XLVI International Symposium on Multiparticle Dynamics

11institutetext: M. Smoluchowski Institute of Physics, Jagiellonian
University, Lojasiewicza 11, 30-348 Krakow, Poland

# Saturation and geometrical scaling

Michal Praszalowicz 11<EMAIL_ADDRESS>
###### Abstract

We discuss emergence of geometrical scaling as a consequence of the non-linear
evolution equations of QCD, which generate a new dynamical scale, known as the
saturation momentum: $Q_{\rm s}$. In the kinematical region where no other
energy scales exist, particle spectra exhibit geometrical scaling (GS), i.e.
they depend on the ratio $p_{\rm T}/Q_{\rm s}$, and the energy dependence
enters solely through the energy dependence of the saturation momentum. We
confront the hypothesis of GS in different systems with experimental data.

##  1 Introduction

In this report we present a concise analysis of GS, slightly extended with
respect to the presentation given at the XLVI International Symposium on
Multiparticle Dynamics. One can find more details in the original publications
Praszalowicz:2012zh – Praszalowicz:2015dta  and in the recent conference
proceedings Praszalowicz:2016jcg , that cover the same topics.

In QCD we have basically two sets of evolution equations that describe the
change of parton densities with decreasing resolution scale $1/Q^{2}$ – DGLAP
equations, or with growing energy (or equivalently with decreasing Bjorken
$x$) – BFKL equation. In both cases the number of partons, or more precisely
the number of gluons, is growing rapidly with the evolution variable. In the
BFKL case however (since the average transverse size of gluons is fixed), we
enter a regime where the partonic system is not dilute and the linear
evolution breaks down. A modified BFKL equation that includes the non-linear
terms is known as the Balitsky-Kovchegov (BK) equation BK . One of the
consequences of the nonlinearities is the emergence of the so called
saturation scale of the form Mueller:2002zm ; Munier:2003vc :

$Q_{\text{s}}^{2}(x)=Q_{0}^{2}({x}/{x_{0}})^{-\lambda}.$ (1)

Munier and Peschanski Munier:2003vc  draw an analogy between the BK equation
and the time evolution of the wave front $u(t,z)$ in one dimensional space
variable $z$:

$\frac{\partial}{\partial t}u(t,z)={\cal O}(\partial/\partial z)$ (2)

where ${\cal O}(\partial/\partial z)$ is a non-linear differential operator
corresponding in QCD to the BK kernel. For a wide class of operators ${\cal
O}$ and initial conditions for $u$, wave front $u$ converges asymptotically to
the traveling wave:

$u(x,z)\rightarrow u(z-v_{\rm c}t),$ (3)

i.e. to the fixed front that moves rigidly with the critical velocity $v_{\rm
c}$. In the QCD context time corresponds to the logarithm of Bjorken $x$:
$t=\ln(x_{0}/x)$ where $x_{0}$ is a constant corresponding to $t=0$, and
$z=\ln(p_{\rm T}^{2}/Q_{0}^{2})$. Translating the argument of the travelling
wave to the QCD variables gives:

$z-v_{\rm c}t=\ln\left(\frac{p_{\rm T}^{2}}{Q_{0}^{2}}\right)-v_{\rm
c}\ln\left(\frac{x_{0}}{x}\right)=\ln\left(\frac{p_{\rm
T}^{2}}{Q_{0}^{2}(x/x_{0})^{-v_{\rm c}}}\right).$ (4)

Therefore the travelling wave corresponds to the scaling solution with the
saturation momentum given by Eq. (1). In the following we shall check whether
GS is present in different pieces of high energy data.

##  2 Deep inelastic scattering (DIS)

Geometrical scaling was first introduced in the context of DIS for
$F_{2}(x)/Q^{2}$ Stasto:2000er . In Fig. 1 we plot $F_{2}(x)/Q^{2}$ as a
function of $Q^{2}$ (left panel) and in terms of the scaling variable
$\tau=Q^{2}/Q_{\rm s}^{2}(x)$ for $\lambda=0.329$ (right panel) for the
combined HERA data HERAdata . Points of different colors correspond to
different Bjorken $x$’s. We see from Fig. 1 that DIS data scale very well with
some exception in the right part of Fig. 1.b. These points, however,
correspond to large Bjorken $x$’s where GS is supposed to break.

Figure 1: Combined DIS data HERAdata  for $F_{2}/Q^{2}$. Different points
forming a wide band as a function of $Q^{2}$ in the left panel correspond to
different Bjorken $x$’s. They fall on a universal curve when plotted in terms
of $\tau$ (right panel). (Figure from the first paper of Ref.
Praszalowicz:2012zh ).

Since the reduced cross-section in ep scattering (which is essentially
proportional to $F_{2}/Q^{2}$) is given as a convolution of the virtual photon
wave function and an unintegrated gluon distribution of the proton
$\varphi_{p}\left({\vec{k}_{\text{T}}^{\,2}},x\right)$, the fact the DIS data
scale implies that
$\varphi_{p}\left({{k}_{\text{T}}^{\,2}},x\right)=\varphi_{p}\left({{k}_{\text{T}}^{\,2}}/Q^{2}_{\rm
s}(x)\right)$.

##  3 Inelastic $p_{\mathrm{T}}$ spectra at the LHC

The cross-section for not too hard gluon production in pp collisions can be
described in the $k_{\text{T}}-$factorization approach by the formula
Gribov:1981kg :

$\frac{d\sigma}{dyd^{2}p_{\text{T}}}=\frac{3\pi\alpha_{\text{s}}}{2p_{\text{T}}^{2}}{\displaystyle\int}{d^{2}\vec{k}_{\text{T}}}\,\varphi_{p}\left({\vec{k}_{\text{T}}^{\,2}},x_{1}\right)\varphi_{p}\left({(\vec{k}-\vec{p}\,)_{\text{T}}^{2},x_{2}}\right)$
(5)

where $\varphi_{p}$ denotes the unintegrated gluon distribution. Here
$x_{1,2}=e^{\pm y}p_{\rm T}/\sqrt{s}$. In the following we shall assume that
produced gluons are in the mid rapidity region ($y\simeq 0$), hence both
Bjorken $x$’s of colliding gluons are equal $x_{1}\simeq x_{2}$ (denoted in
the following as $x$).

Unintegrated gluon distributions have dimension of transverse area. We shall
assume that – following our discussion of DIS – that they depend only on the
ratio of the transverse momentum over the saturation scale

$\varphi_{p}\left({k}_{\text{T}}^{\,2},x\right)=S_{\bot}\phi\left({k}_{\text{T}}^{\,2}/Q_{\rm
s}^{2}(x)\right)$ (6)

where $\phi$ is a dimensionless function. In the case of DIS
$S_{\bot}=\sigma_{0}$ is the dipole-proton cross-section for large dipoles
GolecBiernat:1998js . In the case of heavy ion collisions $S_{\bot}$ is the
transverse size of an overlap of two large nuclei for a given centrality
class. In both cases one can assume that $S_{\bot}$ is energy independent (or
weakly dependent). In that case $d^{2}\vec{k}_{\text{T}}$ integration in (5)
leads to

$\frac{d\sigma}{dyd^{2}p_{\text{T}}}=S_{\bot}^{2}\mathcal{F}(\tau)$ (7)

where $\tau=p_{\rm{T}}^{2}/Q_{\rm s}^{2}(x)$ is a scaling variable and
$\mathcal{F}(\tau)$ is a function related to the integral of $\phi_{p}$’s. We
follow here the parton-hadron duality hypothesis PHD , assuming that the
charged particle spectra are on the average identical to the gluon spectra.

In Fig. 2 we plot ALICE pp data Abelev:2013ala  in terms of $p_{\mathrm{T}}$
(left panel) and in terms of the scaling variable $\tau$ (right panel) for
$\lambda=0.32$. We see that three different curves from the left panel in Fig.
2 overlap over some region if plotted in terms of $\sqrt{\tau}$. The exponent
for which this happens over the largest interval of $\tau$ is $\lambda=0.32$
Praszalowicz:2015dta , which is the value compatible with our model
independent analysis of the DIS data Praszalowicz:2012zh .

Figure 2: Data for pp scattering from ALICE Abelev:2013ala  plotted in terms
of $p_{\mathrm{T}}$ and $\sqrt{\tau}$. Full (black) circles correspond to
$W=7$ TeV, down (red) triangles to 2.76 TeV and up (blue) triangles to 0.9
TeV. (Figure from Ref. Praszalowicz:2015dta .)

In order to illustrate the method of adjusting $\lambda$, we plot in Fig. 3
ratios of the cross-sections at 7 TeV to 2.76 and 0.9 TeV. Proximity of both
ratios to unity for $\lambda=0.32$ is the sign of GS for $p_{\text{T}}$ up to
$4.25$ GeV$/c$ Praszalowicz:2015dta .

Figure 3: Ratios of the cross-sections at 7/2.76 TeV – down (red) triangles
and 7/0.9 TeV – up (blue) triangles, for $\lambda=0$ (left) and 0.32 (right).
(Figure from Ref. Praszalowicz:2015dta .)

It has been argued previously McLerran:2010ex  that GS should hold for
multiplicities, rather than for the cross-sections. This would be true if the
relation between the two was energy independent. This may be the case in HI or
p$A$ collisions where we trigger on some $S_{\bot}$ by selecting the
centrality classes with given number of participants, but it is not true in
the case of the inelastic pp scattering:

$\frac{dN}{dyd^{2}p_{\text{T}}}=\frac{1}{\sigma^{\mathrm{MB}}(W)}\frac{d\sigma}{dyd^{2}p_{\text{T}}}=\frac{S_{\bot}^{2}}{\sigma^{\mathrm{MB}}(W)}\mathcal{F}(\tau)$
(8)

where the minimum bias cross-section $\sigma^{\mathrm{MB}}(W)\neq S_{\bot}$ is
energy-dependent. Nevertheless the multiplicity does scale, however for
different value of the exponent ($\lambda=0.22-0.24$), and over a smaller
range of $p_{\rm T}$. This is discussed in more detail in Ref.
Praszalowicz:2015dta .

##  4 Electron-positron scattering

