%===========================
% Supplementary Information
%===========================

\documentclass[12pt]{article}

%---- Page setup
\usepackage[margin=1in]{geometry}
\usepackage{setspace}
\onehalfspacing

%---- Math & symbols
\usepackage{amsmath, amssymb, amsfonts, bm}

%---- Figures, tables, floats
\usepackage{graphicx}
\usepackage{caption}
\usepackage{subcaption}
\usepackage{booktabs}
\usepackage{float}

%---- Hyperlinks
\usepackage[hidelinks]{hyperref}
\usepackage[super,sort&compress]{natbib}



%---- Numbering with "S" prefix
\renewcommand{\thefigure}{S\arabic{figure}}
\renewcommand{\thetable}{S\arabic{table}}
\renewcommand{\theequation}{S\arabic{equation}}


\usepackage[nomarkers]{endfloat} 
\renewcommand{\efloatseparator}{\mbox{}} % avoid extra separator text



\begin{document}


\section*{Supplementary Information (S1)}

This document provides additional methods, mathematical formulas, hyperparameters, and supporting results for the main text. It includes Supplementary Figures S1--S15 and Supplementary Tables S1--S3, which are cited within this Supplementary Information.


%==========================================================
% PART 3: material from old supplement that can remain in Supplementary
%==========================================================

\section{Supporting analysis of missing-value initialization for the \textit{Drosophila} atlas}
\label{supp:prefill_comparison}

To assess how the initial placeholder choice influences downstream VAE reconstruction, we compared three initialization schemes for missing gene entries in the \textit{Drosophila} atlas: zero-fill, mean-fill, and random-fill. These schemes were evaluated under two VAE objectives: the full objective, which included reconstruction, KL regularization, smooth sparsity, decoder smoothness, and directional alignment; and a reduced objective omitting the decoder-smoothness and directional-alignment terms.

Fig.~\ref{fig:reconsloss} summarizes test-set reconstruction performance for the 99-gene decoder as a function of latent dimensionality. Across both training objectives, random-fill initialization consistently yielded the lowest mean absolute error (MAE), the lowest mean squared error (MSE), and the highest mean coefficient of determination (\(R^2\)), whereas zero-fill performed worst. Based on this comparison, random-fill was used as the default initialization scheme in subsequent \textit{Drosophila} experiments.



\section{Supplementary analyses of teacher-guided mechanistic fitting}
\label{supp:teacher_guidance}

This section provides additional qualitative analyses supporting the teacher-guided versus data-only comparisons reported in the main text. The main quantitative comparison is summarized in Table~1 of the main manuscript, where teacher-guided fits achieve equal or lower one-step forecast error after calibration, with the clearest advantage under stronger sparsity.

Here we focus on complementary visualizations that clarify how the black-box teacher shapes downstream mechanistic fitting. Fig.~\ref{fig:bb_piecewise} shows two representative black-box gene-space realizations over each interval \([t_k,t_{k+1}]\): an anchored pushforward realization and a decoded latent-path realization. For this figure only, the decoded-path realization was obtained by linear interpolation in latent space between consecutive encoded points before decoding, in order to isolate the contribution of the decoder separately from full latent-dynamics integration.

Fig.~\ref{fig:MAE4methods} compares mean absolute prediction error across the 27 fully observed genes between the black-box model, the Hill model, and previously reported erf-weighted LAD baselines. Although the LAD quadratic model achieves the lowest raw error, the Hill model offers the complementary advantage of explicit mechanistic nonlinearity and equation-level editability.

Finally, Fig.~\ref{fig:bb_vs_hill_trajs} overlays black-box and Hill-model trajectories directly in the full 99-dimensional gene space. Solid blue dots denote measured values and hollow blue circles denote VAE-imputed values. Blue curves correspond to black-box teacher trajectories, whereas red curves correspond to the fitted Hill-model trajectories. This figure provides a qualitative view of how well the mechanistic student reproduces the learned black-box dynamics in the full partially observed setting.





\section{Regional robustness of mechanistic fits across the embryo}
\label{supp:universality}

