beam_utils

Utility functions for working with beam maps.

TODO: Make everything radians

apodized_disk(imap, radius, width, modrmap)

Cut out a disk and apodize its edge with a cosine apodization.

Parameters:

Name	Type	Description	Default
imap	Float[ndmap, '... ny nx']	Input map.	required
radius	float	Outer radius of the disk in arcminutes.	required
width	float	Width of the cosine apodization in arcminutes.	required
modrmap	Float[ndmap, 'ny nx']	Radius from the center at each pixel	required

Returns:

Name	Type	Description
omap	Float[ndmap, '... ny nx']	Apodized copy of the input map.

Source code in lat_beams/beam_utils.py

def apodized_disk(
    imap: Float[ndmap, "... ny nx"],
    radius: float,
    width: float,
    modrmap: Float[ndmap, "ny nx"],
) -> Float[ndmap, "... ny nx"]:
    """
    Cut out a disk and apodize its edge with a cosine apodization.

    Parameters
    ----------
    imap : Float[ndmap, "... ny nx"]
        Input map.
    radius : float
        Outer radius of the disk in arcminutes.
    width : float
        Width of the cosine apodization in arcminutes.
    modrmap : Float[ndmap, "ny nx"]
        Radius from the center at each pixel

    Returns
    -------
    omap : Float[ndmap, "... ny nx"]
        Apodized copy of the input map.
    """
    imap = imap.copy()
    imap[..., modrmap > radius] = 0
    if width > 0:
        inner_r = radius - width
        mask = modrmap > inner_r
        width = float(width)
        imap[..., mask] *= 1 + np.cos(
            -np.pi * inner_r / width + np.pi * modrmap[mask] / width
        )
        imap[..., mask] *= 0.5

    return imap

crop_maps(maps, cent, extent)

Crop a list of maps to be smaller. Note that all input maps will be cropped relative to the same pixel.

Parameters:

Name	Type	Description	Default
maps	list[Float[ndmap, 'nx ny']]	List of maps to crop. These should all have the same center.	required
cent	tuple[int, int]	The index of the center pixel.	required
extent	int	The extent of the output map.	required

Returns:

Name	Type	Description
cropped	list[Float[ndmap "2extent 2extent"]]	The cropped maps. Each one will have size (2extent, 2extent) unless that goes outside of the input map's bounding box, in which case the cropped map stops at that box.

Source code in lat_beams/beam_utils.py

def crop_maps(
    maps: list[Float[ndmap, "nx ny"]], cent: tuple[int, int], extent: int
) -> list[Float[ndmap, "2extent 2extent"]]:
    """
    Crop a list of maps to be smaller.
    Note that all input maps will be cropped relative to the same pixel.

    Parameters
    ----------
    maps : list[Float[ndmap, "nx ny"]]
        List of maps to crop.
        These should all have the same center.
    cent : tuple[int, int]
        The index of the center pixel.
    extent : int
        The extent of the output map.

    Returns
    -------
    cropped : list[Float[ndmap "2extent 2extent"]]
        The cropped maps.
        Each one will have size (2*extent, 2*extent)
        unless that goes outside of the input map's bounding box,
        in which case the cropped map stops at that box.
    """
    xmin = max(0, cent[0] - extent)
    xmax = min(maps[0].shape[-2], cent[0] + extent)
    ymin = max(0, cent[1] - extent)
    ymax = min(maps[0].shape[-1], cent[1] + extent)
    slx = slice(int(xmin), int(xmax))
    sly = slice(int(ymin), int(ymax))
    maps = [m[..., slx, sly] for m in maps]
    return maps

estimate_cent(imap, ivar, sigma=5, buf=30, ret_smooth=False)

estimate_cent(imap: Float[np.ndarray, 'nx ny'], ivar: Float[np.ndarray, 'nx ny'], sigma: float = 5, buf: int = 30, peak_radius: int = 10, min_snr: float = 5, ret_smooth: Literal[False] = False) -> tuple[int, int]

estimate_cent(imap: Float[np.ndarray, 'nx ny'], ivar: Float[np.ndarray, 'nx ny'], sigma: float = 5, buf: int = 30, peak_radius: int = 10, min_snr: float = 5, ret_smooth: Literal[True] = True) -> tuple[tuple[int, int], Float[np.ndarray, 'nx ny']]

estimate_cent(imap: Float[np.ndarray, 'nx ny'], ivar: Float[np.ndarray, 'nx ny'], sigma: float = 5, buf: int = 30, peak_radius: int = 10, min_snr: float = 5, ret_smooth: bool = False) -> tuple[int, int] | tuple[tuple[int, int], Float[np.ndarray, 'nx ny']]

Estimate the location of the central pixel of a beam map.

To do this we first construct an inverse-variance weighted SNR map and smooth it with a gaussian of size sigma, then we take the location of the maximum that is farther than buf from the edge of the map. We also require the candidate peak to have sufficient integrated SNR within peak_radius pixels, which helps reject isolated hot pixels.

Parameters:

Name	Type	Description	Default
imap	Float[ndarray, (nx, ny)]	The beam map to look for the center of.	required
ivar	Float[ndarray, (nx, ny)]	The inverse-variance map corresponding to imap. Pixels with non-positive or non-finite inverse variance are not searched.	required
sigma	float	The sigma of the gaussian in pixels to smooth the map by when searching for the max.	5
buf	int	Pixels within buf of the edge of the map will not be searched or used when smoothing. Meant to avoid low hits and edge artifacts.	30
ret_smooth	bool	If True also return the smoothed SNR map.	False

Returns:

Name	Type	Description
cent	tuple[int, int]	The index of the estimated center pixel.

Source code in lat_beams/beam_utils.py