One of the obvious questions concerning pp scattering is whether the observed
scaling is indeed a consequence of gluon saturation or whether it appears due
to the properties of final state radiation that takes place before
hadronization (note that by referring to the parton-hadron duality we have
ignored possible effects of the QCD fragmentation functions). The best way to
address this issue is to perform analogical analysis of the hadronic $p_{\rm
T}$ spectra in e+e- collisions where no GS is expected. To this end we shall
use spectra published by TASSO collaboration Braunschweig:1990yd  at
$W=\sqrt{s}=14$, 22, 35 and 44 GeV.111The author is grateful to J. Chwastowski
for pointing this reference. One might argue that these energies are too small
to reach a firm conclusion on GS in e+e-. However, in Ref. Praszalowicz:2013uu
we have analyzed pp scattering at even lower energies (6 – 17 GeV) finding
positive evidence of GS in the mid rapidity region.

In Fig. 4 ratios of multiplicity spectra

$R(W)=\left.\frac{1}{\sigma_{tot}}\frac{d\sigma}{2\pi p_{\rm T}dp_{\rm
T}}\right|_{44\leavevmode\nobreak\ {\rm
GeV}}{\mbox{\LARGE/}}\left.\frac{1}{\sigma_{tot}}\frac{d\sigma}{2\pi p_{\rm
T}dp_{\rm T}}\right|_{W}$ (9)

for the remaining three values of $W$ are plotted in terms of the scaling
variable $\sqrt{\tau}$ for $\lambda=0$ (note that in this case
$\sqrt{\tau}=p_{\rm T}$) and for $\lambda=0.32$. In the left panel w see
similar behaviour of $R$ as in the pp case (see Fig. 3.a), however in the
right panel where we plot $R$ for $\lambda=0.32$ no GS is seen (compare with
Fig. 3.b).

Figure 4: Ratios $R(W)$ of multiplicity spectra for $W=14$ (down-green
triangles), 22 (up-blue triangles) and 35 GeV (red-full circles) as functions
of $p_{\rm T}$ (left panel) and $\sqrt{\tau}$ for $\lambda=0.32$ (right
panel).

##  5 Summary

In this short note we have argued that geometrical scaling is clearly seen in
the DIS and pp data. In the latter case an observable that scales is the
differential cross-section rather than the multiplicity distribution.
Phenomenological analysis analogical to the one described here has been used
to look for the effects of GS in p$A$ and heavy ion collisions and in $\langle
p_{\rm T}\rangle$ correlation with multiplicity McLerran:2010ex
–Praszalowicz:2016jcg . Most recently effects of the fluctuations of the
saturation momentum in the p$A$ collisions have been studied in Ref.
McLerran:2015lta .

Scaling violations that emerge in the kinematical limit when one of the
Bjorken $x$’s in (5) is close to unity, have been observed in Ref.
Praszalowicz:2013uu . Here, in Sect. 4, we have presented a new evidence that
GS is not present in the spectra of particles produced in e+e- collisions.
Although this has been theoretically expected, conclusions drawn from Fig. 4
support our main hypothesis that geometrical scaling seen in the experimental
data is due to the saturation effects in the initial state of the colliding
hadrons.

## Ackowledgements

The author wants to thank the organizers for an invitation and for putting up
together this successful meeting. This work was supported by the Polish NCN
grant 2014/13/B/ST2/02486.

## References

  * (1) M. Praszalowicz, T. Stebel, JHEP 1303, 090 (2013) and JHEP 1304, 169 (2013).
  * (2) L. McLerran, M. Praszalowicz, Acta Phys. Pol. B 41, 1917 (2010) and Acta Phys. Pol. B 42, 99 (2011).
  * (3) M. Praszalowicz, Acta Phys. Pol. B 42, 1557 (2011) and arXiv:1205.4538 [hep-ph].
  * (4) M. Praszalowicz, A. Francuz, Phys. Rev. D 92, 074036 (2015) .
  * (5) M. Praszalowicz, Acta Phys. Polon. Supp. 9, 395 (2016).
  * (6) I. Balitsky, Nucl. Phys. B 463, 99 (1996);
Y. V. Kovchegov, Phys. Rev. D 60, 034008 (1999) and Phys. Rev. D 61, 074018
(2000).

  * (7) A. H. Mueller and D. N. Triantafyllopoulos, Nucl. Phys. B 640, 331 (2002).
  * (8) S. Munier, R. B. Peschanski, Phys. Rev. Lett. 91, 232001 (2003) and Phys. Rev. D 69, 034008 (2004).
  * (9) A. Stasto, K. Golec-Biernat, J. Kwiecinski, Phys. Rev. Lett. 86, 596 (2001).
  * (10) C. Adloff et al. [H1], Eur. Phys. J. C 21 33 (2001);
S. Chekanov et al. [ZEUS], Eur. Phys. J. C 21, 443 (2001) .

  * (11) L. V. Gribov, E. M. Levin, M. G. Ryskin, Phys. Lett. B 100, 173 (1981).
  * (12) K. J. Golec-Biernat, M. Wusthoff, Phys. Rev. D 59, 014017 (1998) and Phys. Rev. D 60, 114023 (1999).
  * (13) Ya.I. Azimov, Yu.L. Dokshitzer, V.A. Khoze and S.I. Troian, Z. Phys. C 27, 65 (1985);
Yu. L. Dokshitzer, V. A. Khoze and S. I. Troyan, J. Phys. G 17, 1585 (1991) ;
V.A. Khoze and W. Ochs, Int. J. Mod. Phys. A 12, 2949 (1997) .

  * (14) B. B. Abelev et al. [ALICE Collaboration], Eur. Phys. J. C 73, 2662 (2013).
  * (15) W. Braunschweig et al. [TASSO Collaboration], Z. Phys. C 47, 187 (1990). doi:10.1007/BF01552339
  * (16) M. Praszalowicz, Phys. Rev. D 87, 071502 (2013).
  * (17) L. McLerran and M. Praszalowicz, Annals Phys. 372, 215 (2016).


# Compartmental Models for COVID-19 and Control via Policy Interventions

Swapneel Mehta  New York UniversityCenter for Data Science
<EMAIL_ADDRESS>and  Noah Kasmanoff  New York UniversityCenter for Data
Science<EMAIL_ADDRESS>
###### Abstract.

We demonstrate an approach to replicate and forecast the spread of the SARS-
CoV-2 (COVID-19) pandemic using the toolkit of probabilistic programming
languages (PPLs). Our goal is to study the impact of various modeling
assumptions and motivate policy interventions enacted to limit the spread of
infectious diseases. Using existing compartmental models we show how to use
inference in PPLs to obtain posterior estimates for disease parameters. We
improve popular existing models to reflect practical considerations such as
the under-reporting of the true number of COVID-19 cases and motivate the need
to model policy interventions for real-world data. We design an SEI3RD model
as a reusable template and demonstrate its flexibility in comparison to other
models. We also provide a greedy algorithm that selects the optimal series of
policy interventions that are likely to control the infected population
subject to provided constraints. We work within a simple, modular, and
reproducible framework to enable immediate cross-domain access to the state-
of-the-art in probabilistic inference with emphasis on policy interventions.
We are not epidemiologists; the sole aim of this study is to serve as a
exposition of methods, not to directly infer the real-world impact of policy-
making for COVID-19.

##  1\. Introduction

In order to understand and control infectious diseases, it is important to
build realistic models capable of accurately replicating and projecting their
transmission in a region (Atkeson, 2020b; Tang and Wang, 2020; Sameni, 2020).
The underlying assumption being that a model capable of replicating disease
spread has captured the true causal dynamics sufficiently well. Motivated by
this, there has been a spate of research in applying SIR, SEIR, and SEI3RD
models to this problem (COVID et al., 2020; López and Rodó, 2020). These
belong to a class of compartmental models that are underpinned by Lotka-
Volterra dynamics that divide a population into sections or compartments and
describe the probabilistic transitions between them through a set of partial
differential equations. We extend this direction of work with an emphasis on
modeling under-reported cases, decoupling policy interventions from
compartmental transitions, estimating the impact of policy interventions,
selecting a sequence of optimal interventions to control the spread of
diseases, and using the SEI3RD variant as an extension of existing work on
SEIR models since it forms a template for many other extensions of
compartmental models (Giordano et al., 2020; Senapati et al., 2020; Kennedy et
al., 2020; Wolff, 2020; Grimm et al., 2021; Winters, 2020). The authors of
(Hong and Li, 2020), much like us, propose a new statistical tool to visualize
analyses of COVID-19 data. However, our framework makes it much simpler to
inspect the intricacies of modeling assumptions, expand with a custom set of
constraints, and experiment with counterfactual simulations without the need
for extensive compute or data. We have created some demo notebooks which will
be made available publicly
111https://drive.google.com/drive/folders/1Npdn4bS_qlps5EdA6vXlMvRXSljGgYCd?usp=sharing.

In light of probabilistic programming languages (PPLs) (Bingham et al., 2019;
van de Meent et al., 2018; Salvatier et al., 2016; Carpenter et al., 2017)
reaching their ’coming-of-age’ moment, COVID-19 modeling is being explored to
guide and support policy decisions and decision-makers (Wood et al., 2020; de
Witt et al., 2020). In this work, we provide data scientists with a concrete
example of applying probabilistic inference to understand disease spread and
control. To epidemiologists, we offer this manuscript as a guide to design
compartmental models to evaluate the impact of policy interventions (Mandal et
al., 2020; Giordano et al., 2020; Wang et al., 2020b). We include clear
motivation and detailed real-world examples of how to manually explore the
impact of such non-pharmaceutical interventions (NPIs) using synthetically
generated data and provide an algorithm to automatically generate strategies
for implementing governmental policies for disease control.

Our main contributions are as follows:

  * •

Improve existing compartmental models with the ability to deal with the under-
reporting of COVID-19 cases and decoupling NPIs from disease parameters.

  * •

Evaluate our parameter estimates through empirical comparisons with those in
prevalent literature, observed patterns in testing coverage, and highlight the
relative consistency of our predictions compared to anomalies in existing
approaches.

  * •

Model fixed policy interventions and describe an algorithm to select adaptive
policy interventions to aid government efforts to limit the spread of disease.

  * •

Highlight the ease of using PPLs to design, extend, and fit SEI3RD models
using approximate inference to obtain posterior disease parameter estimates;
in addition to having an open-source code-base and our self-contained tutorial
to be released at the time of publication.

##  2\. Susceptible-Infected-Recovered (SIR) Model Dynamics

