nzmodes.py

# Codes for creating and smoothing modes
import numpy as np
def getModes(D, Cn, nModes, threshold=0.001):
    '''Calculate the compression matrix from n to u and the
    modes U that multiply each U.  
    Parameters:
    `D`:      The matrix J^T C_c^{-1} J, where C_c is the covariance
              matrix of the observables c, and J is the Jacobian dc/dn,
              shape=(N,N)
    `Cn`:     Covariance matrix of the n(z) parameter vector n, shape=(N,N)
    `nModes`: Number of modes to retain.
    `threshold`:  The ratio of smallest singular value retained to the 
              largest singular value of an internal matrix during a
              matrix inversion.
    Returns:
    `X`:      The compression matrix, u = X @ n, shape=(nModes,N)
    `U`:      The decoding matrix, n' = U @ u, shape=(N,nModes)
    `dchidq`: The mean chisq shift generated by each mode, shape=(nModes)

    You can safely use fewer than `nModes` of the returned X and U if you 
    think the discarded `dchisq` values are small enough.'''

    Dval, Dvec = np.linalg.eigh(D)
    DeltaD = np.diag(np.sqrt(Dval)) @ Dvec.T # Transposing numpy's eigenvectors

    A = DeltaD @ Cn @ DeltaD.T
    Aval, Avec = np.linalg.eigh(A)
    # Sort eigenvalues to be decreasing
    order = np.argsort(-Aval)
    Aval = Aval[order]
    Avec = Avec[:,order]

    X = Avec.T[:nModes] @ DeltaD   # The encoder from n->u

    # This is the SVD of X given how we made it
    uX = Avec.T
    sX = np.sqrt(Dval)
    vX = Dvec
    
    # Get U from the SVD-controlled inverse of X
    thresh = 0.001
    U = vX @ np.diag(np.where( sX > threshold*np.max(sX), 1/sX, 0.)) @ uX.T[:,:nModes]
    
    return X, U, Aval[:nModes]

class Tz:
    def __init__(self, dz, nz, z0=None):
        '''Class representing sawtooth n(z) kernels (bins) in z.
        First kernel is centered at z0, which defaults to dz if not
        given.  If z0<dz, then its triangle is adjusted to go to zero at
        0, then peak at z0, down to zero at z0+dz
        Arguments:
        `dz`: the step between kernel centers
        `nz`: the number of kernels.
        `z0`: peak of first bin'''
        self.dz = dz
        self.nz = nz
        if z0 is None:
            self.z0 = dz
        else:
            self.z0 = z0
        # Set a flag if we need to treat kernel 0 differently
        self.cut0 = self.z0<dz

    def __call__(self,k,z):
        '''Evaluate dn/dz for the kernel with index k at an array of
        z values.'''
        # Doing duplicative calculations to make this compatible
        # with JAX arithmetic.

        if self.cut0 and k==0:
            # Lowest bin is unusual:
            out = np.where(z>self.z0, 1-(z-self.z0)/self.dz, z/self.z0)
            out = np.maximum(0., out) / ((self.z0+self.dz)/2.)
        else:
            out = np.maximum(0., 1 - np.abs((z-self.z0)/self.dz-k)) / self.dz
        return out
    
    def zbounds(self):
        '''Return lower, upper bounds in z of all the bins in (nz,2) array'''
        zmax = np.arange(1,1+self.nz)*self.dz + self.z0
        zmin = zmax - 2*self.dz
        if self.cut0:
            zmin[0] = 0.
        return np.stack( (zmin, zmax), axis=1)
    
    def dndz(self,coeffs, z):
        '''Calculate dn/dz at an array of z values given set(s) of
        coefficients for the kernels/bins.  The coefficients will
        be normalized to sum to unity, i.e. they will represent the
        fractions within each kernel.
        Arguments:
        `coeffs`:  Array of kernel fractions of shape [...,nz]
        `z`:       Array of redshifts of arbitrary length
        Returns:
        Array of shape [...,len(z)] giving dn/dz at each z for
        each set of coefficients.'''
        
        # Make the kernel coefficients at the z's
        kk = np.array([self(k,z) for k in range(self.nz)])
        return np.einsum('...i,ij->...j',coeffs,kk) / np.sum(coeffs, axis=-1)