Add local subpixel integration

Author: caseyflex

tests/conftest.py

@@ -13,6 +13,7 @@
 from autograd.wrap_util import unary_to_nary
 
 import tidy3d as td
+from tidy3d.config import config
 from tidy3d.log import DEFAULT_LEVEL, set_logging_console, set_logging_level
 
 
@@ -94,6 +95,15 @@ def mpl_config_interactive():
     mpl.use(original_backend)
 
 
+@pytest.fixture(autouse=True)
+def disable_local_subpixel():
+    """Disable local subpixel for the unit tests."""
+    use_local_subpixel = config.use_local_subpixel
+    config.use_local_subpixel = False
+    yield
+    config.use_local_subpixel = use_local_subpixel
+
+
 @pytest.fixture
 def dir_name(request):
     return request.param

tests/test_components/test_packaging.py

@@ -2,7 +2,13 @@
 
 import pytest
 
-from tidy3d.packaging import Tidy3dImportError, check_import, verify_packages_import
+from tidy3d.packaging import (
+    Tidy3dImportError,
+    check_import,
+    supports_local_subpixel,
+    tidy3d_extras,
+    verify_packages_import,
+)
 
 assert check_import("tidy3d") is True
 
@@ -65,5 +71,19 @@ def test_check_import():
     assert mock_check_import("module2") is False
 
 
+def test_tidy3d_extras():
+    import importlib
+
+    has_tidy3d_extras = importlib.util.find_spec("tidy3d_extras") is not None
+    print(f"has_tidy3d_extras = {has_tidy3d_extras}")
+
+    @supports_local_subpixel
+    def get_eps():
+        assert tidy3d_extras["use_local_subpixel"] is False
+        assert tidy3d_extras["mod"] is None
+
+    get_eps()
+
+
 if __name__ == "__main__":
     pytest.main()

tidy3d/components/mode/mode_solver.py

@@ -62,6 +62,7 @@
 from tidy3d.constants import C_0
 from tidy3d.exceptions import SetupError, ValidationError
 from tidy3d.log import log
