Make the four names HamiltonianNeuralNetworks reaches public and documented

CHANGELOG.md

@@ -119,7 +119,7 @@ breaking release).
 
 - **`ForcingLayer`, `ForcingLayerQ`, `ForcingLayerP`, `ForcingLayerQP` and `ForcedSympNet` add a learned non-conservative term to a symplectic integrator.** The five names are `public` and not exported. A forcing layer changes `p` and leaves `q` alone, by adding the output of a `Chain` of its own to `p`; the `Q`, `P` or `QP` of its type says what that network reads — `q`, `p`, or both — and `parameters` fixes the length of the flat vector of system parameters it takes beside the state. `return_parameters = true` makes it return `(output, μ)` instead of the output, so that the layer after it in a `Chain` takes the system parameters too. `dim` is the full phase-space dimension, as for `GradientLayerQ`, and the forcing has `dim ÷ 2` rows; an odd dimension and a forcing type other than `:Q`, `:P` and `:QP` are `ArgumentError`s at construction. `ForcedSympNet(d; upscaling_dimension = 2d, n_layers = 2, n_forcing_layers = 2, activation = tanh, init_upper = true, forcing_type = :P)` is the chain of a `GSympNet` with a forcing layer in front of every layer that changes `p`: it holds `n_layers` gradient layers and one forcing layer per step that changes `p`, it takes no parameters of the system, and it answers `dim` as the `SympNet`s do. A forcing layer holds a `Chain`, so its parameters are a set within the set of the network, which `train!`, `ZygotePullback` and `optimization_step!` carry as they carry any nested set. `scripts/reproduction/damped_oscillator_forced_sympnet.jl` trains one on the damped harmonic oscillator through `TrainingOptimizer` and `train!`. The code is ported from pull request #207 by Benedikt Brantner (relates to #207); there a forcing layer raised "This shouldn't be the case!" for any other tuple, which is a `MethodError` now, the forcing type of the architecture was accepted unchecked and failed at the first call, and the script trained through `Optimizer(AdamOptimizer(), nn)`.
 
-- **`apply_parametric`, `ParametricLoss` and `ParametricDataLoader` train a network whose forward pass takes the parameters of the system beside its input.** `apply_parametric` and `ParametricLoss` are `public` and not exported, so a caller qualifies them or names them in a `using GeometricMachineLearning: …`; `ParametricDataLoader` is `ReducedComplexityModeling`'s and is public there, not here, so a caller takes it from that package or writes `GeometricMachineLearning.ParametricDataLoader`. `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 `μ`, repeated in every column for a matrix and one `μ`, and repeated in every column of every slice for a 3-tensor 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`.
+- **`apply_parametric`, `ParametricLoss` and `ParametricDataLoader` train a network whose forward pass takes the parameters of the system beside its input.** `apply_parametric` and `ParametricLoss` are `public` and not exported, so a caller qualifies them or names them in a `using GeometricMachineLearning: …`; `ParametricDataLoader` is `ReducedComplexityModeling`'s and is public there, not here, so a caller takes it from that package or writes `GeometricMachineLearning.ParametricDataLoader`. `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 `public` `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 `μ`, repeated in every column for a matrix and one `μ`, and repeated in every column of every slice for a 3-tensor 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`.
 
 - **`WideResNetLayer`, `ParametricResNetLayer` and `ParametricResNet` add a ResNet whose hidden layer has a width of its own, and a baseline whose blocks take the parameters of the system.** The three names are `public` and not exported, so a caller qualifies them or names them in a `using GeometricMachineLearning: …`. `WideResNetLayer(dim, width, activation)` computes `x + σ.(W_down σ.(W_up x .+ b_up) .+ b)` for a vector, a matrix, a tensor or a `(q, p)` pair, and `ResNet(dim, n_blocks, width, activation)` and `ResNet(dl, n_blocks, width)` build it wherever `width` differs from the system dimension; a width equal to the dimension keeps the plain `ResNetLayer`, so `ResNet(dim, n_blocks, activation)` and its recorded outputs are unchanged. `ParametricResNet(dim; width = dim, n_blocks = 1, activation = tanh, parameters)` appends the flattened `μ` below the state at the input of every upscaling weight, through `ParametricResNetLayer`s that pass `μ` on to the next layer, and the `NeuralNetwork` built around it gains the functor `nn(x, μ)`, which forwards to `apply_parametric` and so takes one sample with one `μ` or a batch with a vector of `μ`; an input shape that `apply_parametric` does not take raises its `MethodError`. `parameterlength` of both layers counts what `initialparameters` allocates: `width·(dim + 1) + dim·(width + 1)` without system parameters, and `width·(dim + |μ|) + width + dim·(width + 1)` with them. `scripts/reproduction/parametric_resnet.jl` trains the baseline on the analytic data of the harmonic oscillator, with the time as the parameter of the system, through `ParametricDataLoader`, `ParametricLoss` and `train!`. This is the port of pull request #207 by Benedikt Brantner (relates to #207), with `parameterlength(::ParametricResNetLayer)` corrected -- #207's `(width + |μ|)(dim + 1) + dim(width + 1)` counts the entries only at `width = dim + 1` -- and the two defects that #207's own test turned up kept fixed: `ParametricResNet(::DataLoader, …)` forwards `parameters` and `activation` instead of dropping them, and `nn(x, μ)` is the functor of the network and not AbstractNeuralNetworks' generic one.
 