def estimate_cent(
    imap: Float[np.ndarray, "nx ny"],
    ivar: Float[np.ndarray, "nx ny"],
    sigma: float = 5,
    buf: int = 30,
    ret_smooth: bool = False,
) -> tuple[int, int] | tuple[tuple[int, int], Float[np.ndarray, "nx ny"]]:
    """Estimate the location of the central pixel of a beam map.

    To do this we first construct an inverse-variance weighted SNR map and
    smooth it with a gaussian of size `sigma`, then we take the location of
    the maximum that is farther than `buf` from the edge of the map. We also
    require the candidate peak to have sufficient integrated SNR within
    `peak_radius` pixels, which helps reject isolated hot pixels.

    Parameters
    ----------
    imap : Float[np.ndarray, "nx, ny"]
        The beam map to look for the center of.
    ivar : Float[np.ndarray, "nx, ny"]
        The inverse-variance map corresponding to `imap`. Pixels with
        non-positive or non-finite inverse variance are not searched.
    sigma : float
        The sigma of the gaussian in pixels to smooth the map by when
        searching for the max.
    buf : int
        Pixels within `buf` of the edge of the map will not be searched or
        used when smoothing. Meant to avoid low hits and edge artifacts.
    ret_smooth : bool
        If True also return the smoothed SNR map.

    Returns
    -------
    cent : tuple[int, int]
        The index of the estimated center pixel.
    """
    imap = np.asarray(imap, dtype=float)
    ivar = np.asarray(ivar, dtype=float)

    if imap.shape != ivar.shape:
        raise ValueError("imap and ivar must have the same shape")

    valid = np.isfinite(imap) * (imap > 0) * np.isfinite(ivar) * (ivar > 0)
    valid[:buf] = False
    valid[-buf:] = False
    valid[:, :buf] = False
    valid[:, -buf:] = False
    signal = np.where(valid, imap, 0)
    weight = np.where(valid, ivar, 0)
    kern = Gaussian2DKernel(sigma)

    weighted_signal = convolve_fft(
        signal * weight,
        kern,
        normalize_kernel=True,
    )
    smoothed_ivar = convolve_fft(
        weight,
        kern,
        normalize_kernel=True,
    )

    with np.errstate(divide="ignore", invalid="ignore"):
        smoothed = weighted_signal / np.sqrt(smoothed_ivar)

    smoothed[~valid] = -np.inf
    cent = np.unravel_index(
        np.argmax(smoothed, axis=None),
        smoothed.shape,
    )
    best_cent = (int(cent[0]), int(cent[1]))

    if ret_smooth:
        return best_cent, smoothed

    return best_cent

estimate_solid_angle(imap, model, res, data_fwhm, cent, min_sigma)

Estimate the solid angle of a map given a fit model. Here we correct for the bias in our solid angle integration by computing:

\[ \Omega = \tilde{\Omega_{imap}} \frac{\Omega_{model}}{\tilde{\Omega_{imap}}} \]

Where \(\tilde{\Omega}\) implies a solid angle estimated with the solid_angle function from this module.

Parameters:

Name	Type	Description	Default
imap	Float[ndarray, 'nx ny']	The input map to estimate the solid angle of.	required
model	Float[ndarray, 'nx ny']	Model of map to estimate the solid angle of.	required
res	float	The resolution of the map in arcseconds.	required
data_fwhm	float	The FWHM of the map in arcseconds.	required
cent	tuple[int, int]	The index of the center pixel.	required
min_sigma	float	The number of sigma to use as the radius for aperture photometry.	required

Returns:

Name	Type	Description
data_solid_angle_meas	float	\(\tilde{\Omega_{imap}}\) in stradians.
model_solid_angle_meas	float	\(\tilde{\Omega_{model}}\) in stradians.
model_solid_angle_true	float	\(Omega_{model}\) in stradians.
data_solid_angle_corr	float	\(Omega\) in stradians.

Source code in lat_beams/beam_utils.py

def estimate_solid_angle(
    imap: Float[np.ndarray, "nx ny"],
    model: Float[np.ndarray, "nx ny"],
    res: float,
    data_fwhm: float,
    cent: tuple[int, int],
    min_sigma: float,
) -> tuple[float, float, float, float]:
    r"""
    Estimate the solid angle of a map given a fit model.
    Here we correct for the bias in our solid angle integration by computing:

    $$
    \Omega = \tilde{\Omega_{imap}} \frac{\Omega_{model}}{\tilde{\Omega_{imap}}}
    $$

    Where $\tilde{\Omega}$ implies a solid angle estimated with the `solid_angle` function from this module.

    Parameters
    ----------
    imap : Float[np.ndarray, "nx ny"]
        The input map to estimate the solid angle of.
    model : Float[np.ndarray, "nx ny"]
        Model of map to estimate the solid angle of.
    res : float
        The resolution of the map in arcseconds.
    data_fwhm : float
        The FWHM of the map in arcseconds.
    cent : tuple[int, int]
        The index of the center pixel.
    min_sigma : float
        The number of sigma to use as the radius for aperture photometry.

    Returns
    -------
    data_solid_angle_meas : float
        $\tilde{\Omega_{imap}}$ in stradians.
    model_solid_angle_meas : float
        $\tilde{\Omega_{model}}$ in stradians.
    model_solid_angle_true : float
        $Omega_{model}$ in stradians.
    data_solid_angle_corr : float
        $Omega$ in stradians.
    """
    fwhm_pix = int(data_fwhm / res)
    # Get solid angles
    y = np.linspace(-imap.shape[0] * res / 2, imap.shape[0] * res / 2, imap.shape[0])
    x = np.linspace(-imap.shape[1] * res / 2, imap.shape[1] * res / 2, imap.shape[1])
    kern = Gaussian2DKernel((data_fwhm / 2.3548) / res, (data_fwhm / 2.3548) / res)
    imap_smooth = convolve_fft(imap, kern)
    model_smooth = convolve_fft(model, kern)
    norm = np.max(
        imap_smooth[
            max(0, cent[0] - fwhm_pix) : min(imap_smooth.shape[0], cent[0] + fwhm_pix),
            max(0, cent[1] - fwhm_pix) : min(imap_smooth.shape[1], cent[1] + fwhm_pix),
        ]
    )
    data_solid_angle_meas = solid_angle(
        x, y, imap_smooth, cent, min_sigma * (data_fwhm / 2.355), norm
    )
    model_solid_angle_meas = solid_angle(
        x, y, model_smooth, cent, min_sigma * (data_fwhm / 2.355), norm
    )
    model_solid_angle_true = solid_angle(x, y, model, cent, np.inf, np.max(model))
    data_solid_angle_corr = (
        data_solid_angle_meas * model_solid_angle_true / model_solid_angle_meas
    )
    return (
        data_solid_angle_meas,
        model_solid_angle_meas,
        model_solid_angle_true,
        data_solid_angle_corr,
    )

