Source code for waveformtools.waveforms

""" Classes for handling the waveform modes or data defined on spheres.


Classes
-------
spherical_array:    A 2D data-type.
                    Stores and manages two-dimensional data on surfaces of spherical topology.
modes_array: A data-type.
             Handle and work with mode coefficients.
"""

import sys
import numpy as np
import h5py
from waveformtools.waveformtools import message

# from waveformtools import dataIO
from waveformtools.transforms import Yslm_vec
from waveformtools.grids import spherical_grid
from waveformtools import dataIO

#####################
# Units
#####################

G = 6.67 * 10 ** (-11.0)
c = 3.0 * 10**8  # m/s
Msun = 1.988500 * 10**30.0  # g, SI
muc = c**2 / (G * Msun)  # g, NR to SI
tuc = G * Msun / (c**3)  # s, NR to SI
dMpc = 3.0857 * 10 ** (22)  # m


#################################################
# Spherical array class
################################################


# @jitclass(spec_sp)
[docs]class spherical_array: """A class for handling waveforms on a sphere. Attributes ---------- label: str The label of the waveform data. time_axis: 1d array The time axis of the data. frequency_axis: 1d array The frequency axis if the data is represented in frequency domain. grid_info: spherical_grid An instance of the `spherical_grid` class. data_len: int The length of the data along the time axis. Methods ------- delta_t: Fetch the time stepping `delta_t`. to_modes_array: Find the waveform expressed in the SWSH basis. boost: Boost the waveform. supertranslate: Supertranslate the waveform. """ def __init__( self, label=None, time_axis=None, frequency_axis=None, data=None, data_dir=None, file_name=None, grid_info=None, spin_weight=2, ): self.label = label # self.base_dir = base_dir # The base directory containing the self.data = data self.file_name = file_name self.data_dir = data_dir self.time_axis = time_axis self.frequency_axis = frequency_axis self.grid_info = grid_info self.spin_weight = spin_weight
[docs] def delta_t(self, value=None): """Sets and returns the value of time stepping :math:`dt`. Parameters ---------- value : float, optional The value of :math:`dt` to set to the attribute. Returns ------- delta_t: float Sets the attribute. """ # if not self.delta_t: if not value: try: delta_t = self.time_axis[1] - self.time_axis[0] except Exception as ex: print("Please input the value of `delta_t` or supply the `time_axis` to the waveform.", ex) else: delta_t = value return delta_t
@property def data_len(self): """Returns the length of the time/frequency axis. Returns ------- data_len: int The length of the time/frequency axis. """ try: data_len = len(self.time_axis) except Exception as ex: data_len = len(self.frequency_axis) message(ex) return data_len
[docs] def to_modes_array(self, grid_info=None, spin_weight=None, ell_max=8): """Decompose a given spherical array function on a sphere into Spin Weighted Spherical Harmonic modes. Parameters ---------- spin_weight: int, optional The spin weight of the waveform. It defaults to -2 for a gravitational waveform. ell_max: int, optional The maximum value of the :math:`\\ell` polar quantum number. Defaults to 8. grid_info: class instance The class instance that contains the properties of the spherical grid. Returns ------- waveforms_modes: modes_array An instance of the `modes_array` class containing the decomposed modes. Notes ----- 1. Assumes that the sphere on which this decomposition is carried out is so far out that the coordinate system is spherical polar on a round sphere. 2. Assumes that the poper area is the same as its co-ordinate area. 3. Ensure that the label of the input spherical array indicates whether it is a time domain data or frequency domain data. """ if grid_info is None: if self.grid_info is None: message("Please specify the grid specs. Assuming defaults.") grid_info = spherical_grid() self.grid_info = grid_info else: grid_info = self.grid_info if spin_weight is None: if self.spin_weight is None: message("Please specify spin weight. Assuming 2") spin_weight = 2 self.spin_weight = spin_weight else: spin_weight = self.spin_weight spin_weight = abs(spin_weight) # Compute the meshgrid for theta and phi. theta, phi = grid_info.meshgrid # Create a modes array object # Create a modes list. modes_list = construct_mode_list(ell_max, spin_weight=spin_weight) if not self.label: label = "decomposed_time_domain" else: label = self.label # Create a mode array for the decomposed_waveform waveform_modes = modes_array(label=label, ell_max=ell_max, spin_weight=spin_weight) # Inherit the time or frequency axis. if "time" in label: # axis = self.time_axis waveform_modes.time_axis = self.time_axis else: # axis = self.frequency_axis waveform_modes.frequency_axis = self.frequency_axis # Create the modes_array waveform_modes.create_modes_array(ell_max=ell_max, data_len=self.data_len) waveform_modes.modes_list = modes_list # The area element on the sphere # Compute the meshgrid for theta and phi. theta, phi = grid_info.meshgrid # sqrt_met_det = np.sin(theta) sqrt_met_det = np.sqrt(np.power(np.sin(theta), 2)) darea = sqrt_met_det * grid_info.dtheta * grid_info.dphi modes_list = [item for item in modes_list if item[0] >= spin_weight] for mode in modes_list: ell_value, all_emm_values = mode for emm_value in all_emm_values: # m value. # Spin weighted spherical harmonic function at (theta, phi) Ybasis_fun = np.conj(Yslm_vec(spin_weight, ell=ell_value, emm=emm_value, theta_grid=theta, phi_grid=phi)) Ydarea = Ybasis_fun * darea # print(Ydarea.shape) # print(full_integrand) # Using quad multipole_ell_emm = np.tensordot(self.data, Ydarea, axes=((0, 1), (0, 1))) # print(f'l{ell_value}m{emm_value}', multipole_ell_emm) waveform_modes.set_mode_data(ell_value, emm_value, multipole_ell_emm) return waveform_modes
[docs] def boost(self, conformal_factor, grid_info=None): """Boost the waveform given the unboosted waveform and the boost conformal factor. Parameters ---------- self: spherical_array A class instance of `spherical array`. conformal_factor: 2d array The conformal factor for the Lorentz transformation. It may be a single floating point number or an array on a spherical grid. The array will be of dimensions [ntheta, nphi] grid_info: class instance The class instance that contains the properties of the spherical grid. Returns ------- boosted_waveform: sp_array The class instance `sp_array` that contains the boosted waveform. """ from waveformtools.waveforms import spherical_array if grid_info is None: grid_info = self.grid_info # Compute the boosted waveform on the spherical grid on all the elements. boosted_waveform_data = np.transpose(self.data, (2, 0, 1)) * conformal_factor # boosted_waveform_data = np.multiply(self.data, conformal_factor[:, :, None]) # boosted_waveform_data = boosted_waveform_item # boosted_waveform_data = np.array([np.transpose(item)*conformal_factor # for item in np.transpose(self.data)]) # Construct a 2d waveform array object boosted_waveform = spherical_array( grid_info=grid_info, data=np.transpose(np.array(boosted_waveform_data), (1, 2, 0)) ) # Assign other attributes. boosted_waveform.label = "boosted " + self.label boosted_waveform.time_axis = self.time_axis return boosted_waveform
[docs] def supertranslate(self, supertransl_alpha_sp, order=1): """Supertranslate the :math:`\\Psi_{4\\ell m}` waveform modes, give the, the supertranslation parameter and the order. Parameters ---------- supertransl_alpha_modes : modes_array The modes_array containing the supertranslation mode coefficients. grid_info: class instance The class instance that contains the properties of the spherical grid using which the computations are carried out. order: int The number of terms to use to approximate the supertranslation. Returns ------- Psi4_supertransl: modes_array A class instance containg the modes of the supertranslated waveform :math:`\\Psi_4`. """ # Create a spherical_array to hold the supertranslated waveform Psi4_supertransl_sp = spherical_array( grid_info=self.grid_info, label="{} -> supertranslated time".format(self.label) ) delta_t = float(self.delta_t()) # Set the data. data = 0 # data = self.data # Psi4_supertransl_sp.data = self.data # delta = 0 # count = 0 from waveformtools.differentiate import differentiate5_vec_numba for index in range(order): # print(f'{index} loop') dPsidu = self.data for dorder in range(index + 1): # print(f'differentiating {dorder+1} times') dPsidu = differentiate5_vec_numba(dPsidu, delta_t) print("Incrementing...") # delta = np.power(supertransl_alpha_sp.data, index+1) * dPsidu / np.math.factorial(index+1) # print(delta/self.data) data += np.power(supertransl_alpha_sp.data, index + 1) * dPsidu / np.math.factorial(index + 1) # delta data += self.data print("Equal to original waveform?", (data == self.data).all()) Psi4_supertransl_sp.data = data Psi4_supertransl_sp.time_axis = self.time_axis print("Done.") return Psi4_supertransl_sp
[docs] def load_qlm_data(self, data_dir=None, grid_info=None, bh=0, variable="sigma"): """Load the 2D shear data from h5 files. Parameters ---------- file_name: str The name of the file containing data. data_dir: str The name of the directory containing data. grid_info: class instance An instance of the grid_info class. bh: int The black hole number (0, 1 or 2) """ # import sys # import re # import json # if file_name is None: # if self.file_name is None: # print('Please supply the file name!') # else: # file_name = self.file_name # else: # if self.file_name is None: # self.file_name = file_name if data_dir is None: if self.data_dir is None: print("Please supply the data directory!") else: data_dir = self.data_dir else: if self.data_dir is None: self.data_dir = data_dir if grid_info is None: if self.grid_info is None: print("Please supply the grid spec!") else: grid_info = self.grid_info else: if self.grid_info is None: self.grid_info = grid_info # get the full path. file_name = f"qlm_{variable}[{bh}].xy.h5" full_path = self.data_dir + file_name # cflag = 0 nghosts = grid_info.nghosts ntheta = grid_info.ntheta nphi = grid_info.nphi # Open the modes file. with h5py.File(full_path, "r") as wfile: # Get all the mode keys. modes_keys_list = list(wfile.keys()) # modes_keys_list = sorted(modes_keys_list) # Get the mode keys containing the data. modes_keys_list = [item for item in modes_keys_list if "it=" in item] # Get the itaration numbers. iteration_numbers = sorted(get_iteration_numbers_from_keys(modes_keys_list)) # sargs = np.argsort(iteration_numbers) # iteration_numbers = iteration_numbers[sargs] modes_keys_list = sort_keys(modes_keys_list) # Construct the data array. data_array = [] for key in modes_keys_list: # data_item = np.array(wfile[key]) # print(data_item.shape) data_item = np.array(wfile[key])[nghosts : nphi - nghosts, nghosts : ntheta - nghosts] try: data_item = data_item["real"] + 1j * data_item["imag"] except Exception as ex: message(ex) pass data_array.append(data_item) self.data = np.transpose(np.array(data_array), (2, 1, 0)) self.iteration_axis = np.array(iteration_numbers) ######################################################### # Load inv_coords data ######################################################### inv_file_name = f"qlm_inv_z[{bh}].xy.h5" # get the full path. full_path = self.data_dir + inv_file_name # Open the modes file. with h5py.File(full_path, "r") as wfile: # Get all the mode keys. modes_keys_list = list(wfile.keys()) # modes_keys_list = sorted(modes_keys_list) # Get the mode keys containing the data. modes_keys_list = [item for item in modes_keys_list if "it=" in item] modes_keys_list = sort_keys(modes_keys_list) data_array = [] for key in modes_keys_list: data_item = np.array(wfile[key])[nghosts : nphi - nghosts, nghosts : ntheta - nghosts] # data_item = data_item['real'] + 1j*data_item['imag'] data_array.append(data_item) self.invariant_coordinates_data = np.transpose(np.array(data_array), (2, 1, 0)) ######################################################### # Load metric determinant data ######################################################### twometric_qtt_file_name = f"qlm_qtt[{bh}].xy.h5" twometric_qtp_file_name = f"qlm_qtp[{bh}].xy.h5" twometric_qpp_file_name = f"qlm_qpp[{bh}].xy.h5" # set the full path. full_path = self.data_dir + twometric_qtt_file_name # Open the modes file. with h5py.File(full_path, "r") as wfile: # Get all the mode keys. modes_keys_list = list(wfile.keys()) # modes_keys_list = sorted(modes_keys_list) # Get the mode keys containing the data. modes_keys_list = [item for item in modes_keys_list if "it=" in item] modes_keys_list = sort_keys(modes_keys_list) qtt_data_array = [] for key in modes_keys_list: data_item = np.array(wfile[key])[nghosts : nphi - nghosts, nghosts : ntheta - nghosts] # data_item = data_item['real'] + 1j*data_item['imag'] qtt_data_array.append(data_item) qtt_data_array = np.array(qtt_data_array) qtt_data_array = np.transpose(qtt_data_array, (2, 1, 0)) # set the full path. full_path = self.data_dir + twometric_qtp_file_name # Open the modes file. with h5py.File(full_path, "r") as wfile: # Get all the mode keys. modes_keys_list = list(wfile.keys()) # modes_keys_list = sorted(modes_keys_list) # Get the mode keys containing the data. modes_keys_list = [item for item in modes_keys_list if "it=" in item] modes_keys_list = sort_keys(modes_keys_list) qtp_data_array = [] for key in modes_keys_list: data_item = np.array(wfile[key])[nghosts : nphi - nghosts, nghosts : ntheta - nghosts] # data_item = data_item['real'] + 1j*data_item['imag'] qtp_data_array.append(data_item) qtp_data_array = np.array(qtp_data_array) qtp_data_array = np.transpose(qtp_data_array, (2, 1, 0)) # set the full path. full_path = self.data_dir + twometric_qpp_file_name # Open the modes file. with h5py.File(full_path, "r") as wfile: # Get all the mode keys. modes_keys_list = list(wfile.keys()) # modes_keys_list = sorted(modes_keys_list) # Get the mode keys containing the data. modes_keys_list = [item for item in modes_keys_list if "it=" in item] modes_keys_list = sort_keys(modes_keys_list) qpp_data_array = [] for key in modes_keys_list: data_item = np.array(wfile[key])[nghosts : nphi - nghosts, nghosts : ntheta - nghosts] # data_item = data_item['real'] + 1j*data_item['imag'] qpp_data_array.append(data_item) qpp_data_array = np.array(qpp_data_array) qpp_data_array = np.transpose(qpp_data_array, (2, 1, 0)) sqrt_met_det = np.sqrt( np.linalg.det( np.transpose(np.array([[qtt_data_array, qtp_data_array], [qtp_data_array, qpp_data_array]]), (2, 3, 4, 0, 1)) ) ) self.sqrt_met_det_data = sqrt_met_det
[docs] def to_shear_modes_array(self, grid_info=None, spin_weight=None, ell_max=8): """Decompose a given spherical array function on a sphere into Spin Weighted Spherical Harmonic modes. Parameters ---------- spin_weight: int, optional The spin weight of the waveform. It defaults to -2 for a gravitational waveform. ell_max: int, optional The maximum value of the :math:`\\ell` polar quantum number. Defaults to 8. grid_info: class instance The class instance that contains the properties of the spherical grid. Returns ------- waveforms_modes: modes_array An instance of the `modes_array` class containing the decomposed modes. Notes ----- 1. Assumes that the sphere on which this decomposition is carried out is so far out that the coordinate system is spherical polar on a round sphere. 2. Assumes that the poper area is the same as its co-ordinate area. 3. Ensure that the label of the input spherical array indicates whether it is a time domain data or frequency domain data. """ if grid_info is None: if self.grid_info is None: message("Please specify the grid specs. Assuming defaults.") grid_info = spherical_grid() self.grid_info = grid_info else: grid_info = self.grid_info if spin_weight is None: if self.spin_weight is None: message("Please specify spin weight. Assuming 2") spin_weight = 2 self.spin_weight = spin_weight else: spin_weight = self.spin_weight spin_weight = abs(spin_weight) # Compute the meshgrid for theta and phi. theta, phi = grid_info.meshgrid # Create a modes array object # Create a modes list. modes_list = construct_mode_list(ell_max, spin_weight=spin_weight) if not self.label: label = "decomposed_time_domain" else: label = self.label # Create a mode array for the decomposed_waveform waveform_modes = modes_array(label=label, ell_max=ell_max, spin_weight=spin_weight) # Inherit the time or frequency axis. if "time" in label: # axis = self.time_axis waveform_modes.time_axis = self.time_axis else: # axis = self.frequency_axis waveform_modes.frequency_axis = self.frequency_axis # Create the modes_array waveform_modes.time_axis = self.time_axis[:] # sargs = np.argsort(waveform_modes.time_axis) # print(sargs) waveform_modes.time_axis = waveform_modes.time_axis waveform_modes.create_modes_array(ell_max=ell_max, data_len=self.data_len) waveform_modes.modes_list = modes_list # The area element on the sphere # Compute the meshgrid for theta and phi. theta, phi = grid_info.meshgrid phi = np.transpose(np.array([phi for index in range(len(self.time_axis))]), (1, 2, 0)) # sqrt_met_det = np.sin(theta) # sqrt_met_det = np.sqrt(np.power(np.sin(theta), 2)) darea = self.sqrt_met_det_data * grid_info.dtheta * grid_info.dphi theta = np.emath.arccos(self.invariant_coordinates_data) modes_list = [item for item in modes_list if item[0] >= spin_weight] for mode in modes_list: ell_value, all_emm_values = mode for emm_value in all_emm_values: # m value. # print(f'Processing l{ell_value} m{emm_value}') # Spin weighted spherical harmonic function at (theta, phi) Ybasis_fun = np.conj( Yslm_vec(spin_weight=spin_weight, ell=ell_value, emm=emm_value, theta_grid=theta, phi_grid=phi) ) # Ybasis_fun = np.array([np.conj(Yslm_vec(spin_weight=spin_weight, # ell=ell_value, emm=emm_value, theta_grid=theta[:, :, index], # phi_grid=phi[:, :, index])) for index in range(self.data_len)]) # Ybasis_fun = np.transpose(Ybasis_fun, (1, 2, 0)) # print('Ybasis_fun', Ybasis_fun.shape) Ydarea = Ybasis_fun * darea # print('Ydarea', Ydarea.shape) # print(full_integrand) # Using quad # print('self.data', self.data.shape) # multipole_ell_emm = np.tensordot(self.data, Ydarea, axes=((0, 1), (0, 1))) multipole_ell_emm = np.sum(self.data * Ydarea, (0, 1)) # print(f'l{ell_value}m{emm_value}', multipole_ell_emm) # print('multipole_ell_emm', multipole_ell_emm.shape) waveform_modes.set_mode_data(ell_value, emm_value, data=multipole_ell_emm) return waveform_modes
# @jitclass(spec_ma)
[docs]class modes_array: """A class that holds mode array of waveforms Attributes ---------- label: str The label of the modes array. r_ext: float The extraction radius. modes_list: list The list of available modes in the format [l1, [m values], [l2, [m values], ...]] ell_max: int The maximum :math:`\\ell` mode available. modes_data: 3d array The three dimensional array containing modes in time/frequency space. The axis of the array is (:math:`\\ell`, :math:`m`, time/freq). base_dir: str The base directory containing the modes data. data_dir: str The subdirectory in which to look for the data. filename: str The filename containg the modes data. Methods ------- get_metadata: Get the metadata associated with the modes_array. mode: Get the data for the given :math:`\\ell, m` mode. create_modes_array: A private method to create an empty modes_array of given shape. delta_t: Set the attribute `delta_t` and/ or return its value. load_modes: Load the waveform modes from a specified h5 file. save_modes: Save the waveform modes to a specified h5 file. set_mode_data: Set the `mode` data of specified modes. to_frequency_basis: Get the `modes_array` in frequency basis from its time basis representation. to_time_basis: Get the `modes_array` in temporal basis from its frequency basis representation. extrap_to_inf: Extrapolate the modes to infinity. supertranslate: Supertranslate the waveform modes. boost: Boost the waveform modes. """ def __init__( self, data_dir=None, file_name=None, modes_data=None, time_axis=None, frequency_axis=None, key_format=None, ell_max=None, modes_list=None, label=None, r_ext=500, out_file_name=None, maxtime=None, date=None, time=None, key_ex=None, spin_weight=-2, ): self.label = label self.data_dir = data_dir self.file_name = file_name self.modes_data = modes_data self.key_format = key_format self.ell_max = ell_max self.modes_list = modes_list self.r_ext = r_ext self.time_axis = time_axis self.frequency_axis = frequency_axis self.out_file_name = out_file_name self.maxtime = maxtime self.date = date self.time = time self.key_ex = key_ex self.spin_weight = spin_weight
[docs] def get_metadata(self): """Get the metadata associated with the instance. Returns ------- metadata: dict A dictionary of metedata. """ # The metadata dict unnecessary_keys = ["_time_axis", "_modes_data", "freq_axis"] # Get all attributes # metadata = self.__dict__ metadata = {} for key, val in self.__dict__.items(): if key in unnecessary_keys: pass else: metadata.update({key: val}) # for item in unnecessary_keys: # metadata.pop(item, None) # self.metadata = metadata return metadata
[docs] def mode(self, ell, emm): """The modes array data structure to hold waveform modes. Parameters ---------- ell: int The :math:`\\ell` index of the mode. emm: int The :math:`m` index of the mode. Returns ------- mode_data: array The array of the requested mode. """ emm_index = ell + emm return self.modes_data[ell, emm_index, :]
@property def time_axis(self): """The time axis""" return np.array(self._time_axis) @time_axis.setter def time_axis(self, time_axis): self._time_axis = time_axis @property def modes_data(self): """The modes array""" return np.array(self._modes_data) @modes_data.setter def modes_data(self, modes_data): self._modes_data = modes_data
[docs] def create_modes_array(self, ell_max=None, data_len=None): """Create a modes array and initialize it with zeros. Parameters ---------- ell_max: int The maximum :math:`\\ell` value of the modes. data_len: int The number of points along the third (time / frequency) axis. Returns ------- self.modes_array: modes_array sets the `self.modes_array` attribute. """ import time import datetime # Check ell_max if ell_max is None: try: ell_max = self.ell_max except Exception as ex: message(ex) raise NameError("Please supply ell_max") if data_len is None: try: data_len = self.data_len except Exception as ex: message(ex) raise NameError("Please supply data_len") if self.modes_list is None: self.modes_list = construct_mode_list(ell_max=ell_max, spin_weight=self.spin_weight) # self.modes_data = np.zeros([ell_max + 1, 2 * (ell_max + 1) + 1, data_len], dtype=np.complex128) self.modes_data = np.zeros((ell_max + 1, 2 * (ell_max + 1) + 1, data_len), dtype=np.complex128) # Set the time metadata time_now = time.localtime() time_now = time.strftime("%H:%M:%S", time_now) date_now = str(datetime.date.today()) if self.time is None: # Assign time and date stamp if it doesnt exist self.time = time_now self.date = date_now
@property def data_len(self): """Returns the length of the time/frequency axis. Returns ------- data_len: int THe length of the time/frequency axis. """ try: data_len = len(self.time_axis) except Exception as ex: message(ex) data_len = len(self.frequency_axis) return data_len @data_len.setter def data_len(self, data_len): self._data_len = data_len
[docs] def delta_t(self, value=None): """Sets and returns the value of time stepping :math:`dt`. Parameters ---------- value : float, optional The value of :math:`dt` to set to the attribute. Returns ------- self.delta_t: float Sets the attribute. """ if not value: try: delta_t = self.time_axis[1] - self.time_axis[0] except Exception as ex: message("Please input the value of `delta_t` or supply the `time_axis` to the waveform.", ex) else: delta_t = value return delta_t
@property def delta_f(self, value=None): """Sets and returns the value of frequency stepping :math:`df`. Parameters ---------- value : float, optional The value of :math:`df` to set to the attribute. Returns ------- delta_f: float Sets the attribute. """ # if not self.delta_t: if not value: try: delta_f = self.frequency_axis[1] - self.frequency_axis[0] except Exception as ex: message("Please input the value of `delta_f` or supply the `frequency_axis` to the waveform.", ex) else: delta_f = value return delta_f
[docs] def load_modes( self, label=None, data_dir=None, file_name=None, ftype="generic", var_type="Psi4", resam_type="finest", interp_kind="cubic", extrap_order=4, r_ext=None, ell_max=None, pre_key=None, modes_list=None, crop=False, centre=True, key_ex=None, save_as_ma=False, compression_opts=None, r_ext_factor=1, debug=False, ): """Load the waveform mode data from an hdf file. Parameters ---------- extrap_order : int, optional For SpEC waveforms only. This is the order of extrapolation to use. pre_key: str, optional A string containing the key of the group in the HDF file in which the modes` dataset exists. It defaults to `None`. mode_numbers: list The mode numbers to load from the file. Each item in the list is a list that contains two integrer numbers, one for the mode index :math:`\\ell` and the other for the mode index :math:`m`. crop: bool Whether or not to crop the beginning of the input waveform. If yes, the first :math:`t=r_{ext}` length of data will be discarded. Returns ------- waveform_obj: 3d array Sets the three dimensional array `waveform.modes` that contains the required :math:`\\ell, m` modes. Examples -------- # Note # Update this! #>>> from waveformtools.waveforms import waveform #>>> waveform.set_basedir('./') #>>> waveform.set_datadir('data/') #>>> mode_numbers = [[2, 2], [3, 3]] #>>> waveform.load_data(mode_numbers=mode_numbers) """ if debug is False: wfs_nl = 1 if not data_dir: data_dir = self.data_dir else: self.data_dir = data_dir if not file_name: file_name = self.file_name else: self.file_name = file_name if not ell_max: ell_max = self.ell_max else: self.ell_max = ell_max if not label: label = self.label # if self.data_dir is not None: # data_dir = self.data_dir # if self.file_name is not None: # file_name = self.file_name print("Passing", data_dir, file_name) if ftype == "generic": dataIO.load_gen_data_from_disk( self, label, data_dir, file_name, r_ext, ell_max, pre_key, modes_list, crop, centre, key_ex, r_ext_factor ) elif ftype == "RIT": if var_type == "Psi4": dataIO.load_RIT_Psi4_data_from_disk( wfa=self, data_dir=data_dir, file_name=file_name, resam_type=resam_type, interp_kind=interp_kind, r_ext=r_ext, ell_max=ell_max, pre_key=pre_key, modes_list=modes_list, crop=crop, centre=centre, key_ex=key_ex, r_ext_factor=r_ext_factor, ) elif var_type == "Strain": # print(file_name) dataIO.load_RIT_Strain_data_from_disk( self, data_dir=data_dir, file_name=file_name, label=label, resam_type=resam_type, interp_kind=interp_kind, ell_max=ell_max, save_as_ma=save_as_ma, modes_list=modes_list, crop=crop, centre=centre, r_ext_factor=r_ext_factor, debug=debug, ) else: message(f"Data {ftype} {var_type} not supported yet!") sys.exit(0) elif ftype == "SpEC": wfs_nl = dataIO.load_SpEC_data_from_disk( self, label, data_dir, file_name, extrap_order, r_ext, ell_max, centre, modes_list, save_as_ma, resam_type, interp_kind, compression_opts, r_ext_factor, debug, ) elif ftype == "SpECTRE": dataIO.load_SpECTRE_data_from_disk( self, label, data_dir, file_name, r_ext, ell_max, centre, modes_list, save_as_ma, resam_type, interp_kind, compression_opts, r_ext_factor, ) else: message(f"Data {ftype} {var_type} not supported yet!") sys.exit(0) return wfs_nl
[docs] def save_modes( self, ell_max=None, pre_key=None, key_format=None, modes_to_save=None, out_file_name="mp_new_modes.h5", r_ext_factor=None, compression_opts=0, r_ext=None, ): """Save the waveform mode data to an hdf file. Parameters ---------- pre_key: str, optional A string containing the key of the group in the HDF file in which the modes` dataset exists. It defaults to `None`. mode_numbers: list The mode numbers to load from the file. Each item in the list is a list that contains two integrer numbers, one for the mode index :math:`\\ell` and the other for the mode index :math:`m`. Returns ------- waveform_obj: 3d array Sets the three dimensional array `waveform.modes` that contains the required :math:`\\ell, m` modes. Examples -------- >>> from waveformtools.waveforms import waveform >>> waveform.set_basedir('./') >>> waveform.set_datadir('data/') >>> mode_numbers = [[2, 2], [3, 3]] >>> waveform.load_data(mode_numbers=mode_numbers) """ # import dataIO from waveformtools import dataIO dataIO.save_modes_data_to_gen( self, ell_max=None, pre_key=None, key_format=None, modes_to_save=None, out_file_name="mp_new_modes.h5", r_ext_factor=None, compression_opts=0, r_ext=None, )
[docs] def set_mode_data(self, ell_value, emm_value, data): """Set the mode array data for the respective :math:`(\\ell, m)` mode. Parameters ---------- ell_value: int The :math:`\\ell` polar mode number. emm_value: int The :math:`emm` azimuthal mode number. data: 1d array The array consisting of mode data for the requested mode. Returns ------- self.mode_data: modes_data The updated modes data. """ # Compute the emm index given ell. emm_index = emm_value + ell_value # Set the mode data. self._modes_data[ell_value, emm_index] = data
[docs] def to_spherical_array(self, grid_info, meth_info, spin_weight=None): """Obtain the spherical array from the modes array. Parameters ---------- grid_info: cls instance An instance of the "spherical_grid" class to hold the grid info. meth_info : cls instance An instance of the class waveformtools.diagnostics.method_info that provides information on what methods to use for integration. Returns ------- waveform_sp: spherical_array A member of the "spherical_array" class constructed from the given "modes_rray". """ # Create a spherical array. waveform_sp = spherical_array(label=self.label, grid_info=grid_info) if spin_weight is None: if self.spin_weight is not None: spin_weight = self.spin_weight else: spin_weight = -2 self.spin_weight = spin_weight spin_weight = abs(spin_weight) waveform_sp.spin_weight = spin_weight # Set the time-axis try: waveform_sp.time_axis = self.time_axis except Exception as ex: message(ex) waveform_sp.frequency_axis = self.frequency_axis # Get the coordinate meshgrid theta, phi = grid_info.meshgrid s1, s2 = theta.shape s3 = self.data_len sp_data = np.zeros((s1, s2, s3), dtype=np.complex128) modes_list = [item for item in self.modes_list if item[0] >= spin_weight] for item in modes_list: # Get modes. ell, emm_list = item # if ell<spin_weight: # continue for emm in emm_list: # For every l, m sp_data += np.multiply.outer( Yslm_vec(spin_weight, ell=ell, emm=emm, theta_grid=theta, phi_grid=phi), self.mode(ell, emm) ) # print(sp_data) # Set the data of the spherical array. waveform_sp.data = sp_data try: waveform_sp.time_axis = self.time_axis except Exception as ex: message(ex) waveform_sp.frequency_axis = self.frequency_axis return waveform_sp
[docs] def trim(self, trim_upto_time=None): """Trim the modes_array at the beginning and center about the peak of the 2,2 mode. Parameters ---------- time: float The time unit upto which to discard. Returns ------- Re-sets the `time_axis` and `modes_array` data. """ if trim_upto_time is None: trim_upto_time = self.r_ext # Compute the start index start = int(trim_upto_time / self.delta_t()) # Trim the time axis self._time_axis = self.time_axis[start:] # Trim the data self._modes_data = self.modes_data[:, :, start:] # Recenter the time axis max_ind = np.argmax(np.absolute(self.mode(2, 2))) self._time_axis = self.time_axis - self.time_axis[max_ind]
[docs] def to_frequency_basis(self): """Compute the modes in frequency basis. Returns ------- waveform_tilde_modes: modes_array A modes_array containing the modes in frequency basis. """ # Create a new modes array waveform_tilde_modes = modes_array(label=f"{self.label} -> frequency_domain") waveform_tilde_modes.create_modes_array(ell_max=self.ell_max, data_len=self.data_len) from waveformtools.transforms import compute_fft for mode in self.modes_list: # Extrapolate every mode # Ge the ell value ell_value, emm_list = mode for emm_value in emm_list: freq_axis, freq_data = compute_fft(self.mode(ell_value, emm_value), self.delta_t()) waveform_tilde_modes.set_mode_data(ell_value, emm_value, freq_data) waveform_tilde_modes.frequency_axis = freq_axis waveform_tilde_modes.ell_max = self.ell_max waveform_tilde_modes.modes_list = self.modes_list return waveform_tilde_modes
[docs] def to_time_basis(self): """Compute the modes in time basis. Returns ------- waveform_modes : modes_array A modes_array containing the modes in frequency basis. """ # Create a new modes array waveform_modes = modes_array(label=f"{self.label} -> time_domain") waveform_modes.create_modes_array(ell_max=self.ell_max, data_len=self.data_len) from waveformtools.transforms import compute_ifft for mode in self.modes_list: # Extrapolate every mode # Ge the ell value ell_value, emm_list = mode for emm_value in emm_list: time_axis, time_data = compute_ifft(self.mode(ell_value, emm_value), self.delta_f) waveform_modes.set_mode_data(ell_value, emm_value, time_data) try: maxloc = np.argmax(np.absolute(waveform_modes.mode(2, 2))) except Exception as ex: message(ex) maxloc = 0 maxtime = time_axis[maxloc] waveform_modes.time_axis = time_axis - maxtime return waveform_modes
[docs] def extrap_to_inf(self, mass=1, spin=None, modes_list="all", method="SIO", r_ext_factor=1): """Extrapolate the :math:`\\Psi_4` modes to infinity using the perturbative improved second order method. Parameters ---------- mass: float The effective total mass of the system. spin: float The effective spin of the system. modes: modes array, optional The modes to extrapolate. Defaults to `all` if not specified. method: str The method to use for extrapolation. The available methods are: +------------+--------------------------------------+ | Method str | Name | +------------+--------------------------------------+ |'FO' | First order | |'SO' | Second order | |'SIO' | Second improved order | |'NM' | Numerical method (not ready yet) | +------------+--------------------------------------+ Returns ------- waveform_inf_modes: modes array A new modes array that contains the extrapolated modes. """ from functools import partial # Prepare the extrapolating method. if method == "SIO": from waveformtools.extrapolate import waveextract_to_inf_perturbative_twop5_order extrap_method = partial( waveextract_to_inf_perturbative_twop5_order, delta_t=self.delta_t(), areal_radius=self.r_ext, mass=mass, spin=spin, ) if method == "SO": from waveformtools.extrapolate import waveextract_to_inf_perturbative_two_order extrap_method = partial( waveextract_to_inf_perturbative_two_order, delta_t=self.delta_t(), areal_radius=self.r_ext, mass=mass, spin=spin, ) if method == "FO": from waveformtools.extrapolate import waveextract_to_inf_perturbative_one_order extrap_method = partial( waveextract_to_inf_perturbative_one_order, u_ret=self.time_axis, areal_radius=self.r_ext, mass=mass ) if method == "NM": print("This method is not available yet! ") # Prepare the modes to be extrapolated. if modes_list == "all": modes_list = construct_mode_list(self.ell_max) # Create a mode array for the extrapolated waveform. extrap_wf = modes_array( label=f"{self.label} -> rPsi4_inf", modes_list=self.modes_list, ell_max=self.ell_max, r_ext=self.r_ext ) extrap_wf.create_modes_array(ell_max=self.ell_max, data_len=self.data_len) # Retain the time axis. extrap_wf.time_axis = self.time_axis for mode in modes_list: # Extrapolate every mode # Ge the ell value ell_value, emm_list = mode for emm_value in emm_list: # For every emm value print(f"Processing l{ell_value}, m{emm_value}") # Compute rPsi4_lm mode_data = r_ext_factor * self.mode(ell_value, emm_value) # Extrapolate # import ipdb; ipdb.set_trace() extrap_mode_data = extrap_method(rPsi4_rlm=mode_data) # Assign data to new modes array extrap_wf.set_mode_data(ell_value, emm_value, extrap_mode_data) print("Done!") return extrap_wf
[docs] def supertranslate(self, supertransl_alpha_modes, grid_info, order=4): """Supertranslate the :math:`\\Psi_{4\\ell m}` waveform modes, give the, the supertranslation parameter and the order. Parameters ---------- supertransl_alpha_modes : modes_array The modes_array containing the supertranslation mode coefficients. grid_info: class instance The class instance that contains the properties of the spherical grid using which the computations are carried out. order: int The number of terms to use to approximate the supertranslation. Returns ------- Psi4_supertransl: modes_array A class instance containg the modes of the supertranslated waveform :math:`\\Psi_4`. """ import BMS ell_max = self.ell_max # Step 0: Get the grid properties for integrations # Compute the meshgrid for theta and phi. theta, phi = grid_info.meshgrid # theta # Step 1: get the grid function for supertranslation parameter supertransl_alpha_sphere = BMS.compute_supertransl_alpha(supertransl_alpha_modes, theta, phi) # The supertranslation is carried out in frequency space. # Step 2: get the FFT of the mode coefficients of Psi4 Psi4_tilde_modes = self.to_frequency_basis() # Get the 2d angular frequency array omega_axis_2d = Psi4_tilde_modes.omega # Construct the supertranslation factor supertransl_factor = np.sum( np.array([np.power((-1j * omega_axis_2d * supertransl_alpha_sphere), index) for index in range(order)]), axis=0, ) # Multiply with the fourier modes. supertransl_spherical_factor = Psi4_tilde_modes.multiply(supertransl_factor) # Reconstruct the modes # Check! supertransl_spherical_grid = np.zeros(supertransl_spherical_factor.shape, dtype=np.complex128) for ell_value in range(ell_max + 1): for emm_value in range(-ell_value, ell_value + 1): # Multiply with the SWSH basis. supertransl_spherical_grid += supertransl_spherical_factor * Yslm_vec( spin_weight=-2, ell=ell_value, emm=emm_value, theta=theta, phi=phi ) # Step 3: Reconstruct the function on the sphere # Create a spherical_array to hold the supertranslated waveform supertransl_spherical_waveform = spherical_array(grid_info=grid_info) # Set the data. supertransl_spherical_waveform.data = supertransl_spherical_grid # Get modes_array from spherical_array Psi4_supertransl_modes = supertransl_spherical_waveform.to_modes_array(ell_max=ell_max) return Psi4_supertransl_modes
[docs] def boost(self, conformal_factor, grid_info=None): """Boost the waveform given the unboosted waveform and the boost conformal factor. Parameters ---------- conformal_factor: 2d array The conformal factor for the Lorentz transformation. It may be a single floating point number or an array on a spherical grid. The array will be of dimensions [ntheta, nphi] Returns ------- boosted_waveform: spherical_array The class instance `spherical_array` that contains the boosted waveform. """ from waveformtools.grids import spherical_grid # Construct a spherical grid. if grid_info is None: grid_info = spherical_grid() # Get spherical array from modes. unboosted_waveform = self.to_spherical_array(grid_info) boosted_waveform_data = unboosted_waveform.boost(conformal_factor) # Construct a 2d waveform array object boosted_waveform = spherical_array(grid_info=unboosted_waveform.grid_info, data=np.array(boosted_waveform_data)) boosted_waveform.label = "boosted" # Get modes from spherical data. # boosted_waveform_modes = boosted_waveform.to_modes_array() # return boosted_waveform_modes return boosted_waveform
[docs] def get_strain_from_psi4(self, omega0="auto"): """Get the strain `modes_array` from :math:`\\Psi_4` by fixed frequency integration. Parameters ---------- omega0: float, optional The lower cutoff angular frequency for FFI. By default, it computes this from the mode data. Returns ------- strain_waveform: modes_array The computed strain modes. """ # Initialize a mode array for strain. strain_waveform = modes_array( label="{} strain from Psi4".format(self.label), r_ext=self.r_ext, ell_max=8, modes_list=self.modes_list ) strain_waveform.time_axis = self.time_axis strain_waveform.ell_max = self.ell_max data_len = self.data_len strain_waveform.create_modes_array(ell_max=self.ell_max, data_len=data_len) # Integrate, from waveformtools.integrate import fixed_frequency_integrator from waveformtools.waveformtools import get_starting_angular_frequency as sang_f omega_st = omega0 for item in self.modes_list[:]: ell, emm_list = item for emm in emm_list: mode_data = self.mode(ell, emm) if omega0 == "auto": omega_st = abs(sang_f(mode_data, delta_t=self.delta_t())) / 10 strain_mode_data, _ = fixed_frequency_integrator( udata_time=mode_data, delta_t=self.delta_t(), omega0=omega_st, order=2 ) strain_waveform.set_mode_data(ell, emm, data=strain_mode_data) return strain_waveform
[docs] def get_news_from_psi4(self, omega0="auto"): """Get the News `modes_array` from :math:`\\Psi_4` by fixed frequency integration. Parameters ---------- omega0: float, optional The lower cutoff angular frequency for FFI. By default, it computes this as one tenth of the starting frequency of the mode data. Returns ------- news_waveform: modes_array The computed strain modes. """ # Initialize a mode array for strain. # news_waveform = modes_array(label=f'{self.label} news from Psi4', r_ext=500, ell_max=8, modes_list=self.modes_list) news_waveform = modes_array( label="{} news from Psi4".format(self.label), r_ext=500, ell_max=8, modes_list=self.modes_list ) news_waveform.time_axis = self.time_axis news_waveform.ell_max = self.ell_max data_len = self.data_len news_waveform.create_modes_array(ell_max=self.ell_max, data_len=data_len) # Integrate, from waveformtools.integrate import fixed_frequency_integrator from waveformtools.waveformtools import get_starting_angular_frequency as sang_f omega_st = omega0 for item in self.modes_list[:]: ell, emm_list = item for emm in emm_list: mode_data = self.mode(ell, emm) if omega0 == "auto": omega_st = abs(sang_f(mode_data, delta_t=self.delta_t())) / 10 news_mode_data, _ = fixed_frequency_integrator( udata_time=mode_data, delta_t=self.delta_t(), omega0=omega_st, order=1 ) news_waveform.set_mode_data(ell, emm, data=news_mode_data) return news_waveform
[docs] def taper(self, zeros="auto"): """Taper a waveform at both ends and insert zeros if needed Parameters ---------- zeros: int The number of zeros to add at rach end Returns ------- tapered_modes: modes_array Modes array with tapered mode data. """ from waveformtools.waveformtools import taper if zeros == "auto": # Decide the number of zeros data_len = self.data_len nearest_power = int(np.log(data_len) / np.log(2)) req_len = np.power(2, nearest_power + 1) zeros = req_len - data_len print("num_zeros", zeros) # New modes array. tapered_modes = None for item in self.modes_list[:]: ell, emm_list = item for emm in emm_list: input_data_re = self.mode(ell, emm).real input_data_im = self.mode(ell, emm).imag tapered_data_re = taper(input_data_re, zeros=zeros) tapered_data_im = taper(input_data_im, zeros=zeros) # tapered_data_re = taper_tanh(input_data_re, delta_t=self.delta_t()) # tapered_data_im = taper_tanh(input_data_im, delta_t=self.delta_t()) new_data_len = len(tapered_data_re) if tapered_modes is None: tapered_modes = modes_array( label="tapered {}".format(self.label), r_ext=self.r_ext, modes_list=self.modes_list, ell_max=self.ell_max, ) tapered_modes.create_modes_array(ell_max=self.ell_max, data_len=new_data_len) tapered_data = tapered_data_re + 1j * tapered_data_im # print(len(tapered_data_re)) tapered_modes.set_mode_data(ell, emm, data=tapered_data) # Set the time axis new_time_axis = np.arange(0, new_data_len * self.delta_t(), self.delta_t()) tapered_modes.time_axis = new_time_axis # Recenter the modes. tapered_modes.trim(trim_upto_time=0) return tapered_modes
[docs] def taper_tanh(self, time_axis=None, zeros="auto", duration=10, sides="both"): """Taper a waveform at both ends and insert zeros if needed Parameters ---------- zeros: int The number of zeros to add at rach end Returns ------- tapered_modes: modes_array Modes array with tapered mode data. """ from waveformtools.waveformtools import taper_tanh if zeros == "auto": # Decide the number of zeros data_len = self.data_len nearest_power = int(np.log(data_len) / np.log(2)) req_len = np.power(2, nearest_power + 1) zeros = req_len - data_len # print('num_zeros', zeros) # New modes array. tapered_modes = None for item in self.modes_list[:]: ell, emm_list = item for emm in emm_list: input_data_re = self.mode(ell, emm).real input_data_im = self.mode(ell, emm).imag # tapered_data_re = taper(input_data_re, zeros=zeros) # tapered_data_im = taper(input_data_im, zeros=zeros) _, tapered_data_re = taper_tanh(input_data_re, delta_t=self.delta_t(), duration=duration, sides=sides) _, tapered_data_im = taper_tanh(input_data_im, delta_t=self.delta_t(), duration=duration, sides=sides) new_data_len = len(tapered_data_re) if tapered_modes is None: tapered_modes = modes_array( label="tapered {}".format(self.label), r_ext=self.r_ext, modes_list=self.modes_list, ell_max=self.ell_max, ) tapered_modes.create_modes_array(ell_max=self.ell_max, data_len=new_data_len) tapered_data = tapered_data_re + 1j * tapered_data_im # print(len(tapered_data_re)) tapered_modes.set_mode_data(ell, emm, data=tapered_data) # Set the time axis new_time_axis = np.arange(0, new_data_len * self.delta_t(), self.delta_t()) tapered_modes.time_axis = new_time_axis # Recenter the modes. tapered_modes.trim(trim_upto_time=0) return tapered_modes
[docs] def low_cut(self, omega0=0.03, order=2): """Apply the low cut filter from waveformtools.low_cut_filter Parameters ---------- order: int, optional The order of the butterworth filter. omega0: float, optional The cutoff frequency of the butterworth filter. Returns: -------- filtered_modes: modes_array A modes array object containing filtered modes. """ # modes_array for filtered data. filtered_modes = None # Import the filter from waveformtools.waveformtools import low_cut_filter for item in self.modes_list: # Iterate over available modes. ell, emm_list = item for emm in emm_list: if filtered_modes is None: # Create filtered_modes filtered_modes = modes_array( label="lc filtered {}".format(self.label), r_ext=self.r_ext, modes_list=self.modes_list, ell_max=self.ell_max, ) filtered_modes.create_modes_array(ell_max=self.ell_max, data_len=self.data_len) # Get filtered mode data. filtered_data = low_cut_filter(self.mode(ell, emm), self.frequency_axis, omega0=omega0, order=order) # Set the mode data. filtered_modes.set_mode_data(ell, emm, data=filtered_data) # Set the f axis. filtered_modes.frequency_axis = self.frequency_axis return filtered_modes
[docs] def to_td_waveform(self, Mtotal=1, incl_angle=0, distance=1, delta_t=None): """Get the plus and cross polarizations of of the waveform time series by summing the modes. Parameters ---------- incl_angle : float The inclination angle. distance : float Returns ------- taxis, hp, hc : 1d arrays The 1d arrays of the time axis and the polarizations of the waveforms. """ from waveformtools.transforms import Yslm # print(Yslm) th = incl_angle ph = 0 wts = np.zeros(self.data_len, dtype=np.complex128) for ell, emm_list in self.modes_list: for emm in emm_list: Alm = self.mode(ell, emm) Y = Yslm(self.spin_weight, ell, emm, th, ph) # print(Y) # print(Alm) wts = wts + Alm * Y taxis = self.time_axis * tuc * Mtotal wts = Mtotal * wts / (distance * dMpc * muc) if isinstance(delta_t, float): # Resample the waveform new_taxis = np.arange(taxis[0], taxis[-1], delta_t) from waveformtools.waveformtools import interp_resam_wfs wts = interp_resam_wfs(wts, taxis, new_taxis) taxis = new_taxis return taxis, wts.real, wts.imag
####################################################################################################### def _get_modes_list_from_keys(keys_list, r_ext): """Get the modes list from the keys list of an hdf file. Parameters ---------- keys_list: list The list containing all the keys r_ext: float The extraction radius of the data. Returns ------- modes_list: list The list of modes. """ # Sort the keys to ensure a nice # modes list structure. keys_list_orig = sorted(keys_list) if r_ext != -1: keys_list = [item for item in keys_list_orig if f"r{r_ext}" in item] if keys_list == []: print("Got an empty list. Searching for r_ext value in key string") keys_list = [item for item in keys_list_orig if f"{r_ext}" in item] # print('List of keys received', keys_list) # The list of modes. modes_list = [] # Initialize the emm modes sublist. emm_modes_for_ell = [] # Present ell value to # initialize the mode concatenation. ell_old = 0 for key in keys_list: # print(key) # Get the ell value ell_value, emm_value = _get_ell_emm_from_key(key) if ell_value != ell_old: # If the ell value has changed, # update the modes list before moving # onto the next ell value. modes_list.append([ell_old, emm_modes_for_ell]) # The present ell value is the old # ell value. ell_old = ell_value # Reset the ell_mode list. emm_modes_for_ell = [] # Update it with the new emm mode. emm_modes_for_ell.append(emm_value) modes_list.append([ell_value, emm_modes_for_ell]) return modes_list def _get_ell_emm_from_key(key): """Get the :math:`\\ell` and :math:`m` values from a given key string of an hdf file. Parameters ---------- key: str The input key string Returns ------- ell_value: int The :math:`\\ell` value emm_value: int The :math:`m` value. Notes ----- Assumes that the input string has :math:`\\ell` and :math:`m` values in the form `l{int}m{int}`. """ import re str_match = re.search("l\d*", key) ell_str_start = str_match.start() ell_str_end = str_match.end() ell_value = int(key[ell_str_start + 1 : ell_str_end]) str_match = re.search("m-*\d*", key) emm_str_start = str_match.start() emm_str_end = str_match.end() emm_value = int(key[emm_str_start + 1 : emm_str_end]) return ell_value, emm_value
[docs]def get_iteration_numbers_from_keys(keys_list): """Get the iteration number from keys. Parameters ---------- keys_list: list The list of keys. Returns ------- iteration_numbers: list The list containing the iteration numbers. """ import re iteration_numbers = [] for key in keys_list: str_match = re.search(" it=\d* ", key) it_str_start = str_match.start() it_str_end = str_match.end() it_value = int(key[it_str_start + 4 : it_str_end]) iteration_numbers.append(it_value) return iteration_numbers
[docs]def construct_mode_list(ell_max, spin_weight): """ Construct a modes list in the form [[ell1, [emm1, emm2, ...], [ell2, [emm..]],..] given the :math:`\\ell_{max}.` Parameters ---------- spin_weight : int The spin weight of the modes. ell_max : int The :math:`\\ell_{max}` of the modes list. Returns ------- modes_list : list A list containg the mode indices lists. Notes ----- The modes list is the form which the `waveform` object understands. """ # The modes list. modes_list = [] for ell_index in range(abs(spin_weight), ell_max + 1): # Append all emm modes for each ell mode. modes_list.append([ell_index, list(range(-ell_index, ell_index + 1))]) return modes_list
[docs]def sort_keys(modes_keys_list): """Sort the keys in a list based on its iteration number Parameters ---------- modes_keys_list: str The list of keys. Returns ------- sorted_modes_keys_list: str The sorted list. """ iteration_numbers = get_iteration_numbers_from_keys(modes_keys_list) sargs = np.argsort(iteration_numbers) sorted_modes_keys_list = np.array(modes_keys_list)[sargs] return sorted_modes_keys_list