The SIR model dynamics form a template for the increasingly complex
compartmental models we explore. These dynamics are encapsulated in the
following differential equations following the Lotka-Volterra Model:

$\frac{dS}{dt}=-\beta S(t)i(t)\quad,\qquad\frac{dI}{dt}=\beta S(t)i(t)-\gamma
I(t)\quad,\qquad\frac{dR}{dt}=\gamma I(t)$

SIR$\beta$$\gamma$ Figure 1. The SIR Compartmental Model
SEIR$\beta$$\sigma$$\gamma$ Figure 2. Add an ’Exposed’ compartment to the SIR
model to obtain the SEIR model

##  3\. Modeling Noisy Observations

The population-level disease parameters we discussed for the SIR model
(transmission probability $\beta$ and recovery rate $\gamma$) are usually
estimated in terms of measurable quantities from observed data such as the
mean recovery time. The response rate $\rho$ indicates the proportion of
observed infections in reality because we typically cannot expect to observe
every single case of infection. To be clear, it is often the case that people
may not realize they have been infected in the absence of extensive testing.
This has been the case in most countries at least for the initial few months
of the pandemic and some strategies attempt to remedy it in different ways
(Khan et al., 2021; Lachmann, 2020; Jagodnik et al., 2020). Since the real-
world data we intend to use overlaps significantly with this duration, it
makes sense to add this variable to our model. Importantly, the same may be
exacerbated due to socioeconomic and political factors, so $\rho$ allows us to
model this under-reporting of COVID-19 cases in practice. We start by making a
guess about the range for the reproduction number
$R_{0}=\dfrac{\beta}{\gamma}$ and response rate $\rho$. This ’guess’ is
equivalent to placing empirically informed (from past outbreaks, for instance)
priors on these latent variables and estimate their posterior distributions.

###  3.1. Using PPLs for Modeling and Inference

PPLs are a natural candidate among the tools we considered for this problem.
We would like to define a time-series probabilistic model with some underlying
dynamics that encodes our assumptions about the data-generating process; then
we want to fit this model to the data through stochastic variational inference
(Hoffman et al., 2013; Ranganath et al., 2014), and use the posterior
distributions over its parameters to make predictions about future time-steps.
As a sanity check, we also want to test how well it replicates the historical
trajectory of the disease through simulations. Pyro (Bingham et al., 2019), a
deep universal probabilistic programming language, is an excellent candidate
for this because of the following reasons:

  * •

It comprises of thin wrappers around PyTorch (Paszke et al., 2019)
distributions allowing us to write complex generative models interweaving
stochastic and deterministic control flow while still within a familiar and
popular machine learning (ML) framework.

  * •

It offers general-purpose inference algorithms out of the box that allows us
to shift the focus from designing custom inference algorithms to building
expressive models.

  * •

It recently extended support in terms of an API for epidemiological models. We
perform our experiments in Pyro since it is a much more accessible framework
due to the community around it in comparison to modern alternatives like
Pyprob (Baydin et al., 2019). We can confirm this empirically since we spent
some time reproducing work along similar lines in (Wood et al., 2020) and
found much less time required to do the same in Pyro. We do still appreciate
the relative ease of reproducing it compared to other machine learning
research for healthcare.

We consider a partially observed population and want to understand the
putative controls that could achieve a goal defined as ”control the spread of
the disease”, ”reduce the death rate”, or other such equivalent outcomes.
While some of the parameters that define disease spread are controllable,
there are certain non-controllable parameters including properties of the
disease that we will infer. From a computational standpoint, we demonstrate
the capability of the universal probabilistic programming language, Pyro, to
combine a system of equations that define a simulator and perform inference
over the latent variables within the simulation to obtain a posterior
distribution over the model parameters. In addition to an inference engine,
Pyro offers the capability to intervene on the variables within this
simulation in order to obtain potential outcomes of policy changes that can be
expressed within the language as fixing the values of certain stochastic
parameters.

(a) Posterior distribution for $R0$

(b) Posterior distribution for $\rho$

Figure 4. Predicting Future Infections for a Simulated Dataset in a Highly
Infectious Setting

##  4\. Experimental Evaluation of Compartmental Models in Real-world Settings

For our simulation-based study, we used models $M={SIR,SEIR,SEI3RD}$ and
refined versions of each, $M(i)$, wherein we start with an infected population
of 0.01%. We conduct some simulation studies for instance by setting the
$R_{0}=3.0$ and $\rho=0.85$. We perform variational inference (Hoffman et al.,
2013; Ranganath et al., 2014) to obtain posterior parameter estimates close to
the true values. These are indicated by the vertical black line in 3(a) and
3(b). Since the posterior point-estimates are close to the true values for the
simulation, we can see why we are able to accurately replicate the disease
spread shown in 4. We plot the daily infections versus time and the plot
resembles a single wave of infected patients. The shaded portion indicates a
$90\%$ confidence interval for the model’s predictions which closely tracks
the true values of disease spread. While it is no surprise that SIR, SEIR, and
SEI3RD models perform well on simulated data, there are extensive studies in
applying the SEIR model in practice. However, we believe that the SEIR model
also has its failures (see Qualitative Analysis of Successes and Failures)
that often go undetected due to a lack of comprehensive evaluation. We show
this through an extensive set of experiments on real-world data across
multiple geographies. The tabular comparison of the reproduction number or
$R0$ and the $\rho$ estimated via our models are shown in Table 1 for
different time periods which offers the following insights:

  * •

The reported cases for the initial period of January - July 2020 indicate far
less testing than the full periods (upto May 2021 for the USA and January 2021
for the rest of the world, in our dataset).

  * •

The estimated $R0$ seemed high for most regions upto April 2020, and even
though longer-term estimates denoted an $R0$ of a little over 1 for most
regions, the spate of cases underscores the need to reduce the spread of the
infection. At the same time, the reproduction number is not all that matters
and a careful study of mortality rates is warranted as testing increases and
fatalities decrease in order to draw concrete scientific conclusions and
policy recommendations given this evidence.

  * •

We show qualitatively why it is necessary to build more granular, flexible
compartmental models and underscore the need for separately modeling policy
interventions as human influence on these processes.

The references for Comparing $R0$ estimates with literature are drawn from
prevalent literature and expert-curated
resources222https://epiforecasts.io/covid/posts/national/united-states/
333https://covidestim.org/us/GA on the subject (Kamalich Muniz-Rodriguez et
al., [n. d.]; Lau et al., 2020; Prodanov, 2021; Gunzler and Sehgal, 2020). A
certain George E. P. Box would be wont to say ”All $R0$ estimates are wrong
but some may be useful”, and we illustrate this via an empirical and rather
qualitative route. Firstly, in some cases the reference method itself presents
anomalous forecasts like the $R0$ for Germany throughout the year being
predicted as an unusual 22.032, and that of the Netherlands, 9.103 from
(Prodanov, 2021). In comparison, we observe relative consistency at least at
the model level, which might be a strong signal to consider more realistic
modeling choices when using real, noisy data for estimation. Where we lack
consistency, we have a signal in the form of a set of forecasts as well as an
additional estimate of $\rho$. An oddly low $\rho$, for instance, might
convince us of an estimation error despite having high-confidence and a good
fit to the observed data. We observe that the SEI3RD models perform better
than the SEIR and SIR models which might lack the capacity to effectively
model transitions. We also study a refined set of compartmental models which
consider an initial, partially infected (0.01%) population instead of starting
from a single infection (’patient zero’).

Once we successfully model disease progression and forecast for a set of
future time-steps to recover potential new infections. We can examine
counterfactual questions of the nature ’What would have happened to the number
of cases if the government enacted $X$ steps at $Y$ time?’ by introducing
policy interventions $u$ at this stage and figure out the global minimally
invasive intervention to perform that will remain within the desired
thresholds for the infected populace. An adaptive algorithm to determine the
optimal policy interventions is developed in this paper as we continue to
explore breadth-wise analyses pertaining to COVID-19 control, and compare our
models to the impact of real-world interventions. This discussion, albeit a
core contribution, is relegated to the appendix.

##  5\. The Reproduction Number

As responsible data scientists attempting to provide a tool for
epidemiological analysis, it is important not to overstate the relevance of
estimating the correct $R0$ from the data since that is not the only parameter
of interest. Many numbers have been bandied about in the news under the
assumption that controlling $R0$ implies controlling the pandemic. To a
certain degree this is true, however there are important caveats to this
notion that some have expanded upon (Hébert-Dufresne et al., 2020; Maruotti et
al., 2021). The summary of the discussion is that $R0$ must be used as a tool
to paint a partial picture of a disease spread, in conjunction with multiple
factors. In this regard, our consideration of $\rho$, modeling of infected
fractions of populations, and consideration of policy interventions are
significant steps taken to provide a comprehensive idea of the state of a
pandemic.

