Add docstrings to enhance clarity across benchmarks and examples

Author: QuentinWach

benchmarks/compare_mode_solver_fixtures.py

@@ -1,3 +1,5 @@
+"""CLI and helpers for comparing committed mode-solver fixtures against MicroMode."""
+
 from __future__ import annotations
 
 import argparse
@@ -28,6 +30,7 @@
 
 
 def parse_args() -> argparse.Namespace:
+    """Parse fixture-comparison CLI options."""
     parser = argparse.ArgumentParser(description="Inspect committed mode-solver reference fixtures.")
     parser.add_argument(
         "--suite",
@@ -72,6 +75,7 @@ def parse_args() -> argparse.Namespace:
 
 
 def main() -> None:
+    """Inspect fixtures and optionally run local comparisons."""
     args = parse_args()
     fixture_root = args.fixture_root or (DEFAULT_FIXTURE_ROOT / args.suite)
     manifest = read_json(manifest_path(fixture_root))
@@ -125,6 +129,7 @@ def main() -> None:
 
 
 def _compare_local_case(root: Path, entry: dict) -> dict:
+    """Run one reconstructable fixture recipe through MicroMode and compare outputs."""
     case_id = entry["case_id"]
     try:
         import micromode as sm
@@ -321,6 +326,7 @@ def _compare_local_case(root: Path, entry: dict) -> dict:
 def _solver_edges_from_field_coords(
     edges: tuple[np.ndarray, np.ndarray], recipe: dict
 ) -> tuple[np.ndarray, np.ndarray]:
+    """Derive solver edge coordinates from fixture field coordinates."""
     dmin_pmc = tuple(bool(value) for value in recipe.get("dmin_pmc", (False, False)))
     trim_edges = tuple(recipe.get("trim_edges", ((0, 0), (0, 0))))
     out = []
@@ -354,6 +360,7 @@ def _solve_recipe(
     normal_dim: str,
     normal_coord: float,
 ):
+    """Solve all frequencies described by a local fixture recipe."""
     freqs = tuple(float(freq) for freq in ref_n.coords["f"].values)
     if recipe.get("solve_each_frequency"):
         rows = []
@@ -424,6 +431,7 @@ def _solve_recipe_for_freq(
     normal_dim: str,
     normal_coord: float,
 ):
+    """Solve one fixture recipe for one frequency or frequency tuple."""
     freqs = (float(freq),) if np.isscalar(freq) else tuple(float(value) for value in freq)
     centers = tuple((axis_edges[:-1] + axis_edges[1:]) / 2 for axis_edges in edges)
     eps_xx, eps_yy, eps_zz = _eps_components_from_recipe(recipe, edges, centers, tangent_dims, freqs[0], sm)
@@ -460,6 +468,7 @@ def _eps_components_from_recipe(
     freq: float,
     sm,
 ) -> tuple[np.ndarray, np.ndarray, np.ndarray]:
+    """Rasterize diagonal epsilon components at Yee sample locations."""
     if recipe.get("yee_staggered", True):
         return (
             _eps_from_recipe(recipe, (coords[0], edges[1][:-1]), tangent_dims, freq, sm),
@@ -477,6 +486,7 @@ def _eps_from_recipe(
     freq: float,
     sm,
 ) -> np.ndarray:
+    """Rasterize primitive boxes and circles onto one component grid."""
     grids = np.meshgrid(*coords, indexing="ij")
     eps = np.full(tuple(len(coord) for coord in coords), recipe.get("clad_eps", 1.0), dtype=np.complex128)
     for box in recipe.get("boxes", ()):
@@ -502,24 +512,28 @@ def _eps_from_recipe(
 
 
 def _reorder_modes(values: np.ndarray, recipe: dict) -> np.ndarray:
+    """Apply fixture-specific mode ordering before comparison."""
     if recipe.get("sort_order") != "ascending":
         return values
     order = np.argsort(values.real, axis=1)
     return np.take_along_axis(values, order, axis=1)
 
 
 def _reorder_field_modes(values: np.ndarray, recipe: dict) -> np.ndarray:
+    """Apply fixture-specific field ordering for diagnostic overlap checks."""
     if recipe.get("sort_order") != "ascending":
         return values
     # Field reordering is only used for a coarse overlap diagnostic; n sorting is authoritative.
     return values[..., ::-1]
 
 
 def _status(status: str, summary: str, **details) -> dict:
+    """Build a normalized status dictionary for report output."""
     return {"status": status, "failed": status == "fail", "summary": summary, **details}
 
 
 def _n_tolerance(entry: dict, recipe: dict | None = None) -> float:
+    """Resolve the effective-index tolerance for one fixture case."""
     if recipe is not None:
         recipe_tolerance = recipe.get("n_tolerance")
         if recipe_tolerance is not None:

benchmarks/compare_tidy3d_backends.py

@@ -1,3 +1,5 @@
+"""Benchmark MicroMode against equivalent Tidy3D mode-solver setups."""
+
 from __future__ import annotations
 
 import argparse
@@ -13,6 +15,8 @@
 
 @dataclass(frozen=True)
 class BenchmarkCase:
+    """Configuration for one backend comparison problem."""
+
     case_id: str
     description: str
     ny: int
@@ -53,6 +57,7 @@ class BenchmarkCase:
 
 
 def main() -> None:
+    """Run selected backend-comparison cases and print a markdown table."""
     args = parse_args()
     cases = list(PRESETS[args.preset])
     rows = [run_case(case, profile_source=args.profile_source) for case in cases]
@@ -64,6 +69,7 @@ def main() -> None:
 
 
 def parse_args() -> argparse.Namespace:
+    """Parse backend benchmark CLI options."""
     parser = argparse.ArgumentParser(description="Compare MicroMode SciPy and Tidy3D local solves.")
     parser.add_argument("--preset", choices=tuple(PRESETS), default="quick")
     parser.add_argument(
@@ -77,6 +83,7 @@ def parse_args() -> argparse.Namespace:
 
 
 def run_case(case: BenchmarkCase, *, profile_source: str) -> dict[str, object]:
+    """Execute one benchmark case for MicroMode and Tidy3D."""
     tidy3d_solver = make_tidy3d_solver(case)
     materials = (
         micromode_materials_from_tidy3d_solver(tidy3d_solver, case)
@@ -111,6 +118,7 @@ def run_case(case: BenchmarkCase, *, profile_source: str) -> dict[str, object]:
 
 
 def time_micromode(case: BenchmarkCase, materials: mm.Materials) -> tuple[float, np.ndarray]:
+    """Time the MicroMode solve for a prepared material grid."""
     start = time.perf_counter()
     data = mm.solve_modes(
         material_grid=materials,
@@ -124,12 +132,14 @@ def time_micromode(case: BenchmarkCase, materials: mm.Materials) -> tuple[float,
 
 
 def time_tidy3d(solver) -> tuple[float, np.ndarray]:
+    """Time a Tidy3D mode solve."""
     start = time.perf_counter()
     data = solver.solve()
     return time.perf_counter() - start, np.asarray(data.n_eff.values[0], dtype=float)
 
 
 def make_tidy3d_solver(case: BenchmarkCase):
+    """Construct a Tidy3D mode solver for one benchmark case."""
     try:
         import tidy3d as td
         from tidy3d.plugins.mode import ModeSolver
@@ -157,6 +167,7 @@ def make_tidy3d_solver(case: BenchmarkCase):
 
 
 def micromode_materials_from_tidy3d_solver(solver, case: BenchmarkCase) -> mm.Materials:
+    """Rasterize Tidy3D solver materials into a MicroMode grid."""
     eps = np.asarray(solver._solver_eps(tidy3d_frequency()), dtype=np.complex128)
     grid = solver._solver_grid
     return mm.Materials.from_components(
@@ -176,12 +187,14 @@ def micromode_materials_from_tidy3d_solver(solver, case: BenchmarkCase) -> mm.Ma
 
 
 def tidy3d_frequency() -> float:
+    """Return the benchmark frequency in Hz."""
     import tidy3d as td
 
     return float(td.C_0 / WAVELENGTH_UM)
 
 
 def micromode_materials(case: BenchmarkCase) -> mm.Materials:
+    """Build a direct MicroMode material grid for one benchmark case."""
     y_edges = np.linspace(-0.5 * WIDTH_Y, 0.5 * WIDTH_Y, case.ny + 1)
     z_edges = np.linspace(-0.5 * WIDTH_Z, 0.5 * WIDTH_Z, case.nz + 1)
     y = 0.5 * (y_edges[:-1] + y_edges[1:])
@@ -200,6 +213,7 @@ def micromode_materials(case: BenchmarkCase) -> mm.Materials:
 
 
 def tidy3d_structures(td, problem: str):
+    """Return Tidy3D geometry structures for one benchmark problem."""
     structures = [
         td.Structure(
             geometry=td.Box(center=(0.0, 0.0, -0.5001), size=(td.inf, td.inf, 1.0002)),
@@ -227,13 +241,15 @@ def tidy3d_structures(td, problem: str):
 
 
 def max_abs_delta(left: np.ndarray, right: np.ndarray) -> float:
+    """Return the maximum absolute difference between sorted mode arrays."""
     count = min(left.size, right.size)
     if count == 0:
         return float("nan")
     return float(np.max(np.abs(left[:count] - right[:count])))
 
 
 def markdown_table(rows: list[dict[str, object]]) -> str:
+    """Format benchmark rows as a markdown table."""
     header = (
         "| Problem | Grid | MicroMode SciPy (s) | Tidy3D local (s) | max abs Δn_eff SciPy/Tidy3D |\n"
         "|---|---:|---:|---:|---:|"

benchmarks/micromode_solver_benchmark.py

@@ -1,3 +1,5 @@
+"""Time MicroMode solves and record sparse-operator diagnostics."""
+
 from __future__ import annotations
 
 import argparse
@@ -11,6 +13,7 @@
 
 
 def main() -> None:
+    """Run timing cases and write a JSON benchmark report."""
     args = parse_args()
     rows = []
     grids = args.grid or ["20x14", "32x22", "48x32"]
@@ -55,6 +58,7 @@ def main() -> None:
 
 
 def parse_args() -> argparse.Namespace:
+    """Parse solver benchmark CLI options."""
     parser = argparse.ArgumentParser(description="Benchmark MicroMode sparse solves over grid sizes.")
     parser.add_argument(
         "--grid",
@@ -73,6 +77,7 @@ def parse_args() -> argparse.Namespace:
 
 
 def parse_grid_sizes(values: list[str]) -> list[tuple[int, int]]:
+    """Parse grid-size strings like 20x14 into integer pairs."""
     sizes = []
     for value in values:
         left, sep, right = value.lower().partition("x")
@@ -83,6 +88,7 @@ def parse_grid_sizes(values: list[str]) -> list[tuple[int, int]]:
 
 
 def strip_materials(*, nx: int, ny: int) -> mm.Materials:
+    """Build a simple strip-waveguide material grid for timing."""
     x_edges = np.linspace(-1.2, 1.2, nx + 1)
     y_edges = np.linspace(-0.8, 0.8, ny + 1)
     x = 0.5 * (x_edges[:-1] + x_edges[1:])

benchmarks/mode_solver/fixtures.py

@@ -1,3 +1,5 @@
+"""Shared helpers for reading committed mode-solver reference fixtures."""
+
 from __future__ import annotations
 
 import hashlib
@@ -15,22 +17,27 @@
 
 
 def case_dir(root: Path, case_id: str) -> Path:
+    """Return the directory containing one fixture case."""
     return root / case_id
 
 
 def data_path(root: Path, case_id: str) -> Path:
+    """Return the HDF5 path for one fixture case."""
     return case_dir(root, case_id) / "mode_data.hdf5"
 
 
 def summary_path(root: Path, case_id: str) -> Path:
+    """Return the JSON summary path for one fixture case."""
     return case_dir(root, case_id) / "summary.json"
 
 
 def manifest_path(root: Path) -> Path:
+    """Return the manifest path for a fixture suite."""
     return root / "manifest.json"
 
 
 def sha256_file(path: Path) -> str:
+    """Return the SHA-256 digest for a file."""
     digest = hashlib.sha256()
     with path.open("rb") as f:
         for chunk in iter(lambda: f.read(1024 * 1024), b""):
@@ -39,14 +46,17 @@ def sha256_file(path: Path) -> str:
 
 
 def read_json(path: Path) -> dict[str, Any]:
+    """Read a JSON file into a dictionary."""
     return json.loads(path.read_text())
 
 
 def write_json(path: Path, payload: dict[str, Any]) -> None:
+    """Write a dictionary as stable, pretty JSON."""
     path.write_text(json.dumps(payload, indent=2, sort_keys=True) + "\n")
 
 
 def iter_manifest_entries(root: Path) -> tuple[dict[str, Any], ...]:
+    """Return all case entries from a suite manifest."""
     return tuple(read_json(manifest_path(root))["cases"])
 
 
@@ -65,6 +75,7 @@ def load_data_array(path: Path, name: str) -> xr.DataArray:
 
 
 def _read_xarray_group(group: Any) -> xr.DataArray:
+    """Read a legacy xarray-style HDF5 group."""
     if _XARRAY_VALUE_NAME not in group:
         raise KeyError(f"HDF5 group {group.name!r} is missing {_XARRAY_VALUE_NAME!r}")
     values = group[_XARRAY_VALUE_NAME][()]
@@ -75,6 +86,7 @@ def _read_xarray_group(group: Any) -> xr.DataArray:
 
 
 def _infer_dims(shape: tuple[int, ...], coords: dict[str, np.ndarray]) -> tuple[str, ...]:
+    """Infer dimension names from HDF5 coordinate datasets."""
     dims = tuple(dim for dim in _PREFERRED_DIMS if dim in coords)
     if len(dims) == len(shape) and all(len(coords[dim]) == size for dim, size in zip(dims, shape, strict=True)):
         return dims
@@ -86,6 +98,7 @@ def _infer_dims(shape: tuple[int, ...], coords: dict[str, np.ndarray]) -> tuple[
 
 
 def phase_aligned_relative_error(golden: np.ndarray, actual: np.ndarray) -> tuple[float, float]:
+    """Compare complex arrays after removing a global phase offset."""
     g = np.asarray(golden).reshape(-1)
     a = np.asarray(actual).reshape(-1)
     norm_g = float(np.linalg.norm(g))