Source code for pycsamt.ai.inversion.run_config

# Author: LKouadio <etanoyau@gmail.com>
# License: LGPL-3.0
"""End-to-end run configuration for pyCSAMT AI inversion.

:class:`RunConfig` is a thin container that bundles a
:class:`~pycsamt.forward.config.ForwardConfig` and an
:class:`~pycsamt.ai.inversion.config.InversionConfig` into **one**
source-of-truth file.  A single file then covers the complete experiment:
synthetic dataset generation, network architecture, and training.

Supported file formats: Python (``.py``), JSON (``.json``), YAML (``.yml``).

Quick start
-----------
Generate a default template, edit it, run::

    from pycsamt.ai.inversion.run_config import RunConfig

    # 1 — write a single annotated file that covers both forward and inversion
    RunConfig.write_template("experiment_01.yml")

    # 2 — edit experiment_01.yml …

    # 3 — load, validate, and run
    run = RunConfig.from_file("experiment_01.yml")
    run.validate()

    from pycsamt.forward.batch import generate_dataset

    ds = generate_dataset(**run.to_dataset_kwargs())

    inv = run.to_inverter()
    inv.fit(ds, **run.to_fit_kwargs())
    inv.save(run.checkpoint_path())

Cross-config validation
-----------------------
:meth:`RunConfig.validate` checks that the forward and inversion configs are
consistent:

* ``forward.solver`` must equal ``inversion.solver``.
* ``forward.include_phase`` must equal ``inversion.include_phase``.
* When the layer count is fixed (``n_layers_min == n_layers_max``),
  ``forward.n_layers_min`` must equal ``inversion.n_layers``.
"""

from __future__ import annotations

import ast
import json
from dataclasses import asdict, dataclass, field
from pathlib import Path
from textwrap import wrap
from typing import Any

from ...forward.config import (
    _FORWARD_CONFIG_SCHEMA,
    ForwardConfig,
)
from ...models.config_io import (
    _comment_lines,
    _groups,
    _py_value,
    _schema_map,
    _yaml_value,
)
from .config import _INVERSION_CONFIG_SCHEMA, InversionConfig

__all__ = ["RunConfig"]


# ═══════════════════════════════════════════════════════════════════════════════
# Internal serializers
# ═══════════════════════════════════════════════════════════════════════════════

# ── Python ──────────────────────────────────────────────────────────────────


def _write_run_py(
    path: Path,
    fwd_vals: dict[str, Any],
    inv_vals: dict[str, Any],
    name: str,
    description: str,
    title: str,
) -> None:
    """Write a commented two-dict Python run configuration file."""
    by_f = _schema_map(_FORWARD_CONFIG_SCHEMA)
    by_i = _schema_map(_INVERSION_CONFIG_SCHEMA)

    def _block(
        var: str, values: dict[str, Any], by_name, section_title: str
    ) -> list[str]:
        lines: list[str] = [
            f"# ── {section_title} {'─' * max(0, 54 - len(section_title))}",
            f"{var} = {{",
        ]
        for group, names in _groups(values, list(by_name.values())):
            lines.append(f"    # ---- {group} ----")
            for nm in names:
                entry = by_name.get(nm)
                if entry is not None:
                    for cl in _comment_lines(entry.description, "    # "):
                        lines.append(cl)
                lines.append(f"    {nm!r}: {_py_value(values[nm])},")
                lines.append("")
        if lines[-1] == "":
            lines.pop()
        lines.append("}")
        return lines

    header = [
        '"""Source-of-truth run configuration generated by PyCSAMT."""',
        "",
        f"# {title}",
    ]
    if name:
        header.append(f"# Experiment : {name}")
    if description:
        for line in wrap(description, width=70):
            header.append(f"# Description: {line}")
    header += [
        "# Edit FORWARD and INVERSION, then load with RunConfig.from_file().",
        "",
        "",
    ]

    fwd_block = _block("FORWARD", fwd_vals, by_f, "Forward modelling")
    inv_block = _block("INVERSION", inv_vals, by_i, "Inversion")

    lines = (
        header
        + fwd_block
        + [
            "",
            "",
        ]
        + inv_block
    )
    path.write_text("\n".join(lines) + "\n", encoding="utf-8")


def _read_run_py(path: Path) -> tuple[dict[str, Any], dict[str, Any]]:
    """Read ``FORWARD`` and ``INVERSION`` dicts from a Python run file."""
    tree = ast.parse(path.read_text(encoding="utf-8"))
    results: dict[str, dict] = {}
    for node in tree.body:
        if isinstance(node, ast.Assign):
            for target in node.targets:
                if isinstance(target, ast.Name) and target.id in (
                    "FORWARD",
                    "INVERSION",
                ):
                    data = ast.literal_eval(node.value)
                    if isinstance(data, dict):
                        results[target.id] = data
    for key in ("FORWARD", "INVERSION"):
        if key not in results:
            raise ValueError(
                f"Python run config {path} must define a {key} dictionary."
            )
    return results["FORWARD"], results["INVERSION"]


