Merge pull request #1 from NanoComp/SSP_original

Set up Python package, example optimization, and SSP first-order tests

Author: stevengj

.gitignore

@@ -0,0 +1,19 @@
+# macOS
+.DS_Store
+
+# Python bytecode and caches
+__pycache__/
+*.py[cod]
+
+# Packaging/build artifacts
+*.egg-info/
+build/
+dist/
+
+# Test/tool caches
+.pytest_cache/
+.mypy_cache/
+
+# Virtual environments
+.venv/
+venv/

examples/python/ssp1_bilinear_example.py

@@ -0,0 +1,167 @@
+"""
+This is a simple example that demonstrates how to use the original "SSP1" algorithm
+for subpixel-smoothed projection, combined with bilinear interpolation and conic smoothing.
+We set up a simplistic gradient-based optimization problem that attempts
+to drive the mean of the output to zero. We set β=∞ and show the gradient is nonzero.
+"""
+
+import time
+
+import nlopt
+import numpy as np
+from jax import grad, jit, value_and_grad
+from jax import numpy as jnp
+from matplotlib import pyplot as plt
+
+from ssp import conic_filter, get_conic_radius_from_eta_e, ssp1_bilinear
+
+
+def figure_of_merit(x: jnp.ndarray) -> float:
+    """A simple, convex reduction mean as the FOM.
+
+    FOM = 1/(n*m)ΣΣ|x|^2
+    """
+    return jnp.mean((jnp.abs(x) ** 2).flatten())
+
+
+def full_system(x: jnp.ndarray, beta, eta_i, resolution) -> float:
+    """Include the projection and FOM"""
+    return figure_of_merit(ssp1_bilinear(x, beta, eta_i, resolution))
+
+
+def main():
+    """Run the example and optimization."""
+    
+    # --------------------------------------------- #
+    # Visualize the SSP transformation
+    # --------------------------------------------- #
+
+    # First set up a random initial condition. We'll use dimensionless units for everything.
+    lx = 2.0
+    ly = 2.0
+    resolution = 50
+    eta_i = 0.5
+    eta_e = 0.75
+    lengthscale = 0.25
+    filter_radius = get_conic_radius_from_eta_e(lengthscale, eta_e)
+    nx = int(np.round(lx * resolution) + 1)
+    ny = int(np.round(ly * resolution) + 1)
+
+    np.random.seed(42)
+    rho = np.random.rand(nx, ny)
+    beta = np.inf
+
+    rho_filtered = conic_filter(rho, filter_radius, lx, ly, resolution)
+    rho_projected = ssp1_bilinear(rho_filtered, beta, eta_i, resolution)
+
+    plt.figure(figsize=(6, 3))
+    plt.subplot(1, 2, 1)
+    plt.imshow(rho_filtered, vmin=0, vmax=1, cmap="binary")
+    plt.title("Input")
+    plt.colorbar()
+    plt.subplot(1, 2, 2)
+    plt.imshow(rho_projected, vmin=0, vmax=1, cmap="binary")
+    plt.colorbar()
+    plt.title("SSP Output")
+    plt.tight_layout()
+    plt.savefig("projection.png")
+
+    # --------------------------------------------- #
+    # Visualize the SSP gradient for β=∞
+    # --------------------------------------------- #
+
+    d_fom = grad(full_system)
+    df_drho = d_fom(rho_filtered, beta, eta_i, resolution)
+    max_val = jnp.max(df_drho)
+    min_val = jnp.min(df_drho)
+    vmax_vmin = max([abs(max_val), abs(min_val)])
+
+    plt.figure(figsize=(6, 3))
+    plt.subplot(1, 2, 1)
+    plt.imshow(rho_projected, vmin=0, vmax=1, cmap="binary")
+    plt.title("SSP Output")
+    plt.colorbar()
+    plt.subplot(1, 2, 2)
+    plt.imshow(df_drho, vmin=-vmax_vmin, vmax=vmax_vmin, cmap="RdBu")
+    plt.colorbar()
+    plt.title("Gradient")
+    plt.tight_layout()
+    plt.savefig("gradient.png")
+
+
+    # --------------------------------------------- #
+    # Shape optimization via nlopt
+    # --------------------------------------------- #
+
+    def optimization_objective(rho_flat: jnp.ndarray) -> jnp.ndarray:
+        rho_design = rho_flat.reshape((nx, ny))
+        rho_design_filtered = conic_filter(rho_design, filter_radius, lx, ly, resolution)
+        rho_design_projected = ssp1_bilinear(
+            rho_design_filtered, beta, eta_i, resolution
+        )
+        return figure_of_merit(rho_design_projected)
+
+    objective_and_grad = jit(value_and_grad(optimization_objective))
+
+    iteration_history = []
+    time_history = []
+    fom_history = []
+    start_time = time.perf_counter()
+
+    def nlopt_objective(x: np.ndarray, grad_out: np.ndarray) -> float:
+        fom_value, gradient = objective_and_grad(jnp.asarray(x))
+
+        if grad_out.size > 0:
+            grad_out[:] = np.asarray(gradient, dtype=float)
+
+        iteration = len(fom_history) + 1
+        elapsed = time.perf_counter() - start_time
+        fom_scalar = float(fom_value)
+
+        iteration_history.append(iteration)
+        time_history.append(elapsed)
+        fom_history.append(fom_scalar)
+
+        print(f"iter={iteration:03d} time={elapsed:8.3f}s FOM={fom_scalar:.6e}")
+        return fom_scalar
+
+    opt = nlopt.opt(nlopt.LD_CCSAQ, nx * ny)
+    opt.set_lower_bounds(np.zeros(nx * ny))
+    opt.set_upper_bounds(np.ones(nx * ny))
+    opt.set_min_objective(nlopt_objective)
+    opt.set_maxeval(25)
+
+    x_opt = opt.optimize(rho.ravel())
+    final_fom = opt.last_optimum_value()
+
+    rho_opt = x_opt.reshape((nx, ny))
+    rho_opt_filtered = conic_filter(rho_opt, filter_radius, lx, ly, resolution)
+    rho_opt_projected = np.asarray(
+        ssp1_bilinear(rho_opt_filtered, beta, eta_i, resolution)
+    )
+
+    print("\nOptimization complete")
+    print(f"status={opt.last_optimize_result()} final_fom={final_fom:.6e}")
+
+    print("\nIteration log (time in seconds):")
+    for iteration, elapsed, fom in zip(iteration_history, time_history, fom_history):
+        print(f"iter={iteration:03d} time={elapsed:8.3f}s FOM={fom:.6e}")
+
+    plt.figure(figsize=(6, 3))
+    plt.subplot(1, 2, 1)
+    plt.semilogy(iteration_history, fom_history, marker="o", linewidth=1.5)
+    plt.xlabel("Iteration")
+    plt.ylabel("FOM")
+    plt.title("FOM vs Iteration")
+
+    plt.subplot(1, 2, 2)
+    im = plt.imshow(rho_opt_projected, vmin=0, vmax=1, cmap="binary")
+    plt.title("Final SSP Projected Design")
+    plt.colorbar(im)
+
+    plt.tight_layout()
+    plt.savefig("optimization.png")
+
+
+if __name__ == "__main__":
+    main()