get_corr_noise(ps2d, lmap, lmin, lmax)

Estimate the amplitude of correlated noise from a two-dimensional power spectrum.

Parameters:

Name	Type	Description	Default
ps2d	Float[ndarray, 'ny nx']	Two-dimensional power spectrum.	required
lmap	Float[ndarray, '2 ny nx']	Multipole coordinate map containing ell_y and ell_x for each pixel in ps2d.	required
lmin	float	Minimum multipole used to identify the vertical feature in the power spectrum.	required
lmax	float	Maximum multipole used to identify the vertical feature in the power spectrum.	required

Returns:

Name	Type	Description
corr_noise	float	Estimated amplitude of the correlated noise.

Source code in lat_beams/beam_utils.py

def get_corr_noise(
    ps2d: Float[np.ndarray, "ny nx"],
    lmap: Float[np.ndarray, "2 ny nx"],
    lmin: float,
    lmax: float,
) -> float:
    """
    Estimate the amplitude of correlated noise from a two-dimensional power spectrum.

    Parameters
    ----------
    ps2d : Float[np.ndarray, "ny nx"]
        Two-dimensional power spectrum.
    lmap : Float[np.ndarray, "2 ny nx"]
        Multipole coordinate map containing ell_y and ell_x for each pixel
        in `ps2d`.
    lmin : float
        Minimum multipole used to identify the vertical feature in the
        power spectrum.
    lmax : float
        Maximum multipole used to identify the vertical feature in the
        power spectrum.

    Returns
    -------
    corr_noise : float
        Estimated amplitude of the correlated noise.
    """
    # Mask that includes the vertical feature in the PSD.
    c_ell_mask_corr = (np.abs(lmap[1]) <= lmin) & (np.abs(lmap[0]) < lmax)
    corr_noise = np.mean(ps2d[c_ell_mask_corr])

    return float(corr_noise)

get_fit_vec(all_fits, name, fall_back=None)

Get a fit value from all fits in a structured array.

Parameters:

Name	Type	Description	Default
all_fits	Shaped[ndarray, nfits]	The fits to get values from. See load_beam_fits_from_jobs for details on the structure.	required
name	str	The name of the field to load from the AxisManagers in all_fits["aman"].	required
fall_back	str	Field to load in name is not found.	None

Returns:

Name	Type	Description
fit_vec	Quantity	The loaded values. Will have length nfits.

Source code in lat_beams/beam_utils.py

def get_fit_vec(
    all_fits: Shaped[np.ndarray, "nfits"], name: str, fall_back: Optional[str] = None
) -> u.Quantity:
    """
    Get a fit value from all fits in a structured array.

    Parameters
    ----------
    all_fits : Shaped[np.ndarray, "nfits"]
        The fits to get values from.
        See `load_beam_fits_from_jobs` for details on the structure.
    name : str
        The name of the field to load from the AxisManagers in
        `all_fits["aman"]`.
    fall_back : str
        Field to load in `name` is not found.

    Returns
    -------
    fit_vec : u.Quantity
        The loaded values.
        Will have length `nfits`.
    """
    if fall_back is not None:
        dat = u.Quantity(
            [
                aman[name] if name in aman else aman[fall_back]
                for aman in all_fits["aman"]
            ]
        )
    else:
        dat = u.Quantity([aman[name] for aman in all_fits["aman"]])

    if dat.unit == u.Unit(3):  # type: ignore
        dat = dat.value * u.pW
    return dat

get_fwhm_radial_bins(r, y, interpolate=False, frac=0.5)

Estimate FWHM from a radial profile.

Parameters:

Name	Type	Description	Default
r	Float[ndarray, nr]	The radial position at each point.	required
y	Float[ndarray, nr]	The value of the profile at each point.	required
interpolate	bool	If True then interpolate the input profile on an evenly spaced grid of 100 points before estimating the FWHM.	False
frac	float	The fraction of the peak to get the width at. By default this is 0.5 which is the FWHM, but other values can be passed if needed.	0.5

Returns:

Name	Type	Description
fwhm	float	The estimated FWHM in the same units as r.

Source code in lat_beams/beam_utils.py

def get_fwhm_radial_bins(
    r: Float[np.ndarray, "nr"],
    y: Float[np.ndarray, "nr"],
    interpolate: bool = False,
    frac: float = 0.5,
) -> float:
    """
    Estimate FWHM from a radial profile.

    Parameters
    ----------
    r : Float[np.ndarray, "nr"]
        The radial position at each point.
    y : Float[np.ndarray, "nr"]
        The value of the profile at each point.
    interpolate : bool, default: False
        If True then interpolate the input profile on an evenly spaced
        grid of 100 points before estimating the FWHM.
    frac : float, default: 0.5
        The fraction of the peak to get the width at.
        By default this is 0.5 which is the FWHM,
        but other values can be passed if needed.

    Returns
    -------
    fwhm : float
        The estimated FWHM in the same units as `r`.
    """
    half_point = np.max(y) * frac

    if interpolate:
        r_diff = r[1] - r[0]
        interp_func = interp1d(r, y)
        r_interp = np.arange(np.min(r), np.max(r) - r_diff, r_diff / 100)
        y_interp = interp_func(r_interp)
        r, y = (r_interp, y_interp)
    d = y - half_point
    inds = np.where(d > 0)[0]
    if len(inds) < 0.1 * len(r):
        return np.nan
    fwhm = 2 * (r[inds[-1]])
    return cast(float, fwhm)

get_map_noise(imap_centered, ivar, r, rprof, radius_rad, lmin, lmax, fwhm, opt_ang=False)

