Add the parametric core: apply_parametric, ParametricLoss, ParametricDataLoader

CHANGELOG.md

@@ -99,6 +99,8 @@ breaking release).
 
 ### Added
 
+- **`apply_parametric`, `ParametricLoss` and `ParametricDataLoader` train a network whose forward pass takes the parameters of the system beside its input.** The three names are `public` and not exported, so a caller qualifies them or names them in a `using GeometricMachineLearning: …`. `apply_parametric(model, x, μ, ps)` applies a `Chain` to one sample — a vector or a `(q, p)` pair of vectors — with one `μ`, or to a batch — a `2d × 1 × k` tensor or a `(q, p)` pair of `d × 1 × k` tensors — with a vector of `k` values of `μ`, sample by sample; the output has the shape of the input. The `Chain` receives the pair `(x, μ)`, and the internal `concatenate_array_with_parameters` puts the flattened `μ` below the state in the element type of the state, side by side for a matrix and a vector of `μ`, and repeated in every column for a matrix and one `μ`. `ParametricLoss()(model, ps, input, output, μ)` is the relative norm of the error of `apply_parametric`, and `ZygotePullback` takes the triple `(input, output, μ)` with one `μ` or a vector of them. `ParametricDataLoader(data, parameters)` takes a `2d × nt × m` array or a `(q, p)` pair of `d × nt × m` arrays of one array type with a vector of `m` named tuples, and `ParametricDataLoader(ensemble_solution)` takes an `EnsembleSolution` of `(q, p)` states with the parameters of its problem. A `Batch(b)` draws each pair of a start position and a trajectory once per epoch, as for `DataLoader`, and a minibatch is the state, the state one time step later and the `μ` of each sample. `train!` takes a `ParametricDataLoader` and trains on `ParametricLoss()` when no `loss` is given. A `Batch` with a sequence length or a prediction window other than one, a parameter vector whose length is not the number of trajectories, no trajectory, too few time steps for the batch, a `μ` whose element type is not the data's, and a `μ` whose keys, their order or the sizes of its leaves differ from the first `μ`'s raise an `ArgumentError`; a `(q, p)` pair of two array types is a `MethodError`, as for `DataLoader` (K20). The code is ported from pull request #207 by Benedikt Brantner (relates to #207); there `concatenate_array_with_parameters` promoted a `Float32` state to `Float64`, rejected a matrix of more than one column with one `μ`, and the loader and the batch checked their counts with `@assert`.
+
 - **`KernelAbstractions.get_backend` is defined for `PoissonTensor`.** GeometricOptimizers 0.9 asks every operand of a checked product for its backend, so `𝕁 * Y` with a `StiefelManifold` raised `ArgumentError: Implement KernelAbstractions.get_backend(::PoissonTensor...)` without it.
 
 - **The test suite runs on an Apple GPU in its `metal` group.** `test/devices/metal.jl` checks the two device paths that were measured by hand and not tested before: the products of a `PoissonTensor` with a wrapped device array, `𝕁 * view(A, [1, 2, 3, 4], :)` and `𝕁 * B'`, which `ext/GPUArraysCoreExt.jl` carries; and, in the testset "_mul and tensor_tensor_mul on Metal match the host", `_mul` with a dense and a `SymmetricMatrix` weight from the left, a dense one from the right, and `tensor_tensor_mul` on `MetalBackend()`. Each result must stay on the device and match the host in `Float32`, with scalar indexing off. `Pkg.test(test_args = ["metal"])` runs them without the subject tests, and a missing device records a skip, which happens inside a sandbox. Metal (1.10 or later) is a test dependency on every platform; it installs and precompiles on Linux and Windows, and nothing there loads it. On an Apple M4 Max with Metal 1.11.1 and Julia 1.13.1, both testsets pass in a cold process: 4 of 4 assertions in the `PoissonTensor` testset and 8 of 8 in the "_mul and tensor_tensor_mul on Metal match the host" testset.

KNOWN_ISSUES.md

@@ -129,6 +129,8 @@ What is known to be wrong in GeometricMachineLearning and is not fixed.
   The same call with `q = A` gives a `DataLoader` with `input_dim == 4`. The docstring at
   `src/data_loader/data_loader.jl:205` states the one-type signature. Closing it means widening the
   three signatures to `NamedTupleOfArrays{T, (:q, :p)}`, and the docstring with them.
