Features/cuda

CHANGELOG.md

@@ -0,0 +1,27 @@
+# Changelog
+
+All notable changes to this project are documented in this file.
+
+This project follows [Keep a Changelog](https://keepachangelog.com/en/1.0.0/)
+and uses [Semantic Versioning](https://semver.org/).
+
+---
+
+## [Unreleased]
+
+### Added
+- MIT license/authorship headers to all `.py` files ([features/cuda])
+- CUDA feature branch for GPU support
+
+### Changed
+- Restructured branch strategy: `main`, `dev`, `features/cuda`
+- Clarified setup instructions in `README.md`
+
+---
+
+## [0.1.0] – 2025-05-14
+
+### Added
+- Initial codebase and Git setup
+- Created branches: `main`, `dev`, and `feature/cuda`
+- Basic Python project scaffold

spinstep/__init.py__

@@ -1,3 +1,7 @@
+# __init__.py — MIT License
+# Author: Eraldo B. Marques <eraldo.bernardo@gmail.com> — Created: 2025-05-14
+# See LICENSE.txt for full terms. This header must be retained in redistributions.
+
 from .node import Node
 from .traversal import QuaternionDepthIterator
 from .discrete import DiscreteOrientationSet

spinstep/demo.py

@@ -1,3 +1,7 @@
+# demo.py — MIT License
+# Author: Eraldo Marques <eraldo.bernardo@gmail.com> — Created: 2025-05-14
+# See LICENSE.txt for full terms. This header must be retained in redistributions.
+
 from spinstep.node import Node
 from spinstep.traversal import QuaternionDepthIterator
 from spinstep.quaternion_utils import quaternion_from_euler
@@ -7,8 +11,6 @@
 # .Applies a quaternion-based depth-first traversal.
 # .Only visits nodes that lie within a given angular threshold (like aiming a "cone" of rotation).
 
-    Only visits nodes that lie within a given angular threshold (like aiming a "cone" of rotation).
-
 def build_demo_tree():
     """Build a small tree with varied 3D orientations (yaw, pitch, roll)."""
     root = Node("root", [0, 0, 0, 1], [

spinstep/demo1_tree_traversal.py

@@ -1,3 +1,7 @@
+# demo1_tree_traversal.py — MIT License
+# Author: Eraldo Marques <eraldo.bernardo@gmail.com> — Created: 2025-05-14
+# See LICENSE.txt for full terms. This header must be retained in redistributions.
+
 import numpy as np
 from scipy.spatial.transform import Rotation as R
 

spinstep/demo2_full_depth_traversal.py

@@ -1,3 +1,7 @@
+# demo2_full_depth_traversal.py — MIT License
+# Author: Eraldo Marques <eraldo.bernardo@gmail.com> — Created: 2025-05-14
+# See LICENSE.txt for full terms. This header must be retained in redistributions.
+
 import numpy as np
 from scipy.spatial.transform import Rotation as R
 

spinstep/demo3_spatial_traversal.py

@@ -1,3 +1,7 @@
+# demo3_spatial_traversal.py — MIT License
+# Author: Eraldo Marques <eraldo.bernardo@gmail.com> — Created: 2025-05-14
+# See LICENSE.txt for full terms. This header must be retained in redistributions.
+
 import numpy as np
 from scipy.spatial.transform import Rotation as R
 

spinstep/demo4_discrete_traversal.py

@@ -1,3 +1,7 @@
+# demo4_discrete_traversal.py — MIT License
+# Author: Eraldo Marques <eraldo.bernardo@gmail.com> — Created: 2025-05-14
+# See LICENSE.txt for full terms. This header must be retained in redistributions.
+
 import numpy as np
 from scipy.spatial.transform import Rotation as R
 from spinstep.node import Node
@@ -25,4 +29,4 @@
 visited_nodes = []
 for node in iterator:
     visited_nodes.append(node.name)
-    print("Visited:", node.name)
+    print("Visited:", node.name)

spinstep/discrete.py

@@ -1,36 +1,60 @@
-import numpy as np
-from scipy.spatial.transform import Rotation as R
-from sklearn.neighbors import BallTree
+# discrete.py — MIT License
+# Author: Eraldo B. Marques <eraldo.bernardo@gmail.com> — Created: 2025-05-14
+# See LICENSE.txt for full terms. This header must be retained in redistributions.
+
+from spinstep.utils.array_backend import get_array_module
 
 class DiscreteOrientationSet:
-    def __init__(self, orientations):
-        arr = np.array(orientations)
+    def __init__(self, orientations, use_cuda=False):
+        xp = get_array_module(use_cuda)
+        arr = xp.array(orientations)
         if arr.ndim != 2 or arr.shape[1] != 4:
             raise ValueError("Each orientation must be a quaternion [x, y, z, w]")
-        norms = np.linalg.norm(arr, axis=1)
-        if np.any(norms < 1e-8):
+        norms = xp.linalg.norm(arr, axis=1)
+        if xp.any(norms < 1e-8):
             raise ValueError("Zero or near-zero quaternion in orientation set")
         arr = arr / norms[:, None]
         self.orientations = arr
+        self.xp = xp
+        self.use_cuda = use_cuda
 
-        # Precompute rotation vectors for BallTree
-        self.rotvecs = R.from_quat(arr).as_rotvec()
-        if len(arr) > 100:
-            self._balltree = BallTree(self.rotvecs)
-        else:
-            self._balltree = None
+        # Only for CPU/NumPy mode: BallTree for fast queries
+        self._balltree = None
+        if not use_cuda:
+            from scipy.spatial.transform import Rotation as R
+            from sklearn.neighbors import BallTree
+            self.rotvecs = R.from_quat(arr).as_rotvec()
+            if len(arr) > 100:
+                self._balltree = BallTree(self.rotvecs)
+            else:
+                self._balltree = None
 
     def query_within_angle(self, quat, angle):
         """Return indices of orientations within the given angle of quat."""
-        rv = R.from_quat(quat).as_rotvec().reshape(1, -1)
-        if self._balltree is not None:
-            # BallTree in rotation vector space, Euclidean distance ≈ angle for small rotations
-            inds = self._balltree.query_radius(rv, r=angle)[0]
+        if self.use_cuda:
+            # Brute force: batch math on GPU
+            # Convert quat to rotvec on CPU, then broadcast to GPU
+            import numpy as np
+            from scipy.spatial.transform import Rotation as R
+            rv = R.from_quat(np.array(quat)).as_rotvec().reshape(1, -1)
+            rv_gpu = self.xp.array(rv)
+            dists = self.xp.linalg.norm(self.orientations - rv_gpu, axis=1)
+            inds = self.xp.where(dists < angle)[0]
+            return inds
         else:
-            # Brute force for small sets
-            dists = np.linalg.norm(self.rotvecs - rv, axis=1)
-            inds = np.where(dists < angle)[0]
-        return inds
+            from scipy.spatial.transform import Rotation as R
+            rv = R.from_quat(quat).as_rotvec().reshape(1, -1)
+            if self._balltree is not None:
+                inds = self._balltree.query_radius(rv, r=angle)[0]
+            else:
+                dists = self.xp.linalg.norm(self.rotvecs - rv, axis=1)
+                inds = self.xp.where(dists < angle)[0]
+            return inds
+
+    def as_numpy(self):
+        if self.use_cuda:
+            return self.xp.asnumpy(self.orientations)
+        return self.orientations
 
     # ... rest unchanged ...
 

spinstep/discrete_iterator.py

@@ -1,3 +1,7 @@
+# discrete_iterator.py — MIT License
+# Author: Eraldo Marques <eraldo.bernardo@gmail.com> — Created: 2025-05-14
+# See LICENSE.txt for full terms. This header must be retained in redistributions.
+
 import numpy as np
 from scipy.spatial.transform import Rotation as R
 

spinstep/node.py

@@ -1,3 +1,7 @@
+# node.py — MIT License
+# Author: Eraldo Marques <eraldo.bernardo@gmail.com> — Created: 2025-05-14
+# See LICENSE.txt for full terms. This header must be retained in redistributions.
+
 import numpy as np
 
 class Node:

spinstep/quaternion_utils.py

@@ -1,3 +1,7 @@
+# quaternion_utils.py — MIT License
+# Author: Eraldo Marques <eraldo.bernardo@gmail.com> — Created: 2025-05-14
+# See LICENSE.txt for full terms. This header must be retained in redistributions.
+
 import numpy as np
 from scipy.spatial.transform import Rotation as R
 

spinstep/traversal.py

@@ -1,3 +1,7 @@
+# traversal.py — MIT License
+# Author: Eraldo Marques <eraldo.bernardo@gmail.com> — Created: 2025-05-14
+# See LICENSE.txt for full terms. This header must be retained in redistributions.
+
 from scipy.spatial.transform import Rotation as R
 
 class QuaternionDepthIterator:

spinstep/utils/array_backend.py

@@ -0,0 +1,13 @@
+# array_backend.py — MIT License
+# Author: Eraldo B. Marques <eraldo.bernardo@gmail.com> — Created: 2025-05-14
+# See LICENSE.txt for full terms. This header must be retained in redistributions.
+
+def get_array_module(use_cuda=False):
+    if use_cuda:
+        try:
+            import cupy as cp
+            return cp
+        except ImportError:
+            print("[SpinStep] CuPy not found, falling back to NumPy.")
+    import numpy as np
+    return np

spinstep/utils/quaternion_math.py

@@ -0,0 +1,16 @@
+# quaternion_math.py — MIT License
+# Author: Eraldo B. Marques <eraldo.bernardo@gmail.com> — Created: 2025-05-14
+# See LICENSE.txt for full terms. This header must be retained in redistributions.
+
+def batch_quaternion_angle(qs1, qs2, xp):
+    """
+    qs1: (N, 4) array
+    qs2: (M, 4) array
+    xp: array module (np or cp)
+    Returns (N, M) array of angular distances.
+    """
+    # Quaternion inner product: angle = 2*arccos(|dot(q1, q2)|)
+    dots = xp.abs(xp.dot(qs1, qs2.T))
+    dots = xp.clip(dots, -1.0, 1.0)
+    angles = 2 * xp.arccos(dots)
+    return angles