Full PE scripts

The scripts on this page are full parameter estimation examples for a massive black hole binary merger.

Note

We have not yet successfully obtained reliable posteriors sampling with flowMC and jaxns in our tests. The configurations provided in the example scripts should be regarded as preliminary and have not been fully validated.

1. Using the sampler interface of bilby and the likelihood computation of taichi-lang backend

1.1 A single LISA detector with PhenomXAS waveform

full_pe_ti_lisa_xas.py
  1import os
  2# The jax backend is unused on this taichi path; keep jax off the GPU
  3os.environ['JAX_PLATFORMS']='cpu'
  4# Select the GPU for this MPI rank when the scheduler provides these env variables
  5local_rank_id = int(os.environ.get('MPI_LOCALRANKID', 0))
  6gpus_list = [g.strip() for g in os.environ.get('GPU_DEVICE_ORDINAL', '').split(',') if g.strip()]
  7if gpus_list:
  8    os.environ['CUDA_VISIBLE_DEVICES'] = gpus_list[local_rank_id % len(gpus_list)]
  9
 10import bilby
 11import lal
 12import numpy as np
 13import taichi as ti
 14ti.init(arch=ti.gpu, default_fp=ti.f64, cpu_max_num_threads=1, offline_cache=False)
 15# ti.init(arch=ti.cpu, default_fp=ti.f64, cpu_max_num_threads=1, offline_cache=False)
 16
 17from pespace.detector.antenna import InterferometerAntenna, FDResponseModelMarset2018
 18from pespace.detector.tdi import TDIChannelData, FDMichelsonConstantEqualArm
 19from pespace.detector.orbit import available_orbit_models
 20from pespace.inference.interface_bilby import LikelihoodBilbyInterfaceTaichi
 21from tiwave.waveforms import IMRPhenomXAS
 22
 23########################################################################################
 24# set injection
 25tdi_gen = "2.0"
 26tdi_chan = ("A", "E", "T")
 27
 28dt = 10
 29f_min = 1e-4
 30f_max = 0.5*(1/dt)
 31f_ref = f_min
 32t_start = 0.0
 33
 34num_tsamples = 2**np.ceil(np.log2(7*lal.DAYJUL_SI/dt))
 35duration = num_tsamples * dt
 36before_tc = 0.8 * duration
 37after_tc = 0.2 * duration
 38tc = t_start + before_tc
 39
 40params = dict(
 41    total_mass=3e6,
 42    mass_ratio=0.6,
 43    chi_1=0.75,
 44    chi_2=0.62,
 45    luminosity_distance=56000.0,
 46    inclination=0.4,
 47    reference_phase=1.3,
 48    ecliptic_longitude=1.375,
 49    ecliptic_latitude=-1.2108,
 50    polarization=2.659,
 51    coalescence_time=tc,
 52)
 53params = bilby.gw.conversion.generate_mass_parameters(params)
 54
 55tdi_data = TDIChannelData()
 56tdi_data.set_fd_data_from_zero(
 57    channels=tdi_chan, 
 58    duration=duration, 
 59    delta_time=dt,
 60    start_time=t_start,
 61    minimum_frequency=f_min,
 62    maximum_frequency=f_max,
 63    )
 64tdi_data.set_fd_noise_power_density_from_model("LISA_SciRDv1", tdi_generation=tdi_gen)
 65noise = tdi_data.get_fd_noise_realization()
 66tdi_data.add_into_fd_data(noise)
 67
 68orbit_model = available_orbit_models['LISA_analytic']
 69response_model = FDResponseModelMarset2018()
 70tdi_combination = FDMichelsonConstantEqualArm(generation=tdi_gen, orthogonal=True)
 71lisa = InterferometerAntenna(
 72    name="lisa",
 73    tdi_data=tdi_data,
 74    orbit_model=orbit_model,
 75    response_model=response_model,
 76    tdi_combination=tdi_combination,
 77)
 78
 79wf_xas = IMRPhenomXAS(tdi_data.frequency_samples, f_ref, parameter_conversion=bilby.gw.conversion.generate_component_masses)
 80wf_xas.update_waveform(params)
 81lisa.inject_signal(
 82    wf_xas.waveform_container,
 83    params["ecliptic_longitude"],
 84    params["ecliptic_latitude"],
 85    params["polarization"],
 86    params["coalescence_time"],
 87)
 88
 89########################################################################################
 90# set sampling
 91likelihood = LikelihoodBilbyInterfaceTaichi(
 92    waveform=wf_xas,
 93    detector=lisa,
 94    channels=('A', 'E', 'T')
 95    )
 96
 97priors = {}
 98priors["chi_1"] = bilby.core.prior.Uniform(name="chi_1", minimum=-0.99, maximum=0.99)
 99priors["chi_2"] = bilby.core.prior.Uniform(name="chi_2", minimum=-0.99, maximum=0.99)
