"""Generate a standalone, editable Python script from a saved session.
``croak script options.toml`` emits a small script that reproduces the retrieval
by calling the public croak API stage by stage. Every parameter is rendered as an
explicit dataclass literal so the script is transparent and hand-editable, while
the pipeline steps call the same :mod:`croak.session.pipeline` functions the GUI
and the ``replay`` command use, so a generated script never drifts from the real
mapping. Keep the input files and software version with the script.
"""
from __future__ import annotations
from dataclasses import fields, is_dataclass
import numpy as np
from .session import SessionOptions, retarget
__all__ = ["generate_script"]
_HEADER = '''#!/usr/bin/env python3
"""Auto-generated by ``croak script`` — replays a retrieval session.
Edit the parameters below and rerun (``python {script_name}``). The pipeline
section mirrors :func:`croak.run_session`; tweak it freely.
"""
import numpy as np
import croak
from croak.session.dispersion import compose_dispersion
from croak.session.params import (
{params_imports})
from croak.session.pipeline import (
{pipeline_imports})
SEED = {seed}
OUTPUT = {output!r}
'''
#: Loader-dependent head of the pipeline section, keyed by the session's entry
#: (with the simulated "skip filtering & regrid" option as its own variant): how
#: ``data`` and ``tracedata`` are built. Everything after it is shared.
_LOAD_STEPS: dict[str, str] = {
"experimental": (
"data = assemble_load_data(load)\ntracedata = build_tracedata(preproc, data)"
),
"simulated": (
"data = assemble_simulated_load_data(simulated)\n"
"tracedata = build_tracedata(preproc, data)"
),
# Keep the lightweight display/truth bundle alongside the native-grid trace:
# ``data['truth']`` carries the complex ModelPNPS seed when requested.
"simulated-raw": (
"data = assemble_simulated_load_data(simulated)\n"
"tracedata = assemble_simulated_tracedata(simulated)"
),
}
#: Extra names each variant needs from ``croak.session.params`` / ``.pipeline``.
_LOAD_IMPORTS: dict[str, tuple[tuple[str, ...], tuple[str, ...]]] = {
"experimental": (("LoadParams",), ("assemble_load_data", "build_tracedata")),
"simulated": (
("SimulatedLoadParams",),
("assemble_simulated_load_data", "build_tracedata"),
),
"simulated-raw": (
("SimulatedLoadParams",),
("assemble_simulated_load_data", "assemble_simulated_tracedata"),
),
}
#: Names every variant imports, whatever the loader.
_COMMON_PARAMS = (
"BeamPathMirror",
"DispersionParams",
"PreprocParams",
"RetrieveParams",
)
_COMMON_PIPELINE = ("apply_post_filter", "process", "run_retrieval")
_PIPELINE_HEADER = (
"\n# --- pipeline (mirrors croak.run_session: load → preproc → retrieve"
" → dispersion) ---\n"
)
_PIPELINE = """
result = run_retrieval(
retrieve,
tracedata,
truth=data.get("truth"),
rng=np.random.default_rng(SEED),
)
result = apply_post_filter(retrieve, result, tracedata)
final = compose_dispersion(dispersion, result)
processed = process(final, tracedata)
croak.save_result(
final,
OUTPUT,
processed=processed,
measured=tracedata.trace,
truth=data.get("truth"),
force=True,
)
print(f"R = {result.error:.4%} FWHM = {processed.fwhm_retr / 1e-15:.2f} fs")
"""
def _render_value(value: object) -> str:
"""Render one field value as a Python literal (lists of dataclasses too)."""
if isinstance(value, np.generic): # numpy scalar → plain Python literal
return repr(value.item())
if isinstance(value, list):
items = ", ".join(_render_value(v) for v in value)
return f"[{items}]"
if is_dataclass(value) and not isinstance(value, type):
inner = ", ".join(
f"{f.name}={_render_value(getattr(value, f.name))}" for f in fields(value)
)
return f"{type(value).__name__}({inner})"
return repr(value)
def _render_dataclass(name: str, obj: object) -> str:
"""Render ``name = ClassName(field=value, …)`` for the non-default fields.
Only fields that differ from the dataclass default are emitted, so the script
shows what the session actually customised (an unmodified stage renders as an
empty constructor).
"""
default = type(obj)()
lines = [f"{name} = {type(obj).__name__}("]
body = [
f" {f.name}={_render_value(getattr(obj, f.name))},"
for f in fields(obj) # type: ignore[arg-type]
if getattr(obj, f.name) != getattr(default, f.name)
]
if not body:
return f"{name} = {type(obj).__name__}()"
lines.extend(body)
lines.append(")")
return "\n".join(lines)
def _variant(options: SessionOptions) -> str:
"""Key into the loader-dependent tables: the entry, refined by ``raw_direct``.
Raises
------
ValueError
For a synthetic session, whose generator settings are not serialised, so
there is nothing for a script to reload.
"""
if options.entry == "synthetic":
raise ValueError(
"cannot script a synthetic session: the generator's settings are not "
"saved with the options"
)
if options.entry == "simulated":
return "simulated-raw" if options.simulated.raw_direct else "simulated"
return "experimental"
def _import_block(names: tuple[str, ...]) -> str:
"""Render sorted `` Name,`` lines for one parenthesised import."""
return "".join(f" {name},\n" for name in sorted(names))
[docs]
def generate_script(
options: SessionOptions,
*,
base_dir: str | None = None,
output: str = "result.h5",
seed: int = 0,
script_name: str = "replay.py",
) -> str:
"""Render an editable Python script that replays ``options``.
The loader section follows the session's ``entry``: an experimental session
calls :func:`~croak.session.pipeline.assemble_load_data`, a simulated one
:func:`~croak.session.pipeline.assemble_simulated_load_data` (or
:func:`~croak.session.pipeline.assemble_simulated_tracedata` for its raw
"skip filtering & regrid" option), matching
:func:`croak.session.engine.run_session`.
Parameters
----------
options : SessionOptions
The session to turn into a script.
base_dir : str, optional
If given, the script's load paths are retargeted to this dataset folder
(calibration paths are kept fixed), as ``croak replay --base-dir`` does.
output : str, optional
The ``result.h5`` path the generated script writes to.
seed : int, optional
RNG seed embedded in the script for a reproducible run.
script_name : str, optional
The script's own filename, used only in its docstring.
Returns
-------
str
The script source.
Raises
------
ValueError
If the session came from the synthetic generator (nothing to reload).
"""
if base_dir is not None:
options = retarget(options, base_dir)
variant = _variant(options)
params_names, pipeline_names = _LOAD_IMPORTS[variant]
loader = (
_render_dataclass("load", options.load)
if variant == "experimental"
else _render_dataclass("simulated", options.simulated)
)
parts = [
_HEADER.format(
seed=seed,
output=output,
script_name=script_name,
params_imports=_import_block(_COMMON_PARAMS + params_names),
pipeline_imports=_import_block(_COMMON_PIPELINE + pipeline_names),
),
loader,
"",
_render_dataclass("preproc", options.preproc),
"",
_render_dataclass("retrieve", options.retrieve),
"",
_render_dataclass("dispersion", options.dispersion),
_PIPELINE_HEADER + _LOAD_STEPS[variant] + _PIPELINE,
]
return "\n".join(parts)