Using float32

Many modern GPUs have specialized hardware units optimized for float32 operations. Leveraging the floating-point precision of 32-bit may fully exploit the hardware acceleration and gain better performance and memory-efficiency. However, we need to scale responses and PSDs to avoid the underflow. A factor of PC_SI / (MRSUN_SI * MTSUN_SI) is multiplied to responses, and the square of this factor is multiplied to PSDs.

import taichi as ti
ti.init(
    arch=ti.cpu,
    default_fp=ti.f32,
    cpu_max_num_threads=1,
    offline_cache=False,
)
    
import numpy as np
from matplotlib import pyplot as plt
import lal
import bilby

from pespace.detector.antenna import InterferometerAntenna, FDResponseModelMarset2018
from pespace.detector.orbit import available_orbit_models
from pespace.detector.noise import available_noise_models
from pespace.detector.tdi import TDIChannelData, FDMichelsonConstantEqualArm
from tiwave.waveforms import IMRPhenomXAS

[Taichi] version 1.7.4, llvm 15.0.4, commit b4b956fd, linux, python 3.10.19
[Taichi] Starting on arch=x64

[I 02/06/26 22:34:50.713 221985] [shell.py:_shell_pop_print@23] Graphical python shell detected, using wrapped sys.stdout
/tmp/ipykernel_221985/995760217.py:11: UserWarning: Wswiglal-redir-stdio:

SWIGLAL standard output/error redirection is enabled in IPython.
This may lead to performance penalties. To disable locally, use:

with lal.no_swig_redirect_standard_output_error():
    ...

To disable globally, use:

lal.swig_redirect_standard_output_error(False)

Note however that this will likely lead to error messages from
LAL functions being either misdirected or lost when called from
Jupyter notebooks.

To suppress this warning, use:

import warnings
warnings.filterwarnings("ignore", "Wswiglal-redir-stdio")
import lal

  import lal

tdi_gen = "2.0"
tdi_chan = ("A", "E", "T")
dt = 5.0
f_min = 1e-5
f_max = 0.5*(1/dt)
f_ref = f_min
t_start = 0.0
num_tsamples = 2**np.ceil(np.log2(7*lal.DAYJUL_SI/dt))
duration = num_tsamples * dt
params = dict(
    total_mass=3e6,
    mass_ratio=0.6,
    chi_1=0.75,
    chi_2=0.62,
    luminosity_distance=56000.0,
    inclination=0.4,
    reference_phase=1.3,
    ecliptic_longitude=1.375,
    ecliptic_latitude=-1.2108,
    polarization=2.659,
    coalescence_time=0.0,
)
params = bilby.gw.conversion.generate_mass_parameters(params)
params

{'total_mass': 3000000.0,
 'mass_ratio': 0.6,
 'chi_1': 0.75,
 'chi_2': 0.62,
 'luminosity_distance': 56000.0,
 'inclination': 0.4,
 'reference_phase': 1.3,
 'ecliptic_longitude': 1.375,
 'ecliptic_latitude': -1.2108,
 'polarization': 2.659,
 'coalescence_time': 0.0,
 'mass_1': 1875000.0,
 'mass_2': 1125000.0,
 'chirp_mass': 1256226.717491785,
 'symmetric_mass_ratio': np.float64(0.234375)}

tdi_data = TDIChannelData(scaling=True)
tdi_data.set_fd_data_from_zero(
    channels=tdi_chan, 
    duration=duration, 
    delta_time=dt,
    start_time=t_start,
    minimum_frequency=f_min,
    maximum_frequency=f_max,
    )
tdi_data.set_fd_noise_power_density_from_model(available_noise_models["LISA_SciRDv1"], tdi_generation=tdi_gen)
noise = tdi_data.get_fd_noise_realization()
tdi_data.add_into_fd_data(noise)
lisa = InterferometerAntenna(
    name="lisa",
    tdi_data=tdi_data,
    orbit_model=available_orbit_models['LISA_analytic'],
    response_model=FDResponseModelMarset2018(),
    tdi_combination=FDMichelsonConstantEqualArm(generation=tdi_gen, orthogonal=True),
)
wf = IMRPhenomXAS(
    tdi_data.frequency_samples, 
    f_ref, 
    scaling=True, 
    )
wf.update_waveform(params)
lisa.inject_signal(
    wf.waveform_container,
    params["ecliptic_longitude"],
    params["ecliptic_latitude"],
    params["polarization"],
    params["coalescence_time"],
)

/home/nrui/disk_ext/workspace/tiwave/tiwave/waveforms/base_waveform.py:69: UserWarning: check_parameters is disable, make sure all parameters passed in are valid.
  warnings.warn(

freqs_np = tdi_data.data_info.frequency_samples_array
obs_data_np = lisa.tdi_data.fd_data_numpy
tdi_resp_np = lisa.tdi_response_numpy

tab20_colors = plt.cm.tab20.colors
fig, ax = plt.subplots()
line_data_A, = ax.loglog(freqs_np, np.abs(obs_data_np['A']), color=tab20_colors[1], label='data A')
line_data_E, = ax.loglog(freqs_np, np.abs(obs_data_np['E']), color=tab20_colors[3], label='data E')
line_data_T, = ax.loglog(freqs_np, np.abs(obs_data_np['T']), color=tab20_colors[5], label='data T')

line_resp_A, = ax.loglog(freqs_np, np.abs(tdi_resp_np['A']), color=tab20_colors[0], linestyle='dashed', label='response A')
line_resp_E, = ax.loglog(freqs_np, np.abs(tdi_resp_np['E']), color=tab20_colors[2], linestyle='dashed', label='response E')
line_resp_T, = ax.loglog(freqs_np, np.abs(tdi_resp_np['T']), color=tab20_colors[4], linestyle='dashed', label='response T')

ax.set_xlim(f_min, f_max)
ax.set_xlabel('Frequency (Hz)')
ax.set_ylabel('Amplitude (rescaled)')

legend1 = ax.legend(
    handles=[line_data_A, line_data_E, line_data_T],
    loc='upper center',
    bbox_to_anchor=(0.12, 1.0),
)

legend2 = ax.legend(
    handles=[line_resp_A, line_resp_E, line_resp_T],
    loc='upper center',
    bbox_to_anchor=(0.36, 1.0),
)

ax.add_artist(legend1)

<matplotlib.legend.Legend at 0x7f1ed75b0280>