Permutation-Symmetric Particle Layers

CanonicalSortLayer{N, M, Axis}()
CanonicalSortLayer(M::Integer, d::Integer, axis::Integer = 1)

Make an instance of a cotangent-lifted canonical Weyl-chamber particle sorting layer for M particles in $d$ spatial dimensions, sorting along coordinate row Axis (1 <= Axis <= d, default 1).

The first type parameter is the phase-space dimension N = 2Md, which input_dimension and output_dimension return. Given a (q, p) NamedTuple with q, p of shape (d, M) (single snapshot) or (d, M, B) (batch of $B$ snapshots), CanonicalSortLayer{N, M, Axis}() sorts the M particles along coordinate row Axis of q using a stable merge sort (MergeSort), applies the identical permutation $\sigma_b \in S_M$ to both q and p in each batch sample $b$, and returns (q = q_sorted, p = p_sorted, perm = σ). On a flattened phase-space vector of length $2Md$ or snapshot matrix of shape $(2Md, B)$, it reshapes each half to (d, M) or (d, M, B) and returns the sorted flattened state.

The sort is discontinuous in p where two particles cross along the sorting axis, because the permutation changes there, so the Jacobian identity $(\nabla\Psi)^T \mathbb{J}_{2Md} \nabla\Psi = \mathbb{J}_{2Md}$ holds only away from those crossings. Sorting suits systems that keep their particle order, not a collisionless plasma, whose sorted snapshots are not low-rank.

InvertCanonicalSortLayer{N, M}()
InvertCanonicalSortLayer(M::Integer, d::Integer)

Make an instance of the analytic inverse of CanonicalSortLayer for M particles in $d$ spatial dimensions; the first type parameter is the phase-space dimension N = 2Md.

Given (q = q̂, p = p̂, perm = σ), InvertCanonicalSortLayer{N, M}() applies $\sigma_b^{-1}$ (invperm(σ_b)) to both q̂ and p̂ in each batch sample $b$ and returns (q = q_unsorted, p = p_unsorted). When called on (q, p) or a flattened array without a stored perm, it acts as the identity on the canonical Weyl chamber.

PairwiseSymplecticLayer{M, U, TA, C, Sign} <: SympNetLayer{M, M}

Permutation-equivariant pairwise symplectic shear layer on $T^*(\mathbb{R}^d)^{M_p} \cong \mathbb{R}^{M}$ with total phase-space dimension M = 2 * M_p * d and upscaled per-particle feature dimension U = upscale * d.

See PairwiseSymplecticLayerQ, PairwiseSymplecticLayerP, InversePairwiseSymplecticLayerQ, and InversePairwiseSymplecticLayerP.

PairwiseSymplecticLayerQ(M, d, upscale, activation = tanh)

Make an instance of an $S_M$-equivariant pairwise symplectic layer that updates canonical momenta $p$ via self and pairwise relative potential forces:

\[p_i^{\mathrm{new}} = p_i - K_{\mathrm{self}}^T \sigma(K_{\mathrm{self}} q_i + b_{\mathrm{self}}) - \frac{1}{M}\sum_{j=1}^M K_{\mathrm{pair}}^T \tau\bigl(K_{\mathrm{pair}}(q_i - q_j)\bigr),\]

where $\tau(x) = \sigma(x) - \sigma(-x)$ is oddized so that the pairwise force $f_{ij} = -f_{ji}$ is skew-symmetric in $(i, j)$ and conserves total canonical momentum $\sum_{i=1}^M p_i$ when $K_{\mathrm{self}} = 0$.

PairwiseSymplecticLayerP(M, d, upscale, activation = tanh)

Make an instance of an $S_M$-equivariant pairwise symplectic layer that updates canonical positions $q$ via self and pairwise relative velocity differences $p_i - p_j$:

\[q_i^{\mathrm{new}} = q_i + K_{\mathrm{self}}^T \sigma(K_{\mathrm{self}} p_i + b_{\mathrm{self}}) + \frac{1}{M}\sum_{j=1}^M K_{\mathrm{pair}}^T \tau\bigl(K_{\mathrm{pair}}(p_i - p_j)\bigr),\]

where $\tau(x) = \sigma(x) - \sigma(-x)$.

InversePairwiseSymplecticLayerQ(M, d, upscale, activation = tanh)

Analytic symplectic inverse of PairwiseSymplecticLayerQ, flipping the sign of the momentum force update while sharing the identical parameter NamedTuple ps.

InversePairwiseSymplecticLayerP(M, d, upscale, activation = tanh)

Analytic symplectic inverse of PairwiseSymplecticLayerP, flipping the sign of the position displacement update while sharing the identical parameter NamedTuple ps.

CanonicalSlotLayer{N, NOut, M, NSlots}()
CanonicalSlotLayer(M::Integer, d::Integer, NSlots::Integer)

Make an instance of a canonical slot reduction layer mapping M particles in $d$ spatial dimensions to NSlots canonical slots (M >= 2, 1 <= NSlots <= M).

The first two type parameters are the phase-space dimensions N = 2Md of the particles and NOut = 2 NSlots d of the slots, which input_dimension and output_dimension return. The layer first sorts the particles along coordinate row 1 of q with CanonicalSortLayer, so that its output is invariant under a relabelling of particles with distinct first coordinates. It then applies the orthogonal Helmert / Jacobi transformation separating the center-of-mass mode

