# A simple analytical model of magnetic jets

Andrzej A. Zdziarski1, Łukasz Stawarz2, Marek Sikora1 and Krzysztof Nalewajko1
1Nicolaus Copernicus Astronomical Center, Polish Academy of Sciences, Bartycka
18, PL-00-716 Warszawa, Poland
2Astronomical Observatory, Jagiellonian University, Orla 171, PL-30-244
Kraków, Poland

###### Abstract

We propose a simple analytical jet model of magnetic jets, utilizing
conservation laws and some results of published GRMHD jet simulations, and
consider radially-averaged profiles of the physical quantities. We take into
account conversion of the magnetic energy flux to bulk acceleration in jets
formed around rotating black holes assuming the mass continuity equation and
constant jet power, leading to the Bernoulli equation. For assumed profiles of
the bulk Lorentz factor and the radius, this gives the profile of toroidal
magnetic field component along the jet. We then consider the case where the
poloidal field component is connected to a rotating black hole surrounded by
an accretion disc. The formalism recovers the standard formula for the power
extracted from a rotating black hole. We find that the poloidal field strength
dominates over the toroidal one in the comoving frame up to large distances,
which means that jets should be more stable to current-driven kink modes. The
resulting magnetic field profiles can then be used to calculate the jet
synchrotron emission.

###### keywords:

acceleration of particles – black hole physics – radiation mechanisms: non-
thermal – galaxies: active – galaxies: jets

††pagerange: A simple analytical model of magnetic jets–References††pubyear:
2022

##  1 Introduction

Accreting black holes (BHs) of all sizes produce relativistic jets. The
physical mechanisms producing the jets have been subject of intense research.
The most promising mechanisms for the jet production appear to be those linked
to magnetic fields (Blandford & Znajek, 1977; Blandford & Payne, 1982). These
models also predict the jet power to scale with the BH mass (e.g., Heinz &
Sunyaev 2003; Tchekhovskoy et al. 2009, 2010, 2011).

The observed power of extragalactic jets is often found to be very large,
$\gtrsim\dot{M}_{\rm accr}c^{2}$, where $\dot{M}_{\rm accr}$ is the accretion
rate (e.g., Ghisellini et al. 2014; Zamaninasab et al. 2014; Pjanka et al.
2017). The only model capable to account for those powers appears to be that
with electromagnetic extraction of the spin power (Blandford & Znajek, 1977)
from a BH surrounded by a magnetically arrested disc (MAD; Bisnovatyi-Kogan &
Ruzmaikin 1974; Narayan et al. 2003), as demonstrated using GRMHD numerical
simulations by Tchekhovskoy et al. (2011) and McKinney et al. (2012).

The structure of relativistic jets undergoing simultaneous acceleration and
collimation has been investigated by complex semi-analytical relativistic MHD
models, see Li et al. (1992), Vlahakis & Königl (2003), Beskin & Nokhrina
(2006), Lyubarsky (2009). Numerical simulations with relativistic MHD have
been performed by Komissarov et al. (2007, 2009). Then, numerical GRMHD
simulations have been done in a number of papers, see Tchekhovskoy (2015) and
Davis & Tchekhovskoy (2020) for reviews. Here, we provide a simple analytical
formalism reproducing main aspects of those models. We assume the continuity
and energy conservation equations, and then we use the relationship between
the toroidal and poloidal components of the magnetic field for the case of
spin-power extraction based on the numerical simulations of Tchekhovskoy et
al. (2009). However, we do not require the presence of a MAD.

Our model predicts the profiles of the radially-averaged toroidal and poloidal
field components along the jet. We find the poloidal component to dominate
over the toroidal one in the comoving frame up to large distances, which
follows from the conversion of the toroidal field to the bulk acceleration
taking place in the acceleration and collimation zone (ACZ). These profiles
can then be used to calculate self-absorbed synchrotron emission from extended
jets to model their flat radio spectra. Such models are generally based on the
works of Blandford & Königl (1979) and Königl (1981), with a later formulation
(Zdziarski et al., 2019) of the synchrotron emission and self-absorption for
arbitrary profiles of the bulk Lorentz factor, $\Gamma(z)$, the jet radius,
$r(z)$, and the magnetic field strength, where $z$ is the distance from the BH
centre.

##  2 The jet structure

We assume the jet is launched from close vicinity of the BH, and that it first
goes through the ACZ. Details of these processes are poorly understood. The
jet radial profile depends sensitively on the profile of the external
pressure, which is usually uncertain. Still, those profiles can be derived
(e.g., Vlahakis & Königl, 2003), but the resulting formalism is rather
complex. Here, instead, we assume the profiles of both $\Gamma(z)$ and the jet
radius, $r(z)$. In our model, $\Gamma$ increases until it reaches its terminal
bulk Lorentz factor, $\Gamma_{\rm max}$, at which point the jet also becomes
conical. We stress, however, that our general results below do not depend on
these assumptions, and the profiles given below can be modified without
affecting our formalism. Only determination of the specific profiles of the
magnetic field strengths require the explicit specification of $\Gamma(z)$ and
$r(z)$ in the ACZ, see Fig. 1 below. We first assume,

$\Gamma(z)=\begin{cases}a_{\Gamma}\left(z/r_{\rm H}\right)^{1/{q_{1}}},&z_{\rm
m}\equiv a_{\Gamma}^{-q_{1}}r_{\rm H}<z\leq z_{\rm t}\equiv(\Gamma_{\rm
max}/a_{\Gamma})^{q_{1}}r_{\rm H};\cr\Gamma_{\rm max},&z>z_{\rm
t},\cr\end{cases}$ (1)

where $r_{\mathrm{H}}=[1+(1-a_{*}^{2})^{1/2}]r_{\rm g}$ is the BH event
horizon, $r_{\rm g}=GM/c^{2}$ is the gravitational radius, $a_{*}$ is the spin
parameter, and $q_{1}\sim 2$ and $a_{\Gamma}\leq 1$ are constants. The jet
starts to be accelerated at some distance, $z_{\rm m}$, above the stagnation
surface. Its radius is assumed as

$\displaystyle r(z)$ $\displaystyle=\begin{cases}a_{r}\left(z/r_{\rm
H}\right)^{1/q_{2}}r_{\rm H},&z_{\rm m}\leq z\leq z_{\rm
t};\cr\Theta_{0}\left[z+(q_{2}-1)z_{\rm t}\right],&z>z_{\rm t},\cr\end{cases}$
(2) $\displaystyle\Theta_{0}$
$\displaystyle=\frac{a_{r}}{q_{2}}\left(\frac{a_{\Gamma}}{\Gamma_{\rm
max}}\right)^{q_{1}(1-1/q_{2})},$ (3)

where $a_{r}\sim 1$, $1<q_{2}<2$, and $\Theta_{0}$ is the asymptotic opening
half-angle (assuming $\Theta_{0}\ll 1$; thus
$\Theta_{0}\approx\tan\Theta_{0}$). The above formulae follow from the
assumption that the transitions in both $\Gamma$ and $r$ take place at the
same distance, $z_{\rm t}$, and from the continuity in both $r(z)$ and ${\rm
d}r/{\rm d}z$ at $z_{\rm t}$. The latter, in particular, assures the
continuity in the adiabatic loss rate. This differs from, e.g., the model of
Ghisellini & Tavecchio (2009), in which the jet was assumed to have the radius
$=\Theta_{0}z$ in the conical part. We further require the product
$\Theta_{0}\Gamma_{\rm max}$ to be given by a constant, as expected
theoretically (Tchekhovskoy et al., 2009). Then, Jorstad et al. (2005),
Pushkarev et al. (2009) and Clausen-Brown et al. (2013) found the average
values observed in extragalactic jets of $\Theta_{0}\Gamma_{\rm max}\approx
0.1$–0.2. The requirement of a constant $\Theta_{0}\Gamma_{\rm max}$ leads
then to

$q_{1}=\frac{q_{2}}{q_{2}-1},\quad\Theta_{0}\Gamma_{\rm
max}=\frac{a_{r}a_{\Gamma}}{q_{2}},$ (4)

which we will assume hereafter. The former relationship is also given by
Komissarov et al. (2009). On the other hand, their approximation for the
radius has an additional dependence of the normalization on $q_{2}$ (denoted
as $b$ in that paper).

We assume the conservation of the rest mass-flow rate,

$\dot{M}_{\rm j}=2\pi r^{2}\rho^{\prime}c\beta\Gamma,$ (5)

