implement a fref

measureEccentricity/eccDefinition.py

@@ -97,7 +97,8 @@ def get_default_extra_kwargs(self):
             "num_orbits_to_exclude_before_merger": 1,
             "extrema_finding_kwargs": {},   # Gets overriden in methods like
                                             # eccDefinitionUsingAmplitude
-            "debug": True
+            "debug": True,
+            "omega_averaging_method": "average_between_extrema"
             }
         return default_extra_kwargs
 
@@ -157,7 +158,7 @@ def interp_extrema(self, extrema_type="maxima"):
                 f"Sufficient number of {extrema_type} are not found."
                 " Can not create an interpolator.")
 
-    def measure_ecc(self, tref_in):
+    def measure_ecc(self, tref_in=None, fref_in=None):
         """Measure eccentricity and mean anomaly at reference time.
 
         parameters:
@@ -189,33 +190,41 @@ def measure_ecc(self, tref_in):
         mean_ano_ref:
             Measured mean anomaly at tref_out.
         """
-        tref_in = np.atleast_1d(tref_in)
-        omega_peaks_interp, self.peaks_location = self.interp_extrema("maxima")
-        omega_troughs_interp, self.troughs_location = self.interp_extrema("minima")
+        self.omega_peaks_interp, self.peaks_location = self.interp_extrema("maxima")
+        self.omega_troughs_interp, self.troughs_location = self.interp_extrema("minima")
 
         t_peaks = self.t[self.peaks_location]
         t_troughs = self.t[self.troughs_location]
-        t_max = min(t_peaks[-1], t_troughs[-1])
-        t_min = max(t_peaks[0], t_troughs[0])
+        self.t_max = min(t_peaks[-1], t_troughs[-1])
+        self.t_min = max(t_peaks[0], t_troughs[0])
+        if tref_in is not None:
+            tref_in = np.atleast_1d(tref_in)
+        elif fref_in is not None:
+            fref_in = np.atleast_1d(fref_in)
+            # get the tref_in from fref_in
+            tref_in = self.compute_tref_in_from_fref_in(fref_in)
+        else:
+            raise KeyError("Atleast one of tref_in or fref_in should be"
+                           " provided")
         # We measure eccentricity and mean anomaly from t_min to t_max.
         # Note that here we do not include the t_max. This is because
         # the mean anomaly computation requires to looking
         # for a peak before and after the ref time to calculate the current
         # period.
         # If ref time is t_max, which could be equal to the last peak, then
         # there is no next peak and that would cause a problem.
-        self.tref_out = tref_in[np.logical_and(tref_in < t_max,
-                                               tref_in >= t_min)]
+        self.tref_out = tref_in[np.logical_and(tref_in < self.t_max,
+                                               tref_in >= self.t_min)]
 
         # Sanity checks
         # Check if tref_out is reasonable
         if len(self.tref_out) == 0:
