foscat 2025.9.5__py3-none-any.whl → 2025.11.1__py3-none-any.whl

foscat/FoCUS.py

@@ -6,7 +6,7 @@ import numpy as np
 import foscat.HealSpline as HS
 from scipy.interpolate import griddata
 
-TMPFILE_VERSION = "V9_0"
+TMPFILE_VERSION = "V10_0"
 
 
 class FoCUS:
@@ -28,15 +28,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.09.5"
+        self.__version__ = "2025.11.1"
         # P00 coeff for normalization for scat_cov
         self.TMPFILE_VERSION = TMPFILE_VERSION
         self.P1_dic = None
@@ -50,6 +50,7 @@ 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 = {}
@@ -89,13 +90,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
 
@@ -673,13 +667,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 +712,26 @@ 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 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
@@ -1488,7 +1502,7 @@ class FoCUS:
                     if l_kernel == 5:
                         pw = 0.5
                         pw2 = 0.5
-                        threshold = 2e-4
+                        threshold = 2e-5
 
                     elif l_kernel == 3:
                         pw = 1.0 / np.sqrt(2)
@@ -1498,7 +1512,7 @@ class FoCUS:
                     elif l_kernel == 7:
                         pw = 0.5
                         pw2 = 0.25
-                        threshold = 4e-5
+                        threshold = 2e-5
                         
                     import foscat.SphericalStencil as hs
                     import torch
@@ -1517,14 +1531,19 @@ class FoCUS:
                                                 n_gauges=self.NORIENT,
                                                 gauge_type='cosmo')
                     
-                    xx=np.tile(np.arange(self.KERNELSZ)-self.KERNELSZ//2,self.KERNELSZ).reshape(self.KERNELSZ*self.KERNELSZ)
+                    xx=np.tile(np.arange(self.KERNELSZ)-self.KERNELSZ//2,self.KERNELSZ).reshape(self.KERNELSZ,self.KERNELSZ)
                     
-                    wwr=hconvol.to_tensor((np.exp(-pw2*(xx**2+(xx.T)**2))*np.cos(pw*xx*np.pi)).reshape(1,1,self.KERNELSZ*self.KERNELSZ))
+                    wwr=(np.exp(-pw2*(xx**2+(xx.T)**2))*np.cos(pw*xx*np.pi)).reshape(1,1,self.KERNELSZ*self.KERNELSZ)
                     wwr-=wwr.mean()
-                    wwi=hconvol.to_tensor((np.exp(-pw2*(xx**2+(xx.T)**2))*np.sin(pw*xx*np.pi)).reshape(1,1,self.KERNELSZ*self.KERNELSZ))
+                    wwi=(np.exp(-pw2*(xx**2+(xx.T)**2))*np.sin(pw*xx*np.pi)).reshape(1,1,self.KERNELSZ*self.KERNELSZ)
                     wwi-=wwi.mean()
-                    wwr/=(abs(wwr+1J*wwi)).sum()
-                    wwi/=(abs(wwr+1J*wwi)).sum()
+                    amp=np.sum(abs(wwr+1J*wwi))
+                    
+                    wwr/=amp
+                    wwi/=amp
+                    
+                    wwr=hconvol.to_tensor(wwr)
+                    wwi=hconvol.to_tensor(wwi)
                     
                     wavr,indice,mshape=hconvol.make_matrix(wwr)
                     wavi,indice,mshape=hconvol.make_matrix(wwi)
@@ -2180,17 +2199,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:
@@ -2199,22 +2222,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)
@@ -2226,11 +2253,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:
@@ -2239,35 +2271,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:
@@ -2276,24 +2315,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)

foscat/Plot.py

@@ -959,8 +959,9 @@ def conjugate_gradient_normal_equation(data, x0, www, all_idx,
 
         rs_new = np.dot(r, r)
 
-        if verbose and i % 50 == 0:
-            print(f"Iter {i:03d}: residual = {np.sqrt(rs_new):.3e}")
+        if verbose and i % 10 == 0:
+            v=np.mean((LP(p, www, all_idx)-data)**2)
+            print(f"Iter {i:03d}: residual = {np.sqrt(rs_new):.3e},{np.sqrt(v):.3e}")
 
         if np.sqrt(rs_new) < tol:
             if verbose:
@@ -1155,7 +1156,7 @@ def plot_wave(wave,title="spectrum",unit="Amplitude",cmap="viridis"):
     plt.xlabel(r"$k_x$ [cycles / km]")
     plt.ylabel(r"$k_y$ [cycles / km]")
     plt.title(title)
-
+    
 def lonlat_edges_from_ref(shape, ref_lon, ref_lat, dlon, dlat, anchor="center"):
     """
     Build lon/lat *edges* (H+1, W+1) for a regular, axis-aligned grid.

foscat/healpix_unet_torch.py

@@ -982,6 +982,9 @@ def fit(
         n_epoch: int = 10,
         view_epoch: int = 10,
         batch_size: int = 16,
+        x_valid: Union[torch.Tensor, np.ndarray, List[Union[torch.Tensor, np.ndarray]]]= None,
+        y_valid: Union[torch.Tensor, np.ndarray, List[Union[torch.Tensor, np.ndarray]]]= None,
+        save_model: bool = False,
         lr: float = 1e-3,
         weight_decay: float = 0.0,
         clip_grad_norm: Optional[float] = None,
@@ -1005,6 +1008,11 @@ def fit(
     device = model.runtime_device if hasattr(model, "runtime_device") else (torch.device("cuda") if torch.cuda.is_available() else torch.device("cpu"))
     model.to(device)
 
+    if save_model:
+        assert x_valid is None, "If save_mode=True x_valid should not be None"
+        assert y_valid is None, "If save_mode=True y_valid should not be None"
+        best_valid=1E30
+
     # Detect variable-length mode
     varlen_mode = isinstance(x_train, (list, tuple))
 
@@ -1197,6 +1205,14 @@ def fit(
             history.append(epoch_loss)
             # print every view_epoch logical step
             if verbose and ((len(history) % view_epoch == 0) or (len(history) == 1)):
-                print(f"[epoch {len(history)}] loss={epoch_loss:.6f}")
+                if x_valid is not None:
+                    preds=model.predict(model.to_tensor(x_valid)).cpu().numpy()
+                    valid_loss=np.mean((preds-y_valid)**2)
+                    if save_model:
+                        if best_valid>valid_loss:
+                            torch.save({"model": self.state_dict(), "cfg": CFG}, os.path.join(CFG["save_dir"], "best.pt"))
+                    print(f"[epoch {len(history)}] loss={epoch_loss:.4f} loss_valid={valid_loss:.4f}")
+                else:
+                    print(f"[epoch {len(history)}] loss={epoch_loss:.4f}")
 
     return {"loss": history}