where $\rho^{\prime}$ is the mass density, and hereafter all primed quantities
are in the jet frame. Thus, we assume the jet mass loading from the moment of
its formation, which remains uncertain, see, e.g., O’ Riordan et al. (2018).
We also assume the conservation of the jet power, $P_{\rm j}$; thus, we assume
that the radiated power is $\ll\\!P_{\rm j}$. The jet power equals the flux of
the enthalpy, with the contributions from the ion rest energy, the internal
enthalpy of both electrons and ions (consisting of their kinetic energy
density, $u^{\prime}$, and pressure, $p^{\prime}$), and the Poynting flux
integrated over the jet cross section,

$P_{\rm
j}=2\pi\rho^{\prime}r^{2}c^{3}\beta\Gamma^{2}+2\pi(u^{\prime}+p^{\prime})r^{2}c\beta\Gamma^{2}+\frac{{B^{\prime}_{\phi}}^{2}}{2}r^{2}c\beta\Gamma^{2},$
(6)

where the square of the toroidal field strength in the comoving frame,
${B^{\prime}_{\phi}}^{2}$, is averaged over the jet cross section. Here, we
included the power associated with the particle rest energy in the comoving
frame. If most of the rest energy is in ions, it is supplied to the jet by the
accretion flow from large distances. Then, the expression of the usable power
would have the factor $\Gamma(\Gamma-1)$ instead of $\Gamma^{2}$ in the first
term above111We note that Lucchini et al. (2021) and Kantzas et al. (2022)
used a formula for the jet power in their bljet model that is
$\propto\Gamma^{1}$, which appears to be incorrect.. If $\Gamma\gg 1$, this
difference is negligible.

The Poynting flux in an ideal-MHD jet (i.e., with the electric vector
perpendicular to the magnetic one, ${\mathbf{E}}\cdot{\mathbf{B}}=0$) is
$(1/4\pi){\mathbf{B}}\times({\mathbf{v}}\times{\mathbf{B}})$. In a given point
of the jet, this yields the energy-flow rate component parallel to the flow as
$(1/4\pi)(\beta cB_{\phi}^{2}-v_{\phi}B_{\rm p}B_{\phi})$ where $B_{\phi}$ and
$B_{\rm p}$ are the toroidal and poloidal field components in the local
observer frame (hereafter the BH frame), respectively, and $v_{\phi}$ is the
particle rotation velocity taken at an average jet radius. Hereafter, we
neglect the effect of jet rotation, given it is slow beyond the Alfvénic
region, and the radiation contribution from that region is at most small, and
thus neglect the (negative) component of the jet power due to the poloidal
field. Given that

$B_{\phi}=\Gamma B^{\prime}_{\phi},\quad B_{\rm p}=B^{\prime}_{\rm p},$ (7)

we obtain equation (6). Then, the jet magnetisation parameter, defined as the
ratio of the Poynting flux to the rest-energy flux, is

$\sigma\equiv\frac{{B^{\prime}_{\phi}}^{2}}{4\pi\rho^{\prime}c^{2}}.$ (8)

By dividing the expression for the jet power by that for the rest-mass flow,
we obtain the Bernoulli equation, expressing conversion of the magnetic energy
into the bulk motion. Hereafter, we assume the contribution from the internal
enthalpy of the particles, i.e., the second term in the right-hand side of
equation (6), to be small. This is often found to be satisfied in modelling
individual jet sources, see, e.g., Ghisellini et al. (2014), Zdziarski et al.
(2022b). This would also justify the neglect of the internal energy density
term in the denominator of equation (8). Thus, the expression for the
magnetisation parameter following from the Bernoulli equation becomes

$\sigma(z)=\frac{\Gamma_{\rm max}\left(1+\sigma_{\rm
min}\right)}{\Gamma(z)}-1,$ (9)

where a decreasing $\sigma$ corresponds to a conversion of the Poynting flux
into the jet acceleration, and $\sigma_{\rm min}$ is achieved at
$\Gamma=\Gamma_{\rm max}$. The jet power is then,

$P_{\rm j}=\dot{M}_{\rm j}c^{2}\left(1+\sigma_{\rm min}\right)\Gamma_{\rm
max}.$ (10)

Note that this includes the $\dot{M}_{\rm j}c^{2}$ term supplied by the
accretion flow. We can then use equations (5), (8) and (9) to calculate the
(square-averaged) toroidal magnetic field strength along the jet,

${B^{\prime}_{\phi}}^{2}(z)=\frac{2\dot{M}_{\rm
j}c\sigma(z)}{r(z)^{2}\beta(z)\Gamma(z)}.$ (11)

Note that since the jet is accelerated at the expense of the enthalpy flux of
the toroidal component of the magnetic field, ${B^{\prime}_{\phi}}$ decreases
in the ACZ faster than $\propto r(z)^{-1}$, by $(\sigma/\Gamma)^{1/2}$. The
formulation of equations (5–11) is quite general. The value of
$B^{\prime}_{\phi}$ follows uniquely from the assumed profiles of $\Gamma(z)$
and $r(z)$ and the assumption that the acceleration is at the expense of the
magnetic energy, and it is independent of the details of this process. An
additional assumption is that the internal energy is small compared to the
rest energy, but in most actual models it is at most comparable.

The minimum value of the jet magnetisation parameter, $\sigma_{\rm min}$, is
achieved in the conical part of the jet. We consider here two assumptions
about this parameter,

$\sigma_{\rm min}=\begin{cases}(\Theta_{0}\Gamma_{\rm
max}/s)^{2}=\left[a_{r}a_{\Gamma}/(sq_{2})\right]^{2},&{\rm case\ I};\cr{\rm
free},&{\rm case\ II},\cr\end{cases}\ $ (12)

Case I assumes the validity of the approximate relation between the
magnetisation, Lorentz factor and opening angle obtained by Tchekhovskoy et
al. (2009) and Komissarov et al. (2009), where $s$ is constrained by causality
to $\lesssim 1$. In this case, we can take $\Theta_{0}\Gamma_{\rm max}$ to be
a free parameter instead of $\sigma_{\rm min}$, which (for a given $s$)
implies the value of $\sigma_{\rm min}$ (and $a_{r}a_{\Gamma}$). However, case
I is satisfied only if the jet acceleration and collimation are dominated by
ideal MHD processes. While this is likely for $\sigma\gtrsim 1$, conversion of
the Poynting flux to the kinetic energy leading to $\sigma\ll 1$ can be driven
by dissipative processes (see, e.g., in Zdziarski et al. 2015 and Pjanka et
al. 2017). Therefore, we also consider the case of a free $\sigma_{\rm min}$.

In order to find the evolution of the jet poloidal magnetic field,
$B^{\prime}_{\rm p}$, along the outflow, we note the conservation of the jet
magnetic flux, $\Phi_{B}\equiv 2\pi\int_{0}^{r(z)}{\rm
d}r^{\prime}B^{\prime}_{\rm p}(r^{\prime})$, where $r^{\prime}$ denotes the
running radius of a jet cross section, $r^{\prime}\leq r(z)$. Then, the
relationship between $B^{\prime}_{\phi}$ and $B^{\prime}_{\rm p}$ depends on
the angular velocity of the field lines, $\Omega_{\rm f}$. One possibility is
the jet formation by magnetic fields attached to the accretion disc (Blandford
& Payne, 1982), in which case $\Omega_{\rm f}$ is related to the disc
rotation. See, e.g., Sądowski & Sikora (2010) for an analysis of the required
inclination of the field lines. Here, we consider magnetic fields attached to
a rotating BH (a split monopole magnetosphere), resulting in energy extraction
from the BH rotation (Blandford & Znajek, 1977). In this case, $\Omega_{\rm
f}=ca_{*}\ell/2r_{\rm H}$, where $\ell\sim 1/2$ is the ratio of $\Omega_{\rm
f}$ to the BH angular frequency, $\Omega_{\rm H}$. Then,

$B^{\prime}_{\rm p}(r^{\prime})=\frac{c\beta\Gamma}{\Omega_{\rm
f}r^{\prime}}B^{\prime}_{\phi}(r^{\prime}),$ (13)

see, e.g., equation (22) of Tchekhovskoy et al. (2009), where the ideal MHD is
assumed. In this case, conversion of the Poynting flux to the bulk
acceleration causes differential redistribution of poloidal field lines. As a
result, both poloidal and toroidal components of the magnetic field are not
uniform along the jet radius. Their radial dependencies based on the numerical
simulations of Tchekhovskoy et al. (2009) can be approximated as, cf.
equations (17–20) of Zdziarski et al. (2015),

$\frac{B^{\prime}_{\rm p}(r^{\prime})}{B^{\prime}_{\rm
p}(z)}\approx\frac{\sigma}{1+\sigma}\left[\frac{r^{\prime}}{r(z)}\right]^{\frac{-2}{1+\sigma}}\\!\\!\\!,\quad\frac{B^{\prime}_{\phi}(r^{\prime})}{B^{\prime}_{\phi}(z)}\approx\left(\frac{2\sigma}{1+\sigma}\right)^{\frac{1}{2}}\left[\frac{r^{\prime}}{r(z)}\right]^{\frac{\sigma-1}{1+\sigma}},$
(14)

where $B^{\prime}_{\rm p}(z)=\langle{B^{\prime}_{\rm
p}(r^{\prime})}\rangle\equiv\Phi_{B}/(\pi r^{2})$, and
${B^{\prime}_{\phi}}(z)=\langle[B^{\prime}_{\phi}(r^{\prime})]^{2}\rangle^{1/2}$
(as noted above). This implies, using equations (9–11) and (13–14),

$B^{\prime}_{\rm p}(z)\approx B^{\prime}_{\phi}(z)\frac{2^{\frac{3}{2}}r_{\rm
H}\beta\Gamma}{a_{*}\ell
r}\left(\frac{1+\sigma}{\sigma}\right)^{\frac{1}{2}}=\frac{4r_{\rm
H}\left(\beta P_{\rm j}/c\right)^{\frac{1}{2}}}{a_{*}\ell r^{2}}$ (15)

In order to calculate the total field strength in the comoving frame
accurately we would need to add in quadrature the toroidal and poloidal
contributions of equation (14) and then average over the cross section. Here,
we simply approximate it as ${B^{\prime}}^{2}\approx{B^{\prime}_{\rm
p}}^{2}+{B^{\prime}_{\phi}}^{2}$.

Assuming the validity of equations (1–4) in the conical part, we have

$\frac{B^{\prime}_{\rm p}(z)}{B^{\prime}_{\phi}(z)}\approx\frac{2^{3/2}r_{\rm
H}\Gamma_{\rm max}^{2}q_{2}}{a_{*}\ell a_{r}a_{\Gamma}\sigma_{\rm
min}^{1/2}z},\quad z_{\rm eq}\approx\frac{2^{3/2}r_{\rm H}\Gamma_{\rm
max}^{2}q_{2}}{a_{*}\ell a_{r}a_{\Gamma}\sigma_{\rm min}^{1/2}},$ (16)

where we assumed $\sigma_{\rm min}<1$. The two components become equal at
$z_{\rm eq}$. We see that $z_{\rm eq}\gg r_{\rm H}$ generally. In the case of
$q_{2}=2$, $z_{\rm eq}/z_{\rm t}\approx 2^{5/2}a_{\Gamma}/(a_{*}\ell
a_{r}\sigma_{\rm min}^{1/2})$. This can be either greater or less than unity
depending on the values of $\sigma_{\rm min}$, $a_{\Gamma}/a_{r}$ and $a_{*}$.
We note that Vlahakis & Königl (2003) and Komissarov et al. (2009) found
$z_{\rm eq}/z_{\rm t}<1$. Apart from the dependence on the constants, this
difference may also be related to our definition of $z_{\rm t}$, which is
specific to the assumption about $\Gamma(z)$ of equation (1).

Figure 1: Examples of the dependence of the magnetic field components on the
distance for $P_{\rm j}=0.1L_{\rm E}$, $\ell=1/2$, $a_{*}=a_{r}=1$, ,
$a_{\Gamma}=0.4$, $q_{2}=1.7$ (implying $\Gamma_{\rm max}\Theta_{0}\approx
0.24$, $z_{\rm m}\approx 9.3r_{\rm g}$), $q_{1}=q_{2}/(q_{2}-1)$ and (a)
$M=10{\rm M}_{\sun}$, $\Gamma_{\rm max}=4$, $\sigma_{\rm min}=10^{-2}$,
$\Theta_{0}\approx 3.4\degr$, (b) $M=10^{8}{\rm M}_{\sun}$, $\Gamma_{\rm
max}=10$, $\sigma_{\rm min}=(\Gamma_{\rm max}\Theta_{0})^{2}$,
$\Theta_{0}\approx 1.4\degr$, and (c) $M=10^{9}{\rm M}_{\sun}$, $\Gamma_{\rm
max}=30$, $\sigma_{\rm min}=0.1$, $\Theta_{0}\approx 0.45\degr$. The poloidal,
comoving-frame toroidal and total magnetic field strengths are shown by the
blue dotted, red dashed and black solid curve, respectively, and the BH frame
toroidal field is shown by the magenta dot-dashed curve. All components of $B$
scale $\propto(P_{\rm j}/M)^{1/2}M^{-1/2}$. The jet becomes conical above
$z_{\rm t}=\Gamma_{\rm max}^{q_{1}}r_{\rm H}$, where we see kinks in the
dependencies for the toroidal field, which are artefacts of the assumed
$\Gamma(z)$ profile. The solid cyan curves show the growing amplitudes of
linear comoving-frame kink perturbation calculated from equation (24).

Alternatively, we could have used the fact that the average poloidal and
toroidal magnetic field components are expected to be close to equipartition
(in the BH frame) at the Alfvén surface, located at some distance $z_{\rm A}$
from the centre, meaning $B_{\rm p}(z_{\rm A})=B_{\phi}(z_{\rm A})$. Then, we
assume that at the Alfvén surface the jet is strongly magnetised,
$\sigma(z_{\rm A})\gg 1$. Under this condition, the Alfvén surface approaches
the light cylinder (Appl & Camenzind, 1993; Lyubarsky, 2010). This approach
yields the scaling similar to, but more complex than, that of equation (15),
and somewhat lower values of $B^{\prime}_{\rm p}$.

The conserved poloidal magnetic flux, $\Phi_{B}=\pi B^{\prime}_{\rm p}r^{2}$,
is then [where we set $\beta=1$ consistent with our assumption of the
negligible rest energy flow contribution to the jet power requiring $\Gamma\gg
1$, equation (6)],

$\Phi_{B}=\frac{4\pi r_{\rm H}\left[(1+\sigma_{\rm min})\Gamma_{\rm
max}\dot{M}_{\rm j}c\right]^{1/2}}{a_{*}\ell}=\frac{4\pi r_{\rm H}\left(P_{\rm
j}/c\right)^{1/2}}{a_{*}\ell}.$ (17)

The power of spin extraction from the BH is $\propto\Omega_{\rm f}(\Omega_{\rm
H}-\Omega_{\rm f})\Phi_{B}^{2}\propto\ell(1-\ell)\Omega_{\rm
H}^{2}\Phi_{B}^{2}$ (Blandford & Znajek, 1977), which is maximized at
$\ell=1/2$, which is usually assumed. Here, we give it following equation
(3.6) of Tchekhovskoy (2015), but expressing it for an arbitrary value of
$\ell$,

$P_{\rm BZ}\approx\frac{\ell(1-\ell)c}{24\pi^{2}r_{\rm
H}^{2}}\Phi_{B}^{2}a_{*}^{2}=\frac{2\ell(1-\ell)}{3\ell^{2}}P_{\rm j},$ (18)

where the first equality is accurate for $a_{*}\lesssim 0.95$ (Tchekhovskoy et
al., 2010), while for $a_{*}\gtrsim 0.95$, the multiplicative correction
factor is given by equation (3.8) of Tchekhovskoy (2015). In the second
equality, we used equation (17). Noting that $P_{\rm j}=P_{\rm
BZ}+\dot{M}_{\rm j}c^{2}$ and for relativistic jets $P_{\rm BZ}\gg\dot{M}_{\rm
j}c^{2}$, we find that equation (18) is satisfied for $\ell=\sqrt{2/5}\approx
0.6$, which is close to 1/2. This shows the overall self-consistency of our
determination of the poloidal flux. The value $\neq$1/2 appears to be due to
both the numerical coefficient in equation 17 being approximate, and the
numerical coefficient in equation (18) depending on the geometry of the
magnetic field (Tchekhovskoy, 2015).

In order to relate $P_{\rm j}$ and $\Phi_{B}$ to the accretion rate,
$\dot{M}_{\rm accr}$, we use the standard definition of a dimensionless
magnetic flux threading the BH on one hemisphere,

$\phi_{\rm BH}\equiv\Phi_{B}/(\dot{M}_{\rm accr}r_{\rm g}^{2}c)^{1/2},$ (19)

(e.g., Tchekhovskoy et al. 2011; McKinney et al. 2012; Davis & Tchekhovskoy
2020), where the magnetic flux conservation implies that $\Phi_{B}=\Phi_{\rm
BH}$. GRMHD simulations yield the maximum value of $\phi_{\rm max}\approx
70(1-0.38a_{*}r_{\rm g}/r_{\rm H})h_{0.3}^{1/2}$, where $h_{0.3}$ is defined
by the half-thickness of the disc being $h_{\rm disc}=r_{\rm disc}0.3h_{\rm
0.3}$ (e.g., Davis & Tchekhovskoy 2020). The maximum $\phi_{\rm BH}$ is
achieved for a spinning BH surrounded by a MAD. With equation (18) at
$\ell=1/2$, we have

