FastQuat 1.0b1__tar.gz → 1.0b2__tar.gz

{fastquat-1.0b1 → fastquat-1.0b2}/PKG-INFO

@@ -1,6 +1,6 @@
-Metadata-Version: 2.4
+Metadata-Version: 2.5
 Name: FastQuat
-Version: 1.0b1
+Version: 1.0b2
 Summary: High-performance quaternions with JAX support
 Project-URL: homepage, https://fastquat.readthedocs.io
 Project-URL: repository, https://github.com/CMBSciPol/fastquat
@@ -55,11 +55,26 @@ from fastquat import Quaternion
 # Create quaternions
 q1 = Quaternion(1)  # Identity quaternion
 q2 = Quaternion(0.7071, 0.7071, 0.0, 0.0)  # 90° rotation around x-axis
-
-# Quaternion operations
-q3 = q1 * q2  # Multiplication
-q_inv = 1 / q1           # Inverse, or q1 ** -1
-q_norm = q1.normalize()  # Normalization
+q = Quaternion(1.0, 0.1, 0.2, 0.3)
+p = 2
+
+# Arithmetic
+q_sum = q1 + q2
+q_diff = q1 - q2
+q_product = q1 * q2
+q_power = q**p
+
+# Normalization
+norm = abs(q)  # Quaternion norm
+q_unit = q.normalize()  # Unit quaternion
+
+# Conjugation and inverse
+q_conj = q.conj()  # Conjugate
+q_inv = 1 / q  # Inverse, or q ** -1
+
+# Other operations
+q_log = q.log()
+q_exp = q.exp()
 
 # Rotate vectors
 vector = jnp.array([1.0, 0.0, 0.0])

{fastquat-1.0b1 → fastquat-1.0b2}/README.md

@@ -32,11 +32,26 @@ from fastquat import Quaternion
 # Create quaternions
 q1 = Quaternion(1)  # Identity quaternion
 q2 = Quaternion(0.7071, 0.7071, 0.0, 0.0)  # 90° rotation around x-axis
-
-# Quaternion operations
-q3 = q1 * q2  # Multiplication
-q_inv = 1 / q1           # Inverse, or q1 ** -1
-q_norm = q1.normalize()  # Normalization
+q = Quaternion(1.0, 0.1, 0.2, 0.3)
+p = 2
+
+# Arithmetic
+q_sum = q1 + q2
+q_diff = q1 - q2
+q_product = q1 * q2
+q_power = q**p
+
+# Normalization
+norm = abs(q)  # Quaternion norm
+q_unit = q.normalize()  # Unit quaternion
+
+# Conjugation and inverse
+q_conj = q.conj()  # Conjugate
+q_inv = 1 / q  # Inverse, or q ** -1
+
+# Other operations
+q_log = q.log()
+q_exp = q.exp()
 
 # Rotate vectors
 vector = jnp.array([1.0, 0.0, 0.0])

fastquat-1.0b2/docs/README.md

@@ -0,0 +1,27 @@
+# Documentation Build Instructions
+
+This documentation uses **PEP 735 dependency groups** instead of separate requirements files to avoid duplication.
+
+
+## Building Documentation Locally
+
+The dependency group `docs` is included in the group `dev` which is installed by default when using  ̀uv` (recommended).
+No extra step other than setting up and activating the virtual environment is required:
+```bash
+# From the docs/ directory
+uv sync --locked --group cuda12
+source .venv/bin/activate
+```
+
+If you're not using uv, you can install the docs dependencies manually:
+```bash
+# Install the project with docs dependencies (pip >= 25.1)
+pip install --upgrade pip
+pip install -e . --group dev
+```
+
+Then build the docs
+```bash
+cd docs
+make html
+```

{fastquat-1.0b1 → fastquat-1.0b2}/docs/source/index.md

@@ -19,26 +19,30 @@ FastQuat provides optimized quaternion operations with full JAX compatibility, f
 ## Quick Start
 
 ```python
-import jax.numpy as jnp
 from fastquat import Quaternion
 
 # Create quaternions
 q1 = Quaternion(1.0)  # Identity quaternion
 q2 = Quaternion(0.7071, 0.7071, 0.0, 0.0)  # 90° rotation around x-axis
