examples/0-basic_pipeline/basic_pipeline.py - main - bari/HPLSE-poster-code

Репозиторий помещен в архив 16 июня 2026. Вся функциональность доступна только для чтения.

examples/0-basic_pipeline/basic_pipeline.py

67 строк2 KB

bari

Add time-domain and overdense plasma examples

13 апр 2026, 01:26

13 апр 2026, 01:26a33d224

"""Simplest end-to-end pipeline."""

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, result_metadata, save_figure, save_metadata
from helpinn import MonochromaticPulse, SampledProfile, TmmSolver


def main() -> None:
    """Run a simple profile -> solver -> plots pipeline."""
    z = np.linspace(0.0, 5.0, 200)
    profile = SampledProfile(
        z=z,
        eps=1.0 + 0.2 * np.exp(-((z - 2.5) / 0.8) ** 2),
    )
    pulse = MonochromaticPulse(
        amplitude=1.0 + 0.0j,
        wavelength=0.8,
        angle_rad=0.0,
        polarization="TE",
    )
    result = TmmSolver(profile, pulse).solve()
    run_dir = create_run_dir(pathlib.Path(__file__).resolve().parent)
    fig_profile, _ = profile.plot()
    fig_fields, _ = result.plot_fields()
    fig_fields_mode, _ = result.plot_fields(mode_index=0)
    fig_energy, _ = result.plot_energy()
    fig_energy_mode, _ = result.plot_energy(mode_index=0)
    fig_coefficients, _ = result.plot_coefficients()
    fig_coefficients_mode, _ = result.plot_coefficients(mode_index=0)
    fig_summary, _ = result.plot_summary()
    fig_summary_mode, _ = result.plot_summary(mode_index=0)
    save_figure(fig_profile, run_dir / "profile.png")
    save_figure(fig_fields, run_dir / "pulse_fields.png")
    save_figure(fig_fields_mode, run_dir / "mode0_fields.png")
    save_figure(fig_energy, run_dir / "pulse_energy.png")
    save_figure(fig_energy_mode, run_dir / "mode0_energy.png")
    save_figure(fig_coefficients, run_dir / "pulse_coefficients.png")
    save_figure(fig_coefficients_mode, run_dir / "mode0_coefficients.png")
    save_figure(fig_summary, run_dir / "pulse_summary.png")
    save_figure(fig_summary_mode, run_dir / "mode0_summary.png")
    save_metadata(
        run_dir / "metadata.json",
        {
            "example": "basic_pipeline",
            "solver": "tmm",
            "cumulative_R": result.cumulative_R,
            "cumulative_T": result.cumulative_T,
            "polarization": pulse.polarization,
            "pulse": result_metadata(result),
            "mode0": result_metadata(result, mode_index=0),
        },
    )
    print("results saved to", run_dir)


if __name__ == "__main__":
    main()