fitter

Master fitting and map making script. You typically want to run the witcher command instead of this.

aic_refit(cfg, metamodel, start, nonpara=False)

Maximum likelihood under the initial noise model, for comparing models with the AIC.

The fit re-estimates the noise from each model's own residuals, so the chi-squared of different models can't be compared. Here the noise is re-estimated from the data alone (as at load time), the free parameters are refit with LM starting from the current parameters (the MCMC medians, or the LM best fit without MCMC), and the result is saved as aic.yaml and par_results_aic.dill in the output directory. The fit's own noise is restored afterwards, nothing else is changed.

Parameters:

Name	Type	Description	Default
cfg	dict	The config, cfg["outdir"] is where the results are written and fitting.maxiter / fitting.chitol set the LM refit.	required
metamodel	MetaModel	The fitted metamodel, its free parameters for the current round are refit.	required
start	str	Description of the starting parameters, saved in aic.yaml.	required
nonpara	bool	If True the outputs get a _nonpara suffix.	False

Returns:

Name	Type	Description
metamodel	MetaModel	The input metamodel with its parameters and fit noise unchanged.

Source code in witch/fitter.py

def aic_refit(cfg, metamodel, start, nonpara=False):
    """
    Maximum likelihood under the initial noise model, for comparing models with the AIC.

    The fit re-estimates the noise from each model's own residuals, so the chi-squared of
    different models can't be compared. Here the noise is re-estimated from the data alone
    (as at load time), the free parameters are refit with LM starting from the current
    parameters (the MCMC medians, or the LM best fit without MCMC), and the result is saved
    as `aic.yaml` and `par_results_aic.dill` in the output directory. The fit's own noise
    is restored afterwards, nothing else is changed.

    Parameters
    ----------
    cfg : dict
        The config, `cfg["outdir"]` is where the results are written and
        `fitting.maxiter` / `fitting.chitol` set the LM refit.
    metamodel : MetaModel
        The fitted metamodel, its free parameters for the current round are refit.
    start : str
        Description of the starting parameters, saved in `aic.yaml`.
    nonpara : bool, default: False
        If True the outputs get a `_nonpara` suffix.

    Returns
    -------
    metamodel : MetaModel
        The input metamodel with its parameters and fit noise unchanged.
    """
    print_once("Refitting under the initial noise model for the AIC")
    suffix = "_nonpara" * nonpara
    start_pars = jnp.array(comm.bcast(np.array(metamodel.parameters), root=0))
    free = np.array(metamodel.to_fit)

    # Noise models live in pytree aux data, so drop jit caches keyed on the old ones
    _initial_noise(metamodel)
    jax.clear_caches()
    start_mm = copy(metamodel).update(
        start_pars, metamodel.errs, metamodel.cov, metamodel.chisq
    )
    chisq_start = float(joint_objective(start_mm, True, False, False)[0])
    fitted, iterations, delta_chisq, lmd = run_lmfit(
        start_mm,
        eval(str(cfg["fitting"].get("maxiter", "10"))),
        eval(str(cfg["fitting"].get("chitol", "1e-5"))),
    )
    chisq_refit = float(joint_objective(fitted, True, False, False)[0])
    _save_outputs(cfg, fitted, f"aic{suffix}")

    if comm.Get_rank() == 0:
        k = int(free.sum())
        pars, errs = np.array(fitted.parameters), np.array(fitted.errs)
        result = {
            "noise": "initial: estimated from the data with no model subtracted",
            "start": start,
            "k": k,
            "chisq_start": chisq_start,
            "chisq_refit": chisq_refit,
            "aic": chisq_refit + 2 * k,
            "lm": {
                "iterations": int(iterations),
                "delta_chisq": float(delta_chisq),
                "lambda": float(lmd),
            },
            # A list in parameter order, names alone can repeat across structures
            "parameters": [
                {
                    "index": i,
                    "name": name,
                    "start": float(start_pars[i]),
                    "refit": float(pars[i]),
                    "err": float(errs[i]),
                }
                for i, name in enumerate(fitted.par_names)
                if free[i]
            ],
        }
        path = _ours(os.path.join(cfg["outdir"], f"aic{suffix}.yaml"))
        with open(path, "w") as file:
            yaml.dump(result, file, sort_keys=False)
        print_once(
            f"AIC chisq: {chisq_start:.2f} at the start, {chisq_refit:.2f} after "
            f"{iterations} LM iterations (lambda {float(lmd):.3g}), saved to {path}"
        )

    # Back to the fit's noise model for the rest of the pipeline
    metamodel = _reestimate_noise(metamodel)
    jax.clear_caches()
    return metamodel

