Skip to content

cfg_loader

Model_xfer dataclass

Bases: Model

Source code in witch/containers/model.py 
701
702
703
704
705
706
707
708
709
710
711
712
713
714
715
716
717
718
719
720
721
722
723
724
725
726
727
728
729
730
731
732
733
734
735
736
737
738
739
740
741
742
743
744
745
746
747
748
749
750
751
752
753
754
755
756
757
758
759
760
761
762
763
764
765
766
767
768
769
770
771
772
773
774
775
776
777
778
@dataclass
class Model_xfer(Model):

    ks: jax.Array = field(default_factory=jnp.array([0]).copy)  # scalar float array
    xfer_vals: jax.Array = field(
        default_factory=jnp.array([1]).copy
    )  # scalar float array

    def __post_init__(self):
        pass

    @classmethod
    def from_parent(cls, parent, xfer_str) -> Self:
        xfer = load_xfer(xfer_str)

        pixsize = np.abs(parent.xyz[1][0][1] - parent.xyz[1][0][0])
        ks = xfer[0, 0:] * pixsize  # This picks up an extra dim?
        xfer_vals = xfer[1, 0:]

        my_dict = {}
        for key in parent.__dataclass_fields__.keys():
            if parent.__dataclass_fields__[key].init:
                my_dict[key] = deepcopy(parent.__dict__[key])

        return cls(**my_dict, ks=ks.ravel(), xfer_vals=xfer_vals)

    @cached_property
    def model(self) -> jax.Array:
        cur_map = core.model(
            self.xyz,
            tuple(self.n_struct),
            tuple(self.n_rbins),
            self.dz,
            *self.pars,
        )
        # Code from JMP, whoever that is, by way of Charles
        farr = np.fft.fft2(cur_map)
        nx, ny = cur_map.shape
        kx = np.outer(np.fft.fftfreq(nx), np.zeros(ny).T + 1.0)
        ky = np.outer(np.zeros(nx).T + 1.0, np.fft.fftfreq(ny))
        k = np.sqrt(kx * kx + ky * ky)

        filt = self.table_filter_2d(k)
        farr *= filt

        return np.real(np.fft.ifft2(farr))

    def table_filter_2d(self, k) -> jax.Array:
        f = interpolate.interp1d(self.ks, self.xfer_vals)
        kbin_min = self.ks.min()
        kbin_max = self.ks.max()

        filt = k * 0.0
        filt[(k >= kbin_min) & (k <= kbin_max)] = f(
            k[(k >= kbin_min) & (k <= kbin_max)]
        )
        filt[(k < kbin_min)] = self.xfer_vals[self.ks == kbin_min]
        filt[(k > kbin_max)] = self.xfer_vals[self.ks == kbin_max]

        return filt

    @cached_property
    def model_grad(self) -> None:
        """
        The evaluated model and its gradient, see `core.model_grad` for details.
        Note that this is cached, but is automatically reset whenever
        `update` is called or `cur_round` changes. Currently computing
        grad for models with transfer function is not supported.

        Returns
        -------
        None
        """

        raise TypeError(
            "Error; Grad cannot currently be computed on Models with transfer function"
        )
        return None  # Shouldnt get here

model_grad cached property

The evaluated model and its gradient, see core.model_grad for details. Note that this is cached, but is automatically reset whenever update is called or cur_round changes. Currently computing grad for models with transfer function is not supported.

Returns:

Type Description
None

deep_merge(a, b)

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

Source code in witch/fitter.py 
688
689
690
691
692
693
694
695
696
697
698
699
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

joint_objective(metamodel, do_loglike=True, do_grad=True, do_curve=True)

Compute the objective for multiple datasets in an MPI aware way.

Parameters:

Name Type Description Default
metamodel MetaModel

The MetaModel we are trying to fit.

 required
do_loglike bool

If True then we will compute the log-likelihood between the model and the data.

 True
do_grad bool

If True then compute the gradient of chi-squared with respect to the model parameters.

 True
do_curve bool

If True than compute the curvature of chi-squared with respect to the model parameters.

 True

Returns:

Name Type Description
loglike Array

The log-likelihood between the model and data. If do_loglike is False then this is jnp.array(0).
grad Array

The gradient of the parameters at there current values. If do_grad is False then this is an array of zeros. This is a (npar,) array.
curve Array

The curvature of the parameter space at the current values. If do_curve is False then this is an array of zeros. This is a (npar, npar) array.
Source code in witch/objective.py 
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
@partial(jax.jit, static_argnames=("do_loglike", "do_grad", "do_curve"))
def joint_objective(
    metamodel: "MetaModel",
    do_loglike: bool = True,
    do_grad: bool = True,
    do_curve: bool = True,
) -> tuple[jax.Array, jax.Array, jax.Array]:
    """
    Compute the objective for multiple datasets in an MPI aware way.


    Parameters
    ----------
    metamodel : MetaModel
        The `MetaModel` we are trying to fit.
    do_loglike : bool, default: True
        If True then we will compute the log-likelihood between
        the model and the data.
    do_grad : bool, default: True
        If True then compute the gradient of chi-squared with
        respect to the model parameters.
    do_curve : bool, default: True
        If True than compute the curvature of chi-squared with
        respect to the model parameters.

    Returns
    -------
    loglike : jax.Array
        The log-likelihood between the model and data.
        If `do_loglike` is `False` then this is `jnp.array(0)`.
    grad : jax.Array
        The gradient of the parameters at there current values.
        If `do_grad` is `False` then this is an array of zeros.
        This is a `(npar,)` array.
    curve : jax.Array
        The curvature of the parameter space at the current values.
        If `do_curve` is `False` then this is an array of zeros.
        This is a `(npar, npar)` array.
    """
    global_comm = metamodel.global_comm
    npar = len(metamodel.parameters)
    loglike = jnp.array(0)
    grad = jnp.zeros(npar)
    curve = jnp.zeros((npar, npar))
    for i in range(len(metamodel.datasets)):
        _loglike, _grad, _curve = metamodel.datasets[i].objective(
            metamodel,
            i,
            do_loglike,
            do_grad,
            do_curve,
        )
        loglike += _loglike
        grad = grad.at[:].add(_grad)
        curve = curve.at[:].add(_curve)
    mpi4jax.barrier(comm=global_comm)
    if do_loglike:
        loglike = mpi4jax.allreduce(loglike, MPI.SUM, comm=global_comm)
    if do_grad:
        grad = mpi4jax.allreduce(grad, MPI.SUM, comm=global_comm)
    if do_curve:
        curve = mpi4jax.allreduce(curve, MPI.SUM, comm=global_comm)

    return loglike, grad, curve

