MLGW is a Machine Learning model to compute the gravitational waves generated by a Binary Black Hole coalescence. It is part of a thesis project at Università di Pisa under the supervision of prof. Walter Del Pozzo.
The model is released as a Python package mlgw
in the PyPI repository: https://pypi.org/project/mlgw/.
You can install the package with
pip install mlgw
The model outputs the waveform when given the two BHs masses and spins. It implements also the dependence of the waveform on the spherical harmonics.
Version 1 outputs only the 22 dominant mode. The model is presented in this paper (also online as arXiv:2011.01958 ), where we present its details, we assess its accuracy and we employ it for analysing the whole GWTC-1, the first catalog of GW transients. Version 2 and 3 are suited to deal with an arbitrary numbers of modes and at the same time they keeps full compatibility with the previous version interface. Version 2 make use of the usual MoE regression, while version 3 implements a neural network regression. Version 3 is presented in this paper.
To generate a WF:
import mlgw
import numpy as np
generator = mlgw.GW_generator() #creating an istance of the generator
theta = np.array([20,10,0.5,-0.3, 1.43, 1.3, 2.3]) #physical parameters [m1,m2,s1,s2, d_L, iota, phi]
times = np.linspace(-8,0.02, 100000) #time grid at which waves shall be evaluated
h_p, h_c = generator.get_WF(theta, times) #returns amplitude and phase of the wave
You can read much more details about the model in the thesis.
To install the latest dev version, you can do
pip install git+https://github.com/stefanoschmidt1995/MLGW.git
The latest version of the package will be released on the PyPI repository. In that case, you can do:
pip install mlgw
The model can be trained with some ready to use commands. If you installed the packaged normally, they should be already available in your shell.
First, you will need to generate a dataset of waveforms; for example:
mlgw_generate_dataset --n-wfs 50000 --n-grid 2000 --modes 22 21 33 44 55 --basefilename dataset/test_dataset --t-coal 2. --t-step 1e-4 --alpha 0.5 --approximant IMRPhenomTPHM --q-range 1 10 --s1-range -0.9 0.9 --s2-range -0.9 0.9
To generate a NN you will need to create first a PCA model. To create a PCA dataset for the 22 mode:
mlgw_fit_NN --pca-dataset pca_datasets/22 --waveform-dataset dataset/test_dataset.22 --n-comp-amp 4 --n-comp-ph 6 --train-frac 0.85
After, you can train a NN regression for the amplitude:
mlgw_fit_NN --pca-dataset pca_datasets/22 --model-directory model/amp/22 --quantity amp --components 0 1 2 3 4 --max-epoch 100 --units 1 --n-layers 20 --polynomial-order 2
Since there are a lots of models, this is a lot of manual work. You can put all the network architectures, as well as the locations of your models, inside a single yaml config.yml
:
modes: 22 21 33 44 55
pca-dataset: pca_datasets/
model-dir: model/
networks:
amp:
quantity: amp
components: 0 1 2 3
max-epoch: 10000
units: 35
n-layers: 1
polynomial-order: 1
features: mc_chieff
learning-rate: 0.001
ph_01:
quantity: ph
components: 0 1
max-epoch: 10000
units: 50
n-layers: 2
polynomial-order: 3
features: mc_eta_logq_chieff
learning-rate: 0.00003
ph_2345:
quantity: ph
components: 2 3 4 5
max-epoch: 10000
units: 50
n-layers: 1
polynomial-order: 1
features: mc_chieff
learning-rate: 0.001
ph_01_residual:
quantity: ph
components: 0 1
max-epoch: 10000
units: 50
n-layers: 2
polynomial-order: 3
features: mc_eta_logq_chieff
learning-rate: 0.0005
residual:
With this, an handy script will create a condor workflow to perform the training in different, completely parallel jobs:
mbank_write_training_dag config.yml
This will create a dag
file and a sub
file.
To launch the dag:
condor_submit_dag train_NN.dag
All the files of the resulting models will be stored in folder model
in the appropriate structure, ready to be used by mlgw. However, there won't be the PCA dataset inside. To add them, you can use the script cp_pca_model.sh
generated by mbank_write_training_dag
.
And you are all done! To load your model, just type inside a python console:
import mlgw
generator = mlgw.GW_generator('path/to/model/')
and you are good to go!
Two papers describing the methods and different versions of the model have been published:
- S. Schmidt et al., Machine learning gravitational waves from binary black hole mergers (link)
- T. Grimbergen et al., Generating Higher Order Modes from Binary Black Hole mergers with Machine Learning (link)
You can also read two thesis, which described version 1 and version 3 of the code respectively:
- S. Schmidt, Gravitational wave modelling with machine learning (link)
- T. Grimbergen, Deriving the Einstein Field Equations and predicting its Gravitational Wave Solutions using Neural Networks (link)
For more information, you can contact me at [email protected]