$P_{\rm j}\approx 1.7\left(\frac{\phi_{\rm BH}r_{\rm g}}{40r_{\rm
H}}\right)^{2}a_{*}^{2}\dot{M}_{\rm accr}c^{2},$ (20)

where we have normalized it to the value of $\phi_{\rm max}$ at $a_{*}=1$. The
numerical coefficient is slightly larger compared to, e.g., equation (13) in
Davis & Tchekhovskoy (2020), where it is 1.3 [which difference appears to
result from the numerical coefficients in equations (15) and (18) being
approximate]. We stress, however, that all of the formulae in our paper do not
require the presence of a MAD. In our formulation, the jet power, $P_{\rm j}$,
is a free parameter, and thus the dimensionless magnetic flux can assume any
value $\leq\phi_{\rm max}(\dot{M}_{\rm accr}r_{\rm g}^{2}c)^{1/2}$. Then,
$\dot{M}_{\rm accr}$ can be obtained if we know the accretion efficiency,
$\epsilon_{\rm accr}$, and the accretion luminosity, $L_{\rm accr}$, by
$L_{\rm acc}\equiv\epsilon_{\rm accr}\dot{M}c^{2}$. We can also express
$L_{\rm accr}$ in terms of the Eddington luminosity, $L_{\rm E}=4\pi GMm_{\rm
p}c/\sigma_{\rm T}$, which is given here for pure hydrogen, and where $m_{\rm
p}$ is the proton mass and $\sigma_{\rm T}$ is the Thomson cross section. We
can also define the jet production efficiency, $\epsilon_{\rm j}$, by $P_{\rm
j}\equiv\epsilon_{\rm j}\dot{M}_{\rm accr}c^{2}$.

Fig. 1 shows some examples of the dependencies of the magnetic field strengths
on the distance along the jet assuming the validity of equations (1–3). We
assume $a_{*}=a_{r}=1$, $a_{\Gamma}=0.4$, $q_{2}=1.7$,
$q_{1}=q_{2}/(q_{2}-1)$. We see that $B^{\prime}_{\rm p}$ dominates over (or
is similar to) $B^{\prime}_{\phi}$ up to relatively large distances. As noted
after equation (11), this follows from the fast decrease of
$B^{\prime}_{\phi}\propto(\sigma/\Gamma)^{1/2}r^{-1}$ in the ACZ. We have then
$B^{\prime}_{\phi}\mathrel{\hbox{\raise 2.15277pt\hbox{$\propto$}\hbox
to0.0pt{\hss\lower 2.15277pt\hbox{$\sim$}}}}z^{-(1/q_{1}+1/q_{2})}$,
$B_{\phi}\mathrel{\hbox{\raise 2.15277pt\hbox{$\propto$}\hbox
to0.0pt{\hss\lower 2.15277pt\hbox{$\sim$}}}}z^{-1/q_{2}}$, $B^{\prime}_{\rm
p}\propto z^{-2/q_{2}}$ in the ACZ, and $B^{\prime}_{\phi}\propto
B_{\phi}\propto z^{-1}$, $B^{\prime}_{\rm p}\propto z^{-2}$ in the conical
part. All components of $B$ are $\propto(P_{\rm j}/M)^{1/2}M^{-1/2}$. The kink
in the dependencies for $B_{\phi}$ and $B^{\prime}_{\phi}$ is artefact of the
assumption that the jet stops its acceleration sharply at $z_{\rm t}$, which
results in the derivative of $\Gamma(z)$ being not continuous. Similar plots
are shown in Vlahakis & Königl (2003) and Komissarov et al. (2009) for
parameters relevant to gamma-ray bursts. As we noted above, they found
$B^{\prime}_{\rm p}=B^{\prime}_{\phi}$ to take place in the ACZ, which
difference is related to our specific assumptions of equations (1–4) as well
as their assumed $\sigma_{\rm min}\approx 1$. We also note that equation (4)
gives the opening angle of $\Theta_{0}=a_{r}a_{\Gamma}/(q_{2}\Gamma_{\rm
max})$, which can be large for a low $\Gamma_{\rm max}$. This is also in
contrast to the observations of accreting BHs in stellar binary systems, where
low values of both $\Gamma_{\rm max}$ and $\Theta_{0}$ are commonly observed
(Miller-Jones et al., 2006; Zdziarski et al., 2022a). This issue can be
resolved for a low enough $a_{\Gamma}$.

We then compare our obtained magnetic field in the jet close to the horizon to
that in a MAD disc. We follow the estimate of the MAD magnetospheric radius of
equation (1) of Xie & Zdziarski (2019),

$B_{z}^{2}\approx\frac{r_{\rm g}\dot{M}_{\rm accr}c^{2}}{r_{\rm
disc}^{3}v_{r}h_{\rm disc}/r_{\rm disc}}\approx\frac{r_{\rm g}P_{\rm
j}}{r_{\rm disc}^{3}v_{r}a_{*}^{2}0.3h_{0.3}^{2}},$ (21)

where in the second equality we used the jet power of equation (20) at the
maximum flux and neglecting the numerical coefficient of 1.7. We note that
$v_{r}$ will not be equal that for a usual accretion disc, given the ‘choked’
nature of MAD accretion. We then compare this estimate at $r_{\rm
disk}=r=r_{\rm H}$ to the $B^{\prime}_{\rm p}$ of equation (15). We find that
their equality is obtained for

$(v_{r}/c)h_{0.3}^{2}\approx(\ell/5)(r_{\rm H}/r_{\rm g}).$ (22)

A similar estimate was done in Tchekhovskoy (2015). We see that the equality
requires $h_{0.3}^{2}v_{r}/c\sim 0.1$.

##  3 Summary and discussion

We have presented a simple analytical model of magnetized jets. We parametrize
the ACZ by profiles of the bulk Lorentz factor and radius [$\Gamma(z)$ and
$r(z)$, equations (1) and (2), respectively), and assume the jet is conical
and at a constant speed beyond the ACZ. However, our general results do not
depend on the specific form of $\Gamma(z)$ and $r(z)$. We then assume a
constant mass flow rate through the jet, equation (5), which is equivalent to
assuming the jet is mass-loaded at its base. Next, we assume the jet power is
conserved, which corresponds to the radiated power being $\ll P_{\rm j}$,
which is often satisfied. Given the above assumptions, we obtain the profiles
of the radius-averaged magnetization parameter and the toroidal component of
the magnetic field strength, with an additional assumption of the internal
particle density being less than the rest energy density. We then relate the
poloidal component of the magnetic field to the toroidal one assuming the
rotation of the magnetic field lines related to the BH angular velocity, and
take into account the radial profiles of the field components being not
uniform. This gives us the conserved magnetic flux. This flux, when used in
the formula for the BH spin-extraction power, yields then the initially
assumed jet power, showing the self-consistency of our approach. The main free
parameters of our model are $P_{\rm j}$, $\Gamma_{\rm max}$ and $\sigma_{\rm
min}$.

Using the above formulation, we find that the poloidal component of the
magnetic field can dominate up to one order of magnitude beyond the length,
$z_{\rm t}$, of the ACZ, see equation (16) and Fig. 1. The cause of this
effect is the conversion of the toroidal magnetic field component into the
bulk acceleration in the ACZ, which results in $B^{\prime}_{\phi}\propto
r^{-1}(\sigma/\Gamma)^{1/2}$, which yields in turn the decline with the radius
a significantly faster than $r^{-1}$. This effect was already noticed in
Vlahakis & Königl (2003) and Komissarov et al. (2009). Furthermore, we stress
that the detailed profiles depend on the profiles of $\Gamma(z)$ and $r(z)$,
and the $B^{\prime}_{\rm p}$–$B^{\prime}_{\phi}$ equipartition may occur at
lower values of $z$ than those shown in Fig. 1. Still, we note the overall
dependence is relatively weak. At the base of the jet, we have $\Gamma\sim 1$
and $r\sim r_{\rm H}$ while in the conical part we have $\Gamma_{\rm max}\gg
1$ and $\sigma_{\rm min}<1$. The profiles of $B^{\prime}_{\phi}$ and
$B^{\prime}_{\rm p}$ have to then smoothly join the two regimes.