load_cfg(cfg, args=None)

Set up datasets and the metamodel from a loaded config, exactly as the fitter does before fitting.

Parameters:

Name Type Description Default
cfg dict

Config as returned by witch.fitter.load_config. Modified in place.

 required
args Namespace

Parsed command line arguments. If given, wnoise, nosub and nofit override the config.

 None

Returns:

Name Type Description
datasets tuple[DataSet]

Processed datasets living on this MPI task.
outdir str

Output directory.
dset_names list[str]

Names of all datasets in the config.
metamodel MetaModel

The metamodel. The (possibly split) MPI communicator is stored in witch.fitter.comm.
Source code in witch/cfg_loader.py 
 14
 15
 16
 17
 18
 19
 20
 21
 22
 23
 24
 25
 26
 27
 28
 29
 30
 31
 32
 33
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
def load_cfg(cfg, args=None):
    """
    Set up datasets and the metamodel from a loaded config, exactly as the
    fitter does before fitting.

    Parameters
    ----------
    cfg : dict
        Config as returned by `witch.fitter.load_config`. Modified in place.
    args : argparse.Namespace, optional
        Parsed command line arguments. If given, `wnoise`, `nosub` and
        `nofit` override the config.

    Returns
    -------
    datasets : tuple[DataSet]
        Processed datasets living on this MPI task.
    outdir : str
        Output directory.
    dset_names : list[str]
        Names of all datasets in the config.
    metamodel : MetaModel
        The metamodel. The (possibly split) MPI communicator is stored in
        `witch.fitter.comm`.
    """
    if "MJ_TODROOT" in os.environ:
        if "WITCH_DATROOT" not in os.environ:
            print_once(
                "You are using MJ_TODROOT, this is depreciated in favor of WITCH_DATROOT. Setting WITCH_DATROOT to the value of MJ_TODROOT..."
            )
            os.environ["WITCH_DATROOT"] = os.environ["MJ_TODROOT"]
        else:
            print_once(
                "Both WITCH_DATROOT and MJ_TODROOT provided! MJ_TODROOT is depreciated, using WITCH_DATROOT..."
            )

    if "MJ_OUTROOT" in os.environ:
        if "WITCH_OUTROOT" not in os.environ:
            print_once(
                "You are using MJ_OUTROOT, this is depreciated in favor of WITCH_OUTROOT. Setting WITCH_OUTROOT to the value of MJ_OUTROOT..."
            )
            os.environ["WITCH_OUTROOT"] = os.environ["MJ_OUTROOT"]
        else:
            print_once(
                "Both WITCH_OUTROOT and MJ_OUTROOT provided! MJ_OUTROOT is depreciated, using WITCH_OUTROOT..."
            )

    # TODO: Serialize cfg to a data class (pydantic?)
    cfg["fit"] = cfg.get("fit", "model" in cfg)
    cfg["sim"] = cfg.get("sim", False)
    cfg["wnoise"] = cfg.get("wnoise", False)
    cfg["map"] = cfg.get("map", True)
    cfg["sub"] = cfg.get("sub", True)
    cfg["mem_starved"] = cfg.get("mem_starved", False)
    if args is not None:
        if args.wnoise:
            cfg["wnoise"] = True
        if args.nosub:
            cfg["sub"] = False
        if args.nofit:
            cfg["fit"] = False

    # Do imports
    for module, name in cfg.get("imports", {}).items():
        mod = import_module(module)
        if isinstance(name, str):
            globals()[name] = mod
        elif isinstance(name, list):
            for n in name:
                locals()[n] = getattr(mod, n)
        else:
            raise TypeError("Expect import name to be a string or a list")

    # Get the functions needed to work with our dataset
    dset_names = list(cfg["datasets"].keys())
    datasets = []
    outdir = None
    # First lets plan how to split things up among mpi tasks
    fnames = {}
    for dset_name in dset_names:
        get_files = eval(cfg["datasets"][dset_name]["funcs"]["get_files"])
        fnames[dset_name] = get_files(dset_name, cfg)
    # Update the fitter's comm so print_once and the fit loop use the split comm
    fnames, fitter.comm, comms_local = _mpi_fsplit(fnames, fitter.comm)
    comm = fitter.comm

    for dset_name in dset_names:
        # If this dataset doesn't live in this proc then skip
        if len(fnames[dset_name]) == 0:
            continue

        if "load_tods" in cfg["datasets"][dset_name]["funcs"]:
            cfg["datasets"][dset_name]["funcs"]["load"] = cfg["datasets"][dset_name][
                "funcs"
            ]["load_tods"]
        elif "load_maps" in cfg["datasets"][dset_name]["funcs"]:
            cfg["datasets"][dset_name]["funcs"]["load"] = cfg["datasets"][dset_name][
                "funcs"
            ]["load_maps"]
        get_files = eval(cfg["datasets"][dset_name]["funcs"]["get_files"])
        load = eval(cfg["datasets"][dset_name]["funcs"]["load"])
        get_info = eval(cfg["datasets"][dset_name]["funcs"]["get_info"])
        make_metadata = eval(cfg["datasets"][dset_name]["funcs"]["make_metadata"])
        preproc = eval(cfg["datasets"][dset_name]["funcs"]["preproc"])
        postproc = eval(cfg["datasets"][dset_name]["funcs"]["postproc"])
        postfit = eval(cfg["datasets"][dset_name]["funcs"]["postfit"])

        # Get data, info, and metadata
        datavec = load(dset_name, cfg, fnames[dset_name], comms_local[dset_name])
        info = get_info(dset_name, cfg, datavec)
        metadata = make_metadata(dset_name, cfg, info)

        # Setup noise
        noise_class = eval(str(cfg["datasets"][dset_name]["noise"]["class"]))
        noise_args = tuple(eval(str(cfg["datasets"][dset_name]["noise"]["args"])))
        noise_kwargs = eval(str(cfg["datasets"][dset_name]["noise"]["kwargs"]))

        # Init dataset
        dataset = DataSet(
            dset_name,
            get_files,
            load,
            get_info,
            make_metadata,
            preproc,
            postproc,
            postfit,
            comm,
            info,
        )

        # Init the data container
        dataset.data = DataSetData(
            dset_name,
            info["mode"],
            info["objective"],
            datavec,
            metadata,
            noise_class,
            noise_args,
            noise_kwargs,
        )

        if "base" in cfg.keys():
            del cfg[
                "base"
            ]  # We've collated to the cfg files so no need to keep the base
        datasets += [dataset]

    datasets = tuple(datasets)
    # Define the model and get stuff setup fitting
    if "metamodel" in cfg:
        metamodel = MetaModel.from_config(
            comm,
            cfg,
            tuple(dataset.data for dataset in datasets),
        )
    else:
        if "model" in cfg:
            raise ValueError(
                "No MetaModel defined but Model definition found. This appears to be a legacy configuration! Please update!"
            )
        metamodel = MetaModel(
            comm,
            tuple(),
            tuple(dataset.data for dataset in datasets),
            tuple(),
            tuple(),
            tuple(),
            jnp.zeros(0),
            jnp.zeros(0),
            jnp.zeros(0),
            jnp.zeros(0),
        )
        print_once("No metamodel defined, setting fit, sim, and sub to False")
        cfg["fit"] = False
        cfg["sim"] = False
        cfg["sub"] = False
    outdir = get_outdir(cfg, metamodel)
    cfg["outdir"] = outdir

    # Now process
    for dataset in datasets:
        os.makedirs(os.path.join(outdir, dataset.name), exist_ok=True)
        dataset.info["outdir"] = outdir
        comm.barrier()
        dataset.preproc(dataset, cfg, metamodel)
        comm.barrier()
        if dataset.mode == "tod":
            dataset.data.datavec = process_tods(cfg, dataset, metamodel)
        elif dataset.mode == "map":
            dataset.data.datavec = process_maps(cfg, dataset, metamodel)
        comm.barrier()
        dataset = jax.block_until_ready(dataset)
        dataset.postproc(dataset, cfg, metamodel)
    metamodel.datasets = tuple(dataset.data for dataset in datasets)

    return datasets, outdir, dset_names, metamodel

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 
702
703
704
705
706
707
708
709
710
711
712
713
714
715
716
717
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

