Custom Layer, Differential Equation as Activation Function. about lux.jl HOT 4 CLOSED

tspan = (0, 200)
time_steps = collect(1:1:10)

function hopf_oscillator(du, u, p, t)
	@unpack W = p
	ω_h = 1.
	μ = 1.
	ω_ext = 1
	I0 = 0.1
	w_rand = rand(5)
	signal = [cos.(i*t) for i in 1:5] 
	Iext = W' * signal
	du[1] = u[1] * (μ - u[1]^2) + I0 * Iext * cos(u[2])
	du[2] = ω_h - I0 * Iext * sin(u[2])/u[1]
end

u0_h = [.1, .1]
p_h = ComponentArray(W = rand(5))

prob_h = ODEProblem(hopf_oscillator, u0_h, dt = 0.1, tspan, p_h)

function gradients_hopf_with_loss_function(p)
	_prob = remake(prob_h, u0 = u0_h, tspan = (0.0, 20.0), p = p)
	sol = solve(_prob, Rosenbrock23(), saveat = time_steps)
	theta = sol[2,:]
        D =  cos.(time_steps)
	r = sol[1, end]
	x = r*cos.(theta) .- D
	sum(abs2, x)  #loss function
end

for i in 1:50
	dp_h = Zygote.gradient(gradients_hopf_with_loss_function, p_h)
	for j in 1:length(p_h)
		p_h[j] = p_h[j] - 0.1 * dp_h[1][j]
	end
end

using Lux, DifferentialEquations, Zygote, SciMLSensitivity, Random, ComponentArrays, Parameters, DataInterpolations, Plots


rng = Random.default_rng()
Random.seed!(rng, 0)

struct Modified_Dense{M <: Lux.AbstractExplicitLayer} <: Lux.AbstractExplicitContainerLayer{(:model,)}
    model::M
    u0
    time_steps 
    solver
    sensealg
    tspan
    kwargs    
end

function Modified_Dense(model::Lux.AbstractExplicitLayer; u0 = 0.1 .* ones(model.out_dims, 2), time_steps = (0.1:0.1:100), solver=Tsit5(), tspan=(0.0f0, 100.0f0), 
    sensealg=InterpolatingAdjoint(; autojacvec=ZygoteVJP()), kwargs...)
    return Modified_Dense(model, u0, time_steps, solver, sensealg, tspan, kwargs)
end

diffeqsol_to_array(x::ODESolution) = mapreduce(permutedims, vcat, x.u)

function (n::Modified_Dense)(x, ps, st)
    function ODE_Activation(du, u, p, t)
        @unpack ps = p
        for k in 1:n.model.out_dims
            du[k, 1] = u[k, 1] * (1 - u[k, 1]^2) + (ps.weight[k]*x(t) + ps.bias[k]) * cos(u[k, 2])
            du[k, 2] = k - (ps.weight[k]*x(t) + ps.bias[k]) * sin(u[k, 2]) / u[k, 1]
        end
    end

    prob = ODEProblem(ODE_Activation, n.u0, n.tspan, ps)
    sol = solve(prob, n.solver, saveat = n.time_steps)
    u1 = sol[:,1,:]
    u2 = sol[:,2,:]
    out = u1.*sin.(u2)
    println(size(out))
    return out, st
end

model = Modified_Dense(Dense(1,10, identity))

ps, st = Lux.setup(rng, model)
ps = ComponentArray(ps=ps)
dev = gpu_device()

x = CubicSpline(sin.(4 * π * model.time_steps), model.time_steps)

model(x, ps, st)

gradient(ps -> sum(first(model(x, ps, st))), ps)