+from tidy3d.packaging import supports_local_subpixel, tidy3d_extras
 
 # Importing the local solver may not work if e.g. scipy is not installed
 IMPORT_ERROR_MSG = """Could not import local solver, 'ModeSolver' objects can still be constructed
@@ -316,7 +317,7 @@ def _get_solver_grid(
             _, plane_inds = Box.pop_axis([0, 1, 2], normal_axis)
             for dim, sym in enumerate(solver_symmetry):
                 if sym != 0:
-                    span_inds[plane_inds[dim], 0] += np.diff(span_inds[plane_inds[dim]]) // 2
+                    span_inds[plane_inds[dim], 0] += np.diff(span_inds[plane_inds[dim]])[0] // 2
 
         return simulation._subgrid(span_inds=span_inds)
 
@@ -1355,18 +1356,21 @@ def _diagonal_material_profile_modal_plane_tranform(
 
     def _solver_eps(self, freq: float) -> ArrayComplex4D:
         """Diagonal permittivity in the shape needed by solver, with normal axis rotated to z."""
-
         # Get diagonal epsilon components in the plane
         eps_tensor = self._get_epsilon(freq)
         # tranformation
         return self._tensorial_material_profile_modal_plane_tranform(eps_tensor, self.normal_axis)
 
+    @supports_local_subpixel
     def _solve_all_freqs(
         self,
         coords: tuple[ArrayFloat1D, ArrayFloat1D],
         symmetry: tuple[Symmetry, Symmetry],
     ) -> tuple[list[float], list[dict[str, ArrayComplex4D]], list[EpsSpecType]]:
         """Call the mode solver at all requested frequencies."""
+        if tidy3d_extras["use_local_subpixel"]:
+            subpixel_ms = tidy3d_extras["mod"].SubpixelModeSolver.from_mode_solver(self)
+            return subpixel_ms._solve_all_freqs(coords=coords, symmetry=symmetry)
 
         fields = []
         n_complex = []
@@ -1380,13 +1384,19 @@ def _solve_all_freqs(
             eps_spec.append(eps_spec_freq)
         return n_complex, fields, eps_spec
 
+    @supports_local_subpixel
     def _solve_all_freqs_relative(
         self,
         coords: tuple[ArrayFloat1D, ArrayFloat1D],
         symmetry: tuple[Symmetry, Symmetry],
         basis_fields: list[dict[str, ArrayComplex4D]],
     ) -> tuple[list[float], list[dict[str, ArrayComplex4D]], list[EpsSpecType]]:
         """Call the mode solver at all requested frequencies."""
+        if tidy3d_extras["use_local_subpixel"]:
+            subpixel_ms = tidy3d_extras["mod"].SubpixelModeSolver.from_mode_solver(self)
+            return subpixel_ms._solve_all_freqs_relative(
+                coords=coords, symmetry=symmetry, basis_fields=basis_fields
+            )
 
         fields = []
         n_complex = []
@@ -1452,22 +1462,26 @@ def _solve_single_freq(
         )
         return n_complex, fields, eps_spec
 
-    def _rotate_field_coords_inverse(self, field: FIELD) -> FIELD:
+    @classmethod
+    def _rotate_field_coords_inverse(
+        cls, field: FIELD, normal_axis: Axis, plane: MODE_PLANE_TYPE
+    ) -> FIELD:
         """Move the propagation axis to the z axis in the array."""
-        f_x, f_y, f_z = np.moveaxis(field, source=1 + self.normal_axis, destination=3)
-        f_n, f_ts = self.plane.pop_axis((f_x, f_y, f_z), axis=self.normal_axis)
-        return np.stack(self.plane.unpop_axis(f_n, f_ts, axis=2), axis=0)
+        f_x, f_y, f_z = np.moveaxis(field, source=1 + normal_axis, destination=3)
+        f_n, f_ts = plane.pop_axis((f_x, f_y, f_z), axis=normal_axis)
+        return np.stack(plane.unpop_axis(f_n, f_ts, axis=2), axis=0)
 
-    def _postprocess_solver_fields_inverse(self, fields):
+    @classmethod
+    def _postprocess_solver_fields_inverse(cls, fields, normal_axis: Axis, plane: MODE_PLANE_TYPE):
         """Convert ``fields`` to ``solver_fields``. Doesn't change gauge."""
         E = [fields[key] for key in ("Ex", "Ey", "Ez")]
         H = [fields[key] for key in ("Hx", "Hy", "Hz")]
 
-        (Ex, Ey, Ez) = self._rotate_field_coords_inverse(E)
-        (Hx, Hy, Hz) = self._rotate_field_coords_inverse(H)
+        (Ex, Ey, Ez) = cls._rotate_field_coords_inverse(E, normal_axis=normal_axis, plane=plane)
+        (Hx, Hy, Hz) = cls._rotate_field_coords_inverse(H, normal_axis=normal_axis, plane=plane)
 
         # apply -1 to H fields if a reflection was involved in the rotation
-        if self.normal_axis == 1:
+        if normal_axis == 1:
             Hx *= -1
             Hy *= -1
             Hz *= -1
