Bilinear and Low-Rank Stiefel Layers

BilinearStiefelLayer(Nx, Nv, mx, mv; metric_factors = nothing)

Construct a bilinear Stiefel encoder layer that maps a phase-space grid field $F \in \mathbb{R}^{N_x \times N_v}$ (or a batch $F \in \mathbb{R}^{N_x \times N_v \times B}$) to canonical coordinates $(q, p)$ of half-dimension $n = m_x m_v$ via a velocity-parity split and two-sided Stiefel projection.

Given a phase-space array $F$, the layer splits the velocity axis into time-reversal even and odd components on the positive velocity half-grid $1:(N_v \div 2)$,

\[F^+ = \frac{1}{\sqrt{2}}\bigl(F_{:, 1:N_v/2} + F_{:, N_v:-1:N_v/2+1}\bigr), \qquad F^- = \frac{1}{\sqrt{2}}\bigl(F_{:, 1:N_v/2} - F_{:, N_v:-1:N_v/2+1}\bigr),\]

and projects onto orthonormal bases $\Phi_x \in \mathrm{St}(m_x, N_x)$ and $\Phi_v \in \mathrm{St}(m_v, N_v / 2)$:

\[Q_{\mathrm{red}} = \Phi_x^T F^+ \Phi_v \in \mathbb{R}^{m_x \times m_v}, \qquad P_{\mathrm{red}} = \Phi_x^T F^- \Phi_v \in \mathbb{R}^{m_x \times m_v},\]

returning (q = vec(Q_red), p = vec(P_red)) (or matrices of shape (mx * mv, B) for 3D batched input). A pre-split NamedTuple (q = F_plus, p = F_minus) of shape (Nx, Nv ÷ 2) (or (Nx, Nv ÷ 2, B)) is also accepted directly.

When metric_factors = (Rx, Rv) is supplied with Cholesky or square-root mass factors Rx of size (Nx, Nx) and Rv of size (Nv ÷ 2, Nv ÷ 2) or (Nv, Nv), the field is mapped into the Euclidean metric by (Rx, Rv) prior to Stiefel projection.

The coordinates $(q, p)$ are canonical only for the structure $\mathbb{J}_{2n}$ that the layer and LowRankStiefelLayer declare on the parity coordinates. A frontend that is linear in $F$ is not a Poisson map for the Vlasov bracket: with $L^2$-orthonormal even and odd bases, the bracket linearised at $f_0$ gives a $\{Q, P\}$ that is not the identity and can be singular.

See also TransposeBilinearStiefelLayer, LowRankStiefelLayer, and inverse_layer.

TransposeBilinearStiefelLayer(Nx, Nv, mx, mv; metric_factors = nothing)

Construct the right-inverse decoder layer corresponding to BilinearStiefelLayer.

Given canonical coordinates (q = q, p = p) of length $n = m_x m_v$ (or shape (mx * mv, B) for batched input, or as a 2-tuple (Q_red, P_red) of shape (mx, mv)), the layer reshapes $(q, p)$ to $Q_{\mathrm{red}}, P_{\mathrm{red}} \in \mathbb{R}^{m_x \times m_v}$, reconstructs the even and odd velocity components

\[F^+ = \Phi_x Q_{\mathrm{red}} \Phi_v^T, \qquad F^- = \Phi_x P_{\mathrm{red}} \Phi_v^T,\]

applies $(R_x^{-1}, R_v^{-1})$ when metric_factors = (Rx, Rv) is present, and recombines the full velocity grid

\[F_{:, 1:N_v/2} = \frac{1}{\sqrt{2}}(F^+ + F^-), \qquad F_{:, N_v:-1:N_v/2+1} = \frac{1}{\sqrt{2}}(F^+ - F^-).\]

See also BilinearStiefelLayer and inverse_layer.

solve!(l::BilinearStiefelLayer, ps::NamedTuple, F_snapshots::AbstractArray{T, 3})
solve!(nn::NeuralNetwork, F_snapshots::AbstractArray{T, 3})

Initialize the Stiefel bases ps.Φx and ps.Φv of a BilinearStiefelLayer via truncated higher-order singular value decomposition (Tucker-2 HOSVD) on a 3D snapshot array F_snapshots of shape (Nx, Nv, n_snapshots):

1. Splits F_snapshots into weighted velocity-even and velocity-odd tensors $F^+, F^- \in \mathbb{R}^{N_x \times (N_v / 2) \times n_{\mathrm{snapshots}}}$.
2. Sets $\Phi_x \in \mathrm{St}(m_x, N_x)$ to the top-$m_x$ left singular vectors of the spatial mode-1 unfolding $[F_{(1)}^+\;\; F_{(1)}^-]$.
3. Sets $\Phi_v \in \mathrm{St}(m_v, N_v / 2)$ to the top-$m_v$ left singular vectors of the parity-velocity mode-2 unfolding $[F_{(2)}^+\;\; F_{(2)}^-]$.

LowRankStiefelLayer(Nx, Nv, mx, mv; r_out = mx * mv, metric_factors = nothing)

