Skip to content

Commit

Permalink
Merge pull request #4 from cosmodesi/develop
Browse files Browse the repository at this point in the history 
Develop
  • Loading branch information
@otavioalonso
otavioalonso authored Jun 27, 2023
2 parents 0e50809 + ab49849 commit a0872e5
Show file tree
Hide file tree
Showing 4 changed files with 83 additions and 39 deletions.
4 changes: 2 additions & 2 deletions setup.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -2,7 +2,7 @@

setup(
name='thecov',
version='0.1.0',
version='0.1.1',
description='A python package to compute analytic covariance matrices based on Wadekar & Scoccimarro 2019 (arXiv:1910.02914).',
url='https://github.com/cosmodesi/thecov',
author='Otavio Alves',
Expand All @@ -12,7 +12,7 @@
install_requires=['numpy',
'scipy',
'tqdm',
'dask',
# 'dask',
# 'nbodykit',
],

Expand Down
18 changes: 9 additions & 9 deletions thecov/covariance.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -64,9 +64,9 @@ def set_pk(self, pk, ell, has_shotnoise=False):
Whether the power spectrum has shotnoise included or not.
'''

if ell == 0 and not has_shotnoise:
print(f'Adding shotnoise = {self.shotnoise} to ell = 0.')
self._pk[ell] = pk + self.shotnoise
if ell == 0 and has_shotnoise:
print(f'Removing shotnoise = {self.shotnoise} from ell = 0.')
self._pk[ell] = pk - self.shotnoise
else:
self._pk[ell] = pk

Expand All @@ -85,9 +85,9 @@ def get_pk(self, ell, force_return=False, remove_shotnoise=True):

if ell in self._pk.keys():
pk = self._pk[ell]
if remove_shotnoise and ell == 0:
print(f'Removing shotnoise = {self.shotnoise} from ell = 0.')
return pk - self.shotnoise
if (not remove_shotnoise) and ell == 0:
print(f'Adding shotnoise = {self.shotnoise} to ell = 0.')
return pk + self.shotnoise
return pk
elif force_return:
return np.zeros(self.kbins)
Expand Down Expand Up @@ -165,7 +165,7 @@ def _compute_covariance_survey(self):
kbins = self.kbins

# If the power spectrum for a given ell is not set, use a zero array instead
P0 = self.get_pk(0, force_return=True)
P0 = self.get_pk(0, force_return=True, remove_shotnoise=True)
P4 = self.get_pk(2, force_return=True)
P2 = self.get_pk(4, force_return=True)

Expand All @@ -187,13 +187,13 @@ def _compute_covariance_survey(self):
WinKernel[ki, delta_k, 6]*P4[ki]*P0[kj] + \
WinKernel[ki, delta_k, 7]*P4[ki]*P2[kj] + \
WinKernel[ki, delta_k, 8]*P4[ki]*P4[kj] + \
1.01*(
(1 + self.geometry.alpha)*(
WinKernel[ki, delta_k, 9]*(P0[ki] + P0[kj])/2. +
WinKernel[ki, delta_k, 10]*P2[ki] + WinKernel[ki, delta_k, 11]*P4[ki] +
WinKernel[ki, delta_k, 12]*P2[kj] +
WinKernel[ki, delta_k, 13]*P4[kj]
) + \
1.01**2 * WinKernel[ki, delta_k, 14]
(1 + self.geometry.alpha)**2 * WinKernel[ki, delta_k, 14]

self.set_ell_cov(0, 0, cov[:, :, 0])
self.set_ell_cov(2, 2, cov[:, :, 1])
Expand Down
78 changes: 57 additions & 21 deletions thecov/geometry.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -32,7 +32,18 @@


class Geometry(ABC):
pass

def save_attributes(self, filename, attrs):
np.savez(filename if filename.strip()
[-4:] == '.npz' else f'{filename}.npz', **{a: getattr(self, a) for a in attrs})

def load_attributes(self, filename, attrs=None):
with np.load(filename, mmap_mode='r') as data:
if attrs is None:
attrs = data.files
for a in attrs:
setattr(self, a, data[a])



class BoxGeometry(Geometry):
Expand Down Expand Up @@ -178,7 +189,7 @@ def set_nz(self, zedges, nz, *args, **kwargs):
self.nbar = np.average(self.nz, weights=bin_volume)
print(f'Estimated nbar: {self.nbar:.3e} (Mpc/h)^-3')

def set_randoms(self, randoms, alpha=1.0):
def set_randoms(self, randoms, alpha=1.0, fsky=None):
'''Estimates the effective volume and number density of the box based on a
provided catalog of randoms.
Expand All @@ -189,15 +200,20 @@ def set_randoms(self, randoms, alpha=1.0):
alpha : float, optional
Factor to multiply the number density of the randoms. Default is 1.0.
'''
import healpy as hp

from nbodykit.algorithms import zhist

nside = 512
hpixel = hp.ang2pix(nside, randoms['RA'], randoms['DEC'], lonlat=True)
unique_hpixels = np.unique(hpixel)
self.fsky = len(unique_hpixels.compute())/hp.nside2npix(nside)
if fsky is None:
import healpy as hp

print(f'fsky estimated from randoms: {self.fsky:.3f}')
nside = 512
hpixel = hp.ang2pix(nside, randoms['RA'], randoms['DEC'], lonlat=True)
unique_hpixels = np.unique(hpixel)
self.fsky = len(unique_hpixels.compute())/hp.nside2npix(nside)

print(f'fsky estimated from randoms: {self.fsky:.3f}')
else:
self.fsky = fsky

nz_hist = zhist.RedshiftHistogram(
randoms, self.fsky, self.cosmo, redshift='Z')
Expand Down Expand Up @@ -274,7 +290,7 @@ class SurveyGeometry(Geometry, base.FourierBinned):
.. [1] https://arxiv.org/abs/1910.02914
'''

def __init__(self, random_catalog=None, Nmesh=None, BoxSize=None, kmax_mask=0.05, delta_k_max=3, kmodes_sampled=250, alpha=1.0, mesh_kwargs=None, tqdm=shell_tqdm):
def __init__(self, random_catalog=None, Nmesh=None, BoxSize=None, kmax_mask=0.05, delta_k_max=3, kmodes_sampled=400, alpha=1.0, mesh_kwargs=None, tqdm=shell_tqdm):

assert Nmesh is None or Nmesh % 2 == 1, 'Please, use an odd integer for Nmesh.'

Expand Down Expand Up @@ -304,16 +320,12 @@ def __init__(self, random_catalog=None, Nmesh=None, BoxSize=None, kmax_mask=0.05

survey_center = (pos_max + pos_min)/2
print(f'Survey center is at xyz = {survey_center}.')

# print(f'Adding {-pos_min} to xyz coordinates to make them positive.')
# self._randoms['OriginalPosition'] = self._randoms['Position']
# self._randoms['Position'] -= survey_center
else:
self._ngals = None

if BoxSize is None and random_catalog is not None:
# Estimate BoxSize from random catalog
self.BoxSize = max(pos_max - pos_min)
self.BoxSize = max(pos_max - pos_min)*1.1
print(f'Using BoxSize = {self.BoxSize}.')
elif np.ndim(BoxSize) == 0:
self.BoxSize = BoxSize
Expand All @@ -326,7 +338,7 @@ def __init__(self, random_catalog=None, Nmesh=None, BoxSize=None, kmax_mask=0.05

if Nmesh is None:
# Pick odd value that will give at least k_mask = kmax_mask in the FFTs
self.Nmesh = int(kmax_mask*self.BoxSize/np.pi)//2 * 2 + 1
self.Nmesh = int(np.ceil(kmax_mask*self.BoxSize/np.pi))
print(f'Using Nmesh = {self.Nmesh} for the window FFTs.')
else:
self.Nmesh = Nmesh
Expand Down Expand Up @@ -549,6 +561,10 @@ def randoms(self):
@property
def ngals(self):
return self._ngals

@ngals.setter
def ngals(self, ngals):
self._ngals = ngals

@property
def kbins(self):
Expand Down Expand Up @@ -634,8 +650,17 @@ def save_window_kernels(self, filename):
----------
filename : str
Name of the file to save the window kernels.'''
np.savez(filename if filename.strip()
[-4:] == '.npz' else f'{filename}.npz', WinKernel=self.WinKernel, kmin=self.kmin, kmax=self.kmax, dk=self.dk)
self.save_attributes(filename, ['alpha',
'delta_k_max',
'kmodes_sampled',
'shotnoise',
'ngals',
'BoxSize',
'Nmesh',
'dk',
'kmax',
'kmin',
'WinKernel'])

def load_window_kernels(self, filename):
'''Load the window kernels from a file.
Expand All @@ -645,9 +670,20 @@ def load_window_kernels(self, filename):
filename : str
Name of the file to load the window kernels from.
'''
with np.load(filename, mmap_mode='r') as data:
self.WinKernel = data['WinKernel']
self.set_kbins(data['kmin'], data['kmax'], data['dk'])
self.load_attributes(filename)

@classmethod
def from_window_kernels_file(cls, filename):
'''Create geometry object from window kernels file.
Parameters
----------
filename : str
Name of the file to load the window kernels from.
'''
geometry = cls()
geometry.load_window_kernels(filename)
return geometry

@staticmethod
def _compute_window_kernel_row(args):
Expand Down Expand Up @@ -1010,7 +1046,7 @@ def shotnoise(self):
if self._shotnoise is None:
return self.I('12')/self.I('22')
else:
return self.I('12')/self.I('22')
return self._shotnoise

@shotnoise.setter
def shotnoise(self, shotnoise):
Expand Down
22 changes: 15 additions & 7 deletions thecov/utils.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -311,7 +311,7 @@ def plot_cov(cova, covb=None, kmax=None, label_a=None, label_b=None, vmin=-1, vm
'''
return plot_cov_array(cova=cova.cov, covb=covb.cov if covb is not None else None, k=cova.kmid, kmax=kmax, label_a=None, label_b=None, vmin=-1, vmax=1, num_ticks=5, **kwargs)

def plot_cov_diag(cov, k=None, label=None, klim=None, colors=['k', 'r', 'g', 'b'], portrait=False):
def plot_cov_diag(cov, k=None, label=None, klim=None, colors=['k', 'r', 'g', 'b'], portrait=False, logplot=True):
'''Plot the diagonal of a MultipoleCovariance object.
Parameters
Expand Down Expand Up @@ -370,24 +370,32 @@ def plot_cov_diag(cov, k=None, label=None, klim=None, colors=['k', 'r', 'g', 'b'
k = c.kmid
break

p2 = cov[0].get_pk(0)**2 if hasattr(cov[0], 'get_pk') else 1.0

p2 = 1
div_by_p2 = False
for c in cov:
if hasattr(c, 'get_pk'):
p2 = c.get_pk(0)**2
div_by_p2 = True
break

for (l1, l2), ax1, ax2 in zip([(0,0), (2,2), (4,4), (0,2), (0,4), (2,4)], axes1, axes2):

for c,l,color in zip(cov,label,colors):

a = np.diag(c.get_ell_cov(l1,l2).cov)/p2

ax1.semilogy(k, a, label=l, c=color)
ax1.semilogy(k, -a, c=color, ls='dashed')
if logplot:
ax1.semilogy(k, a, label=l, c=color)
ax1.semilogy(k, -a, c=color, ls='dashed')
else:
ax1.plot(k, a, label=l, c=color)

for c,l,color in zip(cov[1:], label[1:], colors[1:]):
a = np.diag(c.get_ell_cov(l1,l2).cov)/p2
b = np.diag(cov[0].get_ell_cov(l1,l2).cov)/p2

ax2.plot(k, a/b-1, label='frac. diff', c=color)

ax1.set_ylabel(r"$C_{l1l2}(k,k)/P_0(k)^2$".replace('l1', str(l1)).replace('l2', str(l2)))
ax1.set_ylabel(f"$C_{{{l1}{l2}}}(k,k){r'/P_0(k)^2' if div_by_p2 else ''}$")
ax2.set_xlabel('k [h/Mpc]')

if len(cov) > 1:
Expand Down

0 comments on commit a0872e5

Please sign in to comment.