para_to_non_para(model, n_rounds, to_copy, rmax, struct_num, sig_params, default, mm_cfg)

Function which approximately converts cluster profiles into a non-parametric form. Note this is only approximate and should be fit afterwords.

Parameters:

Name Type Description Default
model Model

The parametric model to start from.

 required
n_rounds Optional[int]

Number of rounds to fit for output model. If none, copy from self

 required
to_copy list[str]

List of structures, by name, to copy.

 required
rmax float

Maximum radius of the rbins

 required
struct_num int

Structure within model to calculate rbins on

 required
sig_params list[int]

Parameters to consider for computing significance. Only first match will be used.

 required
default tuple[int, ...]

Default rbins to be returned if generation fails.

 required
mm_cfg dict

MetaModel config, parameter_map will be modified if need be.

 required

Returns:

Name Type Description
Model Model

Model with a non-parametric representation of input model

Raises:

Type Description
ValueError

If there are no models to copy
Source code in witch/nonparametric.py 
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
def para_to_non_para(
    model,
    n_rounds: Optional[int],
    to_copy: list[str],
    rmax: float,
    struct_num: int,
    sig_params: list[int],
    default: tuple[int, ...],
    mm_cfg: dict,
) -> Model:
    """
    Function which approximately converts cluster profiles into a non-parametric form. Note this is
    only approximate and should be fit afterwords.

    Parameters
    ----------
    model : Model
        The parametric model to start from.
    n_rounds: Optional int | None
        Number of rounds to fit for output model. If none, copy from self
    to_copy : list[str]
        List of structures, by name, to copy.
    rmax : float
        Maximum radius of the rbins
    struct_num : int
        Structure within model to calculate rbins on
    sig_params: list[str]
        Parameters to consider for computing significance.
        Only first match will be used.
    default: tuple[int, ...]
        Default rbins to be returned if generation fails.
    mm_cfg: dict
        MetaModel config, parameter_map will be modified if need be.
    Returns
    -------
    Model : Model
        Model with a non-parametric representation of input model
    Raises
    ------
    ValueError
        If there are no models to copy
    """
    cur_model = deepcopy(
        model
    )  # Make a copy of model, we don't want to lose structures
    i = 0  # Make sure we keep at least one struct
    to_remove = []
    for structure in cur_model.structures:
        if structure.structure not in to_copy:
            to_remove.append(structure.name)

    if len(to_remove) == len(cur_model.structures):
        raise ValueError("Error: no model structures in {}".format(to_copy))
    for struct in to_remove:
        cur_model.remove_struct(struct)
    params = jnp.array(cur_model.pars)
    params = jnp.ravel(params)
    pressure = core.model3D(
        cur_model.xyz, tuple(cur_model.n_struct), tuple(cur_model.n_rbins), params
    )
    pressure = pressure[
        ..., int(pressure.shape[2] / 2)
    ]  # Take middle slice. Close enough is good enough here, dont care about rounding

    pixsize = np.abs(cur_model.xyz[1][0][1] - cur_model.xyz[1][0][0]).item()

    rs, bin1d, _ = wu.bin_map(pressure, pixsize)

    rbins = get_rbins(cur_model, rmax, struct_num, sig_params, default)
    rbins = np.append(rbins, np.array([np.amax(rs)]))

    condlist = [
        jnp.array((rbins[i] <= rs) & (rs < rbins[i + 1]))
        for i in range(len(rbins) - 2, -1, -1)
    ]

    amps, pows, c = profile_to_broken_power(rs, bin1d, condlist, rbins)

    priors = (-1 * np.inf, np.inf)
    if n_rounds is None:
        n_rounds = model.n_rounds
    if not isinstance(n_rounds, int):
        raise ValueError("Non int n_rounds")
    parameters = [
        Parameter(
            "rbins",
            tuple([False] * n_rounds),
            jnp.atleast_1d(jnp.array(rbins[:-1], dtype=float)),  # Drop last bin
            jnp.zeros_like(jnp.atleast_1d(jnp.array(rbins[:-1])), dtype=float),
            jnp.array(priors, dtype=float),
        ),
        Parameter(
            "amps",
            tuple([True] * n_rounds),
            jnp.atleast_1d(jnp.array(amps, dtype=float)),
            jnp.zeros_like(jnp.atleast_1d(jnp.array(amps)), dtype=float),
            jnp.array(priors, dtype=float),
        ),
        Parameter(
            "pows",
            tuple([True] * n_rounds),
            jnp.atleast_1d(jnp.array(pows, dtype=float)),
            jnp.zeros_like(jnp.atleast_1d(jnp.array(pows)), dtype=float),
            jnp.array(priors, dtype=float),
        ),
        Parameter(
            "dx",  # TODO: miscentering
            tuple([False] * n_rounds),
            jnp.atleast_1d(jnp.array(0, dtype=float)),
            jnp.zeros_like(jnp.atleast_1d(jnp.array(0)), dtype=float),
            jnp.array(priors, dtype=float),
        ),
        Parameter(
            "dy",  # TODO: miscentering
            tuple([False] * n_rounds),
            jnp.atleast_1d(jnp.array(0, dtype=float)),
            jnp.zeros_like(jnp.atleast_1d(jnp.array(0)), dtype=float),
            jnp.array(priors, dtype=float),
        ),
        Parameter(
            "dz",  # TODO: miscentering
            tuple([False] * n_rounds),
            jnp.atleast_1d(jnp.array(0, dtype=float)),
            jnp.zeros_like(jnp.atleast_1d(jnp.array(0)), dtype=float),
            jnp.array(priors, dtype=float),
        ),
        Parameter(
            "c",
            tuple([True] * n_rounds),
            jnp.atleast_1d(jnp.array(c, dtype=float)),
            jnp.zeros_like(jnp.atleast_1d(jnp.array(0)), dtype=float),
            jnp.array(priors, dtype=float),
        ),
    ]

    structure = Structure(
        "nonpara_power", "nonpara_power", parameters, n_rbins=len(rbins) - 1
    )
    structures = list(model.structures)
    if "parameter_map" in mm_cfg:
        mm_cfg["parameter_map"] = [
            [p for p in pm if f".{structures[struct_num].name}." in p]
            for pm in mm_cfg["parameter_map"]
        ]
    structures[struct_num] = structure

    return Model(
        name=model.name,
        structures=tuple(structures),
        xyz=model.xyz,
        dz=model.dz,
        n_rounds=n_rounds,
        cur_round=0,
        to_run=model.to_run,
    )

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 
36
37
38
39
40
41
42
43
44
45
46
47
48
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()