Estimate white and correlated noise from residuals of a beam map.

The radial profile is projected onto the map and subtracted before identifying and masking hot pixels. The residual map is then weighted by the inverse-variance map, masked to an inner and outer radius, and transformed to a two-dimensional power spectrum. The white and correlated noise amplitudes are estimated from this power spectrum.

Parameters:

Name	Type	Description	Default
imap_centered	Float[ndmap, '... ny nx']	Centered input map.	required
ivar	Float[ndmap, '... ny nx']	Inverse-variance map corresponding to imap_centered. Pixels identified as hot are assigned zero inverse variance.	required
r	Float[ndarray, nr]	Radial coordinates corresponding to rprof.	required
rprof	Float[ndarray, nr]	Radial profile of the map. Note that this needs to have the same normalizations applied to it as imap_centered.	required
radius_rad	float	Outer radius of the region used for the noise estimate, in radians. You probably want this to be the mapmaker radius.	required
lmin	float	Minimum multipole used to separate white and correlated noise.	required
lmax	float	Maximum multipole used to estimate the noise.	required
fwhm	float	Beam full width at half maximum, in radians. Used to define the inner radius excluded from the noise estimate and the width of the apodization.	required
opt_ang	bool	If True, optimize the orientation of the correlated-noise stripe in Fourier space. The optimization is restricted to +/- pi/8 around the nominal orientation.	False

Returns:

Name	Type	Description
white_noise	float	Estimated white-noise level.
corr_noise	float	Estimated amplitude of the correlated noise.
ps2d	Float[ndmap, 'ny nx']	Two-dimensional power spectrum of the residual map.
lmap	Float[ndmap, '2 ny nx']	Two-dimensional multipole coordinate maps corresponding to ps2d. The first element contains the y-direction multipoles and the second contains the x-direction multipoles.
hot	Bool[ndarray, 'ny nx']	Boolean mask identifying pixels flagged as hot-pixel outliers.

Source code in lat_beams/beam_utils.py

def get_map_noise(
    imap_centered: Float[ndmap, "... ny nx"],
    ivar: Float[ndmap, "... ny nx"],
    r: Float[np.ndarray, "nr"],
    rprof: Float[np.ndarray, "nr"],
    radius_rad: float,
    lmin: float,
    lmax: float,
    fwhm: float,
    opt_ang: bool = False,
) -> tuple[
    float,
    float,
    Float[ndmap, "ny nx"],
    Float[ndmap, "2 ny nx"],
    Bool[np.ndarray, "ny nx"],
]:
    """
    Estimate white and correlated noise from residuals of a beam map.

    The radial profile is projected onto the map and subtracted before
    identifying and masking hot pixels. The residual map is then weighted
    by the inverse-variance map, masked to an inner and outer radius, and
    transformed to a two-dimensional power spectrum. The white and
    correlated
    noise amplitudes are estimated from this power spectrum.

    Parameters
    ----------
    imap_centered : Float[ndmap, "... ny nx"]
        Centered input map.
    ivar : Float[ndmap, "... ny nx"]
        Inverse-variance map corresponding to `imap_centered`. Pixels
        identified as hot are assigned zero inverse variance.
    r : Float[np.ndarray, "nr"]
        Radial coordinates corresponding to `rprof`.
    rprof : Float[np.ndarray, "nr"]
        Radial profile of the map.
        Note that this needs to have the same normalizations applied
        to it as `imap_centered`.
    radius_rad : float
        Outer radius of the region used for the noise estimate, in radians.
        You probably want this to be the mapmaker radius.
    lmin : float
        Minimum multipole used to separate white and correlated noise.
    lmax : float
        Maximum multipole used to estimate the noise.
    fwhm : float
        Beam full width at half maximum, in radians. Used to define the
        inner radius excluded from the noise estimate and the width of
        the apodization.
    opt_ang : bool, optional
        If True, optimize the orientation of the correlated-noise stripe
        in Fourier space. The optimization is restricted to +/- pi/8
        around the nominal orientation.

    Returns
    -------
    white_noise : float
        Estimated white-noise level.
    corr_noise : float
        Estimated amplitude of the correlated noise.
    ps2d : Float[ndmap, "ny nx"]
        Two-dimensional power spectrum of the residual map.
    lmap : Float[ndmap, "2 ny nx"]
        Two-dimensional multipole coordinate maps corresponding to `ps2d`.
        The first element contains the y-direction multipoles and the
        second contains the x-direction multipoles.
    hot : Bool[np.ndarray, "ny nx"]
        Boolean mask identifying pixels flagged as hot-pixel outliers.
    """

    # Project profile and subtract
    modrmap = enmap.modrmap(imap_centered.shape, imap_centered.wcs)
    br_interp = interp1d(r, rprof, fill_value=0, bounds_error=False)
    brmap = enmap.samewcs(
        br_interp(modrmap.ravel()).reshape(modrmap.shape),
        modrmap,
    )
    imap_sub = cast(ndmap, imap_centered - brmap)

    imap_hot = imap_sub.copy()
    imap_hot[modrmap < 3 * fwhm] = np.nan
    med = np.nanmedian(imap_hot)
    mad = np.nanmedian(np.abs(imap_hot - med))
    z = (imap_sub - med) / (1.4826 * mad)
    hot = np.abs(z) > 5
    hot[modrmap < 3 * fwhm] = False
    if np.sum(hot) > 0:
        labeled, _ = label(hot)  # type: ignore
        labeled = np.asarray(labeled, dtype=np.intp)
        sizes = np.bincount(labeled.ravel())
        small = sizes[labeled] < 4
        hot[small] = False
        hot = binary_dilation(hot)
    hot = np.asarray(hot, bool)
    if np.sum(hot) > 0:
        imap_clean = imap_centered.copy()
        imap_clean[hot] = median_filter(imap_clean, size=15)[hot]
        result = enmap.rbin(imap_clean)
        br = result[0]
        radii = result[1]
        br_interp = interp1d(radii, br, fill_value=0, bounds_error=False)
        brmap = enmap.samewcs(
            br_interp(modrmap.ravel()).reshape(modrmap.shape),
            modrmap,
        )
        imap_sub = cast(ndmap, imap_clean - brmap)
        ivar[hot] = 0
    w = cast(ndmap, np.sqrt(ivar))

    # Cut out residual.
    pmask = cast(ndmap, imap_sub.copy() * 0 + 1)
    pmask -= apodized_disk(pmask, 3 * fwhm, 0.1 * fwhm, modrmap=modrmap)
    w *= pmask

    # Apodize outer edge.
    w = apodized_disk(w, radius_rad, 0.1 * radius_rad, modrmap=modrmap)

    # Compute 2d power spectrum.
    lmap = enmap.lmap(imap_sub.shape, imap_sub.wcs)
    ps2d = enmap.calc_ps2d(enmap.map2harm(w * imap_sub))
    norm = np.mean(w[w > 0] ** 2)
    ps2d /= norm

    lx = lmap[1]
    ly = lmap[0]

    def rotated_lx(theta):
        c, s = np.cos(theta), np.sin(theta)
        return c * lx + s * ly

    def objective(theta):
        lx_rot = rotated_lx(theta)
        stripe = np.abs(lx_rot) < lmin
        return -np.mean(ps2d[stripe])

    theta = 0
    if opt_ang:
        result = minimize_scalar(
            objective,
            bounds=(-np.pi / 2, np.pi / 2),
            method="bounded",
        )
        theta = result.x
        c, s = np.cos(theta), np.sin(theta)
        lx_rot = c * lx + s * ly
        ly_rot = -s * lx + c * ly
        lmap[0][:] = ly_rot
        lmap[1][:] = lx_rot

    white_noise = get_white_noise(ps2d, lmap, lmin, lmax)
    corr_noise = get_corr_noise(ps2d, lmap, lmin, lmax)

    return white_noise, corr_noise, ps2d, lmap, hot, theta