100priors["ecliptic_longitude"] = bilby.core.prior.Uniform(name="ecliptic_longitude", minimum=0, maximum=2 * lal.PI, boundary="periodic")
101priors["ecliptic_latitude"] = bilby.core.prior.Cosine(name="ecliptic_latitude")
102priors["inclination"] = bilby.core.prior.Sine(name="inclination")
103priors["polarization"] = bilby.core.prior.Uniform(name="polarization", minimum=0, maximum=lal.PI, boundary="periodic")
104priors["reference_phase"] = bilby.core.prior.Uniform(name="reference_phase", minimum=0, maximum=2 * lal.PI, boundary="periodic")
105priors["coalescence_time"] = bilby.core.prior.Uniform(name="coalescence_time", minimum=tc-20, maximum=tc+20)
106priors["chirp_mass"] = bilby.gw.prior.UniformInComponentsChirpMass(name="chirp_mass", minimum=params['chirp_mass']-5e3, maximum=params['chirp_mass']+5e3)
107priors["mass_ratio"] = bilby.gw.prior.UniformInComponentsMassRatio(name="mass_ratio", minimum=0.2, maximum=0.9)
108priors["luminosity_distance"] = bilby.core.prior.Uniform(name="luminosity_distance", minimum=4e4, maximum=8e4)
109priors = bilby.core.prior.PriorDict(dictionary=priors)
110
111label = "inj_MBHB_LISA"
112outdir = f"outdir_{label}"
113result = bilby.run_sampler(
114    likelihood=likelihood,
115    priors=priors,
116    outdir=outdir,
117    label=label,
118    sampler='pymultinest',
119    npoints=2048,
120    )

1.2 The LISA-Taiji-Tianqin network with PhenomXHM waveform