\[(Q_{\mathrm{cm}}, P_{\mathrm{cm}}) = \frac{1}{\sqrt{M}}\sum_{i=1}^M (q_i, p_i)\]

from the $M - 1$ zero-sum relative modes $(q_{\mathrm{rel}}, p_{\mathrm{rel}}) \in \mathbb{R}^{d \times (M - 1)}$, sorts the $M - 1$ relative modes by descending squared phase-space norm $\|q_{\mathrm{rel}, k}\|_2^2 + \|p_{\mathrm{rel}, k}\|_2^2$ (sortperm with rev = true, alg = MergeSort, applying the identical permutation $\sigma_{\mathrm{rel}} \in S_{M - 1}$ to both $q_{\mathrm{rel}}$ and $p_{\mathrm{rel}}$), and projects the sorted relative modes onto $N_{\mathrm{slots}} - 1$ relative slots via a learned StiefelManifold parameter $\Phi_{\mathrm{rel}} \in \mathrm{St}(N_{\mathrm{slots}} - 1, M - 1)$ while keeping $(Q_{\mathrm{cm}}, P_{\mathrm{cm}})$ intact as slot 1. The particles are already sorted, so this second sort removes no symmetry, and it makes the output discontinuous where two relative-mode norms cross.

The layer is a linear projection of sorted states for a fixed particle count M; it takes no variable particle count and has no Löwdin-orthogonalised affinity slots. The sort of the particles is discontinuous in p where two particles cross along coordinate row 1, so the identities of the layer's Jacobian hold only away from those crossings. The layer suits systems that keep their particle order, not a collisionless plasma, whose sorted snapshots are not low-rank.

InverseCanonicalSlotLayer{N, NOut, M, NSlots}()
InverseCanonicalSlotLayer(M::Integer, d::Integer, NSlots::Integer)

Analytic symplectic decoder lift (inverse_layer(CanonicalSlotLayer(M, d, NSlots))) mapping NSlots canonical slots back to M particles via $\Phi_{\mathrm{rel}} \in \mathrm{St}(N_{\mathrm{slots}} - 1, M - 1)$ and the inverse Helmert matrix.

It takes the type parameters of CanonicalSlotLayer, so its input dimension is NOut = 2 NSlots d and its output dimension is N = 2Md. For any $\Phi_{\mathrm{rel}} \in \mathrm{St}(N_{\mathrm{slots}} - 1, M - 1)$, the decoder lift satisfies $(\nabla \Psi_{\mathrm{slot}}^{\mathrm{dec}})^T \mathbb{J}_{2Md} \nabla \Psi_{\mathrm{slot}}^{\mathrm{dec}} = \mathbb{J}_{2 N_{\mathrm{slots}} d}$ and preserves $(Q_{\mathrm{cm}}, P_{\mathrm{cm}})$.

SymmetricSymplecticAutoencoder(M, d, reduced_dim; frontend = :pairwise, kwargs...)

Make an instance of a symplectic autoencoder with a particle frontend for M particles in $d$ spatial dimensions (ambient dimension $2Md$) and latent dimension reduced_dim ($2n$).

SymmetricSymplecticAutoencoder combines a particle frontend (frontend = :sort, :pairwise, or :slot) with the weight-tied symplectic core WeightTiedSymplecticAutoencoder so that encoder(nn) and decoder(nn) share a single parameter tree.

With frontend = :sort or frontend = :slot, the latent state is invariant under a relabelling of particles with distinct coordinates along the sorting axis (sort_axis for :sort, coordinate 1 for :slot). With frontend = :pairwise, only the frontend is permutation-equivariant: the core is not, so the latent state is not invariant. The sort of :sort and :slot is discontinuous in p where two particles cross along the sorting axis, so the Jacobian identity of the frontend holds only away from those crossings. :sort and :slot suit systems that keep their particle order, not a collisionless plasma, whose sorted snapshots are not low-rank.

The latent right-inversion $\mathrm{encoder}(nn)(\mathrm{decoder}(nn)(\xi)) = \xi$ holds to machine precision for every $\xi \in \mathbb{R}^{2n}$ with frontend = :pairwise. With frontend = :sort and frontend = :slot, it holds for $\xi$ whose decoded state lies in the canonical chamber, where the sorting permutations on the decoded state are the identity; the decoder does not map into this chamber in general.

Arguments

Besides M::Integer, d::Integer, and even integer reduced_dim::Integer:

• frontend::Symbol = :pairwise: Particle frontend (:sort, :pairwise, or :slot).
• sort_axis::Integer = 1: Coordinate axis (1 <= sort_axis <= d) used by CanonicalSortLayer when frontend == :sort.
• n_pairwise_layers::Integer = 2: Number of alternating PairwiseSymplecticLayerQ / PairwiseSymplecticLayerP shears when frontend == :pairwise.
• pairwise_upscale::Integer = 3: Per-particle upscaling factor in the pairwise symplectic layers.
• n_slots::Integer = max(1, M - 1): Number of canonical slots (1 <= n_slots <= M) in CanonicalSlotLayer when frontend == :slot.
• n_blocks::Integer = 2, n_layers::Integer = 3, n_ambient_layers::Integer = 2, sympnet_upscale::Integer = 5, activation = tanh, init_q::Bool = true: Forwarded to the WeightTiedSymplecticAutoencoder core.