diff --git a/src/Algorithms/Algorithms.jl b/src/Algorithms/Algorithms.jl index a1090829..07cc734e 100644 --- a/src/Algorithms/Algorithms.jl +++ b/src/Algorithms/Algorithms.jl @@ -2,7 +2,7 @@ module Algorithms using Nonconvex, ..Functions, Parameters, Requires using Nonconvex.NonconvexCore: AbstractModel -using ..TopOpt: whichdevice, AbstractDevice, CPU,PENALTY_BEFORE_INTERPOLATION +using ..TopOpt: whichdevice, AbstractDevice, CPU, PENALTY_BEFORE_INTERPOLATION using TimerOutputs, Setfield, StaticArrays using Parameters: @unpack, @pack! using ..Utilities, Ferrite diff --git a/src/Functions/stress_tensor.jl b/src/Functions/stress_tensor.jl index 512cf44c..36c9dd79 100644 --- a/src/Functions/stress_tensor.jl +++ b/src/Functions/stress_tensor.jl @@ -12,7 +12,9 @@ function StressTensor(solver) n = ndofs_per_cell(dh) global_dofs = zeros(Int, n) cellvalues = solver.elementinfo.cellvalues - return StressTensor(problem, solver, global_dofs, cellvalues, collect(CellIterator(dh)), 0.0) + return StressTensor( + problem, solver, global_dofs, cellvalues, collect(CellIterator(dh)), 0.0 + ) end function Ferrite.reinit!(s::StressTensor, cellidx) @@ -49,7 +51,7 @@ function ChainRulesCore.rrule(::typeof(reinit!), st::ElementStressTensor, cellid return reinit!(st, cellidx), _ -> (NoTangent(), NoTangent(), NoTangent()) end -function (f::ElementStressTensor)(u; element_dofs = false) +function (f::ElementStressTensor)(u; element_dofs=false) st = f.stress_tensor reinit!(f, f.cellidx) if element_dofs @@ -62,10 +64,12 @@ function (f::ElementStressTensor)(u; element_dofs = false) dim = TopOptProblems.getdim(st.problem) problem = st.problem E, ν = problem.E, problem.ν - return sum(map(1:n_basefuncs, 1:n_quad) do a, q_point - _u = cellu[dim*(a-1) .+ (1:dim)] - return tensor_kernel(f, q_point, a)(_u) - end) + return sum( + map(1:n_basefuncs, 1:n_quad) do a, q_point + _u = cellu[dim * (a - 1) .+ (1:dim)] + return tensor_kernel(f, q_point, a)(_u) + end, + ) end @params struct ElementStressTensorKernel{T} <: AbstractFunction{T} @@ -85,15 +89,20 @@ function (f::ElementStressTensorKernel)(_u) c2 = E * ν * (1 + ν) return c1 * I + c2 * ϵ end -function ChainRulesCore.rrule( - f::ElementStressTensorKernel, x::AbstractVector, -) +function ChainRulesCore.rrule(f::ElementStressTensorKernel, x::AbstractVector) v, (∇,) = AD.value_and_jacobian(AD.ForwardDiffBackend(), x -> vec(f(x)), x) return reshape(v, f.dim, f.dim), Δ -> (NoTangent(), ∇' * vec(Δ)) end function tensor_kernel(f::StressTensor, quad, basef) - return ElementStressTensorKernel(f.problem.E, f.problem.ν, quad, basef, f.cellvalues, TopOptProblems.getdim(f.problem)) + return ElementStressTensorKernel( + f.problem.E, + f.problem.ν, + quad, + basef, + f.cellvalues, + TopOptProblems.getdim(f.problem), + ) end function tensor_kernel(f::ElementStressTensor, quad, basef) return tensor_kernel(f.stress_tensor, quad, basef) diff --git a/test/Functions/test_common_fns.jl b/test/Functions/test_common_fns.jl index 9619febb..7d4a13a4 100644 --- a/test/Functions/test_common_fns.jl +++ b/test/Functions/test_common_fns.jl @@ -1,4 +1,5 @@ -using TopOpt, Zygote, FiniteDifferences, LinearAlgebra, Test, Random, SparseArrays, ForwardDiff +using TopOpt, + Zygote, FiniteDifferences, LinearAlgebra, Test, Random, SparseArrays, ForwardDiff const FDM = FiniteDifferences using TopOpt: ndofs using Ferrite: ndofs_per_cell, getncells