Merge pull request #112 from invrs-io/metric

Add metrics to meta-atom library challenge

docs/notebooks/library_challenge.ipynb

@@ -230,7 +230,7 @@
    "cell_type": "markdown",
    "metadata": {},
    "source": [
-    "The agreement between the phases reported by Chen et al. and those found here are in good agreement, particularly for the first six nanostructures and at the central wavelength of 550 nm. The last two nanostructures show some difference in phase at the 450 nm and 650 nm wavelengths."
+    "The phases reported by Chen et al. and those found here are in good agreement, particularly for the first six nanostructures and at the central wavelength of 550 nm. The last two nanostructures show some difference in phase at the 450 nm and 650 nm wavelengths."
    ]
   },
   {

src/invrs_gym/challenges/library/challenge.py

@@ -16,9 +16,12 @@
 from invrs_gym.challenges.library import component as library_component
 from invrs_gym.utils import initializers, materials
 
-METAGRATING_EFFICIENCY = "metagrating_efficiency"
-METAGRATING_RELATIVE_EFFICIENCY = "metagrating_relative_efficiency"
-DISTANCE_TO_TARGET = "distance_to_target"
+METAGRATING_EFFICIENCY_RHCP = "metagrating_efficiency_rhcp"
+METAGRATING_EFFICIENCY_LHCP = "metagrating_efficiency_lhcp"
+METAGRATING_RELATIVE_EFFICIENCY_RHCP = "metagrating_relative_efficiency_rhcp"
+METAGRATING_RELATIVE_EFFICIENCY_LHCP = "metagrating_relative_efficiency_lhcp"
+METAGRATING_AVERAGE_EFFICIENCY = "metagrating_average_efficiency"
+METAGRATING_AVERAGE_RELATIVE_EFFICIENCY = "metagrating_average_relative_efficiency"
 
 
 @dataclasses.dataclass
@@ -42,30 +45,50 @@ def metrics(
         params: library_component.Params,
         aux: base.AuxDict,
     ) -> base.AuxDict:
-        """Compute challenge metrics."""
+        """Compute challenge metrics.
+
+        Several challenge metrics relate to the performance of a metagrating assembled
+        from the meta-atom library.
+
+        Args:
+            response: The response of the meta-atom library.
+            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:
+                - Per-wavelength, per-polarization metagrating efficiency for RHCP
+                  excitation.
+                - Per-wavelength, per-polarization metagrating efficiency for LHCP
+                  excitation.
+                - Per-wavelength, per-polarization metagrating relative efficiency for
+                  RHCP excitation.
+                - Per-wavelength, per-polarization metagrating relative efficiency for
+                  LHCP excitation.
+                - Average metagrating efficiency.
+                - Average metagrating relative efficiency.
+        """
         metrics = super().metrics(response, params, aux)
         (
-            metagrating_efficiency,
-            metagrating_relative_efficiency,
+            (efficiency_rhcp, efficiency_lhcp),
+            (relative_efficiency_rhcp, relative_efficiency_lhcp),
         ) = _metagrating_efficiency(response, self.component.spec)
         metrics.update(
             {
-                METAGRATING_EFFICIENCY: metagrating_efficiency,
-                METAGRATING_RELATIVE_EFFICIENCY: metagrating_relative_efficiency,
-                DISTANCE_TO_TARGET: self.distance_to_target(response),
+                METAGRATING_EFFICIENCY_RHCP: efficiency_rhcp,
+                METAGRATING_EFFICIENCY_LHCP: efficiency_lhcp,
+                METAGRATING_RELATIVE_EFFICIENCY_RHCP: relative_efficiency_rhcp,
+                METAGRATING_RELATIVE_EFFICIENCY_LHCP: relative_efficiency_lhcp,
+                METAGRATING_AVERAGE_EFFICIENCY: jnp.mean(
+                    jnp.asarray([efficiency_rhcp, efficiency_lhcp])
+                ),
+                METAGRATING_AVERAGE_RELATIVE_EFFICIENCY: jnp.mean(
+                    jnp.asarray([relative_efficiency_rhcp, relative_efficiency_lhcp])
+                ),
             }
         )
         return metrics
 
-    def distance_to_target(
-        self,
-        response: library_component.LibraryResponse,
-    ) -> jnp.ndarray:
-        """Compute distance from the component `response` to the challenge target."""
-        del response
-        # TODO: implement a distance, or strip distance from the challenge object.
-        return jnp.ones(())
-
 
 def _metagrating_efficiency(
     response: library_component.LibraryResponse,
@@ -157,7 +180,7 @@ def library_density_initializer(
 
 LIBRARY_SIM_PARAMS = library_component.LibrarySimParams(
     wavelength=jnp.asarray([0.45, 0.55, 0.65]),
-    approximate_num_terms=300,
+    approximate_num_terms=200,
     formulation=fmm.Formulation.JONES_DIRECT_FOURIER,
     truncation=basis.Truncation.CIRCULAR,
 )

tests/challenges/library/test_reference_nanostructures.py

@@ -105,24 +105,24 @@ def get_nanostructures(specs):
 
 ATOL_CONSERVED = (
     (0.1, 0.1, 0.1),
-    (0.1, 0.2, 0.4),
-    (0.1, 0.2, 0.3),
-    (0.1, 0.2, 0.3),
-    (0.2, 0.2, 0.5),
-    (0.1, 0.2, 0.4),
+    (0.1, 0.1, 0.1),
+    (0.1, 0.2, 0.2),
+    (0.1, 0.1, 0.3),
+    (0.3, 0.1, 0.4),
+    (0.2, 0.2, 0.4),
     # Final two nanostructures have larger error for some unknown reason.
-    (0.6, 0.2, 0.9),
-    (0.7, 0.2, 1.1),
+    (0.6, 0.1, 0.9),
+    (0.7, 0.1, 1.1),
 )
 ATOL_CONVERTED = (
     (0.1, 0.1, 0.1),
     (0.1, 0.1, 0.1),
     (0.1, 0.2, 0.2),
     (0.2, 0.2, 0.1),
-    (0.1, 0.2, 0.2),
-    (0.2, 0.4, 0.4),
-    (0.8, 0.2, 1.3),
-    (1.2, 0.3, 0.6),
+    (0.1, 0.1, 0.2),
+    (0.3, 0.3, 0.4),
+    (0.9, 0.3, 1.2),
+    (0.9, 0.3, 0.7),
 )
 
 
@@ -147,7 +147,7 @@ def test_nanostructure_phase_matches_expected(self):
             primitive_lattice_vectors=basis.LatticeVectors(
                 u=basis.X * spec.pitch, v=basis.Y * spec.pitch
             ),
-            approximate_num_terms=400,
+            approximate_num_terms=200,
             truncation=basis.Truncation.CIRCULAR,
         )
         response, _ = component.simulate_library(