Detecting Gravitational-waves from Extreme Mass Ratio Inspirals using Convolutional Neural Networks

Galactic nuclei are extreme environments for stellar dynamics Amaro-Seoane (2018). The stellar densities can be as high as $10^{6}\ \rm M_{\odot}\ pc^{-3}$ and its velocity scatter can exceed $1000\ \rm km\ s^{-1}$. Within a galactic nucleus, a large number of COs, including stellar-mass Black Holes, White Dwarfs, and Neutron Stars, are expected to exist. They are considered to be within the “loss cone” if their orbit would lead to their capture by the central MBH, leading to the formation of EMRIs. COs outside the loss cone can fill the loss cone through dynamical processes like two-body relaxation Amaro-Seoane (2018).

In addition to this traditional channel, some new formation channels of EMRI systems have also been proposed. It has been conjectured that a CO can be captured by the accretion disk around central MBH( i.e., Active Galactic Nuclei (AGN) disk). Effects like wind effects, density wave generation, and dynamic friction can assist the inward migration of COs, which further boosts the formation of EMRIs Pan and Yang (2021); Pan et al. (2021).

Although EMRIs are speculated to exist, no direct observational evidence is available yet (c.f.Arcodia et al. (2021)). Therefore, currently, the population properties of EMRIs are derived from theoretical modeling. However, components within the theoretical modeling, like the mass and spin distribution of the central MBH, the erosion of the stellar cusp, the M-$\sigma$ relation, etc. are uncertain, which lead to a very uncertain prediction on EMRI population Babak et al. (2017b). Indeed, different assumptions on EMRI population lead to orders of magnitude difference in terms of detection rate by TianQin Fan et al. (2020). Throughout this analysis, an model is adopted to simulate EMRI population. To be exact, we adopt the model M12 from Babak et al. (2017b).

Theoretical modelling of EMRI waveforms is essential for the successful detection of GWs from these systems. To accurately compute the EMRI waveform, one has to consider the effects of the gravitational self-force, which leads to a deviation of CO orbit from the test particle geodesics under the Kerr metric of the MBH Barack (2009). Unfortunately, waveform simulations with numerical relativity cannot be used for EMRI due to prohibitively high computational costs caused by the very high mass ratio of such systems. Although some progress has been made in adding self-force into the waveform generation Barack (2009); Van De Meent (2018); Miller and Pound (2021); Hughes et al. (2021); McCart et al. (2021); Lynch et al. (2021); Isoyama et al. (2021), a fully realistic solution for EMRI is still yet to be achieved. Nevertheless, there are methods that can generate EMRI waveforms. For example, the Teukolsky-based waveform Hughes et al. (2005) and Numerical Kludge (NK) waveform Babak et al. (2007) solves the Teukolsky equation with the geodesic under Kerr metric. Such methods suffer from low efficiency, while the Analytic Kludge (AK) method that adopts the post-Newtonian (PN) formula can achieve waveform generation in a much faster way Barack and Cutler (2004). The Augmented Analytic Kludge (AAK) waveform Chua and Gair (2015a) can reach an accuracy similar to NK with the generating speed of AK. For our study, since a large number of simulated waveforms are needed for the training, testing, and validation of our CNN method, we focus on the analytic family of waveform models and concentrate on both AK and AAK waveforms. We remark that during the writing of this manuscript, a new and fast waveform has been implemented Katz et al. (2021), which could be used to improve efficiency in the future, but we do not expect any significant difference in our findings. Besides, Katz et al. (2021) adopts the Schwarzschild Black Hole (BH) assumption, while we would like to also study the impact of MBH spin.

The AK method first calculates the orbit of CO through post-Newtonian equations, considering the radiation back-reaction and the Lense-Thirring effect. The overall evolution of the CO orbit is governed by three fundamental frequencies, namely the orbital angular frequency $f_{\rm o}$, the periapsis precession frequency $f_{\rm p}$, and the Lense-Thirring precession frequency $f_{\rm LT}$. The waveform is then calculated based on quadrupolar approximation. Notice that the overall setup of a waveform is not self-consistent, hence the name of a “kludge”. However, this choice of formulation allows a fast way of producing waveforms at the cost of its accuracy.

