foscat 2025.10.2__py3-none-any.whl → 2026.1.1__py3-none-any.whl

foscat/FoCUS.py

@@ -5,8 +5,11 @@ import healpy as hp
 import numpy as np
 import foscat.HealSpline as HS
 from scipy.interpolate import griddata
+from foscat.SphereDownGeo import SphereDownGeo
+from foscat.SphereUpGeo import SphereUpGeo
+import torch
 
-TMPFILE_VERSION = "V10_0"
+TMPFILE_VERSION = "V12_0"
 
 
 class FoCUS:
@@ -28,15 +31,15 @@ class FoCUS:
             use_2D=False,
             use_1D=False,
             return_data=False,
-            JmaxDelta=0,
             DODIV=False,
+            use_median=False,
             InitWave=None,
             silent=True,
             mpi_size=1,
             mpi_rank=0
     ):
 
-        self.__version__ = "2025.10.2"
+        self.__version__ = "2026.01.1"
         # P00 coeff for normalization for scat_cov
         self.TMPFILE_VERSION = TMPFILE_VERSION
         self.P1_dic = None
@@ -50,13 +53,15 @@ class FoCUS:
         self.mpi_rank = mpi_rank
         self.return_data = return_data
         self.silent = silent
+        self.use_median = use_median
 
         self.kernel_smooth = {}
         self.padding_smooth = {}
         self.kernelR_conv = {}
         self.kernelI_conv = {}
         self.padding_conv = {}
-
+        self.down = {}
+        self.up = {}
         if not self.silent:
             print("================================================")
             print("          START FOSCAT CONFIGURATION")
@@ -89,13 +94,6 @@ class FoCUS:
         self.nlog = 0
         self.padding = padding
 
-        if JmaxDelta != 0:
-            print(
-                "OPTION JmaxDelta is not avialable anymore after version 3.6.2. Please use Jmax option in eval function"
-            )
-            return None
-
-        self.OSTEP = JmaxDelta
         self.use_2D = use_2D
         self.use_1D = use_1D
 
@@ -654,6 +652,7 @@ class FoCUS:
         return rim
 
     # --------------------------------------------------------
+    
     def ud_grade_2(self, im, axis=0, cell_ids=None, nside=None,max_poll=False):
 
         if self.use_2D:
@@ -673,13 +672,20 @@ class FoCUS:
             tim = self.backend.bk_reshape(
                 self.backend.bk_cast(im), [ndata, npix, npiy, 1]
             )
