Merge pull request #15 from invrs-io/extractor

Add photon extractor challenge

src/invrs_gym/challenge/__init__.py

@@ -9,6 +9,7 @@
     "wdm_challenge",
     "metagrating",
     "diffractive_splitter",
+    "photon_extractor",
 ]
 
 from invrs_gym.challenge.ceviche.challenge import (
@@ -23,3 +24,4 @@
 )
 from invrs_gym.challenge.diffract.metagrating_challenge import metagrating
 from invrs_gym.challenge.diffract.splitter_challenge import diffractive_splitter
+from invrs_gym.challenge.extractor.challenge import photon_extractor

src/invrs_gym/challenge/diffract/common.py

@@ -167,7 +167,7 @@ def grating_efficiency(
         density_array: Defines the pattern of the grating layer.
         thickness: The thickness of the grating layer. This overrides the grating
             layer thickness given in `spec`.
-        spec: Defines the physical specifcation of the metagrating.
+        spec: Defines the physical specifcation of the grating.
         wavelength: The wavelength of the excitation.
         polarization: The polarization of the excitation, TE or TM.
         expansion: Defines the Fourier expansion for the calculation.

src/invrs_gym/challenge/diffract/splitter_challenge.py

@@ -42,8 +42,8 @@ def __init__(
         """Initializes the grating component.
 
         Args:
-            spec: Defines the physical specification of the grating.
-            sim_params: Defines simulation parameters for the grating.
+            spec: Defines the physical specification of the splitter.
+            sim_params: Defines simulation parameters for the splitter.
             thickness_initializer: Callable which returns the initial thickness for
                 the grating layer from a random key and a bounded array with value
                 equal the thickness from `spec`.

src/invrs_gym/challenge/extractor/challenge.py

@@ -0,0 +1,183 @@
+"""Defines the photon extractor challenge."""
+
+import dataclasses
+from typing import Any, Callable, Dict, Tuple
+
+import jax
+from fmmax import basis, fmm  # type: ignore[import]
+from jax import numpy as jnp
+from jax import tree_util
+from totypes import symmetry, types  # type: ignore[import,attr-defined,unused-ignore]
+
+from invrs_gym.challenge.extractor import component as extractor_component
+
+AuxDict = Dict[str, Any]
+DensityInitializer = Callable[[jax.Array, types.Density2DArray], types.Density2DArray]
+
+
+ENHANCEMENT_FLUX = "enhancement_flux"
+ENHANCEMENT_FLUX_MEAN = "enhancement_flux_mean"
+ENHANCEMENT_DOS = "enhancement_dos"
+ENHANCEMENT_DOS_MEAN = "enhancement_dos_mean"
+DISTANCE_TO_WINDOW = "distance_to_window"
+
+
+@dataclasses.dataclass
+class PhotonExtractorChallenge:
+    """Defines the photon extractor challenge.
+
+    The challenge is based on "Inverse-designed photon extractors for optically
+    addressable defect qubits" by Chakravarthi et al. It involves optimizing a GaP
+    patterned layer on diamond substrate above an implanted nitrogen vacancy defect.
+    An oxide hard mask used to pattern the GaP is left in place after the etch.
+
+    The goal of the optimization is to maximize extraction of 637 nm emission, i.e.
+    to maximize the power coupled from the defect to the ambient above the extractor.
+
+    https://opg.optica.org/optica/fulltext.cfm?uri=optica-7-12-1805
+
+    Attributes:
+        component: The component to be designed.
+        bare_substratee_emitted_power: The power emitted by a nitrogen vacancy defect
+            in a bare diamond substrate, i.e. without the GaP extractor structure.
+        bare_substrate_collected_power: The power collected from a nitrogen vacancy
+            defect in a bare diamond structure.
+        flux_enhancement_lower_bound: Scalar giving the minimum target for flux
+            enhancement. When the flux enhancement exceeds the lower bound, the
+            challenge is considered solved.
+    """
+
+    component: extractor_component.ExtractorComponent
+    bare_substrate_emitted_power: jnp.ndarray
+    bare_substrate_collected_power: jnp.ndarray
+    flux_enhancement_lower_bound: float
+
+    def loss(self, response: extractor_component.ExtractorResponse) -> jnp.ndarray:
+        """Compute a scalar loss from the component `response`."""
+        # The response should have a length-3 trailing axis, corresponding to x, y,
+        # and z-oriented dipoles.
+        assert response.collected_power.shape[-1] == 3
+        return -jnp.mean(response.collected_power)
+
+    def metrics(
+        self,
+        response: extractor_component.ExtractorResponse,
+        params: types.Density2DArray,
+        aux: AuxDict,
+    ) -> AuxDict:
+        """Compute challenge metrics.
+
+        Args:
+            response: The response of the extractor component.
+            params: The parameters where the response was evaluated.
+            aux: The auxilliary quantities returned by the component response method.
+
+        Returns:
+            The metrics dictionary, with the following quantities:
+                - mean enhancement of collected flux
+                - mean enhancement of dipole density of states
+                - the distance to the target flux enhancement
+        """
+        del params, aux
+        enhancement_flux = (
+            response.collected_power / self.bare_substrate_collected_power
+        )
+        enhancement_dos = response.emitted_power / self.bare_substrate_emitted_power
+        return {
+            ENHANCEMENT_FLUX: enhancement_flux,
+            ENHANCEMENT_FLUX_MEAN: jnp.mean(enhancement_flux),
+            ENHANCEMENT_DOS: enhancement_dos,
+            ENHANCEMENT_DOS_MEAN: jnp.mean(enhancement_dos),
+            DISTANCE_TO_WINDOW: jnp.maximum(
+                self.flux_enhancement_lower_bound - jnp.mean(enhancement_flux), 0.0
+            ),
+        }
+
+
+EXTRACTOR_SPEC = extractor_component.ExtractorSpec(
+    permittivity_ambient=(1.0 + 0.0j) ** 2,
+    permittivity_resist=(1.46 + 0.0j) ** 2,
+    permittivity_extractor=(3.31 + 0.0j) ** 2,
+    permittivity_substrate=(2.4102 + 0.0j) ** 2,
+    thickness_ambient=1.0,
+    thickness_resist=0.13,
+    thickness_extractor=0.25,
+    thickness_substrate_before_source=0.1,
+    thickness_substrate_after_source=0.9,
+    width_design_region=1.5,
+    width_padding=0.25,
+    width_pml=0.4,
+    fwhm_source=0.05,
+    offset_monitor_source=0.025,
+    offset_monitor_ambient=0.4,
+    width_monitor_ambient=1.5,
+)
+
+EXTRACTOR_SIM_PARAMS = extractor_component.ExtractorSimParams(
+    grid_spacing=0.01,
+    wavelength=0.637,
+    formulation=fmm.Formulation.JONES_DIRECT,
+    approximate_num_terms=1200,
+    truncation=basis.Truncation.CIRCULAR,
+)
+
+SYMMETRIES: Tuple[str, ...] = (
+    symmetry.REFLECTION_N_S,
+    symmetry.REFLECTION_E_W,
+    symmetry.REFLECTION_NE_SW,
+    symmetry.REFLECTION_NW_SE,
+)
+
+# Minimum width and spacing are 50 nm for the default dimensions.
+MINIMUM_WIDTH = 5
+MINIMUM_SPACING = 5
+
+# Reference power values used to calculate the enhancement. These were computed
+# by `compute_reference_resposne` with 1600 terms in the Fourier expansion.
+BARE_SUBSTRATE_COLLECTED_POWER = jnp.asarray([1.8094667, 1.8083396, 0.10765882])
+BARE_SUBSTRATE_EMITTED_POWER = jnp.asarray([58.385864, 58.383484, 67.01958])
+
+# Target is to achieve flux enhancement of 15 times or greater.
+FLUX_ENHANCEMENT_LOWER_BOUND = 15.0
+
+
+def photon_extractor(
+    minimum_width: int = MINIMUM_WIDTH,
+    minimum_spacing: int = MINIMUM_SPACING,
+    density_initializer: DensityInitializer = extractor_component.identity_initializer,
+    bare_substrate_emitted_power: jnp.ndarray = BARE_SUBSTRATE_EMITTED_POWER,
+    bare_substrate_collected_power: jnp.ndarray = BARE_SUBSTRATE_COLLECTED_POWER,
+    flux_enhancement_lower_bound: float = FLUX_ENHANCEMENT_LOWER_BOUND,
+    spec: extractor_component.ExtractorSpec = EXTRACTOR_SPEC,
+    sim_params: extractor_component.ExtractorSimParams = EXTRACTOR_SIM_PARAMS,
+    symmetries: Tuple[str, ...] = SYMMETRIES,
+) -> PhotonExtractorChallenge:
+    """Photon extractor with 1.5 x 1.5 um design region."""
+    return PhotonExtractorChallenge(
+        component=extractor_component.ExtractorComponent(
+            spec=spec,
+            sim_params=sim_params,
+            density_initializer=density_initializer,
+            minimum_width=minimum_width,
+            minimum_spacing=minimum_spacing,
+            symmetries=symmetries,
+        ),
+        bare_substrate_emitted_power=bare_substrate_emitted_power,
+        bare_substrate_collected_power=bare_substrate_collected_power,
+        flux_enhancement_lower_bound=flux_enhancement_lower_bound,
+    )
+
+
+def bare_substrate_response(
+    spec: extractor_component.ExtractorSpec = EXTRACTOR_SPEC,
+    sim_params: extractor_component.ExtractorSimParams = EXTRACTOR_SIM_PARAMS,
+) -> extractor_component.ExtractorResponse:
+    """Computes the response of the nitrogen vacancy in a bare diamond substrate."""
+    component = extractor_component.ExtractorComponent(
+        spec=spec,
+        sim_params=sim_params,
+        density_initializer=lambda _, d: tree_util.tree_map(jnp.zeros_like, d),
+    )
+    params = component.init(jax.random.PRNGKey(0))
+    response, _ = component.response(params)
+    return response