Update create_abd_from_h5 to include other post-processing options; update documentation

docs/tutorial_abd.rst

@@ -277,3 +277,76 @@ These components of a BMS transformation can also all be stored in the
 :meth:`scri.bms_transformations.BMSTransformation` class, which allows for things like
 reording the components of the BMS transformation, inverting BMS transformations, and
 composing BMS transformations. For more, see :ref:`bms_transformations`. 
+
+==========
+BMS Frames
+==========
+
+All waveforms at future null infinity (and all waveforms more generally) are functions or coordinates.
+Therefore, there are certain "frames" which may be more useful than others, like that of a rest frame.
+For waveforms at future null infinity, the number of coordinates freedoms, i.e., the symmetries, that they
+exhibit is infinite and is summarized by a group known as the BMS group. This controls the types of frames
+that one may map waveforms to. Because GR is covariant, there is no preferred frame. However, for performing
+analysis on waveforms or building waveform models, it turns out that there are certain frames that are
+more useful than others. In particular, within GR one can extend the notion of a rest frame to something called
+a "superrest frame" (see arXiv:2405.08868 or arXiv:2208.04356 for more details), which typically yields waveforms
+that are easier to understand/analyze. Effectively, mapping to this frame amounts to mapping the system to be
+in the center-of-mass frame, with no instananeous memory, and it's angular velocity in the z-direction. For example,
+for a remnant black hole, this corresponds to making the coordinates match those of the usual Kerr metric and is
+therefore incredibly useful (and necessary) for fitting QNMs to NR waveforms.
+
+The function ``scri.asymptotic_bondi_data.map_to_superrest_frame`` maps to this exact frame.
+In particular, it takes as input:
+
+* ``t_0``, the time at which to map to the superrest frame;
+
+* ``target_PsiM_input``, the target Moreschi supermomentum; this should be ``None`` to map to the superrest frame,
+  but to map to the PN BMS frame one should input the PN Moreschi supermomentum (see arXiv:2208.04356).
+
+* ``target_strain_input``, the target strain; this should be ``None`` to map to the superrest frame,
+  but to map to the PN BMS frame one should input the PN strain (see arXiv:2208.04356).
+
+* ``padding_time``, the time window about ``t_0`` to be used when finding the BMS transformation to the superrest frame.
+
+========================================================
+Creating an AsymptoticBondiData Object from a CCE Output
+========================================================
+
+For processing the output of SpECTRE CCE, one may use the function ``scri.SpEC.file_io.create_abd_from_h5``.
+This function takes as input the path to SpECTRE CCE's output file (via the option ``file_name``) and
+creates an ``abd`` object from said file.
+It also can perform a number of other important post-processing steps, such as:
+
+* time translate the time array of the waveforms by the radius of the worldtube; this ensures that the CCE waveforms
+  are more closely aligned (in time) with extrapolated waveform. This is performed via the ``radius`` option.
+
+* scale out the total Christoudoulou mass of the system from each waveform. This is performed via the ``ch_mass`` option.
+
+* interpolate to a coarser time array, such as the time array of the worldtube. This is performed via the ``t_interpolate`` option.
+
+* map to the superrest BMS frame at some time. This is performed via the ``t_0_superrest`` and ``padding_time`` options.
+  E.g., to make reasonable-looking waveforms, one should map to the superrest frame at some time after junk;
+  ``t_0_superrest`` is the time at which to map to this frame, and ``padding_time`` is the window around ``t_0_superrest``
+  that is used when computing this BMS transformation. ``t_0_superrest - padding_time`` should be after junk radiation.
+  ``padding_time`` should be a few hundred :math:`h=2\overline{\sigma}`, e.g., two orbital periods.
+  The function used to do this is ``abd.map_to_superrest_frame`` (see the "BMS Frames" section).
+
+We recommend including all of these post-processing steps when processing a SpECTRE CCE output.
+
+To obtain the strain :math:`h` from the ``abd`` object, one can use the function ``scri.asymptotic_bondi_data.map_to_superrest_frame.MT_to_WM`` via ``h = MT_to_WM(2.0*abd.sigma.bar)``. 
+This is because the strain :math:`h` is related to the shear :math:`\sigma` via :math:`h=2\overline{\sigma}`.
+
+Example usage of this function could be:
+
+.. code-block:: python
+
+  >>> import scri
+  >>> abd = scri.SpEC.file_io.create_abd_from_h5(
+        file_name='CharacteristicExtractVolume_R0292.h5',
+        file_format="spectrecce",
+        ch_mass=1.0,
+        t_interpolate=t_worldtube,
+        t_0_superrest=1600,
+        padding_time=200
+      )
+  >>> h = scri.asymptotic_bondi_data.map_to_superrest_frame.MT_to_WM(2.0*abd.sigma.bar)

