validation.generate_injections does not support IMRPhenomXODE

[JonathanMushkin]

Error description:

In cogwheel/validation/example.py, change APPROXIMANT value to IMRPhenomXODE
Run cells in 1-generate_injections.ipynb, starting from the top, until I reach the cell:
generate_injections.generate_injections_from_config(config, n_cores=8)

See full error output below.

The issue disappears if I keep the approximant at 'IMRPhenomXPHM', or if I change n_cores to 1.
I guess the issue is between IMRPhenomXODE and multiprocessing.

#############################################################################

Setup:
Machine: MacBook Pro 32GB Sequoia 15.0 MacOs.
main branch COGWHEEL, and performing setup required to use IMRPhenomXODE.

###########################################################################

Full error message:

{
"name": "KeyError",
"message": "'IMRPhenomXODE'",
"stack": "---------------------------------------------------------------------------
RemoteTraceback Traceback (most recent call last)
RemoteTraceback:
"""
Traceback (most recent call last):
File "/Users/jonatahm/miniconda3/envs/gwias/lib/python3.12/multiprocessing/pool.py", line 125, in worker
result = (True, func(*args, **kwds))
^^^^^^^^^^^^^^^^^^^
File "/Users/jonatahm/miniconda3/envs/gwias/lib/python3.12/multiprocessing/pool.py", line 51, in starmapstar
return list(itertools.starmap(args[0], args[1]))
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/Users/jonatahm/Work/GW/cogwheel/cogwheel/validation/../../cogwheel/validation/generate_injections.py", line 174, in _get_h_h_rb
return np.sum(likelihood._get_dh_hh_no_asd_drift(par_dic)[1])
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/Users/jonatahm/Work/GW/cogwheel/cogwheel/validation/../../cogwheel/likelihood/relative_binning.py", line 530, in _get_dh_hh_no_asd_drift
d_h, h_h = self._get_dh_hh_complex_no_asd_drift(par_dic) # mpd, mpPd
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/Users/jonatahm/Work/GW/cogwheel/cogwheel/validation/../../cogwheel/likelihood/relative_binning.py", line 553, in _get_dh_hh_complex_no_asd_drift
d_h_mpd, h_h_mpd = self._get_dh_hh_by_m_polarization_detector(
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/Users/jonatahm/Work/GW/cogwheel/cogwheel/validation/../../cogwheel/likelihood/relative_binning.py", line 602, in _get_dh_hh_by_m_polarization_detector
= self.waveform_generator.get_hplus_hcross_at_detectors(
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/Users/jonatahm/Work/GW/cogwheel/cogwheel/validation/../../cogwheel/waveform.py", line 431, in get_hplus_hcross_at_detectors
hplus_hcross = self.get_hplus_hcross(f, waveform_par_dic, by_m)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/Users/jonatahm/Work/GW/cogwheel/cogwheel/validation/../../cogwheel/waveform.py", line 493, in get_hplus_hcross
= APPROXIMANTS[self.approximant].hplus_hcross_by_mode_func(
~~~~~~~~~~~~^^^^^^^^^^^^^^^^^^
KeyError: 'IMRPhenomXODE'
"""

The above exception was the direct cause of the following exception:

KeyError Traceback (most recent call last)
Cell In[8], line 1
----> 1 generate_injections.generate_injections_from_config(config, n_cores=8)

File ~/Work/GW/cogwheel/cogwheel/validation/../../cogwheel/validation/generate_injections.py:47, in generate_injections_from_config(config, n_cores, overwrite)
43 raise FileExistsError(f'{injections_path.as_posix()} exists, '
44 'pass overwrite=True to overwrite')
46 print(f'Generating injections in {injections_path.as_posix()}...')
---> 47 injections_above_threshold = _generate_injections_above_threshold(config,
48 n_cores)
49 injections_above_threshold.to_feather(injections_path)

File ~/Work/GW/cogwheel/cogwheel/validation/../../cogwheel/validation/generate_injections.py:70, in _generate_injections_above_threshold(config, n_cores)
67 likelihood.asd_drift = None # Sets it to 1
69 while len(injs_above_threshold) < config.N_INJECTIONS:
---> 70 batch = _batch_of_injections_above_threshold(config, batch_size,
71 likelihood, n_cores)
72 injs_above_threshold = pd.concat((injs_above_threshold, batch),
73 ignore_index=True)
75 return injs_above_threshold[:config.N_INJECTIONS]

File ~/Work/GW/cogwheel/cogwheel/validation/../../cogwheel/validation/generate_injections.py:102, in _batch_of_injections_above_threshold(config, batch_size, likelihood, n_cores)
100 injection_prior = config.INJECTION_PRIOR_CLS(**config.PRIOR_KWARGS)
101 injections = injection_prior.generate_random_samples(batch_size)
--> 102 h_h_1mpc = _compute_h_h_1mpc(injections, likelihood, n_cores)
104 # Choose physical units for distance, so that the if the loudest
105 # signal was placed at d_ref it would barely make the threshold
106 d_ref = np.sqrt(np.max(h_h_1mpc) / config.H_H_MIN)

File ~/Work/GW/cogwheel/cogwheel/validation/../../cogwheel/validation/generate_injections.py:144, in _compute_h_h_1mpc(injections, likelihood, n_cores, rtol)
142 # Attempt relative binning and estimate accuracy
143 pn_phase_tol = .01
--> 144 h_h_1mpc = _compute_h_h_1mpc_rb(injections, pn_phase_tol, likelihood,
145 n_cores)
146 h_h_1mpc_lowres = _compute_h_h_1mpc_rb(injections, pn_phase_tol*2,
147 likelihood, n_cores)
149 # Recompute inaccurate ones exactly

File ~/Work/GW/cogwheel/cogwheel/validation/../../cogwheel/validation/generate_injections.py:169, in _compute_h_h_1mpc_rb(injections, pn_phase_tol, likelihood, n_cores)
166 par_dics = ({**row, 'd_luminosity': 1} for _, row in injections.iterrows())
167 args = ((likelihood_rb, par_dic) for par_dic in par_dics)
--> 169 h_h_1mpc = _starmap(_get_h_h_rb, args, n_cores)
170 return np.array(h_h_1mpc)

File ~/Work/GW/cogwheel/cogwheel/validation/../../cogwheel/validation/generate_injections.py:198, in _starmap(func, iterable, n_cores)
196 if n_cores != 1:
197 with multiprocessing.Pool(n_cores) as pool:
--> 198 return pool.starmap(func, iterable)
199 else:
200 return [func(*args) for args in iterable]

File ~/miniconda3/envs/gwias/lib/python3.12/multiprocessing/pool.py:375, in Pool.starmap(self, func, iterable, chunksize)
369 def starmap(self, func, iterable, chunksize=None):
370 '''
371 Like map() method but the elements of the iterable are expected to
372 be iterables as well and will be unpacked as arguments. Hence
373 func and (a, b) becomes func(a, b).
374 '''
--> 375 return self._map_async(func, iterable, starmapstar, chunksize).get()

File ~/miniconda3/envs/gwias/lib/python3.12/multiprocessing/pool.py:774, in ApplyResult.get(self, timeout)
772 return self._value
773 else:
--> 774 raise self._value

KeyError: 'IMRPhenomXODE'"
}

[jroulet]

Hi Jonathan, thanks for reporting this.

I don't have a very good idea of how to fix this. I also wasn't able to get multiprocessing and IMRPhenomXODE to work together although I didn't get the same error as you; in my case it hanged up forever. I just used 1 core and manually disabled multiprocessing to get around this.

A hint given your error message is that the cogwheel.waveform.APPROXIMANTS['IMRPhenomXODE'] key should be populated when cogwheel.waveform_models.xode is imported.
Maybe multiprocessing is somehow resetting waveform.APPROXIMANTS to its original state and not reimporting xode for every task or worker?
That's how far I got though, and the few things I've tried to solve this have not worked.