Table 1.  Comparing $R0$ estimates with literature Region | Model | Jan - April, ’20 | Jan - Dec, ’20
---|---|---|---
|  | $R0$ | $\rho$ | Ref. $R0$ | $R0$ | $\rho$ | Ref. $R0$
Italy | SIR | 1.48 $\sigma=.01$ | 0.187 | 2.676 | 1.34 | 0.125 | 1.8024
SIR(i) | 1.712 | 0.76 $\sigma=.03$ | - | 0.608 | 0.506 | -
SEIR | 2.27 $\sigma=.02$ | 0.11 | - | 1.49 | 0.14 | -
SEIR(i) | 0.403 | 0.56 | - | 1.96 | 0.178 | -
SEI3RD | 3.03 $\sigma=.07$ | 0.557 | - | 5.78 | 0.001 | -
SEI3RD(i) | 1.247 | 0.532 | - | 0.215 | 0.505 | -
Netherlands | SIR | 1.69 | 0.36 $\sigma=.01$ | 1.9962 | 1.724 $\sigma=0.84$ | 0.71 | 9.103
SIR(i) | 1.699 $\sigma=.01$ | 0.54 $\sigma=.015$ | - | 0.490 | 0.5079 | -
SEIR | 4.51 $\sigma=.011$ | 3.24 $\sigma=.01$ | - | 1.76 | 0.45 | -
SEIR(i) | 10.048 | 0.71 | - | 0.236 | 0.50 | -
SEI3RD | 4.23 $\sigma=.04$ | 0.735 | - | 2.78 | 0.154 | -
SEI3RD(i) | 3.763 $\sigma=.03$ | 0.70 | - | 3.51 | 0.108 | -
New York | SIR | 1.432 | 0.54 $\sigma=.016$ | 1.21 | 1.475 | 0.027 | 0.81
SIR(i) | 1.643 | 0.594 | - | 2.078 | 0.619 | -
SEIR | 0.88 | 0.50 | - | 0.81 | 0.50 | -
SEIR(i) | 4.09 $\sigma=.01$ | 0.04 | - | 4.98 | 0.14 | -
SEI3RD | 1.53 | 0.504 | - | 1.67 | 0.503 | -
SEI3RD(i) | 10.79 $\sigma=.006$ | 0.0795 | - | 12.295 | 0.253 | -
Germany | SIR | 0.108 | 0.5186 | 1.637 | 1.610 | 0.120 | 22.032
SIR(i) | 1.643 | 0.594 | - | 2.078 | 0.619 | -
SEIR | 2.51 $\sigma=.013$ | 0.76 $\sigma=.01$ | - | 2.93 $\sigma=.1$ | 0.154 | -
SEIR(i) | 1.973 | 0.213 | - | 5.02 $\sigma=.47$ | 0.57 | -
SEI3RD | 0.572 | 0.55 | - | 5.56 $\sigma=.08$ | 0.01 | -
SEI3RD(i) | 1.37 | 0.52 | - | 4.7 | 0.124 | -
Georgia | SIR | 3.74 | 0.303 | 1.45 | 3.296 | 0.343 | 0.84
SIR(i) | 4.075 | 0.302 | - | 3.661 | 0.039 | -
SEIR | 5.3 $\sigma=.17$ | 0.015 | - | 3.47 | 0.56 | -
SEIR(i) | 4.97 $\sigma=.078$ | 0.014 | - | 5.36 | 0.096 | -
SEI3RD | 4.09 $\sigma=.17$ | 0.5483 | - | 1.32 | 0.501 | -
SEI3RD(i) | 6.8 | 0.049 | - | 5.62 | 0.205 | -

###### Acknowledgements.

We would like to acknowledge guidance and support from Kyle Cranmer and Rajesh
Ranganath that laid the foundations for this research project.

## References

  * (1)
  * Atkeson (2020a) Andrew Atkeson. 2020a.  _How deadly is COVID-19? Understanding the difficulties with estimation of its fatality rate_.  Technical Report. National Bureau of Economic Research.
  * Atkeson (2020b) Andrew Atkeson. 2020b.  _What will be the economic impact of COVID-19 in the US? Rough estimates of disease scenarios_.  Technical Report. National Bureau of Economic Research.
  * Atkeson (2021) Andrew Atkeson. 2021\.  _A Parsimonious Behavioral SEIR Model of the 2020 COVID Epidemic in the United States and the United Kingdom_.  Technical Report. National Bureau of Economic Research.
  * Baydin et al. (2019) Atilim Güneş Baydin, Lei Shao, Wahid Bhimji, Lukas Heinrich, Lawrence Meadows, Jialin Liu, Andreas Munk, Saeid Naderiparizi, Bradley Gram-Hansen, Gilles Louppe, et al. 2019\.  Etalumis: Bringing probabilistic programming to scientific simulators at scale. In _Proceedings of the international conference for high performance computing, networking, storage and analysis_. 1–24.
  * Bertozzi et al. (2020) Andrea L Bertozzi, Elisa Franco, George Mohler, Martin B Short, and Daniel Sledge. 2020.  The challenges of modeling and forecasting the spread of COVID-19.  _Proceedings of the National Academy of Sciences_ 117, 29 (2020), 16732–16738.
  * Bingham et al. (2019) Eli Bingham, Jonathan P Chen, Martin Jankowiak, Fritz Obermeyer, Neeraj Pradhan, Theofanis Karaletsos, Rohit Singh, Paul Szerlip, Paul Horsfall, and Noah D Goodman. 2019\.  Pyro: Deep universal probabilistic programming.  _The Journal of Machine Learning Research_ 20, 1 (2019), 973–978.
  * Buckman et al. (2020) Shelby R Buckman, Reuven Glick, Kevin J Lansing, Nicolas Petrosky-Nadeau, and Lily M Seitelman. 2020\.  Replicating and projecting the path of COVID-19 with a model-implied reproduction number.  _Infectious Disease Modelling_ 5 (2020), 635–651.
  * Carpenter et al. (2017) Bob Carpenter, Andrew Gelman, Matthew D Hoffman, Daniel Lee, Ben Goodrich, Michael Betancourt, Marcus A Brubaker, Jiqiang Guo, Peter Li, and Allen Riddell. 2017\.  Stan: a probabilistic programming language.  _Grantee Submission_ 76, 1 (2017), 1–32.
  * Chen et al. (2020) Haonan Chen, Jing He, Wenhui Song, Lianchao Wang, Jiabao Wang, and Yijin Chen. 2020\.  Modeling and interpreting the COVID-19 intervention strategy of China: A human mobility view.  _PloS one_ 15, 11 (2020), e0242761.
  * Chu et al. (2020) Derek K Chu, Elie A Akl, Stephanie Duda, Karla Solo, Sally Yaacoub, Holger J Schünemann, Amena El-harakeh, Antonio Bognanni, Tamara Lotfi, Mark Loeb, et al. 2020\.  Physical distancing, face masks, and eye protection to prevent person-to-person transmission of SARS-CoV-2 and COVID-19: a systematic review and meta-analysis.  _The Lancet_ 395, 10242 (2020), 1973–1987.
  * COVID et al. (2020) Team IHME COVID, RC Reiner, RM Barber, and JK Collins. 2020\.  Modeling COVID-19 scenarios for the United States.  _Nature medicine_ (2020).
  * de Witt et al. (2020) Christian Schroeder de Witt, Bradley Gram-Hansen, Nantas Nardelli, Andrew Gambardella, Rob Zinkov, Puneet Dokania, N Siddharth, Ana Belen Espinosa-Gonzalez, Ara Darzi, Philip Torr, et al. 2020\.  Simulation-Based Inference for Global Health Decisions.  _arXiv preprint arXiv:2005.07062_ (2020).
  * Giordano et al. (2020) Giulia Giordano, Franco Blanchini, Raffaele Bruno, Patrizio Colaneri, Alessandro Di Filippo, Angela Di Matteo, and Marta Colaneri. 2020\.  Modelling the COVID-19 epidemic and implementation of population-wide interventions in Italy.  _Nature medicine_ 26, 6 (2020), 855–860.
  * Giudici and Raffinetti (2020) Paolo Giudici and Emanuela Raffinetti. 2020.  Monitoring Covid-19 Policy Interventions.  _Frontiers in Public Health_ 8 (2020).
  * Grimm et al. (2021) Veronika Grimm, Friederike Mengel, and Martin Schmidt. 2021\.  Extensions of the SEIR model for the analysis of tailored social distancing and tracing approaches to cope with COVID-19.  _Scientific Reports_ 11, 1 (2021), 1–16.
  * Gunzler and Sehgal (2020) Douglas D. Gunzler and Ashwini R. Sehgal. 2020.  Time-Varying COVID-19 Reproduction Number in the United States.  _medRxiv_ (2020).  https://doi.org/10.1101/2020.04.10.20060863 arXiv:https://www.medrxiv.org/content/early/2020/04/15/2020.04.10.20060863.full.pdf
  * Hébert-Dufresne et al. (2020) Laurent Hébert-Dufresne, Benjamin M. Althouse, Samuel V. Scarpino, and Antoine Allard. 2020.  Beyond R0: the importance of contact tracing when predicting epidemics.  _medRxiv_ (2020).  https://doi.org/10.1101/2020.02.10.20021725 arXiv:https://www.medrxiv.org/content/early/2020/02/12/2020.02.10.20021725.full.pdf
  * Hoffman et al. (2013) Matthew D Hoffman, David M Blei, Chong Wang, and John Paisley. 2013\.  Stochastic variational inference.  _Journal of Machine Learning Research_ 14, 5 (2013).
  * Hong and Li (2020) Hyokyoung G Hong and Yi Li. 2020.  Estimation of time-varying reproduction numbers underlying epidemiological processes: A new statistical tool for the COVID-19 pandemic.  _PloS one_ 15, 7 (2020), e0236464.
  * Hoseinpour Dehkordi et al. (2020) Amirhoshang Hoseinpour Dehkordi, Majid Alizadeh, Pegah Derakhshan, Peyman Babazadeh, and Arash Jahandideh. 2020.  Understanding epidemic data and statistics: A case study of COVID-19.  _Journal of medical virology_ 92, 7 (2020), 868–882.
  * International Monetary Fund ([n. d.]) International Monetary Fund. [n. d.].  Policy Responses to COVID-19.  https://www.usa.gov/coronavirus
  * Jagodnik et al. (2020) Kathleen M Jagodnik, Forest Ray, Federico M Giorgi, and Alexander Lachmann. 2020.  Correcting under-reported COVID-19 case numbers: estimating the true scale of the pandemic.  _medRxiv_ (2020).
  * Kamalich Muniz-Rodriguez et al. ([n. d.]) DrPH Kamalich Muniz-Rodriguez, Gerardo Chowell, Jessica S Schwind, Randall Ford, Sylvia K Ofori, Chigozie A Ogwara, Margaret R Davies, Terrence Jacobs, Chi-Hin Cheung, Logan T Cowan, et al. [n. d.].  Time-varying Reproduction Numbers of COVID-19 in Georgia, USA, March 2, 2020 to November 20, 2020.  ([n. d.]).
  * Kennedy et al. (2020) Deanna M Kennedy, Gustavo José Zambrano, Yiyu Wang, and Osmar Pinto Neto. 2020.  Modeling the effects of intervention strategies on COVID-19 transmission dynamics.  _Journal of Clinical Virology_ 128 (2020), 104440.
  * Khan et al. (2021) Soban Qadir Khan, Imran Alam Moheet, Faraz Ahmed Farooqi, Muhanad Alhareky, and Faisal Alonaizan. 2021.  Under-reported COVID-19 cases in South Asian countries.  _F1000Research_ 10, 88 (2021), 88.
  * Kucharski et al. (2020) Adam J Kucharski, Timothy W Russell, Charlie Diamond, Yang Liu, John Edmunds, Sebastian Funk, Rosalind M Eggo, Fiona Sun, Mark Jit, James D Munday, et al. 2020\.  Early dynamics of transmission and control of COVID-19: a mathematical modelling study.  _The lancet infectious diseases_ 20, 5 (2020), 553–558.
  * Lachmann (2020) Alexander Lachmann. 2020\.  Correcting under-reported COVID-19 case numbers.  _MedRxiv_ (2020).
  * Lau et al. (2020) Max SY Lau, Bryan Grenfell, Michael Thomas, Michael Bryan, Kristin Nelson, and Ben Lopman. 2020\.  Characterizing superspreading events and age-specific infectiousness of SARS-CoV-2 transmission in Georgia, USA.  _Proceedings of the National Academy of Sciences_ 117, 36 (2020), 22430–22435.
  * Leung et al. (2020) Nancy HL Leung, Daniel KW Chu, Eunice YC Shiu, Kwok-Hung Chan, James J McDevitt, Benien JP Hau, Hui-Ling Yen, Yuguo Li, Dennis KM Ip, JS Malik Peiris, et al. 2020\.  Respiratory virus shedding in exhaled breath and efficacy of face masks.  _Nature medicine_ 26, 5 (2020), 676–680.
  * Li et al. (2020) Ruiyun Li, Sen Pei, Bin Chen, Yimeng Song, Tao Zhang, Wan Yang, and Jeffrey Shaman. 2020.  Substantial undocumented infection facilitates the rapid dissemination of novel coronavirus (SARS-CoV-2).  _Science_ 368, 6490 (2020), 489–493.
  * Liang et al. (2020) Mingming Liang, Liang Gao, Ce Cheng, Qin Zhou, John Patrick Uy, Kurt Heiner, and Chenyu Sun. 2020.  Efficacy of face mask in preventing respiratory virus transmission: A systematic review and meta-analysis.  _Travel medicine and infectious disease_ 36 (2020), 101751.
  * López and Rodó (2020) Leonardo López and Xavier Rodó. 2020.  The end of social confinement and COVID-19 re-emergence risk.  _Nature Human Behaviour_ 4, 7 (2020), 746–755.
  * Manchein et al. (2020) Cesar Manchein, Eduardo L Brugnago, Rafael M da Silva, Carlos FO Mendes, and Marcus W Beims. 2020\.  Strong correlations between power-law growth of COVID-19 in four continents and the inefficiency of soft quarantine strategies¡? A3B2 show [editpick]?¿.  _Chaos: An Interdisciplinary Journal of Nonlinear Science_ 30, 4 (2020), 041102.
  * Mandal et al. (2020) Sandip Mandal, Tarun Bhatnagar, Nimalan Arinaminpathy, Anup Agarwal, Amartya Chowdhury, Manoj Murhekar, Raman R Gangakhedkar, and Swarup Sarkar. 2020.  Prudent public health intervention strategies to control the coronavirus disease 2019 transmission in India: A mathematical model-based approach.  _The Indian journal of medical research_ 151, 2-3 (2020), 190\.
  * Maruotti et al. (2021) Antonello Maruotti, Massimo Ciccozzi, and Fabio Divino. 2021\.  On the misuse of the reproduction number in the COVID-19 surveillance system in Italy.  _Journal of Medical Virology_ (2021).
  * Ndaïrou et al. (2020) Faïçal Ndaïrou, Iván Area, Juan J Nieto, and Delfim FM Torres. 2020.  Mathematical modeling of COVID-19 transmission dynamics with a case study of Wuhan.  _Chaos, Solitons & Fractals_ 135 (2020), 109846.
  * Paszke et al. (2019) Adam Paszke, Sam Gross, Francisco Massa, Adam Lerer, James Bradbury, Gregory Chanan, Trevor Killeen, Zeming Lin, Natalia Gimelshein, Luca Antiga, et al. 2019\.  Pytorch: An imperative style, high-performance deep learning library.  _arXiv preprint arXiv:1912.01703_ (2019).
  * Prodanov (2021) Dimiter Prodanov. 2021\.  Analytical parameter estimation of the SIR epidemic model. Applications to the COVID-19 pandemic.  _Entropy_ 23, 1 (2021), 59.
  * Ranganath et al. (2014) Rajesh Ranganath, Sean Gerrish, and David Blei. 2014\.  Black box variational inference. In _Artificial intelligence and statistics_. PMLR, 814–822.
  * Salvatier et al. (2016) John Salvatier, Thomas V Wiecki, and Christopher Fonnesbeck. 2016\.  Probabilistic programming in Python using PyMC3.  _PeerJ Computer Science_ 2 (2016), e55.
  * Sameni (2020) Reza Sameni. 2020\.  Mathematical modeling of epidemic diseases; a case study of the COVID-19 coronavirus.  _arXiv preprint arXiv:2003.11371_ (2020).
  * Senapati et al. (2020) Abhishek Senapati, Sourav Rana, Tamalendu Das, and Joydev Chattopadhyay. 2020.  Impact of intervention on the spread of COVID-19 in India: A model based study.  _arXiv preprint arXiv:2004.04950_ (2020).
  * Tang and Wang (2020) Yuanji Tang and Shixia Wang. 2020.  Mathematic modeling of COVID-19 in the United States.  _Emerging microbes & infections_ 9, 1 (2020), 827–829.
  * Thompson (2020) Robin N Thompson. 2020\.  Epidemiological models are important tools for guiding COVID-19 interventions.  _BMC medicine_ 18, 1 (2020), 1–4.
  * United States Government ([n. d.]) United States Government. [n. d.].  Government Response to Coronavirus, COVID-19.  https://www.usa.gov/coronavirus
  * van de Meent et al. (2018) Jan-Willem van de Meent, Brooks Paige, Hongseok Yang, and Frank Wood. 2018.  An introduction to probabilistic programming.  _arXiv preprint arXiv:1809.10756_ (2018).
  * Vasconcelos et al. (2020) Giovani L Vasconcelos, Antônio MS Macêdo, Raydonal Ospina, Francisco AG Almeida, Gerson C Duarte-Filho, Arthur A Brum, and Inês CL Souza. 2020.  Modelling fatality curves of COVID-19 and the effectiveness of intervention strategies.  _PeerJ_ 8 (2020), e9421.
  * Verity et al. (2020) R Verity, LC Okell, and I Dorigatti. 2020.  Estimates of the severity of coronavirus disease 2019: a model-based analysis (vol 20, pg 669, 2020).  (2020).
  * Wang et al. (2020a) Chaolong Wang, Li Liu, Xingjie Hao, Huan Guo, Qi Wang, Jiao Huang, Na He, Hongjie Yu, Xihong Lin, An Pan, et al. 2020a.  Evolving epidemiology and impact of non-pharmaceutical interventions on the outbreak of coronavirus disease 2019 in Wuhan, China.  _MedRxiv_ (2020).
  * Wang et al. (2020b) Lili Wang, Yiwang Zhou, Jie He, Bin Zhu, Fei Wang, Lu Tang, Michael Kleinsasser, Daniel Barker, Marisa C Eisenberg, and Peter XK Song. 2020b.  An epidemiological forecast model and software assessing interventions on the COVID-19 epidemic in China.  _Journal of Data Science_ 18, 3 (2020), 409–432.
  * Winters (2020) Jack M Winters. 2020\.  A Novel Model for Simulating COVID-19 Dynamics Through Layered Infection States that Integrate Concepts from Epidemiology, Biophysics and Medicine: SEI3R2S-Nrec.  _medRxiv_ (2020).
  * Wolff (2020) Michael Wolff. 2020\.  On build-up of epidemiologic models—Development of a SEI3RSD model for the spread of SARS-CoV-2.  _ZAMM-Journal of Applied Mathematics and Mechanics/Zeitschrift für Angewandte Mathematik und Mechanik_ 100, 11 (2020), e202000230.
  * Wood et al. (2020) Frank Wood, Andrew Warrington, Saeid Naderiparizi, Christian Weilbach, Vaden Masrani, William Harvey, Adam Scibior, Boyan Beronov, and Ali Nasseri. 2020.  Planning as inference in epidemiological models.  _arXiv preprint arXiv:2003.13221_ (2020).

