Set up Julia package and example SSP2 optimizations (#2)

* Set up Julia package and example SSP2 optimizations

* Delete examples/julia/Manifest.toml

* fix underflow in tanh projection

* use evalpoly

* fix projection typos/bugs

* refactor projection into a loop over target points

* refactor interpolation into another module

* update example to use a higher resolution target grid and reduce allocations

* correct the usage of cubic_adjoint deriv api

* correct the comparison to finite differences because arrays are reused

* avoid NaN in tanh projection when x == eta at beta = Inf

* change definition of ProjectionProblem and SSP2

* add grid as an optional argument to PaddingProblem

* add init! to interface

* fix typo

* add support for SSP1

* add linear interpolation and use bc in cubic interpolation

* rename R_smoothing_factor to smoothing_radius like python api

* add pythonic API

* add ChainRulesCore rrules for pythonic api

* run checks to make sure different apis match

* document apis of padding module

* document public api in kernels module

* document public api of convolution module

* Document public api of interpolation module

* document the public api of the projection module

* document public api of ssp package

* add package tests

* add README.md

* remove revise dep

* use rffts, precompute conv kernel, use perturb in arb direction, and make adjoint_solve use a structural tangent

* Apply suggestion from @stevengj

---------

Co-authored-by: Steven G. Johnson <stevenj@alum.mit.edu>

Author: lxvm

examples/julia/Project.toml

@@ -0,0 +1,6 @@
+[deps]
+CairoMakie = "13f3f980-e62b-5c42-98c6-ff1f3baf88f0"
+NLopt = "76087f3c-5699-56af-9a33-bf431cd00edd"
+Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
+SSP = "e5b5d2ee-15bb-40cc-a0da-b305b842b7a8"
+Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f"

examples/julia/ssp2_example.jl

@@ -0,0 +1,201 @@
+using SSP: init, solve!, adjoint_solve!
+using SSP: Kernel, Pad, Convolve, Project
+using .Kernel: conickernel
+using .Pad: FillPadding, BoundaryPadding, Inner, PaddingProblem, DefaultPaddingAlgorithm
+using .Convolve: DiscreteConvolutionProblem, FFTConvolution
+using .Project: ProjectionProblem, SSP1_linear, SSP1, SSP2
+
+using Random
+using CairoMakie
+using CairoMakie: colormap
+using NLopt
+
+
+Nx = Ny = 100
+grid = (
+    range(-1, 1, length=Nx),
+    range(-1, 1, length=Ny),
+)
+# Random.seed!(42)
+# design_vars = rand(Nx, Ny)
+# design_vars = [sinpi(x) * sinpi(y) for (x, y) in Iterators.product(grid...)]
+design_vars = let a = 0.5, b = 0.499
+    # Cassini oval
+    [((x^2 + y^2)^2 - 2a^2 * (x^2 - y^2) + a^4 - b^4) + 0.5 for (x, y) in Iterators.product(grid...)]
+end
+radius = 0.1
+
+kernel = conickernel(grid, radius)
+
+padprob = PaddingProblem(;
+    data = design_vars,
+    boundary = BoundaryPadding(size(kernel) .- 1, size(kernel) .- 1),
+    # boundary = FillPadding(1.0, size(kernel) .- 1, size(kernel) .- 1),
+)
+padalg = DefaultPaddingAlgorithm()
+padsolver = init(padprob, padalg)
+padsol = solve!(padsolver)
+
+convprob = DiscreteConvolutionProblem(;
+    data = padsol.value,
+    kernel,
+)
+
+convalg = FFTConvolution()
+convsolver = init(convprob, convalg)
+convsol = solve!(convsolver)
+
+depadprob = PaddingProblem(;
+    data = convsol.value,
+    boundary = Inner(size(kernel) .- 1, size(kernel) .- 1),
+)
+depadalg = DefaultPaddingAlgorithm()
+depadsolver = init(depadprob, depadalg)
+depadsol = solve!(depadsolver)
+
+filtered_design_vars = depadsol.value
+
+# projection points need not be the same as design variable grid
+target_grid = (
+    range(-1, 1, length=Nx * 1),
+    range(-1, 1, length=Ny * 1),
+)
+target_points = vec(collect(Iterators.product(target_grid...)))
+projprob = ProjectionProblem(;
+    rho_filtered=filtered_design_vars,
+    grid,
+    target_points,
+    beta = Inf,
+    eta = 0.5,
+)
+# projalg = SSP1_linear()
+# projalg = SSP1()
+projalg = SSP2()
+projsolver = init(projprob, projalg)
+projsol = solve!(projsolver)
+
+projected_design_vars = projsol.value
+
+let
+    fig = Figure()
+    ax1 = Axis(fig[1,1]; title = "design variables", aspect=DataAspect())
+    h1 = heatmap!(grid..., design_vars; colormap=colormap("grays"))
+    Colorbar(fig[1,2], h1)
+
+    ax2 = Axis(fig[1,3]; title = "SSP2 output", aspect=DataAspect())
+    h2 = heatmap!(target_grid..., reshape(projected_design_vars, length.(target_grid)); colormap=colormap("grays"))
+    Colorbar(fig[1,4], h2)
+    save("design.png", fig)
+end
+
+function fom(data, grid)
+    return sum(abs2, data) / length(data)
+end
+obj = fom(projected_design_vars, grid)
+
+function adjoint_fom(adj_fom, data, grid)
+    adjoint_fom!(similar(data), adj_fom, data, grid)
+end
+function adjoint_fom!(adj_data, adj_fom, data, grid)
+    adj_data .= (adj_fom / length(data)) .* 2 .* data
+    return adj_data
+end
+
+adj_rho_projected = adjoint_fom(1.0, projected_design_vars, grid)
+
+adj_projsol = (; value=adj_rho_projected)
+adj_projprob = adjoint_solve!(projsolver, adj_projsol, projsol.tape)
+adj_depadsol = (; value=adj_projprob.rho_filtered)
+adj_depadprob = adjoint_solve!(depadsolver, adj_depadsol, depadsol.tape)
+adj_convsol = (; value=adj_depadprob.data)
+adj_convprob = adjoint_solve!(convsolver, adj_convsol, convsol.tape)
+adj_padsol = (; value=adj_convprob.data)
+adj_padprob = adjoint_solve!(padsolver, adj_padsol, padsol.tape)
+adj_design_vars = adj_padprob.data
+
+let
+    fig = Figure()
+    ax1 = Axis(fig[1,1]; title = "SSP2 output", aspect=DataAspect())
+    h1 = heatmap!(ax1, target_grid..., reshape(projected_design_vars, length.(target_grid)); colormap=colormap("grays"))
+    Colorbar(fig[1,2], h1)
+
+    ax2 = Axis(fig[1,3]; title = "design variables gradient", aspect=DataAspect())
+    h2 = heatmap!(ax2, grid..., adj_design_vars; colormap=colormap("RdBu"))
+    Colorbar(fig[1,4], h2)
+    save("design_gradient.png", fig)
+end
+
+fom_withgradient = let grid=grid, padsolver=padsolver, convsolver=convsolver, depadsolver=depadsolver, projsolver=projsolver, adj_rho_projected=adj_rho_projected
+    function (design_vars)
+
+        padsolver.data = design_vars
+        padsol = solve!(padsolver)
+        convsolver.data = padsol.value
+        convsol = solve!(convsolver)
+        depadsolver.data = convsol.value
+        depadsol = solve!(depadsolver)
+        projsolver.rho_filtered = depadsol.value
+        projsol = solve!(projsolver)
+
+        _fom = fom(projsol.value, grid)
+        adjoint_fom!(adj_rho_projected, 1.0, projsol.value, grid)
+
+        adj_projsol = (; value=adj_rho_projected)
+        adj_projprob = adjoint_solve!(projsolver, adj_projsol, projsol.tape)
+        adj_depadsol = (; value=adj_projprob.rho_filtered)
+        adj_depadprob = adjoint_solve!(depadsolver, adj_depadsol, depadsol.tape)
+        adj_convsol = (; value=adj_depadprob.data)
+        adj_convprob = adjoint_solve!(convsolver, adj_convsol, convsol.tape)
+        adj_padsol = (; value=adj_convprob.data)
+        adj_padprob = adjoint_solve!(padsolver, adj_padsol, padsol.tape)
+        adj_design_vars = adj_padprob.data
+        return _fom, adj_design_vars
+    end
+end
+
+h = 1e-6
+Random.seed!(0)
+perturb = h * randn(size(design_vars))
+fom_ph, = fom_withgradient(design_vars + perturb)
+fom_mh, = fom_withgradient(design_vars - perturb)
+dfomdh_fd = (fom_ph - fom_mh) / 2h
+
+fom_val, adj_design_vars = fom_withgradient(design_vars)
+dfomdh = 2sum(adj_design_vars .* perturb) / 2h
+@show dfomdh_fd dfomdh
+
+opt = NLopt.Opt(:LD_CCSAQ, length(design_vars))
+evaluation_history = Float64[]
+my_objective_fn = let fom_withgradient=fom_withgradient, evaluation_history=evaluation_history, design_vars=design_vars
+    function (x, grad)
+        val, adj_design = fom_withgradient(reshape(x, size(design_vars)))
+        if !isempty(grad)
+            copy!(grad, vec(adj_design))
+        end
+        push!(evaluation_history, val)
+        return val
+    end
+end
+NLopt.min_objective!(opt, my_objective_fn)
+NLopt.maxeval!(opt, 50)
+fmax, xmax, ret = NLopt.optimize(opt, vec(design_vars))
+
+let
+    padsolver.data = reshape(xmax, size(design_vars))
+    padsol = solve!(padsolver)
+    convsolver.data = padsol.value
+    convsol = solve!(convsolver)
+    depadsolver.data = convsol.value
+    depadsol = solve!(depadsolver)
+    projsolver.rho_filtered = depadsol.value
+    projsol = solve!(projsolver)
+
+    fig = Figure()
+    ax1 = Axis(fig[1,1]; title = "Objective history", yscale=log10, limits = (nothing, (1e-16, 1e1)))
+    h1 = scatterlines!(ax1, evaluation_history)
+
+    ax2 = Axis(fig[1,2]; title = "Final SSP2 design", aspect=DataAspect())
+    h2 = heatmap!(target_grid..., reshape(projsol.value, length.(target_grid)); colormap=colormap("grays"))
+    Colorbar(fig[1,3], h2)
+    save("optimization.png", fig)
+end

examples/julia/ssp_comparison.jl

@@ -0,0 +1,97 @@
+using SSP: conic_filter, ssp1_linear, ssp1, ssp2
+
+using Random
+using CairoMakie
+using CairoMakie: colormap
+using NLopt
+using Zygote
+
+
+Nx = Ny = 100
+grid = (
+    range(-1, 1, length=Nx),
+    range(-1, 1, length=Ny),
+)
+# Random.seed!(42)
+# design_vars = rand(Nx, Ny)
+# design_vars = [sinpi(x) * sinpi(y) for (x, y) in Iterators.product(grid...)]
+design_vars = let a = 0.5, b = 0.499
+    # Cassini oval
+    [((x^2 + y^2)^2 - 2a^2 * (x^2 - y^2) + a^4 - b^4) + 0.5 for (x, y) in Iterators.product(grid...)]
+end
+radius = 0.1
+beta = Inf
+eta = 0.5
+ssp_algs = (ssp1_linear, ssp1, ssp2)
+
+ssp_projections = map(ssp_algs) do ssp
+    rho_filtered = conic_filter(design_vars, radius, grid)
+    rho_projected = ssp(rho_filtered, beta, eta, grid)
+    return rho_projected
+end
+
+let
+    fig = Figure(size = (1200, 400))
+    for (i, (ssp, rho_projected)) in enumerate(zip(ssp_algs, ssp_projections))
+        ax = Axis(fig[1,2i-1]; title = "$(string(nameof(ssp))) projection", aspect=DataAspect())
+        h = heatmap!(grid..., rho_projected; colormap=colormap("grays"))
+        Colorbar(fig[1,2i], h)
+    end
+    save("projection_comparison.png", fig)
+end
+
+function figure_of_merit(rho_projected)
+    sum(abs2, rho_projected) / length(rho_projected)
+end
+
+ssp_projection_gradients = map(ssp_algs) do ssp
+    design_vars_gradient = Zygote.gradient(design_vars) do design_vars
+        rho_filtered = conic_filter(design_vars, radius, grid)
+        rho_projected = ssp(rho_filtered, beta, eta, grid)
+        return figure_of_merit(rho_projected)
+    end
+    return design_vars_gradient[1]
+end
+
+let
+    fig = Figure(size = (1200, 400))
+    for (i, (ssp, rho_projected_gradient)) in enumerate(zip(ssp_algs, ssp_projection_gradients))
+        ax = Axis(fig[1,2i-1]; title = "$(string(nameof(ssp))) projection gradient", aspect=DataAspect())
+        h = heatmap!(grid..., rho_projected_gradient; colormap=colormap("RdBu"))
+        Colorbar(fig[1,2i], h)
+    end
+    save("projection_gradient_comparison.png", fig)
+end
+
+ssp_optimization_histories = map(ssp_algs) do ssp
+    opt = NLopt.Opt(:LD_CCSAQ, length(design_vars))
+    evaluation_history = Float64[]
+    my_objective_fn = let evaluation_history=evaluation_history, design_vars=design_vars
+        function (x, grad)
+            fom, adj_design = Zygote.withgradient(x) do x
+                rho_filtered = conic_filter(reshape(x, size(design_vars)), radius, grid)
+                rho_projected = ssp(rho_filtered, beta, eta, grid)
+                return figure_of_merit(rho_projected)
+            end
+            if !isempty(grad)
+                copy!(grad, vec(adj_design[1]))
+            end
+            push!(evaluation_history, fom)
+            return fom
+        end
+    end
+    NLopt.min_objective!(opt, my_objective_fn)
+    NLopt.maxeval!(opt, 50)
+    fmax, xmax, ret = NLopt.optimize(opt, vec(design_vars))
+    return evaluation_history
+end
+
+let
+    fig = Figure()
+    ax = Axis(fig[1,1]; title = "Optimization history", yscale=log10, limits = (nothing, (1e-16, 1e1)))
+    for (i, (ssp, evaluation_history)) in enumerate(zip(ssp_algs, ssp_optimization_histories))
+        scatterlines!(ax, evaluation_history; label=string(nameof(ssp)))
+    end
+    Legend(fig[1,2], ax)
+    save("evaluation_history_comparison.png", fig)
+end

src/julia/SSP/Project.toml

@@ -0,0 +1,20 @@
+name = "SSP"
+uuid = "e5b5d2ee-15bb-40cc-a0da-b305b842b7a8"
+version = "0.1.0"
+authors = ["Lorenzo Van Munoz <lorenzo@vanmunoz.com>"]
+
+[deps]
+FFTW = "7a1cc6ca-52ef-59f5-83cd-3a7055c09341"
+FastInterpolations = "9ea80cae-fc13-4c00-8066-6eaedb12f34b"
+
+[weakdeps]
+ChainRulesCore = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4"
+
+[extensions]
+SSPChainRulesCoreExt = "ChainRulesCore"
+
+[compat]
+ChainRulesCore = "1"
+FFTW = "1.10.0"
+FastInterpolations = "0.4.4"
+julia = "1.11"

src/julia/SSP/README.md

@@ -0,0 +1,63 @@
+# SSP
+
+A Smoothed Subpixel Projection (SSP) package for topology optimization in Julia.
+Supports N-dimensional data and reverse-mode automatic differentiation with minimal allocations.
+
+## Usage
+
+This package provides a high-level API nearly identical to its python cousin.
+It provides a `conic_filter` routine for smoothing design variables as well as
+three projections: `ssp1_linear`, `ssp1`, and `ssp2`.
+
+### Example
+
+```julia
+using SSP: conic_filter, ssp2
+
+Nx = Ny = 100
+grid = (
+    range(-1, 1, length=Nx),
+    range(-1, 1, length=Ny),
+)
+
+design_vars = let a = 0.5, b = 0.499
+    # Cassini oval
+    [((x^2 + y^2)^2 - 2a^2 * (x^2 - y^2) + a^4 - b^4) + 0.5 for (x, y) in Iterators.product(grid...)]
+end
+radius = 0.1
+beta = Inf
+eta = 0.5
+
+mapping = let radius=radius, grid=grid, beta=beta, eta=eta
+    function (design_vars)
+        rho_filtered = conic_filter(design_vars, radius, grid)
+        rho_projected = ssp2(rho_filtered, beta, eta, grid)
+        return rho_projected
+    end
+end
+
+mapping(design_vars)
+```
+
+This API also supports reverse-mode automatic differentaion through packages
+that rely on ChainRulesCore.jl, such as Zygote.jl. Therefore, calculating
+gradients is straightforward:
+
+```julia
+using Zygote
+
+fom = let mapping=mapping
+    function (x)
+        rho_projected = mapping(x)
+        return sum(abs2, rho_projected) / length(rho_projected)
+    end
+end
+Zygote.withgradient(fom, design_vars)
+```
+
+### Low-level API
+
+We provide a SciML-like API that allows reduced allocations and finer control
+over algorithm details such as boundary conditions and padding for interpolation
+and the selection of projection points. See `SSP/examples/julia/ssp2_example.jl`
+for how to use this API.