Fix penalties and projections (#162)

* Documenter warnonly = true * fix penalties and projections * remove redundant rrule * use Documenter 1 syntax
JuliaTopOpt · Oct 22, 2023 · b0aa456 · b0aa456
1 parent 6444ff1
commit b0aa456
Show file tree

Hide file tree

Showing 4 changed files with 15 additions and 12 deletions.
diff --git a/docs/make.jl b/docs/make.jl
@@ -13,11 +13,12 @@ GENERATED_EXAMPLES = [
 ]
 
 bib = CitationBibliography(joinpath(@__DIR__, "biblio", "ref.bib"))
-makedocs(
-    bib;
+makedocs(;
     sitename="TopOpt.jl",
     format=Documenter.HTML(; prettyurls=get(ENV, "CI", nothing) == "true"),
     # doctest = false,
+    warnonly = true,
+    plugins = [bib],
     pages=[
         "Home" => "index.md",
         "Problem types" => "examples/problem.md",

diff --git a/src/TopOpt.jl b/src/TopOpt.jl
@@ -5,6 +5,9 @@ using Requires, Reexport, ChainRulesCore
 
 @reexport using Nonconvex, NonconvexMMA, NonconvexSemidefinite, NonconvexPercival
 
+# I: interpolated
+# P: penalized
+# F: filtered
 struct PseudoDensities{I,P,F,T,N,A<:AbstractArray{T,N}} <: AbstractArray{T,N}
     x::A
 end
@@ -16,6 +19,9 @@ function PseudoDensities(x::A) where {T,N,A<:AbstractArray{T,N}}
 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 
+function ChainRulesCore.rrule(::Type{PseudoDensities{I, P, F, T, N, A}}, x) where {I, P, F, T, N, A <: AbstractArray{T, N}}
+    PseudoDensities{I, P, F, T, N, A}(x), Δ -> (NoTangent(), Δ isa Tangent ? Δ.x : Δ)
 end
 
 Base.BroadcastStyle(::Type{T}) where {T<:PseudoDensities} = Broadcast.ArrayStyle{T}()
@@ -30,13 +36,6 @@ function Base.similar(
     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}
-    px = PseudoDensities{I,P,F,T,N,A}(x)
-    return px, Δ -> (NoTangent(), Δ isa Matrix ? Δ : Δ.x)
-end
-
 Base.length(x::PseudoDensities) = length(x.x)
 Base.size(x::PseudoDensities, i...) = size(x.x, i...)
 Base.getindex(x::PseudoDensities, i...) = x.x[i...]

diff --git a/src/Utilities/penalties.jl b/src/Utilities/penalties.jl
@@ -5,7 +5,7 @@ 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))
+    return PseudoDensities{I,true,F}(map(P, x.x))
 end
 
 mutable struct PowerPenalty{T} <: AbstractCPUPenalty{T}
@@ -33,17 +33,18 @@ end
 @inline (P::ProjectedPenalty)(x::Real) = P.penalty(P.proj(x))
 @forward_property ProjectedPenalty penalty
 
+(P::AbstractProjection)(x::PseudoDensities{I,T,F}) where {I,T,F} = PseudoDensities{I,T,F}(P(x.x))
+(P::AbstractProjection)(x::AbstractArray) = map(P, x)
+
 mutable struct HeavisideProjection{T} <: AbstractProjection
     β::T
 end
 @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::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)

diff --git a/test/Functions/test_common_fns.jl b/test/Functions/test_common_fns.jl
@@ -14,6 +14,7 @@ get_pen_T(::PseudoDensities{<:Any,T,<:Any}) where {T} = T
             x = PseudoDensities{T1,T2,T3}(rand(4))
             @test typeof(proj(x)) === typeof(x)
             @test typeof(proj.(x)) === typeof(x)
+            Zygote.gradient(sum ∘ proj ∘ PseudoDensities, x)
         end
     end
 
@@ -22,6 +23,7 @@ get_pen_T(::PseudoDensities{<:Any,T,<:Any}) where {T} = T
             x = PseudoDensities{T1,T2,T3}(rand(4))
             @test get_pen_T(pen(x)) === true
             @test get_pen_T(ProjectedPenalty(pen)(x)) === true
+            Zygote.gradient(sum ∘ pen ∘ PseudoDensities, x)
         end
     end
 end