##  Appendix A Appendix

##  Appendix B Qualitative Analysis of Successes and Failures

We have attempted to fit many different kinds of models and obtained certain
parameter estimates which we compare to other, equally noisy parameter
estimates from the literature. The fact is, each country has dealt with
COVID-19 in different ways. While China seemed to be able to take stringent
measures and limit the spread, others like Italy were not able to follow the
same approach. There was a strong focus on the fast-rising numbers of infected
people in the United States (September, 2020), quickly overshadowed by a worse
state of affairs in India (May, 2021). It is crucial to analyse how well each
model is able to not only fit the data but also forecast the spread of disease
accurately. Since we have trained some models on partial amounts of data, let
us take a look at which of the models can identify upcoming ’waves’ in the
progress of the disease.

For instance, in the figures 5 (vertical bars mark the beginning of
forecasts), the SEIR and the SEI3RD models are able to capture the second
wave, with the latter being able to, impressively enough, predict the second
one having only observed the first wave. Similarly, in 6 we can see that both
the SIR and SEIR models seem to perform poorly at data fitting and
forecasting. Note that in both cases, the notion of the ’best model’ as
indicated by comparing estimated $R0$ does not seem to hold, emphasising the
need to consider the least restrictive models when conducting such studies. Of
course, the SEI3RD model is not a panacea. We offer two additional examples of
its forecasts 8 that convince us we need to model an external factor that
allows us to modify the second, larger wave (peak) in a manner similar to how
government interventions delayed the infections from successively peaking as
per the predictions of the SEI3RD model. Also in 8 is the availability of
confidence intervals for the forecast. In general, while the SEI3RD model
performs well at capturing multiple waves of COVID-19 infections, we are
making a fundamentally incorrect modeling assumption by ignoring the fact that
disease evolution was hindered by human intervention, in particular via
government policies designed to limit its spread. In fact, we can utilise
these ’vanilla’ SEI3RD models that are reasonably good at predicting the
potential spread of a disease to help us figure out how best to limit it! This
motivates us to study Policy Interventions.

