Skip to content

forward_modeling

Functions for performing forward modeling

get_chis(m, idx, idy, rhs, v, weight, dd=None)

A faster, but more importantly much less memory intensive, way to get chis. The idea is χ2=(d−Am)TN−1(d−Am). Previously we would calculate the residuals d−Am and calculate directly. However d−Am has shape [ndet, nsamp], which is very big. m has shape [nx, ny], much smaller. A, the pointing reconstruction, can be encapsulated into a few pars. Therefore if we can do everthing in map shape, we save a lot on the ammount of stuff we need to put on the GPU. We can expand χ2 as

dTN−1d−2dTN−1Am+mTATN−1Am=dd−2(m⋅rhs)+mm

the first term only has to do with the data. If we care about the absolute value of χ2, which we do at the end, then we can include it in calculation. For MCMC however, we only care about the relative delta chi2 between models. So we can drop that term. For the other terms

mm=mTATN−1Am

This term is essentially what we've been doing before, except that m is now in map shape, whereas before m was in tod shape so we essentially had Am. So we need to do Am, but this is in general fast and Jon has a very fast way of doing it. Finally rhs need only be computed once, and while it is an additional thing to put on the gpu it is small, map shape.

Parameters:

Name Type Description Default
m NDArray[floating]

The model evaluated at all the map pixels

 required
idx NDArray[floating]

tod.info["model_idx"], the x index output by tod_to_index

 required
idy NDArray[floating]

tod.info["model_idy"], the y index output by tod_to_index

 required
rhs NDArray[floating]

The map output of todvec.make_rhs. Note this is how the data enters into the chi2 calc.

 required
v NDArray[floating]

The right singular vectors for the noise SVD. These rotate the data into the basis of the SVD.

 required
weight NDArray[floating]

The noise weights, in fourier space, SVD decomposed.

 required
dd None | floating

Optional chi2 from dd. Not necessary for MCMC but is important for evaluating goodness of fit.

 None
Outputs

chi2 : np.floating The chi2 of the model m to the data.

Source code in witch/forward_modeling.py 
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
@jax.jit
def get_chis(m, idx, idy, rhs, v, weight, dd=None):
    r"""
    A faster, but more importantly much less memory intensive, way to get chis.
    The idea is ${\chi}^{2} = (d-Am)^T N^{-1} (d-Am)$. Previously we would calculate the residuals $d-Am$
    and calculate directly. However $d-Am$ has shape [ndet, nsamp], which is very big. $m$ has shape
    [nx, ny], much smaller. $A$, the pointing reconstruction, can be encapsulated into a few pars.
    Therefore if we can do everthing in map shape, we save a lot on the ammount of stuff we
    need to put on the GPU. We can expand ${\chi}^{2}$ as

    $$
    d^T N^{-1} d -2d^T N^{-1} A m + m^T A^T N^{-1} A m = dd - 2(m \cdot rhs) + mm
    $$

    the first term only has to do with the data. If we care about the absolute value of ${\chi}^2$,
    which we do at the end, then we can include it in calculation. For MCMC however, we only
    care about the relative delta chi2 between models. So we can drop that term. For the other
    terms

    $$
    mm = m^T A^T N^{-1} A m
    $$

    This term is essentially what we've been doing before, except that m is now in map shape,
    whereas before m was in tod shape so we essentially had $Am$. So we need to do $Am$, but this is
    in general fast and Jon has a very fast way of doing it. Finally rhs need only be computed
    once, and while it is an additional thing to put on the gpu it is small, map shape.

    Parameters
    ----------
    m : NDArray[np.floating]
        The model evaluated at all the map pixels
    idx : NDArray[np.floating]
        tod.info["model_idx"], the x index output by tod_to_index
    idy : NDArray[np.floating]
        tod.info["model_idy"], the y index output by tod_to_index
    rhs : NDArray[np.floating]
        The map output of todvec.make_rhs. Note this is how the data enters into the chi2 calc.
    v : NDArray[np.floating]
        The right singular vectors for the noise SVD. These rotate the data into the basis of
        the SVD.
    weight : NDArray[np.floating]
        The noise weights, in fourier space, SVD decomposed.
    dd : None | np.floating
        Optional chi2 from dd. Not necessary for MCMC but is important for evaluating goodness
        of fit.

    Outputs
    -------
    chi2 : np.floating
        The chi2 of the model m to the data.
    """

    model = m.at[idy.astype(int), idx.astype(int)].get(mode="fill", fill_value=0)

    # model = model.at[:,0].set((jnp.sqrt(0.5)*model)[:,0]) #This doesn't actually do anything
    # model = model.at[:,-1].set((jnp.sqrt(0.5)*model)[:,-1])
    model_rot = jnp.dot(v, model)
    tmp = jnp.hstack(
        [model_rot, jnp.fliplr(model_rot[:, 1:-1])]
    )  # mirror pred so we can do dct of first kind
    predft = jnp.real(jnp.fft.rfft(tmp, axis=1))
    nn = predft.shape[1]

    chisq = (
        jnp.sum(weight[:, :nn] * predft**2) - 2 * jnp.dot(rhs.ravel(), m.ravel()) / 2
    )  # Man IDK about this factor of 2

    return chisq

sample(model_params, xyz, beam, params, tods)

Generate a model realization and compute the chis of that model to data.

Arguements:

tods: Array of tod parameters. See prep tods

params: model parameters

model_params: number of each model componant

xyz: grid to evaluate model at

beam: Beam to smooth by

Returns:

chi2: the chi2 difference of the model to the tods
Source code in witch/forward_modeling.py 
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
def sample(model_params, xyz, beam, params, tods):  # , model_params, xyz, beam):
    """
    Generate a model realization and compute the chis of that model to data.

    Arguements:

        tods: Array of tod parameters. See prep tods

        params: model parameters

        model_params: number of each model componant

        xyz: grid to evaluate model at

        beam: Beam to smooth by

    Returns:

        chi2: the chi2 difference of the model to the tods

    """
    log_like = 0
    n_iso, n_gnfw, n_gauss, n_egauss, n_uni, n_expo, n_power, n_power_cos = model_params

    m = model(
        xyz,
        n_iso,
        n_gnfw,
        n_gauss,
        n_egauss,
        n_uni,
        n_expo,
        n_power,
        n_power_cos,
        -2.5e-05,
        beam,
        params,
    )

    for i, tod in enumerate(tods):
        x, y, rhs, v, weight, norm = tod  # unravel tod

        log_like += jget_chis(m, x, y, rhs, v, weight) / norm

    return log_like