changes to neural

Author: mohamed82008

src/FEA/direct_displacement_solver.jl

@@ -2,17 +2,17 @@ abstract type AbstractFEASolver end
 
 abstract type AbstractDisplacementSolver <: AbstractFEASolver end
 
-@params mutable struct DirectDisplacementSolver{T,dim,TP<:AbstractPenalty{T}} <:
+mutable struct DirectDisplacementSolver{T,dim,TP1<:AbstractPenalty{T},TP2<:StiffnessTopOptProblem{dim,T}, TG <: GlobalFEAInfo{T}, TE <: ElementFEAInfo{dim, T}, Tu <: AbstractVector{T}} <:
                        AbstractDisplacementSolver
-    problem::StiffnessTopOptProblem{dim,T}
-    globalinfo::GlobalFEAInfo{T}
-    elementinfo::ElementFEAInfo{dim,T}
-    u::AbstractVector{T}
-    lhs::AbstractVector{T}
-    rhs::AbstractVector{T}
-    vars::AbstractVector{T}
-    penalty::TP
-    prev_penalty::TP
+    problem::TP2
+    globalinfo::TG
+    elementinfo::TE
+    u::Tu
+    lhs::Tu
+    rhs::Tu
+    vars::Tu
+    penalty::TP1
+    prev_penalty::TP1
     xmin::T
     qr::Bool
 end

src/Functions/Functions.jl

@@ -34,8 +34,10 @@ export Volume,
     apply_boundary_with_zerodiag!,
     apply_boundary_with_meandiag!,
     NeuralNetwork,
-    PredictFunction,
     TrainFunction,
+    PredictFunction,
+    NNParams,
+    Coordinates,
     StressTensor,
     ElementStressTensor,
     MaterialInterpolation,

src/Functions/neural.jl

@@ -1,3 +1,10 @@
+struct Coordinates{C}
+    coords::C
+end
+struct NNParams{W}
+    p::W
+end
+
 function getcentroids(problem::AbstractTopOptProblem)
     dh = problem.ch.dh
     return map(CellIterator(dh)) do cell
@@ -15,20 +22,29 @@ abstract type AbstractMLModel end
     init_params::Any
     params_to_out::Any
     in_to_out::Any
+    centroids::Any
 end
 function NeuralNetwork(nn_model, input_coords::AbstractVector)
     f = x -> nn_model(x)[1]
     @assert all(0 .<= f.(input_coords) .<= 1)
     p, re = Flux.destructure(nn_model)
     return NeuralNetwork(
-        nn_model, Float64.(p), p -> getindex.(re(p).(input_coords), 1), nn_model
+        nn_model,
+        Float64.(p),
+        p -> getindex.(re(p).(input_coords), 1),
+        nn_model,
+        input_coords,
     )
 end
-function NeuralNetwork(nn_model, problem::AbstractTopOptProblem)
+function NeuralNetwork(nn_model, problem::AbstractTopOptProblem; scale = true)
     centroids = getcentroids(problem)
-    m, s = mean(centroids), std(centroids)
-    scentroids = map(centroids) do c
-        (c .- m) ./ s
+    if scale
+        m, s = mean(centroids), std(centroids)
+        scentroids = map(centroids) do c
+            (c .- m) ./ s
+        end
+    else
+        scentroids = centroids
     end
     return NeuralNetwork(nn_model, scentroids)
 end
@@ -46,3 +62,9 @@ end
 function (tf::TrainFunction)(p)
     return PseudoDensities(tf.model.params_to_out(p))
 end
+
+function (ml::NeuralNetwork)(x::AbstractVector{<:Coordinates})
+    return PredictFunction(ml).(x)
+end
+(ml::NeuralNetwork)(x::Coordinates) = PredictFunction(ml)(x.coords)
+(ml::NeuralNetwork)(x::NNParams) = TrainFunction(ml)(x.p)

test/examples/neural.jl

@@ -24,7 +24,7 @@ comp = Compliance(solver)
 volfrac = Volume(solver)
 
 m = 20
