Merge pull request #9 from cosmodesi/develop

Develop

thecov/base.py

@@ -137,6 +137,45 @@ def shape(self):
         '''
 
         return self.cov.shape
+
+    @property
+    def eig(self):
+        '''Compute the eigenvalues and right eigenvectors of the covariance.
+
+        Returns
+        -------
+        A namedtuple with the following attributes:
+            eigenvalues
+            (..., M) array
+                The eigenvalues, each repeated according to its multiplicity.
+                The eigenvalues are not necessarily ordered. The resulting
+                array will be of complex type, unless the imaginary part is
+                zero in which case it will be cast to a real type. When a is
+                real the resulting eigenvalues will be real (0 imaginary
+                part) or occur in conjugate pairs
+
+            eigenvectors
+            (...), M, M) array
+                The normalized (unit “length”) eigenvectors, such that the
+                column eigenvectors[:,i] is the eigenvector corresponding to
+                the eigenvalue eigenvalues[i].
+        '''
+
+        return np.linalg.eig(self.cov)
+
+    @property
+    def eigvals(self):
+        '''Compute the eigenvalues of the covariance.
+
+        Returns
+        -------
+        (..., M,) ndarray
+            The eigenvalues, each repeated according to its multiplicity.
+            They are not necessarily ordered, nor are they necessarily
+            real for real matrices.
+        '''
+
+        return np.linalg.eigvals(self.cov)
 
     def save(self, filename):
         '''Saves the covariance as a .npz file with a specified filename.
@@ -598,4 +637,79 @@ def nmodes(self, nmodes):
         return nmodes
 
 class MultipoleFourierCovariance(MultipoleCovariance, FourierBinned):
-    pass
+
+    def __init__(self):
+        MultipoleCovariance.__init__(self)
+        FourierBinned.__init__(self)
+
+    @property
+    def kmid_matrices(self):
+        kfull = np.concatenate([self.kmid for _ in self.ells])
+        k1 = np.einsum('i,j->ij', kfull, np.ones_like(kfull))
+        k2 = k1.T
+
+        return k1, k2
+
+    @property
+    def ell_matrices(self):
+        ell_array = np.einsum('i,j->ij', self.ells, np.ones(self.kbins)).flatten()
+
+        ell1 = np.einsum('i,j->ij', ell_array, np.ones_like(ell_array))
+        ell2 = ell1.T
+
+        return ell1, ell2
+
+
+    def save_table(self, filename, ells_both_ways=False, fmt=['%.d', '%.d', '%.4f', '%.4f', '%.8e']):
+        k1, k2 = self.kmid_matrices
+        ell1, ell2 = self.ell_matrices
+
+        cov = self.cov
+
+        mask = ell1 <= ell2 if not ells_both_ways else np.ones_like(ell1, dtype=bool)
+
+        np.savetxt(filename, np.concatenate([ell1[mask].reshape(-1, 1),
+                                             ell2[mask].reshape(-1, 1),
+                                               k1[mask].reshape(-1, 1),
+                                               k2[mask].reshape(-1, 1),
+                                              cov[mask].reshape(-1, 1)], axis=1), fmt=fmt, header='ell1 ell2 kmid1 kmid2 cov')
+    def load_table(self, filename):
+        ell1, ell2, k1, k2, value = np.loadtxt(filename).T
+
+        k = np.unique(k1)
+        kbins = len(k)
+
+        assert np.allclose(k, np.unique(k2)), "k1 and k2 are not consistent"
+
+        dk = np.mean(np.diff(k))
+        kmin = k.min() - dk/2
+        kmax = k.max() + dk/2
+
+        ells = np.unique(ell1)
+        assert np.allclose(ells, np.unique(ell2)), "ell1 and ell2 are not consistent"
+
+        ells_both_ways = len(value) == (len(ells)*kbins)**2
+        ells_one_way   = len(value) == (len(ells)**2 + len(ells))/2 * kbins**2
+
+        assert ells_one_way or ells_both_ways, 'length of covariance file doesn\'nt match'
+
+        self.set_kbins(kmin, kmax, dk)
+
+        assert np.allclose(np.unique(k1), self.kmid), "k bins are not linearly spaced"
+
+        kmid_matrix = np.einsum('i,j->ij', k, np.ones_like(k))
+
+        for l1, l2 in itt.combinations_with_replacement(ells, r=2):
+            block_mask = (ell1 == l1) & (ell2 == l2)
+            assert np.allclose(k1[block_mask].reshape(kmid_matrix.shape),   kmid_matrix)
+            assert np.allclose(k2[block_mask].reshape(kmid_matrix.T.shape), kmid_matrix.T)
+            c = value[block_mask].reshape(kbins, kbins)
+            self.set_ell_cov(l1, l2, c)
+
+        return self
+
+    @classmethod
+    def from_table(cls, filename):
+        cov = cls()
+        cov.load_table(filename)
+        return cov