+  `ParametricDataLoader`'s `(q, p)` constructor has the same one-type restriction and widens with
+  it.
 
 ## Upstream
 
@@ -145,6 +147,47 @@ What is known to be wrong in GeometricMachineLearning and is not fixed.
   The blocks are not test failures: both runs are green, and every testset of ea0c397 has the
   same count in both.
 
+### K28 · AbstractNeuralNetworks' loss norm of a `(q, p)` named tuple promotes `Float32` to `Float64`.
+
+- **location:** AbstractNeuralNetworks 0.9.3, `src/losses.jl:73`
+- **kind:** upstream
+- **found:** 2026-10-25, by the critics of the parametric core
+- **evidence:**
+  `_norm(dx::NamedTuple) = sum(map(_norm, values(dx))) / √length(dx)` divides by the square root of
+  an `Int`, which is a `Float64`. So `ParametricLoss` returns a `Float64` on `Float32` `(q, p)`
+  data, and `train!` on a `Float32` `(q, p)` `ParametricDataLoader` returns a `Vector{Float64}`
+  history. Array data keep `Float32`.
+
+  ```julia
+  import AbstractNeuralNetworks
+  x = (q = rand(Float32, 2, 1, 3), p = rand(Float32, 2, 1, 3))
+  typeof(AbstractNeuralNetworks._norm(x))    # Float64
+  ```
+
+### K30 · The gradient of AbstractNeuralNetworks' loss of a `(q, p)` named tuple infers as `Union{Nothing, NetworkParameters}`.
+
+- **location:** AbstractNeuralNetworks 0.9.3, `src/losses.jl:86`
+- **kind:** upstream
+- **found:** #357, review of 2026-10-26; the same inference shows on `origin/main` at `a664f5c9`
+- **evidence:**
+  On Julia 1.13.2 and `a664f5c9`, `Base.return_types` of the pullback of
+  `ps -> _compute_loss((q = ps.L1.W * q, p = ps.L1.b .* p), (q = yq, p = yp))` gives
+  `Tuple{Union{Nothing, NetworkParameters}}` in `Float32` and `Float64`, and the same function
+  with `vcat` in place of the named tuples gives the concrete `NetworkParameters{T, (:L1,), …}`. So
+  `ZygotePullback(ParametricLoss())` on `(q, p)` data infers its gradient as that `Union`, for one
+  sample and for a batch, and a plain sum of squares over the `(q, p)` output of a batch infers
+  concretely. The run-time gradient is a `NetworkParameters`.
+
+  ```julia
+  using GeometricMachineLearning, Zygote
+  import AbstractNeuralNetworks
+  ps = NetworkParameters((L1 = (W = rand(1, 1), b = rand(1)),))
+  q, p, yq, yp = rand(1), rand(1), rand(1), rand(1)
+  v, pb = Zygote.pullback(ps -> AbstractNeuralNetworks._compute_loss(
+      (q = ps.L1.W * q, p = ps.L1.b .* p), (q = yq, p = yp)), ps)
+  Base.return_types(pb, (typeof(v),))    # Tuple{Union{Nothing, NetworkParameters}}
+  ```
+
 ### K4 · Metal.jl 1.11.1 needs a Metal device to precompile, so a sandboxed run with no Metal cache fails before the tests
 
 - **location:** `test/Project.toml:16`

docs/src/data_loader/data_loader.md

@@ -165,4 +165,5 @@ Batch(::Int, ::Int, ::Int)
 GeometricMachineLearning.number_of_batches
 GeometricMachineLearning.convert_input_and_batch_indices_to_array
 GeometricMachineLearning.accuracy
+GeometricMachineLearning.ParametricDataLoader
 ```

docs/src/reduced_order_modeling/losses.md

@@ -46,6 +46,8 @@ where ``\mathbf{x}^{(t)}`` is the solution of the FOM at point ``t`` and ``\math
 TransformerLoss
 AutoEncoderLoss
 ReducedLoss
+GeometricMachineLearning.ParametricLoss
+GeometricMachineLearning.apply_parametric
 projection_error
 reduction_error
 ```

src/GeometricMachineLearning.jl

