improve projection of PseudoDensities (#159)

* improve projection of PseudoDensities

* bump version

* more exports and tests

* formatter toml and formatting

.JuliaFormatter.toml

@@ -0,0 +1 @@
+style = "blue"

Project.toml

@@ -1,7 +1,7 @@
 name = "TopOpt"
 uuid = "53a1e1a5-51bb-58a9-8a02-02056cc81109"
 authors = ["mohamed82008 <[email protected]>", "yijiangh <[email protected]>"]
-version = "0.7.3"
+version = "0.8.0"
 
 [deps]
 AbstractDifferentiation = "c29ec348-61ec-40c8-8164-b8c60e9d9f3d"

src/TopOpt.jl

@@ -8,13 +8,28 @@ using Requires, Reexport, ChainRulesCore
 struct PseudoDensities{I,P,F,T,N,A<:AbstractArray{T,N}} <: AbstractArray{T,N}
     x::A
 end
+function Base.setindex!(A::PseudoDensities, x, inds...)
+    return A.x[inds...] = x
+end
 function PseudoDensities(x::A) where {T,N,A<:AbstractArray{T,N}}
     return PseudoDensities{false,false,false,T,N,A}(x)
 end
 function PseudoDensities{I,P,F}(x::A) where {I,P,F,T,N,A<:AbstractArray{T,N}}
     return PseudoDensities{I,P,F,T,N,A}(x)
 end
 
+Base.BroadcastStyle(::Type{T}) where {T<:PseudoDensities} = Broadcast.ArrayStyle{T}()
+function Base.similar(
+    bc::Broadcast.Broadcasted{Broadcast.ArrayStyle{T}}, ::Type{ElType}
+) where {T,ElType}
+    return similar(T, axes(bc))
+end
+function Base.similar(
+    ::Type{<:TV}, axes::Tuple{Union{Integer,Base.OneTo},Vararg{Union{Integer,Base.OneTo}}}
+) where {I,P,F,T,N,A,TV<:PseudoDensities{I,P,F,T,N,A}}
+    return PseudoDensities{I,P,F}(similar(A, axes))
+end
+
 function ChainRulesCore.rrule(
     ::Type{PseudoDensities{I,P,F,T,N,A}}, x::Matrix
 ) where {I,P,F,T,N,A}
@@ -104,6 +119,6 @@ export TopOpt,
     MMA87,
     MMA02,
     HeavisideProjection,
-    ProjectedPenalty,
-    PowerPenalty
+    SigmoidProjection,
+    ProjectedPenalty
 end

src/Utilities/Utilities.jl

@@ -5,6 +5,7 @@ using ForwardDiff, Ferrite, IterativeSolvers, Requires
 export AbstractPenalty,
     PowerPenalty,
     RationalPenalty,
+    SinhPenalty,
     HeavisideProjection,
     SigmoidProjection,
     ProjectedPenalty,

src/Utilities/penalties.jl

@@ -1,21 +1,27 @@
+import ..TopOpt: PseudoDensities
+
 abstract type AbstractPenalty{T} end
 abstract type AbstractCPUPenalty{T} <: AbstractPenalty{T} end
 abstract type AbstractProjection end
 
+function (P::AbstractCPUPenalty)(x::PseudoDensities{I,<:Any,F}) where {I,F}
+    return PseudoDensities{I,true,F}(P.(x.x))
+end
+
 mutable struct PowerPenalty{T} <: AbstractCPUPenalty{T}
     p::T
 end
-@inline (P::PowerPenalty)(x) = x^(P.p)
+(P::PowerPenalty)(x::Real) = x^(P.p)
 
 mutable struct RationalPenalty{T} <: AbstractCPUPenalty{T}
     p::T
 end
-@inline (R::RationalPenalty)(x) = x / (1 + R.p * (1 - x))
+(R::RationalPenalty)(x::Real) = x / (1 + R.p * (1 - x))
 
 mutable struct SinhPenalty{T} <: AbstractCPUPenalty{T}
     p::T
 end
-@inline (R::SinhPenalty)(x) = sinh(R.p * x) / sinh(R.p)
+(R::SinhPenalty)(x::Real) = sinh(R.p * x) / sinh(R.p)
 
 struct ProjectedPenalty{T,Tpen<:AbstractPenalty{T},Tproj} <: AbstractCPUPenalty{T}
     penalty::Tpen
@@ -24,17 +30,20 @@ end
 function ProjectedPenalty(penalty::AbstractPenalty{T}) where {T}
     return ProjectedPenalty(penalty, HeavisideProjection(10 * one(T)))
 end
-@inline (P::ProjectedPenalty)(x) = P.penalty(P.proj(x))
+@inline (P::ProjectedPenalty)(x::Real) = P.penalty(P.proj(x))
 @forward_property ProjectedPenalty penalty
 
 mutable struct HeavisideProjection{T} <: AbstractProjection
     β::T
 end
-@inline (P::HeavisideProjection)(x) = 1 - exp(-P.β * x) + x * exp(-P.β)
+@inline (P::HeavisideProjection)(x::Real) = 1 - exp(-P.β * x) + x * exp(-P.β)
+(P::HeavisideProjection)(x::AbstractArray) = P.(x)
+
 mutable struct SigmoidProjection{T} <: AbstractProjection
     β::T
 end
-@inline (P::SigmoidProjection)(x) = 1 / (1 + exp((P.β + 1) * (-x + 0.5)))
+@inline (P::SigmoidProjection)(x::Real) = 1 / (1 + exp((P.β + 1) * (-x + 0.5)))
+(P::SigmoidProjection)(x::AbstractArray) = P.(x)
 
 import Base: copy
 copy(p::TP) where {TP<:AbstractPenalty} = TP(p.p)

test/Functions/test_common_fns.jl

@@ -6,11 +6,31 @@ using Ferrite: ndofs_per_cell, getncells
 
 Random.seed!(1)
 
+get_pen_T(::PseudoDensities{<:Any,T,<:Any}) where {T} = T
+
+@testset "Projections and penalties" begin
+    for proj in (HeavisideProjection(5.0), SigmoidProjection(4.0))
+        for T1 in (true, false), T2 in (true, false), T3 in (true, false)
+            x = PseudoDensities{T1,T2,T3}(rand(4))
+            @test typeof(proj(x)) === typeof(x)
+            @test typeof(proj.(x)) === typeof(x)
+        end
+    end
+
+    for pen in (PowerPenalty(3.0), RationalPenalty(3.0), SinhPenalty(3.0))
+        for T1 in (true, false), T2 in (true, false), T3 in (true, false)
+            x = PseudoDensities{T1,T2,T3}(rand(4))
+            @test get_pen_T(pen(x)) === true
+            @test get_pen_T(ProjectedPenalty(pen)(x)) === true
+        end
+    end
+end
+
 @testset "Compliance" begin
     nels = (2, 2)
     problem = HalfMBB(Val{:Linear}, nels, (1.0, 1.0), 1.0, 0.3, 1.0)
     for p in (1.0, 2.0, 3.0)
-        solver = FEASolver(Direct, problem; xmin=0.01, penalty=TopOpt.PowerPenalty(p))
+        solver = FEASolver(Direct, problem; xmin=0.01, penalty=PowerPenalty(p))
         comp = Compliance(solver)
         f = x -> comp(PseudoDensities(x))
         for i in 1:3
@@ -29,7 +49,7 @@ end
     nels = (2, 2)
     problem = HalfMBB(Val{:Linear}, nels, (1.0, 1.0), 1.0, 0.3, 1.0)
     for p in (1.0, 2.0, 3.0)
-        solver = FEASolver(Direct, problem; xmin=0.01, penalty=TopOpt.PowerPenalty(p))
+        solver = FEASolver(Direct, problem; xmin=0.01, penalty=PowerPenalty(p))
         dp = Displacement(solver)
         u = dp(PseudoDensities(solver.vars))
         for _ in 1:3
@@ -50,7 +70,7 @@ end
     nels = (2, 2)
     problem = HalfMBB(Val{:Linear}, nels, (1.0, 1.0), 1.0, 0.3, 1.0)
     for p in (1.0, 2.0, 3.0)
-        solver = FEASolver(Direct, problem; xmin=0.01, penalty=TopOpt.PowerPenalty(p))
+        solver = FEASolver(Direct, problem; xmin=0.01, penalty=PowerPenalty(p))
         vol = Volume(solver)
         constr = x -> vol(PseudoDensities(x)) - 0.3
         for i in 1:3
@@ -69,8 +89,8 @@ end
     nels = (2, 2)
     problem = HalfMBB(Val{:Linear}, nels, (1.0, 1.0), 1.0, 0.3, 1.0)
     for p in (1.0, 2.0, 3.0)
-        solver = FEASolver(Direct, problem; xmin=0.01, penalty=TopOpt.PowerPenalty(p))
-        filter = TopOpt.DensityFilter(solver; rmin=4.0)
+        solver = FEASolver(Direct, problem; xmin=0.01, penalty=PowerPenalty(p))
+        filter = DensityFilter(solver; rmin=4.0)
         for i in 1:3
             x = rand(prod(nels))
             v = rand(prod(nels))
@@ -114,7 +134,7 @@ end
     end
     problem = MultiLoad(base_problem, F)
     for p in (1.0, 2.0, 3.0)
-        solver = FEASolver(Direct, problem; xmin=0.01, penalty=TopOpt.PowerPenalty(p))
+        solver = FEASolver(Direct, problem; xmin=0.01, penalty=PowerPenalty(p))
         exact_svd_block = BlockCompliance(problem, solver; method=:exact)
         constr = Nonconvex.FunctionWrapper(
             x -> exact_svd_block(x) .- 1000.0,
@@ -138,7 +158,7 @@ end
     nels = (2, 2)
     problem = HalfMBB(Val{:Linear}, nels, (1.0, 1.0), 1.0, 0.3, 1.0)
     for p in (1.0, 2.0, 3.0)
-        solver = FEASolver(Direct, problem; xmin=0.01, penalty=TopOpt.PowerPenalty(p))
+        solver = FEASolver(Direct, problem; xmin=0.01, penalty=PowerPenalty(p))
         st = StressTensor(solver)
         # element stress tensor - element 1
         est = st[1]