From f9515f806572672d3e9df851c7a8292bb4e077fe Mon Sep 17 00:00:00 2001 From: Michael Kraus Date: Fri, 30 Oct 2026 02:30:40 +0100 Subject: [PATCH 1/8] Add MultiModalSymplecticAutoencoder, multi-modal losses, and HRedSys VJP and reference shift --- CHANGELOG.md | 2 + .../architectures/symplectic_autoencoder.md | 4 + docs/src/reduced_order_modeling/losses.md | 3 + src/GeometricMachineLearning.jl | 4 + .../multimodal_symplectic_autoencoder.jl | 458 +++++++++++++++++ src/loss/multimodal_losses.jl | 317 ++++++++++++ src/reduced_system/reduced_system.jl | 282 +++++++++-- .../multimodal_symplectic_autoencoder.jl | 463 ++++++++++++++++++ .../reduced_system/reference_shift_and_vjp.jl | 258 ++++++++++ test/runtests.jl | 2 + 10 files changed, 1751 insertions(+), 42 deletions(-) create mode 100644 src/architectures/multimodal_symplectic_autoencoder.jl create mode 100644 src/loss/multimodal_losses.jl create mode 100644 test/architectures/multimodal_symplectic_autoencoder.jl create mode 100644 test/reduced_system/reference_shift_and_vjp.jl diff --git a/CHANGELOG.md b/CHANGELOG.md index 7b17d4492..c1b111662 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -150,6 +150,8 @@ breaking release). - **`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`. +- **`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 whose canonical output dimension matches `core.full_dim` with a shared `SymplecticCompression` core; `encode(nn, modality, x)` 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``. `HRedSys`, `build_v_reduced`, and `build_f_reduced` accept optional reference states `q_ref` and `p_ref` (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 `:jacobian` uses preallocated buffers (`ξ_buf`, `v_intermediate`, `f_intermediate`, `jac_result`, and a per-element-type `IdDict{DataType, Any}` cache for `ForwardDiff.Dual` Newton states). The tests are `test/architectures/multimodal_symplectic_autoencoder.jl` and `test/reduced_system/reference_shift_and_vjp.jl`. + ### Changed - **`SymplecticAttentionQ` and `SymplecticAttentionP` accept a single sample.** A `(q, p)` `NamedTuple` of vectors, and a plain vector, are reshaped to one column, passed through the matrix method and returned as vectors, so the forward layer accepts the same inputs as its inverse. On `main` a `NamedTuple` of vectors and a plain vector raised `MethodError: no method matching _custom_transpose(::Vector{Float64})`. The test is `test/layers/symplectic_attention.jl`. diff --git a/docs/src/architectures/symplectic_autoencoder.md b/docs/src/architectures/symplectic_autoencoder.md index de1b0036a..71fbe9857 100644 --- a/docs/src/architectures/symplectic_autoencoder.md +++ b/docs/src/architectures/symplectic_autoencoder.md @@ -136,6 +136,10 @@ We also see that the intermediate dimension in the decoder is `54` for the spec ```@docs SymplecticAutoencoder WeightTiedSymplecticAutoencoder +MultiModalSymplecticAutoencoder +encode +decode +translate ``` ```@raw latex diff --git a/docs/src/reduced_order_modeling/losses.md b/docs/src/reduced_order_modeling/losses.md index 10d0d15e1..0abdaa6d1 100644 --- a/docs/src/reduced_order_modeling/losses.md +++ b/docs/src/reduced_order_modeling/losses.md @@ -46,6 +46,9 @@ where ``\mathbf{x}^{(t)}`` is the solution of the FOM at point ``t`` and ``\math TransformerLoss AutoEncoderLoss ReducedLoss +CharbonnierRelativeLoss +IsometricRolloutLoss +CrossModalCycleLoss GeometricMachineLearning.ParametricLoss GeometricMachineLearning.apply_parametric projection_error diff --git a/src/GeometricMachineLearning.jl b/src/GeometricMachineLearning.jl index c0d02ac7f..26873bd03 100644 --- a/src/GeometricMachineLearning.jl +++ b/src/GeometricMachineLearning.jl @@ -255,6 +255,8 @@ include("architectures/linear_symplectic_transformer.jl") include("architectures/symplectic_transformer.jl") include("architectures/weight_tied_symplectic_autoencoder.jl") include("architectures/symmetric_symplectic_autoencoder.jl") +include("architectures/multimodal_symplectic_autoencoder.jl") +include("loss/multimodal_losses.jl") export LinearSymplecticAttention, LinearSymplecticAttentionQ, LinearSymplecticAttentionP export LinearSymplecticTransformer @@ -273,6 +275,8 @@ export BilinearStiefelLayer, TransposeBilinearStiefelLayer export LowRankStiefelLayer, DecodeLowRankStiefelLayer export FastDLRATraceLoss export SymmetricSymplecticAutoencoder +export MultiModalSymplecticAutoencoder, encode, decode, translate +export CharbonnierRelativeLoss, IsometricRolloutLoss, CrossModalCycleLoss include("map_to_cpu.jl") diff --git a/src/architectures/multimodal_symplectic_autoencoder.jl b/src/architectures/multimodal_symplectic_autoencoder.jl new file mode 100644 index 000000000..b3ccf29a1 --- /dev/null +++ b/src/architectures/multimodal_symplectic_autoencoder.jl @@ -0,0 +1,458 @@ +function _frontend_output_dimension( + ::BilinearStiefelLayer{Nx, Nv, Mx, Mv}) where {Nx, Nv, Mx, Mv} + 2 * Mx * Mv +end + +function _frontend_output_dimension( + ::LowRankStiefelLayer{Nx, Nv, Mx, Mv}) where {Nx, Nv, Mx, Mv} + 2 * Mx * Mv +end + +function _frontend_output_dimension(c::Chain) + _frontend_output_dimension(last(c.layers)) +end + +function _frontend_output_dimension(l::AbstractExplicitLayer) + output_dimension(l) +end + +function _check_modality(frontends::NamedTuple{names}, modality::Symbol) where {names} + modality in names || throw(ArgumentError( + "Unknown modality :$modality. Available modalities: $names.")) + nothing +end + +ChainRulesCore.@non_differentiable _check_modality(::Any, ::Any) + +@doc raw""" + MultiModalSymplecticAutoencoder(frontends::NamedTuple, core::SymplecticCompression) + +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`: +- [`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``. + +Given `nn = NeuralNetwork(arch, backend, T)`: +- [`encode`](@ref)`(nn, modality::Symbol, x)` applies `frontends[modality]` followed by `encoder(core)` to map modality input `x` into ``(\tilde{q}, \tilde{p}) \in \mathbb{R}^{2n}``; +- [`decode`](@ref)`(nn, modality::Symbol, ξ)` applies `decoder(core)` followed by `inverse_layer(frontends[modality])` to reconstruct in `modality`; +- [`translate`](@ref)`(nn, from::Symbol, to::Symbol, x)` computes `decode(nn, to, encode(nn, from, x))`. +""" +struct MultiModalSymplecticAutoencoder{ + FT <: NamedTuple, + CT <: SymplecticCompression +} <: SymplecticCompression + frontends::FT + core::CT + full_dim::Int + reduced_dim::Int + + function MultiModalSymplecticAutoencoder( + frontends::NamedTuple{names}, core::SymplecticCompression) where {names} + isempty(names) && throw(ArgumentError( + "MultiModalSymplecticAutoencoder requires a non-empty frontends NamedTuple.")) + N2 = core.full_dim + n2 = core.reduced_dim + for mod in names + fe = getfield(frontends, mod) + fe_dim = _frontend_output_dimension(fe) + fe_dim == N2 || throw(ArgumentError( + "Frontend :$mod has output dimension $fe_dim, expected core.full_dim = $N2.")) + end + new{typeof(frontends), typeof(core)}(frontends, core, N2, n2) + end +end + +function AbstractNeuralNetworks.Chain(::MultiModalSymplecticAutoencoder) + throw(ArgumentError( + "A MultiModalSymplecticAutoencoder has no single Chain: build it with NeuralNetwork(arch) and use encode(nn, modality, x), decode(nn, modality, ξ), and translate(nn, from, to, x).")) +end + +struct MultiModalSymplecticLayer{ + M, + FT <: NamedTuple, + IFT <: NamedTuple, + CMT, + CET, + CDT +} <: AbstractExplicitLayer{M, M} + frontends::FT + inverse_frontends::IFT + core_model::CMT + core_encoder::CET + core_decoder::CDT + full_dim::Int + reduced_dim::Int +end + +function parameterlength(l::MultiModalSymplecticLayer) + fe_len = sum(mod -> parameterlength(getfield(l.frontends, mod)), keys(l.frontends); init = 0) + fe_len + parameterlength(l.core_model) +end + +function _init_frontend_params( + rng::AbstractRNG, initializer::AbstractNeuralNetworks.Initializer, + fe::AbstractExplicitLayer, backend::NeuralNetworkBackend, ::Type{T}) where {T} + initialparameters(rng, initializer, fe, backend, T) +end + +function _init_frontend_params( + rng::AbstractRNG, initializer::AbstractNeuralNetworks.Initializer, + fe::Chain, backend::NeuralNetworkBackend, ::Type{T}) where {T} + initialparameters(rng, initializer, fe, backend, T) +end + +function NeuralNetwork( + arch::MultiModalSymplecticAutoencoder, + backend::NeuralNetworkBackend, + ::Type{T}; + rng::AbstractRNG = Random.default_rng(), + initializer::AbstractNeuralNetworks.Initializer = GlorotUniform(), + kwargs...) where {T <: Number} + names = keys(arch.frontends) + fe_ps_vals = map(names) do mod + fe = getfield(arch.frontends, mod) + _init_frontend_params(rng, initializer, fe, backend, T) + end + fe_ps_nt = NamedTuple{names}(fe_ps_vals) + inv_fe_vals = map(names) do mod + inverse_layer(getfield(arch.frontends, mod)) + end + inv_frontends = NamedTuple{names}(inv_fe_vals) + + core_nn = NeuralNetwork(arch.core, backend, T; rng = rng, kwargs...) + core_enc_nn = encoder(core_nn) + core_dec_nn = decoder(core_nn) + + M = arch.full_dim + FT = typeof(arch.frontends) + IFT = typeof(inv_frontends) + CMT = typeof(core_nn.model) + CET = typeof(core_enc_nn.model) + CDT = typeof(core_dec_nn.model) + layer = MultiModalSymplecticLayer{M, FT, IFT, CMT, CET, CDT}( + arch.frontends, + inv_frontends, + core_nn.model, + core_enc_nn.model, + core_dec_nn.model, + arch.full_dim, + arch.reduced_dim + ) + ps = NetworkParameters((frontends = fe_ps_nt, core = params(core_nn.params))) + NeuralNetwork(arch, layer, ps, backend) +end + +function encoder(nn::NeuralNetwork{<:MultiModalSymplecticAutoencoder}) + core_nn = NeuralNetwork( + nn.architecture.core, nn.model.core_model, + NetworkParameters(params(nn).core), networkbackend(nn)) + encoder(core_nn) +end + +function decoder(nn::NeuralNetwork{<:MultiModalSymplecticAutoencoder}) + core_nn = NeuralNetwork( + nn.architecture.core, nn.model.core_model, + NetworkParameters(params(nn).core), networkbackend(nn)) + decoder(core_nn) +end + +inverse_layer(c::Chain) = Chain(reverse(map(inverse_layer, c.layers))...) + +_leaf_eltype(a::AbstractArray{T}) where {T <: Number} = T +_leaf_eltype(nt::NamedTuple) = _leaf_eltype(first(values(nt))) + +_fe_output_to_core_vec(::Any, z_fe::AbstractVecOrMat) = z_fe + +function _fe_output_to_core_vec( + ::Union{BilinearStiefelLayer, LowRankStiefelLayer}, z_fe::NamedTuple) + vcat(z_fe.q, z_fe.p) +end + +function _fe_output_to_core_vec(::Any, z_fe::NamedTuple) + if ndims(z_fe.q) <= 2 + vcat(vec(z_fe.q), vec(z_fe.p)) + else + N = size(z_fe.q, 1) * size(z_fe.q, 2) + B = size(z_fe.q, 3) + vcat(reshape(z_fe.q, N, B), reshape(z_fe.p, N, B)) + end +end + +function _narrow_core_vec(z_vec, x::NamedTuple, ::Type{T}) where {T} + if ndims(first(values(x))) <= 2 + return z_vec::Vector{T} + else + return z_vec::Matrix{T} + end +end + +_narrow_core_vec(z_vec, ::AbstractVector, ::Type{T}) where {T} = z_vec::Vector{T} +function _narrow_core_vec(z_vec, ::AbstractArray{<:Any, 3}, ::Type{T}) where {T} + z_vec::Matrix{T} +end +_narrow_core_vec(z_vec, ::Any, ::Type{T}) where {T} = z_vec + +function _core_vec_to_fe_inv_input( + ::Union{BilinearStiefelLayer, LowRankStiefelLayer}, + z_mid::AbstractVector, full_dim::Int) + N = full_dim ÷ 2 + (q = z_mid[1:N], p = z_mid[(N + 1):full_dim]) +end + +function _core_vec_to_fe_inv_input( + ::Union{BilinearStiefelLayer, LowRankStiefelLayer}, + z_mid::AbstractMatrix, full_dim::Int) + N = full_dim ÷ 2 + (q = z_mid[1:N, :], p = z_mid[(N + 1):full_dim, :]) +end + +function _core_vec_to_fe_inv_input( + ::Any, z_mid::AbstractVecOrMat, ::Int) + z_mid +end + +function _apply_inv_frontend( + ::Chain, inv_fe::Chain, z_in, ps_fe::NamedTuple) + inv_fe(z_in, _reverse_params(ps_fe)) +end + +function _apply_inv_frontend( + ::AbstractExplicitLayer, inv_fe::AbstractExplicitLayer, z_in, ps_fe::NamedTuple) + inv_fe(z_in, ps_fe) +end + +function _core_encode_vec( + ::Union{WeightTiedAutoencoderLayer, SymmetricAutoencoderLayer}, + core_encoder::Chain, ps_core::NamedTuple, z_vec::AbstractVecOrMat) + core_encoder(z_vec, _reverse_params(ps_core)) +end + +function _core_decode_vec( + ::Union{WeightTiedAutoencoderLayer, SymmetricAutoencoderLayer}, + core_decoder::Chain, ps_core::NamedTuple, ξ_vec::AbstractVecOrMat) + core_decoder(ξ_vec, ps_core) +end + +function _core_encode_vec( + ::Any, core_encoder::Chain, ps_core::NamedTuple, z_vec::AbstractVecOrMat) + names = keys(ps_core) + M = length(core_encoder.layers) + enc_names = names[1:M] + ps_enc = NamedTuple{enc_names}(map(k -> getfield(ps_core, k), enc_names)) + core_encoder(z_vec, ps_enc) +end + +function _core_decode_vec( + ::Any, core_decoder::Chain, ps_core::NamedTuple, ξ_vec::AbstractVecOrMat) + names = keys(ps_core) + K = length(core_decoder.layers) + M = length(names) - K + dec_names = names[1:K] + src_names = names[(M + 1):end] + ps_dec = NamedTuple{dec_names}(map(k -> getfield(ps_core, k), src_names)) + core_decoder(ξ_vec, ps_dec) +end + +_latent_to_vec(ξ::AbstractVecOrMat) = ξ +_latent_to_vec(ξ::NamedTuple) = vcat(ξ.q, ξ.p) + +function _encode_active( + fe, ps_fe::NamedTuple, core_model, core_encoder::Chain, ps_core::NamedTuple, x) + z_fe = fe(x, ps_fe) + T = promote_type(_leaf_eltype(x), _leaf_eltype(ps_core)) + z_vec = _narrow_core_vec(_fe_output_to_core_vec(fe, z_fe), x, T) + _core_encode_vec(core_model, core_encoder, ps_core, z_vec) +end + +function _decode_active( + fe, inv_fe, ps_fe::NamedTuple, core_model, core_decoder::Chain, + ps_core::NamedTuple, full_dim::Int, ξ) + ξ_vec = _latent_to_vec(ξ) + z_mid = _core_decode_vec(core_model, core_decoder, ps_core, ξ_vec) + z_in = _core_vec_to_fe_inv_input(fe, z_mid, full_dim) + _apply_inv_frontend(fe, inv_fe, z_in, ps_fe) +end + +_zero_param_tree(a::AbstractArray) = zero(parent(a)) +_zero_param_tree(nt::NamedTuple) = map(_zero_param_tree, nt) + +_resolve_param_grad(g::ChainRulesCore.AbstractZero, p_ref) = _zero_param_tree(p_ref) +_resolve_param_grad(::Nothing, p_ref) = _zero_param_tree(p_ref) +_resolve_param_grad(g::AbstractArray, ::AbstractArray) = g +function _resolve_param_grad( + g::Union{ChainRulesCore.Tangent, NamedTuple}, + p_ref::NamedTuple{names}) where {names} + NamedTuple{names}(map( + k -> _resolve_param_grad(getproperty(g, k), getfield(p_ref, k)), names)) +end + +@doc raw""" + encode(nn::NeuralNetwork{<:MultiModalSymplecticAutoencoder}, modality::Symbol, x) + encode(model::MultiModalSymplecticLayer, ps, modality::Symbol, x) + +Apply `frontends[modality]` followed by `encoder(core)` to map an input `x` in modality `modality` into the shared canonical latent space ``(\tilde{q}, \tilde{p}) \in \mathbb{R}^{2n}``. + +Throws an `ArgumentError` naming `modality` and `keys(frontends)` if `modality` is not present in `frontends`. +""" +Base.@constprop :aggressive function encode( + model::MultiModalSymplecticLayer, ps::NamedTuple, modality::Symbol, x) + _check_modality(model.frontends, modality) + fe = getfield(model.frontends, modality) + ps_fe = getfield(ps.frontends, modality) + _encode_active(fe, ps_fe, model.core_model, model.core_encoder, ps.core, x) +end + +@noinline function _rrule_encode_active( + config::ChainRulesCore.RuleConfig{>:ChainRulesCore.HasReverseMode}, + fe, ps_fe::NamedTuple, core_model, core_encoder::Chain, ps_core::NamedTuple, x) + ChainRulesCore.rrule_via_ad( + config, _encode_active, fe, ps_fe, core_model, core_encoder, ps_core, x) +end + +@noinline function _rrule_decode_active( + config::ChainRulesCore.RuleConfig{>:ChainRulesCore.HasReverseMode}, + fe, inv_fe, ps_fe::NamedTuple, core_model, core_decoder::Chain, + ps_core::NamedTuple, full_dim::Int, ξ) + ChainRulesCore.rrule_via_ad( + config, _decode_active, fe, inv_fe, ps_fe, core_model, core_decoder, + ps_core, full_dim, ξ) +end + +function ChainRulesCore.rrule( + config::ChainRulesCore.RuleConfig{>:ChainRulesCore.HasReverseMode}, + ::typeof(encode), + model::MultiModalSymplecticLayer, + ps::NamedTuple, + modality::Symbol, + x) + _check_modality(model.frontends, modality) + fe = getfield(model.frontends, modality) + ps_fe = getfield(ps.frontends, modality) + y, pb = _rrule_encode_active( + config, fe, ps_fe, model.core_model, model.core_encoder, ps.core, x) + function encode_pullback(Δy) + _, _, d_ps_fe, _, _, d_ps_core, d_x = pb(ChainRulesCore.unthunk(Δy)) + fe_keys = keys(ps.frontends) + d_frontends = NamedTuple{fe_keys}(map( + k -> k === modality ? + _resolve_param_grad( + ChainRulesCore.unthunk(d_ps_fe), getfield(ps.frontends, k)) : + _zero_param_tree(getfield(ps.frontends, k)), + fe_keys)) + d_core = _resolve_param_grad(ChainRulesCore.unthunk(d_ps_core), ps.core) + d_ps = (frontends = d_frontends, core = d_core) + return ( + ChainRulesCore.NoTangent(), + ChainRulesCore.NoTangent(), + d_ps, + ChainRulesCore.NoTangent(), + ChainRulesCore.unthunk(d_x) + ) + end + return y, encode_pullback +end + +Base.@constprop :aggressive function encode( + model::MultiModalSymplecticLayer, ps::NetworkParameters, modality::Symbol, x) + encode(model, params(ps), modality, x) +end + +Base.@constprop :aggressive function encode( + nn::NeuralNetwork{<:MultiModalSymplecticAutoencoder}, modality::Symbol, x) + encode(nn.model, params(nn), modality, x) +end + +@doc raw""" + decode(nn::NeuralNetwork{<:MultiModalSymplecticAutoencoder}, modality::Symbol, ξ) + decode(model::MultiModalSymplecticLayer, ps, modality::Symbol, ξ) + +Apply `decoder(core)` followed by `inverse_layer(frontends[modality])` to reconstruct a latent state ``\xi`` in modality `modality`. + +Throws an `ArgumentError` naming `modality` and `keys(frontends)` if `modality` is not present in `frontends`. +""" +Base.@constprop :aggressive function decode( + model::MultiModalSymplecticLayer, ps::NamedTuple, modality::Symbol, ξ) + _check_modality(model.frontends, modality) + fe = getfield(model.frontends, modality) + inv_fe = getfield(model.inverse_frontends, modality) + ps_fe = getfield(ps.frontends, modality) + _decode_active( + fe, inv_fe, ps_fe, model.core_model, model.core_decoder, + ps.core, model.full_dim, ξ) +end + +function ChainRulesCore.rrule( + config::ChainRulesCore.RuleConfig{>:ChainRulesCore.HasReverseMode}, + ::typeof(decode), + model::MultiModalSymplecticLayer, + ps::NamedTuple, + modality::Symbol, + ξ) + _check_modality(model.frontends, modality) + fe = getfield(model.frontends, modality) + inv_fe = getfield(model.inverse_frontends, modality) + ps_fe = getfield(ps.frontends, modality) + y, pb = _rrule_decode_active( + config, fe, inv_fe, ps_fe, model.core_model, model.core_decoder, + ps.core, model.full_dim, ξ) + function decode_pullback(Δy) + _, _, _, d_ps_fe, _, _, d_ps_core, _, d_ξ = pb(ChainRulesCore.unthunk(Δy)) + fe_keys = keys(ps.frontends) + d_frontends = NamedTuple{fe_keys}(map( + k -> k === modality ? + _resolve_param_grad( + ChainRulesCore.unthunk(d_ps_fe), getfield(ps.frontends, k)) : + _zero_param_tree(getfield(ps.frontends, k)), + fe_keys)) + d_core = _resolve_param_grad(ChainRulesCore.unthunk(d_ps_core), ps.core) + d_ps = (frontends = d_frontends, core = d_core) + return ( + ChainRulesCore.NoTangent(), + ChainRulesCore.NoTangent(), + d_ps, + ChainRulesCore.NoTangent(), + ChainRulesCore.unthunk(d_ξ) + ) + end + return y, decode_pullback +end + +Base.@constprop :aggressive function decode( + model::MultiModalSymplecticLayer, ps::NetworkParameters, modality::Symbol, ξ) + decode(model, params(ps), modality, ξ) +end + +Base.@constprop :aggressive function decode( + nn::NeuralNetwork{<:MultiModalSymplecticAutoencoder}, modality::Symbol, ξ) + decode(nn.model, params(nn), modality, ξ) +end + +@doc raw""" + translate(nn::NeuralNetwork{<:MultiModalSymplecticAutoencoder}, from::Symbol, to::Symbol, x) + translate(model::MultiModalSymplecticLayer, ps, from::Symbol, to::Symbol, x) + +Map an input `x` in modality `from` into modality `to` through the shared canonical latent space via `decode(nn, to, encode(nn, from, x))`. +""" +Base.@constprop :aggressive function translate( + model::MultiModalSymplecticLayer, ps::NamedTuple, from::Symbol, to::Symbol, x) + decode(model, ps, to, encode(model, ps, from, x)) +end + +Base.@constprop :aggressive function translate( + model::MultiModalSymplecticLayer, ps::NetworkParameters, + from::Symbol, to::Symbol, x) + translate(model, params(ps), from, to, x) +end + +Base.@constprop :aggressive function translate( + nn::NeuralNetwork{<:MultiModalSymplecticAutoencoder}, + from::Symbol, to::Symbol, x) + translate(nn.model, params(nn), from, to, x) +end + +function (model::MultiModalSymplecticLayer)(x, ps::NamedTuple) + mod = first(keys(model.frontends)) + decode(model, ps, mod, encode(model, ps, mod, x)) +end diff --git a/src/loss/multimodal_losses.jl b/src/loss/multimodal_losses.jl new file mode 100644 index 000000000..06ed63a55 --- /dev/null +++ b/src/loss/multimodal_losses.jl @@ -0,0 +1,317 @@ +_state_eltype(::AbstractArray{T}) where {T} = T +_state_eltype(nt::NamedTuple{(:q, :p)}) = promote_type(eltype(nt.q), eltype(nt.p)) +function _state_eltype(nt::NamedTuple{(:X, :S, :V)}) + promote_type(eltype(nt.X), eltype(nt.S), eltype(nt.V)) +end + +_frob_sq_norm(x::AbstractArray) = sum(abs2, x) +_frob_sq_norm(x::NamedTuple{(:q, :p)}) = sum(abs2, x.q) + sum(abs2, x.p) +function _frob_sq_norm(x::NamedTuple{(:X, :S, :V)}) + G_X = transpose(x.X) * x.X + G_V = transpose(x.V) * x.V + S_GV = x.S * G_V + T = _state_eltype(x) + max(zero(T), tr(G_X * S_GV * transpose(x.S))) +end + +_frob_sq_diff(x̂::AbstractArray, x::AbstractArray) = sum(abs2, x̂ .- x) +function _frob_sq_diff(x̂::NamedTuple{(:q, :p)}, x::NamedTuple{(:q, :p)}) + sum(abs2, x̂.q .- x.q) + sum(abs2, x̂.p .- x.p) +end +function _frob_sq_diff(x̂::AbstractVector, x::NamedTuple{(:q, :p)}) + sum(abs2, x̂ .- vcat(vec(x.q), vec(x.p))) +end +function _frob_sq_diff(x̂::NamedTuple{(:q, :p)}, x::AbstractVector) + _frob_sq_diff(x, x̂) +end +function _frob_sq_diff(x̂::AbstractMatrix, x::NamedTuple{(:q, :p)}) + if ndims(x.q) <= 2 + sum(abs2, vec(x̂) .- vcat(vec(x.q), vec(x.p))) + else + N = size(x.q, 1) * size(x.q, 2) + B = size(x.q, 3) + sum(abs2, x̂ .- vcat(reshape(x.q, N, B), reshape(x.p, N, B))) + end +end +function _frob_sq_diff(x̂::NamedTuple{(:q, :p)}, x::AbstractMatrix) + _frob_sq_diff(x, x̂) +end +function _frob_sq_diff(x̂::NamedTuple{(:X, :S, :V)}, x::NamedTuple{(:X, :S, :V)}) + FastDLRATraceLoss()(x̂, x) +end + +@doc raw""" + CharbonnierRelativeLoss(; ε = 1e-3, λ_lat = 1e-2) + +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}, +``` + +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. + +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`). +""" +struct CharbonnierRelativeLoss{T <: Real} <: NetworkLoss + ε::T + λ_lat::T + + function CharbonnierRelativeLoss(; ε::Real = 1e-3, λ_lat::Real = 1e-2) + (ε > 0 && λ_lat >= 0) || throw(ArgumentError( + "CharbonnierRelativeLoss requires ε > 0 and λ_lat >= 0, got ε = $ε, λ_lat = $λ_lat.")) + T = promote_type(typeof(float(ε)), typeof(float(λ_lat))) + new{T}(T(ε), T(λ_lat)) + end +end + +function _charbonnier_eval( + loss::CharbonnierRelativeLoss, x̂, x, z_hat::AbstractArray, z::AbstractArray) + T = promote_type(_state_eltype(x̂), _state_eltype(x), eltype(z_hat), eltype(z)) + eps_T = T(loss.ε) + lam_T = T(loss.λ_lat) + eps_sq = eps_T * eps_T + sq_diff = _frob_sq_diff(x̂, x) + denom_sq = _frob_sq_norm(x) + eps_sq + num_rec = sqrt(sq_diff + eps_sq) - eps_T + term_rec = num_rec / sqrt(denom_sq) + term_lat = lam_T * _frob_sq_diff(z_hat, z) / denom_sq + (term_rec + term_lat)::T +end + +Base.@constprop :aggressive function (loss::CharbonnierRelativeLoss)( + model::MultiModalSymplecticLayer, ps::NamedTuple, modality::Symbol, x̂, x) + 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 +end + +Base.@constprop :aggressive function (loss::CharbonnierRelativeLoss)( + model::MultiModalSymplecticLayer, ps::NetworkParameters, modality::Symbol, x̂, x) + loss(model, params(ps), modality, x̂, x) +end + +Base.@constprop :aggressive function (loss::CharbonnierRelativeLoss)( + nn::NeuralNetwork{<:MultiModalSymplecticAutoencoder}, modality::Symbol, x̂, x) + loss(nn.model, params(nn), modality, x̂, x) +end + +Base.@constprop :aggressive function (loss::CharbonnierRelativeLoss)( + model::MultiModalSymplecticLayer, ps::NamedTuple, modality::Symbol, x) + T = promote_type(_state_eltype(x), _leaf_eltype(ps.core)) + 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 +end + +Base.@constprop :aggressive function (loss::CharbonnierRelativeLoss)( + model::MultiModalSymplecticLayer, ps::NetworkParameters, modality::Symbol, x) + loss(model, params(ps), modality, x) +end + +Base.@constprop :aggressive function (loss::CharbonnierRelativeLoss)( + nn::NeuralNetwork{<:MultiModalSymplecticAutoencoder}, modality::Symbol, x) + loss(nn.model, params(nn), modality, x) +end + +Base.@constprop :aggressive function (loss::CharbonnierRelativeLoss)( + model::MultiModalSymplecticLayer, ps::NamedTuple, x) + mod = first(keys(model.frontends)) + loss(model, ps, mod, x) +end + +Base.@constprop :aggressive function (loss::CharbonnierRelativeLoss)( + model::MultiModalSymplecticLayer, ps::NetworkParameters, x) + loss(model, params(ps), x) +end + +Base.@constprop :aggressive function (loss::CharbonnierRelativeLoss)( + nn::NeuralNetwork{<:MultiModalSymplecticAutoencoder}, x) + loss(nn.model, params(nn), x) +end + +@doc raw""" + IsometricRolloutLoss(integrator; α_lat = 1.0, α_rec = 1.0, α_H = 0.1, H = nothing) + +Multi-step latent rollout, decoded trajectory reconstruction, and latent Hamiltonian conservation loss: + +```math +\mathcal{L}_{\mathrm{roll}} = +\frac{1}{K - 1}\sum_{k=2}^{K} +\left( +\alpha_{\mathrm{lat}}\,\|\hat{z}_k - z_k\|_F^2 ++ \alpha_{\mathrm{rec}}\,\|\mathrm{decode}(nn, m, \hat{z}_k) - x_k\|_F^2 ++ \alpha_H\,\bigl|H(\hat{z}_k) - H(\hat{z}_1)\bigr|^2 +\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`. +""" +struct IsometricRolloutLoss{IT, HT, T <: Real} <: NetworkLoss + integrator::IT + H::HT + α_lat::T + α_rec::T + α_H::T + + function IsometricRolloutLoss( + integrator; + α_lat::Real = 1.0, + α_rec::Real = 1.0, + α_H::Real = 0.1, + 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.")) + 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) + +function _check_traj_len(K::Int) + K >= 2 || throw(ArgumentError( + "IsometricRolloutLoss requires at least 2 trajectory steps, got $K.")) + nothing +end + +ChainRulesCore.@non_differentiable _check_traj_len(::Any) + +Base.@constprop :aggressive function (loss::IsometricRolloutLoss)( + model::MultiModalSymplecticLayer, ps::NamedTuple, + modality::Symbol, xs::AbstractVector) + K = length(xs) + _check_traj_len(K) + T = promote_type(_state_eltype(xs[1]), _leaf_eltype(ps.core)) + z1 = encode(model, ps, modality, xs[1]) + H1 = _eval_latent_H(loss.H, z1) + z_curr = z1 + err_lat = zero(T) + err_rec = zero(T) + err_H = zero(T) + for k in 2:K + z_k = encode(model, ps, modality, xs[k]) + z_curr = loss.integrator(z_curr) + 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) + err_H = err_H + abs2(Hk - H1) + end + inv_steps = inv(T(K - 1)) + res = T(loss.α_lat) * err_lat * inv_steps + + T(loss.α_rec) * err_rec * inv_steps + + T(loss.α_H) * err_H * inv_steps + res::T +end + +Base.@constprop :aggressive function (loss::IsometricRolloutLoss)( + model::MultiModalSymplecticLayer, ps::NetworkParameters, + modality::Symbol, xs::AbstractVector) + loss(model, params(ps), modality, xs) +end + +Base.@constprop :aggressive function (loss::IsometricRolloutLoss)( + nn::NeuralNetwork{<:MultiModalSymplecticAutoencoder}, + modality::Symbol, xs::AbstractVector) + loss(nn.model, params(nn), modality, xs) +end + +Base.@constprop :aggressive function (loss::IsometricRolloutLoss)( + model::MultiModalSymplecticLayer, ps::NamedTuple, xs::AbstractVector) + mod = first(keys(model.frontends)) + loss(model, ps, mod, xs) +end + +Base.@constprop :aggressive function (loss::IsometricRolloutLoss)( + model::MultiModalSymplecticLayer, ps::NetworkParameters, xs::AbstractVector) + loss(model, params(ps), xs) +end + +Base.@constprop :aggressive function (loss::IsometricRolloutLoss)( + nn::NeuralNetwork{<:MultiModalSymplecticAutoencoder}, xs::AbstractVector) + loss(nn.model, params(nn), xs) +end + +@doc raw""" + CrossModalCycleLoss(mod_a::Symbol, mod_b::Symbol; β_align = 1.0, β_cycle = 0.1) + +Cross-modal paired latent alignment and cycle consistency 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, +``` + +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`. +""" +struct CrossModalCycleLoss{T <: Real} <: NetworkLoss + mod_a::Symbol + mod_b::Symbol + β_align::T + β_cycle::T + + function CrossModalCycleLoss( + 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)) + end +end + +Base.@constprop :aggressive function _cross_modal_eval( + loss::CrossModalCycleLoss, model::MultiModalSymplecticLayer, + ps::NamedTuple, x_a, x_b) + T = promote_type(_state_eltype(x_a), _state_eltype(x_b), _leaf_eltype(ps.core)) + 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) + res = T(loss.β_align) * _frob_sq_diff(z_a, z_b) + + T(loss.β_cycle) * _frob_sq_diff(z_a_cyc, z_a) + res::T +end + +Base.@constprop :aggressive function (loss::CrossModalCycleLoss)( + 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)( + 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)( + 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)( + 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)( + 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)( + nn::NeuralNetwork{<:MultiModalSymplecticAutoencoder}, pair::Tuple{Any, Any}) + loss(nn.model, params(nn), pair) +end diff --git a/src/reduced_system/reduced_system.jl b/src/reduced_system/reduced_system.jl index 935fde95e..cbc9afa8e 100644 --- a/src/reduced_system/reduced_system.jl +++ b/src/reduced_system/reduced_system.jl @@ -1,3 +1,50 @@ +struct ReducedVFCache{T} + ξ_buf::Vector{T} + v_intermediate::Vector{T} + f_intermediate::Vector{T} + z_bar_buf::Vector{T} + jac_result::DiffResults.MutableDiffResult{1, Vector{T}, Tuple{Matrix{T}}} + dual_cache::IdDict{DataType, Any} +end + +function ReducedVFCache(::Type{T}, N2::Int, n2::Int) where {T} + ξ_buf = zeros(T, 2 * n2) + v_int = zeros(T, N2) + f_int = zeros(T, N2) + z_bar = zeros(T, 2 * N2) + jac_res = DiffResults.JacobianResult(zeros(T, 2 * N2), ξ_buf) + ReducedVFCache{T}( + ξ_buf, v_int, f_int, z_bar, jac_res, IdDict{DataType, Any}()) +end + +function _get_vf_cache( + cache::ReducedVFCache{T0}, ::Type{T}, N2::Int, n2::Int) where {T0, T} + if T === T0 + return cache + end + get!(cache.dual_cache, T) do + ReducedVFCache(T, N2, n2) + end::ReducedVFCache{T} +end + +function _check_pullback_mode(pullback_mode::Symbol) + (pullback_mode === :jacobian || pullback_mode === :vjp) || + throw(ArgumentError( + "pullback_mode must be :jacobian or :vjp, got :$pullback_mode.")) + nothing +end + +function _add_ref!( + z::AbstractVector, q_ref::AbstractVector, p_ref::AbstractVector, N2::Int) + @inbounds for i in 1:N2 + z[i] += q_ref[i] + z[N2 + i] += p_ref[i] + end + z +end + +_add_ref!(z::AbstractVector, ::Nothing, ::Nothing, ::Int) = z + @doc raw""" HRedSys @@ -10,7 +57,7 @@ It can be called using two different constructors. The standard constructor is: ```julia -HRedSys(N, n, encoder, decoder, v_full, f_full, h_full, timespan, timestep, ics; integrator, parameters, v_reduced, f_reduced, h_reduced) +HRedSys(N, n, encoder, decoder, v_full, f_full, h_full, timespan, timestep, ics; integrator, parameters, q_ref, p_ref, pullback_mode, v_reduced, f_reduced, h_reduced) ``` where - `encoder`: a function ``\mathbb{R}^{2N}\mapsto{}\mathbb{R}^{2n}`` @@ -21,6 +68,8 @@ where - `f_reduced`: a (differentiable) mapping defined the same way as in GeometricIntegrators. - `integrator`: is used for integrating the reduced system. It is a required keyword argument, here and in the two `HRedSys(odeensemble, encoder, decoder)` constructors. - `parameters`: the parameters (a NamedTuple or `NullParameters()`) that parametrize the vector fields (the same for full_vector_field and reduced_vector_field). It defaults to `NullParameters()`; the two constructors that take a `HODEProblem` or a `HODEEnsemble` take it from the problem. +- `q_ref`, `p_ref`: optional reference configuration and momentum vectors in ``\mathbb{R}^N`` (defaulting to `zero(ics.q)` and `zero(ics.p)`). The decoder lifts ``\xi = (\tilde{q}, \tilde{p})`` as ``(q, p) = (q_{\mathrm{ref}}, p_{\mathrm{ref}}) + \Psi^{\mathrm{dec}}(\tilde{q}, \tilde{p})`` and the encoder encodes ``(q - q_{\mathrm{ref}}, p - p_{\mathrm{ref}})``. +- `pullback_mode`: `:jacobian` (default) or `:vjp`, forwarded to [`build_v_reduced`](@ref) and [`build_f_reduced`](@ref). - `timespan`: a tuple `(t₀, tₗ)` that specifies start and end point of the time interval over which integration is performed. - `timestep`: the time step - `ics`: the initial condition for the big system. @@ -28,7 +77,7 @@ where The other, and more convenient, constructor is: ```julia -HRedSys(odeensemble, encoder, decoder; integrator) +HRedSys(odeensemble, encoder, decoder; integrator, q_ref, p_ref, pullback_mode) ``` where `odeensemble` can be a `HODEEnsemble` or a `HODEProblem`. `encoder` and `decoder` have to be neural networks. The keyword argument `integrator` is required here as in the standard constructor: no constructor picks one for the caller. @@ -47,7 +96,9 @@ struct HRedSys{ InT, # type of integrator PT, # type of parameters (of the vector field/hamiltonian) IT <: NamedTuple{(:q, :p), Tuple{ST, ST}} where ST <: StateVariable{T}, - TT <: Tuple # type of timespan + TT <: Tuple, # type of timespan + RT <: AbstractVector{T}, + JR } N::Int n::Int @@ -64,24 +115,59 @@ struct HRedSys{ timespan::TT timestep::T ics::IT + q_ref::RT + p_ref::RT + pullback_mode::Symbol + ξ_buf::Vector{T} + v_intermediate::Vector{T} + f_intermediate::Vector{T} + jac_result::JR + dual_cache::IdDict{DataType, Any} end function HRedSys(N::Integer, n::Integer, encoder::NeuralNetwork{<:SymplecticEncoder}, decoder::NeuralNetwork{<:SymplecticDecoder}, v_full, f_full, - h_full, timespan::Tuple, timestep::T, ics; parameters = NullParameters(), - v_reduced = build_v_reduced(v_full, f_full, decoder), - f_reduced = build_f_reduced(v_full, f_full, decoder), - h_reduced = build_h_reduced(h_full, decoder), + h_full, timespan::Tuple, timestep::T, ics; + parameters = NullParameters(), + q_ref::AbstractVector = zeros(eltype(ics.q), length(ics.q)), + p_ref::AbstractVector = zeros(eltype(ics.p), length(ics.p)), + pullback_mode::Symbol = :jacobian, + cache::ReducedVFCache = ReducedVFCache( + eltype(ics.q), decoder.architecture.full_dim ÷ 2, + decoder.architecture.reduced_dim ÷ 2), + v_reduced = build_v_reduced( + v_full, f_full, decoder; + pullback_mode = pullback_mode, q_ref = q_ref, p_ref = p_ref, cache = cache), + f_reduced = build_f_reduced( + v_full, f_full, decoder; + pullback_mode = pullback_mode, q_ref = q_ref, p_ref = p_ref, cache = cache), + h_reduced = build_h_reduced(h_full, decoder; q_ref = q_ref, p_ref = p_ref), integrator) where {T <: Real} + _check_pullback_mode(pullback_mode) + N_ics = length(ics.q) + (length(q_ref) == N_ics && length(p_ref) == N_ics) || + throw(DimensionMismatch( + "q_ref and p_ref must have length N = $N_ics, got $(length(q_ref)) and $(length(p_ref)).")) + T_state = eltype(ics.q) + q_ref_vec = Vector{T_state}(q_ref) + p_ref_vec = Vector{T_state}(p_ref) HRedSys{typeof(timestep), typeof(encoder), typeof(decoder), typeof(v_full), typeof(f_full), typeof(h_full), typeof(v_reduced), typeof(f_reduced), - typeof(h_reduced), typeof(integrator), typeof(parameters), typeof(ics), typeof(timespan)}( + typeof(h_reduced), typeof(integrator), typeof(parameters), typeof(ics), + typeof(timespan), typeof(q_ref_vec), typeof(cache.jac_result)}( N, n, encoder, decoder, v_full, f_full, h_full, v_reduced, f_reduced, - h_reduced, integrator, parameters, timespan, timestep, ics) + h_reduced, integrator, parameters, timespan, timestep, ics, + q_ref_vec, p_ref_vec, pullback_mode, + cache.ξ_buf, cache.v_intermediate, cache.f_intermediate, + cache.jac_result, cache.dual_cache) end function HRedSys(odeproblem::HODEProblem, encoder::NeuralNetwork{<:SymplecticEncoder}, - decoder::NeuralNetwork{<:SymplecticDecoder}; integrator) + decoder::NeuralNetwork{<:SymplecticDecoder}; + integrator, + q_ref::AbstractVector = zeros(eltype(odeproblem.ics.q), length(odeproblem.ics.q)), + p_ref::AbstractVector = zeros(eltype(odeproblem.ics.p), length(odeproblem.ics.p)), + pullback_mode::Symbol = :jacobian) N = encoder.architecture.full_dim n = encoder.architecture.reduced_dim v_eq = odeproblem.equation.v @@ -89,11 +175,22 @@ function HRedSys(odeproblem::HODEProblem, encoder::NeuralNetwork{<:SymplecticEnc h_eq = odeproblem.equation.hamiltonian HRedSys( N, n, encoder, decoder, v_eq, f_eq, h_eq, odeproblem.timespan, odeproblem.timestep, - odeproblem.ics; parameters = odeproblem.parameters, integrator = integrator) + odeproblem.ics; + parameters = odeproblem.parameters, + q_ref = q_ref, + p_ref = p_ref, + pullback_mode = pullback_mode, + integrator = integrator) end function HRedSys(odeensemble::HODEEnsemble, encoder::NeuralNetwork{<:SymplecticEncoder}, - decoder::NeuralNetwork{<:SymplecticDecoder}; integrator) + decoder::NeuralNetwork{<:SymplecticDecoder}; + integrator, + q_ref::AbstractVector = zeros( + eltype(odeensemble.ics[1].q), length(odeensemble.ics[1].q)), + p_ref::AbstractVector = zeros( + eltype(odeensemble.ics[1].p), length(odeensemble.ics[1].p)), + pullback_mode::Symbol = :jacobian) N = encoder.architecture.full_dim n = encoder.architecture.reduced_dim v_eq = odeensemble.equation.v @@ -101,30 +198,74 @@ function HRedSys(odeensemble::HODEEnsemble, encoder::NeuralNetwork{<:SymplecticE h_eq = odeensemble.equation.hamiltonian HRedSys(N, n, encoder, decoder, v_eq, f_eq, h_eq, odeensemble.timespan, odeensemble.timestep, odeensemble.ics[1]; - parameters = odeensemble.parameters[1], integrator = integrator) + parameters = odeensemble.parameters[1], + q_ref = q_ref, + p_ref = p_ref, + pullback_mode = pullback_mode, + integrator = integrator) end # The full vector fields at the decoded point `Ψ(ξ)`, `ξ = (q̃, p̃)`, and the Jacobian `∇Ψ` at `ξ`. # `ForwardDiff.jacobian!` gives both `Ψ(ξ)` and `∇Ψ` from the passes of the decoder in dual numbers: # one pass where `length(ξ)` is at most ForwardDiff's chunk threshold (12 by default), and one # pass per chunk above it. -function evaluate_vf_and_compute_∇Ψ(t, q̃::AbstractVector{T}, p̃::AbstractVector{T}, - parameters, decoder, v_full, f_full) where {T} +function evaluate_vf_and_compute_∇Ψ(t, q̃::AbstractVector{Tq}, p̃::AbstractVector{Tp}, + parameters, decoder, v_full, f_full, + cache::ReducedVFCache = ReducedVFCache( + promote_type(Tq, Tp), + decoder.architecture.full_dim ÷ 2, + decoder.architecture.reduced_dim ÷ 2), + q_ref = nothing, p_ref = nothing) where {Tq, Tp} + T = promote_type(Tq, Tp) N2 = decoder.architecture.full_dim ÷ 2 - ξ = vcat(q̃, p̃) - result = ForwardDiff.jacobian!( - DiffResults.JacobianResult(similar(ξ, 2 * N2), ξ), decoder, ξ) + n2 = decoder.architecture.reduced_dim ÷ 2 + c = _get_vf_cache(cache, T, N2, n2) + ξ_buf = c.ξ_buf + v_intermediate = c.v_intermediate + f_intermediate = c.f_intermediate + copyto!(@view(ξ_buf[1:n2]), q̃) + copyto!(@view(ξ_buf[(n2 + 1):(2 * n2)]), p̃) + result = ForwardDiff.jacobian!(c.jac_result, decoder, ξ_buf) z = DiffResults.value(result) - q, p = @views z[1:N2], z[(N2 + 1):end] - v_intermediate = zeros(T, N2) - f_intermediate = zeros(T, N2) + _add_ref!(z, q_ref, p_ref, N2) + q, p = @views z[1:N2], z[(N2 + 1):(2 * N2)] + fill!(v_intermediate, zero(T)) + fill!(f_intermediate, zero(T)) v_full(v_intermediate, t, q, p, parameters) f_full(f_intermediate, t, q, p, parameters) v_intermediate, f_intermediate, DiffResults.jacobian(result) end +function evaluate_vf_and_vjp(t, q̃::AbstractVector{Tq}, p̃::AbstractVector{Tp}, + parameters, decoder, v_full, f_full, + cache::ReducedVFCache, q_ref = nothing, p_ref = nothing) where {Tq, Tp} + T = promote_type(Tq, Tp) + N2 = decoder.architecture.full_dim ÷ 2 + n2 = decoder.architecture.reduced_dim ÷ 2 + c = _get_vf_cache(cache, T, N2, n2) + ξ_buf = c.ξ_buf + v_intermediate = c.v_intermediate + f_intermediate = c.f_intermediate + z_bar_buf = c.z_bar_buf + copyto!(@view(ξ_buf[1:n2]), q̃) + copyto!(@view(ξ_buf[(n2 + 1):(2 * n2)]), p̃) + z_raw, back = Zygote.pullback(decoder, ξ_buf) + copyto!(z_bar_buf, z_raw) + _add_ref!(z_bar_buf, q_ref, p_ref, N2) + q, p = @views z_bar_buf[1:N2], z_bar_buf[(N2 + 1):(2 * N2)] + fill!(v_intermediate, zero(T)) + fill!(f_intermediate, zero(T)) + v_full(v_intermediate, t, q, p, parameters) + f_full(f_intermediate, t, q, p, parameters) + @inbounds for i in 1:N2 + z_bar_buf[i] = -f_intermediate[i] + z_bar_buf[N2 + i] = v_intermediate[i] + end + first(back(z_bar_buf)) +end + @doc raw""" - build_v_reduced(v_full, f_full, decoder) + build_v_reduced(v_full, f_full, decoder; pullback_mode = :jacobian, q_ref = nothing, p_ref = nothing, cache = ...) Builds the reduced vector field (``q`` part) based on the full vector field for a Hamiltonian system. @@ -138,17 +279,33 @@ We then get ``` `build_v_reduced` outputs the first half of the entries of this vector field. +When `pullback_mode = :jacobian` (default), `ForwardDiff.jacobian!` with `DiffResults.MutableDiffResult` computes ``\Psi^{\mathrm{dec}}(\xi)`` and ``\nabla_\xi \Psi^{\mathrm{dec}}`` in one pass. +When `pullback_mode = :vjp`, `(v(z), f(z))` is evaluated once at ``z = \Psi^{\mathrm{dec}}(\xi)`` and ``(\nabla_\xi \Psi^{\mathrm{dec}})^T \mathbb{J}_{2N}^T [v(z); f(z)]`` is computed via `Zygote.pullback(decoder, ξ)`. """ -function build_v_reduced(v_full, f_full, decoder::NeuralNetwork{<:SymplecticDecoder}) +function build_v_reduced(v_full, f_full, decoder::NeuralNetwork{<:SymplecticDecoder}; + pullback_mode::Symbol = :jacobian, + q_ref = nothing, + p_ref = nothing, + cache::ReducedVFCache = ReducedVFCache( + Float64, decoder.architecture.full_dim ÷ 2, + decoder.architecture.reduced_dim ÷ 2)) + _check_pullback_mode(pullback_mode) N2 = decoder.architecture.full_dim ÷ 2 n2 = decoder.architecture.reduced_dim ÷ 2 function v_reduced( - v, t, q̃::AbstractVector{T}, p̃::AbstractVector{T}, parameters) where {T} - v_intermediate, f_intermediate, ∇Ψ = evaluate_vf_and_compute_∇Ψ( - t, q̃, p̃, parameters, decoder, v_full, f_full) - ∇₂Ψ₁ = @view ∇Ψ[1:N2, (n2 + 1):(2 * n2)] - ∇₂Ψ₂ = @view ∇Ψ[(N2 + 1):(2 * N2), (n2 + 1):(2 * n2)] - v .= -∇₂Ψ₁' * f_intermediate + ∇₂Ψ₂' * v_intermediate + v, t, q̃::AbstractVector{Tq}, p̃::AbstractVector{Tp}, parameters) where {Tq, Tp} + T = promote_type(Tq, Tp) + if pullback_mode === :vjp && (T <: AbstractFloat) + g_ξ = evaluate_vf_and_vjp( + t, q̃, p̃, parameters, decoder, v_full, f_full, cache, q_ref, p_ref) + v .= @view g_ξ[(n2 + 1):(2 * n2)] + else + v_intermediate, f_intermediate, ∇Ψ = evaluate_vf_and_compute_∇Ψ( + t, q̃, p̃, parameters, decoder, v_full, f_full, cache, q_ref, p_ref) + ∇₂Ψ₁ = @view ∇Ψ[1:N2, (n2 + 1):(2 * n2)] + ∇₂Ψ₂ = @view ∇Ψ[(N2 + 1):(2 * N2), (n2 + 1):(2 * n2)] + v .= -∇₂Ψ₁' * f_intermediate + ∇₂Ψ₂' * v_intermediate + end nothing end @@ -156,7 +313,7 @@ function build_v_reduced(v_full, f_full, decoder::NeuralNetwork{<:SymplecticDeco end @doc raw""" - build_f_reduced(v_full, f_full, decoder) + build_f_reduced(v_full, f_full, decoder; pullback_mode = :jacobian, q_ref = nothing, p_ref = nothing, cache = ...) Builds the reduced vector field (``p`` part) based on the full vector field for a Hamiltonian system. @@ -164,31 +321,72 @@ Builds the reduced vector field (``p`` part) based on the full vector field for See [`build_v_reduced`](@ref) for more information. """ -function build_f_reduced(v_full, f_full, decoder::NeuralNetwork{<:SymplecticDecoder}) +function build_f_reduced(v_full, f_full, decoder::NeuralNetwork{<:SymplecticDecoder}; + pullback_mode::Symbol = :jacobian, + q_ref = nothing, + p_ref = nothing, + cache::ReducedVFCache = ReducedVFCache( + Float64, decoder.architecture.full_dim ÷ 2, + decoder.architecture.reduced_dim ÷ 2)) + _check_pullback_mode(pullback_mode) N2 = decoder.architecture.full_dim ÷ 2 n2 = decoder.architecture.reduced_dim ÷ 2 function f_reduced( - f, t, q̃::AbstractVector{T}, p̃::AbstractVector{T}, parameters) where {T} - v_intermediate, f_intermediate, ∇Ψ = evaluate_vf_and_compute_∇Ψ( - t, q̃, p̃, parameters, decoder, v_full, f_full) - ∇₁Ψ₁ = @view ∇Ψ[1:N2, 1:n2] - ∇₁Ψ₂ = @view ∇Ψ[(N2 + 1):(2 * N2), 1:n2] - f .= ∇₁Ψ₁' * f_intermediate - ∇₁Ψ₂' * v_intermediate + f, t, q̃::AbstractVector{Tq}, p̃::AbstractVector{Tp}, parameters) where {Tq, Tp} + T = promote_type(Tq, Tp) + if pullback_mode === :vjp && (T <: AbstractFloat) + g_ξ = evaluate_vf_and_vjp( + t, q̃, p̃, parameters, decoder, v_full, f_full, cache, q_ref, p_ref) + f .= .-(@view g_ξ[1:n2]) + else + v_intermediate, f_intermediate, ∇Ψ = evaluate_vf_and_compute_∇Ψ( + t, q̃, p̃, parameters, decoder, v_full, f_full, cache, q_ref, p_ref) + ∇₁Ψ₁ = @view ∇Ψ[1:N2, 1:n2] + ∇₁Ψ₂ = @view ∇Ψ[(N2 + 1):(2 * N2), 1:n2] + f .= ∇₁Ψ₁' * f_intermediate - ∇₁Ψ₂' * v_intermediate + end nothing end f_reduced end -function build_h_reduced(h_full, decoder::NeuralNetwork{<:SymplecticDecoder}) +function build_h_reduced( + h_full, decoder::NeuralNetwork{<:SymplecticDecoder}; + q_ref = nothing, p_ref = nothing) function h_reduced(t, q, p, params) - h_full(t, decoder(q, p)..., params) + dec = decoder(q, p) + if q_ref === nothing || p_ref === nothing + h_full(t, dec.q, dec.p, params) + else + h_full(t, dec.q .+ q_ref, dec.p .+ p_ref, params) + end end h_reduced end +function encode(rs::HRedSys, qp::NamedTuple{(:q, :p)}) + rs.encoder((q = qp.q .- rs.q_ref, p = qp.p .- rs.p_ref)) +end + +function encode(rs::HRedSys, z::AbstractVector) + N2 = length(rs.q_ref) + rs.encoder(vcat(@view(z[1:N2]) .- rs.q_ref, @view(z[(N2 + 1):(2 * N2)]) .- rs.p_ref)) +end + +function decode(rs::HRedSys, qp̃::NamedTuple{(:q, :p)}) + dec = rs.decoder(qp̃) + (q = dec.q .+ rs.q_ref, p = dec.p .+ rs.p_ref) +end + +function decode(rs::HRedSys, ξ::AbstractVector) + z = rs.decoder(ξ) + N2 = length(rs.q_ref) + vcat(@view(z[1:N2]) .+ rs.q_ref, @view(z[(N2 + 1):(2 * N2)]) .+ rs.p_ref) +end + function integrate_reduced_system(rs::HRedSys) - ics_reduced_nt = rs.encoder(rs.ics) + ics_reduced_nt = encode(rs, rs.ics) # convert to StateVariable format ics_reduced = (q = StateVariable(ics_reduced_nt.q), p = StateVariable(ics_reduced_nt.p)) hode = HODEProblem(rs.v_reduced, rs.f_reduced, rs.h_reduced, rs.timespan, @@ -220,7 +418,7 @@ function reduction_error(rs::HRedSys, sol_full = integrate_full_system(rs), sol_reduced = integrate_reduced_system(rs)) sol_full_matrices = data_tensors_from_geometric_solution(sol_full) sol_reduced_matrices = data_tensors_from_geometric_solution(sol_reduced) - _norm(_diff(rs.decoder(sol_reduced_matrices), sol_full_matrices)) / + _norm(_diff(decode(rs, sol_reduced_matrices), sol_full_matrices)) / _norm(sol_full_matrices) end @@ -240,6 +438,6 @@ This saves the cost of again integrating the systems. """ function projection_error(rs::HRedSys, sol_full = integrate_full_system(rs)) sol_full_matrices = data_tensors_from_geometric_solution(sol_full) - _norm(_diff(rs.decoder(rs.encoder(sol_full_matrices)), sol_full_matrices)) / + _norm(_diff(decode(rs, encode(rs, sol_full_matrices)), sol_full_matrices)) / _norm(sol_full_matrices) end diff --git a/test/architectures/multimodal_symplectic_autoencoder.jl b/test/architectures/multimodal_symplectic_autoencoder.jl new file mode 100644 index 000000000..5065c90f4 --- /dev/null +++ b/test/architectures/multimodal_symplectic_autoencoder.jl @@ -0,0 +1,463 @@ +using GeometricMachineLearning +using GeometricMachineLearning: SymplecticCompression, AutoEncoder, Architecture, + NetworkLoss, params +using GeometricMachineLearning.AbstractNeuralNetworks: GlorotUniform +using KernelAbstractions +using LinearAlgebra +using Test +import Random +import Zygote + +include("../helpers/eltypes.jl") + +Random.seed!(42) + +# Helper to flatten and unflatten nested NamedTuple/Array parameters for finite differences +function _flatten_params(x::AbstractArray{T}) where {T <: Real} + raw = parent(x) + vec_copy = copy(vec(raw)) + sz = size(raw) + is_stiefel = x isa StiefelManifold + rebuild = v -> begin + A = reshape(v, sz) + is_stiefel ? StiefelManifold(A) : A + end + return vec_copy, rebuild +end + +function _flatten_params(nt::NamedTuple{names}) where {names} + vecs = Vector{Float64}[] + rebuilders = Any[] + lengths = Int[] + for k in names + v, r = _flatten_params(getfield(nt, k)) + push!(vecs, Float64.(v)) + push!(rebuilders, r) + push!(lengths, length(v)) + end + flat = isempty(vecs) ? Float64[] : reduce(vcat, vecs) + rebuild = w -> begin + offset = 0 + vals = map(eachindex(names)) do i + len = lengths[i] + sub = w[(offset + 1):(offset + len)] + offset += len + rebuilders[i](sub) + end + NamedTuple{names}(Tuple(vals)) + end + return flat, rebuild +end + +function _flatten_params(ps::NetworkParameters) + flat, rebuild_nt = _flatten_params(params(ps)) + return flat, w -> NetworkParameters(rebuild_nt(w)) +end + +function _extract_grad_vec(g::NamedTuple{names}) where {names} + vecs = Float64[] + for k in names + v = getfield(g, k) + if v isa NamedTuple + append!(vecs, _extract_grad_vec(v)) + elseif v isa AbstractArray + append!(vecs, vec(Float64.(Array(v)))) + end + end + return vecs +end + +_extract_grad_vec(g::NetworkParameters) = _extract_grad_vec(params(g)) + +@testset "Clause 1: MultiModalSymplecticAutoencoder construction and edge cases" begin + @test isdefined(GeometricMachineLearning, :MultiModalSymplecticAutoencoder) + @test isdefined(GeometricMachineLearning, :encode) + @test isdefined(GeometricMachineLearning, :decode) + @test isdefined(GeometricMachineLearning, :translate) + @test MultiModalSymplecticAutoencoder <: SymplecticCompression + @test MultiModalSymplecticAutoencoder <: AutoEncoder + @test MultiModalSymplecticAutoencoder <: Architecture + + core = WeightTiedSymplecticAutoencoder( + 8, 4; n_blocks = 1, n_layers = 1, n_ambient_layers = 0, sympnet_upscale = 2) + + re_empty = r"^(?:ArgumentError: )?MultiModalSymplecticAutoencoder requires a non-empty frontends NamedTuple\.$" + @test_throws re_empty MultiModalSymplecticAutoencoder(NamedTuple(), core) + + fe_bad = CanonicalSortLayer(3, 2) # output dim = 12 != 8 + re_dim = r"^(?:ArgumentError: )?Frontend :particles has output dimension 12, expected core\.full_dim = 8\.$" + @test_throws re_dim MultiModalSymplecticAutoencoder((particles = fe_bad,), core) + + fe_bad_grid = BilinearStiefelLayer(6, 6, 3, 2) # output dim = 2 * 3 * 2 = 12 != 8 + re_dim_grid = r"^(?:ArgumentError: )?Frontend :grid has output dimension 12, expected core\.full_dim = 8\.$" + @test_throws re_dim_grid MultiModalSymplecticAutoencoder((grid = fe_bad_grid,), core) + + fe_part = CanonicalSortLayer(2, 2) # output dim = 8 + fe_grid = BilinearStiefelLayer(4, 4, 2, 2) # output dim = 8 + fe_dlra = LowRankStiefelLayer(4, 4, 2, 2; r_out = 4) # output dim = 8 + arch = MultiModalSymplecticAutoencoder( + (particles = fe_part, grid = fe_grid, dlra = fe_dlra), core) + @test arch.full_dim == 8 + @test arch.reduced_dim == 4 + + nn = NeuralNetwork(arch, CPU(), Float64) + re_unknown = r"^(?:ArgumentError: )?Unknown modality :unknown\. Available modalities: \(:particles, :grid, :dlra\)\.$" + @test_throws re_unknown encode(nn, :unknown, zeros(8)) + @test_throws re_unknown decode(nn, :unknown, zeros(4)) + @test_throws re_unknown translate(nn, :unknown, :grid, zeros(8)) + @test_throws re_unknown translate(nn, :grid, :unknown, zeros(4, 4)) +end + +@testset "Clause 1: encode, decode, translate and right-invertible round-trip ($T)" for T in REAL_ELTYPES + rng = Random.MersenneTwister(101) + nx, nv = 5, 4 + mx, mv = 2, 2 + N = mx * mv # 4 + full_dim = 2 * N # 8 + reduced_dim = 4 + + # Upper-triangular Cholesky metric factors for the spline modality + A_x = randn(rng, T, nx, nx) + A_v = randn(rng, T, nv, nv) + R_x = Matrix{T}(cholesky(Symmetric(transpose(A_x) * A_x + T(2) * I)).U) + R_v = Matrix{T}(cholesky(Symmetric(transpose(A_v) * A_v + T(2) * I)).U) + + fe_sort = CanonicalSortLayer(2, 2) + fe_slot_cm = CanonicalSlotLayer(4, 4, 1) # NOut = 2 * 1 * 4 = 8 + fe_slot = CanonicalSlotLayer(3, 2, 2) # NOut = 2 * 2 * 2 = 8 + fe_spline = BilinearStiefelLayer(nx, nv, mx, mv; metric_factors = (R_x, R_v)) + fe_grid = BilinearStiefelLayer(nx, nv, mx, mv) + fe_dlra = LowRankStiefelLayer(nx, nv, mx, mv; r_out = mx * mv) + + core = WeightTiedSymplecticAutoencoder( + full_dim, reduced_dim; + n_blocks = 1, + n_layers = 1, + n_ambient_layers = 0, + sympnet_upscale = 2 + ) + frontends = ( + particles = fe_sort, + particles_slot_cm = fe_slot_cm, + spline = fe_spline, + grid = fe_grid, + dlra = fe_dlra + ) + arch = MultiModalSymplecticAutoencoder(frontends, core) + nn = NeuralNetwork(arch, CPU(), T; rng = rng) + + # 1. Verify encode, decode and translate match explicit frontend + core composition + q_p = randn(rng, T, 2, 2) + p_p = randn(rng, T, 2, 2) + x_part = (q = q_p, p = p_p) + F_grid = randn(rng, T, nx, nv) + F_spline = randn(rng, T, nx, nv) + r_in = 3 + X_in = Matrix(qr(randn(rng, T, nx, r_in)).Q[:, 1:r_in]) + V_in = Matrix(qr(randn(rng, T, nv, r_in)).Q[:, 1:r_in]) + S_in = randn(rng, T, r_in, r_in) + x_dlra = (X = X_in, S = S_in, V = V_in) + + _enc_grid(net, x) = encode(net, :grid, x) + _enc_spline(net, x) = encode(net, :spline, x) + _dec_grid(net, ξ) = decode(net, :grid, ξ) + + ξ_part = @inferred encode(nn, :particles, x_part) + ξ_grid = @inferred _enc_grid(nn, F_grid) + ξ_spline = @inferred _enc_spline(nn, F_spline) + ξ_dlra = @inferred encode(nn, :dlra, x_dlra) + + @test eltype(ξ_part) === T && length(ξ_part) == reduced_dim + @test eltype(ξ_grid) === T && length(ξ_grid) == reduced_dim + @test eltype(ξ_spline) === T && length(ξ_spline) == reduced_dim + @test eltype(ξ_dlra) === T && length(ξ_dlra) == reduced_dim + + # Explicit frontend followed by core encoder matches encode(nn, modality, x) + z_grid_fe = fe_grid(F_grid, params(nn).frontends.grid) + ξ_grid_manual = encoder(nn)(vcat(z_grid_fe.q, z_grid_fe.p)) + @test isapprox(ξ_grid, ξ_grid_manual; atol = 8 * eps(T), rtol = 8 * eps(T)) + + # Explicit core decoder followed by inverse_layer(frontend) matches decode(nn, modality, ξ) + z_dec_manual = decoder(nn)(ξ_grid) + F_grid_manual = inverse_layer(fe_grid)( + (q = z_dec_manual[1:N], p = z_dec_manual[(N + 1):full_dim]), + params(nn).frontends.grid + ) + F_grid_dec = @inferred _dec_grid(nn, ξ_grid) + @test isapprox(F_grid_dec, F_grid_manual; atol = 8 * eps(T), rtol = 8 * eps(T)) + + # NamedTuple (q, p) latent input to decode + F_grid_dec_nt = decode( + nn, :grid, + (q = ξ_grid[1:(reduced_dim ÷ 2)], p = ξ_grid[(reduced_dim ÷ 2 + 1):reduced_dim])) + @test isapprox(F_grid_dec_nt, F_grid_dec; atol = 8 * eps(T), rtol = 8 * eps(T)) + + # translate(nn, from, to, x) == decode(nn, to, encode(nn, from, x)) + @test translate(nn, :particles, :grid, x_part) == + decode(nn, :grid, encode(nn, :particles, x_part)) + @test translate(nn, :grid, :dlra, F_grid) == + decode(nn, :dlra, encode(nn, :grid, F_grid)) + @test translate(nn, :dlra, :spline, x_dlra) == + decode(nn, :spline, encode(nn, :dlra, x_dlra)) + + # 2. Exact right-inversion encode(nn, mod, decode(nn, mod, ξ)) ≈ ξ to 50 eps(T) + # Global right-invertible frontends on arbitrary ξ ∈ R^{2n}: + for ξ_test in ( + randn(rng, T, reduced_dim), + zeros(T, reduced_dim), + fill(floatmin(T) * T(16), reduced_dim) + ) + for mod in (:grid, :spline, :dlra, :particles_slot_cm) + x_rec = decode(nn, mod, ξ_test) + ξ_rt = encode(nn, mod, x_rec) + @test eltype(ξ_rt) === T + @test norm(ξ_rt - ξ_test) <= 50 * eps(T) * max(one(T), norm(ξ_test)) + end + end + + # Batched latent matrix ξ_batch ∈ R^{2n × B} for :grid, :spline, :particles_slot_cm + ξ_batch = randn(rng, T, reduced_dim, 3) + for mod in (:grid, :spline, :particles_slot_cm) + x_rec_b = decode(nn, mod, ξ_batch) + ξ_rt_b = encode(nn, mod, x_rec_b) + @test eltype(ξ_rt_b) === T && size(ξ_rt_b) == size(ξ_batch) + @test norm(ξ_rt_b - ξ_batch) <= 50 * eps(T) * max(one(T), norm(ξ_batch)) + end + + # Particles with :sort and :slot (NSlots > 1) on the canonical Weyl chamber: + core_ch = WeightTiedSymplecticAutoencoder( + full_dim, 2; + n_blocks = 1, + n_layers = 0, + n_ambient_layers = 0, + sympnet_upscale = 2 + ) + arch_ch = MultiModalSymplecticAutoencoder( + (particles = fe_sort, slot = fe_slot), core_ch) + nn_ch = NeuralNetwork(arch_ch, CPU(), T; rng = rng) + + # Set core Stiefel leaf Φ_core ∈ St(1, 4) so decoded state for ξ_q > 0 lies in the + # canonical chamber for both :sort (increasing particle row-1 q) and :slot (slot 2 active) + Φ_core = zeros(T, N, 1) + Φ_core[3, 1] = one(T) + copyto!(parent(first(values(params(nn_ch).core)).weight), Φ_core) + H_rel_3 = GeometricMachineLearning._helmert_relative_matrix(T, 3) + Φ_rel_mat = reshape(transpose(H_rel_3) * T[-2.0, 0.8, 1.2], 2, 1) + Φ_rel_mat ./= norm(Φ_rel_mat) + copyto!(parent(params(nn_ch).frontends.slot.Φ_rel), Φ_rel_mat) + + _dec_slot(net, ξ) = decode(net, :slot, ξ) + _dec_part(net, ξ) = decode(net, :particles, ξ) + + ξ_slot_test = T[0.8, -0.5] + x_slot_rec = @inferred _dec_slot(nn_ch, ξ_slot_test) + ξ_slot_rt = @inferred encode(nn_ch, :slot, x_slot_rec) + @test norm(ξ_slot_rt - ξ_slot_test) <= 50 * eps(T) * norm(ξ_slot_test) + + # For :sort (M = 2, d = 2), particle 1 has row-1 index 1 and particle 2 has row-1 index 3: + # Φ_core[1, 1] = -1/sqrt(2), Φ_core[3, 1] = +1/sqrt(2) puts q_1 < q_2 for ξ_1 > 0 + Φ_core_sort = T[-inv(sqrt(T(2))), zero(T), inv(sqrt(T(2))), zero(T)] + copyto!(parent(first(values(params(nn_ch).core)).weight), reshape(Φ_core_sort, N, 1)) + ξ_sort_test = T[0.9, -0.4] + x_sort_rec = @inferred _dec_part(nn_ch, ξ_sort_test) + ξ_sort_rt = @inferred encode(nn_ch, :particles, x_sort_rec) + @test norm(ξ_sort_rt - ξ_sort_test) <= 50 * eps(T) * norm(ξ_sort_test) +end + +@testset "Clause 2: Multi-modal losses, zero-input edge case, M1, and Zygote vs FD gradients" begin + @test isdefined(GeometricMachineLearning, :CharbonnierRelativeLoss) + @test isdefined(GeometricMachineLearning, :IsometricRolloutLoss) + @test isdefined(GeometricMachineLearning, :CrossModalCycleLoss) + @test CharbonnierRelativeLoss <: NetworkLoss + @test IsometricRolloutLoss <: NetworkLoss + @test CrossModalCycleLoss <: NetworkLoss + + # Default keyword checks + l_charb_def = CharbonnierRelativeLoss() + @test l_charb_def.ε == 1e-3 + @test l_charb_def.λ_lat == 1e-2 + + step_id = z -> z + 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_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 + + # Edge validation for loss hyperparameters + re_charb_eps = r"^(?:ArgumentError: )?CharbonnierRelativeLoss requires ε > 0 and λ_lat >= 0, got ε = 0\.0, λ_lat = 0\.01\.$" + re_charb_lam = r"^(?:ArgumentError: )?CharbonnierRelativeLoss requires ε > 0 and λ_lat >= 0, got ε = 0\.001, λ_lat = -0\.1\.$" + @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\.$" + @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) + + @testset "Forward evaluation, @inferred, zero edge case, and M1 ($T)" for T in REAL_ELTYPES + rng = Random.MersenneTwister(202) + nx, nv = 4, 4 + mx, mv = 2, 1 + fe_part = CanonicalSlotLayer(3, 1, 2) # output dim = 4 + fe_grid = BilinearStiefelLayer(nx, nv, mx, mv) # output dim = 4 + core = WeightTiedSymplecticAutoencoder( + 4, 2; n_blocks = 1, n_layers = 1, n_ambient_layers = 0, sympnet_upscale = 2) + arch = MultiModalSymplecticAutoencoder( + (particles = fe_part, grid = fe_grid), core) + nn = NeuralNetwork(arch, CPU(), T; rng = rng) + + loss_c = CharbonnierRelativeLoss(; ε = T(1e-2), λ_lat = T(5e-2)) + F = randn(rng, T, nx, nv) + F_hat = decode(nn, :grid, encode(nn, :grid, F)) + + # Manual formula check + z = encode(nn, :grid, F) + 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) + + lam_T * sum(abs2, z_hat .- z) / (sum(abs2, F) + eps_T^2) + + 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)) + + # Independent perturbation F_pert != F_hat so z_hat != z and the latent penalty term is active (catches O3) + F_pert = F .+ T(0.5) .* ones(T, nx, nv) + z_pert = encode(nn, :grid, F_pert) + expected_c_pert = (sqrt(sum(abs2, F_pert .- F) + eps_T^2) - eps_T) / + sqrt(sum(abs2, F) + eps_T^2) + + lam_T * sum(abs2, z_pert .- z) / (sum(abs2, F) + eps_T^2) + val_c_pert = @inferred loss_c(nn, :grid, F_pert, F) + @test isapprox(val_c_pert, expected_c_pert; atol = 10 * eps(T), rtol = 10 * eps(T)) + + # Mutant M1 check: when x̂ == x != 0, CharbonnierRelativeLoss MUST be exact zero(T)! + # Omitting `- ε` in the numerator yields ε / sqrt(||x||_F^2 + ε^2) > 0. + val_identical = @inferred loss_c(nn, :grid, F, F) + @test val_identical === zero(T) + + # Decided at the edges: when x = 0 and x̂ = 0, returns exact zero(T) (not NaN) + F_zero = zeros(T, nx, nv) + val_zero = @inferred loss_c(nn, :grid, F_zero, F_zero) + @test !isnan(val_zero) && val_zero === zero(T) + + # IsometricRolloutLoss forward & @inferred + J2 = [zero(T) one(T); -one(T) zero(T)] + dt = T(0.05) + step_lin = ξ -> ξ .+ dt .* (J2 * ξ) + H_lat = ξ -> T(0.5) * sum(abs2, ξ) + loss_iso = IsometricRolloutLoss( + step_lin; α_lat = T(1.0), α_rec = T(0.5), α_H = T(0.2), H = H_lat) + + re_len = r"^(?:ArgumentError: )?IsometricRolloutLoss requires at least 2 trajectory steps, got 1\.$" + @test_throws re_len loss_iso(nn, :grid, [F]) + + F_traj = [randn(rng, T, nx, nv) for _ in 1:4] + val_iso = @inferred loss_iso(nn, :grid, F_traj) + @test eltype(val_iso) === T && val_iso > zero(T) + + # Manual formula check for IsometricRolloutLoss + z_enc = [encode(nn, :grid, F_k) for F_k in F_traj] + z_roll = Vector{typeof(z_enc[1])}(undef, 4) + z_roll[1] = z_enc[1] + for k in 1:3 + z_roll[k + 1] = step_lin(z_roll[k]) + end + err_lat = sum(sum(abs2, z_roll[k] .- z_enc[k]) for k in 2:4) / T(3) + err_rec = sum(sum(abs2, decode(nn, :grid, z_roll[k]) .- F_traj[k]) for k in 2:4) / + T(3) + err_H = sum(abs2(H_lat(z_roll[k]) - H_lat(z_roll[1])) for k in 2:4) / T(3) + 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)) + x_p = (q = randn(rng, T, 1, 3), p = randn(rng, T, 1, 3)) + val_cyc = @inferred loss_cyc(nn, x_p, F) + @test eltype(val_cyc) === T && val_cyc > zero(T) + + 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) + @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 (catches O5) + loss_cyc_rev = CrossModalCycleLoss( + :grid, :particles; β_align = T(1.2), β_cycle = 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) + @test isapprox(val_cyc_rev, expected_cyc_rev; atol = 10 * eps(T), rtol = 10 * + eps(T)) + end + + @testset "Analytical Zygote gradients match central finite differences in Float64" begin + rng = Random.MersenneTwister(303) + T = Float64 + nx, nv = 4, 4 + mx, mv = 2, 1 + # Full dim = 2 * 2 * 1 = 4, reduced_dim = 2 + fe_part = CanonicalSlotLayer(3, 1, 2) # M = 3, d = 1, NSlots = 2 -> NOut = 4 + fe_grid = BilinearStiefelLayer(nx, nv, mx, mv) # NOut = 4 + core = WeightTiedSymplecticAutoencoder( + 4, 2; n_blocks = 1, n_layers = 1, n_ambient_layers = 0, sympnet_upscale = 2) + arch = MultiModalSymplecticAutoencoder((particles = fe_part, grid = fe_grid), core) + nn = NeuralNetwork(arch, CPU(), T; rng = rng) + + F = randn(rng, T, nx, nv) + F_traj = [randn(rng, T, nx, nv) for _ in 1:3] + # Particle state with well-separated q[1, :] and relative mode norms so sorting permutations are locally constant + x_p = (q = T[-2.0 0.1 2.5], p = T[0.4 -1.3 0.9]) + + loss_c = CharbonnierRelativeLoss(; ε = 1e-2, λ_lat = 5e-2) + F_zero = zeros(T, nx, nv) + g_zero = Zygote.gradient(x_hat -> loss_c(nn, :grid, x_hat, F_zero), F_zero)[1] + @test all(isfinite, g_zero) && all(iszero, g_zero) + g_same = Zygote.gradient(x_hat -> loss_c(nn, :grid, x_hat, F), F)[1] + @test all(isfinite, g_same) && all(iszero, g_same) + + step_fn = ξ -> [ξ[1] + 0.1 * ξ[2] + 0.05 * tanh(ξ[2]), ξ[2] - 0.1 * ξ[1]] + 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) + + flat0, rebuild_ps = _flatten_params(nn.params) + h = 1e-6 + tol = 10 * sqrt(eps(Float64)) + + 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)) + ) + g_zyg = Zygote.gradient(f_ps, nn.params)[1] + g_vec = _extract_grad_vec(g_zyg) + @test length(g_vec) == length(flat0) + + g_fd = zeros(Float64, length(flat0)) + w_plus = copy(flat0) + w_minus = copy(flat0) + for i in eachindex(flat0) + w_plus[i] = flat0[i] + h + w_minus[i] = flat0[i] - h + g_fd[i] = (f_ps(rebuild_ps(w_plus)) - f_ps(rebuild_ps(w_minus))) / (2 * h) + w_plus[i] = flat0[i] + w_minus[i] = flat0[i] + end + @test isapprox(g_vec, g_fd; atol = tol, rtol = tol) + end + end +end diff --git a/test/reduced_system/reference_shift_and_vjp.jl b/test/reduced_system/reference_shift_and_vjp.jl new file mode 100644 index 000000000..56ac7d574 --- /dev/null +++ b/test/reduced_system/reference_shift_and_vjp.jl @@ -0,0 +1,258 @@ +using GeometricMachineLearning +using GeometricMachineLearning: HRedSys, build_v_reduced, build_f_reduced, + integrate_reduced_system, evaluate_vf_and_compute_∇Ψ, + ForwardDiff, StateVariable +using GeometricIntegrators: ImplicitMidpoint +using KernelAbstractions +using LinearAlgebra +using Test +import Random + +include("../helpers/eltypes.jl") + +Random.seed!(42) + +@testset "Clause 3: HRedSys reference-state shifting, pullback_mode = :jacobian | :vjp, and work buffers" begin + model_chk = WeightTiedSymplecticAutoencoder( + 6, 2; n_blocks = 1, n_layers = 1, n_ambient_layers = 0, sympnet_upscale = 2) + nn_chk = NeuralNetwork(model_chk, CPU(), Float64) + v_dummy!(v, t, q, p, ps) = (v .= p; nothing) + f_dummy!(f, t, q, p, ps) = (f .= .-q; nothing) + h_dummy(t, q, p, ps) = 0.5 * (sum(abs2, q) + sum(abs2, p)) + + re_mode = r"^(?:ArgumentError: )?pullback_mode must be :jacobian or :vjp, got :invalid\.$" + @test_throws re_mode build_v_reduced( + v_dummy!, f_dummy!, decoder(nn_chk); pullback_mode = :invalid) + @test_throws re_mode build_f_reduced( + v_dummy!, f_dummy!, decoder(nn_chk); pullback_mode = :invalid) + + re_ref_dim = r"^(?:DimensionMismatch: )?q_ref and p_ref must have length N = 3, got 2 and 3\.$" + ics_chk = (q = StateVariable(zeros(3)), p = StateVariable(zeros(3))) + @test_throws re_ref_dim HRedSys( + 6, 2, encoder(nn_chk), decoder(nn_chk), v_dummy!, f_dummy!, h_dummy, + (0.0, 0.1), 0.05, ics_chk; + integrator = ImplicitMidpoint(), + q_ref = zeros(2), + p_ref = zeros(3) + ) + + @testset "2N = 12, 2n = 4 VJP vs Jacobian agreement and equilibrium integration ($T)" for T in REAL_ELTYPES + rng = Random.MersenneTwister(404) + N = 6 + n = 2 + full_dim = 2 * N # 12 + reduced_dim = 2 * n # 4 + + arch = WeightTiedSymplecticAutoencoder( + full_dim, reduced_dim; + n_blocks = 1, + n_layers = 1, + n_ambient_layers = 0, + sympnet_upscale = 2 + ) + nn = NeuralNetwork(arch, CPU(), T; rng = rng) + + q_ref = randn(rng, T, N) + p_ref = randn(rng, T, N) + + # Shifted quadratic + anharmonic Hamiltonian vector field: + # H(q, p) = 1/2 ||q - q_ref||^2 + 1/2 ||p - p_ref||^2 + 1/4 sum((q - q_ref).^4) + function v_field!(v, t, q, p, params) + @inbounds for i in eachindex(v, p) + v[i] = p[i] - p_ref[i] + end + nothing + end + function f_field!(f, t, q, p, params) + @inbounds for i in eachindex(f, q) + dq = q[i] - q_ref[i] + f[i] = -dq - dq^3 + end + nothing + end + function h_field(t, q, p, params) + T(0.5) * (sum(abs2, q .- q_ref) + sum(abs2, p .- p_ref)) + + T(0.25) * sum(x -> x^4, q .- q_ref) + end + + ics_eq = (q = StateVariable(copy(q_ref)), p = StateVariable(copy(p_ref))) + + rs_jac = HRedSys( + full_dim, reduced_dim, encoder(nn), decoder(nn), + v_field!, f_field!, h_field, + (zero(T), T(1.0)), T(0.05), ics_eq; + integrator = ImplicitMidpoint(), + q_ref = q_ref, + p_ref = p_ref, + pullback_mode = :jacobian + ) + rs_vjp = HRedSys( + full_dim, reduced_dim, encoder(nn), decoder(nn), + v_field!, f_field!, h_field, + (zero(T), T(1.0)), T(0.05), ics_eq; + integrator = ImplicitMidpoint(), + q_ref = q_ref, + p_ref = p_ref, + pullback_mode = :vjp + ) + + # 1. Reference-state shifting (Proposition 13.1): + # For WeightTiedSymplecticAutoencoder, Ψ^dec(0) = 0 and Ψ^enc(0) = 0, so with + # reference shift (q_ref, p_ref), ξ = (0, 0) lifts identically to (q_ref, p_ref) + # and (q_ref, p_ref) encodes identically to ξ = (0, 0). + ξ_zero = zeros(T, reduced_dim) + z_lifted_zero = decode(rs_jac, ξ_zero) + @test eltype(z_lifted_zero) === T + @test isapprox( + z_lifted_zero, vcat(q_ref, p_ref); atol = 10 * eps(T), rtol = 10 * eps(T)) + qp_lifted_zero = decode(rs_jac, (q = zeros(T, n), p = zeros(T, n))) + @test isapprox(qp_lifted_zero.q, q_ref; atol = 10 * eps(T), rtol = 10 * eps(T)) + @test isapprox(qp_lifted_zero.p, p_ref; atol = 10 * eps(T), rtol = 10 * eps(T)) + ξ_enc_ref = encode(rs_jac, vcat(q_ref, p_ref)) + @test eltype(ξ_enc_ref) === T + @test norm(ξ_enc_ref) <= 10 * eps(T) + qp_enc_ref = encode(rs_jac, (q = q_ref, p = p_ref)) + @test norm(qp_enc_ref.q) <= 10 * eps(T) && norm(qp_enc_ref.p) <= 10 * eps(T) + + # Perturbation δz around (q_ref, p_ref) + δq = randn(rng, T, N) + δp = randn(rng, T, N) + ξ_pert = encode(rs_jac, vcat(q_ref .+ δq, p_ref .+ δp)) + @test isapprox( + ξ_pert, encoder(nn)(vcat(δq, δp)); atol = 10 * eps(T), rtol = 10 * eps(T)) + qp_pert = encode(rs_jac, (q = q_ref .+ δq, p = p_ref .+ δp)) + qp_pert_exp = encoder(nn)((q = δq, p = δp)) + @test isapprox(qp_pert.q, qp_pert_exp.q; atol = 10 * eps(T), rtol = 10 * eps(T)) + @test isapprox(qp_pert.p, qp_pert_exp.p; atol = 10 * eps(T), rtol = 10 * eps(T)) + + # 2. VJP vs Jacobian agreement to 50 eps(T) for 2N = 12, 2n = 4 in Float32 and Float64 (catches M2!) + for trial in 1:5 + q̃ = randn(rng, T, n) + p̃ = randn(rng, T, n) + v_jac = zeros(T, n) + f_jac = zeros(T, n) + v_vjp = zeros(T, n) + f_vjp = zeros(T, n) + + rs_jac.v_reduced(v_jac, zero(T), q̃, p̃, nothing) + rs_jac.f_reduced(f_jac, zero(T), q̃, p̃, nothing) + rs_vjp.v_reduced(v_vjp, zero(T), q̃, p̃, nothing) + rs_vjp.f_reduced(f_vjp, zero(T), q̃, p̃, nothing) + + # Verify non-zero vector field so a sign flip (M2) or zeroing is caught + @test norm(v_jac) > T(1e-2) + @test norm(f_jac) > T(1e-2) + @test norm(v_vjp - v_jac) <= 50 * eps(T) * max(one(T), norm(v_jac)) + @test norm(f_vjp - f_jac) <= 50 * eps(T) * max(one(T), norm(f_jac)) + end + + # Also test standalone build_v_reduced and build_f_reduced directly + v_red_j = build_v_reduced( + v_field!, f_field!, decoder(nn); + pullback_mode = :jacobian, q_ref = q_ref, p_ref = p_ref) + f_red_j = build_f_reduced( + v_field!, f_field!, decoder(nn); + pullback_mode = :jacobian, q_ref = q_ref, p_ref = p_ref) + v_red_v = build_v_reduced( + v_field!, f_field!, decoder(nn); + pullback_mode = :vjp, q_ref = q_ref, p_ref = p_ref) + f_red_v = build_f_reduced( + v_field!, f_field!, decoder(nn); + pullback_mode = :vjp, q_ref = q_ref, p_ref = p_ref) + + q̃_test = randn(rng, T, n) + p̃_test = randn(rng, T, n) + vj_out = zeros(T, n) + fj_out = zeros(T, n) + vv_out = zeros(T, n) + fv_out = zeros(T, n) + v_red_j(vj_out, zero(T), q̃_test, p̃_test, nothing) + f_red_j(fj_out, zero(T), q̃_test, p̃_test, nothing) + v_red_v(vv_out, zero(T), q̃_test, p̃_test, nothing) + f_red_v(fv_out, zero(T), q̃_test, p̃_test, nothing) + @test norm(vv_out - vj_out) <= 50 * eps(T) * max(one(T), norm(vj_out)) + @test norm(fv_out - fj_out) <= 50 * eps(T) * max(one(T), norm(fj_out)) + + # 3. Preallocated work buffers & ForwardDiff.Dual cache in HRedSys: + @test hasfield(typeof(rs_jac), :ξ_buf) + @test hasfield(typeof(rs_jac), :v_intermediate) + @test hasfield(typeof(rs_jac), :f_intermediate) + @test hasfield(typeof(rs_jac), :jac_result) + @test hasfield(typeof(rs_jac), :dual_cache) + @test rs_jac.dual_cache isa IdDict{DataType, Any} + + # Evaluate with ForwardDiff.Dual to populate dual_cache and verify agreement + J_v_dual = ForwardDiff.jacobian( + q -> begin + out = similar(q, n) + rs_jac.v_reduced(out, zero(T), q, p̃_test, nothing) + out + end, + q̃_test + ) + @test size(J_v_dual) == (n, n) && all(isfinite, J_v_dual) + @test !isempty(rs_jac.dual_cache) + + # 4. Quadratic Hamiltonian H(q, p) = 1/2 ||(q - q_ref)^2 + (p - p_ref)^2|| + # with equilibrium (q_ref, p_ref) and integrator = ImplicitMidpoint(): + # integrate_reduced_system starting at ics = (q = q_ref, p = p_ref) keeps + # ξ(t) = (0, 0) and (q(t), p(t)) = (q_ref, p_ref) to 10 eps(T) for 20 steps (catches M3!). + function v_quad!(v, t, q, p, params) + @inbounds for i in eachindex(v, p) + v[i] = p[i] - p_ref[i] + end + nothing + end + function f_quad!(f, t, q, p, params) + @inbounds for i in eachindex(f, q) + f[i] = -(q[i] - q_ref[i]) + end + nothing + end + function h_quad(t, q, p, params) + T(0.5) * (sum(abs2, q .- q_ref) + sum(abs2, p .- p_ref)) + end + + rs_eq = HRedSys( + full_dim, reduced_dim, encoder(nn), decoder(nn), + v_quad!, f_quad!, h_quad, + (zero(T), T(1.0)), T(0.05), # 20 steps + (q = StateVariable(copy(q_ref)), p = StateVariable(copy(p_ref))); + integrator = ImplicitMidpoint(), + q_ref = q_ref, + p_ref = p_ref + ) + sol_red = integrate_reduced_system(rs_eq) + @test length(sol_red.q) == 21 # initial state + 20 steps + for step_idx in 0:20 + q̃_k = Vector(sol_red.q[step_idx]) + p̃_k = Vector(sol_red.p[step_idx]) + @test norm(q̃_k) <= 10 * eps(T) + @test norm(p̃_k) <= 10 * eps(T) + lifted_k = decode(rs_eq, (q = q̃_k, p = p̃_k)) + @test norm(lifted_k.q - q_ref) <= 10 * eps(T) * max(one(T), norm(q_ref)) + @test norm(lifted_k.p - p_ref) <= 10 * eps(T) * max(one(T), norm(p_ref)) + end + + # Non-equilibrium trajectory check: verify iterate moves when ics != (q_ref, p_ref) + q_init = q_ref .+ T(0.2) .* randn(rng, T, N) + p_init = p_ref .+ T(0.2) .* randn(rng, T, N) + rs_neq = HRedSys( + full_dim, reduced_dim, encoder(nn), decoder(nn), + v_quad!, f_quad!, h_quad, + (zero(T), T(0.2)), T(0.05), + (q = StateVariable(q_init), p = StateVariable(p_init)); + integrator = ImplicitMidpoint(), + q_ref = q_ref, + p_ref = p_ref, + pullback_mode = :jacobian + ) + sol_red_neq = integrate_reduced_system(rs_neq) + @test norm(Vector(sol_red_neq.q[1]) - Vector(sol_red_neq.q[0])) > T(1e-4) + lifted_0 = decode(rs_neq, ( + q = Vector(sol_red_neq.q[0]), p = Vector(sol_red_neq.p[0]))) + lifted_1 = decode(rs_neq, ( + q = Vector(sol_red_neq.q[1]), p = Vector(sol_red_neq.p[1]))) + @test norm(lifted_1.q - lifted_0.q) > T(1e-4) + end +end diff --git a/test/runtests.jl b/test/runtests.jl index 6e05e2574..4521ea2fc 100644 --- a/test/runtests.jl +++ b/test/runtests.jl @@ -63,6 +63,7 @@ if "core" in GROUPS @safetestset "PSD tests" include("architectures/psd_architecture_tests.jl") @safetestset "SymplecticAutoencoder tests" include("architectures/symplectic_autoencoder_tests.jl") @safetestset "WeightTiedSymplecticAutoencoder" include("architectures/weight_tied_symplectic_autoencoder.jl") + @safetestset "MultiModalSymplecticAutoencoder and multimodal losses" include("architectures/multimodal_symplectic_autoencoder.jl") @safetestset "Check if autoencoder error is lower than PSD error" include("architectures/sae_error_lower_than_psd_error.jl") @safetestset "Architecture forward passes against the recorded outputs" include("architectures/architecture_pins.jl") @safetestset "Architecture keywords and chain builders" include("architectures/architecture_structure.jl") @@ -70,6 +71,7 @@ if "core" in GROUPS @safetestset "ParametricResNet" include("architectures/parametric_resnet.jl") @safetestset "The forward pass of every exported architecture infers" include("architectures/forward_inference.jl") @safetestset "Check reduced model" include("reduced_system/reduced_system.jl") + @safetestset "HRedSys reference-state shifting and VJP pullback mode" include("reduced_system/reference_shift_and_vjp.jl") @safetestset "_norm keeps the element type of its argument" include("utils.jl") @safetestset "Test NetworkLoss + Optimizer" include("integration/losses_and_optimization.jl") @safetestset "ReducedLoss through train!" include("integration/reduced_loss_optimization.jl") From 7f21517dc39d3cde00ddba55961b146a46f1c0c2 Mon Sep 17 00:00:00 2001 From: Michael Kraus Date: Fri, 30 Oct 2026 04:46:28 +0100 Subject: [PATCH 2/8] Fix Chain frontend params, canonical NamedTuple encoding, SymmetricSymplecticAutoencoder core, and build_h_reduced --- .../multimodal_symplectic_autoencoder.jl | 95 ++++++++--- src/reduced_system/reduced_system.jl | 15 +- .../multimodal_symplectic_autoencoder.jl | 148 +++++++++++++++++- .../reduced_system/reference_shift_and_vjp.jl | 37 ++++- 4 files changed, 255 insertions(+), 40 deletions(-) diff --git a/src/architectures/multimodal_symplectic_autoencoder.jl b/src/architectures/multimodal_symplectic_autoencoder.jl index b3ccf29a1..ff5076517 100644 --- a/src/architectures/multimodal_symplectic_autoencoder.jl +++ b/src/architectures/multimodal_symplectic_autoencoder.jl @@ -50,23 +50,32 @@ struct MultiModalSymplecticAutoencoder{ function MultiModalSymplecticAutoencoder( frontends::NamedTuple{names}, core::SymplecticCompression) where {names} - isempty(names) && throw(ArgumentError( - "MultiModalSymplecticAutoencoder requires a non-empty frontends NamedTuple.")) + isempty(names) && throw( + ArgumentError( + "MultiModalSymplecticAutoencoder requires a non-empty frontends NamedTuple." + ), + ) N2 = core.full_dim n2 = core.reduced_dim for mod in names fe = getfield(frontends, mod) fe_dim = _frontend_output_dimension(fe) - fe_dim == N2 || throw(ArgumentError( - "Frontend :$mod has output dimension $fe_dim, expected core.full_dim = $N2.")) + fe_dim == N2 || throw( + ArgumentError( + "Frontend :$mod has output dimension $fe_dim, expected core.full_dim = $N2." + ), + ) end new{typeof(frontends), typeof(core)}(frontends, core, N2, n2) end end function AbstractNeuralNetworks.Chain(::MultiModalSymplecticAutoencoder) - throw(ArgumentError( - "A MultiModalSymplecticAutoencoder has no single Chain: build it with NeuralNetwork(arch) and use encode(nn, modality, x), decode(nn, modality, ξ), and translate(nn, from, to, x).")) + throw( + ArgumentError( + "A MultiModalSymplecticAutoencoder has no single Chain: build it with NeuralNetwork(arch) and use encode(nn, modality, x), decode(nn, modality, ξ), and translate(nn, from, to, x)." + ), + ) end struct MultiModalSymplecticLayer{ @@ -87,7 +96,9 @@ struct MultiModalSymplecticLayer{ end function parameterlength(l::MultiModalSymplecticLayer) - fe_len = sum(mod -> parameterlength(getfield(l.frontends, mod)), keys(l.frontends); init = 0) + fe_len = sum( + mod -> parameterlength(getfield(l.frontends, mod)), keys(l.frontends); init = 0 + ) fe_len + parameterlength(l.core_model) end @@ -100,7 +111,7 @@ end function _init_frontend_params( rng::AbstractRNG, initializer::AbstractNeuralNetworks.Initializer, fe::Chain, backend::NeuralNetworkBackend, ::Type{T}) where {T} - initialparameters(rng, initializer, fe, backend, T) + params(initialparameters(rng, initializer, fe, backend, T)) end function NeuralNetwork( @@ -194,6 +205,21 @@ function _narrow_core_vec(z_vec, ::AbstractArray{<:Any, 3}, ::Type{T}) where {T} end _narrow_core_vec(z_vec, ::Any, ::Type{T}) where {T} = z_vec +function _narrow_fe_vec(z_vec, ::Any, x::NamedTuple, ::Type{T}) where {T} + _narrow_core_vec(z_vec, x, T) +end +_narrow_fe_vec(z_vec, ::Any, ::AbstractVector, ::Type{T}) where {T} = z_vec::Vector{T} +function _narrow_fe_vec( + z_vec, ::BilinearStiefelLayer, ::AbstractMatrix, ::Type{T} +) where {T} + z_vec::Vector{T} +end +_narrow_fe_vec(z_vec, ::Any, ::AbstractMatrix, ::Type{T}) where {T} = z_vec::Matrix{T} +function _narrow_fe_vec(z_vec, ::Any, ::AbstractArray{<:Any, 3}, ::Type{T}) where {T} + z_vec::Matrix{T} +end +_narrow_fe_vec(z_vec, ::Any, ::Any, ::Type{T}) where {T} = z_vec + function _core_vec_to_fe_inv_input( ::Union{BilinearStiefelLayer, LowRankStiefelLayer}, z_mid::AbstractVector, full_dim::Int) @@ -213,6 +239,12 @@ function _core_vec_to_fe_inv_input( z_mid end +_apply_frontend(fe::AbstractExplicitLayer, x, ps_fe::NamedTuple) = fe(x, ps_fe) +_apply_frontend(fe::Chain, x::AbstractVecOrMat, ps_fe::NamedTuple) = fe(x, ps_fe) +function _apply_frontend(fe::Chain, x::NamedTuple{(:q, :p)}, ps_fe::NamedTuple) + fe(_fe_output_to_core_vec(fe, x), ps_fe) +end + function _apply_inv_frontend( ::Chain, inv_fe::Chain, z_in, ps_fe::NamedTuple) inv_fe(z_in, _reverse_params(ps_fe)) @@ -225,13 +257,15 @@ end function _core_encode_vec( ::Union{WeightTiedAutoencoderLayer, SymmetricAutoencoderLayer}, - core_encoder::Chain, ps_core::NamedTuple, z_vec::AbstractVecOrMat) + core_encoder::Union{Chain, AbstractExplicitLayer}, + ps_core::NamedTuple, z_vec::AbstractVecOrMat) core_encoder(z_vec, _reverse_params(ps_core)) end function _core_decode_vec( ::Union{WeightTiedAutoencoderLayer, SymmetricAutoencoderLayer}, - core_decoder::Chain, ps_core::NamedTuple, ξ_vec::AbstractVecOrMat) + core_decoder::Union{Chain, AbstractExplicitLayer}, + ps_core::NamedTuple, ξ_vec::AbstractVecOrMat) core_decoder(ξ_vec, ps_core) end @@ -259,15 +293,18 @@ _latent_to_vec(ξ::AbstractVecOrMat) = ξ _latent_to_vec(ξ::NamedTuple) = vcat(ξ.q, ξ.p) function _encode_active( - fe, ps_fe::NamedTuple, core_model, core_encoder::Chain, ps_core::NamedTuple, x) - z_fe = fe(x, ps_fe) + fe, ps_fe::NamedTuple, core_model, + core_encoder::Union{Chain, AbstractExplicitLayer}, ps_core::NamedTuple, x) + z_fe = _apply_frontend(fe, x, ps_fe) T = promote_type(_leaf_eltype(x), _leaf_eltype(ps_core)) - z_vec = _narrow_core_vec(_fe_output_to_core_vec(fe, z_fe), x, T) - _core_encode_vec(core_model, core_encoder, ps_core, z_vec) + z_vec = _narrow_fe_vec(_fe_output_to_core_vec(fe, z_fe), fe, x, T) + ξ = _core_encode_vec(core_model, core_encoder, ps_core, z_vec) + _narrow_fe_vec(ξ, fe, x, T) end function _decode_active( - fe, inv_fe, ps_fe::NamedTuple, core_model, core_decoder::Chain, + fe, inv_fe, ps_fe::NamedTuple, core_model, + core_decoder::Union{Chain, AbstractExplicitLayer}, ps_core::NamedTuple, full_dim::Int, ξ) ξ_vec = _latent_to_vec(ξ) z_mid = _core_decode_vec(core_model, core_decoder, ps_core, ξ_vec) @@ -284,8 +321,10 @@ _resolve_param_grad(g::AbstractArray, ::AbstractArray) = g function _resolve_param_grad( g::Union{ChainRulesCore.Tangent, NamedTuple}, p_ref::NamedTuple{names}) where {names} - NamedTuple{names}(map( - k -> _resolve_param_grad(getproperty(g, k), getfield(p_ref, k)), names)) + NamedTuple{names}( + map( + k -> _resolve_param_grad(getproperty(g, k), getfield(p_ref, k)), names), + ) end @doc raw""" @@ -301,19 +340,23 @@ Base.@constprop :aggressive function encode( _check_modality(model.frontends, modality) fe = getfield(model.frontends, modality) ps_fe = getfield(ps.frontends, modality) - _encode_active(fe, ps_fe, model.core_model, model.core_encoder, ps.core, x) + T = promote_type(_leaf_eltype(x), _leaf_eltype(ps.core)) + ξ = _encode_active(fe, ps_fe, model.core_model, model.core_encoder, ps.core, x) + _narrow_core_vec(ξ, x, T) end @noinline function _rrule_encode_active( config::ChainRulesCore.RuleConfig{>:ChainRulesCore.HasReverseMode}, - fe, ps_fe::NamedTuple, core_model, core_encoder::Chain, ps_core::NamedTuple, x) + fe, ps_fe::NamedTuple, core_model, + core_encoder::Union{Chain, AbstractExplicitLayer}, ps_core::NamedTuple, x) ChainRulesCore.rrule_via_ad( config, _encode_active, fe, ps_fe, core_model, core_encoder, ps_core, x) end @noinline function _rrule_decode_active( config::ChainRulesCore.RuleConfig{>:ChainRulesCore.HasReverseMode}, - fe, inv_fe, ps_fe::NamedTuple, core_model, core_decoder::Chain, + fe, inv_fe, ps_fe::NamedTuple, core_model, + core_decoder::Union{Chain, AbstractExplicitLayer}, ps_core::NamedTuple, full_dim::Int, ξ) ChainRulesCore.rrule_via_ad( config, _decode_active, fe, inv_fe, ps_fe, core_model, core_decoder, @@ -335,12 +378,14 @@ function ChainRulesCore.rrule( function encode_pullback(Δy) _, _, d_ps_fe, _, _, d_ps_core, d_x = pb(ChainRulesCore.unthunk(Δy)) fe_keys = keys(ps.frontends) - d_frontends = NamedTuple{fe_keys}(map( + d_frontends = NamedTuple{fe_keys}( + map( k -> k === modality ? _resolve_param_grad( ChainRulesCore.unthunk(d_ps_fe), getfield(ps.frontends, k)) : _zero_param_tree(getfield(ps.frontends, k)), - fe_keys)) + fe_keys), + ) d_core = _resolve_param_grad(ChainRulesCore.unthunk(d_ps_core), ps.core) d_ps = (frontends = d_frontends, core = d_core) return ( @@ -400,12 +445,14 @@ function ChainRulesCore.rrule( function decode_pullback(Δy) _, _, _, d_ps_fe, _, _, d_ps_core, _, d_ξ = pb(ChainRulesCore.unthunk(Δy)) fe_keys = keys(ps.frontends) - d_frontends = NamedTuple{fe_keys}(map( + d_frontends = NamedTuple{fe_keys}( + map( k -> k === modality ? _resolve_param_grad( ChainRulesCore.unthunk(d_ps_fe), getfield(ps.frontends, k)) : _zero_param_tree(getfield(ps.frontends, k)), - fe_keys)) + fe_keys), + ) d_core = _resolve_param_grad(ChainRulesCore.unthunk(d_ps_core), ps.core) d_ps = (frontends = d_frontends, core = d_core) return ( diff --git a/src/reduced_system/reduced_system.jl b/src/reduced_system/reduced_system.jl index cbc9afa8e..cb166a186 100644 --- a/src/reduced_system/reduced_system.jl +++ b/src/reduced_system/reduced_system.jl @@ -29,8 +29,10 @@ end function _check_pullback_mode(pullback_mode::Symbol) (pullback_mode === :jacobian || pullback_mode === :vjp) || - throw(ArgumentError( - "pullback_mode must be :jacobian or :vjp, got :$pullback_mode.")) + throw( + ArgumentError( + "pullback_mode must be :jacobian or :vjp, got :$pullback_mode."), + ) nothing end @@ -146,8 +148,11 @@ function HRedSys(N::Integer, n::Integer, encoder::NeuralNetwork{<:SymplecticEnco _check_pullback_mode(pullback_mode) N_ics = length(ics.q) (length(q_ref) == N_ics && length(p_ref) == N_ics) || - throw(DimensionMismatch( - "q_ref and p_ref must have length N = $N_ics, got $(length(q_ref)) and $(length(p_ref)).")) + throw( + DimensionMismatch( + "q_ref and p_ref must have length N = $N_ics, got $(length(q_ref)) and $(length(p_ref))." + ), + ) T_state = eltype(ics.q) q_ref_vec = Vector{T_state}(q_ref) p_ref_vec = Vector{T_state}(p_ref) @@ -355,7 +360,7 @@ function build_h_reduced( h_full, decoder::NeuralNetwork{<:SymplecticDecoder}; q_ref = nothing, p_ref = nothing) function h_reduced(t, q, p, params) - dec = decoder(q, p) + dec = decoder((q = q, p = p)) if q_ref === nothing || p_ref === nothing h_full(t, dec.q, dec.p, params) else diff --git a/test/architectures/multimodal_symplectic_autoencoder.jl b/test/architectures/multimodal_symplectic_autoencoder.jl index 5065c90f4..c6c8af048 100644 --- a/test/architectures/multimodal_symplectic_autoencoder.jl +++ b/test/architectures/multimodal_symplectic_autoencoder.jl @@ -1,12 +1,13 @@ using GeometricMachineLearning -using GeometricMachineLearning: SymplecticCompression, AutoEncoder, Architecture, +using GeometricMachineLearning: + SymplecticCompression, AutoEncoder, Architecture, NetworkLoss, params using GeometricMachineLearning.AbstractNeuralNetworks: GlorotUniform using KernelAbstractions using LinearAlgebra using Test -import Random -import Zygote +using Random: Random +using Zygote: Zygote include("../helpers/eltypes.jl") @@ -381,7 +382,9 @@ end @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)) + loss_cyc = CrossModalCycleLoss( + :particles, :grid; β_align = T(1.2), β_cycle = 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) @test eltype(val_cyc) === T && val_cyc > zero(T) @@ -399,8 +402,10 @@ end 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) - @test isapprox(val_cyc_rev, expected_cyc_rev; atol = 10 * eps(T), rtol = 10 * - eps(T)) + @test isapprox( + val_cyc_rev, expected_cyc_rev; atol = 10 * eps(T), rtol = 10 * + eps(T) + ) end @testset "Analytical Zygote gradients match central finite differences in Float64" begin @@ -460,4 +465,135 @@ end @test isapprox(g_vec, g_fd; atol = tol, rtol = tol) end end + + @testset "Chain particle frontends, SymmetricSymplecticAutoencoder core, and batched @inferred" begin + rng = Random.MersenneTwister(404) + T = Float64 + nx, nv = 6, 6 + mx, mv = 2, 1 + + # 1. Exact-inverse Chain frontend (PairwiseSymplecticLayerQ + PairwiseSymplecticLayerP) + fe_pw_chain = Chain( + PairwiseSymplecticLayerQ(2, 1; hidden_dim = 4), + PairwiseSymplecticLayerP(2, 1; hidden_dim = 4) + ) + fe_grid = BilinearStiefelLayer(nx, nv, mx, mv) + fe_dlra = LowRankStiefelLayer(nx, nv, 2, 1) + core_id = SymplecticAutoencoder( + 4, 4; n_encoder_layers = 1, n_decoder_layers = 1, upscale = 2) + arch_pw = MultiModalSymplecticAutoencoder( + (particle = fe_pw_chain, grid = fe_grid, dlra = fe_dlra), core_id) + nn_pw = NeuralNetwork(arch_pw, CPU(), T; rng = rng) + for (_, lyr_p) in pairs(nn_pw.params.params.core) + for (_, arr) in pairs(lyr_p) + fill!(arr, zero(T)) + end + end + + @test params(nn_pw).frontends.particle isa NamedTuple + x_qp2 = (q = randn(rng, T, 1, 2), p = randn(rng, T, 1, 2)) + x_vec2 = vcat(vec(x_qp2.q), vec(x_qp2.p)) + z_pw_qp = @inferred encode(nn_pw, :particle, x_qp2) + z_pw_vec = @inferred encode(nn_pw, :particle, x_vec2) + @test isapprox(z_pw_qp, z_pw_vec; atol = 10 * eps(T), rtol = 10 * eps(T)) + x_pw_rec = decode(nn_pw, :particle, z_pw_qp) + @test isapprox(x_pw_rec, x_vec2; atol = 10 * eps(T), rtol = 10 * eps(T)) + + # 2. Canonical 4-layer Chain frontend (CanonicalSortLayer -> Pairwise Q/P -> CanonicalSlotLayer) + fe_can_chain = Chain( + CanonicalSortLayer(3, 1, 1), + PairwiseSymplecticLayerQ(3, 1; hidden_dim = 4), + PairwiseSymplecticLayerP(3, 1; hidden_dim = 4), + CanonicalSlotLayer(3, 1, 2) + ) + core_wt = WeightTiedSymplecticAutoencoder( + 4, 2; n_blocks = 1, n_layers = 1, n_ambient_layers = 0, sympnet_upscale = 2) + arch_can = MultiModalSymplecticAutoencoder( + (particle = fe_can_chain, grid = fe_grid, dlra = fe_dlra), core_wt) + nn_can = NeuralNetwork(arch_can, CPU(), T; rng = rng) + + @test params(nn_can).frontends.particle isa NamedTuple + x_qp3 = (q = T[-2.0 0.1 2.5], p = T[0.4 -1.3 0.9]) + x_vec3 = vcat(vec(x_qp3.q), vec(x_qp3.p)) + z_can_qp = @inferred encode(nn_can, :particle, x_qp3) + z_can_vec = @inferred encode(nn_can, :particle, x_vec3) + @test isapprox(z_can_qp, z_can_vec; atol = 10 * eps(T), rtol = 10 * eps(T)) + x_can_rec = @inferred _dec_part(nn_can, z_can_qp) + @test length(x_can_rec) == 6 + + # Batched @inferred for encode, decode, translate across all modalities + B = 3 + x_qp3_batch = (q = randn(rng, T, 1, 3, B), p = randn(rng, T, 1, 3, B)) + x_mat3_batch = vcat(reshape(x_qp3_batch.q, 3, B), reshape(x_qp3_batch.p, 3, B)) + F_batch = randn(rng, T, nx, nv, B) + x_dlra = ( + X = Matrix(qr(randn(rng, T, nx, 2)).Q), + S = randn(rng, T, 2, 2), + V = Matrix(qr(randn(rng, T, nv, 2)).Q) + ) + X_b = zeros(T, nx, 2, B) + S_b = randn(rng, T, 2, 2, B) + V_b = zeros(T, nv, 2, B) + for b in 1:B + X_b[:, :, b] .= Matrix(qr(randn(rng, T, nx, 2)).Q) + V_b[:, :, b] .= Matrix(qr(randn(rng, T, nv, 2)).Q) + end + x_dlra_batch = (X = X_b, S = S_b, V = V_b) + + z_b_qp = @inferred encode(nn_can, :particle, x_qp3_batch) + z_b_mat = @inferred _enc_part(nn_can, x_mat3_batch) + @test size(z_b_qp) == (2, B) + @test isapprox(z_b_qp, z_b_mat; atol = 10 * eps(T), rtol = 10 * eps(T)) + @test size(@inferred _dec_part(nn_can, z_b_qp)) == (6, B) + + z_b_grid = @inferred encode(nn_can, :grid, F_batch) + @test size(z_b_grid) == (2, B) + @test size(@inferred _dec_grid(nn_can, z_b_grid)) == (nx, nv, B) + + z_dlra_1 = @inferred encode(nn_can, :dlra, x_dlra) + @test length(z_dlra_1) == 2 + z_b_dlra = @inferred encode(nn_can, :dlra, x_dlra_batch) + @test size(z_b_dlra) == (2, B) + dec_b_dlra = @inferred _dec_dlra(nn_can, z_b_dlra) + @test size(dec_b_dlra.X) == (nx, 2, B) && size(dec_b_dlra.S) == (2, 2, B) && + size(dec_b_dlra.V) == (nv, 2, B) + + F_from_p_b = @inferred _trans_p2g(nn_can, x_mat3_batch) + @test size(F_from_p_b) == (nx, nv, B) + + # Loss + Zygote gradients on :particle ((q, p) NamedTuple on Chain) and :dlra ((X, S, V) NamedTuple) + loss_c = CharbonnierRelativeLoss(; ε = 1e-2, λ_lat = 5e-2) + val_c_part = @inferred loss_c(nn_can, :particle, x_qp3) + @test isfinite(val_c_part) && val_c_part > zero(T) + g_part = Zygote.gradient( + ps -> loss_c(nn_can.model, ps, :particle, x_qp3), nn_can.params)[1] + @test all(isfinite, _extract_grad_vec(g_part)) + + val_c_dlra = @inferred loss_c(nn_can, :dlra, x_dlra) + @test isfinite(val_c_dlra) && val_c_dlra > zero(T) + g_dlra = Zygote.gradient( + ps -> loss_c(nn_can.model, ps, :dlra, x_dlra), nn_can.params)[1] + @test all(isfinite, _extract_grad_vec(g_dlra)) + + # 3. SymmetricSymplecticAutoencoder as core + core_sym = SymmetricSymplecticAutoencoder( + 4, 2, 1; n_Encoder = 1, n_Decoder = 1, n_sympnet = 1, + hidden_dim = 4, sympnet_upscale = 2) + arch_sym = MultiModalSymplecticAutoencoder( + (particle = fe_can_chain, grid = fe_grid), core_sym) + nn_sym = NeuralNetwork(arch_sym, CPU(), T; rng = rng) + + F_single = randn(rng, T, nx, nv) + z_sym_p = @inferred encode(nn_sym, :particle, x_qp3) + z_sym_g = @inferred _enc_grid(nn_sym, F_single) + @test length(z_sym_p) == 2 && length(z_sym_g) == 2 + F_sym_rec = @inferred _dec_grid(nn_sym, z_sym_g) + @test size(F_sym_rec) == (nx, nv) + F_sym_tr = @inferred _trans_p2g(nn_sym, x_qp3) + @test size(F_sym_tr) == (nx, nv) + + g_sym = Zygote.gradient( + ps -> loss_c(nn_sym.model, ps, :grid, F_single), nn_sym.params)[1] + @test all(isfinite, _extract_grad_vec(g_sym)) + end end diff --git a/test/reduced_system/reference_shift_and_vjp.jl b/test/reduced_system/reference_shift_and_vjp.jl index 56ac7d574..f63d0458f 100644 --- a/test/reduced_system/reference_shift_and_vjp.jl +++ b/test/reduced_system/reference_shift_and_vjp.jl @@ -6,7 +6,7 @@ using GeometricIntegrators: ImplicitMidpoint using KernelAbstractions using LinearAlgebra using Test -import Random +using Random: Random include("../helpers/eltypes.jl") @@ -249,10 +249,37 @@ Random.seed!(42) ) sol_red_neq = integrate_reduced_system(rs_neq) @test norm(Vector(sol_red_neq.q[1]) - Vector(sol_red_neq.q[0])) > T(1e-4) - lifted_0 = decode(rs_neq, ( - q = Vector(sol_red_neq.q[0]), p = Vector(sol_red_neq.p[0]))) - lifted_1 = decode(rs_neq, ( - q = Vector(sol_red_neq.q[1]), p = Vector(sol_red_neq.p[1]))) + lifted_0 = decode( + rs_neq, ( + q = Vector(sol_red_neq.q[0]), p = Vector(sol_red_neq.p[0])) + ) + lifted_1 = decode( + rs_neq, ( + q = Vector(sol_red_neq.q[1]), p = Vector(sol_red_neq.p[1])) + ) @test norm(lifted_1.q - lifted_0.q) > T(1e-4) + + # Verify build_h_reduced evaluates decoder((q = q̃, p = p̃)) with and without reference shift + q̃_test = Vector(sol_red_neq.q[1]) + p̃_test = Vector(sol_red_neq.p[1]) + @test isapprox( + rs_neq.h_reduced(zero(T), q̃_test, p̃_test, nothing), + h_quad(zero(T), lifted_1.q, lifted_1.p, nothing); + atol = 10 * eps(T), rtol = 10 * eps(T) + ) + rs_unshifted = HRedSys( + full_dim, reduced_dim, encoder(nn), decoder(nn), + v_quad!, f_quad!, h_quad, + (zero(T), T(0.2)), T(0.05), + (q = StateVariable(q_init), p = StateVariable(p_init)); + integrator = ImplicitMidpoint(), + pullback_mode = :jacobian + ) + dec_unshifted = decoder(nn)((q = q̃_test, p = p̃_test)) + @test isapprox( + rs_unshifted.h_reduced(zero(T), q̃_test, p̃_test, nothing), + h_quad(zero(T), dec_unshifted.q, dec_unshifted.p, nothing); + atol = 10 * eps(T), rtol = 10 * eps(T) + ) end end From a436a12fbf4bcf3f5d6fdb09754d6e555a486f3e Mon Sep 17 00:00:00 2001 From: Michael Kraus Date: Fri, 30 Oct 2026 06:53:54 +0100 Subject: [PATCH 3/8] Redesign Chain frontend evaluation, single-slot core dispatch, Val modality methods, and HRedSys defaults --- CHANGELOG.md | 2 +- .../multimodal_symplectic_autoencoder.jl | 186 +++++++++++----- src/reduced_system/reduced_system.jl | 33 ++- .../multimodal_symplectic_autoencoder.jl | 200 +++++++++--------- .../reduced_system/reference_shift_and_vjp.jl | 40 ++++ 5 files changed, 295 insertions(+), 166 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index c1b111662..1fd3a878b 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -150,7 +150,7 @@ breaking release). - **`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`. -- **`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 whose canonical output dimension matches `core.full_dim` with a shared `SymplecticCompression` core; `encode(nn, modality, x)` 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``. `HRedSys`, `build_v_reduced`, and `build_f_reduced` accept optional reference states `q_ref` and `p_ref` (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 `:jacobian` uses preallocated buffers (`ξ_buf`, `v_intermediate`, `f_intermediate`, `jac_result`, and a per-element-type `IdDict{DataType, Any}` cache for `ForwardDiff.Dual` Newton states). 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 `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``. `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 `:jacobian` uses preallocated buffers (`ξ_buf`, `v_intermediate`, `f_intermediate`, `jac_result`, and a per-element-type `IdDict{DataType, Any}` cache for `ForwardDiff.Dual` Newton states). The tests are `test/architectures/multimodal_symplectic_autoencoder.jl` and `test/reduced_system/reference_shift_and_vjp.jl`. ### Changed diff --git a/src/architectures/multimodal_symplectic_autoencoder.jl b/src/architectures/multimodal_symplectic_autoencoder.jl index ff5076517..ec47983b9 100644 --- a/src/architectures/multimodal_symplectic_autoencoder.jl +++ b/src/architectures/multimodal_symplectic_autoencoder.jl @@ -205,6 +205,9 @@ function _narrow_core_vec(z_vec, ::AbstractArray{<:Any, 3}, ::Type{T}) where {T} end _narrow_core_vec(z_vec, ::Any, ::Type{T}) where {T} = z_vec +_narrow_latent_decode(z_mid, ::AbstractVector, ::Type{T}) where {T} = z_mid::Vector{T} +_narrow_latent_decode(z_mid, ::AbstractMatrix, ::Type{T}) where {T} = z_mid::Matrix{T} + function _narrow_fe_vec(z_vec, ::Any, x::NamedTuple, ::Type{T}) where {T} _narrow_core_vec(z_vec, x, T) end @@ -239,62 +242,85 @@ function _core_vec_to_fe_inv_input( z_mid end -_apply_frontend(fe::AbstractExplicitLayer, x, ps_fe::NamedTuple) = fe(x, ps_fe) -_apply_frontend(fe::Chain, x::AbstractVecOrMat, ps_fe::NamedTuple) = fe(x, ps_fe) -function _apply_frontend(fe::Chain, x::NamedTuple{(:q, :p)}, ps_fe::NamedTuple) - fe(_fe_output_to_core_vec(fe, x), ps_fe) +@inline _strip_perm(z::NamedTuple{(:q, :p, :perm)}) = (q = z.q, p = z.p) +@inline _strip_perm(z) = z + +@inline _apply_chain_layers(::Tuple{}, z, ::Tuple) = z +@inline function _apply_chain_layers(layers::Tuple, z, ps::Tuple) + z_next = _strip_perm(first(layers)(z, first(ps))) + return _apply_chain_layers(Base.tail(layers), z_next, Base.tail(ps)) +end +@inline function _apply_chain_layers(layers::Tuple, z, ps::NamedTuple) + _apply_chain_layers(layers, z, values(ps)) +end + +@inline _drop_encoder_params(::Tuple{}, ps::Tuple) = ps +@inline function _drop_encoder_params(enc_layers::Tuple, ps::Tuple) + _drop_encoder_params(Base.tail(enc_layers), Base.tail(ps)) +end +@inline function _drop_encoder_params(enc_layers::Tuple, ps::NamedTuple) + _drop_encoder_params(enc_layers, values(ps)) +end +@inline function _drop_encoder_params( + ::NTuple{N, Any}, ::NTuple{N, Any}, ps::Tuple) where {N} + ps +end +@inline function _drop_encoder_params( + full_layers::Tuple, dec_layers::Tuple, ps::Tuple) + _drop_encoder_params(Base.tail(full_layers), dec_layers, Base.tail(ps)) +end +@inline function _drop_encoder_params( + full_layers::Tuple, dec_layers::Tuple, ps::NamedTuple) + _drop_encoder_params(full_layers, dec_layers, values(ps)) +end + +@inline _apply_frontend(fe::AbstractExplicitLayer, x, ps_fe::NamedTuple) = fe(x, ps_fe) +@inline function _apply_frontend(fe::Chain, x::NamedTuple{(:q, :p)}, ps_fe::NamedTuple) + _apply_chain_layers(fe.layers, x, ps_fe) +end +@inline function _apply_frontend(fe::Chain, x::AbstractArray, ps_fe::NamedTuple) + _apply_chain_layers(fe.layers, x, ps_fe) end -function _apply_inv_frontend( +@inline function _apply_inv_frontend( ::Chain, inv_fe::Chain, z_in, ps_fe::NamedTuple) - inv_fe(z_in, _reverse_params(ps_fe)) + _apply_chain_layers(inv_fe.layers, z_in, _reverse_params(ps_fe)) end -function _apply_inv_frontend( +@inline function _apply_inv_frontend( ::AbstractExplicitLayer, inv_fe::AbstractExplicitLayer, z_in, ps_fe::NamedTuple) inv_fe(z_in, ps_fe) end -function _core_encode_vec( +@inline function _core_encode_vec( ::Union{WeightTiedAutoencoderLayer, SymmetricAutoencoderLayer}, - core_encoder::Union{Chain, AbstractExplicitLayer}, - ps_core::NamedTuple, z_vec::AbstractVecOrMat) + core_encoder, ps_core::NamedTuple, z_vec::AbstractVecOrMat) core_encoder(z_vec, _reverse_params(ps_core)) end -function _core_decode_vec( +@inline function _core_decode_vec( ::Union{WeightTiedAutoencoderLayer, SymmetricAutoencoderLayer}, - core_decoder::Union{Chain, AbstractExplicitLayer}, - ps_core::NamedTuple, ξ_vec::AbstractVecOrMat) + core_decoder, ps_core::NamedTuple, ξ_vec::AbstractVecOrMat) core_decoder(ξ_vec, ps_core) end -function _core_encode_vec( - ::Any, core_encoder::Chain, ps_core::NamedTuple, z_vec::AbstractVecOrMat) - names = keys(ps_core) - M = length(core_encoder.layers) - enc_names = names[1:M] - ps_enc = NamedTuple{enc_names}(map(k -> getfield(ps_core, k), enc_names)) - core_encoder(z_vec, ps_enc) +@inline function _core_encode_vec( + ::Chain, core_encoder, ps_core::NamedTuple, z_vec::AbstractVecOrMat) + _apply_chain_layers(core_encoder.layers, z_vec, ps_core) end -function _core_decode_vec( - ::Any, core_decoder::Chain, ps_core::NamedTuple, ξ_vec::AbstractVecOrMat) - names = keys(ps_core) - K = length(core_decoder.layers) - M = length(names) - K - dec_names = names[1:K] - src_names = names[(M + 1):end] - ps_dec = NamedTuple{dec_names}(map(k -> getfield(ps_core, k), src_names)) - core_decoder(ξ_vec, ps_dec) +@inline function _core_decode_vec( + core_model::Chain, core_decoder, ps_core::NamedTuple, ξ_vec::AbstractVecOrMat) + ps_dec = _drop_encoder_params(core_model.layers, core_decoder.layers, ps_core) + _apply_chain_layers(core_decoder.layers, ξ_vec, ps_dec) end _latent_to_vec(ξ::AbstractVecOrMat) = ξ _latent_to_vec(ξ::NamedTuple) = vcat(ξ.q, ξ.p) -function _encode_active( +@inline function _encode_active( fe, ps_fe::NamedTuple, core_model, - core_encoder::Union{Chain, AbstractExplicitLayer}, ps_core::NamedTuple, x) + core_encoder, ps_core::NamedTuple, x) z_fe = _apply_frontend(fe, x, ps_fe) T = promote_type(_leaf_eltype(x), _leaf_eltype(ps_core)) z_vec = _narrow_fe_vec(_fe_output_to_core_vec(fe, z_fe), fe, x, T) @@ -302,12 +328,13 @@ function _encode_active( _narrow_fe_vec(ξ, fe, x, T) end -function _decode_active( +@inline function _decode_active( fe, inv_fe, ps_fe::NamedTuple, core_model, - core_decoder::Union{Chain, AbstractExplicitLayer}, - ps_core::NamedTuple, full_dim::Int, ξ) + core_decoder, ps_core::NamedTuple, full_dim::Int, ξ) ξ_vec = _latent_to_vec(ξ) - z_mid = _core_decode_vec(core_model, core_decoder, ps_core, ξ_vec) + T = promote_type(_leaf_eltype(ξ_vec), _leaf_eltype(ps_core)) + z_mid = _narrow_latent_decode( + _core_decode_vec(core_model, core_decoder, ps_core, ξ_vec), ξ_vec, T) z_in = _core_vec_to_fe_inv_input(fe, z_mid, full_dim) _apply_inv_frontend(fe, inv_fe, z_in, ps_fe) end @@ -329,13 +356,24 @@ end @doc raw""" encode(nn::NeuralNetwork{<:MultiModalSymplecticAutoencoder}, modality::Symbol, x) + encode(nn::NeuralNetwork{<:MultiModalSymplecticAutoencoder}, ::Val{modality}, x) encode(model::MultiModalSymplecticLayer, ps, modality::Symbol, x) + encode(model::MultiModalSymplecticLayer, ps, ::Val{modality}, x) Apply `frontends[modality]` followed by `encoder(core)` to map an input `x` in modality `modality` into the shared canonical latent space ``(\tilde{q}, \tilde{p}) \in \mathbb{R}^{2n}``. Throws an `ArgumentError` naming `modality` and `keys(frontends)` if `modality` is not present in `frontends`. """ -Base.@constprop :aggressive function encode( +@inline Base.@constprop :aggressive function encode( + model::MultiModalSymplecticLayer, ps::NamedTuple, ::Val{modality}, x +) where {modality} + _check_modality(model.frontends, modality) + fe = getfield(model.frontends, modality) + ps_fe = getfield(ps.frontends, modality) + _encode_active(fe, ps_fe, model.core_model, model.core_encoder, ps.core, x) +end + +@inline Base.@constprop :aggressive function encode( model::MultiModalSymplecticLayer, ps::NamedTuple, modality::Symbol, x) _check_modality(model.frontends, modality) fe = getfield(model.frontends, modality) @@ -348,7 +386,7 @@ end @noinline function _rrule_encode_active( config::ChainRulesCore.RuleConfig{>:ChainRulesCore.HasReverseMode}, fe, ps_fe::NamedTuple, core_model, - core_encoder::Union{Chain, AbstractExplicitLayer}, ps_core::NamedTuple, x) + core_encoder, ps_core::NamedTuple, x) ChainRulesCore.rrule_via_ad( config, _encode_active, fe, ps_fe, core_model, core_encoder, ps_core, x) end @@ -356,8 +394,7 @@ end @noinline function _rrule_decode_active( config::ChainRulesCore.RuleConfig{>:ChainRulesCore.HasReverseMode}, fe, inv_fe, ps_fe::NamedTuple, core_model, - core_decoder::Union{Chain, AbstractExplicitLayer}, - ps_core::NamedTuple, full_dim::Int, ξ) + core_decoder, ps_core::NamedTuple, full_dim::Int, ξ) ChainRulesCore.rrule_via_ad( config, _decode_active, fe, inv_fe, ps_fe, core_model, core_decoder, ps_core, full_dim, ξ) @@ -399,25 +436,51 @@ function ChainRulesCore.rrule( return y, encode_pullback end -Base.@constprop :aggressive function encode( - model::MultiModalSymplecticLayer, ps::NetworkParameters, modality::Symbol, x) +function ChainRulesCore.rrule( + config::ChainRulesCore.RuleConfig{>:ChainRulesCore.HasReverseMode}, + ::typeof(encode), + model::MultiModalSymplecticLayer, + ps::NamedTuple, + ::Val{modality}, + x) where {modality} + ChainRulesCore.rrule(config, encode, model, ps, modality, x) +end + +@inline Base.@constprop :aggressive function encode( + model::MultiModalSymplecticLayer, ps::NetworkParameters, + modality::Union{Symbol, Val}, x) encode(model, params(ps), modality, x) end -Base.@constprop :aggressive function encode( - nn::NeuralNetwork{<:MultiModalSymplecticAutoencoder}, modality::Symbol, x) +@inline Base.@constprop :aggressive function encode( + nn::NeuralNetwork{<:MultiModalSymplecticAutoencoder}, + modality::Union{Symbol, Val}, x) encode(nn.model, params(nn), modality, x) end @doc raw""" decode(nn::NeuralNetwork{<:MultiModalSymplecticAutoencoder}, modality::Symbol, ξ) + decode(nn::NeuralNetwork{<:MultiModalSymplecticAutoencoder}, ::Val{modality}, ξ) decode(model::MultiModalSymplecticLayer, ps, modality::Symbol, ξ) + decode(model::MultiModalSymplecticLayer, ps, ::Val{modality}, ξ) Apply `decoder(core)` followed by `inverse_layer(frontends[modality])` to reconstruct a latent state ``\xi`` in modality `modality`. Throws an `ArgumentError` naming `modality` and `keys(frontends)` if `modality` is not present in `frontends`. """ -Base.@constprop :aggressive function decode( +@inline Base.@constprop :aggressive function decode( + model::MultiModalSymplecticLayer, ps::NamedTuple, ::Val{modality}, ξ +) where {modality} + _check_modality(model.frontends, modality) + fe = getfield(model.frontends, modality) + inv_fe = getfield(model.inverse_frontends, modality) + ps_fe = getfield(ps.frontends, modality) + _decode_active( + fe, inv_fe, ps_fe, model.core_model, model.core_decoder, + ps.core, model.full_dim, ξ) +end + +@inline Base.@constprop :aggressive function decode( model::MultiModalSymplecticLayer, ps::NamedTuple, modality::Symbol, ξ) _check_modality(model.frontends, modality) fe = getfield(model.frontends, modality) @@ -466,36 +529,51 @@ function ChainRulesCore.rrule( return y, decode_pullback end -Base.@constprop :aggressive function decode( - model::MultiModalSymplecticLayer, ps::NetworkParameters, modality::Symbol, ξ) +function ChainRulesCore.rrule( + config::ChainRulesCore.RuleConfig{>:ChainRulesCore.HasReverseMode}, + ::typeof(decode), + model::MultiModalSymplecticLayer, + ps::NamedTuple, + ::Val{modality}, + ξ) where {modality} + ChainRulesCore.rrule(config, decode, model, ps, modality, ξ) +end + +@inline Base.@constprop :aggressive function decode( + model::MultiModalSymplecticLayer, ps::NetworkParameters, + modality::Union{Symbol, Val}, ξ) decode(model, params(ps), modality, ξ) end -Base.@constprop :aggressive function decode( - nn::NeuralNetwork{<:MultiModalSymplecticAutoencoder}, modality::Symbol, ξ) +@inline Base.@constprop :aggressive function decode( + nn::NeuralNetwork{<:MultiModalSymplecticAutoencoder}, + modality::Union{Symbol, Val}, ξ) decode(nn.model, params(nn), modality, ξ) end @doc raw""" translate(nn::NeuralNetwork{<:MultiModalSymplecticAutoencoder}, from::Symbol, to::Symbol, x) + translate(nn::NeuralNetwork{<:MultiModalSymplecticAutoencoder}, ::Val{from}, ::Val{to}, x) translate(model::MultiModalSymplecticLayer, ps, from::Symbol, to::Symbol, x) + translate(model::MultiModalSymplecticLayer, ps, ::Val{from}, ::Val{to}, x) Map an input `x` in modality `from` into modality `to` through the shared canonical latent space via `decode(nn, to, encode(nn, from, x))`. """ -Base.@constprop :aggressive function translate( - model::MultiModalSymplecticLayer, ps::NamedTuple, from::Symbol, to::Symbol, x) +@inline Base.@constprop :aggressive function translate( + model::MultiModalSymplecticLayer, ps::NamedTuple, + from::Union{Symbol, Val}, to::Union{Symbol, Val}, x) decode(model, ps, to, encode(model, ps, from, x)) end -Base.@constprop :aggressive function translate( +@inline Base.@constprop :aggressive function translate( model::MultiModalSymplecticLayer, ps::NetworkParameters, - from::Symbol, to::Symbol, x) + from::Union{Symbol, Val}, to::Union{Symbol, Val}, x) translate(model, params(ps), from, to, x) end -Base.@constprop :aggressive function translate( +@inline Base.@constprop :aggressive function translate( nn::NeuralNetwork{<:MultiModalSymplecticAutoencoder}, - from::Symbol, to::Symbol, x) + from::Union{Symbol, Val}, to::Union{Symbol, Val}, x) translate(nn.model, params(nn), from, to, x) end diff --git a/src/reduced_system/reduced_system.jl b/src/reduced_system/reduced_system.jl index cb166a186..c029ad251 100644 --- a/src/reduced_system/reduced_system.jl +++ b/src/reduced_system/reduced_system.jl @@ -137,13 +137,9 @@ function HRedSys(N::Integer, n::Integer, encoder::NeuralNetwork{<:SymplecticEnco cache::ReducedVFCache = ReducedVFCache( eltype(ics.q), decoder.architecture.full_dim ÷ 2, decoder.architecture.reduced_dim ÷ 2), - v_reduced = build_v_reduced( - v_full, f_full, decoder; - pullback_mode = pullback_mode, q_ref = q_ref, p_ref = p_ref, cache = cache), - f_reduced = build_f_reduced( - v_full, f_full, decoder; - pullback_mode = pullback_mode, q_ref = q_ref, p_ref = p_ref, cache = cache), - h_reduced = build_h_reduced(h_full, decoder; q_ref = q_ref, p_ref = p_ref), + v_reduced = nothing, + f_reduced = nothing, + h_reduced = nothing, integrator) where {T <: Real} _check_pullback_mode(pullback_mode) N_ics = length(ics.q) @@ -156,12 +152,27 @@ function HRedSys(N::Integer, n::Integer, encoder::NeuralNetwork{<:SymplecticEnco T_state = eltype(ics.q) q_ref_vec = Vector{T_state}(q_ref) p_ref_vec = Vector{T_state}(p_ref) + v_red = isnothing(v_reduced) ? + build_v_reduced( + v_full, f_full, decoder; + pullback_mode = pullback_mode, q_ref = q_ref_vec, p_ref = p_ref_vec, + cache = cache) : + v_reduced + f_red = isnothing(f_reduced) ? + build_f_reduced( + v_full, f_full, decoder; + pullback_mode = pullback_mode, q_ref = q_ref_vec, p_ref = p_ref_vec, + cache = cache) : + f_reduced + h_red = isnothing(h_reduced) ? + build_h_reduced(h_full, decoder; q_ref = q_ref_vec, p_ref = p_ref_vec) : + h_reduced HRedSys{typeof(timestep), typeof(encoder), typeof(decoder), typeof(v_full), - typeof(f_full), typeof(h_full), typeof(v_reduced), typeof(f_reduced), - typeof(h_reduced), typeof(integrator), typeof(parameters), typeof(ics), + typeof(f_full), typeof(h_full), typeof(v_red), typeof(f_red), + typeof(h_red), typeof(integrator), typeof(parameters), typeof(ics), typeof(timespan), typeof(q_ref_vec), typeof(cache.jac_result)}( - N, n, encoder, decoder, v_full, f_full, h_full, v_reduced, f_reduced, - h_reduced, integrator, parameters, timespan, timestep, ics, + N, n, encoder, decoder, v_full, f_full, h_full, v_red, f_red, + h_red, integrator, parameters, timespan, timestep, ics, q_ref_vec, p_ref_vec, pullback_mode, cache.ξ_buf, cache.v_intermediate, cache.f_intermediate, cache.jac_result, cache.dual_cache) diff --git a/test/architectures/multimodal_symplectic_autoencoder.jl b/test/architectures/multimodal_symplectic_autoencoder.jl index c6c8af048..4e32a0a79 100644 --- a/test/architectures/multimodal_symplectic_autoencoder.jl +++ b/test/architectures/multimodal_symplectic_autoencoder.jl @@ -159,14 +159,13 @@ end S_in = randn(rng, T, r_in, r_in) x_dlra = (X = X_in, S = S_in, V = V_in) - _enc_grid(net, x) = encode(net, :grid, x) - _enc_spline(net, x) = encode(net, :spline, x) - _dec_grid(net, ξ) = decode(net, :grid, ξ) - ξ_part = @inferred encode(nn, :particles, x_part) - ξ_grid = @inferred _enc_grid(nn, F_grid) - ξ_spline = @inferred _enc_spline(nn, F_spline) + ξ_grid = @inferred encode(nn, Val(:grid), F_grid) + ξ_spline = @inferred encode(nn, Val(:spline), F_spline) ξ_dlra = @inferred encode(nn, :dlra, x_dlra) + @test isapprox(ξ_grid, encode(nn, :grid, F_grid); atol = 8 * eps(T), rtol = 8 * eps(T)) + @test isapprox( + ξ_spline, encode(nn, :spline, F_spline); atol = 8 * eps(T), rtol = 8 * eps(T)) @test eltype(ξ_part) === T && length(ξ_part) == reduced_dim @test eltype(ξ_grid) === T && length(ξ_grid) == reduced_dim @@ -184,7 +183,9 @@ end (q = z_dec_manual[1:N], p = z_dec_manual[(N + 1):full_dim]), params(nn).frontends.grid ) - F_grid_dec = @inferred _dec_grid(nn, ξ_grid) + F_grid_dec = @inferred decode(nn, Val(:grid), ξ_grid) + @test isapprox( + F_grid_dec, decode(nn, :grid, ξ_grid); atol = 8 * eps(T), rtol = 8 * eps(T)) @test isapprox(F_grid_dec, F_grid_manual; atol = 8 * eps(T), rtol = 8 * eps(T)) # NamedTuple (q, p) latent input to decode @@ -196,6 +197,8 @@ end # translate(nn, from, to, x) == decode(nn, to, encode(nn, from, x)) @test translate(nn, :particles, :grid, x_part) == decode(nn, :grid, encode(nn, :particles, x_part)) + @test (@inferred translate(nn, Val(:particles), Val(:grid), x_part)) == + translate(nn, :particles, :grid, x_part) @test translate(nn, :grid, :dlra, F_grid) == decode(nn, :dlra, encode(nn, :grid, F_grid)) @test translate(nn, :dlra, :spline, x_dlra) == @@ -247,11 +250,8 @@ end Φ_rel_mat ./= norm(Φ_rel_mat) copyto!(parent(params(nn_ch).frontends.slot.Φ_rel), Φ_rel_mat) - _dec_slot(net, ξ) = decode(net, :slot, ξ) - _dec_part(net, ξ) = decode(net, :particles, ξ) - ξ_slot_test = T[0.8, -0.5] - x_slot_rec = @inferred _dec_slot(nn_ch, ξ_slot_test) + x_slot_rec = @inferred decode(nn_ch, :slot, ξ_slot_test) ξ_slot_rt = @inferred encode(nn_ch, :slot, x_slot_rec) @test norm(ξ_slot_rt - ξ_slot_test) <= 50 * eps(T) * norm(ξ_slot_test) @@ -260,7 +260,7 @@ end Φ_core_sort = T[-inv(sqrt(T(2))), zero(T), inv(sqrt(T(2))), zero(T)] copyto!(parent(first(values(params(nn_ch).core)).weight), reshape(Φ_core_sort, N, 1)) ξ_sort_test = T[0.9, -0.4] - x_sort_rec = @inferred _dec_part(nn_ch, ξ_sort_test) + x_sort_rec = @inferred decode(nn_ch, :particles, ξ_sort_test) ξ_sort_rt = @inferred encode(nn_ch, :particles, x_sort_rec) @test norm(ξ_sort_rt - ξ_sort_test) <= 50 * eps(T) * norm(ξ_sort_test) end @@ -469,131 +469,131 @@ end @testset "Chain particle frontends, SymmetricSymplecticAutoencoder core, and batched @inferred" begin rng = Random.MersenneTwister(404) T = Float64 - nx, nv = 6, 6 - mx, mv = 2, 1 - - # 1. Exact-inverse Chain frontend (PairwiseSymplecticLayerQ + PairwiseSymplecticLayerP) + nx, nv, mx, mv, B = 6, 6, 2, 1, 3 + _grad_vec(x::AbstractArray) = vec(Float64.(parent(x))) + _grad_vec(::Nothing) = Float64[] + _grad_vec(nt) = isempty(nt) ? Float64[] : mapreduce(_grad_vec, vcat, values(nt)) fe_pw_chain = Chain( - PairwiseSymplecticLayerQ(2, 1; hidden_dim = 4), - PairwiseSymplecticLayerP(2, 1; hidden_dim = 4) - ) + PairwiseSymplecticLayerQ(2, 1, 4), PairwiseSymplecticLayerP(2, 1, 4)) fe_grid = BilinearStiefelLayer(nx, nv, mx, mv) - fe_dlra = LowRankStiefelLayer(nx, nv, 2, 1) - core_id = SymplecticAutoencoder( - 4, 4; n_encoder_layers = 1, n_decoder_layers = 1, upscale = 2) - arch_pw = MultiModalSymplecticAutoencoder( - (particle = fe_pw_chain, grid = fe_grid, dlra = fe_dlra), core_id) - nn_pw = NeuralNetwork(arch_pw, CPU(), T; rng = rng) - for (_, lyr_p) in pairs(nn_pw.params.params.core) - for (_, arr) in pairs(lyr_p) - fill!(arr, zero(T)) - end - end - + fe_dlra = LowRankStiefelLayer(nx, nv, 2, 1; r_out = 2) + core_id = WeightTiedSymplecticAutoencoder( + 4, 4; n_blocks = 1, n_layers = 1, n_ambient_layers = 0, sympnet_upscale = 2) + nn_pw = NeuralNetwork( + MultiModalSymplecticAutoencoder( + (particle = fe_pw_chain, grid = fe_grid, dlra = fe_dlra), core_id), + CPU(), T; rng = rng) @test params(nn_pw).frontends.particle isa NamedTuple x_qp2 = (q = randn(rng, T, 1, 2), p = randn(rng, T, 1, 2)) x_vec2 = vcat(vec(x_qp2.q), vec(x_qp2.p)) z_pw_qp = @inferred encode(nn_pw, :particle, x_qp2) z_pw_vec = @inferred encode(nn_pw, :particle, x_vec2) @test isapprox(z_pw_qp, z_pw_vec; atol = 10 * eps(T), rtol = 10 * eps(T)) - x_pw_rec = decode(nn_pw, :particle, z_pw_qp) - @test isapprox(x_pw_rec, x_vec2; atol = 10 * eps(T), rtol = 10 * eps(T)) - - # 2. Canonical 4-layer Chain frontend (CanonicalSortLayer -> Pairwise Q/P -> CanonicalSlotLayer) + x_pw_rec = @inferred decode(nn_pw, Val(:particle), z_pw_qp) + @test isapprox(x_pw_rec, x_vec2; atol = 50 * eps(T), rtol = 50 * eps(T)) + core_sa = 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) + nn_sa_part = NeuralNetwork( + MultiModalSymplecticAutoencoder((particle = fe_pw_chain,), core_sa), + CPU(), T; rng = rng) + z_sa_vec = @inferred encode(nn_sa_part, :particle, x_vec2) + @test z_sa_vec isa Vector{T} && length(z_sa_vec) == 2 + @test (@inferred encode(nn_sa_part, :particle, randn(rng, T, 4, 3))) isa Matrix{T} + @test (@inferred decode(nn_sa_part, :particle, z_sa_vec)) isa Vector{T} + @test (@inferred translate(nn_sa_part, :particle, :particle, x_vec2)) isa Vector{T} fe_can_chain = Chain( - CanonicalSortLayer(3, 1, 1), - PairwiseSymplecticLayerQ(3, 1; hidden_dim = 4), - PairwiseSymplecticLayerP(3, 1; hidden_dim = 4), - CanonicalSlotLayer(3, 1, 2) - ) + CanonicalSortLayer(3, 1, 1), PairwiseSymplecticLayerQ(3, 1, 4), + PairwiseSymplecticLayerP(3, 1, 4), CanonicalSlotLayer(3, 1, 2)) core_wt = WeightTiedSymplecticAutoencoder( 4, 2; n_blocks = 1, n_layers = 1, n_ambient_layers = 0, sympnet_upscale = 2) - arch_can = MultiModalSymplecticAutoencoder( - (particle = fe_can_chain, grid = fe_grid, dlra = fe_dlra), core_wt) - nn_can = NeuralNetwork(arch_can, CPU(), T; rng = rng) - + nn_can = NeuralNetwork( + MultiModalSymplecticAutoencoder( + (particle = fe_can_chain, grid = fe_grid, dlra = fe_dlra), core_wt), + CPU(), T; rng = rng) @test params(nn_can).frontends.particle isa NamedTuple x_qp3 = (q = T[-2.0 0.1 2.5], p = T[0.4 -1.3 0.9]) x_vec3 = vcat(vec(x_qp3.q), vec(x_qp3.p)) z_can_qp = @inferred encode(nn_can, :particle, x_qp3) z_can_vec = @inferred encode(nn_can, :particle, x_vec3) @test isapprox(z_can_qp, z_can_vec; atol = 10 * eps(T), rtol = 10 * eps(T)) - x_can_rec = @inferred _dec_part(nn_can, z_can_qp) - @test length(x_can_rec) == 6 - - # Batched @inferred for encode, decode, translate across all modalities - B = 3 + x_can_rec = @inferred decode(nn_can, Val(:particle), z_can_qp) + @test length(x_can_rec) == 6 && + isapprox(x_can_rec, decode(nn_can, :particle, z_can_qp); atol = 10 * eps(T)) + fe_can_d2 = Chain( + CanonicalSortLayer(3, 2, 1), PairwiseSymplecticLayerQ(3, 2, 4), + PairwiseSymplecticLayerP(3, 2, 4), CanonicalSlotLayer(3, 2, 2)) + core_wt_d2 = WeightTiedSymplecticAutoencoder( + 8, 4; n_blocks = 1, n_layers = 1, n_ambient_layers = 0, sympnet_upscale = 2) + nn_d2 = NeuralNetwork( + MultiModalSymplecticAutoencoder((particle = fe_can_d2,), core_wt_d2), + CPU(), T; rng = rng) + x_qp_d2 = ( + q = T[3.0 -1.0 2.0; -2.0 4.0 0.5], p = T[0.2 -0.5 0.8; 1.1 -0.3 0.4]) + x_vec_d2 = vcat(vec(x_qp_d2.q), vec(x_qp_d2.p)) + ps_d2 = params(nn_d2).frontends.particle + s1_d2 = fe_can_d2.layers[1](x_qp_d2, ps_d2.L1) + s2_d2 = fe_can_d2.layers[2]((q = s1_d2.q, p = s1_d2.p), ps_d2.L2) + s4_d2 = fe_can_d2.layers[4](fe_can_d2.layers[3](s2_d2, ps_d2.L3), ps_d2.L4) + z_exp_d2 = encoder(nn_d2)(vcat(vec(s4_d2.q), vec(s4_d2.p))) + tol = 10 * eps(T) + @test isapprox(@inferred(encode(nn_d2, :particle, x_qp_d2)), z_exp_d2; atol = tol) + @test isapprox(@inferred(encode(nn_d2, :particle, x_vec_d2)), z_exp_d2; atol = tol) x_qp3_batch = (q = randn(rng, T, 1, 3, B), p = randn(rng, T, 1, 3, B)) x_mat3_batch = vcat(reshape(x_qp3_batch.q, 3, B), reshape(x_qp3_batch.p, 3, B)) F_batch = randn(rng, T, nx, nv, B) x_dlra = ( X = Matrix(qr(randn(rng, T, nx, 2)).Q), S = randn(rng, T, 2, 2), - V = Matrix(qr(randn(rng, T, nv, 2)).Q) - ) - X_b = zeros(T, nx, 2, B) - S_b = randn(rng, T, 2, 2, B) - V_b = zeros(T, nv, 2, B) + V = Matrix(qr(randn(rng, T, nv, 2)).Q)) + X_b, S_b, V_b = zeros(T, nx, 2, B), randn(rng, T, 2, 2, B), zeros(T, nv, 2, B) for b in 1:B X_b[:, :, b] .= Matrix(qr(randn(rng, T, nx, 2)).Q) V_b[:, :, b] .= Matrix(qr(randn(rng, T, nv, 2)).Q) end x_dlra_batch = (X = X_b, S = S_b, V = V_b) - z_b_qp = @inferred encode(nn_can, :particle, x_qp3_batch) - z_b_mat = @inferred _enc_part(nn_can, x_mat3_batch) - @test size(z_b_qp) == (2, B) - @test isapprox(z_b_qp, z_b_mat; atol = 10 * eps(T), rtol = 10 * eps(T)) - @test size(@inferred _dec_part(nn_can, z_b_qp)) == (6, B) - + z_b_mat = @inferred encode(nn_can, Val(:particle), x_mat3_batch) + @test size(z_b_qp) == (2, B) && isapprox(z_b_qp, z_b_mat; atol = tol) + @test size(@inferred decode(nn_can, Val(:particle), z_b_qp)) == (6, B) z_b_grid = @inferred encode(nn_can, :grid, F_batch) @test size(z_b_grid) == (2, B) - @test size(@inferred _dec_grid(nn_can, z_b_grid)) == (nx, nv, B) - + @test size(@inferred decode(nn_can, Val(:grid), z_b_grid)) == (nx, nv, B) z_dlra_1 = @inferred encode(nn_can, :dlra, x_dlra) - @test length(z_dlra_1) == 2 z_b_dlra = @inferred encode(nn_can, :dlra, x_dlra_batch) - @test size(z_b_dlra) == (2, B) - dec_b_dlra = @inferred _dec_dlra(nn_can, z_b_dlra) - @test size(dec_b_dlra.X) == (nx, 2, B) && size(dec_b_dlra.S) == (2, 2, B) && - size(dec_b_dlra.V) == (nv, 2, B) - - F_from_p_b = @inferred _trans_p2g(nn_can, x_mat3_batch) + dec_b_dlra = @inferred decode(nn_can, Val(:dlra), z_b_dlra) + @test length(z_dlra_1) == 2 && size(z_b_dlra) == (2, B) && + size(dec_b_dlra.X) == (nx, 2, B) + F_from_p_b = @inferred translate(nn_can, Val(:particle), Val(:grid), x_mat3_batch) @test size(F_from_p_b) == (nx, nv, B) - - # Loss + Zygote gradients on :particle ((q, p) NamedTuple on Chain) and :dlra ((X, S, V) NamedTuple) loss_c = CharbonnierRelativeLoss(; ε = 1e-2, λ_lat = 5e-2) - val_c_part = @inferred loss_c(nn_can, :particle, x_qp3) - @test isfinite(val_c_part) && val_c_part > zero(T) - g_part = Zygote.gradient( - ps -> loss_c(nn_can.model, ps, :particle, x_qp3), nn_can.params)[1] - @test all(isfinite, _extract_grad_vec(g_part)) - - val_c_dlra = @inferred loss_c(nn_can, :dlra, x_dlra) - @test isfinite(val_c_dlra) && val_c_dlra > zero(T) - g_dlra = Zygote.gradient( - ps -> loss_c(nn_can.model, ps, :dlra, x_dlra), nn_can.params)[1] - @test all(isfinite, _extract_grad_vec(g_dlra)) - - # 3. SymmetricSymplecticAutoencoder as core + @test (@inferred loss_c(nn_can, :particle, x_qp3)) > zero(T) + @test (@inferred loss_c(nn_can, :dlra, x_dlra)) > zero(T) + g_enc_pw = GeometricMachineLearning.Zygote.gradient( + ps -> sum(encode(nn_pw.model, ps, :particle, x_qp2)), params(nn_pw))[1] + g_dec_pw = GeometricMachineLearning.Zygote.gradient( + ps -> sum(decode(nn_pw.model, ps, :particle, z_pw_qp)), params(nn_pw))[1] + x_dlra_hat = @inferred decode(nn_can, Val(:dlra), z_dlra_1) + g_dlra = GeometricMachineLearning.Zygote.gradient( + ps -> loss_c(nn_can.model, ps, :dlra, x_dlra_hat, x_dlra), params(nn_can))[1] + @test all(isfinite, _grad_vec(g_enc_pw)) && all(isfinite, _grad_vec(g_dec_pw)) && + all(isfinite, _grad_vec(g_dlra)) core_sym = SymmetricSymplecticAutoencoder( - 4, 2, 1; n_Encoder = 1, n_Decoder = 1, n_sympnet = 1, - hidden_dim = 4, sympnet_upscale = 2) - arch_sym = MultiModalSymplecticAutoencoder( - (particle = fe_can_chain, grid = fe_grid), core_sym) - nn_sym = NeuralNetwork(arch_sym, CPU(), T; rng = rng) - + 2, 1, 2; n_pairwise_layers = 1, pairwise_upscale = 2, + n_blocks = 1, n_layers = 1, n_ambient_layers = 0, sympnet_upscale = 2) + nn_sym = NeuralNetwork( + MultiModalSymplecticAutoencoder( + (particle = fe_can_chain, grid = fe_grid), core_sym), + CPU(), T; rng = rng) F_single = randn(rng, T, nx, nv) z_sym_p = @inferred encode(nn_sym, :particle, x_qp3) - z_sym_g = @inferred _enc_grid(nn_sym, F_single) - @test length(z_sym_p) == 2 && length(z_sym_g) == 2 - F_sym_rec = @inferred _dec_grid(nn_sym, z_sym_g) - @test size(F_sym_rec) == (nx, nv) - F_sym_tr = @inferred _trans_p2g(nn_sym, x_qp3) - @test size(F_sym_tr) == (nx, nv) - - g_sym = Zygote.gradient( - ps -> loss_c(nn_sym.model, ps, :grid, F_single), nn_sym.params)[1] - @test all(isfinite, _extract_grad_vec(g_sym)) + z_sym_g = @inferred encode(nn_sym, Val(:grid), F_single) + F_sym_rec = @inferred decode(nn_sym, Val(:grid), z_sym_g) + F_sym_tr = @inferred translate(nn_sym, Val(:particle), Val(:grid), x_qp3) + @test length(z_sym_p) == 2 && size(F_sym_rec) == size(F_sym_tr) == (nx, nv) + g_sym = GeometricMachineLearning.Zygote.gradient( + ps -> sum(encode(nn_sym.model, ps, :grid, F_single)), params(nn_sym))[1] + @test all(isfinite, _grad_vec(g_sym)) end end diff --git a/test/reduced_system/reference_shift_and_vjp.jl b/test/reduced_system/reference_shift_and_vjp.jl index f63d0458f..2516dcb14 100644 --- a/test/reduced_system/reference_shift_and_vjp.jl +++ b/test/reduced_system/reference_shift_and_vjp.jl @@ -281,5 +281,45 @@ Random.seed!(42) h_quad(zero(T), dec_unshifted.q, dec_unshifted.p, nothing); atol = 10 * eps(T), rtol = 10 * eps(T) ) + + if T === Float32 + q_ref_f64 = Float64.(q_ref) + p_ref_f64 = Float64.(p_ref) + h_quad_generic(t, q, p, params) = eltype(q)(0.5) * + (sum(abs2, q .- rs_neq.q_ref) + + sum(abs2, p .- rs_neq.p_ref)) + rs_mixed_ref = HRedSys( + full_dim, reduced_dim, encoder(nn), decoder(nn), + v_quad!, f_quad!, h_quad_generic, + (zero(T), T(0.2)), T(0.05), + (q = StateVariable(q_init), p = StateVariable(p_init)); + integrator = ImplicitMidpoint(), + q_ref = q_ref_f64, + p_ref = p_ref_f64, + pullback_mode = :vjp + ) + @test eltype(rs_mixed_ref.q_ref) === Float32 + @test eltype(rs_mixed_ref.p_ref) === Float32 + h_val = @inferred rs_mixed_ref.h_reduced( + zero(Float32), q̃_test, p̃_test, nothing) + @test h_val isa Float32 + @test isapprox( + h_val, + rs_neq.h_reduced(zero(Float32), q̃_test, p̃_test, nothing); + atol = 10 * eps(Float32), rtol = 10 * eps(Float32) + ) + v_mixed = zeros(Float32, n) + f_mixed = zeros(Float32, n) + rs_mixed_ref.v_reduced(v_mixed, zero(Float32), q̃_test, p̃_test, nothing) + rs_mixed_ref.f_reduced(f_mixed, zero(Float32), q̃_test, p̃_test, nothing) + v_expected = zeros(Float32, n) + f_expected = zeros(Float32, n) + rs_neq.v_reduced(v_expected, zero(Float32), q̃_test, p̃_test, nothing) + rs_neq.f_reduced(f_expected, zero(Float32), q̃_test, p̃_test, nothing) + @test isapprox( + v_mixed, v_expected; atol = 50 * eps(Float32), rtol = 50 * eps(Float32)) + @test isapprox( + f_mixed, f_expected; atol = 50 * eps(Float32), rtol = 50 * eps(Float32)) + end end end From 432509bf5f05491cff042c364d80d3e5725b2315 Mon Sep 17 00:00:00 2001 From: Michael Kraus Date: Fri, 30 Oct 2026 11:07:55 +0100 Subject: [PATCH 4/8] Fix review findings on the multimodal autoencoder and HRedSys cache - ReducedVFCache takes its Jacobian result type as a parameter instead of naming DiffResults' non-public MutableDiffResult, which failed the ExplicitImports check on every release job. - CharbonnierRelativeLoss accepts batched DLRA input (3-D X, S, V factors). - NeuralNetwork(::MultiModalSymplecticAutoencoder, ...) passes its initializer keyword to the core network as well as to the frontends. - Tests for the batched loss and for the initializer reaching the core. Co-Authored-By: Claude Opus 5.5 --- .../multimodal_symplectic_autoencoder.jl | 3 ++- src/loss/multimodal_losses.jl | 9 +++++++ src/reduced_system/reduced_system.jl | 10 ++++---- .../multimodal_symplectic_autoencoder.jl | 24 ++++++++++++++++++- 4 files changed, 39 insertions(+), 7 deletions(-) diff --git a/src/architectures/multimodal_symplectic_autoencoder.jl b/src/architectures/multimodal_symplectic_autoencoder.jl index ec47983b9..21270ad12 100644 --- a/src/architectures/multimodal_symplectic_autoencoder.jl +++ b/src/architectures/multimodal_symplectic_autoencoder.jl @@ -132,7 +132,8 @@ function NeuralNetwork( end inv_frontends = NamedTuple{names}(inv_fe_vals) - core_nn = NeuralNetwork(arch.core, backend, T; rng = rng, kwargs...) + core_nn = NeuralNetwork( + arch.core, backend, T; rng = rng, initializer = initializer, kwargs...) core_enc_nn = encoder(core_nn) core_dec_nn = decoder(core_nn) diff --git a/src/loss/multimodal_losses.jl b/src/loss/multimodal_losses.jl index 06ed63a55..52cda7c9f 100644 --- a/src/loss/multimodal_losses.jl +++ b/src/loss/multimodal_losses.jl @@ -13,6 +13,15 @@ function _frob_sq_norm(x::NamedTuple{(:X, :S, :V)}) T = _state_eltype(x) max(zero(T), tr(G_X * S_GV * transpose(x.S))) end +# A batch of DLRA states, one per slice along axis 3: the sum of the slices' squared norms. +function _frob_sq_norm(x::NamedTuple{(:X, :S, :V), <:NTuple{3, AbstractArray{<:Any, 3}}}) + T = _state_eltype(x) + X, S, V = T.(x.X), T.(x.S), T.(x.V) + G_X = tensor_transpose_tensor_mul(X, X) + G_V = tensor_transpose_tensor_mul(V, V) + M = tensor_tensor_mul(tensor_transpose_tensor_mul(S, G_X), S) + sum(max.(zero(T), sum(M .* permutedims(G_V, (2, 1, 3)); dims = (1, 2)))) +end _frob_sq_diff(x̂::AbstractArray, x::AbstractArray) = sum(abs2, x̂ .- x) function _frob_sq_diff(x̂::NamedTuple{(:q, :p)}, x::NamedTuple{(:q, :p)}) diff --git a/src/reduced_system/reduced_system.jl b/src/reduced_system/reduced_system.jl index c029ad251..a373fe9f8 100644 --- a/src/reduced_system/reduced_system.jl +++ b/src/reduced_system/reduced_system.jl @@ -1,9 +1,9 @@ -struct ReducedVFCache{T} +struct ReducedVFCache{T, JR} ξ_buf::Vector{T} v_intermediate::Vector{T} f_intermediate::Vector{T} z_bar_buf::Vector{T} - jac_result::DiffResults.MutableDiffResult{1, Vector{T}, Tuple{Matrix{T}}} + jac_result::JR dual_cache::IdDict{DataType, Any} end @@ -13,7 +13,7 @@ function ReducedVFCache(::Type{T}, N2::Int, n2::Int) where {T} f_int = zeros(T, N2) z_bar = zeros(T, 2 * N2) jac_res = DiffResults.JacobianResult(zeros(T, 2 * N2), ξ_buf) - ReducedVFCache{T}( + ReducedVFCache{T, typeof(jac_res)}( ξ_buf, v_int, f_int, z_bar, jac_res, IdDict{DataType, Any}()) end @@ -242,14 +242,14 @@ function evaluate_vf_and_compute_∇Ψ(t, q̃::AbstractVector{Tq}, p̃::Abstract copyto!(@view(ξ_buf[1:n2]), q̃) copyto!(@view(ξ_buf[(n2 + 1):(2 * n2)]), p̃) result = ForwardDiff.jacobian!(c.jac_result, decoder, ξ_buf) - z = DiffResults.value(result) + z = DiffResults.value(result)::Vector{T} _add_ref!(z, q_ref, p_ref, N2) q, p = @views z[1:N2], z[(N2 + 1):(2 * N2)] fill!(v_intermediate, zero(T)) fill!(f_intermediate, zero(T)) v_full(v_intermediate, t, q, p, parameters) f_full(f_intermediate, t, q, p, parameters) - v_intermediate, f_intermediate, DiffResults.jacobian(result) + v_intermediate, f_intermediate, DiffResults.jacobian(result)::Matrix{T} end function evaluate_vf_and_vjp(t, q̃::AbstractVector{Tq}, p̃::AbstractVector{Tp}, diff --git a/test/architectures/multimodal_symplectic_autoencoder.jl b/test/architectures/multimodal_symplectic_autoencoder.jl index 4e32a0a79..0fc9192dc 100644 --- a/test/architectures/multimodal_symplectic_autoencoder.jl +++ b/test/architectures/multimodal_symplectic_autoencoder.jl @@ -2,7 +2,8 @@ using GeometricMachineLearning using GeometricMachineLearning: SymplecticCompression, AutoEncoder, Architecture, NetworkLoss, params -using GeometricMachineLearning.AbstractNeuralNetworks: GlorotUniform +using GeometricMachineLearning.AbstractNeuralNetworks: GlorotUniform, Initializer, + initialparameters using KernelAbstractions using LinearAlgebra using Test @@ -13,6 +14,15 @@ include("../helpers/eltypes.jl") Random.seed!(42) +# A Glorot-uniform initializer that counts its calls, to see which networks it reaches. +struct CountingInitializer <: Initializer + calls::Base.RefValue{Int} +end +function (c::CountingInitializer)(rng, x) + c.calls[] += 1 + GlorotUniform()(rng, x) +end + # Helper to flatten and unflatten nested NamedTuple/Array parameters for finite differences function _flatten_params(x::AbstractArray{T}) where {T <: Real} raw = parent(x) @@ -570,6 +580,18 @@ end loss_c = CharbonnierRelativeLoss(; ε = 1e-2, λ_lat = 5e-2) @test (@inferred loss_c(nn_can, :particle, x_qp3)) > zero(T) @test (@inferred loss_c(nn_can, :dlra, x_dlra)) > zero(T) + @test isapprox(GeometricMachineLearning._frob_sq_norm(x_dlra_batch), + sum(b -> sum(abs2, X_b[:, :, b] * S_b[:, :, b] * V_b[:, :, b]'), 1:B); + rtol = 8 * eps(T)) + @test loss_c(nn_can, :dlra, x_dlra_batch) > zero(T) + calls_mm, calls_fe, calls_core = Ref(0), Ref(0), Ref(0) + NeuralNetwork( + MultiModalSymplecticAutoencoder((grid = fe_grid,), core_wt), CPU(), T; + rng = rng, initializer = CountingInitializer(calls_mm)) + initialparameters(rng, CountingInitializer(calls_fe), fe_grid, CPU(), T) + NeuralNetwork( + core_wt, CPU(), T; rng = rng, initializer = CountingInitializer(calls_core)) + @test calls_core[] > 0 && calls_mm[] == calls_fe[] + calls_core[] g_enc_pw = GeometricMachineLearning.Zygote.gradient( ps -> sum(encode(nn_pw.model, ps, :particle, x_qp2)), params(nn_pw))[1] g_dec_pw = GeometricMachineLearning.Zygote.gradient( From 775e13fad7c7858ee75191d85e0e4e38a8cb8238 Mon Sep 17 00:00:00 2001 From: Michael Kraus Date: Fri, 30 Oct 2026 12:25:42 +0100 Subject: [PATCH 5/8] Reject mixed precision, return no cache memory, record unread HRedSys fields MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - IsometricRolloutLoss and CrossModalCycleLoss raise an ArgumentError that names both element types when the data, the integrator's output or H has another element type than the network's parameters. - build_v_reduced and build_f_reduced default their cache to the decoder's parameter element type; a floating-point state of another element type than the cache's raises an ArgumentError. Dual states keep their own cache. - evaluate_vf_and_compute_∇Ψ returns copies, so no caller holds memory of the cache. - KNOWN_ISSUES.md K40: HRedSys fields that no code in src/ reads. - Tests for each error and for results that do not alias the cache. Co-Authored-By: Claude Opus 5.5 --- CHANGELOG.md | 2 +- KNOWN_ISSUES.md | 14 ++++++++++ src/loss/multimodal_losses.jl | 22 +++++++++++++-- src/reduced_system/reduced_system.jl | 14 +++++++--- .../multimodal_symplectic_autoencoder.jl | 13 +++++++++ .../reduced_system/reference_shift_and_vjp.jl | 27 +++++++++++++++++++ 6 files changed, 86 insertions(+), 6 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 1fd3a878b..75e9f3d24 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -150,7 +150,7 @@ breaking release). - **`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`. -- **`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``. `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 `:jacobian` uses preallocated buffers (`ξ_buf`, `v_intermediate`, `f_intermediate`, `jac_result`, and a per-element-type `IdDict{DataType, Any}` cache for `ForwardDiff.Dual` Newton states). 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 `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` has the element type of the decoder's parameters, and 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 a145adab3..6dfa1dd8d 100644 --- a/KNOWN_ISSUES.md +++ b/KNOWN_ISSUES.md @@ -62,6 +62,20 @@ What is known to be wrong in GeometricMachineLearning and is not fixed. entry means deciding what of an untrained network is worth testing at all — training the two networks first would make the comparison mean something, at a cost the suite does not pay. +### K40 · `HRedSys` carries a cache's buffers and its `pullback_mode` in fields that no code in `src/` reads. + +- **location:** `src/reduced_system/reduced_system.jl:93` +- **kind:** dead code +- **found:** #370 +- **evidence:** + The constructor sets `pullback_mode`, `ξ_buf`, `v_intermediate`, `f_intermediate`, `jac_result` + and `dual_cache`, and the type parameter `JR` carries the type of `jac_result`. `v_reduced` and + `f_reduced` hold the `ReducedVFCache` and the mode in their closures, and every read of + `.ξ_buf`, `.v_intermediate`, `.f_intermediate`, `.jac_result` and `.dual_cache` in `src/` is a read + of a `ReducedVFCache`, never of an `HRedSys`. `test/reduced_system/reference_shift_and_vjp.jl` + 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. + ## Found late ### K17 · `assign_q_and_p` drops the rows from `2N + 1` on without an error. diff --git a/src/loss/multimodal_losses.jl b/src/loss/multimodal_losses.jl index 52cda7c9f..34c45ac8b 100644 --- a/src/loss/multimodal_losses.jl +++ b/src/loss/multimodal_losses.jl @@ -160,6 +160,7 @@ Multi-step latent rollout, decoded trajectory reconstruction, and latent Hamilto 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`. """ struct IsometricRolloutLoss{IT, HT, T <: Real} <: NetworkLoss integrator::IT @@ -193,25 +194,39 @@ end ChainRulesCore.@non_differentiable _check_traj_len(::Any) +function _check_eltype(loss_name::String, what::String, x, ::Type{T}) where {T} + S = x isa Number ? typeof(x) : _state_eltype(x) + S === T || throw(ArgumentError( + "$loss_name: $what has element type $S, but the network has element type $T.")) + nothing +end + +ChainRulesCore.@non_differentiable _check_eltype(::Any, ::Any, ::Any, ::Any) + Base.@constprop :aggressive function (loss::IsometricRolloutLoss)( model::MultiModalSymplecticLayer, ps::NamedTuple, modality::Symbol, xs::AbstractVector) K = length(xs) _check_traj_len(K) - T = promote_type(_state_eltype(xs[1]), _leaf_eltype(ps.core)) + 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) z_curr = z1 err_lat = zero(T) err_rec = zero(T) err_H = zero(T) for k in 2:K + _check_eltype("IsometricRolloutLoss", "the data", xs[k], T) z_k = encode(model, ps, modality, xs[k]) z_curr = loss.integrator(z_curr) + _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) end inv_steps = inv(T(K - 1)) @@ -262,6 +277,7 @@ Cross-modal paired latent alignment and cycle consistency loss between modalitie 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`. +`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 mod_a::Symbol @@ -282,7 +298,9 @@ end Base.@constprop :aggressive function _cross_modal_eval( loss::CrossModalCycleLoss, model::MultiModalSymplecticLayer, ps::NamedTuple, x_a, x_b) - T = promote_type(_state_eltype(x_a), _state_eltype(x_b), _leaf_eltype(ps.core)) + T = _leaf_eltype(ps.core) + _check_eltype("CrossModalCycleLoss", "x_a", x_a, T) + _check_eltype("CrossModalCycleLoss", "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) diff --git a/src/reduced_system/reduced_system.jl b/src/reduced_system/reduced_system.jl index a373fe9f8..5e981e34e 100644 --- a/src/reduced_system/reduced_system.jl +++ b/src/reduced_system/reduced_system.jl @@ -22,6 +22,11 @@ function _get_vf_cache( if T === T0 return cache end + T <: AbstractFloat && throw( + ArgumentError( + "The reduced vector field got a state of element type $T, but its cache has element type $T0." + ), + ) get!(cache.dual_cache, T) do ReducedVFCache(T, N2, n2) end::ReducedVFCache{T} @@ -249,7 +254,9 @@ function evaluate_vf_and_compute_∇Ψ(t, q̃::AbstractVector{Tq}, p̃::Abstract fill!(f_intermediate, zero(T)) v_full(v_intermediate, t, q, p, parameters) f_full(f_intermediate, t, q, p, parameters) - v_intermediate, f_intermediate, DiffResults.jacobian(result)::Matrix{T} + # Copies, so that no caller holds memory of the cache. + copy(v_intermediate), copy(f_intermediate), + copy(DiffResults.jacobian(result)::Matrix{T}) end function evaluate_vf_and_vjp(t, q̃::AbstractVector{Tq}, p̃::AbstractVector{Tp}, @@ -297,13 +304,14 @@ We then get `build_v_reduced` outputs the first half of the entries of this vector field. When `pullback_mode = :jacobian` (default), `ForwardDiff.jacobian!` with `DiffResults.MutableDiffResult` computes ``\Psi^{\mathrm{dec}}(\xi)`` and ``\nabla_\xi \Psi^{\mathrm{dec}}`` in one pass. When `pullback_mode = :vjp`, `(v(z), f(z))` is evaluated once at ``z = \Psi^{\mathrm{dec}}(\xi)`` and ``(\nabla_\xi \Psi^{\mathrm{dec}})^T \mathbb{J}_{2N}^T [v(z); f(z)]`` is computed via `Zygote.pullback(decoder, ξ)`. +The default `cache` has the element type of the decoder's parameters. A state of another floating-point element type than the cache's raises an `ArgumentError`. """ function build_v_reduced(v_full, f_full, decoder::NeuralNetwork{<:SymplecticDecoder}; pullback_mode::Symbol = :jacobian, q_ref = nothing, p_ref = nothing, cache::ReducedVFCache = ReducedVFCache( - Float64, decoder.architecture.full_dim ÷ 2, + _leaf_eltype(params(params(decoder))), decoder.architecture.full_dim ÷ 2, decoder.architecture.reduced_dim ÷ 2)) _check_pullback_mode(pullback_mode) N2 = decoder.architecture.full_dim ÷ 2 @@ -342,7 +350,7 @@ function build_f_reduced(v_full, f_full, decoder::NeuralNetwork{<:SymplecticDeco q_ref = nothing, p_ref = nothing, cache::ReducedVFCache = ReducedVFCache( - Float64, decoder.architecture.full_dim ÷ 2, + _leaf_eltype(params(params(decoder))), decoder.architecture.full_dim ÷ 2, decoder.architecture.reduced_dim ÷ 2)) _check_pullback_mode(pullback_mode) N2 = decoder.architecture.full_dim ÷ 2 diff --git a/test/architectures/multimodal_symplectic_autoencoder.jl b/test/architectures/multimodal_symplectic_autoencoder.jl index 0fc9192dc..7c5bd78fe 100644 --- a/test/architectures/multimodal_symplectic_autoencoder.jl +++ b/test/architectures/multimodal_symplectic_autoencoder.jl @@ -416,6 +416,19 @@ end val_cyc_rev, expected_cyc_rev; atol = 10 * eps(T), 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) + @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, ξ))) + @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)) end @testset "Analytical Zygote gradients match central finite differences in Float64" begin diff --git a/test/reduced_system/reference_shift_and_vjp.jl b/test/reduced_system/reference_shift_and_vjp.jl index 2516dcb14..cfbcbcc17 100644 --- a/test/reduced_system/reference_shift_and_vjp.jl +++ b/test/reduced_system/reference_shift_and_vjp.jl @@ -1,6 +1,7 @@ using GeometricMachineLearning using GeometricMachineLearning: HRedSys, build_v_reduced, build_f_reduced, integrate_reduced_system, evaluate_vf_and_compute_∇Ψ, + evaluate_vf_and_vjp, ReducedVFCache, ForwardDiff, StateVariable using GeometricIntegrators: ImplicitMidpoint using KernelAbstractions @@ -193,6 +194,32 @@ Random.seed!(42) @test size(J_v_dual) == (n, n) && all(isfinite, J_v_dual) @test !isempty(rs_jac.dual_cache) + # A state of another floating-point type than the default cache's (the decoder's + # element type) raises an ArgumentError that names both types. + S = T === Float32 ? Float64 : Float32 + msg_cache = "got a state of element type $S, but its cache has element type $T." + @test_throws msg_cache v_red_j( + zeros(S, n), zero(S), S.(q̃_test), S.(p̃_test), nothing) + @test_throws msg_cache f_red_v( + zeros(S, n), zero(S), S.(q̃_test), S.(p̃_test), nothing) + + # No result holds memory of the cache: mutating a result does not change the next + # call, and a later call does not change an earlier result. + cache_a = ReducedVFCache(T, N, n) + q̃_other, p̃_other = randn(rng, T, n), randn(rng, T, n) + for helper in (evaluate_vf_and_compute_∇Ψ, evaluate_vf_and_vjp) + call(q, p) = helper(zero(T), q, p, nothing, decoder(nn), v_field!, f_field!, + cache_a, q_ref, p_ref) + r1 = call(q̃_test, p̃_test) + r1_saved = deepcopy(r1) + foreach(x -> fill!(x, T(NaN)), r1 isa Tuple ? r1 : (r1,)) + r2 = call(q̃_test, p̃_test) + @test all(map((a, b) -> isapprox(a, b; rtol = 8 * eps(T)), + r2 isa Tuple ? r2 : (r2,), r1_saved isa Tuple ? r1_saved : (r1_saved,))) + call(q̃_other, p̃_other) + @test all(x -> all(isnan, x), r1 isa Tuple ? r1 : (r1,)) + end + # 4. Quadratic Hamiltonian H(q, p) = 1/2 ||(q - q_ref)^2 + (p - p_ref)^2|| # with equilibrium (q_ref, p_ref) and integrator = ImplicitMidpoint(): # integrate_reduced_system starting at ics = (q = q_ref, p = p_ref) keeps From 1dd22384c546783a8909656b8b3c5d3cef100706 Mon Sep 17 00:00:00 2001 From: Michael Kraus Date: Fri, 30 Oct 2026 15:55:16 +0100 Subject: [PATCH 6/8] Address review nits of #370: document :vjp cost, test the Dual-path Jacobian, drop dead branch Document in build_v_reduced and build_f_reduced that pullback_mode = :vjp allocates more and runs slower than :jacobian for small problems, with measured figures. Compare the Dual-path Jacobian of the reduced field with central finite differences under a non-zero reference shift, reject a transposed and a sign-flipped Jacobian, and test reduction_error and projection_error against a loop-based value with a non-zero reference shift. Remove the unreachable nothing branch of build_h_reduced. Document that nn(x) of MultiModalSymplecticAutoencoder uses the first modality only. Rename the plan labels in the two new test files to descriptive names. Co-Authored-By: Claude Sonnet 5.5 --- .../multimodal_symplectic_autoencoder.jl | 3 + src/reduced_system/reduced_system.jl | 14 ++- .../multimodal_symplectic_autoencoder.jl | 14 +-- .../reduced_system/reference_shift_and_vjp.jl | 119 +++++++++++++++++- 4 files changed, 132 insertions(+), 18 deletions(-) diff --git a/src/architectures/multimodal_symplectic_autoencoder.jl b/src/architectures/multimodal_symplectic_autoencoder.jl index 21270ad12..989a34c5d 100644 --- a/src/architectures/multimodal_symplectic_autoencoder.jl +++ b/src/architectures/multimodal_symplectic_autoencoder.jl @@ -38,6 +38,9 @@ Given `nn = NeuralNetwork(arch, backend, T)`: - [`encode`](@ref)`(nn, modality::Symbol, x)` applies `frontends[modality]` followed by `encoder(core)` to map modality input `x` into ``(\tilde{q}, \tilde{p}) \in \mathbb{R}^{2n}``; - [`decode`](@ref)`(nn, modality::Symbol, ξ)` applies `decoder(core)` followed by `inverse_layer(frontends[modality])` to reconstruct in `modality`; - [`translate`](@ref)`(nn, from::Symbol, to::Symbol, x)` computes `decode(nn, to, encode(nn, from, x))`. + +!!! note "`nn(x)` uses the first modality only" + The call `nn(x)` on a network of this architecture makes a round trip through the first modality of `frontends` only: it computes `decode(nn, m, encode(nn, m, x))` for `m = first(keys(frontends))`. Use `encode`, `decode` and `translate` to select another modality. """ struct MultiModalSymplecticAutoencoder{ FT <: NamedTuple, diff --git a/src/reduced_system/reduced_system.jl b/src/reduced_system/reduced_system.jl index 5e981e34e..42cf893a6 100644 --- a/src/reduced_system/reduced_system.jl +++ b/src/reduced_system/reduced_system.jl @@ -305,6 +305,9 @@ We then get When `pullback_mode = :jacobian` (default), `ForwardDiff.jacobian!` with `DiffResults.MutableDiffResult` computes ``\Psi^{\mathrm{dec}}(\xi)`` and ``\nabla_\xi \Psi^{\mathrm{dec}}`` in one pass. When `pullback_mode = :vjp`, `(v(z), f(z))` is evaluated once at ``z = \Psi^{\mathrm{dec}}(\xi)`` and ``(\nabla_\xi \Psi^{\mathrm{dec}})^T \mathbb{J}_{2N}^T [v(z); f(z)]`` is computed via `Zygote.pullback(decoder, ξ)`. The default `cache` has the element type of the decoder's parameters. A state of another floating-point element type than the cache's raises an `ArgumentError`. + +!!! note "`:vjp` is not faster than `:jacobian` for small problems" + Each of the two fields, `v_reduced` and `f_reduced`, has its own pullback, so `:vjp` runs the decoder forward and backward once per field, where `:jacobian` pays one forward pass in dual numbers per call. For `2n = 4` and a decoder of `WeightTiedSymplecticAutoencoder(2N, 4; n_blocks = 1, n_layers = 1, n_ambient_layers = 0, sympnet_upscale = 2)` in `Float64`, one call of `v_reduced` allocates 4448 B (`:jacobian`) against 27968 B (`:vjp`) for `2N = 12`, 8160 B against 29376 B for `2N = 40`, and 29728 B against 37248 B for `2N = 200`. A call takes 0.7 μs against 99 μs for `2N = 12`, 0.8 μs against 104 μs for `2N = 40`, and 2.5 μs against 106 μs for `2N = 200`. The allocation of both modes grows with `N`, and the figures are minima of warm calls in a fresh process with `--check-bounds=auto` on one machine. Use `:jacobian` unless a measurement of your own problem says otherwise. """ function build_v_reduced(v_full, f_full, decoder::NeuralNetwork{<:SymplecticDecoder}; pullback_mode::Symbol = :jacobian, @@ -344,6 +347,9 @@ Builds the reduced vector field (``p`` part) based on the full vector field for `build_f_reduced` outputs the second half of the entries of this vector field. See [`build_v_reduced`](@ref) for more information. + +!!! note "`:vjp` is not faster than `:jacobian` for small problems" + Each field has its own pullback, so `pullback_mode = :vjp` allocates more and runs slower than `:jacobian` for small problems. For `2n = 4` in `Float64`, one call of `f_reduced` allocates 4272 B (`:jacobian`) against 27968 B (`:vjp`) for `2N = 12`, 7744 B against 29376 B for `2N = 40`, and 28048 B against 37248 B for `2N = 200`, and takes 0.7 μs against 99 μs, 0.8 μs against 98 μs, and 2.6 μs against 105 μs. [`build_v_reduced`](@ref) states the setup of these measurements. """ function build_f_reduced(v_full, f_full, decoder::NeuralNetwork{<:SymplecticDecoder}; pullback_mode::Symbol = :jacobian, @@ -377,14 +383,10 @@ end function build_h_reduced( h_full, decoder::NeuralNetwork{<:SymplecticDecoder}; - q_ref = nothing, p_ref = nothing) + q_ref, p_ref) function h_reduced(t, q, p, params) dec = decoder((q = q, p = p)) - if q_ref === nothing || p_ref === nothing - h_full(t, dec.q, dec.p, params) - else - h_full(t, dec.q .+ q_ref, dec.p .+ p_ref, params) - end + h_full(t, dec.q .+ q_ref, dec.p .+ p_ref, params) end h_reduced end diff --git a/test/architectures/multimodal_symplectic_autoencoder.jl b/test/architectures/multimodal_symplectic_autoencoder.jl index 7c5bd78fe..90f6d60d7 100644 --- a/test/architectures/multimodal_symplectic_autoencoder.jl +++ b/test/architectures/multimodal_symplectic_autoencoder.jl @@ -80,7 +80,7 @@ end _extract_grad_vec(g::NetworkParameters) = _extract_grad_vec(params(g)) -@testset "Clause 1: MultiModalSymplecticAutoencoder construction and edge cases" begin +@testset "MultiModalSymplecticAutoencoder construction and edge cases" begin @test isdefined(GeometricMachineLearning, :MultiModalSymplecticAutoencoder) @test isdefined(GeometricMachineLearning, :encode) @test isdefined(GeometricMachineLearning, :decode) @@ -119,7 +119,7 @@ _extract_grad_vec(g::NetworkParameters) = _extract_grad_vec(params(g)) @test_throws re_unknown translate(nn, :grid, :unknown, zeros(4, 4)) end -@testset "Clause 1: encode, decode, translate and right-invertible round-trip ($T)" for T in REAL_ELTYPES +@testset "encode, decode, translate and right-invertible round-trip ($T)" for T in REAL_ELTYPES rng = Random.MersenneTwister(101) nx, nv = 5, 4 mx, mv = 2, 2 @@ -275,7 +275,7 @@ end @test norm(ξ_sort_rt - ξ_sort_test) <= 50 * eps(T) * norm(ξ_sort_test) end -@testset "Clause 2: Multi-modal losses, zero-input edge case, M1, and Zygote vs FD gradients" begin +@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) @@ -312,7 +312,7 @@ end 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) - @testset "Forward evaluation, @inferred, zero edge case, and M1 ($T)" for T in REAL_ELTYPES + @testset "Forward evaluation, @inferred, zero edge case, and exact-zero loss ($T)" for T in REAL_ELTYPES rng = Random.MersenneTwister(202) nx, nv = 4, 4 mx, mv = 2, 1 @@ -343,7 +343,7 @@ end @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)) - # Independent perturbation F_pert != F_hat so z_hat != z and the latent penalty term is active (catches O3) + # 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) expected_c_pert = (sqrt(sum(abs2, F_pert .- F) + eps_T^2) - eps_T) / @@ -352,7 +352,7 @@ end val_c_pert = @inferred loss_c(nn, :grid, F_pert, F) @test isapprox(val_c_pert, expected_c_pert; atol = 10 * eps(T), rtol = 10 * eps(T)) - # Mutant M1 check: when x̂ == x != 0, CharbonnierRelativeLoss MUST be exact zero(T)! + # Exact-zero check: when x̂ == x != 0, CharbonnierRelativeLoss MUST be exact zero(T)! # Omitting `- ε` in the numerator yields ε / sqrt(||x||_F^2 + ε^2) > 0. val_identical = @inferred loss_c(nn, :grid, F, F) @test val_identical === zero(T) @@ -405,7 +405,7 @@ end expected_cyc = T(1.2) * sum(abs2, z_a .- z_b) + T(0.3) * sum(abs2, z_a_cyc .- z_a) @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 (catches O5) + # 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)) val_cyc_rev = @inferred loss_cyc_rev(nn, F, x_p) diff --git a/test/reduced_system/reference_shift_and_vjp.jl b/test/reduced_system/reference_shift_and_vjp.jl index cfbcbcc17..ad18bd380 100644 --- a/test/reduced_system/reference_shift_and_vjp.jl +++ b/test/reduced_system/reference_shift_and_vjp.jl @@ -1,6 +1,7 @@ using GeometricMachineLearning using GeometricMachineLearning: HRedSys, build_v_reduced, build_f_reduced, - integrate_reduced_system, evaluate_vf_and_compute_∇Ψ, + integrate_reduced_system, integrate_full_system, + evaluate_vf_and_compute_∇Ψ, evaluate_vf_and_vjp, ReducedVFCache, ForwardDiff, StateVariable using GeometricIntegrators: ImplicitMidpoint @@ -13,7 +14,7 @@ include("../helpers/eltypes.jl") Random.seed!(42) -@testset "Clause 3: HRedSys reference-state shifting, pullback_mode = :jacobian | :vjp, and work buffers" begin +@testset "HRedSys reference-state shifting, pullback_mode = :jacobian | :vjp, and work buffers" begin model_chk = WeightTiedSymplecticAutoencoder( 6, 2; n_blocks = 1, n_layers = 1, n_ambient_layers = 0, sympnet_upscale = 2) nn_chk = NeuralNetwork(model_chk, CPU(), Float64) @@ -126,7 +127,7 @@ Random.seed!(42) @test isapprox(qp_pert.q, qp_pert_exp.q; atol = 10 * eps(T), rtol = 10 * eps(T)) @test isapprox(qp_pert.p, qp_pert_exp.p; atol = 10 * eps(T), rtol = 10 * eps(T)) - # 2. VJP vs Jacobian agreement to 50 eps(T) for 2N = 12, 2n = 4 in Float32 and Float64 (catches M2!) + # 2. VJP vs Jacobian agreement to 50 eps(T) for 2N = 12, 2n = 4 in Float32 and Float64 (catches a sign error or a swapped block in the pullback) for trial in 1:5 q̃ = randn(rng, T, n) p̃ = randn(rng, T, n) @@ -140,7 +141,7 @@ Random.seed!(42) rs_vjp.v_reduced(v_vjp, zero(T), q̃, p̃, nothing) rs_vjp.f_reduced(f_vjp, zero(T), q̃, p̃, nothing) - # Verify non-zero vector field so a sign flip (M2) or zeroing is caught + # Verify non-zero vector field so a sign flip or zeroing is caught @test norm(v_jac) > T(1e-2) @test norm(f_jac) > T(1e-2) @test norm(v_vjp - v_jac) <= 50 * eps(T) * max(one(T), norm(v_jac)) @@ -194,6 +195,45 @@ Random.seed!(42) @test size(J_v_dual) == (n, n) && all(isfinite, J_v_dual) @test !isempty(rs_jac.dual_cache) + # The Jacobian of the reduced field ξ = (q̃, p̃) ↦ (v, f) through the Dual path, with the + # non-zero reference shift (q_ref, p_ref) above, equals the central finite differences of + # that field. The finite differences run on Float64 states only, which the Float32 cache of + # `rs_jac` does not take; they use `h = cbrt(eps)`, the step that balances the truncation + # error `O(h^2)` (measured: about 0.2 h^2 relative at h = 1e-4) and the round-off error + # `O(eps / h)` (about 4e-11 relative). The mismatch is relative to `max(1, ‖J‖)`, and + # `1e-9` is 25 times the largest of 5 mismatches that were measured with this setup (3.8e-11), + # and 9 orders below the mismatch of a transposed or a sign-flipped Jacobian (at least 1.6). + if T === Float64 + rng_fd = Random.MersenneTwister(505) + function reduced_field(ξ) + out_v = zeros(eltype(ξ), n) + out_f = zeros(eltype(ξ), n) + rs_jac.v_reduced( + out_v, zero(eltype(ξ)), ξ[1:n], ξ[(n + 1):(2 * n)], nothing) + rs_jac.f_reduced( + out_f, zero(eltype(ξ)), ξ[1:n], ξ[(n + 1):(2 * n)], nothing) + vcat(out_v, out_f) + end + h_fd = cbrt(eps(T)) + for trial in 1:5 + ξ_fd = randn(rng_fd, T, reduced_dim) + J_dual = ForwardDiff.jacobian(reduced_field, ξ_fd) + J_fd = similar(J_dual) + for j in 1:reduced_dim + e_j = zeros(T, reduced_dim) + e_j[j] = h_fd + J_fd[:, j] = (reduced_field(ξ_fd + e_j) - reduced_field(ξ_fd - e_j)) / + (2 * h_fd) + end + mismatch(J) = norm(J - J_fd) / max(one(T), norm(J_fd)) + @test size(J_dual) == (reduced_dim, reduced_dim) + @test mismatch(J_dual) <= 1e-9 + # The comparison can fail: a transposed and a sign-flipped Jacobian are rejected. + @test mismatch(transpose(J_dual)) > 1e-3 + @test mismatch(-J_dual) > 1e-3 + end + end + # A state of another floating-point type than the default cache's (the decoder's # element type) raises an ArgumentError that names both types. S = T === Float32 ? Float64 : Float32 @@ -223,7 +263,7 @@ Random.seed!(42) # 4. Quadratic Hamiltonian H(q, p) = 1/2 ||(q - q_ref)^2 + (p - p_ref)^2|| # with equilibrium (q_ref, p_ref) and integrator = ImplicitMidpoint(): # integrate_reduced_system starting at ics = (q = q_ref, p = p_ref) keeps - # ξ(t) = (0, 0) and (q(t), p(t)) = (q_ref, p_ref) to 10 eps(T) for 20 steps (catches M3!). + # ξ(t) = (0, 0) and (q(t), p(t)) = (q_ref, p_ref) to 10 eps(T) for 20 steps (catches a lift or an encoding that ignores the reference shift). function v_quad!(v, t, q, p, params) @inbounds for i in eachindex(v, p) v[i] = p[i] - p_ref[i] @@ -350,3 +390,72 @@ Random.seed!(42) end end end + +@testset "reduction_error and projection_error with a non-zero reference shift ($T)" for T in REAL_ELTYPES + rng = Random.MersenneTwister(606) + 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) + q_ref = randn(rng, T, N) + p_ref = randn(rng, T, N) + z_ref = vcat(q_ref, p_ref) + + v_quad!(v, t, q, p, params) = (v .= p .- p_ref; nothing) + f_quad!(f, t, q, p, params) = (f .= .-(q .- q_ref); nothing) + h_quad(t, q, p, params) = T(0.5) * (sum(abs2, q .- q_ref) + sum(abs2, p .- p_ref)) + + ics = (q = StateVariable(q_ref .+ T(0.3) .* randn(rng, T, N)), + p = StateVariable(p_ref .+ T(0.3) .* randn(rng, T, N))) + rs = HRedSys( + 2 * N, 2 * n, encoder(nn), decoder(nn), v_quad!, f_quad!, h_quad, + (zero(T), T(0.25)), T(0.05), ics; + integrator = ImplicitMidpoint(), q_ref = q_ref, p_ref = p_ref) + sol_full = integrate_full_system(rs) + sol_red = integrate_reduced_system(rs) + + # 5 steps store 6 states, as many as N = 6: a shift broadcast along the time axis instead of + # the axis of the degrees of freedom runs without a dimension error, so only the value shows it. + @test length(sol_full.q) == N + @test length(sol_red.q) == N + + # The expected errors come from a loop over the stored states, with the shift added to and + # subtracted from each state as a vector, not from the broadcasts of `decode` and `encode`. + # `_norm` of a `(q, p)` pair is `(‖q‖ + ‖p‖) / √2`, and the `√2` cancels in the ratio. + function relative_error(reconstruct, shift) + dq, dp, nq, np = zero(T), zero(T), zero(T), zero(T) + for k in 0:(N - 1) + q_k = Vector(sol_full.q[k]) + p_k = Vector(sol_full.p[k]) + ẑ = reconstruct(k, shift) + dq += sum(abs2, ẑ[1:N] .- q_k) + dp += sum(abs2, ẑ[(N + 1):(2 * N)] .- p_k) + nq += sum(abs2, q_k) + np += sum(abs2, p_k) + end + (sqrt(dq) + sqrt(dp)) / (sqrt(nq) + sqrt(np)) + end + projected(k, shift) = decoder(nn)(encoder(nn)( + vcat(Vector(sol_full.q[k]), Vector(sol_full.p[k])) .- shift)) .+ shift + reduced(k, shift) = decoder(nn)(vcat(Vector(sol_red.q[k]), Vector(sol_red.p[k]))) .+ + shift + + proj_expected = relative_error(projected, z_ref) + red_expected = relative_error(reduced, z_ref) + @test projection_error(rs, sol_full) isa T + @test reduction_error(rs, sol_full, sol_red) isa T + @test isapprox(projection_error(rs, sol_full), proj_expected; rtol = 50 * eps(T)) + @test isapprox(reduction_error(rs, sol_full, sol_red), red_expected; rtol = 50 * eps(T)) + # The one-argument forms integrate the systems themselves and agree. + @test isapprox(projection_error(rs), proj_expected; rtol = 50 * eps(T)) + @test isapprox(reduction_error(rs), red_expected; rtol = 50 * eps(T)) + + # The check can fail: the errors are not trivial, and the errors without the shift differ + # from the shifted ones, so a dropped or misplaced shift changes the result. + @test proj_expected > T(1e-2) + @test red_expected > T(1e-2) + @test abs(relative_error(projected, zero(z_ref)) - proj_expected) > + T(1e-3) * proj_expected + @test abs(relative_error(reduced, zero(z_ref)) - red_expected) > T(1e-3) * red_expected +end From 0387deaadca4a0e292062f237d4d4b3ac34ec47f Mon Sep 17 00:00:00 2001 From: Michael Kraus Date: Fri, 30 Oct 2026 16:20:03 +0100 Subject: [PATCH 7/8] Address CodeRabbit review of #370: reuse factors, check each state component, one-sided references - _frob_sq_norm of a batch of DLRA states reuses a factor whose element type is already T. - _check_eltype checks every component of a state, so a mixed-precision particle state raises the ArgumentError. - _add_ref! shifts only the supplied component when one reference is given. - CHANGELOG: state the default cache element type for HRedSys and for the standalone builders. Co-Authored-By: Claude Sonnet 5.5 --- CHANGELOG.md | 2 +- src/loss/multimodal_losses.jl | 16 +++++-- src/reduced_system/reduced_system.jl | 14 ++++++ .../multimodal_symplectic_autoencoder.jl | 31 ++++++++++++ .../reduced_system/reference_shift_and_vjp.jl | 47 +++++++++++++++++++ 5 files changed, 105 insertions(+), 5 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 75e9f3d24..583baa706 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -150,7 +150,7 @@ breaking release). - **`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`. -- **`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` has the element type of the decoder's parameters, and 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 `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`. ### Changed diff --git a/src/loss/multimodal_losses.jl b/src/loss/multimodal_losses.jl index 34c45ac8b..fac169b3f 100644 --- a/src/loss/multimodal_losses.jl +++ b/src/loss/multimodal_losses.jl @@ -16,7 +16,9 @@ end # A batch of DLRA states, one per slice along axis 3: the sum of the slices' squared norms. function _frob_sq_norm(x::NamedTuple{(:X, :S, :V), <:NTuple{3, AbstractArray{<:Any, 3}}}) T = _state_eltype(x) - X, S, V = T.(x.X), T.(x.S), T.(x.V) + X = eltype(x.X) === T ? x.X : T.(x.X) + S = eltype(x.S) === T ? x.S : T.(x.S) + V = eltype(x.V) === T ? x.V : T.(x.V) G_X = tensor_transpose_tensor_mul(X, X) G_V = tensor_transpose_tensor_mul(V, V) M = tensor_tensor_mul(tensor_transpose_tensor_mul(S, G_X), S) @@ -194,10 +196,16 @@ end ChainRulesCore.@non_differentiable _check_traj_len(::Any) +_component_eltypes(x::Number) = (typeof(x),) +_component_eltypes(x::AbstractArray) = (eltype(x),) +_component_eltypes(x::NamedTuple{(:q, :p)}) = (eltype(x.q), eltype(x.p)) +_component_eltypes(x::NamedTuple{(:X, :S, :V)}) = (eltype(x.X), eltype(x.S), eltype(x.V)) + function _check_eltype(loss_name::String, what::String, x, ::Type{T}) where {T} - S = x isa Number ? typeof(x) : _state_eltype(x) - S === T || throw(ArgumentError( - "$loss_name: $what has element type $S, but the network has element type $T.")) + for S in _component_eltypes(x) + S === T || throw(ArgumentError( + "$loss_name: $what has element type $S, but the network has element type $T.")) + end nothing end diff --git a/src/reduced_system/reduced_system.jl b/src/reduced_system/reduced_system.jl index 42cf893a6..b91fb2c65 100644 --- a/src/reduced_system/reduced_system.jl +++ b/src/reduced_system/reduced_system.jl @@ -50,6 +50,20 @@ function _add_ref!( z end +function _add_ref!(z::AbstractVector, q_ref::AbstractVector, ::Nothing, N2::Int) + @inbounds for i in 1:N2 + z[i] += q_ref[i] + end + z +end + +function _add_ref!(z::AbstractVector, ::Nothing, p_ref::AbstractVector, N2::Int) + @inbounds for i in 1:N2 + z[N2 + i] += p_ref[i] + end + z +end + _add_ref!(z::AbstractVector, ::Nothing, ::Nothing, ::Int) = z @doc raw""" diff --git a/test/architectures/multimodal_symplectic_autoencoder.jl b/test/architectures/multimodal_symplectic_autoencoder.jl index 90f6d60d7..8235a81bb 100644 --- a/test/architectures/multimodal_symplectic_autoencoder.jl +++ b/test/architectures/multimodal_symplectic_autoencoder.jl @@ -429,6 +429,19 @@ end @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)) + + # A particle state with one component of another precision than the network's raises + # the same ArgumentError, whichever component it is. + x_p_qS = (q = x_p.q, p = S.(x_p.p)) + x_p_pS = (q = S.(x_p.q), p = x_p.p) + for x_mixed in (x_p_qS, x_p_pS) + @test_throws msg("IsometricRolloutLoss", "the data") loss_iso( + 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)) + end end @testset "Analytical Zygote gradients match central finite differences in Float64" begin @@ -596,6 +609,24 @@ end @test isapprox(GeometricMachineLearning._frob_sq_norm(x_dlra_batch), sum(b -> sum(abs2, X_b[:, :, b] * S_b[:, :, b] * V_b[:, :, b]'), 1:B); rtol = 8 * eps(T)) + # Factors that already have the element type of the batch are reused, not modified: + # value and gradient agree with the sum over the slices, which takes the 2D path. + for Tb in (Float32, Float64) + Xb, Sb, Vb = randn(rng, Tb, nx, 2, B), randn(rng, Tb, 2, 2, B), + randn(rng, Tb, nv, 2, B) + Xb0, Sb0, Vb0 = copy(Xb), copy(Sb), copy(Vb) + slice(b, X, S, V) = (X = X[:, :, b], S = S[:, :, b], V = V[:, :, b]) + f_batch(X, S, V) = GeometricMachineLearning._frob_sq_norm((X = X, S = S, V = V)) + f_slices(X, S, V) = sum( + b -> GeometricMachineLearning._frob_sq_norm(slice(b, X, S, V)), 1:B) + val_b = f_batch(Xb, Sb, Vb) + @test val_b isa Tb + @test isapprox(val_b, f_slices(Xb, Sb, Vb); rtol = 32 * eps(Tb)) + g_b = Zygote.gradient(f_batch, Xb, Sb, Vb) + g_s = Zygote.gradient(f_slices, Xb, Sb, Vb) + @test all(isapprox.(g_b, g_s; rtol = 64 * eps(Tb))) + @test Xb == Xb0 && Sb == Sb0 && Vb == Vb0 + end @test loss_c(nn_can, :dlra, x_dlra_batch) > zero(T) calls_mm, calls_fe, calls_core = Ref(0), Ref(0), Ref(0) NeuralNetwork( diff --git a/test/reduced_system/reference_shift_and_vjp.jl b/test/reduced_system/reference_shift_and_vjp.jl index ad18bd380..5bc386f57 100644 --- a/test/reduced_system/reference_shift_and_vjp.jl +++ b/test/reduced_system/reference_shift_and_vjp.jl @@ -459,3 +459,50 @@ end T(1e-3) * proj_expected @test abs(relative_error(reduced, zero(z_ref)) - red_expected) > T(1e-3) * red_expected end + +@testset "build_v_reduced and build_f_reduced with one reference only ($mode, $which)" for mode in (:jacobian, :vjp), + which in (:q_ref, :p_ref) + + rng = Random.MersenneTwister(707) + T = Float64 + 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) + ref = randn(rng, T, N) + # Both fields depend on q and on p, so a shift of the wrong component changes the result. + v_nl!(v, t, q, p, params) = (v .= p .+ T(0.5) .* q .^ 2; nothing) + f_nl!(f, t, q, p, params) = (f .= .-q .- p .^ 2; nothing) + + q̃, p̃ = randn(rng, T, n), randn(rng, T, n) + ξ = vcat(q̃, p̃) + # The expected values come from the decoder, its Jacobian from ForwardDiff, and a loop that + # adds the reference to the component it belongs to. + z = Vector(decoder(nn)(ξ)) + ∇Ψ = ForwardDiff.jacobian(x -> decoder(nn)(x), ξ) + for i in 1:N + which === :q_ref ? (z[i] += ref[i]) : (z[N + i] += ref[i]) + end + v_z, f_z = zeros(T, N), zeros(T, N) + v_nl!(v_z, zero(T), z[1:N], z[(N + 1):(2 * N)], nothing) + f_nl!(f_z, zero(T), z[1:N], z[(N + 1):(2 * N)], nothing) + v_expected = ∇Ψ[(N + 1):(2 * N), (n + 1):(2 * n)]' * v_z .- + ∇Ψ[1:N, (n + 1):(2 * n)]' * f_z + f_expected = ∇Ψ[1:N, 1:n]' * f_z .- ∇Ψ[(N + 1):(2 * N), 1:n]' * v_z + + kw = (; which => ref) + v_red = build_v_reduced(v_nl!, f_nl!, decoder(nn); pullback_mode = mode, kw...) + f_red = build_f_reduced(v_nl!, f_nl!, decoder(nn); pullback_mode = mode, kw...) + v_out, f_out = zeros(T, n), zeros(T, n) + v_red(v_out, zero(T), q̃, p̃, nothing) + f_red(f_out, zero(T), q̃, p̃, nothing) + @test isapprox(v_out, v_expected; rtol = 1e-10, atol = 1e-12) + @test isapprox(f_out, f_expected; rtol = 1e-10, atol = 1e-12) + + # The shift is not trivial: the unshifted fields differ from the shifted ones. + v_red0 = build_v_reduced(v_nl!, f_nl!, decoder(nn); pullback_mode = mode) + v_out0 = zeros(T, n) + v_red0(v_out0, zero(T), q̃, p̃, nothing) + @test norm(v_out0 .- v_expected) > 1e-3 * norm(v_expected) +end From c0030b4fd8d83d879397c9836d01c841e04308eb Mon Sep 17 00:00:00 2001 From: Michael Kraus Date: Fri, 30 Oct 2026 16:27:01 +0100 Subject: [PATCH 8/8] Add K41 to KNOWN_ISSUES.md: CrossModalCycleLoss with a mixed-component NamedTuple raises a MethodError Co-Authored-By: Claude Sonnet 5.5 --- KNOWN_ISSUES.md | 23 +++++++++++++++++++++++ 1 file changed, 23 insertions(+) diff --git a/KNOWN_ISSUES.md b/KNOWN_ISSUES.md index 6dfa1dd8d..975a238e4 100644 --- a/KNOWN_ISSUES.md +++ b/KNOWN_ISSUES.md @@ -76,6 +76,29 @@ 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`. + +- **location:** `src/loss/multimodal_losses.jl:323` +- **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 + `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: + + ```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 + `_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 + method accepts, and the pull request does not make it. + ## Found late ### K17 · `assign_q_and_p` drops the rows from `2N + 1` on without an error.