New extension for level set reinitialization

Project.toml

@@ -1,7 +1,7 @@
 name = "LevelSetMethods"
 uuid = "ca3df947-0575-4ae6-b9a3-78df393a7451"
 authors = ["Luiz M. Faria and Nicolas Lebbe"]
-version = "0.1.4"
+version = "0.1.5"
 
 [deps]
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
@@ -12,18 +12,21 @@ Interpolations = "a98d9a8b-a2ab-59e6-89dd-64a1c18fca59"
 MMG_jll = "86086c02-e288-5929-a127-40944b0018b7"
 Makie = "ee78f7c6-11fb-53f2-987a-cfe4a2b5a57a"
 MarchingCubes = "299715c1-40a9-479a-aaf9-4a633d36f717"
+NearestNeighbors = "b8a86587-4115-5ab1-83bc-aa920d37bbce"
 
 [extensions]
 InterpolationsExt = "Interpolations"
 MMGSurfaceExt = ["MMG_jll", "MarchingCubes"]
 MMGVolumeExt = "MMG_jll"
 MakieExt = "Makie"
+ReinitializationExt = ["Interpolations", "NearestNeighbors"]
 
 [compat]
 Interpolations = "0.15, 0.16"
 LinearAlgebra = "1"
 MMG_jll = "5"
 Makie = "0.21, 0.22, 0.23, 0.24"
 MarchingCubes = "0.1.10"
+NearestNeighbors = "0.4"
 StaticArrays = "1"
 julia = "1.9"

docs/Project.toml

@@ -8,6 +8,7 @@ LevelSetMethods = "ca3df947-0575-4ae6-b9a3-78df393a7451"
 LiveServer = "16fef848-5104-11e9-1b77-fb7a48bbb589"
 MMG_jll = "86086c02-e288-5929-a127-40944b0018b7"
 MarchingCubes = "299715c1-40a9-479a-aaf9-4a633d36f717"
+NearestNeighbors = "b8a86587-4115-5ab1-83bc-aa920d37bbce"
 StaticArrays = "90137ffa-7385-5640-81b9-e52037218182"
 
 [sources]

docs/make.jl

@@ -4,26 +4,20 @@ using GLMakie
 using MMG_jll
 using MarchingCubes
 using Interpolations
+using NearestNeighbors
 using StaticArrays
 using DocumenterCitations
 
 bib = CitationBibliography(joinpath(@__DIR__, "src", "refs.bib"); style = :numeric)
 
-DocMeta.setdocmeta!(
-    LevelSetMethods,
-    :DocTestSetup,
-    :(using LevelSetMethods);
-    recursive = true,
-)
-
 const page_rename = Dict("developer.md" => "Developer docs") # Without the numbers
 const numbered_pages = [
     file for file in readdir(joinpath(@__DIR__, "src")) if
         file != "index.md" && splitext(file)[2] == ".md" && occursin(r"^\d", file)
 ]
 
 modules = [LevelSetMethods]
-for extension in [:MakieExt, :MMGSurfaceExt, :MMGVolumeExt, :InterpolationsExt]
+for extension in [:MakieExt, :MMGSurfaceExt, :MMGVolumeExt, :InterpolationsExt, :ReinitializationExt]
     ext = Base.get_extension(LevelSetMethods, extension)
     isnothing(ext) && "error loading $ext"
     push!(modules, ext)