# ── JSON ─────────────────────────────────────────────────────────────────────


def _write_run_json(
    path: Path,
    fwd_vals: dict[str, Any],
    inv_vals: dict[str, Any],
    name: str,
    description: str,
    title: str,
) -> None:
    """Write a nested JSON run configuration file."""
    by_f = _schema_map(_FORWARD_CONFIG_SCHEMA)
    by_i = _schema_map(_INVERSION_CONFIG_SCHEMA)

    def _schema_block(vals, by_name):
        return {
            nm: {
                "group": by_name[nm].group if nm in by_name else "General",
                "description": by_name[nm].description
                if nm in by_name
                else "",
            }
            for nm in vals
        }

    payload = {
        "_meta": {
            "title": title,
            "name": name,
            "description": description,
            "note": (
                "Edit values under 'forward' and 'inversion'. "
                "The '_meta' and '_schema' blocks are documentation "
                "and are ignored by RunConfig.from_file()."
            ),
        },
        "_schema": {
            "forward": _schema_block(fwd_vals, by_f),
            "inversion": _schema_block(inv_vals, by_i),
        },
        "forward": fwd_vals,
        "inversion": inv_vals,
    }
    path.write_text(
        json.dumps(payload, indent=2, sort_keys=False, ensure_ascii=False)
        + "\n",
        encoding="utf-8",
    )


def _read_run_json(path: Path) -> tuple[dict[str, Any], dict[str, Any]]:
    """Read ``forward`` and ``inversion`` from a JSON run file."""
    data = json.loads(path.read_text(encoding="utf-8"))
    for key in ("forward", "inversion"):
        if key not in data or not isinstance(data[key], dict):
            raise ValueError(
                f"JSON run config {path} must contain a '{key}' mapping."
            )
    return data["forward"], data["inversion"]


# ── YAML ─────────────────────────────────────────────────────────────────────


def _write_run_yaml(
    path: Path,
    fwd_vals: dict[str, Any],
    inv_vals: dict[str, Any],
    name: str,
    description: str,
    title: str,
) -> None:
    """Write a commented nested YAML run configuration file."""
    by_f = _schema_map(_FORWARD_CONFIG_SCHEMA)
    by_i = _schema_map(_INVERSION_CONFIG_SCHEMA)

    def _section(
        section_key: str, vals, by_name, section_title: str
    ) -> list[str]:
        lines = [
            f"# ── {section_title} {'─' * max(0, 54 - len(section_title))}",
            f"{section_key}:",
        ]
        for group, names in _groups(vals, list(by_name.values())):
            lines.append(f"  # ---- {group} ----")
            for nm in names:
                entry = by_name.get(nm)
                if entry is not None:
                    for cl in _comment_lines(entry.description, "  # "):
                        lines.append(cl)
                lines.append(f"  {nm}: {_yaml_value(vals[nm])}")
                lines.append("")
        if lines and lines[-1] == "":
            lines.pop()
        return lines

    header = [
        f"# {title}",
    ]
    if name:
        header.append(f"# Experiment : {name}")
    if description:
        for line in wrap(description, width=70):
            header.append(f"# Description: {line}")
    header += [
        "# Edit values, then load with RunConfig.from_file().",
        "",
    ]

    fwd_lines = _section("forward", fwd_vals, by_f, "Forward modelling")
    inv_lines = _section("inversion", inv_vals, by_i, "Inversion")

    lines = header + fwd_lines + ["", ""] + inv_lines
    path.write_text("\n".join(lines).rstrip() + "\n", encoding="utf-8")


def _read_run_yaml(path: Path) -> tuple[dict[str, Any], dict[str, Any]]:
    """Read ``forward`` and ``inversion`` from a YAML run file."""
    try:
        import yaml
    except ImportError as exc:
        raise ImportError(
            "Reading YAML run config files requires PyYAML."
        ) from exc
    data = yaml.safe_load(path.read_text(encoding="utf-8"))
    if not isinstance(data, dict):
        raise ValueError(f"YAML run config {path} must contain a mapping.")
    for key in ("forward", "inversion"):
        if key not in data or not isinstance(data[key], dict):
            raise ValueError(
                f"YAML run config {path} must contain a '{key}' mapping."
            )
    return data["forward"], data["inversion"]


# ── Dispatcher ───────────────────────────────────────────────────────────────