Under the transverse and traceless gauge, two polarizations of AK waveform are defined via a $n$-harmonic waveform which is a function of physical parameters ¹¹1The details of the parameters adopted are presented in Table 2:

where $\hat{L}$ is the orbital angular momentum, $\hat{n}$ is the position of the source, and $\gamma$ is an azimuthal angle measuring the direction of pericenter. The coefficients ($a_{n},b_{n},c_{n}$) with $n$ ranges from 1 to 5, are determined by the eccentricity and mean anomaly orbital phase Barack and Cutler (2004). The orbital phase is determined by three fundamental frequencies $f_{\rm o},\ f_{\rm p},$ and $f_{\rm LT}$. More details on waveform generation can be found in Barack and Cutler (2004); Fan et al. (2020). Notice that since the self-force can not be properly incorporated by the AK waveform, the spin parameter for the CO is not included.

The AK waveform quickly deviates from physical systems, due to intrinsic inconsistency. Under the AK model, instead of using the physical parameters, one can adopt the effective parameters to evolve the orbit and produce a waveform similar to the NK model. The NK model is more accurate as it combines PN orbital evolution with Kerr geodesics. Chua et al. implemented the AAK waveform by establishing the link between the physical parameter and the effective parameter. To obtain such a link, one can evolve the NK waveform for a short period and determine the mapping relationship to derive the correct fundamental frequencies $f_{\rm o}$, $f_{\rm p}$, and $f_{\rm LT}$ Chua and Gair (2015b). In this way, the AAK waveform can achieve both the accuracy of NK and the speed of AK.

For a given gravitational wave detector, one can use the antenna pattern to describe its response to a source at a given location, polarization, and frequency Cutler (1998); Cornish and Larson (2001); Cornish and Rubbo (2003). A typical EMRI signal lasts for a long time. If we observe from the point of view of the detector, then a given source would change in location as well as in polarisation. One needs to know the orbit of the spacecraft to derive the overall evolution of the signal as seen at the detector. In this work, we focus our attention on the proposed TianQin mission Luo et al. (2016); Zhang et al. (2020) for the following analysis.

The TianQin mission is designed to be a constellation with three satellites, which accommodate test masses in free fall and use drag-free control to follow these test masses. The satellites share an Earth-central orbit and form a regular triangle. They are connected by laser links and can form Michelson interferometers to detect GW. The geocentric nature of the TianQin orbit makes it adopts a “3-month on 3-month off” work scheme, that the detectors would operate continuously for three months followed by three months of break. Each pair of neighboring laser arms can form interference on the vertex satellite, whose noise can be described by the power spectral density (PSD):

where $f$ is the frequency, $S_{a}=1\times 10^{-30}\ {\rm m^{2}\ s^{-4}/\text{Hz}}$ describes the acceleration noise, $S_{x}=1\times 10^{-24}\ \rm m^{2}/\rm Hz$ describes the positional noise, and $f_{*}=\frac{c}{2\pi L_{\rm arm}}$ is the transfer frequency with $c$ being the speed of light and $L_{\rm arm}=\sqrt{3}\times 10^{8}\ {\rm m}$ is the arm length.

For the low frequency, where $f<f_{*}$, the three interferometers can be combined into two orthogonal channels, corresponding to the two GW polarisations. So the recorded GW signal can be expressed as:

where $F_{I,II}^{+}$ and $F_{I,II}^{\times}$ are the antenna pattern for both plus and cross polarisations, and

Here $\theta_{S}$ and $\phi_{S}$ are the sky location of the source, and $\psi_{S}$ is the polarization angle. All quantities are defined in the detector frame, which makes it time-dependent Liu et al. (2020).

The two effective channels differ only in polarisation angles by $\pi/4$, which means

Finally, a correction term is added to account for the Doppler effect induced by the orbital motion of the detectors Liu et al. (2020).

