Fix x-axis during superrest rotation and minor changes

scri/asymptotic_bondi_data/map_to_superrest_frame.py

@@ -404,19 +404,46 @@ def com_transformation_to_map_to_superrest_frame(abd, N_itr_max=10, rel_err_tol=
     return best_CoM_transformation, rel_errs
 
 
-def rotation_from_spin_charge(chi, t):
+def rotation_from_spin_charge(chi, t, fix_xz_plane=False, fix_yz_plane=False):
     """Obtain the rotation from the remnant BH's spin vector.
-    This finds the rotation that aligns the z-component of the spin vector with the z-axis.
+
+    This finds the rotation that aligns the z-component of the spin vector with the z-axis,
+    i.e., this finds the quaternion q such that q * quaternion.z * q.inverse() = chi_f(t).
+
     Parameters
     ----------
     chi: ndarray, real, shape (..., 3)
         Remnant BH's spin vector.
     t: ndarray, real
         Time array corresponding to the size of the spin vector.
+    fix_xz_plane: bool
+        Whether or not to modify the quaternion by a rotation about z
+        such that the x unit vector remains in the x-z plane
+        after the full rotation.
+        Default is False.
+    fix_yz_plane: bool
+        Whether or not to modify the quaternion by a rotation about z
+        such that the y unit vector remains in the y-z plane
+        after the full rotation.
+        Default is False.
     """
     chi_f = quaternion.quaternion(*chi[np.argmin(abs(t))]).normalized()
-    q = (1 - chi_f * quaternion.z).normalized().components
-    return scri.bms_transformations.BMSTransformation(frame_rotation=q)
+    q = (1 - chi_f * quaternion.z).normalized()
+
+    if fix_xz_plane:
+        y_rotated = quaternion.as_vector_part(q * quaternion.y * q.inverse())
+        phase = np.angle((y_rotated[0] + 1j*y_rotated[1])) - np.pi/2
+        q = q * quaternion.from_rotation_vector(phase * np.array([0, 0, 1]))
+        if (q * quaternion.x * q.inverse()).components[1] < 0:
+            q = q * quaternion.from_rotation_vector(np.pi * np.array([0, 0, 1]))
+    elif fix_yz_plane:
+        x_rotated = quaternion.as_vector_part(q * quaternion.x * q.inverse())
+        phase = np.angle((x_rotated[0] + 1j*x_rotated[1]))
+        q = q * quaternion.from_rotation_vector(phase * np.array([0, 0, 1]))
+        if (q * quaternion.y * q.inverse()).components[2] < 0:
+            q = q * quaternion.from_rotation_vector(np.pi * np.array([0, 0, 1]))
+
+    return scri.bms_transformations.BMSTransformation(frame_rotation=q.components)
 
 
 def rotation_from_vectors(vector, target_vector, t=None):
@@ -438,7 +465,7 @@ def rotation_from_vectors(vector, target_vector, t=None):
     return scri.bms_transformations.BMSTransformation(frame_rotation=q.components)
 
 
-def rotation_to_map_to_superrest_frame(abd, target_strain=None, N_itr_max=10, rel_err_tol=1e-12, print_conv=False):
+def rotation_to_map_to_superrest_frame(abd, target_strain=None, N_itr_max=10, rel_err_tol=1e-12, fix_xz_plane=False, fix_yz_plane=False, print_conv=False):
     """Determine the rotation needed to map an abd object to the superrest frame.
 
     This is found through an iterative solve; e.g., compute the transformation needed to align
@@ -464,6 +491,16 @@ def rotation_to_map_to_superrest_frame(abd, target_strain=None, N_itr_max=10, re
         First value is minimum relative error tolerance between transformation iterations; second value is the
         minimum relative error tolerance between the NR angular velocity and the target angular velocity.
         Default is 1e-12.
+    fix_xz_plane: bool
+        Whether or not to modify the quaternion by a rotation about z
+        such that the x unit vector remains in the x-z plane
+        after the full rotation.
+        Default is False.
+    fix_yz_plane: bool
+        Whether or not to modify the quaternion by a rotation about z
+        such that the y unit vector remains in the y-z plane
+        after the full rotation.
+        Default is False.
     print_conv: bool, defaults to False
         Whether or not to print the termination criterion. Default is False.
     """
@@ -534,12 +571,12 @@ def rotation_to_map_to_superrest_frame(abd, target_strain=None, N_itr_max=10, re
                 chi_prime = chi_prime / np.linalg.norm(chi_prime, axis=-1)[:, None]
 
             rotation_transformation = (
-                rotation_from_spin_charge(chi_prime, abd_prime.t) * rotation_transformation
+                rotation_from_spin_charge(chi_prime, abd_prime.t, fix_xz_plane, fix_yz_plane) * rotation_transformation
             ).reorder(["supertranslation", "frame_rotation", "boost_velocity"])
             # remove supertranslation and CoM components
             rotation_transformation.supertranslation *= 0
             rotation_transformation.boost_velocity *= 0
-
+            
             abd_prime = abd.transform(frame_rotation=rotation_transformation.frame_rotation.components)
 
             chi_prime = abd_prime.bondi_dimensionless_spin()
@@ -683,8 +720,9 @@ def map_to_superrest_frame(
     order=["supertranslation", "rotation", "CoM_transformation"],
     ell_max=None,
     alpha_ell_max=None,
-    fix_time_phase_freedom=False,
     modes=None,
+    fix_xz_plane=False,
+    fix_yz_plane=False,
     print_conv=False,
 ):
     """Transform an abd object to the superrest frame.
@@ -739,18 +777,26 @@ def map_to_superrest_frame(
     order : list, optional
         Order in which to solve for the BMS transformations.
         Default is ["rotation", "CoM_transformation", "supertranslation"].
+        If "time_phase" is included, then a time/phase optimization is performed.
     ell_max : int, optional
         Maximum ell to use for SWSH/Grid transformations.
         Default is self.ell_max.
     alpha_ell_max : int, optional
         Maximum ell of the supertranslation to use.
         Default is self.ell_max.
-    fix_time_phase_freedom : bool, optional
-        Whether or not to fix the time and phase freedom using a 2d minimization scheme.
-        Default is True.
     modes : list, optional
         List of modes to include when performing the 2d alignment.
         Default is every mode.
+    fix_xz_plane: bool
+        Whether or not to modify the quaternion by a rotation about z
+        such that the x unit vector remains in the x-z plane
+        after the full rotation.
+        Default is False.
+    fix_yz_plane: bool
+        Whether or not to modify the quaternion by a rotation about z
+        such that the y unit vector remains in the y-z plane
+        after the full rotation.
+        Default is False.
     print_conv: bool, defaults to False
         Whether or not to print the termination criterion. Default is False.
 
@@ -765,6 +811,9 @@ def map_to_superrest_frame(
     """
     abd = self.copy()
 
+    if order == []:
+        return abd, scri.bms_transformations.BMSTransformation(), None
+
     if target_strain_input is not None:
         target_strain = target_strain_input.copy()
         target_strain.t -= t_0
@@ -832,6 +881,8 @@ def map_to_superrest_frame(
                     target_strain=target_strain,
                     N_itr_max=N_itr_maxes["rotation"],
                     rel_err_tol=rel_err_tols["rotation"],
+                    fix_xz_plane=fix_xz_plane,
+                    fix_yz_plane=fix_yz_plane,
                     print_conv=print_conv,
                 )
             elif transformation == "CoM_transformation":
@@ -874,7 +925,6 @@ def map_to_superrest_frame(
             )
 
         if target_strain is not None and order[-1] == "time_phase":
-            # rel_err is obtained from align2d, so do nothing
             pass
         else:
             rel_err = rel_err_for_abd_in_superrest(abd_interp_prime, target_PsiM, target_strain)