\section{MOTION}
\label{MOTION}
\begin{ipifield}{}%
{Allow chosing the type of calculation to be performed. Holds all the information that is calculation specific, such as geometry optimization parameters, etc.}%
{}%
{\ipiitem{mode}%
{How atoms should be moved at each step in the simulatio. 'replay' means that a simulation is restarted from a previous simulation.}%
{data type: string; options: `vibrations', `minimize', `replay', `neb', `dynamics', `alchemy', `instanton', `dummy', `multi'; }%
}
\begin{ipifield}{\hyperref[GEOP]{optimizer}}%
{Option for geometry optimization}%
{}%
{\ipiitem{mode}%
{The geometry optimization algorithm to be used}%
{default: `lbfgs'; data type: string; options: `sd', `cg', `bfgs', `bfgstrm', `lbfgs'; }%
}
\end{ipifield}
\begin{ipifield}{fixatoms}%
{Indices of the atmoms that should be held fixed.}%
{default:  [ ] ; data type: integer; }%
{\ipiitem{shape}%
{The shape of the array.}%
{default:  (0,) ; data type: tuple; }%
\ipiitem{mode}%
{If 'mode' is 'manual', then the array is read from the content of 'cell' takes a 9-elements vector containing the cell matrix (row-major). If 'mode' is 'abcABC', then 'cell' takes an array of 6 floats, the first three being the length of the sides of the system parallelopiped, and the last three being the angles (in degrees) between those sides. Angle A corresponds to the angle between sides b and c, and so on for B and C. If mode is 'abc', then this is the same as for 'abcABC', but the cell is assumed to be orthorhombic. 'pdb' and 'chk' read the cell from a PDB or a checkpoint file, respectively.}%
{default: `manual'; data type: string; options: `manual', `file'; }%
}
\end{ipifield}
\begin{ipifield}{\hyperref[ALCHEMY]{alchemy}}%
{Option for alchemical exchanges}%
{}%
{\ipiitem{mode}%
{ }%
{default: `dummy'; data type: string; options: `dummy'; }%
}
\end{ipifield}
\begin{ipifield}{\hyperref[PHONONS]{vibrations}}%
{Option for phonon computation}%
{}%
{\ipiitem{mode}%
{The algorithm to be used: finite differences (fd), normal modes finite differences (nmfd), and energy-scaled normal mode finite differences (enmfd).}%
{default: `fd'; data type: string; options: `fd', `nmfd', `enmfd'; }%
}
\end{ipifield}
\begin{ipifield}{\hyperref[INSTANTON]{instanton}}%
{Option for Instanton optimization}%
{}%
{\ipiitem{mode}%
{Use the full or half of the ring polymer during the optimization}%
{default: `rate'; data type: string; options: `rate', `splitting'; }%
}
\end{ipifield}
\begin{ipifield}{\hyperref[INITFILE]{file}}%
{This describes the location to read a trajectory file from.}%
{default: `'; data type: string; }%
{\ipiitem{bead}%
{The index of the bead for which the value will be set. If a negative value is specified, then all beads are assumed.}%
{default:  -1 ; data type: integer; }%
\ipiitem{cell\_units}%
{The units for the cell dimensions.}%
{default: `'; data type: string; }%
\ipiitem{mode}%
{The input data format. 'xyz' and 'pdb' stand for xyz and pdb input files respectively. 'chk' stands for initialization from a checkpoint file.}%
{default: `xyz'; data type: string; options: `xyz', `pdb', `chk'; }%
}
\end{ipifield}
\begin{ipifield}{\hyperref[NEB]{neb\_optimizer}}%
{Option for geometry optimization}%
{}%
{\ipiitem{mode}%
{The geometry optimization algorithm to be used}%
{default: `lbfgs'; data type: string; options: `sd', `cg', `bfgs', `lbfgs'; }%
}
\end{ipifield}
\begin{ipifield}{\hyperref[DYNAMICS]{dynamics}}%
{Option for (path integral) molecular dynamics}%
{}%
{\ipiitem{mode}%
{The ensemble that will be sampled during the simulation. }%
{default: `nve'; data type: string; options: `nve', `nvt', `npt', `nst', `sc', `scnpt'; }%
\ipiitem{splitting}%
{The Louiville splitting used for sampling the target ensemble. }%
{default: `obabo'; data type: string; options: `obabo', `baoab'; }%
}
\end{ipifield}
\begin{ipifield}{fixcom}%
{This describes whether the centre of mass of the particles is fixed.}%
{default:  True ; data type: boolean; }%
{}
\end{ipifield}
\begin{ipifield}{\hyperref[MOTION]{motion}}%
{A motion class that can be included as a member of a 'multi' integrator.}%
{}%
{\ipiitem{mode}%
{How atoms should be moved at each step in the simulatio. 'replay' means that a simulation is restarted from a previous simulation.}%
{data type: string; options: `vibrations', `minimize', `replay', `neb', `dynamics', `alchemy', `instanton', `dummy', `multi'; }%
}
\end{ipifield}
\end{ipifield}


\section{Methods}
The computational methodology described herein has been implemented in the open-source C\_MPI program\footnote{MPI is the abbreviation of message passing interface.} \texttt{fitoog}.\autocite{mccluskey_arm61/fitoog_2019}

\subsection{Simulation Methodology}
%
\begin{marginfigure}
    \centering
    \includegraphics[width=\linewidth]{smallangle/director}
    \caption{A graphical description of the severe coarse-graining applied to the MARTINI description of the \emph{n}-decyltrimethylammonium surfactant molecule for the use of the particle swarm algorithm.}
    \label{fig:director}
\end{marginfigure}
%
The \texttt{fitoog} software takes a series of input files that define the molecules, intramolecular interactions, and scattering lengths for the constituent particles.
The molecule input file is a space-separated file consisting of an index, particle name, \emph{x}-coordinate, \emph{y}-coordinate, \emph{z}-coordinate, and scattering length.
When \texttt{fitoog} reads in a molecular input file, a \texttt{differences} object is created, this stores the differential between the atomistic, or near-atomistic\footnote{Such as the MARTINI-style beading use in this work.} description and the more severe coarse-grained description that is used in the PSO.

Inspired by the coarse-graining methodology applied to directional colloid self-assembly by Law \emph{et al.},\autocite{law_coarse-grained_2016} a severe coarse-graining methodology was developed for use on surfactant molecules in \texttt{fitoog}.
This allowed for a significantly reduced parameter dimensionality to which the PSO\footnote{Which is described in Section \ref{sec:partswarm}.} could be applied.
The severe coarse-graining reduced the surfactant to a ``director'' description; where each surfactant molecule is defined by a position and a direction, shown pictorially in Figure~\ref{fig:director}.
This reduced the parameter dimensionality to just six variables per molecule; three of which described the centre-of-mass position of the molecule and three that describe the angular orientation of the surfactant in space.\footnote{Given the symbols, $a$, $b$, $c$ for the positions and $\phi$, $\omega$, and $\kappa$ for the angles.}
Additionally, if the molecule was a single particle in length, as is the case for the \ce{NO3-} anion in the Section~\ref{sec:real_data}, then the dimensionality is just three, as the direction is arbitrary.

The methodology for the particle swarm in this work followed the implementation discussed in Section~\ref{sec:partswarm}.
Where the parameters to be optimised where the position and orientation of the input molecules.\footnote{The $a$, $b$, $c$, $\phi$, $\omega$, and $\kappa$ mentioned above.}
Following each PSO iteration, the molecule director is expanded from the position variable\footnote{Using the \texttt{differences} object mentioned above.}.
This representation is then rotated based on a rotation matrix, in this work the rotation matrix was constructed by first rotating the rotation axis by $-\phi$ and $-\omega$, then rotating by $\kappa$ in around the $z$-axis, before rotating the axis back to the original position by $\omega$ and $\phi$.\sidecite[Figure~\ref{fig:rot} defines these angles]{evans_rotations_2001}
Following the expansion and reorientation, the scattering profile is then calculated using the Debye equation,\autocite{debye_zerstreuung_1915} this was used over the more efficient Golden Vectors\autocite{watson_rapid_2013} or Fibonacci sequence\autocite{svergun_solution_1994} methods as the aim of this work was to assess the application of the PSO method and efficiency was not the initial goal.
The agreement between the calculated scattering profile and the experimental input scattering was used as a figure of merit, $\zeta$, that was to be optimised by the particle swarm algorithm.
This $\zeta$ was a simple $\chi^2$ value found as follows,
%
\begin{equation}
\zeta = \chi^2 = \sum\frac{(I_{\text{exp}}(q) - I_{\text{calc}}(q))^2}{\text{d}I_{\text{exp}}(q)},
\end{equation}
%
where $I_{\text{exp}}(q)$ is the experimental scattering intensity, $I_{\text{calc}}(q)$ is the calculated scattering intensity, and $\text{d}I_{\text{exp}}(q)$ is the uncertainty in the experimental scattering intensity, all at a given $q$-vector.
%
\begin{marginfigure}
    \centering
\includegraphics[width=\linewidth]{smallangle/rotdia}
    \caption{The definitation of the polar angles used in the coarse grained representation of the surfactant molecule.}
    \label{fig:rot}
\end{marginfigure}
%

Throughout this work, an interia weight for the PSO of 0.4 was used.
Generally, values of between 0 and 2 are used for the global and personal acceleration coefficients\footnote{These are standard values for the PSO method following \cite{sun_study_2010}.}.
For the test case discussed below a value of 2 was used for both acceleration coefficients, however, in the real case discussed later, the global acceleration coefficient was reduced to 1.
This was chosen to reduce the acceleration toward the global best and improve the ability for the PSO to search the parameter space available in this much larger problem.

\subsection{Parallelisation}
\label{sec:para}
The use of a population-based optimisation method, such as the PSO, allowed for easy access to highly parallel simulation.
Parallelisation was achieved by spreading the population evenly across the cores that were available to the simulation.
Inter-core messaging was performed using the MPI libraries, and to ensure efficiency only the figures of merit and the best possible structure were shared across the cores.
This means that during a given \texttt{fitoog} run, the only serial component was the determination of the lowest figure of merit.
The efficiency of the parallelisation was defined by considering the strong and weak scaling of the software.
It was possible to determine the percentage of serial, $s$, and parallel, $p$, components of the software by fitting the speedup\footnote{The time taken for a job to run on a single core divided by the time taken on multiple cores.} with Amdahl's law,\autocite{amdahl_validity_1967}
%
\begin{equation}
\text{speedup} = \frac{1}{s + \sfrac{p}{N}},
\end{equation}
%
where, $N$ is the number of cores in the parallel job, and $s + p = 1$.
While more sophisticated methodologies could be used to further reduce the serial component, such as having a core-level best population that is only occasionally communicated with the entire swarm, this implementation was shown to be highly parallelised and useful for assessing the utility of the PSO method.


\section{NORMALMODES}
\label{NORMALMODES}
\begin{ipifield}{}%
{Deals with the normal mode transformations, including the adjustment of bead masses to give the desired ring polymer normal mode frequencies if appropriate. Takes as arguments frequencies, of which different numbers must be specified and which are used to scale the normal mode frequencies in different ways depending on which 'mode' is specified.}%
{}%
{\ipiitem{transform}%
{Specifies whether to calculate the normal mode transform using a fast Fourier transform or a matrix multiplication. For small numbers of beads the matrix multiplication may be faster.}%
{default: `fft'; data type: string; options: `fft', `matrix'; }%
}
\begin{ipifield}{\hyperref[NMFREQUENCIES]{frequencies}}%
{Specifies normal mode frequencies for a (closed path) calculation}%
{dimension: frequency; default:  [ ] ; data type: float; }%
{\ipiitem{units}%
{The units the input data is given in.}%
{default: `automatic'; data type: string; }%
\ipiitem{shape}%
{The shape of the array.}%
{default:  (0,) ; data type: tuple; }%
\ipiitem{style}%
{Specifies the technique to be used to calculate the dynamical masses.
                                                'rpmd' simply assigns the bead masses the physical mass.
                                                'manual' sets all the normal mode frequencies except the centroid normal mode manually.
                                                'pa-cmd' takes an argument giving the frequency to set all the non-centroid normal modes to.
                                                'wmax-cmd' is similar to 'pa-cmd', except instead of taking one argument it takes two
                                                      ([wmax,wtarget]). The lowest-lying normal mode will be set to wtarget for a
                                                      free particle, and all the normal modes will coincide at frequency wmax. }%
{default: `rpmd'; data type: string; options: `pa-cmd', `wmax-cmd', `manual', `rpmd'; }%
\ipiitem{mode}%
{If 'mode' is 'manual', then the array is read from the content of 'cell' takes a 9-elements vector containing the cell matrix (row-major). If 'mode' is 'abcABC', then 'cell' takes an array of 6 floats, the first three being the length of the sides of the system parallelopiped, and the last three being the angles (in degrees) between those sides. Angle A corresponds to the angle between sides b and c, and so on for B and C. If mode is 'abc', then this is the same as for 'abcABC', but the cell is assumed to be orthorhombic. 'pdb' and 'chk' read the cell from a PDB or a checkpoint file, respectively.}%
{default: `manual'; data type: string; options: `manual', `file'; }%
}
\end{ipifield}
\begin{ipifield}{open\_paths}%
{Indices of the atmoms whose path should be opened (zero-based).}%
{default:  [ ] ; data type: integer; }%
{\ipiitem{shape}%
{The shape of the array.}%
{default:  (0,) ; data type: tuple; }%
\ipiitem{mode}%
{If 'mode' is 'manual', then the array is read from the content of 'cell' takes a 9-elements vector containing the cell matrix (row-major). If 'mode' is 'abcABC', then 'cell' takes an array of 6 floats, the first three being the length of the sides of the system parallelopiped, and the last three being the angles (in degrees) between those sides. Angle A corresponds to the angle between sides b and c, and so on for B and C. If mode is 'abc', then this is the same as for 'abcABC', but the cell is assumed to be orthorhombic. 'pdb' and 'chk' read the cell from a PDB or a checkpoint file, respectively.}%
{default: `manual'; data type: string; options: `manual', `file'; }%
}
\end{ipifield}
\end{ipifield}


% !TEX root = ../zeth-protocol-specification.tex

\section{\zeth~statement after primitive instantiation}\label{instantiation:statement}

After instantiating the various primitives and providing security proofs to justify that they comply with the security requirements listed in previous sections, $\RELCIRC$ now becomes:

\begin{itemize}
    \item For each $i \in [\jsin]$:
    \begin{enumerate}
        \item $ \auxinputs.\jsins{i}.\znote.\apk = \blake{2s}{\taggedaddr \concat \pad{0}{\blakeCompLen}}$ \\ with $\taggedaddr$ defined in~\cref{instantiation:prf-comm-crh:prf}
        \item $\auxinputs.\jsins{i}.\nf{} = \blake{2s}{\taggednf \concat \auxinputs.\jsins{i}.\znote.\rho}$ \\ with $\taggednf$ defined in~\cref{instantiation:prf-comm-crh:prf}
        \item $\auxinputs.\jsins{i}.\cm{} = \blake{2s}{\auxinputs.\jsins{i}.\znote.\noter{} \concat \msg}$ \\ with $\msg = \auxinputs.\jsins{i}.\znote.\apk \concat \auxinputs.\jsins{i}.\znote.\rrho \concat \auxinputs.\jsins{i}.\znote.\notev$
        \item $\auxinputs.\htags{i} = \blake{2s}{ \taggedpk \concat \priminputs.\hsig}$ (malleability fix, see~\cref{appendix:trnm}) with $\taggedpk$ defined in~\cref{instantiation:prf-comm-crh:prf}
        \item $(\auxinputs.\jsins{i}.\znote.\notev) \cdot (1 - e) = 0$ is satisfied for the boolean value $e$ set such that if $\auxinputs.\jsins{i}.\znote.\notev > 0$ then $e = 1$.
        \item The Merkle root $\mkroot'$ used to check the Merkle authentication path $\auxinputs.\jsins{i}.\mkpath$ of commitment $\auxinputs.\jsins{i}.\cm{}$, with $\mimcMPPrime{\rCURVE}$, equals $\priminputs.\mkroot$ if $e = 1$.
        \item $\priminputs.\nfs{i}$ \\ $= \indexedset{\pack{\slice{\auxinputs.\jsins{i}.\nf{}}{k \cdot \fieldBitCap}{(k+1) \cdot \fieldBitCap}}{\FFx{\rCURVE}}}{k \in [\floor{\prfNfOutLen/\fieldBitCap}]}$
        \item $\priminputs.\htags{i}$ \\ $= \indexedset{\pack{\slice{\auxinputs.\htags{i}}{k \cdot \fieldBitCap}{(k+1) \cdot \fieldBitCap}}{\FFx{\rCURVE}}}{k \in [\floor{\prfPkOutLen/\fieldBitCap}]}$
    \end{enumerate}
    \item For each $j \in [\jsout]$:
    \begin{enumerate}
        \item $\auxinputs.\znotes{j}.\rrho = \blake{2s}{ \taggedrho \concat \priminputs.\hsig}$ (malleability fix, see~\cref{appendix:trnm}) with $\taggedrho$ defined in~\cref{instantiation:prf-comm-crh:prf}
        \item $\priminputs.\cms{j} = \blake{2s}{\auxinputs.\znotes{j}.\noter{} \concat \msg }$ \\ with $\msg = \auxinputs.\znotes{j}.\apk \concat \auxinputs.\znotes{j}.\rrho \concat \auxinputs.\znotes{j}.\notev$
    \end{enumerate}
    \item $\priminputs.\hsig = \indexedset{\pack{\slice{\auxinputs.\hsig}{k \cdot \fieldBitCap}{(k+1) \cdot \fieldBitCap}}{\FFx{\rCURVE}}}{k \in [\floor{\crhhsigOutLen/\fieldBitCap}]}$
    \item $\priminputs.\resbits = \packResBits{\indexedset{\auxinputs.\jsins{i}.\nf{}}{i \in [\jsin]}, \auxinputs.\vin, \auxinputs.\vout, \auxinputs.\hsig, \indexedset{\auxinputs.\htags{i}}{i \in [\jsin]}}$
    \item Check that the ``\gls{joinsplit} is balanced'', i.e.~check that the \gls{joinsplit-eq} holds:
    \begin{align*}
        &\pack{\auxinputs.\vin}{\FFx{\rCURVE}} + \sum_{i \in [\jsin]} \pack{\auxinputs.\jsins{i}.\znote.\notev}{\FFx{\rCURVE}} \\
        & = \sum_{j \in [\jsout]} \pack{\auxinputs.\znotes{j}.\notev}{\FFx{\rCURVE}} + \pack{\auxinputs.\vout}{\FFx{\rCURVE}}
    \end{align*}
\end{itemize}

\begin{remark}
    For higher security, we could use \blake{2b}{} with 32-byte output instead of \sha{256}. In fact, since a precompiled contract computing the \blake{2}{}~compression function~\cite{blakecompietf} has been added to the Istanbul release of \ethereum~(EIP 152~\cite{blake-eip}), it could be possible to write a small wrapper on the smart contracts, in order to hash with \blake{2b}{} with any parameter.
\end{remark}

\subsection{Instantiating the packing functions}\label{instantiation:statement:pack}

As we consider SNARKs based on arithmetic circuits defined over a prime field, all variables in the constraint system are interpreted as field elements. Nevertheless, as illustrated in~\cref{zeth-protocol:statement}, part of the statement consists of functions whose co-domains are sets of binary strings (which may be longer than the bit representation of elements of the finite field). While a bit (i.e.~$\{0, 1\}$) is an element of $\FF_p$ ($p$ prime), it is important to minimize the number of gates in the arithmetic circuit (for proof generation efficiency), and to minimize the number of input wires (to improve verification time). This can be done by representing fragments of binary strings as the base 2 decomposition of field elements, thereby ``packing'' binary strings into multiple elements. Converting binary strings into field elements requires the addition of some arithmetic gates (extending the statement to be proven), but reduces the number of primary inputs (reducing the complexity of the SNARK verification carried out on-chain).
The cost of $\groth$ zk-SNARK~\cite{groth2016size} proof verification is linear in the number of primary inputs, since each input acts as a scalar in a costly scalar multiplication of a curve point in $\gset_1$. Hence, while packing slightly increases the prover cost -- by adding constraints to the circuit -- it simplifies the verifier's work.

In this section, we detail the method by which we encode (resp.~decode) a set of binary strings to (resp.~from) sets of field elements. In the rest of this section, the notion of \emph{packing policy} refers to the set of \emph{packing} and \emph{unpacking} functions.

The set of primary inputs is composed of the input nullifiers, the output commitments, the public values (see~\cite[Section 3.4.3]{zethpaper}) along with the signature hash and the authentication tags for security (malleability fix, see~\cref{appendix:trnm}). The complete description of the public inputs is represented in~\cref{instantiation:eq:primary-inputs}.
%
\begin{equation}\label{instantiation:eq:primary-inputs}
    (\indexedset{\priminputs.\nf{i}}{i \in [\jsin]}, \indexedset{\priminputs.\cms{j}}{j \in [\jsout]}, \vin, \vout, \hsig, \indexedset{\priminputs.\htags{i}}{i \in [\jsin]})
\end{equation}

The primary inputs that consist of binary strings are: the nullifiers $\nfs{}$, the public values $\vin$ and $\vout$, the signature hash $\hsig$ and the authentication tags $\htags{}$.

For a binary string $x$, let $\alpha_x = \ceil{\len{x} / \fieldBitCap}$ be the number of field elements required to completely encode $x$ and let $\beta_x = \floor{\len{x} / \fieldBitCap}$ be the number of field elements whose capacity is fully used. Let $\gamma_x = \len{x} \pmod{\fieldBitCap}$ be the number of ``residual'' bits remaining after fully using $\beta_x$ field elements.

\begin{example}\label{instantiation:eg:packing}
    Consider binary strings $A \in \bin^7$ of length 7, to be encoded over the field $\FFx{41}$. This field has a capacity of $5$ bits, and therefore $\alpha_A = 2$, $\beta_A = 1$, and $\gamma_A = 2$. That is, $A$ can be represented as 2 field elements, or as 1 field element with 2 ``residual'' bits.

    Consider $A = (1111011)$. \cref{instantiation:fig:packingA} illustrates how $A$ can be packed as field elements. Note that the 2 residual bits are taken from the ``beginning'' of the bit string, that is, the highest order bits.
\end{example}

\begin{figure}[ht]
    \centering
    \includegraphics[width=1\textwidth]{images/bit-packing-stringA.png}
    \caption{Packing of string $A$ (see~\cref{instantiation:eg:packing})}\label{instantiation:fig:packingA}
\end{figure}

We now consider strategies to pack all primary inputs that are binary strings. A naive approach is to encode each binary string $x$ as $\alpha_x$ field elements. In general, this results in significant waste (and consequently more field elements than necessary), especially when the number of residual bits is small compared to $\fieldBitCap$ (see~\cref{instantiation:fig:packing-multi-strings}). An alternative strategy could be to concatenate all binary strings into a single string $y$ and pack this string into $\alpha_y$ field elements. While this approach minimizes the set of unused bits, each unpack operation would require different shift and mask operations over 2 or 3 field elements. This significantly increases the complexity of the unpacking operation that must be performed on-chain, resulting in a higher gas cost (due to extra logic) or more contract code (if each unpack operation is hard-coded).

The \zeth~protocol requires that each binary string variable $x$ is packed into $\beta_x$ field elements, and the residual bits from all binary strings, along with the public values $\vin$ and $\vout$, are aggregated into a variable $\resbits$. Let $\resBitsBLen$ be the total number of residual bits, and $\resBitsFLen$ be the number of field elements required to represent $\resbits$. We assume that $\zvalueLen < \fieldBitCap$, and define the notation $\gamma_\notev = \zvalueLen$ for the bit lengths of public values $\vin$ and $\vout$. Thus $\resBitsBLen$ is given by
\[
    \resBitsBLen = \gamma_\hsig + 2 \cdot \gamma_\notev + \jsin \cdot (\gamma_{\nf{}} + \gamma_{\htag{}} )
\]
and the lengths, in field elements, of each of the correponding public inputs are
\begin{align*}
    \nfFLen &= \beta_{\nf{}} \\
    \hsigFLen &= \beta_{\hsig} \\
    \htagFLen &= \beta_{\htag{}} \\
    \resBitsFLen &= \ceil{\resBitsBLen / \fieldBitCap}
\end{align*}

\begin{figure}[ht]
    \centering
    \includegraphics[width=1\textwidth]{images/bit-packing-multi-strings.png}
    \caption{Packing of multiple strings. Observe that, by carefully arranging the bits of the input strings, it is possible to output fewer field elements}\label{instantiation:fig:packing-multi-strings}
\end{figure}

The residual bits $\resbits$ are formatted as follows:
\[
    \widetilde{\hsig} \concat \widetilde{\nfs{}} \concat \widetilde{\htags{}} \concat \vin \concat \vout
\]
where $\widetilde{\hsig}, \widetilde{\nfs{}}, \widetilde{\htags{}}$ are, respectively, the $\gamma_\hsig, \gamma_\nf{}, \gamma_\htag{}$ bits.

Note that the public values are packed into the ``last'', or lowest order, $2 \cdot \gamma_\notev$ bits of the resulting field element(s). In this way, their unpack functions are independent of the values $\jsin$ and $\jsout$ and of the number of residual bits required for each bit string (and consequently, independent of the finite field used).

To format the unpacked primary inputs into field elements, we define the following functions.
Given a bit string of length less than $\fieldBitCap$, the algorithm \pack{}{} (see~\cref{zeth-protocol:fig:packing-alg}) returns a field element. Given the nullifiers, public values and authentication tags, the algorithm \packResBits{}{} (see~\cref{zeth-protocol:fig:packing-resbits-alg}) outputs the residual bits. Given a set of packed field elements and the residual bits, the algorithm \unpack{}{} returns the variables reassembled as binary strings. In particular, we have that $\unpack{\priminputs.\nfs{}, \resbits}{\nf{}} = \indexedset{\auxinputs.\jsins{i}.\nf{}}{i \in [\jsin]}$.
\begin{align*}
    &\pack{}{\FFx{\rCURVE}} : \BB^{\leq \fieldBitCap} \to \FFx{\rCURVE} \\
    &\packResBits{} : {(\BB^\prfNfOutLen)}^{\jsin} \times {(\BB^\zvalueLen)}^2 \times \BB^\crhhsigOutLen \times {(\BB^\prfPkOutLen)}^{\jsin} \to (\FFx{\rCURVE})^\resBitsFLen \\
    &\unpack{}{} : \FFx{\rCURVE}^* \times (\FFx{\rCURVE})^{\resBitsFLen} \to \BB^* \\
\end{align*}

The $\unpack{}{}$ functions for nullifiers, public values and signature hash are represented as follows.
\begin{align*}
    &\unpack{}{}_{\hsig} : (\FFx{\rCURVE})^\hsigFLen \times (\FFx{\rCURVE})^{\resBitsFLen} \to \BB^\crhhsigOutLen \\
    &\unpack{}{}_{\nf{}} : (\FFx{\rCURVE})^\nfFLen \times (\FFx{\rCURVE})^{\resBitsFLen} \to \BB^\prfNfOutLen \\
    &\unpack{}{}_{\vin} : \FFx{\rCURVE}^0 \times (\FFx{\rCURVE})^{\resBitsFLen} \to \BB^\zvalueLen \\
    &\unpack{}{}_{\vout} : \FFx{\rCURVE}^0 \times (\FFx{\rCURVE})^{\resBitsFLen} \to \BB^\zvalueLen
\end{align*}

\begin{figure*}
    \begin{minipage}[t]{.4\textwidth}
        \centering
        \procedure{\pack{x}{\FFx{\rCURVE}}}{%
            out \gets 0_{\FFx{\rCURVE}}; \\
            \pcfor i \in [\len{x}] \pcdo: \\
            \t \pcif x[i] = 1 \pcdo: \\
            \t \t out \gets out +_{\FFx{\rCURVE}} 2^{\len{x}-1-i} \\
            \pcreturn out;
        }
        \caption{Algorithm to pack bits into a field element.}\label{zeth-protocol:fig:packing-alg}
    \end{minipage}%
    \begin{minipage}[t]{.6\textwidth}
        \centering
        \procedure{\packResBits{\nfs{}, \vin, \vout, \hsig, \htags{}}}{%
            out \gets []; r \gets \epsilon; \\
            r \gets \vout; \\
            r \gets \vin \concat r; \\
            %
            \pcfor i \in [\jsin] \pcdo:\\
            \t r \gets \slice{\htags{i}}{\beta_{\htags{i}} \cdot \fieldBitCap}{} \concat r; \\
            %
            \pcfor i \in [\jsin] \pcdo:\\
            \t r \gets \slice{\nfs{i}}{\beta_{\nfs{i}} \cdot \fieldBitCap}{} \concat r; \\
            %
            r \gets \slice{\hsig}{\beta_{\hsig} \cdot \fieldBitCap}{} \concat r; \\
            %
            \pcfor i \in [ \ceil{\len{r}/ \fieldBitCap} ] \pcdo: \\
            \t out[i] \gets \pack{\slice{r}{i \cdot \fieldBitCap}{ (i+1) \cdot \fieldBitCap}}{\FFx{\rBN}}; \\
            \pcreturn out;
        }
        \caption{Algorithm to pack residual bits.}\label{zeth-protocol:fig:packing-resbits-alg}
    \end{minipage}
\end{figure*}

\subsubsection{Packing Policy Security}
\begin{proposition}[Packing security]
    For a binary string $x$, it holds that $\unpack{\pack{x}{}}{} = x$ and $\unpack{\packResBits{x}}{} = x$.
\end{proposition}

\subsubsection{Packing Policy Example}

In the case where $\jsin=\jsout=2$, the \BNCurve~is being used (in which field elements hold $\bnFieldBitCap$ bits) and all PRFs and $\crhhsig{}$ output bit-strings of length 256, the unpacked primary inputs are 2167-bit long.
The packing parameters are therefore:
\begin{align*}
    & \resBitsBLen = 5 \times 3 + 64 + 64 = 143 \\
    & \nfFLen = \hsigFLen = \htagFLen = \resBitsFLen = 1
\end{align*}
The packed primary inputs are 2277 bits long, corresponding to a small space overhead of $\approx 5\%$ unused bits. Moreover, the 143-bit residual bits can be packed into a single field element. As such, the primary inputs are encoded as 9 field elements.
Finally, the residual bits are formatted as follows,
\[
    \underbrace{padding}_{\text{113 bits}}
    \concat \underbrace{\hsig}_{3\ bits}
    \concat \underbrace{\nf{1}}_{3\ bits}
    \concat \underbrace{\nf{0}}_{3\ bits}
    \concat \underbrace{\htag{1}}_{3\ bits}
    \concat \underbrace{\htag{0}}_{3\ bits}
    \concat \underbrace{\vin\vphantom{g}}_{64\ bits}
    \concat \underbrace{\vout\vphantom{g}}_{64\ bits}
\]

For the analogous case using \BLSCurve (in which field elements hold $\blsFieldBitCap$ bits), the packing parameters are:
\begin{align*}
    & \resBitsBLen = 5 \times 4 + 64 + 64 = 148 \\
    & \nfFLen = \hsigFLen = \htagFLen = \resBitsFLen = 1
\end{align*}
The residual bits can be packed into a single field element of the form
\[
    \underbrace{padding}_{\text{108 bits}} \concat \underbrace{\hsig}_{4\ bits} \concat \underbrace{\nf{0}}_{4\ bits} \concat \underbrace{\nf{1}}_{4\ bits} \concat \underbrace{\htag{0}}_{4\ bits} \concat \underbrace{\htag{1}}_{4\ bits} \concat \underbrace{\vin\vphantom{g}}_{64\ bits} \concat \underbrace{\vout\vphantom{g}}_{64\ bits}
\]
and the primary inputs are again encoded as 9 field elements.


% JuliaCon proceedings template
\documentclass{juliacon}
\setcounter{page}{1}
\usepackage{graphicx}

\hypersetup{colorlinks=true}
\begin{document}

\input{header}

\maketitle

\begin{abstract}

Flux \cite{innes:2018} is a machine learning framework, written using the numerical computing
language Julia\cite{bezanson2017julia}. The framework makes writing layers as simple as
writing mathematical formulae, and it's advanced AD,
Zygote \cite{DBLP:journals/corr/abs-1810-07951} , applies automatic differentiation (AD) to
calculate derivatives and train the model. It makes heavy use of Julia's language and
compiler features to carry out code analysis and make optimisations.  For example, Julia's
GPU compilation support \cite{besard:2017} can be used to JIT-compile custom GPU kernels for
model layers \cite{CuArrays.jl}. Flux also supports a number of a hardware options, from CPUs,
GPUs and even TPUs via XLA.jl, that compiles Julia code to XLA: an advanced compiler for
linear algebra that is capable of greatly optimizing speed and memory usage in large
deep learning models.

ONNX.jl is an Open Neural Network Exchange backend for the Flux.jl deep learning framework.
ONNX.jl supports directly importing high quality ONNX standard models into Flux, thus saving
time and reducing the need for additional computation resources. This paper aims at introducing
ONNX.jl and explaining how it fits into the bigger picture: How we can use the Julia Language,
specifically Flux.jl and ONNX.jl as a starting for high quality transfer learning of large deep
learning models.

\end{abstract}

\section{Introduction}

The Julia language was introduced to solve the two language problem: In simple words, languages that
are simple to write (high-level) are very slow but those which are difficult to use (low-level) are
way faster. This is because most of the high-level languages weren't written to process a large
amount of data. Thus engineers, researchers and developers have a hard time developing a lot of
high performance languages. At the moment, the common protocol is to write the core of the software
in a low-level language (C/C++/Fortran) and wrap it in a high-level language (Python). This results
in optimized performance and ease of use. The Julia language aims to make best of both worlds. It
provides a high level syntax but manages to perform as fast as C (sometimes even faster).

Flux.jl is a library for implementing machine learning models, written completely in the Julia
programming language. At the heart of Flux.jl lies Zygote.jl: A source-to-source automatic
differentiation (AD) library that makes complete use of the Julia language compiler to generate
backward pass during training phase of a neural network, with complete support for control flow,
recursion, closures and data structures.
Implementing models in Flux.jl is as simple as writing regular Julia code. Implementing models
is as simple as writing the formulae for those, and Zygote.jl will compute the derivatives
seamlessly.

Flux.jl also provides support for other hardware options using external packages such as CuArrays.jl
and CLArrays.jl. CuArrays is written completely in Julia, making implementing GPU kernels very simple.
Making a model run on GPU can be done in a hassle-free manner: It is as simple as calling a few functions
to transfer data to GPU. Flux.jl also has support for running models on Google's Tensor Processing
Unit (TPU). TPUs help in very fast linear algebra computation. Running Flux models on TPUs is possible
through XLA.jl that compiles Julia code to XLA.

The FluxML ecosystem provides a number of supporting packages that provide additional functionalities , some of them being (apart from the aforementioned Flux.jl, Zygote.jl and XLA.jl);

\begin{itemize}
    \item ONNX.jl : Open Neural Network eXchange backend for Flux.jl
    \item Metalhead.jl \cite{Metalhead.jl}: Simple plug and play pretrained Flux.jl computer vision models.
    \item Torch.jl \cite{Torch.jl}: This package aims at exposing Torch.tensor types in Julia.
    \item IRTools.jl \cite{IRTools.jl} : Provides an IR format that is easy to manipulate.
    \item FluxJS.jl \cite{FluxJS.jl} : Runs Flux models in the browser, via tensorflow.js
    \item model-zoo \cite{model-zoo}: Collection of implementation of various Flux deep learning models.
\end{itemize}

\section{Open Neural Network eXchange (ONNX)}
Open Neural Network Exchange (ONNX) \cite{bai2019} is an open ecosystem that empowers AI developers to
choose the right tools as their project evolves. ONNX provides an open source format for AI
models, both deep learning and traditional machine learning. ONNX defines the computation graph for a deep learning model along with various operators used in the model.
It provides a set of specifications to convert a model to a basic ONNX format, and
another set of specifications to get the model back from this ONNX form.   At a high
level, ONNX is designed to allow framework inter-operability. There are many excellent
machine learning libraries in various languages : PyTorch \cite{bai2019} , TensorFlow \cite{abadi2016tensorflow} , MXNet \cite{chen2015mxnet} , and Caffe \cite{jia2014caffe} are
just a few that have become very popular in recent years, but there are many others as well.
Machine learning models can be converted to a serialized ONNX format which can then be run
on a number devices. ONNX Runtime is an inference engine written in C++ framework used to
deploy ONNX format models into production. It works on diverse hardware and support both
deep learning as well as traditional machine learning models.

\subsection{Where does ONNX come in?}
ONNX is a format for representing deep learning models, which can be further run on numerous devices
without worrying much about the implementation. This helps researchers, developers and engineer to
focus on the problem in hand without worrying much about the peripherals, such as the framework
to use, the ability to run a model trained using this particular framework on specialized hardware.
ONNX is is usable anywhere from small mobile devices to large server farms, across chipsets and vendors,
and with extensive runtimes and tools support. ONNX reduces the friction of moving trained AI
models among your favorite tools and frameworks and platforms. A simple example of how ONNX is ideal
for ML is the case when large deep learning models need to be deployed.     Consider
the simple case of deploying a Deep Learning model to an iOS application. This particular model can
be implemented in any framework : TensorFlow, PyTorch, MXNet just to name a few. However, iOS
applications expect to use CoreML inside the application. Up until now, developers have been porting
large models to different frameworks,
which is a waste of time and energy, better spent somewhere else. This is also retraining the entire model from scratch, which isn't efficient. This makes the entire process cumbersome and impractical. ONNX exists to solve this very problem : By connecting the common dots from
different frameworks, ONNX makes it possible to express a model of type A to type B, thus saving time and the need to train the model again.

\section{ONNX backend in Julia}

ONNX.jl is an ONNX backend written in Julia for the Flux.jl machine learning framework. It can be
used for high-quality inference of pretrained ONNX format machine learning models. At the heart of it, ONNX.jl solves a compiler problem by dealing with intermediate code representations to generate readable graphs. While doing this, ONNX operators are mapped to corresponding Flux layers, thus tracing out the model's computation graph at the end. This graph can then be travelled to generate the Julia code for the model.

The first step towards reading any ONNX model in Julia is to have access to all the data structures in the model,
that essentially hold all the required information to load the model. This information includes everything that
is used to completely define the model: Hyper-parameters and parameters in the model. So in the case of a simple
Convolutional Neural Network, this may contain information such as the number of layers, number of units in each
layer, strides, padding, kernel size, number of filters as hyper-parameters and the trained values associated
with each layer as parameters.

Since ONNX models as protobuf serialized, we need a way to read this serialized data into specific data
structures. ProtoBuf.jl is a Julia implementation of protocol buffers that solves this very issue (covered in
section 3.1). It is used to read the ONNX model into the generated Julia data structures. Once we have the
entire model present as a complex Julia structure, we need to read through this structure and map ONNX operators
to corresponding Flux layers/operations (covered in section 3.2). At the same time, model weights or
parameters are separately stored and saved externally as BSON serialized file. Once the model has been loaded,
we end up with two files: `model.jl`: The Julia code for the machine learning model and `weights.bson`: The
weights associated with the layers defined in the `model.jl` file (section 3.3). In the further sections we'll walk through the internals of these individual processes

\subsection{ProtoBuf.jl}

ProtoBuf.jl \cite{Protobuf.jl} is a Julia implementation of Protocol Buffers. It provides a way to read and write data to and
from Julia types from I/O streams. What ProtoBuf.jl gives from onnx.proto3 is the Julian definition of various
data structures that in themselves have all the required attributes to load any ONNX serialized model. As an
example, as simple message defined in onnx.proto3 as:

\begin{lstlisting}[language=julia]
message NodeProto {
  repeated string input = 1;    // namespace Value
  repeated string output = 2;   // namespace Value

  // An optional identifier for this node in a graph.
  // This field MAY be absent in ths version of the
  // IR.
  string name = 3;     // namespace Node

  // The symbolic identifier of the Operator to
  // execute.
  string op_type = 4;  // namespace Operator
  string domain = 7;   // namespace Domain

  // Additional named attributes.
  repeated AttributeProto attribute = 5;

  string doc_string = 6;
}
\end{lstlisting}{}
\textit{(Code snippet taken from onnx/onnx.proto3)}

Results in the corresponding Julia definition of the model as :

\begin{lstlisting}[language=julia]
mutable struct NodeProto <: ProtoType
    input::Vector{AbstractString}
    output::Vector{AbstractString}
    name::AbstractString
    op_type::AbstractString
    domain::AbstractString
    attribute::Vector{AttributeProto}
    doc_string::AbstractString
    NodeProto(; kwargs...) = (o=new(); fillunset(o); isempty(kwargs) || ProtoBuf._protobuild(o, kwargs); o)
end
\end{lstlisting}

Since ONNX tries to inherit properties from diverse frameworks, ONNX serialized models can be large and complicated. While there are a number of complex generated data structures, three of those are essential towards
understanding how data is stored internally:

\begin{itemize}
    \item ModelProto: A very high-level struct that holds all the information. ONNX models are read directly into
            this structure.
    \item GraphProto: This structure captures the entire computation graph of the model.
    \item NodeProto and TensorProto: Information regarding individual nodes in the graph (inputs, outputs and finer attributes) and weights associated with the nodes.
\end{itemize}

\subsection{ModelProto}

ModelProto structure is the structure that holds all the information needed to load a model. Internally, it holds
data such as the version information, model version, docstring, producer details and most importantly: the
computation graph. \newpage

\begin{lstlisting}[language=julia]
mutable struct ModelProto <: ProtoType
    ir_version::Int64
    opset_import::Vector{OperatorSetIdProto}
    producer_name::AbstractString
    producer_version::AbstractString
    domain::AbstractString
    model_version::Int64
    doc_string::AbstractString
    graph::GraphProto
    metadata_props::Vector{StringStringEntryProto}
end #mutable struct ModelProto
\end{lstlisting}

An ONNX model, once read using ProtoBuf.jl is loaded into this ModelProto object before extracting the graph
details. Naturally, at the heart of this is the \textit{graph::GraphProto} attribute that stores the computation
graph of the model.

\subsection{GraphProto}

The GraphProto structure stores information about particular nodes in the graph. This includes
the node metadata, name, input, output and the pre-trained parameters in the \textit{initializer} attribute.

\begin{lstlisting}[language=julia]
mutable struct GraphProto <: ProtoType
    node::Vector{NodeProto}
    name::AbstractString
    initializer::Vector{TensorProto}
    doc_string::AbstractString
    input::Vector{ValueInfoProto}
    output::Vector{ValueInfoProto}
    value_info::Vector{ValueInfoProto}
end #mutable struct GraphProto
\end{lstlisting}

\subsection{TensorProto}
This is the main structure that holds the raw model parameters. For example, in comvolutional layers, the weights
associated with the kernel as available as dense vectors in the \textit{raw\_data} attribute. During graph
traversal, these weights are extracted and reshaped according the shape that is available as a node attribute.

\begin{lstlisting}[language=julia]
mutable struct TensorProto <: ProtoType
    dims::Vector{Int64}
    data_type::Int32
    segment::TensorProto_Segment
    float_data::Vector{Float32}
    int32_data::Vector{Int32}
    string_data::Vector{Array{UInt8,1}}
    int64_data::Vector{Int64}
    name::AbstractString
    doc_string::AbstractString
    raw_data::Array{UInt8,1}
    double_data::Vector{Float64}
    uint64_data::Vector{UInt64}
end #mutable struct TensorProto
\end{lstlisting}

This is most of the information needed to build the model. In the next section, we discuss how we use DataFlow.jl to
travel this graph and extract model parameters and other relevant information.

\subsection{Graph operations via DataFlow.jl}

Once we have the entire model data present as a ModelProto object, the next step is to travel the computation
graph and capture all the operation being done in the graph while mapping those simultaneously to the
corresponding Flux operators.

DataFlow.jl \cite{DataFlow.jl} is a code intermediate representation format, representing Julia code as an expression graph. It provides functions for graph re-structuring , even on cyclic graphs. Graphs can then also be used to generate Julia expression. It can be efficiently used to traverse our \textit{ModelProto.graph} object. However,
during this traversal we want to map ONNX operators to Flux layers and functions. In DataFlow.jl, this becomes
equivalent to creating a new vertex for the required operator and calling in with appropriate Flux functions
, which are inferred from the ONNX operator itself. As an example, let's consider the simple case of the
\textit{BatchNorm} operator in ONNX. Relu is a commonly used activation function in neural networks that can be
expressed as :

\[ relu(x) = max(x, 0) \]

It basically just turns negative neurons off (sets them to 0) and bypasses positive neurons. The Relu operator
in ONNX performs the same operation on an entire vector elementwise. It takes in a single parameter: The input
vector and returns a single value: the result of applying relu on the input. Using DataFlow.jl, this
operator is mapped as :
\newline \newline
\begin{lstlisting}[language=julia]
ops[:Relu] = function (params, x)
  vcall(broadcast, :relu, x)
end
\end{lstlisting}

The definition of \textit{relu} here is defined in Flux: Once the model is written to an external model.jl file,
we can include the file directly after importing Flux and all definitions should be ready for use.

Other complex layers such as Convolution have more complicated implementations, but the essence remains the same,
to collect all inputs and call them with the corresponding Flux function. During this process, model weights
also computed and stored in a dictionary, mapping layer name to the parameters. At the end of the graph traversal,
we have both the required values: the DataFlow graph containing Flux layers and the model weights corresponding to
each of these layers. The DataFlow graph can converted to Julia code using \textit{DataFlow.syntax} that also
assigns variable names as and when needed. This Julia code is then written to an external \textit{model.jl} file.
For saving the weights, we use the BSON.jl package. BSON stands for Binary JSON, a binary encoded serialization
of JSON objects. BSON.jl \cite{BSON.jl} can be used to store and load such structures, our dictionary containing the model weights being one of them.
\subsection{Interface and Design:}
At a top level, ONNX.jl provides a minimal interface for the user; it is just a tool for loading ONNX format
models. Once the model and weight file has been successfully generated, ONNX.jl provides no further functionality.
Users can then treat the resultant model as any other Flux model.

\begin{lstlisting}[language=Julia]
using ONNX, Flux
ONNX.load_model("path_to_onnx_file")
weights = ONNX.load_weights("weights.bson")
model = include("model.jl")
\end{lstlisting}{}

\textit{ONNX.load\_model} here generates the required model and weights file. Internally, it carries out all the
above mentioned graph operations. \textit{model} above can be treated as any other Flux model.

The significant advantage the ONNX.jl provides is that is treats a compiler problem as a graph problem. It
generates the Flux code for the model, which makes it very easy and intuitive to use the same model for
further applications, such as fine-tuning or even replacing existing layers for some other use case. This
is ideal in the case of applications such as neural style transfer, where it is very common to use a pre-trained
network and modify it a bit as a starting point. The generated code can also be helpful for finer debugging of
the model.

Overall, the entire process from the ONNX serialized file to generation of model and weight file can be
summarized as:

\includegraphics[width=8.8cm]{onnx-3.png}

\begin{center}
    \textit{Fig. 1: Flow diagram of ONNX.jl}
\end{center}

Additionally, ONNX.jl also provides helper functions for inspecting the model before loading it. \textit{ONNX.layers} reads an ONNX file and returns a list of all the layers in the model. With the growing
interest around more complicated and deep models, it is possible that an ONNX model might have layers that Flux
itself doesn't support. For handling these, ONNX.jl leaves a \textit{hook} for the users to implement
additional functionality. A \textit{hook} is a function that doesn't have an existing implementation: one
would have to write an implementation for it themselves. However any operator that also has a corresponding
implementation in Flux is completely recognized by ONNX.jl at the moment.

\section{Usage Scenarios}
The ONNX format and ONNX.jl can be used for transfer learning in Flux, where we store knowledge while training a model and
use this knowledge for some other task. The idea is that rather than random initialization of parameters for training a neural network, it's better to take an already trained model, since
it leads to faster convergence. In transfer learning, we take a pretrained model and train it on another
dataset, which might also have a different class distribution. Fine tuning is an approach to transfer learning
where we train on a subset of training data with a smaller learning rate. Transfer learning learning has shown
tremendous results in image classification, object detection, simulations, sentiment and NLP based classification
in recent past.
This is also pretty common when talking about tasks such as neural style transfer where we want to change the
style of an image in accordance with the style of another image. Generative Adversarial Networks (GANs) have shown
to deliver high quality results when trained on top of a pre-trained model. StyleGAN \cite{DBLP:journals/corr/abs-1812-04948} , for example, can use a
pre-trained model to train a custom model to deliver high quality super-resolution results.

\section{Related Work}

In recent times several projects have come up that solve similar issue. One of the most notable project is
TensorFlow's mlir (Multi-Level Intermediate Representation) \cite{lattner2020mlir} . mlir is an evolution of LLVM
\cite{LLVM:CGO04} that defines a
common Intermediate Representation (IR) format, which can be used to represent any DataFlow graph. This common
format unifies machine learning models in TensorFlow or other frameworks.
Other noteworthy approaches in this direction are PFA and NNVM. PFA \cite{10.1145/2939672.2939731} or Portable Format for Analytics is a common
language that aims as easing the transition from development to production. It can be expressed within
the common JSON format and has functionalities such as control structures, loops and user-defined functions.

NNVM\cite{nnvm} is an end-to-end compiler for AI frameworks. It aims at solving the challenges posed by using different
diverse machine learning frameworks. It consist if two major components: NNVM (Neural Network Virtual Machine)
and TVM (Tensor Virtual Machine) \cite{article} . NNVM defines a common computational graph intermediate representation
format and TVM implements the operators used in these computation graphs while optimizing them for the backend hardware.

\section{Future Work}

As ONNX.jl becomes the beginning point for various researchers interested in using Julia for their research,
it is important to note that it also has certain shortcomings. The most significant is that a model can't
be completely loaded unless there's an equivalent implementation of the operator in Flux.jl. An example of this
is Grouped Convolutions. These variants of Convolutional layers were used in AlexNet \cite{NIPS2012_4824} and showed amazing
results. However, since Flux doesn't support these at the moment, the users will need to have an implementation
ready if they choose to import an ONNX model with this particular layer into Flux using ONNX.jl. On the plus side,
a lot of the most commonly used layers are available in Flux and can be readily used. Another thing to note is
that to run ONNX.jl generated code in some other hardware, one might need to do a little restructuring. The model
should work directly on the CPU.

Another challenge moving forward is that we need to constantly update ONNX.jl to support Flux's latest API
changes. This also applies for the other way round: As ONNX operators are updated, we would have to update these
corresponding implementations in ONNX.jl to support the newer models using these updated specifications. Over
the long run, we'd need to constantly keep an eye out for such changes and adapt ONNX.jl to those. Moreover,
there are subtle differences in the way Flux implements operators as compared to other frameworks. As an example,
consider the simple \textit{AveragePool} layer. The way Flux implements this is by padding the input tensor
appropriately and then performing the pooling operation. However, Keras-tensorflow does this by pooling and then
padding. This leads to subtle changes along the edges of the tensor in the output. Such differences occur due to
the way most frameworks deal with such layers, and the only way to avoid this is to check for such
discrepancies.

In recent past, DataFlow.jl has been superseded by another intermediate representation format tool: IRTools.jl.
It provides the ability to work with both lowered and typed Julia code and can be used together with
metaprogramming tools such as Cassette.jl.

There has also been some talk about splitting ONNX.jl into two packages: The first one would do the code
generation and DataFlow related functions while the other would be solely responsible for implementation of the
ONNX operators. This would be greater control and ease while implementing layers or debugging a loaded model. This
should also make implementation pretty straight-forward wherever missing.
For the moment, all this continues to be done by a single package.

\section{Conclusion}

Developing ONNX.jl has been tremendous learning experience. From studying about Intermediate Representation formats for deep
learning models with millions of parameters to loading them in just a couple of lines of code, ONNX.jl has
made it very easy and straight-forward to use a high quality trained model as a starting point for many projects.
Once such example I'd like to point out is DenseNet-121 model. This is deep convolutional network that has
multiple Convolutional, Dense and Pooling blocks. Naturally, implementing this in any framework is going to be
a challenging task. However, thanks to ONNX, we can now use an earlier implementation to import this model into
any other framework. Importing this model (train in Caffe2) via ONNX.jl into Flux can be done in ~3 lines of code.

I was also able to load several large computer vision models loaded from ONNX format at the time of actively
developing the package. Most of these have been added to FluxML/Metalhead.jl for a direct plug-and-play use. These
included:

\begin{itemize}
    \item SqueezeNet \cite{i2016squeezenet}
    \item DenseNet 121  \cite{huang2016densely}
    \item ResNet 50     \cite{he2015deep}
    \item GoogleNet     \cite{szegedy2014going}
\end{itemize}

ONNX.jl serves as a entry point for people looking to use Flux for their research, but want quick results. It
combines the power of the Julia language, the elegance of Flux and the availability of a vast number of pre-trained models. This enables researchers to spend time focusing on the real issues, rather than model portability. \newpage

\input{bib.tex}

\end{document}

% Inspired by the International Journal of Computer Applications template


% !TeX root = ./paper.tex
\documentclass[modern]{aastex62}

% Load the corTeX style definitions
\input{cortex}

% Bibliography stuff
\bibliographystyle{aasjournal}

% Begin!
\begin{document}

% Title
%\title{Inferring a time dependent map of Io's surface from occultations and phase curves.}
\title{Occultation mapping of Io's surface in the near-infrared II: Inferring time variable maps using latent Nonnegative Matrix Factorization}

% Author list
\author{Fran Bartoli\'c}
\email{fbartolic@flatironinstitute.org}
\affil{Center~for~Computational~Astrophysics, Flatiron~Institute, New~York, NY}
\affil{Centre for Exoplanet Science, SUPA, School of Physics and Astronomy, University of St. Andrews, St. Andrews, UK}
\author{Rodrigo Luger}
\author{Daniel Foreman-Mackey}
\affil{Center~for~Computational~Astrophysics, Flatiron~Institute, New~York, NY}
%

\begin{abstract}

\end{abstract}

%
\section{Introduction}
\section{Data}
\label{sec:data}

\begin{figure}[h!]
    \begin{centering}
    \includegraphics[width=0.5\linewidth]{figures/irtf_max_flux.pdf}
    \oscaption{irtf_dataset}{%
        Maximum flux for each curve taken using NASA's IRTF telescope during an
        occultation of Io by Jupiter.
       \label{fig:irtf_max_flux}
    }
    \end{centering}
\end{figure}

The dataset we use to infer maps of Io's surface consists of 112 observations of Io in the near infrared taken from 1996 until 2018 using various instruments at NASA's IRTF observatory.
Figure~\ref{fig:irtf_max_flux} shows the maximum flux in $\mathrm{GW}/\mu \mathrm{m}/\mathrm{sr}$ for each light curve plotted as a function of time which approximately corresponds to the disc integrated flux of Io at the beginning of ingress or the end of egress prior to it being occulted by Jupiter.
Two features of this plot are apparent.
First although the observations span decades the cadence is non-uniform with only a few observations taken between 2008 and 2016.
Second, the baseline brightness is varying stochastically by a large amount and there are several notable events of increased volcanic activity.
All of the observations were taken while Io was \emph{in eclipse}, meaning that it was in Jupiter's shadow and all of the observed emission from the surface is due to thermal radiation.
Observations of Io \emph{in sunlight} on the other hand probe both the thermal (volcanic) emission and the surface albedo variations in the near infrared.

\begin{figure}[h!]
    \begin{centering}
    \includegraphics[width=0.5\linewidth]{figures/irtf_sample_lightcurves.pdf}
    \oscaption{irtf_dataset}{%
        A selection of sample light curves taken during occultations of Io by Jupiter in our dataset.
        The step-like morphology of the light curves is due to bright volcanoes on Io's surface coming in or out of view during an occultation.
        These light curves thus visibly encode information about the features on the surface.
        \label{fig:irtf_sample_lightcurves}
    }
    \end{centering}
\end{figure}

In Figure~\ref{fig:irtf_sample_lightcurves} we display a random subset of all the light curves
in the dataset.
All occultations last for $\sim4$ min and the cadence for each light varies but it on the order of a second.
It follows that over the course of a single exposure Jupiter's limb move by about 15km which provides a lower bound for the size of the features on the surface that we can reliably estimate.
The shapes of light curves strongly deviate from the smooth variability one would expect assuming a homogeneous distribution of thermal emission on the surface.
Especially prominent are light curves with clear step-like features which are present when bright spots come in or out of view during the course of an occultation.
The fact that these features are so clearly visible means that even individual light curves encode a wealth of information about the surface.

The photometric quality varies from year to year because multiple instruments were used over the years.
An additional issue is that estimated errorbars aren't provided.
As with all ground based photometry, the observations are influenced by atmospheric variability which induces correlated noise in the light curves.
Because of this the flux isn't always monotonically increasing or decreasing as one would expect.
All of these issues need to be accounted for in the final model.

\section{Nonnegative Matrix Factorization}
\label{sec:nmf}
In the case when we cannot assume a static map, the situation is considerably more complicated.
In principle, the surface map is different for each \emph{data point}.
Of course, fitting one map per data point is intractable so we might instead fit a single map per light curve, generated by Eq.~\ref{eq:linear_model_pix}.
Even then, we would need to fit on the order of 1k parameters!
This is both becuse the number of spherical coefficients scales as $(l+1)^2$ and because when fitting in pixel space we need many more pixels than spherical harmonics to ensure positivity.
Although fitting a model of such high dimensionality for a single IRTF light curve of high signal to noise ratio is tractable thanks to exact gradients provided by autodifferentiation in \textsf{starry}, we don't want to do that for a hundred or so separate light curves.
Such a model would be very difficult to fit and it would require strong regularization between successive maps in time.
It also wouldn't directly provide much physical insight into the volcanic activity, we would have to conduct a separate analysis on the inferred maps to constrain the time variability of the volcanoes.
Instead, we need an approach which reduces the dimensionality of the problem.
One way of accomplishing that is to expand the spherical harmonic coefficients (or pixels) in a Taylor series in time about a certain point.
This is the approach \cite{luger_tess_2019} took to model a map of Earth's albedo in reflected light from stray Earthshine in the aperture of the TESS space telescope.
This issue with this approach and similar series expansions such as the Fourier series is that with the very long baseline of IRTF data and the observations being scattered throughout the years sporadically, it is unclear which point should be the origin of the expansion.
More importantly, there likely doesn't exist a well defined global time-averaged map of Io so requiring too much smoothness between successive maps wouldn't work.
Another way of reducing the dimensionality of the problem would be a parametric approach where where we place a few spots on the sphere and then fit only for the locations and intensities of each spot.
This is assuming that we know a priori that the surface features are spot like and roughly how many we should expect.
While such a strong assumption might work for studying individual spots in the context of Io, it certainly isn't justified for exoplanets.
In Section~\ref{sec:inverse_problem} we do assume stronger priors on the locations of and shapes of surface features to better constrain the variability of individual volcanoes, but we still fit for pixels, at least in principle allowing the data to override our assumptions.

The approach approach opted for is to treat the problem of inferring a single map per light curve as a probabilistic matrix factorization problem.
The idea is to assume that a map for any given light curve can be expressed as a linear combination of $K$ "basis maps".
Consider $L$ lightcurves, the model for the $l$-th light curve is
\begin{equation}
    \mathbf{f}_l=\mathbf{A}_l\,\mathbf{P}^\dagger\,\mathbf{p}_l'
\end{equation}
We can stack the column vectors $\mathbf{p}_l$ into a matrix $\mathbf{Y}$ of shape $(N_p, L)$ where $N_p$ is the number of pixels for each light curve.
We then assume that the $\mathbf{Y}$ can be decomposed into a product of two matrices $\mathbf{B}$ and $\mathbf{Q}$ as
\begin{equation}
    \mathbb{Y}=\mathbf{B}\,\mathbf{Q}
    \label{eq:nmf}
\end{equation}
where $\mathbf{B}$ has shape $(N_p, K)$ and $\mathbf{Q}$ has shape $(K, L)$.
The model for all light curves can be written as
\begin{equation}
    \mathbb{f}=\mathbb{A}\,\mathrm{vec}(\mathbf{\mathbb{Y}})
    \label{eq:model_all_lcs}
\end{equation}
Where $\mathbf{f}$ is a tall column vector consisting of predictions for all light curves stacked together, $\mathbb{A}$ is a block diagonal matrix with matrices $\mathbf{A}_l\,\mathbf{P}^\dagger$ on the diagonal for $l=1\dots, L$ and the $\mathrm{vec}$ operator stacks the columns of $\mathbb{Y}$ into a tall vector.
The model is \emph{bilinear} which means that is linear in the matrices $\mathbf{B}$ and $\mathbf{Q}$ separately.
The interpretation of Equation~\ref{eq:nmf} is simple, each of the $K$ columns of $\mathbf{B}$ represents pixels of a basis map and the columns of $\mathbf{Q}$ determine how those maps add together to produce a final map for the $l$-th light curve.
For this reason we call $\mathbf{B}$ the \emph{basis matrix} and the matrix $\mathbf{Q}$ the \emph{encoding matrix}.
Ideally, we want the basis maps to be physically meaningful and the coefficients encode the time variability.
A nice feature of this model is that there is no requirement that successive maps are smoothly varying between different light curves.

The matrix factorization problem as written in Equation~\ref{eq:nmf} is highly degenerate and practically intractable a probabilistic framework with a small dataset.
Since each of the columns of $\mathbf{B}$ are pixels representing emitted light or albedo, we can substantially reduce the ambiguity in the decomposition by requiring that both matrices are strictly positive.
This matrix factorization problem is known as \emph{Nonnegative Matrix Factorization} (NMF) \citep{paatero_positive_1994,lee_algorithms_2001} and it has a rich history across many different fields.
It is commonly used for decomposing physical signals in the spectral or the time domain.
Even though NMF is more tractable than unconstrained matrix factorization, it is still an NP hard problem \citep{vavasis_complexity_2009} and it requires additional constraints or priors to be tractable.
Simultaneously transforming $\mathbf{B}\leftarrow \mathbf{B}\,\mathbf{S}^{-1}$ and
$\mathbf{Q}\leftarrow \,\mathbf{S}\mathbf{Q}$ with a nonsingular matrix $\mathbf{S}$ doesn't change the value of the objective function.
The objective function is also invariant to a simultaneous permutation of the columns of $\mathbf{B}$ and rows of $\mathbf{Q}$ and rescaling either of the two matrices by a scalar.
Nevertheless, getting around these degeneracies in practice is possible, especially when we have a lot of prior knowledge on the problem.
For a recent review article with a detalied analysis of these degeneracies and common algorithms, see \cite{fu_nonnegative_2019}.

Broadly, there are two approaches to fitting NMF models, the constrained optimization approach and the probabilistic approach.
In the optimization approach we optimize $\mathbb{Y}$ in a least-squares sense under a set of constraints on either one or both of the matrices \citep[see][for recent examples from the astronomical literature]{acosta-pulido_new_2017, melchior_scarlet_2018,ren_non-negative_2018,ren_using_2020} with the goal of obtaining a point estimate for the two matrices.
The probabilistic approach introduces priors and the inference is usually done with variational inference in order to obtain an approximate posterior distribution over the matrices.
In this work we opt for the probabilistic approach because we care about the uncertainties on the map features a lot.

One difference between most applications of NMF in the literature and NMF within the context of inferring surface maps from 1D light curves is that
in the majority of the work in the literature \citep[except][]{kawahara_global_2020} the matrix $\mathbb{Y}$ is assumed to be directly observed  up to a simple noise term.
In our case the problem is considerably more challenging because we observe the light curves
$\mathbb{A}\,\mathrm{vec}(\mathbf{\mathbb{Y}})$ instead of $\mathbb{Y}$ directly, a process in which some information is lost because, depending on the geometry of the occultations, certain combinations of spherical harmonic coefficients will be in the nullspace.
\cite{kawahara_global_2020} are using NMF to model phase variations in directly imaged exoplanets with the goal of simultaneously inferring a surface map and a spectral decomposition of the map into several components, a case in which the
matrix which is to be decomposed is also unobserved.
Although we don't attempt to solve the problem of spectrally decomposing the maps, our approach is directly applicable to that problem as well. (TODO: should I add more details here?)
One could also imagine inferring a multi spectral component and time dependent map by solving two NMF problems simultaneously.
We leave this application for future work.

\subsection{The likelihood}
\label{ssec:likelihood}
Given Eq.~\ref{eq:nmf} we can write down the likelihood for our data and in Section~\ref{sec:inverse_problem} we dicuss in detail the priors which make the problems of inferring static and dynamic maps tractable.
Assuming a Gaussian noise process for the data with a dense covariance matrix, the (log) likelihood is
\begin{equation}
    \ln\mathcal{L}=-\frac{1}{2}\left[\mathbb{f}_\mathrm{obs}-(\mathbb{f} + \mathbb{b})\right]^{\top} \boldsymbol{\Sigma}^{-1}\left[\mathbb{f}_\mathrm{obs}-(\mathbb{f} + \mathbb{b})\right]
\end{equation}
where $\mathbb{f}_\mathrm{obs}$ is the vector of stacked observed light curves, $\boldsymbol{\Sigma}$ is the data covariance matrix and $\mathbb{b}$ is a fixed flux offset per light curve accounting for stray flux not attributed to Io (mostly due to Jupiter).
We model the data covariance with a Gaussian Process using the fast Celerite method\citep{foreman-mackey_fast_2017} as implemented in the \textsf{exoplanet} package (TODO: cite) to capute correlations imposed by the seeing plus an additional white noise term which is different for each light curve.
The IRTF data isn't provided with any uncertainties so we estimate the uncertainties by filtering the data with the Savitzky Golay filter as implemented in \textsf{SciPy} (TODO: CITE) and estimating the variance of the residuals.
We assume identical errors for all data points in a given light curve but allow for a rescaling factor when fitting the model.
In the next section, we add priors to this likelihood and fit the model on simulated data using optimization and Variational Inference (VI).

\section{The inverse problem}
\label{sec:inverse_problem}

\subsection{Fitting a dynamic map}
\label{ssec:dynamic_map}
% Priors for NMF. Optimization vs. bayesian approach to NMF. Results on simulated data. Comparison of results to matrix factorization without the positivity constraint and clever priors.

\section{Results}
\label{sec:results}

\subsection{Individual events}
\label{ssec:individual}
% Fits of individual events. In particular, those light curves that were analyzed in previous papers.

\subsection{The time-variable map}
\label{ssec:time_variable_map}
%

\subsection{Variability of known hotspots}
\label{ssec:variability_hotspots}
% Plot inferred intensities of known hotspots as function of time, reference previous work.

\section{Mapping volcanic exoplanets}
\label{sec:exoplanets}
% Application to exoplanets. Fitting mock JWST observations.

\section{Conclusions}
\label{sec:conclusions}

% Bibliography
\bibliography{bib}
\end{document}