Figure 5. Fitting SIR (top left), SEI3RD (bottom right) and SEIR models to
data from Italy Figure 6. Fitting SIR (top left), SEIR (top right) and SEI3RD
models to data from the Netherlands Figure 7. Fitting SIR (top left), SEIR
(top right) and SEI3RD models to data from New York Figure 8. Fitting SEI3RD
models to data from Germany and Georgia (right)

##  Appendix C Compartmental Model Extensions in Pyro

There is a vast body of literature on compartmental models for epidemiology.
In particular, the SEIR model seems the most popular and widely applied for
real-world correspondence possibly owing to a trade-off between simplicity and
effectiveness (Buckman et al., 2020). In (COVID et al., 2020) the team fits an
SEIR model to mortality data in an effort to examine possible trajectories of
COVID-19 infections at the state level. In particular, they conduct a review
of COVID-19 model on a state-wise basis with particular emphasis on charting
out potential scenarios in terms of the number of fatalities in the presence
of various non-pharmaceutical interventions. They forecast the best and worst
case outcomes in terms of these numbers and make recommendations to ensure the
safety of the US population in case of epidemic resurgences in many states.
The authors of (López and Rodó, 2020) use an SEIR model to conduct a study of
the recurrence of the COVID-19 pandemic via different post-confinement
scenarios. Their work highlights the importance of non-pharmaceutical
interventions due to the re-emergence risk from the time-decay of acquired
immunity and lack of effective pharmaceutical interventions.

There is however significant work that has gone into building more expressive,
realistic, informed models some of which are SIDARTHE (Giordano et al., 2020),
SEI3HR in India (Senapati et al., 2020), SUEIHCDR in Brazil (Kennedy et al.,
2020), SEI3RSD (Wolff, 2020), SEI3Q3RD (Grimm et al., 2021), SEI3R2S (Winters,
2020), and others (Ndaïrou et al., 2020). SEIR models for the spread of
disease (Buckman et al., 2020; COVID et al., 2020; Buckman et al., 2020) are
prevalent in the existing literature with multitudinous global, national, and
regional studies conducted through studying the disease transmission dynamics
modeled by its compartmental transitions. We conjecture that the SEIR model
can be improved by reducing the modeling assumptions implicit in its
definition such as the explicit separation of the recovered individuals from
the fatalities out of all those in the final ’removed’ compartment. In the
framework offered by Pyro, it is straightforward to modify the compartmental
model structure which makes it extremely useful as a toolkit for practitioners
interested in an experimental perspective on epidemiological modeling.

As a template, we consider the idea of three-layered infection states
corresponding to increasingly infectious ’spreaders’ motivated by the need to
tie epidemiology with medical physiology through modeling causal dynamic
processes in time (Winters, 2020) for defining an SEI3RD model as shown in 9.

Figure 9. The SEI3RD Model from (Wood et al., 2020) along with our policy
intervention parameter $u$

For researchers, an excellent approach to reproducibly expand upon the
existing literature could be following the approach of (Ndaïrou et al., 2020).
They examine an eight-compartment model to obtain disease parameter estimates
of a COVID-19 variant in Wuhan, China. They then conduct a sensitivity
analysis of their model to examine the variance with respect to each parameter
and compare their results with a numerical simulation to examine its
suitability. Our motivation with the SEI3RD model is, in a similar vein, to
expand upon the line of work leading to the SEIR model but with the difference
of offering an open-source template to introduce new variants of compartmental
models that can differ regionally.

Figure 10. Visualizing the disease progression with time (along x-axis) and
infected population fraction (along the y-axis) under varying levels of the
policy intervention parameter $u$ for SIR. Figure 11. Visualizing the disease
progression with time (along x-axis) and infected population fraction (along
the y-axis) under varying levels of the policy intervention parameter $u$ for
the SEIR model.

###  C.1. Policy Interventions

The use of surgical face masks and face shields by healthcare and non-
healthcare workers alike has been shown to significantly reduce or prevent the
transmission of human coronaviruses and influenza viruses through respiratory
droplets from symptomatic individuals in confined spaces (Leung et al., 2020;
Liang et al., 2020; Chu et al., 2020). This is an example of a type of
intervention that falls into the class of non-pharmaceutical interventions
(NPIs) which are important to model in light of the time it requires to
implement substantial pharmaceutical interventions (PIs) such as vaccines. Our
simulation, therefore, focuses on the modeling of NPIs, or what we term policy
interventions, to allow us to build world-models that are reflective of
disease spread in the absence of effective PIs. While vaccines may be the most
effective long-term solution, there is significant impact of short term NPIs
such as isolation and contact tracing (Grimm et al., 2021; Bertozzi et al.,
2020). For this reason, our work focuses extensively on a framework to explore
intervention strategies without the need for human-supervised search.

When governments deal with diseases, they may take certain measures that
result in limiting the exposure of the population to the disease. These
measures may range from mild as in washing hands to stringent as in enforcing
a complete lockdown (United States Government, [n. d.]; International Monetary
Fund, [n. d.]). While some lines of work focus on how to monitor their impact
(Giudici and Raffinetti, 2020; Vasconcelos et al., 2020), we propose an
algorithm to explore new strategies for control. The goal of epidemiological
modeling, particularly the spread of pandemics is to be able to infer the
actions necessary to limit their spread. Every policy to address this is
designed to intervene on the rate of spread through natural or artificial
means, for a short or long term duration. However, most modeling approaches
treat it as affecting the transmission in a deterministic manner whereas in
reality, the impact of government policies varies with time. For this reason,
we decouple modeling of the policy intervention $u$ from the transmission
probability $\beta$. This allows us to introduce our greedy search algorithm
described in (also see 12).

Result: Sequence of policy interventions

initialize sequence;

while  _$t\leq T$_ do

simulate compartmental transitions;

for  _$u=u_{i}$ , increasing from $0$ to $1$_ do

simulate a trajectory with $R_{0}$, $\rho$;

if  _infections $\leq$ threshold_ then

add $u_{i}$ to sequence;

break;

end if

end for

end while

Algorithm 1 Adaptive Interventions through Greedy Search

Modelling policy interventions for COVID-19 has been the focus of a large body
of work ((Giudici and Raffinetti, 2020; Vasconcelos et al., 2020; Giordano et
al., 2020; Chen et al., 2020)). In our work, we start with the similar
implementation of a policy intervention defined by the parameter $u$ which
modifies the SIR compartmental model as per the equation below:

$\displaystyle\beta_{1}$ $\displaystyle=(1-u)\beta$ $\displaystyle\beta$
$\displaystyle=\beta_{1}$

Recall that $\beta$ was the transmission probability of an individual from the
pool of susceptible individuals to the next compartment which is model-
dependent. Changing the values of this policy intervention parameter manually,
as in $u$ taking on a deterministic set of values, and simulating the
resulting trajectories results in the plot shown in 10 for the SIR model and
11 for the SEIR model. The horizontal dashed line indicates the $10\%$
threshold of infections. This approach allows us to monitor what could have
been the impact of governmental policies such as a lockdown implemented over a
long time period. However, $u$ may not be the same at all points in time.
Thus, we introduce the algorithm C.1 to consider a sequence of policy
interventions that is analogous to periodic policy updates or a variance in
the impact of the same policies. A visual analogy of the algorithm is offered
by illustrating the best and worst case outcomes of selecting different policy
interventions at two time-steps in 12 with a description below.

Figure 12. Adaptively Changing the Policy Intervention

The current fraction of infected population is represented in bright red as a
solid line. At each timestep, we perform a simulation of possible trajectories
for different values of the policy intervention $u$. For example, at a time
$t=50$ days we simulate the possible trajectories (dashed lines). The plots
show the best and worst-case options for the progress of the disease depending
on the magnitude of the policy intervention parameter $u$. This is repeated at
each time step $t$, with the illustration of $t=100$ offered in a separate
plot.

A greedy solution, as demonstrated in the algorithm C.1, is to pick the $u$ in
a short-term optimal manner such that it keeps the infected fraction just
within the threshold. However, this might still result in a sequence of large
interventions for a highly infectious disease (long-term lockdowns). There are
better strategies that enact more stringent measures early on so that the
overall impact of interventions across the time series is not as high.
Alternatively, we might weight the outcomes of infected individuals exceeding
a threshold by the amount of time it would take to breach the threshold and
thereby determine a reasonable intervention to perform. There are different
ways of picking the optimal policy intervention depending on the choice of
utility function and optimisation technique, including neural networks. We
believe that the choice should vary subject to demographic factors,
effectiveness of government policies, durations for which governments can
implement stringent measures and acceptable thresholds of infected
individuals.

###  C.2. Limitations

We study the spread of infectious diseases in terms of the COVID-19 pandemic
in an attempt to unify the different models to highlight the utility of PPLs.
For this reason, while we do study the predictions of our models by evaluating
the parameter estimates against those described in the existing literature
(Verity et al., 2020; Hoseinpour Dehkordi et al., 2020; Li et al., 2020;
Kucharski et al., 2020; Wang et al., 2020a) in some experiments, it is more
interesting for us to focus, in this work, on our technique and its extensions
than the fact that we achieve similar results to those in related literature.
We provide a complete comparison of some region-wise parameter estimates with
real-world data in 1 but emphasize that the novelty is in the ease of
application of this technique across the table, its adaptability to
incorporate modeling assumptions, and flexibility of inference across a
spectrum of compartmental models.

It is challenging to model a disease with convenient mathematical assumptions
implicit in compartmental models, that largely ignore individual-level
differences. While there are some inferences that are plausible to make, it is
clear that we must not jump to immediate conclusions in disease modeling due
to the fact that there is almost always non-trivial uncertainty associated
with parameter estimates in that multiple models might reflect similar initial
trajectories (Atkeson, 2020a). Furthermore, certain types of interventions
might not be as effective as others (Manchein et al., 2020). However, modeling
interventions is one way of translating theory into practical advice for
making policy decisions (Thompson, 2020) and it is promising that SEIR models
perform well over time in multiple demographies with regards to predicting the
spread of the disease (Atkeson, 2021). This strengthens the argument that we
would like to further reduce our modeling assumptions and therefore,
uncertainty in parameter estimates, in order to make more confident policy
recommendations eventually. This motivates the use of the SEI3RD model (among
other variants in Compartmental Model Extensions in Pyro) which is more
powerful than the class of SEIR models since it is a derivative that allows
for more granular transmission dynamics between compartments. However, even
with the SEI3RD model, we need to add in external interventions, consider
jointly conditioning on observed fatalities, and inductive biases on
transition probabilities based on estimated incubation period, infection
severity, stratification by demographic factors, and consider utility


