Merge pull request #144 from invrs-io/fields

Add compute fields functionality to diffract challenges

.bumpversion.toml

@@ -1,5 +1,5 @@
 [tool.bumpversion]
-current_version = "v1.3.1"
+current_version = "v1.4.0"
 commit = true
 commit_args = "--no-verify"
 tag = true

README.md

@@ -1,5 +1,5 @@
 # invrs-gym
-`v1.3.1`
+`v1.4.0`
 
 ## Overview
 The `invrs_gym` package is an open-source gym containing a diverse set of photonic design challenges, which are relevant for a wide range of applications such as AR/VR, optical networking, LIDAR, and others.

pyproject.toml

@@ -1,7 +1,7 @@
 [project]
 
 name = "invrs_gym"
-version = "v1.3.1"
+version = "v1.4.0"
 description = "A collection of inverse design challenges"
 keywords = ["topology", "optimization", "jax", "inverse design"]
 readme = "README.md"

src/invrs_gym/__init__.py

@@ -3,7 +3,7 @@
 Copyright (c) 2023 The INVRS-IO authors.
 """
 
-__version__ = "v1.3.1"
+__version__ = "v1.4.0"
 __author__ = "Martin F. Schubert <[email protected]>"
 
 from invrs_gym import challenges as challenges

src/invrs_gym/challenges/diffract/common.py

@@ -26,6 +26,10 @@
 THICKNESS_SPACER = "thickness_spacer"
 DENSITY = "density"
 
+EFIELD = "efield"
+HFIELD = "hfield"
+FIELD_COORDINATES = "field_coordinates"
+
 
 @dataclasses.dataclass
 class GratingSpec:
@@ -55,9 +59,11 @@ class GratingSpec:
     permittivity_spacer: complex
     permittivity_substrate: complex
 
+    thickness_ambient: float
     thickness_cap: float | jnp.ndarray | types.BoundedArray
     thickness_grating: float | jnp.ndarray | types.BoundedArray
     thickness_spacer: float | jnp.ndarray | types.BoundedArray
+    thickness_substrate: float
 
     period_x: float
     period_y: float
@@ -193,6 +199,7 @@ def response(
         *,
         wavelength: Optional[Union[float, jnp.ndarray]] = None,
         expansion: Optional[basis.Expansion] = None,
+        compute_fields: bool = False,
     ) -> Tuple[GratingResponse, base.AuxDict]:
         """Computes the response of the grating.
 