This section provides additional analyses of how spatial training coverage affects mechanistic generalization across the embryo. We partitioned the embryo either along the anterior--posterior axis into anterior, middle, and posterior subsets, or along the dorsal--ventral axis into dorsal, medial, and ventral subsets. These regional analyses parallel earlier region-restricted studies, but here are applied to Hill fields fitted from latent-derived teacher velocities.

For consistency across regions, we did not retrain a separate latent neural ODE for each case. Instead, regional Hill models were trained from a common source of black-box-derived gene-space velocities obtained from a fixed VAE pushforward with trivial latent interpolation. This design ensured that differences across regional fits reflected spatial training coverage rather than differences in upstream representation learning.





As a baseline, we also trained random-30\% models by uniformly sampling 30\% of bins from the training pool, matching the effective training size of each regional partition. To reduce variance, the random-30\% sampling was repeated with 10 different seeds and the resulting errors were averaged. Prediction quality was summarized by the relative error
\[
\frac{|g-\hat g|}{|g+\hat g|+\varepsilon},
\]
where \(g\) is the ground-truth expression and \(\hat g\) is the Hill-model prediction. This normalization reduced domination by highly expressed genes and made error magnitudes more comparable across genes.




Fig.~\ref{fig:universality_amp_dmv_l11en4} and Fig.~\ref{fig:universality_amp_dmv_l11en5} show that balanced spatial coverage improves generalization, with the benefit becoming clearer as sparsity is relaxed. Under stronger regularization, all regional models perform broadly similarly, but the random-30\% baseline still provides a small advantage. Under weaker regularization, that advantage becomes larger, indicating that broad spatial sampling is especially helpful when the fitted network retains more couplings.






The spatial overlays in Fig.~\ref{fig:spatial_amp_dmv} show that region-restricted supervision produces non-uniform error patterns across the embryo. By contrast, the random-30\% baseline gives a more spatially uniform and consistently lower error profile, particularly in posterior and ventral domains where single-region fits tend to overfit. This improvement is not confined to a few genes or locations, but extends across the 27-gene set and the full embryo, supporting the value of balanced spatial coverage during mechanistic distillation.




%%%%%%%%%%==============Supp Figures
\begin{figure}[t]
  \centering
  \includegraphics[width=0.9\textwidth]{FigureS1.png}
  \caption{\textbf{Universal VAE reconstruction for a 1\,s EEG segment (subject S001, run R01).}
  Each panel shows one of the 64 EEG channels (blue: recorded signal; orange: VAE reconstruction of the center sample in each 9-point window).}
  \label{fig:eeg_univvae_64ch_full}
\end{figure}


\begin{figure}[t]
  \centering
  \includegraphics[width=0.9\textwidth]{FigureS2.png}
  \caption{\textbf{Zoomed view of four EEG channels from Fig.~\ref{fig:eeg_univvae_64ch_full}.}
  Blue dots show measured EEG samples and blue curves show universal VAE center reconstructions, illustrating preservation of both slower trends and fast fluctuations.}
  \label{fig:eeg_univvae_4ch_zoom}
\end{figure}


\begin{figure}[t]
  \centering
  \includegraphics[width=0.7\textwidth]{FigureS3.png}
  \caption{\textbf{In silico run-control perturbation in EEG.}
  For subject S033, we compare baseline run control R01 with a control-direction perturbation toward R03 while keeping the same initial states as in Fig. 4 of the main text. %.~\ref{fig:eeg_traj_short}.
  We compute the change in final predicted amplitude \(\Delta x_{\text{final}} = x_{\text{alt}} - x_{\text{base}}\) after a 9-step rollout and aggregate across channels and windows.
  The distribution is shifted below zero (vertical line at 0), indicating a systematic decrease in predicted amplitude under the counterfactual control.}
  \label{fig:eeg_hist_delta}
\end{figure}


\begin{figure}[t]
  \centering
  \includegraphics[width=0.9\textwidth]{Figure_S4_NPJSBA.pdf}
  \caption{\textbf{AirQualityUCI VAE reconstruction scatter on a held-out test subset.}
  Each panel shows a pollutant variable.
  Points compare the measured center value (x-axis) to the VAE-reconstructed center value (y-axis) from 3-hour windows.
  The dashed line indicates \(y=x\).
  Titles report coefficient of determination (\(R^2\)) values computed using measured entries only.}
  \label{fig:aq_vae_recon_scatter_test2_r2}