-act = relu
+act = leakyrelu
 nn = Chain(Dense(2, m, act), Dense(m, m, act), Dense(m, m, act), Dense(m, 1, sigmoid))
 nn_model = NeuralNetwork(nn, problem)
 tf = TrainFunction(nn_model)

test/examples/neural2.jl

@@ -0,0 +1,99 @@
+using TopOpt, Zygote, Flux
+
+E = 1.0 # Young’s modulus
+v = 0.3 # Poisson’s ratio
+f = 1.0 # downward force
+els = (160, 40)
+
+problem = PointLoadCantilever(Val{:Linear}, els, (1.0, 1.0), E, v, f)
+
+# problem settings
+V = 0.5
+xmin = 1e-6
+rmin = 3.0
+
+# SIMP penalty
+p = 1.0
+delta_p = 0.01
+p_max = 5.0
+
+penalty = TopOpt.PowerPenalty(p)
+solver = FEASolver(Direct, problem; xmin, penalty)
+cheqfilter = DensityFilter(solver; rmin)
+comp = Compliance(solver)
+volfrac = Volume(solver)
+
+# constraint aggregation penalty
+alpha = 0.1
+delta_alpha = 0.05
+alpha_max = 100
+
+# neural network
+m = 20
+act = leakyrelu
+nn = NeuralNetwork(
+    Chain(
+        Dense(2, m, act, init = Flux.glorot_normal),
+        Dense(m, m, act, init = Flux.glorot_normal),
+        Dense(m, m, act, init = Flux.glorot_normal),
+        Dense(m, m, act, init = Flux.glorot_normal),
+        Dense(m, m, act, init = Flux.glorot_normal),
+        Dense(m, m, act, init = Flux.glorot_normal),
+        softmax,
+        x -> [x[1]],
+    ),
+    problem,
+    scale = true,
+)
+w0 = nn.init_params
+# w0 ./= 10
+
+C0 = comp(cheqfilter(PseudoDensities(fill(V, TopOpt.getncells(problem)))))
+
+# optimization
+alg = Flux.Optimise.Adam(0.1)
+clip_alg = Flux.Optimise.ClipValue(1.0)
+w = copy(w0)
+Δ = copy(w)
+proj = HeavisideProjection(0.0)
+
+# termination criteria
+eps = Inf
+eps_star = 0.05
+maxiter = 100
+epoch = 1
+constr_tol = 0.01
+violation = Inf
+todensities(w; filter = true) = filter ? PseudoDensities(proj.(cheqfilter(nn(NNParams(w))).x)) : PseudoDensities(proj.(nn(NNParams(w)).x))
+while true
+    epoch > maxiter && break
+    eps < eps_star && violation < constr_tol && break
+    global penalty = TopOpt.PowerPenalty(p)
+    global solver = FEASolver(Direct, problem; xmin, penalty)
+    global cheqfilter = DensityFilter(solver; rmin)
+    global comp = Compliance(solver)
+    global volfrac = Volume(solver)
+
+    global obj = w -> comp(todensities(w, filter = true)) / C0
+    global constr = w -> volfrac(todensities(w, filter = false)) / V - 1
+    global combined_obj = w -> obj(w) + alpha * constr(w)^2
+
+    global Δ = Zygote.gradient(combined_obj, w)[1]
+    @info "grad norm: $(norm(Δ))"
+    Flux.Optimise.apply!(clip_alg, w, Δ)
+    Flux.Optimise.apply!(alg, w, Δ)
+    global w = w - Δ
+    violation = constr(w)
+    global alpha = min(alpha_max, alpha + delta_alpha)
+    global p = min(p_max, p + delta_p)
+    global epoch += 1
+    global x = todensities(w; filter = false)
+    global eps = sum(0.05 .< x.x .< 0.95) / length(x.x)
+    @info "eps = $eps"
+    @info "obj = $(comp(todensities(w; filter = true)))"
+    @info "constr = $(volfrac(x) - V)"
+    @show alpha
+end
+
+using Images, ImageInTerminal
+reshape(Gray.(1 .- x), els...)'