Merge pull request #223 from vijayvarma392/rational_fit

add rational fits

Author: vijayvarma392

README.md

@@ -65,6 +65,7 @@ All of these can be installed through pip or conda.
 * [lalsuite](https://pypi.org/project/lalsuite)
 * [sxs](https://github.com/sxs-collaboration/sxs)
 * [scri](https://github.com/moble/scri)
+* [polyrat](https://pypi.org/project/polyrat/)
 
 # Usage
 See the example notebook [here](https://github.com/vijayvarma392/gw_eccentricity/blob/main/examples/gw_eccentricity_demo.ipynb) for a demo.

examples/gw_eccentricity_demo.ipynb

@@ -347,12 +347,47 @@ def measure_eccentricity(tref_in=None,
         Default: 0.
 
     extra_kwargs: A dict of any extra kwargs to be passed. Allowed kwargs are:
-        spline_kwargs:
+        omega_gw_extrema_interpolation_method : str, default="rational_fit"
+            Specifies the method used to build the interpolations for
+            `omega_gw_pericenters_interp(t)` or
+            `omega_gw_apocenters_interp(t)`. The available options are:
+
+            - `spline`: Uses `scipy.interpolate.InterpolatedUnivariateSpline`.
+                - Best suited for cleaner data, such as when waveform modes
+                  are generated using models like SEOB or TEOB.
+                - Faster to construct and evaluate.
+                - Since it fits through every data point, it may exhibit
+                  oscillatory behavior, particularly near the merger,
+                  especially for noisy NR data.
+
+            - `rational_fit`: Uses `polyrat.StabilizedSKRationalApproximation`.
+                - Can handle both clean and noisy data, e.g., waveform
+                    modes from numerical simulations.
+                - Better monotonic behaviour, particularly near the merger.
+                - Can be slower compared to the `spline` method.
+                  This is because finding optimal numerator and
+                  denominator degree may need several
+                  iterations. See under `get_rational_fit_for_extrema`
+                - Can suppress pathologies in the waveform that might be
+                  seen with `spline`.
+
+            Default value: `"rational_fit"`.
+
+        special_interp_kwargs_for_omega_gw_extrema: dict
+            A dictionary of kwargs for `omega_gw_extrema_interpolation_method`.
+            The defaults are set according to the value of `omega_gw_extrema_interpolation_method`:
+
+            - if omega_gw_extrema_interpolation_method is "spline", default kwargs are set
+              using `utils.get_default_spline_kwargs`.
+            - if omega_gw_extrema_interpolation_method is "rational_fit", default kwargs are set
+              using `utils.get_default_rational_fit_kwargs`.
+
+        general_interp_kwargs: dict
             Dictionary of arguments to be passed to the spline interpolation
             routine (scipy.interpolate.InterpolatedUnivariateSpline) used to
-            compute quantities like omega_gw_pericenters(t) and
-            omega_gw_apocenters(t).
-            Defaults are set using utils.get_default_spline_kwargs
+            interpolate data other than the omega_gw extrema.
+            
+            Defaults are set using `utils.get_default_spline_kwargs`
 
         extrema_finding_kwargs:
             Dictionary of arguments to be passed to the extrema finder,

examples/load_waveform_demo.ipynb

@@ -1019,7 +1019,7 @@ def check_sxs_data_dir(origin, **kwargs):
         f"file named `Strain_N{kwargs['extrap_order']}` since `extrap_order` "
         f"is {kwargs['extrap_order']}."}
     if any([kwargs["include_zero_ecc"], kwargs["include_params_dict"],
-            not kwargs["keep_memory"]]):
+            not kwargs.get("keep_memory", True)]):
         # In newer versions of sxscatalog format, metadata.txt files are
         # replaced by metadata.json file.
         required_metadata_file = "metadata.json" if origin == "SXSCatalog" else "metadata.txt"
@@ -1085,7 +1085,7 @@ def make_return_dict_for_sxs_catalog_format(t, modes_dict, horizon_file_exits,
     dataDict = {"t": t - tpeak,
                 "hlm": modes_dict}
     if any([kwargs["include_zero_ecc"], kwargs["include_params_dict"],
-            not kwargs["keep_memory"]]):
+            not kwargs.get("keep_memory", True)]):
         if os.path.exists(os.path.join(kwargs["data_dir"], "metadata.txt")):
             params_dict = get_params_dict_from_sxs_metadata(
                 os.path.join(kwargs["data_dir"], "metadata.txt"))

gw_eccentricity/eccDefinition.py

@@ -1,6 +1,7 @@
 """Useful functions for the project."""
 import numpy as np
 import argparse
+from polyrat import StabilizedSKRationalApproximation
 from scipy.interpolate import InterpolatedUnivariateSpline
 from scipy.interpolate import PchipInterpolator
 import warnings
@@ -317,6 +318,44 @@ def get_interpolant(oldX,
     return interpolant
 
 
+def get_default_rational_fit_kwargs():
+    """Get default kwargs for rational fit."""
+    default_rational_fit_kwargs = {
+        "num_degree": None,
+        "denom_degree": None,
+        "norm": 2,
+        "maxiter": 20,
+        "verbose": False,
+        "xtol": 1e-07,
+    }
+    return default_rational_fit_kwargs
+
+
+def get_rational_fit(x, y, rational_fit_kwargs=None, check_kwargs=True):
+    """Get rational fit for the data set (x, y).
+
+    We use `polyrat.StabilizedSKRationalApproximation` to obtain
+    rational fit to the data.
+    """
+    if check_kwargs:
+        rational_fit_kwargs = check_kwargs_and_set_defaults(
+            rational_fit_kwargs,
+            get_default_rational_fit_kwargs(),
+            "rational_fit_kwargs",
+            "utils.get_default_rational_fit_kwargs"
+        )
+    # We use a rational approximation based on Stabilized Sanathanan-Koerner Iteration
+    # described in arXiv:2009.10803 and implemented in `polyrat.StabilizedSKRationalApproximation`.
+    rat = StabilizedSKRationalApproximation(**rational_fit_kwargs)
+    # The input x-axis data must be 2-dimensional
+    rat.fit(x.reshape(-1, 1), y)
+    # Using a lambda function to change the input 1-d data to
+    # float64 and reshape to 2-d as required by the fit.
+    # Ensure also that the input and output types are the same.
+    return lambda t: (rat(np.atleast_2d(t).astype('float64').reshape(-1, 1)).item() if np.isscalar(t)\
+                      else rat(np.atleast_2d(t).astype('float64').reshape(-1, 1)))
+
+
 def debug_message(message, debug_level, important=True,
                   point_to_verbose_output=False):
     """Show message based on debug_level.

gw_eccentricity/gw_eccentricity.py

@@ -0,0 +1,3 @@
+[pytest]  
+filterwarnings =  
+    ignore::DeprecationWarning:polyrat.*

gw_eccentricity/load_data.py

@@ -31,7 +31,8 @@ def get_version():
         'h5py',
         'lalsuite',
         'sxs',
-        'scri'
+        'scri',
+        'polyrat'
     ],
     classifiers=[
         "Intended Audience :: Science/Research",

gw_eccentricity/utils.py

@@ -17,10 +17,15 @@
     type=str,
     required=True,
     help="EccDefinition method to save the regression data for.")
+parser.add_argument(
+    "--interp_method",
+    type=str,
+    required=True,
+    help="omega_gw_extrema_interpolation_method to save the regression data for.")
 args = parser.parse_args()
 
 
-def generate_regression_data(method):
+def generate_regression_data(method, interp_method):
     """Generate data for regression test using a method."""
     # Load test waveform
     lal_kwargs = {"approximant": "EccentricTD",
@@ -42,7 +47,7 @@ def generate_regression_data(method):
         raise Exception(f"method {method} is not available. Must be one of "
                         f"{available_methods}")
 
-    extra_kwargs = {}
+    extra_kwargs = {"omega_gw_extrema_interpolation_method": interp_method}
     user_kwargs = extra_kwargs.copy()
     regression_data.update({"extra_kwargs": extra_kwargs})
     # Try evaluating at an array of times
@@ -56,9 +61,9 @@ def generate_regression_data(method):
     meanano_ref = gwecc_dict["mean_anomaly"]
     # We save the measured data 3 reference times
     n = len(tref_out)
-    dict_tref =  {"time": [tref_out[0], tref_out[n//4], tref_out[n//2]],
-                  "eccentricity": [ecc_ref[0], ecc_ref[n//4], ecc_ref[n//2]],
-                  "mean_anomaly": [meanano_ref[0], meanano_ref[n//4], meanano_ref[n//2]]}
+    dict_tref =  {"time": [tref_out[n//8], tref_out[n//4], tref_out[n//2]],
+                  "eccentricity": [ecc_ref[n//8], ecc_ref[n//4], ecc_ref[n//2]],
+                  "mean_anomaly": [meanano_ref[n//8], meanano_ref[n//4], meanano_ref[n//2]]}
 
     # Try evaluating at an array of frequencies
     gwecc_dict = measure_eccentricity(
@@ -70,18 +75,18 @@ def generate_regression_data(method):
     ecc_ref = gwecc_dict["eccentricity"]
     meanano_ref = gwecc_dict["mean_anomaly"]
     n = len(fref_out)
-    dict_fref = {"frequency": [fref_out[0], fref_out[n//4], fref_out[n//2]],
-                 "eccentricity": [ecc_ref[0], ecc_ref[n//4], ecc_ref[n//2]],
-                 "mean_anomaly": [meanano_ref[0], meanano_ref[n//4], meanano_ref[n//2]]}
+    dict_fref = {"frequency": [fref_out[n//8], fref_out[n//4], fref_out[n//2]],
+                 "eccentricity": [ecc_ref[n//8], ecc_ref[n//4], ecc_ref[n//2]],
+                 "mean_anomaly": [meanano_ref[n//8], meanano_ref[n//4], meanano_ref[n//2]]}
     regression_data.update({"tref": dict_tref,
                             "fref": dict_fref})
 
     if not os.path.exists(data_dir):
         os.mkdir(data_dir)
     # save to a json file
-    fl = open(f"{data_dir}/{method}_regression_data.json", "w")
+    fl = open(f"{data_dir}/{method}_{interp_method}_regression_data.json", "w")
     json.dump(regression_data, fl)
     fl.close()
 
 # generate regression data
-generate_regression_data(args.method)
+generate_regression_data(args.method, args.interp_method)