Generating_waveforms.ipynb Notebook Failure

[thompsonphys]

Hi!

I tried to run one of the example notebooks (Generating_waveforms.ipynb) and ran into a couple of issues.

Firstly, it seems that all calls of IMRPhenomD.gen_IMRPhenomD need updating to take in f_ref. Once this issue is resolved, the actual function IMRPhenomD.gen_IMRPhenomD doesn't like the way the frequencies are passed in when defining h0_simple. Using jax 0.4.10 the shapes of fs[0] and fs[:1] are () and (1,0) resp., and the former raises an IndexError:

---------------------------------------------------------------------------
IndexError                                Traceback (most recent call last)
Cell In[33], line 10
      7 h0_simple = lambda theta: IMRPhenomD.gen_IMRPhenomD(fs[0], theta, fs[0]).real
      9 h0_grad = grad(h0_simple)
---> 10 print(h0_grad(theta_ripple_h0))

    [... skipping hidden 10 frame]

Cell In[33], line 7, in (theta)
      1 # Finally, lets take some derivatives!
      2 # We start by writing a simple lambda function to make something that only depends
      3 # on the intrinsic parameters
      4 
      5 # Note here that JAX is expecting a scalar output to compute derivatives
      6 # We therefore just select a single frequency point to make this easy
----> 7 h0_simple = lambda theta: IMRPhenomD.gen_IMRPhenomD(fs[0], theta, f_ref).real
      9 h0_grad = grad(h0_simple)
     10 print(h0_grad(theta_ripple_h0))

File [~/miniconda3/envs/systematics/lib/python3.10/site-packages/ripple/waveforms/IMRPhenomD.py:582](https://file+.vscode-resource.vscode-cdn.net/Users/jthompson/projects/systematics/systematics-git/notebooks/~/miniconda3/envs/systematics/lib/python3.10/site-packages/ripple/waveforms/IMRPhenomD.py:582), in gen_IMRPhenomD(f, params, f_ref)
    579 theta_extrinsic = jnp.array([params[4], params[5], params[6]])
    581 coeffs = get_coeffs(theta_intrinsic)
--> 582 h0 = _gen_IMRPhenomD(f, theta_intrinsic, theta_extrinsic, coeffs, f_ref)
    583 return h0
...
   4537       f"non-None[/Ellipsis](https://file+.vscode-resource.vscode-cdn.net/Ellipsis) indices for dim {arr_ndim}.")
   4538 ellipses = (i for i, elt in enumerate(idx) if elt is Ellipsis)
   4539 ellipsis_index = next(ellipses, None)

IndexError: Too many indices for array: 1 non-None[/Ellipsis](https://file+.vscode-resource.vscode-cdn.net/Ellipsis) indices for dim 0.

Changing to fs[:1] resolves this but now the function outputs an array (of length 1) that breaks grad. In the end I got things to work by changing h0_simple to this (and propagating the changes to the rest of the notebook):

h0_simple = lambda theta: IMRPhenomD.gen_IMRPhenomD(fs[:1], theta, f_ref).real[0]

Also because of the way vmap passes the frequency array by individual value into the polarization functions, the following hack was also needed for the polarization gradient functions, e.g.

hp_ripple_real = lambda theta, f: IMRPhenomD.gen_IMRPhenomD_polar(jnp.array([f]), theta, f_ref)[0].real[0]
hp_ripple_imag = lambda theta, f: IMRPhenomD.gen_IMRPhenomD_polar(jnp.array([f]), theta, f_ref)[0].imag[0]

Just wanted to raise awareness in case this change to the passed frequency array was unexpectedly caused by upstream changes to the code. There's probably a more elegant solution to fix this error!

[tedwards2412]

Thanks for catching this! I pushed a small fix with your changes so that the notebook now runs. At the moment I'm not sure there is a more elegant way round this without changing the phenomD code. I'll look into if there is something we can do in the next release though.