run_lmfit(metamodel, maxiter=10, chitol=1e-05)

Fit a set of models to datasets jointly. This uses a modified Levenberg–Marquardt fitter with flat priors. This function is MPI aware.

Parameters:

Name Type Description Default
metamodel MetaModel

The MetaModel that describes the models and datasets.

 required
maxiter int

The maximum number of iterations to fit.

 10
chitol float

The delta chisq to use as the convergence criteria.

 1e-5

Returns:

Name Type Description
models tuple[Model, ...]

Model with the final set of fit parameters, errors, and chisq.
final_iter int

The number of iterations the fitter ran for.
delta_chisq float

The final delta chisq.
curve Array

The final covariance matrix.
Source code in witch/fitting.py 
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
def run_lmfit(
    metamodel: MetaModel,
    maxiter: int = 10,
    chitol: float = 1e-5,
) -> tuple[MetaModel, int, jax.Array]:
    """
    Fit a set of models to datasets jointly.
    This uses a modified Levenberg–Marquardt fitter with flat priors.
    This function is MPI aware.

    Parameters
    ----------
    metamodel : MetaModel
        The MetaModel that describes the models and datasets.
    maxiter : int, default: 10
        The maximum number of iterations to fit.
    chitol : float, default: 1e-5
        The delta chisq to use as the convergence criteria.

    Returns
    -------
    models : tuple[Model, ...]
        Model with the final set of fit parameters, errors, and chisq.
    final_iter : int
        The number of iterations the fitter ran for.
    delta_chisq : float
        The final delta chisq.
    curve : jax.Array
        The final covariance matrix.
    """
    zero = jnp.array(0.0, jnp.float32)
    chitol = jnp.float32(chitol)
    tf = np.where(np.array(metamodel.to_fit))[0]

    # Every rank has to make the same choice here since the chunked path reduces
    # with mpi4py and the jitted path with mpi4jax. In a joint fit the map ranks
    # follow the TOD ranks onto the chunked path (one chunk per map).
    local_multiple_tods = any(
        dataset.mode == "tod" and len(dataset.datavec.tods) > 1
        for dataset in metamodel.datasets
    )
    if isinstance(metamodel.global_comm, NullComm):
        use_tod_chunks = local_multiple_tods
    else:
        use_tod_chunks = (
            metamodel.global_comm.allreduce(int(local_multiple_tods), op=MPI.MAX) > 0
        )
    if use_tod_chunks:
        return _run_lmfit_tod_chunks(metamodel, maxiter, chitol)

    @jax.jit
    def _cond_func(val):
        i, delta_chisq, lmd, *_ = val
        iterbool = jax.lax.lt(i, maxiter)
        chisqbool = jax.lax.ge(delta_chisq, chitol) + jax.lax.gt(lmd, zero)
        return iterbool * chisqbool

    @jax.jit
    def _body_func(val):
        i, delta_chisq, lmd, metamodel, curve, grad = val
        curve_use = curve.at[:].add(lmd * jnp.diag(jnp.diag(curve)))
        # Get the step
        step = jnp.dot(
            invscale(curve_use.at[tf, :].get().at[:, tf].get(), do_invsafe=True),
            grad.at[tf].get(),
        )
        new_pars, to_fit = _prior_pars_fit(
            metamodel.priors,
            metamodel.parameters.at[tf].add(step),
            jnp.array(metamodel.to_fit),
        )
        # Get errs
        full_cov = invscale(curve_use, do_invsafe=True)
        errs = jnp.where(to_fit, jnp.sqrt(jnp.diag(full_cov)), 0)
        cov = jnp.where(to_fit, full_cov, 0)
        # Now lets get an updated model
        new_metamodel = copy(metamodel).update(
            pars=new_pars, errs=errs, cov=cov, chisq=metamodel.chisq
        )
        new_chisq, new_grad, new_curve = joint_objective(
            new_metamodel, True, True, True
        )
        new_metamodel = copy(new_metamodel).update(
            pars=new_pars, errs=errs, cov=cov, chisq=new_chisq
        )

        new_delta_chisq = jnp.astype(metamodel.chisq - new_metamodel.chisq, jnp.float32)
        metamodel, grad, curve, delta_chisq, lmd = jax.lax.cond(
            new_delta_chisq > 0,
            _success,
            _failure,
            metamodel,
            new_metamodel,
            grad,
            new_grad,
            curve,
            new_curve,
            delta_chisq,
            new_delta_chisq,
            lmd,
        )
        stalled = _stalled(new_delta_chisq, chitol)
        lmd = jnp.where(stalled, zero, lmd)
        delta_chisq = jnp.where(stalled, 0.0, delta_chisq)

        return (i + 1, delta_chisq, lmd, metamodel, curve, grad)

    pars, _ = _prior_pars_fit(
        metamodel.priors, metamodel.parameters, jnp.array(metamodel.to_fit)
    )
    # Recompute chisq at the starting pars so that it matches the current noise model,
    # the stored value may be from before the noise was reestimated
    chisq, grad, curve = joint_objective(metamodel, True, True, True)
    jax.block_until_ready((chisq, grad, curve))
    metamodel = metamodel.update(
        pars=pars, errs=metamodel.errs, cov=metamodel.cov, chisq=chisq
    )
    i, delta_chisq, lmd, metamodel, *_ = jax.lax.while_loop(
        _cond_func,
        _body_func,
        (0, jnp.astype(jnp.inf, jnp.float32), zero.copy(), metamodel, curve, grad),
    )

    return metamodel, i, delta_chisq, lmd