thecov/covariance.py

@@ -23,7 +23,7 @@
 
 __all__ = ['GaussianCovariance']
 
-class GaussianCovariance(base.MultipoleCovariance, base.FourierBinned):
+class GaussianCovariance(base.MultipoleFourierCovariance):
     '''Gaussian covariance matrix of power spectrum multipoles in a given geometry.
 
     Attributes
@@ -33,8 +33,7 @@ class GaussianCovariance(base.MultipoleCovariance, base.FourierBinned):
     '''
 
     def __init__(self, geometry):
-        base.MultipoleCovariance.__init__(self)
-        base.FourierBinned.__init__(self)
+        base.MultipoleFourierCovariance.__init__(self)
 
         self.geometry = geometry
 
@@ -52,8 +51,9 @@ def shotnoise(self):
         -------
         float
             Shotnoise value.'''
-
-        return self.geometry.shotnoise
+
+        return (1 + self.alphabar)*self.geometry.I('12')/self.geometry.I('22')
+        # return self.geometry.shotnoise
 
     def set_shotnoise(self, shotnoise):
         '''Set shotnoise to specified value. Also scales alphabar so that (1 + alphabar)*I12/I22
@@ -67,7 +67,7 @@ def set_shotnoise(self, shotnoise):
 
         print(f'Estimated shotnoise was {self.geometry.shotnoise}')
         print(f'Setting shotnoise to {shotnoise}.')
-        self.geometry.shotnoise = shotnoise
+        # self.geometry.shotnoise = shotnoise
         self.alphabar = shotnoise*self.geometry.I('22')/self.geometry.I('12') - 1
         print(f'Setting alphabar to {self.alphabar} based on given shotnoise value.')
 
@@ -159,6 +159,9 @@ def _compute_covariance_box(self):
         for l1, l2 in itt.combinations_with_replacement(self.ells, r=2):
             self.set_ell_cov(l1, l2, 2/self.nmodes * np.diag(cov[l1, l2]))
 
+        if (self.eigvals < 0).any():
+                    print('WARNING: Covariance matrix is not positive definite.')
+
         return self
 
     def _compute_covariance_survey(self):
@@ -222,5 +225,11 @@ def _compute_covariance_survey(self):
         self.set_ell_cov(0, 4, cov[:, :, 4])
         self.set_ell_cov(2, 4, cov[:, :, 5])
 
+        if (self.eigvals < 0).any():
+            print('WARNING: Covariance matrix is not positive definite.')
+
+        if not np.allclose(self.cov, self.cov.T):
+            print('WARNING: Covariance matrix is not symmetric.')
+
         return self
 

thecov/geometry.py

@@ -355,6 +355,8 @@ def __init__(self, random_catalog=None, Nmesh=None, BoxSize=None, boxsize_factor
             'interlaced':  True,
             'compensated': True,
             'resampler':   'tsc',
+            'position': 'BoxPosition',
+            'weight': 'WEIGHT',
         }
 
         if mesh_kwargs is not None:
@@ -398,7 +400,7 @@ def I(self, W):
         w = W.lower().replace("i", "").replace("w", "")
         if w not in self._I:
             self.W_cat(w)