pyproject.toml

@@ -0,0 +1,30 @@
+[build-system]
+requires = ["setuptools>=68", "wheel"]
+build-backend = "setuptools.build_meta"
+
+[project]
+name = "ssp"
+version = "0.1.0"
+description = "Smoothed subpixel projection (SSP) for topology optimization"
+readme = "README.md"
+requires-python = ">=3.10"
+license = { file = "LICENSE" }
+authors = [
+  { name = "SSP Contributors" }
+]
+dependencies = [
+  "numpy>=1.24",
+  "jax>=0.4",
+]
+
+[project.optional-dependencies]
+examples = [
+  "matplotlib>=3.8",
+  "nlopt>=2.8",
+]
+
+[tool.setuptools]
+include-package-data = true
+
+[tool.setuptools.packages.find]
+where = ["src/python"]

src/python/ssp/__init__.py

@@ -0,0 +1,11 @@
+"""Public Python API for smoothed subpixel projection (SSP)."""
+
+from .core import ssp1_bilinear
+from .utils import conic_filter, get_conic_radius_from_eta_e, tanh_projection
+
+__all__ = [
+    "ssp1_bilinear",
+    "conic_filter",
+    "get_conic_radius_from_eta_e",
+    "tanh_projection",
+]

src/python/ssp/core.py