deep_merge(a, b)

Based on https://gist.github.com/angstwad/bf22d1822c38a92ec0a9?permalink_comment_id=3517209

Source code in witch/fitter.py

def deep_merge(a: dict, b: dict) -> dict:
    """
    Based on https://gist.github.com/angstwad/bf22d1822c38a92ec0a9?permalink_comment_id=3517209
    """
    result = deepcopy(a)
    for bk, bv in b.items():
        av = result.get(bk)
        if isinstance(av, dict) and isinstance(bv, dict):
            result[bk] = deep_merge(av, bv)
        else:
            result[bk] = deepcopy(bv)
    return result

load_config(start_cfg, cfg_path)

We want to load a config and if it has the key "base", load that as well and merge them. We only want to take things from base that are not in the original config so we merge the original into the newly loaded one.

Source code in witch/fitter.py

def load_config(start_cfg, cfg_path):
    """
    We want to load a config and if it has the key "base",
    load that as well and merge them.
    We only want to take things from base that are not in the original config
    so we merge the original into the newly loaded one.
    """
    with open(cfg_path) as file:
        new_cfg = yaml.safe_load(file)
    cfg = deep_merge(new_cfg, start_cfg)
    if "base" in new_cfg:
        base_path = new_cfg["base"]
        if not os.path.isabs(base_path):
            base_path = os.path.join(os.path.dirname(cfg_path), base_path)
        return load_config(cfg, base_path)
    return cfg

print_once(*args)

Helper function to print only once when running with MPI. Only the rank 0 process will print.

Parameters:

Name	Type	Description	Default
*args	Unpack[tuple[Any, ...]]	Arguments to pass to print.	()

Source code in witch/fitter.py

def print_once(*args: Unpack[tuple[Any, ...]]):
    """
    Helper function to print only once when running with MPI.
    Only the rank 0 process will print.

    Parameters
    ----------
    *args : Unpack[tuple[Any, ...]]
        Arguments to pass to print.
    """
    if comm.Get_rank() == 0:
        print(*args)
        sys.stdout.flush()

process_maps(cfg, dataset, metamodel, sim_noise_kwargs=None)

Prepare the maps of a dataset: simulate them if requested and compute the noise model.

Parameters:

Name	Type	Description	Default
cfg	dict	The config.	required
dataset	DataSet	The map dataset.	required
metamodel	MetaModel	The model, only needed when simulating.	required
sim_noise_kwargs	dict	Overrides for the noise kwargs when drawing the simulated noise, so a sim can use a different noise model than the fit (e.g. another CMB cosmology through dl). Only used for noise models with a sample method when wnoise is False.	None

Returns:

Name	Type	Description
mapset	SolutionSet	The maps with their noise models computed.

Source code in witch/fitter.py