get_split_vec(fits, split, ctx, round_to=2, metasplits={})

Get an array of metadata to split fits by.

Parameters:

Name	Type	Description	Default
fits	Shaped[ndarray, nfits]	The fits to get values from. See load_beam_fits_from_jobs for details on the structure.	required
split	str	List of collumns in fits or the obsdb to split by, should be one string with + to seperate collumn names.	required
ctx	Context	Context used to lookup values from the obsdb.	required
round_to	int	How many decimal places to round numeric collumns to.	2
metasplits	dict[str, tuple[str, list[str | float]]]	A method of defining a collumn that matches against values from a normal split collumn. Each entry should have some name as the key, which then maps to a tuple. The first element of the tuple should be the split we are matching against and the second should be a list of values to match. In the output split_vec anything that matches will have name in the split and anything that does not match will have NOMATCH. If you want to match against a numeric range instaed then the list should be a two element float list that defines the range [low, high).	{}

Returns:

Name	Type	Description
split_vec	Shaped[ndarray, nfits]	Array of strings containing the values from the split collumns. Values are in the same order as split and are seperated by +s.

Source code in lat_beams/beam_utils.py

def get_split_vec(
    fits: Shaped[np.ndarray, "nfits"],
    split: str,
    ctx: Context,
    round_to: int = 2,
    metasplits: dict[str, tuple[str, list[str]]] = {},
) -> Shaped[np.ndarray, "nfits"]:
    """
    Get an array of metadata to split fits by.

    Parameters
    ----------
    fits : Shaped[np.ndarray, "nfits"]
        The fits to get values from.
        See `load_beam_fits_from_jobs` for details on the structure.
    split : str
        List of collumns in `fits` or the obsdb to split by,
        should be one string with `+` to seperate collumn names.
    ctx : Context
        Context used to lookup values from the obsdb.
    round_to : int
        How many decimal places to round numeric collumns to.
    metasplits : dict[str, tuple[str, list[str | float]]]
        A method of defining a collumn that matches against values
        from a normal split collumn. Each entry should have some `name`
        as the key, which then maps to a tuple. The first element of
        the tuple should be the split we are matching against and the
        second should be a list of values to match. In the output split_vec
        anything that matches will have `name` in the split and anything
        that does not match will have `NOMATCH`.
        If you want to match against a numeric range instaed then
        the `list` should be a two element float list that defines the range
        `[low, high)`.

    Returns
    -------
    split_vec : Shaped[np.ndarray, "nfits"]
        Array of strings containing the values from the split collumns.
        Values are in the same order as `split` and are seperated by `+`s.
    """
    if fits.dtype.names is None:
        raise ValueError("Unknown fits format!")
    if ctx.obsdb is None:
        raise ValueError("obsdb is None!")
    split_vecs = []
    for spl in split.split("+"):
        if spl in metasplits:
            # Structure of metasplits name -> (spl, (vals))
            split_vec = _get_vec(metasplits[spl][0], fits, ctx, round_to)
            split_vec = split_vec.astype(
                f"U{max(len(spl), 7, int(split_vec.dtype.itemsize/4))}"
            )
            if (
                len(metasplits[spl][1]) == 2
                and np.array(metasplits[spl][1]).dtype == float
            ):
                svf = split_vec.astype(float)
                msk = (svf >= metasplits[spl][1][0]) * (svf < metasplits[spl][1][1])
            else:
                msk = np.isin(split_vec, metasplits[spl][1])
            split_vec[msk] = spl
            split_vec[~msk] = "NOMATCH"
        else:
            split_vec = _get_vec(spl, fits, ctx, round_to)
        split_vecs += [split_vec.astype(str)]
    split_vecs = np.column_stack(split_vecs)

    return np.array(["+".join(v) for v in split_vecs])

get_white_noise(ps2d, lmap, lmin, lmax)

Estimate the amplitude of white noise from a two-dimensional power spectrum.

Parameters:

Name	Type	Description	Default
ps2d	Float[ndarray, 'ny nx']	Two-dimensional power spectrum.	required
lmap	Float[ndarray, '2 ny nx']	Multipole coordinate map containing ell_y and ell_x for each pixel in ps2d.	required
lmin	float	Minimum multipole of the white-noise region.	required
lmax	float	Maximum multipole of the white-noise region.	required