def _format_from_path(path: Path, fmt: str | None) -> str:
    if fmt is not None:
        f = fmt.lower().lstrip(".")
        return "yml" if f == "yaml" else f
    suffix = path.suffix.lower()
    if suffix in {".json", ".yml", ".yaml"}:
        return suffix.lstrip(".")
    return "py"


def _target_path(path: Path, fmt: str) -> Path:
    if path.suffix:
        return path
    return path.with_suffix(".yml" if fmt == "yml" else f".{fmt}")


def _write_run(
    path: Path,
    fwd_vals: dict[str, Any],
    inv_vals: dict[str, Any],
    name: str,
    description: str,
    title: str,
    fmt: str,
) -> None:
    if fmt == "py":
        _write_run_py(path, fwd_vals, inv_vals, name, description, title)
    elif fmt == "json":
        _write_run_json(path, fwd_vals, inv_vals, name, description, title)
    else:
        _write_run_yaml(path, fwd_vals, inv_vals, name, description, title)


def _read_run(path: Path) -> tuple[dict[str, Any], dict[str, Any]]:
    suffix = path.suffix.lower()
    if suffix == ".py":
        return _read_run_py(path)
    if suffix == ".json":
        return _read_run_json(path)
    if suffix in {".yml", ".yaml"}:
        return _read_run_yaml(path)
    raise ValueError(
        f"Unsupported run config suffix {suffix!r}. "
        "Use .py, .json, .yml, or .yaml."
    )


# ── Helper — filter unknown keys before constructing sub-configs ──────────────


def _safe_fields(cls, raw: dict[str, Any], strict: bool) -> dict[str, Any]:
    from dataclasses import fields as dc_fields

    allowed = {f.name for f in dc_fields(cls)}
    unknown = sorted(
        set(raw) - allowed - {k for k in raw if k.startswith("_")}
    )
    if unknown and strict:
        raise ValueError(
            f"Unknown {cls.__name__} parameter(s): {', '.join(unknown)}"
        )
    return {k: v for k, v in raw.items() if k in allowed}


# ═══════════════════════════════════════════════════════════════════════════════
# RunConfig
# ═══════════════════════════════════════════════════════════════════════════════


[docs] @dataclass class RunConfig: """Bundle a :class:`~pycsamt.forward.config.ForwardConfig` and an :class:`~pycsamt.ai.inversion.config.InversionConfig` into one source-of-truth experiment file. Parameters ---------- forward : ForwardConfig Dataset generation and solver settings. inversion : InversionConfig Network architecture and training settings. name : str Short experiment identifier written into the file header. description : str Optional free-text note describing the experiment. Notes ----- :meth:`validate` checks internal consistency between the two sub-configs: * ``forward.solver == inversion.solver`` * ``forward.include_phase == inversion.include_phase`` * Fixed layer count (``n_layers_min == n_layers_max``) must match ``inversion.n_layers``. Examples -------- Default run (MT1D, ResNet, 5 layers):: >>> run = RunConfig() >>> run.forward.solver 'mt1d' >>> run.inversion.arch 'resnet' Custom experiment:: >>> run = RunConfig( ... forward=ForwardConfig(solver="mt1d", n_samples=20_000, seed=1), ... inversion=InversionConfig(arch="resnet", n_layers=5, epochs=200), ... name="mt1d_resnet_20k", ... ) >>> run.validate() Write a template, edit it, reload:: >>> path = RunConfig.write_template("experiment_01.yml") >>> run = RunConfig.from_file(path) """ forward: ForwardConfig = field(default_factory=ForwardConfig) inversion: InversionConfig = field(default_factory=InversionConfig) name: str = "" description: str = "" # ───────────────────────────────────────────────────────────────────────── # Validation # ─────────────────────────────────────────────────────────────────────────
[docs] def validate(self) -> None: """Validate both sub-configs and their mutual consistency. Raises ------ ValueError Descriptive message pointing to the offending parameter or the cross-config inconsistency. """ self.forward.validate() self.inversion.validate() if self.forward.solver != self.inversion.solver: raise ValueError( f"forward.solver ({self.forward.solver!r}) does not match " f"inversion.solver ({self.inversion.solver!r}). " "Both must target the same EM method." ) if self.forward.include_phase != self.inversion.include_phase: raise ValueError( f"forward.include_phase ({self.forward.include_phase}) " f"does not match inversion.include_phase " f"({self.inversion.include_phase}). " "The feature vector shape must match between dataset " "generation and training." ) fixed = self.forward.n_layers_min == self.forward.n_layers_max if fixed and self.forward.n_layers_min != self.inversion.n_layers: raise ValueError( f"forward has a fixed layer count of " f"{self.forward.n_layers_min} but inversion.n_layers is " f"{self.inversion.n_layers}. They must agree when the " "layer count is fixed (n_layers_min == n_layers_max)." )
# ───────────────────────────────────────────────────────────────────────── # Convenience pass-throughs # ─────────────────────────────────────────────────────────────────────────
[docs] def to_dataset_kwargs(self) -> dict[str, Any]: """Return kwargs for :func:`~pycsamt.forward.batch.generate_dataset`. Delegates to :meth:`ForwardConfig.to_dataset_kwargs`. """ return self.forward.to_dataset_kwargs()
[docs] def to_inverter(self): """Return a configured, untrained :class:`~pycsamt.ai.inversion.inv1d.EMInverter1D`. Delegates to :meth:`InversionConfig.to_inverter`. """ return self.inversion.to_inverter()
[docs] def to_fit_kwargs(self) -> dict[str, Any]: """Return kwargs for :meth:`~pycsamt.ai.inversion.inv1d.EMInverter1D.fit`. Delegates to :meth:`InversionConfig.to_fit_kwargs`. """ return self.inversion.to_fit_kwargs()
[docs] def checkpoint_path(self) -> Path | None: """Return the checkpoint file path, or ``None`` if disabled. Delegates to :meth:`InversionConfig.checkpoint_path`. """ return self.inversion.checkpoint_path()
# ───────────────────────────────────────────────────────────────────────── # Config file I/O # ─────────────────────────────────────────────────────────────────────────
[docs] def to_template( self, path: str | Path = "run_config.py", *, fmt: str | None = None, ) -> Path: """Write this run configuration to an annotated source-of-truth file. Parameters ---------- path : path-like, default "run_config.py" Destination. The suffix selects the format (``.py``, ``.json``, ``.yml``). fmt : {"py", "json", "yml", "yaml"}, optional Explicit format override. Returns ------- pathlib.Path """ p = Path(path) fmt_out = _format_from_path(p, fmt) out = _target_path(p, fmt_out) out.parent.mkdir(parents=True, exist_ok=True) fwd_vals = asdict(self.forward) inv_vals = asdict(self.inversion) label = self.name or "PyCSAMT end-to-end experiment" _write_run( out, fwd_vals, inv_vals, name=self.name, description=self.description, title=label, fmt=fmt_out, ) return out
[docs] @classmethod def write_template( cls, path: str | Path = "run_config.py", *, fmt: str | None = None, name: str = "", description: str = "", ) -> Path: """Generate a documented source-of-truth run configuration file. Writes a single file covering both dataset generation and network training with default parameter values and an inline comment for every parameter. Parameters ---------- path : path-like, default "run_config.py" Destination. fmt : {"py", "json", "yml", "yaml"}, optional Explicit format override. name : str Experiment name written into the file header. description : str Free-text description written into the file header. Returns ------- pathlib.Path Examples -------- >>> from pycsamt.ai.inversion.run_config import RunConfig >>> path = RunConfig.write_template("experiment_01.yml") >>> path.suffix '.yml' """ return cls(name=name, description=description).to_template( path, fmt=fmt )
[docs] @classmethod def from_file( cls, path: str | Path, *, strict: bool = True, ) -> RunConfig: """Load a run configuration from a source-of-truth file. Parameters ---------- path : path-like Python, JSON, YML, or YAML file generated by :meth:`write_template` or following the same structure. strict : bool, default True If ``True``, unknown parameter keys raise :class:`ValueError`. If ``False``, unknown keys are silently ignored. Returns ------- RunConfig Examples -------- >>> RunConfig.write_template("run.json") PosixPath('run.json') >>> run = RunConfig.from_file("run.json") >>> run.forward.solver 'mt1d' >>> run.inversion.arch 'resnet' """ p = Path(path) raw_fwd, raw_inv = _read_run(p) fwd = ForwardConfig(**_safe_fields(ForwardConfig, raw_fwd, strict)) inv = InversionConfig(**_safe_fields(InversionConfig, raw_inv, strict)) return cls(forward=fwd, inversion=inv)
#: Alias — matches the convention used by ModEmConfig, OccamConfig, ForwardConfig. read = from_file # ───────────────────────────────────────────────────────────────────────── # repr / summary # ─────────────────────────────────────────────────────────────────────────
[docs] def summary(self) -> str: """Return a human-readable multi-line summary of the full run config.""" lines = ["RunConfig"] if self.name: lines.append(f" name : {self.name}") if self.description: lines.append(f" description : {self.description}") lines.append("") # forward section (indent each line by 2) for line in self.forward.summary().splitlines(): lines.append(" " + line) lines.append("") # inversion section for line in self.inversion.summary().splitlines(): lines.append(" " + line) return "\n".join(lines)
def __repr__(self) -> str: return self.summary()