full_pe_ti_network_xhm.py
  1import os
  2# The jax backend is unused on this taichi path; keep jax off the GPU
  3os.environ['JAX_PLATFORMS']='cpu'
  4# Select the GPU for this MPI rank when the scheduler provides these env variables
  5local_rank_id = int(os.environ.get('MPI_LOCALRANKID', 0))
  6gpus_list = [g.strip() for g in os.environ.get('GPU_DEVICE_ORDINAL', '').split(',') if g.strip()]
  7if gpus_list:
  8    os.environ['CUDA_VISIBLE_DEVICES'] = gpus_list[local_rank_id % len(gpus_list)]
  9
 10import bilby
 11import lal
 12import numpy as np
 13import taichi as ti
 14ti.init(arch=ti.gpu, default_fp=ti.f64, cpu_max_num_threads=1, offline_cache=False)
 15# ti.init(arch=ti.cpu, default_fp=ti.f64, cpu_max_num_threads=1, offline_cache=False)
 16
 17from pespace.detector.antenna import InterferometerAntenna, FDResponseModelMarset2018
 18from pespace.detector.tdi import TDIChannelData, FDMichelsonConstantEqualArm
 19from pespace.detector.orbit import available_orbit_models
 20from pespace.inference.interface_bilby import LikelihoodBilbyInterfaceTaichi
 21from tiwave.waveforms import IMRPhenomXHM
 22
 23########################################################################################
 24# common settings
 25# set injection
 26tdi_gen = "2.0"
 27tdi_chan = ("A", "E", "T")
 28
 29dt = 10
 30f_min = 1e-4
 31f_max = 0.5*(1/dt)
 32f_ref = f_min
 33t_start = 0.0
 34
 35num_tsamples = 2**np.ceil(np.log2(7*lal.DAYJUL_SI/dt))
 36duration = num_tsamples * dt
 37before_tc = 0.8 * duration
 38after_tc = 0.2 * duration
 39tc = t_start + before_tc
 40
 41params = dict(
 42    total_mass=3e6,
 43    mass_ratio=0.6,
 44    chi_1=0.75,
 45    chi_2=0.62,
 46    luminosity_distance=56000.0,
 47    inclination=0.4,
 48    reference_phase=1.3,
 49    ecliptic_longitude=1.375,
 50    ecliptic_latitude=-1.2108,
 51    polarization=2.659,
 52    coalescence_time=tc,
 53)
 54params = bilby.gw.conversion.generate_mass_parameters(params)
 55
 56# detector settings
 57tdi_data_lisa = TDIChannelData()
 58tdi_data_lisa.set_fd_data_from_zero(
 59    channels=tdi_chan, 
 60    duration=duration, 
 61    delta_time=dt,
 62    start_time=t_start,
 63    minimum_frequency=f_min,
 64    maximum_frequency=f_max,
 65    )
 66tdi_data_lisa.set_fd_noise_power_density_from_model("LISA_SciRDv1", tdi_generation=tdi_gen)
 67noise_lisa = tdi_data_lisa.get_fd_noise_realization()
 68tdi_data_lisa.add_into_fd_data(noise_lisa)
 69lisa = InterferometerAntenna(
 70    name="lisa",
 71    tdi_data=tdi_data_lisa,
 72    orbit_model=available_orbit_models['LISA_analytic'],
 73    response_model=FDResponseModelMarset2018(),
 74    tdi_combination=FDMichelsonConstantEqualArm(generation=tdi_gen, orthogonal=True),
 75)
 76
 77tdi_data_taiji = TDIChannelData()
 78tdi_data_taiji.set_fd_data_from_zero(
 79    channels=tdi_chan, 
 80    duration=duration, 
 81    delta_time=dt,
 82    start_time=t_start,
 83    minimum_frequency=f_min,
 84    maximum_frequency=f_max,
 85    )
 86tdi_data_taiji.set_fd_noise_power_density_from_model("Taiji_TDC", tdi_generation=tdi_gen)
 87noise_taiji = tdi_data_taiji.get_fd_noise_realization()
 88tdi_data_taiji.add_into_fd_data(noise_taiji)
 89taiji = InterferometerAntenna(
 90    name="taiji",
 91    tdi_data=tdi_data_taiji,
 92    orbit_model=available_orbit_models['Taiji_analytic'],
 93    response_model=FDResponseModelMarset2018(),
 94    tdi_combination=FDMichelsonConstantEqualArm(generation=tdi_gen, orthogonal=True),
 95)
 96
 97tdi_data_tianqin = TDIChannelData()
 98tdi_data_tianqin.set_fd_data_from_zero(
 99    channels=tdi_chan, 
100    duration=duration, 
101    delta_time=dt,
102    start_time=t_start,
103    minimum_frequency=f_min,
104    maximum_frequency=f_max,
105    )
106tdi_data_tianqin.set_fd_noise_power_density_from_model("Tianqin_GWSpace", tdi_generation=tdi_gen)
107noise_tianqin = tdi_data_tianqin.get_fd_noise_realization()
108tdi_data_tianqin.add_into_fd_data(noise_tianqin)
109tianqin = InterferometerAntenna(
110    name="tianqin",
111    tdi_data=tdi_data_tianqin,
112    orbit_model=available_orbit_models['Tianqin_analytic'],
113    response_model=FDResponseModelMarset2018(),
114    tdi_combination=FDMichelsonConstantEqualArm(generation=tdi_gen, orthogonal=True),
115)
116
117common_freqs = tdi_data_lisa.data_info.frequency_samples_array
118wf_xhm = IMRPhenomXHM(common_freqs, f_ref, parameter_conversion=bilby.gw.conversion.generate_component_masses)
119wf_xhm.update_waveform(params)
120lisa.inject_signal(
121    wf_xhm.waveform_container,
122    params["ecliptic_longitude"],
123    params["ecliptic_latitude"],
124    params["polarization"],
125    params["coalescence_time"],
126)
127taiji.inject_signal(
128    wf_xhm.waveform_container,
129    params["ecliptic_longitude"],
130    params["ecliptic_latitude"],
131    params["polarization"],
132    params["coalescence_time"],
133)
134tianqin.inject_signal(
135    wf_xhm.waveform_container,
136    params["ecliptic_longitude"],
137    params["ecliptic_latitude"],
138    params["polarization"],
139    params["coalescence_time"],
140)
141
142########################################################################################
143# sampling settings
144likelihood = LikelihoodBilbyInterfaceTaichi(
145    waveform=wf_xhm,
146    detector=(lisa, taiji, tianqin,),
147    channels=('A', 'E', 'T')
148    )
149
150priors = {}
151priors["chi_1"] = bilby.core.prior.Uniform(name="chi_1", minimum=-0.99, maximum=0.99)
152priors["chi_2"] = bilby.core.prior.Uniform(name="chi_2", minimum=-0.99, maximum=0.99)
153priors["ecliptic_longitude"] = bilby.core.prior.Uniform(name="ecliptic_longitude", minimum=0, maximum=2 * lal.PI, boundary="periodic")
154priors["ecliptic_latitude"] = bilby.core.prior.Cosine(name="ecliptic_latitude")
155priors["inclination"] = bilby.core.prior.Sine(name="inclination")
156priors["polarization"] = bilby.core.prior.Uniform(name="polarization", minimum=0, maximum=lal.PI, boundary="periodic")
157priors["reference_phase"] = bilby.core.prior.Uniform(name="reference_phase", minimum=0, maximum=2 * lal.PI, boundary="periodic")
158priors["coalescence_time"] = bilby.core.prior.Uniform(name="coalescence_time", minimum=tc-20, maximum=tc+20)
159priors["chirp_mass"] = bilby.gw.prior.UniformInComponentsChirpMass(name="chirp_mass", minimum=params['chirp_mass']-5e3, maximum=params['chirp_mass']+5e3)
160priors["mass_ratio"] = bilby.gw.prior.UniformInComponentsMassRatio(name="mass_ratio", minimum=0.2, maximum=0.9)
161priors["luminosity_distance"] = bilby.core.prior.Uniform(name="luminosity_distance", minimum=4e4, maximum=8e4)
162priors = bilby.core.prior.PriorDict(dictionary=priors)
163
164label = "inj_MBHB_network_HM"
165outdir = f"outdir_{label}"
166result = bilby.run_sampler(
167    likelihood=likelihood,
168    priors=priors,
169    outdir=outdir,
170    label=label,
171    sampler='pymultinest',
172    npoints=2048,
173    )

