kazewong / jim Goto Github PK
View Code? Open in Web Editor NEWGravitational-wave data analysis tools in Jax
License: MIT License
Gravitational-wave data analysis tools in Jax
License: MIT License
In lines 384 and 385 in jim/src/jimgw /detector.py, the variance of the real part or the imaginary part of the noise should have half of the full variance [1].
Therefore, the line should be
noise_real = jax.random.normal(key, shape=freqs.shape) * jnp.sqrt(var / 2.)
noise_imag = jax.random.normal(subkey, shape=freqs.shape) * jnp.sqrt(var / 2.)
[1] https://en.wikipedia.org/wiki/Complex_normal_distribution
Currently, the optimization is done through scipy differential evolution by calling the jax kernel probably sequentially, which is pretty slow.
Ideally we want something completely Jaxxed, such as a vmap lbfgs to solve for the optimum
I discovered PyDantic when I was playing with langchain, which I think is beneficial for building application in python in general.
That said, this is one more dependency in our list, which might be an overkill for an astro software anyway.
Maybe worth thinking whether we want this in the future.
Currently, JIM computes the density change in the parameter space due to transformation by calculating the Jacobian determinant. The absolute value of the Jacobian determinant should be used; otherwise, feeding it into the calculation of the log Jacobian determinant may result in nan values.
It might be beneficial to work with whiten data such that the value and the dynamic range of data won't be present in the analysis. This could be useful to leverage lower precision format like float16/float32
Here is some things we might want to add to the run manager:
Currently, Jim only supports 2G detectors (LIGO and VIRGO). It would be beneficial to add 3G detectors such as Einstein Telescope and Cosmic Explorer for a lot of future applications with Jim, and test the robustness of relative binning for these scenarios.
Currently the user need to go into the python code to set up their own run. For future production, there should be a wrapper interfacing with the core and a configuration files, which the user can just specify the configuration they want for the run (Probably in a somewhat restrictive way),
On top of this, it would be good to have some sort of batch job options that the code will consume an array of configuration files then run on them, such that we don't need to pay the compilation overhead repeatedly.
Currently jim has minimal amount of unit test. Since there are more and more devs joining the development, it is better to add unit tests now than later.
Here is a non-exhaustive list of unit tests that are needed:
Currently there is only uniform prior with a guard outside the domain available.
This create issue of broken waveform returning nan when the gradient is large in HMC, which shoot the test point outside the domain and raising nan.
Here is a list of priors that could be useful to have:
Script:
from jimgw.single_event.runManager import SingleEventPERunManager, SingleEventRun
import jax.numpy as jnp
import jax
import os
outdir = os.path.dirname(__file__)
label = os.path.splitext(os.path.basename(__file__))[0]
jax.config.update("jax_enable_x64", True)
mass_matrix = jnp.eye(15)
mass_matrix = mass_matrix.at[1, 1].set(1e-3)
mass_matrix = mass_matrix.at[9, 9].set(1e-3)
mass_matrix = mass_matrix * 3e-3
local_sampler_arg = {"step_size": mass_matrix}
bounds = jnp.array(
[
[10.0, 40.0],
[0.125, 1.0],
[0, jnp.pi],
[0, 2*jnp.pi],
[0.0, 1.0],
[0, jnp.pi],
[0, 2*jnp.pi],
[0.0, 1.0],
[0.0, 2000.0],
[-0.05, 0.05],
[0.0, 2 * jnp.pi],
[-1.0, 1.0],
[0.0, jnp.pi],
[0.0, 2 * jnp.pi],
[-1.0, 1.0],
]
)
run = SingleEventRun(
seed=0,
path='',
detectors=["H1", "L1"],
priors={
"M_c": {"name": "Uniform", "xmin": 10.0, "xmax": 80.0},
"q": {"name": "MassRatio"},
"s1": {"name": "Sphere"},
"s2": {"name": "Sphere"},
"d_L": {"name": "Uniform", "xmin": 0.0, "xmax": 2000.0},
"t_c": {"name": "Uniform", "xmin": -0.05, "xmax": 0.05},
"phase_c": {"name": "Uniform", "xmin": 0.0, "xmax": 2 * jnp.pi},
"cos_iota": {"name": "CosIota"},
"psi": {"name": "Uniform", "xmin": 0.0, "xmax": jnp.pi},
"ra": {"name": "Uniform", "xmin": 0.0, "xmax": 2 * jnp.pi},
"sin_dec": {"name": "SinDec"},
},
waveform_parameters={"name": "RippleIMRPhenomPv2", "f_ref": 20.0},
jim_parameters={
"n_loop_training": 10,
"n_loop_production": 10,
"n_local_steps": 150,
"n_global_steps": 150,
"n_chains": 500,
"n_epochs": 50,
"learning_rate": 0.001,
"max_samples": 45000,
"momentum": 0.9,
"batch_size": 50000,
"use_global": True,
"keep_quantile": 0.0,
"train_thinning": 1,
"output_thinning": 10,
"local_sampler_arg": local_sampler_arg,
},
likelihood_parameters={"name": "HeterodynedTransientLikelihoodFD", "bounds": bounds},
injection=True,
injection_parameters={
"M_c": 28.6,
"eta": 0.24,
"s1_x": 0.05,
"s1_y": -0.05,
"s1_z": 0.05,
"s2_x": -0.05,
"s2_y": 0.05,
"s2_z": 0.05,
"d_L": 440.0,
"t_c": 0.0,
"phase_c": 0.0,
"iota": 0.5,
"psi": 0.7,
"ra": 1.2,
"dec": 0.3,
},
data_parameters={
"trigger_time": 1126259462.4,
"duration": 4,
"post_trigger_duration": 2,
"f_min": 20.0,
"f_max": 1024.0,
"tukey_alpha": 0.2,
"f_sampling": 4096.0,
},
)
run_manager = SingleEventPERunManager(run=run)
run_manager.jim.sample(jax.random.PRNGKey(42))
samples = run_manager.jim.get_samples()
run_manager.save(outdir+'/'+label)
jnp.save(outdir+'/'+label+"_samples.npy", samples)
run_manager.jim.print_summary()
Output:
Run instance provided. Loading from instance.
Initializing detectors.
Injection mode. Need to wait until waveform model is loaded.
Injection mode. Need to wait until waveform model is loaded.
Initializing waveform.
Grabbing GWTC-2 PSD for H1
For detector H1:
The injected optimal SNR is 29.713682403670045
The injected match filter SNR is (30.11265963611454-0.563276239330949j)
Grabbing GWTC-2 PSD for L1
For detector L1:
The injected optimal SNR is 51.04914603748342
The injected match filter SNR is (50.302015972725656-0.13345233070825305j)
Initializing heterodyned likelihood..
No reference parameters are provided, finding it...
Starting the optimizer
Using Adam optimization
Warning: Optimization accessed infinite or NaN log-probabilities.
The reference parameters are {'M_c': nan, 'eta': nan, 's1_x': nan, 's1_y': nan, 's1_z': nan, 's2_x': nan, 's2_y': nan, 's2_z': nan, 'd_L': nan, 't_c': nan, 'phase_c': nan, 'iota': nan, 'psi': nan, 'ra': nan, 'dec': nan}
Constructing reference waveforms..
jax.errors.SimplifiedTraceback: For simplicity, JAX has removed its internal frames from the traceback of the following exception. Set JAX_TRACEBACK_FILTERING=off to include these.
The above exception was the direct cause of the following exception:
Traceback (most recent call last):
File "/home/user/ckng/project/jim_testing/sky_location_frame/ra_dec.py", line 103, in <module>
run_manager = SingleEventPERunManager(run=run)
File "/home/user/ckng/.conda/envs/jim/lib/python3.10/site-packages/jimgw/single_event/runManager.py", line 127, in __init__
local_likelihood = self.initialize_likelihood(local_prior)
File "/home/user/ckng/.conda/envs/jim/lib/python3.10/site-packages/jimgw/single_event/runManager.py", line 185, in initialize_likelihood
return likelihood_presets[name](
File "/home/user/ckng/.conda/envs/jim/lib/python3.10/site-packages/jimgw/single_event/likelihood.py", line 298, in __init__
f_max = jnp.max(f_valid)
File "/home/user/ckng/.conda/envs/jim/lib/python3.10/site-packages/jax/_src/numpy/reductions.py", line 268, in max
return _reduce_max(a, axis=_ensure_optional_axes(axis), out=out,
File "/home/user/ckng/.conda/envs/jim/lib/python3.10/site-packages/jax/_src/numpy/reductions.py", line 260, in _reduce_max
return _reduction(a, "max", np.max, lax.max, -np.inf, has_identity=False,
File "/home/user/ckng/.conda/envs/jim/lib/python3.10/site-packages/jax/_src/numpy/reductions.py", line 115, in _reduction
raise ValueError(f"zero-size array to reduction operation {name} which has no identity")
ValueError: zero-size array to reduction operation max which has no identity
Write a Jim specific tuning guide
Currently Jim is still in its adolescent stage, which it is usable and energetic but it is not very social and user friendly. To make Jim to be usable and more importantly hired by the community, we need to refactor Jim to have a cleaner API with some features useful for productions runs in the future.
Note that on the top level, while we want to be as user-friendly and modular as possible, this should not come with a cost to performance.
Here is a preliminary list of features needed for the refactoring and next release:
This should be the development main branch where the aforementioned features can be branched off from.
Since we have added boundToUnbound transforms in sample transform, I think it is no longer necessary to build UniformPrior on top of a base distribution of LogisticDistribution. The only things from the prior class that are currently accessed in JIM are sample() and log_prob(), so as long as these two functions provide the correct output, the implementation details of the prior class under the hood are no longer important. We could simplify the implementation.
flowMC has released a new version with some internal API changes. Jim should sync up with flowMC
Adding compatibility of population inference to Jim
Dear,
Peter Pang and I are currently exploring the use of jim to perform parameter estimation on BNS events, for which we wish to use the TaylorF2 waveform model that we have currently implemented in ripple and will push to the main ripple repository once we have cleaned up the code and finalized all checks against the lalsuite implementation. In order to be able to make full use of jim, we are kindly requesting to have a fully working and updated example that shows how to analyze GW170817 with jim with the IMRPhenomD waveform, in order to reproduce the PE reported in the jim paper. We have been trying to get a good run set up but have been struggling with adapting the hyperparameters starting from the provided GW150914 run to get satisfactory results for GW170817 that agree with the results from Bilby or those mentioned in the paper.
Thanks in advance for any help provided!
There is a report on the following breaking mode:
When analyzing GW170817 as shown in the example, if the prior of cos(iota) is changed to [-1,-0.9], then choose a number of chain = 800 will give non-sensible result, such as nan in the log_prob.
Further investigation shows number of chains = 601, 801 will also give the same behavior. But 100, 500, 501 performs normally.
This is very weird, which I think it has to do with the training of the normalizing flow and thinning.
Add Periodic BC for better sampling
In a recent PR, there were two updated examples about GW170817, GW170817_PhenomD_NRTv2.py
and GW170817_TaylorF2.py
. Currently they have some paths that are pointing to a local driver of one of the contributers @ThibeauWouters.
@ThibeauWouters Can you update that two scripts?
Run manager takes bounds as a required likelihood parameter. It would be better to set the default bounds to be the same as the prior bounds. This could reduce the setting that users need to input.
Test all the new implementation of the reparametrization (mass, spin, sky location). Related to #77.
Adding summary and plots output after sampling.
Use autotune for MALA in RunManager can allow the user not to input a mass matrix for the MALA local sampler, which could simplify the setup.
It would be good to have a more generic transformation function/interface than the currently available 1-to-1 mapping function in jim
.
That would be particularly useful for nuclear equation-of-state study, in which one would need to sample over the masses (2 parameters) and the EOS parameters (1 to N), then transform them into the tidal deformabilities (2 parameters) by solving the TOV equations.
SingleEventRun requires path for injection, which is not necessary.
Currently the uniform prior doesn't return -inf
outside the prior bound.
Need to add the guard rail in
Spotted by @AlexanderVerhaeghe
The main focus of jim has been on individual event parameter estimation. One obvious extension that can leverage the infrastructure in jim and flowMC is the addition of population analysis, at least in the way that is commonly done with hierarchical Bayesian analysis.
This would require the following steps:
In order to rehaul the likelihood construction to be more flexible and user-friendly, we plan to have a likelihood class to construct kernels that will eventually be sampled over.
Here is a non-exhaustive list
Tagging @maxisi @tedwards2412
Currently, the optimization routine uses a population adam optimizer to try to compute the reference parameters, there are some cases this optimizer might return nan
, usually due to some boundary issue in the prior.
Without modifying the code, using the unconstrained_uniform
prior should alleviate if not solving the problem.
As a first patch it would be good to able to let the user decide what kind of optimizer they want to use. But more important, I think we should think about setting hard boundary in the prior. Because the sampler relies on gradient heavily, perhaps we should discourage such behavior in general
Marginalizing over extrinsic parameters should alleviate our computational load, but more important it should smooth out the likelihood surface so HMC will have a better time traversing the likelihood surface.
Three marginalization that can be done readily is time of coalescence, phase of coalescence (which is only valid when 22 mode is the dominant mode), and distance.
References:
https://arxiv.org/pdf/1809.02293.pdf
Marginalization over time: https://dcc.ligo.org/public/0114/T1400460/002/margtime.pdf
Marginalization over phase: https://dcc.ligo.org/LIGO-T1300326/public
Marginalization over distance: https://journals.aps.org/prd/abstract/10.1103/PhysRevD.93.024013
It would be nice to allow users input injection parameters or prior parameters in any parametrization, while keeping options for what parameters to be sampled in. Similarly, allowing users to access the posteriors in any parameters would be handy.
Hello, I am very interested in your project. I would like to ask where the files injection_WithParser.py, RealDataAnalysis.py, and utils_plotting.py are located in the project. Could you please provide them. Thank you.
The examples don't necessarily run after simply installing everything from pip.
We need to pin the jax version and put the necessary dependencies into setup.cfg
Currently tracking the naming of parameters is done within the prior class, which limits the flexibility and legibility of parameters within Jim.
It is probably more ergonomic to define parameters and track their naming within Jim, and specify transform only at the interface with the likelihood object and prior object.
This should solve the problem specified in #89 .
To facilitate best practice in maintaining the code base, we should starting rolling in support to adhere to pre-commit checks, which check for styles and typing before commiting
The HeterodynedLikelihood function is not working. Changing from TransientLikelihood to HeterodynedLikelihood in the GW150914.py leads to an error.
Currently the main way to run an inference run is to go through the examples, write a script using the Jim library, and run something along the line of python my_script.py
.
This is good for exploration, but in production, it would be nice to have some more infrastructure that support programmatically generating config files, then run the inference.
Since the prior system described in #108 seems to be working, we should write a guide about the concept and a tutorial for using the new prior system before we merge it into the dev branch
Match_filter_SNR with IMRPhenomC drop extremely fast with slight change in mass, which makes learning the inverse mass matrix very difficult, hence wrong in inference.
A declarative, efficient, and flexible JavaScript library for building user interfaces.
๐ Vue.js is a progressive, incrementally-adoptable JavaScript framework for building UI on the web.
TypeScript is a superset of JavaScript that compiles to clean JavaScript output.
An Open Source Machine Learning Framework for Everyone
The Web framework for perfectionists with deadlines.
A PHP framework for web artisans
Bring data to life with SVG, Canvas and HTML. ๐๐๐
JavaScript (JS) is a lightweight interpreted programming language with first-class functions.
Some thing interesting about web. New door for the world.
A server is a program made to process requests and deliver data to clients.
Machine learning is a way of modeling and interpreting data that allows a piece of software to respond intelligently.
Some thing interesting about visualization, use data art
Some thing interesting about game, make everyone happy.
We are working to build community through open source technology. NB: members must have two-factor auth.
Open source projects and samples from Microsoft.
Google โค๏ธ Open Source for everyone.
Alibaba Open Source for everyone
Data-Driven Documents codes.
China tencent open source team.