\end{figure}


\begin{figure}[t]
  \centering
  \includegraphics[width=0.9\textwidth]{Figure_S5_NPJSBA.pdf}
  \caption{\textbf{Local comparison of linear pre-imputation versus post-VAE values within a missing gap.}
  For each variable (\texttt{C6H6\_GT\_}, \texttt{PT08\_S3\_NOx\_}, \texttt{PT08\_S5\_O3\_}), the plot shows a local time window centered around the largest missing gap (shaded).
  Blue: measured values.
  Black: linear-imputed values in the missing interval.
  Red: post-VAE values in the missing interval.}
  \label{fig:aq_linear_vs_postvae_local_windows}
\end{figure}


\begin{figure}[t]
  \centering
  \includegraphics[width=0.3\linewidth]{FigureS6L.png}
  \hspace{5mm}
  \includegraphics[width=0.55\linewidth]{FigureS6R.png}
  \caption{\textbf{Embryo layout and UMAP embedding of latent states.}
  Left: schematic of the blastoderm embryo with 6{,}078 spatial bins. 
  Right: UMAP visualization of latent encodings across all bins, shown separately for the six developmental time points ($t=0$ to $t=5$, left to right). Red corresponds to posterior bins of the virtual embryo and blue corresponds to anterior bins.}
  \label{fig:vae_umap}
\end{figure}



\begin{figure}[t]
  \centering
  \includegraphics[width=\linewidth]{FigureS7.png}
  \caption{\textbf{VAE reconstruction accuracy versus latent dimension on the testing set.}
  Rows correspond to two training objectives: the full five-term loss (top) and a reduced objective omitting the decoder-smoothness and directional-alignment terms (bottom). Columns show mean absolute error (MAE), mean squared error (MSE), and mean coefficient of determination (\(R^2\)). Curves compare zero-fill, mean-fill, and random-fill initialization. Across latent dimensionalities, random-fill consistently gives the lowest error and highest \(R^2\). The zero-initialization point at latent dimension 8 is an outlier and is omitted from the plot.}
  \label{fig:reconsloss}
\end{figure}



\begin{figure}[t]
  \centering
  \includegraphics[width=\linewidth]{FigureS8.png}
  \caption{\textbf{Two black-box trajectory realizations in the 27-dimensional gene space for one representative cell.}
  The underlying VAE was trained with $\lambda_{\ell_1}=10^{-4}$ (see Table 1 of the main text). %~\ref{tab:teacher_vs_not}).
  Blue dots are observed expression at discrete times.
  Solid blue curves integrate in gene space starting from the observed data, eliminating reconstruction error at the segment start.
  Dashed red curves decode latent linear interpolations between neighboring encoded states, incurring reconstruction error at both ends of the curves.
  Gene names are shown above each subplot; segments are not connected across boundaries.}
  \label{fig:bb_piecewise}
\end{figure}



\begin{figure}[t]
  \centering
  \includegraphics[width=\linewidth]{FigureS9.png}
  \caption{\textbf{Mean absolute prediction error on the testing set.}
  Comparison of mean absolute errors for our black-box model and the Hill model, alongside results reported in \cite{han2021mechanistic}.
  The connecting lines are included solely for visual guidance and do not carry interpretive meaning.}
  \label{fig:MAE4methods}
\end{figure}

\begin{figure}[t]
  \centering
  \includegraphics[width=\linewidth]{FigureS10.png}
  \caption{\textbf{Black-box versus Hill-field trajectories in 99-dimensional gene space.}
  Solid blue dots mark observed gene-expression measurements; hollow blue circles indicate values imputed by the VAE.
  Blue curves are the decoder applied to the latent neural ODE trajectories.
  Red curves are integrations of the corresponding Hill model trained from this black-box teacher.
  Subplot titles shown in green denote the 27 genes that are fully observed across all six time points.}
  \label{fig:bb_vs_hill_trajs}