We have made several assumptions, listed above. While some of them correspond
to the standard magnetic jet model, e.g., Tchekhovskoy et al. (2009), the
presented formalism provides a framework in which these assumptions can be
modified. E.g., we can modify equation (5) to allow for mass loading along the
jet, or modify equation (6) to account for radiative losses. We can also
modify the angular velocity of the field lines to that of the model of the jet
powered by the magnetic fields attached to the accretion disc (Blandford &
Payne, 1982). Furthermore, one can assume the geometric profile of $r(z)$ and
calculate the profile of $\Gamma(z)$ using the momentum conservation along a
representative streamline (e.g., assuming self-similarity). However, this
would involve divergence of tensors, and can be done only for individual
stream lines requiring the knowledge of the values of $B^{\prime}_{\phi}$ and
$B^{\prime}_{\rm p}$ along them. This would significantly complicate the
presented formalism (cf., e.g., Vlahakis & Königl 2003).

We then consider the jet stability for the $m=1$ kink modes driven by poloidal
electric current, see Appl et al. (2000). In the force-free case with the
magnetic pitch, $\mathcal{P}^{\prime}\equiv r^{\prime}B^{\prime}_{\rm
p}(r^{\prime})/B^{\prime}_{\phi}(r^{\prime})$, independent of $r^{\prime}$ (as
assured by equation 13), the growth time scale for a linear perturbation is
$\tau_{\rm kink}^{\prime}\simeq 7.5\mathcal{P}^{\prime}/v_{\rm A}^{\prime}$,
where $v_{\rm A}^{\prime}=c\sqrt{\sigma_{\rm p}/(1+\sigma_{\rm p})}$ is the
comoving-frame Alfvén speed at the jet axis and $\sigma_{\rm
p}={B^{\prime}_{\rm p}}^{2}(r^{\prime}=0)/(4\pi\rho^{\prime}c^{2})$. With
equation (13), $\mathcal{P}^{\prime}\approx 2\beta\Gamma r_{\rm
H}/(a_{*}\ell)$, and then

$\tau^{\prime}_{\rm kink}(z)\approx\frac{15\beta\Gamma r_{\rm H}}{a_{*}\ell
c\beta^{\prime}_{\rm A,p}}.$ (23)

We can then calculate the growth of the perturbation amplitude,
$E^{\prime}_{1}$, of the electric field along the jet,

$\int{\rm d}\ln E^{\prime}_{1}=\int_{z_{\rm m}}^{z}{\rm
d}z^{\prime}/[\Gamma(z^{\prime})\beta(z^{\prime})c\tau_{\rm
kink}^{\prime}(z^{\prime})].$ (24)

We note that $\tau_{\rm kink}^{\prime}$ depends on the value of $B_{\rm
p}^{\prime}(r^{\prime}=0)$, which may be divergent according to the profiles
adopted in Eq. (14). In such case, we approximate $\beta_{\rm
A,p}^{\prime}\approx 1$, although we note that in reality the profile of
$B_{\rm p}^{\prime}(r^{\prime})$ should approach a finite value at
$r^{\prime}=0$, which would lower $\beta_{\rm A,p}^{\prime}$ and increase
$\tau_{\rm kink}^{\prime}$. The numerical solutions to equation (24), which we
include in Fig. 1, correspond to the fastest growth of perturbation amplitude
allowed by the functions $\mathcal{P}^{\prime}(z)$ and $\Gamma(z)$. We find
that the perturbation amplitude shows two phases of growth – an initial phase
of slow growth followed by a phase of fast growth. The transition between
these two phases corresponds roughly to the distance at which
$B_{\phi}^{\prime}$ starts to dominate over $B_{\rm p}^{\prime}$, which is
usually somewhat beyond $z_{\rm t}$.

In M87, Asada & Nakamura (2012) and Asada et al. (2014) found $q_{2}\approx
1.7$ and $q_{1}\approx q_{2}/(q_{2}-1)\approx 2.4$, respectively, and $z_{\rm
t}\approx 2\times 10^{5}r_{\rm g}$. However, the initial acceleration is in
the non-relativistic regime, and the above power-law dependence of $\Gamma$ is
reproduced only for $z\gtrsim 2\times 10^{4}r_{\rm g}$. This indicates further
complexity in the observed sources. This, however, affects only our equation
(1) and the relation between $z_{\rm t}$ and $\Gamma_{\rm max}$, and could be
accounted for in future work. In general, the jet radial profile results from
the dependence of the external pressure on the distance.

## ACKNOWLEDGEMENTS

We thank Arieh Königl, Serguei Komissarov and Sasha Veledina for valuable
discussions. We acknowledge support from the Polish National Science Center
under the grants 2019/35/B/ST9/03944 and 2021/41/B/ST9/04306.

## Data Availability

There are no new data associated with this article.

## References

  * Appl & Camenzind (1993) Appl S., Camenzind M., 1993, A&A, 270, 71
  * Appl et al. (2000) Appl S., Lery T., Baty H., 2000, A&A, 355, 818
  * Asada & Nakamura (2012) Asada K., Nakamura M., 2012, ApJ, 745, L28
  * Asada et al. (2014) Asada K., Nakamura M., Doi A., Nagai H., Inoue M., 2014, ApJ, 781, L2
  * Beskin & Nokhrina (2006) Beskin V. S., Nokhrina E. E., 2006, MNRAS, 367, 375
  * Bisnovatyi-Kogan & Ruzmaikin (1974) Bisnovatyi-Kogan G. S., Ruzmaikin A. A., 1974, Ap&SS, 28, 45
  * Blandford & Königl (1979) Blandford R. D., Königl A., 1979, ApJ, 232, 34
  * Blandford & Payne (1982) Blandford R. D., Payne D. G., 1982, MNRAS, 199, 883
  * Blandford & Znajek (1977) Blandford R. D., Znajek R. L., 1977, MNRAS, 179, 433
  * Clausen-Brown et al. (2013) Clausen-Brown E., Savolainen T., Pushkarev A. B., Kovalev Y. Y., Zensus J. A., 2013, A&A, 558, A144
  * Davis & Tchekhovskoy (2020) Davis S. W., Tchekhovskoy A., 2020, ARA&A, 58, 407
  * Ghisellini & Tavecchio (2009) Ghisellini G., Tavecchio F., 2009, MNRAS, 397, 985
  * Ghisellini et al. (2014) Ghisellini G., Tavecchio F., Maraschi L., Celotti A., Sbarrato T., 2014, Nature, 515, 376
  * Heinz & Sunyaev (2003) Heinz S., Sunyaev R. A., 2003, MNRAS, 343, L59
  * Jorstad et al. (2005) Jorstad S. G., et al., 2005, AJ, 130, 1418
  * Kantzas et al. (2022) Kantzas D., Markoff S., Lucchini M., Ceccobello C., Grinberg V., Connors R. M. T., Uttley P., 2022, MNRAS, 510, 5187
  * Komissarov et al. (2007) Komissarov S. S., Barkov M. V., Vlahakis N., Königl A., 2007, MNRAS, 380, 51
  * Komissarov et al. (2009) Komissarov S. S., Vlahakis N., Königl A., Barkov M. V., 2009, MNRAS, 394, 1182
  * Königl (1981) Königl A., 1981, ApJ, 243, 700
  * Li et al. (1992) Li Z.-Y., Chiueh T., Begelman M. C., 1992, ApJ, 394, 459
  * Lucchini et al. (2021) Lucchini M., Russell T. D., Markoff S. B., Vincentelli F., Gardenier D., Ceccobello C., Uttley P., 2021, MNRAS, 501, 5910
  * Lyubarsky (2009) Lyubarsky Y., 2009, ApJ, 698, 1570
  * Lyubarsky (2010) Lyubarsky Y. E., 2010, MNRAS, 402, 353
  * McKinney et al. (2012) McKinney J. C., Tchekhovskoy A., Blandford R. D., 2012, MNRAS, 423, 3083
  * Miller-Jones et al. (2006) Miller-Jones J. C. A., Fender R. P., Nakar E., 2006, MNRAS, 367, 1432
  * Narayan et al. (2003) Narayan R., Igumenshchev I. V., Abramowicz M. A., 2003, PASJ, 55, L69
  * O’ Riordan et al. (2018) O’ Riordan M., Pe’er A., McKinney J. C., 2018, ApJ, 853, 44
  * Pjanka et al. (2017) Pjanka P., Zdziarski A. A., Sikora M., 2017, MNRAS, 465, 3506
  * Pushkarev et al. (2009) Pushkarev A. B., Kovalev Y. Y., Lister M. L., Savolainen T., 2009, A&A, 507, L33
  * Sądowski & Sikora (2010) Sądowski A., Sikora M., 2010, A&A, 517, A18
  * Tchekhovskoy (2015) Tchekhovskoy A., 2015, Launching of Active Galactic Nuclei Jets, ASSL. Springer, p. 45, doi:10.1007/978-3-319-10356-3_3
  * Tchekhovskoy et al. (2009) Tchekhovskoy A., McKinney J. C., Narayan R., 2009, ApJ, 699, 1789
  * Tchekhovskoy et al. (2010) Tchekhovskoy A., Narayan R., McKinney J. C., 2010, ApJ, 711, 50
  * Tchekhovskoy et al. (2011) Tchekhovskoy A., Narayan R., McKinney J. C., 2011, MNRAS, 418, L79
  * Vlahakis & Königl (2003) Vlahakis N., Königl A., 2003, ApJ, 596, 1080
  * Xie & Zdziarski (2019) Xie F.-G., Zdziarski A. A., 2019, ApJ, 887, 167
  * Zamaninasab et al. (2014) Zamaninasab M., Clausen-Brown E., Savolainen T., Tchekhovskoy A., 2014, Nature, 510, 126
  * Zdziarski et al. (2015) Zdziarski A. A., Sikora M., Pjanka P., Tchekhovskoy A., 2015, MNRAS, 451, 927
  * Zdziarski et al. (2019) Zdziarski A. A., Stawarz Ł., Sikora M., 2019, MNRAS, 485, 1210
  * Zdziarski et al. (2022a) Zdziarski A. A., Tetarenko A. J., Sikora M., 2022a, ApJ, 925, 189
  * Zdziarski et al. (2022b) Zdziarski A. A., Phuravhathu D. G., Sikora M., Böttcher M., Chibueze J. O., 2022b, ApJ, 928, L9


