Hand the structured types' backend and section methods back to their owner

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

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"

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`.
 #

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)

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)
+

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

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

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}