@@ -1489,7 +1503,9 @@ def _solve_single_freq_relative(
         if not LOCAL_SOLVER_IMPORTED:
             raise ImportError(IMPORT_ERROR_MSG)
 
-        solver_basis_fields = self._postprocess_solver_fields_inverse(basis_fields)
+        solver_basis_fields = self._postprocess_solver_fields_inverse(
+            fields=basis_fields, normal_axis=self.normal_axis, plane=self.plane
+        )
 
         solver_fields, n_complex, eps_spec = compute_modes(
             eps_cross=self._solver_eps(freq),

tidy3d/components/mode/simulation.py

@@ -250,6 +250,9 @@ def run_local(self):
         """Run locally."""
         from .data.sim_data import ModeSimulationData
 
+        # repeat the calculation every time, in case use_local_subpixel changed
+        self._invalidate_solver_cache()
+
         modes_raw = self._mode_solver.data_raw
         return ModeSimulationData(simulation=self, modes_raw=modes_raw)
 

tidy3d/components/simulation.py

@@ -22,6 +22,7 @@
 from tidy3d.constants import C_0, SECOND, fp_eps, inf
 from tidy3d.exceptions import SetupError, Tidy3dError, Tidy3dImportError, ValidationError
 from tidy3d.log import log
+from tidy3d.packaging import supports_local_subpixel, tidy3d_extras
 from tidy3d.updater import Updater
 
 from .base import cached_property, skip_if_fields_missing
@@ -1461,6 +1462,7 @@ def epsilon(
         sub_grid = self.discretize(box)
         return self.epsilon_on_grid(grid=sub_grid, coord_key=coord_key, freq=freq)
 
+    @supports_local_subpixel
     def epsilon_on_grid(
         self,
         grid: Grid,
@@ -1505,6 +1507,10 @@ def epsilon_on_grid(
                 "Epsilon calculation may be slow."
             )
 
+        if tidy3d_extras["use_local_subpixel"]:
+            subpixel_sim = tidy3d_extras["mod"].SubpixelSimulation.from_simulation(self)
+            return subpixel_sim.epsilon_on_grid(grid=grid, coord_key=coord_key, freq=freq)
+
         def get_eps(structure: Structure, frequency: float, coords: Coords):
             """Select the correct epsilon component if field locations are requested."""
             if coord_key[0] != "E":
@@ -1955,6 +1961,10 @@ def subsection(
         # 2) Assemble the full simulation without validation
         return new_sim.updated_copy(structures=new_structures, deep=deep_copy, validate=False)
 
+    def _invalidate_solver_cache(self) -> None:
+        """Clear cached attributes that become stale when subpixel changes."""
+        self._cached_properties.pop("_mode_solver", None)
+
 
 class Simulation(AbstractYeeGridSimulation):
     """

tidy3d/config.py

@@ -2,6 +2,8 @@
 
 from __future__ import annotations
 
+from typing import Optional
+
 import pydantic.v1 as pd
 
 from .log import DEFAULT_LEVEL, LogLevel, set_log_suppression, set_logging_level
@@ -35,6 +37,13 @@ class Config:
         "for several elements.",
     )
 
+    use_local_subpixel: Optional[bool] = pd.Field(
+        None,
+        title="Whether to use local subpixel averaging. If 'None', local subpixel "
+        "averaging will be used if 'tidy3d-extras' is installed and not used otherwise. "
+        "NOTE: This feature is not yet supported.",
+    )
+
     @pd.validator("logging_level", pre=True, always=True)
     def _set_logging_level(cls, val):
         """Set the logging level if logging_level is changed."""

tidy3d/packaging.py

@@ -12,6 +12,7 @@
 
 import numpy as np
 
+from .config import config
 from .exceptions import Tidy3dImportError
 
 vtk = {
@@ -22,6 +23,8 @@
     "numpy_to_vtk": None,
 }
 
+tidy3d_extras = {"mod": None, "use_local_subpixel": False}
+
 
 def check_import(module_name: str) -> bool:
     """
@@ -175,3 +178,37 @@ def get_numpy_major_version(module=np):
     major_version = int(module_version.split(".")[0])
 
     return major_version
+
+
+def supports_local_subpixel(fn):
+    """When decorating a method, checks that 'tidy3d-extras' is available,
+    conditioned on 'config.use_local_subpixel'."""
+
+    @functools.wraps(fn)
+    def _fn(*args, **kwargs):
+        if config.use_local_subpixel is False:
+            tidy3d_extras["use_local_subpixel"] = False
+            tidy3d_extras["mod"] = None
+        else:
+            # first try to import the module
+            if tidy3d_extras["mod"] is None:
+                try:
+                    import tidy3d_extras as tidy3d_extras_mod
+
+                    tidy3d_extras["mod"] = tidy3d_extras_mod
+                    tidy3d_extras["use_local_subpixel"] = True
+                except ImportError as exc:
+                    tidy3d_extras["mod"] = None
+                    tidy3d_extras["use_local_subpixel"] = False
+                    if config.use_local_subpixel is True:
+                        raise Tidy3dImportError(
+                            "The package 'tidy3d-extras' is required for this "
+                            "operation when 'config.use_local_subpixel' is 'True'. "
+                            "Please install the 'tidy3d-extras' package using, for "
+                            "example, 'pip install tidy3d-extras'. NOTE: This "
+                            "feature is not yet supported."
+                        ) from exc
+
+        return fn(*args, **kwargs)
+
+    return _fn