@@ -154,6 +154,8 @@ breaking release).
 
 - **The tutorial *Multi-Modal Symplectic Autoencoders* and the script `scripts/reproduction/multimodal_symplectic_autoencoder.jl` train a `MultiModalSymplecticAutoencoder` with a grid, a DLRA and a particle frontend in three stages, on synthetic 1D1V data, in `Float32` and `Float64`.** Stage I sets the bases of the grid frontend and the symplectic lift of the core with `solve!`, and the bases of the DLRA frontend start at the grid frontend's plus 0.1 times a random matrix, orthonormalized. Stage II trains the `CharbonnierRelativeLoss` of the grid modality, `CrossModalLoss(:dlra, :grid)` and the error of the particle states translated from the grid modality, with `CompositeMethod(; manifold = FactoredStiefelAdam(), array = Adam())`. Stage III integrates the reduced system of a pendulum Hamiltonian on the frontend coordinates with `HRedSys`, `ImplicitMidpoint()`, `pullback_mode = :vjp` and the background as `q_ref, p_ref`. The script asserts that the projection error after Stage I is below `1e-2`; that Stage II halves the sum of the reconstruction and `CrossModalLoss` terms and keeps every Stiefel leaf orthonormal to `50 eps(T)`; that halving the time step divides the largest energy error by 3 to 5; and that the decoded initial field equals the snapshot to the Stage I projection error. `test/integration/multimodal_pipeline.jl`, in the group `slow`, runs the stages and four defects of them, each of which one of these assertions catches, and checks that the cross-reconstruction term of `CrossModalLoss` alone halves its value where a cycle term in its place does not. The data test the pipeline, not the physics. The script needs `GeometricOptimizers` 0.9.2 for `FactoredStiefelAdam`. The tutorial links `FactoredStiefelAdam` into the `GeometricOptimizers` manual, whose committed inventory, at v0.9.2, holds it.
 
