test/example/train_rectangular_spectrum.py - master - bari/PlasmaTorch

test/example/train_rectangular_spectrum.py
301 строка11 KB
bari
Update plasma state equation API
18 июн 2026, 12:45
18 июн 2026, 12:458b10769
from __future__ import annotations

import os
from pathlib import Path
import sys
import tempfile
import warnings

os.environ.setdefault("MPLCONFIGDIR", str(Path(tempfile.gettempdir()) / "plast_mplconfig"))
os.environ.setdefault("XDG_CACHE_HOME", str(Path(tempfile.gettempdir()) / "plast_cache"))
Path(os.environ["MPLCONFIGDIR"]).mkdir(parents=True, exist_ok=True)
Path(os.environ["XDG_CACHE_HOME"]).mkdir(parents=True, exist_ok=True)

import matplotlib

matplotlib.use("Agg")

import matplotlib.pyplot as plt
import torch

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

from plast import Plasma, Pulse, Solver, drude_model


DTYPE = torch.float64
DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")

ITERATIONS = 500
SAVE_EVERY = 10
LEARNING_RATE = 0.003

CARRIER_WAVELENGTH = 0.80
CARRIER_FREQUENCY = 1.0 / CARRIER_WAVELENGTH
INPUT_WIDTH = 0.065
TARGET_HALF_WIDTH = 0.24
TARGET_EDGE_WIDTH = 0.008
ANGLE_INCIDENCE = 0.24

DOMAIN_LENGTH = 1.00
Z_COUNT = 192
FREQUENCY_COUNT = 141
RBF_COUNT = 64

EPS = 1e-12
MIN_BRIGHTNESS = 1e-3
BRIGHTNESS_WEIGHT = 1e6
MIN_RBF_WIDTH = 0.004
DENSITY_SOFTNESS = 0.020


class GaussianPlasmaFrequency(torch.nn.Module):
    def __init__(self, z: torch.Tensor) -> None:
        super().__init__()
        centers = torch.linspace(z[0] + 0.03, z[-1] - 0.03, RBF_COUNT, dtype=z.dtype, device=z.device)
        normalized_centers = (centers - z[0]) / (z[-1] - z[0])
        self.center_raw = torch.nn.Parameter(torch.log(normalized_centers / (1.0 - normalized_centers)))
        self.width_raw = torch.nn.Parameter(inverse_softplus(torch.full_like(centers, 0.016) - MIN_RBF_WIDTH))
        self.height_raw = torch.nn.Parameter(inverse_softplus(initial_density_ratio(centers) / 0.70) + 2.0)
        self.background_raw = torch.nn.Parameter(inverse_softplus(torch.as_tensor(0.010, dtype=z.dtype, device=z.device)))

    def forward(self, z: torch.Tensor) -> torch.Tensor:
        centers = z[0] + (z[-1] - z[0]) * torch.sigmoid(self.center_raw)
        widths = MIN_RBF_WIDTH + torch.nn.functional.softplus(self.width_raw)
        heights = torch.nn.functional.softplus(self.height_raw)
        basis = torch.exp(-((z[:, None] - centers[None, :]) / widths[None, :]) ** 2)
        density_ratio = torch.nn.functional.softplus(self.background_raw) + torch.sum(basis * heights[None, :], dim=1)
        density_ratio = DENSITY_SOFTNESS * torch.nn.functional.softplus(density_ratio / DENSITY_SOFTNESS)
        density_ratio = edge_window(z) * density_ratio
        return CARRIER_FREQUENCY * torch.sqrt(density_ratio + EPS)


def inverse_softplus(values: torch.Tensor) -> torch.Tensor:
    return values + torch.log(-torch.expm1(-values))


def initial_density_ratio(z: torch.Tensor) -> torch.Tensor:
    ratio = torch.full_like(z, 0.040)
    centers = torch.as_tensor(
        [0.90, 1.18, 1.52, 1.82, 2.06, 2.28, 2.48, 2.66, 2.86, 3.08, 3.34, 3.62],
        dtype=z.dtype,
        device=z.device,
    ) / 4.0
    widths = torch.as_tensor(
        [0.12, 0.08, 0.07, 0.06, 0.05, 0.07, 0.06, 0.055, 0.07, 0.08, 0.09, 0.13],
        dtype=z.dtype,
        device=z.device,
    ) / 4.0
    heights = torch.as_tensor(
        [0.10, 0.16, 0.28, 0.60, 0.78, 0.62, 0.74, 0.52, 0.66, 0.44, 0.26, 0.18],
        dtype=z.dtype,
        device=z.device,
    )
    for center, width, height in zip(centers, widths, heights, strict=True):
        ratio = ratio + height * torch.exp(-((z - center) / width) ** 2)
    return ratio


