Bugfix and enhancement in heterodyned likelihood

src/jimgw/core/single_event/likelihood.py

@@ -438,17 +438,8 @@ def __init__(
         if reference_waveform is None:
             reference_waveform = waveform
 
-        # Get the original frequency grid
-        frequency_original = self.frequencies
-        # Get the grid of the relative binning scheme (contains the final endpoint)
-        # and the center points
-        freq_grid, self.freq_grid_center = self.make_binning_scheme(
-            jnp.array(frequency_original), n_bins
-        )
-        self.freq_grid_low = freq_grid[:-1]
-
         if ref_params:
-            self.ref_params = ref_params
+            self.ref_params = ref_params.copy()
             logging.info(f"Reference parameters provided, which are {self.ref_params}")
         elif prior:
             logging.info("No reference parameters are provided, finding it...")
@@ -484,12 +475,21 @@ def __init__(
         self.B0_array = {}
         self.B1_array = {}
 
+        # Get the original frequency grid
+        frequency_original = self.frequencies
+        # Get the grid of the relative binning scheme (contains the final endpoint)
+        # and the center points
+        freq_grid, self.freq_grid_center = self.make_binning_scheme(
+            jnp.array(frequency_original), n_bins
+        )
+        self.freq_grid_low = freq_grid[:-1]
+
         h_sky = reference_waveform(frequency_original, self.ref_params)
 
         # Get frequency masks to be applied, for both original
         # and heterodyne frequency grid
         h_amp = jnp.sum(
-            jnp.array([jnp.abs(h_sky[key]) for key in h_sky.keys()]), axis=0
+            jnp.array([jnp.abs(h_sky[pol]) for pol in h_sky.keys()]), axis=0
         )
         f_valid = frequency_original[jnp.where(h_amp > 0)[0]]
         f_max = jnp.max(f_valid)
@@ -532,7 +532,6 @@ def __init__(
                 freq_grid,
                 self.freq_grid_center,
             )
-
             self.A0_array[detector.name] = A0[mask_heterodyne_center]
             self.A1_array[detector.name] = A1[mask_heterodyne_center]
             self.B0_array[detector.name] = B0[mask_heterodyne_center]
@@ -556,7 +555,7 @@ def _likelihood(self, params: dict[str, Float], data: dict) -> Float:
                 frequencies_low, waveform_sky_low, params
             )
             waveform_center = detector.fd_response(
-                frequencies_low, waveform_sky_center, params
+                frequencies_center, waveform_sky_center, params
             )
 
             r0 = waveform_center / self.waveform_center_ref[detector.name]
@@ -634,9 +633,8 @@ def make_binning_scheme(
             The bin centers.
         """
         phase_diff_array = self.max_phase_diff(freqs, freqs[0], freqs[-1], chi=chi)  # type: ignore
-        bin_f = interp1d(phase_diff_array, freqs)
         phase_diff = jnp.linspace(phase_diff_array[0], phase_diff_array[-1], n_bins + 1)
-        f_bins = bin_f(phase_diff)
+        f_bins = interp1d(phase_diff_array, freqs)(phase_diff)
         f_bins_center = (f_bins[:-1] + f_bins[1:]) / 2
         return jnp.array(f_bins), jnp.array(f_bins_center)
 
@@ -653,15 +651,15 @@ def compute_coefficients(data, h_ref, psd, freqs, f_bins, f_bins_center):
         for i in range(len(f_bins) - 1):
             f_index = jnp.where((freqs >= f_bins[i]) & (freqs < f_bins[i + 1]))[0]
             freq_shift = freqs[f_index] - f_bins_center[i]
-            A0_array.append(4 * jnp.sum(data_prod[f_index]) * df)
-            A1_array.append(4 * jnp.sum(data_prod[f_index] * freq_shift) * df)
-            B0_array.append(4 * jnp.sum(self_prod[f_index]) * df)
-            B1_array.append(4 * jnp.sum(self_prod[f_index] * freq_shift) * df)
-
-        A0_array = jnp.array(A0_array)
-        A1_array = jnp.array(A1_array)
-        B0_array = jnp.array(B0_array)
-        B1_array = jnp.array(B1_array)
+            A0_array.append(jnp.sum(data_prod[f_index]))
+            A1_array.append(jnp.sum(data_prod[f_index] * freq_shift))
+            B0_array.append(jnp.sum(self_prod[f_index]))
+            B1_array.append(jnp.sum(self_prod[f_index] * freq_shift))
+
+        A0_array = 4 * df * jnp.array(A0_array)
+        A1_array = 4 * df * jnp.array(A1_array)
+        B0_array = 4 * df * jnp.array(B0_array)
+        B1_array = 4 * df * jnp.array(B1_array)
         return A0_array, A1_array, B0_array, B1_array
 
     def maximize_likelihood(
@@ -676,15 +674,15 @@ def maximize_likelihood(
         for transform in sample_transforms:
             parameter_names = transform.propagate_name(parameter_names)
 
+        super_obj = super(HeterodynedTransientLikelihoodFD, self)
+
         def y(x: Float[Array, " n_dims"], data: dict) -> Float:
             named_params = dict(zip(parameter_names, x))
             for transform in reversed(sample_transforms):
                 named_params = transform.backward(named_params)
             for transform in likelihood_transforms:
                 named_params = transform.forward(named_params)
-            return -super(HeterodynedTransientLikelihoodFD, self).evaluate(
-                named_params, data
-            )
+            return -super_obj.evaluate(named_params, data)
 
         print("Starting the optimizer")
 
@@ -720,7 +718,7 @@ def y(x: Float[Array, " n_dims"], data: dict) -> Float:
                 non_finite_index[:common_length]
             ].set(guess[:common_length])
 
-        rng_key, best_fit, log_prob = optimizer.optimize(
+        _, best_fit, log_prob = optimizer.optimize(
             jax.random.PRNGKey(12094), y, initial_position, {}
         )