Skip to content

gauss

Gaussian and multipole expansion beam fitting utilities.

The Gaussian beam model is fit directly to a map using weighted least squares. A fitted Gaussian can then be used as the base beam for a multipole expansion, whose amplitudes are obtained by a weighted linear decomposition. See the individual fitting and model-conversion functions for details.

fit_gauss_map(imap, ivar, posmap, guess, map_units='pW', force_sym=False, mask_size=-1)

Fit a 2D Gaussian beam to a map.

The fit minimizes the inverse-variance-weighted residual between imap and a Gaussian beam. The Gaussian can optionally be constrained to be symmetric, and the fit can optionally be restricted to a circular region around the initial beam center.

Parameters:

Name Type Description Default
imap Float[ndmap, 'ny nx']

Input beam map.

 required
ivar Float[ndmap, 'ny nx']

Inverse-variance map.

 required
posmap Float[ndmap, '2 ny nx']

Beam coordinates. The first component is eta and the second is xi.

 required
guess AxisManager

Initial Gaussian parameters. Must contain xi0, eta0, off, amp, fwhm_xi, fwhm_eta, and phi.

 required
map_units str

Unit of the map amplitude and offset parameters.

 "pW"
force_sym bool

If True, constrain the fitted Gaussian to have equal FWHM in xi and eta and zero rotation angle.

 False
mask_size float

If positive, only pixels within this radius of the initial beam center are included in the fit.

 -1

Returns:

Name Type Description
fit_params Optional[AxisManager]

Fitted Gaussian parameters. Positional parameters are stored in radians. None if the fit fails.
model Optional[Float[ndmap, 'ny nx']]

The fitted Gaussian evaluated on posmap. None if the fit fails.
Source code in lat_beams/fitting/map/gauss.py 
 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