+- **`flatten_system_parameters`, `concatenate_array_with_parameters`, `assign_q_and_p` and `check_forcing_type` are `public` and documented, and are not exported.** [HamiltonianNeuralNetworks.jl](https://github.com/JuliaGML/HamiltonianNeuralNetworks.jl) calls the four functions by their qualified names, so they are part of the API of this package. Each has a docstring that names its methods, what they return and what they throw, in the *Library Functions* section of the neural network integrators, SympNet and tensors pages; `using GeometricMachineLearning` does not bring them in.
+
 ### 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`.

docs/src/architectures/neural_network_integrators.md

@@ -61,6 +61,8 @@ GeometricMachineLearning.WideResNetLayer
 GeometricMachineLearning.ResNetLayer
 GeometricMachineLearning.ParametricResNet
 GeometricMachineLearning.ParametricResNetLayer
+GeometricMachineLearning.flatten_system_parameters
+GeometricMachineLearning.concatenate_array_with_parameters
 integrate(::NeuralNetwork{<:NeuralNetworkIntegrator}, ::NamedTuple{(:q, :p), Tuple{AT, AT}}) where {T, AT<:AbstractVector{T}}
 TransformerIntegrator
 integrate(::NeuralNetwork{<:TransformerIntegrator}, ::NamedTuple{(:q, :p), Tuple{AT, AT}}) where {T, AT<:AbstractMatrix{T}}

docs/src/architectures/sympnet.md

@@ -247,6 +247,7 @@ LASympNet
 GSympNet
 GeometricMachineLearning.ForcedSympNet
 GeometricMachineLearning.ForcingLayer
+GeometricMachineLearning.check_forcing_type
 GeometricMachineLearning.ForcingLayerQ
 GeometricMachineLearning.ForcingLayerP
 GeometricMachineLearning.ForcingLayerQP

docs/src/arrays/tensors.md

@@ -6,4 +6,5 @@ We typically store training data as *tensors with three axes* in `GeometricMachi
 
 ```@docs
 GeometricMachineLearning._mul
+GeometricMachineLearning.assign_q_and_p
 ```

src/GeometricMachineLearning.jl

@@ -236,6 +236,10 @@ include("architectures/parametric_resnet.jl")
 
 # The training of a network that takes the parameters of the system beside its input.
 public apply_parametric, ParametricLoss
+# `HamiltonianNeuralNetworks` calls these four functions by their qualified names, so they are
+# public and documented. They are not exported.
+public flatten_system_parameters, concatenate_array_with_parameters, assign_q_and_p,
+       check_forcing_type
 
 include("reduced_system/reduced_system.jl")
 

src/architectures/forced_sympnet.jl

@@ -64,7 +64,7 @@ struct ForcedSympNet{FT, AT} <: NeuralNetworkIntegrator
             init_upper = g_init_upper_default,
             forcing_type::Symbol = :P)
         _check_forcing_dimension(dim)
-        _check_forcing_type(forcing_type)
+        check_forcing_type(forcing_type)
         new{forcing_type, typeof(activation)}(
             Int(dim), Int(upscaling_dimension), Int(n_layers), Int(n_forcing_layers),
             activation, init_upper)

src/kernels/assign_q_and_p.jl

@@ -1,6 +1,18 @@
-# Split `x` along its first axis into the `(q, p)` `NamedTuple` of its first `N` rows and the `N`
-# rows after them. Rows from `2N + 1` on are ignored. Slicing works on every backend, and Zygote
-# differentiates it, so the split needs neither a kernel nor an `rrule`.
+"""
+    assign_q_and_p(x::AbstractArray, N::Integer)
+
+Split `x` along its first axis, and return the `(q, p)` `NamedTuple` of the rows at the indices
+`1:N` and at the indices `N + 1:2N`, as indexing `x` returns them: copies for an `Array`, and
+ranges for a range. The indices are these whatever the first axis of `x` is, so an offset axis is
+split at index `1`, not at its first row. Rows at other indices are ignored, and the other axes are
+kept. For `N ≤ 0` both halves are empty.
+
+Throws a `BoundsError` when the first axis of `x` does not hold the indices `1:2N`, and an
+`ArgumentError` when `x` has no axis.
+
+Slicing works on every backend, and Zygote differentiates it, so the split needs neither a kernel
+nor an `rrule`.
+"""
 function assign_q_and_p(x::AbstractArray, N::Integer)
     rest = ntuple(Returns(:), ndims(x) - 1)
     (q = x[1:N, rest...], p = x[(N + 1):(2N), rest...])

src/layers/forcing_layer.jl

@@ -5,7 +5,13 @@ const f_activation_default = tanh
 # dissipation enters; the `Q` and `P` suffixes name the component the forcing *reads*.
 const FORCING_TYPES = (:Q, :P, :QP)
 
-function _check_forcing_type(type::Symbol)
+"""
+    check_forcing_type(type::Symbol)
+
+Return `type` when it is one of the forcing types `:Q`, `:P` and `:QP`, and throw an `ArgumentError`
+otherwise. The type names the component that the forcing reads; every type adds to `p`.
+"""
+function check_forcing_type(type::Symbol)
     type in FORCING_TYPES || throw(ArgumentError(
         "the forcing type is `:$(type)`; it is `:Q`, `:P` or `:QP`."))
     type
@@ -76,10 +82,10 @@ function ForcingLayer(dim::Integer, width::Integer = dim,
         parameters::OptionalParameters = NullParameters(),
         return_parameters::Bool = false, type::Symbol = :P)
     _check_forcing_dimension(dim)
-    _check_forcing_type(type)
+    check_forcing_type(type)
     # The parameters of the system fix the length of the flat vector the layer takes beside its
     # state. They are not stored: the `μ` the layer is called with is the one that enters.
-    parameter_length = length(first(_flatten_system_parameters(parameters)))
+    parameter_length = length(first(flatten_system_parameters(parameters)))
     model = _forcing_model(
         Int(dim), Int(width), Int(nhidden), parameter_length, activation, type)
     ForcingLayer{Int(dim), Int(dim), type, return_parameters, typeof(model)}(

src/layers/parametric_resnet_layer.jl

@@ -36,7 +36,7 @@ end
 
 function ParametricResNetLayer(dim::Integer, width::Integer, activation = identity;
         parameters::OptionalParameters = NullParameters(), return_parameters::Bool = false)
-    flat_parameters, layout = _flatten_system_parameters(parameters)
+    flat_parameters, layout = flatten_system_parameters(parameters)
     ParametricResNetLayer{dim, dim, typeof(activation), typeof(layout), return_parameters}(
         width, activation, length(flat_parameters), layout)
 end