\end{figure}


\begin{figure}[t]
  \centering
  \includegraphics[width=0.9\linewidth]{FigureS11.png}
  \caption{\textbf{Spatial partitioning of the embryo into anterior--middle--posterior (AMP) and dorsal--medial--ventral (DMV) subsets.}
  Cells are colored according to their regional assignment: anterior/posterior or dorsal/ventral in deep green, medial in orange.
  Black dots indicate the fixed 10\% test cells that were excluded from training.
  These partitions define the subsets of cells used for region-specific Hill-model fitting.
  As a baseline, we also trained models on a random selection of 30\% of bins uniformly distributed across the embryo (random-30\%), matching the number of bins used in each regional subset.}
  \label{fig:partitions_amp_dmv}
\end{figure}



\begin{figure}[t]
  \centering
  \includegraphics[width=0.48\linewidth]{FigureS12L.png}
  \includegraphics[width=0.48\linewidth]{FigureS12R.png}
  \caption{\textbf{Universality of Hill-model fitting across regions for $\lambda=10^{-4}$ (stronger regularization).}
  Mean relative error per gene for models trained on subsets of cells from either the anterior--middle--posterior (AMP, left) or dorsal--medial--ventral (DMV, right) partition, compared to models trained on randomly selected 30\% of bins across the embryo (random-30\%).
  Under stronger sparsity, the random-30\% baseline already provides a modest advantage, suggesting that balanced spatial coverage improves generalization even when couplings are heavily pruned.}
  \label{fig:universality_amp_dmv_l11en4}
\end{figure}


\begin{figure}[t]
  \centering
  \includegraphics[width=0.48\linewidth]{FigureS13L.png}
  \includegraphics[width=0.48\linewidth]{FigureS13R.png}
  \caption{\textbf{Universality of Hill-model fitting across regions for $\lambda=10^{-5}$ (weaker regularization).}
  Mean relative error per gene for models trained on subsets of cells from either the AMP (left) or DMV (right) partition, compared to random-30\%.
  With weaker sparsity, the advantage of random-30\% becomes more pronounced: balanced coverage consistently outperforms any single-region training, indicating that broader sampling is especially valuable when the network retains more couplings.}
  \label{fig:universality_amp_dmv_l11en5}
\end{figure}









\begin{figure}[t]
  \centering
  \includegraphics[width=0.48\linewidth]{FigureS14L.png}
  \includegraphics[width=0.48\linewidth]{FigureS14R.png}
  \caption{\textbf{Spatial distribution of Hill-model errors for different regional partitions.}
  Relative errors for the 27 fully measured genes are visualized on the embryo surface at latent dimension \(10\) with \(\lambda=10^{-4}\).
  Panels (a--c) show models trained on individual regions (anterior, middle, posterior for AMP; dorsal, medial, ventral for DMV).
  Panel (d) shows the random-30\% baseline.
  In both partitions, random-30\% yields a more spatially uniform and consistently lower error distribution, especially in posterior and ventral domains where regional training tends to overfit.}
  \label{fig:spatial_amp_dmv}
\end{figure}



\begin{figure}[t]
  \centering
  \includegraphics[width=\linewidth]{FigureS15.png}
  \caption{\textbf{Temperature-swap sensitivity of pushforward velocities (January \(\rightarrow\) July).}
  Histograms show the change in instantaneous pushforward slope \(\Delta \dot{x}\) per channel when replacing the temperature control in January with the July temperature at the same day and hour while keeping all other controls fixed.
  Vertical lines indicate \(\Delta \dot{x}=0\).}
  \label{fig:aq_temp_swap_hist}
\end{figure}




%%%%%+++============SuppFigures end here
%%%%%+++============SuppFigures end here
%%%%%+++============SuppFigures end here

