cosmotron_mcp.tools.inference

build_cosmosis_inference

build_cosmosis_inference(session_dir: str, sacc_path: str | None = None, sampler: str | None = None, firecrown_config: str | None = None, priors: dict | None = None, sampler_options: dict | None = None, scale_cuts: dict | None = None, camb_accuracy: str = 'default') -> dict

Assemble a CosmoSIS + Firecrown (YAML-factory) inference config for the session SACC.

PREFER the spec chain — draft_inference_spec → update_inference_spec → validate_inference_spec → compile_inference, then run_inference(config_dir=...). This auto-build path (and run_inference with config_dir=None, which calls it internally) is the legacy route: it skips the human spec gate and the deterministic dry-run.

Resolves (explicit arg › task text) the firecrown factory, the priors, the sampler (+ hyperparameters), and inference scale cuts (ℓ-ranges → firecrown data_source.filters; these are INFERENCE-only and never touch the NaMaster measurement lmax). SAVES the experiment YAML + an HDF5 SACC under {session_dir}/inference/; writes values.ini (ranged min default max for prior'd params → SAMPLED, scalar otherwise → FIXED), a priors.ini for gaussian priors, and a firecrown-only pipeline.ini. sampler defaults to test when none is requested.

δz photo-z shift priors: a task ## nz shifts: block (per-bin src0: gaussian(0.0, 0.011) lines, or nz_shifts: <two-column file>) is resolved into pipeline_config.nz_shifts at bootstrap. When present, this ADDS a per-bin PhotoZShiftFactory to the auto-generated firecrown factory for the named probe(s) and a <tracer>_delta_z gaussian prior (→ sampled) — the tool owns the YAML/ini, the model only names the tracer + prior. A user-supplied firecrown config owns its own systematics and is left untouched. No shifts declared → factory unchanged.

If priors are requested from the literature ("same priors as DESI DR2"), returns needs_literature and writes NO chain config — the orchestrator must fetch + human-confirm them first. Config-only — launch with run_inference.

camb_accuracy: "default" (science-grade HMcode2020+feedback) | "fast" (takahashi halofit, coarse k/z) | "linear" (camb linear-only, CCL adds halofit — SUPER fast). The two test profiles are for smoke tests only; pair with sampler="test" or a small n_live.

Parameters:

Name	Type	Description	Default
session_dir	str	Session directory (must have a SACC data vector).	required
sacc_path	str | None	Explicit SACC path; defaults to the session's results/data_vector_sacc.fits.	None
sampler	str | None	CosmoSIS sampler family token (see the cosmosis-sampler skill); defaults to "test".	None
firecrown_config	str | None	Path to a user-supplied Firecrown YAML config; omit to auto-generate one.	None
priors	dict | None	Prior specs, e.g. {"Omega_c": [min, default, max]} or {"sigma8": "gaussian(mean, sigma)"}; resolved from the task text when omitted.	None
sampler_options	dict | None	Sampler hyperparameter overrides.	None
scale_cuts	dict | None	Per-tracer/pair ℓ-range filters, e.g. {"lens0": [20, 150]}; inference-only, never the NaMaster lmax.	None
camb_accuracy	str	"default" (science-grade), "fast", or "linear" — smoke-test profiles only, pair with sampler="test" or a small n_live.	'default'

Returns:

Type	Description
dict	The assembled config paths, or needs_literature (no chain
dict	config written) if priors were requested from the literature and
dict	need human confirmation first.

Source code in cosmotron_mcp/server.py

@mcp.tool()
@sync_budget_guard
def build_cosmosis_inference(
    session_dir: str,
    sacc_path: str | None = None,
    sampler: str | None = None,
    firecrown_config: str | None = None,
    priors: dict | None = None,
    sampler_options: dict | None = None,
    scale_cuts: dict | None = None,
    camb_accuracy: str = "default",
) -> dict:
    """Assemble a CosmoSIS + Firecrown (YAML-factory) inference config for the session SACC.

    PREFER the spec chain — `draft_inference_spec` → `update_inference_spec`
    → `validate_inference_spec` → `compile_inference`, then
    `run_inference(config_dir=...)`. This auto-build path (and
    `run_inference` with `config_dir=None`, which calls it internally) is
    the legacy route: it skips the human spec gate and the deterministic
    dry-run.

    Resolves (explicit arg › task text) the firecrown factory, the priors, the sampler
    (+ hyperparameters), and inference scale cuts (ℓ-ranges → firecrown
    `data_source.filters`; these are INFERENCE-only and never touch the NaMaster
    measurement lmax). SAVES the experiment YAML + an HDF5 SACC under
    `{session_dir}/inference/`; writes `values.ini` (ranged `min default max` for prior'd
    params → SAMPLED, scalar otherwise → FIXED), a `priors.ini` for gaussian priors, and a
    firecrown-only `pipeline.ini`. `sampler` defaults to `test` when none is requested.

    δz photo-z shift priors: a task `## nz shifts:` block (per-bin
    `src0: gaussian(0.0, 0.011)` lines, or `nz_shifts: <two-column file>`) is resolved
    into `pipeline_config.nz_shifts` at bootstrap. When present, this ADDS a per-bin
    `PhotoZShiftFactory` to the auto-generated firecrown factory for the named probe(s)
    and a `<tracer>_delta_z` gaussian prior (→ sampled) — the tool owns the YAML/ini,
    the model only names the tracer + prior. A user-supplied firecrown config owns its
    own systematics and is left untouched. No shifts declared → factory unchanged.

    If priors are requested *from the literature* ("same priors as DESI DR2"), returns
    `needs_literature` and writes NO chain config — the orchestrator must fetch +
    human-confirm them first. Config-only — launch with `run_inference`.

    camb_accuracy: "default" (science-grade HMcode2020+feedback) | "fast" (takahashi
    halofit, coarse k/z) | "linear" (camb linear-only, CCL adds halofit — SUPER fast). The
    two test profiles are for smoke tests only; pair with sampler="test" or a small n_live.

    Args:
        session_dir: Session directory (must have a SACC data vector).
        sacc_path: Explicit SACC path; defaults to the session's
            `results/data_vector_sacc.fits`.
        sampler: CosmoSIS sampler family token (see the `cosmosis-sampler`
            skill); defaults to ``"test"``.
        firecrown_config: Path to a user-supplied Firecrown YAML config;
            omit to auto-generate one.
        priors: Prior specs, e.g. ``{"Omega_c": [min, default, max]}`` or
            ``{"sigma8": "gaussian(mean, sigma)"}``; resolved from the task
            text when omitted.
        sampler_options: Sampler hyperparameter overrides.
        scale_cuts: Per-tracer/pair ℓ-range filters, e.g.
            ``{"lens0": [20, 150]}``; inference-only, never the NaMaster lmax.
        camb_accuracy: ``"default"`` (science-grade), ``"fast"``, or
            ``"linear"`` — smoke-test profiles only, pair with
            ``sampler="test"`` or a small `n_live`.

    Returns:
        The assembled config paths, or ``needs_literature`` (no chain
        config written) if priors were requested from the literature and
        need human confirmation first.
    """
    _require_human_gates(session_dir)
    result = _build_cosmosis_inference(
        session_dir, sacc_path=sacc_path, sampler=sampler, firecrown_config=firecrown_config,
        priors=priors, sampler_options=sampler_options, scale_cuts=scale_cuts,
        camb_accuracy=camb_accuracy)
    if not result.get("needs_literature"):
        _log_tool_call(session_dir, "build_cosmosis_inference", {
            "sacc_path": sacc_path, "sampler": sampler, "firecrown_config": firecrown_config,
            "priors": priors, "sampler_options": sampler_options, "scale_cuts": scale_cuts,
            "camb_accuracy": camb_accuracy,
        })
    return result

run_inference

run_inference(session_dir: str, config_dir: str | None = None, firecrown_config: str | None = None, sampler: str | None = None, priors: dict | None = None, sampler_options: dict | None = None, scale_cuts: dict | None = None, ncores: int | None = None, jobspec: dict | None = None, camb_accuracy: str = 'default') -> dict

Launch a CosmoSIS inference run as a background job; return a job_id immediately.

PREFER passing config_dir=<compiled_dir> from the spec chain (draft_inference_spec → update_inference_spec → validate_inference_spec → compile_inference) — that path carries the human spec gate and the deterministic dry-run. Omitting config_dir auto-builds one from the session SACC via the legacy build_cosmosis_inference route, skipping both.

Resolves/builds the config (from the session SACC if config_dir is omitted), then submits cosmosis <pipeline.ini> via submit_job — NON-BLOCKING. Poll with check_job(session_dir, job_id). For a real (non-test) sampler, call summarise_posterior(session_dir, job_id) once the job is done.

A real sampler runs under mpirun -n ncores (parallel pool). ncores resolves from the explicit arg › a count in the task ("on N cores") › otherwise needs_cores is returned (NO job) so the orchestrator ASKS the user — it carries available + suggested. Surface warning (auto factory); needs_literature likewise means no job ran. The test sampler runs serially (no MPI, no gate).

Parameters:

Name	Type	Description	Default
session_dir	str	Session directory.	required
config_dir	str | None	Pre-built inference config directory; omit to build one from the session SACC.	None
firecrown_config	str | None	Path to a user-supplied Firecrown YAML config; forwarded to build_cosmosis_inference when building fresh.	None
sampler	str | None	CosmoSIS sampler family token; forwarded when building fresh.	None
priors	dict | None	Prior specs; forwarded when building fresh.	None
sampler_options	dict | None	Sampler hyperparameter overrides; forwarded when building fresh.	None
scale_cuts	dict | None	Per-tracer/pair ℓ-range filters; forwarded when building fresh.	None
ncores	int | None	MPI core count for a real (non-test) sampler; resolved from the task text if omitted.	None
jobspec	dict | None	Optional resource dict, forwarded to submit_job. Remote (Slurm): {"backend": "slurm", "site": "<site>"} runs the chain on the cluster — the config is built locally + made portable (env-var-rooted paths), the input closure is staged via rsync, and the job is submitted as the headless cosmotron_mcp.cli run-inference verb. Returns the async shape {job_id, status: "submitted", mode: "async", site, next: "check_job"} — poll with check_job, then summarise_posterior locally on the pulled chain. Requires the site's env: to define FIRECROWN_DIR + COSMOSIS_STD_LIB_DIR (submit fails early otherwise, naming sites.yaml). MPI runs under srun -n N cosmosis --mpi; if ranks come up as N independent chains, set slurm.srun_args: ["--mpi=pmix"] in sites.yaml.	None
camb_accuracy	str	Forwarded to build_cosmosis_inference when building fresh.	'default'

Returns:

Name	Type	Description
Local	dict	``{job_id, status, config_dir, pipeline_ini, run_dir, sampler,
	dict	ncores, chain_path, config_source, warning, needs_literature,
	dict	needs_cores}.needs_cores`` (no job launched) means the
	dict	orchestrator must ask the user how many cores to use. Remote (slurm):
	dict	{job_id, status: "submitted", mode: "async", site, next: "check_job", …}.

Source code in cosmotron_mcp/server.py

@mcp.tool()
@sync_budget_guard
def run_inference(
    session_dir: str,
    config_dir: str | None = None,
    firecrown_config: str | None = None,
    sampler: str | None = None,
    priors: dict | None = None,
    sampler_options: dict | None = None,
    scale_cuts: dict | None = None,
    ncores: int | None = None,
    jobspec: dict | None = None,
    camb_accuracy: str = "default",
) -> dict:
    """Launch a CosmoSIS inference run as a background job; return a job_id immediately.

    PREFER passing `config_dir=<compiled_dir>` from the spec chain
    (`draft_inference_spec` → `update_inference_spec` → `validate_inference_spec`
    → `compile_inference`) — that path carries the human spec gate and the
    deterministic dry-run. Omitting `config_dir` auto-builds one from the
    session SACC via the legacy `build_cosmosis_inference` route, skipping
    both.

    Resolves/builds the config (from the session SACC if `config_dir` is omitted), then
    submits `cosmosis <pipeline.ini>` via `submit_job` — NON-BLOCKING. Poll with
    `check_job(session_dir, job_id)`. For a real (non-`test`) sampler, call
    `summarise_posterior(session_dir, job_id)` once the job is `done`.

    A real sampler runs under `mpirun -n ncores` (parallel pool). `ncores` resolves from
    the explicit arg › a count in the task ("on N cores") › otherwise `needs_cores` is
    returned (NO job) so the orchestrator ASKS the user — it carries `available` +
    `suggested`. Surface `warning` (auto factory); `needs_literature` likewise means no job
    ran. The `test` sampler runs serially (no MPI, no gate).

    Args:
        session_dir: Session directory.
        config_dir: Pre-built inference config directory; omit to build one
            from the session SACC.
        firecrown_config: Path to a user-supplied Firecrown YAML config;
            forwarded to `build_cosmosis_inference` when building fresh.
        sampler: CosmoSIS sampler family token; forwarded when building
            fresh.
        priors: Prior specs; forwarded when building fresh.
        sampler_options: Sampler hyperparameter overrides; forwarded when
            building fresh.
        scale_cuts: Per-tracer/pair ℓ-range filters; forwarded when
            building fresh.
        ncores: MPI core count for a real (non-`test`) sampler; resolved
            from the task text if omitted.
        jobspec: Optional resource dict, forwarded to `submit_job`.
            **Remote (Slurm):** `{"backend": "slurm", "site": "<site>"}` runs the
            chain on the cluster — the config is built locally + made portable
            (env-var-rooted paths), the input closure is staged via rsync, and the
            job is submitted as the headless `cosmotron_mcp.cli run-inference` verb.
            Returns the async shape `{job_id, status: "submitted", mode: "async",
            site, next: "check_job"}` — poll with `check_job`, then
            `summarise_posterior` locally on the pulled chain. Requires the site's
            `env:` to define `FIRECROWN_DIR` + `COSMOSIS_STD_LIB_DIR` (submit fails
            early otherwise, naming `sites.yaml`). MPI runs under `srun -n N cosmosis
            --mpi`; if ranks come up as N independent chains, set
            `slurm.srun_args: ["--mpi=pmix"]` in `sites.yaml`.
        camb_accuracy: Forwarded to `build_cosmosis_inference` when
            building fresh.

    Returns:
        Local: ``{job_id, status, config_dir, pipeline_ini, run_dir, sampler,
        ncores, chain_path, config_source, warning, needs_literature,
        needs_cores}``. ``needs_cores`` (no job launched) means the
        orchestrator must ask the user how many cores to use. Remote (slurm):
        ``{job_id, status: "submitted", mode: "async", site, next: "check_job", …}``.
    """
    _require_human_gates(session_dir)
    result = _run_inference(
        session_dir, config_dir=config_dir, firecrown_config=firecrown_config,
        sampler=sampler, priors=priors, sampler_options=sampler_options,
        scale_cuts=scale_cuts, ncores=ncores, jobspec=jobspec,
        camb_accuracy=camb_accuracy)
    if result.get("job_id") and not result.get("needs_cores") and not result.get("needs_literature"):
        _log_tool_call(session_dir, "run_inference", {
            "config_dir": config_dir, "firecrown_config": firecrown_config,
            "sampler": sampler, "priors": priors, "sampler_options": sampler_options,
            "scale_cuts": scale_cuts, "ncores": ncores, "jobspec": jobspec,
            "camb_accuracy": camb_accuracy,
        })
    return result

summarise_posterior

summarise_posterior(session_dir: str, job_id: str, params: list[str] | None = None, param_labels: dict[str, str] | None = None) -> dict

Summarise a finished inference chain: per-parameter constraints + a corner plot.

Reads the chain the sampler wrote (inference/chain.txt), computes mean ± std and 16/50/84th percentiles per SAMPLED parameter, writes inference/posterior_summary.json, and renders plots/posterior_corner.png (getdist). Registers both artefacts. The test sampler has no chain → returns {status: "no_chain"}.

Derived duplicate columns (CosmoSIS appends uppercase copies such as OMEGA_M, SIGMA_8 after the sampled ones) are deduplicated automatically.

Parameters:

Name	Type	Description	Default
session_dir	str	Session directory.	required
job_id	str	Job id of the finished (or running) run_inference job.	required
params	list[str] | None	Subset of parameters to show in the corner plot, matched case-insensitively against the short column name (part after --). None shows all sampled params. Example: ["omega_c", "sigma_8", "lens0_bias"].	None
param_labels	dict[str, str] | None	Custom LaTeX label overrides keyed by short param name (case-insensitive), bare LaTeX with no surrounding $. Example: {"lens0_bias": "b_0"}. Falls back to the built-in COSMOSIS_PARAM_LATEX map, then the raw name.	None

Returns:

Type	Description
dict	{status, constraints, summary_path, corner_plot}. The test
dict	sampler has no chain → {status: "no_chain"}.

Source code in cosmotron_mcp/server.py

@mcp.tool()
@sync_budget_guard
def summarise_posterior(
    session_dir: str,
    job_id: str,
    params: list[str] | None = None,
    param_labels: dict[str, str] | None = None,
) -> dict:
    """Summarise a finished inference chain: per-parameter constraints + a corner plot.

    Reads the chain the sampler wrote (`inference/chain.txt`), computes mean ± std and
    16/50/84th percentiles per SAMPLED parameter, writes
    `inference/posterior_summary.json`, and renders `plots/posterior_corner.png`
    (getdist). Registers both artefacts. The `test` sampler has no chain → returns
    `{status: "no_chain"}`.

    Derived duplicate columns (CosmoSIS appends uppercase copies such as OMEGA_M,
    SIGMA_8 after the sampled ones) are deduplicated automatically.

    Args:
        session_dir: Session directory.
        job_id: Job id of the finished (or running) `run_inference` job.
        params: Subset of parameters to show in the corner plot, matched
            case-insensitively against the short column name (part after
            ``--``). `None` shows all sampled params. Example:
            ``["omega_c", "sigma_8", "lens0_bias"]``.
        param_labels: Custom LaTeX label overrides keyed by short param
            name (case-insensitive), bare LaTeX with no surrounding ``$``.
            Example: ``{"lens0_bias": "b_0"}``. Falls back to the built-in
            `COSMOSIS_PARAM_LATEX` map, then the raw name.

    Returns:
        ``{status, constraints, summary_path, corner_plot}``. The `test`
        sampler has no chain → ``{status: "no_chain"}``.
    """
    result = _summarise_posterior(session_dir, job_id,
                                  params=params, param_labels=param_labels)
    if result.get("status") not in ("no_chain", "error"):
        _log_tool_call(session_dir, "summarise_posterior",
                       {"job_id": job_id, "params": params, "param_labels": param_labels})
    return result