Source code for psc_sim.simulation.dispatcher

"""Unified access to IV and EIS simulation backends."""

from __future__ import annotations

from dataclasses import dataclass, field
from pathlib import Path

import numpy as np

from psc_sim.parameters import CellParameters
from psc_sim.simulation.backends import (
    EISBackend,
    IVBackend,
    LTspiceEISBackend,
    LTspiceIVBackend,
    LTspiceTRANBackend,
    PythonEISBackend,
    PythonIVBackend,
)
from psc_sim.simulation.defaults import DEFAULT_IV_SWEEP
from psc_sim.simulation.eis_grid import default_eis_freqs
from psc_sim.simulation.types import IVResult, SweepSpec, TRANResult, TranSpec




@dataclass
class SimulationDispatcher:
    params: CellParameters
    iv_backend: IVBackend = field(default_factory=PythonIVBackend)
    eis_backend: EISBackend = field(default_factory=PythonEISBackend)
    iv_spec: SweepSpec = DEFAULT_IV_SWEEP

    def iv_result(self, spec: SweepSpec | None = None) -> IVResult:
        s = spec or self.iv_spec
        return self.iv_backend.sweep_iv(self.params, s)

    def eis_impedance(self, freqs: np.ndarray) -> tuple[np.ndarray, np.ndarray]:
        return self.eis_backend.impedance(self.params, freqs)

    def ltspice_iv(
        self,
        spec: SweepSpec | None = None,
        *,
        ltspice_exe: str | Path | None = None,
        workdir: str | Path | None = None,
    ) -> IVResult:
        s = spec or self.iv_spec
        backend = LTspiceIVBackend(ltspice_exe=ltspice_exe, workdir=workdir)
        return backend.sweep_iv(self.params, s)

    def ltspice_eis(
        self,
        freqs: np.ndarray | None = None,
        *,
        ltspice_exe: str | Path | None = None,
        workdir: str | Path | None = None,
    ) -> tuple[np.ndarray, np.ndarray, np.ndarray]:
        grid = default_eis_freqs() if freqs is None else np.asarray(freqs, dtype=float)
        backend = LTspiceEISBackend(ltspice_exe=ltspice_exe, workdir=workdir)
        zre, zim = backend.impedance(self.params, grid)
        return grid, zre, zim

    def ltspice_tran(
        self,
        *,
        ltspice_exe: str | Path | None = None,
        workdir: str | Path | None = None,
        tran_spec: TranSpec | None = None,
    ) -> TRANResult:
        backend = LTspiceTRANBackend(
            ltspice_exe=ltspice_exe,
            workdir=workdir,
            tran_spec=tran_spec,
        )
        return backend.run(self.params)