def process_maps(cfg, dataset, metamodel, sim_noise_kwargs=None):
    """
    Prepare the maps of a dataset: simulate them if requested and compute the noise model.

    Parameters
    ----------
    cfg : dict
        The config.
    dataset : DataSet
        The map dataset.
    metamodel : MetaModel
        The model, only needed when simulating.
    sim_noise_kwargs : dict, optional
        Overrides for the noise kwargs when drawing the simulated noise, so a sim can use a
        different noise model than the fit (e.g. another CMB cosmology through `dl`).
        Only used for noise models with a `sample` method when `wnoise` is False.

    Returns
    -------
    mapset : SolutionSet
        The maps with their noise models computed.
    """
    mapset = dataset.datavec
    info = dataset.info
    sim = cfg.get("sim", False)
    if metamodel is None and sim:
        raise ValueError("metamodel cannot be None when simming!")
    metamodel_cur = metamodel
    xfer = info.get("info", "")
    if xfer:
        metamodel_cur = deepcopy(metamodel)
        metamodel_cur.models = tuple(
            Model_xfer.from_parent(model, xfer) for model in metamodel_cur.models
        )

    dset_ind = metamodel_cur.get_dataset_ind(dataset.name)
    for i, imap in enumerate(mapset):
        if sim:
            if jnp.all(imap.ivar == 0):
                print_once(
                    f"ivar for map {imap.name} is all 0! Filling with a dummy value but you should check your maps!"
                )
                imap.ivar = imap.ivar.at[:].set(cfg.get("default_ivar", 1e8))
            # Pixels with ivar == 0 have no weight in the likelihood, so leave them noiseless
            # rather than letting an inf noise level propagate into the data
            good = imap.ivar > 0
            scale = jnp.where(
                good, 1.0 / jnp.sqrt(jnp.where(good, imap.ivar, 1.0)), 0.0
            )
            print("Map noise: ", float(jnp.mean(scale[good])))
            sample = getattr(dataset.noise_class, "sample", None)
            if sample is not None and not cfg["wnoise"]:
                # The noise model can draw its own correlated noise (e.g. NoiseCMBMap)
                args, kwargs = resolve_noise_args(
                    imap, dataset.noise_args, dataset.noise_kwargs
                )
                # Lets a sim use a different noise model than the fit, to test a mismatch
                kwargs |= sim_noise_kwargs or {}
                key = jax.random.key(cfg.get("sim_seed", 0))
                imap.data = sample(key, imap.data, *args, **kwargs)
            else:
                imap.data = (
                    jax.random.normal(
                        jax.random.key(cfg.get("sim_seed", 0)),
                        shape=imap.data.shape,
                        dtype=imap.data.dtype,
                    )
                    * scale
                )
                if not cfg["wnoise"]:
                    print_once(
                        "Noise model cannot be sampled... simming map with white noise"
                    )

            pred = metamodel_cur.model_proj(dset_ind, i)
            imap.data = imap.data + pred
        print("Map scale: ", jnp.mean(jnp.abs(imap.data)))
        imap.compute_noise(
            dataset.noise_class, None, *dataset.noise_args, **dataset.noise_kwargs
        )
    return mapset

resolve_noise_args(imap, args, kwargs)

Evaluate noise arguments that refer to the map, like compute_noise does.

Note

This is a temporary copy of the self... evaluation inside jitkasi's WCSMap.compute_noise, needed because the sim calls a noise model's sample directly. Once jitkasi exposes that evaluation as a public helper, import it from there and delete this function (and its use in process_maps). Until then keep the two in step: a string argument that starts with self is evaluated with self set to the map.

Parameters:

Name	Type	Description	Default
imap	WCSMap	The map that self refers to in the arguments.	required
args	tuple	Positional noise arguments, strings starting with self are evaluated.	required
kwargs	dict	Keyword noise arguments, string values starting with self are evaluated.	required

Returns:

Name	Type	Description
args	list	The evaluated positional arguments.
kwargs	dict	The evaluated keyword arguments.

Source code in witch/fitter.py

def resolve_noise_args(imap, args: tuple, kwargs: dict) -> tuple[list, dict]:
    """
    Evaluate noise arguments that refer to the map, like `compute_noise` does.

    Note
    ----
    This is a temporary copy of the `self...` evaluation inside jitkasi's
    `WCSMap.compute_noise`, needed because the sim calls a noise model's `sample` directly.
    Once jitkasi exposes that evaluation as a public helper, import it from there and delete
    this function (and its use in `process_maps`). Until then keep the two in step: a string
    argument that starts with `self` is evaluated with `self` set to the map.

    Parameters
    ----------
    imap : WCSMap
        The map that `self` refers to in the arguments.
    args : tuple
        Positional noise arguments, strings starting with `self` are evaluated.
    kwargs : dict
        Keyword noise arguments, string values starting with `self` are evaluated.

    Returns
    -------
    args : list
        The evaluated positional arguments.
    kwargs : dict
        The evaluated keyword arguments.
    """

    def ev(a):
        return eval(a, {"self": imap}) if isinstance(a, str) and a[:4] == "self" else a

    return [ev(a) for a in args], {k: ev(v) for k, v in kwargs.items()}