-
-# Quaternion operations
-q3 = q1 * q2  # Multiplication
-q_inv = 1 / q1  # Inverse
-q_norm = q1.normalize()  # Normalization
-
-# Rotate vectors
-vector = jnp.array([1.0, 0.0, 0.0])
-rotated = q2.rotate_vector(vector)
-
-# Spherical interpolation (SLERP)
-interpolated = q1.slerp(q2, t=0.5)  # Halfway between q1 and q2
 ```
 
+## Quaternion Operations
+
+| Operation | Syntax | Description |
+|-----------|--------|-------------|
+| Addition | `q1 + q2` | Component-wise addition |
+| Subtraction | `q1 - q2` | Component-wise subtraction |
+| Multiplication | `q1 * q2` | Hamilton product |
+| Exponentiation | `q ** p` | Quaternion power |
+| Inverse | `1 / q` | Multiplicative inverse |
+| Norm | `abs(q)` | Quaternion magnitude |
+| Normalization | `q.normalize()` | Unit quaternion |
+| Conjugate | `q.conj()` | Quaternion conjugate |
+| Rotation | `q.rotate_vector(v)` | Rotate 3D vector |
+| SLERP | `q1.slerp(q2, t)` | Spherical interpolation |
+| Log | `q.log()` | Quaternion logarithm |
+| Exp | `q.exp()` | Quaternion exponential |
+
 ```{toctree}
 :maxdepth: 2
 :caption: User Guide

{fastquat-1.0b1 → fastquat-1.0b2}/src/fastquat/quaternion.py

@@ -292,6 +292,12 @@ class Quaternion:
         for i in range(self.shape[0]):
             yield Quaternion.from_array(self.wxyz[i])
 
+    def __getitem__(self, idx: Any) -> Self:
+        """Index or slice the tensor of quaternions."""
+        if not isinstance(idx, tuple):
+            idx = (idx,)
+        return Quaternion.from_array(self.wxyz[(*idx, slice(None))])
+
     def __pos__(self) -> Self:
         """Quaternion positive."""
         return self

fastquat-1.0b2/tests/test_indexing.py