def fit_gauss_map(
    imap: Float[ndmap, "ny nx"],
    ivar: Float[ndmap, "ny nx"],
    posmap: Float[ndmap, "2 ny nx"],
    guess: AxisManager,
    map_units: str = "pW",
    force_sym: bool = False,
    mask_size: float = -1,
) -> tuple[Optional[AxisManager], Optional[Float[ndmap, "ny nx"]]]:
    """Fit a 2D Gaussian beam to a map.

    The fit minimizes the inverse-variance-weighted residual between `imap`
    and a Gaussian beam. The Gaussian can optionally be constrained to be
    symmetric, and the fit can optionally be restricted to a circular region
    around the initial beam center.

    Parameters
    ----------
    imap : Float[ndmap, "ny nx"]
        Input beam map.
    ivar : Float[ndmap, "ny nx"]
        Inverse-variance map.
    posmap : Float[ndmap, "2 ny nx"]
        Beam coordinates. The first component is `eta` and the second is
        `xi`.
    guess : AxisManager
        Initial Gaussian parameters. Must contain `xi0`, `eta0`, `off`,
        `amp`, `fwhm_xi`, `fwhm_eta`, and `phi`.
    map_units : str, default: "pW"
        Unit of the map amplitude and offset parameters.
    force_sym : bool, default: False
        If `True`, constrain the fitted Gaussian to have equal FWHM in
        `xi` and `eta` and zero rotation angle.
    mask_size : float, default: -1
        If positive, only pixels within this radius of the initial beam center
        are included in the fit.

    Returns
    -------
    fit_params : Optional[AxisManager]
        Fitted Gaussian parameters. Positional parameters are stored in
        radians. `None` if the fit fails.
    model : Optional[Float[ndmap, "ny nx"]]
        The fitted Gaussian evaluated on `posmap`. `None` if the fit
        fails.
    """

    y, x = posmap
    x0 = [
        cast(float, guess.xi0),
        cast(float, guess.eta0),
        cast(float, guess.off),
        cast(float, guess.amp),
        cast(float, guess.fwhm_xi),
        cast(float, guess.fwhm_eta),
        cast(float, guess.phi),
    ]
    bounds = [
        [
            np.min(x) - cast(float, guess.fwhm_xi),
            np.min(y) - cast(float, guess.fwhm_eta),
            -5 * np.max(np.abs(imap)),
            0,
            cast(float, guess.fwhm_xi) / 3,
            cast(float, guess.fwhm_eta) / 3,
            0,
        ],
        [
            np.max(x) + cast(float, guess.fwhm_xi),
            np.max(y) + cast(float, guess.fwhm_eta),
            5 * np.max(imap),
            5 * np.max(imap),
            cast(float, guess.fwhm_xi) * 3,
            cast(float, guess.fwhm_eta) * 3,
            2 * np.pi,
        ],
    ]
    map_unit = u.Unit(map_units)
    par_names = [
        "xi0",
        "eta0",
        "off",
        "amp",
        "fwhm_xi",
        "fwhm_eta",
        "phi",
    ]  # , "wing_r0", "wing_amp"]
    par_units = [
        u.radian,
        u.radian,
        map_unit,
        map_unit,
        u.radian,
        u.radian,
        u.radian,
    ]  # , u.radian, map_unit]  # type: ignore
    if force_sym:
        x0 = x0[:-2]
        bounds[0] = bounds[0][:-2]
        bounds[1] = bounds[1][:-2]
    bounds = [(lb, ub) for lb, ub in zip(*bounds)]

    # Mask out things too far from the starting center
    if mask_size > 0:
        r = np.sqrt((x - x0[0]) ** 2 + (y - x0[1]) ** 2)
        msk = r < mask_size

        fit_imap = np.asarray(imap)[msk]
        fit_ivar = np.asarray(ivar)[msk]
        fit_posmap = np.asarray(posmap)[:, msk]
    else:
        fit_imap = np.asarray(imap)
        fit_ivar = np.asarray(ivar)
        fit_posmap = np.asarray(posmap)

    w = np.sqrt(fit_ivar)

    def _to_pars(coeffs):
        dx, dy, off, amp = coeffs[:4]

        if force_sym:
            fwhm_xi = fwhm_eta = coeffs[4]
            phi = 0
        else:
            fwhm_xi, fwhm_eta, phi = coeffs[4:]

        return dx, dy, off, amp, fwhm_xi, fwhm_eta, phi

    def _resid(coeffs):
        dx, dy, off, amp, fwhm_xi, fwhm_eta, phi = _to_pars(coeffs)
        beam = gaussian2d(fit_posmap, amp, dx, dy, fwhm_xi, fwhm_eta, phi, off)

        return (w * (fit_imap - beam)).ravel()

    # def _objective(
    #     coeffs,
    # ):
    #     resid = _resid(coeffs)
    #     chisq = np.nansum(resid**2)
    #     return chisq

    res = least_squares(
        _resid,
        x0,
        bounds=np.array(bounds).T,
        method="trf",
        x_scale="jac",
    )
    if not res.success:
        print(res)
        return None, None

    # Convert to aman
    aman = AxisManager()
    dx, dy, off, amp, fwhm_xi, fwhm_eta, phi = pars = _to_pars(res.x)
    for n, un, v in zip(par_names, par_units, pars):
        aman.wrap(n, v * un)
    model = gaussian2d(posmap, amp, dx, dy, fwhm_xi, fwhm_eta, phi, off)

    return aman, model

fit_multipole_map(imap, ivar, posmap, guess, map_units='pW', base_beam=None, n_multipoles=5)

Fit a multipole expansion of a beam model to a map.

The input map is modeled as a multipole expansion of base_beam. If no base beam is provided, a unit-amplitude Gaussian is constructed from guess. The multipole amplitudes are fit using inverse-variance weighting.

Parameters:

Name Type Description Default
imap Float[ndmap, 'ny nx']

Input beam map.

 required
ivar Float[ndmap, 'ny nx']

Inverse-variance map.

 required
posmap Float[ndmap, '2 ny nx']

Beam coordinates. The first component is eta and the second is xi.

 required
guess AxisManager

Initial Gaussian beam parameters used to construct base_beam when it is not provided.

 required
map_units str

