From c73b23a961df1b57b81322ae228e4efae81b90b4 Mon Sep 17 00:00:00 2001 From: benedict-96 Date: Wed, 19 Aug 2026 18:14:48 +0900 Subject: [PATCH 1/4] Rebase the parametric generalized Hamiltonian neural networks onto main Reconstructs PR #207 on top of 0.5.0. The branch's own 53 commits are collapsed into this one; the three that followed them are not carried over. `cc1a786f` "Fix GML compilation and add generated artifacts" was 123 MB of build output -- `docs/build.zip` (77 MB), `docs/build 2.zip` (47 MB), eleven LaTeX `.aux`/`.log`/`.fls`/`.fdb_latexmk` files and two `go_migration_inspection_*.txt` -- around four lines of source. `.gitignore` already covers all of it bar the inspection logs. `8d243592` "Resolve merge conflicts" was main's own HDF5 migration re-applied by hand, byte-identical blobs, redundant after a rebase. `43d89206` merged main only as far as `d07b4c26`, so the branch never saw the GeometricOptimizers separation, and its conflict resolution left `src/GeometricMachineLearning.jl` half from each side: `go_bridges.jl`, `src/arrays/` and `src/manifolds/` were `include`d again next to main's `import GeometricOptimizers`, which is `Method overwriting is not permitted during Module precompilation` on Julia 1.10. `4c0e2684` is kept -- it is a real seven-line fix to the symbolic `Jacobian` broadcast and two `build_nn_function` calls. Root cause of the repeated bad merges: main reformatted `src/GeometricMachineLearning.jl` from a 4-space-indented module body to column 0 and the branch never did, so every merge conflicted on the whole file. Here main's version is taken as-is and only the branch's `include`s and `export`s are re-applied. What the feature adds `GeneralizedHamiltonianArchitecture` is implemented; it used to be a stub whose constructor threw. It composes `n_integrators` symplectic Euler steps, each differentiating a learned kinetic or potential energy. Around it: `ForcedGeneralizedHamiltonianArchitecture` and `ForcedSympNet` with `ForcingLayer`s, `ParametricDataLoader`, `ParametricLoss`, a `SymbolicPullback` for it, `ParametricResNet`, and `QPT2`/`QPTOAT2`. Adaptation to the current dependencies The branch predates SymbolicNeuralNetworks 0.5. `symbolize!` is gone -- `symbolic_variables` replaces it -- and symbolic variables are scalar `Num`s rather than `Symbolics.Arr`s (SNN#14), so the parametric `SymbolicPullback` is rewritten in the shape of `SymbolicNeuralNetworks.SymbolicPullback`, with `symbolic_derivative` and `ParameterGradient` in place of the hand-rolled closure and `semi_flatten_network_parameters`. Five `Symbolics.Arr` special cases, three of them type piracy, are unreachable now and deleted. `GeometricProblems.default_parameters` is a function since 0.8, so the three tests call it. `src/optimizers/optimizer.jl` keeps main's version: the branch widened `_optimization_step!` and added `rgrad` methods for `NeuralNetworkParameters` and `nothing` gradients, and the 0.5.0 rewrite covers all of it -- `_tree_optim_step!` already skips a `nothing` block. NeuralNetworkParameters instead of ParameterHandling The system parameters are flattened into the network input. The branch used `ParameterHandling` and pirated three `flatten` methods on it; that does not work, because GeometricOptimizers defines `ParameterHandling.flatten(x)` with an unbound type parameter and that method wins -- D6 in NeuralNetworkParameters' PLAN.md, hit in practice. `flatten`/`unflatten` from NeuralNetworkParameters replace it. Since AbstractNeuralNetworks already exports `params`, no method on a foreign type is needed, and the layout stored in a layer is now a value rather than a closure. This makes registering NeuralNetworkParameters a prerequisite for merging: GML supports Julia 1.10, where `[sources]` does not exist. Bugs fixed on the way -- all on the training path, which nothing ran `concatenate_array_with_parameters(::AbstractMatrix, ::AbstractVector)` used `vcat` where it needs `hcat`, collapsing a batch into one long vector. `optimize_for_one_epoch!` called `_unpack_tuple`, which has never existed in this package. And `Zygote` differentiates *through* the `NeuralNetworkParameters` struct, so every nesting level of the gradient comes back wrapped in `(params = ...,)`; these architectures nest, and `_get_params` only unwraps the top, so `_unwrap_gradient` recurses. `test/generalized_hamiltonian_neural_networks/pghnn_training_test.jl` now covers that path. The other two tests are rewritten: the data-loader one asserted which shuffled batch holds which parameters, which depends on the RNG stream of the Julia version, and now asserts the correspondence itself; the pullback one had no `@test` at all and re-defined in the test file what `src/` now provides, and now checks the symbolic gradient against the summed per-sample `Zygote` gradients, which agree to 4e-16. Type piracy Deleted where it was free: `Base.NamedTuple(::NeuralNetworkParameters)` is `params`, and the `ParameterHandling` and `Symbolics.Arr` methods are gone with their callers. What remains -- `applychain`, four `Chain` functors, `SymbolicNeuralNetworks.Jacobian`, `networkbackend(::LazyArrays.ApplyArray)`, `h5save(::HDF5.Group, ::NeuralNetworkParameters, ...)` and the `SymbolicPullback` call operators -- carries a `TODO` naming its proper home. Also: `GeometricProblems` and `Printf` are dropped from `[deps]`, where the branch had duplicated them out of `[extras]` and nothing in `src/` uses them; `SymplecticEuler`, `SymplecticEulerA` and `SymplecticEulerB` are no longer exported, the training methods that used to carry those names being `SymplecticEulerIntegrator*` now; and the training loop takes `Union{DataLoader, ParametricDataLoader}` rather than an untyped argument. Co-Authored-By: Claude Opus 5 (1M context) --- CHANGELOG.md | 40 +++ Project.toml | 3 + docs/src/GeometricMachineLearning.bib | 15 +- .../hamiltonian_neural_network.md | 17 +- docs/src/architectures/sympnet.md | 1 + docs/src/reduced_order_modeling/losses.md | 1 + .../structure_preservation/symplecticity.md | 2 + ext/HDF5Ext.jl | 10 + ...rcedGeneralizedHamiltonianNeuralNetwork.jl | 161 +++++++++ ...ndentHarmonicOscillatorParametricResnet.jl | 150 ++++++++ scripts/Train_DampedOscillator_QP.jl | 114 ++++++ scripts/forcing_layers_parameter_number.jl | 11 + src/GeometricMachineLearning.jl | 27 +- ..._generalized_hamiltonian_neural_network.jl | 35 ++ src/architectures/forced_sympnet.jl | 65 ++++ .../generalized_hamiltonian_neural_network.jl | 335 ++++++++++++++++++ .../hamiltonian_neural_network.jl | 114 ------ src/architectures/parametric_resnet.jl | 37 ++ src/architectures/resnet.jl | 13 +- .../standard_hamiltonian_neural_network.jl | 85 +++++ src/data_loader/batch.jl | 7 +- src/data_loader/optimize.jl | 4 +- src/data_loader/parametric_data_loader.jl | 143 ++++++++ src/layers/forcing_dissipation_layers.jl | 163 +++++++++ src/layers/parametric_resnet_layer.jl | 64 ++++ src/layers/sympnets.jl | 12 +- src/layers/wide_resnet.jl | 32 ++ src/loss/losses.jl | 22 ++ src/pullbacks/symbolic_hnn_pullback.jl | 69 ++++ src/pullbacks/zygote_pullback.jl | 6 + src/training_method/symplectic_euler.jl | 20 +- src/utils.jl | 84 ++++- .../parametric_data_loader_test.jl | 45 +++ ...hnn_symbolic_pullback_single_layer_test.jl | 63 ++++ .../pghnn_training_test.jl | 47 +++ ...alized_hamiltonian_neural_networks_test.jl | 25 ++ test/runtests.jl | 12 + test/train!/test_method.jl | 4 +- test/training_phnn.jl | 8 +- 39 files changed, 1910 insertions(+), 156 deletions(-) create mode 100644 scripts/TimeDependentHarmonicOscillatorForcedGeneralizedHamiltonianNeuralNetwork.jl create mode 100644 scripts/TimeDependentHarmonicOscillatorParametricResnet.jl create mode 100644 scripts/Train_DampedOscillator_QP.jl create mode 100644 scripts/forcing_layers_parameter_number.jl create mode 100644 src/architectures/forced_generalized_hamiltonian_neural_network.jl create mode 100644 src/architectures/forced_sympnet.jl create mode 100644 src/architectures/generalized_hamiltonian_neural_network.jl create mode 100644 src/architectures/parametric_resnet.jl create mode 100644 src/architectures/standard_hamiltonian_neural_network.jl create mode 100644 src/data_loader/parametric_data_loader.jl create mode 100644 src/layers/forcing_dissipation_layers.jl create mode 100644 src/layers/parametric_resnet_layer.jl create mode 100644 src/layers/wide_resnet.jl create mode 100644 test/data_loader/parametric_data_loader_test.jl create mode 100644 test/generalized_hamiltonian_neural_networks/pghnn_symbolic_pullback_single_layer_test.jl create mode 100644 test/generalized_hamiltonian_neural_networks/pghnn_training_test.jl create mode 100644 test/generalized_hamiltonian_neural_networks_test.jl diff --git a/CHANGELOG.md b/CHANGELOG.md index bf48a4243..f4cf98c39 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -108,6 +108,9 @@ breaking release). matrix cotangent through `_matrix_cotangent`, which fixes the rank *and* gives it the array type of the primal. +- `concatenate_array_with_parameters(::AbstractMatrix, ::AbstractVector)` concatenated a batch with + `vcat` rather than `hcat`, collapsing it into a single long vector ([#207](https://github.com/JuliaGNI/GeometricMachineLearning.jl/pull/207)). + ### Changed - **The kernel `rrule`s honour the ChainRules interface for thunked cotangents.** Twelve pullbacks @@ -139,8 +142,45 @@ breaking release). behind — the value of the innermost `h5save`, an implementation detail of the traversal. Returning the path is what `NeuralNetworkParameters.save(filename, ps)` does, so the two now agree. +- **`SymplecticEuler`, `SymplecticEulerA` and `SymplecticEulerB` are no longer exported** ([#207](https://github.com/JuliaGNI/GeometricMachineLearning.jl/pull/207)). The + names now belong to the layer type of the generalized architectures; the *training methods* they + used to name are `SymplecticEulerIntegrator`, `SymplecticEulerIntegratorA` and + `SymplecticEulerIntegratorB`. `SEuler`, `SEulerA` and `SEulerB`, which is how they are constructed, + are unchanged. +- `src/architectures/hamiltonian_neural_network.jl` is split: it keeps the abstract + `HamiltonianArchitecture`, and `StandardHamiltonianArchitecture` moves to + `standard_hamiltonian_neural_network.jl`. `hamiltonian_vector_field` is narrowed from + `::HamiltonianArchitecture` to `::StandardHamiltonianArchitecture` accordingly ([#207](https://github.com/JuliaGNI/GeometricMachineLearning.jl/pull/207)). + ### Added +**Parametric generalized Hamiltonian neural networks (PGHNNs)** ([#207](https://github.com/JuliaGNI/GeometricMachineLearning.jl/pull/207)). A family of architectures +whose forward pass takes the parameters of the *system* alongside the state, so one network covers a +whole parameter range rather than a single problem instance. + +- **`GeneralizedHamiltonianArchitecture`** is implemented. It used to be a stub whose constructor + threw `error("GHNN still has to be implemented!")`. It composes `n_integrators` symplectic Euler + steps, each of which differentiates a learned kinetic or potential energy — + `SymbolicKineticEnergy` and `SymbolicPotentialEnergy`, built into an executable gradient by + `build_gradient`. The system parameters reach the network as extra input components, flattened + with `NeuralNetworkParameters`' `flatten`/`unflatten`. +- **`ForcedGeneralizedHamiltonianArchitecture`** and **`ForcedSympNet`**, which add `ForcingLayer`s + for forcing and dissipation in the `q`, `p` or both coordinates, following the + Lagrange–d'Alembert integrator of [marsden2001discrete](@cite). +- **`ParametricDataLoader`**, which carries one set of system parameters per trajectory and hands + the matching parameters to each sample of a batch. Built from an `EnsembleSolution` whose members + were integrated at different parameters. +- **`ParametricLoss`**, `FeedForwardLoss` with the system parameters threaded through, and a + `SymbolicPullback(nn, ::ParametricLoss, system_params)` that differentiates it symbolically. + Building that pullback refuses `n_integrators > 1`: the symbolic expression grows + *multiplicatively* with the number of integrators, so the build does not finish. See + [#245](https://github.com/JuliaGNI/GeometricMachineLearning.jl/issues/245). +- **`ParametricResNet`** and a widened **`ResNet`**, which now takes a `width` separate from the + system dimension and uses `WideResNetLayer` when the two differ — the non-structure-preserving + baseline the PGHNNs are compared against. +- `QPT2` and `QPTOAT2`: `QPT`/`QPTOAT` with the array rank fixed but the two array *types* allowed to + differ, which is what splitting an input array into `q` and `p` produces. + - **`load(NeuralNetwork, h5, arch, prototype)`** — a parameter set of the right shape to rebuild the structured leaves against. It is the form that needs no registration: `rebuild` has a prototype to take the non-differentiable fields from, so the file's type tags and diff --git a/Project.toml b/Project.toml index 462957bc8..5ea404921 100644 --- a/Project.toml +++ b/Project.toml @@ -27,6 +27,9 @@ Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f" [weakdeps] HDF5 = "f67ccb44-e63f-5c2f-98bd-6dc0ccc4ba2f" +[sources] +NeuralNetworkParameters = {path = "/Users/mkraus/Datashare/Julia/NeuralNetworkParameters"} + [extensions] HDF5Ext = "HDF5" diff --git a/docs/src/GeometricMachineLearning.bib b/docs/src/GeometricMachineLearning.bib index e069aee6c..a4b523263 100644 --- a/docs/src/GeometricMachineLearning.bib +++ b/docs/src/GeometricMachineLearning.bib @@ -746,9 +746,10 @@ @article{bon2024optimal @article{kingma2014adam, title={Adam: a method for stochastic optimization}, - author={Kingma, DP}, + author={Kingma, Diederik P. and Ba, Jimmy Lei}, journal={arXiv preprint arXiv:1412.6980}, - year={2014} + year={2014}, + note={Published as a conference paper at ICLR 2015} } @article{toda1967vibration, @@ -923,6 +924,16 @@ @article{ge1988lie publisher={Elsevier} } +@article{marsden2001discrete, + title={Discrete mechanics and variational integrators}, + author={Marsden, Jerrold E and West, Matthew}, + journal={Acta numerica}, + volume={10}, + pages={357--514}, + year={2001}, + publisher={Cambridge University Press} +} + @article{otto2023learning, title={Learning nonlinear projections for reduced-order modeling of dynamical systems using constrained autoencoders}, author={Otto, Samuel E and Macchio, Gregory R and Rowley, Clarence W}, diff --git a/docs/src/architectures/hamiltonian_neural_network.md b/docs/src/architectures/hamiltonian_neural_network.md index be33ca27f..8c6234419 100644 --- a/docs/src/architectures/hamiltonian_neural_network.md +++ b/docs/src/architectures/hamiltonian_neural_network.md @@ -42,13 +42,28 @@ Here the derivatives (i.e. vector field data) ``\dot{q}_i^{(t)}`` and ``\dot{p}_ ## Library Functions ```@docs -GeometricMachineLearning.hamiltonian_vector_field(::HamiltonianArchitecture) +GeometricMachineLearning.hamiltonian_vector_field(::StandardHamiltonianArchitecture) GeometricMachineLearning.HamiltonianArchitecture GeometricMachineLearning.StandardHamiltonianArchitecture GeometricMachineLearning.HNNLoss GeometricMachineLearning.symbolic_hamiltonian_vector_field(::GeometricMachineLearning.SymbolicNeuralNetwork) GeometricMachineLearning.SymbolicPullback(::HamiltonianArchitecture) +GeometricMachineLearning.SymbolicEnergy +GeometricMachineLearning.SymbolicPotentialEnergy +GeometricMachineLearning.SymbolicKineticEnergy +GeometricMachineLearning.build_gradient +GeometricMachineLearning.SymplecticEulerA +GeometricMachineLearning.SymplecticEulerB GeometricMachineLearning.GeneralizedHamiltonianArchitecture +GeometricMachineLearning.ForcedGeneralizedHamiltonianArchitecture +GeometricMachineLearning.ForcingLayer +GeometricMachineLearning.ForcingLayerQ +GeometricMachineLearning.ForcingLayerP +GeometricMachineLearning.ForcingLayerQP +GeometricMachineLearning.ParametricDataLoader +GeometricMachineLearning.SymbolicPullback(::GeometricMachineLearning.NeuralNetwork, ::GeometricMachineLearning.ParametricLoss, ::GeometricMachineLearning.GeometricBase.OptionalParameters) +GeometricMachineLearning._flatten_system_parameters +GeometricMachineLearning._unwrap_gradient GeometricMachineLearning._processing GeometricMachineLearning._get_contents GeometricMachineLearning._get_params diff --git a/docs/src/architectures/sympnet.md b/docs/src/architectures/sympnet.md index 72cea0d3c..67978d107 100644 --- a/docs/src/architectures/sympnet.md +++ b/docs/src/architectures/sympnet.md @@ -215,6 +215,7 @@ is the predicted state. In the [example section](@ref "SympNets with `GeometricM SympNet LASympNet GSympNet +ForcedSympNet ``` ```@raw latex diff --git a/docs/src/reduced_order_modeling/losses.md b/docs/src/reduced_order_modeling/losses.md index ba933a840..f7398133b 100644 --- a/docs/src/reduced_order_modeling/losses.md +++ b/docs/src/reduced_order_modeling/losses.md @@ -46,6 +46,7 @@ where ``\mathbf{x}^{(t)}`` is the solution of the FOM at point ``t`` and ``\math TransformerLoss AutoEncoderLoss ReducedLoss +ParametricLoss projection_error reduction_error ``` diff --git a/docs/src/structure_preservation/symplecticity.md b/docs/src/structure_preservation/symplecticity.md index 1e46a0eaa..1fe7a3974 100644 --- a/docs/src/structure_preservation/symplecticity.md +++ b/docs/src/structure_preservation/symplecticity.md @@ -115,7 +115,9 @@ It is important to note that symplecticity is a very strong property[^2] that ma ```@docs PoissonTensor GeometricMachineLearning.QPT +GeometricMachineLearning.QPT2 GeometricMachineLearning.QPTOAT +GeometricMachineLearning.QPTOAT2 ``` ```@raw latex diff --git a/ext/HDF5Ext.jl b/ext/HDF5Ext.jl index 0853d4257..fbec87bfe 100644 --- a/ext/HDF5Ext.jl +++ b/ext/HDF5Ext.jl @@ -7,6 +7,16 @@ import AbstractNeuralNetworks: changebackend, NeuralNetworkBackend, Architecture # `AbstractNeuralNetworks` 0.7, which only re-binds them; reach for them where they are defined. import NeuralNetworkParameters: NetworkParameters, params, save, load +# A `NeuralNetworkParameters` nested inside a parameter tree -- the parameter-dependent +# architectures put one per sub-network. AbstractNeuralNetworks has `save(::H5DataStore, +# ::NeuralNetworkParameters)` for the top level only. +# +# TODO: type piracy -- `h5save` and `NeuralNetworkParameters` are both AbstractNeuralNetworks'. +# This belongs in ANN's own `ext/HDF5Ext.jl`, next to `h5save(::H5DataStore, ::NamedTuple, …)`. +function h5save(h5::HDF5.Group, p::NeuralNetworkParameters, path::AbstractString) + h5save(h5, params(p), path) +end + # --------------------------------------------------------------------------- # changebackend — new methods for GML special array types # diff --git a/scripts/TimeDependentHarmonicOscillatorForcedGeneralizedHamiltonianNeuralNetwork.jl b/scripts/TimeDependentHarmonicOscillatorForcedGeneralizedHamiltonianNeuralNetwork.jl new file mode 100644 index 000000000..e9ca22ce1 --- /dev/null +++ b/scripts/TimeDependentHarmonicOscillatorForcedGeneralizedHamiltonianNeuralNetwork.jl @@ -0,0 +1,161 @@ +using HDF5 +using GeometricMachineLearning +using GeometricMachineLearning: QPT, QPT2, Activation, ParametricLoss, SymbolicNeuralNetwork, SymbolicPullback +using CairoMakie +using NNlib: relu + +# PARAMETERS +omega = 1.0 # natural frequency of the harmonic Oscillator +Omega = 3.5 # frequency of the external sinusoidal forcing +F = .9 # amplitude of the external sinusoidal forcing +ni_dim = 10 # number of initial conditions per dimension (so ni_dim^2 total) +T = 2π * 20 +nt = 1000 # number of time steps +dt = T/nt # time step + +# Generating the initial condition array +IC = vec( [(q=q0, p=p0) for q0 in range(-1, 1, ni_dim), p0 in range(-1, 1, ni_dim)] ) + +# Generating the solution array +ni = ni_dim^2 +q = zeros(Float64, ni, nt+1) +p = zeros(Float64, ni, nt+1) +t = collect(dt * range(0, nt, step=1)) + +""" +Turn a vector of numbers into a vector of `NamedTuple`s to be used by `ParametricDataLoader`. +""" +function turn_parameters_into_correct_format(t::AbstractVector, IC::AbstractVector{<:NamedTuple}) + vec_of_params = NamedTuple[] + for time_step ∈ t + time_step == t[end] || push!(vec_of_params, (t = time_step, )) + end + vcat((vec_of_params for _ in axes(IC, 1))...) +end + +for i in 1:nt+1 + for j=1:ni + q[j,i] = ( IC[j].p - Omega*F/(omega^2-Omega^2) )/ omega *sin(omega*t[i]) + IC[j].q*cos(omega*t[i]) + F/(omega^2-Omega^2)*sin(Omega*t[i]) + p[j,i] = -omega^2*IC[j].q*sin(omega*t[i]) + ( IC[j].p - Omega*F/(omega^2-Omega^2) )*cos(omega*t[i]) + Omega*F/(omega^2-Omega^2)*cos(Omega*t[i]) + # q[j,i] = ( IC[j].p - Omega*F/(omega^2-Omega^2) )/ omega *exp(-omega*t[i]) - IC[j].q*exp(-omega*t[i]) + F/(omega^2-Omega^2)*exp(-Omega*t[i]) + # p[j,i] = -omega^2*IC[j].q*exp(-omega*t[i]) + ( IC[j].p + Omega*F/(omega^2-Omega^2) )*exp(-omega*t[i]) - Omega*F/(omega^2-Omega^2)*exp(-Omega*t[i]) + end + +end + +@doc raw""" +Turn a `NamedTuple` of ``(q,p)`` data into two tensors of the correct format. + +This is the tricky part as the structure of the input array(s) needs to conform with the structure of the parameters. + +Here the data are rearranged in an array of size ``(n, 2, t_f - 1)`` where ``[t_0, t_1, \ldots, t_f]`` is the vector storing the time steps. + +If we deal with different initial conditions as well, we still put everything into the third (parameter) axis. + +# Example + +```jldoctest +using GeometricMachineLearning + +q = [1. 2. 3.; 4. 5. 6.] +p = [1.5 2.5 3.5; 4.5 5.5 6.5] +qp = (q = q, p = p) +turn_q_p_data_into_correct_format(qp) + +# output + +(q = [1.0 2.0; 4.0 5.0;;; 2.0 3.0; 5.0 6.0], p = [1.5 2.5; 4.5 5.5;;; 2.5 3.5; 5.5 6.5]) +``` +""" +function turn_q_p_data_into_correct_format(qp::QPT2{T, 2}) where {T} + number_of_time_steps = size(qp.q, 2) - 1 # not counting t₀ + number_of_initial_conditions = size(qp.q, 1) + q_array = zeros(T, 1, 2, number_of_time_steps * number_of_initial_conditions) + p_array = zeros(T, 1, 2, number_of_time_steps * number_of_initial_conditions) + for initial_condition_index ∈ 0:(number_of_initial_conditions - 1) + for time_index ∈ 1:number_of_time_steps + q_array[:, 1, initial_condition_index * number_of_time_steps + time_index] .= qp.q[initial_condition_index + 1, time_index] + q_array[:, 2, initial_condition_index * number_of_time_steps + time_index] .= qp.q[initial_condition_index + 1, time_index + 1] + p_array[:, 1, initial_condition_index * number_of_time_steps + time_index] .= qp.p[initial_condition_index + 1, time_index] + p_array[:, 2, initial_condition_index * number_of_time_steps + time_index] .= qp.p[initial_condition_index + 1, time_index + 1] + end + end + (q = q_array, p = p_array) +end + +# SAVING TO FILE + +# h5 = h5open(path, "w") +# write(h5, "q", q) +# write(h5, "p", p) +# write(h5, "t", t) +# +# attrs(h5)["ni"] = ni +# attrs(h5)["nt"] = nt +# attrs(h5)["dt"] = dt +# +# close(h5) + +""" +This takes time as a single additional parameter (third axis). +""" +function load_time_dependent_harmonic_oscillator_with_parametric_data_loader(qp::QPT{T}, t::AbstractVector{T}, IC::AbstractVector) where {T} + qp_reformatted = turn_q_p_data_into_correct_format(qp) + t_reformatted = turn_parameters_into_correct_format(t, IC) + ParametricDataLoader(qp_reformatted, t_reformatted) +end + +# This sets up the data loader +dl = load_time_dependent_harmonic_oscillator_with_parametric_data_loader((q = q, p = p), t, IC) + +# This sets up the neural network +width::Int = 1 +nhidden::Int = 1 +n_integrators::Int = 2 +# sigmoid_linear_unit(x::T) where {T<:Number} = x / (T(1) + exp(-x)) +arch1 = ForcedGeneralizedHamiltonianArchitecture(2; activation = tanh, width = width, nhidden = nhidden, n_integrators = n_integrators, parameters = turn_parameters_into_correct_format(t, IC)[1], forcing_type = :P) +arch2 = ForcedGeneralizedHamiltonianArchitecture(2; activation = tanh, width = width, nhidden = nhidden, n_integrators = n_integrators, parameters = turn_parameters_into_correct_format(t, IC)[1], forcing_type = :Q) +arch3 = ForcedGeneralizedHamiltonianArchitecture(2; activation = tanh, width = 2width, nhidden = nhidden, n_integrators = n_integrators, parameters = turn_parameters_into_correct_format(t, IC)[1], forcing_type = :QP) +nn1 = NeuralNetwork(arch1) +nn2 = NeuralNetwork(arch2) +nn3 = NeuralNetwork(arch3) + +# This is where training starts +batch_size = 128 +n_epochs = 200 +batch = Batch(batch_size) +o1 = Optimizer(AdamOptimizer(), nn1) +o2 = Optimizer(AdamOptimizer(), nn2) +o3 = Optimizer(AdamOptimizer(), nn3) +loss = ParametricLoss() +_pb = SymbolicPullback(nn1, loss, turn_parameters_into_correct_format(t, IC)[1]); +_pb = SymbolicPullback(nn2, loss, turn_parameters_into_correct_format(t, IC)[1]); +_pb = SymbolicPullback(nn3, loss, turn_parameters_into_correct_format(t, IC)[1]); + +function train_network() + o1(nn1, dl, batch, n_epochs, loss, _pb) + o2(nn2, dl, batch, n_epochs, loss, _pb) + o3(nn3, dl, batch, n_epochs, loss, _pb) +end + +loss_array = train_network() + +trajectory_number = 20 + +# Testing the network +initial_conditions = (q = q[trajectory_number, 1], p = p[trajectory_number, 1]) +n_steps = nt +trajectory = (q = zeros(1, n_steps), p = zeros(1, n_steps)) +trajectory.q[:, 1] .= initial_conditions.q +trajectory.p[:, 1] .= initial_conditions.p +# note that we have to supply the parameters as a named tuple as well here: +for t_step ∈ 0:(n_steps-2) + qp_temporary = nn3.model((q = [trajectory.q[1, t_step+1]], p = [trajectory.p[1, t_step+1]]), (t = t[t_step+1],), nn3.params) + trajectory.q[:, t_step+2] .= qp_temporary.q + trajectory.p[:, t_step+2] .= qp_temporary.p +end + +fig = Figure() +ax = Axis(fig[1,1]) +lines!(ax, trajectory.q[1,:]; label="nn") +lines!(ax, q[trajectory_number,:]; label="analytic") \ No newline at end of file diff --git a/scripts/TimeDependentHarmonicOscillatorParametricResnet.jl b/scripts/TimeDependentHarmonicOscillatorParametricResnet.jl new file mode 100644 index 000000000..5b6e0f31f --- /dev/null +++ b/scripts/TimeDependentHarmonicOscillatorParametricResnet.jl @@ -0,0 +1,150 @@ +using HDF5 +using GeometricMachineLearning +using GeometricMachineLearning: QPT, QPT2, Activation, ParametricLoss, SymbolicNeuralNetwork, SymbolicPullback +using CairoMakie +using NNlib: relu + +# PARAMETERS +omega = 1.0 # natural frequency of the harmonic Oscillator +Omega = 3.5 # frequency of the external sinusoidal forcing +F = .0 # .9 # amplitude of the external sinusoidal forcing +ni_dim = 10 # number of initial conditions per dimension (so ni_dim^2 total) +T = 2π * 5 +nt = 1000 # number of time steps +dt = T/nt # time step + +# Generating the initial condition array +IC = vec( [(q=q0, p=p0) for q0 in range(-1, 1, ni_dim), p0 in range(-1, 1, ni_dim)] ) + +# Generating the solution array +ni = ni_dim^2 +q = zeros(Float64, ni, nt+1) +p = zeros(Float64, ni, nt+1) +t = collect(dt * range(0, nt, step=1)) + +""" +Turn a vector of numbers into a vector of `NamedTuple`s to be used by `ParametricDataLoader`. +""" +function turn_parameters_into_correct_format(t::AbstractVector, IC::AbstractVector{<:NamedTuple}) + vec_of_params = NamedTuple[] + for time_step ∈ t + time_step == t[end] || push!(vec_of_params, (t = time_step, )) + end + vcat((vec_of_params for _ in axes(IC, 1))...) +end + +for i in 1:nt+1 + for j=1:ni + q[j,i] = ( IC[j].p - Omega*F/(omega^2-Omega^2) )/ omega *sin(omega*t[i]) + IC[j].q*cos(omega*t[i]) + F/(omega^2-Omega^2)*sin(Omega*t[i]) + p[j,i] = -omega^2*IC[j].q*sin(omega*t[i]) + ( IC[j].p - Omega*F/(omega^2-Omega^2) )*cos(omega*t[i]) + Omega*F/(omega^2-Omega^2)*cos(Omega*t[i]) + # q[j,i] = ( IC[j].p - Omega*F/(omega^2-Omega^2) )/ omega *exp(-omega*t[i]) - IC[j].q*exp(-omega*t[i]) + F/(omega^2-Omega^2)*exp(-Omega*t[i]) + # p[j,i] = -omega^2*IC[j].q*exp(-omega*t[i]) + ( IC[j].p + Omega*F/(omega^2-Omega^2) )*exp(-omega*t[i]) - Omega*F/(omega^2-Omega^2)*exp(-Omega*t[i]) + end + +end + +@doc raw""" +Turn a `NamedTuple` of ``(q,p)`` data into two tensors of the correct format. + +This is the tricky part as the structure of the input array(s) needs to conform with the structure of the parameters. + +Here the data are rearranged in an array of size ``(n, 2, t_f - 1)`` where ``[t_0, t_1, \ldots, t_f]`` is the vector storing the time steps. + +If we deal with different initial conditions as well, we still put everything into the third (parameter) axis. + +# Example + +```jldoctest +using GeometricMachineLearning + +q = [1. 2. 3.; 4. 5. 6.] +p = [1.5 2.5 3.5; 4.5 5.5 6.5] +qp = (q = q, p = p) +turn_q_p_data_into_correct_format(qp) + +# output + +(q = [1.0 2.0; 4.0 5.0;;; 2.0 3.0; 5.0 6.0], p = [1.5 2.5; 4.5 5.5;;; 2.5 3.5; 5.5 6.5]) +``` +""" +function turn_q_p_data_into_correct_format(qp::QPT2{T, 2}) where {T} + number_of_time_steps = size(qp.q, 2) - 1 # not counting t₀ + number_of_initial_conditions = size(qp.q, 1) + q_array = zeros(T, 1, 2, number_of_time_steps * number_of_initial_conditions) + p_array = zeros(T, 1, 2, number_of_time_steps * number_of_initial_conditions) + for initial_condition_index ∈ 0:(number_of_initial_conditions - 1) + for time_index ∈ 1:number_of_time_steps + q_array[:, 1, initial_condition_index * number_of_time_steps + time_index] .= qp.q[initial_condition_index + 1, time_index] + q_array[:, 2, initial_condition_index * number_of_time_steps + time_index] .= qp.q[initial_condition_index + 1, time_index + 1] + p_array[:, 1, initial_condition_index * number_of_time_steps + time_index] .= qp.p[initial_condition_index + 1, time_index] + p_array[:, 2, initial_condition_index * number_of_time_steps + time_index] .= qp.p[initial_condition_index + 1, time_index + 1] + end + end + (q = q_array, p = p_array) +end + +# SAVING TO FILE + +# h5 = h5open(path, "w") +# write(h5, "q", q) +# write(h5, "p", p) +# write(h5, "t", t) +# +# attrs(h5)["ni"] = ni +# attrs(h5)["nt"] = nt +# attrs(h5)["dt"] = dt +# +# close(h5) + +""" +This takes time as a single additional parameter (third axis). +""" +function load_time_dependent_harmonic_oscillator_with_parametric_data_loader(qp::QPT{T}, t::AbstractVector{T}, IC::AbstractVector) where {T} + qp_reformatted = turn_q_p_data_into_correct_format(qp) + t_reformatted = turn_parameters_into_correct_format(t, IC) + ParametricDataLoader(qp_reformatted, t_reformatted) +end + +# This sets up the data loader +dl = load_time_dependent_harmonic_oscillator_with_parametric_data_loader((q = q, p = p), t, IC) + +# This sets up the neural network +width::Int = 2 +n_blocks::Int = 1 +n_integrators::Int = 1 +# sigmoid_linear_unit(x::T) where {T<:Number} = x / (T(1) + exp(-x)) +arch = ResNet(2, n_blocks=n_blocks, width=width; activation=tanh, parameters=turn_parameters_into_correct_format(t, IC)[1]) +nn = NeuralNetwork(arch) + +# This is where training starts +batch_size = 128 +n_epochs = 200 +batch = Batch(batch_size) +o = Optimizer(AdamOptimizer(), nn) +loss = ParametricLoss() + +function train_network() + o(nn, dl, batch, n_epochs, loss) +end + +loss_array = train_network() + +trajectory_number = 20 + +# Testing the network +initial_conditions = (q = q[trajectory_number, 1], p = p[trajectory_number, 1]) +n_steps = nt +trajectory = (q = zeros(1, n_steps), p = zeros(1, n_steps)) +trajectory.q[:, 1] .= initial_conditions.q +trajectory.p[:, 1] .= initial_conditions.p +# note that we have to supply the parameters as a named tuple as well here: +for t_step ∈ 0:(n_steps-2) + qp_temporary = nn.model((q = [trajectory.q[1, t_step+1]], p = [trajectory.p[1, t_step+1]]), (t = t[t_step+1],), nn.params) + trajectory.q[:, t_step+2] .= qp_temporary.q + trajectory.p[:, t_step+2] .= qp_temporary.p +end + +fig = Figure() +ax = Axis(fig[1,1]) +lines!(ax, trajectory.q[1,:]; label="nn") +lines!(ax, q[trajectory_number,:]; label="analytic") \ No newline at end of file diff --git a/scripts/Train_DampedOscillator_QP.jl b/scripts/Train_DampedOscillator_QP.jl new file mode 100644 index 000000000..1058ef941 --- /dev/null +++ b/scripts/Train_DampedOscillator_QP.jl @@ -0,0 +1,114 @@ +using HDF5 +using GeometricMachineLearning +using GeometricMachineLearning: QPT, QPT2 +using CairoMakie +using JLD2 +using NNlib: relu + +# PARAMETERS +nu = 0.001 # friction force coefficient +ni_dim = 2 # number of initial conditions per dimension (so ni_dim^2 total) +T = 13 +nt = 100 # number of time steps +dt = T/nt # time step +n_epochs = 100000 +n_epochs = 3 +width = 4 # width of the neural network +nhidden = 3 # number of hidden layers in the neural network +batch_size = 5000 # the size of the batch + +path_out = "D:\\RESEARCH - UTWENTE\\GFHNNs\\Damped Oscillator\\network_TEST.jld2" +#path_out = "/home/tyranowskitm/GFHNNs/DampedOscillator/OUTPUTS/network.jld2" + + +# Generating the initial condition array +IC = vec( [(q=q0, p=p0) for q0 in range(-1, 1, ni_dim), p0 in range(-1, 1, ni_dim)] ) + + +# Generating the solution array +ni = ni_dim^2 +omega = sqrt(4-nu^2) / 2 + +q = zeros(Float64, ni, nt+1) +p = zeros(Float64, ni, nt+1) +t = collect(dt*range(0,nt,step=1)) + +for i in 1:nt+1 + + for j=1:ni + q[j,i] = (1/omega)*( IC[j].p + nu/2 *IC[j].q )*exp(-nu*t[i]/2)*sin(omega*t[i]) + IC[j].q*exp(-nu*t[i]/2)*cos(omega*t[i]) + p[j,i] = -(1/omega)*( IC[j].q + nu/2 *IC[j].p )*exp(-nu*t[i]/2)*sin(omega*t[i]) + IC[j].p*exp(-nu*t[i]/2)*cos(omega*t[i]) + end + +end + + + +@doc raw""" +Turn a `NamedTuple` of ``(q,p)`` data into two tensors of the correct format. + +This is the tricky part as the structure of the input array(s) needs to conform with the structure of the parameters. + +Here the data are rearranged in an array of size ``(n, 2, t_f - 1)`` where ``[t_0, t_1, \ldots, t_f]`` is the vector storing the time steps. + +If we deal with different initial conditions as well, we still put everything into the third (parameter) axis. + +# Example + +```jldoctest +using GeometricMachineLearning + +q = [1. 2. 3.; 4. 5. 6.] +p = [1.5 2.5 3.5; 4.5 5.5 6.5] +qp = (q = q, p = p) +turn_q_p_data_into_correct_format(qp) + +# output + +(q = [1.0 2.0; 4.0 5.0;;; 2.0 3.0; 5.0 6.0], p = [1.5 2.5; 4.5 5.5;;; 2.5 3.5; 5.5 6.5]) +``` +""" +function turn_q_p_data_into_correct_format(qp::QPT2{T, 2}) where {T} + number_of_time_steps = size(qp.q, 2) - 1 # not counting t₀ + number_of_initial_conditions = size(qp.q, 1) + q_array = zeros(T, 1, 2, number_of_time_steps * number_of_initial_conditions) + p_array = zeros(T, 1, 2, number_of_time_steps * number_of_initial_conditions) + for initial_condition_index ∈ 0:(number_of_initial_conditions - 1) + for time_index ∈ 1:number_of_time_steps + q_array[:, 1, initial_condition_index * number_of_time_steps + time_index] .= qp.q[initial_condition_index + 1, time_index] + q_array[:, 2, initial_condition_index * number_of_time_steps + time_index] .= qp.q[initial_condition_index + 1, time_index + 1] + p_array[:, 1, initial_condition_index * number_of_time_steps + time_index] .= qp.p[initial_condition_index + 1, time_index] + p_array[:, 2, initial_condition_index * number_of_time_steps + time_index] .= qp.p[initial_condition_index + 1, time_index + 1] + end + end + (q = q_array, p = p_array) +end + + +# This sets up the data loader +dl = DataLoader(turn_q_p_data_into_correct_format((q = q, p = p))) + +# This sets up the neural network +arch = ForcedSympNet(2; upscaling_dimension = width, n_layers = nhidden, forcing_type = :P) +#arch = ForcedSympNet(2; upscaling_dimension = width, n_layers = nhidden, activation=(x-> max(0,x)^2/2)) +nn = NeuralNetwork(arch) + +# This is where training starts +batch = Batch(batch_size) +o = Optimizer(AdamOptimizer(), nn) + +loss_array = o(nn, dl, batch, n_epochs) + + +# Saving the parameters of the network +println("Saving the parameters of the neural network...") +flush(stdout) + +params = GeometricMachineLearning.map_to_cpu(nn.params) + +save(path_out,"parameters", params, "training loss", loss_array, "ni_dim", ni_dim, "T", T, "nt", nt, "n_epochs", n_epochs, "width", width, "nhidden", nhidden, "batch_size", batch_size, "nu", nu) + +println(" ...Done!") +flush(stdout) + + diff --git a/scripts/forcing_layers_parameter_number.jl b/scripts/forcing_layers_parameter_number.jl new file mode 100644 index 000000000..e27376d04 --- /dev/null +++ b/scripts/forcing_layers_parameter_number.jl @@ -0,0 +1,11 @@ +using GeometricMachineLearning +using GeometricMachineLearning: ForcingLayerP, ForcingLayerQP + +forcing_layer_p = ForcingLayerP(2) +forcing_layer_qp = ForcingLayerQP(2) + +nn_p = NeuralNetwork(forcing_layer_p) +nn_qp = NeuralNetwork(forcing_layer_qp) + +println(parameterlength(nn_p)) +println(parameterlength(nn_qp)) \ No newline at end of file diff --git a/src/GeometricMachineLearning.jl b/src/GeometricMachineLearning.jl index 5311be0c9..b86812bd9 100644 --- a/src/GeometricMachineLearning.jl +++ b/src/GeometricMachineLearning.jl @@ -25,8 +25,12 @@ using InteractiveUtils using TimerOutputs import SymbolicNeuralNetworks import SymbolicNeuralNetworks: SymbolicPullback -using SymbolicNeuralNetworks: derivative, SymbolicNeuralNetwork +using SymbolicNeuralNetworks: derivative, SymbolicNeuralNetwork, AbstractSymbolicNeuralNetwork import Symbolics +# The system parameters of a parameter-dependent architecture are flattened into the network's +# input. Only the two conversions are brought in: the module name would clash with +# `AbstractNeuralNetworks.NeuralNetworkParameters`, the *type* GML re-exports. +using NeuralNetworkParameters: flatten, unflatten # The manifolds, the structured matrix types, the global sections and the retractions are # `GeometricOptimizers`' — GML used to carry near-verbatim copies of all eleven types, which Julia @@ -185,8 +189,11 @@ export StiefelManifold, GrassmannManifold, Manifold export rgrad, metric, check include("layers/sympnets.jl") +include("layers/forcing_dissipation_layers.jl") include("layers/bias_layer.jl") include("layers/resnet.jl") +include("layers/wide_resnet.jl") +include("layers/parametric_resnet_layer.jl") include("layers/manifold_layer.jl") include("layers/stiefel_layer.jl") include("layers/grassmann_layer.jl") @@ -295,11 +302,13 @@ export arch include("backends/backends.jl") include("backends/lux.jl") -export NetworkLoss, TransformerLoss, FeedForwardLoss, AutoEncoderLoss, ReducedLoss, HNNLoss +export NetworkLoss, TransformerLoss, FeedForwardLoss, AutoEncoderLoss, ReducedLoss, HNNLoss, + ParametricLoss #INCLUDE ARCHITECTURES include("architectures/neural_network_integrator.jl") include("architectures/resnet.jl") +include("architectures/parametric_resnet.jl") include("architectures/transformer_integrator.jl") include("architectures/standard_transformer_integrator.jl") include("architectures/sympnet.jl") @@ -308,6 +317,8 @@ include("architectures/symplectic_autoencoder.jl") include("architectures/psd.jl") include("architectures/fixed_width_network.jl") include("architectures/hamiltonian_neural_network.jl") +include("architectures/standard_hamiltonian_neural_network.jl") +include("architectures/generalized_hamiltonian_neural_network.jl") include("architectures/lagrangian_neural_network.jl") include("architectures/variable_width_network.jl") include("architectures/transformer_neural_network.jl") @@ -321,7 +332,8 @@ export ClassificationTransformer, ClassificationLayer export VolumePreservingFeedForward export SymplecticAutoencoder, PSDArch export HamiltonianArchitecture, StandardHamiltonianArchitecture, - GeneralizedHamiltonianArchitecture + GeneralizedHamiltonianArchitecture, ForcedGeneralizedHamiltonianArchitecture +export ForcedSympNet export solve!, encoder, decoder @@ -339,13 +351,18 @@ export AbstractPullback, ZygotePullback, SymbolicPullback include("pullbacks/zygote_pullback.jl") include("pullbacks/symbolic_hnn_pullback.jl") -export DataLoader +export DataLoader, ParametricDataLoader export Batch, optimize_for_one_epoch! include("data_loader/tensor_assign.jl") include("data_loader/matrix_assign.jl") include("data_loader/batch.jl") +# before `optimize.jl`, whose training loop takes either data loader +include("data_loader/parametric_data_loader.jl") include("data_loader/optimize.jl") +include("architectures/forced_sympnet.jl") +include("architectures/forced_generalized_hamiltonian_neural_network.jl") + # INCLUDE TRAINING parameters export TrainingParameters @@ -393,8 +410,6 @@ export train! include("training/train.jl") -export SymplecticEuler -export SymplecticEulerA, SymplecticEulerB export SEuler, SEulerA, SEulerB include("training_method/symplectic_euler.jl") diff --git a/src/architectures/forced_generalized_hamiltonian_neural_network.jl b/src/architectures/forced_generalized_hamiltonian_neural_network.jl new file mode 100644 index 000000000..2bb97632e --- /dev/null +++ b/src/architectures/forced_generalized_hamiltonian_neural_network.jl @@ -0,0 +1,35 @@ +const N_FORCING_LAYERS_DEFAULT = 2 + +""" + ForcedGeneralizedHamiltonianArchitecture <: HamiltonianArchitecture + +A version of [`GeneralizedHamiltonianArchitecture`](@ref) that includes forcing/dissipation terms. Also compare this to [`ForcedSympNet`](@ref). +""" +struct ForcedGeneralizedHamiltonianArchitecture{FT, AT, PT <: OptionalParameters} <: HamiltonianArchitecture{AT} + dim::Int + width::Int + nhidden::Int + n_forcing_layers::Int + n_integrators::Int + parameters::PT + activation::AT + + function ForcedGeneralizedHamiltonianArchitecture(dim; width=dim, nhidden=HNN_nhidden_default, n_forcing_layers=N_FORCING_LAYERS_DEFAULT, n_integrators::Integer=1, activation=HNN_activation_default, parameters=NullParameters(), forcing_type::Symbol=:P) + forcing_type == :P || forcing_type == :Q || forcing_type == :QP || error("Forcing has to be either :Q or :P. It is $(forcing_type).") + activation = (typeof(activation) <: Activation) ? activation : Activation(activation) + new{forcing_type, typeof(activation), typeof(parameters)}(dim, width, nhidden, n_forcing_layers, n_integrators, parameters, activation) + end +end + +function Chain(arch::ForcedGeneralizedHamiltonianArchitecture{FT}) where {FT} + layers = () + kinetic_energy = SymbolicKineticEnergy(arch.dim, arch.width, arch.nhidden, arch.activation; parameters=arch.parameters) + potential_energy = SymbolicPotentialEnergy(arch.dim, arch.width, arch.nhidden, arch.activation; parameters=arch.parameters) + for i ∈ 1:arch.n_integrators + layers = (layers..., SymplecticEulerA(kinetic_energy; return_parameters = true)) + layers = (layers..., ForcingLayer(arch.dim, arch.width, arch.n_forcing_layers, arch.activation; parameters=arch.parameters, return_parameters=true, type=FT)) + _return_parameters = !(i == arch.n_integrators) + layers = (layers..., SymplecticEulerB(potential_energy; return_parameters = _return_parameters)) + end + Chain(layers...) +end \ No newline at end of file diff --git a/src/architectures/forced_sympnet.jl b/src/architectures/forced_sympnet.jl new file mode 100644 index 000000000..582acca44 --- /dev/null +++ b/src/architectures/forced_sympnet.jl @@ -0,0 +1,65 @@ +@doc raw""" + ForcedSympNet <: NeuralNetworkIntegrator + +`ForcedSympNet`s are based on [`SympNet`](@ref)s [jin2020sympnets](@cite) and include [`ForcingLayer`](@ref)s. They are based on [`GSympNet`](@ref)s. + +# Constructor + +```julia +ForcedSympNet(d) +``` + +Make a forced SympNet with dimension ``d.`` + +# Arguments + +Keyword arguments are: +- `upscaling_dimension::Int = 2d`: The *upscaling dimension* of the gradient layer. See the documentation for [`GradientLayerQ`](@ref) and [`GradientLayerP`](@ref) for further explanation. +- `n_layers::Int""" * "$(g_n_layers_default)`" * raw""": The number of layers (i.e. the total number of [`GradientLayerQ`](@ref) and [`GradientLayerP`](@ref)). +- `activation""" * "$(g_activation_default)`" * raw""": The activation function that is applied. +- `init_upper::Bool""" * "$(g_init_upper_default)`" * raw""": Initialize the gradient layer so that it first modifies the $q$-component. +""" +struct ForcedSympNet{FT, AT} <: NeuralNetworkIntegrator + dim::Int + upscaling_dimension::Int + n_layers::Int + n_forcing_layers::Int + act::AT + init_upper::Bool + + function ForcedSympNet(dim::Integer; + upscaling_dimension = 2 * dim, + n_layers = g_n_layers_default, + n_forcing_layers = N_FORCING_LAYERS_DEFAULT, + activation = g_activation_default, + init_upper = g_init_upper_default, + forcing_type::Symbol = :P) + new{forcing_type, typeof(activation)}(dim, upscaling_dimension, n_layers, n_forcing_layers, activation, init_upper) + end + + function ForcedSympNet(dl::DataLoader; + upscaling_dimension = 2 * dl.input_dim, + n_layers = g_n_layers_default, + n_forcing_layers = N_FORCING_LAYERS_DEFAULT, + activation = g_activation_default, + init_upper = g_init_upper_default, + forcing_type::Symbol = :P) + new{forcing_type, typeof(activation)}(dl.input_dim, upscaling_dimension, n_layers, n_forcing_layers, activation, init_upper) + end +end + +function Chain(arch::ForcedSympNet{FT}) where {FT} + layers = () + is_upper_criterion = arch.init_upper ? isodd : iseven + for i in 1:arch.n_layers + layers = + if is_upper_criterion(i) + (layers..., GradientLayerQ(arch.dim, arch.upscaling_dimension, arch.act)) + else + (layers..., + ForcingLayer(arch.dim, arch.upscaling_dimension, arch.n_forcing_layers, arch.act; return_parameters=false, type=FT), + GradientLayerP(arch.dim, arch.upscaling_dimension, arch.act)) + end + end + Chain(layers...) +end \ No newline at end of file diff --git a/src/architectures/generalized_hamiltonian_neural_network.jl b/src/architectures/generalized_hamiltonian_neural_network.jl new file mode 100644 index 000000000..447a9cd3a --- /dev/null +++ b/src/architectures/generalized_hamiltonian_neural_network.jl @@ -0,0 +1,335 @@ +""" + SymbolicEnergy + +See [`SymbolicPotentialEnergy`](@ref) and [`SymbolicKineticEnergy`](@ref). +""" +struct SymbolicEnergy{AT <: Activation, PT, Kinetic} + dim::Int + width::Int + nhidden::Int + parameter_length::Int + parameter_layout::PT + activation::AT + + function SymbolicEnergy(dim, width, nhidden, activation; parameters::OptionalParameters=NullParameters(), type) + @assert iseven(dim) "The input dimension must be an even integer!" + flat_parameters, layout = _flatten_system_parameters(parameters) + _activation = Activation(activation) + new{typeof(_activation), typeof(layout), type}(dim, width, nhidden, length(flat_parameters), layout, _activation) + end +end + +""" + SymbolicPotentialEnergy + +A `const` derived from [`SymbolicEnergy`](@ref). + +# Constructors + +```jldoctest; setup=:(using GeometricMachineLearning; using GeometricMachineLearning: Activation) +julia> params, dim, width, nhidden, activation = (m = 1., ω = π / 2), 2, 2, 1, tanh +((m = 1.0, ω = 1.5707963267948966), 2, 2, 1, tanh) + +julia> se = GeometricMachineLearning.SymbolicPotentialEnergy(dim, width, nhidden, activation; parameters = params); + +``` + +In practice we use `SymbolicPotentialEnergy` (and [`SymbolicKineticEnergy`](@ref)) together with [`build_gradient(::SymbolicEnergy)`](@ref). + +# Parameter Dependence +""" +const SymbolicPotentialEnergy{AT, PT} = SymbolicEnergy{AT, PT, :potential} + +""" + SymbolicKineticEnergy + +A `const` derived from [`SymbolicEnergy`](@ref). + +# Constructors + +See [`SymbolicPotentialEnergy`](@ref). +""" +const SymbolicKineticEnergy{AT, PT} = SymbolicEnergy{AT, PT, :kinetic} + +SymbolicPotentialEnergy(args...; kwargs...) = SymbolicEnergy(args...; type = :potential, kwargs...) +SymbolicKineticEnergy(args...; kwargs...) = SymbolicEnergy(args...; type = :kinetic, kwargs...) + +function Chain(se::SymbolicEnergy) + inner_layers = Tuple( + [Dense(se.width, se.width, se.activation) for _ in 1:se.nhidden] + ) + + Chain( + Dense(se.dim÷2 + se.parameter_length, se.width, se.activation), + inner_layers..., + Linear(se.width, 1; use_bias = false) + ) +end + +# Jacobian with respect to the *first* `dim2` input variables only: for a parameter-dependent +# network the remaining input components are the system parameters, which are not differentiated. +# +# TODO: type piracy -- `Jacobian` and `AbstractSymbolicNeuralNetwork` are both +# SymbolicNeuralNetworks'. The restricted-Jacobian variant belongs there. +function SymbolicNeuralNetworks.Jacobian(f, nn::AbstractSymbolicNeuralNetwork, dim2::Integer) + # make differential of input variables (not of parameters) + Dx = SymbolicNeuralNetworks.symbolic_differentials(nn.input)[1:dim2] + + # Evaluation of gradient + s∇f = hcat([SymbolicNeuralNetworks.expand_derivatives.(dx.(SymbolicNeuralNetworks.Symbolics.scalarize(f))) for dx in Dx]...) + + SymbolicNeuralNetworks.Jacobian(f, s∇f, nn) +end + +function SymbolicNeuralNetworks.Jacobian(nn::AbstractSymbolicNeuralNetwork, dim2::Integer) + + # Evaluation of the symbolic output + soutput = nn.model(nn.input, params(nn)) + + SymbolicNeuralNetworks.Jacobian(soutput, nn, dim2) +end + +""" + build_gradient(se) + +Build a gradient function from a [`SymbolicEnergy`](@ref) `se`. + +# Examples + +```jldoctest; setup=:(using GeometricMachineLearning; using GeometricMachineLearning: SymbolicPotentialEnergy, build_gradient, concatenate_array_with_parameters, OneInitializer; using GeometricMachineLearning.GeometricBase: OptionalParameters) +params, dim, width, nhidden, activation = (m = 1., ω = π / 2), 4, 2, 1, tanh + +se = SymbolicPotentialEnergy(dim, width, nhidden, activation; parameters = params) + +# `OneInitializer` rather than the default random one, so that the output below does not depend on +# the random number stream of the Julia version +network_params = NeuralNetwork(Chain(se); initializer = OneInitializer()).params + +built_grad = build_gradient(se) +grad(qp::AbstractArray, problem_params::OptionalParameters, params::NeuralNetworkParameters) = built_grad(concatenate_array_with_parameters(qp, problem_params), params) + +grad([0.5, 0.25], params, network_params) + +# output + +2×1 Matrix{Float64}: + 2.7907683385233434e-5 + 2.7907683385233434e-5 +``` +""" +function build_gradient(se::SymbolicEnergy) + model = Chain(se) + nn = SymbolicNeuralNetwork(model) + □ = SymbolicNeuralNetworks.Jacobian(nn, se.dim÷2) + SymbolicNeuralNetworks.build_nn_function(SymbolicNeuralNetworks.derivative(□)', nn.params, nn.input; + inplace = false) +end + +struct SymplecticEuler{M, N, FT<:Base.Callable, MT<:Chain, type, ReturnParameters} <: AbstractExplicitLayer{M, N} + gradient_function::FT + energy_model::MT +end + +function parameterlength(integrator::SymplecticEuler) + parameterlength(integrator.energy_model) +end + +function initialparameters(rng::Random.AbstractRNG, init_weight::AbstractNeuralNetworks.Initializer, integrator::SymplecticEuler, backend::KernelAbstractions.Backend, ::Type{T}) where {T} + initialparameters(rng, init_weight, integrator.energy_model, backend, T) +end + +const SymplecticEulerA{M, N, FT, AT, ReturnParameters} = SymplecticEuler{M, N, FT, AT, :A, ReturnParameters} +const SymplecticEulerB{M, N, FT, AT, ReturnParameters} = SymplecticEuler{M, N, FT, AT, :B, ReturnParameters} + +""" +Changes ``q`` (based on the kinetic energy). +""" +function SymplecticEulerA(se::SymbolicKineticEnergy; return_parameters::Bool) + gradient_function = build_gradient(se) + c = Chain(se) + SymplecticEuler{se.dim, se.dim, typeof(gradient_function), typeof(c), :A, return_parameters}(gradient_function, c) +end + +""" +Changes ``p`` (based on the potential energy). +""" +function SymplecticEulerB(se::SymbolicPotentialEnergy; return_parameters::Bool) + gradient_function = build_gradient(se) + c = Chain(se) + SymplecticEuler{se.dim, se.dim, typeof(gradient_function), typeof(c), :B, return_parameters}(gradient_function, c) +end + +# A network with no system parameters gets its input unchanged; without this the empty flat vector +# would have to be `vcat`ed on, which loses the element type. +concatenate_array_with_parameters(qp::AbstractArray, ::NullParameters) = qp + +function concatenate_array_with_parameters(qp::AbstractVector, params::NamedTuple) + vcat(qp, _flatten_system_parameters(params)[1]) +end + +function concatenate_array_with_parameters(qp::AbstractMatrix, params::NamedTuple) + @assert size(qp, 2) == 1 + vcat(qp, repeat(_flatten_system_parameters(params)[1], 1, size(qp, 2))) +end + +function concatenate_array_with_parameters(qp::AbstractArray{T, 3}, params::AbstractVector) where {T} + @assert size(qp, 3) == length(params) + matrices = Tuple(concatenate_array_with_parameters(qp[:, :, i], params[i]) for i in axes(qp, 3)) + cat(matrices...; dims = 3) +end + +# function concatenate_array_with_parameters(qp::AbstractMatrix, params::OptionalParameters) +# hcat((concatenate_array_with_parameters(qp[:, i], params) for i in axes(qp, 2))...) +# end + +# One parameter set per column, so the columns are concatenated *side by side*: the result is a +# matrix with `size(qp, 1) + parameter_length` rows, one column per sample. +function concatenate_array_with_parameters(qp::AbstractMatrix, params::AbstractVector) + @assert _size(qp, 2) == length(params) + hcat((concatenate_array_with_parameters(@view(qp[:, i]), params[i]) for i in axes(params, 1))...) +end + +function (integrator::SymplecticEulerA{M, N, FT, AT, false})(qp::QPT2, problem_params::OptionalParameters, params::NeuralNetworkParameters) where {M, N, FT, AT} + input = concatenate_array_with_parameters(qp.p, problem_params) + (q = @view((qp.q + integrator.gradient_function(input, params))[:, 1]), p = qp.p) +end + +function (integrator::SymplecticEulerB{M, N, FT, AT, false})(qp::QPT2, problem_params::OptionalParameters, params::NeuralNetworkParameters) where {M, N, FT, AT} + input = concatenate_array_with_parameters(qp.q, problem_params) + (q = qp.q, p = @view((qp.p - integrator.gradient_function(input, params))[:, 1])) +end + +function (integrator::SymplecticEulerA{M, N, FT, AT, true})(qp::QPT2, problem_params::OptionalParameters, params::NeuralNetworkParameters) where {M, N, FT, AT} + input = concatenate_array_with_parameters(qp.p, problem_params) + ((q = @view((qp.q + integrator.gradient_function(input, params))[:, 1]), p = qp.p), problem_params) +end + +function (integrator::SymplecticEulerB{M, N, FT, AT, true})(qp::QPT2, problem_params::OptionalParameters, params::NeuralNetworkParameters) where {M, N, FT, AT} + input = concatenate_array_with_parameters(qp.q, problem_params) + ((q = qp.q, p = @view((qp.p - integrator.gradient_function(input, params))[:, 1])), problem_params) +end + +function (integrator::SymplecticEuler)(qp_params::Tuple{<:QPTOAT2, <:OptionalParameters}, params::NeuralNetworkParameters) + integrator(qp_params..., params) +end + +function (integrator::SymplecticEuler)(::TT, ::NeuralNetworkParameters) where {TT <: Tuple} + error("The input is of type $(TT). This shouldn't be the case!") +end + +function (integrator::SymplecticEuler{M, N, FT, AT, Type, true})(qp::AbstractArray, problem_params::OptionalParameters, params::NeuralNetworkParameters) where {M, N, FT, AT, Type} + @assert iseven(size(qp, 1)) + n = size(qp, 1)÷2 + qp_split = assign_q_and_p(qp, n) + evaluated = integrator(qp_split, problem_params, params)[1] + (vcat(evaluated.q, evaluated.p), problem_params) +end + +function (integrator::SymplecticEuler{M, N, FT, AT, Type, false})(qp::AbstractArray, problem_params::OptionalParameters, params::NeuralNetworkParameters) where {M, N, FT, AT, Type} + @assert iseven(size(qp, 1)) + n = size(qp, 1)÷2 + qp_split = assign_q_and_p(qp, n) + evaluated = integrator(qp_split, problem_params, params) + vcat(evaluated.q, evaluated.p) +end + +(integrator::SymplecticEuler)(qp::QPTOAT2, params::NeuralNetworkParameters) = integrator(qp, NullParameters(), params) + +""" + GeneralizedHamiltonianArchitecture <: HamiltonianArchitecture + +A realization of generalized Hamiltonian neural networks (GHNNs) as introduced in [horn2025generalized](@cite). + +Also see [`StandardHamiltonianArchitecture`](@ref). + +# Constructor + +The constructor takes the following input arguments: +1. `dim`: system dimension, +2. `width = dim`: width of the hidden layer. By default this is equal to `dim`, +3. `nhidden = $(HNN_nhidden_default)`: the number of hidden layers, +4. `n_integrators`: the number of integrators used in the GHNN. +5. `activation = $(HNN_activation_default)`: the activation function used in the GHNN, +""" +struct GeneralizedHamiltonianArchitecture{AT, PT <: OptionalParameters} <: HamiltonianArchitecture{AT} + dim::Int + width::Int + nhidden::Int + n_integrators::Int + parameters::PT + activation::AT + + function GeneralizedHamiltonianArchitecture(dim; width=dim, nhidden=HNN_nhidden_default, n_integrators::Integer=1, activation=HNN_activation_default, parameters=NullParameters()) + activation = (typeof(activation) <: Activation) ? activation : Activation(activation) + new{typeof(activation), typeof(parameters)}(dim, width, nhidden, n_integrators, parameters, activation) + end +end + +# The parameter-dependent layers pass `(state, system parameters)` down the chain, so `applychain` +# has to accept that tuple as its data argument. +# +# TODO: type piracy -- `applychain` is AbstractNeuralNetworks' and every argument type here is +# `Base`'s. ANN's own `applychain(layers, x, ps::Union{NamedTuple, NeuralNetworkParameters})` is +# already generic in `x`; widening the `@generated` method the same way would remove the need. +@generated function AbstractNeuralNetworks.applychain(layers::Tuple, x::Tuple{<:QPTOAT2, <:OptionalParameters}, ps::Tuple) + N = length(fieldtypes((layers))) + x_symbols = vcat([:x], [gensym() for _ in 1:N]) + calls = [:(($(x_symbols[i + 1])) = layers[$i]($(x_symbols[i]), ps[$i])) for i in 1:N] + push!(calls, :(return $(x_symbols[N + 1]))) + return Expr(:block, calls...) +end + +index_qpt(qp::QPT2{T, 2}, i, j) where {T} = (q = qp.q[i, j], p = qp.p[i, j]) +index_gpt(qp::QPT2{T, 3}, i, j, k) where {T} = (q = qp.q[i, j, k], p = qp.p[i, j, k]) + +function Chain(ghnn_arch::GeneralizedHamiltonianArchitecture) + c = () + kinetic_energy = SymbolicKineticEnergy(ghnn_arch.dim, ghnn_arch.width, ghnn_arch.nhidden, ghnn_arch.activation; parameters=ghnn_arch.parameters) + potential_energy = SymbolicPotentialEnergy(ghnn_arch.dim, ghnn_arch.width, ghnn_arch.nhidden, ghnn_arch.activation; parameters=ghnn_arch.parameters) + + for n in 1:ghnn_arch.n_integrators + c = (c..., SymplecticEulerA(kinetic_energy; return_parameters = true)) + c = n == ghnn_arch.n_integrators ? (c..., SymplecticEulerB(potential_energy; return_parameters=false)) : (c..., SymplecticEulerB(potential_energy; return_parameters=true)) + end + + Chain(c...) +end + +function (nn::NeuralNetwork{GT})(qp::QPTOAT2, problem_params::OptionalParameters) where {GT <: GeneralizedHamiltonianArchitecture} + nn.model(qp, problem_params, params(nn)) +end + +# TODO: type piracy -- `Chain` is AbstractNeuralNetworks' and so is every argument type of the four +# functors below. A `ParametricChain` wrapper owned by GML, or these methods upstream, would fix it. +function (model::Chain)(qp::QPTOAT2, problem_params::OptionalParameters, params::Union{NeuralNetworkParameters, NamedTuple}) + model((qp, problem_params), params) +end + +function (c::Chain)(qp::QPT2{T, 3}, system_params::AbstractVector, ps::Union{NamedTuple, NeuralNetworkParameters})::QPT2{T} where {T} + @assert size(qp.q, 3) == length(system_params) + @assert size(qp.q, 2) == 1 + output_vectorwise = [c(index_gpt(qp, :, 1, i), system_params[i], ps) for i in axes(system_params, 1)] + q_output = hcat([single_output_vectorwise.q for single_output_vectorwise ∈ output_vectorwise]...) + p_output = hcat([single_output_vectorwise.p for single_output_vectorwise ∈ output_vectorwise]...) + (q = reshape(q_output, size(q_output, 1), 1, size(q_output, 2)), p = reshape(p_output, size(p_output, 1), 1, size(p_output, 2))) +end + +function (c::Chain)(qp::AbstractArray{T, 2}, system_params::AbstractVector, ps::Union{NamedTuple, NeuralNetworkParameters}) where {T} + @assert _size(qp, 2) == length(system_params) + qp_reshaped = reshape(qp, size(qp, 1), 1, length(system_params)) + c(qp_reshaped, system_params, ps) +end + +function (c::Chain)(qp::AbstractArray{T, 3}, system_params::AbstractVector, ps::Union{NamedTuple, NeuralNetworkParameters}) where {T} + @assert size(qp, 3) == length(system_params) + @assert size(qp, 2) == 1 + @assert iseven(size(qp, 1)) + n = size(qp, 1)÷2 + qp_split = assign_q_and_p(qp, n) + c_output = c(qp_split, system_params, ps)::QPT + reshape(vcat(c_output.q, c_output.p), 2n, length(system_params)) +end + +# TODO: type piracy -- `networkbackend` is AbstractNeuralNetworks' and `ApplyArray` is LazyArrays'. +# Belongs in ANN, which already dispatches `networkbackend` on array types it does not own either. +AbstractNeuralNetworks.networkbackend(::LazyArrays.ApplyArray) = AbstractNeuralNetworks.CPU() diff --git a/src/architectures/hamiltonian_neural_network.jl b/src/architectures/hamiltonian_neural_network.jl index 2b2ab690d..ea4aaf765 100644 --- a/src/architectures/hamiltonian_neural_network.jl +++ b/src/architectures/hamiltonian_neural_network.jl @@ -13,118 +13,4 @@ function HamiltonianArchitecture(dim::Integer, width::Integer, nhidden::Integer, StandardHamiltonianArchitecture(dim, width, nhidden, activation) end -""" - StandardHamiltonianArchitecture <: HamiltonianArchitecture - -A realization of the standard Hamiltonian neural network (HNN) [greydanus2019hamiltonian](@cite). - -Also see [`GeneralizedHamiltonianArchitecture`](@ref). - -# Constructor - -The constructor takes the following input arguments: -1. `dim`: system dimension, -2. `width = dim`: width of the hidden layer. By default this is equal to `dim`, -3. `nhidden = $(HNN_nhidden_default)`: the number of hidden layers, -4. `activation = $(HNN_activation_default)`: the activation function used in the HNN. -""" -struct StandardHamiltonianArchitecture{AT} <: HamiltonianArchitecture{AT} - dim::Int - width::Int - nhidden::Int - activation::AT - - function StandardHamiltonianArchitecture(dim, width=dim, nhidden=HNN_nhidden_default, activation=HNN_activation_default) - new{typeof(activation)}(dim, width, nhidden, activation) - end -end - -GHNN_integrator_default = nothing - -""" - GeneralizedHamiltonianArchitecture <: HamiltonianArchitecture - -A realization of generalized Hamiltonian neural networks (GHNNs) as introduced in [horn2025generalized](@cite). - -Also see [`StandardHamiltonianArchitecture`](@ref). - -# Constructor - -The constructor takes the following input arguments: -1. `dim`: system dimension, -2. `width = dim`: width of the hidden layer. By default this is equal to `dim`, -3. `nhidden = $(HNN_nhidden_default)`: the number of hidden layers, -4. `activation = $(HNN_activation_default)`: the activation function used in the GHNN, -5. `integrator = $(GHNN_integrator_default)`: the integrator that is used to design the GHNN. -""" -struct GeneralizedHamiltonianArchitecture{AT, IT} <: HamiltonianArchitecture{AT} - dim::Int - width::Int - nhidden::Int - activation::AT - integrator::IT - - function GeneralizedHamiltonianArchitecture(dim, width=dim, nhidden=HNN_nhidden_default, activation=HNN_activation_default, integrator=GHNN_integrator_default) - error("GHNN still has to be implemented!") - new{typeof(activation), typeof(integrator)}(dim, width, nhidden, activation, integrator) - end -end - @inline AbstractNeuralNetworks.dim(arch::HamiltonianArchitecture) = arch.dim - -""" - symbolic_hamiltonian_vector_field(nn::SymbolicNeuralNetwork) - -Get the symbolic expression for the vector field belonging to the HNN `nn`. - -# Implementation - -This is calling `SymbolicNeuralNetworks.Jacobian` and then multiplies the result with a Poisson tensor. -""" -function symbolic_hamiltonian_vector_field(nn::SymbolicNeuralNetwork) - □ = SymbolicNeuralNetworks.Jacobian(nn) - n = input_dimension(nn.model) ÷ 2 - # The Poisson tensor is built from *integers* on purpose: a `Float64` literal in the symbolic - # expression would widen the result of a `Float32` network. `PoissonTensor`, which has the same - # convention, is an `AbstractMatrix{Float64}` and would do exactly that. - 𝕆 = zeros(Int, n, n) - 𝕀 = Matrix(1I, n, n) - 𝕁 = [𝕆 𝕀; -𝕀 𝕆] - # `Jacobian` uses the convention `□[i, j] = ∂fᵢ/∂xⱼ` and the HNN output is scalar, so the one - # row of `derivative(□)` is the gradient of the Hamiltonian. The vector field is built as a - # *vector*, so that the generated function returns what `HNNLoss` compares against: a vector - # for a single sample, and one column per sample for a batch. - ∇H = vec(derivative(□)) - 𝕁 * ∇H -end - -""" - hamiltonian_vector_field(arch::HamiltonianArchitecture) - -Compute an executable expression of the Hamiltonian vector field of a [`HamiltonianArchitecture`](@ref). - -# Implementation - -This first computes a symbolic expression of the vector field using [`symbolic_hamiltonian_vector_field`](@ref). - -The function is built with `inplace = false`: [`HNNLoss`](@ref) wraps it and is differentiated with -`Zygote`, and the in-place kernel `SymbolicNeuralNetworks.build_nn_function` builds by default -*mutates* its result, which `Zygote` does not support. -""" -function hamiltonian_vector_field(arch::HamiltonianArchitecture) - nn = SymbolicNeuralNetwork(arch) - hvf = symbolic_hamiltonian_vector_field(nn) - SymbolicNeuralNetworks.build_nn_function(hvf, nn.params, nn.input; inplace = false) -end - -function Chain(arch::HamiltonianArchitecture) - inner_layers = Tuple( - [Dense(arch.width, arch.width, arch.activation) for _ in 1:arch.nhidden] - ) - - Chain( - Dense(arch.dim, arch.width, arch.activation), - inner_layers..., - Linear(arch.width, 1; use_bias = false) - ) -end \ No newline at end of file diff --git a/src/architectures/parametric_resnet.jl b/src/architectures/parametric_resnet.jl new file mode 100644 index 000000000..5652ef4cc --- /dev/null +++ b/src/architectures/parametric_resnet.jl @@ -0,0 +1,37 @@ +struct ParametricResNet{AT <: Activation, PT <: OptionalParameters} <: NeuralNetworkIntegrator + sys_dim::Int + n_blocks::Int + width::Int + parameters::PT + activation::AT + + function ParametricResNet(dim; width=dim, n_blocks = HNN_nhidden_default, activation=HNN_activation_default, parameters=NullParameters()) + activation = (typeof(activation) <: Activation) ? activation : Activation(activation) + new{typeof(activation), typeof(parameters)}(dim, n_blocks, width, parameters, activation) + end +end + +function ParametricResNet(dl::DataLoader, n_blocks::Integer, width::Integer=dl.input_dim; activation=HNN_activation_default, parameters=NullParameters()) + ParametricResNet(dl.input_dim; width=width, n_blocks=n_blocks, activation) +end + +function ResNet(input_dim::Integer, n_blocks::Integer, width::Integer=input_dim; activation=HNN_activation_default, parameters=NullParameters()) + typeof(parameters) <: NullParameters ? ResNet(input_dim, n_blocks, width, activation) : ParametricResNet(input_dim; n_blocks=n_blocks, width=width, parameters=parameters, activation=activation) +end + +function ResNet(input_dim::Integer; n_blocks::Integer, width::Integer=input_dim, activation=HNN_activation_default, parameters=NullParameters()) + ResNet(input_dim, n_blocks, width; activation=activation, parameters=parameters) +end + +function Chain(arch::ParametricResNet{AT}) where AT + layers = () + for _ in 1:arch.n_blocks + # nonlinear layers + layers = (layers..., ParametricResNetLayer(arch.sys_dim, arch.width, arch.activation; parameters=arch.parameters, return_parameters=true)) + end + + # linear layers for the output + layers = (layers..., ParametricResNetLayer(arch.sys_dim, arch.width, identity; parameters=arch.parameters, return_parameters=false)) + + Chain(layers...) +end \ No newline at end of file diff --git a/src/architectures/resnet.jl b/src/architectures/resnet.jl index cc8e98519..bc8829b60 100644 --- a/src/architectures/resnet.jl +++ b/src/architectures/resnet.jl @@ -23,20 +23,23 @@ where `dl` is an instance of `DataLoader`. See [`iterate`](@ref) for an example of this. """ struct ResNet{AT} <: NeuralNetworkIntegrator - sys_dim::Int - n_blocks::Int + sys_dim::Int + n_blocks::Int + width::Int activation::AT end -function ResNet(dl::DataLoader, n_blocks::Integer; activation = tanh) - ResNet(dl.input_dim, n_blocks, activation) +ResNet(sys_dim::Integer, n_blocks::Integer, activation) = ResNet(sys_dim, n_blocks, sys_dim, activation) + +function ResNet(dl::DataLoader, n_blocks::Integer, width::Integer=dl.input_dim; activation = tanh) + ResNet(dl.input_dim, n_blocks, width, activation) end function Chain(arch::ResNet{AT}) where AT layers = () for _ in 1:arch.n_blocks # nonlinear layers - layers = (layers..., ResNetLayer(arch.sys_dim, arch.activation; use_bias=true)) + layers = (layers..., arch.sys_dim == arch.width ? ResNetLayer(arch.sys_dim, arch.activation; use_bias=true) : WideResNetLayer(arch.sys_dim, arch.width, arch.activation)) end # linear layers for the output diff --git a/src/architectures/standard_hamiltonian_neural_network.jl b/src/architectures/standard_hamiltonian_neural_network.jl new file mode 100644 index 000000000..a18b01062 --- /dev/null +++ b/src/architectures/standard_hamiltonian_neural_network.jl @@ -0,0 +1,85 @@ +""" + StandardHamiltonianArchitecture <: HamiltonianArchitecture + +A realization of the standard Hamiltonian neural network (HNN) [greydanus2019hamiltonian](@cite). + +Also see [`GeneralizedHamiltonianArchitecture`](@ref). + +# Constructor + +The constructor takes the following input arguments: +1. `dim`: system dimension, +2. `width = dim`: width of the hidden layer. By default this is equal to `dim`, +3. `nhidden = $(HNN_nhidden_default)`: the number of hidden layers, +4. `activation = $(HNN_activation_default)`: the activation function used in the HNN. +""" +struct StandardHamiltonianArchitecture{AT} <: HamiltonianArchitecture{AT} + dim::Int + width::Int + nhidden::Int + activation::AT + + function StandardHamiltonianArchitecture(dim::Integer, width=dim, + nhidden=HNN_nhidden_default, activation=HNN_activation_default) + @assert iseven(dim) "The input dimension must be an even integer." + new{typeof(activation)}(dim, width, nhidden, activation) + end +end + +""" + symbolic_hamiltonian_vector_field(nn::SymbolicNeuralNetwork) + +Get the symbolic expression for the vector field belonging to the HNN `nn`. + +# Implementation + +This is calling `SymbolicNeuralNetworks.Jacobian` and then multiplies the result with a Poisson tensor. +""" +function symbolic_hamiltonian_vector_field(nn::SymbolicNeuralNetwork) + □ = SymbolicNeuralNetworks.Jacobian(nn) + n = input_dimension(nn.model) ÷ 2 + # The Poisson tensor is built from *integers* on purpose: a `Float64` literal in the symbolic + # expression would widen the result of a `Float32` network. `PoissonTensor`, which has the same + # convention, is an `AbstractMatrix{Float64}` and would do exactly that. + 𝕆 = zeros(Int, n, n) + 𝕀 = Matrix(1I, n, n) + 𝕁 = [𝕆 𝕀; -𝕀 𝕆] + # `Jacobian` uses the convention `□[i, j] = ∂fᵢ/∂xⱼ` and the HNN output is scalar, so the one + # row of `derivative(□)` is the gradient of the Hamiltonian. The vector field is built as a + # *vector*, so that the generated function returns what `HNNLoss` compares against: a vector + # for a single sample, and one column per sample for a batch. + ∇H = vec(derivative(□)) + 𝕁 * ∇H +end + +""" + hamiltonian_vector_field(arch::StandardHamiltonianArchitecture) + +Compute an executable expression of the Hamiltonian vector field of a +[`StandardHamiltonianArchitecture`](@ref). + +# Implementation + +This first computes a symbolic expression of the vector field using [`symbolic_hamiltonian_vector_field`](@ref). + +The function is built with `inplace = false`: [`HNNLoss`](@ref) wraps it and is differentiated with +`Zygote`, and the in-place kernel `SymbolicNeuralNetworks.build_nn_function` builds by default +*mutates* its result, which `Zygote` does not support. +""" +function hamiltonian_vector_field(arch::StandardHamiltonianArchitecture) + nn = SymbolicNeuralNetwork(arch) + hvf = symbolic_hamiltonian_vector_field(nn) + SymbolicNeuralNetworks.build_nn_function(hvf, nn.params, nn.input; inplace = false) +end + +function Chain(arch::StandardHamiltonianArchitecture) + inner_layers = Tuple( + [Dense(arch.width, arch.width, arch.activation) for _ in 1:arch.nhidden] + ) + + Chain( + Dense(arch.dim, arch.width, arch.activation), + inner_layers..., + Linear(arch.width, 1; use_bias = false) + ) +end diff --git a/src/data_loader/batch.jl b/src/data_loader/batch.jl index c2afecbba..2f3e759fa 100644 --- a/src/data_loader/batch.jl +++ b/src/data_loader/batch.jl @@ -149,18 +149,19 @@ function number_of_batches(dl::DataLoader{T, AT, OT, :RegularData}, batch::Batch Int(ceil(dl.input_time_steps * dl.n_params / batch.batch_size)) end -function batch_over_two_axes(batch::Batch, number_columns::Int, third_dim::Int, dl::DataLoader) +function batch_over_two_axes(batch::Batch, number_columns::Integer, third_dim::Integer, n_batches::Integer) time_indices = shuffle(1:number_columns) parameter_indices = shuffle(1:third_dim) complete_indices = Iterators.product(time_indices, parameter_indices) |> collect |> vec batches = () - n_batches = number_of_batches(dl, batch) for batch_number in 1:(n_batches - 1) batches = (batches..., complete_indices[(batch_number - 1) * batch.batch_size + 1 : batch_number * batch.batch_size]) end (batches..., complete_indices[(n_batches - 1) * batch.batch_size + 1:end]) end +batch_over_two_axes(batch::Batch, number_of_columns::Integer, third_dim::Integer, dl::DataLoader) = batch_over_two_axes(batch, number_of_columns, third_dim, number_of_batches(dl, batch)) + function (batch::Batch)(dl::DataLoader{T, BT, OT, :RegularData}) where {T, AT<:AbstractArray{T, 3}, BT<:Union{AT, NamedTuple{(:q, :p), Tuple{AT, AT}}}, OT} batch_over_two_axes(batch, dl.input_time_steps, dl.n_params, dl) end @@ -195,7 +196,7 @@ end output[i, j, k] = data[i, indices[1, k] + seq_length + j - 1, indices[2, k]] end -# this is neeced if we want to use the vector of tuples in a kernel +# this is needed if we want to use the vector of tuples in a kernel function convert_vector_of_tuples_to_matrix(backend::Backend, batch_indices_tuple::Vector{Tuple{Int, Int}}) _batch_size = length(batch_indices_tuple) diff --git a/src/data_loader/optimize.jl b/src/data_loader/optimize.jl index a479c5bf6..0c47fca70 100644 --- a/src/data_loader/optimize.jl +++ b/src/data_loader/optimize.jl @@ -83,8 +83,8 @@ _copy(qp::QPT) = (q = copy(qp.q), p = copy(qp.p)) _copy(t::Tuple{<:QPTOAT, <:QPTOAT}) = _copy.(t) function (o::Optimizer)(nn::NeuralNetwork, - dl::DataLoader, - batch::Batch, + dl::Union{DataLoader, ParametricDataLoader}, + batch::Batch, n_epochs::Integer, loss::NetworkLoss, _pullback::AbstractPullback = ZygotePullback(loss); show_progress = true) diff --git a/src/data_loader/parametric_data_loader.jl b/src/data_loader/parametric_data_loader.jl new file mode 100644 index 000000000..6a520a7c7 --- /dev/null +++ b/src/data_loader/parametric_data_loader.jl @@ -0,0 +1,143 @@ +""" + ParametricDataLoader + +Very similar to [`DataLoader`](@ref), but can deal with parametric problems. +""" +struct ParametricDataLoader{T, AT<:QPTOAT2, VT<:AbstractVector} + input::AT + input_dim::Int + input_time_steps::Int + parameters::VT + n_params::Int + + function ParametricDataLoader(data::QPTOAT2{T, 3}, parameters::AbstractVector) where {T} + input_dim, input_time_steps, n_params = _size(data) + @assert T == _eltype(parameters) "Provided data and parameters must have the same eltype!" + @assert length(parameters) == _size(data, 3) "The number of provided parameters and the parameter axis of the supplied data do not have the same length!" + + new{T, typeof(data), typeof(parameters)}(data, input_dim, input_time_steps, parameters, n_params) + end +end + +function ParametricDataLoader(input::AbstractMatrix{T}, parameters::AbstractVector) where {T} + ParametricDataLoader(reshape(input, size(input)..., 1), parameters) +end + +# The same solution shape `DataLoader(::EnsembleSolution)` takes: since GeometricSolutions 0.6 a +# solution carries the time series and the vector field alongside `q` and `p`. +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}} + } + + sys_dim = length(ensemble_solution.s[1].q[0]) + input_time_steps = length(ensemble_solution.t) + n_params = length(ensemble_solution.s) + params = ensemble_solution.problem.parameters + + data = (q = zeros(T, sys_dim, input_time_steps, n_params), p = zeros(T, sys_dim, input_time_steps, n_params)) + + for (solution, i) in zip(ensemble_solution.s, axes(ensemble_solution.s, 1)) + for dim in 1:sys_dim + data.q[dim, :, i] = solution.q[:, dim] + data.p[dim, :, i] = solution.p[:, dim] + end + end + + ParametricDataLoader(data, params) +end + +# """ +# rearrange_parameters(parameters) +# +# Rearrange `parameters` such that they can be used by [`ParametricDataLoader`](@ref). +# """ +# function rearrange_parameters(parameters::Vector{<:NamedTuple}) +# parameters_rearranged = zeros(_eltype(parameters), ) +# end + +# function batch_over_two_axes(batch::Batch, number_columns::Int, third_dim::Int, dl::ParametricDataLoader) +# time_indices = shuffle(1:number_columns) +# parameter_indices = shuffle(1:third_dim) +# complete_indices = Iterators.product(time_indices, parameter_indices) |> collect |> vec +# batches = () +# n_batches = number_of_batches(dl, batch) +# for batch_number in 1:(n_batches - 1) +# batches = (batches..., complete_indices[(batch_number - 1) * batch.batch_size + 1 : batch_number * batch.batch_size]) +# end +# (batches..., complete_indices[(n_batches - 1) * batch.batch_size + 1:end]) +# end + +function optimize_for_one_epoch!( opt::Optimizer, + model, + ps::Union{NeuralNetworkParameters, NamedTuple}, + dl::ParametricDataLoader{T}, + batch::Batch, + _pullback::AbstractPullback, + λY) where T + count = 0 + total_error = T(0) + batches = batch(dl) + for batch_indices in batches + count += 1 + # these `copy`s should not be necessary! coming from a Zygote problem! + _input_nt_output_nt_parameter_indices = convert_input_and_batch_indices_to_array(dl, batch, batch_indices) + # input_nt_output_nt = _input_nt_output_nt_parameter_indices[1:2] + loss_value, pullback = _pullback(ps, model, _input_nt_output_nt_parameter_indices) + total_error += loss_value + dp = _unwrap_gradient(_get_contents(pullback(one(loss_value)))) + optimization_step!(opt, λY, ps, dp) + end + total_error / count +end + +function parameter_indices(parameters::AbstractVector, parameter_indices::AbstractVector{Int}) + [parameters[parameter_index] for parameter_index in parameter_indices] +end + +function parameter_indices(parameters::AbstractVector, batch_indices::AbstractMatrix{Int}) + parameter_indices(parameters, batch_indices[2, :]) +end + +function parameter_indices(dl::ParametricDataLoader, indices::AbstractArray{Int}) + parameter_indices(dl.parameters, indices) +end + +function convert_input_and_batch_indices_to_array(dl::ParametricDataLoader{T, BT}, batch::Batch, batch_indices_tuple::Vector{Tuple{Int, Int}}) where {T, AT<:AbstractArray{T, 3}, BT<:NamedTuple{(:q, :p), Tuple{AT, AT}}} + backend = networkbackend(dl.input.q) + + # the batch size is smaller for the last batch + _batch_size = length(batch_indices_tuple) + + batch_indices = convert_vector_of_tuples_to_matrix(backend, batch_indices_tuple) + + q_input = KernelAbstractions.allocate(backend, T, dl.input_dim ÷ 2, batch.seq_length, _batch_size) + p_input = similar(q_input) + + assign_input_from_vector_of_tuples! = assign_input_from_vector_of_tuples_kernel!(backend) + assign_input_from_vector_of_tuples!(q_input, p_input, dl.input, batch_indices, ndrange=(dl.input_dim ÷ 2, batch.seq_length, _batch_size)) + + q_output = KernelAbstractions.allocate(backend, T, dl.input_dim ÷ 2, batch.prediction_window, _batch_size) + p_output = similar(q_output) + + assign_output_from_vector_of_tuples! = assign_output_from_vector_of_tuples_kernel!(backend) + assign_output_from_vector_of_tuples!(q_output, p_output, dl.input, batch_indices, batch.seq_length, ndrange=(dl.input_dim ÷ 2, batch.prediction_window, _batch_size)) + + (q = q_input, p = p_input), (q = q_output, p = p_output), parameter_indices(dl, batch_indices) +end + +function number_of_batches(dl::ParametricDataLoader, batch::Batch) + @assert dl.input_time_steps ≥ (batch.seq_length + batch.prediction_window) "The number of time steps has to be greater than sequence length + prediction window." + Int(ceil((dl.input_time_steps - (batch.seq_length - 1) - batch.prediction_window) * dl.n_params / batch.batch_size)) +end + +function (batch::Batch)(dl::ParametricDataLoader) + batch_over_two_axes(batch, dl.input_time_steps - (batch.seq_length - 1) - batch.prediction_window, dl.n_params, number_of_batches(dl, batch)) +end + +function (o::Optimizer)(nn::NeuralNetwork{<:GeneralizedHamiltonianArchitecture}, dl::ParametricDataLoader, batch::Batch{:FeedForward}, n_epochs::Integer=1, loss::NetworkLoss=ParametricLoss(); kwargs...) + _pullback::AbstractPullback = ZygotePullback(loss) + o(nn, dl, batch, n_epochs, loss, _pullback; kwargs...) +end \ No newline at end of file diff --git a/src/layers/forcing_dissipation_layers.jl b/src/layers/forcing_dissipation_layers.jl new file mode 100644 index 000000000..cc7ea16cd --- /dev/null +++ b/src/layers/forcing_dissipation_layers.jl @@ -0,0 +1,163 @@ +@doc raw""" + ForcingLayer <: AbstractExplicitLayer + +Layers that can learn dissipative or forcing terms, but not conservative ones. + +Use the constructors [`ForcingLayerQ`](@ref) and [`ForcingLayerP`](@ref) for this. + +!!! warn + The forcing is dependent on either ``q`` or ``p``, but always applied to the ``p`` component. + +The forcing layers are inspired by the Lagrange-d'Alembert integrator from [marsden2001discrete; Example 3.2.2](@cite): + +```math +\begin{aligned} + q^{(t+1)} = & q^{(t)} + & hM^{-1}p^{(t)}, \\ + p^{(t+1)} = & p^{(t)} + & -h\nabla{}U(q^{(t+1)}) + hf_H(q^{(t+1)}, p^{(t)}), +\end{aligned} +``` +for a separable Hamiltonian ``H(q, p) = T(p) + U(q) = p^TM^{-1}p + U(q)`` and external forcing ``f_H.`` +""" +struct ForcingLayer{M,N,PT,CT,type,ReturnParameters} <: AbstractExplicitLayer{M,N} + dim::Int + width::Int + nhidden::Int + parameter_length::Int + parameter_layout::PT + model::CT +end + +parameterlength(l::ForcingLayer) = parameterlength(l.model) + +function initialparameters(rng::Random.AbstractRNG, init_weight::AbstractNeuralNetworks.Initializer, integrator::ForcingLayer, backend::KernelAbstractions.Backend, ::Type{T}) where {T} + initialparameters(rng, init_weight, integrator.model, backend, T) +end + +""" + ForcingLayerQ + +A layer that is derived from the more general [`ForcingLayer`](@ref) and the resulting forcing only depends on the ``q`` component. +""" +const ForcingLayerQ{M,N,FT,AT,ReturnParameters} = ForcingLayer{M,N,FT,AT,:Q,ReturnParameters} + +""" + ForcingLayerP + +A layer that is derived from the more general [`ForcingLayer`](@ref) and the resulting forcing only depends on the ``p`` component. +""" +const ForcingLayerP{M,N,FT,AT,ReturnParameters} = ForcingLayer{M,N,FT,AT,:P,ReturnParameters} + +""" + ForcingLayerQP + +A layer that is derived from the more general [`ForcingLayer`](@ref) and the resulting forcing only depends on the ``q`` and the ``p`` component. +""" +const ForcingLayerQP{M,N,FT,AT,ReturnParameters} = ForcingLayer{M,N,FT,AT,:QP,ReturnParameters} + +function build_chain(dim::Integer, width::Integer, nhidden::Integer, parameter_length::Integer, activation, type::Symbol) + inner_layers = Tuple( + [Dense(width, width, activation) for _ in 1:nhidden] + ) + + Chain( + type == :QP ? Dense(dim + parameter_length, width, activation) : Dense(dim ÷ 2 + parameter_length, width, activation), + inner_layers..., + Linear(width, dim ÷ 2; use_bias=false) + ) +end + +function ForcingLayer(dim::Integer, width::Integer, nhidden::Integer, activation; parameters::OptionalParameters=NullParameters(), return_parameters::Bool, type::Symbol) + flat_parameters, layout = _flatten_system_parameters(parameters) + parameter_length = length(flat_parameters) + c = build_chain(dim, width, nhidden, parameter_length, activation, type) + ForcingLayer{dim,dim,typeof(layout),typeof(c),type,return_parameters}(dim, width, nhidden, parameter_length, layout, c) +end + +""" + ForcingLayerQ(dim) + +# Examples + +```julia +ForcingLayerQ(dim, width, nhidden, activation; parameters, return_parameters) +``` +""" +function ForcingLayerQ(dim::Integer, width::Integer=dim, nhidden::Integer=HNN_nhidden_default, activation=HNN_activation_default; parameters::OptionalParameters=NullParameters(), return_parameters::Bool=false) + ForcingLayer(dim, width, nhidden, activation; parameters=parameters, return_parameters=return_parameters, type=:Q) +end + +""" + ForcingLayerP(dim) + +See [`ForcingLayerQ`](@ref). + +# Examples + +```julia +ForcingLayerP(dim, width, nhidden, activation; parameters, return_parameters) +``` +""" +function ForcingLayerP(dim::Integer, width::Integer=dim, nhidden::Integer=HNN_nhidden_default, activation=HNN_activation_default; parameters::OptionalParameters=NullParameters(), return_parameters::Bool=false) + ForcingLayer(dim, width, nhidden, activation; parameters=parameters, return_parameters=return_parameters, type=:P) +end + +function ForcingLayerQP(dim::Integer, width::Integer=dim, nhidden::Integer=HNN_nhidden_default, activation=HNN_activation_default; parameters::OptionalParameters=NullParameters(), return_parameters::Bool=false) + ForcingLayer(dim, width, nhidden, activation; parameters=parameters, return_parameters=return_parameters, type=:QP) +end + +function (integrator::ForcingLayerQ{M,N,FT,AT,false})(qp::QPT2, problem_params::OptionalParameters, params::NeuralNetworkParameters) where {M,N,FT,AT} + input = concatenate_array_with_parameters(qp.q, problem_params) + (q=qp.q, p=qp.p + integrator.model(input, params)) +end + +function (integrator::ForcingLayerP{M,N,FT,AT,false})(qp::QPT2, problem_params::OptionalParameters, params::NeuralNetworkParameters) where {M,N,FT,AT} + input = concatenate_array_with_parameters(qp.p, problem_params) + (q=qp.q, p=qp.p + integrator.model(input, params)) +end + +function (integrator::ForcingLayerQP{M,N,FT,AT,false})(qp::QPT2, problem_params::OptionalParameters, params::NeuralNetworkParameters) where {M,N,FT,AT} + input = concatenate_array_with_parameters(vcat(qp.q, qp.p), problem_params) + (q=qp.q, p=qp.p + integrator.model(input, params)) +end + +function (integrator::ForcingLayerQ{M,N,FT,AT,true})(qp::QPT2, problem_params::OptionalParameters, params::NeuralNetworkParameters) where {M,N,FT,AT} + input = concatenate_array_with_parameters(qp.q, problem_params) + ((q=qp.q, p=qp.p + integrator.model(input, params)), problem_params) +end + +function (integrator::ForcingLayerP{M,N,FT,AT,true})(qp::QPT2, problem_params::OptionalParameters, params::NeuralNetworkParameters) where {M,N,FT,AT} + input = concatenate_array_with_parameters(qp.p, problem_params) + ((q=qp.q, p=qp.p + integrator.model(input, params)), problem_params) +end + +function (integrator::ForcingLayerQP{M,N,FT,AT,true})(qp::QPT2, problem_params::OptionalParameters, params::NeuralNetworkParameters) where {M,N,FT,AT} + input = concatenate_array_with_parameters(vcat(qp.q, qp.p), problem_params) + ((q=qp.q, p=qp.p + integrator.model(input, params)), problem_params) +end + +function (integrator::ForcingLayer)(qp_params::Tuple{<:QPTOAT2,<:OptionalParameters}, params::NeuralNetworkParameters) + integrator(qp_params..., params) +end + +function (integrator::ForcingLayer)(::TT, ::NeuralNetworkParameters) where {TT<:Tuple} + error("The input is of type $(TT). This shouldn't be the case!") +end + +function (integrator::ForcingLayer{M,N,FT,AT,Type,true})(qp::AbstractArray, problem_params::OptionalParameters, params::NeuralNetworkParameters) where {M,N,FT,AT,Type} + @assert iseven(size(qp, 1)) + n = size(qp, 1) ÷ 2 + qp_split = assign_q_and_p(qp, n) + evaluated = integrator(qp_split, problem_params, params)[1] + (vcat(evaluated.q, evaluated.p), problem_params) +end + +function (integrator::ForcingLayer{M,N,FT,AT,Type,false})(qp::AbstractArray, problem_params::OptionalParameters, params::NeuralNetworkParameters) where {M,N,FT,AT,Type} + @assert iseven(size(qp, 1)) + n = size(qp, 1) ÷ 2 + qp_split = assign_q_and_p(qp, n) + evaluated = integrator(qp_split, problem_params, params) + vcat(evaluated.q, evaluated.p) +end + +(integrator::ForcingLayer)(qp::QPTOAT2, params::NeuralNetworkParameters) = integrator(qp, NullParameters(), params) +(integrator::ForcingLayer)(qp::QPTOAT2, params::NamedTuple) = integrator(qp, NeuralNetworkParameters(params)) diff --git a/src/layers/parametric_resnet_layer.jl b/src/layers/parametric_resnet_layer.jl new file mode 100644 index 000000000..149d6c4e0 --- /dev/null +++ b/src/layers/parametric_resnet_layer.jl @@ -0,0 +1,64 @@ +struct ParametricResNetLayer{M, N, F1 <: Activation, PT, ReturnParameters} <: AbstractExplicitLayer{M, N} + width::Int + activation::F1 + parameter_length::Int + parameter_layout::PT +end + +function ParametricResNetLayer(dim::Integer, width::Integer, activation=identity; parameters::OptionalParameters=NullParameters(), return_parameters::Bool) + flat_parameters, layout = _flatten_system_parameters(parameters) + _activation = Activation(activation) + ParametricResNetLayer{dim, dim, typeof(_activation), typeof(layout), return_parameters}(width, _activation, length(flat_parameters), layout) +end + +function initialparameters(rng::Random.AbstractRNG, init_weight::AbstractNeuralNetworks.Initializer, l::ParametricResNetLayer{M, M}, backend::KernelAbstractions.Backend, ::Type{T}; init_bias = ZeroInitializer()) where {M, T} + upscale_weight = KernelAbstractions.allocate(backend, T, l.width, M + l.parameter_length) + upscale_bias = KernelAbstractions.allocate(backend, T, l.width) + downscale_weight = KernelAbstractions.allocate(backend, T, M, l.width) + bias = KernelAbstractions.allocate(backend, T, M) + init_weight(rng, upscale_weight) + init_weight(rng, downscale_weight) + init_bias(rng, upscale_bias) + init_bias(rng, bias) + (upscale_weight=upscale_weight, downscale_weight=downscale_weight, upscale_bias=upscale_bias, bias=bias) +end + +parameterlength(l::ParametricResNetLayer{M, M}) where {M} = (l.width + l.parameter_length) * (M + 1) + M * (l.width + 1) + +function (d::ParametricResNetLayer{M, M, F, PT, false})(x::AbstractVecOrMat, problem_params::OptionalParameters, ps::NamedTuple) where {M, F, PT} + input = concatenate_array_with_parameters(x, problem_params) + x + d.activation.(ps.downscale_weight * d.activation.(ps.upscale_weight * input .+ ps.upscale_bias) .+ ps.bias) +end + +function (d::ParametricResNetLayer{M, M, F, PT, true})(x::AbstractVecOrMat, problem_params::OptionalParameters, ps::NamedTuple) where {M, F, PT} + input = concatenate_array_with_parameters(x, problem_params) + (x + d.activation.(ps.downscale_weight * d.activation.(ps.upscale_weight * input .+ ps.upscale_bias) .+ ps.bias), problem_params) +end + +# function (d::ParametricResNetLayer{M, M, F, PT, false})(x::AbstractArray{T, 3}, problem_params::AbstractVector, ps::NamedTuple) where {M, F, PT, T} +# input = concatenate_array_with_parameters(x, problem_params) +# x + d.activation.(mat_tensor_mul(ps.downscale_weight, d.activation.(mat_tensor_mul(ps.upscale_weight, x) .+ ps.upscale_bias)) .+ ps.bias) +# end +# +# function (d::ParametricResNetLayer{M, M, F, PT, true})(x::AbstractArray{T, 3}, problem_params::AbstractVector, ps::NamedTuple) where {M, F, PT, T} +# input = concatenate_array_with_parameters(x, problem_params) +# (x + d.activation.(mat_tensor_mul(ps.downscale_weight, d.activation.(mat_tensor_mul(ps.upscale_weight, x) .+ ps.upscale_bias)) .+ ps.bias), problem_params) +# end + +(d::ParametricResNetLayer)(input::Tuple, ps::NamedTuple) = length(input) == 2 ? d(input..., ps) : error("The tuple must contain the input array/nt as well as the system parameters.") + +function (d::ParametricResNetLayer{M, M, F, PT, false})(z::QPT, problem_params::OptionalParameters, ps::NamedTuple) where {M, F, PT} + @assert iseven(M) + @assert size(z.q, 1) * 2 == M + N2 = M ÷ 2 + output = d(vcat(z.q, z.p), problem_params, ps) + assign_q_and_p(output, N2) +end + +function (d::ParametricResNetLayer{M, M, F, PT, true})(z::QPT, problem_params::OptionalParameters, ps::NamedTuple) where {M, F, PT} + @assert iseven(M) + @assert size(z.q, 1) * 2 == M + N2 = M ÷ 2 + output = d(vcat(z.q, z.p), problem_params, ps) + (assign_q_and_p(output[1], N2), problem_params) +end \ No newline at end of file diff --git a/src/layers/sympnets.jl b/src/layers/sympnets.jl index 5122467a3..cecbf49cd 100644 --- a/src/layers/sympnets.jl +++ b/src/layers/sympnets.jl @@ -236,36 +236,36 @@ function custom_vec_mul(scale::AbstractVector{T}, x::AbstractArray{T, 3}) where vec_tensor_mul(scale, x) end -@inline function (d::ActivationLayerQ{M, M})(x::NamedTuple, ps) where {M} +@inline function (d::ActivationLayerQ{M, M})(x::QPT2, ps) where {M} size(x.q, 1) == M÷2 || error("Dimension mismatch.") return (q = x.q + custom_vec_mul(ps.scale, d.activation.(x.p)), p = x.p) end -@inline function (d::ActivationLayerP{M, M})(x::NamedTuple, ps) where {M} +@inline function (d::ActivationLayerP{M, M})(x::QPT2, ps) where {M} size(x.q, 1) == M÷2 || error("Dimension mismatch.") return (q = x.q, p = x.p + custom_vec_mul(ps.scale, d.activation.(x.q))) end -@inline function (d::GradientLayerQ{M, M})(x::NamedTuple, ps) where {M} +@inline function (d::GradientLayerQ{M, M})(x::QPT2, ps) where {M} size(x.q, 1) == M÷2 || error("Dimension mismatch.") (q = x.q + custom_mat_mul(ps.weight', (custom_vec_mul(ps.scale, d.activation.(custom_mat_mul(ps.weight, x.p) .+ ps.bias)))), p = x.p) end -@inline function(d::GradientLayerP{M, M})(x::NamedTuple, ps) where {M} +@inline function(d::GradientLayerP{M, M})(x::QPT2, ps) where {M} size(x.q, 1) == M÷2 || error("Dimension mismatch.") (q = x.q, p = x.p + custom_mat_mul(ps.weight', (custom_vec_mul(ps.scale, d.activation.(custom_mat_mul(ps.weight, x.q) .+ ps.bias))))) end -@inline function(d::LinearLayerQ{M, M})(x::NamedTuple, ps) where {M} +@inline function(d::LinearLayerQ{M, M})(x::QPT2, ps) where {M} size(x.q, 1) == M÷2 || error("Dimension mismatch.") (q = x.q + custom_mat_mul(ps.weight, x.p), p = x.p) end -@inline function(d::LinearLayerP{M, M})(x::NamedTuple, ps) where {M} +@inline function(d::LinearLayerP{M, M})(x::QPT2, ps) where {M} size(x.q, 1) == M÷2 || error("Dimension mismatch.") (q = x.q, p = x.p + custom_mat_mul(ps.weight, x.q)) end diff --git a/src/layers/wide_resnet.jl b/src/layers/wide_resnet.jl new file mode 100644 index 000000000..f1a4c4682 --- /dev/null +++ b/src/layers/wide_resnet.jl @@ -0,0 +1,32 @@ +struct WideResNetLayer{M, N, F1} <: AbstractExplicitLayer{M, N} + width::Int + activation::F1 +end + +WideResNetLayer(dim::Integer, width::Integer, activation=identity) = WideResNetLayer{dim, dim, typeof(activation)}(width, activation) + +function initialparameters(rng::Random.AbstractRNG, init_weight::AbstractNeuralNetworks.Initializer, l::WideResNetLayer{M, M}, backend::KernelAbstractions.Backend, ::Type{T}; init_bias = ZeroInitializer()) where {M, T} + upscale_weight = KernelAbstractions.allocate(backend, T, l.width, M) + upscale_bias = KernelAbstractions.allocate(backend, T, l.width) + downscale_weight = KernelAbstractions.allocate(backend, T, M, l.width) + bias = KernelAbstractions.allocate(backend, T, M) + init_weight(rng, upscale_weight) + init_weight(rng, downscale_weight) + init_bias(rng, upscale_bias) + init_bias(rng, bias) + (upscale_weight=upscale_weight, downscale_weight=downscale_weight, upscale_bias=upscale_bias, bias=bias) +end + +parameterlength(l::WideResNetLayer{M, M}) where {M} = l.width * (M + 1) + M * (l.width + 1) + +(d::WideResNetLayer{M, M})(x::AbstractVecOrMat, ps::NamedTuple) where {M} = x + d.activation.(ps.downscale_weight * d.activation.(ps.upscale_weight * x .+ ps.upscale_bias) .+ ps.bias) + +(d::WideResNetLayer{M, M})(x::AbstractArray{T, 3}, ps::NamedTuple) where {M, T} = x + d.activation.(mat_tensor_mul(ps.downscale_weight, d.activation.(mat_tensor_mul(ps.upscale_weight, x) .+ ps.upscale_bias)) .+ ps.bias) + +function (d::WideResNetLayer{M, M})(z::QPT, ps::NamedTuple) where {M} + @assert iseven(M) + @assert size(z.q, 1) * 2 == M + N2 = M ÷ 2 + output = d(vcat(z.q, z.p), ps) + assign_q_and_p(output, N2) +end \ No newline at end of file diff --git a/src/loss/losses.jl b/src/loss/losses.jl index 2354a1551..c814d972a 100644 --- a/src/loss/losses.jl +++ b/src/loss/losses.jl @@ -257,3 +257,25 @@ function (loss::ReducedLoss)(model::Chain, params::NetworkParameters, input::CT, output::CT) where {CT <: QPTOAT} _compute_loss(loss.decoder(model(loss.encoder(input), params)), output) end + +@doc raw""" + ParametricLoss() + +The loss for a network whose forward pass takes the parameters of the system alongside the input, +i.e. the parameter-dependent architectures built on [`GeneralizedHamiltonianArchitecture`](@ref). + +It is `FeedForwardLoss` with the system parameters threaded through: + +```math +L(\mathtt{input}, \mathtt{output}, \mu) = ||\mathcal{NN}(\mathtt{input}, \mu) - \mathtt{output}||. +``` + +This loss does not have any parameters. +""" +struct ParametricLoss <: NetworkLoss end + +function (loss::ParametricLoss)(model::Chain, + params::Union{NamedTuple, NeuralNetworkParameters}, input::CT, output::CT, + system_parameters::Union{NamedTuple, AbstractVector}) where {CT <: QPTOAT} + _compute_loss(model(input, system_parameters, params), output) +end diff --git a/src/pullbacks/symbolic_hnn_pullback.jl b/src/pullbacks/symbolic_hnn_pullback.jl index c5168be93..a812d4be9 100644 --- a/src/pullbacks/symbolic_hnn_pullback.jl +++ b/src/pullbacks/symbolic_hnn_pullback.jl @@ -25,3 +25,72 @@ function SymbolicPullback(arch::HamiltonianArchitecture) soutput; reduce = +) SymbolicPullback(loss, SymbolicNeuralNetworks.ParameterGradient(gradient_function)) end + +@doc raw""" + SymbolicPullback(nn, loss, system_params) + +The `SymbolicPullback` for a network whose forward pass also takes the parameters of the *system*, +i.e. one built on [`GeneralizedHamiltonianArchitecture`](@ref), with a [`ParametricLoss`](@ref). + +# Implementation + +This is `SymbolicNeuralNetworks.SymbolicPullback(nn, loss)` with the system parameters threaded +through. `build_nn_function` generates a function of *one* input array, so the flattened system +parameters are appended to the network input, and the symbolic expression splits them off again with +[`_flatten_system_parameters`](@ref) and `unflatten`. The numeric side does the same concatenation, +in the call operators below. + +`reduce = +`: the loss of a batch is the sum of the losses of its samples, so its gradient is the +sum of the per-sample gradients. +""" +function SymbolicPullback(nn::NeuralNetwork, loss::ParametricLoss, + system_params::OptionalParameters; cse::Bool = true, inplace::Bool = true) + symbolic_system_parameters = SymbolicNeuralNetworks.symbolic_variables(system_params, :S) + symbolic_network_parameters = SymbolicNeuralNetworks.symbolic_variables(params(nn), :W) + + input_dim = input_dimension(nn.model) + _, parameter_layout = _flatten_system_parameters(SymbolicNeuralNetworks.Symbolics.Num, + symbolic_system_parameters) + sinput = Symbolics.variables(:x, 1:(input_dim + length(system_params))) + soutput = Symbolics.variables(:y, 1:output_dimension(nn.model)) + symbolic_system_input = unflatten(parameter_layout, sinput[(input_dim + 1):end]) + + symbolic_loss = loss(nn.model, symbolic_network_parameters, sinput[1:input_dim], soutput, + symbolic_system_input) + differentials = SymbolicNeuralNetworks.symbolic_differentials(symbolic_network_parameters) + gradient = SymbolicNeuralNetworks.symbolic_derivative(symbolic_loss, differentials) + gradient_function = SymbolicNeuralNetworks.build_nn_function( + gradient, symbolic_network_parameters, sinput, soutput; + reduce = +, cse = cse, inplace = inplace) + SymbolicPullback(loss, SymbolicNeuralNetworks.ParameterGradient(gradient_function)) +end + +# TODO: type piracy -- `SymbolicPullback` is `SymbolicNeuralNetworks`', and so is every argument +# type here. These belong upstream, together with a `build_nn_function` that takes more than one +# data argument, which is what would make the concatenation below unnecessary. +# +# The generated pullback takes *one* input array, so the system parameters are appended to the +# network input before it is called; the loss, which knows about them, gets them separately. +function (_pullback::SymbolicPullback)(ps, model, + input_output_params::Tuple{<:AbstractMatrix, <:AbstractMatrix, + <:Union{NamedTuple, AbstractVector}})::Tuple + input, output, system_params = input_output_params + _pullback.loss(model, ps, input, output, system_params), + _pullback.fun(concatenate_array_with_parameters(input, system_params), output, ps) +end + +# A batch with a time axis: the network is applied sample-wise, so the time and parameter axes are +# folded into one before the pullback sees them. +function (_pullback::SymbolicPullback)(ps, model, + input_output_params::Tuple{AT, AT, <:Union{NamedTuple, AbstractVector}})::Tuple where {T, AT <: AbstractArray{T, 3}} + input, output, system_params = input_output_params + _input = reshape(input, size(input, 1), size(input, 2) * size(input, 3)) + _output = reshape(output, size(output, 1), size(output, 2) * size(output, 3)) + _pullback(ps, model, (_input, _output, system_params)) +end + +function (_pullback::SymbolicPullback)(ps, model, + input_output_params::Tuple{<:QPT, <:QPT, <:Union{NamedTuple, AbstractVector}})::Tuple + input, output, system_params = input_output_params + _pullback(ps, model, (vcat(input.q, input.p), vcat(output.q, output.p), system_params)) +end diff --git a/src/pullbacks/zygote_pullback.jl b/src/pullbacks/zygote_pullback.jl index df882e224..ac7fba027 100644 --- a/src/pullbacks/zygote_pullback.jl +++ b/src/pullbacks/zygote_pullback.jl @@ -31,6 +31,12 @@ end ps -> _pullback.loss(model, ps, input_nt), ps) (_pullback::ZygotePullback)(ps, model, input_nt_output_nt::Tuple{<:QPTOAT, <:QPTOAT})::Tuple = Zygote.pullback( ps -> _pullback.loss(model, ps, input_nt_output_nt...), ps) +# The parameter-dependent architectures take the system parameters as a third element of the +# input tuple, either as a `NamedTuple` of parameters or as one vector entry per sample. +(_pullback::ZygotePullback)(ps, model, input_output_params::Tuple{<:QPTOAT, <:QPTOAT, <:NamedTuple})::Tuple = Zygote.pullback( + ps -> _pullback.loss(model, ps, input_output_params...), ps) +(_pullback::ZygotePullback)(ps, model, input_output_params::Tuple{<:QPTOAT, <:QPTOAT, <:AbstractVector})::Tuple = Zygote.pullback( + ps -> _pullback.loss(model, ps, input_output_params...), ps) """ _get_contents(returned_pullback) diff --git a/src/training_method/symplectic_euler.jl b/src/training_method/symplectic_euler.jl index fab5de86b..ea06c5b38 100644 --- a/src/training_method/symplectic_euler.jl +++ b/src/training_method/symplectic_euler.jl @@ -1,27 +1,27 @@ -abstract type SymplecticEuler <: HnnTrainingMethod end +abstract type SymplecticEulerIntegrator <: HnnTrainingMethod end -struct SymplecticEulerA <: SymplecticEuler end -struct SymplecticEulerB <: SymplecticEuler end +struct SymplecticEulerIntegratorA <: SymplecticEulerIntegrator end +struct SymplecticEulerIntegratorB <: SymplecticEulerIntegrator end SEuler(;sqdist = sqeuclidean) = SEulerA(sqdist = sqdist) -SEulerA(;sqdist = sqeuclidean) = TrainingMethod{SymplecticEulerA, PhaseSpaceSymbol, TrajectoryData, typeof(sqdist)}(sqdist) -SEulerB(;sqdist = sqeuclidean) = TrainingMethod{SymplecticEulerB, PhaseSpaceSymbol, TrajectoryData, typeof(sqdist)}(sqdist) +SEulerA(;sqdist = sqeuclidean) = TrainingMethod{SymplecticEulerIntegratorA, PhaseSpaceSymbol, TrajectoryData, typeof(sqdist)}(sqdist) +SEulerB(;sqdist = sqeuclidean) = TrainingMethod{SymplecticEulerIntegratorB, PhaseSpaceSymbol, TrajectoryData, typeof(sqdist)}(sqdist) -function loss_single(::TrainingMethod{SymplecticEulerA}, nn::AbstractNeuralNetwork{<:HamiltonianArchitecture}, qₙ, qₙ₊₁, pₙ, pₙ₊₁, Δt, params = params(nn)) +function loss_single(::TrainingMethod{SymplecticEulerIntegratorA}, nn::AbstractNeuralNetwork{<:HamiltonianArchitecture}, qₙ, qₙ₊₁, pₙ, pₙ₊₁, Δt, params = params(nn)) dH = vectorfield(nn, [qₙ₊₁...,pₙ...], params) sqeuclidean(dH[1],(qₙ₊₁-qₙ)/Δt) + sqeuclidean(dH[2],(pₙ₊₁-pₙ)/Δt) end -function loss_single(::TrainingMethod{SymplecticEulerB}, nn::AbstractNeuralNetwork{<:HamiltonianArchitecture}, qₙ, qₙ₊₁, pₙ, pₙ₊₁, Δt, params = params(nn)) +function loss_single(::TrainingMethod{SymplecticEulerIntegratorB}, nn::AbstractNeuralNetwork{<:HamiltonianArchitecture}, qₙ, qₙ₊₁, pₙ, pₙ₊₁, Δt, params = params(nn)) dH = vectorfield(nn, [qₙ...,pₙ₊₁...], params) sqeuclidean(dH[1],(qₙ₊₁-qₙ)/Δt) + sqeuclidean(dH[2],(pₙ₊₁-pₙ)/Δt) end -get_loss(::TrainingMethod{<:SymplecticEuler}, ::AbstractNeuralNetwork{<:HamiltonianArchitecture}, data::TrainingData{<:DataSymbol{<:PhaseSpaceSymbol}}, args) = (get_data(data,:q, args...), get_data(data,:q, next(args...)...), get_data(data,:p, args...), get_data(data,:p,next(args...)...), get_Δt(data)) +get_loss(::TrainingMethod{<:SymplecticEulerIntegrator}, ::AbstractNeuralNetwork{<:HamiltonianArchitecture}, data::TrainingData{<:DataSymbol{<:PhaseSpaceSymbol}}, args) = (get_data(data,:q, args...), get_data(data,:q, next(args...)...), get_data(data,:p, args...), get_data(data,:p,next(args...)...), get_Δt(data)) -loss(ti::TrainingMethod{<:SymplecticEuler}, nn::AbstractNeuralNetwork{<:HamiltonianArchitecture}, data::TrainingData{<:DataSymbol{<:PhaseSpaceSymbol}}, index_batch = eachindex(ti, data), params = params(nn)) = +loss(ti::TrainingMethod{<:SymplecticEulerIntegrator}, nn::AbstractNeuralNetwork{<:HamiltonianArchitecture}, data::TrainingData{<:DataSymbol{<:PhaseSpaceSymbol}}, index_batch = eachindex(ti, data), params = params(nn)) = mapreduce(args->loss_single(Zygote.ignore_derivatives(ti), nn, get_loss(ti, nn, data, args)..., params),+, index_batch) -min_length_batch(::SymplecticEuler) = 2 \ No newline at end of file +min_length_batch(::SymplecticEulerIntegrator) = 2 \ No newline at end of file diff --git a/src/utils.jl b/src/utils.jl index d4d28a872..dbf0ee1a0 100644 --- a/src/utils.jl +++ b/src/utils.jl @@ -151,7 +151,16 @@ qp = (q = [1, 2], p = [3, 4]) ``` """ -const QPT{T} = NamedTuple{(:q, :p), Tuple{AT, AT}} where {T, AT <: AbstractArray{T}} +const QPT{T} = NamedTuple{(:q, :p), Tuple{AT, AT}} where {T, N, AT <: AbstractArray{T, N}} + +@doc raw""" + QPT2 + +[`QPT`](@ref) with the number of dimensions of the two arrays fixed, but their types allowed to +differ. A `Chain` that splits an input array into `q` and `p` produces views of different types, so +the layers of a parameter-dependent network dispatch on this rather than on `QPT`. +""" +const QPT2{T, N} = NamedTuple{(:q, :p), Tuple{AT₁, AT₂}} where {T, N, AT₁ <: AbstractArray{T, N}, AT₂ <: AbstractArray{T, N}} @doc raw""" QPTOAT @@ -165,9 +174,82 @@ This could be data in ``(q, p)\in\mathbb{R}^{2d}`` form or come from an arbitrar """ const QPTOAT{T} = Union{QPT{T}, AbstractArray{T}} where {T} +@doc raw""" + QPTOAT2 + +[`QPTOAT`](@ref) with the number of dimensions of the arrays fixed: + +```julia +const QPTOAT2 = Union{QPT2, AbstractArray} +``` +""" +const QPTOAT2{T, N} = Union{QPT2{T, N}, AbstractArray{T, N}} where {T, N} + Base.:≈(qp₁::QPT, qp₂::QPT) = (qp₁.q ≈ qp₂.q) & (qp₁.p ≈ qp₂.p) +@doc raw""" + _flatten_system_parameters(parameters) + _flatten_system_parameters(T, parameters) + +Flatten the parameters of the *system* (not of the network) into a vector, together with the +`NeuralNetworkParameters.ParameterLayout` that puts such a vector back into the original shape. + +The parameter-dependent architectures — [`GeneralizedHamiltonianArchitecture`](@ref) and the layers +it is built from — feed the system parameters to the network as extra input components, so they have +to be a vector. `NullParameters` flattens to an empty one, which makes the parameter-free case fall +out of the same code path. + +The layout is a *value*, not a closure, so a layer can store it in a field and stay inferable. +""" +_flatten_system_parameters(parameters::NamedTuple) = flatten(parameters) +_flatten_system_parameters(::NullParameters) = flatten(NamedTuple()) +_flatten_system_parameters(::Type{T}, parameters::NamedTuple) where {T} = flatten(T, parameters) +_flatten_system_parameters(::Type{T}, ::NullParameters) where {T} = flatten(T, NamedTuple()) + +""" + _unwrap_gradient(dp) + +Strip the `NetworkParameters` wrappers and the `(params = …,)` layers out of a gradient, so +that it has the same shape as the parameters it belongs to. + +`Zygote` differentiates *through* the `NetworkParameters` struct, so the gradient of a +parameter set comes back as a `NamedTuple` with a single `params` field. [`_get_params`](@ref) undoes +that at the top level. The parameter-dependent architectures nest — a `SymplecticEuler` layer +holds the parameters of a whole sub-network — so the unwrapping has to recurse. +""" +_unwrap_gradient(dp) = dp +_unwrap_gradient(dp::NetworkParameters) = _unwrap_gradient(params(dp)) +_unwrap_gradient(dp::NamedTuple{(:params,)}) = _unwrap_gradient(dp.params) +_unwrap_gradient(dp::NamedTuple) = map(_unwrap_gradient, dp) + _eltype(x) = eltype(x) _eltype(ps::NamedTuple) = _eltype(ps[1]) _eltype(ps::Tuple) = _eltype(ps[1]) _eltype(ps::NetworkParameters) = _eltype(params(ps)[1]) + +# `ParametricDataLoader` stores one `NamedTuple` of system parameters per trajectory, and they all +# have to agree with the element type of the data. +function _eltype(parameters::AbstractVector{<:NamedTuple}) + T = _eltype(first(parameters)) + for p in parameters + _eltype(p) == T || error("The parameters do not all have the same element type.") + end + T +end + +# `size` that also works on `(q, p)` data, where the first axis is the concatenation of the two. +_size(x) = size(x) +function _size(qp::QPT) + q_size = _size(qp.q) + p_size = _size(qp.p) + @assert q_size == p_size + (2q_size[1], q_size[2:end]...) +end + +_size(x, a::Integer) = size(x, a) +function _size(qp::QPT, a::Integer) + q_size = _size(qp.q, a) + p_size = _size(qp.p, a) + @assert q_size == p_size + a == 1 ? 2q_size : q_size +end diff --git a/test/data_loader/parametric_data_loader_test.jl b/test/data_loader/parametric_data_loader_test.jl new file mode 100644 index 000000000..12c9caa3d --- /dev/null +++ b/test/data_loader/parametric_data_loader_test.jl @@ -0,0 +1,45 @@ +using GeometricMachineLearning +using GeometricMachineLearning: convert_input_and_batch_indices_to_array +using Test +using GeometricProblems.CoupledHarmonicOscillator: hodeensemble, default_parameters +using GeometricIntegrators: ImplicitMidpoint, integrate +using Random: seed! +seed!(123) + +function make_alternative_parameters_by_adding_constant(params::NamedTuple = default_parameters(), + a::Number = 1.) + NamedTuple{keys(params)}(Tuple(value .+ a for value in values(params))) +end + +all_parameters = [default_parameters(), make_alternative_parameters_by_adding_constant()] + +h_ensemble = hodeensemble(; parameters = all_parameters) +sol = integrate(h_ensemble, ImplicitMidpoint()) +dl = ParametricDataLoader(sol) +batch = Batch(2) +batch_indices = batch(dl) + +# Each entry of a batch is a `(time index, parameter index)` pair, and the third element of what +# `convert_input_and_batch_indices_to_array` returns has to be the parameters of *that* trajectory. +# The batches are shuffled, so this asserts the correspondence rather than which batch holds which +# parameters -- pinning the latter makes the test depend on the RNG stream of the Julia version. +function batch_is_consistent(n::Integer) + input, output, parameters = convert_input_and_batch_indices_to_array(dl, batch, batch_indices[n]) + all(enumerate(batch_indices[n])) do (k, (time_index, parameter_index)) + parameters[k] == all_parameters[parameter_index] && + input.q[:, 1, k] == dl.input.q[:, time_index, parameter_index] && + input.p[:, 1, k] == dl.input.p[:, time_index, parameter_index] && + output.q[:, 1, k] == dl.input.q[:, time_index + 1, parameter_index] && + output.p[:, 1, k] == dl.input.p[:, time_index + 1, parameter_index] + end +end + +@test all(batch_is_consistent, eachindex(batch_indices)) + +# Both parameter sets have to turn up somewhere, or the assertion above would also pass on a data +# loader that always returned the first one. +returned_parameters = Set(parameters + for n in eachindex(batch_indices) + for parameters in last(convert_input_and_batch_indices_to_array( + dl, batch, batch_indices[n]))) +@test returned_parameters == Set(all_parameters) diff --git a/test/generalized_hamiltonian_neural_networks/pghnn_symbolic_pullback_single_layer_test.jl b/test/generalized_hamiltonian_neural_networks/pghnn_symbolic_pullback_single_layer_test.jl new file mode 100644 index 000000000..a4c52604e --- /dev/null +++ b/test/generalized_hamiltonian_neural_networks/pghnn_symbolic_pullback_single_layer_test.jl @@ -0,0 +1,63 @@ +# The symbolic pullback of a parameter-dependent network, on the smallest such network there is: a +# single `SymplecticEulerB` layer built from a `SymbolicPotentialEnergy`. +# +# `SymbolicPullback(nn, ::ParametricLoss, system_params)` builds its gradient with `reduce = +`, so +# what it returns is the *sum* of the per-sample gradients -- that is the convention +# `SymbolicNeuralNetworks.SymbolicPullback` uses too. The test compares against exactly that, +# computed with `Zygote` one sample at a time. + +using GeometricMachineLearning +using GeometricMachineLearning: ParametricLoss, SymbolicPotentialEnergy, SymplecticEulerB +using AbstractNeuralNetworks: params +using Random: seed! +using Test +import Zygote + +seed!(1234) + +system_parameters = (m = 1.0, ω = π / 2) +dim, width, nhidden, activation = 2, 2, 1, tanh +n_samples = 10 + +se = SymbolicPotentialEnergy(dim, width, nhidden, activation; parameters = system_parameters) +nn = NeuralNetwork(Chain(SymplecticEulerB(se; return_parameters = false))) + +loss = ParametricLoss() +pullback = SymbolicPullback(nn, loss, system_parameters) + +input = rand(dim, n_samples) +output = rand(dim, n_samples) +# one parameter set per sample, which is the shape `ParametricDataLoader` hands to the optimizer +batch_parameters = fill(system_parameters, n_samples) + +loss_value, gradient = pullback(params(nn), nn.model, (input, output, batch_parameters)) + +@test loss_value ≈ loss(nn.model, params(nn), input, output, batch_parameters) + +symbolic_gradient = gradient(1.0) + +function summed_per_sample_gradient() + total = nothing + for i in 1:n_samples + single = Zygote.gradient( + ps -> loss(nn.model, ps, input[:, i:i], output[:, i:i], [system_parameters]), + params(nn))[1] + block = single.L1.params + total = isnothing(total) ? block : map((a, b) -> map(+, a, b), total, block) + end + total +end + +reference_gradient = summed_per_sample_gradient() + +@test keys(symbolic_gradient) == (:L1,) +for layer in keys(reference_gradient) + for parameter in keys(reference_gradient[layer]) + @test symbolic_gradient.L1[layer][parameter] ≈ reference_gradient[layer][parameter] + end +end + +# A gradient of all zeros would pass the loop above if the reference were zero too, so check that +# the network actually depends on its parameters here. +@test any(any(abs.(block) .> 1e-8) for layer in keys(reference_gradient) + for block in values(reference_gradient[layer])) diff --git a/test/generalized_hamiltonian_neural_networks/pghnn_training_test.jl b/test/generalized_hamiltonian_neural_networks/pghnn_training_test.jl new file mode 100644 index 000000000..e1b4c4200 --- /dev/null +++ b/test/generalized_hamiltonian_neural_networks/pghnn_training_test.jl @@ -0,0 +1,47 @@ +# One epoch of training a `GeneralizedHamiltonianArchitecture` on a `ParametricDataLoader`, which is +# the path that ties the pieces together: the batch splitter, the parametric loss, the `Zygote` +# pullback through the symbolic gradient of the energies, and the optimizer step over a *nested* +# parameter set. + +using GeometricMachineLearning +using AbstractNeuralNetworks: params +using GeometricProblems.CoupledHarmonicOscillator: hodeensemble, default_parameters +using GeometricIntegrators: ImplicitMidpoint, integrate +using Random: seed! +using Test + +seed!(1234) + +function shift_parameters(params::NamedTuple, a::Number) + NamedTuple{keys(params)}(Tuple(value .+ a for value in values(params))) +end + +all_parameters = [default_parameters(), shift_parameters(default_parameters(), 0.5)] + +sol = integrate(hodeensemble(; parameters = all_parameters), ImplicitMidpoint()) +dl = ParametricDataLoader(sol) + +arch = GeneralizedHamiltonianArchitecture(dl.input_dim; parameters = default_parameters()) +nn = NeuralNetwork(arch) +parameters_before = deepcopy(params(nn)) + +n_epochs = 2 +loss_array = Optimizer(AdamOptimizer(), nn)(nn, dl, Batch(200), n_epochs; show_progress = false) + +@test length(loss_array) == n_epochs +@test all(isfinite, loss_array) +@test all(>(0), loss_array) + +# The optimizer has to reach every block of the *nested* parameter set: the architecture is a chain +# of `SymplecticEuler` layers, each of which holds the parameters of a whole sub-network. +function every_block_moved(before, after) + all(keys(before)) do layer + all(keys(before[layer])) do sublayer + all(keys(before[layer][sublayer])) do parameter + before[layer][sublayer][parameter] != after[layer][sublayer][parameter] + end + end + end +end + +@test every_block_moved(parameters_before, params(nn)) diff --git a/test/generalized_hamiltonian_neural_networks_test.jl b/test/generalized_hamiltonian_neural_networks_test.jl new file mode 100644 index 000000000..824ded805 --- /dev/null +++ b/test/generalized_hamiltonian_neural_networks_test.jl @@ -0,0 +1,25 @@ +using GeometricMachineLearning +using GeometricMachineLearning: OptionalParameters, OneInitializer +using GeometricProblems.HarmonicOscillator: odeproblem, default_parameters +using GeometricIntegrators +using Test + +sol = integrate(odeproblem(), ImplicitMidpoint()) +dim = length(sol.problem.ics.q) + +dl = DataLoader(sol) + +function test_ghnn_without_parameters(dim::Integer = dim) + arch = GeneralizedHamiltonianArchitecture(dim) + nn = NeuralNetwork(arch; initializer=OneInitializer()) + @test nn([1., 1.]) ≈ [1.003217200759985, 0.9968055760434815] +end + +function test_ghnn_with_parameters(dim::Integer = dim, parameters::OptionalParameters = default_parameters()) + arch = GeneralizedHamiltonianArchitecture(dim, parameters = parameters) + nn = NeuralNetwork(arch; initializer=OneInitializer()) + @test nn([1., 1.], parameters) ≈ [1.0000350420844089, 0.9999649603746436] +end + +test_ghnn_without_parameters() +test_ghnn_with_parameters() \ No newline at end of file diff --git a/test/runtests.jl b/test/runtests.jl index d8758cd8d..b697aea21 100644 --- a/test/runtests.jl +++ b/test/runtests.jl @@ -37,6 +37,15 @@ end @safetestset "Hamiltonian Neural Network " begin include("hamiltonian_neural_network_tests.jl") end +@safetestset "Generalized Hamiltonian Neural Network " begin + include("generalized_hamiltonian_neural_networks_test.jl") +end +@safetestset "Symbolic pullback for a single-layer PGHNN " begin + include("generalized_hamiltonian_neural_networks/pghnn_symbolic_pullback_single_layer_test.jl") +end +@safetestset "PGHNN training on a ParametricDataLoader " begin + include("generalized_hamiltonian_neural_networks/pghnn_training_test.jl") +end @safetestset "Manifold Neural Network Layers " begin include("layers/manifold_layers.jl") end @@ -128,6 +137,9 @@ end @safetestset "Test data loader for a tensor (q and p data) " begin include("data_loader/draw_batch_for_tensor_test.jl") end +@safetestset "Parametric DataLoader " begin + include("data_loader/parametric_data_loader_test.jl") +end @info "Starting network-loss and kernel tests" @safetestset "Test NetworkLoss + Optimizer " begin diff --git a/test/train!/test_method.jl b/test/train!/test_method.jl index fbce31e80..4df65d03e 100644 --- a/test/train!/test_method.jl +++ b/test/train!/test_method.jl @@ -24,7 +24,7 @@ exacthnn = ExactHnn() sympeuler = SEuler() -@test GeometricMachineLearning.type(sympeuler) == SymplecticEulerA +@test GeometricMachineLearning.type(sympeuler) == SymplecticEulerIntegratorA @test symbols(sympeuler) == PhaseSpaceSymbol @test shape(sympeuler) == TrajectoryData @test min_length_batch(sympeuler) == 2 @@ -64,7 +64,7 @@ midpointlnn = VariaMidPoint() ######################################### @testerror GeometricMachineLearning.type(default_Method(sympnet, tra_pos_data)) -@test GeometricMachineLearning.type(default_method(hnn, tra_ps_data)) == SymplecticEulerA +@test GeometricMachineLearning.type(default_method(hnn, tra_ps_data)) == SymplecticEulerIntegratorA @test GeometricMachineLearning.type(default_method(hnn, sam_dps_data)) == HnnExactMethod @test GeometricMachineLearning.type(default_method(sympnet, tra_ps_data)) == BasicSympNetMethod @test GeometricMachineLearning.type(default_method(lnn, tra_pos_data)) == VariationalMidPointMethod diff --git a/test/training_phnn.jl b/test/training_phnn.jl index 4f30fda1c..0a65aa7b0 100644 --- a/test/training_phnn.jl +++ b/test/training_phnn.jl @@ -5,7 +5,7 @@ using GeometricIntegrators: ImplicitMidpoint, integrate using Random: seed! seed!(123) -function make_alternative_parameters_by_adding_constant(params::NamedTuple=default_parameters, n::Integer=1, a::Number=1.) +function make_alternative_parameters_by_adding_constant(params::NamedTuple=default_parameters(), n::Integer=1, a::Number=1.) _keys = keys(params) values = () for (key, i) in zip(_keys, 1:length(_keys)) @@ -18,13 +18,13 @@ function make_alternative_parameters_by_adding_constant(params::NamedTuple, n::I [make_alternative_parameters_by_adding_constant(params,n, a) for a ∈ a_vals] end -alternative_parameters = make_alternative_parameters_by_adding_constant(default_parameters, 1, Vector(.1:.1:1.)) +alternative_parameters = make_alternative_parameters_by_adding_constant(default_parameters(), 1, Vector(.1:.1:1.)) h_ensemble = hodeensemble(; parameters = alternative_parameters) sol = integrate(h_ensemble, ImplicitMidpoint()) dl = ParametricDataLoader(sol) batch = Batch(100) -arch = GeneralizedHamiltonianArchitecture(4; parameters = default_parameters) +arch = GeneralizedHamiltonianArchitecture(4; parameters = default_parameters()) nn = NeuralNetwork(arch) o = Optimizer(AdamOptimizer(), nn) -o(nn, dl, batch) \ No newline at end of file +o(nn, dl, batch) From 98528bfe23602baf9660886a5eeea14887656722 Mon Sep 17 00:00:00 2001 From: Michael Kraus Date: Wed, 19 Aug 2026 23:18:07 +0900 Subject: [PATCH 2/4] Cover the parametric and forced pieces, and fix what that uncovered codecov put the patch at 55%, with six of the new files at 0%: nothing built a forcing layer, a wide or parametric ResNet layer, a `ParametricResNet`, a `ForcedSympNet` or a `ForcedGeneralizedHamiltonianArchitecture`. Writing a construct-and-evaluate test for each turned up two defects. `ForcedGeneralizedHamiltonianArchitecture` could not be evaluated at all. `(nn::NeuralNetwork{GT})(qp, problem_params)` and the `Optimizer` entry point were both written for `GT <: GeneralizedHamiltonianArchitecture`, and the forced architecture is a *sibling* of that under `HamiltonianArchitecture`, not a subtype. So `nn(x, mu)` fell through to AbstractNeuralNetworks' generic functor, which read the system parameters as the network parameters and reached the first layer as a `Float64`. Both methods are now defined for it too, in its own file; widening to `HamiltonianArchitecture` would be wrong, since `StandardHamiltonianArchitecture` takes no system parameters. `ParametricResNet(::DataLoader, n_blocks, width; parameters = ...)` accepted `parameters` and did not forward it, so that constructor always built a network with no parameter dependence. The test also pins the `ForcingLayer` convention, which is the opposite of what the names suggest: `Q`/`P`/`QP` say what the forcing *depends on*, not what it changes. All three add to `p` and leave `q` alone, which is what a force does to the `p` equation. It perturbs one coordinate at a time and checks the output only moves for a coordinate the layer is named after. Co-Authored-By: Claude Opus 5 (1M context) --- CHANGELOG.md | 7 + ..._generalized_hamiltonian_neural_network.jl | 17 ++- src/architectures/parametric_resnet.jl | 2 +- ...arametric_layers_and_architectures_test.jl | 129 ++++++++++++++++++ test/runtests.jl | 3 + 5 files changed, 156 insertions(+), 2 deletions(-) create mode 100644 test/generalized_hamiltonian_neural_networks/parametric_layers_and_architectures_test.jl diff --git a/CHANGELOG.md b/CHANGELOG.md index f4cf98c39..1d6b9b18c 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -110,6 +110,13 @@ breaking release). - `concatenate_array_with_parameters(::AbstractMatrix, ::AbstractVector)` concatenated a batch with `vcat` rather than `hcat`, collapsing it into a single long vector ([#207](https://github.com/JuliaGNI/GeometricMachineLearning.jl/pull/207)). +- **`ForcedGeneralizedHamiltonianArchitecture` could not be evaluated at all.** The + parameter-dependent `NeuralNetwork` functor and the `Optimizer` entry point were defined for + `GeneralizedHamiltonianArchitecture` only, and the two are siblings under `HamiltonianArchitecture` + rather than sub- and supertype, so `nn(x, μ)` fell through to the generic functor and read the + *system* parameters as the *network* parameters. +- `ParametricResNet(::DataLoader, n_blocks, width; parameters = …)` accepted `parameters` and then + dropped it, silently building a network with no parameter dependence. ### Changed diff --git a/src/architectures/forced_generalized_hamiltonian_neural_network.jl b/src/architectures/forced_generalized_hamiltonian_neural_network.jl index 2bb97632e..9b417c7d6 100644 --- a/src/architectures/forced_generalized_hamiltonian_neural_network.jl +++ b/src/architectures/forced_generalized_hamiltonian_neural_network.jl @@ -32,4 +32,19 @@ function Chain(arch::ForcedGeneralizedHamiltonianArchitecture{FT}) where {FT} layers = (layers..., SymplecticEulerB(potential_energy; return_parameters = _return_parameters)) end Chain(layers...) -end \ No newline at end of file +end + +# `ForcedGeneralizedHamiltonianArchitecture` and `GeneralizedHamiltonianArchitecture` are siblings +# under `HamiltonianArchitecture`, so the parameter-dependent forward pass and training entry point +# defined for the latter do not cover this one. `HamiltonianArchitecture` itself is too wide: +# `StandardHamiltonianArchitecture` takes no system parameters. +function (nn::NeuralNetwork{<:ForcedGeneralizedHamiltonianArchitecture})(qp::QPTOAT2, + problem_params::OptionalParameters) + nn.model(qp, problem_params, params(nn)) +end + +function (o::Optimizer)(nn::NeuralNetwork{<:ForcedGeneralizedHamiltonianArchitecture}, + dl::ParametricDataLoader, batch::Batch{:FeedForward}, n_epochs::Integer = 1, + loss::NetworkLoss = ParametricLoss(); kwargs...) + o(nn, dl, batch, n_epochs, loss, ZygotePullback(loss); kwargs...) +end diff --git a/src/architectures/parametric_resnet.jl b/src/architectures/parametric_resnet.jl index 5652ef4cc..2a4ff022d 100644 --- a/src/architectures/parametric_resnet.jl +++ b/src/architectures/parametric_resnet.jl @@ -12,7 +12,7 @@ struct ParametricResNet{AT <: Activation, PT <: OptionalParameters} <: NeuralNet end function ParametricResNet(dl::DataLoader, n_blocks::Integer, width::Integer=dl.input_dim; activation=HNN_activation_default, parameters=NullParameters()) - ParametricResNet(dl.input_dim; width=width, n_blocks=n_blocks, activation) + ParametricResNet(dl.input_dim; width=width, n_blocks=n_blocks, activation=activation, parameters=parameters) end function ResNet(input_dim::Integer, n_blocks::Integer, width::Integer=input_dim; activation=HNN_activation_default, parameters=NullParameters()) diff --git a/test/generalized_hamiltonian_neural_networks/parametric_layers_and_architectures_test.jl b/test/generalized_hamiltonian_neural_networks/parametric_layers_and_architectures_test.jl new file mode 100644 index 000000000..eaa669edd --- /dev/null +++ b/test/generalized_hamiltonian_neural_networks/parametric_layers_and_architectures_test.jl @@ -0,0 +1,129 @@ +# Construction and a forward pass for each parameter-dependent and forced piece. The other PGHNN +# tests build a plain `GeneralizedHamiltonianArchitecture`, so none of these were evaluated +# anywhere -- which is how `ForcedGeneralizedHamiltonianArchitecture`, exported and documented, +# came to have a forward pass that threw. + +using GeometricMachineLearning +using GeometricMachineLearning: ForcingLayerQ, ForcingLayerP, ForcingLayerQP, + ParametricResNetLayer, WideResNetLayer, ParametricResNet +using AbstractNeuralNetworks: params +using Random: seed! +using Test + +seed!(1234) + +const DIM = 4 +const HALF = DIM ÷ 2 +const WIDTH = 8 +const SYSTEM_PARAMETERS = (m = 1.0, ω = π / 2) + +finite(x::AbstractArray) = all(isfinite, x) +finite(qp::NamedTuple) = finite(qp.q) && finite(qp.p) + +# The `Q`/`P`/`QP` suffix names what the forcing *depends on*, not what it changes: a force enters +# the `ṗ` equation, so all three add to `p` and leave `q` alone. +@testset "ForcingLayer$name" for (name, Layer, depends_on) in ( + ("Q", ForcingLayerQ, (:q,)), ("P", ForcingLayerP, (:p,)), ("QP", ForcingLayerQP, (:q, :p))) + layer = Layer(DIM; parameters = SYSTEM_PARAMETERS) + nn = NeuralNetwork(layer) + @test parameterlength(nn) > 0 + + z = (q = rand(HALF), p = rand(HALF)) + out = layer(z, SYSTEM_PARAMETERS, params(nn)) + @test keys(out) == (:q, :p) + @test size(out.q) == size(z.q) && size(out.p) == size(z.p) + @test finite(out) + @test out.q == z.q + @test out.p != z.p + + # perturb one coordinate at a time: the forcing may only move when a coordinate it is named + # after does + for coordinate in (:q, :p) + perturbed = merge(z, NamedTuple{(coordinate,)}((z[coordinate] .+ 1.0,))) + moved = layer(perturbed, SYSTEM_PARAMETERS, params(nn)).p .- perturbed.p != + out.p .- z.p + @test moved == (coordinate in depends_on) + end + + # the same layer applied to the concatenated array form + array_out = layer(vcat(z.q, z.p), SYSTEM_PARAMETERS, params(nn)) + @test array_out ≈ vcat(out.q, out.p) +end + +@testset "WideResNetLayer" begin + layer = WideResNetLayer(DIM, WIDTH, tanh) + nn = NeuralNetwork(Chain(layer)) + ps = params(nn).L1 + @test parameterlength(layer) == WIDTH * (DIM + 1) + DIM * (WIDTH + 1) + + for input in (rand(DIM), rand(DIM, 3), rand(DIM, 3, 2)) + out = layer(input, ps) + @test size(out) == size(input) + @test finite(out) + end + + z = (q = rand(HALF), p = rand(HALF)) + out = layer(z, ps) + @test keys(out) == (:q, :p) + @test out ≈ (q = layer(vcat(z.q, z.p), ps)[1:HALF], p = layer(vcat(z.q, z.p), ps)[(HALF + 1):DIM]) +end + +@testset "ParametricResNetLayer" begin + layer = ParametricResNetLayer(DIM, WIDTH, tanh; + parameters = SYSTEM_PARAMETERS, return_parameters = false) + nn = NeuralNetwork(Chain(layer)) + ps = params(nn).L1 + + # one `NamedTuple` of system parameters describes one sample, so a matrix input is a single + # column; a batch is a vector of parameter sets, which is what `ParametricResNet` builds + for input in (rand(DIM), rand(DIM, 1)) + out = layer(input, SYSTEM_PARAMETERS, ps) + @test size(out) == size(input) + @test finite(out) + end + @test_throws AssertionError layer(rand(DIM, 3), SYSTEM_PARAMETERS, ps) + + z = (q = rand(HALF), p = rand(HALF)) + @test finite(layer(z, SYSTEM_PARAMETERS, ps)) +end + +@testset "ResNet with a width of its own" begin + # `sys_dim == width` keeps the plain `ResNetLayer`; a different width switches to + # `WideResNetLayer`, which is the path `ParametricResNet` compares against + narrow = NeuralNetwork(ResNet(DIM, 2, DIM)) + wide = NeuralNetwork(ResNet(DIM, 2, WIDTH)) + @test parameterlength(wide) > parameterlength(narrow) + @test size(wide(rand(DIM))) == (DIM,) + @test finite(wide(rand(DIM))) +end + +@testset "ParametricResNet" begin + arch = ParametricResNet(DIM; width = WIDTH, n_blocks = 2, parameters = SYSTEM_PARAMETERS) + nn = NeuralNetwork(arch) + out = nn.model(rand(DIM), SYSTEM_PARAMETERS, params(nn)) + @test size(out) == (DIM,) + @test finite(out) + + # the `DataLoader` constructor used to accept `parameters` and drop it + dl = DataLoader(rand(DIM, 20); suppress_info = true) + @test ParametricResNet(dl, 2, WIDTH; parameters = SYSTEM_PARAMETERS).parameters == + SYSTEM_PARAMETERS +end + +@testset "ForcedSympNet $forcing_type" for forcing_type in (:Q, :P, :QP) + nn = NeuralNetwork(ForcedSympNet(DIM; forcing_type = forcing_type)) + out = nn(rand(DIM)) + @test size(out) == (DIM,) + @test finite(out) +end + +@testset "ForcedGeneralizedHamiltonianArchitecture $forcing_type" for forcing_type in (:Q, :P, :QP) + arch = ForcedGeneralizedHamiltonianArchitecture(DIM; parameters = SYSTEM_PARAMETERS, + forcing_type = forcing_type) + nn = NeuralNetwork(arch) + out = nn(rand(DIM), SYSTEM_PARAMETERS) + @test size(out) == (DIM,) + @test finite(out) +end + +@test_throws ErrorException ForcedGeneralizedHamiltonianArchitecture(DIM; forcing_type = :X) diff --git a/test/runtests.jl b/test/runtests.jl index b697aea21..9539d9de8 100644 --- a/test/runtests.jl +++ b/test/runtests.jl @@ -46,6 +46,9 @@ end @safetestset "PGHNN training on a ParametricDataLoader " begin include("generalized_hamiltonian_neural_networks/pghnn_training_test.jl") end +@safetestset "Parametric and forced layers and architectures " begin + include("generalized_hamiltonian_neural_networks/parametric_layers_and_architectures_test.jl") +end @safetestset "Manifold Neural Network Layers " begin include("layers/manifold_layers.jl") end From cc932436b1a9428517c3112fe1d490c8a31fde82 Mon Sep 17 00:00:00 2001 From: Michael Kraus Date: Thu, 20 Aug 2026 10:36:02 +0900 Subject: [PATCH 3/4] Address @benedict-96's review Dead code he asked about, all three confirmed obsolete rather than merely unused. `rearrange_parameters` was never finished -- `zeros(_eltype(parameters), )` has no dimensions -- and never called. The commented `batch_over_two_axes` for `ParametricDataLoader` is superseded by the `n_batches::Integer` method in `batch.jl`, which the live code calls. The commented three-dimensional `ParametricResNetLayer` methods were also wrong: both compute `input` and then use `x`, so they were not a starting point for anything. The hardcoded `D:\RESEARCH - UTWENTE\...` output path is now `joinpath(get(ENV, "GML_OUTPUT_DIR", @__DIR__), ...)`. The three scripts each carried a verbatim copy of the same dataset helpers -- `turn_q_p_data_into_correct_format` was byte-identical in all three -- and the two forced-oscillator scripts also shared their analytic solution. All of it moves to `scripts/parametric_data_helpers.jl`, about 200 lines less duplication. The damped oscillator keeps its own analytic solution: it looks similar but is a different system, not a duplicate. That file is a staging post; GMLDatasets #5 tracks moving the helpers where they belong. Docstrings for `WideResNetLayer`, `ParametricResNetLayer` and `ParametricResNet`, which had none at all, and they are in the manual now. They passed `missing_docs` only because Documenter checks bindings that already carry a docstring. SymbolicPullback performance ---------------------------- Measured rather than assumed, and it is the opposite way round from what the report suggested. Per call the symbolic pullback is 100-1000x *faster* than Zygote and flat at 0.1 ms, where Zygote goes 10 -> 24 -> 34 -> 140 ms as the network grows. What explodes is the build, multiplicatively in `n_integrators`. At `dim = 4, width = 4, nhidden = 1`, phase by phase: n_integrators = 1: loss 3.4e5 chars, its derivative 1.5e8 chars, build 1.4 s n_integrators = 2: loss 1.4e9 chars, never returns, past 8 GB Each `SymplecticEuler` layer's forward pass calls the executable gradient `build_gradient` produced for its energy network; traced on symbolic `Num`s that inlines the whole gradient expression, so stacking integrators inlines it inside itself. `cse = true` cannot help -- it runs at code generation, long after the expression exists. The real fix is to stop tracing the inlined chain and compose the pullback layer by layer from the gradients each layer has already built, which is upstream-shaped work; issue #245 has the numbers and the argument. Here the build refuses `n_integrators > 1` with an explanation instead of appearing to hang. Not addressed, deliberately: narrative documentation of the forced and parametric architectures, which is waiting on the tikz figures he offered. Co-Authored-By: Claude Opus 5 (1M context) --- CHANGELOG.md | 6 + .../hamiltonian_neural_network.md | 4 + ...rcedGeneralizedHamiltonianNeuralNetwork.jl | 88 +------- ...ndentHarmonicOscillatorParametricResnet.jl | 88 +------- scripts/Train_DampedOscillator_QP.jl | 193 +++++++----------- scripts/parametric_data_helpers.jl | 95 +++++++++ ..._generalized_hamiltonian_neural_network.jl | 2 + src/architectures/parametric_resnet.jl | 21 ++ src/data_loader/parametric_data_loader.jl | 21 -- src/layers/parametric_resnet_layer.jl | 35 +++- src/layers/wide_resnet.jl | 18 ++ src/pullbacks/symbolic_hnn_pullback.jl | 38 ++++ ...hnn_symbolic_pullback_single_layer_test.jl | 12 ++ 13 files changed, 306 insertions(+), 315 deletions(-) create mode 100644 scripts/parametric_data_helpers.jl diff --git a/CHANGELOG.md b/CHANGELOG.md index 1d6b9b18c..23fc3fa5f 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -117,6 +117,12 @@ breaking release). *system* parameters as the *network* parameters. - `ParametricResNet(::DataLoader, n_blocks, width; parameters = …)` accepted `parameters` and then dropped it, silently building a network with no parameter dependence. +- `SymbolicPullback(nn, ::ParametricLoss, μ)` now throws for `n_integrators > 1` instead of appearing + to hang. The symbolic expression grows *multiplicatively* with the number of integrators — measured + at `dim = 4, width = 4, nhidden = 1`, the loss is 3.4 ⋅ 10⁵ characters at one integrator and + 1.4 ⋅ 10⁹ at two, and the build never returns. One integrator builds in ≈1.4 s, and the result + evaluates about 100× faster than the `Zygote` pullback. See + [#245](https://github.com/JuliaGNI/GeometricMachineLearning.jl/issues/245). ### Changed diff --git a/docs/src/architectures/hamiltonian_neural_network.md b/docs/src/architectures/hamiltonian_neural_network.md index 8c6234419..82a1086ec 100644 --- a/docs/src/architectures/hamiltonian_neural_network.md +++ b/docs/src/architectures/hamiltonian_neural_network.md @@ -61,7 +61,11 @@ GeometricMachineLearning.ForcingLayerQ GeometricMachineLearning.ForcingLayerP GeometricMachineLearning.ForcingLayerQP GeometricMachineLearning.ParametricDataLoader +GeometricMachineLearning.ParametricResNet +GeometricMachineLearning.ParametricResNetLayer +GeometricMachineLearning.WideResNetLayer GeometricMachineLearning.SymbolicPullback(::GeometricMachineLearning.NeuralNetwork, ::GeometricMachineLearning.ParametricLoss, ::GeometricMachineLearning.GeometricBase.OptionalParameters) +GeometricMachineLearning._check_symbolic_pullback_is_tractable GeometricMachineLearning._flatten_system_parameters GeometricMachineLearning._unwrap_gradient GeometricMachineLearning._processing diff --git a/scripts/TimeDependentHarmonicOscillatorForcedGeneralizedHamiltonianNeuralNetwork.jl b/scripts/TimeDependentHarmonicOscillatorForcedGeneralizedHamiltonianNeuralNetwork.jl index e9ca22ce1..ec352faaf 100644 --- a/scripts/TimeDependentHarmonicOscillatorForcedGeneralizedHamiltonianNeuralNetwork.jl +++ b/scripts/TimeDependentHarmonicOscillatorForcedGeneralizedHamiltonianNeuralNetwork.jl @@ -4,6 +4,8 @@ using GeometricMachineLearning: QPT, QPT2, Activation, ParametricLoss, SymbolicN using CairoMakie using NNlib: relu +include(joinpath(@__DIR__, "parametric_data_helpers.jl")) + # PARAMETERS omega = 1.0 # natural frequency of the harmonic Oscillator Omega = 3.5 # frequency of the external sinusoidal forcing @@ -18,92 +20,8 @@ IC = vec( [(q=q0, p=p0) for q0 in range(-1, 1, ni_dim), p0 in range(-1, 1, ni_di # Generating the solution array ni = ni_dim^2 -q = zeros(Float64, ni, nt+1) -p = zeros(Float64, ni, nt+1) t = collect(dt * range(0, nt, step=1)) - -""" -Turn a vector of numbers into a vector of `NamedTuple`s to be used by `ParametricDataLoader`. -""" -function turn_parameters_into_correct_format(t::AbstractVector, IC::AbstractVector{<:NamedTuple}) - vec_of_params = NamedTuple[] - for time_step ∈ t - time_step == t[end] || push!(vec_of_params, (t = time_step, )) - end - vcat((vec_of_params for _ in axes(IC, 1))...) -end - -for i in 1:nt+1 - for j=1:ni - q[j,i] = ( IC[j].p - Omega*F/(omega^2-Omega^2) )/ omega *sin(omega*t[i]) + IC[j].q*cos(omega*t[i]) + F/(omega^2-Omega^2)*sin(Omega*t[i]) - p[j,i] = -omega^2*IC[j].q*sin(omega*t[i]) + ( IC[j].p - Omega*F/(omega^2-Omega^2) )*cos(omega*t[i]) + Omega*F/(omega^2-Omega^2)*cos(Omega*t[i]) - # q[j,i] = ( IC[j].p - Omega*F/(omega^2-Omega^2) )/ omega *exp(-omega*t[i]) - IC[j].q*exp(-omega*t[i]) + F/(omega^2-Omega^2)*exp(-Omega*t[i]) - # p[j,i] = -omega^2*IC[j].q*exp(-omega*t[i]) + ( IC[j].p + Omega*F/(omega^2-Omega^2) )*exp(-omega*t[i]) - Omega*F/(omega^2-Omega^2)*exp(-Omega*t[i]) - end - -end - -@doc raw""" -Turn a `NamedTuple` of ``(q,p)`` data into two tensors of the correct format. - -This is the tricky part as the structure of the input array(s) needs to conform with the structure of the parameters. - -Here the data are rearranged in an array of size ``(n, 2, t_f - 1)`` where ``[t_0, t_1, \ldots, t_f]`` is the vector storing the time steps. - -If we deal with different initial conditions as well, we still put everything into the third (parameter) axis. - -# Example - -```jldoctest -using GeometricMachineLearning - -q = [1. 2. 3.; 4. 5. 6.] -p = [1.5 2.5 3.5; 4.5 5.5 6.5] -qp = (q = q, p = p) -turn_q_p_data_into_correct_format(qp) - -# output - -(q = [1.0 2.0; 4.0 5.0;;; 2.0 3.0; 5.0 6.0], p = [1.5 2.5; 4.5 5.5;;; 2.5 3.5; 5.5 6.5]) -``` -""" -function turn_q_p_data_into_correct_format(qp::QPT2{T, 2}) where {T} - number_of_time_steps = size(qp.q, 2) - 1 # not counting t₀ - number_of_initial_conditions = size(qp.q, 1) - q_array = zeros(T, 1, 2, number_of_time_steps * number_of_initial_conditions) - p_array = zeros(T, 1, 2, number_of_time_steps * number_of_initial_conditions) - for initial_condition_index ∈ 0:(number_of_initial_conditions - 1) - for time_index ∈ 1:number_of_time_steps - q_array[:, 1, initial_condition_index * number_of_time_steps + time_index] .= qp.q[initial_condition_index + 1, time_index] - q_array[:, 2, initial_condition_index * number_of_time_steps + time_index] .= qp.q[initial_condition_index + 1, time_index + 1] - p_array[:, 1, initial_condition_index * number_of_time_steps + time_index] .= qp.p[initial_condition_index + 1, time_index] - p_array[:, 2, initial_condition_index * number_of_time_steps + time_index] .= qp.p[initial_condition_index + 1, time_index + 1] - end - end - (q = q_array, p = p_array) -end - -# SAVING TO FILE - -# h5 = h5open(path, "w") -# write(h5, "q", q) -# write(h5, "p", p) -# write(h5, "t", t) -# -# attrs(h5)["ni"] = ni -# attrs(h5)["nt"] = nt -# attrs(h5)["dt"] = dt -# -# close(h5) - -""" -This takes time as a single additional parameter (third axis). -""" -function load_time_dependent_harmonic_oscillator_with_parametric_data_loader(qp::QPT{T}, t::AbstractVector{T}, IC::AbstractVector) where {T} - qp_reformatted = turn_q_p_data_into_correct_format(qp) - t_reformatted = turn_parameters_into_correct_format(t, IC) - ParametricDataLoader(qp_reformatted, t_reformatted) -end +q, p = forced_harmonic_oscillator_solution(t, IC; omega = omega, Omega = Omega, F = F) # This sets up the data loader dl = load_time_dependent_harmonic_oscillator_with_parametric_data_loader((q = q, p = p), t, IC) diff --git a/scripts/TimeDependentHarmonicOscillatorParametricResnet.jl b/scripts/TimeDependentHarmonicOscillatorParametricResnet.jl index 5b6e0f31f..2e068693b 100644 --- a/scripts/TimeDependentHarmonicOscillatorParametricResnet.jl +++ b/scripts/TimeDependentHarmonicOscillatorParametricResnet.jl @@ -4,6 +4,8 @@ using GeometricMachineLearning: QPT, QPT2, Activation, ParametricLoss, SymbolicN using CairoMakie using NNlib: relu +include(joinpath(@__DIR__, "parametric_data_helpers.jl")) + # PARAMETERS omega = 1.0 # natural frequency of the harmonic Oscillator Omega = 3.5 # frequency of the external sinusoidal forcing @@ -18,92 +20,8 @@ IC = vec( [(q=q0, p=p0) for q0 in range(-1, 1, ni_dim), p0 in range(-1, 1, ni_di # Generating the solution array ni = ni_dim^2 -q = zeros(Float64, ni, nt+1) -p = zeros(Float64, ni, nt+1) t = collect(dt * range(0, nt, step=1)) - -""" -Turn a vector of numbers into a vector of `NamedTuple`s to be used by `ParametricDataLoader`. -""" -function turn_parameters_into_correct_format(t::AbstractVector, IC::AbstractVector{<:NamedTuple}) - vec_of_params = NamedTuple[] - for time_step ∈ t - time_step == t[end] || push!(vec_of_params, (t = time_step, )) - end - vcat((vec_of_params for _ in axes(IC, 1))...) -end - -for i in 1:nt+1 - for j=1:ni - q[j,i] = ( IC[j].p - Omega*F/(omega^2-Omega^2) )/ omega *sin(omega*t[i]) + IC[j].q*cos(omega*t[i]) + F/(omega^2-Omega^2)*sin(Omega*t[i]) - p[j,i] = -omega^2*IC[j].q*sin(omega*t[i]) + ( IC[j].p - Omega*F/(omega^2-Omega^2) )*cos(omega*t[i]) + Omega*F/(omega^2-Omega^2)*cos(Omega*t[i]) - # q[j,i] = ( IC[j].p - Omega*F/(omega^2-Omega^2) )/ omega *exp(-omega*t[i]) - IC[j].q*exp(-omega*t[i]) + F/(omega^2-Omega^2)*exp(-Omega*t[i]) - # p[j,i] = -omega^2*IC[j].q*exp(-omega*t[i]) + ( IC[j].p + Omega*F/(omega^2-Omega^2) )*exp(-omega*t[i]) - Omega*F/(omega^2-Omega^2)*exp(-Omega*t[i]) - end - -end - -@doc raw""" -Turn a `NamedTuple` of ``(q,p)`` data into two tensors of the correct format. - -This is the tricky part as the structure of the input array(s) needs to conform with the structure of the parameters. - -Here the data are rearranged in an array of size ``(n, 2, t_f - 1)`` where ``[t_0, t_1, \ldots, t_f]`` is the vector storing the time steps. - -If we deal with different initial conditions as well, we still put everything into the third (parameter) axis. - -# Example - -```jldoctest -using GeometricMachineLearning - -q = [1. 2. 3.; 4. 5. 6.] -p = [1.5 2.5 3.5; 4.5 5.5 6.5] -qp = (q = q, p = p) -turn_q_p_data_into_correct_format(qp) - -# output - -(q = [1.0 2.0; 4.0 5.0;;; 2.0 3.0; 5.0 6.0], p = [1.5 2.5; 4.5 5.5;;; 2.5 3.5; 5.5 6.5]) -``` -""" -function turn_q_p_data_into_correct_format(qp::QPT2{T, 2}) where {T} - number_of_time_steps = size(qp.q, 2) - 1 # not counting t₀ - number_of_initial_conditions = size(qp.q, 1) - q_array = zeros(T, 1, 2, number_of_time_steps * number_of_initial_conditions) - p_array = zeros(T, 1, 2, number_of_time_steps * number_of_initial_conditions) - for initial_condition_index ∈ 0:(number_of_initial_conditions - 1) - for time_index ∈ 1:number_of_time_steps - q_array[:, 1, initial_condition_index * number_of_time_steps + time_index] .= qp.q[initial_condition_index + 1, time_index] - q_array[:, 2, initial_condition_index * number_of_time_steps + time_index] .= qp.q[initial_condition_index + 1, time_index + 1] - p_array[:, 1, initial_condition_index * number_of_time_steps + time_index] .= qp.p[initial_condition_index + 1, time_index] - p_array[:, 2, initial_condition_index * number_of_time_steps + time_index] .= qp.p[initial_condition_index + 1, time_index + 1] - end - end - (q = q_array, p = p_array) -end - -# SAVING TO FILE - -# h5 = h5open(path, "w") -# write(h5, "q", q) -# write(h5, "p", p) -# write(h5, "t", t) -# -# attrs(h5)["ni"] = ni -# attrs(h5)["nt"] = nt -# attrs(h5)["dt"] = dt -# -# close(h5) - -""" -This takes time as a single additional parameter (third axis). -""" -function load_time_dependent_harmonic_oscillator_with_parametric_data_loader(qp::QPT{T}, t::AbstractVector{T}, IC::AbstractVector) where {T} - qp_reformatted = turn_q_p_data_into_correct_format(qp) - t_reformatted = turn_parameters_into_correct_format(t, IC) - ParametricDataLoader(qp_reformatted, t_reformatted) -end +q, p = forced_harmonic_oscillator_solution(t, IC; omega = omega, Omega = Omega, F = F) # This sets up the data loader dl = load_time_dependent_harmonic_oscillator_with_parametric_data_loader((q = q, p = p), t, IC) diff --git a/scripts/Train_DampedOscillator_QP.jl b/scripts/Train_DampedOscillator_QP.jl index 1058ef941..b24ab61a9 100644 --- a/scripts/Train_DampedOscillator_QP.jl +++ b/scripts/Train_DampedOscillator_QP.jl @@ -1,114 +1,79 @@ -using HDF5 -using GeometricMachineLearning -using GeometricMachineLearning: QPT, QPT2 -using CairoMakie -using JLD2 -using NNlib: relu - -# PARAMETERS -nu = 0.001 # friction force coefficient -ni_dim = 2 # number of initial conditions per dimension (so ni_dim^2 total) -T = 13 -nt = 100 # number of time steps -dt = T/nt # time step -n_epochs = 100000 -n_epochs = 3 -width = 4 # width of the neural network -nhidden = 3 # number of hidden layers in the neural network -batch_size = 5000 # the size of the batch - -path_out = "D:\\RESEARCH - UTWENTE\\GFHNNs\\Damped Oscillator\\network_TEST.jld2" -#path_out = "/home/tyranowskitm/GFHNNs/DampedOscillator/OUTPUTS/network.jld2" - - -# Generating the initial condition array -IC = vec( [(q=q0, p=p0) for q0 in range(-1, 1, ni_dim), p0 in range(-1, 1, ni_dim)] ) - - -# Generating the solution array -ni = ni_dim^2 -omega = sqrt(4-nu^2) / 2 - -q = zeros(Float64, ni, nt+1) -p = zeros(Float64, ni, nt+1) -t = collect(dt*range(0,nt,step=1)) - -for i in 1:nt+1 - - for j=1:ni - q[j,i] = (1/omega)*( IC[j].p + nu/2 *IC[j].q )*exp(-nu*t[i]/2)*sin(omega*t[i]) + IC[j].q*exp(-nu*t[i]/2)*cos(omega*t[i]) - p[j,i] = -(1/omega)*( IC[j].q + nu/2 *IC[j].p )*exp(-nu*t[i]/2)*sin(omega*t[i]) + IC[j].p*exp(-nu*t[i]/2)*cos(omega*t[i]) - end - -end - - - -@doc raw""" -Turn a `NamedTuple` of ``(q,p)`` data into two tensors of the correct format. - -This is the tricky part as the structure of the input array(s) needs to conform with the structure of the parameters. - -Here the data are rearranged in an array of size ``(n, 2, t_f - 1)`` where ``[t_0, t_1, \ldots, t_f]`` is the vector storing the time steps. - -If we deal with different initial conditions as well, we still put everything into the third (parameter) axis. - -# Example - -```jldoctest -using GeometricMachineLearning - -q = [1. 2. 3.; 4. 5. 6.] -p = [1.5 2.5 3.5; 4.5 5.5 6.5] -qp = (q = q, p = p) -turn_q_p_data_into_correct_format(qp) - -# output - -(q = [1.0 2.0; 4.0 5.0;;; 2.0 3.0; 5.0 6.0], p = [1.5 2.5; 4.5 5.5;;; 2.5 3.5; 5.5 6.5]) -``` -""" -function turn_q_p_data_into_correct_format(qp::QPT2{T, 2}) where {T} - number_of_time_steps = size(qp.q, 2) - 1 # not counting t₀ - number_of_initial_conditions = size(qp.q, 1) - q_array = zeros(T, 1, 2, number_of_time_steps * number_of_initial_conditions) - p_array = zeros(T, 1, 2, number_of_time_steps * number_of_initial_conditions) - for initial_condition_index ∈ 0:(number_of_initial_conditions - 1) - for time_index ∈ 1:number_of_time_steps - q_array[:, 1, initial_condition_index * number_of_time_steps + time_index] .= qp.q[initial_condition_index + 1, time_index] - q_array[:, 2, initial_condition_index * number_of_time_steps + time_index] .= qp.q[initial_condition_index + 1, time_index + 1] - p_array[:, 1, initial_condition_index * number_of_time_steps + time_index] .= qp.p[initial_condition_index + 1, time_index] - p_array[:, 2, initial_condition_index * number_of_time_steps + time_index] .= qp.p[initial_condition_index + 1, time_index + 1] - end - end - (q = q_array, p = p_array) -end - - -# This sets up the data loader -dl = DataLoader(turn_q_p_data_into_correct_format((q = q, p = p))) - -# This sets up the neural network -arch = ForcedSympNet(2; upscaling_dimension = width, n_layers = nhidden, forcing_type = :P) -#arch = ForcedSympNet(2; upscaling_dimension = width, n_layers = nhidden, activation=(x-> max(0,x)^2/2)) -nn = NeuralNetwork(arch) - -# This is where training starts -batch = Batch(batch_size) -o = Optimizer(AdamOptimizer(), nn) - -loss_array = o(nn, dl, batch, n_epochs) - - -# Saving the parameters of the network -println("Saving the parameters of the neural network...") -flush(stdout) - -params = GeometricMachineLearning.map_to_cpu(nn.params) - -save(path_out,"parameters", params, "training loss", loss_array, "ni_dim", ni_dim, "T", T, "nt", nt, "n_epochs", n_epochs, "width", width, "nhidden", nhidden, "batch_size", batch_size, "nu", nu) - -println(" ...Done!") -flush(stdout) - - +using HDF5 +using GeometricMachineLearning +using GeometricMachineLearning: QPT, QPT2 +using CairoMakie +using JLD2 +using NNlib: relu + +include(joinpath(@__DIR__, "parametric_data_helpers.jl")) + +# PARAMETERS +nu = 0.001 # friction force coefficient +ni_dim = 2 # number of initial conditions per dimension (so ni_dim^2 total) +T = 13 +nt = 100 # number of time steps +dt = T/nt # time step +n_epochs = 100000 +n_epochs = 3 +width = 4 # width of the neural network +nhidden = 3 # number of hidden layers in the neural network +batch_size = 5000 # the size of the batch + +# next to the script, unless GML_OUTPUT_DIR says otherwise -- an absolute path from whoever ran it +# last is no use to anybody else +path_out = joinpath(get(ENV, "GML_OUTPUT_DIR", @__DIR__), "damped_oscillator_network.jld2") + + +# Generating the initial condition array +IC = vec( [(q=q0, p=p0) for q0 in range(-1, 1, ni_dim), p0 in range(-1, 1, ni_dim)] ) + + +# Generating the solution array +ni = ni_dim^2 +omega = sqrt(4-nu^2) / 2 + +q = zeros(Float64, ni, nt+1) +p = zeros(Float64, ni, nt+1) +t = collect(dt*range(0,nt,step=1)) + +for i in 1:nt+1 + + for j=1:ni + q[j,i] = (1/omega)*( IC[j].p + nu/2 *IC[j].q )*exp(-nu*t[i]/2)*sin(omega*t[i]) + IC[j].q*exp(-nu*t[i]/2)*cos(omega*t[i]) + p[j,i] = -(1/omega)*( IC[j].q + nu/2 *IC[j].p )*exp(-nu*t[i]/2)*sin(omega*t[i]) + IC[j].p*exp(-nu*t[i]/2)*cos(omega*t[i]) + end + +end + + + +end + + +# This sets up the data loader +dl = DataLoader(turn_q_p_data_into_correct_format((q = q, p = p))) + +# This sets up the neural network +arch = ForcedSympNet(2; upscaling_dimension = width, n_layers = nhidden, forcing_type = :P) +#arch = ForcedSympNet(2; upscaling_dimension = width, n_layers = nhidden, activation=(x-> max(0,x)^2/2)) +nn = NeuralNetwork(arch) + +# This is where training starts +batch = Batch(batch_size) +o = Optimizer(AdamOptimizer(), nn) + +loss_array = o(nn, dl, batch, n_epochs) + + +# Saving the parameters of the network +println("Saving the parameters of the neural network...") +flush(stdout) + +params = GeometricMachineLearning.map_to_cpu(nn.params) + +save(path_out,"parameters", params, "training loss", loss_array, "ni_dim", ni_dim, "T", T, "nt", nt, "n_epochs", n_epochs, "width", width, "nhidden", nhidden, "batch_size", batch_size, "nu", nu) + +println(" ...Done!") +flush(stdout) + + diff --git a/scripts/parametric_data_helpers.jl b/scripts/parametric_data_helpers.jl new file mode 100644 index 000000000..eda9a15bf --- /dev/null +++ b/scripts/parametric_data_helpers.jl @@ -0,0 +1,95 @@ +# Shared by the parametric/forced training scripts in this directory. These helpers reshape a +# trajectory ensemble into the `(system dimension, time, parameter)` layout `ParametricDataLoader` +# wants, with one `NamedTuple` of system parameters per column. +# +# They are here rather than in each script because all three used to carry a verbatim copy. They do +# not really belong in `scripts/` either -- generating and reshaping data sets is `GMLDatasets`' +# job -- https://github.com/JuliaGNI/GMLDatasets.jl/issues/5 tracks moving them there -- so treat +# this file as a staging post. + +using GeometricMachineLearning +using GeometricMachineLearning: QPT, QPT2 + +""" +Turn a vector of numbers into a vector of `NamedTuple`s to be used by `ParametricDataLoader`. +""" +function turn_parameters_into_correct_format(t::AbstractVector, IC::AbstractVector{<:NamedTuple}) + vec_of_params = NamedTuple[] + for time_step ∈ t + time_step == t[end] || push!(vec_of_params, (t = time_step, )) + end + vcat((vec_of_params for _ in axes(IC, 1))...) +end + +@doc raw""" +Turn a `NamedTuple` of ``(q,p)`` data into two tensors of the correct format. + +This is the tricky part as the structure of the input array(s) needs to conform with the structure of the parameters. + +Here the data are rearranged in an array of size ``(n, 2, t_f - 1)`` where ``[t_0, t_1, \ldots, t_f]`` is the vector storing the time steps. + +If we deal with different initial conditions as well, we still put everything into the third (parameter) axis. + +# Example + +```jldoctest +using GeometricMachineLearning + +q = [1. 2. 3.; 4. 5. 6.] +p = [1.5 2.5 3.5; 4.5 5.5 6.5] +qp = (q = q, p = p) +turn_q_p_data_into_correct_format(qp) + +# output + +(q = [1.0 2.0; 4.0 5.0;;; 2.0 3.0; 5.0 6.0], p = [1.5 2.5; 4.5 5.5;;; 2.5 3.5; 5.5 6.5]) +``` +""" +function turn_q_p_data_into_correct_format(qp::QPT2{T, 2}) where {T} + number_of_time_steps = size(qp.q, 2) - 1 # not counting t₀ + number_of_initial_conditions = size(qp.q, 1) + q_array = zeros(T, 1, 2, number_of_time_steps * number_of_initial_conditions) + p_array = zeros(T, 1, 2, number_of_time_steps * number_of_initial_conditions) + for initial_condition_index ∈ 0:(number_of_initial_conditions - 1) + for time_index ∈ 1:number_of_time_steps + q_array[:, 1, initial_condition_index * number_of_time_steps + time_index] .= qp.q[initial_condition_index + 1, time_index] + q_array[:, 2, initial_condition_index * number_of_time_steps + time_index] .= qp.q[initial_condition_index + 1, time_index + 1] + p_array[:, 1, initial_condition_index * number_of_time_steps + time_index] .= qp.p[initial_condition_index + 1, time_index] + p_array[:, 2, initial_condition_index * number_of_time_steps + time_index] .= qp.p[initial_condition_index + 1, time_index + 1] + end + end + (q = q_array, p = p_array) +end + +""" +This takes time as a single additional parameter (third axis). +""" +function load_time_dependent_harmonic_oscillator_with_parametric_data_loader(qp::QPT{T}, t::AbstractVector{T}, IC::AbstractVector) where {T} + qp_reformatted = turn_q_p_data_into_correct_format(qp) + t_reformatted = turn_parameters_into_correct_format(t, IC) + ParametricDataLoader(qp_reformatted, t_reformatted) +end + +""" + forced_harmonic_oscillator_solution(t, IC; omega, Omega, F) + +The analytic solution of the sinusoidally forced harmonic oscillator, as `(q, p)` arrays of size +`(number of initial conditions, length(t))`. +""" +function forced_harmonic_oscillator_solution(t::AbstractVector, IC::AbstractVector{<:NamedTuple}; + omega::Real, Omega::Real, F::Real) + ni = length(IC) + q = zeros(Float64, ni, length(t)) + p = zeros(Float64, ni, length(t)) + amplitude = F / (omega^2 - Omega^2) + for i in eachindex(t) + for j in 1:ni + q[j, i] = (IC[j].p - Omega * amplitude) / omega * sin(omega * t[i]) + + IC[j].q * cos(omega * t[i]) + amplitude * sin(Omega * t[i]) + p[j, i] = -omega^2 * IC[j].q * sin(omega * t[i]) + + (IC[j].p - Omega * amplitude) * cos(omega * t[i]) + + Omega * amplitude * cos(Omega * t[i]) + end + end + (q = q, p = p) +end diff --git a/src/architectures/forced_generalized_hamiltonian_neural_network.jl b/src/architectures/forced_generalized_hamiltonian_neural_network.jl index 9b417c7d6..1360e7eb1 100644 --- a/src/architectures/forced_generalized_hamiltonian_neural_network.jl +++ b/src/architectures/forced_generalized_hamiltonian_neural_network.jl @@ -48,3 +48,5 @@ function (o::Optimizer)(nn::NeuralNetwork{<:ForcedGeneralizedHamiltonianArchitec loss::NetworkLoss = ParametricLoss(); kwargs...) o(nn, dl, batch, n_epochs, loss, ZygotePullback(loss); kwargs...) end + +_n_symplectic_integrators(arch::ForcedGeneralizedHamiltonianArchitecture) = arch.n_integrators diff --git a/src/architectures/parametric_resnet.jl b/src/architectures/parametric_resnet.jl index 2a4ff022d..67ad36e5e 100644 --- a/src/architectures/parametric_resnet.jl +++ b/src/architectures/parametric_resnet.jl @@ -1,3 +1,24 @@ +@doc raw""" + ParametricResNet(dim; width, n_blocks, activation, parameters) + +A [`ResNet`](@ref) whose blocks also take the parameters of the *system*, built from +[`ParametricResNetLayer`](@ref)s. + +This is the architecture without structure preservation that +[`GeneralizedHamiltonianArchitecture`](@ref) is compared against: it maps +``(x, \mu) \mapsto x'`` with the same information available to it, but nothing in it makes the map +symplectic. + +# Keyword arguments + +- `width = dim`: the width of the hidden layer of each block, +- `n_blocks = $(HNN_nhidden_default)`: the number of blocks, +- `activation = $(HNN_activation_default)`, +- `parameters = NullParameters()`: a `NamedTuple` of system parameters, used for its shape. + +`ResNet(dim; n_blocks, width, parameters)` dispatches here when `parameters` is anything other than +`NullParameters`. +""" struct ParametricResNet{AT <: Activation, PT <: OptionalParameters} <: NeuralNetworkIntegrator sys_dim::Int n_blocks::Int diff --git a/src/data_loader/parametric_data_loader.jl b/src/data_loader/parametric_data_loader.jl index 6a520a7c7..8b2be711b 100644 --- a/src/data_loader/parametric_data_loader.jl +++ b/src/data_loader/parametric_data_loader.jl @@ -49,27 +49,6 @@ function ParametricDataLoader(ensemble_solution::EnsembleSolution{T, T1, Vector{ ParametricDataLoader(data, params) end -# """ -# rearrange_parameters(parameters) -# -# Rearrange `parameters` such that they can be used by [`ParametricDataLoader`](@ref). -# """ -# function rearrange_parameters(parameters::Vector{<:NamedTuple}) -# parameters_rearranged = zeros(_eltype(parameters), ) -# end - -# function batch_over_two_axes(batch::Batch, number_columns::Int, third_dim::Int, dl::ParametricDataLoader) -# time_indices = shuffle(1:number_columns) -# parameter_indices = shuffle(1:third_dim) -# complete_indices = Iterators.product(time_indices, parameter_indices) |> collect |> vec -# batches = () -# n_batches = number_of_batches(dl, batch) -# for batch_number in 1:(n_batches - 1) -# batches = (batches..., complete_indices[(batch_number - 1) * batch.batch_size + 1 : batch_number * batch.batch_size]) -# end -# (batches..., complete_indices[(n_batches - 1) * batch.batch_size + 1:end]) -# end - function optimize_for_one_epoch!( opt::Optimizer, model, ps::Union{NeuralNetworkParameters, NamedTuple}, diff --git a/src/layers/parametric_resnet_layer.jl b/src/layers/parametric_resnet_layer.jl index 149d6c4e0..29abb531f 100644 --- a/src/layers/parametric_resnet_layer.jl +++ b/src/layers/parametric_resnet_layer.jl @@ -1,3 +1,28 @@ +@doc raw""" + ParametricResNetLayer(dim, width, activation; parameters, return_parameters) + +A [`WideResNetLayer`](@ref) whose hidden layer also sees the parameters of the *system*. + +The flattened system parameters are appended to the input of the upscaling weight, so the layer +computes + +```math + x \mapsto x + \sigma(W_\mathrm{down}\sigma(W_\mathrm{up}[x; \mu] + b_\mathrm{up}) + b), +``` + +where ``\mu`` are the system parameters. `W_\mathrm{up}` is therefore +``\mathrm{width}\times(\mathrm{dim} + |\mu|)`` wide. + +# Keyword arguments + +- `parameters = NullParameters()`: a `NamedTuple` of system parameters, used only for its *shape* -- + the layer stores the resulting `NeuralNetworkParameters.ParameterLayout` and the flattened length, + and the values are supplied per call. +- `return_parameters::Bool`: whether to pass the system parameters on to the next layer alongside the + output, which is what lets a `Chain` of these thread them through. + +This is the building block of [`ParametricResNet`](@ref). +""" struct ParametricResNetLayer{M, N, F1 <: Activation, PT, ReturnParameters} <: AbstractExplicitLayer{M, N} width::Int activation::F1 @@ -35,16 +60,6 @@ function (d::ParametricResNetLayer{M, M, F, PT, true})(x::AbstractVecOrMat, prob (x + d.activation.(ps.downscale_weight * d.activation.(ps.upscale_weight * input .+ ps.upscale_bias) .+ ps.bias), problem_params) end -# function (d::ParametricResNetLayer{M, M, F, PT, false})(x::AbstractArray{T, 3}, problem_params::AbstractVector, ps::NamedTuple) where {M, F, PT, T} -# input = concatenate_array_with_parameters(x, problem_params) -# x + d.activation.(mat_tensor_mul(ps.downscale_weight, d.activation.(mat_tensor_mul(ps.upscale_weight, x) .+ ps.upscale_bias)) .+ ps.bias) -# end -# -# function (d::ParametricResNetLayer{M, M, F, PT, true})(x::AbstractArray{T, 3}, problem_params::AbstractVector, ps::NamedTuple) where {M, F, PT, T} -# input = concatenate_array_with_parameters(x, problem_params) -# (x + d.activation.(mat_tensor_mul(ps.downscale_weight, d.activation.(mat_tensor_mul(ps.upscale_weight, x) .+ ps.upscale_bias)) .+ ps.bias), problem_params) -# end - (d::ParametricResNetLayer)(input::Tuple, ps::NamedTuple) = length(input) == 2 ? d(input..., ps) : error("The tuple must contain the input array/nt as well as the system parameters.") function (d::ParametricResNetLayer{M, M, F, PT, false})(z::QPT, problem_params::OptionalParameters, ps::NamedTuple) where {M, F, PT} diff --git a/src/layers/wide_resnet.jl b/src/layers/wide_resnet.jl index f1a4c4682..caad09177 100644 --- a/src/layers/wide_resnet.jl +++ b/src/layers/wide_resnet.jl @@ -1,3 +1,21 @@ +@doc raw""" + WideResNetLayer(dim, width, activation) + +A [`ResNetLayer`](@ref) whose hidden layer is `width` wide, independent of `dim`. + +`ResNetLayer` applies a single ``\mathrm{dim} \times \mathrm{dim}`` weight, so its capacity is tied +to the system dimension. This one upscales to `width`, applies the activation, and downscales again: + +```math + x \mapsto x + \sigma(W_\mathrm{down}\sigma(W_\mathrm{up}x + b_\mathrm{up}) + b), +``` + +with ``W_\mathrm{up}\in\mathbb{R}^{\mathrm{width}\times\mathrm{dim}}`` and +``W_\mathrm{down}\in\mathbb{R}^{\mathrm{dim}\times\mathrm{width}}``. [`ResNet`](@ref) uses it +whenever the `width` it is given differs from the system dimension. + +Also see [`ParametricResNetLayer`](@ref), which additionally takes the parameters of the system. +""" struct WideResNetLayer{M, N, F1} <: AbstractExplicitLayer{M, N} width::Int activation::F1 diff --git a/src/pullbacks/symbolic_hnn_pullback.jl b/src/pullbacks/symbolic_hnn_pullback.jl index a812d4be9..b7ddaf535 100644 --- a/src/pullbacks/symbolic_hnn_pullback.jl +++ b/src/pullbacks/symbolic_hnn_pullback.jl @@ -26,6 +26,43 @@ function SymbolicPullback(arch::HamiltonianArchitecture) SymbolicPullback(loss, SymbolicNeuralNetworks.ParameterGradient(gradient_function)) end +# How many `SymplecticEuler` layers the architecture stacks. Only the generalized architectures +# have more than one; everything else traces as a single pass. +_n_symplectic_integrators(::Any) = 1 +_n_symplectic_integrators(arch::GeneralizedHamiltonianArchitecture) = arch.n_integrators + +@doc raw""" + _check_symbolic_pullback_is_tractable(arch) + +Throw if building a `SymbolicPullback` for `arch` would not finish. + +`SymbolicPullback` traces the whole chain symbolically, and each `SymplecticEuler` layer inlines the +symbolic gradient of its energy network — an expression that is itself already a derivative. Stacking +integrators inlines that expression inside itself, so it grows *multiplicatively* rather than +additively. Measured at `dim = 4, width = 4, nhidden = 1`: + +| `n_integrators` | symbolic loss | its parameter derivative | build time | +|---|---|---|---| +| 1 | 3.4 ⋅ 10⁵ characters | 1.5 ⋅ 10⁸ characters | ≈ 1.4 s | +| 2 | 1.4 ⋅ 10⁹ characters | — | does not finish, past 8 GB | + +So one integrator is fine — and worth it, the built function evaluates about 100 times faster than +the `Zygote` pullback — while two are hopeless. Rather than let that look like a hang, refuse it. + +Removing the limit means not tracing the inlined chain at all: composing the pullback layer by layer +from the gradients each `SymplecticEuler` has already built. See +[issue #245](https://github.com/JuliaGNI/GeometricMachineLearning.jl/issues/245). +""" +function _check_symbolic_pullback_is_tractable(arch) + n = _n_symplectic_integrators(arch) + n > 1 && throw(ArgumentError( + "cannot build a `SymbolicPullback` for an architecture with $(n) integrators: the symbolic " * + "expression grows multiplicatively with `n_integrators`, and already exceeds 10⁹ terms at " * + "two, so the build does not finish. Use `ZygotePullback(loss)`, which is what `Optimizer` " * + "uses by default, or reduce `n_integrators` to 1. See GeometricMachineLearning issue #245.")) + nothing +end + @doc raw""" SymbolicPullback(nn, loss, system_params) @@ -45,6 +82,7 @@ sum of the per-sample gradients. """ function SymbolicPullback(nn::NeuralNetwork, loss::ParametricLoss, system_params::OptionalParameters; cse::Bool = true, inplace::Bool = true) + _check_symbolic_pullback_is_tractable(nn.architecture) symbolic_system_parameters = SymbolicNeuralNetworks.symbolic_variables(system_params, :S) symbolic_network_parameters = SymbolicNeuralNetworks.symbolic_variables(params(nn), :W) diff --git a/test/generalized_hamiltonian_neural_networks/pghnn_symbolic_pullback_single_layer_test.jl b/test/generalized_hamiltonian_neural_networks/pghnn_symbolic_pullback_single_layer_test.jl index a4c52604e..a55b061fc 100644 --- a/test/generalized_hamiltonian_neural_networks/pghnn_symbolic_pullback_single_layer_test.jl +++ b/test/generalized_hamiltonian_neural_networks/pghnn_symbolic_pullback_single_layer_test.jl @@ -61,3 +61,15 @@ end # the network actually depends on its parameters here. @test any(any(abs.(block) .> 1e-8) for layer in keys(reference_gradient) for block in values(reference_gradient[layer])) + +# Building the pullback for a stacked architecture is refused rather than left to hang: the symbolic +# expression grows multiplicatively with `n_integrators` and already exceeds 10⁹ terms at two. +stacked = NeuralNetwork(GeneralizedHamiltonianArchitecture(dim; n_integrators = 2, + parameters = system_parameters)) +@test_throws ArgumentError SymbolicPullback(stacked, loss, system_parameters) + +# one integrator is the supported case, and it builds +single = NeuralNetwork(GeneralizedHamiltonianArchitecture(dim; n_integrators = 1, + parameters = system_parameters)) +@test SymbolicPullback(single, loss, system_parameters) isa SymbolicPullback + From 6db0b980bbf8bd77d1dc5e27bda21d7c9943a081 Mon Sep 17 00:00:00 2001 From: Michael Kraus Date: Mon, 24 Aug 2026 02:30:41 +0900 Subject: [PATCH 4/4] Follow AbstractNeuralNetworks 0.7 and SymbolicNeuralNetworks 0.6 The parameter container moved to `NeuralNetworkParameters` and ANN 0.7 removed the old name outright, so `NeuralNetworkParameters` as a *type* is `NetworkParameters` throughout. The module keeps its name where it is a module. `Project.toml` is now identical to main's: this work needs no dependency main does not already have. The `[sources]` pin and the CI step that went with it are gone -- NeuralNetworkParameters 0.1.1 is registered and main depends on it. A local `path =` entry that had survived from an earlier resolve goes with them. Three pirated methods delete rather than move, because upstream now covers them: * `applychain(::Tuple, ::Tuple{<:QPTOAT2, <:OptionalParameters}, ::Tuple)`. ANN 0.7 leaves `applychain`'s data argument untyped, and its comment gives the reason -- the old signature "forced downstream packages to commit type piracy to push anything else through a `Chain`". The generic method carries the `(state, system parameters)` tuple already. * `networkbackend(::LazyArrays.ApplyArray)`. LazyArrays is not a dependency of GML any more, so this would not have compiled. * `h5save(::HDF5.Group, ::NetworkParameters, ::AbstractString)`. NeuralNetworkParameters' own extension has the `::H5DataStore` method, and `HDF5.Group` is one. GML's was not merely redundant but worse: being more specific it shadowed the upstream one, which writes the `kind` attribute that records what to rebuild a leaf as. The four three-argument `Chain` functors become `apply_parametric`, a function GML owns. That is newly possible for the same reason the first deletion is: with `applychain` generic in its data argument, the tuple threads through the ordinary two-argument functor, so these were only an entry point and never needed to be methods on someone else's type. Two pirated groups are left, both genuinely upstream and both still marked: `SymbolicNeuralNetworks.Jacobian(f, nn, dim2)`, which 0.6 did not add, and the three-tuple `SymbolicPullback` call operators. The `n_integrators > 1` guard stays. SNN 0.6 shipped the layerwise construction (SNN #49) and its changelog names this case, but it does not reach it: the seam is a plain `Vector{Num}`, so a `SymplecticEuler` with `return_parameters = true` -- which passes the system parameters on and returns a `Tuple` -- cannot be seeded. `composes_layerwise` still says the chain decomposes, so `:auto` commits and then raises, where it is documented never to. Filed as SNN #54; measured at `n_integrators = 1`, where `layerwise = false` builds in 3.2 s and `:auto` throws. Co-Authored-By: Claude Opus 5 (1M context) --- Project.toml | 3 - .../hamiltonian_neural_network.md | 1 + ext/HDF5Ext.jl | 10 --- ..._generalized_hamiltonian_neural_network.jl | 2 +- .../generalized_hamiltonian_neural_network.jl | 85 ++++++++++--------- src/data_loader/parametric_data_loader.jl | 2 +- src/layers/forcing_dissipation_layers.jl | 22 ++--- src/loss/losses.jl | 4 +- src/pullbacks/symbolic_hnn_pullback.jl | 15 +++- ...arametric_layers_and_architectures_test.jl | 5 +- 10 files changed, 77 insertions(+), 72 deletions(-) diff --git a/Project.toml b/Project.toml index 5ea404921..462957bc8 100644 --- a/Project.toml +++ b/Project.toml @@ -27,9 +27,6 @@ Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f" [weakdeps] HDF5 = "f67ccb44-e63f-5c2f-98bd-6dc0ccc4ba2f" -[sources] -NeuralNetworkParameters = {path = "/Users/mkraus/Datashare/Julia/NeuralNetworkParameters"} - [extensions] HDF5Ext = "HDF5" diff --git a/docs/src/architectures/hamiltonian_neural_network.md b/docs/src/architectures/hamiltonian_neural_network.md index 82a1086ec..505830dc3 100644 --- a/docs/src/architectures/hamiltonian_neural_network.md +++ b/docs/src/architectures/hamiltonian_neural_network.md @@ -61,6 +61,7 @@ GeometricMachineLearning.ForcingLayerQ GeometricMachineLearning.ForcingLayerP GeometricMachineLearning.ForcingLayerQP GeometricMachineLearning.ParametricDataLoader +GeometricMachineLearning.apply_parametric GeometricMachineLearning.ParametricResNet GeometricMachineLearning.ParametricResNetLayer GeometricMachineLearning.WideResNetLayer diff --git a/ext/HDF5Ext.jl b/ext/HDF5Ext.jl index fbec87bfe..0853d4257 100644 --- a/ext/HDF5Ext.jl +++ b/ext/HDF5Ext.jl @@ -7,16 +7,6 @@ import AbstractNeuralNetworks: changebackend, NeuralNetworkBackend, Architecture # `AbstractNeuralNetworks` 0.7, which only re-binds them; reach for them where they are defined. import NeuralNetworkParameters: NetworkParameters, params, save, load -# A `NeuralNetworkParameters` nested inside a parameter tree -- the parameter-dependent -# architectures put one per sub-network. AbstractNeuralNetworks has `save(::H5DataStore, -# ::NeuralNetworkParameters)` for the top level only. -# -# TODO: type piracy -- `h5save` and `NeuralNetworkParameters` are both AbstractNeuralNetworks'. -# This belongs in ANN's own `ext/HDF5Ext.jl`, next to `h5save(::H5DataStore, ::NamedTuple, …)`. -function h5save(h5::HDF5.Group, p::NeuralNetworkParameters, path::AbstractString) - h5save(h5, params(p), path) -end - # --------------------------------------------------------------------------- # changebackend — new methods for GML special array types # diff --git a/src/architectures/forced_generalized_hamiltonian_neural_network.jl b/src/architectures/forced_generalized_hamiltonian_neural_network.jl index 1360e7eb1..cfc60b992 100644 --- a/src/architectures/forced_generalized_hamiltonian_neural_network.jl +++ b/src/architectures/forced_generalized_hamiltonian_neural_network.jl @@ -40,7 +40,7 @@ end # `StandardHamiltonianArchitecture` takes no system parameters. function (nn::NeuralNetwork{<:ForcedGeneralizedHamiltonianArchitecture})(qp::QPTOAT2, problem_params::OptionalParameters) - nn.model(qp, problem_params, params(nn)) + apply_parametric(nn.model, qp, problem_params, params(nn)) end function (o::Optimizer)(nn::NeuralNetwork{<:ForcedGeneralizedHamiltonianArchitecture}, diff --git a/src/architectures/generalized_hamiltonian_neural_network.jl b/src/architectures/generalized_hamiltonian_neural_network.jl index 447a9cd3a..d1cdfeb17 100644 --- a/src/architectures/generalized_hamiltonian_neural_network.jl +++ b/src/architectures/generalized_hamiltonian_neural_network.jl @@ -106,7 +106,7 @@ se = SymbolicPotentialEnergy(dim, width, nhidden, activation; parameters = param network_params = NeuralNetwork(Chain(se); initializer = OneInitializer()).params built_grad = build_gradient(se) -grad(qp::AbstractArray, problem_params::OptionalParameters, params::NeuralNetworkParameters) = built_grad(concatenate_array_with_parameters(qp, problem_params), params) +grad(qp::AbstractArray, problem_params::OptionalParameters, params::NetworkParameters) = built_grad(concatenate_array_with_parameters(qp, problem_params), params) grad([0.5, 0.25], params, network_params) @@ -189,35 +189,35 @@ function concatenate_array_with_parameters(qp::AbstractMatrix, params::AbstractV hcat((concatenate_array_with_parameters(@view(qp[:, i]), params[i]) for i in axes(params, 1))...) end -function (integrator::SymplecticEulerA{M, N, FT, AT, false})(qp::QPT2, problem_params::OptionalParameters, params::NeuralNetworkParameters) where {M, N, FT, AT} +function (integrator::SymplecticEulerA{M, N, FT, AT, false})(qp::QPT2, problem_params::OptionalParameters, params::NetworkParameters) where {M, N, FT, AT} input = concatenate_array_with_parameters(qp.p, problem_params) (q = @view((qp.q + integrator.gradient_function(input, params))[:, 1]), p = qp.p) end -function (integrator::SymplecticEulerB{M, N, FT, AT, false})(qp::QPT2, problem_params::OptionalParameters, params::NeuralNetworkParameters) where {M, N, FT, AT} +function (integrator::SymplecticEulerB{M, N, FT, AT, false})(qp::QPT2, problem_params::OptionalParameters, params::NetworkParameters) where {M, N, FT, AT} input = concatenate_array_with_parameters(qp.q, problem_params) (q = qp.q, p = @view((qp.p - integrator.gradient_function(input, params))[:, 1])) end -function (integrator::SymplecticEulerA{M, N, FT, AT, true})(qp::QPT2, problem_params::OptionalParameters, params::NeuralNetworkParameters) where {M, N, FT, AT} +function (integrator::SymplecticEulerA{M, N, FT, AT, true})(qp::QPT2, problem_params::OptionalParameters, params::NetworkParameters) where {M, N, FT, AT} input = concatenate_array_with_parameters(qp.p, problem_params) ((q = @view((qp.q + integrator.gradient_function(input, params))[:, 1]), p = qp.p), problem_params) end -function (integrator::SymplecticEulerB{M, N, FT, AT, true})(qp::QPT2, problem_params::OptionalParameters, params::NeuralNetworkParameters) where {M, N, FT, AT} +function (integrator::SymplecticEulerB{M, N, FT, AT, true})(qp::QPT2, problem_params::OptionalParameters, params::NetworkParameters) where {M, N, FT, AT} input = concatenate_array_with_parameters(qp.q, problem_params) ((q = qp.q, p = @view((qp.p - integrator.gradient_function(input, params))[:, 1])), problem_params) end -function (integrator::SymplecticEuler)(qp_params::Tuple{<:QPTOAT2, <:OptionalParameters}, params::NeuralNetworkParameters) +function (integrator::SymplecticEuler)(qp_params::Tuple{<:QPTOAT2, <:OptionalParameters}, params::NetworkParameters) integrator(qp_params..., params) end -function (integrator::SymplecticEuler)(::TT, ::NeuralNetworkParameters) where {TT <: Tuple} +function (integrator::SymplecticEuler)(::TT, ::NetworkParameters) where {TT <: Tuple} error("The input is of type $(TT). This shouldn't be the case!") end -function (integrator::SymplecticEuler{M, N, FT, AT, Type, true})(qp::AbstractArray, problem_params::OptionalParameters, params::NeuralNetworkParameters) where {M, N, FT, AT, Type} +function (integrator::SymplecticEuler{M, N, FT, AT, Type, true})(qp::AbstractArray, problem_params::OptionalParameters, params::NetworkParameters) where {M, N, FT, AT, Type} @assert iseven(size(qp, 1)) n = size(qp, 1)÷2 qp_split = assign_q_and_p(qp, n) @@ -225,7 +225,7 @@ function (integrator::SymplecticEuler{M, N, FT, AT, Type, true})(qp::AbstractArr (vcat(evaluated.q, evaluated.p), problem_params) end -function (integrator::SymplecticEuler{M, N, FT, AT, Type, false})(qp::AbstractArray, problem_params::OptionalParameters, params::NeuralNetworkParameters) where {M, N, FT, AT, Type} +function (integrator::SymplecticEuler{M, N, FT, AT, Type, false})(qp::AbstractArray, problem_params::OptionalParameters, params::NetworkParameters) where {M, N, FT, AT, Type} @assert iseven(size(qp, 1)) n = size(qp, 1)÷2 qp_split = assign_q_and_p(qp, n) @@ -233,7 +233,7 @@ function (integrator::SymplecticEuler{M, N, FT, AT, Type, false})(qp::AbstractAr vcat(evaluated.q, evaluated.p) end -(integrator::SymplecticEuler)(qp::QPTOAT2, params::NeuralNetworkParameters) = integrator(qp, NullParameters(), params) +(integrator::SymplecticEuler)(qp::QPTOAT2, params::NetworkParameters) = integrator(qp, NullParameters(), params) """ GeneralizedHamiltonianArchitecture <: HamiltonianArchitecture @@ -265,20 +265,6 @@ struct GeneralizedHamiltonianArchitecture{AT, PT <: OptionalParameters} <: Hamil end end -# The parameter-dependent layers pass `(state, system parameters)` down the chain, so `applychain` -# has to accept that tuple as its data argument. -# -# TODO: type piracy -- `applychain` is AbstractNeuralNetworks' and every argument type here is -# `Base`'s. ANN's own `applychain(layers, x, ps::Union{NamedTuple, NeuralNetworkParameters})` is -# already generic in `x`; widening the `@generated` method the same way would remove the need. -@generated function AbstractNeuralNetworks.applychain(layers::Tuple, x::Tuple{<:QPTOAT2, <:OptionalParameters}, ps::Tuple) - N = length(fieldtypes((layers))) - x_symbols = vcat([:x], [gensym() for _ in 1:N]) - calls = [:(($(x_symbols[i + 1])) = layers[$i]($(x_symbols[i]), ps[$i])) for i in 1:N] - push!(calls, :(return $(x_symbols[N + 1]))) - return Expr(:block, calls...) -end - index_qpt(qp::QPT2{T, 2}, i, j) where {T} = (q = qp.q[i, j], p = qp.p[i, j]) index_gpt(qp::QPT2{T, 3}, i, j, k) where {T} = (q = qp.q[i, j, k], p = qp.p[i, j, k]) @@ -296,40 +282,59 @@ function Chain(ghnn_arch::GeneralizedHamiltonianArchitecture) end function (nn::NeuralNetwork{GT})(qp::QPTOAT2, problem_params::OptionalParameters) where {GT <: GeneralizedHamiltonianArchitecture} - nn.model(qp, problem_params, params(nn)) + apply_parametric(nn.model, qp, problem_params, params(nn)) end -# TODO: type piracy -- `Chain` is AbstractNeuralNetworks' and so is every argument type of the four -# functors below. A `ParametricChain` wrapper owned by GML, or these methods upstream, would fix it. -function (model::Chain)(qp::QPTOAT2, problem_params::OptionalParameters, params::Union{NeuralNetworkParameters, NamedTuple}) - model((qp, problem_params), params) +@doc raw""" + apply_parametric(model, qp, system_parameters, ps) + +Apply `model` to `qp` with the parameters of the *system* alongside the state, and the network +parameters `ps`. + +`system_parameters` is either one `OptionalParameters` for the whole input, or — for a batch drawn by +[`ParametricDataLoader`](@ref) — one entry per sample. + +# Implementation + +A `Chain` of parameter-dependent layers threads `(state, system parameters)` from layer to layer, so +this hands the pair to the ordinary two-argument `Chain` functor and lets +`AbstractNeuralNetworks.applychain` carry it: since AbstractNeuralNetworks 0.7 that method leaves its +data argument untyped, so a tuple needs no method of its own. + +This is a function of GML's rather than a three-argument functor on `Chain`, which would be type +piracy — `Chain` and every argument type belong to AbstractNeuralNetworks. +""" +function apply_parametric(model::Chain, qp::QPTOAT2, problem_params::OptionalParameters, + ps::Union{NetworkParameters, NamedTuple}) + model((qp, problem_params), ps) end -function (c::Chain)(qp::QPT2{T, 3}, system_params::AbstractVector, ps::Union{NamedTuple, NeuralNetworkParameters})::QPT2{T} where {T} +function apply_parametric(c::Chain, qp::QPT2{T, 3}, system_params::AbstractVector, + ps::Union{NamedTuple, NetworkParameters})::QPT2{T} where {T} @assert size(qp.q, 3) == length(system_params) @assert size(qp.q, 2) == 1 - output_vectorwise = [c(index_gpt(qp, :, 1, i), system_params[i], ps) for i in axes(system_params, 1)] + output_vectorwise = [apply_parametric(c, index_gpt(qp, :, 1, i), system_params[i], ps) + for i in axes(system_params, 1)] q_output = hcat([single_output_vectorwise.q for single_output_vectorwise ∈ output_vectorwise]...) p_output = hcat([single_output_vectorwise.p for single_output_vectorwise ∈ output_vectorwise]...) - (q = reshape(q_output, size(q_output, 1), 1, size(q_output, 2)), p = reshape(p_output, size(p_output, 1), 1, size(p_output, 2))) + (q = reshape(q_output, size(q_output, 1), 1, size(q_output, 2)), + p = reshape(p_output, size(p_output, 1), 1, size(p_output, 2))) end -function (c::Chain)(qp::AbstractArray{T, 2}, system_params::AbstractVector, ps::Union{NamedTuple, NeuralNetworkParameters}) where {T} +function apply_parametric(c::Chain, qp::AbstractArray{T, 2}, system_params::AbstractVector, + ps::Union{NamedTuple, NetworkParameters}) where {T} @assert _size(qp, 2) == length(system_params) qp_reshaped = reshape(qp, size(qp, 1), 1, length(system_params)) - c(qp_reshaped, system_params, ps) + apply_parametric(c, qp_reshaped, system_params, ps) end -function (c::Chain)(qp::AbstractArray{T, 3}, system_params::AbstractVector, ps::Union{NamedTuple, NeuralNetworkParameters}) where {T} +function apply_parametric(c::Chain, qp::AbstractArray{T, 3}, system_params::AbstractVector, + ps::Union{NamedTuple, NetworkParameters}) where {T} @assert size(qp, 3) == length(system_params) @assert size(qp, 2) == 1 @assert iseven(size(qp, 1)) n = size(qp, 1)÷2 qp_split = assign_q_and_p(qp, n) - c_output = c(qp_split, system_params, ps)::QPT + c_output = apply_parametric(c, qp_split, system_params, ps)::QPT reshape(vcat(c_output.q, c_output.p), 2n, length(system_params)) end - -# TODO: type piracy -- `networkbackend` is AbstractNeuralNetworks' and `ApplyArray` is LazyArrays'. -# Belongs in ANN, which already dispatches `networkbackend` on array types it does not own either. -AbstractNeuralNetworks.networkbackend(::LazyArrays.ApplyArray) = AbstractNeuralNetworks.CPU() diff --git a/src/data_loader/parametric_data_loader.jl b/src/data_loader/parametric_data_loader.jl index 8b2be711b..3d9ef84cc 100644 --- a/src/data_loader/parametric_data_loader.jl +++ b/src/data_loader/parametric_data_loader.jl @@ -51,7 +51,7 @@ end function optimize_for_one_epoch!( opt::Optimizer, model, - ps::Union{NeuralNetworkParameters, NamedTuple}, + ps::Union{NetworkParameters, NamedTuple}, dl::ParametricDataLoader{T}, batch::Batch, _pullback::AbstractPullback, diff --git a/src/layers/forcing_dissipation_layers.jl b/src/layers/forcing_dissipation_layers.jl index cc7ea16cd..fc41bc6b3 100644 --- a/src/layers/forcing_dissipation_layers.jl +++ b/src/layers/forcing_dissipation_layers.jl @@ -105,45 +105,45 @@ function ForcingLayerQP(dim::Integer, width::Integer=dim, nhidden::Integer=HNN_n ForcingLayer(dim, width, nhidden, activation; parameters=parameters, return_parameters=return_parameters, type=:QP) end -function (integrator::ForcingLayerQ{M,N,FT,AT,false})(qp::QPT2, problem_params::OptionalParameters, params::NeuralNetworkParameters) where {M,N,FT,AT} +function (integrator::ForcingLayerQ{M,N,FT,AT,false})(qp::QPT2, problem_params::OptionalParameters, params::NetworkParameters) where {M,N,FT,AT} input = concatenate_array_with_parameters(qp.q, problem_params) (q=qp.q, p=qp.p + integrator.model(input, params)) end -function (integrator::ForcingLayerP{M,N,FT,AT,false})(qp::QPT2, problem_params::OptionalParameters, params::NeuralNetworkParameters) where {M,N,FT,AT} +function (integrator::ForcingLayerP{M,N,FT,AT,false})(qp::QPT2, problem_params::OptionalParameters, params::NetworkParameters) where {M,N,FT,AT} input = concatenate_array_with_parameters(qp.p, problem_params) (q=qp.q, p=qp.p + integrator.model(input, params)) end -function (integrator::ForcingLayerQP{M,N,FT,AT,false})(qp::QPT2, problem_params::OptionalParameters, params::NeuralNetworkParameters) where {M,N,FT,AT} +function (integrator::ForcingLayerQP{M,N,FT,AT,false})(qp::QPT2, problem_params::OptionalParameters, params::NetworkParameters) where {M,N,FT,AT} input = concatenate_array_with_parameters(vcat(qp.q, qp.p), problem_params) (q=qp.q, p=qp.p + integrator.model(input, params)) end -function (integrator::ForcingLayerQ{M,N,FT,AT,true})(qp::QPT2, problem_params::OptionalParameters, params::NeuralNetworkParameters) where {M,N,FT,AT} +function (integrator::ForcingLayerQ{M,N,FT,AT,true})(qp::QPT2, problem_params::OptionalParameters, params::NetworkParameters) where {M,N,FT,AT} input = concatenate_array_with_parameters(qp.q, problem_params) ((q=qp.q, p=qp.p + integrator.model(input, params)), problem_params) end -function (integrator::ForcingLayerP{M,N,FT,AT,true})(qp::QPT2, problem_params::OptionalParameters, params::NeuralNetworkParameters) where {M,N,FT,AT} +function (integrator::ForcingLayerP{M,N,FT,AT,true})(qp::QPT2, problem_params::OptionalParameters, params::NetworkParameters) where {M,N,FT,AT} input = concatenate_array_with_parameters(qp.p, problem_params) ((q=qp.q, p=qp.p + integrator.model(input, params)), problem_params) end -function (integrator::ForcingLayerQP{M,N,FT,AT,true})(qp::QPT2, problem_params::OptionalParameters, params::NeuralNetworkParameters) where {M,N,FT,AT} +function (integrator::ForcingLayerQP{M,N,FT,AT,true})(qp::QPT2, problem_params::OptionalParameters, params::NetworkParameters) where {M,N,FT,AT} input = concatenate_array_with_parameters(vcat(qp.q, qp.p), problem_params) ((q=qp.q, p=qp.p + integrator.model(input, params)), problem_params) end -function (integrator::ForcingLayer)(qp_params::Tuple{<:QPTOAT2,<:OptionalParameters}, params::NeuralNetworkParameters) +function (integrator::ForcingLayer)(qp_params::Tuple{<:QPTOAT2,<:OptionalParameters}, params::NetworkParameters) integrator(qp_params..., params) end -function (integrator::ForcingLayer)(::TT, ::NeuralNetworkParameters) where {TT<:Tuple} +function (integrator::ForcingLayer)(::TT, ::NetworkParameters) where {TT<:Tuple} error("The input is of type $(TT). This shouldn't be the case!") end -function (integrator::ForcingLayer{M,N,FT,AT,Type,true})(qp::AbstractArray, problem_params::OptionalParameters, params::NeuralNetworkParameters) where {M,N,FT,AT,Type} +function (integrator::ForcingLayer{M,N,FT,AT,Type,true})(qp::AbstractArray, problem_params::OptionalParameters, params::NetworkParameters) where {M,N,FT,AT,Type} @assert iseven(size(qp, 1)) n = size(qp, 1) ÷ 2 qp_split = assign_q_and_p(qp, n) @@ -151,7 +151,7 @@ function (integrator::ForcingLayer{M,N,FT,AT,Type,true})(qp::AbstractArray, prob (vcat(evaluated.q, evaluated.p), problem_params) end -function (integrator::ForcingLayer{M,N,FT,AT,Type,false})(qp::AbstractArray, problem_params::OptionalParameters, params::NeuralNetworkParameters) where {M,N,FT,AT,Type} +function (integrator::ForcingLayer{M,N,FT,AT,Type,false})(qp::AbstractArray, problem_params::OptionalParameters, params::NetworkParameters) where {M,N,FT,AT,Type} @assert iseven(size(qp, 1)) n = size(qp, 1) ÷ 2 qp_split = assign_q_and_p(qp, n) @@ -159,5 +159,5 @@ function (integrator::ForcingLayer{M,N,FT,AT,Type,false})(qp::AbstractArray, pro vcat(evaluated.q, evaluated.p) end -(integrator::ForcingLayer)(qp::QPTOAT2, params::NeuralNetworkParameters) = integrator(qp, NullParameters(), params) +(integrator::ForcingLayer)(qp::QPTOAT2, params::NetworkParameters) = integrator(qp, NullParameters(), params) (integrator::ForcingLayer)(qp::QPTOAT2, params::NamedTuple) = integrator(qp, NeuralNetworkParameters(params)) diff --git a/src/loss/losses.jl b/src/loss/losses.jl index c814d972a..61d5de238 100644 --- a/src/loss/losses.jl +++ b/src/loss/losses.jl @@ -275,7 +275,7 @@ This loss does not have any parameters. struct ParametricLoss <: NetworkLoss end function (loss::ParametricLoss)(model::Chain, - params::Union{NamedTuple, NeuralNetworkParameters}, input::CT, output::CT, + params::Union{NamedTuple, NetworkParameters}, input::CT, output::CT, system_parameters::Union{NamedTuple, AbstractVector}) where {CT <: QPTOAT} - _compute_loss(model(input, system_parameters, params), output) + _compute_loss(apply_parametric(model, input, system_parameters, params), output) end diff --git a/src/pullbacks/symbolic_hnn_pullback.jl b/src/pullbacks/symbolic_hnn_pullback.jl index b7ddaf535..2a581d5f2 100644 --- a/src/pullbacks/symbolic_hnn_pullback.jl +++ b/src/pullbacks/symbolic_hnn_pullback.jl @@ -50,7 +50,16 @@ So one integrator is fine — and worth it, the built function evaluates about 1 the `Zygote` pullback — while two are hopeless. Rather than let that look like a hang, refuse it. Removing the limit means not tracing the inlined chain at all: composing the pullback layer by layer -from the gradients each `SymplecticEuler` has already built. See +from the gradients each `SymplecticEuler` has already built. + +`SymbolicNeuralNetworks` 0.6 added exactly that construction +([SNN #49](https://github.com/JuliaGNI/SymbolicNeuralNetworks.jl/issues/49)), but it does not reach +this case yet. Its seam is a plain `Vector{Num}`, so it assumes every layer maps an array to an +array; a `SymplecticEuler` built with `return_parameters = true` passes the system parameters on to +the next layer and returns a `Tuple`, which `layer_seed` cannot seed. `composes_layerwise` says the +chain decomposes, so `layerwise = :auto` commits to that path and then raises — which is +[SNN #54](https://github.com/JuliaGNI/SymbolicNeuralNetworks.jl/issues/54). Until the seam can carry +what a layer passes alongside the state, this limit stays. See [issue #245](https://github.com/JuliaGNI/GeometricMachineLearning.jl/issues/245). """ function _check_symbolic_pullback_is_tractable(arch) @@ -59,7 +68,9 @@ function _check_symbolic_pullback_is_tractable(arch) "cannot build a `SymbolicPullback` for an architecture with $(n) integrators: the symbolic " * "expression grows multiplicatively with `n_integrators`, and already exceeds 10⁹ terms at " * "two, so the build does not finish. Use `ZygotePullback(loss)`, which is what `Optimizer` " * - "uses by default, or reduce `n_integrators` to 1. See GeometricMachineLearning issue #245.")) + "uses by default, or reduce `n_integrators` to 1. See GeometricMachineLearning issue #245, " * + "and SymbolicNeuralNetworks issues #49 and #54 for the upstream construction that will " * + "eventually lift this.")) nothing end diff --git a/test/generalized_hamiltonian_neural_networks/parametric_layers_and_architectures_test.jl b/test/generalized_hamiltonian_neural_networks/parametric_layers_and_architectures_test.jl index eaa669edd..c07f16dd4 100644 --- a/test/generalized_hamiltonian_neural_networks/parametric_layers_and_architectures_test.jl +++ b/test/generalized_hamiltonian_neural_networks/parametric_layers_and_architectures_test.jl @@ -5,7 +5,8 @@ using GeometricMachineLearning using GeometricMachineLearning: ForcingLayerQ, ForcingLayerP, ForcingLayerQP, - ParametricResNetLayer, WideResNetLayer, ParametricResNet + ParametricResNetLayer, WideResNetLayer, ParametricResNet, + apply_parametric using AbstractNeuralNetworks: params using Random: seed! using Test @@ -100,7 +101,7 @@ end @testset "ParametricResNet" begin arch = ParametricResNet(DIM; width = WIDTH, n_blocks = 2, parameters = SYSTEM_PARAMETERS) nn = NeuralNetwork(arch) - out = nn.model(rand(DIM), SYSTEM_PARAMETERS, params(nn)) + out = apply_parametric(nn.model, rand(DIM), SYSTEM_PARAMETERS, params(nn)) @test size(out) == (DIM,) @test finite(out)