scri/SpEC/file_io/__init__.py

@@ -33,7 +33,6 @@
 rpxmb = rotating_paired_xor_multishuffle_bzip2
 
 
-
 def translate_data_types_GWFrames_to_waveforms(d):
     if d < 8:
         return {0: UnknownDataType, 1: h, 2: hdot, 3: psi4, 4: psi3, 5: psi2, 6: psi1, 7: psi0}[d]
@@ -149,7 +148,6 @@ def read_from_h5(file_name, **kwargs):
         pass  # Probably couldn't find the metadata.json/metadata.txt file
 
     try:  # Make sure the h5py.File gets closed, even in the event of an exception
-
         # Add the old history to the new history, if found
         try:
             try:
@@ -420,7 +418,48 @@ def write_to_h5(w, file_name, file_write_mode="w", attributes={}, use_NRAR_forma
         f.close()
 
 
-def create_abd_from_h5(file_format, convention="SpEC", **kwargs):
+def make_variable_dimensionless(WM, ch_mass=1.0):
+    """
+    Make variable dimensionless by scaling out Christodoulou mass.
+
+    Parameters:
+    -----------
+    WM: scri.WaveformModes
+        Waveform data.
+    ch_mass: float, optional
+        Total Christodoulou mass of the system.
+        [Default: 1.0].
+
+    """
+
+    if WM.m_is_scaled_out:
+        raise ValueError("Data is already dimensionless!")
+
+    if WM.dataType in [psi4, psi3, psi2, psi1, psi0]:
+        unit_scale_factor = (ch_mass) ** (WM.dataType - 4)
+    elif WM.dataType == h:
+        unit_scale_factor = 1 / ch_mass
+    elif WM.dataType == hdot:
+        unit_scale_factor = 1.0
+    else:
+        raise ValueError("DataType not determined.")
+
+    WM.t = WM.t / ch_mass
+    WM.data = WM.data * unit_scale_factor
+    WM.m_is_scaled_out = True
+
+
+def create_abd_from_h5(
+    file_name,
+    file_format,
+    convention="SpEC",
+    radius=None,
+    ch_mass=None,
+    t_interpolate=None,
+    t_0_superrest=None,
+    padding_time=None,
+    **kwargs,
+):
     """Returns an AsymptoticBondiData object with waveform data loaded from specified H5 files.
 
     The AsymptoticBondiData class internally uses the Moreschi-Boyle conventions, see the following reference:
@@ -429,18 +468,30 @@ def create_abd_from_h5(file_format, convention="SpEC", **kwargs):
 
     Parameters
     ----------
-    file_format : 'SXS', 'CCE', 'SpECTRECCE1', 'RPDMB', or 'RPXMB'
+    file_name : str,
+        Path to .h5 file to be loaded.
+    file_format : 'SXS', 'SpECTRECCE', 'RPDMB', or 'RPXMB'
         The H5 files may be in the one of the following file formats:
           * 'SXS'  - Dimensionless extrapolated waveform files found in the SXS Catalog, also known as NRAR format.
-          * 'CCE'  - Asymptotic waveforms output by SpECTRE CCE. These are not dimensionless.
-          * 'SpECTRECCE1' - Asymptotic waveforms output by SpECTRE CCE's original format.
-          * SpECTRECCE' - Asymptotic waveforms output by SpECTRE CCE's newest format.  (NOTE: this may break compatibility)
+          * 'SpECTRECCE' - Asymptotic waveforms output by SpECTRE CCE.
           * 'RPDMB' - Dimensionless waveforms compressed using the rotating_paired_diff_multishuffle_bzip2 format.
           * 'RPXMB' - Dimensionless waveforms compressed using the rotating_paired_xor_multishuffle_bzip2 format.
     convention : 'SpEC' or 'Moreschi-Boyle'
         The data conventions of the waveform data that will be loaded. This defaults to 'SpEC' since this will be
-        most often used with 'SpEC' convention waveforms.
-
+        most often used with 'SpEC' convention waveforms. The output convention is Moreschi-Boyle.
+    radius : str, optional
+        Worldtube radius used when running CCE; only needed for versions of SpECTRE before PR #5985.
+        The time array of the worldtube is translated by this radius.
+    ch_mass : float, optional
+        Total Christodoulou mass of the system.
+    t_interpolate: float array, optional
+        Time array to interpolate to, e.g., the time array of the worldtube.
+    t_0_superrest: float, optional
+        When to map to the BMS superrest frame.
+        Typically a few hundred M after the junk radiation is sufficient.
+    padding_time: float, optional
+        Time window length around t_0_superrest to use when mapping to the superrest frame.
+        Typically a few hundred M or a few orbits is sufficient.
 
     Keyword Parameters
     ------------------
@@ -466,22 +517,29 @@ def create_abd_from_h5(file_format, convention="SpEC", **kwargs):
     # Load waveform data from H5 files into WaveformModes objects
     WMs = {}
     filenames = {}
-    if file_format == "spectrecce" or file_format == "spectrecce1":
-        file_name =  kwargs.pop("file_name")
+    if file_format == "spectrecce":
         with h5py.File(file_name, "r") as f:
-            cce = f["Cce"]
-            time = cce["Strain.dat"][:, 0]
+            cce_key = [x for x in f.keys() if "Spectre" in x][0]
+            if cce_key == []:
+                cce_key = "Cce"
+                data_label_suffix = ".dat"
+            else:
+                radius = cce_key.split("R")[1][:4]
+                data_label_suffix = ""
+
+            cce = f[cce_key]
+            time = cce[f"Strain{data_label_suffix}"][:, 0]
             indices = monotonic_indices(time)
             time = time[indices]
-            ell_max = int(np.sqrt((cce["Strain.dat"].shape[1]-1) / 2) - 1)
+            ell_max = int(np.sqrt((cce[f"Strain{data_label_suffix}"].shape[1] - 1) / 2) - 1)
             for data_label in ["Psi4", "Psi3", "Psi2", "Psi1", "Psi0", "Strain"]:
                 if data_label != "Strain":
                     dataType = DataNames.index(data_label)
                 else:
                     dataType = DataNames.index("h")
                 WMs[data_label] = WaveformModes(
                     t=time,
-                    data=cce[f"{data_label}.dat"][indices, 1:].view(np.complex128),
+                    data=cce[f"{data_label}{data_label_suffix}"][indices, 1:].view(np.complex128),
                     ell_min=0,
                     ell_max=ell_max,
                     frameType=Inertial,
@@ -508,9 +566,7 @@ def create_abd_from_h5(file_format, convention="SpEC", **kwargs):
                     WMs[data_label] = rotating_paired_xor_multishuffle_bzip2.load(filenames[data_label])[0]
                     WMs[data_label].to_inertial_frame()
                 elif file_format == "rpdmb" or file_format == "rpdm":
-                    WMs[data_label] = WaveformModes.from_sxs(
-                        sxs.rpdmb.load(filenames[data_label])
-                    )
+                    WMs[data_label] = WaveformModes.from_sxs(sxs.rpdmb.load(filenames[data_label]))
                 else:
                     raise ValueError(
                         f"File format '{file_format}' not recognized. "
@@ -525,6 +581,7 @@ def create_abd_from_h5(file_format, convention="SpEC", **kwargs):
     # Sanity check
     if not WMs:
         raise ValueError("No filenames have been provided. The data of at least one waveform quantity is required.")
+
     WM_ref = WMs[list(WMs.keys())[0]]
     for i in WMs:
         if not (WM_ref.t == WMs[i].t).all():
@@ -533,6 +590,15 @@ def create_abd_from_h5(file_format, convention="SpEC", **kwargs):
                 f"for {list(WMs.keys())[i].data_type_string} has a different set of times."
             )
 
+    for i in WMs:
+        # Make waveforms dimensionless (if they already are, does nothing)
+        if not ch_mass is None:
+            make_variable_dimensionless(WMs[i], ch_mass)
+
+        # Adjust time by worldtube radius
+        if file_format == "spectrecce":
+            WMs[i].t -= float(radius)
+
     # Create an instance of AsymptoticBondiData
     abd = AsymptoticBondiData(
         time=WM_ref.t,
@@ -581,4 +647,14 @@ def create_abd_from_h5(file_format, convention="SpEC", **kwargs):
         )
         abd.sigma = abd.sigma.bar
 
+    # Interpolate to finer time array, if specified
+    if not t_interpolate is None:
+        idx1 = np.argmin(abs(t_interpolate - abd.t[0])) + 1
+        idx2 = np.argmin(abs(t_interpolate - abd.t[-1])) + 1 - 1
+        abd = abd.interpolate(t_interpolate[idx1:idx2])
+
+    # Map to superrest frame at some time over some window, if specified
+    if not (t_0_superrest is None or padding_time is None):
+        abd, BMS, _ = abd.map_to_superrest_frame(t_0=t_0_superrest, padding_time=padding_time)
+
     return abd