Unit of the fitted multipole amplitudes.

 "pW"
base_beam Optional[Float[ndmap, 'ny nx']]

Base beam whose angular dependence is expanded. If None, a unit-amplitude Gaussian is constructed from guess.

 None
n_multipoles int

Number of multipoles to fit. 0 fits only the monopole, 1 adds the dipole, and so on.

 5

Returns:

Name Type Description
fit_params AxisManager

Fitted multipole amplitudes stored in amps with shape (n_multipoles, 2). The second dimension contains the cosine and sine amplitudes. The AxisManager also contains multipoles and term axes.
model Float[ndmap, 'ny nx']

Beam model evaluated with the fitted multipole amplitudes.
Source code in lat_beams/fitting/map/gauss.py 
188
189
190
191
192
193
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
def fit_multipole_map(
    imap: Float[ndmap, "ny nx"],
    ivar: Float[ndmap, "ny nx"],
    posmap: Float[ndmap, "2 ny nx"],
    guess: AxisManager,
    map_units: str = "pW",
    base_beam: Optional[Float[ndmap, "ny nx"]] = None,
    n_multipoles: int = 5,
) -> tuple[AxisManager, Float[ndmap, "ny nx"]]:
    """
    Fit a multipole expansion of a beam model to a map.

    The input map is modeled as a multipole expansion of `base_beam`.
    If no base beam is provided, a unit-amplitude Gaussian is constructed
    from `guess`. The multipole amplitudes are fit using inverse-variance
    weighting.

    Parameters
    ----------
    imap : Float[ndmap, "ny nx"]
        Input beam map.
    ivar : Float[ndmap, "ny nx"]
        Inverse-variance map.
    posmap : Float[ndmap, "2 ny nx"]
        Beam coordinates. The first component is `eta` and the second is
        `xi`.
    guess : AxisManager
        Initial Gaussian beam parameters used to construct `base_beam` when
        it is not provided.
    map_units : str, default: "pW"
        Unit of the fitted multipole amplitudes.
    base_beam : Optional[Float[ndmap, "ny nx"]], default: None
        Base beam whose angular dependence is expanded. If `None`, a
        unit-amplitude Gaussian is constructed from `guess`.
    n_multipoles : int, default: 5
        Number of multipoles to fit. `0` fits only the monopole, `1` adds
        the dipole, and so on.

    Returns
    -------
    fit_params : AxisManager
        Fitted multipole amplitudes stored in `amps` with shape
        `(n_multipoles, 2)`. The second dimension contains the cosine and
        sine amplitudes. The `AxisManager` also contains `multipoles` and
        `term` axes.
    model : Float[ndmap, "ny nx"]
        Beam model evaluated with the fitted multipole amplitudes.
    """

    if base_beam is None:
        base_beam = gaussian2d(
            posmap,
            1,
            cast(u.Quantity, guess.xi0).value,
            cast(u.Quantity, guess.eta0).value,
            cast(u.Quantity, guess.fwhm_xi).value,
            cast(u.Quantity, guess.fwhm_eta).value,
            cast(u.Quantity, guess.phi).value,
            0,
        )
    y, x = posmap
    theta = np.arctan2(
        y - cast(u.Quantity, guess.eta0).to(u.radian).value,
        x - cast(u.Quantity, guess.xi0).to(u.radian).value,
    )

    # Compute model
    if n_multipoles == 0:
        amps = np.array([[cast(u.Quantity, guess.amp).value, 0]])
    else:
        amps = multipole_decomp(base_beam, imap, ivar, n_multipoles, theta, True)
    model = imap.copy()
    model[...] = multipole_expansion(base_beam, amps, theta)
    model = cast(
        ndmap,
        multipole_expansion(base_beam, amps, theta),
    )

    # Convert to aman
    m_units = u.Unit(map_units)
    aman = AxisManager()
    mp_ax = IndexAxis("multipoles", n_multipoles)
    sc_ax = LabelAxis("term", ["cos", "sin"])
    aman.wrap("amps", amps * m_units, [(0, mp_ax), (1, sc_ax)])

    return aman, model

