# 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()