@@ -0,0 +1,155 @@
+"""
+These are the core algorithms and routines behind SSP.
+"""
+
+from jax import numpy as jnp
+
+from .utils import ArrayLikeType, tanh_projection
+
+# Naively we might pick a smoothing radius of exactly 0.5 (such that the diameter 
+# of the smoothing kernel spans the full pixel/voxel). However, in the case when 
+# a geometry interface is defined _right_ on the edge of a voxel (which often 
+# happens when the user initializes a geometry with a specific design), no smoothing 
+# occurs. By extending the kernel slightly beyond the edge of the pixel, we account 
+# for these corner cases.
+DEFAULT_SMOOTHING_RADIUS = 0.55
+
+def ssp1_bilinear(
+    rho_filtered: ArrayLikeType,
+    beta: float,
+    eta: float,
+    resolution: float,
+    smoothing_radius: float = DEFAULT_SMOOTHING_RADIUS
+):
+    """Project using the original SSP1 algorithm with bilinear interpolation.
+
+    This technique integrates out the discontinuity within the projection
+    function, allowing the user to smoothly increase β from 0 to ∞ without
+    losing the gradient. Effectively, a level set is created, and from this
+    level set, SSP1 bilinear subpixel smoothing is applied to the interfaces (if
+    any are present).
+
+    In order for this to work, the input array must already be smooth (e.g. by
+    filtering).
+
+    While the original approach involves numerical quadrature, this approach
+    performs a "trick" by assuming that the user is always infinitely projecting
+    (β=∞). In this case, the expensive quadrature simplifies to an analytic
+    fill-factor expression. When to use this fill factor requires some careful
+    logic.
+
+    For one, we want to make sure that the user can indeed project at any level
+    (not just infinity). So in these cases, we simply check if in interface is
+    within the pixel. If not, we revert to the standard filter plus project
+    technique.
+
+    If there is an interface, we want to make sure the derivative remains
+    continuous both as the interface leaves the cell, *and* as it crosses the
+    center. To ensure this, we need to account for the different possibilities.
+
+    Ref: A. M. Hammond, A. Oskooi, I. M. Hammond, M. Chen, S. E. Ralph, and 
+    S. G. Johnson, “Unifying and accelerating level-set and density-based topology 
+    optimization by subpixel-smoothed projection,” Optics Express, vol. 33, 
+    pp. 33620-33642, July 2025. Editor's Pick.
+
+    Args:
+        rho_filtered: The (2D) input design parameters (already filered e.g.
+            with a conic filter).
+        beta: The thresholding parameter in the range [0, inf]. Determines the
+            degree of binarization of the output.
+        eta: The threshold point in the range [0, 1].
+        resolution: resolution of the design grid.
+        smoothing_radius: The smoothing radius of the kernel relative to the 
+            pixel/voxel "width."
+    Returns:
+        The projected and smoothed output.
+
+    Example:
+        >>> Lx = 2; Ly = 2
+        >>> resolution = 50
+        >>> eta_i = 0.5; eta_e = 0.75
+        >>> lengthscale = 0.1
+        >>> filter_radius = get_conic_radius_from_eta_e(lengthscale, eta_e)
+        >>> Nx = onp.round(Lx * resolution) + 1
+        >>> Ny = onp.round(Ly * resolution) + 1
+        >>> rho = onp.random.rand(Nx, Ny)
+        >>> beta = npa.inf
+        >>> rho_filtered = conic_filter(rho, filter_radius, Lx, Ly, resolution)
+        >>> rho_projected = smoothed_projection(rho_filtered, beta, eta_i, resolution)
+    """
+    # TODO [alechammond] The current implementation is ported from meep 
+    # and only supports 2D inputs. We'll want to generalize this to 
+    # arbitrary dimensions.
+
+    # TODO [alechammond] Note that currently, the underlying assumption 
+    # is that the smoothing kernel is a circle, which means dx = dy.
+    dx = dy = 1 / resolution
+    R_smoothing = smoothing_radius * dx
+
+    rho_projected = tanh_projection(rho_filtered, beta=beta, eta=eta)
+
+    # Compute the spatial gradient (using finite differences) of the *filtered*
+    # field, which will always be smooth and is the key to our approach. This
+    # gradient essentially represents the normal direction pointing the the
+    # nearest inteface.
+    rho_filtered_grad = jnp.gradient(rho_filtered)
+    rho_filtered_grad_helper = (rho_filtered_grad[0] / dx) ** 2 + (
+        rho_filtered_grad[1] / dy
+    ) ** 2
+
+    # Note that a uniform field (norm=0) is problematic, because it creates
+    # divide by zero issues and makes backpropagation difficult, so we sanitize
+    # and determine where smoothing is actually needed. The value where we don't
+    # need smoothings doesn't actually matter, since all our computations our
+    # purely element-wise (no spatial locality) and those pixels will instead
+    # rely on the standard projection. So just use 1, since it's well behaved.
+    nonzero_norm = jnp.abs(rho_filtered_grad_helper) > 0
+
+    rho_filtered_grad_norm = jnp.sqrt(
+        jnp.where(nonzero_norm, rho_filtered_grad_helper, 1)
+    )
+    rho_filtered_grad_norm_eff = jnp.where(nonzero_norm, rho_filtered_grad_norm, 1)
+
+    # The distance for the center of the pixel to the nearest interface
+    d = (eta - rho_filtered) / rho_filtered_grad_norm_eff
+
+    # Only need smoothing if an interface lies within the voxel. Since d is
+    # actually an "effective" d by this point, we need to ignore values that may
+    # have been sanitized earlier on.
+    needs_smoothing = nonzero_norm & (jnp.abs(d) < R_smoothing)
+
+    # The fill factor is used to perform the SSP1 bilinear subpixel smoothing.
+    # We use the (2D) analytic expression that comes when assuming the smoothing
+    # kernel is a circle. Note that because the kernel contains some
+    # expressions that are sensitive to NaNs, we have to use the "double where"
+    # trick to avoid the Nans in the backward trace. This is a common problem
+    # with array-based AD tracers, apparently. See here:
+    # https://github.com/google/jax/issues/1052#issuecomment-5140833520
+    d_R = d / R_smoothing
+    F = jnp.where(
+        needs_smoothing, 0.5 - 15 / 16 * d_R + 5 / 8 * d_R**3 - 3 / 16 * d_R**5, 1.0
+    )
+    # F(-d)
+    F_minus = jnp.where(
+        needs_smoothing, 0.5 + 15 / 16 * d_R - 5 / 8 * d_R**3 + 3 / 16 * d_R**5, 1.0
+    )
+
+    # Determine the upper and lower bounds of materials in the current pixel (before projection).
+    rho_filtered_minus = rho_filtered - R_smoothing * rho_filtered_grad_norm_eff * F
+    rho_filtered_plus = (
+        rho_filtered + R_smoothing * rho_filtered_grad_norm_eff * F_minus
+    )
+
+    # Finally, we project the extents of our range.
+    rho_minus_eff_projected = tanh_projection(rho_filtered_minus, beta=beta, eta=eta)
+    rho_plus_eff_projected = tanh_projection(rho_filtered_plus, beta=beta, eta=eta)
+
+    # Only apply smoothing to interfaces
+    rho_projected_smoothed = (
+        1 - F
+    ) * rho_minus_eff_projected + F * rho_plus_eff_projected
+    return jnp.where(
+        needs_smoothing,
+        rho_projected_smoothed,
+        rho_projected,
+    )