-            self._I[w] = (self._randoms[f'W{w}'].sum() * self.alpha).compute()
+            self._I[w] = ((self._randoms[f'W{w}']*self._randoms[f'WEIGHT']).sum() * self.alpha).compute()
         return self._I[w]
 
     def W(self, W):
@@ -450,7 +452,7 @@ def compute_cartesian_fft(self, W):
 
         with self.tqdm(total=22, desc=f'Computing moments of W{w}') as pbar:
             self.set_cartesian_fft(f'W{w}', self._format_fft(self.randoms.to_mesh(
-                value=f'W{w}', position='BoxPosition', **mesh_kwargs).paint(mode='complex')))
+                value=f'W{w}',**mesh_kwargs).paint(mode='complex')))
 
             # Check if zero mode of FFT Wij is consistent with integral Iij
             center = self.Nmesh//2
@@ -470,7 +472,7 @@ def compute_cartesian_fft(self, W):
                 self.randoms[label] = self.randoms[f'W{w}'] * \
                     x[i]*x[j] / (x[0]**2 + x[1]**2 + x[2]**2)
                 self.set_cartesian_fft(label, self._format_fft(self.randoms.to_mesh(
-                    value=label, position='BoxPosition', **mesh_kwargs).paint(mode='complex')))
+                    value=label, **mesh_kwargs).paint(mode='complex')))
 
                 pbar.update(1)
 
@@ -479,7 +481,7 @@ def compute_cartesian_fft(self, W):
                 self.randoms[label] = self.randoms[f'W{w}'] * x[i] * \
                     x[j]*x[k]*x[l] / (x[0]**2 + x[1]**2 + x[2]**2)**2
                 self.set_cartesian_fft(label, self._format_fft(self.randoms.to_mesh(
-                    value=label, position='BoxPosition', **mesh_kwargs).paint(mode='complex')))
+                    value=label, **mesh_kwargs).paint(mode='complex')))
 
                 pbar.update(1)
 
@@ -488,12 +490,12 @@ def _format_fft(self, fourier):
         # PFFT omits modes that are determined by the Hermitian condition. Adding them:
         fourier_full = utils.r2c_to_c2c_3d(fourier)
 