where $f$ is the frequency of GW, $R=1$AU and barred symbols denote parameters in Solar System Barycentre frame, with $\bar{\Phi}(t)=\bar{\phi}_{0}+2\pi t/T,\ T=1\rm year$ is Earth’s orbital period around the Sun, and $\bar{\phi}_{0}$ is the initial location of TianQin at time t = 0. Notice that in higher frequencies when $f\geq f_{*}$, the long-wavelength approximation breaks down and the two channels are no longer uncorrelated, hence an extra correction term is needed in Equation 2 Cornish and Larson (2001); Cornish and Rubbo (2003). Fortunately, EMRI signals we are interested in, are located mostly within the low-frequency range. This is because the highest orbital frequency of CO is determined by the last stable orbit(LSO), where $\nu_{\rm LSO}=1/(2\pi M)\times 6^{3/2}$. Given a typical MBH with $10^{5}M_{\odot}$, $\nu_{LSO}=0.04<f_{*}=0.28$ Hz. Therefore, we stick with the long-wavelength approximation throughout the work.

As an illustration, we show in Figure 1 an example EMRI signal (on both channels) on the TianQin sensitivity curve. This EMRI system consists of a 10 $M_{\odot}$ BH and MBH with 4.5 $\times 10^{6}\ M_{\odot}$ and a spin of 0.97. The Signal-to-Noise Ratio (SNR) of this signal equals 50.

We define the overall SNR as the root sum square of inner products of waveform with itself across two channels of TianQin,

Here we use “ $\langle g|h\rangle$” to represent the inner product between the two frequency domain waveforms $\tilde{g}(f)$ and $\tilde{h}(f)$. Notice that in the last line, the SNR in a channel is simply the integration of squared modulus of whitened signal over the frequency range between $f_{\rm low}=10^{-4}\ \rm Hz$ and $f_{\rm high}=1\ \rm Hz$, which are sensitivity frequency band of TianQin.

For the detection of the signal in noise we are going to use CNNs. In order to train and verify CNN properly, one needs to divide the data into three groups, namely training data, validation data, and testing data. In these datasets, half of them are signal+noise samples, and the rest are noise-only samples. Using training data, the CNN learns how to classify samples. Validation data is incorporated in the training process to verify that the CNN is learning correctly. Finally, testing data can test the efficiency of the trained CNN.

In our case, we aim to accomplish the task of classifying the data into two categories: with signal, or without signal. One can express the data $d$ as the addition of random Gaussian noise $n$ and the GW signal $h$ (in the case it is present in the data):

The raw output of TianQin is three correlated time domain data channels. After some preprocessing they can be reduced to two channels with uncorrelated noise, plus a channel with a response to astrophysical signals highly suppressed. This process is implemented by the Time Delay Interferometry (TDI) Estabrook et al. (2000); Tinto and Armstrong (1999). Out of simplicity, we do not implement TDI, but rather treat TianQin as two orthogonal Michelson interferometers. In Figure 2, we illustrate a data sample that contains an EMRI signal, using the parameters demonstrated in Figure 1.

We prepare the generation of noise and signal separately. Due to the “3-month on 3-month off” work scheme of TianQin, the longest chunk of continuous data lasts only three months. To avoid the complexity introduced by considering the data gap issue, we focus our analysis on data length of three months (or 7864320 seconds). By adopting the long-wavelength approximation, it is also sufficient to assume the sample rate of $1/30$ Hz, which leads to the data size of 262144 samples. The parameters are summarised in Table 1.

The detector noise was assumed to be Gaussian and was generated by sampling from the one-sided PSD given by Equation 1. For the simulation of an EMRI waveform signal, we used both AK and AAK models. For different purposes that we will detail in the following sections, we generate 7 groups of signals, whose properties are explained in Table 4, and the shared default parameter distributions are listed in Table 2. The data are whitened in the frequency domain. When data are converted back to the time domain, we apply a Tukey window function with the coefficient $\alpha=1/4$ to avoid the spectrum leakage.

Notice that Table 2 does not reflect an astrophysical model, and one might argue that an astrophysically motivated distribution for EMRI parameters could be more realistic. However, an astrophysically motivated distribution is not necessarily the best choice for training the CNN. The reason for that is that the astrophysically motivated models are subject to huge uncertainty. Therefore sticking with any specific model runs a risk of biasing the network.

