squarenet 1.3.0__tar.gz → 1.3.2__tar.gz

{squarenet-1.3.0 → squarenet-1.3.2}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: squarenet
-Version: 1.3.0
+Version: 1.3.2
 Summary: Gridification for point clouds in any dimension (fast greedy optimal transport through Cartesian Sort).
 Author-email: ArmanddeCacqueray <armanddecacqueray@sfr.fr>
 Project-URL: Homepage, https://github.com/ArmanddeCacqueray/SquareNet

{squarenet-1.3.0 → squarenet-1.3.2}/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "squarenet"
-version = "1.3.0"
+version = "1.3.2"
 description = "Gridification for point clouds in any dimension (fast greedy optimal transport through Cartesian Sort)."
 authors = [{ name = "ArmanddeCacqueray", email = "armanddecacqueray@sfr.fr" }]
 readme = "README.md"

{squarenet-1.3.0 → squarenet-1.3.2}/src/squarenet/__init__.py

@@ -1,4 +1,4 @@
 from .squarenet import SquareNet 
 
 __all__ = ["SquareNet", "samplepoints", "place_at", "plotpoints", "sqplot", "carthesian_sort"]
-__version__ = "1.3.0"
+__version__ = "1.3.2"

{squarenet-1.3.0 → squarenet-1.3.2}/src/squarenet/neighbormap.py