+            '''
             tim = self.backend.bk_reshape(
                 tim[:, 0 : 2 * (npix // 2), 0 : 2 * (npiy // 2), :],
                 [ndata, npix // 2, 2, npiy // 2, 2, 1],
             )
-
-            res = self.backend.bk_reduce_sum(self.backend.bk_reduce_sum(tim, 4), 2) / 4
-
+            
+            #res = self.backend.bk_reduce_sum(self.backend.bk_reduce_sum(tim, 4), 2) / 4
+            '''
+            
+            if self.use_median:
+                res = self.backend.downsample_median_2x2(tim)
+            else:
+                res = self.backend.downsample_mean_2x2(tim)
+    
             if len(ishape) == 2:
                 return (
                     self.backend.bk_reshape(
@@ -711,13 +717,42 @@ class FoCUS:
                 self.backend.bk_cast(im), [ndata, npix // 2, 2]
             )
 
-            res = self.backend.bk_reduce_mean(tim, -1)
+            if self.use_median:
+                res=self.backend.bk_reduce_median(tim,axis=-1)
+            else:
+                res=self.backend.bk_reduce_mean(tim,axis=-1)
 
             return self.backend.bk_reshape(res, ishape[0:-1] + [npix // 2]), None
 
         else:
             shape = list(im.shape)
-            if max_poll:
+            if nside is None:
+                l_nside=int(np.sqrt(shape[-1]//12))
+            else:
+                l_nside=nside
+            
+            nbatch=1
+            for k in range(len(shape)-1):
+                nbatch*=shape[k]
+            if l_nside not in self.down:
+                print('initialise down', l_nside)
+                self.down[l_nside] = SphereDownGeo(nside_in=l_nside, dtype=self.all_bk_type,mode="smooth", in_cell_ids=cell_ids)
+            
+            res,out_cell=self.down[l_nside](self.backend.bk_reshape(im,[nbatch,1,shape[-1]]))
+            
+            return self.backend.bk_reshape(res,shape[:-1]+[out_cell.shape[0]]),out_cell
+            '''
+            if self.use_median:
+                if cell_ids is not None:
+                    sim, new_cell_ids = self.backend.binned_mean(im, cell_ids,reduce='median')
+                    return sim, new_cell_ids
+            
+                return self.backend.bk_reduce_median(
+                    self.backend.bk_reshape(im, shape[0:-1]+[shape[-1]//4,4]), axis=-1
+                ),None
+            
+            elif max_poll:
+            
                 if cell_ids is not None:
                     sim, new_cell_ids = self.backend.binned_mean(im, cell_ids,reduce='max')
                     return sim, new_cell_ids
@@ -733,6 +768,7 @@ class FoCUS:
                 return self.backend.bk_reduce_mean(
                     self.backend.bk_reshape(im, shape[0:-1]+[shape[-1]//4,4]), axis=-1
                 ),None
+            '''
 
     # --------------------------------------------------------
     def up_grade(self, im, nout,
@@ -822,6 +858,7 @@ class FoCUS:
             else:
                 lout = nside
                 
+            '''
             if (lout,nout) not in self.pix_interp_val or force_init_index:
                 if not self.silent:
                     print("compute lout nout", lout, nout)
@@ -912,12 +949,32 @@ class FoCUS:
                     
                     del w
                     del p
-
-            if lout == nout:
-                imout = im
-            else:
-                # work only on the last column
-
+            '''
+            shape=list(im.shape)
+            nbatch=1
+            for k in range(len(shape)-1):
+                nbatch*=shape[k]
+            
+            im=self.backend.bk_reshape(im,[nbatch,1,shape[-1]])
+            
+            while lout<nout:
+                if lout not in self.up:
+                    if o_cell_ids is None:
+                        l_o_cell_ids=torch.tensor(np.arange(12*(lout**2),dtype='int'),device=im.device)
+                    else:
+                        l_o_cell_ids=o_cell_ids
+                    self.up[lout] = SphereUpGeo(nside_out=lout, 
+                                                dtype=self.all_bk_type,
+                                                 cell_ids_out=l_o_cell_ids,
+                                                 up_norm="col_l1")
+                im, fine_ids = self.up[lout](self.backend.bk_cast(im))
+                lout*=2
+                if lout<nout and o_cell_ids is not None:
+                    o_cell_ids=torch.repeat(fine_ids,4)*4+ \
+                                torch.tile(torch.tensor([0,1,2,3],device=fine_ids.device,dtype=fine_ids.dtype),fine_ids.shape[0])
+            
+            return self.backend.bk_reshape(im,shape[:-1]+[im.shape[-1]])
+            '''
                 ndata = 1
                 for k in range(len(ishape)-1):
                     ndata = ndata * ishape[k]
@@ -946,6 +1003,7 @@ class FoCUS:
                     return self.backend.bk_reshape(
                         imout, ishape[0:-1]+[imout.shape[-1]]
                     )
+            '''
         return imout
 
     # --------------------------------------------------------
@@ -1340,7 +1398,9 @@ class FoCUS:
                     else:
                         l_cell_ids=cell_ids
                         
-                    nvalid=self.KERNELSZ**2
+                    nvalid=4*self.KERNELSZ**2
+                    if nvalid>12*nside**2:
+                        nvalid=12*nside**2
                     idxEB=hconvol.idx_nn[:,0:nvalid]
                     tmpEB=np.zeros([self.NORIENT,4,l_cell_ids.shape[0],nvalid],dtype='complex')       
                     tmpS=np.zeros([4,l_cell_ids.shape[0],nvalid],dtype='float')           
@@ -1486,7 +1546,7 @@ class FoCUS:
                     
                 else:
                     if l_kernel == 5:
-                        pw = 0.5
+                        pw = 0.75
                         pw2 = 0.5
                         threshold = 2e-5
 
@@ -2185,17 +2245,21 @@ class FoCUS:
             #    mtmp = l_mask[:,self.KERNELSZ // 2 : -self.KERNELSZ // 2,self.KERNELSZ // 2 : -self.KERNELSZ // 2,:]
             #    vtmp = l_x[:,self.KERNELSZ // 2 : -self.KERNELSZ // 2,self.KERNELSZ // 2 : -self.KERNELSZ // 2,:]
 
-            v1 = self.backend.bk_reduce_sum(
-                self.backend.bk_reduce_sum(mtmp * vtmp, axis=-1), -1
-            )
-            v2 = self.backend.bk_reduce_sum(
-                self.backend.bk_reduce_sum(mtmp * vtmp * vtmp, axis=-1), -1
-            )
-            vh = self.backend.bk_reduce_sum(
-                self.backend.bk_reduce_sum(mtmp, axis=-1), -1
-            )
+            if self.use_median:
+                res,res2  = self.backend.bk_masked_median_2d_weiszfeld(vtmp, mtmp)
+            else:
+                v1 = self.backend.bk_reduce_sum(
+                    self.backend.bk_reduce_sum(mtmp * vtmp, axis=-1), -1
+                )
+                v2 = self.backend.bk_reduce_sum(
+                    self.backend.bk_reduce_sum(mtmp * vtmp * vtmp, axis=-1), -1
+                )
+                vh = self.backend.bk_reduce_sum(
+                    self.backend.bk_reduce_sum(mtmp, axis=-1), -1
+                )
 
-            res = v1 / vh
+                res = v1 / vh
+                res2= v2 / vh
 
             oshape = [x.shape[0]] + [mask.shape[0]]
             if len(x.shape) > 3:
@@ -2204,22 +2268,26 @@ class FoCUS:
                 oshape = oshape + [1]
                 
             if calc_var:
+                if self.use_median:
+                    vh = self.backend.bk_reduce_sum(
+                        self.backend.bk_reduce_sum(mtmp, axis=-1), -1
+                    )
                 if self.backend.bk_is_complex(vtmp):
                     res2 = self.backend.bk_sqrt(
                         (
                             (
-                                self.backend.bk_real(v2) / self.backend.bk_real(vh)
+                                self.backend.bk_real(res2)
                                 - self.backend.bk_real(res) * self.backend.bk_real(res)
                             )
                             + (
-                                self.backend.bk_imag(v2) / self.backend.bk_real(vh)
+                                self.backend.bk_imag(res2)
                                 - self.backend.bk_imag(res) * self.backend.bk_imag(res)
                             )
                         )
                         / self.backend.bk_real(vh)
                     )
                 else:
-                    res2 = self.backend.bk_sqrt((v2 / vh - res * res) / (vh))
+                    res2 = self.backend.bk_sqrt((res2 - res * res) / (vh))
 
                 res = self.backend.bk_reshape(res, oshape)
                 res2 = self.backend.bk_reshape(res2, oshape)
@@ -2231,11 +2299,16 @@ class FoCUS:
         elif self.use_1D:
             mtmp = l_mask
             vtmp = l_x
-            v1 = self.backend.bk_reduce_sum(l_mask * vtmp, axis=-1)
-            v2 = self.backend.bk_reduce_sum(l_mask * vtmp * vtmp, axis=-1)
-            vh = self.backend.bk_reduce_sum(l_mask , axis=-1)
+            
+            if self.use_median:
+                res,res2  = self.backend.bk_masked_median(l_x, l_mask)
+            else:
+                v1 = self.backend.bk_reduce_sum(l_mask * l_x, axis=-1)
+                v2 = self.backend.bk_reduce_sum(l_mask * l_x * l_x, axis=-1)
+                vh = self.backend.bk_reduce_sum(l_mask, axis=-1)
 
-            res = v1 / vh
+                res = v1 / vh
+                res2= v2 / vh
 
             oshape = [x.shape[0]] + [mask.shape[0]]
             if len(x.shape) > 1:
@@ -2244,35 +2317,42 @@ class FoCUS:
                 oshape = oshape + [1]
 
             if calc_var:
+                if self.use_median:
+                    vh = self.backend.bk_reduce_sum(l_mask, axis=-1)
+                    
                 if self.backend.bk_is_complex(vtmp):
                     res2 = self.backend.bk_sqrt(
                         (
                             (
-                                self.backend.bk_real(v2) / self.backend.bk_real(vh)
+                                self.backend.bk_real(res2)
                                 - self.backend.bk_real(res) * self.backend.bk_real(res)
                             )
                             + (
-                                self.backend.bk_imag(v2) / self.backend.bk_real(vh)
+                                self.backend.bk_imag(res2)
                                 - self.backend.bk_imag(res) * self.backend.bk_imag(res)
                             )
                         )
                         / self.backend.bk_real(vh)
                     )
                 else:
-                    res2 = self.backend.bk_sqrt((v2 / vh - res * res) / (vh))
+                    res2 = self.backend.bk_sqrt((res2 - res * res) / (vh))
+                    
                 res = self.backend.bk_reshape(res, oshape)
                 res2 = self.backend.bk_reshape(res2, oshape)
                 return res, res2
             else:
                 res = self.backend.bk_reshape(res, oshape)
                 return res
-
         else:
-            v1 = self.backend.bk_reduce_sum(l_mask * l_x, axis=-1)
-            v2 = self.backend.bk_reduce_sum(l_mask * l_x * l_x, axis=-1)
-            vh = self.backend.bk_reduce_sum(l_mask, axis=-1)
+            if self.use_median:
+                res,res2  = self.backend.bk_masked_median(l_x, l_mask)
+            else:
+                v1 = self.backend.bk_reduce_sum(l_mask * l_x, axis=-1)
+                v2 = self.backend.bk_reduce_sum(l_mask * l_x * l_x, axis=-1)
+                vh = self.backend.bk_reduce_sum(l_mask, axis=-1)
 
-            res = v1 / vh
+                res = v1 / vh
+                res2= v2 / vh
 
             oshape = []
             if len(shape) > 1:
@@ -2281,24 +2361,27 @@ class FoCUS:
                 oshape = [1]
                 
             oshape = oshape + [mask.shape[0]]
+
             if len(shape) > 2:
                 oshape = oshape + shape[1:-1]
             else:
                 oshape = oshape + [1]
 
             if calc_var:
+                if self.use_median:
+                    vh = self.backend.bk_reduce_sum(l_mask, axis=-1)
                 if self.backend.bk_is_complex(l_x):
                     res2 = self.backend.bk_sqrt(
                         (
-                            self.backend.bk_real(v2) / self.backend.bk_real(vh)
+                            self.backend.bk_real(res2)
                             - self.backend.bk_real(res) * self.backend.bk_real(res)
-                            + self.backend.bk_imag(v2) / self.backend.bk_real(vh)
+                            + self.backend.bk_imag(res2)
                             - self.backend.bk_imag(res) * self.backend.bk_imag(res)
                         )
                         / self.backend.bk_real(vh)
                     )
                 else:
-                    res2 = self.backend.bk_sqrt((v2 / vh - res * res) / (vh))
+                    res2 = self.backend.bk_sqrt((res2 - res * res) / (vh))
 
                 res = self.backend.bk_reshape(res, oshape)
                 res2 = self.backend.bk_reshape(res2, oshape)