SNR of a signal is inversely proportional to the distance. When assessing the performance of a CNN, it is convenient to compare signals with similar SNRs than with similar distances. Therefore, we do not fix distribution of distance $D_{L}$ in Table 2. The choice of SNR range is determined by balancing the hardware requirement between data storage, waveform generation speed, and I/O speed in training.

CNN is a specific realization of Neural Network (NN) which features the inclusion of convolution operations Ng et al. (2013); Heaton (2017). In NN, each “neuron” represents a non-linear mapping or an affine transformation function. The non-linearity is achieved by an activation function. A typical realization of a CNN includes convolution layers, pooling layers, and fully-connected layers (or dense layers) LeCun et al. (2015). It can be treated as a classifier. Such a CNN composed of multiple neurons can be understood as a non-linear function $f_{w}(\cdot)$ that maps the input space of the data to the output space:

where $y_{i}$ is the output probability for the $i$-th data $d_{i}$ to belong to a certain class, $w$ are weights of the CNN.

Based on training data a CNN can learn features of the input data by updating the weights associated with different neurons by minimizing the difference between the output prediction and training data labels. The size of our training data is large, therefore we cannot optimize weights with all the data at once and need to adopt mini-batching. The training process is divided into multiple epochs. For each epoch, the weights are updated by the optimizer algorithm. When the training process converges, the CNN can classify new data into different categories.

To achieve the reliable ability to classify, one needs, first, to define the loss function $L(y_{\rm true},y_{\rm pred})$, which quantifies the difference between the CNN output and labels in the training process. This could be done by adopting the binary cross-entropy function:

where $N$ is the total number of training data samples, for the $j$-th data instance. The noise-only samples are labeled as $y^{j}_{\rm true}=0$, and samples containing signals are labeled as $y^{j}_{\rm true}=1$. The output of CNN $y^{j}_{\rm pred}$ represents the assigned probability for each category. The loss function is minimized when the predicted value for the label class matches the one from the training data with the highest probability. In the training process, the weights are updated through the Nadam algorithm, such that the loss function can be minimized by following the adaptive gradient descent with Nesterov momentum Dozat (2016). During training, validation data is used to test the trained CNN after each epoch and use those result to evaluate the performance of the CNN and its stability towards overfitting.

In a CNN different layers play different roles. In the convolution layer, features can be extracted by applying convolution between the input and the convolution kernel. After multiple convolution layers with various filters, a CNN can learn different features of the input data from lower-level to higher-level. The purpose of the pooling layer is to reduce the size of the convolved features while remaining the key features so that one can reduce computation complexity and avoid over-fitting. In the fully-connected layer, neurons are connected with all activation values to the previous layer, with the final output indicating the predicted class probability. In this work, rectified linear unit (relu) is selected as the activation function to all but the output layer, for which the softmax function is adopted. Relu function is convenient and efficient to avoid vanishing gradient problems when training a CNN Krizhevsky et al. (2012), and softmax function at the last layer achieves a normalization so that the output is constrained in the range of (0,1). In this way, the output can appropriately represent the probability that the input sample belongs to one of the classes. The performance of a CNN depends on many hyper-parameters, like CNN depth, convolutional layer number, number of kernels, and their respective size.

In order for CNN to gain the classification ability, a certain amount of input data is necessary for the training. More samples are needed to train CNN to detect weaker signals. In the actual training, it turns out that the I/O of training data becomes the bottleneck. Therefore, we are limited in the size of total data, which in turn determines the SNR lower limit. After some trial and error, we converged at 50 as the SNR lower limit in our training sample. On the other hand, stronger signals can be detected more easily. We concentrate on signals with SNR comparable with the lower limit, therefore an upper limit of SNR 120 is adopted. The total amount of $2.5\times 10^{5}$ signal plus noise samples are used to train the final CNN.

With the data ready, we need to train the network to maximize the classification ability without over-fitting. In general, a more complicated network is expected to perform better on the training data. However, a too-complicated neural network runs a risk of over-fitting or being too acquainted with the training data so that the performance on training data is not generalizable towards unfamiliar data. Therefore one needs to check the consistency of the CNN performance on both training data and validation data. One can identify the over-fitting if the performance on the training data is significantly better than the validation data. This can be relieved by adopting a less complicated neural network, adding dropout layers, or increasing the size of the training set.