@@ -204,6 +211,8 @@ def response(
                 by the `init` method.
             wavelength: Optional wavelength to override the default in `sim_params`.
             expansion: Optional expansion to override the default `expansion`.
+            compute_fields: If `True`, computes and xz cross section for electric
+                and magnetic fields, which makes the calculation more expensive.
 
         Returns:
             The `(response, aux)` tuple.
@@ -212,14 +221,15 @@ def response(
             expansion = self.expansion
         if wavelength is None:
             wavelength = self.sim_params.wavelength
-        transmission_efficiency, reflection_efficiency = grating_efficiency(
+        (transmission_efficiency, reflection_efficiency), aux = grating_efficiency(
             density=params,
             spec=self.spec,
             wavelength=jnp.asarray(wavelength),
             polar_angle=jnp.asarray(self.sim_params.polar_angle),
             azimuthal_angle=jnp.asarray(self.sim_params.azimuthal_angle),
             expansion=expansion,
             formulation=self.sim_params.formulation,
+            compute_fields=compute_fields,
         )
         response = GratingResponse(
             wavelength=jnp.asarray(wavelength),
@@ -229,7 +239,7 @@ def response(
             reflection_efficiency=reflection_efficiency,
             expansion=expansion,
         )
-        return response, {}
+        return response, aux
 
 
 class GratingWithOptimizableThicknessComponent(base.Component):
@@ -302,6 +312,7 @@ def response(
         *,
         wavelength: Optional[Union[float, jnp.ndarray]] = None,
         expansion: Optional[basis.Expansion] = None,
+        compute_fields: bool = False,
     ) -> Tuple[GratingResponse, base.AuxDict]:
         """Computes the response of the grating component.
 
@@ -313,6 +324,8 @@ def response(
                 the `init` method.
             wavelength: Optional wavelength to override the default in `sim_params`.
             expansion: Optional expansion to override the default `expansion`.
+            compute_fields: If `True`, computes and xz cross section for electric
+                and magnetic fields, which makes the calculation more expensive.
 
         Returns:
             The `(response, aux)` tuple.
@@ -327,14 +340,15 @@ def response(
             thickness_grating=jnp.asarray(params[THICKNESS_GRATING].array),
             thickness_spacer=jnp.asarray(params[THICKNESS_SPACER].array),
         )
-        transmission_efficiency, reflection_efficiency = grating_efficiency(
+        (transmission_efficiency, reflection_efficiency), aux = grating_efficiency(
             density=params[DENSITY],  # type: ignore[arg-type]
             spec=spec,
             wavelength=jnp.asarray(wavelength),
             polar_angle=jnp.asarray(self.sim_params.polar_angle),
             azimuthal_angle=jnp.asarray(self.sim_params.azimuthal_angle),
             expansion=expansion,
             formulation=self.sim_params.formulation,
+            compute_fields=compute_fields,
         )
         response = GratingResponse(
             wavelength=jnp.asarray(wavelength),
@@ -344,7 +358,7 @@ def response(
             reflection_efficiency=reflection_efficiency,
             expansion=expansion,
         )
-        return response, {}
+        return response, aux
 
 
 def seed_density(grid_shape: Tuple[int, int], **kwargs: Any) -> types.Density2DArray:
@@ -401,7 +415,8 @@ def grating_efficiency(
     azimuthal_angle: jnp.ndarray,
     expansion: basis.Expansion,
     formulation: fmm.Formulation,
-) -> Tuple[jnp.ndarray, jnp.ndarray]:
+    compute_fields: bool,
+) -> Tuple[Tuple[jnp.ndarray, jnp.ndarray], base.AuxDict]:
     """Compute the per-order transmission and reflection efficiency for a grating.
 
     The excitation for the calculation are separate a TE- and TM-polarized plane waves
@@ -416,6 +431,7 @@ def grating_efficiency(
         azimuthal_angle: The azimuthal angle of the excitation.
         expansion: Defines the Fourier expansion for the calculation.
         formulation: Defines the FMM formulation to be used.
+        compute_fields: If `True`, computes fields in an xz cross section.
 
     Returns:
         The per-order transmission and reflection efficiency, having shape
@@ -463,14 +479,25 @@ def grating_efficiency(
     # Layer thicknesses for the ambient and substrate are set to zero; these do not
     # affect the result of the calculation.
     layer_thicknesses = (
-        jnp.zeros(()),
+        jnp.asarray(spec.thickness_ambient),
         jnp.asarray(spec.thickness_cap),
         jnp.asarray(spec.thickness_grating),
         jnp.asarray(spec.thickness_spacer),
-        jnp.zeros(()),
+        jnp.asarray(spec.thickness_substrate),
     )
 
-    s_matrix = scattering.stack_s_matrix(layer_solve_results, layer_thicknesses)
+    if compute_fields:
+        # If fields wanted, compute the full set of interior scattering matrices.
+        s_matrices_interior = scattering.stack_s_matrices_interior(
+            layer_solve_results=layer_solve_results,
+            layer_thicknesses=layer_thicknesses,
+        )
+        s_matrix = s_matrices_interior[-1][0]
+    else:
+        s_matrix = scattering.stack_s_matrix(
+            layer_solve_results=layer_solve_results,
+            layer_thicknesses=layer_thicknesses,
+        )
 
     n = expansion.num_terms
     assert tuple(expansion.basis_coefficients[0, :]) == (0, 0)
@@ -511,4 +538,36 @@ def grating_efficiency(
     transmission_efficiency = bwd_flux_ambient / total_incident_flux
     reflection_efficiency = fwd_flux_substrate / total_incident_flux
 
-    return transmission_efficiency, reflection_efficiency
+    # -------------------------------------------------------------------------
+    # Compute fields in an xz cross section.
+    # -------------------------------------------------------------------------
+
+    aux = {}
+    if compute_fields:
+        amplitudes_interior = fields.stack_amplitudes_interior(
+            s_matrices_interior=s_matrices_interior,
+            forward_amplitude_0_start=jnp.zeros_like(bwd_amplitude_substrate_end),
+            backward_amplitude_N_end=bwd_amplitude_substrate_end,
+        )
+        x = jnp.linspace(0, spec.period_x, density.shape[0])
+        y = jnp.ones_like(x) * spec.period_y / 2
+        layer_znum = tuple(
+            [int(jnp.round(t / spec.grid_spacing) + 1) for t in layer_thicknesses]
+        )
+        (ex, ey, ez), (hx, hy, hz), (x, y, z) = fields.stack_fields_3d_on_coordinates(
+            amplitudes_interior=amplitudes_interior,
+            layer_solve_results=layer_solve_results,
+            layer_thicknesses=layer_thicknesses,
+            layer_znum=layer_znum,
+            x=x,
+            y=y,
+        )
+        aux.update(
+            {
+                EFIELD: (ex, ey, ez),
+                HFIELD: (hx, hy, hz),
+                FIELD_COORDINATES: (x, y, z),
+            }
+        )
+
+    return (transmission_efficiency, reflection_efficiency), aux

src/invrs_gym/challenges/diffract/metagrating_challenge.py

@@ -144,9 +144,11 @@ def _value_for_order(
     permittivity_encapsulation=(1.0 + 0.00001j) ** 2,
     permittivity_spacer=(1.45 + 0.0j) ** 2,
     permittivity_substrate=(1.45 + 0.0j) ** 2,
+    thickness_ambient=0.0,
     thickness_cap=0.0,
     thickness_grating=0.325,
     thickness_spacer=0.0,
+    thickness_substrate=0.0,
     period_x=float(1.050 / jnp.sin(jnp.deg2rad(50.0))),
     period_y=0.525,
     grid_spacing=0.0117,  # Yields a grid shape of `(118, 45)`.

src/invrs_gym/challenges/diffract/splitter_challenge.py

@@ -262,9 +262,11 @@ def extract_orders_for_splitting(
     permittivity_encapsulation=(1.0 + 0.00001j) ** 2,
     permittivity_spacer=(1.0 + 0.0j) ** 2,
     permittivity_substrate=(1.0 + 0.0j) ** 2,
+    thickness_ambient=0.0,
     thickness_cap=types.BoundedArray(array=0.0, lower_bound=0.0, upper_bound=0.1),
     thickness_grating=types.BoundedArray(array=0.692, lower_bound=0.5, upper_bound=1.5),
     thickness_spacer=types.BoundedArray(array=0.0, lower_bound=0.0, upper_bound=0.1),
+    thickness_substrate=0.0,
     period_x=7.2,
     period_y=7.2,
     grid_spacing=0.04,  # Yields a grid shape of `(180, 180)`.

tests/challenges/diffract/test_common.py

@@ -21,9 +21,11 @@
     permittivity_encapsulation=(1.0 + 0.00001j) ** 2,
     permittivity_spacer=(1.45 + 0.0j) ** 2,
     permittivity_substrate=(1.45 + 0.0j) ** 2,
+    thickness_ambient=0.0,
     thickness_cap=0.0,
     thickness_grating=0.325,
     thickness_spacer=0.0,
+    thickness_substrate=0.0,
     period_x=float(1.050 / jnp.sin(jnp.deg2rad(50.0))),
     period_y=0.525,
     grid_spacing=0.0117,
@@ -36,9 +38,11 @@
     permittivity_encapsulation=(1.0 + 0.00001j) ** 2,
     permittivity_spacer=(1.45 + 0.0j) ** 2,
     permittivity_substrate=(1.45 + 0.0j) ** 2,
+    thickness_ambient=0.0,
     thickness_cap=types.BoundedArray(array=0.0, lower_bound=0.0, upper_bound=0.1),
     thickness_grating=types.BoundedArray(array=0.6, lower_bound=0.5, upper_bound=1.5),
     thickness_spacer=types.BoundedArray(array=0.0, lower_bound=0.0, upper_bound=0.1),
+    thickness_substrate=0.0,
     period_x=float(1.050 / jnp.sin(jnp.deg2rad(50.0))),
     period_y=0.525,
     grid_spacing=0.0117,
@@ -111,6 +115,19 @@ def test_multiple_wavelengths(self):
             (2, mc.expansion.num_terms, 2),
         )
 
+    def test_compute_fields(self):
+        mc = common.SimpleGratingComponent(
+            spec=SIMPLE_GRATING_SPEC,
+            sim_params=LIGHTWEIGHT_SIM_PARAMS,
+            density_initializer=lambda _, seed_density: seed_density,
+        )
+        params = mc.init(jax.random.PRNGKey(0))
+        response, aux = mc.response(params, compute_fields=True)
+        self.assertSequenceEqual(
+            set(aux.keys()),
+            {common.EFIELD, common.HFIELD, common.FIELD_COORDINATES},
+        )
+
 
 class GratingWithOptimizableThicknessComponentTest(unittest.TestCase):
     def test_can_jit_response(self):
@@ -141,3 +158,17 @@ def test_multiple_wavelengths(self):
             response.transmission_efficiency.shape,
             (2, mc.expansion.num_terms, 2),
         )
+
+    def test_compute_fields(self):
+        mc = common.GratingWithOptimizableThicknessComponent(
+            spec=GRATING_WITH_THICKNESS_SPEC,
+            sim_params=LIGHTWEIGHT_SIM_PARAMS,
+            thickness_initializer=lambda _, thickness: thickness,
+            density_initializer=lambda _, seed_density: seed_density,
+        )
+        params = mc.init(jax.random.PRNGKey(0))
+        response, aux = mc.response(params, compute_fields=True)
+        self.assertSequenceEqual(
+            set(aux.keys()),
+            {common.EFIELD, common.HFIELD, common.FIELD_COORDINATES},
+        )