@@ -11,6 +11,12 @@ using AbstractNeuralNetworks
 # `NamedTuple` of layers is walked with `Base.map`, which needs nothing from anybody.
 import NeuralNetworkParameters: NetworkParameters
 using NeuralNetworkParameters: mapstorage
+# `flatten` puts the parameters of a system into the input of a network that takes them
+# (`src/parametric.jl`); `parameter_eltype` checks their element type against the data's, and
+# `parameterlayout` their shape against each other's. `map_cotangent` and `storage_gradient` turn the
+# cotangent of a parameter set into its gradient in the reverse rule of a batch.
+using NeuralNetworkParameters: flatten, parameter_eltype, parameterlayout, map_cotangent,
+                               storage_gradient
 using ChainRulesCore
 using GeometricBase
 using GeometricSolutions: GeometricSolution, EnsembleSolution, TimeSeries
@@ -205,6 +211,7 @@ export solve!, encoder, decoder
 
 export iterate
 
+include("parametric.jl")
 include("loss/losses.jl")
 
 export AbstractPullback, ZygotePullback
@@ -214,8 +221,12 @@ export DataLoader
 export Batch, train!
 include("data_loader/tensor_assign.jl")
 include("data_loader/batch.jl")
+include("data_loader/parametric_data_loader.jl")
 include("data_loader/train.jl")
 
+# The training of a network that takes the parameters of the system beside its input.
+public apply_parametric, ParametricLoss, ParametricDataLoader
+
 include("reduced_system/reduced_system.jl")
 
 export HRedSys, reduction_error, projection_error, integrate_reduced_system,

src/data_loader/parametric_data_loader.jl

