Feature: Support wavelength-dependent refractive index (material dispersion)

[singer-yang]

Summary

Currently, all refractive indices in DiffTMM are static scalars — they do not vary with wavelength. This makes it impossible to accurately model real dispersive materials (e.g. SiO₂, TiO₂, Si, metals) whose refractive index n(λ) changes significantly across the simulation bandwidth.

This issue tracks adding first-class support for wavelength-dependent (dispersive) refractive indices via two complementary mechanisms:

1. Refractive index lookup table — user-supplied (λ, n) pairs, interpolated at simulation wavelengths
2. Sellmeier equation — analytic dispersion formula with fitted coefficients

---

Background & Current Limitation

Where dispersion is missing today

In film_solver_isotropic.py, layer refractive indices are stored once at construction:

# film_solver_isotropic.py ~line 405
self.refract_idx_layers = n_layers_t.unsqueeze(0).expand(self.n_mirrors, -1).clone()
# shape: (n_mirrors, n_layers)  — no wavelength dimension

In film_solver_anisotropic.py, the same pattern applies for three-axis materials:

# film_solver_anisotropic.py ~line 571
self.refract_idx_layers  # shape: (n_mirrors, n_layers, 3)  — no wavelength dimension

When simulate(theta, wvln) is called, the wavelength wvln is used only to compute the wave number k0 = 2π/λ. The refractive index n is broadcast as a constant across all wavelengths. This is correct for optimization of structural parameters but breaks any use case that requires accurate broadband modelling of real materials.

Where to inject wavelength-dependence

In _compute_isotropic_tmm() (line ~76), n_layers has shape (batch, 1, 1, layer). Adding a wavelength axis would make it (batch, wv, 1, layer), which already broadcasts correctly with the existing k0 shape (batch, wv, 1, 1).

In create_jones_matrix_AOIAz() (line ~370), n_2d has shape (batch, layer, 3). The equivalent dispersive shape would be (batch, wv, layer, 3).

---

Proposed Design

1. RefractiveIndexTable — tabulated dispersion

class RefractiveIndexTable:
    """
    Wavelength-indexed lookup table for n and k.
    Interpolates (linearly or via cubic spline) at query wavelengths.
    """
    def __init__(self, wavelengths_um: Tensor, n: Tensor, k: Tensor | None = None):
        ...

    def __call__(self, wvln: Tensor) -> Tensor:
        """Returns complex n+ik interpolated at each wavelength. Shape: (n_wvlns,)"""
        ...

Usage example:

# e.g., SiO2 data from refractiveindex.info
table = RefractiveIndexTable(
    wavelengths_um=torch.tensor([0.4, 0.5, 0.6, 0.7, 0.8]),
    n=torch.tensor([1.470, 1.462, 1.458, 1.455, 1.453]),
)

2. SellmeierMaterial — analytic dispersion

The Sellmeier equation (single-term form, extendable to N terms):

$$n^2(\lambda) = 1 + \sum_{i=1}^{N} \frac{B_i \lambda^2}{\lambda^2 - C_i}$$

where λ is in µm, and B_i, C_i are material-specific coefficients.

class SellmeierMaterial:
    """
    Sellmeier dispersion formula. Coefficients B and C must match in length.
    λ units: micrometers.
    """
    def __init__(self, B: list[float], C: list[float]):
        ...

    def __call__(self, wvln: Tensor) -> Tensor:
        """Returns real refractive index n(λ). Shape: (n_wvlns,)"""
        ...

Usage example (BK7 glass, 3-term Sellmeier):

bk7 = SellmeierMaterial(
    B=[1.03961212, 0.231792344, 1.01046945],
    C=[6.00069867e-3, 2.00179144e-2, 1.03560653e2],  # in µm²
)

3. Updated solver API

Both IsotropicFilmSolver and AnisotropicFilmSolver should accept a material callable (any object with __call__(wvln) -> Tensor) or a plain scalar/complex value per layer. This keeps full backward compatibility.

solver = IsotropicFilmSolver(
    n_layers_list=[
        1.0,                             # air — static, as today
        SellmeierMaterial(B=[...], C=[...]),   # TiO2 — dispersive
        RefractiveIndexTable(wv, n, k),  # custom material from file
        1.5,                             # substrate — static
    ],
    d_list=[...],
)

When simulate(theta, wvln) is called, each dispersive material callable is evaluated at wvln to produce a (n_wvlns,) tensor, which is then inserted into n_layers at the wavelength dimension before the existing TMM computation.

---

Acceptance Criteria

• RefractiveIndexTable class with linear and cubic-spline interpolation modes
• SellmeierMaterial class supporting N-term Sellmeier equation
• Both classes return torch.complex64 tensors (to support lossy materials via k)
• IsotropicFilmSolver accepts dispersive materials per layer alongside existing static scalars
• AnisotropicFilmSolver accepts dispersive materials per principal axis (nx(λ), ny(λ), nz(λ))
• Dispersive path is fully differentiable through torch.autograd (gradients flow through n(λ))
• Backward compatibility: existing code passing plain scalars or tensors continues to work unchanged
• Unit tests comparing dispersive TMM output against the reference tmm_numpy implementation for a known material (e.g. SiO₂ using tabulated data)
• At least one example notebook or benchmark demonstrating broadband simulation with a real dispersive material

---

Implementation Notes

• Interpolation in RefractiveIndexTable should use torch.searchsorted + linear interpolation to stay on-device and remain differentiable w.r.t. n/k values (useful for fitting tabulated data).
• SellmeierMaterial is already analytic and trivially differentiable.
• The Sellmeier formula can produce n² < 0 below resonance (absorption edge). Handle this by returning the complex square root torch.sqrt(n2.clamp(min=0) + 0j) or raising an error for out-of-range wavelengths.
• For anisotropic materials, the three principal indices may each have independent dispersion (e.g. birefringent crystals like quartz or calcite). The SellmeierMaterial and RefractiveIndexTable classes should therefore be accepted per axis.
• Predefined constants for common materials (SiO₂, TiO₂, Si, BK7, MgF₂, ZnS …) would be a useful follow-up, but are out of scope for this issue.