run_mcmc(metamodel, num_steps=5000, num_leaps=10, step_size=jnp.array(0.02), sample_which=-1, burn_in=0.1, max_tries=20)

Run Hamilonian Monte Carlo to estimate the posterior for our model. Currently this function only support flat priors, but more will be supported down the line. In order to ensure accuracy of the noise model used, it is reccomended that you run at least one round of fit_tods followed by noise reestimation before this function.

This is MPI aware. Eventually this will be replaced with something more jaxy.

Parameters:

Name Type Description Default
metamodel MetaModel

The MetaModel that describes the models and datasets.

 required
num_steps int

The number of steps to run MCMC for.

 5000
num_leaps int

The number of leapfrog steps to take at each sample.

 10
step_size Array

The step size to use in the leapfrog algorithm. This should be tuned to get an acceptance fraction of ~.65.

 .02
sample_which int

Sets which parameters to sample. If this is >= 0 then we will sample which ever parameters were fit in that round of fitting. If this is -1 then we will sample which ever parameters were fit in the last round of fitting. If this is -2 then any parameters that were ever fit will be sampled. If this is <= -3 or >= model.n_rounds then all parameters are sampled.

 -1,
burn_in float

Fractional burn-in period of samples to discard

 0.1
max_tries int

Number of tries to tune step size that will be attemted. If 0 no tuning will be run.

 10

Returns:

Name Type Description
metamodel MetaModel

The metamodel with MCMC estimated parameters and errors. The parameters are estimated as the mean of the samples. The errors are estimated as the standard deviation. This also has the chi-squared of the estimated parameters.
flat_samples Array