@@ -0,0 +1,208 @@
+"""Indexing tests for Quaternion class.
+
+Tests for __getitem__: integers, slices, boolean masks, fancy integer arrays, None
+(newaxis), Ellipsis, and combinations thereof. The component axis (w, x, y, z) must
+never be reachable by these indices -- only the batch shape is indexed.
+"""
+
+import jax
+import jax.numpy as jnp
+import numpy as np
+import pytest
+
+from fastquat.quaternion import Quaternion
+
+
+def _arange_quat(*shape: int) -> Quaternion:
+    """A Quaternion whose wxyz array is `jnp.arange(prod(shape) * 4).reshape(*shape, 4)`."""
+    size = 1
+    for dim in shape:
+        size *= dim
+    return Quaternion.from_array(jnp.arange(size * 4, dtype=jnp.float32).reshape(*shape, 4))
+
+
+@pytest.mark.parametrize(
+    'idx',
+    [0, -1, 2, slice(1, 4), slice(None, None, 2), slice(None, None, -1), slice(2, None)],
+)
+@pytest.mark.parametrize('do_jit', [False, True])
+def test_getitem_int_and_slice(idx, do_jit):
+    """Basic and strided/negative slicing over the leading axis."""
+
+    # `idx` is a Python-level constant closed over by `func`, not a traced argument.
+    def func(q):
+        return q[idx]
+
+    if do_jit:
+        func = jax.jit(func)
+
+    q = _arange_quat(6)
+    result = func(q)
+    expected = q.wxyz[idx]
+    assert result.wxyz.shape == expected.shape
+    assert jnp.allclose(result.wxyz, expected)
+
+
+def test_getitem_boolean_mask():
+    """Boolean-mask indexing preserves the component axis.
+
+    Not parametrized over `do_jit`: a traced boolean mask is non-concrete under `jit`, exactly
+    as it would be for a plain JAX array (`arr[mask]`) -- this is a JAX limitation, not
+    something `Quaternion.__getitem__` needs to work around.
+    """
+    mask = jnp.array([True, False, True, True, False])
+    q = _arange_quat(5)
+    result = q[mask]
+    assert result.shape == (3,)
+    assert jnp.allclose(result.wxyz, q.wxyz[mask])
+
+
+@pytest.mark.parametrize('do_jit', [False, True])
+def test_getitem_fancy_integer_array(do_jit):
+    """Gather with a dynamic integer index array (e.g. batched detector selection)."""
+    idx = jnp.array([3, 0, 0, 4])
+
+    def func(q, i):
+        return q[i]
+
+    if do_jit:
+        func = jax.jit(func)
+
+    q = _arange_quat(5)
+    result = func(q, idx)
+    assert result.shape == (4,)
+    assert jnp.allclose(result.wxyz, q.wxyz[idx])
+
+
+@pytest.mark.parametrize('do_jit', [False, True])
+def test_getitem_fancy_integer_array_out_of_range_clamps(do_jit):
+    """Out-of-range indices clamp like default JAX gather semantics (no IndexError)."""
+    idx = jnp.array([0, 10, -10])
+
+    def func(q, i):
+        return q[i]
+
+    if do_jit:
+        func = jax.jit(func)
+
+    q = _arange_quat(5)
+    result = func(q, idx)
+    assert jnp.allclose(result.wxyz, q.wxyz[idx])
+
+
+@pytest.mark.parametrize(
+    'idx,expected_shape',
+    [
+        ((None, slice(None)), (1, 5)),
+        ((slice(None), None), (5, 1)),
+        ((None, slice(None), None), (1, 5, 1)),
+    ],
+)
+@pytest.mark.parametrize('do_jit', [False, True])
+def test_getitem_none_newaxis(idx, expected_shape, do_jit):
+    """None/newaxis inserts a batch axis without disturbing the components."""
+
+    def func(q):
+        return q[idx]
+
+    if do_jit:
+        func = jax.jit(func)
+
+    q = _arange_quat(5)
+    result = func(q)
+    assert result.shape == expected_shape
+    assert jnp.allclose(result.wxyz, q.wxyz[idx])
+
+
+@pytest.mark.parametrize('do_jit', [False, True])
+def test_getitem_ellipsis_with_none(do_jit):
+    """Ellipsis combined with None on a multi-dimensional batch shape."""
+
+    def func(q):
+        return q[None, ...]
+
+    if do_jit:
+        func = jax.jit(func)
+
+    q = _arange_quat(2, 3)
+    result = func(q)
+    assert result.shape == (1, 2, 3)
+    assert jnp.allclose(result.wxyz, q.wxyz[None, ...])
+
+
+@pytest.mark.parametrize('do_jit', [False, True])
+def test_getitem_multidim_tuple(do_jit):
+    """A tuple index spanning multiple batch dimensions."""
+
+    def func(q):
+        return q[1:, 0]
+
+    if do_jit:
+        func = jax.jit(func)
+
+    q = _arange_quat(3, 4)
+    result = func(q)
+    assert result.shape == (2,)
+    assert jnp.allclose(result.wxyz, q.wxyz[1:, 0])
+
+
+def test_getitem_scalar_index_reduces_ndim():
+    """Indexing with a plain int drops that axis, exactly like a JAX array."""
+    q = _arange_quat(6)
+    result = q[2]
+    assert result.shape == ()
+    assert jnp.allclose(result.wxyz, q.wxyz[2])
+
+
+def test_getitem_never_indexes_component_axis():
+    """A full-rank slice tuple must still leave all 4 components intact."""
+    q = _arange_quat(3)
+    # `q[:]` only has one explicit index; the hidden component axis must stay whole.
+    result = q[:]
+    assert result.wxyz.shape == (3, 4)
+    assert jnp.allclose(result.wxyz, q.wxyz)
+
+
+@pytest.mark.parametrize('do_jit', [False, True])
+def test_getitem_broadcast_multiplication(do_jit):
+    """The pattern furax needs: q1[None, :] * q2[:, None] broadcasts like raw arrays would."""
+
+    def func(q1, q2):
+        return q1[None, :] * q2[:, None]
+
+    if do_jit:
+        func = jax.jit(func)
+
+    q1 = Quaternion.random(jax.random.key(0), shape=(5,))
+    q2 = Quaternion.random(jax.random.key(1), shape=(3,))
+    result = func(q1, q2)
+    assert result.shape == (3, 5)
+
+    # Cross-check against the equivalent reshape-based broadcast (the old workaround).
+    expected = q1.reshape((1, 5)) * q2.reshape((3, 1))
+    assert jnp.allclose(result.wxyz, expected.wxyz)
+
+
+def test_getitem_preserves_dtype():
+    q = Quaternion.from_array(jnp.ones((4, 4), dtype=jnp.float16))
+    assert q[1:3].dtype == jnp.float16
+
+
+def test_getitem_pytree_leaf_indexing_under_vmap():
+    """Quaternion indexing composes with jax.vmap, exercising the pytree registration."""
+    q = _arange_quat(3, 4)
+
+    def first_row(qi: Quaternion) -> Quaternion:
+        return qi[0]
+
+    result = jax.vmap(first_row)(q)
+    assert result.shape == (3,)
+    assert jnp.allclose(result.wxyz, q.wxyz[:, 0])
+
+
+def test_getitem_boolean_mask_matches_numpy_reference():
+    rng = np.random.default_rng(0)
+    mask = rng.random(7) > 0.5
+    q = _arange_quat(7)
+    result = q[jnp.asarray(mask)]
+    assert jnp.allclose(result.wxyz, q.wxyz[mask])

