# SPDX-FileCopyrightText: 2026 John C. Travers and contributors
#
# SPDX-License-Identifier: GPL-3.0-or-later
"""Own the optional TDSE runtime behind the provider-free response seam.
The provider's whole-electron currents and shadow dipoles remain unscaled SI
arrays. Source assembly applies the calibration and density separately. Resource
admission, call accounting, the archive and the validation of every returned record
live here; the provider itself runs in a worker process of its own
(``solprop.driver._source_worker``, with ``solprop.driver._pyroraptor`` as its
adapter, SEP 49), and solprop's process never imports it.
Two atomic models are bound (SEP 12, SEP 20): the single-active-electron (SAE)
response plan, and the helium haCC response plan, either from one of the
provider's qualified presets or from a custom basis. An atomic window drives the
atom over part of the lattice: the binding passes the retained samples multiplied
by the window's gate and pads the returned currents with zeros, so a
:class:`~solprop.translate.tdse.TDSEResult` always spans the full lattice.
"""
import contextlib
import dataclasses
import json
import math
import os
import warnings
from collections.abc import Mapping
from contextlib import ExitStack
from dataclasses import asdict, dataclass
from typing import Any, Literal, Protocol, cast, runtime_checkable
import numpy as np
from jaxtyping import Float
from solprop import exceptions
from solprop._process import thread_budget
from solprop.constants import AU_TIME, EPS0, C
from solprop.driver._pyroraptor import _PRESET_FIELDS, _admit_memory, check_info
from solprop.driver._source_worker import (
CONTROL_LIMIT,
Record,
RemoteError,
Worker,
WorkerError,
worker_environment,
)
from solprop.exceptions import (
BackendError,
MaterialError,
RunTerminated,
SolpropError,
StopRequested,
TDSEArchiveError,
TDSEDomainWarning,
TDSEUnqualifiedModelWarning,
)
from solprop.log import StopFlag
from solprop.nonlinear.tdse import (
HaccBasis,
TDSEBasis,
TDSECalibration,
TDSEExecution,
)
from solprop.output.tdse_archive import (
ArchiveOutputs,
TDSEArchiveWriter,
record_serializer,
writer_memory,
)
from solprop.translate.tdse import (
AtomicCallBudget,
AtomicCallLedger,
AtomicWindowRecord,
TDSEBinding,
TDSEDescriptor,
TDSEPurpose,
TDSEResult,
lattice_times,
waveform_identity,
)
__all__ = [
"DiagnosticSession",
"HaccPreset",
"basis_from_preset",
"estimate_binding_memory",
"hacc_basis",
"hacc_preset",
"prepare_binding",
"provider_revision",
]
_Preset = Literal["production", "high_accuracy", "strong_field"]
_HaccVariant = Literal["standard", "fast"]
# The adapter module the worker runs (SEP 49 section 6.5).
_ADAPTER = "solprop.driver._pyroraptor"
# Native classes never reach this process; what the worker returns are records.
_Native = Any
# The archive role of each evaluation purpose (SEP 20 section 5.8); a predictor or
# corrector of a rejected trial reads back as "rejected".
_ARCHIVE_ROLES: dict[str, str] = {
"startup": "startup",
"predictor": "predictor",
"corrected": "corrector",
"diagnostic": "diagnostic",
}
# A preset's calibration admission band: from 20 eV (hc / 20 eV = 62 nm), the
# presets' below-threshold bound, to 1.6 um, the band the 800 nm qualifying drive and
# run_short's grid occupy (SEP 20 section 5.1). Not an independent qualification.
_PRESET_WAVELENGTH_RANGE = (62e-9, 1.6e-6)
# The preset records a run's provenance carries (the provider's evidence apart).
_PRESET_RECORD = tuple(name for name in _PRESET_FIELDS if name != "source")
@runtime_checkable
class _Channel(Protocol):
"""The worker the binding talks to, a ``_source_worker.Worker``."""
identity: Mapping[str, object]
@property
def alive(self) -> bool:
"""Whether the worker may take a request."""
...
def request(self, kind: str, value: object = None) -> object:
"""Send one request and return its reply."""
...
def close(self) -> None:
"""End the worker between calls."""
...
def kill(self) -> None:
"""End the worker at once."""
...
[docs]
def provider_revision(build: Mapping[str, object]) -> str:
"""Return the provider source revision a build record names.
solprop pins no provider revision (SEP 49 section 6.5): it uses the installed
provider and records the revision its worker reports.
Parameters
----------
build : collections.abc.Mapping
A binding's or diagnostic session's build record (``_build_record``).
Returns
-------
str
The provider's reported source revision, or ``""`` for an empty record.
"""
native = build.get("native")
if not isinstance(native, Mapping):
return ""
return str(cast("Mapping[str, object]", native).get("source_revision", ""))
def _open_provider(
*, stop: StopFlag | None = None, limit: int = CONTROL_LIMIT
) -> _Channel:
"""Start a worker running the provider's adapter and wait for its handshake.
The worker gets the process's whole thread budget (SEP 49 section 6.6) and the
parent's environment less native's own OpenMP wait settings.
- ``stop``: the run's flag, polled while the worker is awaited.
- ``limit``: the bound on any frame, bytes.
"""
environment = worker_environment(
os.environ, {"PYRORAPTOR_THREADS": str(thread_budget().threads)}
)
worker = Worker(_ADAPTER, environment=environment, limit=limit, stop=stop)
try:
worker.start()
except RemoteError as error:
raise _local_error(error, "TDSE provider start failed") from error
return worker
def _local_error(error: RemoteError, context: str) -> SolpropError:
"""Return the worker's exception as this process raises it.
solprop's own exceptions keep their class and message; any other becomes a
:class:`BackendError` naming `context`, the message and its classification.
"""
module, _, name = error.type_name.rpartition(".")
kind = getattr(exceptions, name, None) if module == "solprop.exceptions" else None
if isinstance(kind, type) and issubclass(kind, SolpropError):
try:
local = kind(error.detail)
except TypeError:
pass
else:
local.add_note(
f"raised in the TDSE provider worker:\n{error.remote_traceback}"
)
return local
wrapped = BackendError(f"{context}: {error.detail} [{error.classification}]")
wrapped.add_note(f"the worker's traceback:\n{error.remote_traceback}")
return wrapped
def _build_record(
identity: Mapping[str, object], execution: TDSEExecution | None
) -> dict[str, object]:
"""Refuse an execution-space mismatch the worker reported; return its build.
No provider revision is refused: solprop uses the installed provider, whose
revision, build and execution space the record keeps (SEP 49 section 6.5).
- ``identity``: the handshake's provider identity: its build and distribution.
- ``execution``: required device family, or None when only resolving a preset.
- Returns: serializable native build and installed distribution identities.
"""
build = cast("Mapping[str, object]", identity["native"])
space = str(build["execution_space"]).lower()
devices = {"cpu": {"openmp", "serial", "threads"}, "cuda": {"cuda"}, "hip": {"hip"}}
if execution is not None and space not in devices[execution.device]:
raise BackendError(
f"TDSE requested device={execution.device!r}, but the native build uses "
f"{build['execution_space']!r}"
)
record: dict[str, object] = {
"native": dict(build),
"distribution": identity.get("distribution"),
}
if "direct_url" in identity:
record["direct_url"] = identity["direct_url"]
return record
def _query(kind: str, request: Mapping[str, object]) -> tuple[Any, str]:
"""Ask a short-lived worker one question (a preset, a basis) and end it.
- Returns: the reply, and the provider revision the worker reported.
"""
worker = _open_provider()
try:
revision = provider_revision(_build_record(worker.identity, None))
return worker.request(kind, dict(request)), revision
except RemoteError as error:
raise _local_error(error, "TDSE provider query failed") from error
finally:
worker.close()
[docs]
def basis_from_preset(target: str, preset: _Preset) -> TDSEBasis:
"""Resolve a provider preset into an immutable, serializable SI basis.
Parameters
----------
target : str
Atomic target, currently ``helium`` only.
preset : {'production', 'high_accuracy', 'strong_field'}
Name of the provider's measured basis and spectral cut.
Returns
-------
TDSEBasis
Fully resolved radial, angular, potential and spectral parameters.
The preset's atomic step is deliberately separate from this basis.
Raises
------
MaterialError
If the target or preset is unsupported.
BackendError
If the optional provider is missing.
"""
if target != "helium" or preset not in (
"production",
"high_accuracy",
"strong_field",
):
raise MaterialError("TDSE requires helium and a supported provider preset")
reply, revision = _query("basis_from_preset", {"target": target, "preset": preset})
fields = dict(reply)
provider_source = fields.pop("provider_source")
return TDSEBasis(**fields, source=f"pyroraptor@{revision}: {provider_source}")
# ------------------------------------------------------------------- haCC records
[docs]
@dataclass(frozen=True)
class HaccPreset:
"""Hold a provider's qualified helium haCC preset as solprop records (SEP 20).
Attributes
----------
basis : HaccBasis
Frozen copy of the preset's model and basis, with its provider identity.
variant : {'standard', 'fast'}
The provider's numerics variant: ``'standard'`` (cut 200 hartree, RK4 step
at most 0.0088 au) or ``'fast'`` (cut 125 hartree, step at most 0.02045 au,
about 2.3 times fewer RK4 steps per call), each with its own identity and
measured records.
step_margin : float
The step ceiling over the largest stable RK4 step the provider measured at
the preset's cut on the qualifying field: 0.62 (standard) or 0.9 (fast).
max_atomic_step : float
Largest RK4 step in s, ``basis.max_step_au`` converted with solprop's atomic
unit of time; pass it as ``TDSEGasResponse.max_atomic_step``.
accuracy : float
Largest relative change of the closed total current, ``max_t |dJ| / max_t
|J|``, across the preset's numerical ladders (lmax, finite-element order,
cut, step) on the qualifying waveform, up to 8 fs after its centre. For the
fast variant, the standard value plus the measured largest fast-to-standard
change (1.5e-4 of the peak), a triangle-inequality bound.
box_sensitivity : float
The same measure across the R0 ladder (80, 120 and 160 bohr).
gauge_sensitivity : float
The same measure across the Rg ladder (12, 20 and 28 bohr).
max_peak_intensity : float
Largest qualified peak intensity in W/m^2.
wavelength : float
Carrier wavelength of the qualifying drive in m.
window : float
Atomic-window duration the accuracy ladders drove, in s.
device_admissible : bool
Whether the provider's device route admits the preset's lattice.
identity : str
The provider's preset identity, equal to ``basis.preset``.
source : str
The provider's evidence record.
Notes
-----
The sensitivities are measured on one sampled configuration and waveform; they are
not error bounds on the helium response or on a propagated spectrum.
"""
basis: HaccBasis
variant: _HaccVariant
step_margin: float
max_atomic_step: float
accuracy: float
box_sensitivity: float
gauge_sensitivity: float
max_peak_intensity: float
wavelength: float
window: float
device_admissible: bool
identity: str
source: str
[docs]
def calibration(self) -> TDSECalibration:
"""Return the research calibration record documenting the preset's evidence.
Returns
-------
TDSECalibration
Research status, no uncertainty (the sensitivities are not certified
uncertainties), the 62 nm to 1.6 um admission band, and limitations
naming the window, the intensity bound, the shadow's finite-tail
qualification, the diagnostic-only bound and the measured
sensitivities with their metric and waveform.
"""
metric = (
"max_t |dJ| / max_t |J| of the closed total current on the run_short "
"peak-power waveform (8.0e14 W/cm^2, node 0) up to 8 fs after its centre"
)
return TDSECalibration(
self.identity,
"research",
self.basis.identity,
_PRESET_WAVELENGTH_RANGE,
source=self.source,
uncertainty=None,
limitations=(
f"qualified over an atomic window of {self.window:.4g} s; longer "
"windows accumulate about 1% of the peak per escape burst",
f"qualified up to a peak intensity of {self.max_peak_intensity:.4g} "
f"W/m^2 at a {self.wavelength:.4g} m carrier, spectrum below the "
"helium threshold",
"the order-one shadow equals the closed current's first order only up "
"to the states' amplitudes at the closure surface",
"neglected_current_bound is a loose diagnostic, not an accuracy bound",
f"accuracy {self.accuracy:.4g}, box sensitivity "
f"{self.box_sensitivity:.4g}, gauge sensitivity "
f"{self.gauge_sensitivity:.4g}: {metric}; not error bounds",
f"numerics variant {self.variant!r}: RK4 step at "
f"{self.step_margin:.2g} of the provider's measured stability limit "
"at the preset's cut; a stronger or different drive can move it",
"wavelength_range is an admission band, not independent qualification",
),
)
[docs]
def hacc_preset(ions: Literal[1, 2], variant: _HaccVariant = "standard") -> HaccPreset:
"""Resolve one of the provider's qualified helium haCC presets.
Parameters
----------
ions : {1, 2}
1 for haCC(1) (the He+ 1s channel), 2 for haCC(5) (He+ n <= 2).
variant : {'standard', 'fast'}, optional
The provider's numerics variant (``hacc.helium_preset(ions, variant)``).
``'fast'`` lowers the cut to 125 hartree and raises the RK4 step ceiling
to 0.02045 au, with the same lattice, box, gauge and validity domain: about
2.3 times fewer RK4 steps per call, and a closed current within 1.5e-4 of
the standard variant's peak on the qualifying field. It has its own
identity, so its own consistency cache entry.
Returns
-------
HaccPreset
SI records of the preset's model, accuracy and validity domain. No model
is prepared.
Raises
------
MaterialError
If ``ions`` is not the integer 1 or 2, or ``variant`` is not ``'standard'``
or ``'fast'``.
BackendError
If the optional provider is missing or has no haCC surface, or its
preset's variant disagrees with the request.
Examples
--------
>>> from solprop.driver.tdse import hacc_preset
>>> preset = hacc_preset(1, "fast") # doctest: +SKIP
>>> preset.basis.label, preset.basis.preset # doctest: +SKIP
('haCC(1) fast', 'hacc-helium-1-fast-v1-e38292e1b30729c1')
"""
if type(ions) is not int or ions not in (1, 2):
raise MaterialError("hacc_preset needs ions=1 (haCC(1)) or ions=2 (haCC(5))")
if not isinstance(variant, str) or variant not in ("standard", "fast"):
raise MaterialError("hacc_preset needs variant='standard' or variant='fast'")
reply, revision = _query("hacc_preset", {"ions": ions, "variant": variant})
preset = cast("Mapping[str, Any]", reply["preset"])
if preset["variant"] != variant:
raise BackendError(
f"the provider returned the {preset['variant']!r} haCC preset for the "
f"{variant!r} variant"
)
source = f"pyroraptor@{revision}: {preset['source']}"
basis = HaccBasis(
**cast("dict[str, Any]", reply["fields"]),
source=source,
preset=preset["identity"],
)
if basis.variant != variant:
raise BackendError(
f"the provider's {variant!r} preset identity {preset['identity']!r} does "
"not name that variant"
)
return HaccPreset(
basis=basis,
variant=variant,
step_margin=float(preset["step_margin"]),
max_atomic_step=basis.max_step_au * AU_TIME,
accuracy=float(preset["accuracy"]),
box_sensitivity=float(preset["box_sensitivity"]),
gauge_sensitivity=float(preset["gauge_sensitivity"]),
max_peak_intensity=float(preset["max_peak_intensity"]),
wavelength=float(preset["wavelength"]),
window=float(preset["window"]),
device_admissible=bool(preset["device_admissible"]),
identity=preset["identity"],
source=source,
)
[docs]
def hacc_basis(
atom: object,
basis: object,
gauge: object,
cut_energy: float,
max_step_au: float,
*,
source: str,
) -> HaccBasis:
"""Convert provider haCC objects into a custom, unqualified ``HaccBasis``.
The "accept any provider model" path (SEP 20 section 5.1): the record's identity
is derived from its full definition, and it is warned and recorded as unqualified
when a run prepares it. Build its calibration record yourself, with
``status="research"``, no uncertainty and the limitation that the basis is
unqualified.
Parameters
----------
atom : pyroraptor.hacc.Helium
The atom and its ion count (1 or 2).
basis : pyroraptor.hacc.Basis
Radial axis in m, irECS and angular limit.
gauge : pyroraptor.hacc.MixedGauge or pyroraptor.hacc.VelocityGauge
The gauge; the velocity gauge becomes ``gauge_radius=None``.
cut_energy : float
Spectral cut in J.
max_step_au : float
RK4 step ceiling in atomic units of time.
source : str
Why this basis was chosen; recorded in its identity.
Returns
-------
HaccBasis
Custom record with ``preset=None``.
Raises
------
MaterialError
For argon, an ion count other than 1 or 2, or any field the record cannot
represent.
BackendError
If the optional provider is missing or has no haCC surface.
"""
# The objects travel to the worker as records of their fields, and are checked
# there against the provider's own classes (SEP 49 section 6.5).
fields, _ = _query(
"hacc_basis",
{
"atom": atom,
"basis": basis,
"gauge": gauge,
"cut_energy": cut_energy,
"max_step_au": max_step_au,
},
)
return HaccBasis(**cast("dict[str, Any]", fields), source=source)
# ---------------------------------------------------------------- binding helpers
@dataclass(frozen=True)
class _Chunk:
"""Identify one stable contiguous radial batch.
Attributes
----------
start : int
Inclusive nonnegative row index.
stop : int
Exclusive row index, strictly greater than start.
"""
start: int
stop: int
@property
def rows(self) -> int:
"""Return the number of rows in this contiguous chunk.
- Returns: positive difference between the exclusive stop and inclusive start.
"""
return self.stop - self.start
def _atomic_samples(descriptor: TDSEDescriptor) -> int:
"""Return the samples the provider integrates over.
The atomic window's when there is one, else the lattice's.
"""
window = descriptor.atomic_window
return descriptor.samples if window is None else window.samples
def _estimated_memory(
descriptor: TDSEDescriptor,
chunks: tuple[_Chunk, ...],
native_outputs: str | None = None,
) -> dict[str, int]:
"""Estimate conservative data storage before allocating the native setup.
- ``native_outputs``: the provider's output selection when a full archive widens the
descriptor's; None takes ``descriptor.outputs``.
- Returns: named conservative atomic storage estimates in bytes; vendor setup
workspaces and allocator overhead are not hard bounded. The haCC provider offers
no estimate before preparation (SEP 20, C7), so only its host histories and the
transport copies are estimated; its model and plans are admitted from their
measured ledgers.
"""
samples = _atomic_samples(descriptor)
outputs = descriptor.outputs if native_outputs is None else native_outputs
histories = 48 if outputs == "full" else 20
transport = _transport_bytes(
descriptor, max(chunk.rows for chunk in chunks), outputs
)
host = {
"estimate.host_histories": 8 * histories * chunks[-1].stop * descriptor.samples,
# A call's copies on the pipe, in this process and in the worker (SEP 49).
"estimate.transport": 2 * transport,
}
basis = descriptor.response.basis
if not isinstance(basis, TDSEBasis):
return host
radial = sum(order - 1 for order in basis.radial_orders)
radial += (basis.infinite_order or 0) - 1
channels = basis.lmax + 1
sizes = tuple(sorted({chunk.rows for chunk in chunks}))
substeps = max(1, math.ceil(descriptor.spacing / descriptor.max_atomic_step))
stages = 3 * (samples - 1) * substeps + 1
return {
"estimate.shared_dense_setup": 16 * 12 * channels * radial * radial,
"estimate.plan_states": sum(
16 * 24 * (descriptor.shadow_order + 1) * rows * channels * radial
for rows in sizes
),
"estimate.plan_stages": sum(8 * stages * (4 * rows + 16) for rows in sizes),
**host,
}
def _json_default(value: object) -> object:
"""Convert a provider record or array into JSON-compatible containers."""
if isinstance(value, Record):
# A provider record from the worker: its fields, as the dataclass gave them.
return dict(vars(value))
if isinstance(value, np.ndarray):
return value.tolist()
if isinstance(value, np.generic):
return value.item()
raise TypeError(f"Unsupported TDSE metadata value {type(value).__name__}")
def _encoded(value: object) -> str:
"""Encode finite metadata in a deterministic JSON representation.
- ``value``: nested public records, arrays and finite primitive values.
"""
return json.dumps(value, default=_json_default, sort_keys=True, allow_nan=False)
def _failure_safe(value: object) -> object:
"""Preserve nonfinite failed monitors as explicit strings in valid JSON."""
encoded = json.dumps(value, default=_json_default, sort_keys=True)
return json.loads(encoded, parse_constant=lambda token: f"nonfinite:{token}")
def _input_diagnostics(
fields: Float[np.ndarray, "n_nodes n_to"], spacing: float
) -> object:
"""Measure field endpoints and area without changing the supplied waveform.
- ``fields``: finite SI fields in V/m, shape (n_nodes, n_to).
- ``spacing``: positive sample spacing in s.
- Returns: Per-node peak and endpoint fields in V/m and pulse area in V s/m.
"""
return {
"peak_field_v_per_m": np.max(np.abs(fields), axis=1),
"leading_field_v_per_m": fields[:, 0].copy(),
"trailing_field_v_per_m": fields[:, -1].copy(),
"leading_eight_peak_v_per_m": np.max(np.abs(fields[:, :8]), axis=1),
"trailing_eight_peak_v_per_m": np.max(np.abs(fields[:, -8:]), axis=1),
"pulse_area_v_s_per_m": np.trapezoid(fields, dx=spacing, axis=1),
}
def _drive_diagnostics(
drive: Float[np.ndarray, "n_nodes n_window"], spacing: float
) -> object:
"""Record the tapered drive's area and vector-potential endpoint (SEP 20, 4.6).
Multiplying by the gate need not preserve a zero-area drive; the area is
recorded, never removed.
- ``drive``: the exact gated field passed to the provider, V/m.
- ``spacing``: sample spacing in s.
- Returns: Per-node area in V s/m, which is minus the vector potential at ``t_e``
when ``A(t_s) = 0``, and the endpoint fields in V/m.
"""
area = np.trapezoid(drive, dx=spacing, axis=1)
return {
"drive_area_v_s_per_m": area,
"drive_final_vector_potential_v_s_per_m": -area,
"drive_leading_field_v_per_m": drive[:, 0].copy(),
"drive_trailing_field_v_per_m": drive[:, -1].copy(),
}
def _row_chunks(execution: TDSEExecution, n_nodes: int) -> tuple[_Chunk, ...]:
"""Partition radial rows into stable contiguous batches."""
if type(n_nodes) is not int or n_nodes < 1:
raise BackendError("TDSE binding requires a positive integer n_nodes")
return tuple(
_Chunk(start, min(start + execution.max_batch_rows, n_nodes))
for start in range(0, n_nodes, execution.max_batch_rows)
)
def _classify(error: BaseException) -> str:
"""Classify an evaluation's failure for its record (SEP 20, 5.6).
- Returns: the worker's classification of a provider exception (``native``,
``interface``, ``contract`` or ``unexpected``); ``worker`` for a worker that
exited or broke the protocol; for this process's own validation, ``contract``
for a solprop refusal, ``interface`` for a ``ValueError`` or ``TypeError`` and
``unexpected`` otherwise.
"""
if isinstance(error, RemoteError):
return error.classification
if isinstance(error, WorkerError):
return "worker"
if isinstance(error, SolpropError):
return "contract"
if isinstance(error, ValueError | TypeError):
return "interface"
return "unexpected"
def _described(error: BaseException) -> tuple[str, str]:
"""Return an error's type name and message, the worker's own for a remote one."""
if isinstance(error, RemoteError):
return error.type_name.rpartition(".")[2], error.detail
return type(error).__name__, str(error)
def _frame_limit(descriptor: TDSEDescriptor, chunks: tuple[_Chunk, ...]) -> int:
"""Bound one worker frame by the prepared shapes (SEP 49 section 6.5).
The largest array either way is a chunk's shadow record, three float64 orders
of every row on the lattice; four orders' worth, twice over, is the bound.
"""
rows = max(chunk.rows for chunk in chunks)
return max(CONTROL_LIMIT, 2 * 4 * 8 * rows * descriptor.samples)
def _transport_bytes(descriptor: TDSEDescriptor, rows: int, outputs: str) -> int:
"""Return one process's transport copies of the largest chunk's call.
The drive twice (its frame is built whole before it is written) and the result
once (its arrays are read into their own buffers), in bytes.
"""
histories = 48 if outputs == "full" else 20
return 8 * rows * _atomic_samples(descriptor) * (2 + histories)
def _peak_intensity(fields: Float[np.ndarray, "n_nodes n_samples"]) -> list[float]:
"""Return each row's peak cycle-averaged intensity ``eps0 c E_max^2 / 2``.
- ``fields``: real fields in V/m.
- Returns: peak intensities in W/m^2, one per row.
"""
peak = np.max(np.abs(fields), axis=1)
return [float(value) for value in 0.5 * EPS0 * C * peak**2]
[docs]
def estimate_binding_memory(
descriptor: TDSEDescriptor,
execution: TDSEExecution,
n_nodes: int,
*,
archive: ArchiveOutputs | None = None,
) -> Mapping[str, int]:
"""Estimate atomic storage without importing the optional provider.
Parameters
----------
descriptor : TDSEDescriptor
Resolved basis, fixed time lattice and requested response products.
execution : TDSEExecution
Stable maximum chunk size, determining the persistent plan shapes.
n_nodes : int
Positive total number of radial rows.
archive : {'full', 'required'} or None, optional
The output selection of a raw archive, whose writer buffers are added and
whose full selection widens the provider histories; None without one.
Returns
-------
mapping of str to int
Conservative atomic host and device data estimates in bytes. Optical
host storage is excluded and must be added for a shared run budget.
Vendor allocation overhead and setup workspaces are not hard bounded.
A haCC model has no estimate before preparation; its measured ledgers are
admitted after preparation.
Raises
------
BackendError
If the total row count is invalid.
"""
chunks = _row_chunks(execution, n_nodes)
estimate = _estimated_memory(
descriptor, chunks, "full" if archive == "full" else None
)
if archive is not None:
estimate.update(writer_memory(rows=chunks[0].rows, samples=descriptor.samples))
return estimate
[docs]
def prepare_binding(
descriptor: TDSEDescriptor,
execution: TDSEExecution,
n_nodes: int,
*,
defer_plans: bool = False,
archive: TDSEArchiveWriter | None = None,
stop: StopFlag | None = None,
) -> TDSEBinding:
"""Prepare persistent, stably chunked atomic plans for one response lattice.
The provider runs in a worker process the binding starts here and owns until it
is closed (SEP 49): every wait on it is cancelled by Ctrl-C, SIGTERM and `stop`.
Parameters
----------
descriptor : TDSEDescriptor
Resolved scientific description and fixed fine time lattice in SI.
execution : TDSEExecution
Device, memory limit, stable maximum row count and atomic-call limit.
n_nodes : int
Positive total number of radial rows; fixed for the binding's lifetime.
defer_plans : bool, optional
For a haCC model, prepare and admit the model only, and build the
propagation plans when ``build_plans()`` is called. The run uses this so the
consistency check's temporary one-row plan is built and closed before the
propagation plans exist, and peak memory is the larger of the two, not their
sum. Ignored for SAE, whose plans are always built here.
archive : TDSEArchiveWriter or None, optional
An open archive session (SEP 20 section 5.8). Every provider evaluation is
written to it, input first and raw result before validation; a full
archive builds the plans with the provider's full record. The caller owns
and closes it. Its buffers are admitted with the binding's memory.
stop : solprop.log.StopFlag or None, optional
The run's stop flag, polled while an atomic call is awaited; setting it ends
the call at once, the worker killed, with
:class:`~solprop.exceptions.StopRequested`.
Returns
-------
TDSEBinding
Owning runtime adapter. Close it when the propagation ends. Evaluation
requires a successful reservation of all logical calls in the trial.
Raises
------
BackendError
If the provider, device, scientific identity or prepared state is invalid.
DriverError
If estimated or measured persistent data exceed the declared byte limit.
"""
if type(n_nodes) is not int or n_nodes < 1:
raise BackendError("TDSE binding requires a positive integer n_nodes")
if type(defer_plans) is not bool:
raise BackendError("TDSE defer_plans must be a bool")
if archive is not None and not isinstance(archive, TDSEArchiveWriter):
raise BackendError("TDSE archive must be an open TDSEArchiveWriter")
return _ProviderBinding(
descriptor,
execution,
n_nodes,
defer_plans=defer_plans,
archive=archive,
stop=stop,
)
# ------------------------------------------------------------------ haCC model
@dataclass
class _HaccView:
"""The binding's record of the haCC model the worker prepared (SEP 20, C1).
Attributes
----------
basis : HaccBasis
The model's record.
outputs : {'full', 'total_current'} or None
The plans' native output mode; None for a custom mixed-gauge basis before
its first plan decides it.
closure_known : bool or None
Whether the plans carry the escaped-electron closure; None until known.
preset : mapping or None
The verified preset's records (its identity, variant, accuracy and domain), or
None for a custom basis.
"""
basis: HaccBasis
outputs: str | None
closure_known: bool | None
preset: Mapping[str, Any] | None
@classmethod
def from_reply(cls, basis: HaccBasis, reply: Mapping[str, Any]) -> "_HaccView":
"""Return the view of the model a ``prepare`` reply describes."""
return cls(basis, reply["outputs"], reply["closure"], reply["preset"])
def update(self, reply: Mapping[str, Any]) -> None:
"""Take the output mode and closure a plan decided, from a reply's record."""
self.outputs, self.closure_known = reply["outputs"], reply["closure"]
@property
def closure(self) -> bool:
"""Whether the plans carry the escaped-electron closure."""
return self.closure_known is True
def check_info(
self, raw: _Native, *, rows: int, samples: int, spacing: float
) -> None:
"""Require a result's ``info`` to describe exactly this model."""
check_info(
raw.info,
rows=rows,
samples=samples,
spacing=spacing,
outputs=self.outputs,
cut_energy=self.basis.cut_energy,
closure=self.closure,
preset=self.preset is not None,
model_match=_matches(raw),
)
def record(self) -> dict[str, object]:
"""Return the model's provenance for run metadata.
- Returns: label, preset identity or custom flag, qualification, closure, the
preset's accuracy and domain, and the complement policy.
"""
preset = self.preset
return {
"label": self.basis.label,
"basis_identity": self.basis.identity,
"preset": self.basis.preset,
"custom": self.basis.preset is None,
"qualified": self.basis.qualified,
"closure": self.closure,
"native_outputs": self.outputs,
"complement": "not_applicable",
"preset_record": None
if preset is None
else {name: preset[name] for name in _PRESET_RECORD},
}
def _returned(reply: object) -> _Native:
"""Return an atomic call's record, typed by the worker's name for it.
The record's ``record_type`` becomes ``sae.grid_response``, ``hacc.total_current_
response`` and so on, or ``other:<type>``; the worker's comparison of a haCC
result's atom, basis and gauge with the model rides beside its fields.
"""
answer = cast("Mapping[str, Any]", reply)
raw = answer["raw"]
if not isinstance(raw, Record):
# Nothing, or an object without fields, which the worker sends as its repr.
raise BackendError(f"the TDSE provider returned no result record: {raw!r}")
raw.computed["provider_type"] = raw.record_type.rpartition(".")[2]
raw.record_type = answer["kind"]
raw.computed["model_match"] = answer["model_match"]
return raw
def _kind(raw: object) -> str:
"""Return a returned record's type name (:func:`_returned`)."""
return raw.record_type if isinstance(raw, Record) else "other"
def _matches(raw: object) -> bool:
"""Whether the worker found a haCC result's model to be the prepared one."""
return isinstance(raw, Record) and raw.computed.get("model_match") is True
def _fields_of(value: object) -> dict[str, Any]:
"""Return a returned record's fields."""
return dict(vars(value))
def _component(shadow: _Native, quantity: str, order: int) -> np.ndarray:
"""Return one shadow component, as the provider's ``component`` returned it.
The worker evaluated the provider's method (``_pyroraptor``); a component it
refused raises here as the provider raised it.
"""
# The worker evaluated every quantity at each order the shadow reports, and the
# binding asks only for orders it has checked against those labels.
found = cast("Mapping[str, Any]", shadow.components)[f"{quantity}:{order}"]
if isinstance(found, Mapping):
refusal = cast("Mapping[str, str]", found)
if refusal["type"] == "ValueError":
raise ValueError(refusal["error"])
raise RuntimeError(f"{refusal['type']}: {refusal['error']}")
return cast("np.ndarray", found)
# --------------------------------------------------------------------- binding
class _ProviderBinding:
"""Keep the optional runtime private to the driver and expose only the seam.
- ``defer_plans``: for haCC, build the propagation plans only in
:meth:`build_plans`.
"""
def __init__(
self,
descriptor: TDSEDescriptor,
execution: TDSEExecution,
n_nodes: int,
*,
defer_plans: bool = False,
archive: TDSEArchiveWriter | None = None,
stop: StopFlag | None = None,
) -> None:
self._archive = archive
# A full archive widens the provider's selection, never the descriptor's,
# which keeps naming the products the source policy needs (SEP 20, 6.2).
self._native_outputs: str = (
"full"
if archive is not None and archive.outputs == "full"
else descriptor.outputs
)
self._trial: int | None = None
self._trials_opened = 0
self._trial_evaluations: list[int] = []
self._defer_plans = defer_plans
self._descriptor = descriptor
self._execution = execution
self._n_nodes = n_nodes
self._chunks = _row_chunks(execution, n_nodes)
self._budget = AtomicCallBudget(execution.max_atomic_calls, len(self._chunks))
self._stack = ExitStack()
self._plan_records: dict[str, object] = {}
self._closed = False
self._cached: TDSEResult | None = None
self._cached_fields: Float[np.ndarray, "n_nodes n_to"] | None = None
self._diagnostic_cache_hits = 0
self._last_call: dict[str, object] = {}
self._failures: list[object] = []
self._hacc: _HaccView | None = None
self._window: AtomicWindowRecord | None = descriptor.atomic_window
self._gate = descriptor.source_gate()
self._samples = _atomic_samples(descriptor)
self._domain_reasons: dict[str, object] = {}
self._domain_flags: list[object] = []
self._consistency: object = None
self._estimate = _estimated_memory(
descriptor, self._chunks, self._native_outputs
)
_admit_memory(
self._estimate,
execution.memory_limit,
)
self._memory: dict[str, int] = (
{}
if archive is None
else archive.memory_ledger(
rows=self._chunks[0].rows, samples=descriptor.samples
)
)
# The worker is the binding's from here to close: its stack ends it.
self._stop = stop
self._worker = _open_provider(
stop=stop, limit=_frame_limit(descriptor, self._chunks)
)
self._stack.callback(self._worker.close)
try:
self._build = _build_record(self._worker.identity, execution)
self._prepare()
_admit_memory(
self._memory,
execution.memory_limit,
)
except BaseException as error:
self._worker.kill()
self._stack.close()
raise _preparation_error(error) from error
@property
def native_outputs(self) -> str | None:
"""Return the provider's output selection for this binding's plans.
- Returns: the descriptor's selection, or ``"full"`` for a full archive; None
for a custom haCC basis before its first plan decides it.
"""
if self._hacc is not None:
return self._hacc.outputs
return self._native_outputs
# -- preparation ---------------------------------------------------------------
def _prepare(self) -> None:
"""Prepare the model and its plans in the worker; admit the measured memory."""
basis = self._descriptor.response.basis
reply = cast(
"Mapping[str, Any]",
self._worker.request("prepare", self._prepare_request(basis)),
)
self._absorb(reply)
if isinstance(basis, HaccBasis):
self._hacc = _HaccView.from_reply(basis, reply["hacc"])
if basis.preset is not None:
self._preset_domain_at_start(self._hacc)
if not self._defer_plans:
self._warn_custom()
# Include the simultaneously live provider histories, normalized chunks,
# concatenated result and exact-input diagnostic cache in host admission, and
# a call's copies on the pipe in each process (SEP 49 section 6.5).
self._memory["host.result_and_cache_reserve"] = self._estimate[
"estimate.host_histories"
]
transport = _transport_bytes(
self._descriptor,
max(chunk.rows for chunk in self._chunks),
self._native_outputs,
)
self._memory["transport.solprop"] = transport
self._memory["transport.worker"] = transport
def _prepare_request(self, basis: TDSEBasis | HaccBasis) -> dict[str, object]:
"""Return the plain records the worker builds the model and plans from."""
descriptor, execution = self._descriptor, self._execution
return {
"model": "hacc" if isinstance(basis, HaccBasis) else "sae",
"basis": asdict(basis),
"rows": sorted({chunk.rows for chunk in self._chunks}),
"samples": self._samples,
"spacing": float(descriptor.spacing),
"max_atomic_step": descriptor.max_atomic_step,
"shadow_order": descriptor.shadow_order,
"complement": execution.complement,
"outputs": self._native_outputs,
"archive_outputs": "required"
if self._archive is None
else self._archive.outputs,
"initial": (descriptor.initial_l, descriptor.initial_index),
"growing": descriptor.growing,
"growth_tolerance": descriptor.growth_tolerance,
"window_duration": descriptor.window_duration,
"accept_removed_component": descriptor.accept_removed_component,
"device": execution.device,
"memory": dict(self._memory),
"memory_limit": execution.memory_limit,
"defer_plans": self._defer_plans,
}
def _ensure_worker(self) -> None:
"""Start and prepare a fresh worker when the last one failed (SEP 49, 6.5).
A worker is never used again after a failure; the binding's records, budget
and archive are its own and carry over. The new worker must report the same
provider build.
"""
if self._worker.alive:
return
self._worker = _open_provider(
stop=self._stop, limit=_frame_limit(self._descriptor, self._chunks)
)
self._stack.callback(self._worker.close)
basis = self._descriptor.response.basis
try:
if _build_record(self._worker.identity, self._execution) != self._build:
raise BackendError(
"the TDSE provider's build changed between its workers"
)
request = self._prepare_request(basis)
request["defer_plans"] = not self.plans_built
self._absorb(
cast("Mapping[str, Any]", self._worker.request("prepare", request))
)
except BaseException as error:
self._worker.kill()
raise _preparation_error(error) from error
def _absorb(self, reply: Mapping[str, Any]) -> None:
"""Take the worker's measured memory and plan records."""
self._memory.update(cast("Mapping[str, int]", reply["memory"]))
self._plan_records = dict(cast("Mapping[str, object]", reply["plans"]))
if self._hacc is not None and reply["hacc"] is not None:
self._hacc.update(reply["hacc"])
def _warn_custom(self) -> None:
"""Warn that a custom basis is unqualified, once its plans know the closure."""
model = self._hacc
assert model is not None
basis = model.basis
if basis.preset is not None:
return
missing = (
""
if model.closure
else "; its total current lacks the escaped-electron closure, so the "
"source uses the in-box first order and a total that omits the "
"momentum carried out through R0"
)
warnings.warn(
f"{basis.label} is a custom haCC basis with no qualified accuracy "
f"and no preset domain; it is recorded as unqualified{missing}",
TDSEUnqualifiedModelWarning,
stacklevel=5,
)
@property
def plans_built(self) -> bool:
"""Return whether the propagation plans exist.
- Returns: False only for a haCC binding prepared with ``defer_plans`` before
:meth:`build_plans`.
"""
return self._hacc is None or bool(self._plan_records)
def build_plans(self) -> None:
"""Build and admit a deferred haCC binding's propagation plans.
Call once, after any consistency diagnostic, whose temporary plan is then
already closed. A failure closes every resource, as at preparation.
"""
self._ensure_open()
if self.plans_built:
raise BackendError("the TDSE binding's plans are already built")
try:
reply = self._worker.request(
"build_plans",
{
"memory": dict(self._memory),
"memory_limit": self._execution.memory_limit,
},
)
self._absorb(cast("Mapping[str, Any]", reply))
_admit_memory(
self._memory,
self._execution.memory_limit,
)
except BaseException as error:
self._worker.kill()
self.close()
raise _preparation_error(error) from error
self._warn_custom()
def _preset_domain_at_start(self, model: _HaccView) -> None:
"""Flag a preset run's duration or carrier beyond its qualification."""
preset = model.preset
assert preset is not None
configuration = json.loads(self._descriptor.configuration_json)
lattice = configuration["lattice"]
duration = lattice["atomic_duration"]
if duration > preset["window"]:
self._flag_domain(
"duration",
{"atomic_duration": duration, "preset_window": preset["window"]},
f"the atomic duration {duration:.6g} s exceeds the preset's tested "
f"window {preset['window']:.6g} s",
)
carrier = lattice["reference_wavelength"]
if not math.isclose(carrier, preset["wavelength"], rel_tol=1e-12):
self._flag_domain(
"carrier",
{
"reference_wavelength": carrier,
"preset_wavelength": preset["wavelength"],
},
f"the reference wavelength {carrier:.6g} m differs from the preset's "
f"qualifying carrier {preset['wavelength']:.6g} m",
)
def _flag_domain(self, reason: str, record: object, message: str) -> None:
"""Record a domain reason and warn once per reason per run."""
if reason in self._domain_reasons:
return
self._domain_reasons[reason] = record
assert self._hacc is not None
warnings.warn(
f"{self._hacc.basis.label} preset outside its qualified domain: {message}; "
"the run proceeds and the record is flagged",
TDSEDomainWarning,
stacklevel=4,
)
# -- accounting ----------------------------------------------------------------
@property
def ledger(self) -> AtomicCallLedger:
"""Return attempted-call accounting, including failed evaluations."""
return self._budget.snapshot()
def reserve(self, logical_calls: int) -> None:
"""Reserve every chunk of a complete trial before its first atomic call.
A successful reservation opens the next optical trial, whose evaluations
the archive labels with its identifier (SEP 20 section 5.8).
"""
self._ensure_open()
self._ensure_plans()
self._budget.reserve(logical_calls)
self._trial = self._trials_opened
self._trials_opened += 1
self._trial_evaluations = []
def record_trial(self, *, accepted: bool) -> None:
"""Record the stepper's decision on the latest trial in the archive."""
if self._trial is None:
raise BackendError("a TDSE trial outcome needs a reserved trial")
if self._archive is not None:
self._archive.record_trial(
self._trial, accepted=accepted, evaluations=self._trial_evaluations
)
def release_reservation(self) -> None:
"""Release unused capacity while retaining every attempted-call charge."""
self._budget.release_reservation()
def restore_ledger(self, ledger: AtomicCallLedger) -> None:
"""Restore prior attempts before performing work in a resumed binding."""
self._ensure_open()
self._budget.restore(ledger)
def _ensure_open(self) -> None:
"""Reject use after native resources have been closed."""
if self._closed:
raise BackendError("TDSE binding is closed")
def _ensure_plans(self) -> None:
"""Reject propagation before a deferred binding has built its plans."""
if not self.plans_built:
raise BackendError(
"the TDSE binding's propagation plans are not built; call build_plans"
)
# -- evaluation ----------------------------------------------------------------
def _drive(
self, fields: Float[np.ndarray, "n_nodes n_to"], chunk: _Chunk
) -> Float[np.ndarray, "n_rows n_samples"]:
"""Return the exact float64 field passed to the provider for one chunk.
- ``fields``: Full-lattice nodal fields in V/m.
"""
rows = fields[chunk.start : chunk.stop]
window = self._window
if window is None:
return rows
assert self._gate is not None
span = slice(window.first, window.last + 1)
return np.ascontiguousarray(rows[:, span] * self._gate[span])
def _pad(
self, values: Float[np.ndarray, "n_rows n_samples"]
) -> Float[np.ndarray, "n_rows n_to"]:
"""Place a window's currents on the full lattice with zeros outside."""
window = self._window
if window is None:
return values
full = np.zeros((values.shape[0], self._descriptor.samples), dtype=values.dtype)
full[:, window.first : window.last + 1] = values
return full
def _pad_valid(self, valid: np.ndarray) -> np.ndarray:
"""Place a window's validity on the full lattice; undriven samples are valid."""
window = self._window
if window is None:
return valid
# The provider marks a row invalid at and after its first failure, and the
# padded samples after the window would then read as valid. That cannot hide
# a failure: a failed outcome raises in _normalize and any invalid sample
# inside the window raises in TDSEResult, so this array reaches the source
# only when every driven sample is valid. Outside the window the atom is not
# driven and its padded current is exactly zero.
full = np.ones((valid.shape[0], self._descriptor.samples), dtype=bool)
full[:, window.first : window.last + 1] = valid
return full
def evaluate( # noqa: PLR0912, PLR0915 - one evaluation's complete sequence
self,
fields: Float[np.ndarray, "n_nodes n_to"],
*,
z: float,
purpose: TDSEPurpose,
density: float | None = None,
) -> TDSEResult:
"""Evaluate unchanged node fields from the prepared ground state.
With an archive, the evaluation's input is written before its first
provider call and each chunk's raw result as soon as it returns, before
validation; a write failure stops the evaluation before the next provider
call (SEP 20 section 5.8).
- ``fields``: finite float64 fields in V/m, shape (n_nodes, n_to).
- ``z``: finite optical position in m for provenance and failure reports.
- ``purpose``: reason for the source evaluation; only exact diagnostics may
reuse data.
- ``density``: finite nonnegative number density at `z` in m^-3, an archive
label.
- Returns: owning, read-only, unscaled microscopic SI currents and shadow
dipoles.
"""
self._ensure_open()
self._ensure_plans()
identity = waveform_identity(fields)
if fields.shape != (self._n_nodes, self._descriptor.samples):
raise BackendError(
"TDSE fields changed their prepared row count or sample count"
)
if not math.isfinite(z) or purpose not in (
"startup",
"predictor",
"corrected",
"diagnostic",
):
raise BackendError(
"TDSE evaluation needs a finite z and a declared purpose"
)
if density is not None and (
not isinstance(density, float | int)
or isinstance(density, bool)
or not math.isfinite(density)
or density < 0
):
raise BackendError("TDSE density label must be finite and nonnegative")
if (
purpose == "diagnostic"
and self._cached is not None
and identity == self._cached.input_identity
and self._cached_fields is not None
and np.array_equal(fields, self._cached_fields)
):
self._diagnostic_cache_hits += 1
return self._cached
self._ensure_worker()
archive = self._archive
if archive is not None:
archive.ensure_writable()
drives = [self._drive(fields, chunk) for chunk in self._chunks]
domain = self._call_domain(fields, drives, z=z, purpose=purpose)
evaluation = self._archive_begin(
fields,
drives,
labels={
"trial_id": self._trial,
"role": _ARCHIVE_ROLES[purpose],
"purpose": purpose,
"z": float(z),
"density": None if density is None else float(density),
"input_identity": identity,
},
)
try:
self._budget.begin_evaluation(purpose=purpose)
except BackendError as error:
self._archive_failure(evaluation, None, error, "contract")
raise
self._last_call = {
"input_identity": identity,
"configuration_identity": self._descriptor.identity,
"z": z,
"purpose": purpose,
"input_diagnostics": _input_diagnostics(fields, self._descriptor.spacing),
"chunks": [],
}
if self._window is not None:
self._last_call["drive_diagnostics"] = _drive_diagnostics(
np.concatenate(drives, axis=0), self._descriptor.spacing
)
if domain is not None:
self._last_call["domain"] = domain
results: list[TDSEResult] = []
chunk_records: list[object] = []
self._last_call["chunks"] = chunk_records
for index, (chunk, drive) in enumerate(zip(self._chunks, drives, strict=True)):
# A failed write raises before this point, so no provider call follows
# one (SEP 20 section 5.8); the writer refuses any later evaluation.
self._budget.charge(purpose=purpose, z=z)
raw: _Native = None
archived: TDSEArchiveError | None = None
try:
# Public evaluate always resets every row to the prepared state.
raw = _returned(
self._worker.request(
"evaluate", {"rows": chunk.rows, "drive": drive}
)
)
# Captured as returned, before any validation (SEP 20, 5.8).
archived = self._archive_result(evaluation, index, raw)
self._validate_diagnostics(raw, chunk.rows, self._samples)
chunk_records.append(self._native_record(raw, chunk))
results.append(self._normalize(raw, chunk, identity))
except (KeyboardInterrupt, RunTerminated, StopRequested) as stopping:
# A user's stop during the call: the worker is already killed, the
# call stays charged, and it is no provider failure (SEP 49, 6.7).
self._budget.release_reservation()
self._archive_failure(evaluation, index, stopping, "interrupted")
raise
except Exception as error:
# Whatever failed, this worker is not used again (SEP 49, 6.5): the
# next evaluation starts a fresh one.
self._worker.kill()
self._budget.record_failure()
classification = _classify(error)
name, detail = _described(error)
failure = {
"z": z,
"row_start": chunk.start,
"row_stop": chunk.stop,
"input_identity": identity,
"classification": classification,
"error": f"{name}: {detail}",
"native": None
if raw is None
else self._failed_native_record(raw, chunk),
}
self._failures.append(_failure_safe(failure))
self._budget.release_reservation()
wrapped = BackendError(
f"TDSE failure at z={z:.17g} m, "
f"rows [{chunk.start}, {chunk.stop}), "
f"input={identity}: {detail} [{classification}]"
)
_attach_remote(wrapped, error)
if archived is not None:
wrapped.add_note(
f"the TDSE archive could not store this result: {archived}"
)
self._archive_failure(
evaluation, index, error, classification, note=wrapped
)
raise wrapped from error
if archived is not None:
self._budget.release_reservation()
raise archived
result = self._combine(results, identity)
self._cached = result
self._cached_fields = np.array(fields, dtype=np.float64, order="C", copy=True)
self._cached_fields.setflags(write=False)
if archive is not None and evaluation is not None:
try:
archive.finish_evaluation(evaluation, "returned")
except TDSEArchiveError:
self._budget.release_reservation()
raise
return result
# -- archive -------------------------------------------------------------------
def _archive_begin(
self,
fields: Float[np.ndarray, "n_nodes n_to"],
drives: list[Float[np.ndarray, "_ n_samples"]],
*,
labels: Mapping[str, object],
) -> int | None:
"""Write an evaluation's labels and complete input before any provider call.
- ``fields``: the full-lattice nodal fields, V/m.
- ``drives``: the exact provider input of each chunk, V/m.
- Returns: the archive's evaluation identifier, or None without an archive.
"""
archive = self._archive
if archive is None:
return None
descriptor = self._descriptor
window = self._window
origin = float(lattice_times(descriptor.samples, descriptor.spacing)[0])
identifier = archive.begin_evaluation(
{
**labels,
"configuration_identity": descriptor.identity,
"basis_identity": descriptor.response.basis.identity,
"model": descriptor.model,
"descriptor_outputs": descriptor.outputs,
"native_outputs": self.native_outputs,
"archive_outputs": archive.outputs,
"lattice": "optical",
"lattice_samples": descriptor.samples,
"spacing": float(descriptor.spacing),
"optical_time_origin": origin,
"window_first": 0 if window is None else window.first,
"window_last": descriptor.samples - 1
if window is None
else window.last,
"provider_samples": self._samples,
"provider_time_zero": origin if window is None else window.t_s,
},
prewindow=fields,
drives=[
(chunk.start, chunk.stop, drive)
for chunk, drive in zip(self._chunks, drives, strict=True)
],
gate=None if window is None else self._gate,
binding=self._binding_record(),
)
self._trial_evaluations.append(identifier)
return identifier
def _binding_record(self) -> dict[str, object]:
"""Return the build, execution, plan and chunk records an archive shares.
- Returns: provider plan/build and execution configuration.
"""
return {
"build": self._build,
"execution": asdict(self._execution),
"plans": self._plan_records,
"chunks": [asdict(chunk) for chunk in self._chunks],
"hacc": None if self._hacc is None else self._hacc.record(),
}
def _archive_result(
self, evaluation: int | None, index: int, raw: _Native
) -> TDSEArchiveError | None:
"""Store one returned record, returning a failure instead of raising it.
The binding's own validation runs next either way: a contract failure of
the same result is then the primary error, and the archive failure a note.
- ``evaluation``: the archive's evaluation, or None without an archive.
- Returns: the write failure, if any.
"""
archive = self._archive
if archive is None or evaluation is None:
return None
outputs = self.native_outputs
assert outputs is not None
try:
archive.write_result(
evaluation,
index,
raw,
record_serializer(self._descriptor.model, outputs),
)
except TDSEArchiveError as error:
return error
return None
def _archive_failure(
self,
evaluation: int | None,
index: int | None,
error: BaseException,
classification: str,
*,
note: BaseException | None = None,
) -> None:
"""Record an evaluation's exception, keeping the caller's error primary.
- ``evaluation``: the archive's evaluation, or None without an archive.
- ``note``: the error the caller raises, which gains a note if recording fails.
"""
archive = self._archive
if archive is None or evaluation is None:
return
try:
archive.write_failure(
evaluation, chunk=index, error=error, classification=classification
)
archive.finish_evaluation(evaluation, "raised")
except TDSEArchiveError as failure:
(error if note is None else note).add_note(
f"the TDSE archive could not record this failure: {failure}"
)
def _call_domain(
self,
fields: Float[np.ndarray, "n_nodes n_to"],
drives: list[Float[np.ndarray, "_ n_samples"]],
*,
z: float,
purpose: TDSEPurpose,
) -> dict[str, object] | None:
"""Flag a preset call whose peak input intensity exceeds the preset's bound.
The peaks before and after the drive taper are both recorded, and the flag
uses the larger, so the gate cannot conceal a field above the bound.
- ``fields``: Full-lattice nodal fields in V/m.
- ``z``: optical position in m.
- Returns: the call's domain record for a preset, None otherwise.
"""
model = self._hacc
if model is None or model.preset is None:
return None
bound = float(model.preset["max_peak_intensity"])
before = _peak_intensity(fields)
after = _peak_intensity(np.concatenate(drives, axis=0))
outside = max(before) > bound
record: dict[str, object] = {
"peak_intensity_w_per_m2": before,
"driven_peak_intensity_w_per_m2": after,
"max_peak_intensity_w_per_m2": bound,
"outside_preset_domain": outside,
}
if outside:
self._domain_flags.append({"z": z, "purpose": purpose, **record})
self._flag_domain(
"intensity",
{"max_peak_intensity_w_per_m2": bound},
f"a call's peak input intensity {max(before):.6g} W/m^2 exceeds "
f"{bound:.6g} W/m^2",
)
return record
def _bound(self, raw: _Native) -> _Native:
"""Return a result's neglected-current bound, wherever its record keeps it.
- Returns: the bound array, or None for a haCC model without the closure.
"""
if self._hacc is None and _kind(raw) == "sae.grid_response":
return raw.plasma.neglected_current_bound
return raw.neglected_current_bound
def _failed_native_record(self, raw: _Native, chunk: _Chunk) -> object:
"""Retain diagnostic facts even when the native result is malformed."""
try:
return _failure_safe(
{
"row_start": chunk.start,
"row_stop": chunk.stop,
"monitors": raw.monitors,
"outcomes": raw.outcomes,
"valid": raw.valid,
"neglected_current_bound": self._bound(raw),
}
)
except (AttributeError, TypeError, ValueError) as error:
name = raw.computed["provider_type"] if isinstance(raw, Record) else "None"
return {"unreadable_native_record": name, "error": str(error)}
def _validate_diagnostics(self, raw: _Native, rows: int, samples: int) -> None:
"""Reject malformed diagnostics before reducing or labelling row arrays."""
if self._hacc is not None and (
not isinstance(raw.total_current, np.ndarray)
or raw.total_current.shape != (rows, samples)
or not np.all(np.isfinite(raw.total_current))
):
raise BackendError("haCC total current is malformed or non-finite")
bound = self._bound(raw)
absent_admitted = self._hacc is not None and not self._hacc.closure
if bound is None and absent_admitted:
pass
elif (
not isinstance(bound, np.ndarray)
or bound.dtype != np.dtype(np.float64)
or bound.shape != (rows, samples)
or not np.all(np.isfinite(bound))
or np.any(bound < 0)
):
raise BackendError("TDSE neglected-current bounds are malformed")
# The haCC monitors carry a bool non_finite flag beside float64 norms and
# the int64 step index (SEP 20, C4).
kinds = {"monitored_step": np.int64, "non_finite": np.bool_}
for name, vector in _fields_of(raw.monitors).items():
if vector is None:
continue
dtype = np.dtype(kinds.get(name, np.float64))
if (
not isinstance(vector, np.ndarray)
or vector.dtype != dtype
or vector.shape != (rows,)
or (dtype == np.float64 and not np.all(np.isfinite(vector)))
):
raise BackendError(f"TDSE row monitor {name!r} is malformed")
def _native_record(self, raw: _Native, chunk: _Chunk) -> object:
"""Summarise validated native diagnostics with their original row positions.
- ``raw``: public result whose diagnostic shapes and finiteness were validated.
- Returns: Per-row monitors, outcomes, validity and maximum current bounds. A
haCC record adds the bound's maximum relative to the row's peak total current,
``None`` with ``zero_reference`` for a zero-current row; the bound is a
diagnostic and never compared with a threshold.
"""
bound = self._bound(raw)
record: dict[str, object] = {
"row_start": chunk.start,
"row_stop": chunk.stop,
"monitors": raw.monitors,
"outcomes": raw.outcomes,
"valid": np.all(raw.valid, axis=1),
"neglected_current_bound_max": None
if bound is None
else np.max(bound, axis=1),
}
if self._hacc is not None:
peaks = np.max(np.abs(raw.total_current), axis=1)
maxima = None if bound is None else np.max(bound, axis=1)
record["neglected_current_bound"] = (
None
if maxima is None
else [
{
"relative_to_peak": None if peak == 0 else float(top / peak),
"zero_reference": bool(peak == 0),
}
for top, peak in zip(maxima, peaks, strict=True)
]
)
return record
def _hacc_currents(
self, raw: _Native, rows: int, samples: int
) -> tuple[np.ndarray, np.ndarray]:
"""Validate a haCC result and return its total and first-order currents.
- ``raw``: ``hacc.TotalCurrentResponse`` or ``hacc.GridResponse``.
- Returns: ``j_tot`` and ``j_1`` on the retained samples, C m/s per atom.
"""
model = self._hacc
assert model is not None
model.check_info(
raw,
rows=rows,
samples=samples,
spacing=float(self._descriptor.spacing),
)
if model.outputs == "total_current":
if _kind(raw) != "hacc.total_current_response":
raise BackendError(
"TDSE compact contract returned a different result type"
)
if raw.shadow_orders != (1,) or raw.shadow_total_current is None:
raise BackendError(
"TDSE compact response has incorrect shadow-order labels"
)
if raw.shadow_total_current.shape != (rows, 1, samples):
raise BackendError(
"TDSE compact response has incorrect shadow array extents"
)
return raw.total_current, raw.shadow_total_current[:, 0, :]
if _kind(raw) != "hacc.grid_response" or raw.shadow is None:
raise BackendError("haCC full output requires a response with shadows")
if raw.shadow.orders != (1,):
raise BackendError("TDSE full response has incorrect shadow-order labels")
# With the closure the total's order one is returned; without it the
# in-box order one is the first order of the field-impulse total, whose
# added term is at least third order (SEP 20, 4.3).
first = _component(
raw.shadow, "total_current" if model.closure else "current", 1
)
return raw.total_current, first
def _sae_currents(
self, raw: _Native, chunk: _Chunk
) -> tuple[
np.ndarray, np.ndarray, np.ndarray | None, np.ndarray | None, tuple[int, ...]
]:
"""Validate an SAE result and return its currents and shadow products.
- ``raw``: SAE ``GridResponse`` or ``TotalCurrentResponse``.
- Returns: total and first-order currents, the optional third-order current and
dipole, and the order labels, on the retained samples.
"""
descriptor = self._descriptor
expected_orders = (1,) if descriptor.shadow_order == 1 else (1, 3)
info = raw.info
if (
info.rows != chunk.rows
or info.samples != self._samples
or info.spacing != descriptor.spacing
or info.shadow_order != descriptor.shadow_order
or info.outputs != self._native_outputs
):
raise BackendError(
"TDSE result metadata disagree with the prepared contract"
)
if self._native_outputs == "full":
if _kind(raw) != "sae.grid_response" or raw.shadow is None:
raise BackendError(
"TDSE shadow model requires a full response with shadow products"
)
orders = raw.shadow.orders
if orders != expected_orders:
raise BackendError(
"TDSE full response has incorrect shadow-order labels"
)
# A full archive reads the same total and shadow currents from the full
# record; only a descriptor that needs it takes the third-order dipole.
return (
raw.plasma.total_current,
_component(raw.shadow, "total_current", 1),
_component(raw.shadow, "total_current", 3) if 3 in orders else None,
_component(raw.shadow, "charge_dipole", 3)
if descriptor.outputs == "full"
else None,
orders,
)
if _kind(raw) != "sae.total_current_response":
raise BackendError("TDSE compact contract returned a different result type")
orders = raw.shadow_orders
if orders != expected_orders or raw.shadow_total_current is None:
raise BackendError(
"TDSE compact response has incorrect shadow-order labels"
)
if raw.shadow_total_current.shape != (chunk.rows, len(orders), self._samples):
raise BackendError(
"TDSE compact response has incorrect shadow array extents"
)
third = raw.shadow_total_current[:, orders.index(3), :] if 3 in orders else None
return (
raw.total_current,
raw.shadow_total_current[:, orders.index(1), :],
third,
None,
orders,
)
def _normalize(self, raw: _Native, chunk: _Chunk, identity: str) -> TDSEResult:
"""Validate native labels and copy one chunk into the host seam result."""
descriptor = self._descriptor
if self._hacc is not None:
total, first = self._hacc_currents(raw, chunk.rows, self._samples)
third = dipole = None
orders: tuple[int, ...] = (1,)
else:
total, first, third, dipole, orders = self._sae_currents(raw, chunk)
if len(raw.outcomes) != chunk.rows or any(
not outcome.ok for outcome in raw.outcomes
):
raise BackendError("TDSE provider reported a failed atomic row")
if total.shape != (chunk.rows, self._samples):
raise BackendError("TDSE provider returned incorrect current extents")
# A window admits tdse_native only (it requires the time collar), so no
# third-order product is ever padded.
return TDSEResult(
self._pad(total),
self._pad(first),
input_identity=identity,
configuration_identity=descriptor.identity,
valid=self._pad_valid(raw.valid),
shadow_orders=orders,
third_order_current=third,
third_order_dipole=dipole,
monitors_json=_encoded(self._native_record(raw, chunk)),
outcomes_json=_encoded(
[
{"row": chunk.start + index, **_fields_of(value)}
for index, value in enumerate(raw.outcomes)
]
),
)
def _combine(self, results: list[TDSEResult], identity: str) -> TDSEResult:
"""Concatenate normalized chunks in the original radial row order."""
third = None
dipole = None
if self._descriptor.shadow_order == 3:
currents = [result.third_order_current for result in results]
assert all(value is not None for value in currents)
third = np.concatenate(
[value for value in currents if value is not None], axis=0
)
if self._descriptor.outputs == "full":
dipoles = [result.third_order_dipole for result in results]
assert all(value is not None for value in dipoles)
dipole = np.concatenate(
[value for value in dipoles if value is not None], axis=0
)
return TDSEResult(
np.concatenate([result.total_current for result in results], axis=0),
np.concatenate([result.first_order_current for result in results], axis=0),
input_identity=identity,
configuration_identity=self._descriptor.identity,
valid=np.concatenate([result.valid for result in results], axis=0),
shadow_orders=results[0].shadow_orders,
third_order_current=third,
third_order_dipole=dipole,
monitors_json=_encoded(
{"chunks": [json.loads(result.monitors_json) for result in results]}
),
outcomes_json=_encoded(
[
outcome
for result in results
for outcome in json.loads(result.outcomes_json)
]
),
)
# -- the haCC consistency diagnostic --------------------------------------------
def diagnostic_currents(
self, field: Float[np.ndarray, "1 n_samples"], spacing: float
) -> tuple[Float[np.ndarray, " n_samples"], Float[np.ndarray, " n_samples"]]:
"""Evaluate one diagnostic row on the run's open haCC model (SEP 20, 4.9).
A temporary one-row plan on the protocol lattice is built, admitted with the
propagation plans against the memory limit, charged as one ``diagnostic``
call to this binding's ledger and closed again.
- ``field``: finite float64 field in V/m, shape ``(1, n_samples)``, at ``m *
spacing``.
- ``spacing``: positive protocol spacing in s.
- Returns: ``j_tot`` and ``j_1`` in C m/s per atom on the protocol samples.
"""
self._ensure_open()
model = self._hacc
if model is None:
raise BackendError("the consistency diagnostic needs a haCC binding")
self._ensure_worker()
context = _DiagnosticContext(
model,
self._budget,
self._failures,
self._worker,
self._memory,
self._execution.memory_limit,
archive=self._archive,
labels={
"configuration_identity": self._descriptor.identity,
"binding": self._binding_record(),
},
)
return _diagnostic(context, field, float(spacing))
def record_consistency(self, record: Mapping[str, object]) -> None:
"""Attach the run's consistency provenance to the binding's metadata."""
self._consistency = dict(record)
# -- metadata ------------------------------------------------------------------
def metadata(self) -> Mapping[str, object]:
"""Return independent JSON-compatible build, setup and call diagnostics."""
record: dict[str, object] = {
"configuration_identity": self._descriptor.identity,
"configuration": json.loads(self._descriptor.configuration_json),
"execution": asdict(self._execution),
"build": self._build,
"plans": self._plan_records,
"chunks": [asdict(chunk) for chunk in self._chunks],
"ledger": asdict(self.ledger),
"memory_estimate_bytes": self._estimate,
"memory_bytes": self._memory,
"diagnostic_cache_hits": self._diagnostic_cache_hits,
"last_call": self._last_call,
"failures": self._failures,
"closed": self._closed,
}
if self._window is not None:
record["atomic_window"] = self._window.to_json()
if self._hacc is not None:
record["model"] = "hacc"
record["hacc"] = self._hacc.record()
record["domain"] = {
"reasons": self._domain_reasons,
"flagged_calls": self._domain_flags,
}
record["consistency"] = self._consistency
if self._archive is not None:
record["archive"] = {
**self._archive.reference(),
"native_outputs": self.native_outputs,
"trials": self._trials_opened,
}
return cast("dict[str, object]", _failure_safe(record))
def memory_ledger(self) -> Mapping[str, int]:
"""Return native storage counted once and the admitted host-result reserve.
- Returns: named atomic allocations in bytes, excluding separately owned optics.
"""
return dict(self._memory)
def close(self) -> None:
"""End the worker, which closes the plans before the setup, keeping records.
Metadata and attempted-call counts survive; so does a worker already killed.
"""
if not self._closed:
self._budget.release_reservation()
self._stack.close()
self._cached = None
self._cached_fields = None
self._closed = True
@dataclass(frozen=True)
class _DiagnosticContext:
"""Collect what one charged haCC diagnostic call needs.
Attributes
----------
model : _HaccView
The open prepared model's record.
budget : AtomicCallBudget
Ledger the call is charged to, as ``diagnostic``.
failures : list
Failure records the call appends to.
worker : _Channel
The worker holding the model.
memory : mapping
Memory already held, in bytes.
limit : int or None
Memory limit in bytes.
archive : TDSEArchiveWriter or None
Archive session receiving the call, or None.
labels : mapping
The run's identities for the archive record: configuration identity (None
for a standalone check) and the build and execution records.
"""
model: _HaccView
budget: AtomicCallBudget
failures: list[object]
worker: _Channel
memory: Mapping[str, int]
limit: int | None
archive: TDSEArchiveWriter | None = None
labels: Mapping[str, object] = dataclasses.field(default_factory=dict)
def _diagnostic(
context: _DiagnosticContext,
field: Float[np.ndarray, "1 n_samples"],
spacing: float,
) -> tuple[Float[np.ndarray, " n_samples"], Float[np.ndarray, " n_samples"]]:
"""Run one charged one-row haCC diagnostic on a temporary plan in the worker.
- ``field``: finite float64 field in V/m, shape ``(1, n_samples)``.
- ``spacing``: protocol spacing in s.
- Returns: owning ``j_tot`` and ``j_1`` in C m/s.
"""
model, budget, worker = context.model, context.budget, context.worker
if (
not isinstance(field, np.ndarray)
or field.dtype != np.dtype(np.float64)
or field.ndim != 2
or field.shape[0] != 1
or field.shape[1] < 4
or not np.all(np.isfinite(field))
):
raise BackendError("the consistency field must be finite float64 (1, samples)")
samples = int(field.shape[1])
try:
prepared = cast(
"Mapping[str, Any]",
worker.request("diagnostic_plan", {"samples": samples, "spacing": spacing}),
)
except RemoteError as error:
raise _local_error(error, "TDSE diagnostic preparation failed") from error
try:
model.update(prepared["hacc"])
extra = cast("Mapping[str, int]", prepared["memory"])
# The temporary plan overlaps the propagation plans: admit their sum.
_admit_memory(
{**context.memory, **extra, "diagnostic.host": 8 * 20 * samples},
context.limit,
)
budget.reserve(1, calls=1)
try:
total, first = _diagnostic_call(context, field, spacing)
finally:
budget.release_reservation()
finally:
# A failed or cancelled call has ended the worker already, which then refuses
# the close at once as not running: nothing is left to close.
with contextlib.suppress(RemoteError, WorkerError):
worker.request("diagnostic_close", {})
return (
np.array(total[0], dtype=np.float64, copy=True),
np.array(first[0], dtype=np.float64, copy=True),
)
def _diagnostic_call(
context: _DiagnosticContext,
field: Float[np.ndarray, "1 n_samples"],
spacing: float,
) -> tuple[np.ndarray, np.ndarray]:
"""Make, archive and validate the diagnostic's one charged atomic call."""
model, budget = context.model, context.budget
samples = int(field.shape[1])
drive = np.ascontiguousarray(field)
evaluation = _diagnostic_archive_begin(context, drive, spacing)
budget.begin_evaluation(purpose="diagnostic", calls=1)
budget.charge(purpose="diagnostic")
archived: TDSEArchiveError | None = None
try:
raw = _returned(context.worker.request("diagnostic_evaluate", {"drive": drive}))
archived = _diagnostic_archive_result(context, evaluation, raw)
binding_view = _DiagnosticView(model, spacing)
binding_view.validate(raw, samples)
total, first = binding_view.currents(raw, samples)
except (KeyboardInterrupt, RunTerminated, StopRequested) as stopping:
_diagnostic_archive_failure(
context, evaluation, stopping, "interrupted", stopping
)
raise
except Exception as error:
context.worker.kill()
budget.record_failure()
kind = _classify(error)
name, detail = _described(error)
context.failures.append(
_failure_safe(
{
"purpose": "diagnostic",
"classification": kind,
"error": f"{name}: {detail}",
}
)
)
wrapped = BackendError(f"TDSE consistency diagnostic failed: {detail} [{kind}]")
_attach_remote(wrapped, error)
if archived is not None:
wrapped.add_note(
f"the TDSE archive could not store this result: {archived}"
)
_diagnostic_archive_failure(context, evaluation, error, kind, wrapped)
raise wrapped from error
if archived is not None:
raise archived
if context.archive is not None and evaluation is not None:
context.archive.finish_evaluation(evaluation, "returned")
return total, first
def _diagnostic_archive_begin(
context: _DiagnosticContext,
field: Float[np.ndarray, "1 n_samples"],
spacing: float,
) -> int | None:
"""Write a consistency diagnostic's labels and input before its provider call.
- ``field``: the exact protocol field passed, V/m.
- ``spacing``: protocol spacing in s.
- Returns: the archive's evaluation identifier, or None without an archive.
"""
archive = context.archive
if archive is None:
return None
archive.ensure_writable()
model = context.model
samples = int(field.shape[1])
labels = dict(context.labels)
binding = labels.pop("binding", {})
return archive.begin_evaluation(
{
"trial_id": None,
"role": "diagnostic",
"purpose": "consistency",
"z": None,
"density": None,
"input_identity": waveform_identity(field),
"basis_identity": model.basis.identity,
"model": "hacc",
"descriptor_outputs": "total_current",
"native_outputs": model.outputs,
"archive_outputs": archive.outputs,
# The consistency protocol's own lattice: provider time zero is its
# first sample, which is not a time of the optical coordinate.
"lattice": "consistency_protocol",
"lattice_samples": samples,
"spacing": float(spacing),
"optical_time_origin": None,
"window_first": 0,
"window_last": samples - 1,
"provider_samples": samples,
"provider_time_zero": 0.0,
**labels,
},
prewindow=field,
drives=[(0, 1, field)],
gate=None,
binding=cast("Mapping[str, object]", binding),
)
def _diagnostic_archive_result(
context: _DiagnosticContext, evaluation: int | None, raw: _Native
) -> TDSEArchiveError | None:
"""Store a consistency diagnostic's returned record before its validation.
- Returns: the write failure, if any.
"""
archive = context.archive
if archive is None or evaluation is None:
return None
outputs = context.model.outputs
assert outputs is not None
try:
archive.write_result(evaluation, 0, raw, record_serializer("hacc", outputs))
except TDSEArchiveError as error:
return error
return None
def _diagnostic_archive_failure(
context: _DiagnosticContext,
evaluation: int | None,
error: BaseException,
classification: str,
primary: BaseException,
) -> None:
"""Record a failed consistency diagnostic, keeping the provider error primary.
- ``primary``: the error raised, which gains a note if recording fails.
"""
archive = context.archive
if archive is None or evaluation is None:
return
try:
archive.write_failure(
evaluation, chunk=0, error=error, classification=classification
)
archive.finish_evaluation(evaluation, "raised")
except TDSEArchiveError as failure:
primary.add_note(f"the TDSE archive could not record this failure: {failure}")
class _DiagnosticView:
"""Validate a one-row diagnostic result with the binding's haCC rules.
- ``spacing``: protocol spacing in s.
"""
def __init__(self, model: _HaccView, spacing: float) -> None:
self._model = model
self._spacing = spacing
def validate(self, raw: _Native, samples: int) -> None:
"""Check the bound, monitors, outcomes and validity of one row."""
bound = raw.neglected_current_bound
if bound is None:
if self._model.closure:
raise BackendError("TDSE neglected-current bounds are malformed")
elif (
not isinstance(bound, np.ndarray)
or bound.dtype != np.dtype(np.float64)
or bound.shape != (1, samples)
or not np.all(np.isfinite(bound))
or np.any(bound < 0)
):
raise BackendError("TDSE neglected-current bounds are malformed")
if len(raw.outcomes) != 1 or not raw.outcomes[0].ok or not np.all(raw.valid):
raise BackendError("TDSE provider reported a failed diagnostic row")
def currents(self, raw: _Native, samples: int) -> tuple[np.ndarray, np.ndarray]:
"""Return the validated total and first-order currents.
- Returns: ``(1, samples)`` currents in C m/s.
"""
model = self._model
model.check_info(raw, rows=1, samples=samples, spacing=self._spacing)
if model.outputs == "total_current":
if (
_kind(raw) != "hacc.total_current_response"
or raw.shadow_orders != (1,)
or raw.shadow_total_current is None
or raw.shadow_total_current.shape != (1, 1, samples)
):
raise BackendError("haCC compact diagnostic result is malformed")
return raw.total_current, raw.shadow_total_current[:, 0, :]
if (
_kind(raw) != "hacc.grid_response"
or raw.shadow is None
or raw.shadow.orders != (1,)
):
raise BackendError("haCC full diagnostic result is malformed")
first = _component(
raw.shadow, "total_current" if model.closure else "current", 1
)
return raw.total_current, first
[docs]
class DiagnosticSession:
"""Prepare a haCC model alone for a standalone consistency check.
Used by :func:`solprop.driver.tdse_consistency.check_consistency`, which needs
one diagnostic row and no propagation plans. Close it when done.
Parameters
----------
basis : HaccBasis
Validated preset or custom record.
execution : TDSEExecution
Device, memory limit and the explicit call budget of the check.
archive : TDSEArchiveWriter or None, optional
Archive session receiving the diagnostic call, labelled ``diagnostic``.
stop : solprop.log.StopFlag or None, optional
A flag polled while the provider's worker is awaited.
"""
def __init__(
self,
basis: HaccBasis,
execution: TDSEExecution,
*,
archive: TDSEArchiveWriter | None = None,
stop: StopFlag | None = None,
) -> None:
self._execution = execution
self._archive = archive
self._budget = AtomicCallBudget(execution.max_atomic_calls, 1)
self._failures: list[object] = []
self._worker = _open_provider(stop=stop)
try:
self.build = _build_record(self._worker.identity, execution)
reply = cast(
"Mapping[str, Any]",
self._worker.request(
"prepare",
{
"model": "hacc",
"basis": asdict(basis),
"rows": [],
"complement": execution.complement,
"archive_outputs": "required"
if archive is None
else archive.outputs,
"device": execution.device,
"memory": {},
"memory_limit": execution.memory_limit,
"defer_plans": True,
},
),
)
self._model = _HaccView.from_reply(basis, reply["hacc"])
self._memory = dict(cast("Mapping[str, int]", reply["memory"]))
if archive is not None:
self._memory.update(archive.memory_ledger(rows=1, samples=1))
_admit_memory(self._memory, execution.memory_limit)
except BaseException as error:
self._worker.kill()
raise _preparation_error(error) from error
[docs]
@property
def ledger(self) -> AtomicCallLedger: # numpydoc ignore=RT01
"""The check's attempted-call accounting.
Diagnostic charges, including failures.
"""
return self._budget.snapshot()
[docs]
def diagnostic_currents(
self, field: Float[np.ndarray, "1 n_samples"], spacing: float
) -> tuple[Float[np.ndarray, " n_samples"], Float[np.ndarray, " n_samples"]]:
"""Evaluate one charged diagnostic row on the prepared model.
Parameters
----------
field : numpy.ndarray
Finite float64 field in V/m, shape ``(1, n_samples)``.
spacing : float
Protocol spacing in s.
Returns
-------
tuple of numpy.ndarray
``j_tot`` and ``j_1`` in C m/s per atom.
"""
context = _DiagnosticContext(
self._model,
self._budget,
self._failures,
self._worker,
self._memory,
self._execution.memory_limit,
archive=self._archive,
labels={
"configuration_identity": None,
"binding": {
"build": self.build,
"execution": asdict(self._execution),
"hacc": self._model.record(),
},
},
)
return _diagnostic(context, field, float(spacing))
[docs]
def close(self) -> None:
"""End the worker and its model; repeated calls are harmless."""
self._worker.close()
def _preparation_error(error: BaseException) -> BaseException:
"""Return a preparation failure as the binding raises it.
solprop's own refusals and interruptions as they are; a provider exception the
worker reported as solprop's class or ``TDSE preparation failed: ... [kind]``.
"""
if isinstance(error, RemoteError):
return _local_error(error, "TDSE preparation failed")
if not isinstance(error, Exception) or isinstance(error, SolpropError):
return error
return BackendError(f"TDSE preparation failed: {error} [{_classify(error)}]")
def _attach_remote(wrapped: BaseException, error: BaseException) -> None:
"""Attach a worker's traceback or its standard error to the error raised."""
if isinstance(error, RemoteError):
wrapped.add_note(f"the provider worker's traceback:\n{error.remote_traceback}")
for note in getattr(error, "__notes__", ()):
wrapped.add_note(note)