# Technical report: Improving the properties of molecules generated by LIMO

Vineet Thumuluri
Computer Science and Engineering
University of California San Diego
<EMAIL_ADDRESS>
&Peter Eckmann
Computer Science and Engineering
University of California San Diego
<EMAIL_ADDRESS>
&Michael K. Gilson
Skaggs School of Pharmacy and Pharmaceutical Sciences
University of California San Diego
<EMAIL_ADDRESS>
&Rose Yu
Computer Science and Engineering
University of California San Diego
<EMAIL_ADDRESS>

###### Abstract

This technical report investigates variants of the Latent Inceptionism on
Molecules (LIMO) framework to improve the properties of generated molecules.
We conduct ablative studies of molecular representation, decoder model, and
surrogate model training scheme. The experiments suggest that an autogressive
Transformer decoder with GroupSELFIES achieves the best average properties for
the random generation task.

##  1 LIMO framework

LIMO (Eckmann et al. (2022)) is a molecular generation technique that improves
a given set of properties by mapping molecules to a latent space and uses
inexpensive property surrogates to turn a discrete space optimization problem
into a continuous one. There are multiple components to this framework which
are described below.

###  1.1 Generative model

Given a dataset of molecular strings, a generative model is first trained to
mimic the distribution of SELFIES Krenn et al. (2020) tokens, i.e. smallest
units of a molecule string. In this case, a Variational AutoEncoder (VAE)
Kingma & Welling (2022), is trained to map molecules ($X$) to a latent space
($Z$) using an encoder ($E$) i.e. $E:X\xrightarrow{}Z$ and decoder (D) that
learns an inverse mapping i.e. $D:Z\xrightarrow{}X$, such that it is easy to
reconstruct $X$ from Z using the decoder ($D$) i.e. $D(E(X))$ $\approx$ $X$
(Reconstruction), and the latent space is close to a unit normal distribution
i.e. $D_{KL}(z||\mathcal{N}(0,\,I))\approx 0$ (Regularization).

This defines a generative model $p(x)=\int p(x|z)p(z)$. The VAE parameters are
obtained by minimizing the ELBO loss Kingma & Welling (2022). In this work, we
use a variant that weights the regularization term to avoid a problem during
optimization where the KL term vanishes Higgins et al. (2017).

$L_{\mathrm{BETA}}(\phi,\beta)=-\mathbb{E}_{\mathbf{z}\sim
q_{\phi}(\mathbf{z}\mid\mathbf{x})}\log
p_{\theta}(\mathbf{x}\mid\mathbf{z})+\beta
D_{\mathrm{KL}}\left(q_{\phi}(\mathbf{z}\mid\mathbf{x})\|p_{\theta}(\mathbf{z})\right)$

$\beta$ value of 0.1 was used in LIMO (Eckmann et al. (2022)).

###  1.2 Surrogate Model

LIMO uses a property predictor model to map the molecular representation to
property values Eckmann et al. (2022). The molecular representation can either
be a sample from the latent space or the decoder output. The property
predictor acts as a fully differentiable and faster surrogate model to the
oracle property predictor tools that are often computationally expensive. The
properties are generally scalar values and thus the parameters of the neural
network are learned by minimizing the mean-squared error between the
predictions and oracle property values.

###  1.3 Reverse optimization

The key idea of LIMO is the latent inceptionism technique to reverse optimize
the latent space and generate molecules with desired properties. Specifically,
a large number of random points are sampled according to the unit normal
distribution in the latent space. These serve as the starting points for the
search. Molecules with the optimized properties are generated by local
optimization in the latent space with the gradient approximated using the
trained surrogate model. The overall objective being optimized is a weighted
sum of one or more property values. All of the minima are then decoded to
molecules and the top molecules, measured using the oracle property value
predictor, are chosen.

##  2 Attempted modifications to LIMO

We conduct an extensive ablation study of the LIMO framework, investigating
different variations of LIMO components. These include changes to the
tokenization of the molecular strings, the decoding model from the latent
space to the molecule, and changes to how the property predictor is trained.

###  2.1 Molecular Representation

To represent molecules, LIMO used SELFIES Krenn et al. (2020), a molecular
string representation that ensures the tokens can be combined in any way to
always form valid molecules. The molecular strings can be decomposed into
tokens in many ways. GroupSELFIES Cheng et al. (2022) builds on top of the
chemical validity guarantees of SELFIES by enabling group tokens, thereby
creating additional flexibility to the representation. It extends SELFIES
where the tokens can be larger fragments and depend on the dataset and the
fragmentation technique.

In this work, we use the default fragmentation method from Cheng et al. (2022)
with different datasets:

(1) GS-Paper: GroupSELFIES tokens from the original paper Cheng et al. (2022).

(2) GS-USPTO: GroupSELFIES tokens extracted from USPTO reactions Lowe (2017).

(3) GS-Zinc: GroupSELFIES tokens extracted from ZINC250K molecules Akhmetshin
et al. (2021).

Table 1 shows the total number of tokens, the average and maximum sequence
length for each of the molecular string representations for the training
dataset (ZINC 250K) considered in this study.

Table 1: Statistics of different molecular string representation schemes when tokenizing the training dataset (ZINC250K). | Tokenizer
---|---
| SELFIES | GroupSELFIES
|  | ZINC250K | USPTO | Paper
Total tokens | 108 | 304 | 242 | 248
Max length | 72 | 75 | 77 | 93
Avg. length | 37.43 | 37.07 | 36.61 | 29.86

###  2.2 Decoder model

The decoder in LIMO maps from latent space to a molecular string
representation. Such mapping can be approximated in multiple ways. Here we
discuss two broad strategies i.e. Non-autoregressive (NAR) and Autoregressive
(AR) decoder models. In this work, both of these are implemented using the
transformer architecture Vaswani et al. (2023).

####  2.2.1 Non-autoregressive models

LIMO Eckmann et al. (2022) uses an MLP decoder by modelling the joint
distribution of the target tokens ($y$) as conditionally independent i.e.
$P(y)=\prod_{i=1}^{n}P\left(y_{i}\mid Z\right)$. While this has the advantage
of inference being parallelizable, it suffers from the multi-modality problem
Gu et al. (2018). There have been multiple works that try to mitigate this
issue through iterative refinement such as CMLM Yang et al. (2021), which uses
masked language modeling and multi-step decoding. At inference time, first,
the input is fully masked and in each subsequent iteration, the least likely
tokens are masked and the procedure is repeated. A follow-up work CMLMC Huang
et al. (2022), aims to improve the training objective of CMLM by addressing
the mismatch at inference by including a token denoising loss in addition to
the masking loss. Another advantage of NAR models is the ability to do
conditional generation by constraining on a given scaffold.

####  2.2.2 Autoregressive models

Autoregressive models, such as Transformers, are slower since only one token
is decoded at each step and the cost scales at least linearly with the length
of the target sequence length. However, the full joint distribution can be
modeled i.e. $P(y)=\prod_{i=1}^{n}P\left(y_{i}\mid
y_{1},...,y_{i-1},Z\right)$, and thus does not suffer from the multi-modality
problem. AR models thus have been shown to have superior task performance Ren
et al. (2020). In contrast to NAR models, conditional generation is non-
trivial, thus losing the ability to do scaffold-constrained generation.

###  2.3 Surrogate Model training

