warnings and more edge case handling

Author: yaugenst-flex

tidy3d/plugins/autograd/invdes/parametrizations.py

@@ -1,7 +1,7 @@
 from __future__ import annotations
 
 from collections import deque
-from typing import Callable, Optional, Union
+from typing import Callable, Literal, Optional, Union
 
 import autograd.numpy as np
 import pydantic.v1 as pd
@@ -110,6 +110,7 @@ def initialize_params_from_simulation(
     params0: np.ndarray,
     *,
     freq: Optional[float] = None,
+    outside_handling: Literal["extrapolate", "mask", "nan"] = "mask",
     maxiter: int = 100,
     bounds: tuple[Optional[float], Optional[float]] = (0.0, 1.0),
     rel_improve_tol: float = 1e-3,
@@ -133,6 +134,10 @@ def initialize_params_from_simulation(
       using coordinate-aware interpolation (no per-iteration interpolation).
     - Early stopping uses a single knob ``rel_improve_tol``; optimization stops once the
       relative improvement over a small fixed window falls below this value.
+    - Points outside the base‑epsilon coverage can be handled via ``outside_handling``:
+      ``'extrapolate'`` (use nearest extrapolation, default in earlier versions),
+      ``'mask'`` (ignore outside points using a coverage mask; default), or
+      ``'nan'`` (non‑extended grid; treat outside as NaN and ignore in the loss).
 
     Parameters
     ----------
@@ -148,6 +153,12 @@ def initialize_params_from_simulation(
         Maximum number of L‑BFGS‑B iterations.
     freq : float, optional
         Frequency at which permittivity is evaluated. If ``None``, uses infinite frequency.
+    outside_handling : {"extrapolate", "mask", "nan"} = "mask"
+        Strategy for points where design coordinates fall outside the sampled base epsilon:
+        - "extrapolate": include them using nearest extrapolation.
+        - "mask": include only points within the coverage bounds of ``sim.epsilon`` on the
+          extended subgrid.
+        - "nan": sample on a non‑extended subgrid and ignore points where interpolation returns NaN.
     bounds : tuple[float | None, float | None] = (0.0, 1.0)
         Element‑wise parameter bounds, e.g. ``(0.0, 1.0)`` or ``(-1.0, 1.0)``. Use ``None`` to
         indicate an unbounded side.
@@ -193,7 +204,13 @@ def initialize_params_from_simulation(
     (3, 3)
     """
     structure_init = param_to_structure(params0)
-    subgrid = sim.discretize(structure_init.geometry, extend=True)
+
+    if outside_handling not in ("extrapolate", "mask", "nan"):
+        raise ValueError("'outside_handling' must be one of {'extrapolate', 'mask', 'nan'}.")
+
+    # build base epsilon and interpolate onto design coords
+    extend_flag = outside_handling != "nan"
+    subgrid = sim.discretize(structure_init.geometry, extend=extend_flag)
     eps_base_da = sim.epsilon_on_grid(grid=subgrid, coord_key="centers", freq=freq)
 
     design_eps_da = structure_init.medium.permittivity
@@ -202,18 +219,53 @@ def initialize_params_from_simulation(
         y=np.array(design_eps_da.coords["y"]),
         z=np.array(design_eps_da.coords["z"]),
     )
-    eps_base_interp = design_coords.spatial_interp(
-        array=eps_base_da, interp_method="linear", fill_value="extrapolate"
-    ).data
 
-    denom = np.sqrt(np.sum(eps_base_interp.real**2 + eps_base_interp.imag**2))
+    if outside_handling == "nan":
+        eps_base_interp = design_coords.spatial_interp(
+            array=eps_base_da, interp_method="linear", fill_value=np.nan
+        ).data
+        mask = np.isfinite(eps_base_interp)
+    else:
+        eps_base_interp = design_coords.spatial_interp(
+            array=eps_base_da, interp_method="linear", fill_value="extrapolate"
+        ).data
+        if outside_handling == "mask":
+            # build mask from coverage bounds of base epsilon coordinates
+            xb, yb, zb = [np.array(eps_base_da.coords[d]) for d in ("x", "y", "z")]
+            xd, yd, zd = [np.array(design_eps_da.coords[d]) for d in ("x", "y", "z")]
+            mask_x = (xd >= np.min(xb)) & (xd <= np.max(xb))
+            mask_y = (yd >= np.min(yb)) & (yd <= np.max(yb))
+            mask_z = (zd >= np.min(zb)) & (zd <= np.max(zb))
+            mask = (mask_x[:, None, None]) & (mask_y[None, :, None]) & (mask_z[None, None, :])
+        else:
+            mask = None
+
+    if mask is not None:
+        covered = int(np.sum(mask))
+        total = int(np.prod(mask.shape))
+        frac = covered / max(total, 1)
+        if frac < 0.9:
+            td.log.warning(
+                f"Only {frac:.1%} of design points are covered by base epsilon sampling. "
+                "Consider adding a 'MeshOverrideStructure' or adjusting design coordinates."
+            )
+
+    if mask is None:
+        denom = np.sqrt(np.sum(eps_base_interp.real**2 + eps_base_interp.imag**2))
+    else:
+        denom = np.sqrt(
+            np.sum((eps_base_interp.real[mask]) ** 2 + (eps_base_interp.imag[mask]) ** 2)
+        )
     denom = 1.0 if denom == 0 else denom
 
     def loss_fn(params_vec: np.ndarray) -> float:
         params = params_vec.reshape(params0.shape)
         structure = param_to_structure(params)
         eps_design = structure.medium.permittivity.data
-        res = eps_base_interp - eps_design
+        if mask is None:
+            res = eps_base_interp - eps_design
+            return 0.5 * np.sum(res.real**2 + res.imag**2) / denom
+        res = (eps_base_interp - eps_design)[mask]
         return 0.5 * np.sum(res.real**2 + res.imag**2) / denom
 
     val_and_grad = value_and_grad(loss_fn)