test: isolate numerical autograd artifacts

Author: yaugenst-flex

tests/test_components/autograd/numerical/conftest.py

@@ -0,0 +1,40 @@
+from __future__ import annotations
+
+import os
+import re
+from pathlib import Path
+
+import pytest
+
+ARTIFACT_ENV_VAR = "TIDY3D_NUMERICAL_ARTIFACT_DIR"
+DEFAULT_RELATIVE_DIR = Path("tests/tmp/autograd_numerical")
+
+
+def _sanitize_segment(value: str) -> str:
+    sanitized = re.sub(r"[^\w.-]+", "_", value)
+    sanitized = sanitized.strip("_")
+    return sanitized or "case"
+
+
+def _resolve_artifact_root() -> Path:
+    env_value = os.environ.get(ARTIFACT_ENV_VAR)
+    if env_value:
+        root = Path(env_value).expanduser()
+    else:
+        repo_root = Path(__file__).resolve().parents[4]
+        root = repo_root / DEFAULT_RELATIVE_DIR
+    root.mkdir(parents=True, exist_ok=True)
+    return root
+
+
+@pytest.fixture(scope="session")
+def numerical_artifact_root() -> Path:
+    return _resolve_artifact_root()
+
+
+@pytest.fixture
+def numerical_case_dir(request, numerical_artifact_root: Path) -> Path:
+    safe_nodeid = _sanitize_segment(request.node.nodeid.replace(os.sep, "_"))
+    case_dir = numerical_artifact_root / safe_nodeid
+    case_dir.mkdir(parents=True, exist_ok=True)
+    return case_dir

tests/test_components/autograd/numerical/test_autograd_box_polyslab_numerical.py

@@ -3,6 +3,7 @@
 from __future__ import annotations
 
 import sys
+from pathlib import Path
 
 import autograd.numpy as anp
 import numpy as np
@@ -40,6 +41,12 @@
     sys.stdout = sys.stderr
 
 
+def case_identifier(is_3d: bool, infinite_dim_2d: int | None, shift_box_center: bool) -> str:
+    geometry_tag = "3d" if is_3d else f"2d_infinite_dim_{infinite_dim_2d}"
+    shift_tag = "shifted" if shift_box_center else "centered"
+    return f"box_polyslab_{geometry_tag}_{shift_tag}"
+
+
 def dimension_permutation(infinite_dim: int) -> tuple[int, int]:
     offset_dim = (1 + infinite_dim) % 3
     final_dim = (1 + offset_dim) % 3
@@ -191,11 +198,13 @@ def run_parameter_simulations(
     tag: str,
     base_sim: td.Simulation,
     fom,
-    tmp_path,
+    artifact_dir: Path,
     *,
     local_gradient: bool,
 ):
     simulation_dict = {}
+    output_dir = artifact_dir
+    output_dir.mkdir(parents=True, exist_ok=True)
 
     for idx, param_values in enumerate(parameter_sets):
         geometry = make_geometry(param_values, box_center)
@@ -209,16 +218,19 @@ def run_parameter_simulations(
 
     if len(simulation_dict) == 1:
         key, sim = next(iter(simulation_dict.items()))
+        result_path = output_dir / f"{key}.hdf5"
         sim_data = web.run(
             sim,
             task_name=key,
+            path=str(result_path),
             local_gradient=local_gradient,
             verbose=VERBOSE,
         )
         return fom(sim_data)
 
     sim_data_map = web.run_async(
         simulation_dict,
+        path_dir=str(output_dir),
         local_gradient=local_gradient,
         verbose=VERBOSE,
     )
@@ -232,10 +244,13 @@ def make_objective(
     tag: str,
     base_sim: td.Simulation,
     fom,
-    tmp_path,
+    case_dir: Path,
     *,
     local_gradient: bool,
 ):
+    artifact_dir = case_dir / tag
+    artifact_dir.mkdir(parents=True, exist_ok=True)
+
     def objective(parameters):
         results = run_parameter_simulations(
             parameters,
@@ -244,7 +259,7 @@ def objective(parameters):
             tag,
             base_sim,
             fom,
-            tmp_path,
+            artifact_dir,
             local_gradient=local_gradient,
         )
 
@@ -304,7 +319,9 @@ def squeeze_dimension(array: np.ndarray, is_3d: bool, infinite_dim: int | None)
     ],
 )
 @pytest.mark.parametrize("shift_box_center", (True, False))
-def test_box_and_polyslab_gradients_match(is_3d, infinite_dim_2d, shift_box_center, tmp_path):
+def test_box_and_polyslab_gradients_match(
+    is_3d, infinite_dim_2d, shift_box_center, numerical_case_dir
+):
     """Test that the box and polyslab gradients match for rectangular slab geometries. Allow
     comparison as well to finite difference values."""
 