The property predictor surrogate model training depends on the VAE latent
space i.e. the input to the surrogate predictor used in LIMO is a decoded
representation of the point in the latent space, there are two approaches to
optimizing the whole system.

The first is to perform the training sequentially, in which case, the VAE is
trained to convergence using the molecular string dataset and then its
parameters are frozen. This ensures a fixed latent space. Random molecules are
generated and their property values are computed to create the training
dataset for the neural network-based surrogate.

An alternative strategy is to jointly train both the VAE and the surrogate. In
this case, the property values of the training molecules need to be computed.
An advantage of this approach is that the latent space can be informed of the
desired properties explicitly, however, the generative model optimization
becomes harder as there are more training objectives to balance (empirically
we find the losses to be higher) and a lot more oracle compute is required. In
this study we chose to use a surrogate that directly predicts using the latent
space to simplify the optimization.

##  3 Experiments

###  3.1 Setup

All methods are trained on the ZINC-250K dataset. The parameters are optimized
using the Adam optimizer Kingma & Ba (2017), using a cosine-annealing learning
rate schedule for one cycle Loshchilov & Hutter (2017) for a total of 100K
iterations. To prevent KL from vanishing across models with minimal tuning, a
cyclic beta schedule was used while training the VAE Fu et al. (2019). We
follow the LIMO paper and use the following properties to evaluate the quality
of generated optimized molecules.

  * •

SA: Synthetic-accessibility, computed with SAScorer Ertl & Schuffenhauer
(2009). A score between 1 (easy to make) and 10 (very difficult to make).

  * •

QED: Quantitative Estimate of Drug-likeness, estimated by RDKit Landrum . A
value between 0 (non-drug-like) and 1 (drug-like).

  * •

BA: Binding Affinity to the human estrogen receptor (PDB 1ERR), computed with
AutoDock-GPU Santos-Martins et al. (2021). Outputs a prediction of the $\Delta
G$ in kcal/mol, with lower values indicating stronger binding. We used 0 when
the value could not be computed.

Table 2: Property comparison of all generated molecules from different model variants. Model Variants | Count with valid AutoDock BA | $BA\downarrow$ | $SA\downarrow$ | $QED\uparrow$
---|---|---|---|---
LIMO | Default-Retrained | 9813 | -4.67 | 4.46 | 0.57
GS-Paper | 8729 | -5.44 | 2.89 | 0.62
GS-USPTO | 9947 | -5.56 | 4.46 | 0.61
GS-Zinc | 9868 | -5.57 | 4.63 | 0.59
AR | Default | 9194 | -5.03 | 4.93 | 0.56
Joint-Z | 9229 | -5.57 | 3.56 | 0.76
AR-JointZ | GS-Paper | 9768 | -5.26 | 2.91 | 0.81
GS-USPTO | 9939 | -5.58 | 3.28 | 0.73
GS-Zinc | 9967 | -5.74 | 3.39 | 0.80
CMLMC | 8758 | -5.11 | 3.65 | 0.30
ZINC250K (Training dataset ) | 242591 | -5.35 | 3.05 | 0.73

###  3.2 Property Optimization Results

We apply reverse optimization to generate a set of random molecules from the
trained generative model. For each experiment, the property optimization step
was applied with 10,000 starting points and all of the resulting minima are
considered.

We consider the following model variants (1) LIMO: The original LIMO Eckmann
et al. (2022). (2) AR: LIMO with decoder replaced with an autoregressive
Transformer model. (3) AR-JointZ: AR with surrogate trained simultaneously
with VAE. (4) CMLMC: Non-Autoregressive decoder model from Huang et al. (2022)
For LIMO, AR and AR-JointZ, we further consider different molecular string
representation schemes (see Section 2.1 for more details on the various
molecular representations).

The results in Tables 2 show the number of binding ligands and the average
property scores of those ligands for the human estrogen receptor (PDB 1ERR)
target by considering 10,000 generated molecules each. We also report in
Appendix Table 3, 4, 5 the results for molecules with the top 100 BA, the
molecules with the top 100 SA, and the molecules with the top 100 QED
respectively. The AR VAE decoder with the GroupSELFIES variants has the
molecules with the best generated SA, QED (tokens provided by Cheng et al.
(2022)), and BA (tokens extracted from ZINC250K).

We also include the corresponding property values of the training dataset
(ZINC250K) as a reference. Note that the property values of generated
molecules are even better than those of the training dataset. This
demonstrates the great potential of using LIMO for de novo drug discovery.

Figure 1 compares the distributions of property values for optimized molecules
generated by LIMO, AR-JointZ GS-Zinc, and the training dataset. We can see
that AR-JointZ with GS-ZINC representation improves the original LIMO model
across all metrics. The model even outperforms the training dataset in terms
of BA and QED values.

Figure 1: Distributions of property values for optimized molecules generated
by LIMO, AR-JointZ GS-Zinc, and the training dataset.

###  3.3 Latent space analysis

While Table 2 shows that the trained models are able to match the property
distribution in the training dataset approximately, Table 3, 4, 5 in Appendix
show that the optimized molecules, however, fall short of molecules with the
best property values in the training dataset. This difference in the best
property values is due to the surrogate model-based optimization in the latent
space.

We first investigate the organization of the latent space and identify
features that would indicate the ease of optimization in this space. Figure 2
shows the property value distribution in the learned latent space of LIMO. In
the case of sequentially training the VAE and then the surrogate, the latent
space receives no explicit feedback from the property values. Hence, the
property values are not smoothly distributed, i.e. molecules decoded from the
same neighborhood of the latent space can have large variations in their
property values.

Since the surrogate is a simple approximation of the true latent space to
property mapping, we measure the mean-squared error and the Dirichlet energy
(Equation 1) to quantify the difference in surface smoothness, as well as the
local correlation as measured using the pearson correlation coefficient to
quantify the correctness of gradient directions in both the surfaces.

$\lambda_{y}=\frac{1}{N}y^{\mathrm{T}}Ly$ (1)

Figure 2: Objective plotted along a random subspace of the latent space.
Latent space points were decoded to molecules by uniformly sampling a region
constructed using two orthogonal random directions. The plot shows that
objective values vary greatly even for nearby latent points.

where $y$ is the property value, $L=D-A$, $D$ is the degree matrix, $A$ is the
KNN adjacency matrix, and $N$ is the number of samples considered for
computing the Dirichlet energy. The property value measured in our experiments
is the weighted sum of SA, QED, and BA as used in Eckmann et al. (2022). The
training molecules are encoded into the latent space and 5 nearest neighbors
are found using KNN which are then used to compute the Dirichlet energy.

The mean-squared error was found to decrease when training the surrogate
jointly with the VAE (0.634) vs when training them separately (0.986),
indicating that joint training leads to a better fit.

Additionally, the Dirichlet energies are lower when jointly training (33.31 vs
94.04), indicating that optimizing for properties using the jointly trained
surrogate leads to smoother changes in property values i.e. molecules decoded
from the same neighborhood of the latent space tend to have more similar
property values.

Lastly, the local correlation as measured using the pearson correlation
coefficient decreases when jointly training (0.48 vs 0.76), indicating that
local gradient-based optimization using the jointly trained surrogate is less
accurate.

##  4 Future work

###  4.1 Latent space organization

While jointly training the VAE and the surrogate improves the generated
molecules by explicit feedback making the latent space interpolable with
respect to properties, further improvements to the latent space organization
via topological constraints Moor et al. (2021); Keller & Welling (2022) is a
promising direction.

###  4.2 Semi-autoregressive modelling

There is a tradeoff between controllable generation and sequence modeling
ability by changing the non-autoregressive decoder to an autoregressive one.
Huang et al. (2022) is an attempt at merging the best of both worlds by having
an iterative non-autoregressive decoder. Other such approaches include semi-
autoregressive models such as those used in Han et al. (2023) which generate
blocks of text autoregressively, and order-agnostic autoregressive models such
as those used in Alamdari et al. (2023) which can generate text in an
arbitrary order i.e. not just left to right.

###  4.3 Further analysis of GroupSELFIES tokens

The experiments done in this work show that the performance of the
GroupSELFIES tokens depends greatly on the extracted tokens. While there is no
obvious difference between the tokens extracted from different datasets,
further study is needed to determine the cause of the apparent different in
performance.

Pursuing these directions should improve the properties of molecules generated
by LIMO, as well as maintain the ability to constrain the molecular scaffold.

