From 95ff91660120738f27a559618fe9799d5b3b668e Mon Sep 17 00:00:00 2001 From: Michael Kraus Date: Mon, 24 Aug 2026 14:23:40 +0900 Subject: [PATCH] Hand the structured types' backend and section methods back to their owner Both were piracy of the same shape: `changebackend` is `AbstractNeuralNetworks`' generic and `GlobalSection` is `GeometricOptimizers`', while `StiefelManifold`, `SymmetricMatrix` and the rest are `GeometricOptimizers`' types and `NetworkParameters` is `NeuralNetworkParameters`'. This package owns neither side of either pair. `GeometricOptimizers` 0.5 defines both, so they go from here. The `changebackend` methods also sat inside the HDF5 extension, which has nothing to do with them: `changebackend(GPU(), nn)` on a network with a manifold weight was a `MethodError` unless HDF5 happened to be loaded. Upstream covers the two horizontal lifts as well, which were missing here. `test/hdf5_support.jl` is untouched on purpose -- it imports `changebackend` from `AbstractNeuralNetworks`, so upstream's methods answer by dispatch and the testsets keep their meaning. Nothing under `src/` referenced `changebackend` at all. The compat floor moves to `GeometricOptimizers = "0.5"` accordingly, which is what makes this commit depend on that release rather than merely prefer it. --- CHANGELOG.md | 15 ++ Project.toml | 2 +- ext/HDF5Ext.jl | 35 +---- scripts/enzyme.jl | 39 ++++++ scripts/harmonic_oscillator_sympnet.jl | 58 ++++++++ scripts/sae_script2.jl | 181 +++++++++++++++++++++++++ scripts/zygote.jl | 20 +++ src/optimizers/optimizer.jl | 4 - 8 files changed, 315 insertions(+), 39 deletions(-) create mode 100644 scripts/enzyme.jl create mode 100644 scripts/harmonic_oscillator_sympnet.jl create mode 100644 scripts/sae_script2.jl create mode 100644 scripts/zygote.jl diff --git a/CHANGELOG.md b/CHANGELOG.md index 6cfdd0fe1..5e644e2e7 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -24,8 +24,23 @@ remaining walks. `map_to_cpu` loses six per-type methods, `apply_toNT` turns out **It also makes the optimizer cache immune to a change coming in `GeometricOptimizers`**, which is the half of this release with no visible effect today — see *Fixed*. +**Requires `GeometricOptimizers` 0.5**, which is where the structured types' `changebackend` and +`GlobalSection` methods now live. + ### Removed (breaking) +- **The five `changebackend` methods for `GeometricOptimizers`' types are gone, and so is + `GeometricOptimizers.GlobalSection(::NetworkParameters)`.** Both were type piracy of the same shape: + the generic belongs to one package, the types to another, and this package owns neither. Both now + live in `GeometricOptimizers` 0.5, which is what the compat floor moves for. + + The `changebackend` methods also sat inside the **HDF5 extension**, which had nothing to do with + HDF5 — so `changebackend(GPU(), nn)` on a network with a manifold weight was a `MethodError` unless + HDF5 happened to be loaded. Upstream covers the horizontal lifts too, which were missing here. + + `test/hdf5_support.jl` needs no edit: it imports `changebackend` from `AbstractNeuralNetworks`, and + upstream's methods answer by dispatch. Nothing under `src/` referenced `changebackend` at all. + - **`apply_toNT` is gone from the export list and from the package. It was `Base.map`.** ```julia diff --git a/Project.toml b/Project.toml index 462957bc8..6832e4aaf 100644 --- a/Project.toml +++ b/Project.toml @@ -39,7 +39,7 @@ ForwardDiff = "1" GeometricBase = "0.14" GeometricEquations = "0.21" GeometricIntegrators = "0.18.2" -GeometricOptimizers = "0.4.1" +GeometricOptimizers = "0.5" GeometricSolutions = "0.6" HDF5 = "0.16, 0.17" KernelAbstractions = "0.9" diff --git a/ext/HDF5Ext.jl b/ext/HDF5Ext.jl index 0853d4257..2d74621d0 100644 --- a/ext/HDF5Ext.jl +++ b/ext/HDF5Ext.jl @@ -2,44 +2,11 @@ module HDF5Ext using HDF5 using GeometricMachineLearning -import AbstractNeuralNetworks: changebackend, NeuralNetworkBackend, Architecture +import AbstractNeuralNetworks: NeuralNetworkBackend, Architecture # `save`, `load`, `params` and the parameter container are `NeuralNetworkParameters`' as of # `AbstractNeuralNetworks` 0.7, which only re-binds them; reach for them where they are defined. import NeuralNetworkParameters: NetworkParameters, params, save, load -# --------------------------------------------------------------------------- -# changebackend — new methods for GML special array types -# -# AbstractNeuralNetworks.changebackend handles AbstractArray and NamedTuple. -# Moving a NeuralNetwork between devices fails for parameters that include -# StiefelManifold, SymmetricMatrix, or SkewSymMatrix without these methods. -# -# `changebackend` is `AbstractNeuralNetworks`' and the types are `GeometricOptimizers`', so these -# methods are piracy the same way the `h5save` ones were before `GeometricOptimizers` took over the -# leaf protocol. They belong in a `GeometricOptimizers` extension on `AbstractNeuralNetworks`; that -# is a separate change with its own release chain, so they stay here for now. -# --------------------------------------------------------------------------- - -function changebackend(backend::NeuralNetworkBackend, Y::StiefelManifold) - StiefelManifold(changebackend(backend, Y.A)) -end - -function changebackend(backend::NeuralNetworkBackend, A::SymmetricMatrix) - SymmetricMatrix(changebackend(backend, A.S), A.n) -end - -function changebackend(backend::NeuralNetworkBackend, A::SkewSymMatrix) - SkewSymMatrix(changebackend(backend, A.S), A.n) -end - -function changebackend(backend::NeuralNetworkBackend, A::LowerTriangular) - LowerTriangular(changebackend(backend, A.S), A.n) -end - -function changebackend(backend::NeuralNetworkBackend, A::UpperTriangular) - UpperTriangular(changebackend(backend, A.S), A.n) -end - # --------------------------------------------------------------------------- # save / load — the entry points that dispatch on this package's `NeuralNetwork`. # diff --git a/scripts/enzyme.jl b/scripts/enzyme.jl new file mode 100644 index 000000000..78612cd4a --- /dev/null +++ b/scripts/enzyme.jl @@ -0,0 +1,39 @@ +using BenchmarkTools +using Enzyme +using LinearAlgebra +using Zygote + + +loss(A,x) = norm(A*x) + +# function loss(A,x) +# y = zero(x) +# mul!(y,A,x) +# norm(y) +# end + +function test(n) + A = rand(n,n) + x = rand(n) + + l = a -> loss(a,x) + + dA = zero(A) + + println("\nn = $n") + + println("\nEnzyme (autodiff):") + @btime Enzyme.autodiff(Reverse, $l, Active, Duplicated($A, $dA)) + + println("\nEnzyme (gradient):") + @btime Enzyme.gradient(Reverse, $l, $A) + + println("\nZygote:") + @btime Zygote.gradient($l, $A)[1] + + println("") +end + +test(100) +test(1000) +test(10000) diff --git a/scripts/harmonic_oscillator_sympnet.jl b/scripts/harmonic_oscillator_sympnet.jl new file mode 100644 index 000000000..4f48ec5d6 --- /dev/null +++ b/scripts/harmonic_oscillator_sympnet.jl @@ -0,0 +1,58 @@ +using GeometricMachineLearning +using GeometricIntegrators: ImplicitMidpoint, integrate +import GeometricProblems.HarmonicOscillator as ho + +# the problem is the ODE of the harmonic oscillator +ho_problem = ho.hodeproblem(; tspan = 500) + +# integrate the system +solution = integrate(ho_problem, ImplicitMidpoint()) + +dl_raw = DataLoader(solution; suppress_info = true) + +# specify the data type and the backend +type = Float64 +backend = CPU() + +# we can then make a new instance of `DataLoader` with this backend and type. +dl = DataLoader(dl_raw, backend, type) + + +const upscaling_dimension = 2 +const nhidden = 1 +const activation = tanh +const n_layers = 4 # number of layers for the G-SympNet +const depth = 4 # number of layers in each linear block in the LA-SympNet + +# calling G-SympNet architecture +gsympnet = GSympNet(dl; upscaling_dimension = upscaling_dimension, + n_layers = n_layers, + activation = activation) + +# initialize the networks +g_nn = NeuralNetwork(gsympnet, backend, type) + +# set up optimizer; for this we first need to specify the optimization method +opt_method = AdamOptimizer(type) + +# we then call the optimizer struct which allocates the cache +g_opt = Optimizer(opt_method, g_nn) + +# determine the batch size (the number of samples in one batch) +const batch_size = 16 + +batch = Batch(batch_size) + +# number of training epochs +const nepochs = 100 + +# perform training (returns array that contains the total loss for each training step) +g_loss_array = g_opt(g_nn, dl, batch, nepochs; show_progress = false) + +ics = (q=dl.input.q[:, 1, 1], p=dl.input.p[:, 1, 1]) + +steps_to_plot = 1000 + +#predictions +g_trajectory = iterate(g_nn, ics; n_points = steps_to_plot) + diff --git a/scripts/sae_script2.jl b/scripts/sae_script2.jl new file mode 100644 index 000000000..902a8db61 --- /dev/null +++ b/scripts/sae_script2.jl @@ -0,0 +1,181 @@ +using GeometricIntegrators: integrate, ImplicitMidpoint +using GeometricMachineLearning +import Random # hide +import GeometricProblems.TodaLattice as tl +using JLD2 +using CairoMakie + +sae_dir = "animations" +mkpath(sae_dir) + +N = tl.Ñ # hide +Δx = 1. / (N - 1) # hide +Ω = -0.5 : Δx : 0.5 # hide +tl.μ + +# todo +#pr = tl.hodeproblem(; tspan = (0.0, 8.)) +pr = tl.hodeproblem(; tspan = (0.0, 800.)) +@time "FOM + Implicit Midpoint" sol = integrate(pr, ImplicitMidpoint()) + +dl_cpu = DataLoader(sol; autoencoder = true, suppress_info = true) + +const reduced_dim = 2 + +Random.seed!(123) # hide +sae_arch = SymplecticAutoencoder(dl_cpu.input_dim, reduced_dim; n_encoder_blocks = 4, + n_decoder_blocks = 4, + n_encoder_layers = 2, + n_decoder_layers = 2) + +const mtc = GeometricMachineLearning.map_to_cpu + +sae_trained_parameters = load("../docs/src/tutorials/sae_parameters.jld2")["sae_parameters"] +_nnp(ps::Tuple) = NeuralNetworkParameters{Tuple(Symbol("L$(i)") for i in 1:length(ps))}(ps) +sae_nn_cpu = NeuralNetwork(sae_arch, Chain(sae_arch), _nnp(sae_trained_parameters), CPU()) + +sae_rs = HRedSys(pr, encoder(sae_nn_cpu), decoder(sae_nn_cpu); integrator = ImplicitMidpoint()) + +# @time "FOM + Implicit Midpoint" sol_full = integrate_full_system(sae_rs) # hide +@time "SAE + Implicit Midpoint" sol_sae_reduced = integrate_reduced_system(sae_rs) # hide + + + + + +const T = Float32 +_T(qp::NamedTuple{(:q, :p)}) = (q = T.(qp.q), p = T.(qp.p)) + +dl_reduced = DataLoader(encoder(sae_nn_cpu)(_T(dl_cpu.input))) + +# lines(dl_reduced.input.q[1, :, 1], dl_reduced.input.p[1, :, 1]) + +# sympnet_arch = GSympNet(2; n_layers = 10) +# sympnet_nn = NeuralNetwork(sympnet_arch, T) +# o = Optimizer(AdamOptimizer(), sympnet_nn) +# o(sympnet_nn, dl_reduced, Batch(10), 500) + +morange = RGBf(255 / 256, 127 / 256, 14 / 256) +mred = RGBf(214 / 256, 39 / 256, 40 / 256) +mpurple = RGBf(148 / 256, 103 / 256, 189 / 256) +mblue = RGBf(31 / 256, 119 / 256, 180 / 256) +mgreen = RGBf(44 / 256, 160 / 256, 44 / 256) + +function plot_solution(time_step; theme = :light, framerate = 50) + textcolor = theme == :dark ? :white : :black + fig = Figure(size = (1000, 500), figure_padding = (5,50,5,10), fontsize = 24) + ax = Axis(fig[1, 1], backgroundcolor = :transparent, + bottomspinecolor = textcolor, + topspinecolor = textcolor, + leftspinecolor = textcolor, + rightspinecolor = textcolor, + xtickcolor = textcolor, + ytickcolor = textcolor, + xticklabelcolor = textcolor, + yticklabelcolor = textcolor, + xlabel=L"\omega", + ylabel=L"q", + xlabelcolor = textcolor, + ylabelcolor = textcolor) + lines!(ax, sol.s.q[time_step], label = rich("FOM + Implicit Midpoint"; color = textcolor), color = mblue) + lines!(ax, sae_rs.decoder((q = sol_sae_reduced.s.q[time_step], p = sol_sae_reduced.s.p[time_step])).q, label = rich("SAE + Implicit Midpoint"; color = textcolor), color = mgreen) + axislegend(ax; position = :rt) + xlims!(ax, 0, 200) + ylims!(ax, 0, 1) + fig +end + + +#### Transformer + +const seq_length = 4 +integrator_architecture = StandardTransformerIntegrator(reduced_dim; + transformer_dim = 20, + n_blocks = 3, + n_heads = 5, + L = 3, + upscaling_activation = tanh) + +nn_integrator_parameters = load("../docs/src/tutorials/integrator_parameters.jld2")["integrator_parameters"] # hide +integrator_nn = NeuralNetwork(integrator_architecture, Chain(integrator_architecture), _nnp(nn_integrator_parameters), CPU()) # hide + +# todo +#n_time_steps = 100 +n_time_steps = 10000 + +ics = (q = dl_reduced.input.q[:, 1:seq_length], p = dl_reduced.input.p[:, 1:seq_length]) +time_series = iterate(mtc(integrator_nn), ics; n_points = n_time_steps, prediction_window = seq_length) +function plot_solution2(time_step; theme = :light, framerate = 50) + textcolor = theme == :dark ? :white : :black + fig = Figure(size = (1000, 500), figure_padding = (5,50,5,10), fontsize = 24) + ax = Axis(fig[1, 1], backgroundcolor = :transparent, + bottomspinecolor = textcolor, + topspinecolor = textcolor, + leftspinecolor = textcolor, + rightspinecolor = textcolor, + xtickcolor = textcolor, + ytickcolor = textcolor, + xticklabelcolor = textcolor, + yticklabelcolor = textcolor, + xlabel=L"\omega", + ylabel=L"q", + xlabelcolor = textcolor, + ylabelcolor = textcolor) + lines!(ax, sol.s.q[time_step], label = rich("FOM + Implicit Midpoint"; color = textcolor), color = mblue) + # prediction = (q = time_series.q[:, end], p = time_series.p[:, end]) + prediction = (q = time_series.q[:, time_step], p = time_series.p[:, time_step]) + prediction_big = decoder(sae_nn_cpu)(prediction) + + lines!(ax, prediction_big.q; label = rich("SAE + Transformer"; color = textcolor), color = mpurple) + axislegend(ax; position = :rt) + xlims!(ax, 0, 200) + ylims!(ax, 0, 1) + fig +end + +# ics3 = (q = ics.q[:, 1], p = ics.p[:, 1]) +# +# time_series2 = iterate(sympnet_nn, ics3; n_points = n_time_steps) +# +# function plot_solution3(time_step; theme = :light, framerate = 50) +# textcolor = theme == :dark ? :white : :black +# fig = Figure(size = (1000, 500), figure_padding = (5,50,5,5), fontsize = 24) +# ax = Axis(fig[1, 1], backgroundcolor = :transparent, +# bottomspinecolor = textcolor, +# topspinecolor = textcolor, +# leftspinecolor = textcolor, +# rightspinecolor = textcolor, +# xtickcolor = textcolor, +# ytickcolor = textcolor, +# xticklabelcolor = textcolor, +# yticklabelcolor = textcolor, +# xlabel=L"\omega", +# ylabel=L"q", +# xlabelcolor = textcolor, +# ylabelcolor = textcolor) +# lines!(ax, sol.s.q[time_step], label = rich("FOM + Implicit Midpoint"; color = textcolor), color = mblue) +# time_series = iterate(sympnet_nn, ics3; n_points = time_step) +# # prediction = (q = time_series.q[:, end], p = time_series.p[:, end]) +# prediction = (q = time_series2.q[:, time_step], p = time_series2.p[:, time_step]) +# prediction_big = decoder(sae_nn_cpu)(prediction) +# +# lines!(ax, prediction_big.q; label = rich("SAE + SympNet"; color = textcolor), color = mpurple) +# axislegend(ax; position = :rt) +# xlims!(ax, 0, 200) +# fig +# end + +# todo +#time_steps = 1:5 # axes(time_series.q, 2) +time_steps = 1:500 # axes(time_series.q, 2) + +for time_step in time_steps + fig1 = plot_solution(time_step) + save(sae_dir * "/sae-midpoint-$(string(time_step, pad = 3)).pdf", fig1) + + fig2 = plot_solution2(time_step) + save(sae_dir * "/sae-transformer-$(string(time_step, pad = 3)).pdf", fig2) + +# fig3 = plot_solution3(time_step) +# save(sae_dir * "/sae-sympnet-$(string(time_step, pad = 3)).pdf", fig3) +end diff --git a/scripts/zygote.jl b/scripts/zygote.jl new file mode 100644 index 000000000..cde81e962 --- /dev/null +++ b/scripts/zygote.jl @@ -0,0 +1,20 @@ +using Zygote, Printf, LinearAlgebra + +const number_data_points = 1000 + +const data_input = [[i] for i in 1:number_data_points] + +function_to_be_differentiated(input, A) = norm(A*input) + +function gradient_eval(data, num, A = rand(100000,1)) + input = data[num] + @printf "First one: " + @time Zygote.gradient(A -> function_to_be_differentiated(input, A), A)[1] + @printf "Second one:" + @time Zygote.gradient(A -> function_to_be_differentiated(data[num], A), A)[1] + @printf "\n" +end + +for i in 1:5 + gradient_eval(data_input, Int(ceil(rand()*number_data_points))) +end diff --git a/src/optimizers/optimizer.jl b/src/optimizers/optimizer.jl index c41f5619e..171b9a2c0 100644 --- a/src/optimizers/optimizer.jl +++ b/src/optimizers/optimizer.jl @@ -1,10 +1,6 @@ # Optimizer machinery on top of GeometricOptimizers. # Kept out of `utils.jl` because it dispatches on `Manifold`, which is defined later. -# Extend GlobalSection so it works with NetworkParameters (wraps a NamedTuple). -GeometricOptimizers.GlobalSection(ps::NetworkParameters) = - GeometricOptimizers.GlobalSection(params(ps)) - # Backward-compat alias const AbstractCache{T} = GeometricOptimizers.OptimizerCache{T}