diff --git a/bilby/gw/likelihood/roq.py b/bilby/gw/likelihood/roq.py index 61b0ca95a..60c50d65f 100644 --- a/bilby/gw/likelihood/roq.py +++ b/bilby/gw/likelihood/roq.py @@ -22,18 +22,31 @@ class _WaveformGeneratorWrapper: """ def __init__(self, waveform_generator, update_basis): - object.__setattr__(self, '_wrapped', waveform_generator) - object.__setattr__(self, '_update_basis', update_basis) + self._wrapped = waveform_generator + self._update_basis = update_basis def frequency_domain_strain(self, parameters=None, **kwargs): self._update_basis(parameters) return self._wrapped.frequency_domain_strain(parameters, **kwargs) + def __setstate__(self, d): + # since we overrode __setattr__ just doing self.__dict__ = d + # won't work + for k, v in d.items(): + setattr(self, k, v) + def __getattr__(self, name): return getattr(self._wrapped, name) def __setattr__(self, name, value): - setattr(self._wrapped, name, value) + if name in ["_wrapped", "_update_basis"]: + object.__setattr__(self, name, value) + elif not hasattr(self, "_wrapped"): + raise AttributeError( + "ROQ waveform generator wrapper not properly initialized" + ) + else: + setattr(self._wrapped, name, value) class ROQGravitationalWaveTransient(GravitationalWaveTransient): diff --git a/test/gw/likelihood/roq_test.py b/test/gw/likelihood/roq_test.py index 34fa5504f..d16db2f17 100644 --- a/test/gw/likelihood/roq_test.py +++ b/test/gw/likelihood/roq_test.py @@ -8,6 +8,7 @@ import array_api_compat as aac import bilby +import dill import h5py import numpy as np import pytest @@ -285,6 +286,14 @@ def test_create_roq_weights_fails_due_to_min_frequency(self): priors=self.priors, ) + def test_roq_likelihood_pickles(self): + # the backend waveform generator is not pickleable... + self.roq.waveform_generator = self.roq.waveform_generator.wfg + dill.pickles(self.roq) + self.roq.waveform_generator = BackendWaveformGenerator( + self.roq.waveform_generator, self.xp + ) + def test_create_roq_weights_fails_due_to_duration(self): ifos = bilby.gw.detector.InterferometerList(["H1"]) ifos.set_strain_data_from_power_spectral_densities(