2. Using samplers with jax

jax is not installed by default, please install dependencies additionally:

pip install "pespace[jax]"
pip install "pespace[samplers]"
pip install "rippleGW @ git+https://github.com/nnrui/ripple"

2.1 Adaptive tempering SMC with NUTS and Hessian

full_pe_jax_smc.py
  1from collections import namedtuple
  2import os
  3
  4os.environ["XLA_PYTHON_CLIENT_PREALLOCATE"] = "false"
  5os.environ["JAX_COMPILATION_CACHE_DIR"] = "/home/changfenggroup/nrui/.cache/jax"
  6os.environ["JAX_PERSISTENT_CACHE_MIN_COMPILE_TIME_SECS"] = "10"
  7
  8import jax
  9
 10jax.config.update("jax_platform_name", "gpu")
 11jax.config.update("jax_enable_x64", True)
 12
 13import lal
 14import numpy as np
 15import bilby
 16
 17from pespace.detector.antenna import InterferometerAntenna, FDResponseModelMarset2018
 18from pespace.detector.tdi import TDIChannelData, FDMichelsonConstantEqualArm
 19from pespace.detector.orbit import KaplerianHeliocentric
 20from pespace.inference.likelihood_jax import WhittleLikelihoodJax
 21from pespace.inference.sampler.blackjax_smc import BlackjaxSMCwithNUTS
 22from pespace.waveform.interface_ripple import IMRPhenomXAS
 23from pespace.inference.sampler.base import Cosine, Sine
 24
 25########################################################################################
 26# set injection
 27tdi_gen = "2.0"
 28tdi_chan = ("A", "E", "T")
 29
 30dt = 10
 31f_min = 1e-4
 32f_max = 0.5 * (1 / dt)
 33f_ref = f_min
 34t_start = 0.0
 35
 36num_tsamples = 2 ** np.ceil(np.log2(7 * lal.DAYJUL_SI / dt))
 37duration = num_tsamples * dt
 38before_tc = 0.8 * duration
 39tc = t_start + before_tc
 40
 41params = dict(
 42    total_mass=3e6,
 43    mass_ratio=0.6,
 44    chi_1=0.75,
 45    chi_2=0.62,
 46    luminosity_distance=56000.0,
 47    inclination=0.4,
 48    reference_phase=1.3,
 49    ecliptic_longitude=1.375,
 50    ecliptic_latitude=-1.2108,
 51    polarization=2.659,
 52    coalescence_time=tc,
 53)
 54params = bilby.gw.conversion.generate_mass_parameters(params)
 55
 56tdi_data = TDIChannelData(backend="jax")
 57tdi_data.set_fd_data_from_zero(
 58    channels=tdi_chan,
 59    duration=duration,
 60    delta_time=dt,
 61    start_time=t_start,
 62    minimum_frequency=f_min,
 63    maximum_frequency=f_max,
 64)
 65tdi_data.set_fd_noise_power_density_from_model("LISA_SciRDv1", tdi_generation=tdi_gen)
 66noise = tdi_data.get_fd_noise_realization(seed=2026, output_type="numpy")
 67tdi_data.add_into_fd_data(noise)
 68
 69orbit_model = KaplerianHeliocentric(2.5e9, 0.0, -np.pi / 9, backend="jax")
 70response_model = FDResponseModelMarset2018(backend="jax")
 71tdi_combination = FDMichelsonConstantEqualArm(generation=tdi_gen, orthogonal=True, backend="jax")
 72lisa = InterferometerAntenna(
 73    name="lisa",
 74    tdi_data=tdi_data,
 75    orbit_model=orbit_model,
 76    response_model=response_model,
 77    tdi_combination=tdi_combination,
 78    backend="jax",
 79)
 80
 81wf_xas = IMRPhenomXAS(tdi_data.frequency_samples, f_ref)
 82WaveformParams = namedtuple("WaveformParams", wf_xas.required_params)
 83wf_params_inj = WaveformParams(
 84    **{name: params[name] for name in wf_xas.required_params}
 85)
 86lisa.inject_signal(
 87    wf_xas.get_waveform(wf_params_inj),
 88    params["ecliptic_longitude"],
 89    params["ecliptic_latitude"],
 90    params["polarization"],
 91    params["coalescence_time"],
 92)
 93
 94########################################################################################
 95# set sampling
 96likelihood = WhittleLikelihoodJax(
 97    waveform=wf_xas, detector=lisa, channels=("A", "E", "T")
 98)
 99