-        return fft.fftshift(fourier_full)[::-1, ::-1, ::-1] * self.ngals
+        return fft.fftshift(fourier_full) * self.ngals
 
     def _unformat_fft(self, fourier, window):
 
         fourier_cut = fft.ifftshift(
-            fourier[::-1, ::-1, ::-1])[:, :, :fourier.shape[2]//2+1] / self.ngals
+            fourier)[:, :, :fourier.shape[2]//2+1] / self.ngals
 
         return np.conj(fourier_cut) if window == '12' else fourier_cut
 
@@ -596,11 +598,15 @@ def compute_window_kernels(self):
 
         kfun = 2*np.pi/self.BoxSize
 
-        # sample kmodes from each k1 bin
-        # kmodes = np.array([[utils.sample_from_shell(kmin/kfun, kmax/kfun) for _ in range(
-        #     self.kmodes_sampled)] for kmin, kmax in zip(self.kedges[:-1], self.kedges[1:])])
-
-        kmodes, Nmodes = utils.sample_from_cube(self.kmax/kfun, self.dk/kfun, self.kmodes_sampled)
+        kmodes = np.array([[utils.sample_from_shell(kmin/kfun, kmax/kfun) for _ in range(
+            self.kmodes_sampled)] for kmin, kmax in zip(self.kedges[:-1], self.kedges[1:])])
+
+        Nmodes = utils.nmodes(self.BoxSize**3, self.kedges[:-1], self.kedges[1:])
+
+        # sample kmodes from each k1 bin        
+        # kmodes, Nmodes = utils.sample_kmodes(self.kmax, self.dk, self.BoxSize, self.kmodes_sampled)
+
+        # kmodes, Nmodes = utils.sample_from_cube(self.kmax/kfun, self.dk/kfun, self.kmodes_sampled)
 
         # Note: as the window falls steeply with k, only low-k regions are needed for the calculation.
         # Therefore, high-k modes in the FFTs can be discarded

thecov/tests/test_base.py

@@ -1,5 +1,6 @@
 import numpy as np
 import thecov.base
+import os
 
 def test_multipole_covariance_symmetrization():
     cov00, cov22, cov44, cov02, cov04, cov24 = np.random.rand(6, 100, 100)
@@ -76,10 +77,36 @@ def test_multipole_covariance_addition():
 
     addition = cov1 + cov2
 
-    assert addition.get_ell_cov(0,0).cov == cov1_00 + cov2_00
-    assert addition.get_ell_cov(2,2).cov == cov1_22 + cov2_22
-    assert addition.get_ell_cov(4,4).cov == cov1_44 + cov2_44
+    assert (addition.get_ell_cov(0,0).cov == cov1_00 + cov2_00).all()
+    assert (addition.get_ell_cov(2,2).cov == cov1_22 + cov2_22).all()
+    assert (addition.get_ell_cov(4,4).cov == cov1_44 + cov2_44).all()
 
-    assert addition.get_ell_cov(0,2).cov == cov1_02 + cov2_02
-    assert addition.get_ell_cov(0,4).cov == cov1_04 + cov2_04
-    assert addition.get_ell_cov(2,4).cov == cov1_24 + cov2_24
+    assert (addition.get_ell_cov(0,2).cov == cov1_02 + cov2_02).all()
+    assert (addition.get_ell_cov(0,4).cov == cov1_04 + cov2_04).all()
+    assert (addition.get_ell_cov(2,4).cov == cov1_24 + cov2_24).all()
+
+def test_multipole_fourier_covariance_save_load_table():
+    cov = thecov.base.MultipoleFourierCovariance()
+    cov.set_kbins(0., 0.4, 0.005)
+
+    cov00, cov22, cov44, cov02, cov04, cov24 = np.random.rand(6, cov.kbins, cov.kbins)
+
+    cov.set_ell_cov(0, 0, cov00)
+    cov.set_ell_cov(2, 2, cov22)
+    cov.set_ell_cov(4, 4, cov44)
+
+    cov.set_ell_cov(0, 2, cov02)
+    cov.set_ell_cov(0, 4, cov04)
+    cov.set_ell_cov(4, 2, cov24.T)
+
+    cov.save_table('test1.txt')
+    cov.save_table('test2.txt', ells_both_ways=True)
+
+    cov1 = thecov.base.MultipoleFourierCovariance.from_table('test1.txt')
+    cov2 = thecov.base.MultipoleFourierCovariance.from_table('test2.txt')
+
+    assert np.allclose(cov1.cov, cov.cov)
+    assert np.allclose(cov2.cov, cov.cov)
+
+    os.remove('test1.txt')
+    os.remove('test2.txt')

thecov/utils.py

@@ -139,6 +139,24 @@ def sample_from_cube(rmax, dr, max_modes:int=200):
 
     return modes, Nmodes
 
+def sample_kmodes(kmax, dk, BoxSize, kmodes_sampled, shell_approx_bin=10):
+
+    # Wavelength where spherical shell approximation kicks in
+    k_shell = shell_approx_bin*dk
+
+    kfun = 2*np.pi/BoxSize
+
+    # Uses full cube from k = 0 to k_shell
+    cube_modes, cube_Nmodes = sample_from_cube(k_shell/kfun, dk/kfun, max_modes=kmodes_sampled)
+
+    # Uses spherical shell approximation from k = k_shell to kmax
+    kedges_shell = np.arange(k_shell, kmax + dk, dk)
+    shell_modes = [np.array([sample_from_shell(kmin/kfun, kmax/kfun) for _ in range(
+                    kmodes_sampled)]) for kmin, kmax in zip(kedges_shell[:-1], kedges_shell[1:])]
+    shell_Nmodes = nmodes(BoxSize**3, kedges_shell[:-1], kedges_shell[1:])
+
+    return cube_modes + shell_modes, np.array(cube_Nmodes + list(shell_Nmodes))
+
 def nmodes(volume, kmin, kmax):
     '''Compute the number of modes in a given shell.