Returns:

Name	Type	Description
white_noise	float	Estimated white-noise level.

Source code in lat_beams/beam_utils.py

def get_white_noise(
    ps2d: Float[np.ndarray, "ny nx"],
    lmap: Float[np.ndarray, "2 ny nx"],
    lmin: float,
    lmax: float,
) -> float:
    """
    Estimate the amplitude of white noise from a two-dimensional power spectrum.

    Parameters
    ----------
    ps2d : Float[np.ndarray, "ny nx"]
        Two-dimensional power spectrum.
    lmap : Float[np.ndarray, "2 ny nx"]
        Multipole coordinate map containing ell_y and ell_x for each pixel
        in `ps2d`.
    lmin : float
        Minimum multipole of the white-noise region.
    lmax : float
        Maximum multipole of the white-noise region.

    Returns
    -------
    white_noise : float
        Estimated white-noise level.
    """
    ps2d_white = ps2d.copy()

    # PSD will have strong vertical feature that we want to avoid.
    c_ell_mask = np.abs(lmap[1]) > lmin
    ps2d_white[~c_ell_mask] = 0

    lbin_result = enmap.lbin(ps2d_white)
    c_ell_white = lbin_result[0]
    ells = lbin_result[1]
    ells_cut = ells[ells > lmin]

    # Compute factor that corrects for lost modes from c_ell_mask.
    # This is just the arclen inside of a box of width 2 * lmin at radius ell.
    corr_frac = (2 * np.pi * ells_cut) / (
        2 * np.pi * ells_cut - 4 * ells_cut * np.arcsin(lmin / ells_cut)
    )
    c_ell_white[ells > lmin] *= corr_frac

    white_noise = np.mean(c_ell_white[(ells > lmin) & (ells <= lmax)])

    return white_noise

load_beam_fits_from_jobs(fpath, joblist, jdb=None)

Load beam fits from a list of jobs.

Parameters:

Name	Type	Description	Default
fpath	str	The path to the HDF5 file containing the fits.	required
joblist	list[Job]	List of jobs to load fits for. Jobs should be of jclass fit_map.	required
jdb	Optional[JobManager]	If passed then jobs that couldn't be loaded are reopened.	None

Returns:

Name	Type	Description
all_fits	Shaped[ndarray, nfits]	Loaded fits. This is a numpy structured array with the following collumns: obs_id : str, the obs_id of the fit data wafer_slot : str, the wafer slot of the fit data stream_id : str, the stream id of the fit data array : str, the array of the fit data band : str, the band (ie. f090) of the fit data split: str, the split that was run for this map source : str, the source that was fit time : float, the time of the observation hour : float, what hour of the day the observation was at aman : AxisManager, the loaded fit

Raises:

Type	Description
ValueError	If loaded fits do not all contain the same structure.

Source code in lat_beams/beam_utils.py

def load_beam_fits_from_jobs(
    fpath: str, joblist: list[jobdb.Job], jdb: Optional[jobdb.JobManager] = None
) -> Shaped[np.ndarray, "nfits"]:
    """
    Load beam fits from a list of jobs.

    Parameters
    ----------
    fpath : str
        The path to the HDF5 file containing the fits.
    joblist : list[jobdb.Job]
        List of jobs to load fits for.
        Jobs should be of jclass `fit_map`.
    jdb : Optional[jobdb.JobManager], default: None
        If passed then jobs that couldn't be loaded are reopened.

    Returns
    -------
    all_fits : Shaped[np.ndarray, "nfits"]
        Loaded fits.
        This is a numpy structured array with the following collumns:

        * obs_id : str, the obs_id of the fit data
        * wafer_slot : str, the wafer slot of the fit data
        * stream_id : str, the stream id of the fit data
        * array : str, the array of the fit data
        * band : str, the band (ie. f090) of the fit data
        * split: str, the split that was run for this map
        * source : str, the source that was fit
        * time : float, the time of the observation
        * hour : float, what hour of the day the observation was at
        * aman : AxisManager, the loaded fit

    Raises
    ------
    ValueError
        If loaded fits do not all contain the same structure.
    """
    f = h5py.File(fpath, mode="r")
    obs_ids = np.array([job.tags["obs_id"] for job in joblist])
    times = np.array([float(o.split("_")[1]) for o in obs_ids])
    wafer_slots = np.array([job.tags["wafer_slot"] for job in joblist])
    stream_ids = np.array([job.tags["stream_id"] for job in joblist])
    arrays = np.array([job.tags["array"] for job in joblist])
    splits = np.array([job.tags.get("split", "") for job in joblist])
    bands = np.array([job.tags["band"] for job in joblist])
    sources = np.array([job.tags["source"] for job in joblist])
    dates = np.array([dt.date.fromtimestamp(ct) for ct in times])
    tdelt = (
        np.array(
            [
                ct - dt.datetime(year=d.year, month=d.month, day=d.day).timestamp()
                for ct, d in zip(times, dates)
            ]
        )
        / 3600
    )

    # amans = np.array(
    #     [
    #         AxisManager.load(f[os.path.join(o, s, b, m)])
    #         for o, s, b, m in zip(obs_ids, stream_ids, bands, splits)
    #     ]
    # )
    amans = []
    for job, o, u, b, m in zip(joblist, obs_ids, arrays, bands, splits):
        try:
            aman = AxisManager.load(f[os.path.join(o, u, b, m)])
            amans += [aman]
        except:
            print(os.path.join(o, u, b, m))
            if jdb is None:
                continue
            with jdb.session_scope() as session:
                job.jstate = "open"
                session.merge(job)
                session.commit()
    # check that all fits have the same pars
    par_list = np.sort(list(amans[0]._fields.keys()))
    for aman in amans:
        pars = np.sort(list(aman._fields.keys()))
        if not np.array_equal(par_list, pars):
            raise ValueError("Not all fits have the same pars!")
    f.close()
    dtype = [
        ("obs_id", obs_ids.dtype),
        ("wafer_slot", wafer_slots.dtype),
        ("stream_id", stream_ids.dtype),
        ("array", arrays.dtype),
        ("band", bands.dtype),
        ("split", splits.dtype),
        ("source", sources.dtype),
        ("time", float),
        ("hour", float),
        ("aman", "O"),
    ]
    all_fits = np.fromiter(
        zip(
            obs_ids,
            wafer_slots,
            stream_ids,
            arrays,
            bands,
            splits,
            sources,
            times,
            tdelt,
            amans,
        ),
        dtype,
        count=len(amans),
    )
    return all_fits