## References

  * Akhmetshin et al. (2021) Tagir Akhmetshin, Arkadii I. Lin, Daniyar Mazitov, Evgenii Ziaikin, Timur Madzhidov, and Alexandre Varnek.  ZINC 250K data sets.  12 2021.  doi: 10.6084/m9.figshare.17122427.v1.  URL https://figshare.com/articles/dataset/ZINC_250K_data_sets/17122427.
  * Alamdari et al. (2023) Sarah Alamdari, Nitya Thakkar, Rianne van den Berg, Alex X. Lu, Nicolo Fusi, Ava P. Amini, and Kevin K. Yang.  Protein generation with evolutionary diffusion: sequence is all you need.  September 2023.  doi: 10.1101/2023.09.11.556673.  URL http://dx.doi.org/10.1101/2023.09.11.556673.
  * Cheng et al. (2022) Austin Cheng, Andy Cai, Santiago Miret, Gustavo Malkomes, Mariano Phielipp, and Alán Aspuru-Guzik.  Group selfies: A robust fragment-based molecular string representation.  2022\.  doi: 10.48550/ARXIV.2211.13322.  URL https://arxiv.org/abs/2211.13322.
  * Eckmann et al. (2022) Peter Eckmann, Kunyang Sun, Bo Zhao, Mudong Feng, Michael K Gilson, and Rose Yu.  Limo: Latent inceptionism for targeted molecule generation.  2022\.
  * Ertl & Schuffenhauer (2009) Peter Ertl and Ansgar Schuffenhauer.  Estimation of synthetic accessibility score of drug-like molecules based on molecular complexity and fragment contributions.  _Journal of Cheminformatics_ , 1(1), June 2009.  ISSN 1758-2946.  doi: 10.1186/1758-2946-1-8.  URL http://dx.doi.org/10.1186/1758-2946-1-8.
  * Fu et al. (2019) Hao Fu, Chunyuan Li, Xiaodong Liu, Jianfeng Gao, Asli Celikyilmaz, and Lawrence Carin.  Cyclical annealing schedule: A simple approach to mitigating kl vanishing, 2019.
  * Gu et al. (2018) Jiatao Gu, James Bradbury, Caiming Xiong, Victor O. K. Li, and Richard Socher.  Non-autoregressive neural machine translation, 2018.
  * Han et al. (2023) Xiaochuang Han, Sachin Kumar, and Yulia Tsvetkov.  Ssd-lm: Semi-autoregressive simplex-based diffusion language model for text generation and modular control, 2023.
  * Higgins et al. (2017) Irina Higgins, Loic Matthey, Arka Pal, Christopher Burgess, Xavier Glorot, Matthew Botvinick, Shakir Mohamed, and Alexander Lerchner.  beta-VAE: Learning basic visual concepts with a constrained variational framework.  In _International Conference on Learning Representations_ , 2017.  URL https://openreview.net/forum?id=Sy2fzU9gl.
  * Huang et al. (2022) Xiao Shi Huang, Felipe Perez, and Maksims Volkovs.  Improving non-autoregressive translation models without distillation.  In _International Conference on Learning Representations_ , 2022.  URL https://openreview.net/forum?id=I2Hw58KHp8O.
  * Keller & Welling (2022) T. Anderson Keller and Max Welling.  Topographic vaes learn equivariant capsules, 2022.
  * Kingma & Ba (2017) Diederik P. Kingma and Jimmy Ba.  Adam: A method for stochastic optimization, 2017.
  * Kingma & Welling (2022) Diederik P Kingma and Max Welling.  Auto-encoding variational bayes, 2022.
  * Krenn et al. (2020) Mario Krenn, Florian Häse, AkshatKumar Nigam, Pascal Friederich, and Alan Aspuru-Guzik.  Self-referencing embedded strings (selfies): A 100  _Machine Learning: Science and Technology_ , 1(4):045024, October 2020.  ISSN 2632-2153.  doi: 10.1088/2632-2153/aba947.  URL http://dx.doi.org/10.1088/2632-2153/aba947.
  * (15) Greg Landrum.  Rdkit: Open-source cheminformatics.  URL http://www.rdkit.org.
  * Loshchilov & Hutter (2017) Ilya Loshchilov and Frank Hutter.  Sgdr: Stochastic gradient descent with warm restarts, 2017.
  * Lowe (2017) Daniel Lowe.  Chemical reactions from US patents (1976-Sep2016).  6 2017.  doi: 10.6084/m9.figshare.5104873.v1.  URL https://figshare.com/articles/dataset/Chemical_reactions_from_US_patents_1976-Sep2016_/5104873.
  * Moor et al. (2021) Michael Moor, Max Horn, Bastian Rieck, and Karsten Borgwardt.  Topological autoencoders, 2021.
  * Ren et al. (2020) Yi Ren, Jinglin Liu, Xu Tan, Zhou Zhao, Sheng Zhao, and Tie-Yan Liu.  A study of non-autoregressive model for sequence generation, 2020.
  * Santos-Martins et al. (2021) Diogo Santos-Martins, Leonardo Solis-Vasquez, Andreas F Tillack, Michel F Sanner, Andreas Koch, and Stefano Forli.  Accelerating a¡scp¿uto¡/scp¿d¡scp¿ock¡/scp¿4 with gpus and gradient-based local search.  _Journal of Chemical Theory and Computation_ , 17(2):1060–1073, January 2021.  ISSN 1549-9626.  doi: 10.1021/acs.jctc.0c01006.  URL http://dx.doi.org/10.1021/acs.jctc.0c01006.
  * Vaswani et al. (2023) Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N. Gomez, Lukasz Kaiser, and Illia Polosukhin.  Attention is all you need, 2023.
  * Yang et al. (2021) Ziyi Yang, Yinfei Yang, Daniel Cer, Jax Law, and Eric Darve.  Universal sentence representation learning with conditional masked language model, 2021.

##  Appendix A Appendix

We report the top 100 molecules in terms of SA, BA and QED in the
corresponding tables.

Table 3: Top 100 BA molecules Model Variants | Count with valid AutoDock BA | $BA\downarrow$ | $SA\downarrow$ | $QED\uparrow$
---|---|---|---|---
LIMO | Default-Retrained | 9813 | -7.05 | 4.94 | 0.58
GS-Paper | 8729 | -7.65 | 3.54 | 0.49
GS-USPTO | 9947 | -7.55 | 4.75 | 0.59
GS-Zinc | 9868 | -7.46 | 4.97 | 0.57
AR | Default | 9194 | -6.80 | 5.18 | 0.67
Joint-Z | 9229 | -7.16 | 4.32 | 0.78
AR-JointZ | GS-Paper | 9768 | -6.53 | 3.05 | 0.82
GS-USPTO | 9939 | -7.11 | 3.41 | 0.77
GS-Zinc | 9967 | -7.11 | 3.65 | 0.84
CMLMC |  | 8758 | -7.85 | 4.04 | 0.28
Dataset | ZINC250K | 242591 | -8.22 | 3.08 | 0.51
Table 4: Top 100 SA molecules Model Variants | Count with valid AutoDock BA | $BA\downarrow$ | $SA\downarrow$ | $QED\uparrow$
---|---|---|---|---
LIMO | Default-Retrained | 9813 | -3.24 | 2.19 | 0.52
GS-Paper | 8729 | -3.36 | 1.12 | 0.53
GS-USPTO | 9947 | -3.52 | 2.27 | 0.56
GS-Zinc | 9868 | -3.87 | 2.57 | 0.54
AR | Default | 9194 | -4.38 | 2.14 | 0.68
Joint-Z | 9229 | -5.02 | 1.80 | 0.85
AR-JointZ | GS-Paper | 9768 | -5.10 | 1.72 | 0.74
GS-USPTO | 9939 | -5.19 | 1.83 | 0.85
GS-Zinc | 9967 | -5.37 | 2.01 | 0.85
CMLMC |  | 8758 | -4.40 | 1.74 | 0.22
Dataset | ZINC250K | 242591 | -4.74 | 1.42 | 0.80
Table 5: Top 100 QED molecules Model Variants | Count with valid AutoDock BA | $BA\downarrow$ | $SA\downarrow$ | $QED\uparrow$
---|---|---|---|---
LIMO | Default-Retrained | 9813 | -5.49 | 4.65 | 0.83
GS-Paper | 8729 | -5.47 | 3.02 | 0.91
GS-USPTO | 9947 | -5.83 | 4.51 | 0.89
GS-Zinc | 9868 | -5.93 | 4.60 | 0.86
AR | Default | 9194 | -5.52 | 3.83 | 0.91
| Joint-Z | 9229 | -5.46 | 2.74 | 0.94
AR-JointZ | GS-Paper | 9768 | -5.23 | 2.55 | 0.94
GS-USPTO | 9939 | -5.54 | 2.57 | 0.94
GS-Zinc | 9967 | -5.67 | 2.60 | 0.94
CMLMC |  | 8758 | -5.43 | 4.59 | 0.77
Dataset | ZINC250K | 242591 | -5.38 | 2.75 | 0.95