\section{Method}
\label{section:method}
\subsection{A new empirical gyrochronology relation}

We fit a broken power law to the 650 Myr Praesepe cluster in order to
calibrate a gyrochronology relation that takes advantage of new data available
from the \ktwo\ and \gaia\ surveys.
This relation captures the detailed shape of Praesepe's rotation period-color
relation, and is calibrated to \gaia\ \gcolor\ color.
Praesepe is the ideal calibration cluster because it has the largest number of
members with precisely measured rotation periods, over a large range of colors
\citep[spanning spectral types A0 through M6,][]{rebull2017}, of any open
cluster.
It is relatively tightly clustered on the sky and many of its members were
targeted in a single \ktwo\ campaign, from which rotation periods have been
measured via light curve frequency analysis \citep{douglas2017, rebull2017}.
We compiled rotation periods, \Gaia\ photometry and \gaia\ parallaxes for
members of the 650 Myr Praesepe cluster \citep{fossati2008} by identifying
Praesepe members with measured rotation periods from \citet{douglas2017} in
the \ktwo-\gaia\ crossmatch catalog provided at
\url{https://gaia-kepler.fun/}.
This catalog cross-matched the EPIC catalog \citep{huber2016} with the
Gaia DR2 catalog \citep{brown2018}, using a 1'' search radius.
The result was a sample of 757 stars with rotation periods, parallaxes and
\gaia\ $G$, $G_{BP}$ and $G_{RP}$-band photometry, shown in figure
\ref{fig:praesepe}\footnote{This analysis was performed in a Jupyter notebook
available here: \url{
    https://github.com/RuthAngus/stardate/blob/master/paper/code/Praesepe.ipynb}}.
Although this new model does not perfectly describe the rotation period of
every star, it provides a better representation of rotational evolution than
previous models such as the \citet{angus2015} model (shown as the blue solid
line in figure \ref{fig:praesepe}).
\begin{figure}
  \caption{
    The rotation periods of Praesepe members \citep{douglas2016},
    vs. their \Gaia\ colors (\gcolor) with a broken power law model, fit to
    these data.
    The dashed line and shaded region shows the mean and variance model
    described in equation \ref{eqn:gyro} at 650 Myrs (lower model) and 4.56
    Gyrs (upper model).
    The solid blue line shows the \citep{angus2015} gyrochronology relation at
    650 Myrs (lower model) and 4.56 Gyrs (upper model).
    Shaded regions show the 1$\sigma$ range of the rotation period
    model (equation \ref{eqn:gyro}).
    The rotation periods of stars bluer than around 0.56 dex and redder than
    around 2.7 dex in \gcolor\ are modeled as a broad log-normal distribution
    with a standard deviation of 0.5 dex, added to observational
    uncertainties.
    This figure was generated in a Jupyter notebook available at
    \url{https://github.com/RuthAngus/stardate/blob/master/paper/code/Fitting_Praesepe.ipynb}
}
  \centering
    \includegraphics[width=1.1\textwidth]{Praesepe.pdf}
\label{fig:praesepe}
\end{figure}
In order to fit a relation to Praesepe, we removed rotational outliers bluer
than \gcolor\ = 2.7 via sigma-clipping and fit a 5th-order polynomial to the
remaining FGK and early M stars.
\racomment{We found that a 5th order polynomial provided a substantially
better fit than lower-order polynomials, which were not able to capture the
sharp `elbow’ in the rotation period-color relation. Additional orders
provided either a worse
fit, tending towards extreme values at the boundaries, or diminishing returns
in goodness-of-fit.}
We also fit a straight line to the late M dwarfs (\gcolor\ $>$ 2.7), to
capture the mass-dependent initial rotation periods of low mass stars
\citep{somers2017}.
We fit a separable straight line function to the period-age relation using
the ages of Praesepe and the Sun\footnote{The fitting process was
performed in a Jupyter notebook available at
\url{https://github.com/RuthAngus/stardate/blob/master/paper/code/Fitting_Praesepe.ipynb}.}.
This new Praesepe-calibrated gyrochronology relation is,
\begin{equation}
    \log_{10}(P_\mathrm{rot}) =
    c_A\log_{10}(t) +
    \sum_{n=0}^4 c_n[\log_{10}(G_{BP}-G_{RP})]^n
\label{eqn:fgk_gyro}
\end{equation}
for stars with \gcolor\ $<$ 2.7 and
\begin{equation}
    \log_{10}(P_\mathrm{rot}) =
    c_A\log_{10}(t) +
    \sum_{m=0}^1 b_m[\log_{10}(G_{BP}-G_{RP})]^m
\label{eqn:m_gyro}
\end{equation}
for stars with \gcolor\ $>$ 2.7, where $P_{\mathrm{rot}}$ is rotation period
in days and $t$ is age in years.
Best-fit coefficient values are shown in table \ref{tab:coefficients}.
\begin{table}[h!]
  \begin{center}
      \caption{Coefficient values for equations \ref{eqn:fgk_gyro} and
      \ref{eqn:m_gyro}.}
    \label{tab:coefficients}
    \begin{tabular}{l|c} % <-- Alignments: 1st column left and 2nd middle, with vertical lines in between
      Coefficient & Value  \\
      \hline
      $c_A$ & 0.65 $\pm$ 0.05 \\
      $c_0$ & -4.7 $\pm$ 0.5 \\
      $c_1$ & 0.72 $\pm$ 0.05 \\
      $c_2$ & -4.9 $\pm$ 0.2 \\
      $c_3$ & 29 $\pm$ 2 \\
      $c_4$ & -38 $\pm$ 4 \\
      $b_0$ & 0.9 $\pm$ 0.5 \\
      $b_1$ & -13.6 $\pm$ 0.1 \\
    \end{tabular}
  \end{center}
\end{table}

We calibrated this new relation using Praesepe and the Sun alone, without
including other clusters, because different open clusters have slightly
different period-color relationships \citep{agueros2018, agueros2018b,
curtis2018, curtis2019} and, given that we used an age-color separable
relation, adding in extra clusters was unlikely to improve the calibration.
We did not include asteroseismic stars because most are slightly evolved and
would require a gyrochronology relation that depends on $\log g$.
This new relation, fit to Praesepe and the Sun, does not perfectly predict the
rotation periods of stars at all colors and ages, but it provides several
improvements over previous empirical gyrochronology relations.
Firstly, it uses new \ktwo\ rotation period measurements to model the
period-color relation of Praesepe in detail, secondly, it includes a model for
the rotational behavior of M dwarfs, and thirdly, it is calibrated to \Gaia\
\gcolor\ color: a directly observable quantity and the most widely available
photometric color index.

Equation \ref{eqn:fgk_gyro}, describing the rotational evolution of FGK and
early M stars, is most closely analogous to previously calibrated empirical
gyrochronology relations \citep[\eg][]{barnes2003, barnes2007, mamajek2008,
barnes2010, angus2015}.
It describes stars with Sun-like magnetic dynamos that follow a
`Skumanich-like' magnetic braking law, \ie\ their rotation period increases
with the square root of their age.
It does not describe stars hotter than around 6250 K which have a thin
convective layers and a weak magnetic dynamo, nor does it describe fully
convective stars which take a long time to converge onto the
\citet{skumanich1972} braking law \citep{krishnamurthi1997}.
In addition, stars with Rossby numbers larger than around 2 do not show
Skumanich-like magnetic braking \citep{vansaders2016, vansaders2018}, so
equation \ref{eqn:fgk_gyro} does not describe these stars.
It also does not describe the rotation periods of subgiants or giants, whose
rotation periods are influenced by their expanding radii, changing winds, and
core-to-surface differential rotation \citep[\eg][]{vansaders2013, tayar2018}.
Finally, this relation does not describe the rotation periods of dynamically
or magnetically interacting binaries which often rotate more rapidly than
isolated stars at the same age and color \citep{douglas2016}.
In order to include non-Skumanich type stars in our combined gyrochronal and
isochronal model, we designed a composite gyrochronology relation which
describes a mean and variance model for the rotation period distributions
across the HRD.
Rotation periods are modeled differently for stars of different photometric
color, EEP, Rossby number, age and metallicity.
The rotation period model for each of these groups is described below.
\begin{itemize}
    \item{The rotation periods of late F, GK, and early M dwarfs
        (0.56 $<$ \gcolor\ $<$ 2.7), with Rossby numbers less than 2, are
        modeled as a log-normal distribution, with mean given by equation
        \ref{eqn:fgk_gyro}, and variance given by the squared inverse of their
        observational uncertainties.}
    \item{The rotation periods of fully convective stars (\gcolor\ $>$ 2.7)
        are modeled with a log-normal distribution with mean given by equation
        \ref{eqn:m_gyro}, and an extra standard deviation of
        0.5 dex, added to observational uncertainties.
        This distribution reflects the observed rotation periods of late M
        dwarfs in the Praesepe cluster which span a broad range at every
        color.
        Late M dwarfs with masses $\lesssim$ 0.3 M$_\odot$, temperatures
        $\lesssim$ 3500 and \gcolor $\gtrsim$ 2.7 exhibit weak magnetic
        braking until at least after the age of Praesepe ($\sim$650 million
        years).
    \item{The rotation periods of F-type and hotter stars, with \gcolor\ $<$
        0.56 dex, are modeled as a log-normal distribution with mean,
        $\log_{10}(P_\mathrm{rot})=0.56$,
        and an additional standard deviation, added to observational
        uncertainties, of 0.5 dex.
        Stars more massive than around 1.25 M$_\odot$, with a temperature
        $\gtrsim$ 6250 K and a \gcolor\ $\lesssim$ 0.56 do not spin down
        appreciably over their main-sequence lifetimes because they do not
        have the deep convective envelope needed to generate strong magnetic
        fields.
        This model for the mean and variance of hot star rotation periods is
        based on stars hotter than 6250 K in the \citet{mcquillan2014} sample,
        which have a mean $\log_{10}(P_\mathrm{rot})$ of 0.56 dex and a
        standard deviation of around 0.5 dex.}
        Both hot and cool stars retain rotation periods that are similar to
        their primordial distribution \citep[see \eg][]{matt2012,
        somers2017}.}
    \item{The rotation periods of stars with large Rossby numbers are
        modeled as follows.
        The age at which a star's Rossby number would exceed 2 is calculated
        by inverting equation \ref{eqn:fgk_gyro}.
        If a star's age is greater than this, its mean rotation period is
        given by $P_\mathrm{max} = 2/\tau$, where $\tau$ is the convective
        turnover timescale, calculated via stellar mass using equation 11 from
        \citet{wright2011}.}
    \item{The rotation periods of subgiants are described with a log-normal
        distribution with mean given by equation \ref{eqn:fgk_gyro},
        \ref{eqn:m_gyro} or 0.56, depending on its color, and an additional
        standard deviation of 5 dex.
        This is not an accurate model of subgiant rotation
        periods \citep[see, \eg][]{vansaders2013} but the highly
        inflated variance makes it a weakly constraining one.
        Isochrone fitting provides precise ages for subgiants, so by inflating
        the variance of the gyrochronology relation at large EEP, we allow a
        star's position on the HRD/CMD to dominate the age information over
        its rotation period.
        This is useful because we have not yet built an accurate rotation-age
        relation for subgiants into our model.}
    \item{We model stars younger than around 250 Myrs with a log-normal
        distribution with mean function given by equation \ref{eqn:fgk_gyro},
        \ref{eqn:m_gyro} or 0.56, depending on it whether it has
        an intermediate, red, or blue color respectively and a inflated
        standard deviation of 0.5 dex.
        Rotation periods in young open clusters show a large amount of scatter
        because they have not yet converged onto the \citet{skumanich1972}
        spin-down sequence.}
    \item{Finally, we model stars with very low and high metallicites (-0.2
        $>$ [Fe/H] $>$ 0.2) with a log-normal distribution with mean given by
        equation \ref{eqn:fgk_gyro}, \ref{eqn:m_gyro} or 0.56, depending on
        its color, and a standard deviation of 0.5 dex.
        Gyrochronology is not calibrated at these extreme metallicities due to
        a lack of suitable metal poor and rich calibration stars.
        Rather than assume the same gyrochronology model can be na\"ively
        applied to these stars, we take a more conservative approach and model
        them with a broad Gaussian distribution.}
\end{itemize}
Inflating the variance of the rotation period distribtion for stars with
non-Skumanich magnetic braking behavior has two purposes: 1) in the cases of
hot stars and fully convective stars, it allows the broad distributions of
rotation periods observed in clusters and the field to be matched, and 2) it
down-weights the age-information provided by rotation periods in regions of
the HRD/CMD where rotation periods are not information-rich or the
gyrochronology model is inaccurate or uncalibrated.
If the observational uncertainties on rotation periods are 5\% on average,
which corresponds to 0.05/$\ln(10)$ $\sim$ 0.02 dex, adding 0.5 amounts to
a 25$\times$ increase in standard deviation, or a $\sim$ 600$\times$ increase
in variance.
% The variance on stellar ages inferred via gyrochronology is directly
% proportional to the variance on rotation periods, so a 600$\times$ increase
% in period variance corresponds to a 600$\times$ reduction in the
% age-information provide by gyrochronology, a 600$\times$ increase in
% gyrochronal age variance and a 25$\times$ increase in gyrochronal standard
% deviation.
So, practically speaking, when the standard deviation is inflated to 0.5 dex
or more, ages are almost entirely inferred via isochrone fitting.

We used the following composite gyrochronology model to infer ages from
rotation periods,
\begin{equation}
    \log_{10}(P_\mathrm{rot}) \sim \begin{cases}
        \mathcal{N}\left[(\ref{eqn:fgk_gyro}), (\sigma+\sigma_P)^2\right], & Ro < 2,
        0.56 < G_{BP} - G_{RP} < 2.7 \\
        \mathcal{N}[(\ref{eqn:m_gyro}), (\sigma+\sigma_P)^2], & Ro < 2,
        G_{BP} - G_{RP} > 2.7 \\
        \mathcal{N}[\log_{10}(P_\mathrm{max}), (\sigma+\sigma_P)^2],& Ro \geq 2 \\
        \mathcal{N}[0.56, (\sigma+\sigma_P)^2],& G_{BP} - G_{RP} < 0.56,
    \end{cases}
\label{eqn:gyro}
\end{equation}
where $\sigma$ is the relative period uncertainty, divided by $\ln(10)$, on
individual rotation period measurements and $\sigma_P$ is an additional
scatter that is a function of EEP, age, metallicity and color.
It takes a maximal value of 0.5 for hot stars and fully convective stars, 5
for subgiants and giants, and a minimal value of zero for late F, GK and early
M dwarfs.
The variance model is shown in figure \ref{fig:variance}.
Sigmoid functions were used to provide smooth transitions between regions of
low and high variance.
Sharp changes in variance would produce sharp changes in likelihood, which
would cause the posterior distributions over stellar parameters to be more
difficult to sample.
The sigmoid functions shown in figure \ref{fig:variance} reach half their
maximum values at \gcolor\ = 0.56 and 2.5 for hot and cool stars respectively,
EEP = 454 for subgiants, age = 250 Myrs for young stars, $[$M/H$]$ = -0.2 for
metal poor stars and 0.2 for metal rich stars.
The logistic growth rate, or steepness, of the sigmoid functions is 100
dex$^{-1}$ for both color transitions, .2 EEP$^{-1}$ for the EEP transition,
20 dex$^{-1}$ for the age transition and 5 dex$^{-1}$ for the both metallicity
transitions.
The additional standard deviation is additive, so if a star is \eg\ hot,
evolved and metal poor, the additional standard deviation of its rotation
period rises to 6.

\begin{figure}
  \caption{
    The additional rotation period scatter, $\sigma_P$, added to the
    observational period uncertainties in the model (see equation
    \ref{eqn:gyro}).
    The standard deviation was increased for early F and hotter
    stars (\gcolor $<$ 0.56), late M dwarfs (\gcolor\ $>$ 2.7) and evolved
    stars (EEP $\gtrsim$ 420) in order to down-weight the age-information
    supplied
    by rotation periods and reproduce observed rotation period distributions.
    We also increased the variance for stars younger than around 250 Myrs,
    because the rotation periods of these stars typically have not yet
    converged onto a tight gyrochronology sequence, and for very high and low
    metallicity stars (-0.2 $\gtrsim$ [Fe/H] $\gtrsim$ 0.2) because the
    gyrochronology relations have not been calibrated at these extreme values.
    Down-weighting the gyrochronal likelihood by the inverse variance
    ($1/\sigma^2$) allowed the ages of these stars to be mostly inferred
    via isochrone fitting.
    Sigmoid functions were used to provide smooth transitions between regions
    of low and high variance.
}
  \centering
    \includegraphics[width=1.\textwidth]{variance}
\label{fig:variance}
\end{figure}

\subsection{Simultaneously fitting gyrochronology and isochrones}
The previous part of this section describes the model for the mean and
variance of rotation periods as a function of their ages (and colors, EEPs,
masses and metallicities).
In what follows, we describe how this model was combined with a stellar
evolution model to infer stellar ages (and other parameters) via
gyrochronology and isochrone fitting simultaneously.
Our goal was to infer the age of a star from its observable properties by
estimating the posterior probability density function (PDF) over age,
\begin{equation} \label{eqn:eqn1}
    % p(t|{\bf m_x}, T_{\mathrm{eff}}, \log(g), \hat{F},
    p(t|{\bf m_x}, P_{\mathrm{rot}}, \bar{\pi}),
\end{equation}
where $t$ is age, ${\bf m_x}$ is a vector of
apparent magnitudes in various bandpasses,
\prot\ is the rotation period and \pmega\ is parallax.
Spectroscopic properties (\teff, \logg\ and $\mathrm{[Fe/H]}$) and/or
asteroseismic parameters (\dnu\ and \numax) may also be available for a star,
in which case they would appear to the right of the `$|$' in the above
equation since they are observables.
In order to calculate a posterior PDF over age, other stellar parameters must
be marginalized over.
These parameters are distance ($D$), V-band extinction ($A_V$), the
inferred metallicity, $[M/H]$\footnote{The inferred metallicity, [M/H] is a
model parameter which is different to the {\it observed} metallicity, [Fe/H]
which would appear on the right side of the $`|'$ in equation \ref{eqn:eqn1}.},
and equivalent evolutionary phase (abbreviated to EEP or
$E$).
EEP is a dimensionless number ranging from around 200 for M dwarfs up
to around 1600 for giants and is 355 for the Sun \citep[see][]{dotter2016,
choi2016}.
Stars are defined as subgiants when their EEP exceeds 454.
Mass is uniquely defined by EEP, age and metallicity.
The marginalization involves integrating over these extra parameters,
\begin{eqnarray} \label{eqn:bayes}
    % & p(A|{\bf m_x}, T_{\mathrm{eff}}, \log(g), \hat{F},
    & p(t|{\bf m_x},
    P_{\mathrm{rot}}, \bar{\pi})
\\ \nonumber
    % & \propto \int p({\bf m_x}, T_{\mathrm{eff}}, \log(g), \hat{F},
    & \propto \int p({\bf m_x},
    P_\mathrm{rot}, \bar{\pi}|
    t, E, [M/H], D, A_V)~p(t)\, p(E)\, p([M/H])\, p(D)\, p(A_V)\, \mathrm{d}E\,
    \mathrm{d}[M/H]\, \mathrm{d}D\, \mathrm{d}A_V.
\end{eqnarray}
This equation is a form of Bayes' rule,
\begin{equation} \label{eqn:eqn2}
\mathrm{Posterior} \propto \mathrm{Likelihood} \times \mathrm{Prior},
\end{equation}
where the likelihood of the data given the model is,
\begin{equation} \label{eqn:full_likelihood}
    % p({\bf m_x}, T_{\mathrm{eff}}, \log(g), \hat{F}, \bar{\pi},
    p({\bf m_x}, \bar{\pi}, P_{\mathrm{rot}}|t, E, D, A_V, [M/H]),
\end{equation}
and the prior PDF over parameters is,
\begin{equation} \label{eqn:prior}
    p(t)\, p(E)\, p(D)\, p(A_V)\, p([M/H]).
\end{equation}
The priors we used are described in the appendix.

We assumed that the process of magnetic braking is independent of hydrogen
burning in the core, outside of the dependencies that are captured in the
model.
This assumption allowed us to multiply two separate likelihood
functions together: one computed using an isochronal model and one computed
using a gyrochronal model.
We assumed that the probability of observing the measured observables given
the model parameters was a Gaussian and that the observables were identically
and independently distributed.
The isochronal likelihood function was,
\begin{eqnarray} \label{eqn:isochrones_only_likelihood}
    % & \mathcal{L_{\mathrm{iso}}} = p({\bf m_x}, T_{\mathrm{eff}}, \log(g),
    & \mathcal{L_{\mathrm{iso}}} = p({\bf m_x},
    \bar{\pi}|t, E, [M/H], D,
    A_V) \\ \nonumber
    & = \frac{1}{\sqrt{(2\pi)^n \det(\Sigma)}}
    \exp\left( -\frac{1}{2} ({\bf O_I} - {\bf I})^T \Sigma ^{-1}
    ({\bf O_I} - {\bf I})\right),
\end{eqnarray}
where ${\bf O_I}$ is the vector of n observables: \pmega, ${\bf m_x}$ plus
spectroscopic and/or asteroseismic observables if available, and $\Sigma$ is
the covariance matrix of that set of observables.
${\bf I}$ is the vector of {\it model} observables that correspond to a set of
parameters: $t$, $E$, $[M/H]$, $D$ and $A_V$, calculated using an isochrone model.
We assumed there is no covariance between these observables and so this
covariance matrix consists of individual parameter variances along the
diagonal with zeros everywhere else.
The gyrochronal likelihood function was,
\begin{eqnarray} \label{eqn:gyro_likelihood}
    & \mathcal{L_{\mathrm{gyro}}} = p(P_\mathrm{rot} |t, E, [M/H], D, A_V) \\ \nonumber
    & = \frac{1}{\sqrt{(2\pi) \det(\Sigma_P)}}
    \exp\left( -\frac{1}{2} ({\bf P_O} - {\bf P_P})^T \Sigma_P ^{-1}
    ({\bf P_O} - {\bf P_P})\right),
\end{eqnarray}
where ${\bf P_O}$ is a 1-D vector of observed logarithmic rotation periods,
and ${\bf P_P}$ is the vector of corresponding logarithmic rotation periods,
predicted by the model.
$\Sigma_P$ was comprised of individual rotation period measurement
uncertainties, plus an additional variance that is a function of EEP and
\gcolor\
color, added in quadrature.
This variance accounts for the stochastic nature of the rotation periods of
very hot and very cool stars and allowed us to predominantly use isochrone
fitting to measure the ages of subgiants.
The full likelihood used in our model was the product of these two likelihood
functions,
\begin{equation}
    \mathcal{L}_{\mathrm{full}} = \mathcal{L}_{\mathrm{iso}} \times
    \mathcal{L}_{\mathrm{gyro}}.
\end{equation}

The inference processes proceded as follows.
First, a set of parameters: age, EEP, metallicity, distance and extinction,
and observables for a single star were passed to the isochronal likelihood
function in equation \eqref{eqn:isochrones_only_likelihood}.
Then, a set of observables corresponding to those parameters were generated
from the MIST model grid using {\tt isochrones.py} \citep{isochrones} and
compared to the measured observables via the isochronal likelihood,
$\mathcal{L}_{\mathrm{iso}}$ (also computed using {\tt isochrones.py}).
The parameters were also passed to the gyrochronology model (equation
\ref{eqn:gyro}) where $t$, $E$,
$[M/H]$, $D$ and $A_V$ were used to calculate \gcolor\ color and mass from the
MIST model grid and, in turn, rotation period via the gyrochronology model.
EEP and \gcolor\ were also used to calculate the additional rotation period
variance, added to the individual period uncertainties.
This model rotation period was compared to the measured rotation period using
gyrochronal likelhood function of equation \ref{eqn:gyro_likelihood}.
The gyrochronal log-likelihood was added to the isochronal log-likelihood to
give the full likelihood, which was then added to the log-prior to produce a
single sample from the posterior PDF.

Ages were inferred with \sd\ using Markov Chain Monte Carlo.
The joint posterior PDF over age, mass, metallicity, distance and extinction
was sampled using the affine invariant ensemble sampler, {\tt emcee}
\citep{foreman-mackey2013} with 50 walkers.
Samples were drawn from the posterior PDF until 100 {\it independent} samples
were obtained.
We actively estimated the autocorrelation length, which indicates how many
steps were taken per independent sample, after every 100 steps using the
autocorrelation tool built into {\tt emcee}.
The MCMC concluded when {\it either} 100 times the autocorrelation length was
reached and the change in autocorrelation length over 100 samples was less
than 0.01, {\it or} the maximum of 500,000 samples was obtained.
This method is trivially parallelizable, since the inference process for each
star can be performed on a separate core.
The age of a single star can be inferred in around 1 hour on a laptop
computer.


\section{Quantitative Results}
The system outlined here performs quite well, even on somewhat difficult questions.
For example, for the questions ``Mary earns \$46 cleaning a home. How many homes did she clean, if she made 276 dollars?''
the system is able to correctly identify that ``\$46'' should be multiplied by ``x homes'' (remember that variables are instantiated by ``much-many\_a\_rel'' predicates, and they are associated with the ARG1 entity, here ``homes'').
Here, the system can use only the distance between the entity noun vectors and the verbs associated with each, as no other abstract predications obtain (besides ``much-many\_a\_rel'', common to all variables regardless of operation).

However, the system does still make many significant errors.
For the questions ``There are 8 calories in a candy bar. How many calories are there in 3 candy bars?''
the system erroneously predicts that ``x calories'' should be subtracted from ``3 candy bars'', whereas in a correct equation the latter would be divided by the former.
This error is in part due to the fact that ``3 candy bars'' is not associated with any verbal predication, and ``x calories'' is only associated with a relatively neutral ``\_be\_v\_there\_rel'' verbal predication.
However, according to the parse, the ARG1 of the ``\_in\_p\_rel'' associated with the candy bars is the ARG0 of the verbal predicate associated with the calories.
At present, our system is not capable of exploring such connections between event predicates, but we plan to improve this in future work.

\section{Conclusion}
This paper has presented a method for utilizing Minimal Recursion Semantics for improving a math word problem solving system.
We show that with limited adaptation to the math word problem domain, we can achieve a 3.0\% absolute improvement in operator prediction accuracy compared to a previous state-of-the-art semantic-based technique.
Furthermore, the techniques for relating entities across sentences discussed here are more applicable to other domains besides math word problems due to the general nature of the semantic representation over which they work.
Finally, we make our code available to the public to support further development of applications which take advantage of Minimal Recursion Semantics.


\documentclass[journal, letterpaper]{IEEEtran}
%\documentclass{scrartcl}

\usepackage[ngerman,english]{babel}
%\usepackage[latin1]{inputenc}
\usepackage[utf8]{inputenc}
\usepackage[T1]{fontenc}
\usepackage{amsmath}
\usepackage{amsthm}
\usepackage{amsfonts}
\usepackage{tikz}
\usepackage{verbatim}
\usepackage{subcaption}
\usepackage{algorithm}
\usepackage{algorithmic}
\usepackage[pdftex]{hyperref}

\renewcommand{\algorithmicrequire}{\textbf{Input:}}
\renewcommand{\algorithmiccomment}[1]{\ \ // #1} % C-like // Comments

\hyphenation{render}

% No clubs and widows allowed
\clubpenalty10000
\widowpenalty10000
\displaywidowpenalty=10000

\begin{document}

%\title{Simulating elastic spheres without external forces}
%\subtitle{Project 1 for class CS6491 Computer Graphics}
\title{Swirl\\
	{\large Project 2 for class CS6491 Computer Graphics}}
%\author{Sebastian Weiss}
\author{Sebastian Weiss, Kristian Eberhardson\\ \today}
%\date{\today}

\maketitle

\begin{tikzpicture}[remember picture,overlay]
   \node[anchor=north east,inner sep=0pt] at (current page.north east)
              {\includegraphics[scale=1]{pic}};
\end{tikzpicture}
\begin{tikzpicture}[remember picture,overlay]
   \node[anchor=north west,inner sep=10pt] at (current page.north west)
              {\includegraphics[scale=0.4]{pictures/P4.png}};
\end{tikzpicture}

\section{Objective}
We are provided with a starting and ending frame in 3D, called $A_0$ and $A_1$. We wish to create a continuous, affine and steady motion which interpolates $A_0$ and $A_1$.

\section{Input}
The Frames $A_0$ and $A_1$ consist of a point $P_i$ and a right hand coordinate system of orthogonal vectors $I_i$, $J_i$, and $K_i$, each of the same length ($i \in \{0,1\}$). All vectors are three dimensional.
\subsection{Terms}
1. An \underline{affine} transformation in three dimensions can be expressed as a 4x4 matrix multiplied onto a point or a vector, is reversible, and can be expressed solely as a combination of rotation, translation, scaling and shearing. In this work, we only concentrate on rotation, translation and scaling.
\newline
2. A \underline{continuous} transformation is one in which $\forall \epsilon>0 \\ \exists \delta>0 \; : \; |P_x - P_y|<\epsilon \; \forall x,y \; \text{with} \; |x-y|<\delta$. $P_x$ and $P_y$ denotes the interpolated point at time $x$ and $y$. Intuitively, this means we can draw the motion with a pen in a single stroke.
\newline
3. A \underline{steady} transformation is independent of the time t and the velocities are constant over time.

\section{Related interpolation schemes}

\subsection{Logarithmic spiral in 2D}
In 2D, a frame is given by three two-dimensional vectors. Prof. Rossignac proposes a combination of a linear rotation and an exponential scaling to realize a steady interpolation between two dimensional frames, see Fig. \ref{fig:P6}.
The formula for this approach is the following:
\begin{equation}
 P(t) := F + m^t FP_0 ^{\;\circ} (\alpha t)
\label{eq:logSpiral}
\end{equation}
$FP_0 ^{\;\circ} (\alpha t)$ denotes the normal 2d rotation of the vector $FP_0$ around the origin with the angle $\alpha t$.
\begin{figure}
	\centering
		\includegraphics[scale=0.4]{pictures/P6.png}
	\caption{log spiral in 2D}
	\label{fig:P6}
\end{figure}

\subsection{Linear interpolation in 3D}
For all points in $A_0$, draw a straight line to its corresponding point in $A_1$. Motion is simply moving the frame along these lines.  This is neither an elegant solution nor is it steady.

\section{Proposed model}
We wish to create a swirling motion in 3D related to the logarithmic spiral in 2D. Therefore, we are using a combination of linear rotation around the rotation axis $N$ with the origin in $F$ and an exponential scaling with the factor $m$. We propose the following formula:
\begin{equation}
 P(t) := F + m^t FP_0 ^{\;\circ} (\alpha t; N)
\label{eq:Interpolation}
\end{equation}
Since the rotation and scaling are steady on its own and these operations commute, the final motion is also steady.

\subsection{3D rotation}
A rotation in 2D is very straightforward because there is no choice of the rotation axis. In 3D we are given the rotation axis and rotate around it in the following way:

\begin{equation}
\begin{array}{lcl}
 X^{\circ}(\alpha; N) &:=& W + U^{\circ}(\alpha; N) \\
                        & =& W + \cos \alpha U - \sin \alpha (N \times U) \\
 \multicolumn{3}{l}{\text{with } W := W\angle\underline{N} = (W\bullet\underline{N})\underline{N}} \\
 \multicolumn{3}{l}{\text{and } U = X-W}
\end{array}
\label{eq:Rotation}
\end{equation}

Fig. \ref{tikz:Rot} visualizes the rotation around the axis $N$ in 3D.

\begin{figure}[H]
\centering
\begin{tikzpicture}
	\filldraw[color=white,fill=blue!20,fill opacity=0.5] (-3,-1.5) -- (2,-1.5) -- (3,1.5) -- (-2,1.5) -- cycle;
	\draw[color=orange] (-0.5,0) arc (-180:-61.7:0.5);
	\draw[line width=1mm,color=black!60,->] (0,0) -- (-1.5,0);
	\draw[line width=1mm,color=black!60,->] (0,0) -- (0.7,-1.3);
	\draw[line width=1mm,dotted,color=black!50] (-1.5,0) -- (-1.5,1);
	\draw[line width=1mm,dotted,color=black!50] (0.7,-1.3) -- (0.7,0.5);
	\draw[line width=1.5mm,color=green,->,line cap=round] (0,0) -- (-1.5,1);
	\draw[line width=1.5mm,color=red,->,line cap=round] (0,0) -- (0.7,0.5);
	\draw[line width=1.5mm,color=black,->,line cap=round] (0,0) -- (0,2);
	\node at (0.4,2) {$N$};
	\node at (-0.7,1) {$FP_0$};
	\node at (1.7,0.7) {$FP_1=FP_0 ^{\;\circ} (\alpha; N)$};
	\node at (-1.8,0.5) {$W$};
	\node at (1,-0.4) {$W$};
	\node at (-0.75,-0.3) {$U$};
	\node at (0.1,-0.9) {$U$};
	\node[color=orange] at (-0.35,-0.6) {$\alpha$};
\end{tikzpicture}
\caption{3D-Rotation}
\label{tikz:Rot}
\end{figure}

\subsection{Calculating $N$, $\alpha$ and $m$}
Since $m$ is the scaling factor, it is calculated as:
\begin{equation}
 m := |I_1| / |I_0|
\label{eq:m}
\end{equation}
Because $I_i$, $J_i$, and $K_i$ are of the same length for each $i$, it does not matter which we use.

To calculate the rotation axis $N$, we first define
\begin{equation}
\Delta I:= I_1 - I_0, \ \Delta J := J_1 - J_0, \ \Delta K := K_1 - K_0
\end{equation}
Then we pick the largest vector in the euclidean norm out of
\begin{equation}
N_a := \Delta I \times \Delta J, \ N_b := \Delta I \times \Delta K, \ N_c := \Delta J \times \Delta K
\end{equation}
We call the largest vector $N$. Since $I_i$, $J_i$ and $K_i$ form an orthogonal basis, these vectors all point in the same direction and differ only in the scaling. The reasons for choosing the largest one are numerically issues. Some of these $N$'s can be very small when the rotation axis is almost parallel to the frame axes, they are zero if they are exactly parallel. The special case when all $N$'s are zero, so $\Delta I$, $\Delta J$ and $\Delta K$ are all zero, happens if there is no rotation at all. See section \ref{NoRot} on how to deal with this case. In the later computation, we always need the normalized vector $\underline{N}:=N/|N|$.

Furthermore, we need the rotation angle $\alpha$. To compute this angle, we first project each basis vector into the plane defined by $\underline{N}$ and normalize them:
\begin{equation}
\begin{array}{lcl}
 proj(X) &:=& X - X\angle\underline{N} = X - (X \bullet \underline{N})\underline{N} \\
 norm(X) &:=& X / |X| \\
 B' &:=& norm(proj(B_i)) \ \forall B \in \{I_0,...,K_1\}
\end{array}
\label{eq:ProjBasis}
\end{equation}
The $B'$-vectors correspond to normalized versions of the $U$-vectors in Fig. \ref{tikz:Rot}.
Then we define
\begin{equation}
 \alpha = \max \{ \cos^{-1}(I_0' \bullet I_1'), \cos^{-1}(J_0' \bullet J_1'), \cos^{-1}(K_0' \bullet K_1') \}
\label{eq:}
\end{equation}
We take the maximum angle here because all angles are equal by the definition of $N$. At most one of these angles can be zero, this happens in the case if the rotation axis is exactly orthogonal to that frame axis, see Fig. \ref{fig:P1}.
\begin{figure}
	\centering
		\includegraphics[scale=0.4]{pictures/P1.png}
	\caption{rotation parallel to a frame axis}
	\label{fig:P1}
\end{figure}

\subsection{Calculating $F$}
Now we have $N$, $\alpha$ and $m$ and we can compute $F$ now by solving the equation system
\begin{equation}
 P_1 = P(1) := F + m FP_0 ^{\;\circ} (\alpha; N)
\label{eq:Feq}
\end{equation}
This system is a linear system in the three variables $F_x$, $F_y$ and $F_z$. We obtain an explicit formulation for $F$ by converting it into matrix form, solving it using Cramer's Rule and applying simplification. For purposes of notation, we introduce the following variables for the vector components: $P_0=(P_{0,x},P_{0,y},P_{0,z})^T, P_1=(P_{1,x},P_{1,y},P_{1,z})^T,\\ \underline{N}=(N_x,N_y,N_z)^T$.
\begin{equation}
\begin{array}{lcl}
	F_x &=& (m P_{0,x} + m^3 P_{0,x} - P_{1,x} - m^2 P_{1,x} \\
		&& + m (1+m) (\cos\alpha - 1) (P_{0,x} - P_{1,x}) N_y^2 \\
		&& + m (1+m) (\cos\alpha - 1) (P_{0,x} - P_{1,x}) N_z^2 \\
		&& - m (1+m) (\cos\alpha - 1) (P_{0,y} - P_{1,y}) N_x N_y \\
		&& - m (1+m) (\cos\alpha - 1) (P_{0,z} - P_{1,z}) N_x N_z \\
		&& - m (m-1) \sin\alpha (P_{0,y} - P_{1,y}) N_z \\
		&& + m (m-1) \sin\alpha (P_{0,z} - P_{1,z}) N_y \\
		&& - 2 m^2 P_{0,x} \cos\alpha + 2 m P_{1,x} \cos\alpha ) \\
		&& / (m-1)(1 + m^2 - 2m\cos\alpha)
\end{array}
\label{eq:Fx}
\end{equation}
\begin{equation}
\begin{array}{lcl}
 F_y &=& -(m^2 P_{0,y} - m^3 P_{0,y} + P_{1,y} - m P_{1,y} \\
		&& + m (1+m) (\cos\alpha - 1) (P_{0,x} - P_{1,x}) N_x N_y \\
		&& + m (1+m) (\cos\alpha - 1) (P_{0,y} - P_{1,y}) N_y^2 \\
		&& + m (1+m) (\cos\alpha - 1) (P_{0,z} - P_{1,z}) N_y N_z \\
		&& - m (m-1) \sin\alpha (P_{0,x} - P_{1,x}) N_z \\
		&& + m (m-1) \sin\alpha (P_{0,z} - P_{1,z}) N_x \\
		&& - m P_{0,y} \cos\alpha + m^2 P_{0,y} \cos\alpha \\
		&& - m P_{1,y} \cos\alpha + m^2 P_{1,y} \cos\alpha) \\
		&& / (m-1)(1 + m^2 - 2m\cos\alpha)
\end{array}
\label{eq:Fy}
\end{equation}
\begin{equation}
\begin{array}{lcl}
 F_z &=& -(m^2 P_{0,z} - m^3 P_{0,z} + P_{1,z} - m P_{1,z} \\
		&& + m (1+m) (\cos\alpha - 1) (P_{0,x} - P_{1,x}) N_x N_z \\
		&& + m (1+m) (\cos\alpha - 1) (P_{0,y} - P_{1,y}) N_y N_z \\
		&& + m (1+m) (\cos\alpha - 1) (P_{0,z} - P_{1,z}) N_z^2 \\
		&& - m (m-1) \sin\alpha (P_{0,y} - P_{1,y}) N_x \\
		&& + m (m-1) \sin\alpha (P_{0,x} - P_{1,x}) N_y \\
		&& - m P_{0,z} \cos\alpha + m^2 P_{0,z} \cos\alpha \\
		&& - m P_{1,z} \cos\alpha + m^2 P_{1,z} \cos\alpha) \\
		&& / (m-1)(1 + m^2 - 2m\cos\alpha)
\end{array}
\label{eq:Fz}
\end{equation}
For the implementation, we strongly suggest to further factorize these formulas and to store the common intermediate results.

\subsection{Dealing with special cases}
The proposed model contains two singularities in the computations:
First, $N$ is the zero vector when the rotation between the two input frames is zero.
Second, $F$ goes to infinity when $m$ goes to one, since the translation of the frames rely on the scaling.

\subsubsection{No translation}
Strictly speaking, this case is no special case at all. When there is no translation, $P_0$, $P_1$ and $F$ are all at the same position, see Fig. \ref{fig:P8}. This model handles this case without further additions.
\begin{figure}
	\centering
		\includegraphics[scale=0.4]{pictures/P8.png}
	\caption{no translation}
	\label{fig:P8}
\end{figure}

\subsubsection{No rotation, but scaling}\label{NoRot}
In the case that there is no rotation, the rotation axis $N$ is the zero vector and the rotation angle $\alpha$ is also zero.
The translation is then realized by the scaling, see Fig. \ref{fig:P2}. The fixpoint $F$ is calculated as usual. Intuitively, this is the intersection point of the lines connecting the tips of the arrows and the frame origin from the start frame to the end frame.
\begin{figure}
	\centering
		\includegraphics[scale=0.4]{pictures/P2.png}
	\caption{no rotation, translation by scaling}
	\label{fig:P2}
\end{figure}

\subsubsection{No rotation, no scaling}
If we start with the case presented in the previous section and let $m$ run to 1, the fixpoint $F$ runs against infinity. We therefore observe that the case of no rotation and no scaling degenerates to a linear interpolation with pure translation, see Fig. \ref{fig:P5}.
\begin{figure}
	\centering
		\includegraphics[scale=0.4]{pictures/P5.png}
	\caption{pure translation}
	\label{fig:P5}
\end{figure}

\subsubsection{No scaling, but rotation and translation}
If the scaling factor $m$ is 1, $F$ goes to infinity. The swirl degenerates to a circular movement in the rotation plane plus a linear translation orthogonal to the rotation plane. The swirl is therefore a perfect spiral, see Fig. \ref{fig:P7}.
\begin{figure}
	\centering
		\includegraphics[scale=0.3]{pictures/P7.png}
	\caption{no scaling, but rotation and translation}
	\label{fig:P7}
\end{figure}
This special case can be implemented by projecting the frames in the rotation plane, solving the circular motion there equivalent to the Logarithmic Spiral in 2D and then adding a translation orthogonal to the rotation plane.

\section{Implementation}
We provide a framework for playing around with our Swirl Implementation. You can rotate the camera and scale it. Furthermore, the start and end frame can be modified by clicking and dragging the center of the frames or the tips of the frame axes. Additionally, you can select if you want to display the rotation axis $N$ (gray line) and the fixpoint $F$ (gray sphere on the rotation axis) and if the extrapolating frames should be drawn or not.

A typical swirl produced by our program is shown in Fig. \ref{fig:P3}. You can see how the motion is attracted by the fixpoint $F$.

\begin{figure}
	\centering
		\includegraphics[scale=0.3]{pictures/P3.png}
	\caption{swirl motion is attracted by the fixpoint $F$}
	\label{fig:P3}
\end{figure}

\section{Future work}
As an extension to this work, we could include shearing in the interpolation model. With shearing, every affine transformation would be supported.

\end{document}


\documentclass{article}
\usepackage{tocloft}
\include{common_symbols_and_format}
\renewcommand{\cfttoctitlefont}{\Large\bfseries}

\begin{document}
\logo
\rulename{Percentage Volume Oscillator} %Argument is name of rule
\tblofcontents

\ruledescription{The percentage volume oscillator (PVO) is a momentum indicator that measures the difference between two moving averages as a percentage of the larger moving average. The fast moving average uses a lower period, and is thus more reactive to changes in volume. The slow moving average uses a faster period and is therefore less reactive. The PVO is positive when the shorter volume EMA is above the longer volume EMA and negative when the shorter volume EMA is below. This indicator can be used to define the ups and downs for volume, which can then be used to confirm or refute other signals. A bullish reversal of an asset is identify when the PVO cross above zero line. And a bearish reversal when the PVO cross below the zero line.
}

\howtotrade
{The strategy is to identify Bullish and Bearish Reversal.
Bullish Reversal - when PVO is above zero \&
Bearish Reversal - when PVO is below zero.
}

\ruleparameters %You can include however many arguments (in groups of 4) as you want!
{Short term look back Length}{12}{Short term look back length used to compute EMA.}{$\lookbacklength_{s}$}
{Long term look back Length}{26}{Long term look back length used to compute EMA.}{$\lookbacklength_{l}$}
{Signal look back Length}{9}{Look back length used to generate Signal line.}{$S_{l}$}
\stoptable %must be included or Tex engine runs infinitely

\newpage
\section{Equation}
Below are the equations which govern how this specific trading rule calculates a trading position.

\begin{equation}
    PVO = \frac{EMA(\lookbacklength_{s}) -         EMA(\lookbacklength_{l})}{EMA(\lookbacklength_{l})} \times 100
\end{equation}
\begin{equation}
    Signal = EMA(S_{l})
\end{equation}
\begin{equation}
    PVO_{hist} = PVO - Signal
\end{equation}
\\ % creates some space after equation
where:

$EMA(\lookbacklength_{s})$: is the short term exponentially weighted volume average.

$EMA(\lookbacklength_{l})$: is the long term exponentially weighted volume average.

$EMA(S_{l})$: is the exponentially weighted volume average computed to generate signal line.

$PVO_{hist}$: is the Percentage Volume Oscillator histogram.

\keyterms
\furtherlinks %The footer
\end{document}


\documentclass{article}
\usepackage{color}
\usepackage{listings}
\usepackage[T1]{fontenc}
%\usepackage[dvipsnames,table]{xcolor}

\lstdefinelanguage{MiniZinc}
{	morekeywords={var,morph,let,in,pareto,lex,par,offers,int,solve,output,not,search,satisfy,new,endif,maximize,params,instantiates,with,bool,in,type,PVS,PVSType,minimize,float,constraint,soft,sum,forall,exists,array,of,include,predicate,then,commit,post,set,function,if,else,repeat,next,ann,break},
	sensitive=false,
	morecomment=[l]{\%},
	morecomment=[s]{/*}{*/},
	morestring=[b]",
	morestring=[b]',
}

\lstset
{   stringstyle=\ttfamily\color{blue},
	tabsize=2,
	showstringspaces=false,
	numberfirstline=false,
	%flexiblecolumns=true,
	captionpos=b,
	backgroundcolor=\color{lightlightgray},
	frame=single,
	%xleftmargin=\parindent,
	numberstyle={\ttfamily\fontsize{6pt}{6pt}\selectfont}
}
\lstset{language=MiniZinc}

\def\mystyle{}

\begin{document}

\Css{div.lstlisting .ecit-1000 {font-family: monospace;color:darkred; font-style:normal;}}
\Css{div.lstlisting .ectt-1000 {font-family: monospace;color:black}}
\Css{div.lstlisting .ecss-1000 {font-family: monospace;color:red}}
\Css{div.lstlisting .ecbx-1000 {font-family: monospace;color:forestgreen}}
\Css{div.lstlisting { background-color: \#F0F0F0; border: 1px solid black; padding: 5px;} }

\section{Test!}

\begin{lstlisting}
include "classic_o.mzn"; % output of minibrass
include "soft_constraints/pvs_gen_search.mzn";

% the basic, "classic" CSP
set of int: DOM = 1..3;

var DOM: x; var DOM: y; var DOM: z;

solve
:: int_search([x,y,z], input_order, indomain_min, complete)
search pvs_BAB();

output ["x = \(x); y = \(y); z = \(z)"] ++
       [ "\nValuations:  overall = \(topLevelObjective)\n"];

\end{lstlisting}

\begin{lstlisting}
include "defs.mbr";

PVS: cr1 = new ConstraintPreferences("cr1") {
   soft-constraint c1: 'x + 1 = y';
   soft-constraint c2: 'z = y + 2';
   soft-constraint c3: 'x + y <= 3';

   crEdges : '[| mbr.c2, mbr.c1 | mbr.c3, mbr.c1 |]';
   useSPD: 'false' ;
};

solve cr1;
\end{lstlisting}

\end{document}


\documentclass{article}
\usepackage[utf8]{inputenc}

\usepackage{amsmath}
\usepackage{graphicx}
\usepackage{xcolor}
\usepackage[export]{adjustbox}[2011/08/13]
\usepackage{float}
\usepackage{hyperref}
\usepackage{setspace}
\usepackage{fullpage}

\hypersetup{colorlinks=true}

% shortcuts for covariant basis vectors
\newcommand{\er}{{\mathbf e}_{\rho}}
\newcommand{\ev}{{\mathbf e}_{\vartheta}}
\newcommand{\ez}{{\mathbf e}_{\zeta}}

\bibliographystyle{ieeetr}

\title{APC 524 Design Document}
\author{Daniel Dudt, Dario Panici, Evan Yerger}
\date{December 15, 2020}

\begin{document}

\maketitle

\section{Background}

\subsection{Motivation}

Determining the equilibrium of a plasma is crucial for the design and operation of magnetic confinement fusion reactors.
Equilibrium calculations are used for understanding basic plasma physics, interpreting diagnostic data from experiments, running real-time control systems to stabilize the plasma, and for optimizing the design of future machines.
Most mainstream magnetic confinement concepts for controlled fusion involve toroidal reactors, such as the tokamak and stellarator.
The tokamak is characterized by toroidal symmetry, while the stellarator does not have an ignorable coordinate and is fully ``three-dimensional''.
The nonlinear partial differential equations that describe an equilibrium are computationally challenging to solve, especially for the complicated geometries of stellarators.
Existing codes such as VMEC \cite{Hirshman1983} are expensive to run and do not always converge to the desired result, making them poorly equipped for future progress.
DESC \cite{Dudt2020} is a modern code designed to meet the increasing demands of advanced stellarator performance, and the goal of this project is to make significant contributions to the development of DESC.

\subsection{Theory}
\label{sec:theory}

Ideal magnetohydrodynamics (MHD) is a single-fluid model of plasma.
It describes the equilibrium of a static plasma through a force balance equation, along with Amp\`ere's and Gauss's laws:
%
\begin{subequations}
  \label{eq:equil}
  \begin{align}
    \label{eq:momentum}
    \mathbf{J} \times \mathbf{B} &= \nabla p \\
    \label{eq:Ampere}
    \nabla \times \mathbf{B} &= \mu_0 \mathbf{J} \\
    \label{eq:Gauss}
    \nabla \cdot \mathbf{B} &= 0.
  \end{align}
\end{subequations}
%
Here $\mathbf{J}$ is the current density, $\mathbf{B}$ is the magnetic field, $p$ is the plasma pressure, and $\mu_0$ is the magnetic constant.
Combining the force balance (\ref{eq:momentum}) and  Amp\`ere's Law (\ref{eq:Ampere}), a force balance error $\mathbf{F}$ can be defined as
%
\begin{equation}
  \label{eq:F}
  \mathbf{F} \equiv \frac{1}{\mu_0} \left( \nabla \times \mathbf{B} \right) \times \mathbf{B} - \nabla p = F_\rho \nabla\rho + F_\beta \mathbf{\beta} = \mathbf{0}
\end{equation}
%
where $\mathbf{\beta} \equiv B^\zeta \nabla \vartheta - B^\vartheta \nabla \zeta$.
The flux coordinate system $(\rho,\vartheta,\zeta)$ is a specific choice that makes the magnetic field lines ``appear straight'' to simplify calculations, and is related to the usual toroidal coordinates $(R,\phi,Z)$ as shown in Figure \ref{fig:coords}.
Equation (\ref{eq:F}) is a system of nonlinear partial differential equations (PDEs) and it must be satisfied throughout the entire plasma volume when in equilibrium.
Using Gauss's law (\ref{eq:Gauss}) and assuming that nested flux surfaces exist $\mathbf{B}\cdot\nabla\rho=B^\rho=0$, the magnetic field can be expressed in the following contravariant form of the flux coordinate system $(\rho,\vartheta,\zeta)$:
%
\begin{equation}
  \label{eq:B}
  \mathbf{B} = B^\vartheta \ev + B^\zeta \ez = \frac{\partial_\rho \Psi}{2\pi \sqrt{g}} \left( \iota \ev + \ez \right).
\end{equation}
%
Here $\Psi$ is the toroidal magnetic flux and $\iota\equiv d\vartheta/d\zeta = B^\vartheta/B^\zeta$ is the rotational transform.
The covariant basis vectors $\ev$ and $\ez$ and the jacobian of the coordinate system $\sqrt{g}$ are known from the shapes of the flux surfaces: $R(\rho,\vartheta,\zeta)$ and $Z(\rho,\vartheta,\zeta)$.
Solving the force balance given in (\ref{eq:F}) subject to a magnetic field of the form given in (\ref{eq:B}) will satisfy all of the MHD equilibrium conditions provided in (\ref{eq:equil}).
The actual scalar equations that get minimized are
%
\begin{subequations}
  \label{eq:f}
  \begin{align}
    f_\rho(R,Z) &= F_\rho \lVert\nabla\rho\rVert_2 \sqrt{g} \Delta\rho\Delta\vartheta\Delta\zeta \text{sign}\left(\nabla\rho\cdot\er\right) \\
    f_\beta(R,Z) &= F_\beta \lVert\mathbf{\beta}\rVert_2 \sqrt{g} \Delta\rho\Delta\vartheta\Delta\zeta \text{sign}\left(\mathbf{\beta}\cdot\ev\right) \text{sign}\left(\mathbf{\beta}\cdot\ez\right)
  \end{align}
\end{subequations}
%
along with additional equations to enforce the boundary conditions.
This is known as an ``inverse formulation'' because the toroidal coordinates are solved for in terms of the flux coordinates, which are treated as the independent variables.
Calculating the equilibrium magnetic field is equivalent to determining the transformation between these two coordinate systems.

The PDEs are discretized using pseudospectral methods.
The flux surfaces are represented by the coefficients of a global Fourier-Zernike basis set of the form:
%
\begin{subequations}
	\begin{align}
	\label{eq:R_basis}
	R(\rho,\vartheta,\zeta) &= \sum_{n=-N}^{N} \sum_{m=-M}^{M} \sum_{l\in L} R_{lmn} \mathcal{Z}^{m}_{l}(\rho,\vartheta) \mathcal{F}^{n}(\zeta) \\
	\label{eq:Z_basis}
	Z(\rho,\vartheta,\zeta) &= \sum_{n=-N}^{N} \sum_{m=-M}^{M} \sum_{l\in L} Z_{lmn} \mathcal{Z}^{m}_{l}(\rho,\vartheta) \mathcal{F}^{n}(\zeta)
	\end{align}
\end{subequations}
%
where $\mathcal{Z}^{m}_{l}$ is the Zernike polynomial with radial mode $l$ and azimuthal mode $m$, $\mathcal{F}^{n}$ is a Fourier component with wave number $n$, and $L = |m|, |m|+2, |m|+4, \ldots, 2 M - |m|$.
From the state vector of coefficients $\mathbf{x} = [R_{lmn}, Z_{lmn}]^T$, the values and partial derivatives of $R(\rho,\vartheta,\zeta)$ and $Z(\rho,\vartheta,\zeta)$ are transformed to physical space at a set of collocation points.
All of the nonlinear calculations involved to compute (\ref{eq:f}) are performed in physical space at these nodes, resulting in a residual vector $\mathbf{f} = [f_\rho, f_\beta]^T$.
Starting from an initial guess, DESC computes the equilibrium solution by finding the flux surface geometry that minimize these errors on a given grid of collocation points: $\mathbf{f}(\mathbf{x}) \approx \mathbf{0}$.
This optimization process is performed with a quasi-Newton method, and will be referred to as the ``inner loop'' of the DESC algorithm.
The ``outer loop'' performs a sequence of these optimizations with different input parameters, and can perturb the previous solution to give a good initial guess for the next optimization step with the new inputs.
For example, this outer loop can be used to increase the numerical resolution from a crude initial solution to one with many spectral modes.
Another application is to perform scans in solution space over a physics parameter such as pressure, by starting from a vacuum solution and then solving for the equilibrium at increasingly higher pressures.

\begin{figure}
	\centering
	\includegraphics[width=0.8\linewidth,center]{./figs/coordinates.pdf}
	\caption{Toroidal coordinate system $(R,\phi,Z)$ and the flux coordinates $(\rho,\vartheta,\zeta)$.}
	\label{fig:coords}
\end{figure}

\section{Project Goals}

\subsection{Initial State}

Before undertaking this project, DESC already existed as an open-source Python software package.
The user interface was to create an input file detailing all of the solver parameters, and then pass that input file as a command line argument.
The user could also pass flags to plot the results, but this output was only preset routines without much opportunity for customization.
Internally, the DESC algorithm was implemented sequentially with a functional rather than class structure.
Although the functions were well documented, each had a unique call signature and there was a lack of a common interface between different parts of the code.
New features were added to the code by creating another option in the input file format, and handling the new cases with if-else logic in the main driver script.

DESC relies on other software packages to outsource some functionality.
The NumPy \cite{NumPy} library is used to handle and vectorize all array operations in the system of equations.
Initially, all of the optimization routines for the inner loop were provided by the SciPy \cite{SciPy} optimization library.
Many of these routines rely on information of how the objective function changes with respect to state variables, which is encoded in the Jacobian matrix $\frac{\partial\mathbf{f}}{\partial\mathbf{x}}$.
This information is also needed to perform the outer loop perturbations, so it is important for DESC to compute the Jacobian matrix quickly and accurately.
This is accomplished with the use of JAX \cite{JAX}, an open source machine learning package for Python.
JAX provides automatic differentiation of arbitrary Python functions by overloading NumPy operations, as well as just-in-time (JIT) compilation and optimization for speed-up on repeated function calls.
It also allows for operations to be computed on a GPU, which can greatly accelerate the matrix-vector operations necessary for the optimization algorithm.
The initial state of the code had JAX implemented, but was disorganized about when it should be used over regular NumPy operations.

\subsection{Goals}

While DESC was functional in its initial state, it did not have a unified or high-level class structure organization to it.
It also was not extensively covered by testing, had not been thoroughly profiled for speed optimization, and lacked certain features that would make it more useful to the user.
With those shortcoming in mind, the aim of this project was to improve DESC in the following ways:
%
\begin{enumerate}
\item Refactor the code into a modular class structure
\item Expand the coverage of the existing testing suite
\item Profile the code to identify bottlenecks and explore optimization options
\item Extend the plotting capabilities to provide more solution visualization capabilities
\end{enumerate}

As mentioned in the previous section, the previous version of DESC had very minimal class structure.
This resulted in cumbersome control logic and unclear interfaces between different parts of the code.
Our primary goal of this project is to refactor DESC into a class structure that takes advantage of object-oriented programming.
This improvement will streamline the code and provide a consistent application programming interface (API) for adding new functionality as users expand the software to meet additional applications.
The new modularity will also make it easier to perform unit testing to ensure that the code is implemented properly and returns the desired output.
The existing continuous integration (CI) automated testing suite only covered 22\% of the DESC code, and our goal was to increase this converage to over 50\%.
Computational efficiency is also essential for DESC, since it was developed in the hope of finding stellarator equilibria faster than other codes.
Once the major refactoring was complete, we planned to profile the code to identify the bottlenecks and pursue any opportunities for optimization.
Finally, we also intended to create a class responsible for plotting.
Accurate solutions are only useful if they can be viewed, and the goal of this effort was to provide the user with a more flexible and interactive analysis tool.

\section{Design}

\subsection{Software Architecture}

The final architecture of the refactored DESC code is represented by the Unified Modeling Language (UML) diagram in Figure \ref{fig:DESC_UML}.
At the core of the design is the \texttt{Configuration} class, which acts as a container to store all of the spectral coefficients that describe a plasma state and their corresponding spectral basis sets.
The \texttt{Equilibrium} class inherits from \texttt{Configuration}, and uses the decorator pattern to add more information about how the equilibrium solution was solved such as the objective function that was minimized and the optimization method used.
The \texttt{Equilibrium} attributes get updated when the equilibrium is solved in each inner loop, but it also contains another \texttt{Configuration} attribute to hold the initial guess.
\texttt{EquilibriaFamily}, which inherits from \texttt{MutableSequence}, is used as a container to store the \texttt{Equilibrium} for each outer loop iteration.

\texttt{ObjectiveFunction} is an Abstract Base Class (ABC) that represents the objective function $\mathbf{f}(\mathbf{x})$ to be minimized.
Section \ref{sec:theory} only outlined one such function for the equilibrium force error, but there could be other equivalent definitions for equilibrium (such as minimizing the plasma energy) or other physical quantities of interest that need to be optimized.
The \texttt{ObjectiveFunctionFactory} provides a factory design method to systematically determine which child of \texttt{ObjectiveFunction} to use.
This objective function would be minimized by an optimization object with a very similar structure: an \texttt{OptimizerFactory} will return an instance of an ABC \texttt{Optimizer} class that represents different optimization methods.
This has not been implemented in the code yet because the SciPy optimization library is still being used, but the planned design is included in the UML in red.

Another major component of the design is the \texttt{Transform} class, which is responsible for transforming spectral coefficients to real space and fitting data to a spectral basis.
This is needed whenever the \texttt{Configuration} data is used, such as during the evaluation of an objective function or while plotting the state of the plasma.
A \texttt{Transform} object is defined by a \texttt{Grid} and a \texttt{Basis}.
The \texttt{Basis} is an ABC that represents a set of basis functions such as the Fourier-Zernike basis from (\ref{eq:R_basis}) and (\ref{eq:Z_basis}), a double Fourier series to represent the boundary surface, or a power series to describe the pressure profile.
The \texttt{Grid} represents a set of collocation points, and is intended to act as an abstract class but can be instatiated directly to create a grid of arbitrary nodes.

The outer loop is designed to be handled by the \texttt{EquilibriumSolver} class, which will be responsible for calling the factory methods at each iteration, updating the transforms as the resolution changes, and perturbing the solutions between iterations.
This functionality still lives in a script, but has been streamlined to work with the new interface and will eventually be updated into the proposed class structure.
There are several other classes that are necessary to achieve the desired modularity of the code, and can be seen in Figure \ref{fig:DESC_UML}.

\subsection{Input \& Output}

A new interface for the input and output (IO) of these class objects was also developed.
This design is particularly useful for saving and loading the solutions contained in an \texttt{EquilibriaFamily}, but the interface is intentionally general for any class object.
Any object we want to make savable inherits from the ABC \texttt{IOAble}, which includes methods to save and initialize the object from a file.
Child classes need only specify a method to initialize from initialization arguments (i.e. initializing from scratch), a list of attributes to save and load, and a dictionary of classes it may need to initialize on the fly during a load from file.
This architecture was chosen after our initial idea resulted in circular import errors:
either the objects could depend on the IO in some way or the IO could depend on the objects in some way, but both could not happen while keeping the IO functions and objects in different modules.
We decided to have the objects inherit from the IO, as it would allow for a homogeneous and simple interface when saving and loading objects.

We also altered the IO back end code significantly towards a number of goals: a homogeneous interface across file formats, protecting read/write permissions, and less code reproduction.
Starting with an ABC \texttt{IO}, which includes methods handling how all files will be treated, we specify that files should be closed on garbage collection, and files of the specified type should be opened if the class constructor is passed a file path.
Classes specifying IO operations specific to file-format, like opening files and creating sub directories inherit directly from \texttt{IO}.
Parallel to this, we introduced the abstract base classes \texttt{Reader} and \texttt{Writer}, which specify the interface for any \texttt{IO} class that reads from or writes to file, respectively.
The utility of this inheritance structure, shown in Figure \ref{fig:EquilIO_UML}, is that one only needs to specify a file format wrapper for \texttt{IO} and read or write functions to have a reader or writer that can be hot-swapped for any other.

\subsection{User Interface}

The user interface is generally split into two parts: one computes an equilibrium from command line arguments; the other to plot solutions from saved files (generated by the solver).
When calculating a solution, the user must call the program with the input filename as an argument.
Other optional arguments may be included and are detailed in the existing documentation.
The command line arguments and input file are parsed, and the required instances of the class objects will be created in order to solve.
A solution algorithm is called as a method of the EquilibriaFamily class, which creates new Equilibrium objects inside of it and saves aspects of the solution to the output file, as well as uses the linked EquilibriumSolver object to solve the problem.
Another envisioned use of the code is in an interactive sense, where the user could instantiate an Equilibrium or EquilibriaFamily from an input file, and then use the solve method of these classes to compute the solutions.

Once a solution has been found and an output file saved, the user can then instantiate a \texttt{Plot} class with the name of the output file.
Methods of this class are called to plot specified parts of the solution (magnetic field, flux surfaces, etc) for specified domains (1-D profiles, 2-D plots at a given toroidal cross-section, etc) at a specified iteration of the outer loop.
In this way, the user can use the Plot class to visualize not only the final result but also to see intermediate solutions with different input parameters, which are stored as Equilibrium objects inside of an EquilibriaFamily.

\begin{figure}[!h]
	\includegraphics[width=\textwidth,center]{figs/DESC_UML.pdf}
	\caption{UML diagram of the improved DESC software architecture. Classes in green have already been implemented; classes in red are planned.}
	\label{fig:DESC_UML}
\end{figure}

\begin{figure}[!h]
	\includegraphics[width=0.8\textwidth,center]{figs/EquilibriumIO_UML.pdf}
	\caption{UML Diagram for the new Eqiulibrium IO setup.}
	\label{fig:EquilIO_UML}
\end{figure}

\section{Development Process}

\subsection{Git Workflow}

For our project, we decided to go with the git workflow that uses infinite, parallel development and master branches.
The idea of the workflow is to branch from \texttt{master} to work on feature branches, then merge those into \texttt{dev} where the testing takes place.
Then, once the tests have passed, the feature branch is merged back into \texttt{master}.

This works well when the feature branches are all independent of each other, so merge conflicts on \texttt{dev} are minimized.
However, in our project we were refactoring an existing code base and a lot of our feature branches had some sort of inter-dependence, especially at the start.
A case of this is in the creation of the \texttt{Transform} class that is used as the base class for the coordinates at which we evaluate our objective functions.
We initially each made separate feature branches off of \texttt{master} with a defined refactoring task for each branch.
But, when it came time to merge the branches back into\texttt{master}, we realized that the \texttt{transform} branch, where the \texttt{Transform}, \texttt{Grid}, and \texttt{Basis} classes were created and other code was edited to work with them, touched many parts of the other branches' files.
This created a headache when it came time to merge the \texttt{transform} branch into \texttt{dev}, as not only merge conflicts had to be resolved, but also other working parts from other feature branches were broken.

How we ended up resolving this issue was to first merge the other, less broadly-changing feature branches into \texttt{master}.
Then, from the \texttt{transform} branch, we cherry-picked the latest commits from \texttt{master} into \texttt{transform}.
We could then resolve merge conflicts locally on the \texttt{transform} branch so it would play nice with the other branch commits.
This allowed us to keep the \texttt{dev} branch strictly for testing and not for bug-fixing commits.
So, we were able to modify the other necessary parts of the code that enabled it to run with the \texttt{Transform} classes, and then merged everything back to master without conflicts.

From this experience, we learned that the infinite and parallel \texttt{dev} and \texttt{master} workflow works best when feature branches are smaller in scope and are merged back to master often.
This way, any individual feature branch does not fall behind the master branch, and each feature branch can more easily be independent of each other.
Another thing we learned as we used this workflow was that the \texttt{dev} branch should be used only for testing.
Finding bugs on \texttt{dev} and then fixing them by committing code to \texttt{dev} seemed fine at first.
However, when we would go to make further changes on our feature branches, we would realize that some bug we fixed on \texttt{dev} was still present on the feature branch.
This would then require us to have to re-commit the same bug fix on the feature branch.
Then when we wanted to test the feature branch by merging it into \texttt{dev} before submitting a pull request into \texttt{master}, we would encounter confusing merge conflicts due to having fixed the bug separately on both \texttt{dev} and on the feature branch.
Thus, we learned the hard way that the \texttt{dev} branch should never be advanced alone but rather by merging feature branches, and all work and fixes should be done on feature branches to keep all branches up-to-date on bug fixes.

\subsection{Continuous Integration}

The continuous integration (CI) workflow with automated testing helped us catch bugs as new code was developed.
In our requirements file we had initially listed support for the latest version of JAX, version 0.2.5, as a remnant from the original DESC repository when we forked it.
JAX is still being actively developed and there have been major changes in the past year.
When pushing changes to the new \texttt{Transform} class structure that relies heavily on the use of \texttt{jax.numpy} arrays, the new changes were passing the tests on our local computers without JAX but failing on the virtual machines that were running with JAX.
This problem was puzzling because JAX is intended to replicate most of the usual numpy operations.
Around the same time, an independent user of the code had also reported a problem with running JAX on the original DESC repository.
It turned out that some of the operations we were using were no longer supported in the newer version of JAX, and reducing the installation requirement to version 0.1.77 resolved the issue.
The lesson learned was that maintaining compatibility with other software packages can be laborsome, but frequent testing (from both automated tests and beta users) can help catch problems when they arise.

\section{Profiling \& Optimization}
\label{sec:profiling}

Once a stable version of the refactored code was reached, we performed profiling to measure its performance.
Although DESC is intended to be primarily run with JAX on large clusters that have multiple GPUs, the profiling was performed on a CPU without JAX.
The JIT-ed code is more difficult to profile since the compiled functions are not executed as originally written, so working without JAX allowed us to use conventional python profiler tools.
This represents the slowest-case scenario for running DESC, but is not irrelevant for real applications -- some users may wish to run computations on their personal laptop.
We assumed that any performance gains made on this setup would also benefit other architectures because JAX only adds overhead to the underlying NumPy operations, and GPU parallelization is still limited by the serial processes of each thread.
The wall time to execute DESC depends strongly on the particular inputs chosen, since the dimension of the system of equations scales with the spectral resolution.
For these profiling tests, a stellarator input with 8 poloidal Fourier modes and 2 toroidal Fourier modes were used, with 100 inner loop iterations in a single outer loop step.
For reference, the total run time for this test case was about two minutes.

At a high level, the DESC algorithm has two major components: pre-computation of the spectral transform matrices, and the inner loop optimization.
The transform matrices are built by evaluating basis functions on a given grid, and the optimization loop is repeatedly calling the objective function.
Within the objective function itself, there are two components: the equilibrium force balance errors, and the errors in satisfying the boundary conditions.
At the spectral resolution of the reference case, each of these components consumes a roughly equal portion of the total objective function evaluation time.
cProfile was used to identify which sub-functions were the bottleneck for the overall algorithm, and then the line profiler lprun was used to determine the most expensive operations within those functions.

We discovered that the single function taking up the most computation time was the one responsible for computing the coordinate system jacobian, which is calculated from the triple product $\sqrt{g} = \er\cdot\ev\times\ez$ and gets called each time the objective function is evaluated.
The cross-product, called by \texttt{jax.numpy.cross(a, b)}, is a relatively expensive operation because it involves manipulating multidimensional arrays.
This vector algebra was easily avoided by writing out the triple product explicitly in the code, reducing all of the operations to element-wise multiplication and addition.
This reduced the computation time of the equilibrium force errors by 36\% and the total evaluation time of the objective function by 20\%, as shown in Figure \ref{fig:compute_time_opt}.
Writing out the triple product explicitly does make the code somewhat less human-readable, but it does not add many extra lines of code and is certainly worth doing for the notable speedup.

The next greatest bottleneck after the jacobian terms was Fourier basis evaluation.
This function gets called when building the transform matrices, and also when evaluating the boundary condition errors in the objective function since that transformation cannot be pre-computed.
In the original version of this function, the Fourier series was written in the complex notation $e^{im\theta}(im)^d$, where $d$ is the required order of derivative with respect to $\theta$.
Line profiling revealed that the differentiation step was taking a substantial amount of time, but was not always necessary.
The derivatives are needed in the transform matrices that can be pre-computed, while the objective function that gets called repeatedly only evaluates the Fourier series without differentiating.
This excessive operation was avoided by re-writing the function to evaluate the derivatives recursively, so that the base case of $d=0$ does not waste time differentiating.
Additional performance gains were made by replacing the complex format with all real values, and changing how the two-dimensional arrays were formed.
These edits reduced the computation time of the boundary conditions errors by 19\% and the total evaluation time of the objective function by 10\%, also shown in Figure \ref{fig:compute_time_opt}.
Since the Fourier basis evaluation is also called while pre-computing the transform matrices, this also reduced the time to build the Fourier-Zernike transformations by 36\% as shown in Figure \ref{fig:compile_time_opt}.
The build process is now dominated by the Jacobi polynomial evaluations, which comprise the radial terms of the Zernike basis functions (see Figure \ref{fig:compile_time_rel}).

With both of these optimizations combined, the time to compute the objective function was reduced by 28\% from about 4.8 ms to 3.4 ms.
The equilibrium force balance errors and boundary condition errors still take roughly the same amount of time to compute, as they did in the initial reference version (see Figure \ref{fig:compute_time_rel}).

\begin{figure}[H]
	\centering
	\includegraphics[width=0.6\linewidth,center]{./figs/compute_time_opt.png}
	\caption{Reduction in execution time to compute the objective function, relative to the original code, for different stages of optimization.}
	\label{fig:compute_time_opt}
\end{figure}
%
\begin{figure}[H]
	\centering
	\includegraphics[width=0.6\linewidth,center]{./figs/compile_time_opt.png}
	\caption{Reduction in execution time to evaluate the Fourier-Zernike basis function, relative to the original code, for different stages of optimization.}
	\label{fig:compile_time_opt}
\end{figure}
%
\begin{figure}[H]
	\centering
	\includegraphics[width=0.6\linewidth,center]{./figs/compile_time_rel.png}
	\caption{Relative portion of time to evaluate the Fourier-Zernike basis function spent in each subfunction, for different stages of optimization.}
	\label{fig:compile_time_rel}
\end{figure}
%
\begin{figure}[H]
	\centering
	\includegraphics[width=0.6\linewidth,center]{./figs/compute_time_rel.png}
	\caption{Relative portion of time to compute the objective function spent in each subfunction, for different stages of optimization.}
	\label{fig:compute_time_rel}
\end{figure}

\section{Future Work}

\subsection{Finish Realization of Object-Oriented Design}

The most immediate future work for this project would be to complete the refactoring of the code into an object-oriented design.
The main part of the code that remains in a functional implementation is the outer loop solver.
While the refactoring we completed did increase the legibility of the outer loop solver function, there are still aspects that would benefit from being modularized.
This would come in the form of creating the \texttt{EquilibriumSolver} and \texttt{Optimizer} classes from the UML diagram, which would handle the logic of the inner-outer loop optimization algorithm.
We foresee the most difficulty in this task coming from deciding how to split the necessary steps and control logic across the respective objects.
For example, in the outer loop, a new \texttt{Equilibrium} must be created and perturbed, in order to then be sent into the inner optimization loop.
Should the \texttt{EquilibriumSolver} object be able to create new \texttt{Equilibrium} objects, even though it is not directly connected to an \texttt{Equilibrium}?
Or would this be better handled by \texttt{EquilibriaFamily}?
We have in our UML diagram an outline of what we think is the best to approach, but it is the concrete answers to these questions that we will need to decide in order to finish the implementation of our complete design.

\subsection{Expand Testing Suite}

While we did increase the percentage of code covered by tests by roughly 20\%, there is still a lot of room for improvement.
Increasing the modularity of the code helped us to more easily implement unit tests for the lower level, simpler components.
Though there are still simple classes left for us to cover with unit tests, among what is remaining are the classes that take these simpler objects and use them in a more complex way to return some physically meaningful quantites.
We have two goals to improve testing in the future, given this complexity.

The first is to have a broader integration testing suite, where we test the overall output of our code against analytically solvable solutions.
While this sort of testing suite cannot narrow down issues to specific objects or functions, it can narrow down issues to certain types of problems.
For example, an equilibrium that is axisymmetric is much simpler to compute than a non-axisymmetric equilibrium, as they are different dimensions (2-D versus 3-D).
So, with integration testing of both types of equilibria, we would be able to at least see if problems arise only with certain types of equilibria.
Additionally, this sort of testing suite would instill confidence in the output of the code.

The second is to use the mock object features of the Python unit testing modules to contrive simple examples for these more complex classes.
This would work by allowing us to create mock objects to take the place of objects that would normally be created by reading an input file and performing intermediate calculations.
We would then be able to create simple scenarios where we could more easily check the expected output, when we have control over exactly what the inputs into the function are.

\subsection{Further Profiling \& Optimization}

A preliminary round of profiling and code optimization was performed as mentioned in Section \ref{sec:profiling}, but there is certainly more room for improvement.
Profiling is an ongoing effort, but our time was limited during this project and we decided to prioritize the code refactoring over optimization of existing code.
It would be helpful to profile additional test cases with higher spectral resolutions to understand how the relative bottlenecks scale with the dimensionality of the system.
These same tests then need to be reproduced with JAX in use to check if there are opportunities specific to their overloading of the usual NumPy operations, which will have to be done using the TensorBoard profiler.
We also need to determine when it is advantageous to JIT functions and optimize the code to reduce those compile times.
Another need is to profile on different computer architectures, including with GPUs.
DESC currently has some GPU compatibility and basic parallelization, but this avenue has not been explored in much detail yet.
Finally, work is needed to optimize memory management.
The Jacobian matrix calculations can consume large amounts of memory at high resolutions, and storing this data could be especially problematic when running on GPUs.
We should explore potential solutions to remedy this, such as splitting up the Jacobian and evaluating it on multiple GPUs in parallel.

\bibliography{sources}

\end{document}


\documentclass{article}[12pt]
\usepackage{amsmath}
\usepackage{hyperref}
\usepackage{graphicx}
\graphicspath{ {./} }
\newcommand{\rot}[3]{#3#2#1}
\newcommand{\myskip}{\bigskip\noindent}

\title{Interpreting Neural Networks by Reducing Nonlinearities during Training}
\author{Elod Pal Csirmaz\\
\texttt{\rot{\rot{maz.}{csir}{ep}com}{@}{elod}}}

\date{\today}
\begin{document}

\maketitle

\begin{abstract}
Machine learning and neural networks are used more and more widely,
despite the fact that they mostly function as black boxes, and we
have little insight into how they generate their output.
This makes it difficult to scrutinize and account for decisions made using these systems.

In this paper we present a simple solution that makes it possible to
extract rules from a neural network that employs Parametric Rectified Linear Units (PReLUs).
We introduce a force, applied in parallel to backpropagation, that
aims to reduce PReLUs into the identity function, which then causes
the neural network to collapse into a smaller system of linear functions and inequalities
suitable for review or use by human decision makers.

As this force reduces the capacity of neural networks, it is expected to help avoid overfitting as well.

\myskip
{\bf Keywords:} % (4-6)
machine learning,
neural networks,
rule extraction

\myskip
{\bf MSC classes:} % http://www.ams.org/mathscinet/msc/msc2010.html
62M45,  % Neural nets and related approaches
68T05   % Learning and adaptive systems

% \myskip
% {\bf ACM classes:} % http://www.acm.org/about/class/ccs98-html

\end{abstract}

\section{Introduction}

Machine learning solutions and, more specifically, neural networks are used
in more and more areas of our lives. At the same time, since it is difficult to
grasp at any level how they function and so they can be seen as black boxes,
there is a reluctance in their adoption in the case of companies who need to
remain accountable for business and/or legal reasons. This fact also impacts
the trust the public places in such systems.

Because of this, providing tools that can help us understand how these
models arrive at a certain decision is an area of active research.
The tools range from visualization tools, tools that identify the
most significant input, tools that allow experimenting with the input data
to understand the relative importance the model attributes to different
signals, and, the focus of this paper, significantly reducing the complexity
of a neural network to the point where ``rules'' for decisions could be
extracted from it, which can then be reviewed or even used by human decision-makers.

One form these rules can take is a closed, relatively simple algebraic
expression (for example, a linear function) which has a behavior more
predictable for humans than the full function implemented by a neural
network involving parameters and nonlinear functions in the range of
millions.

In this paper we detail a method of reducing the complexity of a neural
network during its training, which then allows describing its output as a
simple system of linear functions limited to regions defined by linear
inequalities.
In particular,
we investigate feed-forward neural networks
where all nonlinearities are Parametric Rectified Linear Units (PReLUs)
defined as
\[ f(x) = \max(0, x) + a \min(0, x). \]
where $0\leq a\leq 1$.\cite{PR}

Let the input of the network be the vector (or tensor, which we flatten) $I=(i_1, i_2, \ldots, i_n)$
and the output be the vector (or tensor) $O=(o_1, o_2, \ldots, o_m)$.

\section{The Neural Network as a Combination of Linear Functions}

It is easy to see that the output $O$ of such a network can be expressed
as a set of linear combinations of $i_x$, each restricted to a particular region of input values
by a set of linear inequalities on $i_x$ (the conditions).

Intuitively, this is true because the output is a linear function of the input
as long as (that is, inside regions of the possible input vectors in which) none of the inputs of PReLUs
cross from negative to positive or \emph{vice versa}; and because the borders of these regions
are linear as well.

More formally, we can show that this statement is true by constructing a suitable
system of linear combinations and inequalities for any given neural network.
For all PReLUs in the network, let us decide whether their inputs are negative, or nonnegative.
If the number of PReLUs is $\delta$, this means $2^\delta$ possible scenarios.
In each scenario, the neural network becomes completely linear, with the output being a linear function of the input,
and, in fact, with all partial results inside the network being linear functions of the input as well.
This entails that the inputs of PReLUs are also linear functions of the input $I$, yielding
the $\delta$ linear inequalities (the conditions) that mark out the region of the input space where this scenario applies.

Naturally, this system of $2^\delta$ linear combinations yielding the output, each controlled by $\delta$
conditions is likely to be redundant, as the conditions may not be independent,
or may be contradictory, in which case some of the $2^\delta$ regions of the input space will be empty.
But this system will be an accurate description of the output of the neural network nevertheless.

\myskip
While a linear function may seem easy to interpret and to predict for us humans, clearly $2^\delta$
functions will not be so if $\delta$ is in the range of thousands or millions as in most modern neural networks.
We therefore need to construct an approximation of the output of our neural network which is linear in much
larger regions of the input space, and has far less different regions as well.
This is equivalent to trying to find another neural network that approximates our network well,
but contains a small number of PReLUs only.

\section{Reducing PReLUs into Linearities}

We achieve this by continuing training the original neural network while applying a force on each of
its PReLUs that moves its parameter $a$ towards 1.

Note that if $a=1$, then the PReLU degenerates into the identity function $f(x)=x$, and ceases to be a nonlinearity.
In a sense, it disappears, and the neural network ``collapses'' around it, inasmuch as the linear mapping
that precedes the degenerate PReLU and the one after it can now be combined into a single linear map.

This force therefore finds a balance between approximating the training data
and reducing the number of PReLUs, thereby yielding a neural network that is feasible to express
as a set of linear functions and inequalities for human consumption.
By removing nonlinearities, this force also reduces the capacity of the neural network,
and therefore we expect that it can also be helpful to avoid overfitting.

We chose a
force that is independent of $a$ and is applied after each backpropagation step, and moves the parameter by an adjustment rate
$\eta_p$ but never above 1:
\[ a_{t+1} = \max\left(0,\;\min\big(1,\;a_t + \eta_p\,\mathrm{sgn}(1-a_t)\,\big)\right) \]

In our research we considered any PReLU fully linear if
\[ a > 0.9995 \]

We also made $\eta_p$ dependent on the current error exhibited by the network,
to allow the network initially to train without interference, and then
pull the PReLU parameters to 1 more and more aggressively. We chose
\[ \eta_p = \eta_0\,0.01\,\max\big(0,\; -\log_{10}(err)-2\big) \]
where $\eta_0=0.01$ is the learning rate for the whole model, and $err$ is the training error.
This means that $\eta_p$ will start becoming non-zero when the training error
falls below 0.01, and can grow indefinitely as the training error falls
(although the $a$ parameters are clipped anyway).

\section{Example}

Sample code demonstrating the above is available at
\url{https://github.com/csirmaz/trained-linearization}.
The code implements training a small neural net with the force on PReLU parameters,
and contains logic to extract a system of linear inequalities and combinations
from the weights and other parameters of the model.

The neural network itself
has 4 input nodes, 3 output nodes, and
4 hidden layers in between with 5 nodes on each.
These layers are linked by trainable fully-connected layers
with PReLU activation following them, except for the last fully connected layer,
which simply provides the output.

The network is expected to learn the following relationships:
\begin{align}
\mathrm{out}_1 &= \mathrm{in}_1 \;\mathrm{xor}\; \mathrm{in}_2 \\
\mathrm{out}_2 &= \mathrm{in}_3 \;\mathrm{xor}\; \mathrm{in}_4 \\
\mathrm{out}_3 &= \mathrm{in}_1 \cdot \mathrm{in}_2
\end{align}
where all the inputs and outputs are 0 or 1 with some extra noise added to the inputs.
We generate the training data according to the expressions above, for all $2^4$ possible combinations of 0 and 1 for the input.
New noise is generated for each batch of training data.

\begin{figure}[ht]
\centering
\includegraphics[height=10cm]{example}
\caption{A trained network with 4 input nodes (top) and 3 output nodes (bottom). Red lines denote positive, blue lines negative weights.
Lines leading from the right denote bias values. Green boxes mark PReLUs that are considered linear, while red ones mark nonlinear PReLUs.}
\label{fig}
\end{figure}

Figure \ref{fig} visualizes a trained model. Two nonlinearities are enough to solve this problem, and indeed the training
settled on a solution where only two PReLUs were nonlinear.

Although in the visualization the relationships between the inputs and the outputs are far from clear,
once the rest of the PReLUs are considered linear, all the weights collapse into very simple expressions.
See figure \ref{output} for the output of the code listing these expressions for the $2^2$ cases arising from the
two nonlinear PReLUs.
These correspond completely to the relationships in the training data.

\begin{figure}[ht]
\begin{verbatim}
IF +0.00*in1 +0.00*in2 -1.00*in3 -1.00*in4 +1.00 < 0
   (PReLU #3 on level 1 is neg. ln(1-weight)=-2.58)
IF +1.00*in1 +1.00*in2 -0.00*in3 -0.00*in4 -1.00 < 0
   (PReLU #1 on level 4 is neg. ln(1-weight)=-2.60)
THEN
  out1 = +1.01*in1 +1.01*in2 -0.00*in3 -0.00*in4 +0.00
  out2 = -0.00*in1 -0.00*in2 -1.03*in3 -1.03*in4 +2.04
  out3 = -0.00*in1 -0.00*in2 +0.00*in3 +0.00*in4 -0.00

IF +0.00*in1 +0.00*in2 -1.00*in3 -1.00*in4 +1.00 < 0
   (PReLU #3 on level 1 is neg. ln(1-weight)=-2.58)
IF +1.00*in1 +1.00*in2 -0.00*in3 -0.00*in4 -1.00 > 0
   (PReLU #1 on level 4 is pos. ln(1-weight)=-2.60)
THEN
  out1 = -1.01*in1 -1.01*in2 +0.00*in3 +0.00*in4 +2.02
  out2 = -0.00*in1 -0.00*in2 -1.03*in3 -1.03*in4 +2.04
  out3 = +1.00*in1 +1.00*in2 +0.00*in3 +0.00*in4 -1.01

IF +0.00*in1 +0.00*in2 -1.00*in3 -1.00*in4 +1.00 > 0
   (PReLU #3 on level 1 is pos. ln(1-weight)=-2.58)
IF +1.00*in1 +1.00*in2 +0.00*in3 +0.00*in4 -1.00 < 0
   (PReLU #1 on level 4 is neg. ln(1-weight)=-2.60)
THEN
  out1 = +1.01*in1 +1.01*in2 +0.00*in3 +0.00*in4 +0.00
  out2 = -0.00*in1 -0.00*in2 +1.00*in3 +1.00*in4 +0.01
  out3 = -0.00*in1 -0.00*in2 -0.00*in3 -0.00*in4 +0.00

IF +0.00*in1 +0.00*in2 -1.00*in3 -1.00*in4 +1.00 > 0
   (PReLU #3 on level 1 is pos. ln(1-weight)=-2.58)
IF +1.00*in1 +1.00*in2 +0.00*in3 +0.00*in4 -1.00 > 0
   (PReLU #1 on level 4 is pos. ln(1-weight)=-2.60)
THEN
  out1 = -1.01*in1 -1.01*in2 -0.00*in3 -0.00*in4 +2.02
  out2 = -0.00*in1 -0.00*in2 +1.00*in3 +1.00*in4 +0.01
  out3 = +1.00*in1 +1.00*in2 -0.00*in3 -0.00*in4 -1.01
\end{verbatim}
\caption{Sample output of the code}
\label{output}
\end{figure}

\section{Future work}

Future work can include and improved algorithm to extract linear rules from the network
that recognizes dependencies between the inequalities to simplify its output.

Another alternative is to approximate PReLUs that are deemed linear not with $f(x)=x$,
but with e.g. $f(x)=\frac{a+1}{2}x$, which can allow replacing more PReLUs with a fully linear function than we currently do.

\begin{thebibliography}{99}
\bibitem{PR}
Kaiming He and
               Xiangyu Zhang and
               Shaoqing Ren and
               Jian Sun,
``Delving Deep into Rectifiers: Surpassing Human-Level Performance on
               ImageNet Classification,''
Microsoft Research,
6 February 2015.
http://arxiv.org/abs/1502.01852.
\end{thebibliography}

\end{document}