perf(tidy3d): FXC-3721 Speed up test suite

Author: marcorudolphflex

pyproject.toml

@@ -298,9 +298,10 @@ banned-module-level-imports = ["scipy", "matplotlib"]
 
 [tool.pytest.ini_options]
 # TODO: remove --assert=plain when https://github.com/scipy/scipy/issues/22236 is resolved
-addopts = "--cov=tidy3d --doctest-modules -n auto --dist worksteal --assert=plain -m 'not numerical'"
+addopts = "--cov=tidy3d --doctest-modules -n auto --dist worksteal --assert=plain -m 'not numerical and not perf'"
 markers = [
   "numerical: marks numerical tests for adjoint gradients that require running simulations (deselect with '-m \"not numerical\"')",
+  "perf: marks tests which test the runtime of operations (deselect with '-m \"not perf\"')",
 ]
 env = ["MPLBACKEND=Agg", "OMP_NUM_THREADS=1", "TIDY3D_MICROWAVE__SUPPRESS_RF_LICENSE_WARNING=true"]
 doctest_optionflags = "NORMALIZE_WHITESPACE ELLIPSIS"

tests/test_components/autograd/test_autograd.py

@@ -658,6 +658,9 @@ def plot_sim(sim: td.Simulation, plot_eps: bool = True) -> None:
     args = [("polyslab", "mode")]
 
 
+ASYNC_TEST_ARGS = args[:2]
+
+
 def get_functions(structure_key: str, monitor_key: str) -> dict[str, typing.Callable]:
     if structure_key == ALL_KEY:
         structure_keys = structure_keys_
@@ -856,50 +859,55 @@ def objective(*args):
         assert anp.all(grad != 0.0), "some gradients are 0"
 
 
-@pytest.mark.slow
-def _run_autograd_async(use_emulated_run, fn_dicts, use_task_names=True):
-    """Shared core logic for the autograd async objective test (all combos in one batch)."""
+def _compare_async_vs_sync(fn_dicts) -> None:
+    """Compare async vs non-async autograd for a subset of structure/monitor pairs."""
 
-    def objective(*params):
-        """Compute combined objective across all structure/monitor combinations."""
+    # synchronous objective: run() one sim after another
+    def objective_sync(*params):
+        total = 0.0
+        for i, fn_dict in enumerate(fn_dicts):
+            sim = fn_dict["sim"](*params)
+            data = run(sim, task_name=f"autograd_sync_{i}", verbose=False)
+            total += fn_dict["postprocess"](data)
+        return total
+
+    def objective_async(*params):
         sims = {}
         for i, fn_dict in enumerate(fn_dicts):
-            if use_task_names:
-                # simulate different tasks per (structure, monitor)
-                sims[f"sim_{i}_fwd"] = fn_dict["sim"](*params)
-            else:
-                sims[f"sim_{i}"] = fn_dict["sim"](*params)
+            sim = fn_dict["sim"](*params)
+            key = f"autograd_{i}"
+            sims[key] = sim
 
-        # run all simulations asynchronously in one go
-        batch_data = run_async(sims, verbose=False)
+        batch_data = run_async(sims, verbose=False, local_gradient=False)
 
-        # accumulate total objective value
-        total_value = 0.0
+        total = 0.0
         for i, fn_dict in enumerate(fn_dicts):
-            key = f"sim_{i}_fwd" if use_task_names else f"sim_{i}"
-            total_value += fn_dict["postprocess"](batch_data[key])
+            key = f"autograd_{i}"
+            total += fn_dict["postprocess"](batch_data[key])
+        return total
 
-        return total_value
+    val_sync, grad_sync = ag.value_and_grad(objective_sync)(params0)
+    val_async, grad_async = ag.value_and_grad(objective_async)(params0)
 
-    val, grad = ag.value_and_grad(objective)(params0)
-    print(f"use_task_names={use_task_names}", val, grad)
-    assert anp.all(grad != 0.0), "some gradients are 0"
+    val_sync = float(val_sync)
+    val_async = float(val_async)
+    grad_sync = np.asarray(grad_sync)
+    grad_async = np.asarray(grad_async)
 
+    np.testing.assert_allclose(val_async, val_sync, rtol=1e-8, atol=1e-10)
+    np.testing.assert_allclose(grad_async, grad_sync, rtol=1e-6, atol=1e-8)
 
-@pytest.mark.slow
-def test_autograd_async_all(use_emulated_run):
-    """Run async autograd objective for all structure/monitor pairs in one combined batch."""
-    # collect all combinations once
-    fn_dicts = [get_functions(structure_key, monitor_key) for structure_key, monitor_key in args]
 
-    # call the unified helper
-    _run_autograd_async(use_emulated_run, fn_dicts, use_task_names=True)
+@pytest.mark.slow
+def test_autograd_async(use_emulated_run):
+    """Async autograd for a small subset; must match non-async autograd."""
 
+    # only use two structure/monitor combinations to keep this test cheap
+    fn_dicts = [
+        get_functions(structure_key, monitor_key) for structure_key, monitor_key in ASYNC_TEST_ARGS
+    ]
 
-def test_autograd_async_without_taskname(use_emulated_run):
-    """Test autograd objective varying use_task_names only."""
-    fn_dict = get_functions(*args[0])
-    _run_autograd_async(use_emulated_run, [fn_dict, fn_dict], use_task_names=False)
+    _compare_async_vs_sync(fn_dicts)
 
 
 class TestTupleGrads:
@@ -1021,6 +1029,7 @@ def objective(*args):
         grad = ag.grad(objective)(params0)
 
 
+@pytest.mark.perf
 def test_autograd_speed_num_structures(use_emulated_run):
     """Test an objective function through tidy3d autograd."""
 
@@ -1144,31 +1153,14 @@ def objective(*args):
 
 @pytest.mark.slow
 def test_autograd_async_server(use_emulated_run):
-    """Run async autograd objective across all structure/monitor combos in one batch."""
-
-    # gather all structure/monitor function pairs
-    fn_dicts = [get_functions(structure_key, monitor_key) for structure_key, monitor_key in args]
-
-    def objective(*params):
-        """Combined async objective for all cases."""
-        sims = {}
-        for i, fn_dict in enumerate(fn_dicts):
-            sim = fn_dict["sim"](*params)
-            # mimic two tasks per combination (as original test did)
-            sims[f"autograd_{i}"] = sim
-
-        # run everything asynchronously in one go
-        batch_data = run_async(sims, verbose=False, local_gradient=False)
+    """Same comparison, but with alternative task-keying (server-style)."""
 
-        # sum the results from all simulations
-        total_value = 0.0
-        for i, fn_dict in enumerate(fn_dicts):
-            total_value += fn_dict["postprocess"](batch_data[f"autograd_{i}"])
-        return total_value
+    fn_dicts = [
+        get_functions(structure_key, monitor_key) for structure_key, monitor_key in ASYNC_TEST_ARGS
+    ]
 
-    # compute value and gradient once for the entire batch
-    val, grad = ag.value_and_grad(objective)(params0)
-    assert np.all(np.abs(grad) > 0), "some gradients are 0"
+    # here we exercise the alternative key style in the async dict
+    _compare_async_vs_sync(fn_dicts)
 
 
 @pytest.mark.parametrize("structure_key", ("custom_med",))

tests/test_components/test_IO.py

@@ -176,6 +176,7 @@ def test_1a_simulation_load_export2(tmp_path):
     assert SIM2 == SIM3, "original and loaded simulations are not the same"
 
 
+@pytest.mark.perf
 def test_validation_speed(tmp_path):
     sizes_bytes = []
     times_sec = []