adjoint: add exception for NaN data in adjoint pipeline

CHANGELOG.md

@@ -27,6 +27,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 - Added `SimulationMap` and `SimulationDataMap` immutable dictionary-like containers for managing collections of simulations and results.
 - Added `TerminalComponentModelerData`, `ComponentModelerData`, `MicrowaveSMatrixData`, and introduced multiple DataArrays for modeler workflow data structures.
 - Added autograd support for dispersive material models: `Sellmeier`, `Drude`, `Lorentz`, `Debye` and their custom medium variants.
+- Added check and exception for NaN data in the adjoint pipeline to raise issue to user before adjoint source creation failure.
 
 ### Changed
 - Validate mode solver object for large number of grid points on the modal plane.

tests/test_components/autograd/test_autograd.py

@@ -2813,3 +2813,24 @@ def objective(params):
     val, grad = ag.value_and_grad(objective)(params0)
 
     assert anp.all(grad != 0.0), "some gradients are 0 for conductivity-only test"
+
+
+@pytest.mark.parametrize("structure_key, monitor_key", args)
+def test_vjp_nan(use_emulated_run, structure_key, monitor_key):
+    """Test vjp data that has nan in it is flagged as an error."""
+
+    fn_dict = get_functions(structure_key, monitor_key)
+    make_sim = fn_dict["sim"]
+    postprocess = fn_dict["postprocess"]
+
+    def objective(*args):
+        """Objective function."""
+        sim = make_sim(*args)
+        if PLOT_SIM:
+            plot_sim(sim, plot_eps=True)
+        data = run(sim, task_name="autograd_test", verbose=False)
+        value = (postprocess(data) + float("nan")) ** 2
+        return value
+
+    with pytest.raises(AdjointError, match="aN values detected for data field"):
+        grad = ag.grad(objective)(params0)

tidy3d/web/api/autograd/autograd.py

@@ -912,6 +912,14 @@ def setup_adj(
         k: get_static(v) for k, v in data_fields_vjp.items() if not np.allclose(v, 0)
     }
 
+    for k, v in data_fields_vjp.items():
+        if np.any(np.isnan(v)):
+            raise AdjointError(
+                f"NaN values detected for data field {k} in the adjoint pipeline. This may be "
+                f"due to NaN values in the simulation data or the computed value of your "
+                f"objective function."
+            )
+
     # if all entries are zero, there is no adjoint sim to run
     if not data_fields_vjp:
         return []