# The detection, extraction and parameter estimation of extreme-mass-ratio
inspirals with deep learning

Qianyun Yun<EMAIL_ADDRESS>Shanghai Astronomical Observatory, Chinese Academy
of Sciences, Shanghai, China, 200030  Hangzhou Institute for Advanced Study,
University of Chinese Academy of Sciences, Hangzhou 310124, China  School of
Physics and Astronomy, Shanghai Jiao Tong University 800 Dongchuan RD.,Minhang
District, Shanghai, 200240, China  Wen-Biao Han  Corresponding author:
<EMAIL_ADDRESS>Hangzhou Institute for Advanced Study, University of Chinese
Academy of Sciences, Hangzhou 310124, China  Shanghai Astronomical
Observatory, Chinese Academy of Sciences, Shanghai, China, 200030  School of
Astronomy and Space Science, University of Chinese Academy of Sciences,
Beijing, China, 100049  Taiji Laboratory for Gravitational Wave Universe
(Beijing/Hangzhou), University of Chinese Academy of Sciences, Beijing 100049,
China  Yi-Yang Guo  Lanzhou Center for Theoretical Physics, Key Laboratory of
Theoretical Physics of Gansu Province,
and Key Laboratory of Quantum Theory and Applications of MoE,
Lanzhou University, Lanzhou, Gansu 730000, China  Institute of Theoretical
Physics & Research Center of Gravitation, Lanzhou University, Lanzhou 730000,
China  He Wang  International Centre for Theoretical Physics Asia-Pacific
(ICTP-AP), University of Chinese Academy of Sciences (UCAS), Beijing, China
Taiji Laboratory for Gravitational Wave Universe (Beijing/Hangzhou),
University of Chinese Academy of Sciences, Beijing 100049, China  Minghui Du
Center for Gravitational Wave Experiment, National Microgravity Laboratory,
Institute of Mechanics, Chinese Academy of Sciences, Beĳing 100190, China

###### Abstract

One of the primary goals of space-borne gravitational wave detectors is to
detect and analyze extreme-mass-ratio inspirals (EMRIs). This endeavor
presents a significant challenge due to the complex and lengthy EMRI signals,
further compounded by their inherently faint nature. In this letter, we
introduce a 2-layer Convolutional Neural Network (CNN) approach to detect EMRI
signals for space-borne detectors, achieving a true positive rate (TPR) of
96.9 % at a 1 % false positive rate (FPR) for signal-to-noise ratio (SNR) from
50 to 100. Especially, the key intrinsic parameters of EMRIs such as the mass
and spin of the supermassive black hole (SMBH) and the initial eccentricity of
the orbit can be inferred directly by employing a VGG network. The mass and
spin of the SMBH can be determined at 99 % and 92 % respectively. This will
greatly reduce the parameter spaces and computing cost for the following
Bayesian parameter estimation. Our model also has a low dependency on the
accuracy of the waveform model. This study underscores the potential of deep
learning methods in EMRI data analysis, enabling the rapid detection of EMRI
signals and efficient parameter estimation.

Introduction. Since the first detection of GW150914 LIGOScientific:2016dsl,
ground-based gravitational wave (GW) detectors such as LIGO Harry:2010zz and
VIRGO VIRGO:2014yos have made remarkable progress. These detectors have
subsequently detected nearly one hundred GW events of the merger of two black
holes with stellar mass before 2022 LIGOScientific:2021djp. These observations
enable researchers to test Einstein’s general theory of relativity
LIGOScientific:2019fpa and explore the universe through an independent
approach LIGOScientific:2021aug.

Ground-based detectors such as LIGO, Virgo, and Kagra KAGRA:2013rdx focus on
gravitational waves within a frequency range of 10 Hz to a few thousand Hertz,
as determined by their sensitivity curves LIGOScientific:2021djp;
Kaiser:2020tlg. While space-borne interferometers, represented by missions
like LISA LISA:2017pwj and Taiji Hu:2017mde are specifically engineered to
detect low-frequency gravitational waves, covering the frequency from 0.1 mHz
to 1 Hz. One primary objective of these space-borne interferometers is to
detect extreme-mass-ratio inspiral systems Babak:2017tow; Ruan:2018tsw;
Han:2018hby.

An EMRI system consists of a stellar-mass black hole orbiting around a much
heavier black hole with a mass ranging from $10^{4}-10^{7}M_{\odot}$. These
supermassive black holes (SMBHs) are typically located at the centers of
galaxies Babak:2017tow. The inspirals of compact objects around the SMBHs emit
GWs at low frequencies. Investigating EMRI systems provides opportunities to
examine theories of gravity Babak:2017tow; Han:2018hby. This exploration
enables precise mapping mapping of the spacetime around SMBHs and the testing
of the no-hair theory Maselli:2021men; LISA:2022kgy. Moreover, it can help us
comprehend the mass distribution of SMBHs and their relationships with host
galaxies, as well as trace the formation and evolution of them LISA:2022yao.

Analyzing EMRI signals is very challenging due to their prolonged duration and
intricate characteristics. These signals can span for several years, demanding
massive computational resources to generate sufficient and high accurate
waveforms for data analysis Cornish:2008zd; Babak:2009ua. Traditional
gravitational wave data analysis methods, like match filtering and Bayesian
parameter estimation, require a minimum of $10^{40}$ EMRI waveform templates
to cover the 14-parameter space Gair:2004iv. Even with the application of the
F–statistic algorithm Wang:2012xh, a substantial amount of templates are still
required. In addition, the traditional methods heavily depend on template
accuracy in data analysis. Nevertheless, the intricate harmonics and
modulations of EMRI signals can create difficulties in modeling an accurate
template due to factors like self-force effects.

Deep learning techniques have demonstrated their remarkable efficacy in GW
data analysis. They excel in both detecting and characterizing GW signals for
both simulated and real situations George:2017pmj; Wang:2019zaj; Jin:2023ahl.
Additionally, deep learning has already been used for the detection of EMRIs
Zhang:2022xuq; Zhao:2022qob. Moreover, deep learning serves as a valuable tool
for estimating galaxy parameters tuccillo2018deep, which can also be employed
for GW parameter estimation.

In this letter, we use a 2-layer CNN network designed for detecting EMRI
signals in the time-frequency domain. We compute the EMRI signals using
several waveform models and simulate the noise using Taiji’s analytical PSD.
Considering the practical application in space detectors, we incorporate
second-generation Time Delay Interferometry (TDI) Amaro-Seoane:2012vvq. Then
we apply the Q-transform to convert the time series data into time-frequency
data before feeding them into the network. We also employ a UNet network to
extract EMRI signals from noise and estimate their parameters using a VGG net.
Though our research primarily concentrates on Taiji data, our method is
versatile for other space-borne GW detectors like LISA.

Methodology. The initial step in training a deep learning Network is the
creation of the training and testing dataset. Our dataset consists of two
elements: the combination of signals and noise ($d=h+n$), labeled as 1, and
noise alone ($n$), labeled as 0. The demonstration of these two types of data
is shown in Fig. 1.

Figure 1: Left: The CQT plot and the time series plot of signal and noise.
Right: The CQT plot and the time series plot of pure noise.

We employ the AAK model Chua:2017ujo within specified parameter ranges (see
Table 1) to generate EMRI signals. These signals last 0.5 years and are
sampled at 10-second intervals. We utilize second-generation TDI techniques
Amaro-Seoane:2012vvq to get the A, E, and T channel data of the Taiji
detector. TDI is a crucial method for reducing laser frequency noise and
ensuring desired sensitivity in space-borne detector missions. We chose the A
and E channels in this letter. Subsequently, we compute SNR values using the
analytical PSD, maintaining them within the range of 50-100. The mathematical
expression for the PSD of the X channel
($A=\frac{1}{\sqrt{2}}(Z-X),E=\frac{1}{\sqrt{6}}(X-2Y+Z)$) is as follows
Ren:2023yec: $\mathrm{PSD}{X_{2}}=64\sin^{2}(\omega L)\sin^{2}(2\omega
L)\left(P_{\mathrm{oms}}+(3+\cos(2\omega L))P_{\mathrm{acc}}\right)$, where
$\omega=2\pi f/c$, and $P_{\mathrm{oms}}(f)$ represents high-frequency noise
from the optical metrology system, while $P_{\mathrm{acc}}(f)$ results from
test mass acceleration. We utilize this PSD to create the detector’s noise.
Half of these noises are combined with the signals, as illustrated in the
bottom-left panel of Fig.1. The data of pure noise is displayed in the bottom-
right panel of Fig.1. The dataset consists of 6000 samples for both $d$ and
$n$. The dataset is subsequently divided into training and testing sets, with
the former for model training and the latter for evaluation.

Table 1: The parameter ranges of AKK and effective-one-body-Teukolsky waveforms for training and testing data. Parameter |  Range (Uniform distribution) |  Parameter |  Range (Uniform distribution)
---|---|---|---
$M/M_{\odot}$ |  ($10^{5},10^{8}$) |  $\gamma$ |  (0, 0.1)
$\mu/M_{\odot}$ |  (10, 100) ; $M/\mu\geq 10^{4}$ |  $\theta_{S}$ |  (0, $\pi$)
$a$ |  $(10^{-3}$, 0.90) |  $\phi_{S}$ |  (0, $2\pi$ )
$e_{0}$ |  (0.005, 0.6) |  $\theta_{K}$ |  (0, $\pi$)
$p_{0}/M$ |  (10, 12) |  $\phi_{K}$ |  (0, $2\pi$)
$\iota$ |  (0, 0.1) |  |