100priors = {
101    "chi_1": {"distribution": "Uniform", "low": -0.99, "high": 0.99},
102    "chi_2": {"distribution": "Uniform", "low": -0.99, "high": 0.99},
103    "ecliptic_longitude": {"distribution": "Uniform", "low": 0.0, "high": 2.0 * np.pi},
104    "ecliptic_latitude": {"distribution": Cosine},
105    "inclination": {"distribution": Sine},
106    "polarization": {"distribution": "Uniform", "low": 0.0, "high": np.pi},
107    "reference_phase": {"distribution": "Uniform", "low": 0.0, "high": 2.0 * np.pi},
108    "coalescence_time": {
109        "distribution": "Uniform",
110        "low": tc - 10.0,
111        "high": tc + 10.0,
112    },
113    "chirp_mass": {
114        "distribution": "Uniform",
115        "low": params["chirp_mass"] - 5e3,
116        "high": params["chirp_mass"] + 5e3,
117    },
118    "mass_ratio": {"distribution": "Uniform", "low": 0.2, "high": 0.9},
119    "luminosity_distance": {"distribution": "Uniform", "low": 4e4, "high": 8e4},
120
121    "symmetric_mass_ratio": {
122        "transform": lambda mass_ratio: mass_ratio / (1.0 + mass_ratio) ** 2,
123        "from": ("mass_ratio",),
124    },
125}
126
127sampler = BlackjaxSMCwithNUTS(
128    likelihood,
129    priors,
130    label="full_pe_smc",
131    rkey_seed=2026,
132    num_particles=2048,
133    batch_size=1024,
134    hessian_batch_size=64,
135    show_progress=True,
136)
137results = sampler.run()

2.2 jaxns

