SciML compatibility #192
Description
I saw SciML.jl and it has a lot of functionalities. How about using them in TopOpt.jl?
The library helps people to use adjoint sensitivities and gradient based on AD or others.
For example,
using OrdinaryDiffEq, SciMLSensitivity
using Plots, LinearAlgebra, Optimisers
using Optimization, OptimizationPolyalgorithms, Zygote # Zygote still imported if needed
# ------------------------------------------------
# Problem Setup Parameters
# ------------------------------------------------
L = 1.0 # Domain length
N_el = 100 # Number of elements (design variables)
dx = L / N_el # Element length
N_nodes = N_el + 1 # Number of nodes
penal = 3.0 # SIMP penalization factor
ρ_min = 0.001 # Minimum density value
ρ_max = 1.0 # Maximum density value
Vfrac = 0.5 # Target volume fraction (allowed mass fraction)
target_mass = Vfrac * L # Target mass
# Final time for the ODE solve (chosen long enough to reach near–steady state)
T_final = 10.0
# ------------------------------------------------
# Helper Functions: Assemble Stiffness Matrix and Load Vector
# ------------------------------------------------
"""
K_matrix(rho)
Assemble the global stiffness matrix for a 1D rod using a SIMP formulation.
Each element contributes a 2×2 block scaled by (rho[i]^penal)/dx.
The Dirichlet BC at node 1 is enforced by replacing the first row and column.
"""
function K_matrix(rho)
# Assemble contributions from each element into a full matrix (N_nodes x N_nodes)
K = reduce(+, (
(rho[i]^penal)/dx *
[ (j == i && k == i ? 1.0 :
j == i && k == i+1 ? -1.0 :
j == i+1 && k == i ? -1.0 :
j == i+1 && k == i+1 ? 1.0 : 0.0)
for j in 1:N_nodes, k in 1:N_nodes ]
for i in 1:N_el
))
# Enforce Dirichlet BC at node 1: zero out row and column, set K[1,1]=1.
K_bc = [ (i == 1 && j == 1) ? 1.0 :
(i == 1 || j == 1) ? 0.0 :
K[i,j]
for i in 1:N_nodes, j in 1:N_nodes ]
return reshape(K_bc, (N_nodes, N_nodes))
end
"""
f_vector()
Returns the load vector for the 1D rod.
A unit load is applied at the last node while the first node is fixed.
"""
function f_vector()
return [ (i == 1 ? 0.0 : (i == N_nodes ? 1.0 : 0.0)) for i in 1:N_nodes ]
end
# ------------------------------------------------
# ODE Definition for the State Equation
# ------------------------------------------------
"""
ode_state(u, rho, t)
Defines the time–dependent ODE such that its steady state satisfies the equilibrium:
K(ρ)⋅u = f.
We use a simple first–order dynamics: u′ = –K(ρ)*u + f.
Note: The design variable ρ (a vector of length N_el) is passed as the parameter.
"""
function ode_state(du, u, rho, t)
K = K_matrix(rho)
f = f_vector()
du .= -K * u + f
end
# ------------------------------------------------
# Objective Function: Compliance + Mass Penalty
# ------------------------------------------------
"""
compliance_objective(rho)
For a given design vector ρ, we solve the ODE (from t=0 to T_final) to obtain the
steady–state displacement field u. The compliance (taken as the displacement at the
loaded end, u[end]) is computed and augmented with a quadratic penalty if the total mass
(sum(ρ)*dx) exceeds the target.
"""
function compliance_objective(rho)
u0 = zeros(N_nodes) # Initial condition for displacements
tspan = (0.0, T_final)
# Set up the ODEProblem with rho as the parameter.
prob = ODEProblem(ode_state, u0, tspan, rho)
# Solve using Tsit5 and an adjoint sensitivity algorithm (QuadratureAdjoint from DiffEq)
sol = solve(prob, Tsit5(); sensealg= InterpolatingAdjoint(), abstol=1e-8, reltol=1e-8)
uT = sol[end] # Steady state displacement (at t = T_final)
compliance = uT[end] # Compliance: displacement at the loaded (last) node
mass = sum(rho) * dx
penalty_coef = 1e6
penalty = penalty_coef * max(0.0, mass - target_mass)^2
return compliance + penalty
end
# ------------------------------------------------
# Visualization Function
# ------------------------------------------------
function plot_results(rho_opt)
# Solve the ODE with the optimal density to get the displacement field
u0 = zeros(N_nodes)
tspan = (0.0, T_final)
prob = ODEProblem(ode_state, u0, tspan, rho_opt)
sol = solve(prob, Tsit5(); sensealg=InterpolatingAdjoint(), abstol=1e-8, reltol=1e-8)
u_opt = sol[end]
x_nodes = range(0, stop=L, length=N_nodes)
p1 = plot(x_nodes, u_opt, lw=3, label="Displacement",
title="Displacement Field", xlabel="x", ylabel="u")
p2 = plot(1:N_el, rho_opt, lw=3, label="Density",
title="Optimal Density Distribution", xlabel="Element Index", ylabel="ρ")
plot(p1, p2, layout=(1,2), size=(1000,400))
end
# ------------------------------------------------
# Optimization Setup Using Adjoint Sensitivities from SciMLSensitivity.jl
# ------------------------------------------------
# Initial guess: uniform density equal to Vfrac
rho0 = fill(Vfrac, N_el)
# Callback to monitor optimization progress
loss_history = Float64[]
rho_history = []
cb = function (state, loss)
println("Loss: ", loss)
push!(loss_history, loss)
push!(rho_history, copy(state.u))
false # continue optimization
end
# Here we use the adjoint sensitivity approach from SciMLSensitivity.
# (Assuming the wrapper type provided by Optimization.jl is named `SciMLSensitivityAdjoint`.)
adtype = Optimization.AutoZygote() # Use adjoint sensitivities from SciMLSensitivity.jl
optf = Optimization.OptimizationFunction((ρ, p) -> compliance_objective(ρ), adtype)
optprob = Optimization.OptimizationProblem(optf, rho0)
# Solve the optimization problem using PolyOpt
# res = Optimization.solve(optprob, PolyOpt(), callback = cb, maxiters=10)
res = Optimization.solve(optprob, Optimisers.Descent(0.01), callback = cb, maxiters=10)
println("Optimal density distribution: ", res.u)
println("Final objective (compliance + penalty): ", res.f)
# ------------------------------------------------
# Plot the Final Results
# ------------------------------------------------
plot_results(res.u)