process_model(aman, solved, model, noise, min_snr, c, map_units, pixsize, data_fwhm, min_sigma, job, logger)

Convenience function to postproccess a map and it's fit model. This fundtion checks the SNR of the model, computes its radial profile, and computes the solid angle.

Parameters:

Name	Type	Description	Default
aman	AxisManager	AxisManager with the fit parameters. No values are read off this, but it wil be modified in place.	required
solved	Float[ndarray, 'nx ny']	The map that was fit.	required
model	Float[ndarray, 'nx ny']	The model computed on the same grid as the map.	required
noise	float	The noise level of the map.	required
min_snr	float	The minimum SNR of the model. If the SNR is less than this then None is returned.	required
c	tuple[int, int]	The index of the center pixel.	required
map_units	Unit	The units of the map.	required
pixsize	float	The pixel size in arcseconds.	required
data_fwhm	float	The data fwhm in arcseconds.	required
min_sigma	float	See estimate_solid_angle.	required
job	Optional[Job]	The job associated with the fit. Pass None if running without a job.	required
logger	Optional[LoggerLike]	The logger to log with. Pass None if running without a logger.	required

Returns:

Name	Type	Description
aman	Optional[AxisManager]	The input aman modified to add the radial profile of the model and all the solid angles output by estimate_solid_angle. If the SNR chech is not passed then None is returned.

Source code in lat_beams/beam_utils.py

def process_model(
    aman: AxisManager,
    solved: Float[np.ndarray, "nx ny"],
    model: Float[np.ndarray, "nx ny"],
    noise: float,
    min_snr: float,
    c: tuple[int, int],
    map_units: u.Unit,
    pixsize: float,
    data_fwhm: float,
    min_sigma: float,
    job: Optional[jobdb.Job],
    logger: Optional[LoggerLike],
) -> Optional[AxisManager]:
    """
    Convenience function to postproccess a map and it's fit model.
    This fundtion checks the SNR of the model,
    computes its radial profile,
    and computes the solid angle.

    Parameters
    ----------
    aman : AxisManager
        AxisManager with the fit parameters.
        No values are read off this, but it wil be modified in place.
    solved : Float[np.ndarray, "nx ny"]
        The map that was fit.
    model : Float[np.ndarray, "nx ny"]
        The model computed on the same grid as the map.
    noise : float
        The noise level of the map.
    min_snr : float
        The minimum SNR of the model.
        If the SNR is less than this then `None` is returned.
    c : tuple[int, int]
        The index of the center pixel.
    map_units : u.Unit
        The units of the map.
    pixsize : float
        The pixel size in arcseconds.
    data_fwhm : float
        The data fwhm in arcseconds.
    min_sigma : float
        See `estimate_solid_angle`.
    job : Optional[jobdb.Job]
        The job associated with the fit.
        Pass `None` if running without a job.
    logger : Optional[LoggerLike]
        The logger to log with.
        Pass `None` if running without a logger.

    Returns
    -------
    aman : Optional[AxisManager]
        The input `aman` modified to add the radial profile
        of the model and all the solid angles output by
        `estimate_solid_angle`.
        If the SNR chech is not passed then `None` is returned.
    """
    # Check snr
    if np.nanmax(model) / noise < min_snr:
        msg = "Model SNR too low"
        if logger is None:
            print(f"{msg}")
        elif job is None:
            logger.error("%s", msg)
        if job is not None:
            fail(job, ErrCode.SNR_LOW, msg, logger)
        return None

    # Get model profile
    mprof = radial_profile(model, c[::-1])
    aman.wrap("mprof", mprof * map_units)

    # Get solid angle
    try:
        (
            data_solid_angle_meas,
            model_solid_angle_meas,
            model_solid_angle_true,
            data_solid_angle_corr,
        ) = estimate_solid_angle(solved, model, pixsize, data_fwhm.value, c, min_sigma)
    except ValueError as e:
        msg = f"Failed to estimate solid angle with error: {e}"
        if logger is None:
            print(f"{msg}")
        elif job is None:
            logger.error("%s", msg)
        if job is not None:
            fail(job, ErrCode.OMEGA_FAILED, msg, logger)
        return None
    aman.wrap("data_solid_angle_meas", data_solid_angle_meas * u.sr)
    aman.wrap("data_solid_angle_corr", data_solid_angle_corr * u.sr)
    aman.wrap("model_solid_angle_meas", model_solid_angle_meas * u.sr)
    aman.wrap("model_solid_angle_true", model_solid_angle_true * u.sr)

    return aman

radial_profile(data, center, avg=True)

Compute the radial profile of a beam, this is a naive way of doing thing just looking at the pixels.

Parameters:

Name	Type	Description	Default
data	Float[ndarray, 'nx ny']	The input beam.	required
center	tuple[int, int]	The index of the center pixel.	required

Returns:

Name	Type	Description
radialprofile	Float[ndarray, nr]	Radial profile of the input map.

Source code in lat_beams/beam_utils.py

def radial_profile(
    data: Float[np.ndarray, "nx ny"],
    center: tuple[int, int],
    avg=True,
) -> Float[np.ndarray, "nr"]:
    """
    Compute the radial profile of a beam, this is a naive way of doing thing just looking at the pixels.

    Parameters
    ----------
    data : Float[np.ndarray, "nx ny"]
        The input beam.
    center : tuple[int, int]
        The index of the center pixel.

    Returns
    -------
    radialprofile : Float[np.ndarray, "nr"]
        Radial profile of the input map.
    """
    msk = np.isfinite(data.ravel())
    y, x = np.indices(data.shape)
    r = np.sqrt((x - center[0]) ** 2 + (y - center[1]) ** 2)
    r = r.astype(int)

    tbin = np.bincount(r.ravel()[msk], data.ravel()[msk])
    if not avg:
        return tbin
    nr = np.bincount(r.ravel()[msk])
    radialprofile = tbin / nr
    return radialprofile

