Merge pull request #27 from SciML/demos

Updated demos to use new `Lux.PeriodicEmbedding`
SciML · May 14, 2024 · f93a351 · f93a351
2 parents 96f2125 + 1ee202f
commit f93a351
Show file tree

Hide file tree

Showing 4 changed files with 6 additions and 18 deletions.
diff --git a/Project.toml b/Project.toml
@@ -13,7 +13,7 @@ SciMLBase = "0bca4576-84f4-4d90-8ffe-ffa030f20462"
 [compat]
 ForwardDiff = "0.10"
 JuMP = "1"
-Lux = "0.5"
+Lux = "0.5.45"
 ModelingToolkit = "8, 9"
 NLopt = "1"
 NeuralPDE = "5.10"

diff --git a/test/damped_pendulum.jl b/test/damped_pendulum.jl
@@ -45,11 +45,7 @@ dim_state = length(bounds)
 dim_hidden = 15
 dim_output = 2
 chain = [Lux.Chain(
- Lux.WrappedFunction(x -> vcat(
- transpose(sin.(x[1, :])),
- transpose(cos.(x[1, :])),
- transpose(x[2, :])
- )),
+ PeriodicEmbedding([1], [2π]),
  Dense(3, dim_hidden, tanh),
  Dense(dim_hidden, dim_hidden, tanh),
  Dense(dim_hidden, 1, use_bias = false)

diff --git a/test/inverted_pendulum.jl b/test/inverted_pendulum.jl
@@ -36,11 +36,7 @@ dim_phi = 2
 dim_u = 1
 dim_output = dim_phi + dim_u
 chain = [Lux.Chain(
- Lux.WrappedFunction(x -> vcat(
- transpose(sin.(x[1, :])),
- transpose(cos.(x[1, :])),
- transpose(x[2, :])
- )),
+ PeriodicEmbedding([1], [2π]),
  Dense(3, dim_hidden, tanh),
  Dense(dim_hidden, dim_hidden, tanh),
  Dense(dim_hidden, 1, use_bias = false)
@@ -151,7 +147,7 @@ closed_loop_dynamics = ODEFunction(
 
 # Starting still at bottom
 downward_equilibrium = zeros(2)
-ode_prob = ODEProblem(closed_loop_dynamics, downward_equilibrium, [0.0, 35.0], p)
+ode_prob = ODEProblem(closed_loop_dynamics, downward_equilibrium, [0.0, 75.0], p)
 sol = solve(ode_prob, Tsit5())
 # plot(sol)
 
@@ -162,7 +158,7 @@ x_end, y_end = sin(θ_end), -cos(θ_end)
 
 # Starting at a random point
 x0 = lb .+ rand(2) .* (ub .- lb)
-ode_prob = ODEProblem(closed_loop_dynamics, x0, [0.0, 20.0], p)
+ode_prob = ODEProblem(closed_loop_dynamics, x0, [0.0, 75.0], p)
 sol = solve(ode_prob, Tsit5())
 # plot(sol)
 

diff --git a/test/inverted_pendulum_ODESystem.jl b/test/inverted_pendulum_ODESystem.jl
@@ -49,11 +49,7 @@ dim_phi = 2
 dim_u = 1
 dim_output = dim_phi + dim_u
 chain = [Lux.Chain(
- Lux.WrappedFunction(x -> vcat(
- transpose(sin.(x[1, :])),
- transpose(cos.(x[1, :])),
- transpose(x[2, :])
- )),
+ PeriodicEmbedding([1], [2π]),
  Dense(3, dim_hidden, tanh),
  Dense(dim_hidden, dim_hidden, tanh),
  Dense(dim_hidden, 1, use_bias = false)