Parameterizations#

API reference for the design-variable parameterizations used in inverse design. A parameterization maps a small set of design variables onto a field of the state – usually the material density State.rho – in a differentiable way, so that State.resolve() yields a pure function of the design variables that jax.grad / torch.autograd can differentiate.

Class-Reference#

class neuralmag.LevelSetParameterization(num_rbfs_x, num_rbfs_y, *, xlimits=None, ylimits=None, geometry=None, rbf_c=0.0001, gaussian_sigma=0.0, gaussian_radius=None, sigmoid_a=20.0, smooth_max_p=90.0, eikonal=True)#

Differentiable RBF level-set parameterization of the material density.

Parameterizes State.rho over a rectangular design region of a 2D mesh by a grid of radial basis functions whose amplitudes are the design variables. Each RBF contributes a cone

\[\phi_r(\vec{x}) = s_r - \sqrt{(x - x_r)^2 + (y - y_r)^2 + c^2},\]

the cones are combined into one level-set field by a smooth maximum

\[\phi(\vec{x}) = \left( \sum_r \tilde\phi_r(\vec{x})^p \right)^{1/p} - \Delta\phi, \qquad \tilde\phi_r = \phi_r + \Delta\phi,\]

evaluated through logsumexp for numerical stability (the shift \(\Delta\phi = |\min_r \phi_r|\) makes the base positive), optionally Gaussian-filtered to impose a minimum feature size, renormalized to unit slope (eikonal=True, the default)

\[\hat\phi(\vec{x}) = \frac{\phi(\vec{x})}{\sqrt{|\nabla \phi(\vec{x})|^2 + \epsilon}},\]

and finally projected onto \((0, 1)\) by a sigmoid of sharpness \(a\)

\[\rho(\vec{x}) = \operatorname{sigmoid}\big(a \, \hat\phi(\vec{x})\big).\]

The renormalization is what makes \(a\) the inverse interface width: the zero set of \(\phi\) is unchanged (the divisor is strictly positive), but its slope becomes 1 cell^-1 there, so the 10–90 % transition of \(\rho\) spans \(\approx 4.4 / a\) cells regardless of the design. Without it the Gaussian filter and the ridges of the smooth maximum flatten \(|\nabla\phi|\) by a design-dependent factor (0.04–0.31 measured across production designs, with a 2–3x spread within a single design), so the same \(a\) yields different interface widths and a continuation in \(a\) sharpens far less than it appears to. Pass eikonal=False for the raw projection \(\operatorname{sigmoid}(a\phi)\).

Outside the design region \(\rho\) is 1, and the whole field is multiplied by the fixed geometry mask, so a single design region can be embedded in an arbitrary fixed layout. The result is clipped to [state.eps, 1], never to exactly 0 (see State.rho).

Raising a during optimization (a continuation scheme) drives \(\rho\) from a smooth, well-conditioned field toward a binary material distribution. state.lsf_sigmoid_a is a state attribute rather than a captured constant precisely so it can be ramped – or traced as a solver parameter – without re-registering.

Note

Lengths here are cell indices, not metres: xlimits, ylimits, gaussian_sigma and gaussian_radius all count mesh cells, because the design grid is defined against the mesh rather than against physical space.

Warning

On the JAX backend, run in float64 (nm.config.dtype = "float64" before any tensor is created, or NM_DTYPE=float64): reverse-mode gradients through the LLG time integration have been observed to be NaN under float32 while the forward pass looks healthy, which silently breaks an optimization. register() warns if a JAX state is float32. On the torch backend the gradients flow normally under float32; note however that the smooth maximum evaluates p * log(...) with p of order 100, so in float32 the level set loses some accuracy near a cone’s apex on either backend.

Parameters:

num_rbfs_x (int) – Number of RBF centers along x.
num_rbfs_y (int) – Number of RBF centers along y.
xlimits (tuple of int, optional) – (start, stop) cell index range of the design region along x, defaults to the full mesh.
ylimits (tuple of int, optional) – (start, stop) cell index range of the design region along y, defaults to the full mesh.
geometry (array_like, optional) – Fixed mask of shape mesh.n multiplied onto the parameterized density, e.g. 1 in the allowed layout and state.eps in permanent voids. Defaults to all ones.
rbf_c (float, optional) – Regularization \(c\) of the cone tip, keeping the square root differentiable at the RBF center.
gaussian_sigma (float, optional) – Standard deviation of the Gaussian filter in cells. 0 disables filtering.
gaussian_radius (int, optional) – Half-width of the filter kernel in cells, defaults to ceil(3 * gaussian_sigma).
sigmoid_a (float, optional) – Initial projection sharpness \(a\).
eikonal (bool, optional) – Renormalize the level-set field to unit gradient magnitude before the sigmoid (default True), so that sigmoid_a sets the interface width in cells (\(\approx 4.4 / a\)) independently of gaussian_sigma and of the design. False projects the raw field.
smooth_max_p (float, optional) – Exponent \(p\) of the smooth maximum. Larger values approximate the true maximum more closely at the cost of conditioning.
attributes (Registered state)
---------------------------
variables (* state.rbf_amplitudes -- the design) – (num_rbfs_x, num_rbfs_y); initialized to zeros if not already set.
shape – (num_rbfs_x, num_rbfs_y); initialized to zeros if not already set.

:param * state.lsf_sigmoid_a – the projection sharpness \(a\).: :param * state.lsf_geometry – the fixed mask.: :param * state.rho – the resulting CellFunction.:

Examples

>>> import neuralmag as nm
>>> state = nm.State(nm.Mesh((60, 40), (5e-9, 5e-9, 3e-9)))
>>> lsf = nm.LevelSetParameterization(12, 8, xlimits=(10, 50), ylimits=(5, 35))
>>> lsf.register(state)
>>> # pure, differentiable function of the design variables
>>> rho_func = state.resolve("rho", ["rbf_amplitudes"])
>>> rho = rho_func(state.rbf_amplitudes)

property num_rbfs#: The shape (num_rbfs_x, num_rbfs_y) of the design variables.

register(state, name='rho')#

Everything that cannot change during an optimization – the RBF centers, the cone distances, the filter kernel and its gather indices, the constant material outside the design region – is precomputed here and captured, so each evaluation is a handful of tensor operations on the design variables. In particular the RBF distance field is built once instead of per forward pass; it costs num_rbfs_x * num_rbfs_y * nx_design * ny_design values of memory.

Parameters:

state (State) – The state to register with. Must carry a 2D mesh.
name (str, optional) – The attribute the parameterized density is bound to.