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.