examples/3-reference_angles/reference_angles.py - main - bari/HPLSE-poster-code

Репозиторий помещен в архив 16 июня 2026. Вся функциональность доступна только для чтения.
examples/3-reference_angles/reference_angles.py
127 строк5 KB
bari
Add regression tests and runnable examples
12 апр 2026, 23:59
12 апр 2026, 23:59d8ddeab
"""Analytical reference scans for Fresnel, Brewster, and critical angles."""

from __future__ import annotations

import pathlib
import sys

import numpy as np

ROOT = pathlib.Path(__file__).resolve().parents[2]
if str(ROOT) not in sys.path:
    sys.path.insert(0, str(ROOT))

from examples.common import create_run_dir, save_figure, save_metadata
from helpinn import MonochromaticPulse, RiccatiSolver, SampledProfile, SmmSolver, TmmSolver
from helpinn.utils import new_figure


def fresnel_reflectance(eps_left: float, eps_right: float, angle_rad: np.ndarray, polarization: str) -> np.ndarray:
    """Return Fresnel reflectance for one interface."""
    n_left = np.sqrt(eps_left + 0j)
    n_right = np.sqrt(eps_right + 0j)
    sin_theta_t = n_left * np.sin(angle_rad) / n_right
    cos_theta_i = np.cos(angle_rad)
    cos_theta_t = np.sqrt(1.0 - sin_theta_t**2 + 0j)
    if polarization == "TE":
        numerator = n_left * cos_theta_i - n_right * cos_theta_t
        denominator = n_left * cos_theta_i + n_right * cos_theta_t
    else:
        numerator = n_right * cos_theta_i - n_left * cos_theta_t
        denominator = n_right * cos_theta_i + n_left * cos_theta_t
    return np.abs(numerator / denominator) ** 2


def single_interface_profile(eps_left: float, eps_right: float) -> SampledProfile:
    """Build a one-interface profile."""
    return SampledProfile(
        z=np.array([0.0, 1.0]),
        eps=np.array([eps_left, eps_right]),
        eps_in=eps_left,
        eps_out=eps_right,
        discontinuity_mask=np.array([False, True]),
    )


def run_scan(eps_left: float, eps_right: float, polarization: str, angles: np.ndarray) -> dict[str, np.ndarray]:
    """Run one angular scan against all solvers and Fresnel reference."""
    profile = single_interface_profile(eps_left, eps_right)
    results: dict[str, np.ndarray] = {
        "angles": angles,
        "fresnel_R": fresnel_reflectance(eps_left, eps_right, angles, polarization),
    }
    for solver_name, solver_cls in {
        "tmm": TmmSolver,
        "smm": SmmSolver,
        "riccati": RiccatiSolver,
    }.items():
        reflectance = []
        for angle in angles:
            pulse = MonochromaticPulse(wavelength=1.0, angle_rad=float(angle), polarization=polarization)
            result = solver_cls(profile, pulse).solve()
            reflectance.append(result.cumulative_R)
        results[f"{solver_name}_R"] = np.asarray(reflectance, dtype=float)
    return results


def plot_scan(scan: dict[str, np.ndarray], output_path: pathlib.Path, title: str) -> None:
    """Plot one scan with solver curves and Fresnel reference."""
    fig, ax = new_figure()
    angle_deg = np.degrees(scan["angles"])
    ax.plot(angle_deg, scan["fresnel_R"], label="Fresnel", linewidth=2.0)
    ax.plot(angle_deg, scan["tmm_R"], label="TMM")
    ax.plot(angle_deg, scan["smm_R"], label="SMM")
    ax.plot(angle_deg, scan["riccati_R"], label="Riccati")
    ax.set_xlabel("angle [deg]")
    ax.set_ylabel("R")
    ax.set_title(title)
    ax.legend()
    save_figure(fig, output_path)


def plot_difference(scan: dict[str, np.ndarray], output_path: pathlib.Path, title: str) -> None:
    """Plot solver differences against Fresnel reference."""
    fig, ax = new_figure()
    angle_deg = np.degrees(scan["angles"])
    fresnel = scan["fresnel_R"]
    ax.plot(angle_deg, np.abs(scan["tmm_R"] - fresnel), label="|TMM - Fresnel|")
    ax.plot(angle_deg, np.abs(scan["smm_R"] - fresnel), label="|SMM - Fresnel|")
    ax.plot(angle_deg, np.abs(scan["riccati_R"] - fresnel), label="|Riccati - Fresnel|")
    ax.set_xlabel("angle [deg]")
    ax.set_ylabel("absolute error")
    ax.set_title(title)
    ax.legend()
    save_figure(fig, output_path)


def main() -> None:
    """Run analytical reference scans and save plots."""
    run_dir = create_run_dir(pathlib.Path(__file__).resolve().parent)
    angles_air_glass = np.linspace(0.0, np.radians(89.9), 800)
    brewster_scan = run_scan(1.0, 2.25, "TM", angles_air_glass)
    normal_scan = run_scan(1.0, 2.25, "TE", angles_air_glass)
    critical_angle = float(np.arcsin(np.sqrt(1.0 / 2.25)))
    angles_glass_air = np.linspace(0.0, np.radians(89.9), 800)
    critical_scan = run_scan(2.25, 1.0, "TE", angles_glass_air)

    plot_scan(normal_scan, run_dir / "normal_te_scan.png", "Air to glass, TE")
    plot_difference(normal_scan, run_dir / "normal_te_error.png", "Air to glass, TE error")
    plot_scan(brewster_scan, run_dir / "brewster_tm_scan.png", "Air to glass, TM")
    plot_difference(brewster_scan, run_dir / "brewster_tm_error.png", "Air to glass, TM error")
    plot_scan(critical_scan, run_dir / "critical_te_scan.png", "Glass to air, TE")
    plot_difference(critical_scan, run_dir / "critical_te_error.png", "Glass to air, TE error")

    save_metadata(
        run_dir / "metadata.json",
        {
            "example": "reference_angles",
            "brewster_angle_deg": float(np.degrees(np.arctan(np.sqrt(2.25 / 1.0)))),
            "critical_angle_deg": float(np.degrees(critical_angle)),
            "normal_incidence_reference_R": float(fresnel_reflectance(1.0, 2.25, np.array([0.0]), "TE")[0]),
        },
    )
    print("results saved to", run_dir)


if __name__ == "__main__":
    main()