radial_profile_lin(data, posmap, xi0=0.0, eta0=0.0, r=None, n_bins=None, rmax=None)

Compute the azimuthally averaged radial profile using the true beam center. This computes a matrix for the binning operation that can be reused.

Parameters:

Name	Type	Description	Default
data	ndarray	Map to bin.	required
posmap	tuple[ndarray, ndarray]	Beam-coordinate maps (eta, xi).	required
xi0	float	Beam center in the same units as posmap.	0.0
eta0	float	Beam center in the same units as posmap.	0.0
r	ndarray	Radial bin centers. If omitted, n_bins equally spaced bins are generated from zero to rmax.	None
n_bins	int	Number of radial bins when r is not supplied.	None
rmax	float	Maximum radius when generating bins.	None

Returns:

Name	Type	Description
r	ndarray	Radial bin centers.
profile	ndarray	Azimuthally averaged radial profile.
R	csr_matrix	Radial binning operator satisfying profile = R @ data.ravel().

Source code in lat_beams/beam_utils.py

def radial_profile_lin(
    data,
    posmap,
    xi0=0.0,
    eta0=0.0,
    r=None,
    n_bins=None,
    rmax=None,
):
    """
    Compute the azimuthally averaged radial profile using the true beam center.
    This computes a matrix for the binning operation that can be reused.

    Parameters
    ----------
    data : ndarray
        Map to bin.
    posmap : tuple[ndarray, ndarray]
        Beam-coordinate maps `(eta, xi)`.
    xi0, eta0 : float
        Beam center in the same units as `posmap`.
    r : ndarray, optional
        Radial bin centers. If omitted, `n_bins` equally spaced bins are
        generated from zero to `rmax`.
    n_bins : int, optional
        Number of radial bins when `r` is not supplied.
    rmax : float, optional
        Maximum radius when generating bins.

    Returns
    -------
    r : ndarray
        Radial bin centers.
    profile : ndarray
        Azimuthally averaged radial profile.
    R : scipy.sparse.csr_matrix
        Radial binning operator satisfying
        `profile = R @ data.ravel()`.
    """
    eta, xi = posmap
    rmap = np.hypot(xi - xi0, eta - eta0).ravel()
    data = np.asarray(data).ravel()

    if r is None:
        if n_bins is None:
            raise ValueError("Specify either r or n_bins.")
        if rmax is None:
            rmax = rmap.max()
        dr = rmax / n_bins
        r = np.arange(n_bins) * dr
    else:
        r = np.asarray(r)
        if len(r) < 2:
            raise ValueError("r must contain at least two bin centers.")
        dr = np.mean(np.diff(r))

    edges = np.r_[0.0, r[1:] - dr / 2, r[-1] + dr / 2]
    bin_idx = np.digitize(rmap, edges) - 1
    valid = (bin_idx >= 0) & (bin_idx < len(r)) & np.isfinite(data)

    pix = np.flatnonzero(valid)
    bins = bin_idx[valid]
    counts = np.bincount(bins, minlength=len(r))

    weights = 1.0 / np.maximum(counts[bins], 1)
    R = sparse.csr_matrix(
        (weights, (bins, pix)),
        shape=(len(r), len(data)),
    )

    profile = np.asarray(R @ data).ravel()
    return r, profile, R

solid_angle(az, el, beam, cent, r1, norm)

Compute the integrated solid angle of a beam map. This uses aperture photometry to handle bias from the background of the map.

Parameters:

Name	Type	Description	Default
az	Float[ndarray, nx]	The x coordinates of the map in arcseconds.	required
el	Float[ndarray, nx]	The y coordinates of the map in arcseconds.	required
beam	Float[ndarray, 'nx ny']	The beam map to compute the solid angle of.	required
cent	tuple[int, int]	The index of the center pixel.	required
r1	float	The radius of the inner ring in aperture photometry in arcseconds.	required
norm	float	The value to normalize the map by.	required

Returns:

Name	Type	Description
solid_angle	float	the solid angle in stradians.

Source code in lat_beams/beam_utils.py

def solid_angle(
    az: Float[np.ndarray, "nx"],
    el: Float[np.ndarray, "ny"],
    beam: Float[np.ndarray, "nx ny"],
    cent: tuple[int, int],
    r1: float,
    norm: float,
) -> float:
    """
    Compute the integrated solid angle of a beam map.
    This uses aperture photometry to handle bias from the background of the map.

    Parameters
    ----------
    az : Float[np.ndarray, "nx"]
        The x coordinates of the map in arcseconds.
    el : Float[np.ndarray, "nx"]
        The y coordinates of the map in arcseconds.
    beam : Float[np.ndarray, "nx ny"]
        The beam map to compute the solid angle of.
    cent : tuple[int, int]
        The index of the center pixel.
    r1 : float
        The radius of the inner ring in aperture photometry in arcseconds.
    norm : float
        The value to normalize the map by.

    Returns
    -------
    solid_angle : float
        the solid angle in stradians.
    """
    r2 = np.sqrt(2) * r1

    _az, _el = np.meshgrid(az, el)
    r = np.sqrt((_az - _az[cent]) ** 2 + (_el - _el[cent]) ** 2)
    # convert from arcsec to rad
    az = np.deg2rad(az / 3600)
    el = np.deg2rad(el / 3600)

    integrand = beam / norm

    # perform the solid angle integral
    _integrand = integrand.copy()
    _integrand[r > r1] = 0
    integral_inner = trapz(trapz(_integrand, el, axis=0), az, axis=0)

    _integrand = integrand.copy()
    _integrand[(r < r1) + (r > r2)] = 0
    integral_outer = trapz(trapz(_integrand, el, axis=0), az, axis=0)
    return integral_inner - integral_outer