\begin{table}[H]
\centering
\small
\begin{tabular}{@{}p{12.2cm}@{}}
\toprule
$\displaystyle \mathcal{L}_{\mathrm{recon}}
=\frac{1}{B}\sum_{b=1}^{B}
\frac{\big\|m_t^{(b)}\odot(\hat g_t^{(b)}-g_t^{(b)})\big\|_2^2}
{\sum_i m_{t,i}^{(b)}+\varepsilon}$ \\[6pt]
$\displaystyle \mathcal{L}_{\mathrm{KL}}=\tfrac{1}{2}\sum_{j=1}^{d_z}\!\big(\mu_j^2+\sigma_j^2-\log\sigma_j^2-1\big),\qquad
w_{\mathrm{KL}}=\beta\cdot\min\!\Big(1,\frac{\text{epoch}+1}{\text{warmup}}\Big)$ \\[6pt]
$\displaystyle \mathcal{L}_{\mathrm{Jac}}=\mathbb{E}_{v}\left\|J_{\mathrm{Dec}}(z)\,v\right\|_2^2
\;\approx\; \frac{1}{P}\sum_{p=1}^{P}\left\|\frac{\mathrm{Dec}(z+\varepsilon v_p)-\mathrm{Dec}(z-\varepsilon v_p)}{2\varepsilon}\right\|_2^2$ \\[8pt]
$\displaystyle \mathcal{L}_{\mathrm{dir}}=1-\cos\!\big(J_{\mathrm{Dec}}(z)\,r,\;\Delta g\big),\quad
r=\frac{\mu_{t+1}-\mu_t}{\|\mu_{t+1}-\mu_t\|},\ \Delta g=g_{t+1}-g_t$ \\[8pt]
$\displaystyle \mathcal{L}_{\mathrm{sp}}=\sum_{\ell}\sum_{i,j} W^{(\ell)}_{ij}\,\tanh\!\big(W^{(\ell)}_{ij}/\tau\big)$ \\
\bottomrule
\end{tabular}
\caption{Loss components used in the VAE: (1) masked reconstruction $\mathcal{L}_{\mathrm{recon}}$; (2) KL with warm-up $w_{\mathrm{KL}}\mathcal{L}_{\mathrm{KL}}$; 
(3) decoder Jacobian-size proxy $\mathcal{L}_{\mathrm{Jac}}$; (4) optional directional alignment $\mathcal{L}_{\mathrm{dir}}$; (5) smooth sparsity $\mathcal{L}_{\mathrm{sp}}$.
The total objective is $\mathcal{L}=\mathcal{L}_{\mathrm{recon}}+w_{\mathrm{KL}}\mathcal{L}_{\mathrm{KL}}+\lambda_J\mathcal{L}_{\mathrm{Jac}}+\lambda_{\mathrm{dir}}\mathcal{L}_{\mathrm{dir}}+\lambda_{\mathrm{sp}}\mathcal{L}_{\mathrm{sp}}$.
Here \(m_t^{(b)}\) denotes the binary observation mask for sample \(b\), \(\odot\) denotes elementwise multiplication, and \(\varepsilon>0\) is a small constant for numerical stability.}
\label{tab:supp:vae:loss-terms}
\end{table}


\begin{table}[H]
\centering
\small
\begin{tabular}{@{}llp{7.7cm}@{}}
\toprule
Category & Setting & Value / Notes \\ \midrule

\textbf{Inputs/targets}
  & Encoder input & $[g_t \Vert s] \in \mathbb{R}^{G+3}$; nonnegativity of $g_t$ preserved \\
  & Decoder target & Genes only, $\hat g_t\in\mathbb{R}^{G}$; spatial coordinates are not reconstructed \\
  & Masking & Per-row binary mask on $g_t$ (99D: loaded; 27D: all ones).
                 \emph{The same mask is applied to reconstruction, Jacobian, and directional terms so unobserved/pre-imputed entries never affect gradients.} \\

\addlinespace
\textbf{Model}
  & Depth & Compact MLP encoder/decoder (4 fully connected layers each) \\
  & Activations & Smooth nonlinearities (e.g., \texttt{tanh} or \texttt{gelu}) \\
  & Decoder dropout & $\approx 0.10$ between hidden layers (encoder dropout usually $0$) \\
  & Output scaling & Sigmoid output scaled by $3.18$ (matches training range) \\
  & Numeric precision & \texttt{float64} throughout \\