def edge_window(z: torch.Tensor) -> torch.Tensor:
    left = torch.sigmoid((z - z[0] - 0.055) / 0.014)
    right = torch.sigmoid((z[-1] - z - 0.055) / 0.014)
    window = left * right
    return torch.where((z == z[0]) | (z == z[-1]), torch.zeros_like(window), window)


def make_grid() -> tuple[torch.Tensor, torch.Tensor]:
    z_inner = torch.linspace(0.0, DOMAIN_LENGTH, Z_COUNT, dtype=DTYPE, device=DEVICE)
    grid = torch.cat([z_inner[:1], z_inner, z_inner[-1:]])
    return z_inner, grid


def make_pulse() -> tuple[Pulse, torch.Tensor, torch.Tensor]:
    frequency_margin = TARGET_HALF_WIDTH + 8.0 * TARGET_EDGE_WIDTH
    f_min = CARRIER_FREQUENCY - frequency_margin
    f_max = CARRIER_FREQUENCY + frequency_margin
    frequencies = torch.linspace(f_min, f_max, FREQUENCY_COUNT, dtype=DTYPE, device=DEVICE)
    amplitudes = torch.exp(-((frequencies - CARRIER_FREQUENCY) / INPUT_WIDTH) ** 2).to(torch.complex128)
    angles = torch.full_like(frequencies, ANGLE_INCIDENCE)
    pulse = Pulse(
        fft=amplitudes,
        fftfreq=frequencies,
        angle_incidence=angles,
        polarization="TE",
    )
    return pulse, frequencies, amplitudes


def flat_top(frequencies: torch.Tensor) -> torch.Tensor:
    distance = torch.abs(frequencies - CARRIER_FREQUENCY)
    return torch.sigmoid((TARGET_HALF_WIDTH - distance) / TARGET_EDGE_WIDTH)


def solve_profile(
    model: GaussianPlasmaFrequency,
    z_inner: torch.Tensor,
    grid: torch.Tensor,
    pulse: Pulse,
) -> dict[str, torch.Tensor]:
    plasma_frequency = model(z_inner)
    frequency = torch.cat([plasma_frequency[:1], plasma_frequency, plasma_frequency[-1:]])
    plasma = Plasma(
        grid=grid,
        frequency=frequency,
        state_equation=drude_model(gamma=1e-4),
    )
    with warnings.catch_warnings():
        warnings.simplefilter("ignore", UserWarning)
        result = Solver(plasma=plasma, pulse=pulse).solve()
    return {
        "plasma_frequency": plasma_frequency,
        "permittivity": plasma._get_permittivity(pulse.fftfreq),
        "R": result.R,
        "T": result.T,
    }


def spectrum_terms(pulse: Pulse, frequencies: torch.Tensor, solution: dict[str, torch.Tensor]) -> dict[str, torch.Tensor]:
    input_power = torch.abs(pulse.fft) ** 2
    output_amplitude = pulse.fft * solution["T"]
    output_power = torch.abs(output_amplitude) ** 2
    normalized_output = output_power / (torch.max(output_power) + EPS)
    guide = flat_top(frequencies)
    brightness = torch.sum(output_power) / (torch.sum(input_power) + EPS)
    shape_loss = torch.mean((normalized_output - guide) ** 2)
    loss = 4.0 * shape_loss + BRIGHTNESS_WEIGHT * torch.relu(MIN_BRIGHTNESS - brightness) ** 2
    return {
        "loss": loss,
        "shape_loss": shape_loss,
        "brightness": brightness,
        "input_power": input_power,
        "output_power": output_power,
        "normalized_output": normalized_output,
        "guide": guide,
        "output_amplitude": output_amplitude,
    }