@@ -0,0 +1,110 @@
+@doc raw"""
+    ParametricDataLoader(data, parameters)
+    ParametricDataLoader(ensemble_solution)
+
+Make a loader of trajectories of a system together with the parameters ``\mu`` of the system that
+each trajectory belongs to.
+
+`data` is a `2d × nt × m` array, or a `(q, p)` named tuple of two `d × nt × m` arrays of one
+array type: `m` trajectories of `nt` time steps, with `q` above `p` where they are stacked.
+`parameters` is a vector of `m` named tuples, the ``\mu`` of each trajectory. Every ``\mu`` has the
+keys of the first, in the same order and at every level, and the same sizes of its leaves. The data
+and every ``\mu`` have one element type. An `EnsembleSolution` of a problem with `(q, p)` states
+gives its `q` and `p` and the parameters of its problem.
+
+# The batch
+
+A `Batch(batch_size)` draws, in one epoch, each pair of a start position and a trajectory once, in
+minibatches of `batch_size`; the last minibatch is shorter. A drawn minibatch is the triple
+`(input, output, μs)`: the state at the start position, the state one time step later, and the
+``\mu`` of each sample. [`train!`](@ref) on a `ParametricDataLoader` trains on
+[`ParametricLoss`](@ref)`()` when no `loss` is given, so the network is applied with
+[`apply_parametric`](@ref), one time step per sample. A `Batch` with a sequence length or a
+prediction window other than one is an `ArgumentError`.
+
+A problem without system parameters is trained with a [`DataLoader`](@ref).
+"""
+struct ParametricDataLoader{
+    T, AT <: ArrayOrNamedTuple{T}, PT <: AbstractVector{<:NamedTuple}}
+    input::AT
+    parameters::PT
+    input_time_steps::Int
+    n_params::Int
+
+    function ParametricDataLoader(input::AT,
+            parameters::PT) where {T,
+            AT <: Union{AbstractArray{T, 3},
+                NamedTuple{(:q, :p), Tuple{AT3, AT3}} where {AT3 <: AbstractArray{T, 3}}},
+            PT <: AbstractVector{<:NamedTuple}}
+        input_time_steps, n_params = data_size(input)
+        n_params ≥ 1 && length(parameters) == n_params || throw(ArgumentError(
+            "the data hold $(n_params) trajectories and there are $(length(parameters)) " *
+            "system parameters; a `ParametricDataLoader` takes one for each trajectory, " *
+            "and at least one trajectory."))
+        layout = parameterlayout(first(parameters))
+        for (j, μ) in enumerate(parameters)
+            S = parameter_eltype(μ)
+            S === T || S === Union{} ||
+                throw(ArgumentError(
+                    "the data have the element type $(T) and a system parameter has the " *
+                    "element type $(S); a `ParametricDataLoader` takes one element type."))
+            parameterlayout(μ) == layout || throw(ArgumentError(
+                "system parameter $(j) differs from system parameter 1 in its keys, their " *
+                "order or the sizes of its leaves; a `ParametricDataLoader` takes one shape."))
+        end
+        new{T, AT, PT}(input, parameters, input_time_steps, n_params)
+    end
+end
+
+# The number of time steps and the number of trajectories of the data.
+data_size(input::AbstractArray{<:Any, 3}) = (size(input, 2), size(input, 3))
+
+function data_size(input::NamedTuple{(:q, :p)})
+    size(input.q) == size(input.p) || throw(ArgumentError(
+        "`q` has the size $(size(input.q)) and `p` has the size $(size(input.p))."))
+    data_size(input.q)
+end
+
+function ParametricDataLoader(ensemble_solution::EnsembleSolution{T,
+        T1,
+        Vector{ST}}) where {
+        T, T1, TuT, TT <: TimeSeries{T1},
+        ST <: GeometricSolution{T, T1, TT, NamedTuple{(:t, :q, :p, :q̇, :ṗ), TuT}}}
+    ParametricDataLoader(
+        (q = data_tensor([s.q for s in ensemble_solution.s]),
+            p = data_tensor([s.p for s in ensemble_solution.s])),
+        ensemble_solution.problem.parameters)
+end
+
+function networkbackend(dl::ParametricDataLoader{T, <:AbstractArray}) where {T}
+    networkbackend(dl.input)
+end
+function networkbackend(dl::ParametricDataLoader{T, <:NamedTuple}) where {T}
+    networkbackend(dl.input.q)
+end
+
+# The start positions are those with a target `prediction_window` steps later; the network takes one
+# time step per sample and returns one, so the sequence length and the prediction window are one.
+function (batch::Batch)(dl::ParametricDataLoader)
+    batch.seq_length == 1 || throw(ArgumentError(
+        "a `ParametricDataLoader` draws one time step per sample, and the `Batch` asks for a " *
+        "sequence length of $(batch.seq_length)."))
+    batch.prediction_window == 1 || throw(ArgumentError(
+        "a `ParametricDataLoader` draws a target of one time step per sample, and the `Batch` " *
+        "asks for a prediction window of $(batch.prediction_window)."))
+    n = dl.input_time_steps - batch.prediction_window
+    n ≥ 1 || throw(ArgumentError(
+        "the `ParametricDataLoader` has $(dl.input_time_steps) time steps and the `Batch` asks " *
+        "for a prediction window of $(batch.prediction_window), which leaves no minibatch."))
+    batch_over_two_axes(batch, n, dl.n_params)
+end
+
+# The input at the start position, the target `seq_length` steps later, and the `μ` of the trajectory
+# of each sample, which row 2 of the batch indices names.
+function convert_input_and_batch_indices_to_array(dl::ParametricDataLoader, batch::Batch,
+        batch_indices_tuple::Vector{Tuple{Int, Int}})
+    indices = convert_vector_of_tuples_to_matrix(networkbackend(dl), batch_indices_tuple)
+    (gather(dl.input, indices, 0, batch.seq_length),
+        gather(dl.input, indices, batch.seq_length, batch.prediction_window),
+        [dl.parameters[trajectory] for (_, trajectory) in batch_indices_tuple])
+end

src/data_loader/train.jl

@@ -2,7 +2,8 @@
     train!(nn, opt, dl, batch, n_epochs; loss, show_progress = true)
 
 Train the neural network `nn` with the `GeometricOptimizers.TrainingOptimizer` `opt` on the data of
-the [`DataLoader`](@ref) `dl` for `n_epochs` epochs, and return the loss history.
+the [`DataLoader`](@ref) or [`ParametricDataLoader`](@ref) `dl` for `n_epochs` epochs, and return
+the loss history.
 
 Each epoch draws the minibatches of `dl` with the [`Batch`](@ref) `batch` and takes one step per
 minibatch: `loss` computes the loss on the minibatch and its gradient, and
@@ -15,8 +16,9 @@ gradient, a zero leaf included. Entry `i` of the history is the mean of the mini
 trained through a [`ZygotePullback`](@ref). Without it, a [`TransformerIntegrator`](@ref) with a
 `Batch{:Transformer}` trains on [`TransformerLoss`](@ref)`(batch)`, a
 [`NeuralNetworkIntegrator`](@ref) with a `Batch{:FeedForward}` on `FeedForwardLoss()`, and an
