@@ -70,3 +70,125 @@ def tanh_projection(
7070 num = np .tanh (beta * eta ) + np .tanh (beta * (array - eta ))
7171 denom = np .tanh (beta * eta ) + np .tanh (beta * (1 - eta ))
7272 return num / denom
73+
74+
75+ def smoothed_projection (
76+ array : NDArray ,
77+ beta : float = BETA_DEFAULT ,
78+ eta : float = ETA_DEFAULT ,
79+ ) -> NDArray :
80+ """
81+ The details of this projection are described in the paper by Alec Hammond:
82+ https://arxiv.org/pdf/2503.20189
83+
84+ Project using subpixel smoothing, which allows for β→∞.
85+ This technique integrates out the discontinuity within the projection
86+ function, allowing the user to smoothly increase β from 0 to ∞ without
87+ losing the gradient. Effectively, a level set is created, and from this
88+ level set, first-order subpixel smoothing is applied to the interfaces (if
89+ any are present).
90+
91+ In order for this to work, the input array must already be smooth (e.g. by
92+ filtering).
93+
94+ While the original approach involves numerical quadrature, this approach
95+ performs a "trick" by assuming that the user is always infinitely projecting
96+ (β=∞). In this case, the expensive quadrature simplifies to an analytic
97+ fill-factor expression. When to use this fill factor requires some careful
98+ logic.
99+
100+ For one, we want to make sure that the user can indeed project at any level
101+ (not just infinity). So in these cases, we simply check if in interface is
102+ within the pixel. If not, we revert to the standard filter plus project
103+ technique.
104+
105+ If there is an interface, we want to make sure the derivative remains
106+ continuous both as the interface leaves the cell, *and* as it crosses the
107+ center. To ensure this, we need to account for the different possibilities.
108+
109+ Parameters
110+ ----------
111+ array : NDArray
112+ The (2D) input design parameters (already filered e.g. with a conic filter)
113+ beta : float = BETA_DEFAULT
114+ The thresholding parameter in the range [0, inf]. Determines the degree of binarization of the output.
115+ eta : float = ETA_DEFAULT
116+ The threshold point / midpoint of the projection.
117+
118+ Returns
119+ -------
120+ NDArray
121+ The array after applying the smoothed projection.
122+
123+ Example:
124+ >>> import autograd.numpy as np
125+ >>> from tidy3d.plugins.autograd.invdes.filters import ConicFilter
126+ >>> arr = np.random.uniform(size=(50, 50))
127+ >>> filter = ConicFilter(kernel_size=5)
128+ >>> arr_filtered = filter(arr)
129+ >>> eta = 0.5 # center of projection
130+ >>> smoothed = smoothed_projection(arr, beta=np.inf, eta=eta)
131+ """
132+ # sanity checks
133+ assert array .ndim == 2
134+ # Note that currently, the underlying assumption is that the smoothing
135+ # kernel is a circle, which means dx = dy.
136+ dx = dy = 1
137+ R_smoothing = np .sqrt (1 / np .pi ) # unit area of smoothing kernel
138+
139+ rho_projected = tanh_projection (array , beta = beta , eta = eta )
140+
141+ # Compute the spatial gradient (using finite differences) of the *filtered*
142+ # field, which will always be smooth and is the key to our approach. This
143+ # gradient essentially represents the normal direction pointing the the
144+ # nearest inteface.
145+ rho_filtered_grad = np .gradient (array )
146+ rho_filtered_grad_helper = (rho_filtered_grad [0 ] / dx ) ** 2 + (rho_filtered_grad [1 ] / dy ) ** 2
147+
148+ # Note that a uniform field (norm=0) is problematic, because it creates
149+ # divide by zero issues and makes backpropagation difficult, so we sanitize
150+ # and determine where smoothing is actually needed. The value where we don't
151+ # need smoothings doesn't actually matter, since all our computations are
152+ # purely element-wise (no spatial locality) and those pixels will instead
153+ # rely on the standard projection. So just use 1, since it's well behaved.
154+ nonzero_norm = np .abs (rho_filtered_grad_helper ) > 1e-10
155+
156+ rho_filtered_grad_norm = np .sqrt (np .where (nonzero_norm , rho_filtered_grad_helper , 1 ))
157+ rho_filtered_grad_norm_eff = np .where (nonzero_norm , rho_filtered_grad_norm , 1 )
158+
159+ # The distance for the center of the pixel to the nearest interface
160+ d = (eta - array ) / rho_filtered_grad_norm_eff
161+
162+ # Only need smoothing if an interface lies within the voxel. Since d is
163+ # actually an "effective" d by this point, we need to ignore values that may
164+ # have been sanitized earlier on.
165+ needs_smoothing = nonzero_norm & (np .abs (d ) < R_smoothing )
166+
167+ # The fill factor is used to perform simple, first-order subpixel smoothing.
168+ # We use the (2D) analytic expression that comes when assuming the smoothing
169+ # kernel is a circle. Note that because the kernel contains some
170+ # expressions that are sensitive to NaNs, we have to use the "double where"
171+ # trick to avoid the Nans in the backward trace. This is a common problem
172+ # with array-based AD tracers, apparently. See here:
173+ # https://github.com/google/jax/issues/1052#issuecomment-5140833520
174+ d_R = d / R_smoothing
175+ F = np .where (needs_smoothing , 0.5 - 15 / 16 * d_R + 5 / 8 * d_R ** 3 - 3 / 16 * d_R ** 5 , 1.0 )
176+ # F(-d)
177+ F_minus = np .where (needs_smoothing , 0.5 + 15 / 16 * d_R - 5 / 8 * d_R ** 3 + 3 / 16 * d_R ** 5 , 1.0 )
178+
179+ # Determine the upper and lower bounds of materials in the current pixel (before projection).
180+ rho_filtered_minus = array - R_smoothing * rho_filtered_grad_norm_eff * F
181+ rho_filtered_plus = array + R_smoothing * rho_filtered_grad_norm_eff * F_minus
182+
183+ # Finally, we project the extents of our range.
184+ rho_minus_eff_projected = tanh_projection (rho_filtered_minus , beta = beta , eta = eta )
185+ rho_plus_eff_projected = tanh_projection (rho_filtered_plus , beta = beta , eta = eta )
186+
187+ # Only apply smoothing to interfaces
188+ rho_projected_smoothed = (1 - F ) * rho_minus_eff_projected + F * rho_plus_eff_projected
189+ smoothed = np .where (
190+ needs_smoothing ,
191+ rho_projected_smoothed ,
192+ rho_projected ,
193+ )
194+ return smoothed
0 commit comments