foscat 2025.6.1__py3-none-any.whl → 2025.7.1__py3-none-any.whl

foscat/HealSpline.py

@@ -0,0 +1,211 @@
+from scipy.interpolate import interp1d
+import foscat.CircSpline as sc
+import foscat.Spline1D as sc1d
+import numpy as np
+import healpy as hp
+
+class heal_spline:
+    def __init__(
+        self,
+        level):
+        nside=2**level
+        self.nside_store=2**(level//2)
+        self.spline_tree={}
+        
+        self.nside=nside
+        #compute colatitude
+        idx_th=np.zeros([4*nside],dtype='int')
+        n=0
+        d=0
+        for k in range(nside):
+            d+=4
+            idx_th[k]=n
+            n+=d
+            
+        for k in range(2*nside-1):
+            idx_th[k+nside]=n
+            n+=d
+                
+        for k in range(nside):
+            idx_th[k+3*nside-1]=n
+            n+=d
+            d-=4
+        idx_th[-1]=12*nside**2
+    
+        th0_val,ph0_val=hp.pix2ang(self.nside,idx_th[:-1])
+        self.th0_val=th0_val    
+        self.ph0_val=ph0_val    
+        
+        self.idx_th=idx_th
+        
+        #init spline table
+
+        self.spline_lat=sc1d.Spline1D(4*self.nside-1,3)
+
+        #convert colatitude in ring index
+        self.f_interp_th = interp1d(np.concatenate([[0],(th0_val[:-1]+th0_val[1:])/2,[np.pi]],0),
+                                    np.arange(4*self.nside)/(4*self.nside), 
+                                   kind='cubic', fill_value='extrapolate')    
+
+    
+    def ang2weigths(self,th,ph,threshold=1E-2,nest=True): 
+        th0=self.f_interp_th(th).flatten()
+        
+        idx_lat,w_th=self.spline_lat.eval(th0.flatten())
+            
+        www     = np.zeros([4,4,th0.shape[0]])
+        all_idx = np.zeros([4,4,th0.shape[0]],dtype='int')
+            
+        iring_tab=np.unique(idx_lat)
+        for iring in iring_tab:
+            spline_table=sc.CircSpline(self.idx_th[iring+1]-self.idx_th[iring],3)
+            for k in range(4):
+                iii=np.where(idx_lat[k]==iring)[0]
+                idx,w=spline_table.eval((ph[iii]-self.ph0_val[iring])/(2*np.pi))
+                idx=idx+self.idx_th[iring]
+                for m in range(4):
+                    www[k,m,iii]=w[m]*w_th[k,iii]
+                    all_idx[k,m,iii]=idx[m]
+
+        www=www.reshape(16,www.shape[2])
+        all_idx=all_idx.reshape(16,all_idx.shape[2])
+            
+        heal_idx,inv_idx = np.unique(all_idx,
+                                    return_inverse=True)
+        all_idx = inv_idx
+        if nest:
+            heal_idx = hp.ring2nest(self.nside,heal_idx)
+        self.cell_ids = heal_idx
+            
+        hit=np.bincount(all_idx.flatten(),weights=www.flatten())
+        www[hit[all_idx]<threshold]=0.0
+        www=www/np.sum(www,0)[None,:]
+        return www,all_idx,heal_idx
+        
+    def P(self,x,www,all_idx):
+        return np.sum(www*x[all_idx],0)
+    
+    #PT(y) must return a 1D NumPy array of shape (N,)
+    def PT(self,y,www,all_idx,hit):
+        value=np.bincount(all_idx.flatten(),weights=(www*y[None,:]).flatten())
+        return value*hit
+
+    # the data is of dimension M
+    # the x is of dimension N=12*nside**2
+    
+    def conjugate_gradient_normal_equation(self,data, x0, www, all_idx, max_iter=100, tol=1e-8, verbose=True):
+        """
+        Solve (PᵗP)x = Pᵗy using explicit Conjugate Gradient without scipy.cg.
+    
+        Parameters:
+        ----------
+        P       : function(x) → forward operator (ℝⁿ → ℝᵐ)
+        PT      : function(y) → adjoint operator (ℝᵐ → ℝⁿ)
+        data    : array_like, observed data y ∈ ℝᵐ
+        x0      : array_like, initial guess for x ∈ ℝⁿ
+        max_iter: maximum number of iterations
+        tol     : convergence tolerance on relative residual
+        verbose : if True, print convergence info
+    
+        Returns:
+        -------
+        x       : estimated solution ∈ ℝⁿ
+        """
+        x = x0.copy()
+        
+        hit=np.bincount(all_idx.flatten(),weights=www.flatten())
+        hit[hit>0]=1/hit[hit>0]
+        
+        # Compute b = Pᵗ y # This part could be distributed easily
+        b = self.PT(data,www,all_idx,hit)
+        
+        # Compute initial residual r = b - A x = b - Pᵗ P x
+        Ax = self.PT(self.P(x,www,all_idx),www,all_idx,hit)# This part could be distributed easily
+        r = b - Ax
+    
+        # Initial direction
+        p = r.copy()
+        rs_old = np.dot(r, r)
+    
+        for i in range(max_iter):
+            # Apply A p = Pᵗ P p
+            Ap = self.PT(self.P(p,www,all_idx),www,all_idx,hit)# This part could be distributed easily
+    
+            alpha = rs_old / np.dot(p, Ap)
+            x += alpha * p
+            r -= alpha * Ap
+    
+            rs_new = np.dot(r, r)
+            
+            if verbose and i%50==0:
+                print(f"Iter {i:03d}: residual = {np.sqrt(rs_new):.3e}")
+    
+            if np.sqrt(rs_new) < tol:
+                if verbose:
+                    print(f"Converged. Iter {i:03d}: residual = {np.sqrt(rs_new):.3e}")
+                break
+    
+            p = r + (rs_new / rs_old) * p
+            rs_old = rs_new
+    
+        return x
+        
+    def Fit(self,X,th,ph,threshold=1E-2,nest=True, max_iter=100, tol=1e-8):
+        
+        www,all_idx,hidx=self.ang2weigths(th,ph,threshold=threshold,nest=nest)
+
+        self.heal_idx=hidx
+        self.spline=self.conjugate_gradient_normal_equation(X,
+                                                            (self.heal_idx*0).astype('float'),
+                                                            www,
+                                                            all_idx
+                                                            , max_iter=max_iter,
+                                                            tol=tol
+                                                            )
+        scale=(self.nside//self.nside_store)**2
+        h,ih=np.unique(hidx//scale,return_inverse=True)
+        for k in range(h.shape[0]):
+            spl=np.zeros([scale])
+            spl[hidx[ih==k]-scale*h[k]]=self.spline[ih==k]
+            self.spline_tree[h[k]]=spl
+
+    def SetParam(self,nside,spline,heal_idx):
+        
+        self.heal_idx=heal_idx
+        self.nside=nside
+        self.spline=spline
+        self.level=int(np.log2(nside))
+        self.nside_store=2**(int(self.level//2))
+        
+        self.spline_tree={}
+        
+        scale=(self.nside//self.nside_store)**2
+        h,ih=np.unique(heal_idx//scale,return_inverse=True)
+        for k in range(h.shape[0]):
+            spl=np.zeros([scale])
+            spl[heal_idx[ih==k]-scale*h[k]]=self.spline[ih==k]
+            self.spline_tree[h[k]]=spl
+        
+        
+    
+    def Transform(self,th,ph,threshold=1E-2,nest=True):
+        
+        www,all_idx,hidx=self.ang2weigths(th,ph,threshold=threshold,nest=nest)
+
+        x=np.zeros([hidx.shape[0]])
+        scale=(self.nside//self.nside_store)**2
+        h,ih=np.unique(hidx//scale,return_inverse=True)
+        for k in range(h.shape[0]):
+            if h[k] in self.spline_tree:
+                spl=self.spline_tree[h[k]]
+                x[ih==k]=spl[hidx[ih==k]-scale*h[k]]
+        data=self.P(x,www,all_idx)
+        return data
+
+    def FitTransform(self,X,th,ph,threshold=1E-2,nest=True):
+
+        self.Fit(X,th,ph)
+        
+        t,p=hp.pix2ang(self.nside,self.heal_idx,nest=True)
+
+        return self.Transform(t,p)

foscat/heal_NN.py

@@ -7,18 +7,19 @@ import foscat.scat_cov as sc
 class CNN:
 
     def __init__(
-        self,
-        nparam=1,
-        KERNELSZ=3,
-        NORIENT=4,
-        chanlist=[],
-        in_nside=1,
-        n_chan_in=1,
-        SEED=1234,
-        all_type='float32',
-        filename=None,
-        scat_operator=None,
-        BACKEND='tensorflow'
+            self,
+            nparam=1,
+            KERNELSZ=3,
+            NORIENT=4,
+            chanlist=[],
+            in_nside=1,
+            n_chan_in=1,
+            SEED=1234,
+            add_undersample_data=False,
+            all_type='float32',
+            filename=None,
+            scat_operator=None,
+            BACKEND='tensorflow'
     ):
 
         if filename is not None:
@@ -38,6 +39,7 @@ class CNN:
             self.x = self.scat_operator.backend.bk_cast(outlist[6])
             self.out_nside = self.in_nside // (2**(self.nscale+1))
         else:
+            self.add_undersample_data=add_undersample_data
             self.nscale = len(chanlist)-1
             self.npar = nparam
             self.n_chan_in = n_chan_in
@@ -89,9 +91,9 @@ class CNN:
     def get_number_of_weights(self):
         totnchan = 0
         for i in range(self.nscale):
-            totnchan = totnchan + self.chanlist[i] * self.chanlist[i + 1]
+            totnchan = totnchan + (self.chanlist[i]+int(self.add_undersample_data)* self.n_chan_in) * self.chanlist[i + 1]
         return (
-            self.npar * 12 * self.out_nside**2 * self.chanlist[self.nscale]*self.NORIENT
+            self.npar * 12 * self.out_nside**2 * (self.chanlist[self.nscale]+int(self.add_undersample_data)* self.n_chan_in)*self.NORIENT
             + totnchan * self.KERNELSZ * (self.KERNELSZ//2+1)*self.NORIENT*self.NORIENT
             + self.KERNELSZ * (self.KERNELSZ//2+1) * self.n_chan_in * self.chanlist[0]*self.NORIENT
         )
@@ -190,7 +192,7 @@ class CNN:
         m=np.mean(m,0)
         return self.scat_operator.backend.bk_cast(m[None,None,...])
         
-    def eval(self, im, 
+    def eval(self, in_im, 
             indices=None, 
             weights=None, 
             out_map=False, 
@@ -204,14 +206,12 @@ class CNN:
         )
         nn = self.KERNELSZ * (self.KERNELSZ//2+1) * self.n_chan_in * self.chanlist[0]*self.NORIENT
 
-        im = self.scat_operator.healpix_layer(im[:,:,None,:], ww)
+        im = self.scat_operator.healpix_layer(in_im[:,:,None,:], ww)
         
         if first_layer_rot is not None:
             im = self.backend.bk_reshape(im,[im.shape[0],im.shape[1],self.NORIENT,1,im.shape[3]])
             im = self.backend.bk_reduce_sum(im*first_layer_rot,2)
             
-        if out_map:
-            return im
         
         if activation=='relu':
             im = self.backend.bk_relu(im)
@@ -220,6 +220,16 @@ class CNN:
         
         im = self.backend.bk_reduce_sum(self.backend.bk_reshape(im,[im.shape[0],im.shape[1],self.NORIENT,im.shape[3]//4,4]),4)
 
+        if self.add_undersample_data:
+            l_im=self.backend.bk_repeat(self.backend.bk_reshape(in_im,[in_im.shape[0],in_im.shape[1],1,in_im.shape[2]]),self.NORIENT,2)
+            l_im=self.backend.bk_reduce_sum(
+                                           self.backend.bk_reshape(l_im,[in_im.shape[0],in_im.shape[1],self.NORIENT,l_im.shape[3]//4,4]), 4)
+            im=self.backend.bk_concat([im,l_im],1)
+                
+            
+        if out_map:
+            return im
+        
         for k in range(self.nscale):
             ww = self.scat_operator.backend.bk_reshape(
                 x[
@@ -227,17 +237,21 @@ class CNN:
                     + self.KERNELSZ
                     * (self.KERNELSZ//2+1)
                     * self.NORIENT*self.NORIENT
-                    * self.chanlist[k]
+                    * (self.chanlist[k]+int(self.add_undersample_data)* self.n_chan_in)
                     * self.chanlist[k + 1]
                 ],
-                [self.chanlist[k], self.NORIENT, self.KERNELSZ * (self.KERNELSZ//2+1),  self.chanlist[k + 1], self.NORIENT],
+                [self.chanlist[k]+int(self.add_undersample_data)* self.n_chan_in,
+                 self.NORIENT,
+                 self.KERNELSZ * (self.KERNELSZ//2+1),
+                 self.chanlist[k + 1],
+                 self.NORIENT],
             )
             nn = (
                 nn
                 + self.KERNELSZ
                 * (self.KERNELSZ//2+1)
                 * self.NORIENT*self.NORIENT
-                * self.chanlist[k]
+                * (self.chanlist[k]+int(self.add_undersample_data)* self.n_chan_in)
                 * self.chanlist[k + 1]
             )
             if indices is None:
@@ -253,12 +267,17 @@ class CNN:
                 im = self.backend.bk_abs(im)
             im = self.backend.bk_reduce_sum(self.backend.bk_reshape(im,[im.shape[0],im.shape[1],self.NORIENT,im.shape[3]//4,4]),4)
 
+            if self.add_undersample_data:
+                l_im=self.backend.bk_reduce_sum(
+                    self.backend.bk_reshape(l_im,[l_im.shape[0],l_im.shape[1],self.NORIENT,l_im.shape[3]//4,4]), 4)
+                im=self.backend.bk_concat([im,l_im],1)
+
         ww = self.scat_operator.backend.bk_reshape(
             x[
                 nn : nn
-                + self.npar * 12 * self.out_nside**2 * self.chanlist[self.nscale]*self.NORIENT
+                + self.npar * 12 * self.out_nside**2 * (self.chanlist[self.nscale]+int(self.add_undersample_data)* self.n_chan_in)*self.NORIENT
             ],
-            [12 * self.out_nside**2 * self.chanlist[self.nscale]*self.NORIENT, self.npar],
+            [12 * self.out_nside**2 * (self.chanlist[self.nscale]+int(self.add_undersample_data)* self.n_chan_in)*self.NORIENT, self.npar],
         )
 
         im = self.scat_operator.backend.bk_matmul(
@@ -429,4 +448,4 @@ class GCNN:
         im = self.backend.bk_reduce_mean(im[:,:,:,None]*ww,[1,2])
         #im = self.scat_operator.backend.bk_relu(im)
          
-        return im
+        return im