"""Reusable topology-optimization transforms and reverse-tape planning.
The functions in this module are deliberately independent of a particular
photonic device. They provide the common mechanics needed by density-based
inverse design: minimum-feature filtering, projection continuation,
erosion/dilation penalties, deterministic hard binarization, and an automatic
choice of reversible-recorder storage that respects the accelerator budget.
"""
from __future__ import annotations
import math
from collections.abc import Callable, Mapping
from dataclasses import dataclass
from functools import reduce
from operator import mul
from typing import Any, Literal, Protocol
import jax
import jax.numpy as jnp
import numpy as np
import optax
from jax.scipy.signal import convolve2d
from fdtdx.constants import c
from fdtdx.interfaces.modules import DtypeConversion
from fdtdx.interfaces.recorder import Recorder
from fdtdx.interfaces.time_filter import LinearReconstructEveryK
from fdtdx.typing import BackendOption
@dataclass(frozen=True)
class RecorderMemoryPlan:
"""Auditable storage decision for a reversible FDTD boundary tape."""
raw_bytes: int
estimated_bytes: int
budget_bytes: int
device_limit_bytes: int
storage_dtype: str
temporal_stride: int
latent_steps: int
@property
def compression_ratio(self) -> float:
return self.raw_bytes / max(self.estimated_bytes, 1)
[docs]
@dataclass(frozen=True)
class TopologyOptimizerConfig:
"""Device-independent bounded Adam continuation policy."""
iterations: int
continuation_steps: int | None = None
learning_rate: float = 0.1
final_learning_rate_fraction: float = 0.1
beta_start: float = 1.0
beta_end: float = 50.0
lower_bound: float = 0.0
upper_bound: float = 1.0
maximize: bool = True
early_stop_patience: int | None = None
early_stop_min_delta: float = 1e-4
algorithm: Literal["adam", "sgd", "lbfgs"] = "adam"
momentum: float = 0.9
lbfgs_memory_size: int = 10
@dataclass(frozen=True)
class TopologyStepRecord:
"""Small host-side diagnostic record for one objective evaluation."""
step: int
objective: float
beta: float
gradient_norm: float | None
accepted: bool
[docs]
@dataclass(frozen=True)
class TopologyOptimizationResult:
"""Parameters and diagnostics returned by :func:`optimize_topology`."""
final_parameters: np.ndarray
best_parameters: np.ndarray
best_objective: float
best_step: int
records: tuple[TopologyStepRecord, ...]
@dataclass(frozen=True)
class SafeguardedLBFGSDirection:
"""A limited-memory ascent proposal with auditable curvature health.
``relative_predicted_ascent`` compares the selected direction with a
diagonally scaled steepest-ascent direction after both are normalized to
the same peak coordinate motion. This is the normalization used by the
bounded trust steps in the inverse-design campaigns, so the comparison is
invariant to arbitrary L-BFGS scaling.
"""
direction: np.ndarray
method: str
retained_step_history: np.ndarray
retained_gradient_difference_history: np.ndarray
pruned_pairs: int
predicted_ascent: float
reference_ascent: float
relative_predicted_ascent: float
def safeguarded_lbfgs_ascent_direction(
gradient: np.ndarray,
step_history: np.ndarray,
gradient_difference_history: np.ndarray,
*,
diagonal_scale: np.ndarray | None = None,
trainable_mask: np.ndarray | None = None,
parameters: np.ndarray | None = None,
lower_bounds: np.ndarray | None = None,
upper_bounds: np.ndarray | None = None,
minimum_relative_ascent: float = 0.1,
minimum_secant_cosine: float | None = None,
) -> SafeguardedLBFGSDirection:
"""Return a curvature-accelerated direction only while it is healthy.
The routine is deliberately objective- and architecture-independent. It
filters numerically unreliable secant pairs, tests the full L-BFGS model,
and then removes the oldest pairs until the peak-normalized predicted
ascent is a useful fraction of current-gradient ascent. If no suffix is
healthy, it returns the current diagonally scaled gradient. All work is
host-side linear algebra; no objective or gradient evaluation is added.
``gradient_difference_history`` follows the maximization convention
``g_previous - g_current``, so the usual minimization L-BFGS recursion is
applied to ``-gradient``. Box bounds are optional. When supplied, an
outward direction on an active bound is removed before health is scored.
"""
gradient64 = np.asarray(gradient, dtype=np.float64)
if gradient64.ndim != 1:
raise ValueError("gradient must be one-dimensional")
steps = np.asarray(step_history, dtype=np.float64)
changes = np.asarray(gradient_difference_history, dtype=np.float64)
if steps.ndim != 2 or changes.ndim != 2 or steps.shape != changes.shape:
raise ValueError("L-BFGS histories must be equally shaped matrices")
if steps.shape[1] != gradient64.size:
raise ValueError("L-BFGS history width must match the gradient")
if not 0.0 < minimum_relative_ascent <= 1.0:
raise ValueError("minimum_relative_ascent must lie in (0, 1]")
if diagonal_scale is None:
scale = np.ones_like(gradient64)
else:
scale = np.asarray(diagonal_scale, dtype=np.float64)
if scale.shape != gradient64.shape:
raise ValueError("diagonal_scale must match the gradient")
if not np.all(np.isfinite(scale)) or np.any(scale <= 0.0):
raise ValueError("diagonal_scale must be finite and positive")
if trainable_mask is None:
mask = np.ones_like(gradient64)
else:
mask = np.asarray(trainable_mask, dtype=np.float64)
if mask.shape != gradient64.shape:
raise ValueError("trainable_mask must match the gradient")
bounded = parameters is not None or lower_bounds is not None or upper_bounds is not None
if bounded and (parameters is None or lower_bounds is None or upper_bounds is None):
raise ValueError("parameters and both bounds must be supplied together")
if bounded:
current = np.asarray(parameters, dtype=np.float64)
lower = np.asarray(lower_bounds, dtype=np.float64)
upper = np.asarray(upper_bounds, dtype=np.float64)
if current.shape != gradient64.shape or lower.shape != gradient64.shape or upper.shape != gradient64.shape:
raise ValueError("parameters and bounds must match the gradient")
secant_floor = (
float(np.sqrt(np.finfo(np.float32).eps))
if minimum_secant_cosine is None
else float(minimum_secant_cosine)
)
if not 0.0 <= secant_floor < 1.0:
raise ValueError("minimum_secant_cosine must lie in [0, 1)")
tiny = np.finfo(np.float64).tiny
def project(direction: np.ndarray) -> np.ndarray:
projected = np.asarray(direction, dtype=np.float64) * mask
if bounded:
projected = projected.copy()
projected[(current <= lower) & (projected < 0.0)] = 0.0
projected[(current >= upper) & (projected > 0.0)] = 0.0
return projected
def normalized_ascent(direction: np.ndarray) -> tuple[np.ndarray, float]:
projected = project(direction)
peak = float(np.max(np.abs(projected), initial=0.0))
if not np.isfinite(peak) or peak <= tiny:
return projected, 0.0
normalized = projected / peak
return projected, float(np.vdot(gradient64, normalized).real)
fallback, reference_ascent = normalized_ascent(gradient64 * scale)
empty = np.empty((0, gradient64.size), dtype=np.float64)
if not np.isfinite(reference_ascent) or reference_ascent <= tiny:
return SafeguardedLBFGSDirection(
direction=fallback,
method="projected_scaled_gradient_stationary",
retained_step_history=empty,
retained_gradient_difference_history=empty.copy(),
pruned_pairs=int(steps.shape[0]),
predicted_ascent=float(reference_ascent),
reference_ascent=float(reference_ascent),
relative_predicted_ascent=1.0,
)
def valid_pairs(start: int) -> tuple[np.ndarray, np.ndarray]:
kept_steps: list[np.ndarray] = []
kept_changes: list[np.ndarray] = []
for step, change in zip(steps[start:], changes[start:], strict=True):
curvature = float(np.dot(step, change))
norm_product = float(np.linalg.norm(step) * np.linalg.norm(change))
if (
np.all(np.isfinite(step))
and np.all(np.isfinite(change))
and curvature > secant_floor * max(norm_product, tiny)
):
kept_steps.append(step)
kept_changes.append(change)
if not kept_steps:
return empty, empty.copy()
return np.stack(kept_steps), np.stack(kept_changes)
def two_loop(pair_steps: np.ndarray, pair_changes: np.ndarray) -> np.ndarray:
value = -gradient64.copy()
alphas: list[float] = []
inverse_curvatures: list[float] = []
for step, change in zip(
reversed(pair_steps), reversed(pair_changes), strict=True
):
inverse_curvature = 1.0 / float(np.dot(step, change))
inverse_curvatures.append(inverse_curvature)
alpha = inverse_curvature * float(np.dot(step, value))
alphas.append(alpha)
value -= alpha * change
last_step = pair_steps[-1]
last_change = pair_changes[-1]
value *= float(
np.dot(last_step, last_change) / np.dot(last_change, last_change)
)
for step, change, inverse_curvature, alpha in zip(
pair_steps,
pair_changes,
reversed(inverse_curvatures),
reversed(alphas),
strict=True,
):
value += step * (
alpha - inverse_curvature * float(np.dot(change, value))
)
return -value
pair_count = int(steps.shape[0])
for history_start in range(pair_count):
retained_steps, retained_changes = valid_pairs(history_start)
if retained_steps.shape[0] == 0:
continue
candidate, candidate_ascent = normalized_ascent(
two_loop(retained_steps, retained_changes) * scale
)
relative_ascent = candidate_ascent / reference_ascent
if (
np.all(np.isfinite(candidate))
and np.isfinite(relative_ascent)
and relative_ascent >= minimum_relative_ascent
):
return SafeguardedLBFGSDirection(
direction=candidate,
method="safeguarded_lbfgs",
retained_step_history=retained_steps,
retained_gradient_difference_history=retained_changes,
pruned_pairs=pair_count - int(retained_steps.shape[0]),
predicted_ascent=float(candidate_ascent),
reference_ascent=float(reference_ascent),
relative_predicted_ascent=float(relative_ascent),
)
return SafeguardedLBFGSDirection(
direction=fallback,
method="scaled_steepest_ascent_curvature_fallback",
retained_step_history=empty,
retained_gradient_difference_history=empty.copy(),
pruned_pairs=pair_count,
predicted_ascent=float(reference_ascent),
reference_ascent=float(reference_ascent),
relative_predicted_ascent=1.0,
)
@dataclass(frozen=True)
class LowRadiationSubspaceConfig:
"""Configuration for an optional low-radiation direction projector.
The rows supplied to :func:`update_low_radiation_subspace` span the range
of a local radiation curvature factor such as ``J.H @ J``. Projecting a
design direction out of that range leaves the locally dark/null-space
component. In practice the rows can be direct radiation gradients and
gradient secants, since a secant is a matrix-free approximation to one
radiation-Hessian action.
``enabled=False`` is deliberately the default so existing optimizers are
bit-for-bit unaffected until the mechanism is explicitly selected.
"""
enabled: bool = False
maximum_rank: int = 8
relative_singular_value_cutoff: float = 1.0e-5
projection_strength: float = 1.0
def __post_init__(self) -> None:
if not isinstance(self.enabled, bool):
raise TypeError("enabled must be boolean")
if (
not isinstance(self.maximum_rank, int)
or isinstance(self.maximum_rank, bool)
or self.maximum_rank < 1
):
raise ValueError("maximum_rank must be a positive integer")
if (
not np.isfinite(self.relative_singular_value_cutoff)
or not 0.0 <= self.relative_singular_value_cutoff < 1.0
):
raise ValueError(
"relative_singular_value_cutoff must lie in [0, 1)"
)
if (
not np.isfinite(self.projection_strength)
or not 0.0 <= self.projection_strength <= 1.0
):
raise ValueError("projection_strength must lie in [0, 1]")
@dataclass(frozen=True)
class LowRadiationProjection:
"""Projected proposal and diagnostics for one optimizer step."""
direction: np.ndarray
basis: np.ndarray
singular_values: np.ndarray
rank: int
removed_norm_fraction: float
sensitive_norm_before: float
sensitive_norm_after: float
@dataclass(frozen=True)
class LinearizedConstraintProjection:
"""A proposal projected into one linearized feasible half-space."""
direction: np.ndarray
active: bool
unconstrained_derivative: float
constrained_derivative: float
required_derivative: float
correction_multiplier: float
correction_norm_fraction: float
def _normalized_finite_rows(
rows: np.ndarray | jax.Array | None,
*,
width: int,
) -> np.ndarray:
"""Return finite, unit-norm rows without changing their span."""
if rows is None:
return np.empty((0, width), dtype=np.float64)
values = np.asarray(rows, dtype=np.float64)
if values.size == 0:
return np.empty((0, width), dtype=np.float64)
if values.ndim == 1:
values = values[None, :]
values = values.reshape((values.shape[0], -1))
if values.shape[1] != width:
raise ValueError(
f"radiation sensitivity width {values.shape[1]} does not match "
f"the design width {width}"
)
norms = np.linalg.norm(values, axis=1)
threshold = np.finfo(np.float64).eps * math.sqrt(max(width, 1))
keep = np.all(np.isfinite(values), axis=1) & np.isfinite(norms) & (norms > threshold)
if not np.any(keep):
return np.empty((0, width), dtype=np.float64)
return values[keep] / norms[keep, None]
def update_low_radiation_subspace(
basis: np.ndarray | jax.Array | None,
radiation_gradient: np.ndarray | jax.Array,
*,
previous_radiation_gradient: np.ndarray | jax.Array | None = None,
accepted_step: np.ndarray | jax.Array | None = None,
config: LowRadiationSubspaceConfig | None = None,
) -> tuple[np.ndarray, np.ndarray]:
"""Update a compact radiation-sensitive range from gradients and secants.
Near a dark design, inverse Q is locally quadratic and
``grad(loss_new) - grad(loss_old) ~= H_rad @ accepted_step``. Accumulating
these matrix-free Hessian actions learns ``range(H_rad)`` online; its
orthogonal complement is the low-radiation subspace. The current loss
gradient is included because it lies in the same range for the local
quadratic model and protects Q before the first accepted secant exists.
Returns an orthonormal row basis and the retained singular values. No
history is changed when the feature is disabled.
"""
if config is None:
config = LowRadiationSubspaceConfig()
gradient = np.asarray(radiation_gradient, dtype=np.float64)
design_shape = gradient.shape
width = gradient.size
existing = _normalized_finite_rows(basis, width=width)
if not config.enabled:
return existing.reshape((existing.shape[0], *design_shape)), np.empty(
(0,), dtype=np.float64
)
# ``radiation_gradient`` is one design-space vector even when the design
# itself is an image or volume. Flatten it explicitly so a 2-D design is
# not misread as a stack of short basis rows.
candidates = [
existing,
_normalized_finite_rows(gradient.reshape(-1), width=width),
]
if (previous_radiation_gradient is None) != (accepted_step is None):
raise ValueError(
"previous_radiation_gradient and accepted_step must be supplied together"
)
if previous_radiation_gradient is not None and accepted_step is not None:
previous = np.asarray(previous_radiation_gradient, dtype=np.float64)
step = np.asarray(accepted_step, dtype=np.float64)
if previous.shape != design_shape or step.shape != design_shape:
raise ValueError("radiation-gradient secant arrays must match the current gradient")
step_norm = float(np.linalg.norm(step))
if np.isfinite(step_norm) and step_norm > np.finfo(np.float64).eps:
secant_action = (gradient - previous) / step_norm
candidates.append(
_normalized_finite_rows(secant_action.reshape(-1), width=width)
)
matrix = np.concatenate(candidates, axis=0)
if matrix.shape[0] == 0:
return np.empty((0, *design_shape), dtype=np.float64), np.empty(
(0,), dtype=np.float64
)
_, singular_values, right = np.linalg.svd(matrix, full_matrices=False)
cutoff = config.relative_singular_value_cutoff * singular_values[0]
rank = min(
config.maximum_rank,
int(np.count_nonzero(singular_values > cutoff)),
)
retained = right[:rank]
return retained.reshape((rank, *design_shape)), singular_values[:rank]
def project_low_radiation_direction(
direction: np.ndarray | jax.Array,
radiation_basis: np.ndarray | jax.Array | None,
*,
config: LowRadiationSubspaceConfig | None = None,
) -> LowRadiationProjection:
"""Remove the learned radiation-sensitive component of a proposal.
``radiation_basis`` need not already be orthogonal. It is compressed with
an SVD on every call because the intended rank is small (usually 4--12),
while the number of design variables may be very large.
"""
if config is None:
config = LowRadiationSubspaceConfig()
values = np.asarray(direction, dtype=np.float64)
if values.size == 0 or not np.all(np.isfinite(values)):
raise ValueError("direction must be a non-empty finite array")
flat = values.reshape(-1)
rows = _normalized_finite_rows(radiation_basis, width=flat.size)
if not config.enabled or rows.shape[0] == 0 or config.projection_strength == 0.0:
return LowRadiationProjection(
direction=np.asarray(values),
basis=rows.reshape((rows.shape[0], *values.shape)),
singular_values=np.empty((0,), dtype=np.float64),
rank=0,
removed_norm_fraction=0.0,
sensitive_norm_before=0.0,
sensitive_norm_after=0.0,
)
_, singular_values, right = np.linalg.svd(rows, full_matrices=False)
cutoff = config.relative_singular_value_cutoff * singular_values[0]
rank = min(
config.maximum_rank,
int(np.count_nonzero(singular_values > cutoff)),
)
orthonormal = right[:rank]
coefficients = orthonormal @ flat
sensitive = orthonormal.T @ coefficients
projected = flat - config.projection_strength * sensitive
projected_coefficients = orthonormal @ projected
direction_norm = float(np.linalg.norm(flat))
return LowRadiationProjection(
direction=projected.reshape(values.shape),
basis=orthonormal.reshape((rank, *values.shape)),
singular_values=singular_values[:rank],
rank=rank,
removed_norm_fraction=float(
config.projection_strength * np.linalg.norm(sensitive)
/ max(direction_norm, np.finfo(np.float64).tiny)
),
sensitive_norm_before=float(np.linalg.norm(coefficients)),
sensitive_norm_after=float(np.linalg.norm(projected_coefficients)),
)
def project_linearized_constraint_direction(
direction: np.ndarray | jax.Array,
constraint_gradient: np.ndarray | jax.Array,
*,
metric_constraint_direction: np.ndarray | jax.Array | None = None,
minimum_directional_derivative: float = 0.0,
minimum_cosine: float = 0.0,
) -> LinearizedConstraintProjection:
"""Project a proposal into ``grad(constraint) @ direction >= required``.
The constraint is assumed to be feasible when it increases. With the
defaults this is the closest constant-or-improving direction. A positive
``minimum_cosine`` adds a small first-order feasibility margin so that the
linear term can dominate adverse second-order curvature after a finite
step.
``metric_constraint_direction`` may be ``H @ grad(constraint)`` for a
positive-definite inverse-Hessian or other optimizer metric. Omitting it
gives the ordinary Euclidean projection.
"""
proposal = np.asarray(direction, dtype=np.float64)
gradient = np.asarray(constraint_gradient, dtype=np.float64)
if proposal.size == 0 or proposal.shape != gradient.shape:
raise ValueError("direction and constraint_gradient must have one non-empty shape")
if not np.all(np.isfinite(proposal)) or not np.all(np.isfinite(gradient)):
raise ValueError("direction and constraint_gradient must be finite")
if not np.isfinite(minimum_directional_derivative):
raise ValueError("minimum_directional_derivative must be finite")
if not np.isfinite(minimum_cosine) or not 0.0 <= minimum_cosine < 1.0:
raise ValueError("minimum_cosine must lie in [0, 1)")
proposal_flat = proposal.reshape(-1)
gradient_flat = gradient.reshape(-1)
proposal_norm = float(np.linalg.norm(proposal_flat))
gradient_norm = float(np.linalg.norm(gradient_flat))
threshold = np.finfo(np.float64).eps * math.sqrt(max(proposal_flat.size, 1))
if proposal_norm <= threshold:
raise ValueError("direction must have nonzero norm")
if gradient_norm <= threshold:
raise ValueError("constraint_gradient must have nonzero norm")
if metric_constraint_direction is None:
metric_flat = gradient_flat
else:
metric = np.asarray(metric_constraint_direction, dtype=np.float64)
if metric.shape != proposal.shape or not np.all(np.isfinite(metric)):
raise ValueError("metric_constraint_direction must be finite and shape-matched")
metric_flat = metric.reshape(-1)
denominator = float(np.vdot(gradient_flat, metric_flat))
if not np.isfinite(denominator) or denominator <= threshold:
raise ValueError(
"metric_constraint_direction must have positive constraint curvature"
)
unconstrained = float(np.vdot(gradient_flat, proposal_flat))
required = max(
float(minimum_directional_derivative),
float(minimum_cosine * gradient_norm * proposal_norm),
)
if unconstrained >= required:
return LinearizedConstraintProjection(
direction=np.asarray(proposal),
active=False,
unconstrained_derivative=unconstrained,
constrained_derivative=unconstrained,
required_derivative=required,
correction_multiplier=0.0,
correction_norm_fraction=0.0,
)
multiplier = (required - unconstrained) / denominator
correction = multiplier * metric_flat
constrained = proposal_flat + correction
return LinearizedConstraintProjection(
direction=constrained.reshape(proposal.shape),
active=True,
unconstrained_derivative=unconstrained,
constrained_derivative=float(np.vdot(gradient_flat, constrained)),
required_derivative=required,
correction_multiplier=float(multiplier),
correction_norm_fraction=float(
np.linalg.norm(correction)
/ max(proposal_norm, np.finfo(np.float64).tiny)
),
)
DesignActivation = Literal["identity", "sigmoid", "tanh"]
class DesignBasis2D(Protocol):
"""Interface implemented by differentiable two-dimensional design bases.
Basis objects are static Python configuration captured by a JAX objective.
The optimized value is always the array returned by :attr:`parameter_shape`.
``decode`` maps those parameters to a material density on ``shape`` while
``encode`` provides a host-side initializer from an existing density map.
"""
shape: tuple[int, int]
@property
def parameter_shape(self) -> tuple[int, ...]: ...
@property
def parameter_bounds(self) -> tuple[float, float]: ...
def latent(self, parameters: jax.Array) -> jax.Array: ...
def decode(self, parameters: jax.Array) -> jax.Array: ...
def encode(self, density: jax.Array | np.ndarray) -> jax.Array: ...
class DesignBasis1D(Protocol):
"""Interface for differentiable one-dimensional fabrication bases."""
size: int
@property
def parameter_shape(self) -> tuple[int, ...]: ...
@property
def parameter_bounds(self) -> tuple[float, float]: ...
def latent(self, parameters: jax.Array) -> jax.Array: ...
def decode(self, parameters: jax.Array) -> jax.Array: ...
def encode(self, density: jax.Array | np.ndarray) -> jax.Array: ...
def _validate_design_shape(shape: tuple[int, int]) -> tuple[int, int]:
if len(shape) != 2 or any(not isinstance(value, int) or isinstance(value, bool) or value < 2 for value in shape):
raise ValueError("shape must contain two integer dimensions of at least two cells")
return tuple(shape)
def _validate_design_spacing(spacing: tuple[float, float]) -> tuple[float, float]:
if len(spacing) != 2 or any(not np.isfinite(value) or value <= 0 for value in spacing):
raise ValueError("spacing must contain two positive finite values")
return tuple(float(value) for value in spacing)
def _validate_design_size(size: int) -> int:
if not isinstance(size, int) or isinstance(size, bool) or size < 2:
raise ValueError("size must be an integer of at least two cells")
return size
def _activate_design_latent(values: jax.Array, activation: DesignActivation) -> jax.Array:
if activation == "identity":
return values
if activation == "sigmoid":
return jax.nn.sigmoid(values)
if activation == "tanh":
return 0.5 + 0.5 * jnp.tanh(values)
raise ValueError("activation must be 'identity', 'sigmoid', or 'tanh'")
def _inverse_design_activation(
density: jax.Array | np.ndarray,
activation: DesignActivation,
*,
epsilon: float = 1e-4,
) -> np.ndarray:
values = np.asarray(density, dtype=np.float64)
if values.ndim != 2:
raise ValueError("density must be a two-dimensional array")
if not 0 < epsilon < 0.5:
raise ValueError("epsilon must lie strictly between zero and one half")
if not np.all(np.isfinite(values)):
raise ValueError("density must contain only finite values")
if activation == "identity":
return values
clipped = np.clip(values, epsilon, 1.0 - epsilon)
if activation == "sigmoid":
return np.log(clipped) - np.log1p(-clipped)
if activation == "tanh":
return np.arctanh(2.0 * clipped - 1.0)
raise ValueError("activation must be 'identity', 'sigmoid', or 'tanh'")
def _inverse_design_activation_1d(
density: jax.Array | np.ndarray,
activation: DesignActivation,
*,
epsilon: float = 1e-4,
) -> np.ndarray:
values = np.asarray(density, dtype=np.float64)
if values.ndim != 1:
raise ValueError("density must be a one-dimensional array")
if not 0 < epsilon < 0.5:
raise ValueError("epsilon must lie strictly between zero and one half")
if not np.all(np.isfinite(values)):
raise ValueError("density must contain only finite values")
if activation == "identity":
return values
clipped = np.clip(values, epsilon, 1.0 - epsilon)
if activation == "sigmoid":
return np.log(clipped) - np.log1p(-clipped)
if activation == "tanh":
return np.arctanh(2.0 * clipped - 1.0)
raise ValueError("activation must be 'identity', 'sigmoid', or 'tanh'")
def _period_is_selected(
wave_number: float,
*,
minimum_period: float | None,
maximum_period: float | None,
) -> bool:
if wave_number == 0:
return True
period = 2.0 * math.pi / wave_number
if minimum_period is not None and period < minimum_period:
return False
return maximum_period is None or period <= maximum_period
def _validate_period_band(minimum_period: float | None, maximum_period: float | None) -> None:
if minimum_period is not None and (not np.isfinite(minimum_period) or minimum_period <= 0):
raise ValueError("minimum_period must be positive and finite")
if maximum_period is not None and (not np.isfinite(maximum_period) or maximum_period <= 0):
raise ValueError("maximum_period must be positive and finite")
if minimum_period is not None and maximum_period is not None and minimum_period > maximum_period:
raise ValueError("minimum_period must not exceed maximum_period")
@dataclass(frozen=True)
class PixelBasis2D:
"""One independent design parameter per in-plane material pixel."""
shape: tuple[int, int]
activation: DesignActivation = "identity"
def __post_init__(self) -> None:
_validate_design_shape(self.shape)
_activate_design_latent(jnp.asarray(0.0), self.activation)
@property
def parameter_shape(self) -> tuple[int, int]:
return self.shape
@property
def parameter_bounds(self) -> tuple[float, float]:
return (0.0, 1.0) if self.activation == "identity" else (-math.inf, math.inf)
def latent(self, parameters: jax.Array) -> jax.Array:
parameters = jnp.asarray(parameters)
if parameters.shape != self.shape:
raise ValueError(f"pixel parameters must have shape {self.shape}, got {parameters.shape}")
return parameters
def decode(self, parameters: jax.Array) -> jax.Array:
return _activate_design_latent(self.latent(parameters), self.activation)
def encode(self, density: jax.Array | np.ndarray) -> jax.Array:
latent = _inverse_design_activation(density, self.activation)
if latent.shape != self.shape:
raise ValueError(f"density must have shape {self.shape}, got {latent.shape}")
return jnp.asarray(latent, dtype=jnp.float32)
@dataclass(frozen=True)
class BoundaryTopology:
"""Component-and-hole signature used to guard boundary evolution."""
material_components: int
enclosed_voids: int
BoundaryEdge = Literal["min_x", "max_x", "min_y", "max_y"]
_ALL_BOUNDARY_EDGES: tuple[BoundaryEdge, ...] = (
"min_x",
"max_x",
"min_y",
"max_y",
)
def _validate_exterior_edges(
exterior_edges: tuple[BoundaryEdge, ...],
) -> tuple[BoundaryEdge, ...]:
edges = tuple(exterior_edges)
if len(set(edges)) != len(edges) or any(edge not in _ALL_BOUNDARY_EDGES for edge in edges):
raise ValueError(
"exterior_edges must contain unique values from "
"('min_x', 'max_x', 'min_y', 'max_y')"
)
return edges
[docs]
def boundary_topology_signature(
density: jax.Array | np.ndarray,
*,
threshold: float = 0.5,
diagonal_connectivity: bool = False,
exterior_edges: tuple[BoundaryEdge, ...] = _ALL_BOUNDARY_EDGES,
) -> BoundaryTopology:
"""Return the material-component and enclosed-void counts of a 2-D mask.
Air connected to one of ``exterior_edges`` is exterior and is therefore
not counted as a hole. Omit symmetry planes from that tuple: a half-hole
touching a mirror plane then remains an enclosed hole after unfolding.
This host-side diagnostic is intended for accepted candidate checks, not
for differentiation inside a JIT-compiled objective.
"""
from scipy import ndimage
values = np.asarray(density)
if values.ndim != 2 or min(values.shape) < 2:
raise ValueError("density must be a two-dimensional array at least two cells wide")
if not np.all(np.isfinite(values)):
raise ValueError("density must contain only finite values")
if not 0.0 <= threshold <= 1.0:
raise ValueError("threshold must lie in [0, 1]")
exterior_edges = _validate_exterior_edges(exterior_edges)
connectivity = 2 if diagonal_connectivity else 1
structure = ndimage.generate_binary_structure(2, connectivity)
material = values >= threshold
_, material_components = ndimage.label(material, structure=structure)
void_labels, void_components = ndimage.label(~material, structure=structure)
edge_values = {
"min_x": void_labels[0, :],
"max_x": void_labels[-1, :],
"min_y": void_labels[:, 0],
"max_y": void_labels[:, -1],
}
exterior_labels: set[int] = set()
for edge in exterior_edges:
exterior_labels.update(np.unique(edge_values[edge]))
exterior_labels.discard(0)
enclosed_voids = void_components - len(exterior_labels)
return BoundaryTopology(int(material_components), int(enclosed_voids))
def _topology_region_labels(
density: jax.Array | np.ndarray,
*,
threshold: float,
diagonal_connectivity: bool,
exterior_edges: tuple[BoundaryEdge, ...],
) -> tuple[np.ndarray, np.ndarray]:
"""Label material components and non-exterior voids for lineage checks."""
from scipy import ndimage
values = np.asarray(density)
connectivity = 2 if diagonal_connectivity else 1
structure = ndimage.generate_binary_structure(2, connectivity)
material_labels, _ = ndimage.label(values >= threshold, structure=structure)
void_labels, _ = ndimage.label(values < threshold, structure=structure)
edge_values = {
"min_x": void_labels[0, :],
"max_x": void_labels[-1, :],
"min_y": void_labels[:, 0],
"max_y": void_labels[:, -1],
}
exterior_labels: set[int] = set()
for edge in exterior_edges:
exterior_labels.update(np.unique(edge_values[edge]))
exterior_labels.discard(0)
enclosed_labels = np.zeros_like(void_labels)
enclosed_index = 0
for label in range(1, int(void_labels.max()) + 1):
if label not in exterior_labels:
enclosed_index += 1
enclosed_labels[void_labels == label] = enclosed_index
return material_labels, enclosed_labels
def _labels_have_one_to_one_lineage(
reference_labels: np.ndarray,
candidate_labels: np.ndarray,
) -> bool:
"""Return whether every labeled region overlaps exactly one counterpart."""
for source, destination in (
(reference_labels, candidate_labels),
(candidate_labels, reference_labels),
):
for label in range(1, int(source.max()) + 1):
overlaps = np.unique(destination[source == label])
overlaps = overlaps[overlaps > 0]
if overlaps.size != 1:
return False
return True
def signed_distance_level_set(
density: jax.Array | np.ndarray,
*,
spacing: tuple[float, float] = (1.0, 1.0),
threshold: float = 0.5,
) -> np.ndarray:
"""Fit a positive-in-material signed-distance field to a density image.
The returned level set lives on cell *vertices* and therefore has shape
``(density.shape[0] + 1, density.shape[1] + 1)``. Thresholding a grey
topology result before this conversion selects the zero-contour once;
subsequent optimization moves that contour instead of changing material
values throughout the bulk.
"""
from scipy import ndimage
values = np.asarray(density, dtype=np.float64)
if values.ndim != 2 or min(values.shape) < 2:
raise ValueError("density must be a two-dimensional array at least two cells wide")
_validate_design_spacing(spacing)
if not np.all(np.isfinite(values)):
raise ValueError("density must contain only finite values")
if not 0.0 <= threshold <= 1.0:
raise ValueError("threshold must lie in [0, 1]")
# Average the four adjacent cell values onto every vertex. Edge padding
# extends the design normally rather than introducing a fictitious air rim.
padded = np.pad(values, ((1, 1), (1, 1)), mode="edge")
vertex_density = 0.25 * (
padded[:-1, :-1]
+ padded[1:, :-1]
+ padded[:-1, 1:]
+ padded[1:, 1:]
)
material = vertex_density >= threshold
if np.all(material) or not np.any(material):
raise ValueError("density must contain at least one material interface")
inside_distance = ndimage.distance_transform_edt(material, sampling=spacing)
outside_distance = ndimage.distance_transform_edt(~material, sampling=spacing)
level_set = np.where(material, inside_distance, -outside_distance)
return np.asarray(level_set, dtype=np.float32)
def fit_level_set_to_density(
density: jax.Array | np.ndarray,
*,
spacing: tuple[float, float] = (1.0, 1.0),
threshold: float = 0.5,
supersample: int = 4,
interface_width: float | None = None,
fit_band_width: float | None = None,
max_iterations: int = 200,
exterior_edges: tuple[BoundaryEdge, ...] = _ALL_BOUNDARY_EDGES,
) -> np.ndarray:
"""Fit a sign-preserving vertex level set to a two-phase density field.
A thresholded signed-distance initializer establishes the topology. The
magnitudes of vertices near that interface are then fitted to the supplied
subpixel density while every vertex sign remains fixed. This preserves the
selected topology but places the zero contour much more accurately than a
direct cell-to-vertex threshold, which is important for thin photonic
features during a density-to-boundary handoff.
"""
from scipy.optimize import minimize
values = np.asarray(density, dtype=np.float32)
spacing = _validate_design_spacing(spacing)
if values.ndim != 2 or min(values.shape) < 2 or not np.all(np.isfinite(values)):
raise ValueError("density must be a finite two-dimensional array")
if not isinstance(max_iterations, int) or max_iterations < 1:
raise ValueError("max_iterations must be a positive integer")
if not isinstance(supersample, int) or supersample < 1:
raise ValueError("supersample must be a positive integer")
grid_scale = max(spacing)
width = min(spacing) / supersample if interface_width is None else interface_width
if width <= 0.0:
raise ValueError("interface_width must be positive")
band = 2.5 * grid_scale if fit_band_width is None else fit_band_width
if band <= 0.0:
raise ValueError("fit_band_width must be positive")
edges = _validate_exterior_edges(exterior_edges)
initial = signed_distance_level_set(values, spacing=spacing, threshold=threshold)
active = np.abs(initial) <= band
active_indices = np.flatnonzero(active)
if active_indices.size == 0:
return initial
base_flat = jnp.asarray(initial).reshape(-1)
indices = jnp.asarray(active_indices, dtype=jnp.int32)
target = jnp.asarray(values)
vertex_shape = initial.shape
def loss(active_values: jax.Array) -> jax.Array:
level_set = base_flat.at[indices].set(active_values).reshape(vertex_shape)
candidate = level_set_fill_fraction(
level_set,
supersample=supersample,
interface_width=width,
)
return jnp.mean(jnp.square(candidate - target))
value_and_gradient = jax.jit(jax.value_and_grad(loss))
def objective(active_values: np.ndarray) -> tuple[float, np.ndarray]:
value, gradient = value_and_gradient(jnp.asarray(active_values, dtype=jnp.float32))
return float(value), np.asarray(gradient, dtype=np.float64)
initial_active = initial.reshape(-1)[active_indices].astype(np.float64)
epsilon = max(1.0e-6 * grid_scale, np.finfo(np.float32).eps)
magnitude_limit = max(4.0 * grid_scale, band)
bounds = [
(epsilon, magnitude_limit) if value >= 0.0 else (-magnitude_limit, -epsilon)
for value in initial_active
]
result = minimize(
objective,
initial_active,
method="L-BFGS-B",
jac=True,
bounds=bounds,
options={
"maxiter": max_iterations,
"ftol": 1.0e-14,
"gtol": 1.0e-8,
"maxls": 20,
},
)
if not np.all(np.isfinite(result.x)):
raise RuntimeError("level-set density fit produced nonfinite vertices")
fitted = initial.copy().reshape(-1)
fitted[active_indices] = result.x
fitted = fitted.reshape(initial.shape).astype(np.float32)
fitted_density = np.asarray(
level_set_fill_fraction(
jnp.asarray(fitted),
supersample=supersample,
interface_width=width,
)
)
initial_topology = boundary_topology_signature(
values,
threshold=threshold,
exterior_edges=edges,
)
fitted_topology = boundary_topology_signature(
fitted_density,
threshold=threshold,
exterior_edges=edges,
)
if fitted_topology != initial_topology:
raise RuntimeError("sign-preserving level-set fit changed material topology")
return fitted
def level_set_fill_fraction(
level_set: jax.Array,
*,
supersample: int = 4,
interface_width: float,
) -> jax.Array:
"""Return a differentiable subpixel fill fraction for a binary boundary.
Bilinear interpolation locates the zero contour at a regular set of
quadrature points inside each Yee cell. A compact, continuously
differentiable Heaviside maps those samples to material coverage. It is
exactly zero or one outside ``interface_width`` and therefore introduces
no optimizable grey bulk material: fractional values occur only in cells
intersected by the geometric interface.
Crucially, the same smooth coverage is used by both the forward Maxwell
solve and reverse-mode differentiation. This avoids the inconsistent
straight-through construction in which an infinitesimal boundary motion
had a nonzero derivative but left the permittivity seen by FDTD unchanged.
"""
level_set = jnp.asarray(level_set)
if level_set.ndim != 2 or min(level_set.shape) < 3:
raise ValueError("level_set must be a two-dimensional vertex grid at least three points wide")
if not isinstance(supersample, int) or isinstance(supersample, bool) or supersample < 1:
raise ValueError("supersample must be a positive integer")
if interface_width <= 0:
raise ValueError("interface_width must be positive")
coordinate = (jnp.arange(supersample, dtype=level_set.dtype) + 0.5) / supersample
u = coordinate[None, None, :, None]
v = coordinate[None, None, None, :]
phi00 = level_set[:-1, :-1, None, None]
phi10 = level_set[1:, :-1, None, None]
phi01 = level_set[:-1, 1:, None, None]
phi11 = level_set[1:, 1:, None, None]
subcell_level_set = (
(1.0 - u) * (1.0 - v) * phi00
+ u * (1.0 - v) * phi10
+ (1.0 - u) * v * phi01
+ u * v * phi11
)
# C1 compact smoothstep approximation to H(phi). The compact support is
# preferable to a sigmoid here: uniform cells remain exactly binary while
# interface cells change continuously as the physical contour moves.
coordinate = jnp.clip(
0.5 + 0.5 * subcell_level_set / interface_width,
0.0,
1.0,
)
material = coordinate * coordinate * (3.0 - 2.0 * coordinate)
return jnp.mean(material, axis=(2, 3))
[docs]
@dataclass(frozen=True)
class LevelSetBoundary2D:
"""Binary, topology-guarded boundary-normal design parameterization.
``parameters`` are signed physical displacements on the vertex grid.
Positive displacement grows material along its outward normal and negative
displacement retreats it. A cosine narrow-band window makes the Maxwell
derivative exactly zero in uniform bulk regions, so the gradient is a
normal boundary sensitivity rather than a pixel-density sensitivity.
The basis cannot nucleate a remote component because displacement vanishes
outside ``narrow_band_width``. Boundary collision can still merge or erase
existing components, so :meth:`preserves_topology` and :meth:`rebase`
provide an explicit host-side acceptance guard.
"""
reference_level_set: np.ndarray
spacing: tuple[float, float] = (1.0, 1.0)
narrow_band_width: float = 4.0
maximum_displacement: float = 2.0
interface_width: float = 0.25
supersample: int = 4
threshold: float = 0.5
diagonal_connectivity: bool = False
exterior_edges: tuple[BoundaryEdge, ...] = _ALL_BOUNDARY_EDGES
def __post_init__(self) -> None:
level_set = np.asarray(self.reference_level_set, dtype=np.float32)
if level_set.ndim != 2 or min(level_set.shape) < 3:
raise ValueError("reference_level_set must be a two-dimensional vertex grid")
if not np.all(np.isfinite(level_set)):
raise ValueError("reference_level_set must contain only finite values")
if not np.any(level_set < 0.0) or not np.any(level_set >= 0.0):
raise ValueError("reference_level_set must contain a zero interface")
_validate_design_spacing(self.spacing)
if self.narrow_band_width <= 0 or self.maximum_displacement <= 0 or self.interface_width <= 0:
raise ValueError("boundary widths and maximum displacement must be positive")
if self.maximum_displacement > self.narrow_band_width:
raise ValueError("maximum_displacement must not exceed narrow_band_width")
if not isinstance(self.supersample, int) or isinstance(self.supersample, bool) or self.supersample < 1:
raise ValueError("supersample must be a positive integer")
if not 0.0 <= self.threshold <= 1.0:
raise ValueError("threshold must lie in [0, 1]")
object.__setattr__(self, "exterior_edges", _validate_exterior_edges(self.exterior_edges))
frozen_level_set = level_set.copy()
frozen_level_set.setflags(write=False)
object.__setattr__(self, "reference_level_set", frozen_level_set)
[docs]
@classmethod
def from_density(
cls,
density: jax.Array | np.ndarray,
*,
spacing: tuple[float, float] = (1.0, 1.0),
threshold: float = 0.5,
narrow_band_width: float | None = None,
maximum_displacement: float | None = None,
interface_width: float | None = None,
supersample: int = 4,
diagonal_connectivity: bool = False,
exterior_edges: tuple[BoundaryEdge, ...] = _ALL_BOUNDARY_EDGES,
) -> LevelSetBoundary2D:
"""Threshold a grey result once and fit its boundary level set."""
spacing = _validate_design_spacing(spacing)
grid_scale = max(spacing)
narrow_band = 4.0 * grid_scale if narrow_band_width is None else narrow_band_width
maximum = 0.5 * narrow_band if maximum_displacement is None else maximum_displacement
width = min(spacing) / supersample if interface_width is None else interface_width
return cls(
reference_level_set=signed_distance_level_set(
density,
spacing=spacing,
threshold=threshold,
),
spacing=spacing,
narrow_band_width=narrow_band,
maximum_displacement=maximum,
interface_width=width,
supersample=supersample,
threshold=threshold,
diagonal_connectivity=diagonal_connectivity,
exterior_edges=exterior_edges,
)
@property
def shape(self) -> tuple[int, int]:
return (
self.reference_level_set.shape[0] - 1,
self.reference_level_set.shape[1] - 1,
)
@property
def parameter_shape(self) -> tuple[int, int]:
return self.reference_level_set.shape
@property
def parameter_bounds(self) -> tuple[float, float]:
return (-self.maximum_displacement, self.maximum_displacement)
[docs]
def initial_parameters(self, *, dtype: jnp.dtype = jnp.float32) -> jax.Array:
"""Return the zero-displacement start for this reference boundary."""
return jnp.zeros(self.parameter_shape, dtype=dtype)
[docs]
def boundary_window(self, *, dtype: jnp.dtype = jnp.float32) -> jax.Array:
"""Cosine window equal to one on the interface and zero in the bulk."""
distance = jnp.abs(jnp.asarray(self.reference_level_set, dtype=dtype))
phase = jnp.clip(distance / self.narrow_band_width, 0.0, 1.0)
return jnp.where(
distance < self.narrow_band_width,
0.5 * (1.0 + jnp.cos(jnp.pi * phase)),
0.0,
)
def latent(self, parameters: jax.Array) -> jax.Array:
parameters = jnp.asarray(parameters)
if parameters.shape != self.parameter_shape:
raise ValueError(
f"boundary parameters must have shape {self.parameter_shape}, got {parameters.shape}"
)
displacement = jnp.clip(
parameters,
-self.maximum_displacement,
self.maximum_displacement,
)
reference = jnp.asarray(self.reference_level_set, dtype=parameters.dtype)
return reference + self.boundary_window(dtype=parameters.dtype) * displacement
def decode(self, parameters: jax.Array) -> jax.Array:
return level_set_fill_fraction(
self.latent(parameters),
supersample=self.supersample,
interface_width=self.interface_width,
)
def encode(self, density: jax.Array | np.ndarray) -> jax.Array:
"""Fit a same-topology binary density as a boundary displacement."""
values = np.asarray(density)
if values.shape != self.shape:
raise ValueError(f"density must have shape {self.shape}, got {values.shape}")
if not self.preserves_topology(values):
raise ValueError("density topology differs from the reference boundary")
candidate = signed_distance_level_set(
values,
spacing=self.spacing,
threshold=self.threshold,
)
window = np.asarray(self.boundary_window(), dtype=np.float64)
displacement = np.zeros_like(candidate, dtype=np.float64)
active = window > 0.25
displacement[active] = (
candidate[active] - self.reference_level_set[active]
) / window[active]
return jnp.asarray(
np.clip(
displacement,
-self.maximum_displacement,
self.maximum_displacement,
),
dtype=jnp.float32,
)
[docs]
def interface_normal(self, parameters: jax.Array) -> jax.Array:
"""Return the outward material normal on the vertex grid."""
level_set = self.latent(parameters)
grad_x, grad_y = jnp.gradient(level_set, self.spacing[0], self.spacing[1])
magnitude = jnp.sqrt(grad_x**2 + grad_y**2 + 1e-24)
# The level set is positive in material, hence its gradient points in;
# the material-outward normal has the opposite sign.
return -jnp.stack((grad_x / magnitude, grad_y / magnitude), axis=0)
@property
def topology(self) -> BoundaryTopology:
reference_density = np.asarray(self.decode(self.initial_parameters()))
return boundary_topology_signature(
reference_density,
threshold=self.threshold,
diagonal_connectivity=self.diagonal_connectivity,
exterior_edges=self.exterior_edges,
)
def preserves_topology(self, density: jax.Array | np.ndarray) -> bool:
candidate = np.asarray(density)
if candidate.shape != self.shape:
raise ValueError(f"density must have shape {self.shape}, got {candidate.shape}")
if boundary_topology_signature(
candidate,
threshold=self.threshold,
diagonal_connectivity=self.diagonal_connectivity,
exterior_edges=self.exterior_edges,
) != self.topology:
return False
reference = np.asarray(self.decode(self.initial_parameters()))
reference_material, reference_voids = _topology_region_labels(
reference,
threshold=self.threshold,
diagonal_connectivity=self.diagonal_connectivity,
exterior_edges=self.exterior_edges,
)
candidate_material, candidate_voids = _topology_region_labels(
candidate,
threshold=self.threshold,
diagonal_connectivity=self.diagonal_connectivity,
exterior_edges=self.exterior_edges,
)
return _labels_have_one_to_one_lineage(
reference_material,
candidate_material,
) and _labels_have_one_to_one_lineage(reference_voids, candidate_voids)
[docs]
def rebase(
self,
parameters: jax.Array | np.ndarray,
*,
preserve_topology: bool = True,
) -> tuple[LevelSetBoundary2D, jax.Array]:
"""Accept a moved interface and return a fresh signed-distance stage."""
density = np.asarray(self.decode(jnp.asarray(parameters)))
if preserve_topology and not self.preserves_topology(density):
raise ValueError("candidate boundary changes the reference topology")
rebased = type(self).from_density(
density,
spacing=self.spacing,
threshold=self.threshold,
narrow_band_width=self.narrow_band_width,
maximum_displacement=self.maximum_displacement,
interface_width=self.interface_width,
supersample=self.supersample,
diagonal_connectivity=self.diagonal_connectivity,
exterior_edges=self.exterior_edges,
)
if preserve_topology and rebased.topology != self.topology:
raise ValueError("signed-distance rebasing changed the reference topology")
return rebased, rebased.initial_parameters()
def _cubic_reflecting_interpolation_matrix(
output_size: int,
control_size: int,
*,
dtype: jnp.dtype,
) -> jax.Array:
"""Return a uniform cubic B-spline matrix with reflected end knots."""
if output_size < 2 or control_size < 4:
raise ValueError("cubic interpolation requires two outputs and four controls")
coordinate = jnp.linspace(0.0, control_size - 1, output_size, dtype=dtype)
interval = jnp.floor(coordinate).astype(jnp.int32)
fraction = coordinate - interval.astype(dtype)
fraction2 = fraction * fraction
fraction3 = fraction2 * fraction
weights = (
(1.0 - fraction) ** 3 / 6.0,
(3.0 * fraction3 - 6.0 * fraction2 + 4.0) / 6.0,
(-3.0 * fraction3 + 3.0 * fraction2 + 3.0 * fraction + 1.0) / 6.0,
fraction3 / 6.0,
)
period = 2 * (control_size - 1)
rows = jnp.arange(output_size, dtype=jnp.int32)
matrix = jnp.zeros((output_size, control_size), dtype=dtype)
for weight, offset in zip(weights, (-1, 0, 1, 2), strict=True):
wrapped = jnp.mod(interval + offset, period)
indices = jnp.where(wrapped <= control_size - 1, wrapped, period - wrapped)
matrix = matrix.at[rows, indices].add(weight)
return matrix
def _cubic_reflecting_sampling_matrix(
coordinates: jax.Array,
control_size: int,
*,
dtype: jnp.dtype,
) -> jax.Array:
"""Evaluate a reflecting uniform cubic B-spline at arbitrary coordinates."""
if control_size < 4:
raise ValueError("cubic interpolation requires at least four controls")
coordinate = jnp.asarray(coordinates, dtype=dtype)
if coordinate.ndim != 1:
raise ValueError("coordinates must be one-dimensional")
interval = jnp.floor(coordinate).astype(jnp.int32)
fraction = coordinate - interval.astype(dtype)
fraction2 = fraction * fraction
fraction3 = fraction2 * fraction
weights = (
(1.0 - fraction) ** 3 / 6.0,
(3.0 * fraction3 - 6.0 * fraction2 + 4.0) / 6.0,
(-3.0 * fraction3 + 3.0 * fraction2 + 3.0 * fraction + 1.0) / 6.0,
fraction3 / 6.0,
)
period = 2 * (control_size - 1)
rows = jnp.arange(coordinate.size, dtype=jnp.int32)
matrix = jnp.zeros((coordinate.size, control_size), dtype=dtype)
for weight, offset in zip(weights, (-1, 0, 1, 2), strict=True):
wrapped = jnp.mod(interval + offset, period)
indices = jnp.where(wrapped <= control_size - 1, wrapped, period - wrapped)
matrix = matrix.at[rows, indices].add(weight)
return matrix
@dataclass(frozen=True)
class SplinePhaseField2D:
"""Topology-free smooth material interface with cut-cell FDTD fractions.
The optimized values are an *absolute signed level set* on a coarse cubic
B-spline control grid, rather than independent material densities. The
spline is sampled directly at subcell quadrature points and a compact C1
Heaviside converts its sign to material coverage. Thus the underlying
interface is continuous and smooth, uniform cells are exactly zero or one,
and only cells cut by the interface retain fractional fill in the sharp
limit.
Unlike :class:`SplineLevelSetBoundary2D`, this discovery representation
has no immutable reference contour or narrow-band window. Any control may
cross zero, so components and holes may appear or disappear while topology
is still being discovered. A campaign should slowly reduce
``interface_width``, then replay into a topology-guarded boundary chart
once its topology has stabilized.
"""
shape: tuple[int, int]
control_shape: tuple[int, int]
spacing: tuple[float, float] = (1.0, 1.0)
interface_width: float = 1.0
supersample: int = 4
def __post_init__(self) -> None:
_validate_design_shape(self.shape)
_validate_design_spacing(self.spacing)
if len(self.control_shape) != 2 or min(self.control_shape) < 4:
raise ValueError("control_shape must contain two dimensions of at least four")
if any(
control > cells + 1
for control, cells in zip(self.control_shape, self.shape, strict=True)
):
raise ValueError("control_shape cannot exceed the material vertex grid")
if not np.isfinite(self.interface_width) or self.interface_width <= 0.0:
raise ValueError("interface_width must be positive and finite")
if not isinstance(self.supersample, int) or isinstance(self.supersample, bool):
raise TypeError("supersample must be a positive integer")
if self.supersample < 1:
raise ValueError("supersample must be a positive integer")
@property
def parameter_shape(self) -> tuple[int, int]:
return self.control_shape
@property
def parameter_bounds(self) -> tuple[float, float]:
return (-math.inf, math.inf)
def _sampling_matrices(self, dtype: jnp.dtype) -> tuple[jax.Array, jax.Array]:
matrices = []
for cells, controls in zip(self.shape, self.control_shape, strict=True):
physical = (jnp.arange(cells * self.supersample, dtype=dtype) + 0.5) / (
cells * self.supersample
)
coordinates = physical * (controls - 1)
matrices.append(
_cubic_reflecting_sampling_matrix(coordinates, controls, dtype=dtype)
)
return matrices[0], matrices[1]
def latent(self, parameters: jax.Array) -> jax.Array:
"""Return the continuous phase field sampled on material vertices."""
parameters = jnp.asarray(parameters)
if tuple(parameters.shape) != self.control_shape:
raise ValueError(
f"phase-field parameters must have shape {self.control_shape}, got {parameters.shape}"
)
x_matrix = _cubic_reflecting_interpolation_matrix(
self.shape[0] + 1,
self.control_shape[0],
dtype=parameters.dtype,
)
y_matrix = _cubic_reflecting_interpolation_matrix(
self.shape[1] + 1,
self.control_shape[1],
dtype=parameters.dtype,
)
return x_matrix @ parameters @ y_matrix.T
def decode(self, parameters: jax.Array) -> jax.Array:
"""Integrate the smooth material indicator inside every FDTD cell."""
parameters = jnp.asarray(parameters)
if tuple(parameters.shape) != self.control_shape:
raise ValueError(
f"phase-field parameters must have shape {self.control_shape}, got {parameters.shape}"
)
x_matrix, y_matrix = self._sampling_matrices(parameters.dtype)
sampled = x_matrix @ parameters @ y_matrix.T
coordinate = jnp.clip(
0.5 + 0.5 * sampled / self.interface_width,
0.0,
1.0,
)
material = coordinate * coordinate * (3.0 - 2.0 * coordinate)
reshaped = material.reshape(
self.shape[0],
self.supersample,
self.shape[1],
self.supersample,
)
return jnp.mean(reshaped, axis=(1, 3))
def encode(self, density: jax.Array | np.ndarray) -> jax.Array:
"""Least-squares fit a signed-distance initializer to spline controls."""
values = np.asarray(density, dtype=np.float32)
if tuple(values.shape) != self.shape:
raise ValueError(
f"density must have shape {self.shape}, got {values.shape}"
)
level_set = signed_distance_level_set(values, spacing=self.spacing)
x_matrix = np.asarray(
_cubic_reflecting_interpolation_matrix(
self.shape[0] + 1,
self.control_shape[0],
dtype=jnp.float32,
)
)
y_matrix = np.asarray(
_cubic_reflecting_interpolation_matrix(
self.shape[1] + 1,
self.control_shape[1],
dtype=jnp.float32,
)
)
# Solve X C Y.T ~= phi without forming the Kronecker product.
left = np.linalg.lstsq(x_matrix, level_set, rcond=1.0e-8)[0]
controls = np.linalg.lstsq(y_matrix, left.T, rcond=1.0e-8)[0].T
if not np.all(np.isfinite(controls)):
raise RuntimeError("spline phase-field fit produced nonfinite controls")
return jnp.asarray(controls, dtype=jnp.float32)
def fit_to_density(
self,
density: jax.Array | np.ndarray,
*,
maximum_control_magnitude: float | None = None,
max_iterations: int = 500,
) -> jax.Array:
"""Fit the decoded cut-cell density without constraining topology.
:meth:`encode` is a fast signed-distance least-squares initializer.
Coarse spline interpolation can move its zero contour, however, so a
direct fit of the *decoded* cell fractions is preferable at a
density-to-boundary handoff. This host-side fit performs no Maxwell
solves and imposes no sign or component-lineage constraint: small
islands may disappear, holes may close, and nearby components may
fuse whenever that is the closest representable smooth contour.
``maximum_control_magnitude`` is an optional physical bound on the
absolute phase field. Clipping far-away saturated bulk values does
not move the zero contour, and gives an optimizer a finite normalized
coordinate chart after the handoff.
"""
from scipy.optimize import minimize
values = np.asarray(density, dtype=np.float32)
if tuple(values.shape) != self.shape:
raise ValueError(f"density must have shape {self.shape}, got {values.shape}")
if not np.all(np.isfinite(values)):
raise ValueError("density must contain only finite values")
if not isinstance(max_iterations, int) or max_iterations < 1:
raise ValueError("max_iterations must be a positive integer")
if maximum_control_magnitude is not None and (
not np.isfinite(maximum_control_magnitude)
or maximum_control_magnitude <= 0.0
):
raise ValueError("maximum_control_magnitude must be positive and finite")
initial = np.asarray(self.encode(values), dtype=np.float32)
if maximum_control_magnitude is not None:
initial = np.clip(
initial,
-maximum_control_magnitude,
maximum_control_magnitude,
)
target = jnp.asarray(values)
def loss(flat_controls: jax.Array) -> jax.Array:
controls = flat_controls.reshape(self.control_shape)
return jnp.mean(jnp.square(self.decode(controls) - target))
value_and_gradient = jax.jit(jax.value_and_grad(loss))
def objective(flat_controls: np.ndarray) -> tuple[float, np.ndarray]:
value, gradient = value_and_gradient(
jnp.asarray(flat_controls, dtype=jnp.float32)
)
return float(value), np.asarray(gradient, dtype=np.float64)
bounds = None
if maximum_control_magnitude is not None:
bounds = [
(-maximum_control_magnitude, maximum_control_magnitude)
] * initial.size
result = minimize(
objective,
initial.reshape(-1).astype(np.float64),
method="L-BFGS-B",
jac=True,
bounds=bounds,
options={
"maxiter": max_iterations,
"ftol": 1.0e-14,
"gtol": 1.0e-8,
"maxls": 30,
},
)
initial_value = objective(initial.reshape(-1))[0]
fitted = result.x if float(result.fun) <= initial_value else initial.reshape(-1)
controls = fitted.reshape(self.control_shape)
if not np.all(np.isfinite(controls)):
raise RuntimeError(
"spline phase-field density fit produced nonfinite controls"
)
return jnp.asarray(controls, dtype=jnp.float32)
@dataclass(frozen=True)
class SplineLevelSetBoundary2D:
"""Topology-guarded boundary motion on a coarse cubic-spline grid.
The immutable signed-distance field supplies the literal reference
interface. Parameters are physical normal displacements sampled on a
smaller tensor-product cubic B-spline grid. Their interpolated field is
multiplied by the same narrow-band window as :class:`LevelSetBoundary2D`,
so the Maxwell derivative remains a boundary sensitivity while neighboring
boundary segments move coherently.
"""
reference_level_set: np.ndarray
control_shape: tuple[int, int]
spacing: tuple[float, float] = (1.0, 1.0)
narrow_band_width: float = 4.0
maximum_displacement: float = 2.0
interface_width: float = 0.25
supersample: int = 4
threshold: float = 0.5
diagonal_connectivity: bool = False
exterior_edges: tuple[BoundaryEdge, ...] = _ALL_BOUNDARY_EDGES
def __post_init__(self) -> None:
level_set = np.asarray(self.reference_level_set, dtype=np.float32)
if level_set.ndim != 2 or min(level_set.shape) < 3:
raise ValueError("reference_level_set must be a two-dimensional vertex grid")
if len(self.control_shape) != 2 or min(self.control_shape) < 4:
raise ValueError("control_shape must contain two dimensions of at least four")
if any(
control > size
for control, size in zip(self.control_shape, level_set.shape, strict=True)
):
raise ValueError("control_shape cannot exceed the reference vertex grid")
validated = LevelSetBoundary2D(
reference_level_set=level_set,
spacing=self.spacing,
narrow_band_width=self.narrow_band_width,
maximum_displacement=self.maximum_displacement,
interface_width=self.interface_width,
supersample=self.supersample,
threshold=self.threshold,
diagonal_connectivity=self.diagonal_connectivity,
exterior_edges=self.exterior_edges,
)
frozen = np.asarray(validated.reference_level_set).copy()
frozen.setflags(write=False)
object.__setattr__(self, "reference_level_set", frozen)
object.__setattr__(self, "exterior_edges", validated.exterior_edges)
@classmethod
def from_density(
cls,
density: jax.Array | np.ndarray,
*,
control_shape: tuple[int, int],
spacing: tuple[float, float] = (1.0, 1.0),
threshold: float = 0.5,
narrow_band_width: float | None = None,
maximum_displacement: float | None = None,
interface_width: float | None = None,
supersample: int = 4,
diagonal_connectivity: bool = False,
exterior_edges: tuple[BoundaryEdge, ...] = _ALL_BOUNDARY_EDGES,
fit_reference: bool = True,
fit_max_iterations: int = 200,
) -> SplineLevelSetBoundary2D:
spacing = _validate_design_spacing(spacing)
grid_scale = max(spacing)
width = (
min(spacing) / supersample if interface_width is None else interface_width
)
if fit_reference:
reference_level_set = fit_level_set_to_density(
density,
spacing=spacing,
threshold=threshold,
supersample=supersample,
interface_width=width,
max_iterations=fit_max_iterations,
exterior_edges=exterior_edges,
)
else:
reference_level_set = signed_distance_level_set(
density,
spacing=spacing,
threshold=threshold,
)
return cls(
reference_level_set=reference_level_set,
control_shape=control_shape,
spacing=spacing,
narrow_band_width=(
4.0 * grid_scale if narrow_band_width is None else narrow_band_width
),
maximum_displacement=(
2.0 * grid_scale
if maximum_displacement is None
else maximum_displacement
),
interface_width=width,
supersample=supersample,
threshold=threshold,
diagonal_connectivity=diagonal_connectivity,
exterior_edges=exterior_edges,
)
@property
def shape(self) -> tuple[int, int]:
return (
self.reference_level_set.shape[0] - 1,
self.reference_level_set.shape[1] - 1,
)
@property
def parameter_shape(self) -> tuple[int, int]:
return self.control_shape
@property
def parameter_bounds(self) -> tuple[float, float]:
return (-self.maximum_displacement, self.maximum_displacement)
def initial_parameters(self, *, dtype: jnp.dtype = jnp.float32) -> jax.Array:
return jnp.zeros(self.parameter_shape, dtype=dtype)
def _pixel_basis(self) -> LevelSetBoundary2D:
return LevelSetBoundary2D(
reference_level_set=self.reference_level_set,
spacing=self.spacing,
narrow_band_width=self.narrow_band_width,
maximum_displacement=self.maximum_displacement,
interface_width=self.interface_width,
supersample=self.supersample,
threshold=self.threshold,
diagonal_connectivity=self.diagonal_connectivity,
exterior_edges=self.exterior_edges,
)
def _matrices(self, dtype: jnp.dtype) -> tuple[jax.Array, jax.Array]:
return (
_cubic_reflecting_interpolation_matrix(
self.reference_level_set.shape[0],
self.control_shape[0],
dtype=dtype,
),
_cubic_reflecting_interpolation_matrix(
self.reference_level_set.shape[1],
self.control_shape[1],
dtype=dtype,
),
)
def boundary_window(self, *, dtype: jnp.dtype = jnp.float32) -> jax.Array:
"""Return controls whose cubic support intersects the narrow band."""
x_matrix, y_matrix = self._matrices(dtype)
vertex_window = self._pixel_basis().boundary_window(dtype=dtype)
influence = x_matrix.T @ (vertex_window > 0.0).astype(dtype) @ y_matrix
return (influence > 0.0).astype(dtype)
def latent(self, parameters: jax.Array) -> jax.Array:
parameters = jnp.asarray(parameters)
if parameters.shape != self.parameter_shape:
raise ValueError(
f"spline parameters must have shape {self.parameter_shape}, got {parameters.shape}"
)
controls = jnp.clip(parameters, *self.parameter_bounds)
x_matrix, y_matrix = self._matrices(parameters.dtype)
displacement = x_matrix @ controls @ y_matrix.T
pixel = self._pixel_basis()
return jnp.asarray(self.reference_level_set, dtype=parameters.dtype) + (
pixel.boundary_window(dtype=parameters.dtype) * displacement
)
def decode(self, parameters: jax.Array) -> jax.Array:
return level_set_fill_fraction(
self.latent(parameters),
supersample=self.supersample,
interface_width=self.interface_width,
)
@property
def topology(self) -> BoundaryTopology:
return self._pixel_basis().topology
def preserves_topology(self, density: jax.Array | np.ndarray) -> bool:
return self._pixel_basis().preserves_topology(density)
def interface_normal(self, parameters: jax.Array) -> jax.Array:
level_set = self.latent(parameters)
grad_x, grad_y = jnp.gradient(level_set, self.spacing[0], self.spacing[1])
magnitude = jnp.sqrt(grad_x**2 + grad_y**2 + 1e-24)
return -jnp.stack((grad_x / magnitude, grad_y / magnitude), axis=0)
def rebase(
self,
parameters: jax.Array | np.ndarray,
*,
preserve_topology: bool = True,
) -> tuple[SplineLevelSetBoundary2D, jax.Array]:
density = np.asarray(self.decode(jnp.asarray(parameters)))
if preserve_topology and not self.preserves_topology(density):
raise ValueError("candidate spline boundary changes the reference topology")
rebased = type(self).from_density(
density,
control_shape=self.control_shape,
spacing=self.spacing,
threshold=self.threshold,
narrow_band_width=self.narrow_band_width,
maximum_displacement=self.maximum_displacement,
interface_width=self.interface_width,
supersample=self.supersample,
diagonal_connectivity=self.diagonal_connectivity,
exterior_edges=self.exterior_edges,
)
if preserve_topology and rebased.topology != self.topology:
raise ValueError("spline boundary rebasing changed the reference topology")
return rebased, rebased.initial_parameters()
[docs]
def topology_safe_boundary_step(
basis: LevelSetBoundary2D | SplineLevelSetBoundary2D,
parameters: jax.Array | np.ndarray,
normal_velocity: jax.Array | np.ndarray,
*,
step_size: float,
backtracking_scales: tuple[float, ...] = (1.0, 0.5, 0.25, 0.125, 0.0625),
) -> tuple[jax.Array, float]:
"""Take the largest proposed normal step that preserves interface topology.
``normal_velocity`` is normally the gradient of a scalar Maxwell objective
with respect to the basis parameters. The returned scale is zero when all
candidates would merge, split, create, or erase a boundary component. An
objective-aware optimizer should additionally accept only improving steps;
this helper supplies the independent geometric guard for its line search.
This is deliberately a host-side operation because connected-component
topology is discrete and should never be hidden inside a surrogate JAX
gradient.
"""
if not np.isfinite(step_size) or step_size <= 0.0:
raise ValueError("step_size must be positive and finite")
if not backtracking_scales or any(
not np.isfinite(scale) or not 0.0 < scale <= 1.0
for scale in backtracking_scales
):
raise ValueError("backtracking_scales must contain values in (0, 1]")
current = jnp.asarray(parameters)
velocity = jnp.asarray(normal_velocity, dtype=current.dtype)
if current.shape != basis.parameter_shape or velocity.shape != basis.parameter_shape:
raise ValueError(f"parameters and velocity must have shape {basis.parameter_shape}")
if not bool(jnp.all(jnp.isfinite(current))) or not bool(jnp.all(jnp.isfinite(velocity))):
raise ValueError("parameters and velocity must contain only finite values")
active_velocity = velocity * basis.boundary_window(dtype=current.dtype)
peak = float(jnp.max(jnp.abs(active_velocity)))
if peak == 0.0:
return current, 0.0
unit_velocity = active_velocity / peak
lower, upper = basis.parameter_bounds
for scale in backtracking_scales:
candidate = jnp.clip(current + scale * step_size * unit_velocity, lower, upper)
if basis.preserves_topology(basis.decode(candidate)):
return candidate, float(scale)
return current, 0.0
@dataclass(frozen=True)
class CosineBasis1D:
"""Orthonormal DCT-II basis for a finite radial or longitudinal line.
``mode_count`` retains the lowest modes, including the constant bias.
Unlike a periodic Fourier series, the DCT does not identify the two ends
of the line. A complete ``mode_count == size`` basis is an invertible
reparameterization of the sampled line; a truncated basis supplies a
smooth, globally correlated topology prior.
"""
size: int
mode_count: int
spacing: float = 1.0
minimum_period: float | None = None
maximum_period: float | None = None
activation: DesignActivation = "sigmoid"
def __post_init__(self) -> None:
size = _validate_design_size(self.size)
if not isinstance(self.mode_count, int) or isinstance(self.mode_count, bool) or self.mode_count < 1:
raise ValueError("mode_count must be a positive integer")
if self.mode_count > size:
raise ValueError("cosine mode_count cannot exceed the decoded line size")
if not np.isfinite(self.spacing) or self.spacing <= 0:
raise ValueError("spacing must be positive and finite")
_validate_period_band(self.minimum_period, self.maximum_period)
_activate_design_latent(jnp.asarray(0.0), self.activation)
if not self.mode_indices:
raise ValueError("the requested cosine period band contains no modes")
@property
def mode_indices(self) -> tuple[int, ...]:
modes = []
for mode in range(self.mode_count):
wave_number = math.pi * mode / (self.size * self.spacing)
if _period_is_selected(
wave_number,
minimum_period=self.minimum_period,
maximum_period=self.maximum_period,
):
modes.append(mode)
return tuple(modes)
@property
def parameter_shape(self) -> tuple[int]:
return (len(self.mode_indices),)
@property
def parameter_bounds(self) -> tuple[float, float]:
return (-math.inf, math.inf)
def _basis(self, dtype: jnp.dtype = jnp.float32) -> jax.Array:
modes = jnp.asarray(self.mode_indices, dtype=dtype)
coordinate = jnp.arange(self.size, dtype=dtype) + 0.5
normalization = jnp.where(
modes == 0,
jnp.sqrt(1.0 / self.size),
jnp.sqrt(2.0 / self.size),
)
return (
jnp.cos(jnp.pi * coordinate[:, None] * modes[None, :] / self.size)
* normalization[None, :]
)
def latent(self, parameters: jax.Array) -> jax.Array:
parameters = jnp.asarray(parameters)
if parameters.shape != self.parameter_shape:
raise ValueError(
f"cosine parameters must have shape {self.parameter_shape}, got {parameters.shape}"
)
return self._basis(parameters.dtype) @ parameters
def decode(self, parameters: jax.Array) -> jax.Array:
return _activate_design_latent(self.latent(parameters), self.activation)
def encode(self, density: jax.Array | np.ndarray) -> jax.Array:
latent = _inverse_design_activation_1d(density, self.activation)
if latent.shape != (self.size,):
raise ValueError(
f"density must have shape {(self.size,)}, got {latent.shape}"
)
basis = np.asarray(self._basis(jnp.float32), dtype=np.float64)
coefficients = np.einsum("xk,x->k", basis, latent)
return jnp.asarray(coefficients, dtype=jnp.float32)
@dataclass(frozen=True)
class CosineBasis2D:
"""Separable DCT-II material basis, ideal for mirror-reduced domains.
``mode_shape=(mx, my)`` retains the lowest ``mx * my`` cosine modes before
optional physical-period filtering. Cosines have zero normal derivative
at each design-box boundary, making this basis particularly natural on an
x/y mirror-reduced quadrant. The zero mode is always retained as a bias.
"""
shape: tuple[int, int]
mode_shape: tuple[int, int]
spacing: tuple[float, float] = (1.0, 1.0)
minimum_period: float | None = None
maximum_period: float | None = None
activation: DesignActivation = "sigmoid"
def __post_init__(self) -> None:
shape = _validate_design_shape(self.shape)
_validate_design_spacing(self.spacing)
_validate_period_band(self.minimum_period, self.maximum_period)
if len(self.mode_shape) != 2 or any(
not isinstance(value, int) or isinstance(value, bool) or value < 1 for value in self.mode_shape
):
raise ValueError("mode_shape must contain two positive integers")
if any(count > size for count, size in zip(self.mode_shape, shape, strict=True)):
raise ValueError("cosine mode_shape cannot exceed the decoded grid shape")
_activate_design_latent(jnp.asarray(0.0), self.activation)
if len(self.mode_indices) == 0:
raise ValueError("the requested cosine period band contains no modes")
@property
def mode_indices(self) -> tuple[tuple[int, int], ...]:
nx, ny = self.shape
dx, dy = self.spacing
modes = []
for mx in range(self.mode_shape[0]):
for my in range(self.mode_shape[1]):
wave_number = math.hypot(math.pi * mx / (nx * dx), math.pi * my / (ny * dy))
if _period_is_selected(
wave_number,
minimum_period=self.minimum_period,
maximum_period=self.maximum_period,
):
modes.append((mx, my))
return tuple(modes)
@property
def parameter_shape(self) -> tuple[int]:
return (len(self.mode_indices),)
@property
def parameter_bounds(self) -> tuple[float, float]:
return (-math.inf, math.inf)
def _basis(self, dtype: jnp.dtype = jnp.float32) -> jax.Array:
nx, ny = self.shape
modes = jnp.asarray(self.mode_indices, dtype=dtype)
x = jnp.arange(nx, dtype=dtype) + 0.5
y = jnp.arange(ny, dtype=dtype) + 0.5
mx = modes[:, 0]
my = modes[:, 1]
norm_x = jnp.where(mx == 0, jnp.sqrt(1.0 / nx), jnp.sqrt(2.0 / nx))
norm_y = jnp.where(my == 0, jnp.sqrt(1.0 / ny), jnp.sqrt(2.0 / ny))
along_x = jnp.cos(jnp.pi * x[:, None] * mx[None, :] / nx) * norm_x[None, :]
along_y = jnp.cos(jnp.pi * y[:, None] * my[None, :] / ny) * norm_y[None, :]
return along_x[:, None, :] * along_y[None, :, :]
def latent(self, parameters: jax.Array) -> jax.Array:
parameters = jnp.asarray(parameters)
if parameters.shape != self.parameter_shape:
raise ValueError(f"cosine parameters must have shape {self.parameter_shape}, got {parameters.shape}")
return jnp.einsum("xyk,k->xy", self._basis(parameters.dtype), parameters)
def decode(self, parameters: jax.Array) -> jax.Array:
return _activate_design_latent(self.latent(parameters), self.activation)
def encode(self, density: jax.Array | np.ndarray) -> jax.Array:
latent = _inverse_design_activation(density, self.activation)
if latent.shape != self.shape:
raise ValueError(f"density must have shape {self.shape}, got {latent.shape}")
basis = np.asarray(self._basis(jnp.float32), dtype=np.float64)
coefficients = np.einsum("xyk,xy->k", basis, latent)
return jnp.asarray(coefficients, dtype=jnp.float32)
@dataclass(frozen=True)
class FourierBasis2D:
"""General real plane-wave basis with independently learned phase.
Parameters are ordered as one bias followed by cosine amplitudes and sine
amplitudes for every unique reciprocal-grid vector. Unlike
:class:`CosineBasis2D`, this basis does not impose mirror parity.
"""
shape: tuple[int, int]
max_mode_indices: tuple[int, int]
spacing: tuple[float, float] = (1.0, 1.0)
minimum_period: float | None = None
maximum_period: float | None = None
activation: DesignActivation = "sigmoid"
def __post_init__(self) -> None:
shape = _validate_design_shape(self.shape)
_validate_design_spacing(self.spacing)
_validate_period_band(self.minimum_period, self.maximum_period)
if len(self.max_mode_indices) != 2 or any(
not isinstance(value, int) or isinstance(value, bool) or value < 0 for value in self.max_mode_indices
):
raise ValueError("max_mode_indices must contain two non-negative integers")
if any(limit > size // 2 for limit, size in zip(self.max_mode_indices, shape, strict=True)):
raise ValueError("Fourier mode indices cannot exceed the grid Nyquist indices")
_activate_design_latent(jnp.asarray(0.0), self.activation)
@property
def mode_indices(self) -> tuple[tuple[int, int], ...]:
nx, ny = self.shape
dx, dy = self.spacing
modes = []
for mx in range(-self.max_mode_indices[0], self.max_mode_indices[0] + 1):
for my in range(-self.max_mode_indices[1], self.max_mode_indices[1] + 1):
# Keep exactly one vector from each conjugate pair.
if not (mx > 0 or (mx == 0 and my > 0)):
continue
wave_number = math.hypot(2.0 * math.pi * mx / (nx * dx), 2.0 * math.pi * my / (ny * dy))
if _period_is_selected(
wave_number,
minimum_period=self.minimum_period,
maximum_period=self.maximum_period,
):
modes.append((mx, my))
return tuple(modes)
@property
def parameter_shape(self) -> tuple[int]:
return (1 + 2 * len(self.mode_indices),)
@property
def parameter_bounds(self) -> tuple[float, float]:
return (-math.inf, math.inf)
def _basis(self, dtype: jnp.dtype = jnp.float32) -> jax.Array:
nx, ny = self.shape
if not self.mode_indices:
return jnp.ones((*self.shape, 1), dtype=dtype)
modes = jnp.asarray(self.mode_indices, dtype=dtype)
x = (jnp.arange(nx, dtype=dtype) + 0.5) / nx
y = (jnp.arange(ny, dtype=dtype) + 0.5) / ny
phase = 2.0 * jnp.pi * (x[:, None, None] * modes[None, None, :, 0] + y[None, :, None] * modes[None, None, :, 1])
return jnp.concatenate(
(
jnp.ones((*self.shape, 1), dtype=dtype),
jnp.cos(phase),
jnp.sin(phase),
),
axis=2,
)
def latent(self, parameters: jax.Array) -> jax.Array:
parameters = jnp.asarray(parameters)
if parameters.shape != self.parameter_shape:
raise ValueError(f"Fourier parameters must have shape {self.parameter_shape}, got {parameters.shape}")
return jnp.einsum("xyk,k->xy", self._basis(parameters.dtype), parameters)
def decode(self, parameters: jax.Array) -> jax.Array:
return _activate_design_latent(self.latent(parameters), self.activation)
def encode(self, density: jax.Array | np.ndarray) -> jax.Array:
latent = _inverse_design_activation(density, self.activation)
if latent.shape != self.shape:
raise ValueError(f"density must have shape {self.shape}, got {latent.shape}")
matrix = np.asarray(self._basis(jnp.float32), dtype=np.float64).reshape((-1, self.parameter_shape[0]))
coefficients, *_ = np.linalg.lstsq(matrix, latent.reshape(-1), rcond=None)
return jnp.asarray(coefficients, dtype=jnp.float32)
@dataclass(frozen=True)
class FFTGridBasis2D:
r"""Dense real Fourier grid decoded with an :math:`O(N\log N)` FFT.
A real material field has one independent real degree of freedom per grid
cell, even though its Fourier spectrum is complex. This class packs each
self-conjugate reciprocal point as one real value and every other Hermitian
pair as one real and one imaginary value. With no period band,
``parameter_shape == (prod(shape),)`` exactly: there are no redundant or
aliased coefficients.
Use this basis when thousands of Fourier modes are desired. The explicit
:class:`FourierBasis2D` remains preferable for a small hand-selected bank,
because it exposes each cosine/sine wave directly.
"""
shape: tuple[int, int]
spacing: tuple[float, float] = (1.0, 1.0)
minimum_period: float | None = None
maximum_period: float | None = None
activation: DesignActivation = "sigmoid"
def __post_init__(self) -> None:
_validate_design_shape(self.shape)
_validate_design_spacing(self.spacing)
_validate_period_band(self.minimum_period, self.maximum_period)
_activate_design_latent(jnp.asarray(0.0), self.activation)
if self.parameter_shape[0] == 0:
raise ValueError("the requested FFT period band contains no modes")
@staticmethod
def _signed_index(index: int, size: int) -> int:
return index if index <= size // 2 else index - size
def _period_selected(self, index: tuple[int, int]) -> bool:
nx, ny = self.shape
dx, dy = self.spacing
mx = self._signed_index(index[0], nx)
my = self._signed_index(index[1], ny)
wave_number = math.hypot(2.0 * math.pi * mx / (nx * dx), 2.0 * math.pi * my / (ny * dy))
return _period_is_selected(
wave_number,
minimum_period=self.minimum_period,
maximum_period=self.maximum_period,
)
@property
def self_conjugate_indices(self) -> tuple[tuple[int, int], ...]:
nx, ny = self.shape
result = []
for ix in range(nx):
for iy in range(ny):
partner = ((-ix) % nx, (-iy) % ny)
if partner == (ix, iy) and self._period_selected((ix, iy)):
result.append((ix, iy))
return tuple(result)
@property
def paired_indices(self) -> tuple[tuple[int, int], ...]:
nx, ny = self.shape
result = []
for ix in range(nx):
for iy in range(ny):
partner = ((-ix) % nx, (-iy) % ny)
flat = ix * ny + iy
partner_flat = partner[0] * ny + partner[1]
if flat < partner_flat and self._period_selected((ix, iy)):
result.append((ix, iy))
return tuple(result)
@property
def conjugate_indices(self) -> tuple[tuple[int, int], ...]:
nx, ny = self.shape
return tuple(((-ix) % nx, (-iy) % ny) for ix, iy in self.paired_indices)
@property
def parameter_shape(self) -> tuple[int]:
return (len(self.self_conjugate_indices) + 2 * len(self.paired_indices),)
@property
def parameter_bounds(self) -> tuple[float, float]:
return (-math.inf, math.inf)
@property
def reciprocal_wave_count(self) -> int:
"""Number of unique real waves, counting a conjugate pair once."""
return len(self.self_conjugate_indices) + len(self.paired_indices)
def latent(self, parameters: jax.Array) -> jax.Array:
parameters = jnp.asarray(parameters)
if parameters.shape != self.parameter_shape:
raise ValueError(f"FFT-grid parameters must have shape {self.parameter_shape}, got {parameters.shape}")
self_count = len(self.self_conjugate_indices)
pair_count = len(self.paired_indices)
complex_dtype = jnp.complex128 if parameters.dtype == jnp.float64 else jnp.complex64
spectrum = jnp.zeros(self.shape, dtype=complex_dtype)
if self_count:
self_indices = jnp.asarray(self.self_conjugate_indices, dtype=jnp.int32)
spectrum = spectrum.at[self_indices[:, 0], self_indices[:, 1]].set(parameters[:self_count])
if pair_count:
pair_indices = jnp.asarray(self.paired_indices, dtype=jnp.int32)
conjugate_indices = jnp.asarray(self.conjugate_indices, dtype=jnp.int32)
real = parameters[self_count : self_count + pair_count]
imaginary = parameters[self_count + pair_count :]
values = real + 1j * imaginary
spectrum = spectrum.at[pair_indices[:, 0], pair_indices[:, 1]].set(values)
spectrum = spectrum.at[conjugate_indices[:, 0], conjugate_indices[:, 1]].set(jnp.conj(values))
return jnp.real(jnp.fft.ifft2(spectrum, norm="ortho"))
def decode(self, parameters: jax.Array) -> jax.Array:
return _activate_design_latent(self.latent(parameters), self.activation)
def encode(self, density: jax.Array | np.ndarray) -> jax.Array:
latent = _inverse_design_activation(density, self.activation)
if latent.shape != self.shape:
raise ValueError(f"density must have shape {self.shape}, got {latent.shape}")
spectrum = np.fft.fft2(latent, norm="ortho")
self_values = np.asarray([spectrum[index].real for index in self.self_conjugate_indices])
pair_values = np.asarray([spectrum[index] for index in self.paired_indices])
packed = np.concatenate((self_values, pair_values.real, pair_values.imag))
return jnp.asarray(packed, dtype=jnp.float32)
@dataclass(frozen=True)
class RadialCosineBasis2D:
"""Concentric circular or elliptical waves with optional learned phase.
This basis is useful for bullseye cavities and facing curved Bragg mirrors.
``axis_scale`` controls the ellipse metric and ``window_sigma`` optionally
turns the waves into localized radial/Gabor-like features. The parameter
order is bias, cosine amplitudes, then (when enabled) sine amplitudes.
"""
shape: tuple[int, int]
spacing: tuple[float, float]
periods: tuple[float, ...]
center: tuple[float, float] = (0.0, 0.0)
axis_scale: tuple[float, float] = (1.0, 1.0)
window_sigma: tuple[float, float] | None = None
learn_phase: bool = True
activation: DesignActivation = "sigmoid"
def __post_init__(self) -> None:
_validate_design_shape(self.shape)
_validate_design_spacing(self.spacing)
if len(self.periods) == 0 or any(not np.isfinite(value) or value <= 0 for value in self.periods):
raise ValueError("periods must contain positive finite values")
if len(set(self.periods)) != len(self.periods):
raise ValueError("radial periods must be unique")
if len(self.center) != 2 or any(not np.isfinite(value) for value in self.center):
raise ValueError("center must contain two finite coordinates")
if len(self.axis_scale) != 2 or any(not np.isfinite(value) or value <= 0 for value in self.axis_scale):
raise ValueError("axis_scale must contain two positive finite values")
if self.window_sigma is not None and (
len(self.window_sigma) != 2 or any(not np.isfinite(value) or value <= 0 for value in self.window_sigma)
):
raise ValueError("window_sigma must contain two positive finite values")
_activate_design_latent(jnp.asarray(0.0), self.activation)
@property
def parameter_shape(self) -> tuple[int]:
phase_count = 2 if self.learn_phase else 1
return (1 + phase_count * len(self.periods),)
@property
def parameter_bounds(self) -> tuple[float, float]:
return (-math.inf, math.inf)
def _basis(self, dtype: jnp.dtype = jnp.float32) -> jax.Array:
nx, ny = self.shape
dx, dy = self.spacing
x = ((jnp.arange(nx, dtype=dtype) + 0.5) - nx / 2.0) * dx - self.center[0]
y = ((jnp.arange(ny, dtype=dtype) + 0.5) - ny / 2.0) * dy - self.center[1]
radius = jnp.sqrt((x[:, None] / self.axis_scale[0]) ** 2 + (y[None, :] / self.axis_scale[1]) ** 2)
periods = jnp.asarray(self.periods, dtype=dtype)
phase = 2.0 * jnp.pi * radius[:, :, None] / periods[None, None, :]
if self.window_sigma is None:
window = jnp.ones(self.shape, dtype=dtype)
else:
window = jnp.exp(
-0.5 * (x[:, None] / self.window_sigma[0]) ** 2 - 0.5 * (y[None, :] / self.window_sigma[1]) ** 2
)
components = [jnp.ones((*self.shape, 1), dtype=dtype), window[:, :, None] * jnp.cos(phase)]
if self.learn_phase:
components.append(window[:, :, None] * jnp.sin(phase))
return jnp.concatenate(tuple(components), axis=2)
def latent(self, parameters: jax.Array) -> jax.Array:
parameters = jnp.asarray(parameters)
if parameters.shape != self.parameter_shape:
raise ValueError(f"radial parameters must have shape {self.parameter_shape}, got {parameters.shape}")
return jnp.einsum("xyk,k->xy", self._basis(parameters.dtype), parameters)
def decode(self, parameters: jax.Array) -> jax.Array:
return _activate_design_latent(self.latent(parameters), self.activation)
def encode(self, density: jax.Array | np.ndarray) -> jax.Array:
latent = _inverse_design_activation(density, self.activation)
if latent.shape != self.shape:
raise ValueError(f"density must have shape {self.shape}, got {latent.shape}")
matrix = np.asarray(self._basis(jnp.float32), dtype=np.float64).reshape((-1, self.parameter_shape[0]))
coefficients, *_ = np.linalg.lstsq(matrix, latent.reshape(-1), rcond=None)
return jnp.asarray(coefficients, dtype=jnp.float32)
def automatic_pulse_time_window(
*,
spectral_width_hz: float,
path_length_m: float,
max_refractive_index: float,
source_offset_sigmas: float = 6.0,
tail_sigmas: float = 5.0,
transit_passes: float = 2.0,
max_time: float | None = None,
) -> float:
"""Choose a conservative fixed window for reversible pulsed FDTD.
Reversible adjoints must size their boundary tape before a run, so a
data-dependent early-stop cannot reduce their memory allocation. This
device-independent estimate includes the complete Gaussian source window
and a configurable number of worst-index traversals of the scene. A
caller may retain a reference solver's longer run time as ``max_time``;
the returned value never exceeds that cap.
"""
if spectral_width_hz <= 0 or path_length_m <= 0 or max_refractive_index < 1:
raise ValueError(
"spectral_width_hz and path_length_m must be positive and max_refractive_index must be at least one"
)
if source_offset_sigmas < 0 or tail_sigmas <= 0 or transit_passes <= 0:
raise ValueError("pulse-window sigma counts and transit_passes must be positive")
sigma_t = 1.0 / (2.0 * math.pi * spectral_width_hz)
source_window = (source_offset_sigmas + tail_sigmas) * sigma_t
transit_window = transit_passes * path_length_m * max_refractive_index / c
estimate = source_window + transit_window
if max_time is not None:
if max_time <= 0:
raise ValueError("max_time must be positive when provided")
estimate = min(estimate, max_time)
return float(estimate)
def extend_periodic_mirror_to_boundary(
centers: np.ndarray,
radii: np.ndarray,
*,
period: float,
radius: float,
maximum_extent: float,
end_clearance: float = 0.0,
) -> tuple[np.ndarray, np.ndarray]:
"""Complete a one-sided periodic hole mirror without entering a boundary layer.
High-Q cavity optimization is easily limited by an artificially short mirror
rather than by the cavity taper itself. This discrete setup helper appends as
many complete nominal periods as fit before the caller's non-PML extent. The
dielectric beam may and should continue through the PML; the appended holes do
not. Geometry controls can subsequently optimize the returned centers/radii.
All lengths use the caller's consistent unit system. Existing centers must be
positive and strictly increasing, so the function works with metres,
micrometres, or normalized lattice units.
"""
centers = np.asarray(centers, dtype=np.float64)
radii = np.asarray(radii, dtype=np.float64)
if centers.ndim != 1 or radii.ndim != 1 or centers.shape != radii.shape or centers.size == 0:
raise ValueError("centers and radii must be non-empty one-dimensional arrays with equal shapes")
if np.any(centers <= 0) or np.any(np.diff(centers) <= 0):
raise ValueError("centers must be positive and strictly increasing")
if np.any(radii <= 0) or period <= 0 or radius <= 0 or maximum_extent <= 0 or end_clearance < 0:
raise ValueError("radii, period, radius, and maximum_extent must be positive; clearance must be non-negative")
if maximum_extent <= centers[-1] + radii[-1] + end_clearance:
return centers.copy(), radii.copy()
out_centers = centers.tolist()
out_radii = radii.tolist()
next_center = out_centers[-1] + period
while next_center + radius + end_clearance <= maximum_extent:
out_centers.append(next_center)
out_radii.append(radius)
next_center += period
return np.asarray(out_centers, dtype=np.float64), np.asarray(out_radii, dtype=np.float64)
def differentiable_circle_union_fill_fraction(
x_edges: jax.Array,
y_edges: jax.Array,
centers_x: jax.Array,
radii: jax.Array,
*,
supersample: int = 8,
edge_width: float,
) -> jax.Array:
"""Rasterize binary circular geometry with a smooth shape derivative.
The returned *forward* value is the exact supersampled area fraction of the
union of the circles. Its derivative is supplied by a sigmoid approximation
to the same boundary (a straight-through estimator). This prevents an
optimizer from exploiting grey interface cells that are absent when a final
literal geometry is rebuilt, while retaining useful GPU-native shape
derivatives with respect to circle centers and radii.
``x_edges`` and ``y_edges`` are physical cell edges. All circles are centered
at ``(centers_x[i], 0)``; this common nanophotonic case covers hole tapers and
one-dimensional photonic-crystal mirrors. The output has shape
``(len(x_edges) - 1, len(y_edges) - 1)``.
"""
x_edges = jnp.asarray(x_edges)
y_edges = jnp.asarray(y_edges)
centers_x = jnp.asarray(centers_x)
radii = jnp.asarray(radii)
if x_edges.ndim != 1 or y_edges.ndim != 1 or x_edges.size < 2 or y_edges.size < 2:
raise ValueError("x_edges and y_edges must be one-dimensional arrays with at least two entries")
if centers_x.ndim != 1 or radii.ndim != 1 or centers_x.shape != radii.shape or centers_x.size == 0:
raise ValueError("centers_x and radii must be non-empty one-dimensional arrays with equal shapes")
if supersample < 1:
raise ValueError("supersample must be at least one")
if edge_width <= 0:
raise ValueError("edge_width must be positive")
dtype = jnp.result_type(x_edges, y_edges, centers_x, radii)
offsets = (jnp.arange(supersample, dtype=dtype) + 0.5) / supersample
x_sub = x_edges[:-1, None] + offsets[None, :] * jnp.diff(x_edges)[:, None]
y_sub = y_edges[:-1, None] + offsets[None, :] * jnp.diff(y_edges)[:, None]
hard_outside = jnp.ones((x_sub.shape[0], y_sub.shape[0], supersample, supersample), dtype=bool)
soft_outside = jnp.ones(hard_outside.shape, dtype=dtype)
for center, radius in zip(centers_x, radii, strict=True):
dx = x_sub[:, None, :, None] - center
dy = y_sub[None, :, None, :]
distance = jnp.sqrt(dx**2 + dy**2)
hard_inside = distance < radius
soft_inside = jax.nn.sigmoid((radius - distance) / edge_width)
hard_outside = jnp.logical_and(hard_outside, jnp.logical_not(hard_inside))
soft_outside = soft_outside * (1.0 - soft_inside)
hard_fill = jnp.mean(jnp.logical_not(hard_outside), axis=(2, 3), dtype=dtype)
soft_fill = jnp.mean(1.0 - soft_outside, axis=(2, 3))
return soft_fill + jax.lax.stop_gradient(hard_fill - soft_fill)
def fill_fraction_interface_normal(
fill_fraction: jax.Array,
*,
spacing: tuple[float, float],
) -> jax.Array:
"""Return the metric-aware in-plane normal of a 2-D fill fraction."""
fill_fraction = jnp.asarray(fill_fraction)
if fill_fraction.ndim != 2 or min(fill_fraction.shape) < 2:
raise ValueError("fill_fraction must be a two-dimensional array at least two cells wide")
if len(spacing) != 2 or spacing[0] <= 0 or spacing[1] <= 0:
raise ValueError("spacing must contain two positive values")
grad_x, grad_y = jnp.gradient(fill_fraction, spacing[0], spacing[1])
gradient = -jnp.stack((grad_x, grad_y, jnp.zeros_like(fill_fraction)), axis=0)
magnitude_squared = jnp.sum(gradient**2, axis=0)
# sqrt has an undefined derivative at exactly zero. The tiny squared floor
# leaves every nonzero forward normal unchanged at float precision while
# keeping reverse-mode derivatives finite in uniform bulk regions.
safe_magnitude = jnp.sqrt(magnitude_squared + 1e-24)
return jnp.where(
magnitude_squared[None, ...] > 1e-24,
gradient / safe_magnitude,
0.0,
)
def diagonal_subpixel_inverse_permittivity(
fill_fraction: jax.Array,
normal: jax.Array,
*,
background_permittivity: float,
inclusion_permittivity: float,
) -> jax.Array:
"""Apply diagonal Farjadpour smoothing to a binary fill-fraction field."""
fill_fraction = jnp.asarray(fill_fraction)
normal = jnp.asarray(normal)
if normal.shape != (3, *fill_fraction.shape):
raise ValueError(f"normal must have shape {(3, *fill_fraction.shape)}, got {normal.shape}")
if background_permittivity <= 0 or inclusion_permittivity <= 0:
raise ValueError("permittivities must be positive")
eps_background = jnp.asarray(background_permittivity, dtype=fill_fraction.dtype)
eps_inclusion = jnp.asarray(inclusion_permittivity, dtype=fill_fraction.dtype)
eps_bar = fill_fraction * eps_inclusion + (1.0 - fill_fraction) * eps_background
eps_harmonic = 1.0 / (fill_fraction / eps_inclusion + (1.0 - fill_fraction) / eps_background)
delta = eps_bar - eps_harmonic
epsilon = jnp.stack([eps_bar - delta * normal[axis] ** 2 for axis in range(3)], axis=0)
return 1.0 / epsilon
def _shape_elements(shape: tuple[int, ...]) -> int:
return reduce(mul, shape, 1)
def recorder_bytes_per_step(
input_shape_dtypes: Mapping[str, jax.ShapeDtypeStruct],
*,
storage_dtype: jnp.dtype | None = None,
) -> int:
"""Return bytes required for one stored boundary-history sample."""
total = 0
for spec in input_shape_dtypes.values():
dtype = np.dtype(storage_dtype if storage_dtype is not None else spec.dtype)
total += _shape_elements(tuple(spec.shape)) * dtype.itemsize
return total
def accelerator_memory_limit_bytes(*, fallback_bytes: int = 8 * 1024**3) -> int:
"""Return the memory visible to JAX without depending on vendor tooling."""
try:
stats = jax.devices()[0].memory_stats() or {}
limit = int(stats.get("bytes_limit", 0))
except (IndexError, RuntimeError, TypeError, ValueError):
limit = 0
return limit if limit > 0 else fallback_bytes
def plan_reversible_recorder(
input_shape_dtypes: Mapping[str, jax.ShapeDtypeStruct],
max_time_steps: int,
*,
backend: BackendOption = "gpu",
device_limit_bytes: int | None = None,
tape_budget_fraction: float = 0.12,
max_temporal_stride: int = 8,
) -> tuple[Recorder, RecorderMemoryPlan]:
"""Build the highest-fidelity reversible recorder that fits automatically.
The planner first tries an exact float32 tape, then float16 storage, then
progressively wider linear-reconstruction strides. Only the boundary
tape receives this budget; the conservative default leaves 88% of JAX's
visible memory for fields, detector state, compilation, and reverse
workspace. A failure is explicit instead of allowing an opaque GPU OOM.
"""
if max_time_steps < 2:
raise ValueError("max_time_steps must be at least 2")
if not 0.0 < tape_budget_fraction < 1.0:
raise ValueError("tape_budget_fraction must lie strictly between 0 and 1")
if max_temporal_stride < 1:
raise ValueError("max_temporal_stride must be positive")
limit = accelerator_memory_limit_bytes() if device_limit_bytes is None else int(device_limit_bytes)
budget = int(limit * tape_budget_fraction)
raw_per_step = recorder_bytes_per_step(input_shape_dtypes)
raw_bytes = raw_per_step * max_time_steps
candidates: list[tuple[jnp.dtype | None, int]] = [(None, 1), (jnp.float16, 1)]
stride = 2
while stride <= max_temporal_stride:
candidates.append((jnp.float16, stride))
stride *= 2
for storage_dtype, temporal_stride in candidates:
latent_steps = max_time_steps if temporal_stride == 1 else math.ceil((max_time_steps - 1) / temporal_stride) + 1
per_step = recorder_bytes_per_step(input_shape_dtypes, storage_dtype=storage_dtype)
estimated = per_step * latent_steps
if estimated > budget:
continue
modules = []
if storage_dtype is not None:
modules.append(DtypeConversion(dtype=storage_dtype))
if temporal_stride > 1:
modules.append(LinearReconstructEveryK(k=temporal_stride))
recorder = Recorder(modules=modules)
plan = RecorderMemoryPlan(
raw_bytes=raw_bytes,
estimated_bytes=estimated,
budget_bytes=budget,
device_limit_bytes=limit,
storage_dtype=(
np.dtype(next(iter(input_shape_dtypes.values())).dtype).name
if storage_dtype is None
else np.dtype(storage_dtype).name
),
temporal_stride=temporal_stride,
latent_steps=latent_steps,
)
return recorder, plan
minimum = recorder_bytes_per_step(input_shape_dtypes, storage_dtype=jnp.float16) * (
math.ceil((max_time_steps - 1) / max_temporal_stride) + 1
)
raise MemoryError(
"reversible boundary tape does not fit the automatic memory budget: "
f"minimum={minimum / 1024**3:.2f} GiB, "
f"budget={budget / 1024**3:.2f} GiB, "
f"max_temporal_stride={max_temporal_stride}. Use local spatial refinement, "
"a shorter converged time window, or a larger explicitly validated stride."
)
[docs]
def initialize_planned_recorder(
input_shape_dtypes: Mapping[str, jax.ShapeDtypeStruct],
max_time_steps: int,
*,
backend: BackendOption = "gpu",
**plan_kwargs,
):
"""Plan a recorder and allocate its state in one call."""
recorder, plan = plan_reversible_recorder(
input_shape_dtypes,
max_time_steps,
backend=backend,
**plan_kwargs,
)
recorder, state = recorder.init_state(
input_shape_dtypes=dict(input_shape_dtypes),
max_time_steps=max_time_steps,
backend=backend,
)
return recorder, state, plan
def conic_filter_kernel(radius: float, spacing: float) -> jax.Array:
"""Return the normalized conic kernel used by Tidy3D topology examples."""
if radius <= 0 or spacing <= 0:
raise ValueError("radius and spacing must be positive")
radius_pixels = math.ceil(radius / spacing)
coordinates = jnp.linspace(-1.0, 1.0, 2 * radius_pixels + 1, dtype=jnp.float32)
xx, yy = jnp.meshgrid(coordinates, coordinates, indexing="ij")
kernel = jnp.maximum(0.0, 1.0 - jnp.sqrt(xx**2 + yy**2))
return kernel / jnp.sum(kernel)
def conic_filter_kernel_1d(radius: float, spacing: float) -> jax.Array:
"""Return a normalized one-dimensional conic fabrication kernel.
This is the radial analogue of :func:`conic_filter_kernel`. It supports
bodies of revolution whose complete mask is controlled by one radial line
while retaining a physical minimum-feature radius.
"""
if radius <= 0 or spacing <= 0:
raise ValueError("radius and spacing must be positive")
radius_pixels = math.ceil(radius / spacing)
coordinate = jnp.arange(-radius_pixels, radius_pixels + 1, dtype=jnp.float32)
kernel = jnp.maximum(0.0, 1.0 - jnp.abs(coordinate) * spacing / radius)
return kernel / jnp.sum(kernel)
def conic_filter_1d(
parameters: jax.Array,
*,
radius: float,
spacing: float,
padding: str = "reflect",
) -> jax.Array:
"""Apply a conic minimum-feature filter to a one-dimensional design."""
parameters = jnp.asarray(parameters)
if parameters.ndim != 1:
raise ValueError("parameters must be one-dimensional")
if padding not in {"reflect", "constant"}:
raise ValueError("padding must be 'reflect' or 'constant'")
kernel = conic_filter_kernel_1d(radius, spacing)
pad = kernel.shape[0] // 2
padded = jnp.pad(parameters, (pad, pad), mode=padding)
return jnp.convolve(padded, kernel, mode="valid")
[docs]
def conic_filter(
parameters: jax.Array,
*,
radius: float,
spacing: float,
padding: str = "reflect",
) -> jax.Array:
"""Apply a conic minimum-feature filter with explicit boundary padding."""
kernel = conic_filter_kernel(radius, spacing)
pad = kernel.shape[0] // 2
if padding not in {"reflect", "constant"}:
raise ValueError("padding must be 'reflect' or 'constant'")
padded = jnp.pad(parameters, ((pad, pad), (pad, pad)), mode=padding)
return convolve2d(padded, kernel, mode="valid")
[docs]
def tanh_projection(values: jax.Array, *, beta: float | jax.Array, eta: float = 0.5) -> jax.Array:
"""Smoothly project filtered values toward a binary material system."""
beta_array = jnp.asarray(beta, dtype=values.dtype)
eta_array = jnp.asarray(eta, dtype=values.dtype)
return (jnp.tanh(beta_array * eta_array) + jnp.tanh(beta_array * (values - eta_array))) / (
jnp.tanh(beta_array * eta_array) + jnp.tanh(beta_array * (1.0 - eta_array))
)
def filter_and_project(
parameters: jax.Array,
*,
radius: float,
spacing: float,
beta: float | jax.Array,
eta: float = 0.5,
padding: str = "reflect",
) -> jax.Array:
"""Apply the standard differentiable fabrication transform."""
return tanh_projection(
conic_filter(parameters, radius=radius, spacing=spacing, padding=padding),
beta=beta,
eta=eta,
)
def filter_and_project_1d(
parameters: jax.Array,
*,
radius: float,
spacing: float,
beta: float | jax.Array,
eta: float = 0.5,
padding: str = "reflect",
) -> jax.Array:
"""Filter and project a one-dimensional radial or longitudinal design."""
return tanh_projection(
conic_filter_1d(parameters, radius=radius, spacing=spacing, padding=padding),
beta=beta,
eta=eta,
)
def fabrication_density_2d(
parameters: jax.Array,
basis: DesignBasis2D,
*,
beta: float | jax.Array,
eta: float = 0.5,
filter_radius: float | None = None,
filter_spacing: float = 1.0,
trainable_mask: jax.Array | np.ndarray | None = None,
fixed_density: jax.Array | np.ndarray | None = None,
exact_binary: bool = False,
binary_threshold: float = 0.5,
padding: str = "reflect",
) -> jax.Array:
"""Decode a basis and apply one common fabrication-aware density map.
This function is the bridge between any :class:`DesignBasis2D` and an
FDTD material array. Hard-air regions have ``trainable_mask=0`` and
``fixed_density=0``; fixed material has ``trainable_mask=0`` and a nonzero
``fixed_density``. The constraints are imposed both before filtering and
after projection so neither operation can bleed material into a hard gap.
``exact_binary=True`` uses a straight-through estimator: the forward FDTD
geometry is literal binary while its reverse derivative is that of the
smooth projection. Erosion/dilation robustness can be evaluated by
calling this function at several ``eta`` values with the same parameters.
"""
density = basis.decode(parameters)
if density.ndim != 2 or tuple(density.shape) != tuple(basis.shape):
raise ValueError(f"decoded density must have shape {basis.shape}, got {density.shape}")
if filter_spacing <= 0:
raise ValueError("filter_spacing must be positive")
if not 0.0 <= binary_threshold <= 1.0:
raise ValueError("binary_threshold must lie in [0, 1]")
trainable = jnp.ones(basis.shape, dtype=density.dtype) if trainable_mask is None else jnp.asarray(trainable_mask)
fixed = jnp.zeros(basis.shape, dtype=density.dtype) if fixed_density is None else jnp.asarray(fixed_density)
if trainable.shape != basis.shape or fixed.shape != basis.shape:
raise ValueError("trainable_mask and fixed_density must match the decoded basis shape")
trainable = trainable.astype(density.dtype)
fixed = fixed.astype(density.dtype)
pinned = density * trainable + fixed
if filter_radius is None:
filtered = pinned
else:
if filter_radius <= 0:
raise ValueError("filter_radius must be positive when provided")
filtered = conic_filter(
pinned,
radius=filter_radius,
spacing=filter_spacing,
padding=padding,
)
projected = tanh_projection(filtered, beta=beta, eta=eta)
if exact_binary:
binary = (projected >= binary_threshold).astype(projected.dtype)
projected = projected + jax.lax.stop_gradient(binary - projected)
return jnp.clip(projected * trainable + fixed, 0.0, 1.0)
@jax.custom_jvp
def _slab_fill_fractions_traced(
edges: jax.Array,
thickness: jax.Array,
center: jax.Array,
) -> jax.Array:
lower = center - 0.5 * thickness
upper = center + 0.5 * thickness
overlap = jnp.maximum(
0.0, jnp.minimum(edges[1:], upper) - jnp.maximum(edges[:-1], lower)
)
return jnp.clip(overlap / jnp.diff(edges), 0.0, 1.0)
@_slab_fill_fractions_traced.defjvp
def _slab_fill_fractions_jvp(primals, tangents):
"""Use the centered shape derivative when an interface lies on an edge."""
edges, thickness, center = primals
_, thickness_tangent, center_tangent = tangents
fractions = _slab_fill_fractions_traced(edges, thickness, center)
lower = center - 0.5 * thickness
upper = center + 0.5 * thickness
cell_lower = edges[:-1]
cell_upper = edges[1:]
inverse_width = 1.0 / jnp.diff(edges)
def interface_weight(position):
interior = (position > cell_lower) & (position < cell_upper)
on_edge = (position == cell_lower) | (position == cell_upper)
return (interior.astype(edges.dtype) + 0.5 * on_edge.astype(edges.dtype)) * inverse_width
upper_weight = interface_weight(upper)
lower_weight = interface_weight(lower)
thickness_derivative = 0.5 * (upper_weight + lower_weight)
center_derivative = upper_weight - lower_weight
tangent = (
thickness_derivative * thickness_tangent
+ center_derivative * center_tangent
)
return fractions, tangent
def slab_fill_fractions(
z_edges: jax.Array | np.ndarray,
*,
thickness: float | jax.Array,
center: float | jax.Array = 0.0,
) -> jax.Array:
"""Return exact, differentiable cell overlaps for an extruded slab.
``thickness`` and ``center`` may be traced scalar arrays, allowing layer
thicknesses to be optimized directly through an FDTD material raster. A
Python scalar is still validated eagerly; traced callers are responsible
for applying positive bounds (normally with a sigmoid parameterization).
"""
edges = jnp.asarray(z_edges)
if edges.ndim != 1 or edges.shape[0] < 2:
raise ValueError("z_edges must be a one-dimensional array with at least two entries")
thickness_array = jnp.asarray(thickness, dtype=edges.dtype)
center_array = jnp.asarray(center, dtype=edges.dtype)
if thickness_array.ndim != 0 or center_array.ndim != 0:
raise ValueError("thickness and center must be scalars")
if isinstance(thickness, (int, float, np.integer, np.floating)) and (
float(thickness) <= 0 or not np.isfinite(float(thickness))
):
raise ValueError("thickness must be positive and finite")
if isinstance(center, (int, float, np.integer, np.floating)) and not np.isfinite(
float(center)
):
raise ValueError("center must be finite")
return _slab_fill_fractions_traced(edges, thickness_array, center_array)
def extrude_density_2d(
density: jax.Array,
z_fill_fractions: jax.Array | np.ndarray,
) -> jax.Array:
"""Extrude a 2-D material density through fractional slab cells."""
density = jnp.asarray(density)
z_fill = jnp.asarray(z_fill_fractions, dtype=density.dtype)
if density.ndim != 2:
raise ValueError("density must be two-dimensional")
if z_fill.ndim != 1 or z_fill.shape[0] < 1:
raise ValueError("z_fill_fractions must be a non-empty one-dimensional array")
return density[:, :, None] * z_fill[None, None, :]
[docs]
def erosion_dilation_penalty(
parameters: jax.Array,
*,
radius: float,
spacing: float,
beta: float = 10.0,
padding: str = "reflect",
) -> jax.Array:
"""Return the RMS morphological close/open discrepancy."""
def operation(values: jax.Array, eta: float) -> jax.Array:
return filter_and_project(
values,
radius=radius,
spacing=spacing,
beta=beta,
eta=eta,
padding=padding,
)
opened = operation(operation(parameters, 0.99), 0.01)
closed = operation(operation(parameters, 0.01), 0.99)
difference = closed - opened
return jnp.linalg.norm(difference) / jnp.sqrt(difference.size)
[docs]
def binary_density(
parameters: jax.Array,
*,
radius: float,
spacing: float,
threshold: float = 0.5,
padding: str = "reflect",
) -> jax.Array:
"""Create the deterministic fully binary fabrication design.
Thresholding occurs after the same minimum-feature filter used during
optimization. It is intentionally fixed at the projection threshold;
objective-aware threshold sweeps would grant a hidden per-device
optimization stage.
"""
filtered = conic_filter(parameters, radius=radius, spacing=spacing, padding=padding)
return (filtered >= threshold).astype(parameters.dtype)
def binary_density_1d(
parameters: jax.Array,
*,
radius: float,
spacing: float,
threshold: float = 0.5,
padding: str = "reflect",
) -> jax.Array:
"""Create a deterministic binary one-dimensional fabrication profile."""
filtered = conic_filter_1d(
parameters,
radius=radius,
spacing=spacing,
padding=padding,
)
return (filtered >= threshold).astype(parameters.dtype)
def binary_fraction(values: jax.Array, *, tolerance: float = 1e-7) -> jax.Array:
"""Return the fraction of entries exactly at either material endpoint."""
return jnp.mean((values <= tolerance) | (values >= 1.0 - tolerance))
[docs]
def d4_symmetrize(values: jax.Array) -> jax.Array:
"""Average a square material map over the eight D4 transformations.
This is the material symmetry used by four-port crossings: four rotations
and their mirror images. Applying it after filtering and projection
matches the topology transform used by Tidy3D's S-matrix example while
retaining one unconstrained parameter per design pixel.
"""
if values.ndim != 2 or values.shape[0] != values.shape[1]:
raise ValueError("D4 symmetry requires a square two-dimensional array")
rotations = tuple(jnp.rot90(values, index) for index in range(4))
return sum(rotations + tuple(value[:, ::-1] for value in rotations)) / 8.0
def density_from_permittivity(
permittivity: jax.Array | np.ndarray,
*,
low: float,
high: float,
shape: tuple[int, int] | None = None,
) -> jax.Array:
"""Initialize bounded topology parameters from an existing geometry.
The material is normalized between its endpoint permittivities, clipped
to the legal design interval, and optionally resampled to the optimizer
grid. This is the local equivalent of initializing parameters from a
base simulation rather than supplying unpublished optimized pixels.
"""
if high <= low:
raise ValueError("high permittivity must exceed low permittivity")
density = jnp.clip(
(jnp.asarray(permittivity, dtype=jnp.float32) - low) / (high - low),
0.0,
1.0,
)
if density.ndim != 2:
raise ValueError("permittivity initializer requires a two-dimensional map")
if shape is not None and tuple(density.shape) != tuple(shape):
density = jax.image.resize(density, shape, method="linear")
return density
def linear_continuation(
step: int | jax.Array,
num_steps: int,
*,
start: float,
end: float,
) -> jax.Array:
"""Inclusive linear continuation shared by Tidy3D-style optimizers."""
if num_steps < 1:
raise ValueError("num_steps must be positive")
fraction = jnp.asarray(step, dtype=jnp.float32) / max(num_steps - 1, 1)
return start * (1.0 - fraction) + end * fraction
def smooth_minimum(
values: jax.Array,
*,
temperature: float | jax.Array = 1.0,
normalize: bool = True,
) -> jax.Array:
"""Differentiable minimum for bounded multi-objective efficiencies.
``normalize=True`` removes the constant ``-tau*log(N)`` offset so equal
channel efficiencies retain their physical scale. It does not change the
gradient and is convenient for parity metrics.
"""
values = jnp.asarray(values)
temperature = jnp.asarray(temperature, dtype=values.dtype)
result = -temperature * jax.nn.logsumexp(-values / temperature)
if normalize:
result = result + temperature * jnp.log(values.size)
return result
def balanced_smooth_minimum(
values: jax.Array,
fraction: float | jax.Array,
*,
start_temperature: float = 1.0,
end_temperature: float = 0.05,
) -> jax.Array:
"""Smooth-min continuation that prevents a weak objective from dying.
Multi-port photonic optimizations often collapse one channel because a
high-temperature soft minimum behaves almost like an average on the
physical 0--1 efficiency scale. Exponentially lowering the temperature
keeps the broad, stable initial gradient while progressively emphasizing
the worst channel. Final validation should still use the problem's
published metric; this function is an optimizer aid, not a score change.
"""
if start_temperature <= 0 or end_temperature <= 0:
raise ValueError("smooth-min temperatures must be positive")
if end_temperature > start_temperature:
raise ValueError("end_temperature must not exceed start_temperature")
fraction = jnp.clip(jnp.asarray(fraction, dtype=jnp.float32), 0.0, 1.0)
temperature = start_temperature * (end_temperature / start_temperature) ** fraction
return smooth_minimum(values, temperature=temperature, normalize=True)
[docs]
def optimize_topology(
objective: Callable[[jax.Array, jax.Array, jax.Array], tuple[jax.Array, Any]],
initial_parameters: jax.Array | np.ndarray,
config: TopologyOptimizerConfig,
*,
accept: Callable[[Any], bool] | None = None,
callback: Callable[[TopologyStepRecord, np.ndarray, Any], None] | None = None,
) -> TopologyOptimizationResult:
"""Run a bounded Adam optimization with automatic projection continuation.
``objective(parameters, beta, fraction)`` returns ``(value, auxiliary)``.
The scalar ``fraction`` progresses inclusively from zero to one and lets a
device objective schedule additional weights without owning an optimizer
loop. ``accept`` may reject numerically nonphysical candidates (for
example, passive-power violations) from best-design selection.
"""
if config.iterations < 1:
raise ValueError("iterations must be positive")
if config.learning_rate <= 0:
raise ValueError("learning_rate must be positive")
if config.lower_bound >= config.upper_bound:
raise ValueError("lower_bound must be smaller than upper_bound")
if not 0.0 < config.final_learning_rate_fraction <= 1.0:
raise ValueError("final_learning_rate_fraction must lie in (0, 1]")
if config.early_stop_patience is not None and config.early_stop_patience < 1:
raise ValueError("early_stop_patience must be positive when provided")
if config.early_stop_min_delta < 0:
raise ValueError("early_stop_min_delta must be nonnegative")
if config.algorithm not in ("adam", "sgd", "lbfgs"):
raise ValueError("algorithm must be 'adam', 'sgd', or 'lbfgs'")
if not 0.0 <= config.momentum < 1.0:
raise ValueError("momentum must lie in [0, 1)")
if config.lbfgs_memory_size < 1:
raise ValueError("lbfgs_memory_size must be positive")
continuation_steps = config.iterations if config.continuation_steps is None else config.continuation_steps
if continuation_steps < 1 or continuation_steps > config.iterations:
raise ValueError("continuation_steps must lie between 1 and iterations")
parameters = jnp.asarray(initial_parameters, dtype=jnp.float32)
if config.algorithm == "adam":
optimizer = optax.adam(config.learning_rate)
elif config.algorithm == "sgd":
optimizer = optax.sgd(
config.learning_rate,
momentum=config.momentum,
nesterov=config.momentum > 0.0,
)
else:
# A fixed-step limited-memory BFGS transform avoids multiplying the
# cost of an FDTD iteration by a line search. Bounds are enforced after
# the quasi-Newton update, just as for Adam and SGD.
optimizer = optax.chain(
optax.scale_by_lbfgs(memory_size=config.lbfgs_memory_size),
optax.scale(-config.learning_rate),
)
optimizer_state = optimizer.init(parameters)
records: list[TopologyStepRecord] = []
best_parameters = np.asarray(parameters)
best_objective = -np.inf if config.maximize else np.inf
best_step = 0
progress_objective = best_objective
last_progress_step = 0
def optimization_step(parameters, optimizer_state, beta, fraction, learning_rate_scale):
(value, auxiliary), gradient = jax.value_and_grad(
objective,
argnums=0,
has_aux=True,
)(parameters, beta, fraction)
update_gradient = -gradient if config.maximize else gradient
updates, next_optimizer_state = optimizer.update(
update_gradient,
optimizer_state,
parameters,
)
updates = jax.tree.map(
lambda update: update * learning_rate_scale,
updates,
)
next_parameters = jnp.clip(
optax.apply_updates(parameters, updates),
config.lower_bound,
config.upper_bound,
)
return (
next_parameters,
next_optimizer_state,
value,
auxiliary,
jnp.linalg.norm(gradient),
)
compiled_optimization_step = jax.jit(optimization_step)
evaluate_objective = jax.jit(objective)
use_compiled_step = True
def consider(
step: int,
value: jax.Array,
beta: jax.Array,
gradient_norm: float | None,
auxiliary: Any,
) -> None:
nonlocal best_parameters, best_objective, best_step
value_host = float(value)
accepted = True if accept is None else bool(accept(auxiliary))
record = TopologyStepRecord(
step=step,
objective=value_host,
beta=float(beta),
gradient_norm=gradient_norm,
accepted=accepted,
)
records.append(record)
better = value_host > best_objective if config.maximize else value_host < best_objective
if accepted and better:
best_objective = value_host
best_parameters = np.asarray(parameters)
best_step = step
if callback is not None:
callback(record, np.asarray(parameters), auxiliary)
for step in range(config.iterations):
fraction = jnp.asarray(min(step / max(continuation_steps - 1, 1), 1.0), dtype=jnp.float32)
beta = linear_continuation(
min(step, continuation_steps - 1),
continuation_steps,
start=config.beta_start,
end=config.beta_end,
)
if step >= continuation_steps and config.iterations > continuation_steps:
tail_fraction = (step - continuation_steps + 1) / (config.iterations - continuation_steps)
learning_rate_scale = config.final_learning_rate_fraction**tail_fraction
else:
learning_rate_scale = 1.0
step_arguments = (
parameters,
optimizer_state,
beta,
fraction,
jnp.asarray(learning_rate_scale, dtype=parameters.dtype),
)
try:
step_output = (
compiled_optimization_step(*step_arguments) if use_compiled_step else optimization_step(*step_arguments)
)
except jax.errors.ConcretizationTypeError:
# Preserve compatibility with objectives that intentionally perform
# host-side scalar inspection or bookkeeping. Production numerical
# objectives stay on the compiled path.
use_compiled_step = False
step_output = optimization_step(*step_arguments)
next_parameters, optimizer_state, value, auxiliary, gradient_norm = step_output
accepted = True if accept is None else bool(accept(auxiliary))
value_host = float(value)
progress = (
value_host > progress_objective + config.early_stop_min_delta
if config.maximize
else value_host < progress_objective - config.early_stop_min_delta
)
if accepted and progress:
progress_objective = value_host
last_progress_step = step
consider(step, value, beta, float(gradient_norm), auxiliary)
parameters = next_parameters
if (
config.early_stop_patience is not None
and step + 1 >= continuation_steps
and step - last_progress_step >= config.early_stop_patience
):
break
final_fraction = jnp.asarray(1.0, dtype=jnp.float32)
final_beta = jnp.asarray(config.beta_end, dtype=jnp.float32)
if use_compiled_step:
final_value, final_auxiliary = evaluate_objective(parameters, final_beta, final_fraction)
else:
final_value, final_auxiliary = objective(parameters, final_beta, final_fraction)
consider(
config.iterations,
final_value,
final_beta,
None,
final_auxiliary,
)
return TopologyOptimizationResult(
final_parameters=np.asarray(parameters),
best_parameters=best_parameters,
best_objective=best_objective,
best_step=best_step,
records=tuple(records),
)