Time-frequency analysis proves to be a highly effective approach in EMRI data
analysis Wang:2012xh. Dominant harmonics in EMRI signals become clear in the
time-frequency domain. Therefore, we chose to employ the Constant Q Transform
(CQT) for data preprocessing before feeding them into the networks. CQT is a
special case of Variable Q transform(VQT). Both of them are related to complex
Morlet wavelet transform. The definition is: $\delta f_{k}=2^{1/n}\cdot\delta
f_{k-1}=\left(2^{1/n}\right)^{k}\cdot\delta f_{\min}$, where $\delta f_{k}$ is
the bandwidth of the k-th filter, $f_{min}$ is the central frequency of the
lowest filter, and n is the number of filters per octave. Define the quality
factor Q for CQT: $Q=\frac{f_{k}}{\delta f_{k}}$. With this factor, we can
define window length: $N[k]=\frac{f_{s}}{\delta f_{k}}=\frac{f_{s}}{f_{k}}Q$.
Hamming window is one of the most commonly used windows, take this for
example: $W[k,n]=\alpha-(1-\alpha)\cos\frac{2\pi
n}{N[k]-1},\quad\alpha=25/46,\quad 0\leqslant n\leqslant N[k]-1$. With above
equations, define Q transform as:
$X[k]=\frac{1}{N[k]}\sum_{n=0}^{N[k]-1}W[k,n]x[n]e^{\frac{-j2\pi Qn}{N[k]}}$.
We use librosa (https://librosa.org/doc/latest/index.html) to do the process,
store the Q-transform results as gray images, and resize them to (255,255)
matrixes as the input data. The top panel of Fig. 1 shows the training data
post-Q-transform, where the input signal has an SNR of 50. In this figure,
EMRI features are not distinctly visible due to noise, making it challenging
to visually distinguish between $d$ and $n$. To solve this problem, we try to
use CNN to distinguish these two kind of data.

For this target, we adopt a basic framework of a CNN network architecture,
which includes two Conv2d layers with filter sizes of 16 and 32. The
convolutional layers have kernel sizes of $3\times 3$ and $3\times 3$,
respectively. Following each convolutional layer, a Batchnorm layer is
introduced, followed by the application of a standard ReLU activation function
as non-linearities layers, and subsequently, a max-pooling layer is applied.
The convolution layers have a stride of 1, while the pooling layers have a
stride of 2. Following the two convolutional blocks, there are two fully
connected layers with sizes of 128 and 64.

Then we use the simulated data to train our detection network. During the
training, the cross-entropy loss is employed as the loss function, defined as:

$L=\sum_{i=1}^{N}\left(-y_{i_{\text{true}}}\cdot\log(y_{i_{\text{pred}}})-(1-y_{i_{\text{true}}})\cdot\log(1-y_{i_{\text{pred}}})\right).$
(1)

The loss quantifies the difference between our model’s predictions and the
true labels. Through backpropagation, the model updates each parameter to
minimize the loss. The training goes through 200 epochs. The Adam optimizer is
employed with a learning rate set at 0.01.

The performance of the network is assessed by testing data and the results are
visually shown by the Receiver operating characteristics (ROC) curve. The ROC
shows our model’s ability to distinguish $d$ and $n$ by plotting the true
positive rate against the false positive rate. The quantified performance is
determined by the Area Under the ROC Curve (AUC). A higher AUC score, closer
to 1, means that the model has a better capability to distinguish the EMRI
signals.

Figure 2: ROC curves for the trained CNN model to detect EMRI signals
generated by the AAK (red) and EOB (blue) waveform model. The SNRs of the
testing data range in [50,100]. The figure inside zooms in on the TPR from 0
to 0.1.

In Fig. 2, the red curve shows the performance of our model in detecting EMRI
signals simulated by the AAK templates, achieving an AUC of 99 %. The
detection achieves a 96.9% TPR while 1% FPR with SNRs from 50 to 100.

To evaluate the robustness of our network, we also utilize an alternative
waveform model to generate test data. This is crucial since the theoretical
EMRI signals may differ from the true EMRI signals due to the accuracy and the
real scenarios. Traditional GW detection relies heavily on accurate waveform
templates, which may lead to missed detections when true signals differ from
the theoretical ones. We believe that deep learning can reduce this template
dependency. To confirm this, the effective-one-body-Teukolsky (ET)
Shen:2023pje waveform model is used to generate new data to test our detection
network trained with AAK waveforms.

In Fig. 2, the blue ROC curve shows the capability to detect EMRI signals
generated by ET model. The CNN network also performs effectively on this new
test dataset, with an AUC of 99 %. The TPR can be still 90% at 1% FPR. This
result shows the robustness of our network in the EMRI detention. It indicates
that even if the theoretical waveforms are not accurate enough, our network
retains the capability to identify the EMRI signals.

After successfully recognizing EMRI signals, we employ the UNet network to
extract them ronneberger2015u. The extraction process utilizes signal-with-
noise as input data and pure signals as the target
(https://github.com/milesial/Pytorch-UNet/tree/master/unet). This Unet
network’s architecture is illustrated in Fig. 3, which depicts the process of
extracting the final time-frequency signal from the original image.

UNet comprises an encoder and a decoder. The encoder is composed of four
“Down” blocks, reducing the input image’s resolution and extracting the
crucial features. The decoder, composed of four “Up” blocks, performs
upsampling to restore features to the input size, and then generates the final
output. “Skip Connections” is the key process, which can merge features of the
encoder and decoder at the same hierarchical level (four gray parallel lines
with arrows in the top panel of Fig. 3). Details of the blocks are illustrated
at the bottom panel of Fig. 3.

During the training, we utilize Binary Cross-Entropy with Logits as the loss
function:

$\text{L}=\text{max}(y_{\text{pred}},0)-y_{\text{pred}}\cdot
y_{\text{true}}+\log(1+\exp(-\text{abs}(y_{\text{pred}}))).$ (2)

This loss function is commonly utilized in Unet-based segmentation tasks to
evaluate the difference between predicted outputs and actual targets. We train
50 epochs to improve the quality of the output, and Fig. 3 shows how we
extract a time-frequency EMRI signal from a noisy one using the Unet network.

After the extraction, we use a VGG simonyan2014very network for parameter
estimation, inspired by the approach in Ref. tuccillo2018deep, where deep
learning was applied for estimating galaxy parameters. The architecture of our
model is depicted in Fig. 4. In this model, input images go through three
blocks. The initial block consists of two 2D convolution layers with a 4x4
filter, followed by a 2x2 max pooling layer and a dropout layer. The
subsequent two blocks have the same architecture, each comprising 4x4
convolution and max pooling layers.

We then begin to train the VGG network for parameter estimation. Firstly, we
create a new dataset with signals and noise, dividing it into training and
testing sets. Following this, we again use the previous UNet model to extract
signals from the noise and input these extracted signals to the VGG network.
The true parameters of the signals are the targets. The network undergoes 50
epochs of training. We utilize Mean Squared Error (MSE) as the loss function,
defined as:

$\mathrm{L}=\frac{1}{N}\sum_{i=1}^{N}(y_{i_{\text{pred}}}-y_{i_{\text{true}}})^{2}.$
(3)

After training, we input testing data to extract signals and obtain their
parameters. Fig. 5 illustrates the relationship between our predicted values
and the true values.

In the top panel of Fig. 5, each point represents the true and predicted mass
of the SMBH in an EMRI. The accuracy of our predictions is calculated by the
coefficient of determination $R^{2}$, defined as
$R^{2}=1-\frac{\sum_{i}^{n}\left(y_{i}-f_{i}\right)^{2}}{\sum_{i}^{n}\left(y_{i}-\bar{y}\right)^{2}}$.
Here, $f_{i}$ denotes the predicted value, $y_{i}$ is the true value, and
$\bar{y}$ is the mean of the true value. Remarkably, the predicted mass
achieves an exceptionally good $R^{2}$ value of 0.989. In the middle panel of
Fig. 5, we present predictions for the spins of SMBHs, yielding an $R^{2}$ of
0.92. Lastly, the bottom panel of Fig. 5 showcases our predictions for the
initial eccentricity of the orbits, accompanied by an $R^{2}$ of 0.95.

Figure 3: Top: Structural details of our UNet network in this letter. Bottom:
The details of each module for our UNet network in this letter.  Figure 4: The
scheme of Architecture in the parameter estimation work.

Figure 5: The parameter estimation results from the VGG net of the mass
values, spin values, and initial eccentricity values of the center SMBHs.

Conclusion and discussion. Detection and parameter estimation of EMRIs are
still a challenge for space-borne projects. Due to the weakness, complexity
(at least 14 parameters), and long duration of the EMRI signals, the
traditional matched-filtering and Bayesian inference may have difficulty in
detecting EMRIs and obtaining accurate source parameters. The PE based on the
Bayesian method requires very accurate waveform templates and huge amounts of
computing resources. Fortunately, the booming machine learning technologies
offer a new chance to overcome this problem. In this Letter, our CNN-based
model exhibits excellent performance for the detection of EMRIs and achieves a
96.9% TPR at a 1% FPR with SNRs from 50 to 100. Especially, our model also
shows good robustness, if the training and testing data come from different
waveform models, the TPR can be still 91% at 1% FPR.

Furthermore, we also employ UNet to extract EMRI signals from noise and VGG to
estimate their main intrinsic parameters, such as the mass and spin of the
SMBH. The statistical error is approximately 1 % for mass, less than 10 % for
spin, and less than 5 % for orbital eccentricity. These intrinsic parameters
directly determine the physical characteristics of the EMRI sources, and will
greatly reduce the parameter space and the computing cost for further Bayesian
PE with Markov Chain Monte Carlo.

In this letter, we demonstrate the great potential of deep learning in the
area of detection, extraction, and parameter estimation of EMRIs. However, in
the present work, the SNRs of the signals are all larger than 50. Therefore,
we will analyze the weaker EMRI signals in the upcoming work, with SNR as low
as 20. This kind of signals are more popular for space-borne interferometers
and will improve the detection rate.

###### Acknowledgements.

We would like to express our sincere appreciation to Quanfeng Xu (SHAO) and
Yuewei Zhang (SJTU) for their invaluable assistance and guidance about the
deep learning part throughout this research project. Thanks to the Lanzhou
University Summer School on Gravitational Waves where we talk about the
advantages of using time-frequency domain data to analysis EMRI signals with
Xiaobo Zou (LZU) and discuss the CQT with Professor Soumya D. Mohanty (The
University of Texas Rio Grande Valley), lecturer Qunyin Xie (LZU), and her
student Shuzhu Jin. This work was supported by the National Key R&D Program of
China (Grant Nos. 2021YFC2203002), the National Natural Science Foundation of
China (Grant Nos. 12173071). Wen-Biao Han was supported by the CAS Project for
Young Scientists in Basic Research (Grant No. YSBR-006). This work made use of
the High Performance Computing Resource in the Core Facility for Advanced
Research Computing at Shanghai Astronomical Observatory.