better docstrings

Author: flaport

src/meow/arrays.py

@@ -131,19 +131,43 @@ def _assert_dtype(arr: np.ndarray, dtype: str) -> np.ndarray:
 
 
 def Dim(ndim: int, *, coerce: bool = True) -> AfterValidator:  # noqa: N802
-    """Validator to ensure the array has a specific number of dimensions."""
+    """Validator to ensure the array has a specific number of dimensions.
+
+    Args:
+        ndim: the required number of dimensions.
+        coerce: if True, reshape the array to match; if False, raise on mismatch.
+
+    Returns:
+        A pydantic AfterValidator that checks or coerces the array dimensions.
+    """
     f = _coerce_dim if coerce else _assert_dim
     return AfterValidator(partial(f, ndim=ndim))
 
 
 def DType(dtype: str, *, coerce: bool = True) -> AfterValidator:  # noqa: N802
-    """Validator to ensure the array has a specific dtype."""
+    """Validator to ensure the array has a specific dtype.
+
+    Args:
+        dtype: the required data type as a string (e.g. "float64").
+        coerce: if True, cast the array to the dtype; if False, raise on mismatch.
+
+    Returns:
+        A pydantic AfterValidator that checks or coerces the array dtype.
+    """
     f = _coerce_dtype if coerce else _assert_dtype
     return AfterValidator(partial(f, dtype=dtype))
 
 
 def Shape(*shape: int, coerce: bool = True) -> AfterValidator:  # noqa: N802
-    """Validator to ensure the array has a specific shape."""
+    """Validator to ensure the array has a specific shape.
+
+    Args:
+        *shape: the required dimensions of the array.
+        coerce: if True, reshape the array to match; if False, raise on mismatch.
+
+    Returns:
+        A pydantic AfterValidator that checks or coerces the array shape.
+    """
     f = _coerce_shape if coerce else _assert_shape
     return AfterValidator(partial(f, shape=shape))
 

src/meow/base_model.py

@@ -181,7 +181,14 @@ def _visualize(self, **_: Any) -> None:
 
 
 def cache(prop: Callable) -> Callable:
-    """Decorator to cache the result of a property method."""
+    """Decorator to cache the result of a property method.
+
+    Args:
+        prop: the property method whose result should be cached.
+
+    Returns:
+        A wrapped callable that caches and returns the computed value.
+    """
     prop_name = getattr(prop, "__name__", "")
     if not prop_name:
         return prop
@@ -201,7 +208,14 @@ def getter(self):  # noqa: ANN001,ANN202
 
 
 def cached_property(method: Callable):  # noqa: ANN201
-    """Decorator to cache the result of a property method."""
+    """Decorator to cache the result of a property method.
+
+    Args:
+        method: the method to wrap as a cached property.
+
+    Returns:
+        A property descriptor that caches the method's return value.
+    """
     return property(cache(method))
 
 

src/meow/cell.py

@@ -122,7 +122,17 @@ def create_cells(
     Ls: Annotated[NDArray, Dim(1), DType("float64")],
     z_min: float = 0.0,
 ) -> list[Cell]:
-    """Create multiple `Cell` objects with a `Mesh` and a collection of cell lengths."""
+    """Create multiple `Cell` objects with a `Mesh` and a collection of cell lengths.
+
+    Args:
+        structures: the 3D structures shared by all cells.
+        mesh: a single mesh or a list of meshes, one per cell.
+        Ls: a 1D array of cell lengths.
+        z_min: the starting z-coordinate for the first cell.
+
+    Returns:
+        A list of Cell objects spanning the given lengths.
+    """
     Ls = np.asarray(Ls, float)
     if Ls.ndim != 1:
         msg = f"Ls should be 1D. Got shape: {Ls.shape}."

src/meow/eme/cascade.py

@@ -29,7 +29,20 @@ def compute_s_matrix_sax(
     propagations_fn: Callable = compute_propagation_s_matrices,
     **_: Any,
 ) -> sax.SDenseMM:
-    """Calculate the S-matrix for given sets of modes."""
+    """Calculate the S-matrix for given sets of modes.
+
+    Args:
+        modes: Modal basis for each cell in the stack.
+        cells: Cells from which to derive propagation lengths. Either cells
+            or cell_lengths must be provided.
+        cell_lengths: Optional explicit propagation lengths per cell.
+        sax_backend: SAX backend used for circuit evaluation.
+        interfaces_fn: Callable that computes interface S-matrices.
+        propagations_fn: Callable that computes propagation S-matrices.
+
+    Returns:
+        A tuple ``(S, port_map)`` in SAX dense multimode format.
+    """
     propagations = propagations_fn(modes, cells, cell_lengths=cell_lengths)
     interfaces = interfaces_fn(modes)
     net = _get_netlist(propagations, interfaces)
@@ -94,7 +107,15 @@ def _get_netlist(
 
 
 def downselect_s(S: sax.SDenseMM, ports: list[str]) -> sax.SDenseMM:
-    """Downselect the S-matrix to the given ports."""
+    """Downselect the S-matrix to the given ports.
+
+    Args:
+        S: A tuple ``(S_matrix, port_map)`` in SAX dense multimode format.
+        ports: Port names to keep.
+
+    Returns:
+        A new ``(S_matrix, port_map)`` tuple containing only the selected ports.
+    """
     S_matrix, port_map = S
     idxs = [port_map[port] for port in ports]
     S_matrix = S_matrix[idxs, :][:, idxs]

src/meow/eme/interface.py

@@ -221,6 +221,20 @@ def compute_interface_s_matrices(
     The same sharp edges apply here as for the single-interface solve:
     orthonormalization, inner-product choice, TSVD cutoff, and passivity method
     must all be interpreted consistently across the full stack.
+
+    Args:
+        modes: Ordered list of modal bases across the stack.
+        inner_product: Inner-product callable for forming overlap matrices.
+        conjugate: Whether to use the conjugated formulation. Inferred from
+            inner_product if None.
+        tsvd_rcond: Relative singular-value cutoff for the TSVD solve.
+        passivity_method: Method for enforcing S-matrix passivity.
+        enforce_reciprocity: Whether to symmetrize each interface S-matrix.
+        ignore_warnings: Whether to suppress numerical warnings.
+
+    Returns:
+        Dict mapping interface names (e.g. ``"i_0_1"``) to SAX dense
+        multimode S-matrix tuples.
     """
     return {
         f"i_{i}_{i + 1}": compute_interface_s_matrix(

src/meow/eme/propagation.py

@@ -23,6 +23,13 @@ def compute_propagation_s_matrix(modes: Modes, cell_length: float) -> sax.SDictM
     Each mode acquires a phase ``exp(2j * pi * neff / wl * cell_length)`` while
     propagating through the cell. Backward propagation is mirrored by the
     bidirectional port mapping in the returned SAX dictionary.
+
+    Args:
+        modes: Modal basis for the cell.
+        cell_length: Physical length of the cell.
+
+    Returns:
+        SAX dictionary mapping port pairs to complex transmission values.
     """
     s_dict = {
         (f"left@{i}", f"right@{i}"): jnp.exp(
@@ -67,7 +74,15 @@ def compute_propagation_s_matrices(
 
 
 def select_ports(S: sax.SDenseMM, ports: list[str]) -> sax.SDenseMM:
-    """Keep a subset of ports from an S-matrix."""
+    """Keep a subset of ports from an S-matrix.
+
+    Args:
+        S: A tuple ``(s_matrix, port_map)`` in SAX dense multimode format.
+        ports: Port names to retain.
+
+    Returns:
+        A new ``(s_matrix, port_map)`` tuple with only the selected ports.
+    """
     s, pm = S
     idxs = jnp.array([pm[port] for port in ports], dtype=jnp.int32)
     s = s[idxs, :][:, idxs]
@@ -112,7 +127,16 @@ def pi_pairs(
     interfaces: dict[str, sax.SDenseMM],
     sax_backend: sax.Backend,
 ) -> list[sax.STypeMM]:
-    """Return propagation-interface pairs for a full stack."""
+    """Return propagation-interface pairs for a full stack.
+
+    Args:
+        propagations: Propagation S-matrices keyed by cell index.
+        interfaces: Interface S-matrices keyed by adjacent cell pairs.
+        sax_backend: SAX backend used for cascading.
+
+    Returns:
+        List of cascaded propagation-interface S-matrices, one per cell.
+    """
     pairs: list[sax.STypeMM] = []
     for i in range(len(propagations)):
         propagation = propagations[f"p_{i}"]
@@ -128,7 +152,16 @@ def pi_pairs(
 def l2r_matrices(
     pairs: list[sax.STypeMM], identity: sax.SDenseMM, sax_backend: sax.Backend
 ) -> list[sax.STypeMM]:
-    """Return cumulative left-to-right S-matrices."""
+    """Return cumulative left-to-right S-matrices.
+
+    Args:
+        pairs: Propagation-interface pair S-matrices from ``pi_pairs``.
+        identity: Identity-like S-matrix used as the initial accumulator.
+        sax_backend: SAX backend used for cascading.
+
+    Returns:
+        List of cumulative S-matrices from the left boundary up to each cell.
+    """
     matrices: list[sax.STypeMM] = [identity]
     for pair in pairs[:-1]:
         matrices.append(_connect_two(matrices[-1], pair, sax_backend))
@@ -138,7 +171,15 @@ def l2r_matrices(
 def r2l_matrices(
     pairs: list[sax.STypeMM], sax_backend: sax.Backend
 ) -> list[sax.STypeMM]:
-    """Return cumulative right-to-left S-matrices."""
+    """Return cumulative right-to-left S-matrices.
+
+    Args:
+        pairs: Propagation-interface pair S-matrices from ``pi_pairs``.
+        sax_backend: SAX backend used for cascading.
+
+    Returns:
+        List of cumulative S-matrices from each cell to the right boundary.
+    """
     matrices = [pairs[-1]]
     for pair in pairs[-2::-1]:
         matrices.append(_connect_two(pair, matrices[-1], sax_backend))
@@ -150,7 +191,15 @@ def split_square_matrix(
 ) -> tuple[
     tuple[ComplexArray2D, ComplexArray2D], tuple[ComplexArray2D, ComplexArray2D]
 ]:
-    """Split a square matrix into its four block submatrices."""
+    """Split a square matrix into its four block submatrices.
+
+    Args:
+        matrix: Square matrix to split.
+        idx: Row/column index at which to split.
+
+    Returns:
+        Nested tuple ``((top_left, top_right), (bottom_left, bottom_right))``.
+    """
     if matrix.shape[0] != matrix.shape[1]:
         msg = "Matrix has to be square."
         raise ValueError(msg)
@@ -167,7 +216,18 @@ def compute_mode_amplitudes(
     excitation_l: ComplexArray1D,
     excitation_r: ComplexArray1D,
 ) -> tuple[ComplexArray1D, ComplexArray1D]:
-    """Solve for the forward and backward modal amplitudes in one cell."""
+    """Solve for the forward and backward modal amplitudes in one cell.
+
+    Args:
+        u: Cumulative left-to-right S-matrix (dense array).
+        v: Cumulative right-to-left S-matrix (dense array).
+        m: Number of right-side (forward) modes.
+        excitation_l: Left boundary excitation vector.
+        excitation_r: Right boundary excitation vector.
+
+    Returns:
+        Tuple ``(forward, backward)`` of complex amplitude vectors.
+    """
     n = u.shape[0] - m
     _, [u21, u22] = split_square_matrix(u, n)
     [v11, v12], _ = split_square_matrix(v, m)
@@ -185,7 +245,18 @@ def propagate(
     excitation_l: ComplexArray1D,
     excitation_r: ComplexArray1D,
 ) -> tuple[list[ComplexArray1D], list[ComplexArray1D]]:
-    """Propagate boundary excitations through cumulative S-matrices."""
+    """Propagate boundary excitations through cumulative S-matrices.
+
+    Args:
+        l2rs: Cumulative left-to-right S-matrices for each cell.
+        r2ls: Cumulative right-to-left S-matrices for each cell.
+        excitation_l: Left boundary excitation vector.
+        excitation_r: Right boundary excitation vector.
+
+    Returns:
+        Tuple ``(forwards, backwards)`` where each element is a list of
+        complex amplitude vectors, one per cell.
+    """
     forwards = []
     backwards = []
     for l2r, r2l in zip(l2rs, r2ls, strict=False):
@@ -212,7 +283,20 @@ def plot_fields(
     y: float,
     z: FloatArray1D,
 ) -> tuple[ComplexArray2D, FloatArray1D]:
-    """Reconstruct an ``Ex(x, z)`` field slice from propagated modal amplitudes."""
+    """Reconstruct an ``Ex(x, z)`` field slice from propagated modal amplitudes.
+
+    Args:
+        modes: Mode sets for each cell.
+        cells: Cells defining the geometry and mesh.
+        forwards: Forward amplitude vectors per cell.
+        backwards: Backward amplitude vectors per cell.
+        y: Transverse y-coordinate at which to sample the field.
+        z: Global z-grid on which to reconstruct the field.
+
+    Returns:
+        Tuple ``(field, x)`` where ``field`` is the complex ``Ex(z, x)``
+        array and ``x`` is the transverse sampling grid.
+    """
     mesh_y = cells[0].mesh.y
     mesh_x = cells[0].mesh.x
     mesh_x = mesh_x[:-1] + np.diff(mesh_x) / 2
@@ -275,6 +359,14 @@ def track_modes(
     Unmatched modes are appended after the matched subset in their original
     order. This function does not change the physics of the interface solve; it
     only makes the local basis more continuous for reconstruction and plotting.
+
+    Args:
+        modes: Ordered list of modal bases across the stack.
+        inner_product_fn: Inner-product callable used to compute overlaps
+            between neighboring mode sets.
+
+    Returns:
+        Reordered and phase-aligned list of modal bases.
     """
     if not modes:
         return []

src/meow/fde/lumerical.py

@@ -27,7 +27,18 @@ def compute_modes_lumerical(
     post_process: Callable = post_process_modes,
     sim: Sim | None = None,
 ) -> list[Mode]:
-    """Compute ``Modes` for a given ``FdeSpec` (Lumerical backend)."""
+    """Compute ``Modes`` for a given ``CrossSection`` (Lumerical backend).
+
+    Args:
+        cs: the cross-section to solve modes for.
+        num_modes: number of modes to compute.
+        unit: unit scaling factor (default 1e-6 for micrometres).
+        post_process: callable applied to the raw mode list before returning.
+        sim: optional Lumerical simulation object; uses the global one if None.
+
+    Returns:
+        The computed and post-processed list of modes.
+    """
     from lumapi import LumApiError  # type: ignore[reportMissingImports]
 
     sim = get_sim(sim=sim)
@@ -122,7 +133,14 @@ def create_lumerical_geometries(
     env: Environment,
     unit: float,
 ) -> None:
-    """Create Lumerical geometries from a list of structures."""
+    """Create Lumerical geometries from a list of structures.
+
+    Args:
+        sim: the Lumerical simulation object.
+        structures: 3-D structures to add to the simulation.
+        env: the environment containing material/wavelength info.
+        unit: unit scaling factor.
+    """
     sim = get_sim(sim=sim)
     sim.switchtolayout()
     sim.deleteall()
@@ -131,7 +149,15 @@ def create_lumerical_geometries(
 
 
 def get_sim(**kwargs: Any) -> Sim:
-    """Get the Lumerical simulation object."""
+    """Get the Lumerical simulation object.
+
+    Args:
+        **kwargs: keyword arguments; pass ``sim`` to set and return a specific
+            simulation object.
+
+    Returns:
+        The active Lumerical simulation object.
+    """
     global _sim  # noqa: PLW0603
     sim = kwargs.get("sim", None)
     if sim is not None: