Source code for cosmofit.core.run

"""
Running an input end to end.
"""

from __future__ import annotations

from .info import load_info, validate_info
from .model import Model
from .mpi import world
from .output import Output
from .registry import resolve


__all__ = ["run"]




def run(info, seed: int | None = None, callback=None):
    """
    Build the model an input describes, run its sampler, and return
    both.

    Parameters
    ----------
    info : dict, str or Path
        The input; it must have a ``sampler`` block -- or a ``post``
        block, to reweight a finished run (see :mod:`core.post`).
    seed : int, optional
        For the sampler's random draws, unless the input sets one.
    callback : callable, optional
        Called with a dict of progress (``steps``, ``Rminus1`` or
        ``tau``, ``acceptance``, ``seconds``, and the stopping
        targets) at every check of a sampler that checks -- for a
        progress bar. In the first process only.

    Returns
    -------
    (dict, Sampler)
        The validated input and the sampler that ran -- its
        ``products()`` hold the result. For a ``post`` input, the
        :class:`~core.post.Post` that ran instead.
    """

    info = load_info(info)

    if "post" in info:

        from .post import Post

        if world().size > 1:
            raise RuntimeError("Post-processing runs in one process; run it without mpirun.")

        post = Post(info)
        post.run()

        return info, post

    info = validate_info(info)

    if not info["sampler"]:
        raise ValueError(
            "run() needs a 'sampler' block (evaluate, minimize, ...). "
            "To build the model alone, use get_model()."
        )

    model = Model(info)

    (name, options), = info["sampler"].items()

    sampler_cls = resolve("sampler", name, options)

    if not getattr(sampler_cls, "parallel", False) and world().size > 1:
        raise RuntimeError(
            f"The {name!r} sampler runs in one process; under mpirun every "
            f"process would run it into the same output. Run it without "
            f"mpirun, or use mcmc."
        )

    options = {k: v for k, v in options.items() if k != "class"}

    sampler = sampler_cls(options, model, seed=seed, output=Output(info))
    sampler.callback = callback

    sampler.run()

    return info, sampler