-            if tref_in[-1] > t_max:
-                raise Exception(f"tref_in is later than t_max={t_max}, "
+            if tref_in[-1] > self.t_max:
+                raise Exception(f"tref_in is later than t_max={self.t_max}, "
                                 "which corresponds to min(last periastron "
                                 "time, last apastron time).")
-            if tref_in[0] < t_min:
-                raise Exception(f"tref_in is earlier than t_min={t_min}, "
+            if tref_in[0] < self.t_min:
+                raise Exception(f"tref_in is earlier than t_min={self.t_min}, "
                                 "which corresponds to max(first periastron "
                                 "time, first apastron time).")
             raise Exception("tref_out is empty. This can happen if the "
@@ -239,8 +248,8 @@ def measure_ecc(self, tref_in):
         # compute eccentricty from the value of omega_peaks_interp
         # and omega_troughs_interp at tref_out using the fromula in
         # ref. arXiv:2101.11798 eq. 4
-        self.omega_peak_at_tref_out = omega_peaks_interp(self.tref_out)
-        self.omega_trough_at_tref_out = omega_troughs_interp(self.tref_out)
+        self.omega_peak_at_tref_out = self.omega_peaks_interp(self.tref_out)
+        self.omega_trough_at_tref_out = self.omega_troughs_interp(self.tref_out)
         self.ecc_ref = ((np.sqrt(self.omega_peak_at_tref_out)
                          - np.sqrt(self.omega_trough_at_tref_out))
                         / (np.sqrt(self.omega_peak_at_tref_out)
@@ -395,6 +404,95 @@ def compute_res_amp_and_omega(self):
         self.res_omega22 = (self.omega22
                             - self.omega22_zeroecc_interp)
 
+    def compute_orbital_averaged_omega22_at_periastrons(self):
+        """Compute reference frequency by orbital averaging at periastron."""
+        @np.vectorize
+        def orbital_averaged_omega22_at_periastrons(n):
+            """Compute orbital averaged omega22 between n and n+1 peaks."""
+            # integrate omega22 between n and n+1 peaks
+            integ = (self.phase22[self.peaks_location[n+1]]
+                     - self.phase22[self.peaks_location[n]])
+            period = (self.t[self.peaks_location[n+1]]
+                      - self.t[self.peaks_location[n]])
+            return integ / period
+        # get the mid points between the peaks as avg time
+        t_average = (self.t[self.peaks_location][1:]
+                     + self.t[self.peaks_location][:-1]) / 2
+        omega22_average = orbital_averaged_omega22_at_periastrons(
+            np.arange(len(self.peaks_location) - 1))
+        return t_average, omega22_average
+
+    def find_tref_from_orbital_averaged_omega22_at_periastrons(self,
+                                                               omega22_ref):
+        """Find the reference time given a reference frequency."""
+        t_average, omega22_average = self.compute_orbital_averaged_omega22_at_periastrons()
+        t_of_omega = InterpolatedUnivariateSpline(omega22_average, t_average)
+        return t_of_omega(omega22_ref)
+
+    def compute_omega22_average_between_extrema(self):
+        """Find omega22 average by taking mean of peaks_interp and troughs_interp."""
+        t = np.arange(self.t_min, self.t_max, self.t[1] - self.t[0])
+        return t, (self.omega_peaks_interp(t)
+                   + self.omega_troughs_interp(t)) / 2
+
+    def find_tref_from_omega22_average_between_extrema(self, fref):
+        """Find the reference time given a reference frequency."""
+        t, average_omega22 = self.compute_omega22_average_between_extrema()
+        # check if average omega22 is monotonically increasing
+        domega22_dt = np.gradient(average_omega22, t)
+        if any(domega22_dt <= 0):
+            warnings.warn("Omega22 average between extrema is not "
+                          "monotonically increasing.")
+        t_of_omega22 = InterpolatedUnivariateSpline(average_omega22, t)
+        return t_of_omega22(fref)
+
+    def compute_omega22_zeroecc(self):
+        """Find omega22 from zeroecc data."""
+        t = np.arange(self.t_min, self.t_max, self.t[1] - self.t[0])
+        omega22_zeroecc = np.interp(t, self.t_zeroecc, self.omega22_zeroecc)
+        return t, omega22_zeroecc
+
+    def find_tref_from_omega22_zeroecc(self, fref):
+        """Find the reference time given a reference frequency."""
+        t, average_omega22 = self.compute_omega22_zeroecc()
+        # check if average omega22 is monotonically increasing
+        domega22_dt = np.gradient(average_omega22, t)
+        if any(domega22_dt <= 0):
+            warnings.warn("Omega22 average between extrema is not "
+                          "monotonically increasing.")
+        t_of_omega22 = InterpolatedUnivariateSpline(average_omega22, t)
+        return t_of_omega22(fref)
+
+    def get_availabe_omega_averaging_methods(self):
+        """Return available omega averaging methods."""
+        available_methods = {
+            "average_between_extrema": [self.compute_omega22_average_between_extrema, self.find_tref_from_omega22_average_between_extrema],
+            "orbital_average_at_periastron": [self.compute_orbital_averaged_omega22_at_periastrons, self.find_tref_from_orbital_averaged_omega22_at_periastrons],
+            "omega22_zeroecc": [self.compute_omega22_zeroecc, self.find_tref_from_omega22_zeroecc]
+        }
+        return available_methods
+
+    def compute_tref_in_from_fref_in(self, fref_in):
+        """Compute tref_in from fref_in using chosen omega average method."""
+        available_averaging_methods = self.get_availabe_omega_averaging_methods()
+        method = self.extra_kwargs["omega_averaging_method"]
+        if method in available_averaging_methods:
+            self.omega22_average = available_averaging_methods[method][0]()[1]
+            # check if the fref falls within the range of omega22 average
+            if fref_in[0] < self.omega22_average[0]:
+                raise Exception("fref_in is less than minimum available "
+                                "frequency "
+                                f"{self.omega22_average[0]}")
+            if fref_in[-1] > self.omega22_average[-1]:
+                raise Exception("fref_in is greater than maximum available "
+                                "frequency "
+                                f"{self.omega22_average[-1]}")
+            tref_in = available_averaging_methods[method][1](fref_in)
+            return tref_in
+        else:
+            raise KeyError(f"Omega averaging method {method} does not exist. "
+                           f"Must be one of {available_averaging_methods.keys()}")
+
     def make_diagnostic_plots(self, usetex=True, **kwargs):
         """Make dignostic plots for the eccDefinition method.
 

measureEccentricity/measureEccentricity.py

@@ -47,7 +47,8 @@ def get_available_methods():
     return models
 
 
-def measure_eccentricity(tref_in, dataDict, method="Amplitude",
+def measure_eccentricity(tref_in=None, fref_in=None,
+                         dataDict=None, method="Amplitude",
                          return_ecc_method=False,
                          spline_kwargs=None,
                          extra_kwargs=None):
@@ -131,7 +132,8 @@ def measure_eccentricity(tref_in, dataDict, method="Amplitude",
                                                spline_kwargs=spline_kwargs,
                                                extra_kwargs=extra_kwargs)
 
-        tref_out, ecc_ref, mean_ano_ref = ecc_method.measure_ecc(tref_in)
+        tref_out, ecc_ref, mean_ano_ref = ecc_method.measure_ecc(
+            tref_in=tref_in, fref_in=fref_in)
         if not return_ecc_method:
             return tref_out, ecc_ref, mean_ano_ref
         else:

test/test_interface.py

@@ -1,6 +1,7 @@
 import measureEccentricity
 from measureEccentricity import load_data
 from measureEccentricity import measure_eccentricity
+import numpy as np
 
 
 def test_interface():
@@ -37,3 +38,16 @@ def test_interface():
 
         # Make diagnostic plots
         eccMethod.make_diagnostic_plots(usetex=False)
+
+        # Try evaluating at single frequency
+        tref_out, ecc_ref, meanano_ref = measure_eccentricity(
+            fref_in=0.025,
+            method=method,
+            dataDict=dataDict)
+
+        # Try evaluating at an array of frequencies
+        tref_out, ecc_ref, meanano_ref, eccMethod = measure_eccentricity(
+            fref_in=np.arange(0.025, 0.035, 0.001),
+            method=method,
+            dataDict=dataDict,
+            return_ecc_method=True)