Finally, we need to fix the training architecture. We start our trial of convolution kernel size of (1,5) as suggested by previous studies Gabbard et al. (2018); Bayley et al. (2020), and finally settled with a size of (1,34) through trial and error. We apply the max-pooling in the pooling layer, which can retain the maximum value of each patch of the input. The mini-batch size is set to 56, and we train the CNN with 300 epochs. For each mini-batch, 56 simulated datasets will be randomly selected from samples, half of which are signal plus noise samples and the other half are noise-only samples. The wall clock time for each epoch of training is about 2 hours. Strategies like early stopping Yao (2007) are also employed to increase efficiency. In our experiments, the training process will be stopped when either 300 epochs are reached, or the accuracy of validation data no longer increases for 50 epochs. Finally, it takes about 10.5 days to train the final CNN. This trained model can afterwards be used for testing. We summarise the architecture of the CNN in Table 3.

Once a CNN is trained, we test the performance with 7 different groups of testing data, the properties of which are introduced in Table 4. The setup of these groups is motivated by the testing of two main points, namely the validity of the CNN, and the sensitivity of the CNN. For testing the validity, we prepare three groups of signal+noise samples: signals that follow the same distribution as the training data, signals that follow an astrophysically motivated distribution, and signals produced by AAK but with the same distribution as the training data. For testing the sensitivity, we apply the CNN on top of the remaining four groups, in which many intrinsic parameters are fixed, but either mass, redshift, or spin is enumerated over a certain list. Taking group 4 as an example, it replaces the SNR distribution in group 1 with a set of fixed SNRs, taking values $10,\ 20,\ 30,\ \ldots,\ 130$. Similarly, in group 5, it enumerates the mass of the MBH, taking values of $10^{4},\ 10^{4.5},\ \ldots,\ 10^{7}\ M_{\odot}$. In group 6, it enumerates the spin of the MBH in a similar way, taking values of $0.0,\ 0.2,\ 0.4,\ 0.6,\ 0.8,\ 0.98$ with fixed MBH mass is $10^{6}\ M_{\odot}$. In group 7, it focuses on EMRI systems with three kinds of MBHs: 1. $M=10^{5.5}\ M_{\odot}$ and $a=0.98$; 2. $M=10^{6}\ M_{\odot}$ and $a=0.98$; 3. $M=10^{6}\ M_{\odot}$ and $a=0$. And they enumerates the redshift, taking the value of $0.1,\ 0.15,\ 0.2,\ 0.25,\ 0.3,\ 0.35.$

The trained CNN takes the testing data and outputs a probability $y_{\rm pred}$, which predicts the probability that the data contains a signal. Conversely, $1-y_{\rm pred}$ predicts the probability that it is purely Gaussian noise.

To measure the effectiveness of the trained CNN model, we used the testing data of groups 1-3 in Table 4. One can compare the efficiencies among detection methods through the Receiver Operator Characteristics (ROC) curve, where the True Alarm Probability (TAP) is plotted against the False Alarm Probability (FAP).

In the calculation of the ROC, one first chooses a value of $y^{*}$ as a detection threshold. Any value of $y_{\rm pred}>y^{*}$ is regarded as an alarm. At a given $y^{*}$, the FAP and the TAP is constructed as the fraction of noise-only and s͡ignal plus noise samples, respectively, that are reported as an alarm. By varying $y^{*}$, one obtain pairs of TAP-FAPs, and a line that links all these pairs is the ROC curve. By definition, the TAP is a monotonic function of the FAP.

The blue line in Figure 3 represents a ROC curve for group 1. The dotted line is the worst possible performance which indicates a total inability of distinguishing signal from noise. One can observe that for FAP equal to 1%, the corresponding TAP is approximately 92%. This fact indicates that when the CNN is applied to data constructed in a similar way to training data, it has a very good ability to distinguish the signal from noise. We remark that the results from group 1 are expected to be the upper limit among all groups and could be used as a benchmark.

