Machine learning modelling of the gravitational waves generated by black-hole binaries


Keywords
datanalysis, gravitational-waves, machine-learning
License
Other
Install
pip install mlgw==2.0.0

Documentation

MLGW

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.

Installation

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

How to train the model

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!

Info

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 stefanoschmidt1995@gmail.com