diff --git a/CHANGELOG.md b/CHANGELOG.md index 583baa706..0676ce3dc 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -146,11 +146,11 @@ breaking release). - **`scripts/utilities/metal_abort_runs.jl` counts whether one of the two Metal process aborts recurs, one fresh process per counted run, and `scripts/utilities/metal_abort_session.jl` runs the arm that needs two requests to one session.** Two processes died in the first Adam call of a structured network on `MetalBackend()` in `Float32` — `LASympNet(4)` from a command-encoder assertion while a second evaluation read a Julia vector, and `VolumePreservingFeedForward(4)` from host-heap corruption detected in the LLVM JIT — and neither had a reproducer when the runs began. **One death of the session arm has the signature of the `LASympNet(4)` item:** run 18 died with `SIGSEGV` in `AGX::ComputeContext…::endComputePass` under `-[AGXG16XFamilyCommandBuffer commit]` and Metal.jl's `commit_with_queue_key!`, from `flush_batched_queues!` inside a task's `device_synchronize`, with no GeometricMachineLearning frame in the stack. It happened in the load phase, before the overlap stamps and outside the arm's counted window, so it is recorded as a death and not as a reproduction, and the arm stays open. A counted run prints its pid, `MTL_DEBUG_LAYER`, `Threads.nthreads()`, the version and path of every package of the stack, `Metal.functional()` and the device's name before any Metal work; it takes a forward pass and a Zygote pullback on one network and times the *first* training call of a new one with `@timed Base.invokelatest`, so that the compilation is inside the timer, and a `compile_time` of 0 makes the run void. `N` is 59 per arm, and a run's crash report counts only when the report's `"pid"` is the run's own, so the sixteen reports of other Julia processes in `~/Library/Logs/DiagnosticReports` are not counted; `audit` recounts an arm from its logs and the reports directory. Neither arm that a shell process runs aborted in 59 of 59 runs, on Julia 1.13.3 and Metal.jl 1.11.1, the Metal.jl of the aborts, and the third arm, a private session's two `ex` calls, has 16 counted runs, all clean, and the one recorded death. The session script appends `M8-TRAIN-START` and `M8-READ-START` to the run's log file before the timed call can finish, `audit` counts a death of that arm as its abort only when they show the overlap, and a child that dies before it prints its pid is counted from the pid that the parent logs. An issue for `JuliaGPU/Metal.jl` is drafted and not filed, on the hypothesis that `flush_batched_queues!` commits another task's open command buffer with no lock. `KNOWN_ISSUES.md` K35 records the counts, both stacks, the procedure and the draft. -- **`CanonicalSortLayer`, `InvertCanonicalSortLayer`, `PairwiseSymplecticLayer`, `PairwiseSymplecticLayerQ`, `PairwiseSymplecticLayerP`, `InversePairwiseSymplecticLayerQ`, `InversePairwiseSymplecticLayerP`, `CanonicalSlotLayer`, `InverseCanonicalSlotLayer`, and `SymmetricSymplecticAutoencoder` add particle symplectic layers that are invariant or equivariant under relabelling of the particles, and a weight-tied symplectic autoencoder with three particle frontends.** All ten names are exported. `CanonicalSortLayer(M, d, axis = 1)` makes `CanonicalSortLayer{N, M, Axis}` with `N = 2Md`, sorts `M` particles in ``\mathbb{R}^d`` along coordinate row `axis` of `q` using a stable merge sort (`MergeSort`), applies the same permutation ``\sigma \in S_M`` to both `q` and `p`, and returns `(q = q_sorted, p = p_sorted, perm = σ)` on `(q, p)` `NamedTuple` inputs of shape `(d, M)` or `(d, M, B)` (or the sorted state on flattened `2Md` vectors and `(2Md, B)` matrices). `InvertCanonicalSortLayer(M, d)` makes `InvertCanonicalSortLayer{N, M}` and applies ``\sigma^{-1}`` (`invperm`) to restore the original particle ordering; `inverse_layer` pairs it with `CanonicalSortLayer`. `CanonicalSortLayer`, `InvertCanonicalSortLayer`, `CanonicalSlotLayer` and `InverseCanonicalSlotLayer` carry the phase-space dimension of the flattened state as their first type parameter, `N = 2Md`, so `input_dimension` and `output_dimension` return the length of that state, and their inputs are checked against the fixed `d`, with the first dimension of an input exactly its input dimension. `PairwiseSymplecticLayerQ(M, d, upscale, σ)` and `PairwiseSymplecticLayerP(M, d, upscale, σ)` apply ``S_M``-equivariant canonical shears combining a self term ``K_{\mathrm{self}}^T \sigma(K_{\mathrm{self}} u_i + b_{\mathrm{self}})`` and a skew-symmetric pairwise relative potential term ``\frac{1}{M}\sum_{j=1}^M K_{\mathrm{pair}}^T \tau(K_{\mathrm{pair}}(u_i - u_j))`` with ``\tau(x) = \sigma(x) - \sigma(-x)``, conserving total canonical momentum (or position) when ``K_{\mathrm{self}} = 0`` and pairing with `InversePairwiseSymplecticLayerQ` and `InversePairwiseSymplecticLayerP` through `inverse_layer`. `CanonicalSlotLayer(M, d, NSlots)` (`M >= 2`, `1 <= NSlots <= M`, `CanonicalSlotLayer{N, NOut, M, NSlots}` with `NOut = 2 NSlots d`) first sorts the particles along coordinate row 1 of `q` with `CanonicalSortLayer`, so its slots are invariant under a relabelling of particles with distinct first coordinates. It then separates 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 Helmert relative modes, sorts the ``M - 1`` relative modes canonically by descending squared phase-space norm ``\|q_{\mathrm{rel}, k}\|_2^2 + \|p_{\mathrm{rel}, k}\|_2^2`` (`MergeSort`), and projects the sorted relative modes onto ``N_{\mathrm{slots}} - 1`` slots via a learned `StiefelManifold` parameter ``\Phi_{\mathrm{rel}} \in \mathrm{St}(N_{\mathrm{slots}} - 1, M - 1)`` (`NSlots == 1` keeps only the center-of-mass slot with no parameter leaf). `InverseCanonicalSlotLayer(M, d, NSlots)` takes the same type parameters, with input dimension `NOut` and output dimension `N`, and lifts the slots back symplectically via ``\Phi_{\mathrm{rel}}`` and the transpose Helmert matrix. `SymmetricSymplecticAutoencoder(M, d, reduced_dim; frontend = :sort | :pairwise | :slot, ...) <: SymplecticCompression` combines one of the three particle frontends with the weight-tied `WeightTiedSymplecticAutoencoder` core. With `:sort` and `: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 `:pairwise` only the frontend is permutation-equivariant, because the core is not, so the latent state is not invariant. Right-inversion ``\mathrm{encoder}(nn)(\mathrm{decoder}(nn)(\xi)) = \xi`` holds to machine precision for every ``\xi`` with `:pairwise`; with `:sort` and `:slot` it holds only for ``\xi`` whose decoded state lies in the canonical chamber, where the sorting permutations on the decoded state are the identity, because the decoder does not map into that chamber in general. The tests are `test/layers/symmetric_particle_layers.jl`. +- **`CanonicalSortLayer`, `InvertCanonicalSortLayer`, `PairwiseSymplecticLayer`, `PairwiseSymplecticLayerQ`, `PairwiseSymplecticLayerP`, `InversePairwiseSymplecticLayerQ`, `InversePairwiseSymplecticLayerP`, `CanonicalSlotLayer`, `InverseCanonicalSlotLayer`, and `SymmetricSymplecticAutoencoder` add particle symplectic layers that are invariant or equivariant under relabelling of the particles, and a weight-tied symplectic autoencoder with three particle frontends.** All ten names are exported. `CanonicalSortLayer(M, d, axis = 1)` makes `CanonicalSortLayer{N, M, Axis}` with `N = 2Md`, sorts `M` particles in ``\mathbb{R}^d`` along coordinate row `axis` of `q` using a stable merge sort (`MergeSort`), applies the same permutation ``\sigma \in S_M`` to both `q` and `p`, and returns `(q = q_sorted, p = p_sorted, perm = σ)` on `(q, p)` `NamedTuple` inputs of shape `(d, M)` or `(d, M, B)` (or the sorted state on flattened `2Md` vectors and `(2Md, B)` matrices). `InvertCanonicalSortLayer(M, d)` makes `InvertCanonicalSortLayer{N, M}` and applies ``\sigma^{-1}`` (`invperm`) to restore the original particle ordering; `inverse_layer` pairs it with `CanonicalSortLayer`. `CanonicalSortLayer`, `InvertCanonicalSortLayer`, `CanonicalSlotLayer` and `InverseCanonicalSlotLayer` carry the phase-space dimension of the flattened state as their first type parameter, `N = 2Md`, so `input_dimension` and `output_dimension` return the length of that state, and their inputs are checked against the fixed `d`, with the first dimension of an input exactly its input dimension. `PairwiseSymplecticLayerQ(M, d, upscale, σ)` and `PairwiseSymplecticLayerP(M, d, upscale, σ)` apply ``S_M``-equivariant canonical shears combining a self term ``K_{\mathrm{self}}^T \sigma(K_{\mathrm{self}} u_i + b_{\mathrm{self}})`` and a skew-symmetric pairwise relative potential term ``\frac{1}{M}\sum_{j=1}^M K_{\mathrm{pair}}^T \tau(K_{\mathrm{pair}}(u_i - u_j))`` with ``\tau(x) = \sigma(x) - \sigma(-x)``, conserving total canonical momentum (or position) when ``K_{\mathrm{self}} = 0`` and pairing with `InversePairwiseSymplecticLayerQ` and `InversePairwiseSymplecticLayerP` through `inverse_layer`. `CanonicalSlotLayer(M, d, NSlots)` (`M >= 2`, `1 <= NSlots <= M`, `CanonicalSlotLayer{N, NOut, M, NSlots}` with `NOut = 2 NSlots d`) first sorts the particles along coordinate row 1 of `q` with `CanonicalSortLayer`, so its slots are invariant under a relabelling of particles with distinct first coordinates. It then separates 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 Helmert relative modes, sorts the ``M - 1`` relative modes canonically by descending squared phase-space norm ``\|q_{\mathrm{rel}, k}\|_2^2 + \|p_{\mathrm{rel}, k}\|_2^2`` (`MergeSort`), and projects the sorted relative modes onto ``N_{\mathrm{slots}} - 1`` slots via a learned `StiefelManifold` parameter ``\Phi_{\mathrm{rel}} \in \mathrm{St}(N_{\mathrm{slots}} - 1, M - 1)`` (`NSlots == 1` keeps only the center-of-mass slot with no parameter leaf). `InverseCanonicalSlotLayer(M, d, NSlots)` takes the same type parameters, with input dimension `NOut` and output dimension `N`, and lifts the slots back symplectically via ``\Phi_{\mathrm{rel}}`` and the transpose Helmert matrix. `SymmetricSymplecticAutoencoder(M, d, reduced_dim; frontend = :sort | :pairwise | :slot, ...) <: SymplecticCompression` combines one of the three particle frontends with the weight-tied `WeightTiedSymplecticAutoencoder` core. With `:sort` and `: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 `:pairwise` only the frontend is permutation-equivariant, because the core is not, so the latent state is not invariant. Right-inversion ``\mathrm{encoder}(nn)(\mathrm{decoder}(nn)(\xi)) = \xi`` holds to machine precision for every ``\xi`` with `:pairwise`; with `:sort` and `:slot` it holds only for ``\xi`` whose decoded state lies in the canonical chamber, where the sorting permutations on the decoded state are the identity, because the decoder does not map into that chamber in general. Sorting is discontinuous where two particles cross, so the `:sort` and `:slot` frontends suit systems that keep their particle order, and the slot layer is a linear projection of the sorted states for a fixed particle count; the docstrings state both limits. The tests are `test/layers/symmetric_particle_layers.jl`. -- **`BilinearStiefelLayer`, `TransposeBilinearStiefelLayer`, `LowRankStiefelLayer`, `DecodeLowRankStiefelLayer`, `FastDLRATraceLoss`, and `solve!` on `BilinearStiefelLayer` add two-sided velocity-parity Stiefel frontends, a gauge-invariant dynamical low-rank Stiefel frontend and decoder, a fast low-rank trace loss, and Tucker-2 HOSVD initialization.** The five new types are exported. `BilinearStiefelLayer(Nx, Nv, mx, mv; metric_factors = nothing)` (`Nx >= 1`, even `Nv >= 2`, `1 <= mx <= Nx`, `1 <= mv <= Nv ÷ 2`) allocates `(Φx = StiefelManifold(...), Φv = StiefelManifold(...))` with ``\Phi_x \in \mathrm{St}(m_x, N_x)`` and ``\Phi_v \in \mathrm{St}(m_v, N_v / 2)``, splits an input array ``F`` of shape `(Nx, Nv)` or `(Nx, Nv, B)` into velocity-even ``F^+ = \frac{1}{\sqrt{2}}(F_{:, 1:N_v/2} + F_{:, N_v:-1:N_v/2+1})`` and velocity-odd ``F^- = \frac{1}{\sqrt{2}}(F_{:, 1:N_v/2} - F_{:, N_v:-1:N_v/2+1})`` components (or accepts a pre-split `(q, p)` `NamedTuple` of shape `(Nx, Nv ÷ 2)` or `(Nx, Nv ÷ 2, B)`), applies optional Cholesky mass factors `(Rx, Rv)` when `metric_factors` is supplied, and projects ``Q_{\mathrm{red}} = \Phi_x^T F^+ \Phi_v`` and ``P_{\mathrm{red}} = \Phi_x^T F^- \Phi_v`` to return vectorized `(q = vec(Q_red), p = vec(P_red))` of half-dimension ``n = m_x m_v``; `TransposeBilinearStiefelLayer` (`inverse_layer(l)`) is its exact right-inverse. `LowRankStiefelLayer(Nx, Nv, mx, mv; r_out = mx * mv, metric_factors = nothing)` takes a DLRA `NamedTuple` `(X = X, S = S, V = V)` of arbitrary rank ``r_{\mathrm{in}} \ge 1`` and evaluates ``Q_{\mathrm{red}} = (\Phi_x^T X) S ((V^+)^T \Phi_v)`` and ``P_{\mathrm{red}} = (\Phi_x^T X) S ((V^-)^T \Phi_v)`` in ``O(r_{\mathrm{in}}(m_x N_x + m_v N_v))`` operations without forming ``X S V^T``, matching `BilinearStiefelLayer` on ``F = X S V^T`` identically and remaining invariant under ``O(r_{\mathrm{in}}) \times O(r_{\mathrm{in}})`` gauge transformations, while `DecodeLowRankStiefelLayer` (`inverse_layer(l)`) computes the truncated SVD of ``[Q_{\mathrm{red}}\;\; P_{\mathrm{red}}]`` of rank ``r_{\mathrm{eff}} = \min(r_{\mathrm{out}}, m_x, 2 m_v)`` and QR re-orthonormalizes the lifted factors to return `(X = X_hat, S = S_hat, V = V_hat)` on the Stiefel manifold. `FastDLRATraceLoss()` evaluates ``\|X S V^T - \hat{X} \hat{S} \hat{V}^T\|_F^2`` in ``O(r^2(N_x + N_v))`` operations via Gram-matrix traces clamped at `zero(T)`. `solve!(nn, F_snapshots)` and `solve!(l, ps, F_snapshots)` initialize `Φx` and `Φv` on a `BilinearStiefelLayer` from the top-``m_x`` spatial and top-``m_v`` parity-velocity left singular vectors of `F_snapshots`. The tests are `test/layers/bilinear_and_low_rank_stiefel.jl`. +- **`BilinearStiefelLayer`, `TransposeBilinearStiefelLayer`, `LowRankStiefelLayer`, `DecodeLowRankStiefelLayer`, `FastDLRATraceLoss`, and `solve!` on `BilinearStiefelLayer` add two-sided velocity-parity Stiefel frontends, a gauge-invariant dynamical low-rank Stiefel frontend and decoder, a fast low-rank trace loss, and Tucker-2 HOSVD initialization.** The five new types are exported. `BilinearStiefelLayer(Nx, Nv, mx, mv; metric_factors = nothing)` (`Nx >= 1`, even `Nv >= 2`, `1 <= mx <= Nx`, `1 <= mv <= Nv ÷ 2`) allocates `(Φx = StiefelManifold(...), Φv = StiefelManifold(...))` with ``\Phi_x \in \mathrm{St}(m_x, N_x)`` and ``\Phi_v \in \mathrm{St}(m_v, N_v / 2)``, splits an input array ``F`` of shape `(Nx, Nv)` or `(Nx, Nv, B)` into velocity-even ``F^+ = \frac{1}{\sqrt{2}}(F_{:, 1:N_v/2} + F_{:, N_v:-1:N_v/2+1})`` and velocity-odd ``F^- = \frac{1}{\sqrt{2}}(F_{:, 1:N_v/2} - F_{:, N_v:-1:N_v/2+1})`` components (or accepts a pre-split `(q, p)` `NamedTuple` of shape `(Nx, Nv ÷ 2)` or `(Nx, Nv ÷ 2, B)`), applies optional Cholesky mass factors `(Rx, Rv)` when `metric_factors` is supplied, and projects ``Q_{\mathrm{red}} = \Phi_x^T F^+ \Phi_v`` and ``P_{\mathrm{red}} = \Phi_x^T F^- \Phi_v`` to return vectorized `(q = vec(Q_red), p = vec(P_red))` of half-dimension ``n = m_x m_v``; `TransposeBilinearStiefelLayer` (`inverse_layer(l)`) is its exact right-inverse. `LowRankStiefelLayer(Nx, Nv, mx, mv; r_out = mx * mv, metric_factors = nothing)` takes a DLRA `NamedTuple` `(X = X, S = S, V = V)` of arbitrary rank ``r_{\mathrm{in}} \ge 1`` and evaluates ``Q_{\mathrm{red}} = (\Phi_x^T X) S ((V^+)^T \Phi_v)`` and ``P_{\mathrm{red}} = (\Phi_x^T X) S ((V^-)^T \Phi_v)`` in ``O(r_{\mathrm{in}}(m_x N_x + m_v N_v))`` operations without forming ``X S V^T``, matching `BilinearStiefelLayer` on ``F = X S V^T`` identically and remaining invariant under ``O(r_{\mathrm{in}}) \times O(r_{\mathrm{in}})`` gauge transformations, while `DecodeLowRankStiefelLayer` (`inverse_layer(l)`) computes the truncated SVD of ``[Q_{\mathrm{red}}\;\; P_{\mathrm{red}}]`` of rank ``r_{\mathrm{eff}} = \min(r_{\mathrm{out}}, m_x, 2 m_v)`` and QR re-orthonormalizes the lifted factors to return `(X = X_hat, S = S_hat, V = V_hat)` on the Stiefel manifold. `FastDLRATraceLoss()` evaluates ``\|X S V^T - \hat{X} \hat{S} \hat{V}^T\|_F^2`` in ``O(r^2(N_x + N_v))`` operations via Gram-matrix traces clamped at `zero(T)`. `solve!(nn, F_snapshots)` and `solve!(l, ps, F_snapshots)` initialize `Φx` and `Φv` on a `BilinearStiefelLayer` from the top-``m_x`` spatial and top-``m_v`` parity-velocity left singular vectors of `F_snapshots`. The `(q, p)` of the parity frontends is canonical only for the structure that they declare on the parity coordinates, not for the Vlasov bracket, and the docstrings say so. The tests are `test/layers/bilinear_and_low_rank_stiefel.jl`. -- **`MultiModalSymplecticAutoencoder`, `encode`, `decode`, `translate`, `CharbonnierRelativeLoss`, `IsometricRolloutLoss`, and `CrossModalCycleLoss` add multi-modal symplectic compression and multi-modal training losses, and `HRedSys` gains reference-state shifting, reverse-mode VJP vector fields, and preallocated work buffers.** All seven new names are exported. `MultiModalSymplecticAutoencoder(frontends::NamedTuple, core::SymplecticCompression) <: SymplecticCompression` combines a non-empty `NamedTuple` of modality frontends (individual layers or `Chain`s of symplectic particle layers) whose canonical output dimension matches `core.full_dim` with a shared `SymplecticCompression` core; `encode(nn, modality, x)` (with `modality` a `Symbol` or `Val(modality)`) applies `frontends[modality]` followed by `encoder(core)`, `decode(nn, modality, ξ)` applies `decoder(core)` followed by `inverse_layer(frontends[modality])`, and `translate(nn, from, to, x)` computes `decode(nn, to, encode(nn, from, x))`, raising an `ArgumentError` on an empty or dimension-mismatched `frontends` or an unknown modality symbol. `CharbonnierRelativeLoss(; ε = 1e-3, λ_lat = 1e-2)` evaluates the exact-zero Charbonnier relative reconstruction loss ``(\sqrt{\|\hat{x} - x\|_F^2 + \varepsilon^2} - \varepsilon) / \sqrt{\|x\|_F^2 + \varepsilon^2} + \lambda_{\mathrm{lat}} \|\mathrm{encode}(nn, m, \hat{x}) - \mathrm{encode}(nn, m, x)\|_F^2 / (\|x\|_F^2 + \varepsilon^2)`` (supporting low-rank `(X, S, V)` inputs via `FastDLRATraceLoss`), `IsometricRolloutLoss(integrator; α_lat = 1.0, α_rec = 1.0, α_H = 0.1, H = nothing)` evaluates multi-step latent rollout error, decoded trajectory error, and latent Hamiltonian conservation penalty, and `CrossModalCycleLoss(mod_a, mod_b; β_align = 1.0, β_cycle = 0.1)` evaluates paired latent alignment ``\|z_a - z_b\|_F^2`` plus cycle consistency ``\|\mathrm{encode}(nn, b, \mathrm{decode}(nn, b, z_a)) - z_a\|_F^2``. `IsometricRolloutLoss` raises an `ArgumentError` naming both element types when the data, the output of `integrator` or `H` has another element type than the network's parameters, and `CrossModalCycleLoss` does so for `x_a` or `x_b`. `HRedSys`, `build_v_reduced`, and `build_f_reduced` accept optional reference states `q_ref` and `p_ref` (converted to the state element type, shifting the decoder lift by ``(q_{\mathrm{ref}}, p_{\mathrm{ref}})`` and subtracting ``(q_{\mathrm{ref}}, p_{\mathrm{ref}})`` before `rs.encoder` in `integrate_reduced_system`) and `pullback_mode::Symbol = :jacobian | :vjp`, where `:vjp` evaluates ``(\nabla_\xi \Psi^{\mathrm{dec}})^T \mathbb{J}_{2N}^T [v(z); f(z)]`` in one reverse-mode pullback `Zygote.pullback(decoder, ξ)` and both modes use the preallocated buffers of `cache` (`ξ_buf`, `v_intermediate`, `f_intermediate`, and `jac_result` for `:jacobian`, with a per-element-type `IdDict{DataType, Any}` of caches for `ForwardDiff.Dual` Newton states). The default `cache` of `build_v_reduced` and `build_f_reduced` has the element type of the decoder's parameters, and the default `cache` of `HRedSys` has the element type of `ics.q`; a state of another floating-point element type than the cache's raises an `ArgumentError`. The tests are `test/architectures/multimodal_symplectic_autoencoder.jl` and `test/reduced_system/reference_shift_and_vjp.jl`. +- **`MultiModalSymplecticAutoencoder`, `encode`, `decode`, `translate`, `CharbonnierRelativeLoss`, `IsometricRolloutLoss`, and `CrossModalLoss` add multi-modal symplectic compression and multi-modal training losses, and `HRedSys` gains reference-state shifting, reverse-mode VJP vector fields, and preallocated work buffers.** All seven new names are exported. `MultiModalSymplecticAutoencoder(frontends::NamedTuple, core::SymplecticCompression) <: SymplecticCompression` combines a non-empty `NamedTuple` of modality frontends (individual layers or `Chain`s of symplectic particle layers) whose canonical output dimension matches `core.full_dim` with a shared `SymplecticCompression` core; `encode(nn, modality, x)` (with `modality` a `Symbol` or `Val(modality)`) applies `frontends[modality]` followed by `encoder(core)`, `decode(nn, modality, ξ)` applies `decoder(core)` followed by `inverse_layer(frontends[modality])`, and `translate(nn, from, to, x)` computes `decode(nn, to, encode(nn, from, x))`, raising an `ArgumentError` on an empty or dimension-mismatched `frontends` or an unknown modality symbol. `CharbonnierRelativeLoss(; ε = 1e-3, λ_lat = 1e-2)` evaluates, for an external reconstruction `x̂`, the exact-zero Charbonnier relative reconstruction loss ``(\sqrt{\|\hat{x} - x\|_F^2 + \varepsilon^2} - \varepsilon) / \sqrt{\|x\|_F^2 + \varepsilon^2} + \lambda_{\mathrm{lat}} \|\mathrm{encode}(nn, m, \hat{x}) - \mathrm{encode}(nn, m, x)\|_F^2 / (\|x\|_F^2 + \varepsilon^2)``, and without `x̂` the same formula for ``\hat{x} = \mathrm{decode}(nn, m, \mathrm{encode}(nn, m, x))`` (supporting low-rank `(X, S, V)` inputs via `FastDLRATraceLoss`). Both forms of `CharbonnierRelativeLoss`, the cross-reconstruction term of `CrossModalLoss` and the `α_rec` term of `IsometricRolloutLoss` compare fields in the form that `decode` returns, so a pre-split pair `(q = F⁺, p = F⁻)` of a `BilinearStiefelLayer` modality is recombined into the full field first. The two halves of a pre-split pair need one element type: halves of different element types raise an `ArgumentError` naming both types, with the message `A pre-split pair (q = F⁺, p = F⁻) needs one element type, got q of element type Float64 and p of element type Float32.` `IsometricRolloutLoss(integrator; α_lat = 1.0, α_rec = 1.0, α_H = 0.0, H = nothing)` evaluates multi-step latent rollout error, decoded trajectory error, and latent Hamiltonian conservation penalty. `IsometricRolloutLoss(integrator)` has no energy term. `α_H > 0` requires `H`, and the constructor raises an `ArgumentError` without one, because a fixed latent energy such as ``\frac{1}{2}\|\hat{z}\|_2^2`` is not the energy of the system. With `α_H = 0` the loss takes no `H`, and a given `H` raises an `ArgumentError` with the message `IsometricRolloutLoss: H has no effect when α_H = 0, got α_H = 0.0 and a latent Hamiltonian H. Set α_H > 0, or pass no H.` `CrossModalLoss(mod_a, mod_b; β_align = 1.0, β_cross = 1.0)` evaluates paired latent alignment ``\|z_a - z_b\|_F^2`` plus the cross-reconstruction error ``\|\mathrm{decode}(nn, b, z_a) - x_b\|_F^2``. A negative or `NaN` weight raises an `ArgumentError`, and `x_a` and `x_b` whose latent states have different sizes, such as one sample against a batch, raise a `DimensionMismatch`. `IsometricRolloutLoss` raises an `ArgumentError` naming both element types when the data or the output of `integrator` has another element type than the network's parameters, or when `H` does so with `α_H > 0`, and `CrossModalLoss` does so for `x_a` or `x_b`. `HRedSys`, `build_v_reduced`, and `build_f_reduced` accept optional reference states `q_ref` and `p_ref` (converted to the state element type, shifting the decoder lift by ``(q_{\mathrm{ref}}, p_{\mathrm{ref}})`` and subtracting ``(q_{\mathrm{ref}}, p_{\mathrm{ref}})`` before `rs.encoder` in `integrate_reduced_system`) and `pullback_mode::Symbol = :jacobian | :vjp`, where `:vjp` evaluates ``(\nabla_\xi \Psi^{\mathrm{dec}})^T \mathbb{J}_{2N}^T [v(z); f(z)]`` in one reverse-mode pullback `Zygote.pullback(decoder, ξ)` and both modes use the preallocated buffers of `cache` (`ξ_buf`, `v_intermediate`, `f_intermediate`, and `jac_result` for `:jacobian`, with a per-element-type `IdDict{DataType, Any}` of caches for `ForwardDiff.Dual` Newton states). The default `cache` of `build_v_reduced` and `build_f_reduced` has the element type of the decoder's parameters, and the default `cache` of `HRedSys` has the element type of `ics.q`; a state of another floating-point element type than the cache's raises an `ArgumentError`. The tests are `test/architectures/multimodal_symplectic_autoencoder.jl` and `test/reduced_system/reference_shift_and_vjp.jl`. ### Changed diff --git a/KNOWN_ISSUES.md b/KNOWN_ISSUES.md index 975a238e4..ebc725192 100644 --- a/KNOWN_ISSUES.md +++ b/KNOWN_ISSUES.md @@ -76,24 +76,24 @@ What is known to be wrong in GeometricMachineLearning and is not fixed. asserts the fields with `hasfield`. The fields stay because code outside this package may need them; removing them changes the type's parameters, its layout and those tests. -### K41 · `CrossModalCycleLoss(nn, x_a, x_b)` with a mixed-component `NamedTuple` raises a `MethodError`, not the documented `ArgumentError`. +### K41 · `CrossModalLoss(nn, x_a, x_b)` with a mixed-component `NamedTuple` raises a `MethodError`, not the documented `ArgumentError`. -- **location:** `src/loss/multimodal_losses.jl:323` +- **location:** `src/loss/multimodal_losses.jl:382` - **kind:** defect - **found:** #370 - **evidence:** - The three `(nn, x_a, x_b)` methods (lines 323, 329 and 335) annotate `x_a` and `x_b` with + The three `(nn, x_a, x_b)` methods (lines 382, 388 and 394) annotate `x_a` and `x_b` with `ArrayOrNamedTuple`, and AbstractNeuralNetworks defines that as `Union{AbstractArray{T}, NamedTupleOfArrays{T}}`, with one element type `T` for every component. A `NamedTuple` whose `q` and `p` differ in element type matches no `T`, so dispatch fails before - `_check_eltype` (line 204) runs. In a fresh `julia --startup-file=no --project=.` process: + `_check_eltype` (line 251) runs. In a fresh `julia --startup-file=no --project=.` process: ```julia T = typeof((q = zeros(Float32, 2, 3), p = zeros(Float64, 2, 3))) T <: AbstractNeuralNetworks.ArrayOrNamedTuple # false ``` - The pair form `CrossModalCycleLoss(nn, (x_a, x_b))` takes `Tuple{Any, Any}`, reaches + The pair form `CrossModalLoss(nn, (x_a, x_b))` takes `Tuple{Any, Any}`, reaches `_check_eltype`, and is tested. The docstring says that another element type raises an `ArgumentError`. Closing this entry means widening the `x_a` and `x_b` annotations of the three methods, for example to `Union{AbstractArray, NamedTuple}`. That is a design call about what the diff --git a/docs/src/layers/stiefel_layer.md b/docs/src/layers/stiefel_layer.md index 3fa3670f1..8d70e2cee 100644 --- a/docs/src/layers/stiefel_layer.md +++ b/docs/src/layers/stiefel_layer.md @@ -1,9 +1,14 @@ # Bilinear and Low-Rank Stiefel Layers `BilinearStiefelLayer`, `TransposeBilinearStiefelLayer`, `LowRankStiefelLayer`, -`DecodeLowRankStiefelLayer`, and `FastDLRATraceLoss` provide canonical velocity-parity -Stiefel frontends for two-sided phase-space grid fields ``F \in \mathbb{R}^{N_x \times N_v}`` -and dynamical low-rank approximation (DLRA) states ``(X, S, V)``. +`DecodeLowRankStiefelLayer`, and `FastDLRATraceLoss` provide velocity-parity Stiefel +frontends for two-sided phase-space grid fields ``F \in \mathbb{R}^{N_x \times N_v}`` and +dynamical low-rank approximation (DLRA) states ``(X, S, V)``. + +Their ``(q, p)`` is canonical only for the structure ``\mathbb{J}`` that they 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. ## Bilinear Grid and Spline Stiefel Layers diff --git a/docs/src/layers/symmetric_particle_layers.md b/docs/src/layers/symmetric_particle_layers.md index 765da2900..dc6790f62 100644 --- a/docs/src/layers/symmetric_particle_layers.md +++ b/docs/src/layers/symmetric_particle_layers.md @@ -11,7 +11,7 @@ P_\sigma^T \mathbb{J}_{2Md} P_\sigma = \mathbb{J}_{2Md}, \qquad P_\sigma = \begin{bmatrix} \Pi_\sigma \otimes I_d & \mathbb{O} \\ \mathbb{O} & \Pi_\sigma \otimes I_d \end{bmatrix}. ``` -A standard [`PSDLayer`](@ref), [`GradientLayerQ`](@ref), or [`GradientLayerP`](@ref) breaks ``S_M`` symmetry unless the particle state is canonicalized or processed by ``S_M``-equivariant symplectic layers. `GeometricMachineLearning` provides three particle frontends and the composite architecture [`SymmetricSymplecticAutoencoder`](@ref). The sorting and slot frontends are ``S_M``-invariant; the pairwise shears are ``S_M``-equivariant. +A standard [`PSDLayer`](@ref), [`GradientLayerQ`](@ref), or [`GradientLayerP`](@ref) breaks ``S_M`` symmetry unless the particle state is canonicalized or processed by ``S_M``-equivariant symplectic layers. `GeometricMachineLearning` provides three particle frontends and the composite architecture [`SymmetricSymplecticAutoencoder`](@ref). The sorting and slot frontends are invariant under a relabelling of particles with distinct coordinates along the sorting axis; the pairwise shears are ``S_M``-equivariant. Sorting is discontinuous in ``p`` where two particles cross along the sorting axis, so the Jacobian identity of a sorting frontend holds only away from those crossings. The sorting and slot frontends suit systems that keep their particle order, not a collisionless plasma, whose sorted snapshots are not low-rank. ## Canonical Weyl-Chamber Sorting @@ -36,7 +36,7 @@ InversePairwiseSymplecticLayerP ## Canonical Slot Reduction -[`CanonicalSlotLayer`](@ref) sorts the particles along coordinate row 1 of ``q`` with [`CanonicalSortLayer`](@ref), separates the center-of-mass mode ``(Q_{\mathrm{cm}}, P_{\mathrm{cm}})`` from the ``M - 1`` zero-sum relative modes via the orthogonal Helmert / Jacobi transformation, sorts the ``M - 1`` relative modes canonically by descending squared phase-space norm ``\|q_{\mathrm{rel}, k}\|_2^2 + \|p_{\mathrm{rel}, k}\|_2^2``, and projects onto ``N_{\mathrm{slots}} - 1`` relative slots via a learned Stiefel matrix ``\Phi_{\mathrm{rel}} \in \mathrm{St}(N_{\mathrm{slots}} - 1, M - 1)``. Its analytic decoder lift is [`InverseCanonicalSlotLayer`](@ref). +[`CanonicalSlotLayer`](@ref) sorts the particles along coordinate row 1 of ``q`` with [`CanonicalSortLayer`](@ref), separates the center-of-mass mode ``(Q_{\mathrm{cm}}, P_{\mathrm{cm}})`` from the ``M - 1`` zero-sum relative modes via the orthogonal Helmert / Jacobi transformation, sorts the ``M - 1`` relative modes by descending squared phase-space norm ``\|q_{\mathrm{rel}, k}\|_2^2 + \|p_{\mathrm{rel}, k}\|_2^2``, and projects onto ``N_{\mathrm{slots}} - 1`` relative slots via a learned Stiefel matrix ``\Phi_{\mathrm{rel}} \in \mathrm{St}(N_{\mathrm{slots}} - 1, M - 1)``. Its analytic decoder lift is [`InverseCanonicalSlotLayer`](@ref). The particles are already sorted, so the 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. ```@docs CanonicalSlotLayer @@ -45,7 +45,7 @@ InverseCanonicalSlotLayer ## Symmetric Symplectic Autoencoder -[`SymmetricSymplecticAutoencoder`](@ref) combines a chosen particle frontend (`:sort`, `:pairwise`, or `:slot`) with [`WeightTiedSymplecticAutoencoder`](@ref). Only `:sort` and `:slot` give a permutation-invariant latent state; `:pairwise` is equivariant only in its frontend, because the core is not equivariant. +[`SymmetricSymplecticAutoencoder`](@ref) combines a chosen particle frontend (`:sort`, `:pairwise`, or `:slot`) with [`WeightTiedSymplecticAutoencoder`](@ref). Only `:sort` and `:slot` give a latent state that is invariant under a relabelling of particles with distinct coordinates along the sorting axis; `:pairwise` is equivariant only in its frontend, because the core is not equivariant. ```@docs SymmetricSymplecticAutoencoder diff --git a/docs/src/reduced_order_modeling/losses.md b/docs/src/reduced_order_modeling/losses.md index 0abdaa6d1..86b5ae549 100644 --- a/docs/src/reduced_order_modeling/losses.md +++ b/docs/src/reduced_order_modeling/losses.md @@ -48,7 +48,7 @@ AutoEncoderLoss ReducedLoss CharbonnierRelativeLoss IsometricRolloutLoss -CrossModalCycleLoss +CrossModalLoss GeometricMachineLearning.ParametricLoss GeometricMachineLearning.apply_parametric projection_error diff --git a/src/GeometricMachineLearning.jl b/src/GeometricMachineLearning.jl index 26873bd03..b5a186457 100644 --- a/src/GeometricMachineLearning.jl +++ b/src/GeometricMachineLearning.jl @@ -276,7 +276,7 @@ export LowRankStiefelLayer, DecodeLowRankStiefelLayer export FastDLRATraceLoss export SymmetricSymplecticAutoencoder export MultiModalSymplecticAutoencoder, encode, decode, translate -export CharbonnierRelativeLoss, IsometricRolloutLoss, CrossModalCycleLoss +export CharbonnierRelativeLoss, IsometricRolloutLoss, CrossModalLoss include("map_to_cpu.jl") diff --git a/src/architectures/multimodal_symplectic_autoencoder.jl b/src/architectures/multimodal_symplectic_autoencoder.jl index 989a34c5d..62791277f 100644 --- a/src/architectures/multimodal_symplectic_autoencoder.jl +++ b/src/architectures/multimodal_symplectic_autoencoder.jl @@ -29,7 +29,8 @@ ChainRulesCore.@non_differentiable _check_modality(::Any, ::Any) Make an instance of a multi-modal symplectic autoencoder mapping a non-empty `NamedTuple` of modality frontends into a shared canonical latent space ``\xi = (\tilde{q}, \tilde{p}) \in \mathbb{R}^{2n}`` via `core`. -Each entry of `frontends` maps a modality symbol (such as `:particles`, `:spline`, `:grid`, or `:dlra`) to a symplectic or Poisson frontend layer whose canonical output dimension equals `core.full_dim = 2N`: +Each entry of `frontends` maps a modality symbol (such as `:particles`, `:spline`, `:grid`, or `:dlra`) to a frontend layer whose output dimension equals `core.full_dim = 2N`. +The latent state is canonical for the structure that each frontend declares; for the two parity frontends, [`BilinearStiefelLayer`](@ref) and [`LowRankStiefelLayer`](@ref), that structure is the ``\mathbb{J}_{2N}`` that they declare on the parity coordinates, not the Vlasov bracket. The frontends are: - [`CanonicalSortLayer`](@ref), [`CanonicalSlotLayer`](@ref), [`PairwiseSymplecticLayerQ`](@ref), [`PairwiseSymplecticLayerP`](@ref), or a `Chain` of particle layers with output dimension ``2N``; - [`BilinearStiefelLayer`](@ref)`(Nx, Nv, Mx, Mv)` with ``2 M_x M_v = 2N`` (with or without `metric_factors` for spline vs grid modalities); - [`LowRankStiefelLayer`](@ref)`(Nx, Nv, Mx, Mv; r_out = ...)` with ``2 M_x M_v = 2N``. diff --git a/src/architectures/symmetric_symplectic_autoencoder.jl b/src/architectures/symmetric_symplectic_autoencoder.jl index 48e1f459f..ca68b5bbe 100644 --- a/src/architectures/symmetric_symplectic_autoencoder.jl +++ b/src/architectures/symmetric_symplectic_autoencoder.jl @@ -6,6 +6,7 @@ Make an instance of a symplectic autoencoder with a particle frontend for `M` pa `SymmetricSymplecticAutoencoder` combines a particle frontend (`frontend = :sort`, `:pairwise`, or `:slot`) with the weight-tied symplectic core [`WeightTiedSymplecticAutoencoder`](@ref) so that [`encoder`](@ref)`(nn)` and [`decoder`](@ref)`(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. diff --git a/src/layers/bilinear_stiefel_layer.jl b/src/layers/bilinear_stiefel_layer.jl index f3db0e6fe..ef7964750 100644 --- a/src/layers/bilinear_stiefel_layer.jl +++ b/src/layers/bilinear_stiefel_layer.jl @@ -30,6 +30,12 @@ When `metric_factors = (Rx, Rv)` is supplied with Cholesky or square-root mass f `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`](@ref) 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`](@ref), [`LowRankStiefelLayer`](@ref), and [`inverse_layer`](@ref). """ diff --git a/src/layers/canonical_slot_layer.jl b/src/layers/canonical_slot_layer.jl index fe32e6f9e..e6261498a 100644 --- a/src/layers/canonical_slot_layer.jl +++ b/src/layers/canonical_slot_layer.jl @@ -41,7 +41,7 @@ end 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` collective canonical slots (`M >= 2`, `1 <= NSlots <= M`). +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`](@ref), 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 @@ -49,7 +49,12 @@ The layer first sorts the particles along coordinate row 1 of `q` with [`Canonic ```math (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 canonically 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`](@extref GeometricOptimizers GeometricOptimizers.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. +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`](@extref GeometricOptimizers GeometricOptimizers.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. """ struct CanonicalSlotLayer{N, NOut, M, NSlots} <: AbstractExplicitLayer{N, NOut} function CanonicalSlotLayer{N, NOut, M, NSlots}() where {N, NOut, M, NSlots} @@ -181,7 +186,7 @@ end InverseCanonicalSlotLayer{N, NOut, M, NSlots}() InverseCanonicalSlotLayer(M::Integer, d::Integer, NSlots::Integer) -Analytic symplectic decoder lift (`inverse_layer(CanonicalSlotLayer(M, d, NSlots))`) mapping `NSlots` collective canonical slots back to `M` particles via ``\Phi_{\mathrm{rel}} \in \mathrm{St}(N_{\mathrm{slots}} - 1, M - 1)`` and the inverse Helmert matrix. +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`](@ref), 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}})``. diff --git a/src/layers/canonical_sort_layer.jl b/src/layers/canonical_sort_layer.jl index 302ad14db..ba089d52c 100644 --- a/src/layers/canonical_sort_layer.jl +++ b/src/layers/canonical_sort_layer.jl @@ -34,6 +34,9 @@ Make an instance of a cotangent-lifted canonical Weyl-chamber particle sorting l 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. """ struct CanonicalSortLayer{N, M, Axis} <: AbstractExplicitLayer{N, N} function CanonicalSortLayer{N, M, Axis}() where {N, M, Axis} diff --git a/src/layers/low_rank_stiefel_layer.jl b/src/layers/low_rank_stiefel_layer.jl index 587918c71..d55be9904 100644 --- a/src/layers/low_rank_stiefel_layer.jl +++ b/src/layers/low_rank_stiefel_layer.jl @@ -24,6 +24,12 @@ which produces the exact same canonical coordinates `(q = vec(Q_red), p = vec(P_ 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`](@ref) 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`](@ref), [`BilinearStiefelLayer`](@ref), [`FastDLRATraceLoss`](@ref), and [`inverse_layer`](@ref). """ diff --git a/src/loss/multimodal_losses.jl b/src/loss/multimodal_losses.jl index fac169b3f..6b90fcfd6 100644 --- a/src/loss/multimodal_losses.jl +++ b/src/loss/multimodal_losses.jl @@ -51,6 +51,27 @@ function _frob_sq_diff(x̂::NamedTuple{(:X, :S, :V)}, x::NamedTuple{(:X, :S, :V) FastDLRATraceLoss()(x̂, x) end +# The data that a decoded state of a frontend is compared with, in the form that `decode` +# returns for it. `decode` through a `BilinearStiefelLayer` returns the full field, so a pre-split +# pair `(q = F⁺, p = F⁻)` is recombined into that field; every other input is compared as it is. +# The two halves of a pre-split pair must have one element type. +function _check_split_eltypes(q, p) + eltype(q) === eltype(p) || throw(ArgumentError( + "A pre-split pair (q = F⁺, p = F⁻) needs one element type, got q of element type $(eltype(q)) and p of element type $(eltype(p)).")) + nothing +end + +ChainRulesCore.@non_differentiable _check_split_eltypes(::Any, ::Any) + +function _reconstruction_target(::BilinearStiefelLayer, x::NamedTuple{(:q, :p)}) + _check_split_eltypes(x.q, x.p) + _parity_recombine(x.q, x.p) +end +_reconstruction_target(::Any, x) = x +function _reconstruction_target(model::MultiModalSymplecticLayer, modality::Symbol, x) + _reconstruction_target(getfield(model.frontends, modality), x) +end + @doc raw""" CharbonnierRelativeLoss(; ε = 1e-3, λ_lat = 1e-2) @@ -59,10 +80,14 @@ Smooth Charbonnier relative reconstruction and latent consistency loss: ```math \mathcal{L}_{\mathrm{Charb}}(\hat{x}, x) = \frac{\sqrt{\|\hat{x} - x\|_F^2 + \varepsilon^2} - \varepsilon}{\sqrt{\|x\|_F^2 + \varepsilon^2}} -+ \lambda_{\mathrm{lat}}\,\frac{\|\mathrm{encode}(nn, m, \hat{x}) - \mathrm{encode}(nn, m, x)\|_F^2}{\|x\|_F^2 + \varepsilon^2}, ++ \lambda_{\mathrm{lat}}\,\frac{\|\mathrm{encode}(nn, m, \hat{x}) - \mathrm{encode}(nn, m, x)\|_F^2}{\|x\|_F^2 + \varepsilon^2}. ``` -where ``\varepsilon > 0`` (default `1e-3`) regularizes the Frobenius norm at ``\hat{x} = x`` and ``x = 0``, and ``\lambda_{\mathrm{lat}} \ge 0`` (default `1e-2`) weights the latent re-encoding consistency penalty. +Called as `loss(nn, m, x̂, x)`, it takes an external reconstruction ``\hat{x}``. Called as `loss(nn, m, x)` (or `loss(nn, x)` for the first modality), it uses ``\hat{x} = \mathrm{decode}(nn, m, \mathrm{encode}(nn, m, x))``. + +Both forms compare ``\hat{x}`` with ``x`` in the form that `decode` returns: a pre-split pair `(q = F⁺, p = F⁻)` of a [`BilinearStiefelLayer`](@ref) modality, passed as ``x`` or as ``\hat{x}``, is recombined into the full field. The two halves of such a pair have one element type; halves of different element types raise an `ArgumentError`. + +Here ``\varepsilon > 0`` (default `1e-3`) regularizes the Frobenius norm at ``\hat{x} = x`` and ``x = 0``, and ``\lambda_{\mathrm{lat}} \ge 0`` (default `1e-2`) weights the latent re-encoding consistency penalty. When ``x = 0`` and ``\hat{x} = 0``, the loss returns exact `zero(T)` (not `NaN`), and its gradient at ``\hat{x} = x`` is finite and zero (`0 / ε = 0`). """ @@ -97,7 +122,8 @@ Base.@constprop :aggressive function (loss::CharbonnierRelativeLoss)( T = promote_type(_state_eltype(x̂), _state_eltype(x), _leaf_eltype(ps.core)) z = encode(model, ps, modality, x) z_hat = encode(model, ps, modality, x̂) - _charbonnier_eval(loss, x̂, x, z_hat, z)::T + _charbonnier_eval(loss, _reconstruction_target(model, modality, x̂), + _reconstruction_target(model, modality, x), z_hat, z)::T end Base.@constprop :aggressive function (loss::CharbonnierRelativeLoss)( @@ -116,7 +142,7 @@ Base.@constprop :aggressive function (loss::CharbonnierRelativeLoss)( z = encode(model, ps, modality, x) x̂ = decode(model, ps, modality, z) z_hat = encode(model, ps, modality, x̂) - _charbonnier_eval(loss, x̂, x, z_hat, z)::T + _charbonnier_eval(loss, x̂, _reconstruction_target(model, modality, x), z_hat, z)::T end Base.@constprop :aggressive function (loss::CharbonnierRelativeLoss)( @@ -146,7 +172,7 @@ Base.@constprop :aggressive function (loss::CharbonnierRelativeLoss)( end @doc raw""" - IsometricRolloutLoss(integrator; α_lat = 1.0, α_rec = 1.0, α_H = 0.1, H = nothing) + IsometricRolloutLoss(integrator; α_lat = 1.0, α_rec = 1.0, α_H = 0.0, H = nothing) Multi-step latent rollout, decoded trajectory reconstruction, and latent Hamiltonian conservation loss: @@ -160,9 +186,13 @@ Multi-step latent rollout, decoded trajectory reconstruction, and latent Hamilto \right), ``` -where ``z_k = \mathrm{encode}(nn, m, x_k)``, ``\hat{z}_1 = z_1``, ``\hat{z}_{k+1} = \mathrm{integrator}(\hat{z}_k)`` for ``k = 1, \dots, K - 1``, and `H` is a callable latent Hamiltonian (default `nothing`, which uses the canonical quadratic energy ``\frac{1}{2}\|\hat{z}\|_2^2``). -The weights default to `α_lat = 1.0`, `α_rec = 1.0`, and `α_H = 0.1`. -The data, the output of `integrator` and the value of `H` have the element type of the network's parameters; another element type raises an `ArgumentError`. +where ``z_k = \mathrm{encode}(nn, m, x_k)``, ``\hat{z}_1 = z_1``, ``\hat{z}_{k+1} = \mathrm{integrator}(\hat{z}_k)`` for ``k = 1, \dots, K - 1``, and `H` is a callable latent Hamiltonian. +`H(ẑ)` takes the latent state as `encode` and `integrator` give it, a vector of length ``2n`` or a matrix of size ``(2n, B)`` for a batched snapshot, and returns one scalar, one value for a whole batch. +The decoded state is compared with ``x_k`` in the form that `decode` returns: a pre-split pair `(q = F⁺, p = F⁻)` of a [`BilinearStiefelLayer`](@ref) modality is recombined into the full field. +The weights default to `α_lat = 1.0`, `α_rec = 1.0`, and `α_H = 0.0`, so `IsometricRolloutLoss(integrator)` has no energy term. + +`H` has no default: the loss requires a latent Hamiltonian `H` when `α_H > 0`, and the constructor raises an `ArgumentError` without one. A fixed latent energy such as ``\frac{1}{2}\|\hat{z}\|_2^2`` is not the energy of the system. With `α_H = 0`, the loss has no energy term and takes no `H`: a given `H` would have no effect, so the constructor raises an `ArgumentError`. +The data, the output of `integrator` and, with `α_H > 0`, the value of `H` have the element type of the network's parameters; another element type raises an `ArgumentError`. """ struct IsometricRolloutLoss{IT, HT, T <: Real} <: NetworkLoss integrator::IT @@ -175,18 +205,35 @@ struct IsometricRolloutLoss{IT, HT, T <: Real} <: NetworkLoss integrator; α_lat::Real = 1.0, α_rec::Real = 1.0, - α_H::Real = 0.1, + α_H::Real = 0.0, H = nothing) (α_lat >= 0 && α_rec >= 0 && α_H >= 0) || throw(ArgumentError( "IsometricRolloutLoss requires non-negative weights α_lat, α_rec, α_H, got α_lat = $α_lat, α_rec = $α_rec, α_H = $α_H.")) + (iszero(α_H) || H !== nothing) || throw(ArgumentError( + "IsometricRolloutLoss requires a latent Hamiltonian H when α_H > 0, got α_H = $α_H and H = nothing. Pass H, or set α_H = 0.")) + (!iszero(α_H) || H === nothing) || throw(ArgumentError( + "IsometricRolloutLoss: H has no effect when α_H = 0, got α_H = $α_H and a latent Hamiltonian H. Set α_H > 0, or pass no H.")) T = promote_type(typeof(float(α_lat)), typeof(float(α_rec)), typeof(float(α_H))) new{typeof(integrator), typeof(H), T}( integrator, H, T(α_lat), T(α_rec), T(α_H)) end end -_eval_latent_H(::Nothing, z::AbstractArray{T}) where {T} = T(0.5) * sum(abs2, z) -_eval_latent_H(H_fn, z::AbstractArray) = H_fn(z) +# The value of `H` at the first latent state, and the penalty `|H(ẑ) - H(ẑ₁)|²` of a later one. +# `H = nothing` occurs only with `α_H = 0`, and then `H` is never called. +_initial_latent_H(::Nothing, z, ::Type{T}) where {T} = nothing +function _initial_latent_H(H, z, ::Type{T}) where {T} + H1 = H(z) + _check_eltype("IsometricRolloutLoss", "the value of H", H1, T) + H1 +end + +_latent_H_penalty(::Nothing, z, ::Nothing, ::Type{T}) where {T} = zero(T) +function _latent_H_penalty(H, z, H1, ::Type{T}) where {T} + Hk = H(z) + _check_eltype("IsometricRolloutLoss", "the value of H", Hk, T) + abs2(Hk - H1) +end function _check_traj_len(K::Int) K >= 2 || throw(ArgumentError( @@ -219,8 +266,7 @@ Base.@constprop :aggressive function (loss::IsometricRolloutLoss)( T = _leaf_eltype(ps.core) _check_eltype("IsometricRolloutLoss", "the data", xs[1], T) z1 = encode(model, ps, modality, xs[1]) - H1 = _eval_latent_H(loss.H, z1) - _check_eltype("IsometricRolloutLoss", "the value of H", H1, T) + H1 = _initial_latent_H(loss.H, z1, T) z_curr = z1 err_lat = zero(T) err_rec = zero(T) @@ -232,10 +278,9 @@ Base.@constprop :aggressive function (loss::IsometricRolloutLoss)( _check_eltype("IsometricRolloutLoss", "the integrator's output", z_curr, T) err_lat = err_lat + _frob_sq_diff(z_curr, z_k) x_rec_k = decode(model, ps, modality, z_curr) - err_rec = err_rec + _frob_sq_diff(x_rec_k, xs[k]) - Hk = _eval_latent_H(loss.H, z_curr) - _check_eltype("IsometricRolloutLoss", "the value of H", Hk, T) - err_H = err_H + abs2(Hk - H1) + err_rec = err_rec + + _frob_sq_diff(x_rec_k, _reconstruction_target(model, modality, xs[k])) + err_H = err_H + _latent_H_penalty(loss.H, z_curr, H1, T) end inv_steps = inv(T(K - 1)) res = T(loss.α_lat) * err_lat * inv_steps + @@ -273,80 +318,94 @@ Base.@constprop :aggressive function (loss::IsometricRolloutLoss)( end @doc raw""" - CrossModalCycleLoss(mod_a::Symbol, mod_b::Symbol; β_align = 1.0, β_cycle = 0.1) + CrossModalLoss(mod_a::Symbol, mod_b::Symbol; β_align = 1.0, β_cross = 1.0) -Cross-modal paired latent alignment and cycle consistency loss between modalities `mod_a` and `mod_b`: +Cross-modal paired latent alignment and cross-reconstruction loss between modalities `mod_a` and `mod_b`: ```math \mathcal{L}_{\mathrm{cross}}(x_a, x_b) = \beta_{\mathrm{align}}\,\|z_a - z_b\|_F^2 -+ \beta_{\mathrm{cycle}}\,\|\mathrm{encode}(nn, b, \mathrm{decode}(nn, b, z_a)) - z_a\|_F^2, ++ \beta_{\mathrm{cross}}\,\|\mathrm{decode}(nn, b, z_a) - x_b\|_F^2, ``` where ``z_a = \mathrm{encode}(nn, a, x_a)`` and ``z_b = \mathrm{encode}(nn, b, x_b)``. -The weights default to `β_align = 1.0` and `β_cycle = 0.1`. +The pair ``(x_a, x_b)`` is the same physical state in two modalities: the alignment term asks that both encode to one latent state, and the cross-reconstruction term asks that the latent state of ``x_a`` decodes to ``x_b``. +The loss is not symmetric in ``a`` and ``b``: the cross-reconstruction term decodes only into modality `mod_b`, so `CrossModalLoss(mod_b, mod_a)` is another loss. +`decode(nn, b, z_a)` is compared with ``x_b`` in the form that `decode` returns: a pre-split pair `(q = F⁺, p = F⁻)` of a [`BilinearStiefelLayer`](@ref) modality is recombined into the full field. +A batch is the sum of its samples' losses; `x_a` and `x_b` with different numbers of samples raise a `DimensionMismatch`. + +The weights are non-negative and default to `β_align = 1.0` and `β_cross = 1.0`; a negative or `NaN` weight raises an `ArgumentError`. `x_a` and `x_b` have the element type of the network's parameters; another element type raises an `ArgumentError`. """ -struct CrossModalCycleLoss{T <: Real} <: NetworkLoss +struct CrossModalLoss{T <: Real} <: NetworkLoss mod_a::Symbol mod_b::Symbol β_align::T - β_cycle::T + β_cross::T - function CrossModalCycleLoss( + function CrossModalLoss( mod_a::Symbol, mod_b::Symbol; - β_align::Real = 1.0, β_cycle::Real = 0.1) - (β_align >= 0 && β_cycle >= 0) || throw(ArgumentError( - "CrossModalCycleLoss requires non-negative weights β_align and β_cycle, got β_align = $β_align, β_cycle = $β_cycle.")) - T = promote_type(typeof(float(β_align)), typeof(float(β_cycle))) - new{T}(mod_a, mod_b, T(β_align), T(β_cycle)) + β_align::Real = 1.0, β_cross::Real = 1.0) + (β_align >= 0 && β_cross >= 0) || throw(ArgumentError( + "CrossModalLoss requires non-negative weights β_align and β_cross, got β_align = $β_align, β_cross = $β_cross.")) + T = promote_type(typeof(float(β_align)), typeof(float(β_cross))) + new{T}(mod_a, mod_b, T(β_align), T(β_cross)) end end +function _check_sample_count(z_a, z_b) + size(z_a) == size(z_b) || throw(DimensionMismatch( + "CrossModalLoss: the latent states of x_a and x_b have different sizes, $(size(z_a)) and $(size(z_b)).")) + nothing +end + +ChainRulesCore.@non_differentiable _check_sample_count(::Any, ::Any) + Base.@constprop :aggressive function _cross_modal_eval( - loss::CrossModalCycleLoss, model::MultiModalSymplecticLayer, + loss::CrossModalLoss, model::MultiModalSymplecticLayer, ps::NamedTuple, x_a, x_b) T = _leaf_eltype(ps.core) - _check_eltype("CrossModalCycleLoss", "x_a", x_a, T) - _check_eltype("CrossModalCycleLoss", "x_b", x_b, T) + _check_eltype("CrossModalLoss", "x_a", x_a, T) + _check_eltype("CrossModalLoss", "x_b", x_b, T) z_a = encode(model, ps, loss.mod_a, x_a) z_b = encode(model, ps, loss.mod_b, x_b) - x_b_rec = decode(model, ps, loss.mod_b, z_a) - z_a_cyc = encode(model, ps, loss.mod_b, x_b_rec) + _check_sample_count(z_a, z_b) + x̂_b = decode(model, ps, loss.mod_b, z_a) + target = _reconstruction_target(model, loss.mod_b, x_b) res = T(loss.β_align) * _frob_sq_diff(z_a, z_b) + - T(loss.β_cycle) * _frob_sq_diff(z_a_cyc, z_a) + T(loss.β_cross) * _frob_sq_diff(x̂_b, target) res::T end -Base.@constprop :aggressive function (loss::CrossModalCycleLoss)( +Base.@constprop :aggressive function (loss::CrossModalLoss)( model::MultiModalSymplecticLayer, ps::NamedTuple, x_a::ArrayOrNamedTuple, x_b::ArrayOrNamedTuple) _cross_modal_eval(loss, model, ps, x_a, x_b) end -Base.@constprop :aggressive function (loss::CrossModalCycleLoss)( +Base.@constprop :aggressive function (loss::CrossModalLoss)( model::MultiModalSymplecticLayer, ps::NetworkParameters, x_a::ArrayOrNamedTuple, x_b::ArrayOrNamedTuple) _cross_modal_eval(loss, model, params(ps), x_a, x_b) end -Base.@constprop :aggressive function (loss::CrossModalCycleLoss)( +Base.@constprop :aggressive function (loss::CrossModalLoss)( nn::NeuralNetwork{<:MultiModalSymplecticAutoencoder}, x_a::ArrayOrNamedTuple, x_b::ArrayOrNamedTuple) loss(nn.model, params(nn), x_a, x_b) end -Base.@constprop :aggressive function (loss::CrossModalCycleLoss)( +Base.@constprop :aggressive function (loss::CrossModalLoss)( model::MultiModalSymplecticLayer, ps::NamedTuple, pair::Tuple{Any, Any}) _cross_modal_eval(loss, model, ps, pair[1], pair[2]) end -Base.@constprop :aggressive function (loss::CrossModalCycleLoss)( +Base.@constprop :aggressive function (loss::CrossModalLoss)( model::MultiModalSymplecticLayer, ps::NetworkParameters, pair::Tuple{Any, Any}) _cross_modal_eval(loss, model, params(ps), pair[1], pair[2]) end -Base.@constprop :aggressive function (loss::CrossModalCycleLoss)( +Base.@constprop :aggressive function (loss::CrossModalLoss)( nn::NeuralNetwork{<:MultiModalSymplecticAutoencoder}, pair::Tuple{Any, Any}) loss(nn.model, params(nn), pair) end diff --git a/test/architectures/multimodal_symplectic_autoencoder.jl b/test/architectures/multimodal_symplectic_autoencoder.jl index 8235a81bb..962e22b92 100644 --- a/test/architectures/multimodal_symplectic_autoencoder.jl +++ b/test/architectures/multimodal_symplectic_autoencoder.jl @@ -278,10 +278,10 @@ end @testset "Multi-modal losses, zero-input edge case, exact-zero loss, and Zygote vs FD gradients" begin @test isdefined(GeometricMachineLearning, :CharbonnierRelativeLoss) @test isdefined(GeometricMachineLearning, :IsometricRolloutLoss) - @test isdefined(GeometricMachineLearning, :CrossModalCycleLoss) + @test isdefined(GeometricMachineLearning, :CrossModalLoss) @test CharbonnierRelativeLoss <: NetworkLoss @test IsometricRolloutLoss <: NetworkLoss - @test CrossModalCycleLoss <: NetworkLoss + @test CrossModalLoss <: NetworkLoss # Default keyword checks l_charb_def = CharbonnierRelativeLoss() @@ -289,16 +289,35 @@ end @test l_charb_def.λ_lat == 1e-2 step_id = z -> z + H_lat_def = z -> sum(abs2, z) / 2 l_iso_def = IsometricRolloutLoss(step_id) @test l_iso_def.α_lat == 1.0 @test l_iso_def.α_rec == 1.0 - @test l_iso_def.α_H == 0.1 - - l_cyc_def = CrossModalCycleLoss(:particles, :grid) + @test l_iso_def.α_H == 0.0 + @test l_iso_def.H === nothing + @test IsometricRolloutLoss(step_id; α_H = 0.1, H = H_lat_def).H === H_lat_def + + # There is no default latent Hamiltonian: α_H > 0 requires H, and α_H = 0 needs none. + re_iso_H = r"^(?:ArgumentError: )?IsometricRolloutLoss requires a latent Hamiltonian H when α_H > 0, got α_H = 0\.1 and H = nothing\. Pass H, or set α_H = 0\.$" + @test_throws re_iso_H IsometricRolloutLoss(step_id; α_H = 0.1) + @test_throws re_iso_H IsometricRolloutLoss(step_id; α_H = 0.1, H = nothing) + @test IsometricRolloutLoss(step_id; α_H = 0.0).H === nothing + # With α_H = 0 a given H has no effect, so the constructor refuses it. + re_iso_H0 = r"^(?:ArgumentError: )?IsometricRolloutLoss: H has no effect when α_H = 0, got α_H = 0\.0 and a latent Hamiltonian H\. Set α_H > 0, or pass no H\.$" + @test_throws re_iso_H0 IsometricRolloutLoss(step_id; H = H_lat_def) + @test_throws re_iso_H0 IsometricRolloutLoss(step_id; α_H = 0.0, H = H_lat_def) + # A negative zero is zero: it is refused with an H and accepted without one. + re_iso_H0_neg = r"^(?:ArgumentError: )?IsometricRolloutLoss: H has no effect when α_H = 0, got α_H = -0\.0 and a latent Hamiltonian H\. Set α_H > 0, or pass no H\.$" + @test_throws re_iso_H0_neg IsometricRolloutLoss(step_id; α_H = -0.0, H = H_lat_def) + @test IsometricRolloutLoss(step_id; α_H = -0.0).H === nothing + + l_cyc_def = CrossModalLoss(:particles, :grid) @test l_cyc_def.mod_a === :particles @test l_cyc_def.mod_b === :grid @test l_cyc_def.β_align == 1.0 - @test l_cyc_def.β_cycle == 0.1 + @test l_cyc_def.β_cross == 1.0 + @test !hasfield(typeof(l_cyc_def), :β_cycle) + @test_throws MethodError CrossModalLoss(:particles, :grid; β_cycle = 0.1) # Edge validation for loss hyperparameters re_charb_eps = r"^(?:ArgumentError: )?CharbonnierRelativeLoss requires ε > 0 and λ_lat >= 0, got ε = 0\.0, λ_lat = 0\.01\.$" @@ -306,11 +325,15 @@ end @test_throws re_charb_eps CharbonnierRelativeLoss(; ε = 0.0) @test_throws re_charb_lam CharbonnierRelativeLoss(; λ_lat = -0.1) - re_iso = r"^(?:ArgumentError: )?IsometricRolloutLoss requires non-negative weights α_lat, α_rec, α_H, got α_lat = -1\.0, α_rec = 1\.0, α_H = 0\.1\.$" + re_iso = r"^(?:ArgumentError: )?IsometricRolloutLoss requires non-negative weights α_lat, α_rec, α_H, got α_lat = -1\.0, α_rec = 1\.0, α_H = 0\.0\.$" @test_throws re_iso IsometricRolloutLoss(step_id; α_lat = -1.0) - re_cyc = r"^(?:ArgumentError: )?CrossModalCycleLoss requires non-negative weights β_align and β_cycle, got β_align = 1\.0, β_cycle = -0\.1\.$" - @test_throws re_cyc CrossModalCycleLoss(:particles, :grid; β_cycle = -0.1) + re_cross = r"^(?:ArgumentError: )?CrossModalLoss requires non-negative weights β_align and β_cross, got β_align = 1\.0, β_cross = -0\.1\.$" + @test_throws re_cross CrossModalLoss(:particles, :grid; β_cross = -0.1) + re_cross_nan = r"^(?:ArgumentError: )?CrossModalLoss requires non-negative weights β_align and β_cross, got β_align = 1\.0, β_cross = NaN\.$" + @test_throws re_cross_nan CrossModalLoss(:particles, :grid; β_cross = NaN) + re_align = r"^(?:ArgumentError: )?CrossModalLoss requires non-negative weights β_align and β_cross, got β_align = -0\.1, β_cross = 1\.0\.$" + @test_throws re_align CrossModalLoss(:particles, :grid; β_align = -0.1) @testset "Forward evaluation, @inferred, zero edge case, and exact-zero loss ($T)" for T in REAL_ELTYPES rng = Random.MersenneTwister(202) @@ -333,16 +356,46 @@ end z_hat = encode(nn, :grid, F_hat) eps_T = T(1e-2) lam_T = T(5e-2) - expected_c = (sqrt(sum(abs2, F_hat .- F) + eps_T^2) - eps_T) / - sqrt(sum(abs2, F) + eps_T^2) + + expected_c_rec = (sqrt(sum(abs2, F_hat .- F) + eps_T^2) - eps_T) / + sqrt(sum(abs2, F) + eps_T^2) + expected_c = expected_c_rec + lam_T * sum(abs2, z_hat .- z) / (sum(abs2, F) + eps_T^2) + # Both forms have the latent term; the self-reconstruction form uses x̂ = decode(encode(x)). val_c1 = @inferred loss_c(nn, :grid, F) val_c2 = @inferred loss_c(nn, :grid, F_hat, F) @test eltype(val_c1) === T && eltype(val_c2) === T @test isapprox(val_c1, expected_c; atol = 10 * eps(T), rtol = 10 * eps(T)) @test isapprox(val_c2, expected_c; atol = 10 * eps(T), rtol = 10 * eps(T)) + # In the self-reconstruction form, a pre-split field (q = F⁺, p = F⁻) gives the value of + # the full field F: the decoded field is compared with F⁺ and F⁻ recombined. + F_plus = (F[:, 1:(nv ÷ 2)] .+ F[:, nv:-1:(nv ÷ 2 + 1)]) ./ sqrt(T(2)) + F_minus = (F[:, 1:(nv ÷ 2)] .- F[:, nv:-1:(nv ÷ 2 + 1)]) ./ sqrt(T(2)) + F_split = (q = F_plus, p = F_minus) + @test isapprox(@inferred(loss_c(nn, :grid, F_split)), val_c1; + atol = 10 * eps(T), rtol = 10 * eps(T)) + # The form with an external x̂ recombines a pre-split x and a pre-split x̂ in the same way. + @test isapprox(@inferred(loss_c(nn, :grid, F_hat, F_split)), val_c2; + atol = 10 * eps(T), rtol = 10 * eps(T)) + @test @inferred(loss_c(nn, :grid, F_split, F_split)) === zero(T) + # The two halves of a pre-split pair have one element type: halves of different element + # types raise an ArgumentError in both forms, and a pair of one element type gives the + # value of the full field that it recombines to by hand. + if T === Float64 + F_split_mixed = (q = F_plus, p = Float32.(F_minus)) + re_mixed = r"^(?:ArgumentError: )?A pre-split pair \(q = F⁺, p = F⁻\) needs one element type, got q of element type Float64 and p of element type Float32\.$" + @test_throws re_mixed loss_c(nn, :grid, F_hat, F_split_mixed) + @test_throws re_mixed loss_c(nn, :grid, F_split_mixed) + F_minus_same = Float64.(F_split_mixed.p) + F_split_same = (q = F_plus, p = F_minus_same) + F_same_full = hcat((F_plus .+ F_minus_same) ./ sqrt(T(2)), + ((F_plus .- F_minus_same) ./ sqrt(T(2)))[:, end:-1:1]) + val_same = @inferred loss_c(nn, :grid, F_hat, F_split_same) + val_same_full = loss_c(nn, :grid, F_hat, F_same_full) + @test isapprox(val_same, val_same_full; rtol = 10 * eps(T)) + end + # Independent perturbation F_pert != F_hat so z_hat != z and the latent penalty term is active (catches a dropped latent penalty) F_pert = F .+ T(0.5) .* ones(T, nx, nv) z_pert = encode(nn, :grid, F_pert) @@ -391,9 +444,34 @@ end expected_iso = T(1.0) * err_lat + T(0.5) * err_rec + T(0.2) * err_H @test isapprox(val_iso, expected_iso; atol = 15 * eps(T), rtol = 15 * eps(T)) - # CrossModalCycleLoss forward & @inferred - loss_cyc = CrossModalCycleLoss( - :particles, :grid; β_align = T(1.2), β_cycle = T(0.3) + # With α_H = 0 and no H, the loss is the α_lat and α_rec terms alone. + loss_iso_noH = IsometricRolloutLoss(step_lin; α_lat = T(1.0), α_rec = T(0.5), α_H = 0) + val_iso_noH = @inferred loss_iso_noH(nn, :grid, F_traj) + @test eltype(val_iso_noH) === T + @test isapprox(val_iso_noH, T(1.0) * err_lat + T(0.5) * err_rec; + atol = 15 * eps(T), rtol = 15 * eps(T)) + + # IsometricRolloutLoss(integrator) has the weights α_lat = α_rec = 1 and no energy term. + loss_iso_def = IsometricRolloutLoss(step_lin) + val_iso_def = @inferred loss_iso_def(nn, :grid, F_traj) + @test eltype(val_iso_def) === T + @test isapprox(val_iso_def, err_lat + err_rec; + atol = 15 * eps(T), rtol = 15 * eps(T)) + + # A trajectory of pre-split fields gives the value of the trajectory of full fields. + split_traj = [( + q = (F_k[:, 1:(nv ÷ 2)] .+ F_k[:, nv:-1:(nv ÷ 2 + 1)]) ./ + sqrt(T(2)), + p = (F_k[:, 1:(nv ÷ 2)] .- F_k[:, nv:-1:(nv ÷ 2 + 1)]) ./ + sqrt(T(2))) + for F_k in F_traj] + @test isapprox(@inferred(loss_iso(nn, :grid, split_traj)), val_iso; + atol = 15 * eps(T), rtol = 15 * eps(T)) + + # CrossModalLoss forward & @inferred: the alignment term and the + # cross-reconstruction term ‖decode(b, z_a) - x_b‖². + loss_cyc = CrossModalLoss( + :particles, :grid; β_align = T(1.2), β_cross = T(0.3) ) x_p = (q = randn(rng, T, 1, 3), p = randn(rng, T, 1, 3)) val_cyc = @inferred loss_cyc(nn, x_p, F) @@ -401,34 +479,83 @@ end z_a = encode(nn, :particles, x_p) z_b = encode(nn, :grid, F) - z_a_cyc = encode(nn, :grid, decode(nn, :grid, z_a)) - expected_cyc = T(1.2) * sum(abs2, z_a .- z_b) + T(0.3) * sum(abs2, z_a_cyc .- z_a) + expected_cyc = T(1.2) * sum(abs2, z_a .- z_b) + + T(0.3) * sum(abs2, decode(nn, :grid, z_a) .- F) @test isapprox(val_cyc, expected_cyc; atol = 10 * eps(T), rtol = 10 * eps(T)) - # Reverse direction CrossModalCycleLoss(:grid, :particles) where :particles is only chamber-wise right-invertible (checks the cycle term against the explicit encode-decode round trip) - loss_cyc_rev = CrossModalCycleLoss( - :grid, :particles; β_align = T(1.2), β_cycle = T(0.5)) + # Reverse direction CrossModalLoss(:grid, :particles): the particle decoder returns + # the flat state [vec(q); vec(p)]. + loss_cyc_rev = CrossModalLoss( + :grid, :particles; β_align = T(1.2), β_cross = T(0.5)) val_cyc_rev = @inferred loss_cyc_rev(nn, F, x_p) - z_b_cyc = encode(nn, :particles, decode(nn, :particles, z_b)) expected_cyc_rev = T(1.2) * sum(abs2, z_b .- z_a) + - T(0.5) * sum(abs2, z_b_cyc .- z_b) + T(0.5) * sum( + abs2, decode(nn, :particles, z_b) .- vcat(vec(x_p.q), vec(x_p.p))) @test isapprox( - val_cyc_rev, expected_cyc_rev; atol = 10 * eps(T), rtol = 10 * - eps(T) - ) + val_cyc_rev, expected_cyc_rev; atol = 10 * eps(T), rtol = 10 * eps(T)) + + # The three receivers, each with x_a and x_b and with the pair (x_a, x_b). + ps_nt, ps_np = params(nn), nn.params + @test (@inferred loss_cyc(nn, (x_p, F))) == val_cyc + @test (@inferred loss_cyc(nn.model, ps_nt, x_p, F)) == val_cyc + @test (@inferred loss_cyc(nn.model, ps_nt, (x_p, F))) == val_cyc + @test (@inferred loss_cyc(nn.model, ps_np, x_p, F)) == val_cyc + @test (@inferred loss_cyc(nn.model, ps_np, (x_p, F))) == val_cyc + + # A batch is the sum over its samples, with no division by B. + B = 3 + x_p_b = (q = randn(rng, T, 1, 3, B), p = randn(rng, T, 1, 3, B)) + F_b = randn(rng, T, nx, nv, B) + val_cyc_b = @inferred loss_cyc(nn, x_p_b, F_b) + @test eltype(val_cyc_b) === T + val_cyc_sum = sum( + b -> loss_cyc(nn, (q = x_p_b.q[:, :, b], p = x_p_b.p[:, :, b]), F_b[:, :, b]), + 1:B) + @test isapprox(val_cyc_b, val_cyc_sum; rtol = 10 * eps(T)) + @test (@inferred loss_cyc(nn, (x_p_b, F_b))) == val_cyc_b + + # x_a and x_b with different numbers of samples raise a DimensionMismatch. + re_samples = r"^(?:DimensionMismatch: )?CrossModalLoss: the latent states of x_a and x_b have different sizes, \(2,\) and \(2, 3\)\.$" + @test_throws re_samples loss_cyc(nn, x_p, F_b) + re_samples_rev = r"^(?:DimensionMismatch: )?CrossModalLoss: the latent states of x_a and x_b have different sizes, \(2, 3\) and \(2,\)\.$" + @test_throws re_samples_rev loss_cyc(nn, x_p_b, F) + + # A pre-split field (q = F⁺, p = F⁻) for x_b gives the value of the full field F. The + # parity components come from explicit parity matrices: vec(F⁺) = P₊ vec(F). + half = nv ÷ 2 + P_plus = zeros(T, nx * half, nx * nv) + P_minus = zeros(T, nx * half, nx * nv) + for j in 1:half, i in 1:nx + + row = (j - 1) * nx + i + col_top = (j - 1) * nx + i + col_bot = (nv - j) * nx + i + P_plus[row, col_top] = P_plus[row, col_bot] = inv(sqrt(T(2))) + P_minus[row, col_top] = inv(sqrt(T(2))) + P_minus[row, col_bot] = -inv(sqrt(T(2))) + end + split_of(F2) = (q = reshape(P_plus * vec(F2), nx, half), + p = reshape(P_minus * vec(F2), nx, half)) + @test isapprox(@inferred(loss_cyc(nn, x_p, split_of(F))), val_cyc; + rtol = 10 * eps(T)) + F_b_split = (q = cat((split_of(F_b[:, :, b]).q for b in 1:B)...; dims = 3), + p = cat((split_of(F_b[:, :, b]).p for b in 1:B)...; dims = 3)) + @test isapprox(@inferred(loss_cyc(nn, x_p_b, F_b_split)), val_cyc_b; + rtol = 10 * eps(T)) # Another precision than the network's raises an ArgumentError that names both. S = T === Float32 ? Float64 : Float32 msg(name, what) = "$name: $what has element type $S, but the network has element type $T." - loss_iso_int = IsometricRolloutLoss(ξ -> S.(step_lin(ξ)); H = H_lat) + loss_iso_int = IsometricRolloutLoss(ξ -> S.(step_lin(ξ)); α_H = 0.1, H = H_lat) @test_throws msg("IsometricRolloutLoss", "the integrator's output") loss_iso_int( nn, :grid, F_traj) - loss_iso_H = IsometricRolloutLoss(step_lin; H = ξ -> S(0.5) * S(sum(abs2, ξ))) + loss_iso_H = IsometricRolloutLoss( + step_lin; α_H = 0.1, H = ξ -> S(0.5) * S(sum(abs2, ξ))) @test_throws msg("IsometricRolloutLoss", "the value of H") loss_iso_H( nn, :grid, F_traj) @test_throws msg("IsometricRolloutLoss", "the data") loss_iso( nn, :grid, [S.(F_k) for F_k in F_traj]) - @test_throws msg("CrossModalCycleLoss", "x_b") loss_cyc(nn, x_p, S.(F)) + @test_throws msg("CrossModalLoss", "x_b") loss_cyc(nn, x_p, S.(F)) # A particle state with one component of another precision than the network's raises # the same ArgumentError, whichever component it is. @@ -439,8 +566,8 @@ end nn, :particles, [x_mixed, x_p]) @test_throws msg("IsometricRolloutLoss", "the data") loss_iso( nn, :particles, [x_p, x_mixed]) - @test_throws msg("CrossModalCycleLoss", "x_a") loss_cyc(nn, (x_mixed, F)) - @test_throws msg("CrossModalCycleLoss", "x_b") loss_cyc_rev(nn, (F, x_mixed)) + @test_throws msg("CrossModalLoss", "x_a") loss_cyc(nn, (x_mixed, F)) + @test_throws msg("CrossModalLoss", "x_b") loss_cyc_rev(nn, (F, x_mixed)) end end @@ -473,7 +600,7 @@ end H_fn = ξ -> 0.5 * sum(abs2, ξ) + 0.05 * sum(ξ .^ 4) loss_iso = IsometricRolloutLoss( step_fn; α_lat = 1.0, α_rec = 0.8, α_H = 0.25, H = H_fn) - loss_cyc = CrossModalCycleLoss(:particles, :grid; β_align = 1.0, β_cycle = 0.2) + loss_cyc = CrossModalLoss(:particles, :grid; β_align = 1.0, β_cross = 0.2) flat0, rebuild_ps = _flatten_params(nn.params) h = 1e-6 @@ -482,7 +609,7 @@ end for (name, f_ps) in ( ("CharbonnierRelativeLoss", ps -> loss_c(nn.model, ps, :grid, F)), ("IsometricRolloutLoss", ps -> loss_iso(nn.model, ps, :grid, F_traj)), - ("CrossModalCycleLoss", ps -> loss_cyc(nn.model, ps, x_p, F)) + ("CrossModalLoss", ps -> loss_cyc(nn.model, ps, x_p, F)) ) g_zyg = Zygote.gradient(f_ps, nn.params)[1] g_vec = _extract_grad_vec(g_zyg) @@ -663,3 +790,140 @@ end @test all(isfinite, _grad_vec(g_sym)) end end + +# The untied core of the control networks below. +function _untied_core() + SymplecticAutoencoder( + 4, 2; n_encoder_blocks = 1, n_decoder_blocks = 1, n_encoder_layers = 1, + n_decoder_layers = 1, n_decoder_output_layers = 1, sympnet_upscale = 2) +end + +@testset "Cross-reconstruction term against a dense linear oracle ($T)" for T in REAL_ELTYPES + rng = Random.MersenneTwister(505) + Nx, Nv, mx, mv = 4, 4, 2, 1 + half = Nv ÷ 2 + s = inv(sqrt(T(2))) + # Exactly orthonormal leaves, with entries 0, ±1 and ±1/√2. + Φx = T[s 0; -s 0; 0 1; 0 0] + Φv = reshape(T[s, s], half, mv) + Φc = reshape(T[s, -s], 2, 1) + + fe = BilinearStiefelLayer(Nx, Nv, mx, mv) + core = WeightTiedSymplecticAutoencoder( + 4, 2; n_blocks = 1, n_layers = 0, n_ambient_layers = 0, sympnet_upscale = 2) + nn = NeuralNetwork( + MultiModalSymplecticAutoencoder((a = fe, b = fe), core), CPU(), T; rng = rng) + nn_sa = NeuralNetwork( + MultiModalSymplecticAutoencoder((a = fe, b = fe), _untied_core()), CPU(), T; + rng = rng) + for net in (nn, nn_sa), m in (:a, :b) + + copyto!(parent(getfield(params(net).frontends, m).Φx), Φx) + copyto!(parent(getfield(params(net).frontends, m).Φv), Φv) + end + @test length(params(nn).core) == 1 + copyto!(parent(params(nn).core.L1.weight), Φc) + + # Dense encoder and decoder from the leaves and explicit parity matrices, vec(F) x-fastest: + # vec(F⁺) = P₊ vec(F), vec(F⁻) = P₋ vec(F), vec(Φxᵀ G Φv) = kron(Φv, Φx)ᵀ vec(G), and the + # weight-tied core with no SympNet layers is (q, p) ↦ (Φcᵀ q, Φcᵀ p). + P_plus = zeros(T, Nx * half, Nx * Nv) + P_minus = zeros(T, Nx * half, Nx * Nv) + for j in 1:half, i in 1:Nx + + row = (j - 1) * Nx + i + col_top = (j - 1) * Nx + i + col_bot = (Nv - j) * Nx + i + P_plus[row, col_top] = P_plus[row, col_bot] = s + P_minus[row, col_top] = s + P_minus[row, col_bot] = -s + end + K = kron(Φv, Φx) + O = zeros(T, 1, 2) + C_enc = [transpose(Φc) O; O transpose(Φc)] + Enc = C_enc * vcat(transpose(K) * P_plus, transpose(K) * P_minus) + Dec = hcat(transpose(P_plus) * K, transpose(P_minus) * K) * transpose(C_enc) + + x_a = T(3.7) .* randn(rng, T, Nx, Nv) + x_b = T(3.7) .* randn(rng, T, Nx, Nv) + z_a = Enc * vec(x_a) + z_b = Enc * vec(x_b) + @test isapprox(encode(nn, :a, x_a), z_a; rtol = 10 * eps(T)) + @test isapprox(vec(decode(nn, :b, z_a)), Dec * z_a; rtol = 10 * eps(T)) + + β_align, β_cross = T(0.7), T(1.3) + loss = CrossModalLoss(:a, :b; β_align = β_align, β_cross = β_cross) + expected = β_align * sum(abs2, z_a - z_b) + β_cross * sum(abs2, Dec * z_a - vec(x_b)) + val = @inferred loss(nn, x_a, x_b) + @test val isa T + @test isapprox(val, expected; rtol = 10 * eps(T)) + + # With β_align = 0 the loss is the cross term alone, which is not small and depends on x_b. + loss_cross = CrossModalLoss(:a, :b; β_align = 0, β_cross = 1) + cross = sum(abs2, Dec * z_a - vec(x_b)) + @test isapprox(loss_cross(nn, x_a, x_b), cross; rtol = 10 * eps(T)) + @test cross > 1e-2 + x_b2 = x_b .+ T(0.5) + cross2 = sum(abs2, Dec * z_a - vec(x_b2)) + @test isapprox(loss_cross(nn, x_a, x_b2), cross2; rtol = 10 * eps(T)) + @test abs(loss_cross(nn, x_a, x_b2) - loss_cross(nn, x_a, x_b)) > T(1e-2) * cross + + # The cycle term ‖encode(b, decode(b, z_a)) - z_a‖ is zero for the weight-tied core with a + # right-invertible frontend, and not zero for the untied core. + z_a_nn = encode(nn, :a, x_a) + cycle = norm(encode(nn, :b, decode(nn, :b, z_a_nn)) - z_a_nn) + @test cycle <= 10 * eps(T) * max(one(T), norm(z_a_nn)) + z_a_sa = encode(nn_sa, :a, x_a) + cycle_sa = norm(encode(nn_sa, :b, decode(nn_sa, :b, z_a_sa)) - z_a_sa) + @test cycle_sa > T(1e-3) * norm(z_a_sa) + + # The self-reconstruction form of CharbonnierRelativeLoss has the latent term: on the untied + # core, where that term is not zero, the value is the reconstruction term plus the latent + # term, the value of the form with an external x̂ = decode(encode(x)). + ε = T(1e-2) + loss_c = CharbonnierRelativeLoss(; ε = ε, λ_lat = 1) + x̂_a = decode(nn_sa, :a, encode(nn_sa, :a, x_a)) + rec = (sqrt(sum(abs2, x̂_a - x_a) + ε^2) - ε) / sqrt(sum(abs2, x_a) + ε^2) + lat = sum(abs2, encode(nn_sa, :a, x̂_a) - z_a_sa) / (sum(abs2, x_a) + ε^2) + @test lat > T(1e-3) * rec + val_self = @inferred loss_c(nn_sa, :a, x_a) + @test val_self isa T + @test isapprox(val_self, rec + lat; rtol = 10 * eps(T)) + @test isapprox(loss_c(nn_sa, :a, x̂_a, x_a), rec + lat; rtol = 10 * eps(T)) +end + +@testset "Docstrings state the limits of the frontends and the losses" begin + # The source text of every docstring of X, with its line breaks as spaces. `Docs.doc` has + # no method until the REPL standard library loads, so the text comes from the module's + # metadata. + function docstring(X) + binding = Base.Docs.Binding(GeometricMachineLearning, nameof(X)) + multidoc = Base.Docs.meta(GeometricMachineLearning)[binding] + replace(join((join(d.text) for d in values(multidoc.docs)), "\n"), '\n' => ' ') + end + documented = (CanonicalSortLayer, CanonicalSlotLayer, InverseCanonicalSlotLayer, + SymmetricSymplecticAutoencoder, BilinearStiefelLayer, LowRankStiefelLayer, + MultiModalSymplecticAutoencoder, CharbonnierRelativeLoss, IsometricRolloutLoss, + CrossModalLoss) + absent = ("collective canonical slots", "relative modes canonically", + "canonical quadratic energy", "β_cycle", "symplectic or Poisson frontend") + for X in documented + str = docstring(X) + for phrase in absent + @test !occursin(phrase, str) + end + @test !occursin(r"(Proposition|Theorem|Corollary|Section) \d", str) + end + for X in (CanonicalSortLayer, CanonicalSlotLayer, SymmetricSymplecticAutoencoder) + @test occursin("keep their particle order", docstring(X)) + end + @test occursin("fixed particle count", docstring(CanonicalSlotLayer)) + for X in (BilinearStiefelLayer, LowRankStiefelLayer, MultiModalSymplecticAutoencoder) + @test occursin("declare on the parity coordinates", docstring(X)) + end + @test occursin("requires a latent Hamiltonian", docstring(IsometricRolloutLoss)) + @test occursin("same physical state", docstring(CrossModalLoss)) + @test occursin("not symmetric", docstring(CrossModalLoss)) + @test !occursin("skips", docstring(CharbonnierRelativeLoss)) + @test occursin("have one element type", docstring(CharbonnierRelativeLoss)) +end diff --git a/test/reduced_system/reference_shift_and_vjp.jl b/test/reduced_system/reference_shift_and_vjp.jl index 5bc386f57..130d9c4e7 100644 --- a/test/reduced_system/reference_shift_and_vjp.jl +++ b/test/reduced_system/reference_shift_and_vjp.jl @@ -506,3 +506,70 @@ end v_red0(v_out0, zero(T), q̃, p̃, nothing) @test norm(v_out0 .- v_expected) > 1e-3 * norm(v_expected) end + +@testset "HRedSys with a reference that is not an equilibrium, against the unshifted system in shifted coordinates ($T, $mode)" for T in REAL_ELTYPES, + mode in (:jacobian, :vjp) + + rng = Random.MersenneTwister(808) + N = 6 + n = 2 + arch = WeightTiedSymplecticAutoencoder( + 2 * N, 2 * n; n_blocks = 1, n_layers = 1, n_ambient_layers = 0, sympnet_upscale = 2) + nn = NeuralNetwork(arch, CPU(), T; rng = rng) + + # H(q, p) = ½‖p‖² + ½‖q‖² + ¼ Σ qᵢ⁴, centred at the origin. + v_field!(v, t, q, p, params) = (v .= p; nothing) + f_field!(f, t, q, p, params) = (f .= .-q .- q .^ 3; nothing) + h_field(t, q, p, params) = T(0.5) * (sum(abs2, p) + sum(abs2, q)) + + T(0.25) * sum(x -> x^4, q) + + q_ref = randn(rng, T, N) + p_ref = randn(rng, T, N) + z_ref = vcat(q_ref, p_ref) + # The full vector field (∂H/∂p, -∂H/∂q) does not vanish at the reference. + @test norm(vcat(p_ref, q_ref .+ q_ref .^ 3)) >= 0.1 + q0 = q_ref .+ T(0.3) .* randn(rng, T, N) + p0 = p_ref .+ T(0.3) .* randn(rng, T, N) + + rs = HRedSys( + 2 * N, 2 * n, encoder(nn), decoder(nn), v_field!, f_field!, h_field, + (zero(T), one(T)), T(0.05), (q = StateVariable(q0), p = StateVariable(p0)); + integrator = ImplicitMidpoint(), q_ref = q_ref, p_ref = p_ref, + pullback_mode = mode) + + # The oracle: the system without a reference, with the fields translated by the reference + # and the initial state written in shifted coordinates. + v̂!(v, t, y_q, y_p, params) = v_field!(v, t, y_q .+ q_ref, y_p .+ p_ref, params) + f̂!(f, t, y_q, y_p, params) = f_field!(f, t, y_q .+ q_ref, y_p .+ p_ref, params) + ĥ(t, y_q, y_p, params) = h_field(t, y_q .+ q_ref, y_p .+ p_ref, params) + rs_oracle = HRedSys( + 2 * N, 2 * n, encoder(nn), decoder(nn), v̂!, f̂!, ĥ, + (zero(T), one(T)), T(0.05), + (q = StateVariable(q0 .- q_ref), p = StateVariable(p0 .- p_ref)); + integrator = ImplicitMidpoint(), pullback_mode = mode) + + sol = integrate_reduced_system(rs) + sol_oracle = integrate_reduced_system(rs_oracle) + @test length(sol.q) == 21 # initial state + 20 steps + for k in 0:20 + q̃_k, p̃_k = Vector(sol.q[k]), Vector(sol.p[k]) + ξ_k = vcat(q̃_k, p̃_k) + ξ̂_k = vcat(Vector(sol_oracle.q[k]), Vector(sol_oracle.p[k])) + bound = 50 * eps(T) * max(one(T), norm(ξ̂_k)) + @test norm(ξ_k - ξ̂_k) <= bound + @test norm(decode(rs, ξ_k) - (decoder(nn)(ξ̂_k) .+ z_ref)) <= bound + @test abs(rs.h_reduced(zero(T), q̃_k, p̃_k, nothing) - + rs_oracle.h_reduced(zero(T), q̃_k, p̃_k, nothing)) <= bound + end + # The trajectory moves. + @test norm(vcat(Vector(sol.q[20]), Vector(sol.p[20])) - + vcat(Vector(sol.q[0]), Vector(sol.p[0]))) > T(1e-2) + + # The check can fail: at ξ_0 the shifted reduced field differs from the unshifted one. + q̃_0, p̃_0 = Vector(sol.q[0]), Vector(sol.p[0]) + v_shifted, v_unshifted = zeros(T, n), zeros(T, n) + rs.v_reduced(v_shifted, zero(T), q̃_0, p̃_0, nothing) + build_v_reduced(v_field!, f_field!, decoder(nn); pullback_mode = mode)( + v_unshifted, zero(T), q̃_0, p̃_0, nothing) + @test norm(v_shifted - v_unshifted) > 1e-3 * norm(v_shifted) +end