full_pe_jax_jaxns.py
  1from collections import namedtuple
  2import os
  3
  4os.environ["XLA_PYTHON_CLIENT_PREALLOCATE"] = "false"
  5os.environ["JAX_COMPILATION_CACHE_DIR"] = "/home/changfenggroup/nrui/.cache/jax"
  6os.environ["JAX_PERSISTENT_CACHE_MIN_COMPILE_TIME_SECS"] = "10"
  7
  8import jax
  9
 10jax.config.update("jax_platform_name", "gpu")
 11jax.config.update("jax_enable_x64", True)
 12
 13import lal
 14import numpy as np
 15import bilby
 16
 17from pespace.detector.antenna import InterferometerAntenna, FDResponseModelMarset2018
 18from pespace.detector.tdi import TDIChannelData, FDMichelsonConstantEqualArm
 19from pespace.detector.orbit import KaplerianHeliocentric
 20from pespace.inference.likelihood_jax import WhittleLikelihoodJax
 21from pespace.inference.sampler.jaxns import JaxnsSampler
 22from pespace.waveform.interface_ripple import IMRPhenomXAS
 23from pespace.inference.sampler.jaxns import tfpCosine, tfpSine
 24
 25########################################################################################
 26# set injection
 27tdi_gen = "2.0"
 28tdi_chan = ("A", "E", "T")
 29
 30dt = 10
 31f_min = 1e-4
 32f_max = 0.5 * (1 / dt)
 33f_ref = f_min
 34t_start = 0.0
 35
 36num_tsamples = 2 ** np.ceil(np.log2(7 * lal.DAYJUL_SI / dt))
 37duration = num_tsamples * dt
 38before_tc = 0.8 * duration
 39tc = t_start + before_tc
 40
 41params = dict(
 42    total_mass=3e6,
 43    mass_ratio=0.6,
 44    chi_1=0.75,
 45    chi_2=0.62,
 46    luminosity_distance=56000.0,
 47    inclination=0.4,
 48    reference_phase=1.3,
 49    ecliptic_longitude=1.375,
 50    ecliptic_latitude=-1.2108,
 51    polarization=2.659,
 52    coalescence_time=tc,
 53)
 54params = bilby.gw.conversion.generate_mass_parameters(params)
 55
 56tdi_data = TDIChannelData(backend="jax")
 57tdi_data.set_fd_data_from_zero(
 58    channels=tdi_chan,
 59    duration=duration,
 60    delta_time=dt,
 61    start_time=t_start,
 62    minimum_frequency=f_min,
 63    maximum_frequency=f_max,
 64)
 65tdi_data.set_fd_noise_power_density_from_model("LISA_SciRDv1", tdi_generation=tdi_gen)
 66noise = tdi_data.get_fd_noise_realization(seed=2026, output_type="numpy")
 67tdi_data.add_into_fd_data(noise)
 68
 69orbit_model = KaplerianHeliocentric(2.5e9, 0.0, -np.pi / 9, backend="jax")
 70response_model = FDResponseModelMarset2018(backend="jax")
 71tdi_combination = FDMichelsonConstantEqualArm(generation=tdi_gen, orthogonal=True, backend="jax")
 72lisa = InterferometerAntenna(
 73    name="lisa",
 74    tdi_data=tdi_data,
 75    orbit_model=orbit_model,
 76    response_model=response_model,
 77    tdi_combination=tdi_combination,
 78    backend="jax",
 79)
 80
 81wf_xas = IMRPhenomXAS(tdi_data.frequency_samples, f_ref)
 82WaveformParams = namedtuple("WaveformParams", wf_xas.required_params)
 83wf_params_inj = WaveformParams(
 84    **{name: params[name] for name in wf_xas.required_params}
 85)
 86lisa.inject_signal(
 87    wf_xas.get_waveform(wf_params_inj),
 88    params["ecliptic_longitude"],
 89    params["ecliptic_latitude"],
 90    params["polarization"],
 91    params["coalescence_time"],
 92)
 93
 94########################################################################################
 95# set sampling
 96likelihood = WhittleLikelihoodJax(
 97    waveform=wf_xas, detector=lisa, channels=("A", "E", "T")
 98)
 99