gaussian2d_from_aman(posmap, aman)

Evaluate a Gaussian beam from an AxisManager.

Parameters:

Name Type Description Default
posmap Float[ndmap, '2 ny nx']

Beam coordinates. The first component is eta and the second is xi.

 required
aman AxisManager

AxisManager containing the Gaussian parameters. If it contains a gaussian field, that field is used. Otherwise the parameters are read directly from aman.

 required

Returns:

Name Type Description
beam Float[ndmap, 'ny nx']

Gaussian beam evaluated at posmap.
Source code in lat_beams/fitting/map/gauss.py 
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
def gaussian2d_from_aman(
    posmap: Float[ndmap, "2 ny nx"],
    aman: AxisManager,
) -> Float[ndmap, "ny nx"]:
    """Evaluate a Gaussian beam from an AxisManager.

    Parameters
    ----------
    posmap : Float[ndmap, "2 ny nx"]
        Beam coordinates. The first component is `eta` and the second is
        `xi`.
    aman : AxisManager
        AxisManager containing the Gaussian parameters. If it contains a
        `gaussian` field, that field is used. Otherwise the parameters are
        read directly from `aman`.

    Returns
    -------
    beam : Float[ndmap, "ny nx"]
        Gaussian beam evaluated at `posmap`.
    """
    if "gaussian" in aman._fields:
        aman = aman.gaussian

    return gaussian2d(
        posmap,
        cast(u.Quantity, aman.amp).value,
        cast(u.Quantity, aman.xi0).to(u.radian).value,
        cast(u.Quantity, aman.eta0).to(u.radian).value,
        cast(u.Quantity, aman.fwhm_xi).to(u.radian).value,
        cast(u.Quantity, aman.fwhm_eta).to(u.radian).value,
        cast(u.Quantity, aman.phi).to(u.radian).value,
        cast(u.Quantity, aman.off).value,
    )

gaussian2d_multipoles_from_aman(posmap, aman)

Evaluate a Gaussian multipole beam from an AxisManager.

Parameters:

Name Type Description Default
posmap Float[ndmap, '2 ny nx']

Beam coordinates. The first component is eta and the second is xi.

 required
aman AxisManager

AxisManager containing gaussian and gauss_multipole fields. gaussian contains the base Gaussian parameters and gauss_multipole.amps contains the multipole amplitudes.

 required

Returns:

Name Type Description
beam Float[ndmap, 'ny nx']

Gaussian multipole beam evaluated at posmap.
Source code in lat_beams/fitting/map/gauss.py 
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 gaussian2d_multipoles_from_aman(
    posmap: Float[ndmap, "2 ny nx"],
    aman: AxisManager,
) -> Float[ndmap, "ny nx"]:
    """Evaluate a Gaussian multipole beam from an AxisManager.

    Parameters
    ----------
    posmap : Float[ndmap, "2 ny nx"]
        Beam coordinates. The first component is `eta` and the second is
        `xi`.
    aman : AxisManager
        AxisManager containing `gaussian` and `gauss_multipole` fields.
        `gaussian` contains the base Gaussian parameters and
        `gauss_multipole.amps` contains the multipole amplitudes.

    Returns
    -------
    beam : Float[ndmap, "ny nx"]
        Gaussian multipole beam evaluated at `posmap`.
    """
    base_beam = gaussian2d_from_aman(posmap, aman.gaussian)
    base_beam -= cast(u.Quantity, aman.gaussian.off).value
    base_beam /= cast(u.Quantity, aman.gaussian.amp).value

    y, x = posmap
    theta = np.arctan2(
        y - cast(u.Quantity, aman.gaussian.eta0).to(u.radian).value,
        x - cast(u.Quantity, aman.gaussian.xi0).to(u.radian).value,
    )
    amps = (np.array(aman.gauss_multipole.amps.value),)

    model = posmap[0].copy()
    model[...] = multipole_expansion(base_beam, amps, theta)
    model = cast(ndmap, multipole_expansion(base_beam, amps, theta))

    return model