Construct a gauge-invariant dynamical low-rank approximation (DLRA) Stiefel encoder layer that maps factored rank-$r_{\mathrm{in}}$ states $(X = X, S = S, V = V)$ with $X \in \mathbb{R}^{N_x \times r_{\mathrm{in}}}$, $S \in \mathbb{R}^{r_{\mathrm{in}} \times r_{\mathrm{in}}}$, and $V \in \mathbb{R}^{N_v \times r_{\mathrm{in}}}$ (or 3D batched arrays) to canonical coordinates $(q, p)$ of half-dimension $n = m_x m_v$ in $O(r_{\mathrm{in}}(m_x N_x + m_v N_v))$ operations without assembling the dense $N_x \times N_v$ matrix $X S V^T$.

Splitting $V$ into velocity-even and velocity-odd half-grid factors $V^+, V^- \in \mathbb{R}^{(N_v / 2) \times r_{\mathrm{in}}}$ and contracting with $\Phi_x \in \mathrm{St}(m_x, N_x)$ and $\Phi_v \in \mathrm{St}(m_v, N_v / 2)$ gives

\[Q_{\mathrm{red}} = (\Phi_x^T X)\, S\, \bigl((V^+)^T \Phi_v\bigr), \qquad P_{\mathrm{red}} = (\Phi_x^T X)\, S\, \bigl((V^-)^T \Phi_v\bigr),\]

which produces the exact same canonical coordinates (q = vec(Q_red), p = vec(P_red)) as BilinearStiefelLayer applied to $F = X S V^T$ and is strictly invariant under orthogonal gauge transformations $(X, S, V) \mapsto (X O_x, O_x^T S O_v, V O_v)$ for all $O_x, O_v \in O(r_{\mathrm{in}})$.

The coordinates $(q, p)$ are canonical only for the structure $\mathbb{J}_{2n}$ that the layer and BilinearStiefelLayer declare on the parity coordinates. A frontend that is linear in $f = X S V^T$ is not a Poisson map for the Vlasov bracket: with $L^2$-orthonormal even and odd bases, the bracket linearised at $f_0$ gives a $\{Q, P\}$ that is not the identity and can be singular.

See also DecodeLowRankStiefelLayer, BilinearStiefelLayer, FastDLRATraceLoss, and inverse_layer.

DecodeLowRankStiefelLayer(Nx, Nv, mx, mv; r_out = mx * mv, metric_factors = nothing)

Construct the Stiefel-preserving low-rank SVD decoder layer corresponding to LowRankStiefelLayer.

Given canonical coordinates (q = q, p = p) (or (Q_red, P_red)), the layer forms the $m_x \times 2 m_v$ matrix $[Q_{\mathrm{red}}\;\; P_{\mathrm{red}}]$, computes its small singular value decomposition $U_S \Sigma_S W_S^T$ truncated to rank $r_{\mathrm{eff}} = \min(r_{\mathrm{out}}, m_x, 2 m_v)$, lifts $\tilde{X} = \Phi_x U_S$ and $\tilde{V}$ from $W_S$ and $\Phi_v$ (applying $(R_x^{-1}, R_v^{-1})$ when metric_factors is supplied), and QR re-orthonormalizes $\tilde{X} = \hat{X} R_X$ and $\tilde{V} = \hat{V} R_V$ with $\hat{S} = R_X \Sigma_S R_V^T$ so that the returned factors (X = X_hat, S = S_hat, V = V_hat) satisfy $\hat{X}^T \hat{X} = I_{r_{\mathrm{eff}}}$ and $\hat{V}^T \hat{V} = I_{r_{\mathrm{eff}}}$.

See also LowRankStiefelLayer, TransposeBilinearStiefelLayer, and inverse_layer.

FastDLRATraceLoss()

Construct a gauge-invariant low-rank trace loss that evaluates the squared Frobenius reconstruction error $\|X S V^T - \hat{X} \hat{S} \hat{V}^T\|_F^2$ between two factored dynamical low-rank approximation (DLRA) states (X = X, S = S, V = V) and (X = X_hat, S = S_hat, V = V_hat) in $O(r^2 (N_x + N_v))$ operations without forming any dense $N_x \times N_v$ array:

\[\|X S V^T - \hat{X} \hat{S} \hat{V}^T\|_F^2 = \mathrm{tr}\bigl(S^T (X^T X) S (V^T V)\bigr) - 2\,\mathrm{tr}\bigl(S^T (X^T \hat{X}) \hat{S} (\hat{V}^T V)\bigr) + \mathrm{tr}\bigl(\hat{S}^T (\hat{X}^T \hat{X}) \hat{S} (\hat{V}^T \hat{V})\bigr),\]

clamped at zero(T) via max(zero(T), val) to prevent negative roundoff at zero error. Because the Gram matrices $X^T X, V^T V, \hat{X}^T \hat{X}, \hat{V}^T \hat{V}$ are computed explicitly, the identity holds even when the basis factors are only approximately orthonormal.

See also LowRankStiefelLayer and DecodeLowRankStiefelLayer.