def reconstruct_time(frequencies: torch.Tensor, amplitudes: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
    time = torch.linspace(-18.0, 18.0, 900, dtype=frequencies.dtype, device=frequencies.device)
    phase = torch.exp(-2j * torch.pi * frequencies[:, None] * time[None, :])
    return time, torch.real(torch.sum(amplitudes[:, None] * phase, dim=0))


def to_cpu(tensor: torch.Tensor) -> torch.Tensor:
    return tensor.detach().cpu()


def plot_frame(
    frame_path: Path,
    history: list[dict[str, float]],
    z_inner: torch.Tensor,
    frequencies: torch.Tensor,
    pulse: Pulse,
    solution: dict[str, torch.Tensor],
    terms: dict[str, torch.Tensor],
) -> None:
    frame_path.parent.mkdir(parents=True, exist_ok=True)
    fig, axes = plt.subplots(2, 2, figsize=(7.2, 5.2), constrained_layout=True)

    axes[0, 0].plot([item["iteration"] for item in history], [item["loss"] for item in history], color="tab:blue")
    axes[0, 0].set_yscale("log")
    axes[0, 0].set_xlabel("iteration")
    axes[0, 0].set_ylabel("loss")
    axes[0, 0].grid(True, alpha=0.25)

    time, input_trace = reconstruct_time(frequencies, pulse.fft)
    _, output_trace = reconstruct_time(frequencies, terms["output_amplitude"])
    axes[0, 1].plot(to_cpu(time), to_cpu(input_trace), color="tab:blue", label="input")
    output_time_axis = axes[0, 1].twinx()
    output_time_axis.plot(to_cpu(time), to_cpu(output_trace), color="tab:orange", label="output")
    axes[0, 1].set_xlabel("time")
    axes[0, 1].set_ylabel("input field", color="tab:blue")
    output_time_axis.set_ylabel("output field", color="tab:orange")
    axes[0, 1].tick_params(axis="y", labelcolor="tab:blue")
    output_time_axis.tick_params(axis="y", labelcolor="tab:orange")
    axes[0, 1].grid(True, alpha=0.25)

    axes[1, 0].plot(to_cpu(z_inner), to_cpu(solution["plasma_frequency"]), color="tab:blue")
    axes[1, 0].axhline(CARRIER_FREQUENCY, color="tab:red", linestyle=":", linewidth=1.0, label="carrier frequency")
    axes[1, 0].set_xlabel("z")
    axes[1, 0].set_ylabel("plasma frequency")
    axes[1, 0].legend(frameon=False)
    axes[1, 0].grid(True, alpha=0.25)

    order = torch.argsort(frequencies)
    input_norm = terms["input_power"] / (torch.max(terms["input_power"]) + EPS)
    axes[1, 1].plot(to_cpu(frequencies[order]), to_cpu(input_norm[order]), color="tab:blue", label="input")
    output_spectrum_axis = axes[1, 1].twinx()
    output_spectrum_axis.plot(
        to_cpu(frequencies[order]),
        to_cpu(terms["normalized_output"][order]),
        color="tab:orange",
        label="output",
    )
    axes[1, 1].plot(to_cpu(frequencies[order]), to_cpu(terms["guide"][order]), color="tab:green", linestyle="--", label="guide")
    axes[1, 1].set_xlabel("frequency")
    axes[1, 1].set_ylabel("input normalized power", color="tab:blue")
    output_spectrum_axis.set_ylabel("output normalized power", color="tab:orange")
    axes[1, 1].set_ylim(-0.02, 1.05)
    output_spectrum_axis.set_ylim(-0.02, 1.05)
    axes[1, 1].tick_params(axis="y", labelcolor="tab:blue")
    output_spectrum_axis.tick_params(axis="y", labelcolor="tab:orange")
    axes[1, 1].grid(True, alpha=0.25)

    latest = history[-1]
    fig.suptitle(
        f"iteration {latest['iteration']}: "
        f"loss={latest['loss']:.4g}, "
        f"shape={latest['shape_loss']:.4g}, "
        f"brightness={latest['brightness']:.3g}",
        fontsize=10,
    )
    fig.savefig(frame_path, dpi=180)
    plt.close(fig)


def main() -> None:
    torch.manual_seed(37)
    frames_dir = Path(__file__).with_name("frames")
    z_inner, grid = make_grid()
    pulse, frequencies, _ = make_pulse()
    model = GaussianPlasmaFrequency(z_inner).to(device=DEVICE)
    optimizer = torch.optim.Adam(model.parameters(), lr=LEARNING_RATE)
    history: list[dict[str, float]] = []

    for iteration in range(ITERATIONS + 1):
        optimizer.zero_grad()
        solution = solve_profile(model, z_inner, grid, pulse)
        terms = spectrum_terms(pulse, frequencies, solution)
        if iteration < ITERATIONS:
            terms["loss"].backward()
            optimizer.step()
            solution = solve_profile(model, z_inner, grid, pulse)
            terms = spectrum_terms(pulse, frequencies, solution)

        history.append(
            {
                "iteration": float(iteration),
                "loss": float(terms["loss"].detach().cpu()),
                "shape_loss": float(terms["shape_loss"].detach().cpu()),
                "brightness": float(terms["brightness"].detach().cpu()),
            }
        )

        if iteration % SAVE_EVERY == 0 or iteration == ITERATIONS:
            plot_frame(
                frames_dir / f"iteration_{iteration:03d}.png",
                history,
                z_inner,
                frequencies,
                pulse,
                solution,
                terms,
            )


if __name__ == "__main__":
    main()