# Scalable optimal-transport based particle filter: assimilation update in `dapy`

$\newcommand{\range}[2][1]{#1{:}\texttt{#2}}$
The assimilation update is

$$
z^p_t(s_m) =
\sum_{q\in\range{P}} \sum_{b\in\range{B}}
\hat{a}^{p,q}_{t,b} \phi_b(s_m) \vec{z}^q_{t}(s_m).
$$

Due to the requirement that the basis functions $\beta_{\range{M}}$ satisfy $\beta_m(s_m) = 1$ and $\beta_m(s_n) = 0~~\forall n \neq m$ we have that

$$
\vec{z}^p_t(s_m) = \vec{x}^p_{t,m} \quad\text{and}\quad z^p_t(s_m) = x^p_{t,m}.
$$

If we also denote $\phi_{b,m} = \phi(s_m)$ then we have that

$$
x^p_{t,m} =
\sum_{q\in\range{P}}\sum_{b\in\range{B}}
\hat{a}^{p,q}_{t,b} \phi_{b,m} \vec{x}^q_{t,m} =
\sum_{b\in\range{B}}\sum_{q\in\range{P}}
\hat{a}^{p,q}_{t,b} \phi_{b,m} \vec{x}^q_{t,m}.
$$

Translating to the syntax in the implementation in `dapy`, assuming the assimilation update is being performed for time index $\texttt{time_index} = t$ the statement

state_particle_patches = pou.split_into_patches_and_scale(state_particles_mesh)

computes the *non-zero elements* of the pointwise products $\phi_{b,m} \vec{x}^p_{t,m}$ with `state_particles_mesh` a $\texttt{P} \times \texttt{N} \times \texttt{M}$ dimensional array ($\texttt{P}$ number of particles, $\texttt{N}$ dimension of local state at each mesh node, $\texttt{M}$ number of spatial mesh nodes) such that

$$\texttt{state_particles_mesh}[p, :, m] = \vec{x}^p_{t, m}$$

and `state_particle_patches` a $\texttt{P} \times \texttt{N} \times \texttt{B} \times \texttt{C}$ dimensional array ($\texttt{B}$ the number of patches in the POU, $\texttt{C} = | \{m \in \range{M}: \phi_{b,m} \neq 0 \} |$, i.e. the number of non-zero entries in $\phi_{b,\range{M}}$ which is assumed constant for each $b$) such that

$$
\texttt{state_particle_patches}[p, :, b, c] =
\phi_{b,m_b(c)}\vec{x}^p_{t, m_b(c)}
$$

where $m_b(c)$ maps $c \in \range{C}$ to each of the indices $\{m \in \range{M}: \phi_{b,m} \neq 0 \}$.

The statement

post_state_particle_patches = np.einsum(
    "bpq,qnbc->pnbc", per_patch_transform_matrices, state_particle_patches)

computes the linear transformations of $\phi_{b,m} \vec{x}^p_{t,m}$ by the per-patch transform matrices $\hat{a}^{p,q}_{t,b}$ with `per_patch_transform_matrices` a $\texttt{B} \times \texttt{P} \times \texttt{P}$ array such that $\texttt{per_patch_transform_matrices}[b, p, q] = \hat{a}^{p,q}_{t,b}$ and `post_state_particle_patches` a $\texttt{P} \times \texttt{N} \times \texttt{B} \times \texttt{C}$ dimensional array with

$$
\texttt{post_state_particle_patches}[p, :, b, c] =
\sum_{q\in\range{P}}\hat{a}^{p,q}_{t,b} \phi_{b,m_b(c)} \vec{x}^q_{t,m_b(c)}.
$$

The statement

post_state_particles_mesh = pou.combine_patches(post_state_particles_patches)

combines the posterior state particles patches by summing over all patches with non-zero values at a particle mesh node index with `post_state_particles_mesh` a $\texttt{P} \times \texttt{N} \times \texttt{M}$ dimensional array with

$$
\texttt{post_state_particles_mesh}[p, :, m] =
\sum_{b\in\range{B}}\sum_{q\in\range{P}}
\hat{a}^{p,q}_{t,b} \phi_{b,m} \vec{x}^q_{t,m}.
$$