@@ -37,14 +31,13 @@ makedocs(;
     format = Documenter.HTML(;
         canonical = "https://maltezfaria.github.io/LevelSetMethods.jl",
         collapselevel = 2,
-    ),
-    pages = vcat(
+    ), pages = vcat(
         "Home" => "index.md",
         "terms.md",
         "time-integrators.md",
         "boundary-conditions.md",
         "Extensions" =>
-            ["extension-makie.md", "extension-mmg.md", "extension-interpolations.md"],
+            ["extension-makie.md", "extension-mmg.md", "extension-interpolations.md", "extension-reinitialization.md"],
         "Examples" => ["example-zalesak.md", "example-shape-optim.md"],
         numbered_pages,
     ),

docs/src/extension-makie.md

@@ -55,4 +55,4 @@ end
 
 !!! tip "Plotting a `LevelSetEquation`"
     Calling `plot` on a [`LevelSetEquation`](@ref) defaults to plotting the `LevelSet` given by its
-    [`current_state`](@ref); exactly the same as calling `plot(current_state(equation))`.
+    `current_state`; exactly the same as calling `plot(current_state(equation))`.

docs/src/extension-reinitialization.md

@@ -0,0 +1,46 @@
+# Reinitialization
+
+The `ReinitializationExt` extension provides a [`reinitialize!`](@ref) method to
+transform a level set function into a signed distance function by computing the closest
+point on the interface using Newton's method.
+
+## Usage
+
+After loading the required dependencies for this extension (`Interpolations` and
+`NearestNeighbors`), simply call `reinitialize!` on a `LevelSet` or a `LevelSetEquation` to
+reinitialize the level set function in-place:
+
+```julia
+using LevelSetMethods
+using CairoMakie
+using NearestNeighbors
+using Interpolations
+
+grid = CartesianGrid((-1, -1), (1, 1), (100, 100))
+# An ellipse
+sdf = LevelSet(x -> sqrt(x[1]^2 + x[2]^2) - 0.5, grid)
+ϕ = LevelSet(x -> x[1]^2 + x[2]^2 - 0.5^2, grid)
+LevelSetMethods.set_makie_theme!()
+fig = Figure(; size = (800, 300))
+ax1 = Axis(fig[1, 1]; title="Signed Distance")
+ax2 = Axis(fig[1, 2]; title="Before Reinitialization", ylabel = "", yticklabelsvisible = false)
+ax3 = Axis(fig[1, 3]; title="After Reinitialization", ylabel = "", yticklabelsvisible = false)
+
+contour!(ax1, sdf; levels = [0.25, 0, 0.5], labels = true, labelsize = 14)
+contour!(ax2, ϕ; levels = [0.25, 0, 0.5], labels = true, labelsize = 14)
+
+# Reinitialize using Newton's method
+reinitialize!(ϕ)
+contour!(ax3, ϕ; levels = [0.25, 0, 0.5], labels = true, labelsize = 14)
+fig
+```
+
+You can easily check that the reinitialized level set function is indeed a signed distance:
+
+```julia
+max_er = maximum(eachindex(grid)) do i
+  abs(ϕ[i] - sdf[i])
+end
+@test max_er < 1e-10 # hide
+println("Maximum error after reinitialization: $max_er")
+```

docs/src/index.md

@@ -96,8 +96,7 @@ Here is what the `.gif` file looks like:
 
 ![Dumbbell](ls_intro.gif)
 
-For more interesting applications and advanced usage, see the [examples section](@ref
-examples)!
+For more interesting applications and advanced usage, see the examples section!
 
 !!! note "Other resources"
     There is an almost one-to-one correspondance between each of the [`LevelSetTerm`](@ref)s

docs/src/refs.bib

@@ -30,12 +30,23 @@ @article{mitchell2007toolbox
 }
 
 @article{allaire2004structural,
-  title={Structural optimization using sensitivity analysis and a level-set method},
-  author={Allaire, Gr{\'e}goire and Jouve, Fran{\c{c}}ois and Toader, Anca-Maria},
-  journal={Journal of computational physics},
-  volume={194},
-  number={1},
-  pages={363--393},
-  year={2004},
-  publisher={Elsevier}
+  title     = {Structural optimization using sensitivity analysis and a level-set method},
+  author    = {Allaire, Gr{\'e}goire and Jouve, Fran{\c{c}}ois and Toader, Anca-Maria},
+  journal   = {Journal of computational physics},
+  volume    = {194},
+  number    = {1},
+  pages     = {363--393},
+  year      = {2004},
+  publisher = {Elsevier}
+}
+
+@article{saye2014high,
+  title     = {High-order methods for computing distances to implicitly defined surfaces},
+  author    = {Saye, Robert},
+  journal   = {Communications in Applied Mathematics and Computational Science},
+  volume    = {9},
+  number    = {1},
+  pages     = {107--141},
+  year      = {2014},
+  publisher = {Mathematical Sciences Publishers}
 }

ext/MakieExt.jl

@@ -46,10 +46,12 @@ function Makie.convert_arguments(
     return _contour_plot(ϕ)
 end
 
-function Makie.convert_arguments(::Type{<:Volume}, ϕ::LSM.LevelSet)
+function Makie.convert_arguments(P::Type{<:Volume}, ϕ::LSM.LevelSet)
     LSM.dimension(ϕ) == 3 ||
         throw(ArgumentError("Volume plot only supported for 3D level-sets."))
-    return _volume_plot(ϕ)
+    x, y, z, v = _volume_plot(ϕ)
+    # Delegate to Makie's existing conversion for arrays
+    return Makie.convert_arguments(P, x, y, z, v)
 end
 
 function _contour_plot(ϕ::LSM.LevelSet)
@@ -63,9 +65,9 @@ function _volume_plot(ϕ::LSM.LevelSet)
     msh = LSM.mesh(ϕ)
     x, y, z = LSM.xgrid(msh), LSM.ygrid(msh), LSM.zgrid(msh)
     v = LSM.values(ϕ)
-    # NOTE: volume gives a warning when passed an AbstractVector for x,y,z,
-    # and asks for a tuple instead
-    return extrema(x), extrema(y), extrema(z), v
+    # Return extrema as tuples - Makie will handle conversion
+    xlims, ylims, zlims = extrema(x), extrema(y), extrema(z)
+    return xlims, ylims, zlims, v
 end
 
 Makie.@recipe(LevelSetPlot, eq) do scene

ext/ReinitializationExt.jl

@@ -0,0 +1,123 @@
+module ReinitializationExt
+
+import LevelSetMethods as LSM
+using Interpolations
+using NearestNeighbors
+using LinearAlgebra
+using StaticArrays
+
+function __init__()
+    return @info "Loading extension for Newton reinitialization"
+end
+
+function LSM.reinitialize!(eq::LSM.LevelSetEquation; kwargs...)
+    LSM.reinitialize!(LSM.current_state(eq); kwargs...)
+    return eq
+end
+
+function LSM.reinitialize!(
+        ϕ::LSM.LevelSet;
+        upsample::Int = 4,
+        maxiters::Int = 10,
+        xtol::Float64 = 1.0e-8,
+        ftol::Float64 = 1.0e-8,
+    )
+    grid = LSM.mesh(ϕ)
+    vals = ϕ.vals
+    itp = Interpolations.interpolate(ϕ)
+    f(x) = itp(x...)
+    ∇f(x) = Interpolations.gradient(itp, x...)
+    ∇²f(x) = Interpolations.hessian(itp, x...)
+
+    # Sample the interface
+    pts = _sample_interface(grid, f, ∇f, upsample, maxiters, ftol)
+    tree = KDTree(pts)
+
+    for I in eachindex(grid)
+        x = grid[I]
+        # Find closest point in the cloud
+        idx, dist = nn(tree, x)
+        x0 = pts[idx]
+        # Refine with Newton's method
+        cp = _closest_point(f, ∇f, ∇²f, x, x0, maxiters, xtol, ftol)
+        vals[I] = sign(vals[I]) * norm(x - cp)
+    end
+    return ϕ
+end
+
+function _sample_interface(grid, f, ∇f, upsample, maxiter, ftol)
+    pts = Vector{SVector{LSM.dimension(grid), Float64}}()
+    for I in CartesianIndices(LSM.size(grid) .- 1)
+        Ip = CartesianIndex(Tuple(I) .+ 1)
+        lc, hc = grid[I], grid[Ip]
+        samples = (lc .+ (hc .- lc) .* (SVector(j, k) .+ 0.5) ./ upsample for j in 0:(upsample - 1), k in 0:(upsample - 1))
+        # Skip cells where all samples have the same sign
+        all(x -> f(x) > 0, samples) || all(x -> f(x) < 0, samples) && continue
+        # Go over samples and push them to the interface
+        for x in samples
+            pt = _project_to_interface(f, ∇f, x, maxiter, ftol)
+            push!(pts, pt)
+        end
+    end
+    return pts
+end
+
+function _project_to_interface(f, ∇f, x0, maxiter, ftol)
+    x = x0
+    for _ in 1:maxiter
+        val = f(x)
+        val < ftol && break # close enough to the interface
+        grad = ∇f(x)
+        norm_grad = norm(grad)
+        iszero(norm_grad) && break
+        δx = val * grad / norm_grad^2
+        x = x - δx
+    end
+    f(x) > ftol && @warn "projection to interface did not converge to $ftol at x=$x"
+    return x
+end
+
+function _closest_point(f, ∇f, ∇²f, xq::SVector, x0::SVector, maxiters, xtol, ftol)
+    x = x0
+    ∇p_x0 = ∇f(x0)
+    λ = dot(xq - x0, ∇p_x0) / dot(∇p_x0, ∇p_x0)
+
+    converged = false
+    for _ in 1:maxiters
+        px = f(x)
+        ∇p = ∇f(x)
+        ∇²p = ∇²f(x)
+
+        # System for Newton's method
+        # ∇L = [ x - xq + λ∇p ] = 0
+        #      [      p(x)     ]
+        grad_L = vcat(x - xq + λ * ∇p, px)
+
+        # Hessian of the Lagrangian
+        # H_L = [ I + λ∇²p   ∇p ]
+        #       [    ∇p'      0 ]
+        hess_L = hcat(vcat(I + λ * ∇²p, ∇p'), vcat(∇p, 0))
+
+        # Solve for the update
+        δ = -hess_L \ grad_L
+        δx = δ[1:(end - 1)]
+        δλ = δ[end]
+
+        # Update variables
+        α = 1.0 # TODO: reduce step size if not diverging?
+        x = x + α * δx
+        λ = λ + α * δλ
+
+        # Check for convergence
+        if norm(δx) < xtol && norm(f(x)) < ftol
+            converged = true
+            break
+        end
+    end
+
+    converged || @warn "closest point search did not converge at xq=$xq"
+
+    return x
+end
+
+end

src/LevelSetMethods.jl

@@ -30,7 +30,8 @@ export CartesianGrid,
     WENO5,
     LevelSetEquation,
     integrate!,
-    current_time
+    current_time,
+    reinitialize!
 
 """
     makie_theme()
@@ -49,4 +50,39 @@ function set_makie_theme! end
 function export_volume_mesh end
 function export_surface_mesh end
 
+"""
+    reinitialize!(ϕ::LevelSet; upsample=2, maxiters=10, xtol=1e-8)
+
+Reinitializes the level set `ϕ` to a signed distance, modifying it in place.
+
+The method works by first sampling the zero-level set of the interface, and then for each
+grid point, finding the closest point on the interface using a Newton-based method. The
+distance to the closest point is then used as the new value of the level set at that grid
+point, with the sign determined by the original level set value. See [saye2014high](@cite)
+for more details.
+
+## Arguments
+
+  - `ϕ`: The level set to reinitialize.
+
+## Keyword Arguments
+
+  - `upsample`: number of samples to take in each cell when sampling the interface.
+    Higher values yield better initial guesses for the closest point search, but increase
+    the computational cost.
+  - `maxiters`: maximum number of iterations to use in the Newton's method
+    to find the closest point on the interface.
+  - `xtol`: convergence tolerance for the Newton's method. The iterations stop when
+    the change in position is less than `xtol`.
+
+!!! note
+    This functionality is provided by the `ReinitializationExt` module, which
+    requires loading `Interpolations.jl` and `NearestNeighbors.jl`.
+"""
+function reinitialize! end
+
+# tomatoes tomatos ...
+reinitialise!(args...; kwargs...) = reinitialize!(args...; kwargs...)
+
+
 end # module