Skip to content

Custom filters #4

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
wants to merge 3 commits into
base: main
Choose a base branch
from
Open

Custom filters #4

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
f590c13
added options for user-provided cov_ell for filtering anand added eig…
AdriJD Jan 4, 2024
56941f7
added lim and lim0 to wav_noise
AdriJD Feb 21, 2024
c7d595e
minor format changes
zatkins2 Feb 21, 2024
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
73 changes: 54 additions & 19 deletions mnms/filters.py
View file
Open in desktop
Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Did you consider making a separate filter for this new feature? That is probably simpler than modifying each of the other filters to accommodate this new feature. You would only need to write a new model function (registered with model=True). You wouldn't need to write a new simulation filter, you would only need to add a registration on top of existing filters (e.g., analogous to how ivar_basic and raw_ivar_basic share a simulation filters). Maybe there is a downside to this but I think it would be more modular/easier to maintain.

If you prefer to modify the existing filters, I think you missed adding your new feature to iso_harmonic_ivar_scaledep_model

Original file line number Diff line number Diff line change
Expand Up @@ -74,7 +74,7 @@ def identity(inp, *args, **kwargs):
@register('map', 'map', iso_filt_method='harmonic', model=True)
def iso_harmonic_ivar_none_model(imap, mask_est=1, ainfo=None, lmax=None,
post_filt_rel_downgrade=1,
post_filt_downgrade_wcs=None,
post_filt_downgrade_wcs=None, cov_ell=None,
lim=1e-6, lim0=None, **kwargs):
"""Filter a map by an ell-dependent matrix in harmonic space. The
filter is measured as the pseudospectra of the input and returned.
Expand All @@ -96,6 +96,9 @@ def iso_harmonic_ivar_none_model(imap, mask_est=1, ainfo=None, lmax=None,
bandlimits the measured pseudospectra by the same factor.
post_filt_downgrade_wcs : astropy.wcs.WCS, optional
Assign this wcs to the filtered maps, by default None.
cov_ell : (*preshape, *preshape, nell) np.ndarray, optional
If provided, skip the estimation of the ell-dependent matrix and
instead use this matrix for the harmonic filter, by default None.
lim : float, optional
Set eigenvalues smaller than lim * max(eigenvalues) to zero.
lim0 : float, optional
Expand All @@ -118,12 +121,19 @@ def iso_harmonic_ivar_none_model(imap, mask_est=1, ainfo=None, lmax=None,
The matrix square-root is taken over the matrix axes before filtering, as
appropriate for a map.
"""
# measure sqrt_cov_ell and inv_sqrt_cov_ell
sqrt_cov_ell, inv_sqrt_cov_ell = utils.measure_iso_harmonic(
imap, 0.5, -0.5, mask_est=mask_est, ainfo=ainfo, lmax=lmax,
lim=lim, lim0=lim0
if cov_ell is None:
sqrt_cov_ell, inv_sqrt_cov_ell = utils.measure_iso_harmonic(
imap, 0.5, -0.5, mask_est=mask_est, ainfo=ainfo, lmax=lmax,
lim=lim, lim0=lim0
)

else:
preshape = cov_ell.shape[:len(cov_ell.shape[:-1]) // 2]
sqrt_cov_ell, inv_sqrt_cov_ell = [
utils.eigpow_flat(
cov_ell, e, preshape, lim=lim, lim0=lim0, copy=True)
for e in (0.5, -0.5)
]

# also need to downgrade the measured power spectra!
sqrt_cov_ell = sqrt_cov_ell[..., :lmax//post_filt_rel_downgrade+1]

Expand Down Expand Up @@ -166,7 +176,7 @@ def iso_harmonic_ivar_none(alm, sqrt_cov_ell=None, ainfo=None, lmax=None,
@register('map', 'map', iso_filt_method='harmonic', ivar_filt_method='basic', model=True)
def iso_harmonic_ivar_basic_model(imap, sqrt_ivar=1, mask_est=1, ainfo=None,
lmax=None, post_filt_rel_downgrade=1,
post_filt_downgrade_wcs=None,
post_filt_downgrade_wcs=None, cov_ell=None,
lim=1e-6, lim0=None, **kwargs):
"""Filter a map by another map in map space, and then an ell-dependent
matrix in harmonic space. The harmonic filter is measured as the
Expand All @@ -191,6 +201,9 @@ def iso_harmonic_ivar_basic_model(imap, sqrt_ivar=1, mask_est=1, ainfo=None,
bandlimits the measured pseudospectra by the same factor.
post_filt_downgrade_wcs : astropy.wcs.WCS, optional
Assign this wcs to the filtered maps, by default None.
cov_ell : (*preshape, *preshape, nell) np.ndarray, optional
If provided, skip the esimation of the ell-dependent matrix and
instead use this matrix for the harmonic filter, by default None.
lim : float, optional
Set eigenvalues smaller than lim * max(eigenvalues) to zero.
lim0 : float, optional
Expand Down Expand Up @@ -220,14 +233,14 @@ def iso_harmonic_ivar_basic_model(imap, sqrt_ivar=1, mask_est=1, ainfo=None,
return iso_harmonic_ivar_none_model(
imap, mask_est=mask_est, ainfo=ainfo, lmax=lmax,
post_filt_rel_downgrade=post_filt_rel_downgrade,
post_filt_downgrade_wcs=post_filt_downgrade_wcs, lim=lim,
lim0=lim0
post_filt_downgrade_wcs=post_filt_downgrade_wcs,
cov_ell=cov_ell, lim=lim, lim0=lim0
)

@register('map', 'map', iso_filt_method='harmonic_raw', ivar_filt_method='basic', model=True)
def iso_harmonic_raw_ivar_basic_model(imap, sqrt_ivar=1, mask_est=1, ainfo=None,
lmax=None, post_filt_rel_downgrade=1,
post_filt_downgrade_wcs=None,
post_filt_downgrade_wcs=None, cov_ell=None,
lim=1e-6, lim0=None, **kwargs):
"""Filter a map by another map in map space, and then an ell-dependent
matrix in harmonic space. The harmonic filter is measured as the
Expand All @@ -253,6 +266,9 @@ def iso_harmonic_raw_ivar_basic_model(imap, sqrt_ivar=1, mask_est=1, ainfo=None,
bandlimits the measured pseudospectra by the same factor.
post_filt_downgrade_wcs : astropy.wcs.WCS, optional
Assign this wcs to the filtered maps, by default None.
cov_ell : (*preshape, *preshape, nell) np.ndarray, optional
If provided, skip the esimation of the ell-dependent matrix and
instead use this matrix for the harmonic filter, by default None.
lim : float, optional
Set eigenvalues smaller than lim * max(eigenvalues) to zero.
lim0 : float, optional
Expand All @@ -279,11 +295,19 @@ def iso_harmonic_raw_ivar_basic_model(imap, sqrt_ivar=1, mask_est=1, ainfo=None,
# net power effect of var map
sqrt_var = np.reciprocal(sqrt_ivar, where=sqrt_ivar!=0) * (sqrt_ivar!=0)

sqrt_cov_ell, inv_sqrt_cov_ell = utils.measure_iso_harmonic(
imap, 0.5, -0.5, mask_est=mask_est, mask_norm=sqrt_var, ainfo=ainfo, #TODO: is this right in general?
lmax=lmax, lim=lim, lim0=lim0
if cov_ell is None:
sqrt_cov_ell, inv_sqrt_cov_ell = utils.measure_iso_harmonic(
imap, 0.5, -0.5, mask_est=mask_est, mask_norm=sqrt_var, ainfo=ainfo, #TODO: is this right in general?
lmax=lmax, lim=lim, lim0=lim0
)

else:
preshape = cov_ell.shape[:len(cov_ell.shape[:-1]) // 2]
sqrt_cov_ell, inv_sqrt_cov_ell = [
utils.eigpow_flat(
cov_ell, e, preshape, lim=lim, lim0=lim0, copy=True)
for e in (0.5, -0.5)
]

# also need to downgrade the measured power spectra!
sqrt_cov_ell = sqrt_cov_ell[..., :lmax//post_filt_rel_downgrade+1]

Expand Down Expand Up @@ -464,7 +488,7 @@ def iso_harmonic_raw_ivar_scaledep_model(imap, sqrt_ivar=None, ell_lows=None,
ell_highs=None, profile='cosine',
dtype=np.float32, mask_est=1, ainfo=None,
lmax=None, post_filt_rel_downgrade=1,
post_filt_downgrade_wcs=None,
post_filt_downgrade_wcs=None, cov_ell=None,
lim=1e-6, lim0=None, **kwargs):
"""Filter a map by multiple maps in map space, each map for a different
range of ells in harmonic space. Then, filter that map by an ell-dependent
Expand Down Expand Up @@ -501,6 +525,9 @@ def iso_harmonic_raw_ivar_scaledep_model(imap, sqrt_ivar=None, ell_lows=None,
bandlimits the measured pseudospectra by the same factor.
post_filt_downgrade_wcs : astropy.wcs.WCS, optional
Assign this wcs to the filtered maps, by default None.
cov_ell : (*preshape, *preshape, nell) np.ndarray, optional
If provided, skip the esimation of the ell-dependent matrix and
instead use this matrix for the harmonic filter, by default None.
lim : float, optional
Set eigenvalues smaller than lim * max(eigenvalues) to zero.
lim0 : float, optional
Expand All @@ -527,10 +554,18 @@ def iso_harmonic_raw_ivar_scaledep_model(imap, sqrt_ivar=None, ell_lows=None,
# net power effect of var map
sqrt_var = np.reciprocal(sqrt_ivar[-1], where=sqrt_ivar[-1]!=0) * (sqrt_ivar[-1]!=0)

sqrt_cov_ell, inv_sqrt_cov_ell = utils.measure_iso_harmonic(
imap, 0.5, -0.5, mask_est=mask_est, mask_norm=sqrt_var, ainfo=ainfo,
lmax=lmax, lim=lim, lim0=lim0
if cov_ell is None:
sqrt_cov_ell, inv_sqrt_cov_ell = utils.measure_iso_harmonic(
imap, 0.5, -0.5, mask_est=mask_est, mask_norm=sqrt_var, ainfo=ainfo,
lmax=lmax, lim=lim, lim0=lim0
)
else:
preshape = cov_ell.shape[:len(cov_ell.shape[:-1]) // 2]
sqrt_cov_ell, inv_sqrt_cov_ell = [
utils.eigpow_flat(
cov_ell, e, preshape, lim=lim, lim0=lim0, copy=True)
for e in (0.5, -0.5)
]

# also need to downgrade the measured power spectra!
sqrt_cov_ell = sqrt_cov_ell[..., :lmax//post_filt_rel_downgrade+1]
Expand Down Expand Up @@ -813,4 +848,4 @@ def filter_imaps_scaledep(imaps, ell_lows=None, ell_highs=None,
else:
filt_imap += (imap - utils.ell_filter(imap, 1 - prof, lmax=lmaxi))

return filt_imap
return filt_imap
44 changes: 36 additions & 8 deletions mnms/utils.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -575,6 +575,38 @@ def interp1d_bins(bins, y, return_vals=False, **interp1d_kwargs):
else:
return interp1d(x, y, fill_value=fill_value, **interp1d_kwargs)

def eigpow_flat(mat, exp, preshape, **kwargs):
"""Reshape input (*preshape, *preshape, *postshape)-shaped matrix to an
(ncomp, ncomp, *postshape)-shaped matrix and raise to provided power.

Parameters
----------
mat : (preshape) + (preshape) + (postshape) ndarray
Input array
exp : float
Matrix power.
preshape : tuple
tuple denoting the specified preshape.
kwargs : dict
Keyword arguments to eigpow. Cannot provide "axes" keyword argument.

Returns
-------
out : (preshape) + (preshape) + (postshape) ndarray
Input matrix raised to specified power.
"""

if 'axes' in kwargs:
raise ValueError('"axes" keyword argument is not supported.')

ncomp = np.prod(preshape, dtype=int)

mat = mat.reshape(ncomp, ncomp, -1)
mat = eigpow(mat, exp, axes=[0, 1], **kwargs)
mat = mat.reshape(*preshape, *preshape, -1)

return mat

def get_ps_mat(inp, outbasis, exp, lim=1e-6, lim0=None, inbasis='harmonic',
shape=None, wcs=None):
"""Get a power spectrum matrix from input alm's or fft's, raised to a given
Expand Down Expand Up @@ -645,9 +677,7 @@ def get_ps_mat(inp, outbasis, exp, lim=1e-6, lim0=None, inbasis='harmonic',
if preidx1 != preidx2:
mat[(*preidx2, *preidx1)] = mat[(*preidx1, *preidx2)]

mat = mat.reshape(ncomp, ncomp, -1)
mat = eigpow(mat, exp, axes=[0, 1], lim=lim, lim0=lim0)
mat = mat.reshape(*preshape, *preshape, -1)
mat = eigpow_flat(mat, exp, preshape, lim=lim, lim0=lim0)

if outbasis == 'harmonic':
out = mat
Expand Down Expand Up @@ -699,10 +729,8 @@ def get_ps_mat(inp, outbasis, exp, lim=1e-6, lim0=None, inbasis='harmonic',
)
if preidx1 != preidx2:
mat[(*preidx2, *preidx1)] = mat[(*preidx1, *preidx2)]

mat = mat.reshape(ncomp, ncomp, -1)
mat = eigpow(mat, exp, axes=[0, 1], lim=lim, lim0=lim0)
mat = mat.reshape(*preshape, *preshape, -1)

mat = eigpow_flat(mat, exp, preshape, lim=lim, lim0=lim0)

if outbasis == 'fourier':
out = enmap.empty((*preshape, *preshape, shape[-2], shape[-1]//2+1), wcs, inp.real.dtype) # need "real" matrix
Expand Down Expand Up @@ -2977,4 +3005,4 @@ def filter_weighted(imap, ivar, filter, tol=1e-4, ref=0.9):
omap *= -1
omap += imap
omap *= imap != 0
return omap
return omap
12 changes: 7 additions & 5 deletions mnms/wav_noise.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -5,10 +5,8 @@

import numpy as np


# TODO: add lim, lim0 args
def estimate_sqrt_cov_wav_from_enmap(alm, w_ell, shape, wcs, fwhm_fact=2,
verbose=True):
lim=1e-6, lim0=None, verbose=True):
"""Estimate wavelet-based covariance matrix given noise enmap.

Parameters
Expand All @@ -29,6 +27,10 @@ def estimate_sqrt_cov_wav_from_enmap(alm, w_ell, shape, wcs, fwhm_fact=2,
Can also be a function specifying this factor for a given
ell. Function must accept a single scalar ell value and
return one.
lim : float, optional
Set eigenvalues smaller than lim * max(eigenvalues) to zero.
lim0 : float, optional
If max(eigenvalues) < lim0, set whole matrix to zero.
verbose : bool, optional
Print possibly helpful messages, by default True.

Expand Down Expand Up @@ -64,7 +66,7 @@ def estimate_sqrt_cov_wav_from_enmap(alm, w_ell, shape, wcs, fwhm_fact=2,
cov_wav = noise_utils.estimate_cov_wav(alm, ainfo, w_ell, [0, 2], diag=False,
wav_template=wav_template, fwhm_fact=fwhm_fact)
sqrt_cov_wavs = mat_utils.wavmatpow(cov_wav, 0.5, return_diag=True, axes=[[0,1], [2,3]],
inplace=True)
inplace=True, lim=lim, lim0=lim0)

return {'sqrt_cov_mat': sqrt_cov_wavs}

Expand Down Expand Up @@ -159,4 +161,4 @@ def unit_var_wav(minfos, preshape, dtype, seed=None, nthread=0):

wav_uni.add(widx, m_arr, minfo)

return wav_uni
return wav_uni