fastquat-1.0b1/docs/README.md

@@ -1,96 +0,0 @@
-# Documentation Build Instructions
-
-This documentation uses **PEP 735 dependency groups** instead of separate requirements files to avoid duplication.
-
-## Building Documentation Locally
-
-### Using uv (recommended)
-
-```bash
-# From the docs/ directory
-cd docs/
-uv run --group docs make html
-
-# Or from the root directory
-uv run --group docs sphinx-build docs/source docs/build/html
-```
-
-### Using other tools
-
-If you're not using uv, you can install the docs dependencies manually:
-
-```bash
-# Install the project with docs dependencies
-pip install -e .[docs]
-
-# Then build the docs
-cd docs/
-make html
-```
-
-## What Changed
-
-### ✅ Before (duplicated dependencies)
-- Dependencies listed in both `pyproject.toml` and `docs/requirements.txt`
-- ReadTheDocs config pointed to `docs/requirements.txt`
-
-### ✅ After (PEP 735 dependency groups)
-- Dependencies only in `pyproject.toml` under `[dependency-groups.docs]`
-- ReadTheDocs config uses `extra_requirements: [docs]`
-- No more `docs/requirements.txt` file
-
-## PEP 735 Benefits
-
-1. **Single source of truth** - dependencies defined once in `pyproject.toml`
-2. **No duplication** - eliminates sync issues between files
-3. **Modern standard** - follows PEP 735 specification
-4. **Tool support** - works with uv, pip, and ReadTheDocs
-
-## ReadTheDocs Configuration
-
-The `.readthedocs.yaml` configuration now uses:
-
-```yaml
-python:
-  install:
-    - method: pip
-      path: .
-      extra_requirements:
-        - docs
-```
-
-This automatically installs the project with the `docs` dependency group.
-
-## Available Commands
-
-```bash
-# Build HTML documentation
-uv run --group docs make html
-
-# Clean and rebuild
-uv run --group docs make clean html
-
-# Auto-rebuild on changes (requires sphinx-autobuild)
-uv run --group docs sphinx-autobuild source build/html
-
-# View available make targets
-uv run --group docs make help
-```
-
-## Dependencies in the `docs` Group
-
-The documentation dependencies are defined in `pyproject.toml`:
-
-```toml
-[dependency-groups]
-docs = [
-    "sphinx>=7.0.0",
-    "sphinx-rtd-theme>=1.3.0",
-    "myst-parser>=2.0.0",
-    "nbsphinx>=0.9.0",
-    "jupyter>=1.0.0",
-    "matplotlib>=3.7.0",
-    "numpy>=1.24.0",
-    "pillow>=9.0.0",
-]
-```