@@ -14,14 +14,19 @@ def neighbormap(grided_points, mapidx, criterion = "rank", thresholdcut = 1,
     d = len(gpts.shape) - 1
     wr = windowradius
     ws = 2*wr+1
+    N = np.prod(gpts.shape[:-1])
+    D = gpts.shape[-1]
 
     selection = None
 
-    if np.prod(gpts.shape)*(ws**2) >= sample_size:
-        N = np.prod(gpts.shape[:-1])
+    if N*D*(ws**2) >= sample_size:
         sample_size_x = min(sample_size//(ws**2), N)
         selection = np.random.choice(N, replace = False, size = sample_size_x)
         selection = mapidx(selection)
+    else:
+        selection = mapidx(np.arange(N))
+    mask = np.isfinite(grided_points[selection]).all(axis=-1)
+    selection = tuple(dim_coords[mask] for dim_coords in selection)
         
     pad_width = [(0, 0)]*(d+1)
     for gax in projection:
@@ -39,9 +44,6 @@ def neighbormap(grided_points, mapidx, criterion = "rank", thresholdcut = 1,
         axis=projection
     )
 
-
-
-
     dists = kernel(gpts[selection], gview[selection])
 
     if criterion =="value":

{squarenet-1.3.0 → squarenet-1.3.2}/src/squarenet/squarenet.py

@@ -58,7 +58,8 @@ class SquareNet:
         self.max_iter = max_iter
         self.verbose = verbose
         self.backend = backend
-        self.torchdevice = "cpu"
+        self.device = "cpu"
+        self._torchdevice = "cpu"
 
         if backend not in {"numpy", "jax", "torch"}:
             raise ValueError(f"Unknown backend '{backend}' Should be 'numpy' 'torch' or 'jax'")
@@ -67,14 +68,29 @@ class SquareNet:
             self.xp = np
 
         elif backend == "jax":
+            import jax
             import jax.numpy as jnp
+            import shutil
             self.xp = jnp
+            self.device = jax.default_backend()
+            if shutil.which("nvidia-smi") is not None:
+                try:
+                    import torch
+                    if torch.cuda.is_available():
+                        self._torchdevice = torch.device("cuda")
+                except ImportError:
+                    if warnings_ == True:
+                        warn(
+                            "SquareNet tryed to import pytorch because a GPU was detected."\
+                            "Pytorch was not found so SquareNet will default to cpu for the fit"
+                        )
+
 
         elif backend == "torch":
             import torch
             self.xp = torch
             if torch.cuda.is_available():
-                self.torchdevice = torch.device("cuda")
+                self.device = torch.device("cuda")
 
         self.warnings_ = warnings_
         if stencil is not None:
@@ -88,7 +104,7 @@ class SquareNet:
     # Internal utilities
     # ------------------------------------------------------------------
     def _to_backend(self, x):
-        return to_backend(x, backend = self.backend, device = self.torchdevice, warnings_ = self.warnings_)
+        return to_backend(x, backend = self.backend, device = self.device, warnings_ = self.warnings_)
     
     def _reset_state(self):
         """Reset internal state to initial configuration."""
@@ -96,7 +112,7 @@ class SquareNet:
            self.grid = self.xp.arange(
                 self.N,
                 dtype=self.xp.int32,
-                device=self.torchdevice,
+                device=self.device,
             ).reshape(self.gridshape)
 
         else:
@@ -118,6 +134,7 @@ class SquareNet:
         if isinstance(points, str):
             points = samplepoints(method=points, size=(self.N, self.D))
 
+        points = points.clone() if hasattr(points, "clone") else points.copy()
         points = self._to_backend(points)
         N, D = points.shape
 
@@ -133,7 +150,7 @@ class SquareNet:
 
         return points
     
-    def _make_stencil(self, n_target, dtype=float, max_iter=1000):
+    def _make_stencil(self, n_target, dtype=float, max_iter=100):
         """Create a grid stencil based on a p-holder norm boundary. see utils.make_stencil
         stencil allow to fit n_target points in a grid with shape bigger than n_target"""
         stencil = make_stencil(self.gridshape, n_target, dtype, max_iter)
@@ -164,6 +181,7 @@ class SquareNet:
         Padding is performed such that points in the middle of the grid are true points
         and points in the corner are fictive points with +- inf coordinates, using the stencil.
         """
+        n_target = len(points)
         if self.stencil is None:
             if self.verbose + verbose >= 2:
                 print(
@@ -174,8 +192,14 @@ class SquareNet:
                     "    >>> saved_stencil = sn.stencil\n"
                     "    >>> sn = SquareNet(..., stencil=saved_stencil)\n"
                 )
-            self._make_stencil(n_target=len(points), dtype=points.dtype)
-        
+            self._make_stencil(n_target=n_target, dtype=points.dtype)
+        elif self.xp.isfinite(self.stencil).sum() != self.D*n_target:
+            if self.verbose + verbose >= 2:
+                print(
+                    "\n[SquareNet] Update the stencil to match new number of points...\n"
+                )
+            self._make_stencil(n_target=n_target, dtype=points.dtype)
+
         return fill_in(points, self.stencil, self.xp)
     # ------------------------------------------------------------------
     # Core API
@@ -204,13 +228,24 @@ class SquareNet:
         see ``squarenet.core``
         """
         old_backend = self.backend
-        if self.backend == "torch" and self.torchdevice != "cpu":
-            fit_with_numpy = False 
+        old_device = self.device
+
+        if self._torchdevice != "cpu":
+            #just for the fit, will be updated after
+            self.backend = "torch"
+            self.device = self._torchdevice
         else: 
-            fit_with_numpy = True
             #just for the fit, will be updated after
             self.backend = "numpy"
-
+            self.device = "cpu"
+
+        true_g = np.array(self.gridshape)
+        true_g = true_g[true_g > 1]
+        D = len(true_g)
+        if D >= 5:
+            #method "fast" is as good as the other methods in 5D and more,
+            #robust and ultimate are too slow anyway in those cases
+            method = "fast"
         max_iter = self.max_iter
         verbose = self.verbose
         def _log(msg):
@@ -222,7 +257,7 @@ class SquareNet:
                 print(f"\n=== {title} ===")
 
         points = self._validate_points(points)
-        self.points = points.clone() if hasattr(points, "clone") else points.copy()
+        grid = self._to_backend(self.grid)
 
         _log("Starting gridification... available method [fast, robust, ultimate]")
         _log(f"selected {method}")
@@ -236,8 +271,8 @@ class SquareNet:
         if method == "ultimate":
             _log(f"(max iter: 2 x {mi} + 4 x {mi} + {mi})")
 
-        self.grid, lc = carthesian_sort(self.grid, points, max_iter=max_iter, method = method,
-                                        backend = self.backend, verbose = self.verbose)
+        grid, lc = carthesian_sort(grid, points, max_iter=max_iter, method = method,
+                                    backend = self.backend, verbose = self.verbose)
 
         # --------------------------------------------------
         # Finalization
@@ -267,10 +302,10 @@ class SquareNet:
                         stacklevel=2,
                     )
         
-        if fit_with_numpy:
-            self.backend = old_backend
-            self.grid = to_backend(self.grid, backend = self.backend, device = self.torchdevice)
-            self.points = to_backend(self.points, backend = self.backend, device = self.torchdevice)
+        self.backend = old_backend
+        self.device = old_device
+        self.grid = self._to_backend(grid)
+        self.points = self._to_backend(points)
         # --------------------------------------------------
         # Update mappings
         # --------------------------------------------------

{squarenet-1.3.0 → squarenet-1.3.2}/src/squarenet/utils.py

@@ -8,7 +8,7 @@ def index_identity(shape):
     """index_identy[i, j, k, ...] = [i, j, k, ...]"""
     return np.moveaxis(np.indices(shape), 0, -1)
 
-def make_stencil(gridshape, n_target, dtype=float, max_iter = 1000):
+def make_stencil(gridshape, n_target, dtype=float, max_iter = 100):
     """
     Create a grid stencil based on a p-holder norm boundary.
     Uses binary search to find a smooth convexe subdomain of the hyperrectangular
@@ -52,7 +52,7 @@ def make_stencil(gridshape, n_target, dtype=float, max_iter = 1000):
 
     log_x = np.log(x_abs)
     
-    low_p, high_p = -1.0, 1000.0
+    low_p, high_p = 2.0, 1000.0
     
     for it in range(max_iter):  
         p = (low_p + high_p) / 2.0
@@ -216,50 +216,78 @@ def project(gridpoints, feature_axes=(0, 1), index=0):
 
 def from_backend(x):
     """
-    Convert torch/jax/numpy array to numpy safely.
+    Safely convert Torch, JAX, or array-like objects back to Numpy.
     """
+    # 1. Already a Numpy array
     if isinstance(x, np.ndarray):
         return x
 
-    # torch
+    # 2. PyTorch tensor 
     if hasattr(x, "detach"):
         return x.detach().cpu().numpy()
 
-    # jax
-    if hasattr(x, "__array__"):
-        return np.asarray(x)
-
+    # 3. JAX arrays or any object with __array__ protocol (Pandas, etc.)
     return np.asarray(x)
 
-def to_backend(x, backend="numpy", device = "cpu", warnings_ = True):
+def to_backend(x, backend="numpy", device="cpu", warnings_=True):
     """
-    Convert numpy array to target backend.
+    Convert arrays/tensors between Numpy, Torch, and JAX.
+    Utilizes DLPack for zero-copy transfers across frameworks when possible.
     """
+    # Identify input framework without heavy imports
+    in_type = type(x).__module__.split('.')[0]
+
+    # ---------------------------------------------------------
+    # 1. Target: NUMPY
+    # ---------------------------------------------------------
     if backend == "numpy":
-        return np.asarray(x)
+        return from_backend(x)
 
-    if backend == "torch":
+    # ---------------------------------------------------------
+    # 2. Target: TORCH
+    # ---------------------------------------------------------
+    elif backend == "torch":
         import torch
-        device = torch.device(device)
-        x_device = getattr(x, "device", None)
-        if (
-            warnings_
-            and x_device is not None
-            and hasattr(x_device, "type")
-            and x_device.type == "cpu"
-            and device.type != "cpu"
-        ):
-            warn(
-                "Downgrading tensor as initial device was GPU but asked device is CPU"
-            )
-
-        return torch.as_tensor(x, device=device)
-
-    if backend == "jax":
+        tgt_dev = torch.device(device)
+        
+        if in_type == "torch":
+            if warnings_ and x.device.type != "cpu" and tgt_dev.type == "cpu":
+                warn("Downgrading tensor: initial device was GPU but asked for CPU.")
+            return x.to(tgt_dev)
+            
+        elif in_type in ["jax", "jaxlib"]:
+            # Zero-copy JAX to Torch via DLPack
+            return torch.from_dlpack(x).to(tgt_dev)
+            
+        return torch.as_tensor(x, device=tgt_dev)
+
+    # ---------------------------------------------------------
+    # 3. Target: JAX
+    # ---------------------------------------------------------
+    elif backend == "jax":
+        import jax
         import jax.numpy as jnp
-        return jnp.asarray(x)
-
-    raise ValueError(f"Unknown backend: {backend}")
+        
+        # JAX uses "gpu" instead of "cuda"
+        device = "cpu" if "cpu" in str(device) else "gpu"
+        try:
+            tgt_dev = jax.devices(device)[0]
+        except RuntimeError:
+            tgt_dev = None
+
+        if in_type == "torch":
+            import torch.utils.dlpack
+            # DLPack requires contiguous memory
+            x = x if x.is_contiguous() else x.contiguous()
+            j_arr = jax.dlpack.from_dlpack(torch.utils.dlpack.to_dlpack(x))
+            return jax.device_put(j_arr, tgt_dev) if tgt_dev else j_arr
+            
+        elif in_type in ["jax", "jaxlib"]:
+            return jax.device_put(x, tgt_dev) if tgt_dev else x
+            
+        return jax.device_put(x, tgt_dev) if tgt_dev else jnp.asarray(x)
+
+    raise ValueError(f"Unknown backend: '{backend}'. Expected 'numpy', 'torch', or 'jax'.")
 
 def progress_bar(it, total, bar_length=30):
     progress = it / total

{squarenet-1.3.0 → squarenet-1.3.2}/src/squarenet.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: squarenet
-Version: 1.3.0
+Version: 1.3.2
 Summary: Gridification for point clouds in any dimension (fast greedy optimal transport through Cartesian Sort).
 Author-email: ArmanddeCacqueray <armanddecacqueray@sfr.fr>
 Project-URL: Homepage, https://github.com/ArmanddeCacqueray/SquareNet