100priors = {
101    "chi_1": {"distribution": "Uniform", "low": -0.99, "high": 0.99},
102    "chi_2": {"distribution": "Uniform", "low": -0.99, "high": 0.99},
103    "ecliptic_longitude": {"distribution": "Uniform", "low": 0.0, "high": 2.0 * np.pi},
104    "ecliptic_latitude": {"distribution": tfpCosine},
105    "inclination": {"distribution": tfpSine},
106    "polarization": {"distribution": "Uniform", "low": 0.0, "high": np.pi},
107    "reference_phase": {"distribution": "Uniform", "low": 0.0, "high": 2.0 * np.pi},
108    "coalescence_time": {
109        "distribution": "Uniform",
110        "low": tc - 10.0,
111        "high": tc + 10.0,
112    },
113    "chirp_mass": {
114        "distribution": "Uniform",
115        "low": params["chirp_mass"] - 5e3,
116        "high": params["chirp_mass"] + 5e3,
117    },
118    "mass_ratio": {"distribution": "Uniform", "low": 0.2, "high": 0.9},
119    "luminosity_distance": {"distribution": "Uniform", "low": 4e4, "high": 8e4},
120
121    "symmetric_mass_ratio": {
122        "transform": lambda mass_ratio: mass_ratio / (1.0 + mass_ratio) ** 2,
123        "from": ("mass_ratio",),
124    },
125}
126
127sampler = JaxnsSampler(
128    likelihood,
129    priors,
130    label="full_pe_jaxns",
131    rkey_seed=2026,
132    # c=4096,
133    shell_fraction=0.1,
134    s=500,
135    # k=11,
136    difficult_model=True, 
137    parameter_estimation=True,
138    verbose=True,
139)
140
141results = sampler.run()

2.3 flowMC

full_pe_jax_flowmc.py
  1import os
  2os.environ["XLA_PYTHON_CLIENT_PREALLOCATE"] = "false"
  3os.environ["CUDA_VISIBLE_DEVICES"]="1"
  4from collections import namedtuple
  5
  6import jax
  7jax.config.update("jax_enable_x64", True)
  8jax.config.update("jax_platform_name", "gpu")
  9
 10import numpy as np
 11import lal
 12import bilby
 13
 14from pespace.detector.antenna import (
 15    FDResponseModelMarset2018,
 16    InterferometerAntenna,
 17)
 18from pespace.detector.orbit import KaplerianHeliocentric
 19from pespace.detector.tdi import (
 20    FDMichelsonConstantEqualArm,
 21    TDIChannelData,
 22)
 23from pespace.inference.likelihood_jax import WhittleLikelihoodJax
 24from pespace.inference.sampler.flowmc import FlowMCSampler
 25from pespace.inference.sampler.base import Cosine, Sine
 26from pespace.waveform.interface_ripple import IMRPhenomXAS
 27
 28###############################################################################
 29# set injection
 30tdi_gen = "2.0"
 31tdi_chan = ("A", "E", "T")
 32
 33dt = 10
 34f_min = 1e-4
 35f_max = 0.5 * (1 / dt)
 36f_ref = f_min
 37t_start = 0.0
 38
 39num_tsamples = 2 ** np.ceil(np.log2(7 * lal.DAYJUL_SI / dt))
 40duration = num_tsamples * dt
 41before_tc = 0.8 * duration
 42after_tc = 0.2 * duration
 43tc = t_start + before_tc
 44
 45params = dict(
 46    total_mass=3e6,
 47    mass_ratio=0.6,
 48    chi_1=0.75,
 49    chi_2=0.62,
 50    luminosity_distance=56000.0,
 51    inclination=0.4,
 52    reference_phase=1.3,
 53    ecliptic_longitude=1.375,
 54    ecliptic_latitude=-1.2108,
 55    polarization=2.659,
 56    coalescence_time=tc,
 57)
 58
 59tdi_data = TDIChannelData(backend="jax")
 60tdi_data.set_fd_data_from_zero(
 61    channels=tdi_chan,
 62    duration=duration,
 63    delta_time=dt,
 64    start_time=t_start,
 65    minimum_frequency=f_min,
 66    maximum_frequency=f_max,
 67)
 68tdi_data.set_fd_noise_power_density_from_model("LISA_SciRDv1", tdi_generation=tdi_gen)
 69noise = tdi_data.get_fd_noise_realization(seed=2026, output_type="numpy")
 70tdi_data.add_into_fd_data(noise)
 71
 72orbit_model = KaplerianHeliocentric(2.5e9, 0.0, -np.pi / 9, backend="jax")
 73response_model = FDResponseModelMarset2018(backend="jax")
 74tdi_combination = FDMichelsonConstantEqualArm(generation=tdi_gen, orthogonal=True, backend="jax")
 75lisa = InterferometerAntenna(
 76    name="lisa",
 77    tdi_data=tdi_data,
 78    orbit_model=orbit_model,
 79    response_model=response_model,
 80    tdi_combination=tdi_combination,
 81    backend="jax",
 82)
 83
 84wf_xas = IMRPhenomXAS(tdi_data.frequency_samples, f_ref)
 85
 86params_inj = bilby.gw.conversion.generate_mass_parameters(params)
 87WaveformParams = namedtuple("WaveformParams", wf_xas.required_params)
 88wf_params_inj = WaveformParams(**{name: params_inj[name] for name in wf_xas.required_params})
 89
 90lisa.inject_signal(
 91    wf_xas.get_waveform(wf_params_inj),
 92    params["ecliptic_longitude"],
 93    params["ecliptic_latitude"],
 94    params["polarization"],
 95    params["coalescence_time"],
 96)
 97
 98###############################################################################
 99# set sampling