Array of samples from running MCMC.
Source code in witch/fitting.py 
565
566
567
568
569
570
571
572
573
574
575
576
577
578
579
580
581
582
583
584
585
586
587
588
589
590
591
592
593
594
595
596
597
598
599
600
601
602
603
604
605
606
607
608
609
610
611
612
613
614
615
616
617
618
619
620
621
622
623
624
625
626
627
628
629
630
631
632
633
634
635
636
637
638
639
640
641
642
643
644
645
646
647
648
649
650
651
652
653
654
655
656
657
658
659
660
661
662
663
664
665
666
667
668
669
670
671
672
673
674
675
676
677
678
679
680
681
682
683
684
685
686
687
688
689
690
691
692
693
694
695
696
697
698
699
700
701
702
703
704
705
706
707
708
709
710
711
712
713
714
715
716
717
718
719
720
721
722
723
724
725
726
727
728
729
730
731
732
733
734
735
736
737
738
739
740
741
742
743
744
745
746
747
748
749
750
751
752
753
754
755
756
757
758
759
760
761
762
763
764
765
766
767
768
769
770
771
772
773
774
775
776
777
778
779
780
781
782
783
784
785
786
787
788
789
790
791
792
793
794
795
796
797
798
799
800
801
802
803
804
805
806
807
808
809
810
811
812
813
814
815
816
817
818
819
820
821
822
823
824
825
826
def run_mcmc(
    metamodel: MetaModel,
    num_steps: int = 5000,
    num_leaps: int = 10,
    step_size: jax.Array = jnp.array(0.02),
    sample_which: int = -1,
    burn_in: float = 0.1,
    max_tries: int = 20,
) -> tuple[MetaModel, jax.Array]:
    """
    Run Hamilonian Monte Carlo to estimate the posterior for our model.
    Currently this function only support flat priors, but more will be supported
    down the line. In order to ensure accuracy of the noise model used, it is
    reccomended that you run at least one round of `fit_tods` followed by noise
    reestimation before this function.

    This is MPI aware.
    Eventually this will be replaced with something more jaxy.

    Parameters
    ----------
    metamodel : MetaModel
        The MetaModel that describes the models and datasets.
    num_steps : int, default: 5000
        The number of steps to run MCMC for.
    num_leaps : int, default: 10
        The number of leapfrog steps to take at each sample.
    step_size : lax.Array, default: .02
        The step size to use in the leapfrog algorithm.
        This should be tuned to get an acceptance fraction of ~.65.
    sample_which : int, default: -1,
        Sets which parameters to sample.
        If this is >= 0 then we will sample which ever parameters were
        fit in that round of fitting.
        If this is -1 then we will sample which ever parameters were fit
        in the last round of fitting.
        If this is -2 then any parameters that were ever fit will be sampled.
        If this is <= -3 or >= `model.n_rounds` then all parameters are sampled.
    burn_in : float, default: 0.1
        Fractional burn-in period of samples to discard
    max_tries : int, default: 10
        Number of tries to tune step size that will be attemted.
        If 0 no tuning will be run.

    Returns
    -------
    metamodel : MetaModel
        The metamodel with MCMC estimated parameters and errors.
        The parameters are estimated as the mean of the samples.
        The errors are estimated as the standard deviation.
        This also has the chi-squared of the estimated parameters.
    flat_samples : jax.Array
        Array of samples from running MCMC.

    """
    global_comm = metamodel.global_comm
    mpi4jax.barrier(comm=global_comm)
    rank = global_comm.Get_rank()
    step_size = jnp.array(step_size)

    if burn_in >= 1.0 or burn_in < 0:
        raise ValueError("Error: burn_in must be in range [0, 1)")

    if sample_which >= 0 and sample_which < metamodel.n_rounds:
        metamodel.set_round(sample_which)
        to_fit = metamodel.to_fit
    elif sample_which == -1:
        to_fit = metamodel.to_fit
    elif sample_which == -2:
        to_fit = metamodel.to_fit_ever
    else:
        to_fit = jnp.ones_like(metamodel.to_fit_ever, dtype=bool)
    to_fit = jnp.array(to_fit)
    metamodel = metamodel.add_round(jnp.array(to_fit))

    init_pars = jnp.array(metamodel.parameters)
    init_errs = jnp.zeros_like(init_pars)
    init_cov = jnp.zeros((len(init_pars), len(init_pars)))
    final_pars = init_pars.copy()
    final_errs = init_errs.copy()

    prior_l, prior_u = metamodel.priors
    scale = (jnp.abs(prior_l) + jnp.abs(prior_u)) / 2.0
    scale = jnp.where(jnp.isfinite(scale), scale, init_pars)
    scale = jnp.where(scale == 0, 1, scale)
    init_pars = init_pars.at[:].multiply(1.0 / scale)

    @jax.jit
    def _log_prob_and_grad(pars, metamodel=metamodel, init_pars=init_pars):
        full_pars = init_pars.at[to_fit].set(pars)
        full_pars = full_pars.at[:].multiply(scale)
        _, in_bounds = _prior_pars_fit(
            metamodel.priors, full_pars, jnp.array(metamodel.to_fit)
        )
        log_prior = jnp.sum(
            jnp.where(
                jnp.array(jnp.where(metamodel.to_fit, in_bounds, True)), 0, -1 * jnp.inf
            )
        )
        temp_metamodel = copy(metamodel).update(
            full_pars, init_errs, init_cov, jnp.array(0.0)
        )
        chisq, grad, _ = joint_objective(temp_metamodel, True, True, False)
        log_like = -0.5 * chisq + log_prior
        grad = grad.at[:].multiply(1.0 / scale)
        log_like_grad = grad.at[to_fit].get().ravel()
        log_like_grad = log_like_grad.at[:].add(log_prior)
        return log_like, log_like_grad

    def _get_step_size(step_size, key, comm, max_tries):
        mpi4jax.barrier(comm=global_comm)

        def _can_interp(probs):
            probs = jnp.array(probs)
            msk = jnp.isfinite(probs) * probs > 0
            if jnp.sum(msk) < 2:
                return False
            dprob = (probs[msk] - 0.63) / 0.13
            for n in range(1, 4):
                upper, lower = (
                    jnp.sum((dprob <= n) * (dprob >= 0)),
                    jnp.sum((dprob >= -1 * n) * (dprob < 0)),
                )
                if upper >= 1 and lower >= 1 and upper + lower >= n + 1:
                    return True
            return False

        def _hmc(n_steps, step_size):
            _, chain_probs = hmc(
                init_pars.at[to_fit].get().ravel(),
                _log_prob_and_grad,
                num_steps=n_steps,
                num_leaps=num_leaps,
                step_size=step_size,
                comm=comm,
                key=key,
            )
            prob = 0.0
            if rank == 0:
                prob = jnp.mean(chain_probs)
            prob = jnp.array(mpi4jax.bcast(prob, 0, comm=comm))

            return prob

        if max_tries == 0:
            return step_size
        step_fac = 1
        n_steps = 1
        i = 0
        step_sizes = []
        probs = []
        for i in range(max_tries):
            mpi4jax.barrier(comm=global_comm)
            step_size *= step_fac
            if rank == 0:
                print(f"Trying step size {step_size}")
            prob = _hmc(n_steps, step_size)
            step_sizes += [step_size]
            probs += [prob]

            if prob == 0 or not jnp.isfinite(prob):
                step_fac = 1e-1
                n_steps = 1
            elif prob > 0.6 and prob < 0.66:
                break
            elif _can_interp(probs):
                break
            elif prob <= 0.1:
                step_fac = 1 / 5
                n_steps = 3
            elif prob >= 0.9:
                step_fac = 5
                n_steps = 3
            elif prob <= 0.5:
                step_fac = 1 / 3
                n_steps = 5
            elif prob >= 0.76:
                step_fac = 3
                n_steps = 5
            elif prob > 0.5 and prob < 0.63:
                step_fac = 1 / 1.5
                n_steps = 10
            elif prob < 0.76 and prob > 0.63:
                step_fac = 1 / 1.5
                n_steps = 10
            else:
                raise ValueError(f"Unclear what to do with probability {prob}")
        if probs[-1] > 0.60 and probs[-1] < 0.66:
            if rank == 0:
                print(
                    f"Final step size of {step_size} with an accept probability of {probs[-1]}, took {i} tries to find."
                )
                sys.stdout.flush()
            return step_sizes[-1]
        mpi4jax.barrier(comm=global_comm)
        probs = jnp.array(probs)
        step_sizes = jnp.array(step_sizes)
        msk = jnp.isfinite(probs) & (probs > 0)
        srt = jnp.argsort(probs[msk])
        if jnp.sum(msk) == 0:
            raise ValueError(
                "No finite non-zero step sizes found! Manual intervention is needed!"
            )
        step_size_interp = jnp.interp(
            jnp.array([0.63]), probs[msk][srt], step_sizes[msk][srt]
        )
        step_size = step_size_interp[0].item()
        prob = _hmc(5, step_size)
        if rank == 0:
            print(
                f"Interpolated to get a step size of {step_size} with an acceptance probability {prob:.2%}"
            )
            sys.stdout.flush()
        return step_size

    key = jax.random.PRNGKey(0)
    key = jnp.array(mpi4jax.bcast(key, 0, comm=global_comm))

    step_size = _get_step_size(
        step_size=step_size,
        key=key,
        comm=global_comm,
        max_tries=max_tries,
    )

    chain, _ = hmc(
        init_pars.at[to_fit].get().ravel(),
        _log_prob_and_grad,
        num_steps=num_steps,
        num_leaps=num_leaps,
        step_size=step_size,
        comm=global_comm,
        key=key,
    )
    flat_samples = chain.at[:].multiply(scale.at[to_fit].get())

    fit_idx = np.flatnonzero(np.array(to_fit))
    final_cov = jnp.zeros((len(final_pars), len(final_pars)))
    if rank == 0:
        burn_in = int(num_steps * burn_in)
        flat_samples = flat_samples[burn_in:]
        final_pars = final_pars.at[to_fit].set(jnp.median(flat_samples, axis=0).ravel())
        final_errs = final_errs.at[to_fit].set(jnp.std(flat_samples, axis=0).ravel())
        # The sample covariance only covers the fit parameters, embed it in the full matrix
        final_cov = final_cov.at[fit_idx[:, None], fit_idx[None, :]].set(
            jnp.atleast_2d(jnp.cov(flat_samples, rowvar=False))
        )
    final_pars = jnp.array(mpi4jax.bcast(final_pars, 0, comm=global_comm))
    final_errs = jnp.array(mpi4jax.bcast(final_errs, 0, comm=global_comm))
    final_cov = jnp.array(mpi4jax.bcast(final_cov, 0, comm=global_comm))
    # Placeholder cov and chisq so we can evaluate the chisq at the final parameters
    metamodel = metamodel.update(
        final_pars.block_until_ready(),
        final_errs.block_until_ready(),
        jnp.zeros_like(final_cov),
        jnp.array(0.0),
    )
    chisq, *_ = joint_objective(metamodel, True, False, False)
    metamodel = metamodel.update(final_pars, final_errs, final_cov, chisq)
    jax.block_until_ready(metamodel)

    return metamodel, flat_samples

