fix interpolation

Author: yaugenst-flex

tidy3d/plugins/autograd/invdes/parametrizations.py

@@ -7,11 +7,11 @@
 import pydantic.v1 as pd
 from autograd import value_and_grad
 from numpy.typing import NDArray
-from scipy.ndimage import zoom
 from scipy.optimize import minimize
 
 import tidy3d as td
 from tidy3d.components.base import Tidy3dBaseModel
+from tidy3d.components.grid.grid import Coords
 from tidy3d.plugins.autograd.constants import BETA_DEFAULT, ETA_DEFAULT
 from tidy3d.plugins.autograd.types import KernelType, PaddingType
 
@@ -109,6 +109,7 @@ def initialize_params_from_simulation(
     param_to_structure: Callable[[np.ndarray], td.Structure],
     params0: np.ndarray,
     *,
+    freq: Optional[float] = None,
     maxiter: int = 100,
     bounds: tuple[Optional[float], Optional[float]] = (0.0, 1.0),
     rel_improve_tol: float = 1e-3,
@@ -127,6 +128,9 @@ def initialize_params_from_simulation(
     - The callable ``param_to_structure`` controls the parameterization and must return
       a :class:`.Structure` (typically with a :class:`.CustomMedium`). Its permittivity
       dataset (and coordinates) defines the grid used for comparison.
+    - The base simulation permittivity is sampled once on an extended subgrid covering
+      the design geometry and then interpolated onto the design region coordinates
+      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.
 
@@ -188,12 +192,19 @@ def initialize_params_from_simulation(
     >>> params.shape
     (3, 3)
     """
-
     structure_init = param_to_structure(params0)
-    eps_base = sim.epsilon(box=structure_init.geometry).data
-    eps_param = structure_init.medium.permittivity.data
-    factors = np.array(eps_param.shape) / np.array(eps_base.shape)
-    eps_base_interp = zoom(eps_base, factors, order=1)
+    subgrid = sim.discretize(structure_init.geometry, extend=True)
+    eps_base_da = sim.epsilon_on_grid(grid=subgrid, coord_key="centers", freq=freq)
+
+    design_eps_da = structure_init.medium.permittivity
+    design_coords = Coords(
+        x=np.array(design_eps_da.coords["x"]),
+        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))
     denom = 1.0 if denom == 0 else denom
@@ -214,9 +225,8 @@ def loss_fn(params_vec: np.ndarray) -> float:
         "best_x": params0.ravel().copy(),
     }
 
-    def callback(intermediate_result):
-        xk = intermediate_result.x
-        val = intermediate_result.fun
+    def callback(xk: np.ndarray):
+        val = loss_fn(xk)
         if val < state["best_val"]:
             state["best_val"] = val
             state["best_x"] = xk.copy()