The purple line in Figure 3 corresponds to the ROC curve for group 2. Compared with the blue line, the purple line is lower: for a FAP of 1%, the TAP is about 75%. Since group 2 is an astrophysically motivated population, the majority of signals are expected to be associated with relatively low SNRs. Due to the choice of SNR threshold of 50 for the training data, we observe clustering of SNRs just above 50. This significant decrease in TAP can be explained by such clustering. On the other hand, even though the training data differ in the parameter distribution to the testing data, the trained CNN still has a good ability to detect signals, showing a good generalization ability.

The red line in Figure 3 corresponds to the ROC curve of group 3. In this group, the parameter distribution is the same as in group 1, but a different set of waveform models, namely the AAK waveform, is adopted. This group could be used to study the robustness of the CNN against choices of EMRI waveforms. The similarity between the blue line and the red line indicates that a CNN trained on AK waveforms can achieve good detection ability on AAK waveforms. One should be cautious not to take this conclusion for granted. Although both AK and AAK adopt similar equations, the setup of the waveform calculation makes them dephase quickly after a timescale of around an hour. This observation gives us confidence that the CNN method could be robust, and even if a fully realistic waveform would not be available when space-borne GW detectors are operating, one can still rely on CNN trained on approximate waveforms to achieve good detection potential up to certain SNR values.

One can present the sensitivity of the CNN over parameters like SNR, MBH mass, MBH spin, or redshift, using the distribution of signal parameters group 4-7 in Table 4. We have adopted fixed FAPs of 0.1 and 0.01. The results are presented in Figure 4 in the form of efficiency curves.

Figure 4(a) displays the TAP versus the SNR for signals from group 4. It shows that two efficiency curves are consistent with the expectation that the CNN exhibits higher sensitivity towards stronger signal, and for SNR of higher than about 100, the CNN both can guarantee a certain detection at false alarm probabilities of 0.01 and 0.1.

We study the effectiveness of the CNN over MBH masses with group 5. The MBH spin is fixed as $a=0.98$, and redshifts are varied through $z=0.1,\ 0.2,\ 0.3$. It can be seen from Figure 4(b) that EMRI systems containing a $10^{5.5}M_{\odot}$ MBH are more easily detected by the trained CNN. For such an EMRI with a redshift of $0.1$, the CNN can achieve a TAP of 78% conditioned on a FAP of 0.1. The feature of a sensitivity peak with regards to MBH masses stems from two factors: the sensitivity curve, and the GW signal frequency distribution. Larger MBH masses lead to larger amplitudes and lower frequencies, while the detector is most sensitive only for a limited frequency band. This result is also consistent with horizon distance analysis in a previous study Fan et al. (2020).

EMRI signals contained in group 6 have a fixed MBH mass ($M=10^{6}M_{\odot}$) and fixed redshifts ($z=0.1,\ 0.2,\ 0.3$). From Figure 4(c) one can observe that there is almost no apparent dependency of detection ability on MBH spin. This is a piece of indirect evidence that the detection ability is most sensitive to the SNR. Further investigation shows that EMRI systems with different spins share similar SNRs. This can be roughly explained as follows: the EMRI systems have a smooth evolution of overall GW amplitude before the final merger. According to our setup, the merger is fixed at the end of the three months observation period. Therefore, the spin parameter has a minor impact on the overall SNR.

We present the efficiency curve for group 7 in Figure 4(d). The SNR of a signal is roughly inversely proportional to luminosity distance, which for the nearby universe is approximately proportional to redshift. Therefore, on the vertical axis we plot the TAP multiplied by the redshift. If the TAP is linearly correlated with the SNR, then such multiplication should remain constant. This is exactly what we observed in the blue solid line ($M=10^{5.5}M_{\odot}$ and spin =0.98, FAP equals 0.1). The distribution of SNR within this redshift range (0.1-0.35) roughly spans the range of 20-60, which shows good linearity between SNR and TAP in Figure 4(a). The prediction is violated to the biggest degree for the purple dashed line ($M=10^{6}M_{\odot}$ and $a=0.98$, FAP equals 0.01). Closer events are gaining more detection efficiency, which can be expected from the non-linearity of the dashed line in the similar SNR range in Figure 4(a).

From the above analyses, we can conclude that among all factors, the SNR has the strongest influence on the expected detection efficiency. Changes in other parameters can also lead to a different performance in TAP, but such differences can be mostly explained by the different SNRs.