foscat 2025.6.3__py3-none-any.whl → 2025.7.2__py3-none-any.whl

foscat/HealSpline.py

@@ -0,0 +1,216 @@
+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,
+            gamma=1,
+    ):
+        nside=2**level
+        self.nside_store=2**(level//2)
+        self.spline_tree={}
+        self.gamma=gamma
+        
+        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
+        if self.gamma!=1:
+            www=www**self.gamma
+        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)