Interpolator + DA ¶ 1. Train NerF on Observations
Assume we have a dataset of independent observations, y n y_n y n , on the observation domain ( Ω y , T y ) (\Omega_y,\mathcal{T}_y) ( Ω y , T y ) .
D = { ( x n , t n ) , y n } n = 1 N , x ∈ Ω y ⊆ R D s t ∈ T y ⊆ R + \begin{aligned}
\mathcal{D}=\{(x_n,t_n),y_n\}_{n=1}^N, && &&
x\in\Omega_y\subseteq\mathbb{R}^{D_s} &&
t\in\mathcal{T}_y\subseteq\mathbb{R}^+
\end{aligned} D = {( x n , t n ) , y n } n = 1 N , x ∈ Ω y ⊆ R D s t ∈ T y ⊆ R + We train a neural field, f θ f_\theta f θ , to interpolate the observation wrt the spatiotemporal coordinate values.
y n = f θ ( x n , t n ) + ε n \begin{aligned}
y_n &= f_\theta(x_n, t_n) + \varepsilon_n
\end{aligned} y n = f θ ( x n , t n ) + ε n 2. Train Strong Constrained DA with NerF
u 0 ∼ N ( u b , Σ b ) u t = f θ ( u t − 1 , t ) y t = h θ ( u t , t ) + ε y , ε y ∼ N ( 0 , Σ y ) \begin{aligned}
\boldsymbol{u}_0 &\sim \mathcal{N}(\boldsymbol{u}_b,\boldsymbol{\Sigma}_b)\\
\boldsymbol{u}_t &= \boldsymbol{f_\theta}\left( \boldsymbol{u}_{t-1},t\right)\\
\boldsymbol{y}_t &= \boldsymbol{h_\theta}(\boldsymbol{u}_t, t) +
\boldsymbol{\varepsilon}_y, && &&
\boldsymbol{\varepsilon}_y \sim \mathcal{N}(0,\boldsymbol{\Sigma_y})
\end{aligned} u 0 u t y t ∼ N ( u b , Σ b ) = f θ ( u t − 1 , t ) = h θ ( u t , t ) + ε y , ε y ∼ N ( 0 , Σ y ) We can estimate the state by minimizing the objective function
J ( u ) = ∑ t = 1 T ∣ ∣ y θ ( t ) − h θ ( u t , t ) ∣ ∣ Σ y − 1 2 + ∣ ∣ u 0 − u b ∣ ∣ Σ b − 1 2 \mathcal{J}(u) =
\sum_{t=1}^T||\boldsymbol{y_\theta}(t) - h_\theta(u_t,t)||_{\boldsymbol{\Sigma_y}^{-1}}^2 +
||\boldsymbol{u}_0 - \boldsymbol{u}_b||_{\boldsymbol{\Sigma}_b^{-1}}^2 J ( u ) = t = 1 ∑ T ∣∣ y θ ( t ) − h θ ( u t , t ) ∣ ∣ Σ y − 1 2 + ∣∣ u 0 − u b ∣ ∣ Σ b − 1 2 3. Train NerF
Assume we have a dataset of independent reanalysis points, u n ∗ u_n^* u n ∗ , on the observation domain ( Ω y , T u ) (\Omega_y,\mathcal{T}_u) ( Ω y , T u ) .
D = { ( x n , t n ) , u n ∗ } n = 1 N , x ∈ Ω z ⊆ R D s t ∈ T z ⊆ R + \begin{aligned}
\mathcal{D}=\{(x_n,t_n),u_n^*\}_{n=1}^N, && &&
x\in\Omega_z\subseteq\mathbb{R}^{D_s} &&
t\in\mathcal{T}_z\subseteq\mathbb{R}^+
\end{aligned} D = {( x n , t n ) , u n ∗ } n = 1 N , x ∈ Ω z ⊆ R D s t ∈ T z ⊆ R + We train a neural field, f θ f_\theta f θ , to interpolate the observation wrt the spatiotemporal coordinate values.
u ∗ ( x n , t n ) = f θ ( x n , t n ) + ε n , x ∈ Ω z ⊆ R D s \begin{aligned}
\boldsymbol{u}^*(x_n, t_n) &= f_\theta(x_n, t_n) + \varepsilon_n,
&& &&
\mathbf{x}\in\Omega_z\subseteq\mathbb{R}^{D_s}
\end{aligned} u ∗ ( x n , t n ) = f θ ( x n , t n ) + ε n , x ∈ Ω z ⊆ R D s 5. Train Weak-Constrained DA
u 0 ∼ N ( u b , Σ b ) u t = f θ ( u t − 1 , t ) , ε u ∼ N ( 0 , Σ u ) y t = h θ ( u t , t ) + ε y , ε y ∼ N ( 0 , Σ y ) \begin{aligned}
\boldsymbol{u}_0 &\sim \mathcal{N}(\boldsymbol{u}_b,\boldsymbol{\Sigma}_b)\\
\boldsymbol{u}_t &= \boldsymbol{f_\theta}\left( \boldsymbol{u}_{t-1},t\right), && &&
\boldsymbol{\varepsilon}_u \sim \mathcal{N}(0,\boldsymbol{\Sigma_u}) \\
\boldsymbol{y}_t &= \boldsymbol{h_\theta}(\boldsymbol{u}_t, t) +
\boldsymbol{\varepsilon}_y, && &&
\boldsymbol{\varepsilon}_y \sim \mathcal{N}(0,\boldsymbol{\Sigma_y})
\end{aligned} u 0 u t y t ∼ N ( u b , Σ b ) = f θ ( u t − 1 , t ) , = h θ ( u t , t ) + ε y , ε u ∼ N ( 0 , Σ u ) ε y ∼ N ( 0 , Σ y ) We can estimate the state by minimizing the objective function
J ( u ) = ∑ t = 1 T ∣ ∣ y θ ( t ) − h θ ( u t , t ) ∣ ∣ Σ y − 1 2 + ∑ t = 1 T ∣ ∣ u θ ( t ) − f θ ( u t − 1 , t ) ∣ ∣ Σ u − 1 2 + ∣ ∣ u 0 − u b ∣ ∣ Σ b − 1 2 \mathcal{J}(u) =
\sum_{t=1}^T||\boldsymbol{y_\theta}(t) - h_\theta(u_t,t)||_{\boldsymbol{\Sigma_y}^{-1}}^2 +
\sum_{t=1}^T||\boldsymbol{u_\theta}(t) - \boldsymbol{f_\theta}\left( \boldsymbol{u}_{t-1},t\right)||_{\boldsymbol{\Sigma_u}^{-1}}^2 +
||\boldsymbol{u}_0 - \boldsymbol{u}_b||_{\boldsymbol{\Sigma}_b^{-1}}^2 J ( u ) = t = 1 ∑ T ∣∣ y θ ( t ) − h θ ( u t , t ) ∣ ∣ Σ y − 1 2 + t = 1 ∑ T ∣∣ u θ ( t ) − f θ ( u t − 1 , t ) ∣ ∣ Σ u − 1 2 + ∣∣ u 0 − u b ∣ ∣ Σ b − 1 2 Interpolator + Foundational Models ¶ 2. Train Embedding on NerF
Assume we have a dataset of sequential, independent observations, y t y_t y t , which is given by the neural field, f θ f_\theta f θ .
However, we query the functa on the latent domain, ( Ω z , T z ) (\Omega_z, \mathcal{T}_z) ( Ω z , T z ) .
y θ ( t ) = f θ ( X z , t ) , X z ∈ R D Ω z t ∈ T z ⊆ R + \begin{aligned}
\boldsymbol{y_\theta}(t)&=\boldsymbol{f_\theta}(\mathbf{X}_z,t), && &&
\mathbf{X}_z\in\mathbb{R}^{D_{\Omega_z}} &&
t\in\mathcal{T}_z\subseteq\mathbb{R}^+
\end{aligned} y θ ( t ) = f θ ( X z , t ) , X z ∈ R D Ω z t ∈ T z ⊆ R + where X z = { x ∈ Ω z ∈ R D s } \mathbf{X}_z = \{ \mathbf{x}\in\Omega_z\in\mathbb{R}^{D_s}\} X z = { x ∈ Ω z ∈ R D s } .
We can create a dataset by (quasi-)randomly selecting points
D = { y θ ( t ) } t = 1 T \mathcal{D}=\{ \boldsymbol{y_\theta}(t) \}_{t=1}^T
\hspace{10mm} D = { y θ ( t ) } t = 1 T We train an embedding on the latent domain, z z z , using the Neural Field.
We can also apply a random mask, m \boldsymbol{m} m , to help augment the data by randomly masking pixels.
L ( θ ) = 1 ∣ D ∣ ∑ t ∈ D ∣ ∣ y θ ( t ) − T D ∘ T E ∘ m ∘ y θ ( t ) ∣ ∣ 2 2 \mathcal{L}(\theta) =
\frac{1}{|\mathcal{D}|}\sum_{t\in\mathcal{D}}
||\boldsymbol{y_\theta}(t) - T_D\circ T_E\circ \boldsymbol{m}\circ \boldsymbol{y_\theta}(t)||^2_2 L ( θ ) = ∣ D ∣ 1 t ∈ D ∑ ∣∣ y θ ( t ) − T D ∘ T E ∘ m ∘ y θ ( t ) ∣ ∣ 2 2 Latent Variable ¶ Train (Masked) AutoEncoder on Simulations
PnP for Real Observations
Train AutoEncoder on Sparse Observations
Train Variational AutoEncoder (Probabilistic Reconstruction)
Train U-Net (DEQ)