From a8b221ec356a6115406d21497371bd9dec3294cc Mon Sep 17 00:00:00 2001 From: Alec Hammond Date: Fri, 15 Dec 2023 10:30:23 -0800 Subject: [PATCH 1/3] add subpixel smoothing routine for density TO --- python/adjoint/filters.py | 144 ++++++++++++++++++++++++++++++++++++++ 1 file changed, 144 insertions(+) diff --git a/python/adjoint/filters.py b/python/adjoint/filters.py index 1b03c0b7f..3ad1fe93b 100644 --- a/python/adjoint/filters.py +++ b/python/adjoint/filters.py @@ -1099,3 +1099,147 @@ def gray_indicator(x): density-based topology optimization. Archive of Applied Mechanics, 86(1-2), 189-218. """ return npa.mean(4 * x.flatten() * (1 - x.flatten())) * 100 + +def hybrid_levelset( + x: ArrayLikeType, + beta: float, + eta: float, + Lx: float, + Ly: float, + resolution: float, + filter_radius: float, + periodic_axes: ArrayLikeType = None, +): + """Filter and project using subpixel smoothing, which allows for β→∞. + + 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, first-order subpixel smoothing is applied to the interfaces (if + any are present). + + 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. + + Args: + x: The (2D) input design parameters. + 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]. + Lx: The length of design region in X direction (in Meep units). + Ly: The length of design region in Y direction (in Meep units). + resolution: The resolution of the design grid (not the Meep grid + resolution). + filter_radius: The filter radius (in Meep units). + periodic_axes: The list of axes (x, y = 0, 1) that are to be treated as + periodic. Default is None (all axes are non-periodic). + 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 + >>> A = onp.random.rand(Nx, Ny) + >>> beta = npa.inf + >>> A_smoothed = hybrid_levelset(A, beta, eta_i, Lx, Ly, resolution, filter_radius) + """ + # Note that currently, the underlying assumption is that the smoothing + # kernel is a circle, which means dx = dy. + dx = dy = 1 / resolution + pixel_radius = dx / 2 + + # Perform the standard filter and projection common in density-based TO + x_filtered = conic_filter( + x, + radius=filter_radius, + Lx=Lx, + Ly=Ly, + resolution=resolution, + periodic_axes=periodic_axes, + ) + x_projected = npa.clip(tanh_projection(x_filtered, beta=beta, eta=eta), 0, 1) + + # 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. + x_grad = npa.gradient(x_filtered) + x_grad_norm = npa.sqrt((x_grad[0] / dx) ** 2 + (x_grad[1] / dy) ** 2) + + # The distance for the center of the pixel to the nearest interface + d = (eta - x_filtered) / x_grad_norm + + # The fill factor is used to perform simple, first-order subpixel smoothing. + # We use the (2D) analytic expression that comes when assuming the smoothing + # kernel is a cylinder. 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 + fill_factor = npa.where( + (d <= pixel_radius) & (npa.abs(d / pixel_radius) <= 1), # domain of arccos() and sqrt() + (1 / (npa.pi * pixel_radius**2)) + * ( + pixel_radius**2 + * npa.arccos( + # "double where" trick + npa.where( + npa.abs(d / pixel_radius) < 1, + d / pixel_radius, + 0.0, + ) + ) + - d + * npa.sqrt( + # "double where" trick + npa.where( + (pixel_radius**2 - d**2) > 0, pixel_radius**2 - d**2, 0 + ) + ) + ), + 1, + ) + + # Determine the upper and lower bounds of materials in the current pixel. + x_minus = x_projected - x_grad_norm * pixel_radius + x_plus = x_projected + x_grad_norm * pixel_radius + + # Create an "effective" set of materials that will ensure everything is + # piecewise differentiable, even if an interface moves out of a pixel, or + # through the pixel center. + x_minus_eff = npa.where( + d > 0, + (x_filtered * d + x_minus * (pixel_radius - d)) / pixel_radius, + x_minus, + ) + x_plus_eff = npa.where( + d > 0, + x_plus, + (-x_filtered * d + x_plus * (pixel_radius + d)) / pixel_radius, + ) + + # Only apply smoothing to interfaces + return npa.where( + npa.abs(d) > pixel_radius, + x_projected, + (1 - fill_factor) * x_minus_eff + (fill_factor) * x_plus_eff, + ) \ No newline at end of file From 16702c5a0e4bffc2488945f3f3f59ceeb4655163 Mon Sep 17 00:00:00 2001 From: Alec Hammond Date: Fri, 15 Dec 2023 10:46:26 -0800 Subject: [PATCH 2/3] lint --- python/adjoint/filters.py | 8 +++++--- 1 file changed, 5 insertions(+), 3 deletions(-) diff --git a/python/adjoint/filters.py b/python/adjoint/filters.py index 3ad1fe93b..99d44ab99 100644 --- a/python/adjoint/filters.py +++ b/python/adjoint/filters.py @@ -1100,6 +1100,7 @@ def gray_indicator(x): """ return npa.mean(4 * x.flatten() * (1 - x.flatten())) * 100 + def hybrid_levelset( x: ArrayLikeType, beta: float, @@ -1147,7 +1148,7 @@ def hybrid_levelset( periodic. Default is None (all axes are non-periodic). Returns: The projected and smoothed output. - + Example: >>> Lx = 2 >>> Ly = 2 @@ -1196,7 +1197,8 @@ def hybrid_levelset( # with array-based AD tracers, apparently. See here: # https://github.com/google/jax/issues/1052#issuecomment-5140833520 fill_factor = npa.where( - (d <= pixel_radius) & (npa.abs(d / pixel_radius) <= 1), # domain of arccos() and sqrt() + (d <= pixel_radius) + & (npa.abs(d / pixel_radius) <= 1), # domain of arccos() and sqrt() (1 / (npa.pi * pixel_radius**2)) * ( pixel_radius**2 @@ -1242,4 +1244,4 @@ def hybrid_levelset( npa.abs(d) > pixel_radius, x_projected, (1 - fill_factor) * x_minus_eff + (fill_factor) * x_plus_eff, - ) \ No newline at end of file + ) From 7fb6c6b3213c844f8d4ff711b733eca4bb009554 Mon Sep 17 00:00:00 2001 From: Alec Hammond Date: Thu, 11 Jan 2024 16:31:52 -0800 Subject: [PATCH 3/3] fix backprop and norm --- python/adjoint/filters.py | 293 +++++++++++++++++++------------------- 1 file changed, 147 insertions(+), 146 deletions(-) diff --git a/python/adjoint/filters.py b/python/adjoint/filters.py index 99d44ab99..12d406643 100644 --- a/python/adjoint/filters.py +++ b/python/adjoint/filters.py @@ -700,6 +700,153 @@ def tanh_projection(x: np.ndarray, beta: float, eta: float) -> np.ndarray: ) +def smoothed_projection( + x_smoothed: ArrayLikeType, + beta: float, + eta: float, + resolution: float, +): + """Project using subpixel smoothing, which allows for β→∞. + + 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, first-order 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. + + Args: + x: The (2D) input design parameters. + 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 (not the Meep grid + resolution). + 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 + >>> A = onp.random.rand(Nx, Ny) + >>> beta = npa.inf + >>> A_smoothed = conic_filter(A, filter_radius, Lx, Ly, resolution) + >>> A_projected = smoothed_projection(A_smoothed, beta, eta_i, resolution) + """ + # Note that currently, the underlying assumption is that the smoothing + # kernel is a circle, which means dx = dy. + dx = dy = 1 / resolution + pixel_radius = dx / 2 + + x_projected = tanh_projection(x_smoothed, 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. + x_grad = npa.gradient(x_smoothed) + x_grad_helper = (x_grad[0] / dx) ** 2 + (x_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 = npa.abs(x_grad_helper) > 0 + + x_grad_norm = npa.sqrt(npa.where(nonzero_norm, x_grad_helper, 1)) + x_grad_norm_eff = npa.where(nonzero_norm, x_grad_norm, 1) + + # The distance for the center of the pixel to the nearest interface + d = (eta - x_smoothed) / x_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 & (npa.abs(d) <= pixel_radius) + + # The fill factor is used to perform simple, first-order 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 + + arccos_term = pixel_radius**2 * npa.arccos( + npa.where( + needs_smoothing, + d / pixel_radius, + 0.0, + ) + ) + + sqrt_term = d * npa.sqrt( + npa.where( + needs_smoothing, + pixel_radius**2 - d**2, + 1, + ) + ) + + fill_factor = npa.where( + needs_smoothing, + (1 / (npa.pi * pixel_radius**2)) * (arccos_term - sqrt_term), + 1, + ) + + # Determine the upper and lower bounds of materials in the current pixel. + x_minus = x_smoothed - x_grad_norm * pixel_radius + x_plus = x_smoothed + x_grad_norm * pixel_radius + + # Create an "effective" set of materials that will ensure everything is + # piecewise differentiable, even if an interface moves out of a pixel, or + # through the pixel center. + x_minus_eff_pert = (x_smoothed * d + x_minus * (pixel_radius - d)) / pixel_radius + x_minus_eff = npa.where( + (d > 0), + x_minus_eff_pert, + x_minus, + ) + x_plus_eff_pert = (-x_smoothed * d + x_plus * (pixel_radius + d)) / pixel_radius + x_plus_eff = npa.where( + (d > 0), + x_plus, + x_plus_eff_pert, + ) + + # Only apply smoothing to interfaces + x_projected_smoothed = (1 - fill_factor) * x_minus_eff + (fill_factor) * x_plus_eff + return npa.where( + needs_smoothing, + x_projected_smoothed, + x_projected, + ) + + def heaviside_projection(x, beta, eta): """Projection filter that thresholds the input parameters between 0 and 1. @@ -1099,149 +1246,3 @@ def gray_indicator(x): density-based topology optimization. Archive of Applied Mechanics, 86(1-2), 189-218. """ return npa.mean(4 * x.flatten() * (1 - x.flatten())) * 100 - - -def hybrid_levelset( - x: ArrayLikeType, - beta: float, - eta: float, - Lx: float, - Ly: float, - resolution: float, - filter_radius: float, - periodic_axes: ArrayLikeType = None, -): - """Filter and project using subpixel smoothing, which allows for β→∞. - - 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, first-order subpixel smoothing is applied to the interfaces (if - any are present). - - 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. - - Args: - x: The (2D) input design parameters. - 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]. - Lx: The length of design region in X direction (in Meep units). - Ly: The length of design region in Y direction (in Meep units). - resolution: The resolution of the design grid (not the Meep grid - resolution). - filter_radius: The filter radius (in Meep units). - periodic_axes: The list of axes (x, y = 0, 1) that are to be treated as - periodic. Default is None (all axes are non-periodic). - 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 - >>> A = onp.random.rand(Nx, Ny) - >>> beta = npa.inf - >>> A_smoothed = hybrid_levelset(A, beta, eta_i, Lx, Ly, resolution, filter_radius) - """ - # Note that currently, the underlying assumption is that the smoothing - # kernel is a circle, which means dx = dy. - dx = dy = 1 / resolution - pixel_radius = dx / 2 - - # Perform the standard filter and projection common in density-based TO - x_filtered = conic_filter( - x, - radius=filter_radius, - Lx=Lx, - Ly=Ly, - resolution=resolution, - periodic_axes=periodic_axes, - ) - x_projected = npa.clip(tanh_projection(x_filtered, beta=beta, eta=eta), 0, 1) - - # 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. - x_grad = npa.gradient(x_filtered) - x_grad_norm = npa.sqrt((x_grad[0] / dx) ** 2 + (x_grad[1] / dy) ** 2) - - # The distance for the center of the pixel to the nearest interface - d = (eta - x_filtered) / x_grad_norm - - # The fill factor is used to perform simple, first-order subpixel smoothing. - # We use the (2D) analytic expression that comes when assuming the smoothing - # kernel is a cylinder. 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 - fill_factor = npa.where( - (d <= pixel_radius) - & (npa.abs(d / pixel_radius) <= 1), # domain of arccos() and sqrt() - (1 / (npa.pi * pixel_radius**2)) - * ( - pixel_radius**2 - * npa.arccos( - # "double where" trick - npa.where( - npa.abs(d / pixel_radius) < 1, - d / pixel_radius, - 0.0, - ) - ) - - d - * npa.sqrt( - # "double where" trick - npa.where( - (pixel_radius**2 - d**2) > 0, pixel_radius**2 - d**2, 0 - ) - ) - ), - 1, - ) - - # Determine the upper and lower bounds of materials in the current pixel. - x_minus = x_projected - x_grad_norm * pixel_radius - x_plus = x_projected + x_grad_norm * pixel_radius - - # Create an "effective" set of materials that will ensure everything is - # piecewise differentiable, even if an interface moves out of a pixel, or - # through the pixel center. - x_minus_eff = npa.where( - d > 0, - (x_filtered * d + x_minus * (pixel_radius - d)) / pixel_radius, - x_minus, - ) - x_plus_eff = npa.where( - d > 0, - x_plus, - (-x_filtered * d + x_plus * (pixel_radius + d)) / pixel_radius, - ) - - # Only apply smoothing to interfaces - return npa.where( - npa.abs(d) > pixel_radius, - x_projected, - (1 - fill_factor) * x_minus_eff + (fill_factor) * x_plus_eff, - )