Skip to content

test: isolate numerical autograd artifacts #2981

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 1 commit into from
Nov 12, 2025
Merged

test: isolate numerical autograd artifacts #2981

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
40 changes: 40 additions & 0 deletions tests/test_components/autograd/numerical/conftest.py
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -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
43 changes: 32 additions & 11 deletions tests/test_components/autograd/numerical/test_autograd_box_polyslab_numerical.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -3,6 +3,7 @@
from __future__ import annotations

import sys
from pathlib import Path

import autograd.numpy as anp
import numpy as np
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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)
Expand All @@ -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,
)
Expand All @@ -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,
Expand All @@ -244,7 +259,7 @@ def objective(parameters):
tag,
base_sim,
fom,
tmp_path,
artifact_dir,
local_gradient=local_gradient,
)

Expand Down Expand Up @@ -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."""

Expand All @@ -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(
Expand All @@ -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,
)

Expand All @@ -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(
Expand All @@ -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,
)

Expand Down Expand Up @@ -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)
Expand Down
48 changes: 22 additions & 26 deletions tests/test_components/autograd/numerical/test_autograd_conductivity_numerical.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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.

Expand All @@ -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:
Expand Down Expand Up @@ -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,
Expand Down Expand Up @@ -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:
Expand All @@ -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)
38 changes: 18 additions & 20 deletions tests/test_components/autograd/numerical/test_autograd_mode_polyslab_numerical.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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."""

Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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:
Expand All @@ -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)
Loading
Loading