Enhance documentation and validation across core modules

- Re-export public API in `__init__.py` for user convenience.
- Add detailed comments in `constants.py`, `models.py`, `raster.py`, `result.py`, `scipy_reference.py`, and `sweep.py` to clarify functionality and improve maintainability.
- Ensure explicit handling of public interfaces and validation checks in data classes and solver functions.

Author: QuentinWach

python/micromode/__init__.py

@@ -2,12 +2,16 @@
 
 from __future__ import annotations
 
+# Re-export the small public API from package root so users do not need to know
+# the internal module split.
 from .constants import C_0, EPSILON_0
 from .models import BoundarySpec, Grid, Materials, PmlSpec, Spec
 from .raster import solve_grid, solve_modes, solve_slice
 from .result import Result, overlap
 from .sweep import Sweep, track_modes_by_overlap
 
+# Keep __all__ explicit so documentation and static analysis show the intended
+# public surface.
 __all__ = [
     "C_0",
     "EPSILON_0",

python/micromode/constants.py

@@ -1,4 +1,8 @@
 """Physical constants used by MicroMode."""
 
+# Speed of light in microns per second, matching the coordinate unit used by the
+# raster API.
 C_0 = 2.997_924_58e14
+
+# Vacuum permittivity in solver units.
 EPSILON_0 = 8.854_187_812_800_384e-18

python/micromode/models.py

@@ -14,6 +14,8 @@
 BoundaryCondition = Literal["pec", "pmc"]
 
 
+# These dataclasses are intentionally lightweight: they validate user-facing
+# arrays and options before the numerical code receives them.
 @dataclass(frozen=True)
 class PmlSpec:
     """Perfectly matched layer settings for mode solves."""
@@ -26,18 +28,24 @@ class PmlSpec:
 
     def __post_init__(self) -> None:
         """Normalize constructor inputs and enforce model invariants."""
+        # First normalize the two edge counts because they control PML indexing.
         if len(self.num_cells) != 2:
             raise ValueError("num_cells must contain two non-negative integers")
         num_cells = (
             _coerce_integral("num_cells", self.num_cells[0], minimum=0),
             _coerce_integral("num_cells", self.num_cells[1], minimum=0),
         )
         object.__setattr__(self, "num_cells", num_cells)
+
+        # Stretch parameters must be positive scalars; zero would remove the
+        # damping profile or create invalid divisions in the PML factors.
         for name in ("sigma_max", "kappa_min", "kappa_max"):
             value = float(getattr(self, name))
             if not np.isfinite(value) or value <= 0.0:
                 raise ValueError(f"{name} must be finite and positive")
             object.__setattr__(self, name, value)
+
+        # The outer PML stretch should never be weaker than the inner stretch.
         if self.kappa_max < self.kappa_min:
             raise ValueError("kappa_max must be greater than or equal to kappa_min")
         object.__setattr__(self, "order", _coerce_integral("order", self.order, minimum=1))
@@ -75,9 +83,13 @@ class BoundarySpec:
 
     def __post_init__(self) -> None:
         """Normalize constructor inputs and enforce model invariants."""
+        # Keep only the low-edge symmetry flags; high edges are handled by the
+        # open-domain derivative/PML construction.
         if len(self.low) != 2:
             raise ValueError("low must contain two boundary conditions")
         normalized = tuple(str(value).lower() for value in self.low)
+
+        # Reject misspellings early so they do not silently become PEC walls.
         unknown = set(normalized).difference({"pec", "pmc"})
         if unknown:
             raise ValueError("boundary conditions must be 'pec' or 'pmc'")
@@ -114,12 +126,20 @@ class Grid:
 
     def __post_init__(self) -> None:
         """Normalize constructor inputs and enforce model invariants."""
+        # Store immutable float tuples so downstream validation and metadata
+        # serialization see one consistent representation.
         x_edges = tuple(float(value) for value in self.x_edges)
         y_edges = tuple(float(value) for value in self.y_edges)
         object.__setattr__(self, "x_edges", x_edges)
         object.__setattr__(self, "y_edges", y_edges)
+
+        # The normal axis determines how local solver fields map back to global
+        # x/y/z components, so it must be one of the Cartesian axes.
         if self.normal_axis not in {0, 1, 2}:
             raise ValueError("normal_axis must be 0, 1, or 2")
+
+        # Edge arrays define cell widths; non-monotonic values would produce
+        # negative or zero finite-difference spacings.
         for name, values in {"x_edges": x_edges, "y_edges": y_edges}.items():
             if len(values) < 2:
                 raise ValueError(f"{name} must contain at least two values")
@@ -148,9 +168,13 @@ class Materials:
 
     def __post_init__(self) -> None:
         """Normalize constructor inputs and enforce model invariants."""
+        # Epsilon is required and must already be sampled on the solver cells.
         eps_tensor = np.asarray(self.eps_tensor, dtype=np.complex128)
         if eps_tensor.shape != (3, 3, *self.grid.shape):
             raise ValueError("eps_tensor must have shape (3, 3, nx, ny) matching the grid")
+
+        # Permeability defaults to identity so non-magnetic material grids stay
+        # concise at call sites while the solver always receives full tensors.
         if self.mu_tensor is None:
             mu_tensor = np.zeros_like(eps_tensor)
             for axis in range(3):
@@ -159,6 +183,9 @@ def __post_init__(self) -> None:
             mu_tensor = np.asarray(self.mu_tensor, dtype=np.complex128)
             if mu_tensor.shape != eps_tensor.shape:
                 raise ValueError("mu_tensor must have the same shape as eps_tensor")
+
+        # NaNs and infinities poison sparse assembly, so reject them at the data
+        # model boundary instead of deep inside ARPACK.
         if not np.all(np.isfinite(eps_tensor)) or not np.all(np.isfinite(mu_tensor)):
             raise ValueError("material tensors must contain finite values")
         object.__setattr__(self, "eps_tensor", eps_tensor)
@@ -180,12 +207,16 @@ def from_diagonal(
         normal_coordinate: float = 0.0,
     ) -> Materials:
         """Build a material grid from scalar or diagonal tensor components."""
+        # Build and validate grid metadata before interpreting any component
+        # arrays, because the grid shape drives every component check.
         grid = Grid(
             tuple(float(value) for value in x_edges),
             tuple(float(value) for value in y_edges),
             normal_axis=normal_axis,
             normal_coordinate=normal_coordinate,
         )
+
+        # Missing yy/zz components inherit xx, matching isotropic scalar input.
         eps_diag = _stack_diagonal_components("eps", grid.shape, eps_xx, eps_yy, eps_zz)
         mu_diag = _stack_diagonal_components(
             "mu",
@@ -194,6 +225,9 @@ def from_diagonal(
             np.ones(grid.shape, dtype=np.complex128) if mu_yy is None else mu_yy,
             np.ones(grid.shape, dtype=np.complex128) if mu_zz is None else mu_zz,
         )
+
+        # Store the result in the same full-tensor layout used by tensorial
+        # inputs; diagonal detection later chooses the faster solver path.
         return cls(
             grid=grid, eps_tensor=_diagonal_to_full_tensor(eps_diag), mu_tensor=_diagonal_to_full_tensor(mu_diag)
         )
@@ -226,12 +260,17 @@ def from_components(
         normal_coordinate: float = 0.0,
     ) -> Materials:
         """Build a material grid from diagonal and off-diagonal tensor components."""
+        # Full component construction shares the same grid validation as the
+        # diagonal constructor, then overlays optional off-diagonal entries.
         grid = Grid(
             tuple(float(value) for value in x_edges),
             tuple(float(value) for value in y_edges),
             normal_axis=normal_axis,
             normal_coordinate=normal_coordinate,
         )
+
+        # Start with diagonal tensors so scalar and anisotropic grids follow one
+        # predictable path before optional coupling components are assigned.
         eps_diag = _stack_diagonal_components("eps", grid.shape, eps_xx, eps_yy, eps_zz)
         eps_tensor = _diagonal_to_full_tensor(eps_diag)
         _assign_tensor_offdiagonal(
@@ -245,6 +284,9 @@ def from_components(
             zx=eps_zx,
             zy=eps_zy,
         )
+
+        # Magnetic coupling components are optional; absent mu entries represent
+        # identity permeability on the diagonal and zero off-diagonal coupling.
         mu_diag = _stack_diagonal_components(
             "mu",
             grid.shape,
@@ -264,6 +306,9 @@ def from_components(
             zx=mu_zx,
             zy=mu_zy,
         )
+
+        # The dataclass constructor performs the final finite-value and shape
+        # validation on the assembled full tensors.
         return cls(grid=grid, eps_tensor=eps_tensor, mu_tensor=mu_tensor)
 
     @classmethod
@@ -306,6 +351,9 @@ def from_slice(
 
         axis_index = _normalize_slice_axis(axis)
         edge_values = tuple(float(value) for value in coord_edges)
+
+        # A slice is still represented as a 2D mode plane, so one axis gets the
+        # user coordinates and the other gets a single finite-width cell.
         if len(edge_values) < 2:
             raise ValueError("coord_edges must contain at least two values")
         if invariant_width <= 0.0 or not np.isfinite(invariant_width):
@@ -320,16 +368,24 @@ def from_slice(
 
         def expand(label: str, values: np.ndarray | None) -> np.ndarray | None:
             """Expand a one-dimensional component onto the padded mode-plane grid."""
+            # Optional components stay absent so from_components can apply its
+            # diagonal/permeability defaults in one place.
             if values is None:
                 return None
             array = np.asarray(values, dtype=np.complex128)
             if array.shape != (cell_count,):
                 raise ValueError(f"{label} must have shape {(cell_count,)} for a one-dimensional slice")
+
+            # Insert the singleton invariant axis on the side implied by the
+            # requested varying axis.
             return array[:, None] if axis_index == 0 else array[None, :]
 
         expanded_eps_xx = expand("eps_xx", eps_xx)
         if expanded_eps_xx is None:
             raise ValueError("eps_xx is required")
+
+        # Delegate to the full component constructor so slice and grid inputs
+        # share tensor assembly and validation behavior.
         return cls.from_components(
             x_edges=x_edges,
             y_edges=y_edges,
@@ -383,6 +439,9 @@ def from_subpixel_diagonal(
         x_edge_values = tuple(float(value) for value in x_edges)
         y_edge_values = tuple(float(value) for value in y_edges)
         shape = (len(x_edge_values) - 1, len(y_edge_values) - 1)
+
+        # Average each supplied high-resolution component independently before
+        # reusing the ordinary diagonal tensor constructor.
         averaged = {
             "eps_xx": cls.average_subpixels(eps_xx, shape=shape, subpixel_shape=subpixel_shape, method=average),
             "eps_yy": None
@@ -427,18 +486,26 @@ def average_subpixels(
 
         nx, ny = (int(shape[0]), int(shape[1]))
         sx, sy = (int(subpixel_shape[0]), int(subpixel_shape[1]))
+
+        # Validate the target grouping before reshaping so shape errors are
+        # reported in terms of solver cells and subpixel samples.
         if nx <= 0 or ny <= 0:
             raise ValueError("shape must contain positive cell counts")
         if sx <= 0 or sy <= 0:
             raise ValueError("subpixel_shape must contain positive sample counts")
         array = np.asarray(values, dtype=np.complex128)
+
+        # Accept either a dense raster or an already grouped raster; both become
+        # (nx, ny, sx, sy) before the selected averaging rule is applied.
         if array.shape == (nx * sx, ny * sy):
             grouped = array.reshape(nx, sx, ny, sy).transpose(0, 2, 1, 3)
         elif array.shape == (nx, ny, sx, sy):
             grouped = array
         else:
             raise ValueError(f"subpixel values must have shape {(nx * sx, ny * sy)} or {(nx, ny, sx, sy)}")
 
+        # Choose the averaging rule explicitly so interface-sensitive callers can
+        # request harmonic/geometric behavior without changing solver code.
         if method == "arithmetic":
             return grouped.mean(axis=(2, 3))
         if method == "harmonic":
@@ -463,6 +530,8 @@ def shape(self) -> tuple[int, int]:
     @property
     def is_diagonal(self) -> bool:
         """Return whether epsilon and mu contain only diagonal components."""
+        # A boolean mask lets the check stay independent of grid shape and
+        # catches off-diagonal entries across every solver cell.
         off_diagonal = np.ones((3, 3), dtype=bool)
         np.fill_diagonal(off_diagonal, False)
         mu_tensor = self._resolved_mu_tensor()
@@ -497,24 +566,33 @@ def _stack_diagonal_components(
     zz: np.ndarray | None,
 ) -> np.ndarray:
     """Validate and stack x/y/z diagonal tensor components."""
+    # xx is the required scalar/diagonal component; yy and zz inherit it unless
+    # anisotropic values are provided.
     xx_array = np.asarray(xx, dtype=np.complex128)
     if xx_array.shape != shape:
         raise ValueError(f"{label}_xx must have shape {shape}")
     yy_array = xx_array if yy is None else np.asarray(yy, dtype=np.complex128)
     zz_array = xx_array if zz is None else np.asarray(zz, dtype=np.complex128)
+
+    # All diagonal components must be cell-centered arrays matching the grid.
     if yy_array.shape != shape or zz_array.shape != shape:
         raise ValueError(f"{label}_xx, {label}_yy, and {label}_zz must have shape {shape}")
     return np.stack([xx_array, yy_array, zz_array], axis=0)
 
 
 def _coerce_integral(name: str, value: object, *, minimum: int) -> int:
     """Coerce index-like values while rejecting floats and booleans."""
+    # bool implements the integer protocol, but accepting it would make option
+    # typos like num_cells=(True, False) look valid.
     if isinstance(value, (bool, np.bool_)):
         raise ValueError(f"{name} must contain integers")
     try:
         integer = index(cast(SupportsIndex, value))
     except TypeError as exc:
         raise ValueError(f"{name} must contain integers") from exc
+
+    # Keep the minimum check here so every PML/order caller reports consistent
+    # validation messages.
     if integer < minimum:
         if minimum == 0:
             raise ValueError(f"{name} must contain non-negative integers")
@@ -533,6 +611,8 @@ def _normalize_slice_axis(axis: SliceAxis) -> int:
 
 def _diagonal_to_full_tensor(diagonal: np.ndarray) -> np.ndarray:
     """Expand diagonal components into a full 3x3 tensor grid."""
+    # The solver consumes a full tensor layout even for diagonal materials; all
+    # off-diagonal entries remain zero.
     tensor = np.zeros((3, 3, *diagonal.shape[1:]), dtype=np.complex128)
     for axis in range(3):
         tensor[axis, axis, :, :] = diagonal[axis]
@@ -552,6 +632,8 @@ def _assign_tensor_offdiagonal(
     zy: np.ndarray | None,
 ) -> None:
     """Validate and assign optional off-diagonal tensor components."""
+    # Iterate through physical tensor suffixes instead of open-coding each
+    # assignment, which keeps validation messages and row/column mapping aligned.
     for (row, col), suffix, values in [
         ((0, 1), "xy", xy),
         ((0, 2), "xz", xz),
@@ -563,6 +645,9 @@ def _assign_tensor_offdiagonal(
         if values is None:
             continue
         array = np.asarray(values, dtype=np.complex128)
+
+        # Off-diagonal components are cell-centered just like the diagonal
+        # components; no broadcasting is allowed because that can hide mistakes.
         if array.shape != shape:
             raise ValueError(f"{label}_{suffix} must have shape {shape}")
         tensor[row, col, :, :] = array
@@ -583,10 +668,14 @@ class Spec:
 
     def __post_init__(self) -> None:
         """Normalize constructor inputs and enforce model invariants."""
+        # Validate simple scalar controls before expanding nested model objects.
         if self.num_modes <= 0:
             raise ValueError("num_modes must be positive")
         if self.target_neff is not None and self.target_neff <= 0:
             raise ValueError("target_neff must be positive")
+
+        # Accept compact tuple inputs at the API boundary, but store canonical
+        # model objects so solve_modes can unpack a Spec without extra branches.
         pml = self.pml
         pml = PmlSpec() if pml is None else pml
         if not isinstance(pml, PmlSpec):
@@ -596,6 +685,9 @@ def __post_init__(self) -> None:
         if not isinstance(boundary, BoundarySpec):
             boundary = BoundarySpec(low=boundary)
         object.__setattr__(self, "boundary", boundary)
+
+        # Bend transforms need both a radius and an axis because the Jacobian
+        # depends on which transverse coordinate measures curvature.
         if self.bend_radius is not None and np.isclose(self.bend_radius, 0.0):
             raise ValueError("bend_radius magnitude must be larger than 0")
         if self.bend_radius is not None and self.bend_axis is None: