# Custom layer for a good guess of system dynamics

**URL:** <https://discourse.julialang.org/t/custom-layer-for-a-good-guess-of-system-dynamics/88728>\
**Category:** General Usage\
**Tags:** diffeqflux\
**Created:** [October 14, 2022, 1:15pm UTC](https://discourse.julialang.org/t/custom-layer-for-a-good-guess-of-system-dynamics/88728 "2022-10-14T13:15:43Z")\
**Posts on this page:** 10\
**Page:** 1

<div class="post-metadata">

**Author:** ![Antonino\_Gandolfo](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/antonino_gandolfo/32/18856_2.png) [@Antonino\_Gandolfo](https://discourse.julialang.org/u/Antonino_Gandolfo)\
**Post date:** [October 14, 2022, 1:15pm UTC](https://discourse.julialang.org/t/custom-layer-for-a-good-guess-of-system-dynamics/88728/1 "2022-10-14T13:15:44Z")

</div>

Hi, I’m trying to create a NeuralODE for the forcasting of a dynamic system.  
I did a NN for the forcasting of the system under different initial conditions and it’s work very well.

Now I would like to create NN capable to learn the infuences of a possible variation of a input value, to do that I’m trying to incorporate in the first layer of the NN a customized layer that take cares of the dynamic of the system I want to reproduce.

The problem is that I don’t understand how can I bring the parameters of this custom layer together with the other parameters of the whole NN.

```julia
function law(u,p,t)
    
    F = Float32(5.0)
    α = p

    z1 = u[2]
    z2 = α*F .- u[2] - u[1]
    
    return[z1,z2]
end

#nn
NN = Lux.Chain(x -> law(x,p,t),
                Lux.Dense(2,20,tanh),
                Lux.Dense(30,10,tanh),
                Lux.Dense(10,2))
p, st = Lux.setup(rng, NN)

prob_neuralode = NeuralODE(NN, tspan, Tsit5(), saveat = t)
p_init = Lux.ComponentArray(p)

```

the output is always an empty first layer

```julia
ComponentVector{Float32}(layer_1 = Float32[], layer_2 = (weight = Float32[0.567042 0.4250619; 0.042177025 0.5778097; … ; -0.30559936 -0.09352723; 0.20134257 -0.40837312], bias = Float32[0.0; 0.0; … ; 0.0; 0.0]), layer_3 = (weight = Float32[-0.1402396 0.52604 … -0.17828122 0.22984686; -0.20939533 0.24967496 … 0.2779108 -0.4940759; … ; -0.44822726 0.3582243 … -0.25572324 -0.05863302; 0.43971282 0.09462769 … -0.12614584 -0.41752148], bias = Float32[0.0; 0.0; … ; 0.0; 0.0]), layer_4 = (weight = Float32[0.36662766 0.41478822 … -0.37065506 -0.43711254; -0.15670164 -0.54979855 … -0.58355767 0.009892202], bias = Float32[0.0; 0.0]))

```

---

<div class="post-metadata">

**Author:** ![avikpal](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/avikpal/32/6550_2.png) [@avikpal](https://discourse.julialang.org/u/avikpal)\
**Post date:** [October 14, 2022, 5:34pm UTC](https://discourse.julialang.org/t/custom-layer-for-a-good-guess-of-system-dynamics/88728/2 "2022-10-14T17:34:58Z")

</div>

The problem here is that `Lux.setup` doesn’t know how to initialize `p` for your custom layer so it assumes that your layer has no parameters. See [Lux Interface - Lux.jl](http://lux.csail.mit.edu/stable/manual/interface/) for how to make custom layers work.

---

<div class="post-metadata">

**Author:** ![ChrisRackauckas](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/chrisrackauckas/32/77_2.png) [@ChrisRackauckas](https://discourse.julialang.org/u/ChrisRackauckas)\
**Post date:** [October 14, 2022, 5:50pm UTC](https://discourse.julialang.org/t/custom-layer-for-a-good-guess-of-system-dynamics/88728/3 "2022-10-14T17:50:35Z")

</div>

What’s the error message thrown on that? This seems something that could be handled by posting exactly that when such a case is seen?

---

<div class="post-metadata">

**Author:** ![avikpal](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/avikpal/32/6550_2.png) [@avikpal](https://discourse.julialang.org/u/avikpal)\
**Post date:** [October 14, 2022, 7:14pm UTC](https://discourse.julialang.org/t/custom-layer-for-a-good-guess-of-system-dynamics/88728/4 "2022-10-14T19:14:42Z")

</div>

The issue with the definition OP posted is that it is a perfectly valid model declaration (`Chain` sees that it is a function with 1 argument, so it is automatically converted to a layer with no parameter/state internally), So an error can only occur at runtime when it sees that `p` and `t` are undefined.

---

<div class="post-metadata">

**Author:** ![Antonino\_Gandolfo](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/antonino_gandolfo/32/18856_2.png) [@Antonino\_Gandolfo](https://discourse.julialang.org/u/Antonino_Gandolfo)\
**Post date:** [October 17, 2022, 7:14am UTC](https://discourse.julialang.org/t/custom-layer-for-a-good-guess-of-system-dynamics/88728/5 "2022-10-17T07:14:19Z")

</div>

There is no error in this problem formulation. I can’t figure out how can I pass the parameters of the approximated dynamic layer `law(u,p,t)` in the NeuralODE and soo on in the optimization problem.

Any idea?

---

<div class="post-metadata">

**Author:** ![avikpal](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/avikpal/32/6550_2.png) [@avikpal](https://discourse.julialang.org/u/avikpal)\
**Post date:** [October 18, 2022, 3:17pm UTC](https://discourse.julialang.org/t/custom-layer-for-a-good-guess-of-system-dynamics/88728/6 "2022-10-18T15:17:43Z")

</div>

```julia
struct LawLayer <: Lux.AbstractExplicitLayer
end

Lux.initialparameters(rng, l::LawLayer) = (p = define how to initialize p,)

function (::LawLayer)(u, ps, st)
    
    F = Float32(5.0)
    α = ps.p

    z1 = u[2]
    z2 = α*F .- u[2] - u[1]
    
    return ([z1,z2], st)
end

```

See [Training a Simple LSTM - Lux.jl](http://lux.csail.mit.edu/stable/examples/generated/beginner/SimpleRNN/main/#creating-a-classifier)

---

<div class="post-metadata">

**Author:** ![Antonino\_Gandolfo](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/antonino_gandolfo/32/18856_2.png) [@Antonino\_Gandolfo](https://discourse.julialang.org/u/Antonino_Gandolfo)\
**Post date:** [October 22, 2022, 10:13am UTC](https://discourse.julialang.org/t/custom-layer-for-a-good-guess-of-system-dynamics/88728/7 "2022-10-22T10:13:53Z")

</div>

Really thank you for your support now I’ve clarified several doubts but I miss something yet, please be patient.  
I tried to follow this tutorial [Lux Interface](http://lux.csail.mit.edu/stable/manual/interface/).

```julia
struct LawLayer{F1} <: Lux.AbstractExplicitLayer
  parameters::F1
end

function LawLayer(; parameters = Lux.glorot_uniform)
return LawLayer{typeof(parameters)}(parameters)
end

LawLayer()

Lux.initialparameters(rng, l::LawLayer) = (p = l.parameters(rng,1,1))
Lux.initialstates(::AbstractRNG, ::LawLayer) = NamedTuple()
Lux.parameterlength(::LawLayer) = 1
Lux.statelength(::LawLayer) = 0

function (::LawLayer)(u, ps, st)
    
    F = Float32(5.0)
    α = ps.p

    z1 = u[2]
    z2 = α*F .- u[2] - u[1]
    
    return ([z1,z2], st)
end

```

but I recive the same result:

```julia
Chain(
    layer_1 = WrappedFunction(#8),
    layer_2 = Dense(2 => 30, tanh_fast), # 90 parameters
    layer_3 = Dense(30 => 30, tanh_fast), # 930 parameters
    layer_4 = Dense(30 => 2), # 62 parameters
) # Total: 1_082 parameters,
          # plus 0 states, summarysize 48 bytes.

((layer_1 = NamedTuple(), layer_2 = (weight = Float32[0.10518762 0.30059525; -0.27172458 -0.36132228; … ; 0.061005663 0.044980295; -0.2633441 0.023604192], bias = Float32[0.0; 0.0; … ; 0.0; 0.0]), layer_3 = (weight = Float32[0.16688053 -0.042258434 … 0.07908167 -0.19162428; 0.19909461 0.16200048 … 0.11714643 -0.03535033; … ; -0.13860366 0.009593734 … -0.066989586 -0.0815194; 0.21782668 -0.23353975 … -0.15897466 0.22407342], bias = Float32[0.0; 0.0; … ; 0.0; 0.0]), layer_4 = (weight = Float32[0.215354 0.35274163 … 0.15960616 -0.2909302; 0.055675678 0.3207232 … -0.34296218 -0.34948933], bias = Float32[0.0; 0.0])), (layer_1 = NamedTuple(), layer_2 = NamedTuple(), layer_3 = NamedTuple(), layer_4 = NamedTuple()))

```

what I’m wrong?

---

<div class="post-metadata">

**Author:** ![avikpal](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/avikpal/32/6550_2.png) [@avikpal](https://discourse.julialang.org/u/avikpal)\
**Post date:** [October 22, 2022, 11:54am UTC](https://discourse.julialang.org/t/custom-layer-for-a-good-guess-of-system-dynamics/88728/8 "2022-10-22T11:54:27Z")

</div>

How are you defining the Chain? Seems like `layer_1` is a function and not `LawLayer()`

---

<div class="post-metadata">

**Author:** ![Antonino\_Gandolfo](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/antonino_gandolfo/32/18856_2.png) [@Antonino\_Gandolfo](https://discourse.julialang.org/u/Antonino_Gandolfo)\
**Post date:** [October 22, 2022, 1:55pm UTC](https://discourse.julialang.org/t/custom-layer-for-a-good-guess-of-system-dynamics/88728/9 "2022-10-22T13:55:07Z")

</div>

Sorry, here the whole code:

```julia
rng = Random.default_rng()

in_layer = 30
hidden_layer = 30 

struct LawLayer{F1} <: Lux.AbstractExplicitLayer
  parameters::F1
end

function LawLayer(; parameters = Lux.glorot_uniform)
return LawLayer{typeof(parameters)}(parameters)
end

LawLayer()

Lux.initialparameters(rng, l::LawLayer) = (p = l.parameters(rng,1,1))
Lux.initialstates(::AbstractRNG, ::LawLayer) = NamedTuple()
Lux.parameterlength(::LawLayer) = 1
Lux.statelength(::LawLayer) = 0

function (::LawLayer)(u, ps, st)
    
    F = Float32(5.0)
    α = ps.p

    z1 = u[2]
    z2 = α*F .- u[2] - u[1]
    
    return ([z1,z2], st)
end

#nn
NN = Lux.Chain(u -> LawLayer(u,ps,st),
                Lux.Dense(2,in_layer,tanh),
                Lux.Dense(in_layer,hidden_layer,tanh),
                Lux.Dense(hidden_layer,2))
                      
p, st = Lux.setup(rng, NN)

prob_neuralode = NeuralODE(NN, tspan, Tsit5(), saveat = t)
p_init = Lux.ComponentArray(p)

```

---

<div class="post-metadata">

**Author:** ![avikpal](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/avikpal/32/6550_2.png) [@avikpal](https://discourse.julialang.org/u/avikpal)\
**Post date:** [October 22, 2022, 4:25pm UTC](https://discourse.julialang.org/t/custom-layer-for-a-good-guess-of-system-dynamics/88728/10 "2022-10-22T16:25:09Z")

</div>

Change `u -> LawLayer(u,ps,st)` to `LawLayer()`