run_mh(metamodel, num_samples=10000, num_chains=4, burn_in=0.1, bound=2, seed=0, prior_type='uniform', prior_err_scale=1.0, joint=True, dataset_ind=0, progress=True, proposal='cov')

Run Metropolis-Hastings chains and summarize them into the metamodel.

The chains start at the current metamodel.parameters (e.g. the Levenberg-Marquardt best fit) and only parameters with to_fit set are varied. Mirrors the interface of witch.fitting.run_mcmc.

Parameters:

Name Type Description Default
metamodel MetaModel

Model to sample.

 required
num_samples int

Samples per chain.

 10000
num_chains int

Number of chains.

 4
burn_in float

Fraction of each chain discarded.

 0.1
bound float

Scale for unbounded uniform proposals.

 2
seed int

Base seed; chain i uses seed + i.

 0
prior_type (uniform, normal)

Prior and proposal type.

 "uniform"
prior_err_scale float

Multiplier on metamodel.errs for proposal width and normal prior.

 1.0
joint bool

Use the MPI-aware joint likelihood (chains run serially). If False chains run in threads on a single dataset.

 True
dataset_ind int

Dataset used when joint is False.

 0
progress bool

Show progress bars.

 True
proposal (cov, diag)

cov proposes steps from the LM covariance in metamodel.cov (falls back to diag with the chosen prior_type if unusable); diag uses independent steps.

 "cov"