100likelihood = WhittleLikelihoodJax(
101    waveform=wf_xas,
102    detector=lisa,
103    channels=("A", "E", "T"),
104)
105
106prior_dict = {
107    "chi_1": {"distribution": "Uniform", "low": -0.99, "high": 0.99},
108    "chi_2": {"distribution": "Uniform", "low": -0.99, "high": 0.99},
109    "ecliptic_longitude": {"distribution": "Uniform", "low": 0.0, "high": 2.0 * np.pi},
110    "ecliptic_latitude": {"distribution": Cosine},
111    "inclination": {"distribution": Sine},
112    "polarization": {"distribution": "Uniform", "low": 0.0, "high": np.pi},
113    "reference_phase": {"distribution": "Uniform", "low": 0.0, "high": 2.0 * np.pi},
114    "coalescence_time": {"distribution": "Uniform", "low": tc-10.0, "high": tc+10.0},
115    "chirp_mass": {"distribution": "Uniform", "low": params_inj['chirp_mass']-5e3, "high": params_inj['chirp_mass']+5e3},
116    "mass_ratio": {"distribution": "Uniform", "low": 0.2, "high": 0.9},
117    "luminosity_distance": {"distribution": "Uniform", "low": 4e4, "high": 8e4},
118
119    "symmetric_mass_ratio": {
120        "transform": lambda mass_ratio: mass_ratio / (1.0 + mass_ratio) ** 2,
121        "from": ("mass_ratio",),
122    },
123}
124
125    
126bundle_config = dict(
127    n_local_steps=100,
128    n_global_steps=1000,
129    n_training_loops=20,
130    n_production_loops=10,
131    n_epochs=20,
132    rq_spline_hidden_units=(128, 128),
133    rq_spline_n_bins=16,
134    rq_spline_n_layers=10,
135    n_NFproposal_batch_size=1,
136    n_max_examples=50000,
137    global_thinning=100,
138    early_stopping=True,
139    early_stopping_tolerance=0.1,
140    verbose=True,
141    chain_batch_size=64,
142    mala_step_size=np.array(
143        [1e-2, 1e-2, 1e-3, 2e-3, 1e-1,
144         6e-1, 6e-1, 8e-2, 4e-2, 5e-3, 6e-2]),
145    adapt_step_size=True, 
146    adapt_step_size_per_dim=False,
147    history_window=500,
148    n_temperatures=8, 
149    max_temperature=20,
150)
151
152label="full_pe_MBHB_LISA_flowmc"
153n_chains = 64
154
155sampler = FlowMCSampler(
156    likelihood,
157    prior_dict,
158    label=label,
159    outdir=f"outdir_{label}",
160    rkey_seed=2026,
161    resource_strategy_bundles="RQSpline_MALA_PT",
162    bundle_config=bundle_config,
163    n_chains=n_chains,
164    likelihood_batch_size=64,
165    checkpoint_interval=600.0,
166
167)
168sampler.run()
169