-autoencoder with a `Batch{:FeedForward}` on [`AutoEncoderLoss`](@ref)`()`; any other combination
-needs `loss`, and `train!` raises an `ArgumentError` without it.
+autoencoder with a `Batch{:FeedForward}` on [`AutoEncoderLoss`](@ref)`()`, and any network on a
+`ParametricDataLoader` on [`ParametricLoss`](@ref)`()`; any other combination needs `loss`, and
+`train!` raises an `ArgumentError` without it.
 
 A network of HamiltonianNeuralNetworks.jl trains on vector-field data with
 `loss = SymbolicPullback(arch)` of that package, and on consecutive states with a loss such as its
@@ -29,9 +31,9 @@ with the same `opt` resumes the training where the first one stopped; the only r
 `train!` draws are those of the minibatch shuffle. `n_epochs < 1` is an `ArgumentError`, raised
 before any step.
 """
-function train!(nn::NeuralNetwork, opt::TrainingOptimizer, dl::DataLoader, batch::Batch,
-        n_epochs::Integer;
-        loss::Union{NetworkLoss, AbstractPullback} = NetworkLoss(nn.architecture, batch),
+function train!(nn::NeuralNetwork, opt::TrainingOptimizer,
+        dl::Union{DataLoader, ParametricDataLoader}, batch::Batch, n_epochs::Integer;
+        loss::Union{NetworkLoss, AbstractPullback} = default_loss(nn, dl, batch),
         show_progress::Bool = true)
     n_epochs ≥ 1 ||
         throw(ArgumentError("`train!` needs at least one epoch; `n_epochs` is $(n_epochs)."))
@@ -66,6 +68,13 @@ end
 _as_pullback(l::NetworkLoss) = ZygotePullback(l)
 _as_pullback(p::AbstractPullback) = p
 
+# A `DataLoader` takes the default loss of the architecture and the batch; a `ParametricDataLoader`
+# draws the system parameters of each sample, which `ParametricLoss` hands to the network.
+function default_loss(nn::NeuralNetwork, ::DataLoader, batch::Batch)
+    NetworkLoss(nn.architecture, batch)
+end
+default_loss(::NeuralNetwork, ::ParametricDataLoader, ::Batch) = ParametricLoss()
+
 # The default loss of an architecture and a batch. The two-argument methods live here and not beside
 # `NetworkLoss(::AutoEncoder)` in `src/loss/losses.jl`, because `Batch` is included after `src/loss/`.
 function AbstractNeuralNetworks.NetworkLoss(

src/loss/losses.jl

@@ -226,3 +226,29 @@ function (loss::ReducedLoss)(model::Chain, ps::NetworkParameters,
         input::ArrayOrNamedTuple, output::ArrayOrNamedTuple)
     _compute_loss(loss.decoder(model(loss.encoder(input), ps)), output)
 end
+
+@doc raw"""
+    ParametricLoss()
+
+Make an instance of the loss of a network whose forward pass takes the parameters ``\mu`` of the
+system beside its input.
+
+`ParametricLoss()(model, ps, input, output, μ)` applies the network with
+[`apply_parametric`](@ref) and compares the result with `output` in the relative ``L_2`` norm:
+
+```math
+\mathrm{loss}(\mathcal{NN}, \mathrm{input}, \mathrm{output}, \mu) = ||\mathcal{NN}(\mathrm{input}, \mu) - \mathrm{output}|| / ||\mathrm{output}||.
+```
+
+This is `FeedForwardLoss` with the system parameters passed to the network. `μ` is one
+`OptionalParameters` for one sample, or a vector of them, one per sample of a batch, as a
+[`ParametricDataLoader`](@ref) draws it. `train!` on a `ParametricDataLoader` trains on this loss
+when no `loss` is given.
+"""
+struct ParametricLoss <: NetworkLoss end
+
+function (loss::ParametricLoss)(model::Chain, ps::NetworkParameters,
+        input::ArrayOrNamedTuple, output::ArrayOrNamedTuple,
+        μ::Union{OptionalParameters, AbstractVector})
+    _compute_loss(apply_parametric(model, input, μ, ps), output)
+end