Returns:

Name Type Description
metamodel MetaModel

Metamodel with median parameters, std errors, covariance and chisq.
flat_samples ndarray

Post burn-in samples of the fitted parameters, (nsamp, nfit).
chains ndarray

All chains, (num_chains, num_samples, num_parameters).
acceptance_rates ndarray

Acceptance rate of each chain.
Source code in witch/metropolis_hastings.py 
513
514
515
516
517
518
519
520
521
522
523
524
525
526
527
528
529
530
531
532
533
534
535
536
537
538
539
540
541
542
543
544
545
546
547
548
549
550
551
552
553
554
555
556
557
558
559
560
561
562
563
564
565
566
567
568
569
570
571
572
573
574
575
576
577
578
579
580
581
582
583
584
585
586
587
588
589
590
591
592
593
594
595
596
597
598
599
600
601
602
603
604
605
606
607
608
609
610
611
612
613
614
615
def run_mh(
    metamodel,
    num_samples=10000,
    num_chains=4,
    burn_in=0.1,
    bound=2,
    seed=0,
    prior_type="uniform",
    prior_err_scale=1.0,
    joint=True,
    dataset_ind=0,
    progress=True,
    proposal="cov",
):
    """
    Run Metropolis-Hastings chains and summarize them into the metamodel.

    The chains start at the current ``metamodel.parameters`` (e.g. the
    Levenberg-Marquardt best fit) and only parameters with ``to_fit`` set are
    varied. Mirrors the interface of ``witch.fitting.run_mcmc``.

    Parameters
    ----------
    metamodel : witch.containers.MetaModel
        Model to sample.
    num_samples : int, default=10000
        Samples per chain.
    num_chains : int, default=4
        Number of chains.
    burn_in : float, default=0.1
        Fraction of each chain discarded.
    bound : float, default=2
        Scale for unbounded uniform proposals.
    seed : int, default=0
        Base seed; chain ``i`` uses ``seed + i``.
    prior_type : {"uniform", "normal"}, default="uniform"
        Prior and proposal type.
    prior_err_scale : float, default=1.0
        Multiplier on ``metamodel.errs`` for proposal width and normal prior.
    joint : bool, default=True
        Use the MPI-aware joint likelihood (chains run serially). If False
        chains run in threads on a single dataset.
    dataset_ind : int, default=0
        Dataset used when ``joint`` is False.
    progress : bool, default=True
        Show progress bars.
    proposal : {"cov", "diag"}, default="cov"
        ``cov`` proposes steps from the LM covariance in ``metamodel.cov``
        (falls back to ``diag`` with the chosen prior_type if unusable); ``diag`` uses independent steps.

    Returns
    -------
    metamodel : witch.containers.MetaModel
        Metamodel with median parameters, std errors, covariance and chisq.
    flat_samples : np.ndarray
        Post burn-in samples of the fitted parameters, ``(nsamp, nfit)``.
    chains : np.ndarray
        All chains, ``(num_chains, num_samples, num_parameters)``.
    acceptance_rates : np.ndarray
        Acceptance rate of each chain.
    """
    if joint:
        results = run_chains_serial(
            metamodel,
            num_samples,
            num_chains,
            bound,
            seed,
            prior_type,
            prior_err_scale,
            progress,
            proposal,
        )
    else:
        results = run_chains_parallel(
            metamodel,
            num_samples,
            num_chains,
            bound,
            seed,
            dataset_ind,
            prior_type,
            prior_err_scale,
            proposal,
        )
    chains = np.array([r[0] for r in results])
    acceptance_rates = np.array([r[1] for r in results])

    to_fit = np.array(metamodel.to_fit)
    start = int(burn_in * num_samples)
    flat_samples = chains[:, start:, :][:, :, to_fit].reshape(-1, int(to_fit.sum()))

    pars = np.array(metamodel.parameters)
    errs = np.zeros_like(pars)
    cov = np.zeros((len(pars), len(pars)))
    pars[to_fit] = np.median(flat_samples, axis=0)
    errs[to_fit] = np.std(flat_samples, axis=0)
    cov[np.ix_(to_fit, to_fit)] = np.atleast_2d(np.cov(flat_samples.T))
    chisq = -2 * calc_like_joint(metamodel, jnp.array(pars))
    metamodel = metamodel.update(
        pars=jnp.array(pars), errs=jnp.array(errs), cov=jnp.array(cov), chisq=chisq
    )
    return metamodel, flat_samples, chains, acceptance_rates