@@ -330,13 +347,17 @@ def test_box_and_polyslab_gradients_match(is_3d, infinite_dim_2d, shift_box_cent
             box_center[infinite_dim_2d] = 0.5 * INFINITE_DIM_SIZE_UM
             box_center[final_dim_2d] = 0.5 * PERIODS_UM[0]
 
+    case_id = case_identifier(is_3d, None if is_3d else infinite_dim_2d, shift_box_center)
+    case_dir = numerical_case_dir / case_id
+    case_dir.mkdir(parents=True, exist_ok=True)
+
     box_objective = make_objective(
         make_box_geometry,
         box_center,
         "box",
         base_sim,
         fom,
-        tmp_path,
+        case_dir,
         local_gradient=LOCAL_GRADIENT,
     )
     polyslab_objective = make_objective(
@@ -345,7 +366,7 @@ def test_box_and_polyslab_gradients_match(is_3d, infinite_dim_2d, shift_box_cent
         "polyslab",
         base_sim,
         fom,
-        tmp_path,
+        case_dir,
         local_gradient=LOCAL_GRADIENT,
     )
 
@@ -355,7 +376,7 @@ def test_box_and_polyslab_gradients_match(is_3d, infinite_dim_2d, shift_box_cent
         "box_fd",
         base_sim,
         fom,
-        tmp_path,
+        case_dir,
         local_gradient=False,
     )
     polyslab_objective_fd = make_objective(
@@ -364,7 +385,7 @@ def test_box_and_polyslab_gradients_match(is_3d, infinite_dim_2d, shift_box_cent
         "polyslab_fd",
         base_sim,
         fom,
-        tmp_path,
+        case_dir,
         local_gradient=False,
     )
 
@@ -396,10 +417,10 @@ def test_box_and_polyslab_gradients_match(is_3d, infinite_dim_2d, shift_box_cent
     }
 
     if SAVE_OUTPUT_DATA:
-        np.savez(
-            f"test_diff_init_{'3' if is_3d else '2'}d_infinite_dim_{infinite_dim_2d}.npz",
-            **test_data,
+        npz_path = case_dir / (
+            f"test_diff_init_{'3' if is_3d else '2'}d_infinite_dim_{infinite_dim_2d}.npz"
         )
+        np.savez(npz_path, **test_data)
 
     def angled_overlap_deg(v1, v2):
         norm_v1 = np.linalg.norm(v1)

tests/test_components/autograd/numerical/test_autograd_conductivity_numerical.py

@@ -37,7 +37,7 @@
 SAVE_ADJ_LOC = 1
 LOCAL_GRADIENT = True
 VERBOSE = False
-NUMERICAL_RESULTS_DATA_DIR = "./numerical_conductivity_test/"
+NUMERICAL_RESULTS_SUBDIR = "numerical_conductivity_test"
 SHOW_PRINT_STATEMENTS = False
 
 RMS_THRESHOLD = 0.6
@@ -325,17 +325,9 @@ def flux(sim_data):
 
 
 @pytest.mark.numerical
-@pytest.mark.parametrize(
-    "conductivity_data_test_parameters, dir_name",
-    zip(
-        conductivity_data_test_parameters,
-        ([NUMERICAL_RESULTS_DATA_DIR] if SAVE_FD_ADJ_DATA else [None])
-        * len(conductivity_data_test_parameters),
-    ),
-    indirect=["dir_name"],
-)
+@pytest.mark.parametrize("conductivity_data_test_parameters", conductivity_data_test_parameters)
 def test_finite_difference_conductivity_data(
-    conductivity_data_test_parameters, rng, tmp_path, create_directory
+    conductivity_data_test_parameters, rng, numerical_case_dir
 ):
     """Test autograd conductivity gradients by comparing to numerical finite difference.
 
@@ -356,17 +348,14 @@ def test_finite_difference_conductivity_data(
         Test parameters including wavelengths, monitor configuration, etc.
     rng : numpy.random.Generator
         Random number generator for creating perturbation patterns
-    tmp_path : pathlib.Path
-        Temporary directory for simulation files
-    create_directory : fixture
-        Pytest fixture for creating directories
+    numerical_case_dir : pathlib.Path
+        Case-specific artifact directory for simulation and result files
     """
 
     # Create directory for plots if plotting is enabled
+    results_dir = numerical_case_dir / NUMERICAL_RESULTS_SUBDIR
     if PLOT_FD_ADJ_COMPARISON or SAVE_FD_ADJ_DATA:
-        import os
-
-        os.makedirs(NUMERICAL_RESULTS_DATA_DIR, exist_ok=True)
+        results_dir.mkdir(parents=True, exist_ok=True)
 
     num_tests = 0
     for monitor_size_wvl in monitor_sizes_3d_wvl:
@@ -414,8 +403,8 @@ def test_finite_difference_conductivity_data(
 
     eval_fns, eval_fn_names = make_eval_fns(monitor_size_wvl)
 
-    sim_path_dir = tmp_path / f"test{test_number}"
-    sim_path_dir.mkdir()
+    sim_path_dir = numerical_case_dir / "simulations" / f"test{test_number}"
+    sim_path_dir.mkdir(parents=True, exist_ok=True)
 
     objective = create_objective_function(
         block,
@@ -494,11 +483,21 @@ def test_finite_difference_conductivity_data(
     print("-" * 20)
     print("\n" * 3)
 
-    assert rms_error < RMS_THRESHOLD * fd_mag, "RMS error magnitude too large"
-
     test_results[SAVE_FD_LOC, :] = fd_grad
     test_results[SAVE_ADJ_LOC, :] = pattern_dot_adj_gradient
 
+    save_idx = test_number + 1
+    save_path = None
+    if SAVE_FD_ADJ_DATA:
+        results_dir.mkdir(parents=True, exist_ok=True)
+        save_path = results_dir / f"results_{save_idx}.npy"
+
+    try:
+        assert rms_error < RMS_THRESHOLD * fd_mag, "RMS error magnitude too large"
+    finally:
+        if save_path is not None:
+            np.save(save_path, test_results)
+
     test_number += 1
 
     if PLOT_FD_ADJ_COMPARISON:
@@ -512,12 +511,9 @@ def test_finite_difference_conductivity_data(
         plt.grid(True, alpha=0.3)
 
         # Save the plot
-        plot_filename = f"{NUMERICAL_RESULTS_DATA_DIR}/gradient_comparison_test_{test_number}_{eval_fn_name}.png"
+        plot_filename = results_dir / (f"gradient_comparison_test_{test_number}_{eval_fn_name}.png")
         plt.savefig(plot_filename, dpi=150, bbox_inches="tight")
         print(f"Plot saved to: {plot_filename}")
 
         plt.show()
         plt.close()
-
-    if SAVE_FD_ADJ_DATA:
-        np.save(f"{NUMERICAL_RESULTS_DATA_DIR}/results_{test_number}.npy", test_results)

tests/test_components/autograd/numerical/test_autograd_mode_polyslab_numerical.py

@@ -20,7 +20,7 @@
 SAVE_ADJ_LOC = 1
 LOCAL_GRADIENT = True
 VERBOSE = False
-NUMERICAL_RESULTS_DATA_DIR = "./numerical_mode_polyslab_test/"
+NUMERICAL_RESULTS_SUBDIR = "numerical_mode_polyslab_test"
 SHOW_PRINT_STATEMENTS = False
 
 NUM_MODE_MONITOR_FREQUENCIES = 4
@@ -301,18 +301,8 @@ def objective(vertices):
 
 
 @pytest.mark.numerical
-@pytest.mark.parametrize(
-    "mode_data_test_parameters, dir_name",
-    zip(
-        mode_data_test_parameters,
-        ([NUMERICAL_RESULTS_DATA_DIR] if SAVE_FD_ADJ_DATA else [None])
-        * len(mode_data_test_parameters),
-    ),
-    indirect=["dir_name"],
-)
-def test_finite_difference_mode_data_polyslab(
-    mode_data_test_parameters, rng, tmp_path, create_directory
-):
+@pytest.mark.parametrize("mode_data_test_parameters", mode_data_test_parameters)
+def test_finite_difference_mode_data_polyslab(mode_data_test_parameters, rng, numerical_case_dir):
     """Test a variety of autograd permittivity gradients for ModeData in combination with polyslab by"""
     """comparing them to numerical finite difference."""
 
@@ -351,8 +341,8 @@ def test_finite_difference_mode_data_polyslab(
         size=(np.inf, np.inf, MODE_LAYER_HEIGHT_WVL * mesh_wvl_um + mesh_wvl_um),
     )
 
-    sim_path_dir = tmp_path / f"test{test_number}"
-    sim_path_dir.mkdir()
+    sim_path_dir = numerical_case_dir / "simulations" / f"test{test_number}"
+    sim_path_dir.mkdir(parents=True, exist_ok=True)
 
     # Weights for creating a random objective function over multiple frequencies by
     # summing their contributions by random weights. This helps verify gradient errors
@@ -471,11 +461,22 @@ def eval_fn(sim_data):
     print("-" * 20)
     print("\n" * 3)
 
-    assert rms_error < RMS_THRESHOLD * fd_mag, "RMS error magnitude too large"
-
     test_results[SAVE_FD_LOC, :] = fd_grad
     test_results[SAVE_ADJ_LOC, :] = pattern_dot_adj_gradient
 
+    save_idx = test_number + 1
+    save_path = None
+    if SAVE_FD_ADJ_DATA:
+        results_dir = numerical_case_dir / NUMERICAL_RESULTS_SUBDIR
+        results_dir.mkdir(parents=True, exist_ok=True)
+        save_path = results_dir / f"results_{save_idx}.npy"
+
+    try:
+        assert rms_error < RMS_THRESHOLD * fd_mag, "RMS error magnitude too large"
+    finally:
+        if save_path is not None:
+            np.save(save_path, test_results)
+
     test_number += 1
 
     if PLOT_FD_ADJ_COMPARISON:
@@ -487,6 +488,3 @@ def eval_fn(sim_data):
         plt.ylabel("Gradient value")
         plt.legend()
         plt.show()
-
-    if SAVE_FD_ADJ_DATA:
-        np.save(f"{NUMERICAL_RESULTS_DATA_DIR}/results_{test_number}.npy", test_results)