From 8368761320ec06cc6595d4f1febddb3a48532355 Mon Sep 17 00:00:00 2001 From: benedict-96 Date: Sun, 8 Dec 2024 19:27:43 +0100 Subject: [PATCH 1/5] Added test for reshape and create_array routine. --- src/SymbolicNeuralNetworks.jl | 2 ++ src/utils/build_function.jl | 11 ++++++++-- src/utils/build_function2.jl | 12 +++++++++-- src/utils/create_array.jl | 3 +++ test/reshape_test.jl | 40 +++++++++++++++++++++++++++++++++++ test/runtests.jl | 9 ++++---- 6 files changed, 69 insertions(+), 8 deletions(-) create mode 100644 src/utils/create_array.jl create mode 100644 test/reshape_test.jl diff --git a/src/SymbolicNeuralNetworks.jl b/src/SymbolicNeuralNetworks.jl index c7d8294..898f843 100644 --- a/src/SymbolicNeuralNetworks.jl +++ b/src/SymbolicNeuralNetworks.jl @@ -20,6 +20,8 @@ module SymbolicNeuralNetworks export symbolize include("utils/symbolize.jl") + include("utils/create_array.jl") + export AbstractSymbolicNeuralNetwork export SymbolicNeuralNetwork, SymbolicModel export HamiltonianSymbolicNeuralNetwork, HNNLoss diff --git a/src/utils/build_function.jl b/src/utils/build_function.jl index 33d5469..c33baae 100644 --- a/src/utils/build_function.jl +++ b/src/utils/build_function.jl @@ -25,9 +25,16 @@ end function build_nn_function(eq::EqT, sparams::NeuralNetworkParameters, sinput::Symbolics.Arr) gen_fun = _build_nn_function(eq, sparams, sinput) gen_fun_returned(x, ps) = mapreduce(k -> gen_fun(x, ps, k), hcat, axes(x, 2)) - gen_fun_returned(x::Union{AbstractVector, Symbolics.Arr}, ps) = gen_fun_returned(reshape(x, length(x), 1), ps) + function gen_fun_returned(x::Union{AbstractVector, Symbolics.Arr}, ps) + output_not_reshaped = gen_fun_returned(reshape(x, length(x), 1), ps) + # for vectors we do not reshape, as the output may be a matrix + output_not_reshaped + end # check this! (definitely not correct in all cases!) - gen_fun_returned(x::AbstractArray{<:Number, 3}, ps) = reshape(gen_fun_returned(reshape(x, size(x, 1), size(x, 2) * size(x, 3)), ps), size(x, 1), size(x, 2), size(x, 3)) + function gen_fun_returned(x::AbstractArray{<:Number, 3}, ps) + output_not_reshaped = gen_fun_returned(reshape(x, size(x, 1), size(x, 2) * size(x, 3)), ps) + reshape(output_not_reshaped, size(output_not_reshaped, 1), size(x, 2), size(x, 3)) + end gen_fun_returned end diff --git a/src/utils/build_function2.jl b/src/utils/build_function2.jl index 77e1853..86d5f92 100644 --- a/src/utils/build_function2.jl +++ b/src/utils/build_function2.jl @@ -19,8 +19,16 @@ end function build_nn_function(eq::EqT, sparams::NeuralNetworkParameters, sinput::Symbolics.Arr, soutput::Symbolics.Arr) gen_fun = _build_nn_function(eq, sparams, sinput, soutput) gen_fun_returned(input, output, ps) = mapreduce(k -> gen_fun(input, output, ps, k), +, axes(input, 2)) - gen_fun_returned(input::AT, output::AT, ps) where {AT <: Union{AbstractVector, Symbolics.Arr}} = gen_fun_returned(reshape(input, length(input), 1), reshape(output, length(output), 1), ps) - gen_fun_returned(input::AT, output::AT, ps) where {T, AT <: AbstractArray{T, 3}} = gen_fun_returned(reshape(input, size(input, 1), size(input, 2) * size(input, 3)), reshape(output, size(output, 1), size(output, 2) * size(output, 3)), ps) + function gen_fun_returned(x::AT, y::AT, ps) where {AT <: Union{AbstractVector, Symbolics.Arr}} + output_not_reshaped = gen_fun_returned(reshape(x, length(x), 1), reshape(y, length(y), 1), ps) + # for vectors we do not reshape, as the output may be a matrix + output_not_reshaped + end + # check this! (definitely not correct in all cases!) + function gen_fun_returned(x::AT, y::AT, ps) where {AT <: AbstractArray{<:Number, 3}} + output_not_reshaped = gen_fun_returned(reshape(x, size(x, 1), size(x, 2) * size(x, 3)), reshape(y, size(y, 1), size(y, 2) * size(y, 3)), ps) + reshape(output_not_reshaped, size(output_not_reshaped, 1), size(x, 2), size(x, 3)) + end gen_fun_returned end diff --git a/src/utils/create_array.jl b/src/utils/create_array.jl new file mode 100644 index 0000000..c016f0c --- /dev/null +++ b/src/utils/create_array.jl @@ -0,0 +1,3 @@ +function Symbolics.SymbolicUtils.Code.create_array(::Type{<:Base.ReshapedArray{T, N, P}}, S, nd::Val, d::Val, elems...) where {T, N, P} + Symbolics.SymbolicUtils.Code.create_array(P, S, nd, d, elems...) +end \ No newline at end of file diff --git a/test/reshape_test.jl b/test/reshape_test.jl new file mode 100644 index 0000000..25bd340 --- /dev/null +++ b/test/reshape_test.jl @@ -0,0 +1,40 @@ +using SymbolicNeuralNetworks +using AbstractNeuralNetworks +using Symbolics +using Test + +function set_up_network() + c = Chain(Dense(2, 3)) + nn = SymbolicNeuralNetwork(c) + soutput = nn.model(nn.input, nn.params) + nn_cpu = NeuralNetwork(c) + nn, soutput, nn_cpu +end + +function test_for_input() + nn, soutput, nn_cpu = set_up_network() + input = rand(2, 5) + input2 = reshape((@view input[:, 1:2]), 2, 1, 2) + built_function = build_nn_function(soutput, nn.params, nn.input) + outputs = built_function(input2, nn_cpu.params) + for i in 1:2 + @test nn.model(input[:, i], nn_cpu.params) ≈ outputs[:, 1, i] + end +end + +function test_for_input_and_output() + nn, soutput2, nn_cpu = set_up_network() + input = rand(2, 5) + output = rand(3, 5) + input2 = reshape((@view input[:, 1:2]), 2, 1, 2) + output2 = reshape((@view input[:, 1:3]), 3, 1, 2) + @variables soutput[1:3] + built_function = build_nn_function((soutput - soutput2).^2, nn.params, nn.input, soutput) + outputs = built_function(input2, output2, nn_cpu.params) + for i in 1:2 + @test (nn.model(input[:, i], nn_cpu.params) - output[:, i]).^2 ≈ outputs[:, 1, i] + end +end + +test_for_input() +test_for_input_and_output() \ No newline at end of file diff --git a/test/runtests.jl b/test/runtests.jl index 8daf81c..833b34b 100644 --- a/test/runtests.jl +++ b/test/runtests.jl @@ -3,7 +3,8 @@ using SafeTestsets using Test @safetestset "Docstring tests. " begin include("doctest.jl") end -@safetestset "Symbolic gradient " begin include("symbolic_gradient.jl") end -@safetestset "Symbolic Neural network " begin include("neural_network_derivative.jl") end -@safetestset "Symbolic Params " begin include("test_params.jl") end -# @safetestset "HNN Loss " begin include("test_hnn_loss_pullback.jl") end \ No newline at end of file +@safetestset "Symbolic gradient " begin include("symbolic_gradient.jl") end +@safetestset "Symbolic Neural network " begin include("neural_network_derivative.jl") end +@safetestset "Symbolic Params " begin include("test_params.jl") end +# @safetestset "HNN Loss " begin include("test_hnn_loss_pullback.jl") end +@safetestset "Check if reshape works in the correct way with the generated functions. " begin include("reshape_test.jl") end \ No newline at end of file From d68b2e812b23ac824fc5a519ed7cce6373c93a10 Mon Sep 17 00:00:00 2001 From: benedict-96 Date: Sun, 8 Dec 2024 20:02:53 +0100 Subject: [PATCH 2/5] + -> hcat (was a typo\!) --- src/utils/build_function2.jl | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/utils/build_function2.jl b/src/utils/build_function2.jl index 86d5f92..1f9f3cd 100644 --- a/src/utils/build_function2.jl +++ b/src/utils/build_function2.jl @@ -18,7 +18,7 @@ end function build_nn_function(eq::EqT, sparams::NeuralNetworkParameters, sinput::Symbolics.Arr, soutput::Symbolics.Arr) gen_fun = _build_nn_function(eq, sparams, sinput, soutput) - gen_fun_returned(input, output, ps) = mapreduce(k -> gen_fun(input, output, ps, k), +, axes(input, 2)) + gen_fun_returned(input, output, ps) = mapreduce(k -> gen_fun(input, output, ps, k), hcat, axes(input, 2)) function gen_fun_returned(x::AT, y::AT, ps) where {AT <: Union{AbstractVector, Symbolics.Arr}} output_not_reshaped = gen_fun_returned(reshape(x, length(x), 1), reshape(y, length(y), 1), ps) # for vectors we do not reshape, as the output may be a matrix From 7d09283f19e9acd31f704f289ed99accd34f5f60 Mon Sep 17 00:00:00 2001 From: benedict-96 Date: Sun, 8 Dec 2024 20:03:05 +0100 Subject: [PATCH 3/5] Fixed typo. --- test/reshape_test.jl | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/test/reshape_test.jl b/test/reshape_test.jl index 25bd340..fde5766 100644 --- a/test/reshape_test.jl +++ b/test/reshape_test.jl @@ -27,7 +27,7 @@ function test_for_input_and_output() input = rand(2, 5) output = rand(3, 5) input2 = reshape((@view input[:, 1:2]), 2, 1, 2) - output2 = reshape((@view input[:, 1:3]), 3, 1, 2) + output2 = reshape((@view output[:, 1:2]), 3, 1, 2) @variables soutput[1:3] built_function = build_nn_function((soutput - soutput2).^2, nn.params, nn.input, soutput) outputs = built_function(input2, output2, nn_cpu.params) From 1e3a10e92c3876a06dc4726176f5b5b7bb49c0a3 Mon Sep 17 00:00:00 2001 From: benedict-96 Date: Sun, 8 Dec 2024 23:57:35 +0100 Subject: [PATCH 4/5] Reduce has to be optional (depending on whether we compute derivatives of neural network parameters or other expressions. --- src/pullback.jl | 2 +- src/utils/build_function.jl | 4 ++-- src/utils/build_function2.jl | 15 ++++++++++++--- src/utils/build_function_arrays.jl | 16 ++++++++-------- 4 files changed, 23 insertions(+), 14 deletions(-) diff --git a/src/pullback.jl b/src/pullback.jl index 31da2b3..5ba4050 100644 --- a/src/pullback.jl +++ b/src/pullback.jl @@ -93,7 +93,7 @@ function SymbolicPullback(nn::SymbolicNeuralNetwork, loss::NetworkLoss) @variables soutput[1:output_dimension(nn.model)] symbolic_loss = loss(nn.model, nn.params, nn.input, soutput) symbolic_pullbacks = symbolic_pullback(symbolic_loss, nn) - pbs_executable = build_nn_function(symbolic_pullbacks, nn.params, nn.input, soutput) + pbs_executable = build_nn_function(symbolic_pullbacks, nn.params, nn.input, soutput; reduce = +) function pbs(input, output, params) pullback(::Union{Real, AbstractArray{<:Real}}) = _get_contents(_get_params(pbs_executable(input, output, params))) pullback diff --git a/src/utils/build_function.jl b/src/utils/build_function.jl index c33baae..d7dfcc6 100644 --- a/src/utils/build_function.jl +++ b/src/utils/build_function.jl @@ -22,9 +22,9 @@ function build_nn_function(eq::EqT, nn::AbstractSymbolicNeuralNetwork) build_nn_function(eq, nn.params, nn.input) end -function build_nn_function(eq::EqT, sparams::NeuralNetworkParameters, sinput::Symbolics.Arr) +function build_nn_function(eq::EqT, sparams::NeuralNetworkParameters, sinput::Symbolics.Arr; reduce = hcat) gen_fun = _build_nn_function(eq, sparams, sinput) - gen_fun_returned(x, ps) = mapreduce(k -> gen_fun(x, ps, k), hcat, axes(x, 2)) + gen_fun_returned(x, ps) = mapreduce(k -> gen_fun(x, ps, k), reduce, axes(x, 2)) function gen_fun_returned(x::Union{AbstractVector, Symbolics.Arr}, ps) output_not_reshaped = gen_fun_returned(reshape(x, length(x), 1), ps) # for vectors we do not reshape, as the output may be a matrix diff --git a/src/utils/build_function2.jl b/src/utils/build_function2.jl index 1f9f3cd..8989218 100644 --- a/src/utils/build_function2.jl +++ b/src/utils/build_function2.jl @@ -16,9 +16,9 @@ function build_nn_function(eqs, nn::AbstractSymbolicNeuralNetwork, soutput) build_nn_function(eqs, nn.params, nn.input, soutput) end -function build_nn_function(eq::EqT, sparams::NeuralNetworkParameters, sinput::Symbolics.Arr, soutput::Symbolics.Arr) +function build_nn_function(eq::EqT, sparams::NeuralNetworkParameters, sinput::Symbolics.Arr, soutput::Symbolics.Arr; reduce = hcat) gen_fun = _build_nn_function(eq, sparams, sinput, soutput) - gen_fun_returned(input, output, ps) = mapreduce(k -> gen_fun(input, output, ps, k), hcat, axes(input, 2)) + gen_fun_returned(input, output, ps) = mapreduce(k -> gen_fun(input, output, ps, k), reduce, axes(input, 2)) function gen_fun_returned(x::AT, y::AT, ps) where {AT <: Union{AbstractVector, Symbolics.Arr}} output_not_reshaped = gen_fun_returned(reshape(x, length(x), 1), reshape(y, length(y), 1), ps) # for vectors we do not reshape, as the output may be a matrix @@ -27,11 +27,20 @@ function build_nn_function(eq::EqT, sparams::NeuralNetworkParameters, sinput::Sy # check this! (definitely not correct in all cases!) function gen_fun_returned(x::AT, y::AT, ps) where {AT <: AbstractArray{<:Number, 3}} output_not_reshaped = gen_fun_returned(reshape(x, size(x, 1), size(x, 2) * size(x, 3)), reshape(y, size(y, 1), size(y, 2) * size(y, 3)), ps) - reshape(output_not_reshaped, size(output_not_reshaped, 1), size(x, 2), size(x, 3)) + # if arrays are added together then don't reshape! + optional_reshape(output_not_reshaped, reduce, x) end gen_fun_returned end +function optional_reshape(output_not_reshaped::AbstractVecOrMat, ::typeof(+), ::AbstractArray{<:Number, 3}) + output_not_reshaped +end + +function optional_reshape(output_not_reshaped::AbstractVecOrMat, ::typeof(hcat), input::AbstractArray{<:Number, 3}) + reshape(output_not_reshaped, size(output_not_reshaped, 1), size(input, 2), size(input, 3)) +end + """ _build_nn_function(eq, params, sinput, soutput) diff --git a/src/utils/build_function_arrays.jl b/src/utils/build_function_arrays.jl index d53a6ee..6c5f2b9 100644 --- a/src/utils/build_function_arrays.jl +++ b/src/utils/build_function_arrays.jl @@ -29,8 +29,8 @@ funcs_evaluated = funcs(input, ps) (true, true, true) ``` """ -function build_nn_function(eqs::AbstractArray{<:Union{NamedTuple, NeuralNetworkParameters}}, sparams::NeuralNetworkParameters, sinput::Symbolics.Arr...) - ps_semi = [function_valued_parameters(eq, sparams, sinput...) for eq in eqs] +function build_nn_function(eqs::AbstractArray{<:Union{NamedTuple, NeuralNetworkParameters}}, sparams::NeuralNetworkParameters, sinput::Symbolics.Arr...; reduce = hcat) + ps_semi = [function_valued_parameters(eq, sparams, sinput...; reduce = reduce) for eq in eqs] _pbs_executable(ps_functions, params, input...) = apply_element_wise(ps_functions, params, input...) __pbs_executable(input, params) = _pbs_executable(ps_semi, params, input) @@ -72,8 +72,8 @@ funcs_evaluated = funcs(input, ps) Internally this is using [`function_valued_parameters`](@ref) and [`apply_element_wise`](@ref). """ -function build_nn_function(eqs::Union{NamedTuple, NeuralNetworkParameters}, sparams::NeuralNetworkParameters, sinput::Symbolics.Arr...) - ps = function_valued_parameters(eqs, sparams, sinput...) +function build_nn_function(eqs::Union{NamedTuple, NeuralNetworkParameters}, sparams::NeuralNetworkParameters, sinput::Symbolics.Arr...; reduce = hcat) + ps = function_valued_parameters(eqs, sparams, sinput...; reduce = reduce) _pbs_executable(ps::Union{NamedTuple, NeuralNetworkParameters}, params::NeuralNetworkParameters, input::AbstractArray...) = apply_element_wise(ps, params, input...) __pbs_executable(input::AbstractArray, params::NeuralNetworkParameters) = _pbs_executable(ps, params, input) # return this one if sinput & soutput are supplied @@ -110,13 +110,13 @@ b = c(input, ps).^2 (true, true) ``` """ -function function_valued_parameters(eqs::NeuralNetworkParameters, sparams::NeuralNetworkParameters, sinput::Symbolics.Arr...) - vals = Tuple(build_nn_function(eqs[key], sparams, sinput...) for key in keys(eqs)) +function function_valued_parameters(eqs::NeuralNetworkParameters, sparams::NeuralNetworkParameters, sinput::Symbolics.Arr...; reduce = hcat) + vals = Tuple(build_nn_function(eqs[key], sparams, sinput...; reduce = reduce) for key in keys(eqs)) NeuralNetworkParameters{keys(eqs)}(vals) end -function function_valued_parameters(eqs::NamedTuple, sparams::NeuralNetworkParameters, sinput::Symbolics.Arr...) - vals = Tuple(build_nn_function(eqs[key], sparams, sinput...) for key in keys(eqs)) +function function_valued_parameters(eqs::NamedTuple, sparams::NeuralNetworkParameters, sinput::Symbolics.Arr...; reduce = hcat) + vals = Tuple(build_nn_function(eqs[key], sparams, sinput...; reduce = reduce) for key in keys(eqs)) NamedTuple{keys(eqs)}(vals) end From 0ed585f969e22677d1ac88437e797ce9f1e07650 Mon Sep 17 00:00:00 2001 From: benedict-96 Date: Mon, 16 Dec 2024 09:49:56 +0100 Subject: [PATCH 5/5] Removed SHNN file. This will be included in GML. --- docs/src/hamiltonian_neural_network.md | 98 -------------------------- 1 file changed, 98 deletions(-) delete mode 100644 docs/src/hamiltonian_neural_network.md diff --git a/docs/src/hamiltonian_neural_network.md b/docs/src/hamiltonian_neural_network.md deleted file mode 100644 index f3581c4..0000000 --- a/docs/src/hamiltonian_neural_network.md +++ /dev/null @@ -1,98 +0,0 @@ -# Hamiltonian Neural Network - -Here we build a Hamiltonian neural network as a symbolic neural network. - -```julia hnn -using SymbolicNeuralNetworks -using GeometricMachineLearning -using AbstractNeuralNetworks: Dense, initialparameters, UnknownArchitecture, Model -using LinearAlgebra: norm -using ChainRulesCore -using KernelAbstractions -import Symbolics -import Latexify - -input_dim = 2 -c = Chain(Dense(input_dim, 4, tanh), Dense(4, 1, identity; use_bias = false)) - -nn = HamiltonianSymbolicNeuralNetwork(c) -x_hvf = SymbolicNeuralNetworks.vector_field(nn) -x = x_hvf.x -hvf = x_hvf.hvf -hvf |> Latexify.latexify -``` - -We can now train this Hamiltonian neural network based on vector field data. As a Hamiltonian we take that of a harmonic oscillator: - -```julia hnn -H(z::Array{T}) where T = sum(z.^2) / T(2) -𝕁 = PoissonTensor(input_dim) -hvf_analytic(z) = 𝕁(z) - -const T = Float64 -n_points = 2000 -z_data = randn(T, 2, n_points) -nothing # hide -``` - -We now specify a pullback [`HamiltonianSymbolicNeuralNetwork`](@ref): - -```julia hnn -_pullback = SymbolicPullback(nn) -nothing # hide -``` - -We can now train the network: - -```julia hnn -ps = NeuralNetworkParameters(initialparameters(c, T)) -dl = DataLoader(z_data, hvf_analytic(z_data)) -o = Optimizer(AdamOptimizer(.01), ps) -batch = Batch(200) -const n_epochs = 150 -nn_dummy = NeuralNetwork(UnknownArchitecture(), c, ps, CPU()) -o(nn_dummy, dl, batch, n_epochs, _pullback.loss, _pullback; show_progress = true) -nothing # hide -``` - -We now integrate the vector field: - -```julia hnn -using GeometricIntegrators -hvf_closure(input) = build_nn_function(hvf, x, nn)(input, nn_dummy.params) -function v(v, t, q, params) - v .= hvf_closure(q) -end -pr = ODEProblem(v, (0., 500.), 0.1, [1., 0.]) -sol = integrate(pr, ImplicitMidpoint()) -``` - -```julia hnn -using CairoMakie - -fig = Figure() -ax = Axis(fig[1, 1]) -lines!(ax, [sol.q[i][1] for i in axes(sol.t, 1)].parent, [sol.q[i][2] for i in axes(sol.t, 1)].parent) -``` - -We also train a non-Hamiltonian vector field on the same data for comparison: - -```julia hnn -c_nh = Chain(Dense(2, 10, tanh), Dense(10, 4, tanh), Dense(4, 2, identity; use_bias = false)) -nn_nh = NeuralNetwork(c_nh, CPU()) -o = Optimizer(AdamOptimizer(T), nn_nh) -o(nn_nh, dl, batch, n_epochs * 10, FeedForwardLoss()) # we train for times as long as before -``` - -We now integrate the vector field and plot the solution: - -```julia hnn -vf_closure(input) = c_nh(input, nn_nh.params) -function v_nh(v, t, q, params) - v .= vf_closure(q) -end -pr = ODEProblem(v_nh, (0., 500.), 0.1, [1., 0.]) -sol_nh = integrate(pr, ImplicitMidpoint()) - -lines!(ax, [sol_nh.q[i][1] for i in axes(sol_nh.t, 1)].parent, [sol_nh.q[i][2] for i in axes(sol_nh.t, 1)].parent) -``` \ No newline at end of file