\addlinespace
\textbf{Optimization}
  & Optimizer & AdamW \\
  & Learning-rate schedule & Short warmup to the peak learning rate, then cosine decay to \textbf{5\%} of the peak \\
  & Batch / epochs & Batch size $=500$; up to $3000$ epochs with early stopping \\
  & Checkpointing & Save the model with the best validation loss \\

\addlinespace
\textbf{Defaults}
  & Hidden widths & $(N_1,N_2,N_3)=(1024,\,512,\,256)$ \\
  & Latent prior noise & Standard-normal with scale $0.10$ \\
  & KL regularization & Weight $\beta=10^{-5}$; warmup $\approx 300$ epochs \\
  & Smooth sparsity & Weight decay surrogate: threshold $10^{-5}$; weight in $[0,\,10^{-6}]$ (gentle) \\
  & Decoder Jacobian control & Penalty weight $10^{-4}$; finite-difference step $0.10$; Hutchinson probes $=1$ \\
  & Directional alignment (optional) & Weight $\approx 0$ by default (off); activates only when $\|\Delta g\|>0.02$ \\
  & Peak LR and decoder LR ratio & Peak LR $=10^{-3}$; decoder LR $=0.5\times$ encoder LR \\
  & Warmup length & $1500$ epochs to peak before cosine decay \\

\bottomrule
\end{tabular}
\caption{Exact VAE settings used in our experiments. Values are defaults unless stated otherwise in the main text.}
\label{tab:supp:vae:settings}
\end{table}

\begin{table}[H]
\centering
\small
\begin{tabular}{@{}llp{8.5cm}@{}}
\toprule
Category & Setting & Value / Notes \\ \midrule

\textbf{Inputs}
  & Data source & Full spatiotemporal atlas after imputation \\
  & Pairing & Build $(\mu_t,\mu_{t+1},\sigma_{t+1})$ from consecutive times $t\to t{+}1$ \\

\addlinespace
\textbf{Model}
  & Latent dim $d_z$ & Inherited from VAE (typically $6$–$14$) \\
  & Velocity net & 3 hidden layers, width $N_s$=256 \\
  & Activation & \texttt{tanh} nonlinearity \\
  & Output & $\dot z \in \mathbb{R}^{d_z}$ \\

\addlinespace
\textbf{ODE integration}
  & Solver & Runge–Kutta (\texttt{dopri5}) via \texttt{torchdiffeq} \\
  & Endpoints & Either only $t=0,1$ or $N_{\text{step}}$ points in $[0,1]$ \\
  & Tolerances & Relative tol $10^{-12}$, absolute tol $10^{-12}$ \\
  & Adjoint & Gradient via adjoint method by default \\

\addlinespace
\textbf{Loss}
  & Objective & Relative squared error scaled by $\sigma_{t+1}^2$ \\
  & Formula & $\mathcal{L}_{\text{ODE}}=\tfrac{1}{d_z}\sum_j (z_{t+1,j}-\mu_{t+1,j})^2 / (\sigma_{t+1,j}^2+\varepsilon)$ \\
  & Scale $\varepsilon$ & $10^{-8}$ (stability) \\

\addlinespace
\textbf{Optimization}
  & Optimizer & Adam with maximum LR = $10^{-3}$ (typical) \\
  & LR schedule & Sigmoid-like warmup to epoch \texttt{ep\_maxlr}, then slow decay \\
  & Gradient clipping & Norm clip at $1.0$ \\
  & Batch / epochs & Batch size = 1024, up to 3000 epochs \\
  & Checkpoints & Save model with best validation loss \\

\bottomrule
\end{tabular}
\caption{Neural ODE settings. Symbols: $\mu_t$ latent mean at time $t$, $\sigma_{t+1}$ encoder-predicted posterior standard deviation at time $t+1$, $d_z$ latent dimension, $N_s$ hidden width, $N_{\text{step}}$ number of integration points.}
\label{tab:supp:neuralode:settings}
\end{table}


\clearpage
\section*{Supplemental Figures and Tables}

\processdelayedfloats 







%===========================
% References
%===========================

\clearpage
\section*{Supplemental References}
\bibliographystyle{plainnat}
\bibliography{supp-ref}
\end{document}



