FastQuat 1.0b2__tar.gz → 1.0b4__tar.gz

fastquat-1.0b4/.github/workflows/ci.yml

@@ -0,0 +1,71 @@
+name: CI
+
+permissions: {}
+
+on:
+  push:
+    branches:
+    - main
+  pull_request:
+  workflow_dispatch:
+
+concurrency:
+  group: ${{ github.workflow }}-${{ github.ref_name }}-${{ github.event.pull_request.number || github.sha }}
+  cancel-in-progress: true
+
+jobs:
+
+  pre-commit:
+    name: Pre-commit
+    runs-on: ubuntu-latest
+    timeout-minutes: 5
+    steps:
+    - uses: actions/checkout@v7.0.1
+    - name: Install uv
+      uses: astral-sh/setup-uv@v10.0.1
+      with:
+        version: '0.12.18'
+    - name: Cache prek
+      uses: actions/cache@v6.1.0
+      with:
+        path: ~/.cache/prek
+        key: prek-${{ hashFiles('.pre-commit-config.yaml') }}
+        restore-keys: |
+          prek-
+    # ty resolves third-party imports from the environment, so jax has to be installed.
+    - name: Install dependencies
+      run: uv sync --locked --no-default-groups --group lint
+    - name: Run pre-commit
+      run: uv run --no-sync prek run --all-files --show-diff-on-failure --color always
+
+  test:
+    name: Run tests on Python ${{ matrix.python-version }}
+    runs-on: ubuntu-latest
+    timeout-minutes: 15
+    strategy:
+      fail-fast: false
+      matrix:
+        python-version: ['3.10', '3.11', '3.12', '3.13', '3.14']
+
+    steps:
+      - name: Checkout code
+        uses: actions/checkout@v7.0.1
+        with:
+          fetch-depth: 0
+
+      - name: Install uv
+        uses: astral-sh/setup-uv@v10.0.1
+        with:
+          version: '0.12.18'
+          python-version: ${{ matrix.python-version }}
+
+      - name: Build wheel
+        run: uv build --wheel
+
+      - name: Install project's wheel and dependencies
+        run: |
+          uv sync --locked --no-install-project
+          uv pip install dist/fastquat-*.whl
+
+      - name: Run tests
+        run: uv run --no-sync pytest -v

{fastquat-1.0b2 → fastquat-1.0b4}/.pre-commit-config.yaml

@@ -7,9 +7,9 @@ repos:
     - --all
 
 - repo: https://github.com/astral-sh/ruff-pre-commit
-  rev: "v0.14.7"
+  rev: "v0.16.9"
   hooks:
-    - id: ruff
+    - id: ruff-check
       name: ruff linting
     - id: ruff-format
       name: ruff formatting
@@ -31,21 +31,13 @@ repos:
   - id: check-merge-conflict
 
 - repo: https://github.com/kynan/nbstripout
-  rev: 0.6.1
+  rev: "0.9.1"
   hooks:
   - id: nbstripout
     name: notebook stripping
     args: [--extra-keys=metadata.language_info.version]
 
-- repo: https://github.com/pre-commit/mirrors-mypy
-  rev: 'v1.19.0'
+- repo: https://github.com/astral-sh/ty-pre-commit
+  rev: v0.0.84
   hooks:
-  - id: mypy
-    additional_dependencies:
-    - jax
-    args:
-    - --strict
-    - --show-error-codes
-    - --enable-error-code=ignore-without-code
-    - --allow-untyped-calls
-    files: ^fastquat/
+  - id: ty

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.5
 Name: FastQuat
-Version: 1.0b2
+Version: 1.0b4
 Summary: High-performance quaternions with JAX support
 Project-URL: homepage, https://fastquat.readthedocs.io
 Project-URL: repository, https://github.com/CMBSciPol/fastquat
@@ -89,7 +89,7 @@ interpolated = q1.slerp(q2, t=0.5)  # Halfway between q1 and q2
 ### Core Operations
 - **Quaternion arithmetic**: Addition, multiplication, conjugation, inverse, power, exponentiation, logarithm
 - **Normalization**: Efficient unit quaternion computation
-- **Conversion**: To/from rotation matrices, Euler angles
+- **Conversion**: To/from rotation matrices, axis-angle, Euler angles
 - **Vector rotation**: Direct vector transformation
 
 ### Advanced Interpolation
@@ -124,11 +124,13 @@ from fastquat import Quaternion
 key = jax.random.PRNGKey(42)
 q_batch = Quaternion.random(key, shape=(1000,))
 
+
 # JIT-compiled batch operations
 @jax.jit
 def batch_rotate(quaternions, vectors):
     return quaternions.rotate_vector(vectors)
 
+
 vectors = jax.random.normal(key, (1000, 3))
 rotated_batch = batch_rotate(q_batch, vectors)
 ```

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

@@ -66,7 +66,7 @@ interpolated = q1.slerp(q2, t=0.5)  # Halfway between q1 and q2
 ### Core Operations
 - **Quaternion arithmetic**: Addition, multiplication, conjugation, inverse, power, exponentiation, logarithm
 - **Normalization**: Efficient unit quaternion computation
-- **Conversion**: To/from rotation matrices, Euler angles
+- **Conversion**: To/from rotation matrices, axis-angle, Euler angles
 - **Vector rotation**: Direct vector transformation
 
 ### Advanced Interpolation
@@ -101,11 +101,13 @@ from fastquat import Quaternion
 key = jax.random.PRNGKey(42)
 q_batch = Quaternion.random(key, shape=(1000,))
 
+
 # JIT-compiled batch operations
 @jax.jit
 def batch_rotate(quaternions, vectors):
     return quaternions.rotate_vector(vectors)
 
+
 vectors = jax.random.normal(key, (1000, 3))
 rotated_batch = batch_rotate(q_batch, vectors)
 ```

{fastquat-1.0b2 → fastquat-1.0b4}/docs/source/api/quaternion.md

@@ -18,6 +18,7 @@ compilation, automatic differentiation, and vectorization.
 .. automethod:: fastquat.Quaternion.from_array
 .. automethod:: fastquat.Quaternion.from_scalar_vector
 .. automethod:: fastquat.Quaternion.from_rotation_matrix
+.. automethod:: fastquat.Quaternion.from_axis_angle
 .. automethod:: fastquat.Quaternion.zeros
 .. automethod:: fastquat.Quaternion.ones
 .. automethod:: fastquat.Quaternion.full

{fastquat-1.0b2 → fastquat-1.0b4}/docs/source/development.md

@@ -112,20 +112,21 @@ from jax.typing import ArrayLike
 
 
 class Quaternion:
-   ...
-   def new_method(self, parameter: ArrayLike) -> Self:
-       """Brief description of what the method does.
-
-       Args:
-           parameter: Description of the parameter
-
-       Returns:
-           Description of the return value
-       """
-       # Implementation using JAX operations
-       parameter = jnp.asarray(parameter)
-       result = jnp.some_operation(self.wxyz, parameter)
-       return Quaternion.from_array(result)
+    ...
+
+    def new_method(self, parameter: ArrayLike) -> Self:
+        """Brief description of what the method does.
+
+        Args:
+            parameter: Description of the parameter
+
+        Returns:
+            Description of the return value
+        """
+        # Implementation using JAX operations
+        parameter = jnp.asarray(parameter)
+        result = jnp.some_operation(self.wxyz, parameter)
+        return Quaternion.from_array(result)
 ```
 
 Then add tests:
@@ -138,15 +139,11 @@ import pytest
 from fastquat import Quaternion
 
 
-@pytest.mark.parametrize(
-   'parameter, expected_values',
-   [
-      (..., ...),
-   ]
-)
+@pytest.mark.parametrize('parameter, expected_values', [(..., ...)])
 @pytest.mark.parametrize('do_jit', [False, True])
 def test_new_method(parameter, expected_values, do_jit):
     """Test the new method."""
+
     def test_fn(q_, parameter_):
         return q_.new_method(parameter_)
 

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

@@ -39,6 +39,7 @@ q2 = Quaternion(0.7071, 0.7071, 0.0, 0.0)  # 90° rotation around x-axis
 | Normalization | `q.normalize()` | Unit quaternion |
 | Conjugate | `q.conj()` | Quaternion conjugate |
 | Rotation | `q.rotate_vector(v)` | Rotate 3D vector |
+| Axis-angle | `Quaternion.from_axis_angle(axis, angle)` | Rotation about an axis |
 | SLERP | `q1.slerp(q2, t)` | Spherical interpolation |
 | Log | `q.log()` | Quaternion logarithm |
 | Exp | `q.exp()` | Quaternion exponential |

{fastquat-1.0b2 → fastquat-1.0b4}/docs/source/user-guide/installation.md

@@ -52,12 +52,12 @@ from fastquat import Quaternion
 
 # Create a simple quaternion
 q = Quaternion(1.0)
-print(f"Identity quaternion: {q}")
+print(f'Identity quaternion: {q}')
 
 # Test SLERP functionality
 q2 = Quaternion(0.7071, 0.7071, 0.0, 0.0)
 interpolated = q.slerp(q2, 0.5)
-print(f"SLERP result: {interpolated}")
+print(f'SLERP result: {interpolated}')
 ```
 
 If this runs without errors, FastQuat is properly installed!

{fastquat-1.0b2 → fastquat-1.0b4}/pyproject.toml

@@ -34,6 +34,7 @@ repository = 'https://github.com/CMBSciPol/fastquat'
 [dependency-groups]
 dev = [
     {include-group = 'docs'},
+    {include-group = 'lint'},
     'pytest>=9.0',
     'pytest-notebook>=0.10',
     'visu-hlo',
@@ -48,6 +49,7 @@ docs = [
     'numpy>=1.24.0',
     'pillow>=9.0.0',
 ]
+lint = ['prek>=0.5.4']
 cuda12 = ['jax[cuda12]>=0.4']
 cuda13 = ['jax[cuda13]>=0.7; python_version>="3.11"']
 
@@ -83,5 +85,14 @@ select = [
 [tool.ruff.format]
 quote-style = 'single'
 
+[tool.ty.rules]
+dynamic-function-decorator-return = "error"
+missing-type-argument = "error"
+possibly-unresolved-reference = "warn"
+unused-ignore-comment = "error"
+
+[tool.ty.src]
+include = ["src"]
+
 [tool.uv]
 cache-keys = [{ git = true }]

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

@@ -118,6 +118,32 @@ class Quaternion:
 
         return cls.from_array(jnp.stack([w, x, y, z], axis=-1))
 
+    @classmethod
+    def from_axis_angle(cls, axis: ArrayLike, angle: ArrayLike) -> Self:
+        """Create the unit quaternion of a rotation about an axis.
+
+        The rotation follows the right-hand rule: a positive angle rotates counterclockwise when
+        looking from the tip of the axis towards the origin.
+
+        Args:
+            axis: Array of shape (..., 3) for the rotation axis. It does not need to be normalized.
+            angle: Array of shape (...) for the rotation angle, in radians.
+
+        Returns:
+            Quaternion of shape broadcast(axis.shape[:-1], angle.shape).
+        """
+        axis = jnp.asarray(axis)
+        angle = jnp.asarray(angle)
+        if axis.shape[-1:] != (3,):
+            raise ValueError(f'Axis must have shape (..., 3), got {axis.shape}')
+        dtype = jnp.result_type(axis, angle, float)
+        unit_axis = axis / jnp.linalg.norm(axis, axis=-1, keepdims=True)
+        half_angle = 0.5 * angle.astype(dtype)
+        scalar = jnp.cos(half_angle)
+        vector = jnp.sin(half_angle)[..., None] * unit_axis.astype(dtype)
+        scalar, vector = jnp.broadcast_arrays(scalar[..., None], vector)
+        return cls.from_scalar_vector(scalar[..., 0], vector)
+
     @classmethod
     def zeros(cls, shape: tuple[int, ...], dtype: DTypeLike | None = None) -> Self:
         """Create quaternions with all components set to 0.
@@ -167,7 +193,7 @@ class Quaternion:
 
     @classmethod
     def random(
-        cls, key: jax.random.PRNGKey, shape: tuple[int, ...] = (), dtype: DTypeLike | None = None
+        cls, key: Array, shape: tuple[int, ...] = (), dtype: DTypeLike | None = None
     ) -> Self:
         """Generate normalized random quaternions.
 
@@ -180,7 +206,7 @@ class Quaternion:
             Normalized Quaternion.
         """
         data = jax.random.normal(key, shape + (4,), dtype=dtype)
-        return Quaternion.from_array(data).normalize()
+        return cls.from_array(data).normalize()
 
     @property
     def w(self) -> Array:
@@ -214,13 +240,13 @@ class Quaternion:
         returns the quaternion [NaN, NaN, NaN, NaN].
         """
         norm = abs(self)
-        return Quaternion.from_array(self.wxyz / jnp.expand_dims(norm, axis=-1))
+        return self.from_array(self.wxyz / jnp.expand_dims(norm, axis=-1))
 
     def _inverse(self) -> Self:
         """Quaternion inverse (private method - use 1/q instead)."""
         conj = self.conj()
         norm_sq = jnp.sum(self.wxyz**2, axis=-1)
-        return Quaternion.from_array(conj.wxyz / jnp.expand_dims(norm_sq, axis=-1))
+        return self.from_array(conj.wxyz / jnp.expand_dims(norm_sq, axis=-1))
 
     def to_components(self) -> tuple[Array, Array, Array, Array]:
         return self.w, self.x, self.y, self.z
@@ -290,13 +316,13 @@ class Quaternion:
         if self.ndim == 0:
             raise TypeError('iteration over a 0-d quaternion')
         for i in range(self.shape[0]):
-            yield Quaternion.from_array(self.wxyz[i])
+            yield self.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))])
+        return self.from_array(self.wxyz[(*idx, slice(None))])
 
     def __pos__(self) -> Self:
         """Quaternion positive."""
@@ -304,12 +330,12 @@ class Quaternion:
 
     def __neg__(self) -> Self:
         """Quaternion negation."""
-        return Quaternion.from_array(-self.wxyz)
+        return self.from_array(-self.wxyz)
 
     def __add__(self, other: Any) -> Self:
         """Quaternion addition."""
         if isinstance(other, Quaternion):
-            return Quaternion.from_array(self.wxyz + other.wxyz)
+            return self.from_array(self.wxyz + other.wxyz)
 
         try:
             other = jnp.asarray(other)
@@ -319,7 +345,7 @@ class Quaternion:
         if jnp.iscomplexobj(other):
             raise NotImplementedError('Quaternion and complex addition is not implemented.')
 
-        return Quaternion.from_scalar_vector(self.w + other, self.vector)
+        return self.from_scalar_vector(self.w + other, self.vector)
 
     def __radd__(self, other: Any) -> Self:
         """Quaternion addition."""
@@ -328,7 +354,7 @@ class Quaternion:
     def __sub__(self, other: Any) -> Self:
         """Quaternion subtraction."""
         if isinstance(other, Quaternion):
-            return Quaternion.from_array(self.wxyz - other.wxyz)
+            return self.from_array(self.wxyz - other.wxyz)
 
         try:
             other = jnp.asarray(other)
@@ -338,7 +364,7 @@ class Quaternion:
         if jnp.iscomplexobj(other):
             raise NotImplementedError('Quaternion and complex subtraction is not implemented.')
 
-        return Quaternion.from_scalar_vector(self.w - other, self.vector)
+        return self.from_scalar_vector(self.w - other, self.vector)
 
     def __rsub__(self, other: Any) -> Self:
         """Quaternion subtraction."""
@@ -350,7 +376,7 @@ class Quaternion:
         if jnp.iscomplexobj(other):
             raise NotImplementedError('Quaternion and complex subtraction is not implemented.')
 
-        return Quaternion.from_scalar_vector(other - self.w, -self.vector)
+        return self.from_scalar_vector(other - self.w, -self.vector)
 
     def __mul__(self, other: Any) -> Self:
         """Quaternion multiplication."""
@@ -363,7 +389,7 @@ class Quaternion:
             y = w1 * y2 - x1 * z2 + y1 * w2 + z1 * x2
             z = w1 * z2 + x1 * y2 - y1 * x2 + z1 * w2
 
-            return Quaternion(w, x, y, z)
+            return self.from_array(jnp.stack([w, x, y, z], axis=-1))
 
         try:
             other = jnp.asarray(other)
@@ -373,7 +399,7 @@ class Quaternion:
         if jnp.iscomplexobj(other):
             raise NotImplementedError('Quaternion and complex multiplication is not implemented.')
 
-        return Quaternion.from_array(self.wxyz * jnp.expand_dims(other, axis=-1))
+        return self.from_array(self.wxyz * jnp.expand_dims(other, axis=-1))
 
     def __rmul__(self, other: Any) -> Self:
         """Quaternion multiplication."""
@@ -385,7 +411,7 @@ class Quaternion:
         if jnp.iscomplexobj(other):
             raise NotImplementedError('Quaternion and complex multiplication is not implemented.')
 
-        return Quaternion.from_array(jnp.expand_dims(other, axis=-1) * self.wxyz)
+        return self.from_array(jnp.expand_dims(other, axis=-1) * self.wxyz)
 
     def __truediv__(self, other: Any) -> Self:
         """Quaternion division."""
@@ -400,7 +426,7 @@ class Quaternion:
         if jnp.iscomplexobj(other):
             raise NotImplementedError('Quaternion and complex division is not implemented.')
 
-        return Quaternion.from_array(self.wxyz / jnp.expand_dims(other, axis=-1))
+        return self.from_array(self.wxyz / jnp.expand_dims(other, axis=-1))
 
     def __rtruediv__(self, other: Any) -> Self:
         """Quaternion division."""
@@ -437,7 +463,7 @@ class Quaternion:
             elif exponent == -1:
                 return self._inverse()
             elif exponent == 0:
-                return Quaternion.ones(self.shape, self.dtype)
+                return self.ones(self.shape, self.dtype)
             elif exponent == 1:
                 return self
             elif exponent == 2:
@@ -445,8 +471,9 @@ class Quaternion:
             return (exponent * self.log()).exp()
 
         # General case: q^n = exp(n * log(q))
+        exponent = jnp.asarray(exponent)
         result = (exponent * self.log()).exp().wxyz
-        return Quaternion.from_array(
+        return self.from_array(
             jnp.where(
                 exponent[..., None] == 0, jnp.array([1.0, 0.0, 0.0, 0.0], dtype=self.dtype), result
             )
@@ -455,63 +482,61 @@ class Quaternion:
     def log(self) -> Self:
         """Compute quaternion logarithm.
 
-        For a quaternion q = ‖q‖ * (cos(θ) + sin(θ)v), the logarithm is:
-        log(q) = log(‖q‖) + θ * v
+        For a quaternion q = |q| * (cos(θ) + sin(θ)v), the logarithm is:
+        log(q) = log(|q|) + θ * v
 
         For the zero quaternion, returns (-inf, 0, 0, 0).
 
         Returns:
             The logarithm of the quaternion
         """
-        q_norm = abs(self)
-
-        # Normalize manually to handle zero quaternion (returns 0 instead of NaN)
-        safe_norm = jnp.where(q_norm == 0, 1.0, q_norm)
-        unit_wxyz = self.wxyz / jnp.expand_dims(safe_norm, axis=-1)
-
-        # For unit quaternion q = cos(θ) + sin(θ)v, compute θ and v
-        # θ = arccos(w) and v = vector/|vector|
-        unit_w = unit_wxyz[..., 0]
-        unit_vector = unit_wxyz[..., 1:]
-        theta = jnp.arccos(jnp.clip(unit_w, -1.0, 1.0))
-        vector_norm = jnp.linalg.norm(unit_vector, axis=-1)
-
-        # Handle case where vector is zero (real quaternion)
-        inv_vector_norm = jnp.where(vector_norm == 0, 0.0, 1 / vector_norm)
-        unit_vector = unit_vector * inv_vector_norm[..., None]
-
-        # log(q) = log(|q|) + θ * v
-        log_norm = jnp.log(q_norm)
-        log_q_vector = theta[..., None] * unit_vector
+        scalar_part = self.w
+        vector_part = self.vector
+        vector_norm_sq = jnp.sum(vector_part**2, axis=-1)
+        is_real = vector_norm_sq == 0
+
+        # log(q) = log(|q|) + θ * v/|v|, with θ = atan2(|v|, s).
+        # The where guards keep the gradients finite for real quaternions (|v| = 0).
+        log_norm = 0.5 * jnp.log(scalar_part**2 + vector_norm_sq)
+        safe_vector_norm = jnp.sqrt(jnp.where(is_real, 1.0, vector_norm_sq))
+        safe_scalar_part = jnp.where(scalar_part == 0, 1.0, scalar_part)
+        # θ/|v| tends to 1/s when |v| → 0 (s > 0)
+        theta_over_vector_norm = jnp.where(
+            is_real,
+            1 / safe_scalar_part,
+            jnp.arctan2(safe_vector_norm, scalar_part) / safe_vector_norm,
+        )
+        log_q_vector = theta_over_vector_norm[..., None] * vector_part
 
-        return Quaternion.from_scalar_vector(log_norm, log_q_vector)
+        return self.from_scalar_vector(log_norm, log_q_vector)
 
     def exp(self) -> Self:
         """Compute quaternion exponential.
 
         For a quaternion q = s + v, the exponential is:
-        exp(q) = exp(s) * (cos(‖v‖) + sin(‖v‖) * v/‖v‖)
+        exp(q) = exp(s) * (cos(|v|) + sin(|v|) * v/|v|)
 
         Returns:
             The exponential of the quaternion
         """
         scalar_part = self.w
         vector_part = self.vector
-        vector_norm = jnp.linalg.norm(vector_part, axis=-1)
+        vector_norm_sq = jnp.sum(vector_part**2, axis=-1)
+        is_real = vector_norm_sq == 0
 
-        # exp(s + v) = exp(s) * (cos(|v|) + sin(|v|) * v/|v|)
+        # The where guard keeps the gradients finite for real quaternions (|v| = 0).
+        vector_norm = jnp.where(is_real, 0.0, jnp.sqrt(jnp.where(is_real, 1.0, vector_norm_sq)))
         exp_scalar = jnp.exp(scalar_part)
-        cos_vnorm = jnp.cos(vector_norm)
-        sin_vnorm = jnp.sin(vector_norm)
-
-        # Handle case where |v| = 0 (real quaternion)
-        inv_vector_norm = jnp.where(vector_norm == 0, 0.0, 1 / vector_norm)
-        unit_v = vector_part * jnp.expand_dims(inv_vector_norm, -1)
+        # sin(|v|)/|v| and cos(|v|) = 1 - |v|²/2 (sin(|v|/2)/(|v|/2))², written with sinc
+        # so that the first and second derivatives are exact at |v| = 0
+        sinc_vnorm = jnp.sinc(vector_norm / jnp.pi)
+        sinc_half_vnorm = jnp.sinc(vector_norm / (2 * jnp.pi))
+        cos_vnorm = 1 - 0.5 * vector_norm_sq * sinc_half_vnorm**2
 
         result_w = exp_scalar * cos_vnorm
-        result_vector = exp_scalar * jnp.expand_dims(sin_vnorm, -1) * unit_v
+        result_vector = jnp.expand_dims(exp_scalar * sinc_vnorm, -1) * vector_part
 
-        return Quaternion.from_scalar_vector(result_w, result_vector)
+        return self.from_scalar_vector(result_w, result_vector)
 
     @property
     def nbytes(self) -> int:
@@ -564,12 +589,12 @@ class Quaternion:
 
     def squeeze(self, axis=None) -> Self:
         """Supprime les dimensions de taille 1"""
-        return Quaternion.from_array(jnp.squeeze(self.wxyz, axis=axis))
+        return self.from_array(jnp.squeeze(self.wxyz, axis=axis))
 
     def conjugate(self) -> Self:
         """Quaternion conjugate."""
         sign = jnp.array([1, -1, -1, -1])
-        return Quaternion.from_array(self.wxyz * sign)
+        return self.from_array(self.wxyz * sign)
 
     def conj(self) -> Self:
         """Quaternion conjugate."""
@@ -617,7 +642,7 @@ class Quaternion:
 
         # Linear interpolation case
         result_linear = q1.wxyz + jnp.expand_dims(t * (1 - t), -1) * (q2_corrected - q1.wxyz)
-        result_linear = Quaternion.from_array(result_linear).normalize()
+        result_linear = self.from_array(result_linear).normalize()
 
         # Spherical interpolation case
         theta = jnp.arccos(jnp.clip(dot, 0.0, 1.0))
@@ -632,9 +657,9 @@ class Quaternion:
         result_slerp = (
             jnp.expand_dims(factor1, -1) * q1.wxyz + jnp.expand_dims(factor2, -1) * q2_corrected
         )
-        result_slerp = Quaternion.from_array(result_slerp)
+        result_slerp = self.from_array(result_slerp)
 
         # Choose between linear and spherical interpolation
         result = jnp.where(jnp.expand_dims(use_linear, -1), result_linear.wxyz, result_slerp.wxyz)
 
-        return Quaternion.from_array(result)
+        return self.from_array(result)

{fastquat-1.0b2 → fastquat-1.0b4}/tests/test_pow.py

@@ -1,5 +1,6 @@
 import jax
 import jax.numpy as jnp
+import jax.test_util
 import pytest
 
 from fastquat import Quaternion
@@ -386,3 +387,50 @@ def test_pow_zero_quaternion(cast_type, do_jit):
     # 0^n for n > 0 should give 0
     result_0_2 = func(zero_q, cast_type(2.0))
     assert jnp.allclose(result_0_2.wxyz, zero_q.wxyz)
+
+
+# Gradients of logarithm and exponential
+@pytest.mark.parametrize(
+    'wxyz',
+    [
+        [0.0, 0.0, 0.0, 0.0],  # pure zero: exp at the identity
+        [0.3, 0.0, 0.0, 0.0],  # real
+        [0.0, 0.1, -0.2, 0.3],  # pure imaginary
+        [0.5, 0.1, -0.2, 0.3],
+    ],
+)
+def test_exp_grads(wxyz, enable_x64):
+    """Test exp derivatives up to second order, including at a zero vector part."""
+    jax.test_util.check_grads(
+        lambda a: Quaternion.from_array(a).exp().wxyz, (jnp.array(wxyz),), order=2
+    )
+
+
+@pytest.mark.parametrize(
+    'wxyz',
+    [
+        [1.0, 0.0, 0.0, 0.0],  # identity
+        [2.0, 0.0, 0.0, 0.0],  # real
+        [0.0, 0.1, -0.2, 0.3],  # pure imaginary
+        [0.5, 0.1, -0.2, 0.3],
+    ],
+)
+def test_log_grads(wxyz, enable_x64):
+    """Test log derivatives, including at a zero vector part."""
+    jax.test_util.check_grads(
+        lambda a: Quaternion.from_array(a).log().wxyz, (jnp.array(wxyz),), order=1
+    )
+
+
+def test_exp_jacobian_at_zero():
+    """Test that d exp(v)/dv at v = 0 maps the vector part onto itself."""
+    jac = jax.jacobian(lambda v: Quaternion.from_scalar_vector(0.0, v).exp().wxyz)(jnp.zeros(3))
+    expected = jnp.concatenate([jnp.zeros((1, 3)), jnp.eye(3)])
+    assert jnp.allclose(jac, expected)
+
+
+def test_log_jacobian_at_identity():
+    """Test that d log(1 + v)/dv at v = 0 maps the vector part onto itself."""
+    jac = jax.jacobian(lambda v: Quaternion.from_scalar_vector(1.0, v).log().wxyz)(jnp.zeros(3))
+    expected = jnp.concatenate([jnp.zeros((1, 3)), jnp.eye(3)])
+    assert jnp.allclose(jac, expected)

{fastquat-1.0b2 → fastquat-1.0b4}/tests/test_rotation.py

@@ -79,6 +79,80 @@ def test_from_rotation_matrix_wrong_shape(do_jit):
         func(wrong_matrix)
 
 
+# from_axis_angle
+@pytest.mark.parametrize(
+    'axis, angle, expected',
+    [
+        ([0.0, 0.0, 1.0], 0.0, [1.0, 0.0, 0.0, 0.0]),
+        ([1.0, 0.0, 0.0], jnp.pi, [0.0, 1.0, 0.0, 0.0]),
+        ([0.0, 2.0, 0.0], jnp.pi / 2, [jnp.sqrt(0.5), 0.0, jnp.sqrt(0.5), 0.0]),
+        ([0.0, 0.0, 1.0], -jnp.pi / 2, [jnp.sqrt(0.5), 0.0, 0.0, -jnp.sqrt(0.5)]),
+    ],
+)
+@pytest.mark.parametrize('do_jit', [False, True])
+def test_from_axis_angle(axis, angle, expected, do_jit):
+    """Test from_axis_angle against known quaternions (axis need not be normalized)."""
+    func = Quaternion.from_axis_angle
+    if do_jit:
+        func = jax.jit(func)
+
+    q = func(jnp.array(axis), jnp.array(angle))
+    assert jnp.allclose(q.wxyz, jnp.array(expected), atol=1e-6)
+
+
+@pytest.mark.parametrize('do_jit', [False, True])
+def test_from_axis_angle_right_hand_rule(do_jit):
+    """A positive rotation about z maps x to y."""
+
+    def func(angle):
+        return Quaternion.from_axis_angle(jnp.array([0.0, 0.0, 1.0]), angle).rotate_vector(
+            jnp.array([1.0, 0.0, 0.0])
+        )
+
+    if do_jit:
+        func = jax.jit(func)
+
+    assert jnp.allclose(func(jnp.pi / 2), jnp.array([0.0, 1.0, 0.0]), atol=1e-6)
+
+
+def test_from_axis_angle_consistency_with_matrix():
+    """The rotation matrix of from_axis_angle is Rodrigues' formula."""
+    axis = jnp.array([1.0, -2.0, 0.5])
+    angle = 0.7
+    n = axis / jnp.linalg.norm(axis)
+    k = jnp.array([[0, -n[2], n[1]], [n[2], 0, -n[0]], [-n[1], n[0], 0]])
+    expected = jnp.eye(3) + jnp.sin(angle) * k + (1 - jnp.cos(angle)) * k @ k
+
+    q = Quaternion.from_axis_angle(axis, angle)
+    assert jnp.allclose(q.to_rotation_matrix(), expected, atol=1e-6)
+
+
+def test_from_axis_angle_broadcast():
+    """Axes and angles broadcast against each other."""
+    axes = jnp.eye(3)  # (3, 3)
+    angles = jnp.array([[0.1], [0.2]])  # (2, 1)
+    q = Quaternion.from_axis_angle(axes, angles)
+    assert q.shape == (2, 3)
+    assert jnp.allclose(abs(q), 1.0, atol=1e-6)
+    expected = Quaternion.from_axis_angle(axes[2], angles[1, 0])
+    assert jnp.allclose(q[1, 2].wxyz, expected.wxyz, atol=1e-6)
+
+
+def test_from_axis_angle_grad_at_zero():
+    """The derivative with respect to the angle is finite and exact at zero."""
+
+    def func(angle):
+        return Quaternion.from_axis_angle(jnp.array([0.0, 1.0, 0.0]), angle).wxyz
+
+    jac = jax.jacfwd(func)(0.0)
+    assert jnp.allclose(jac, jnp.array([0.0, 0.0, 0.5, 0.0]))
+
+
+def test_from_axis_angle_wrong_shape():
+    with pytest.raises(ValueError, match='Axis must have shape'):
+        Quaternion.from_axis_angle(jnp.array([1.0, 0.0]), 0.1)
+
+
 # to_rotation_matrix
 @pytest.mark.parametrize('do_jit', [False, True])
 def test_to_rotation_matrix_identity(do_jit):

{fastquat-1.0b2 → fastquat-1.0b4}/uv.lock

@@ -666,6 +666,7 @@ dev = [
     { name = "numpy", version = "2.2.6", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.11'" },
     { name = "numpy", version = "2.4.2", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.11'" },
     { name = "pillow" },
+    { name = "prek" },
     { name = "pytest" },
     { name = "pytest-notebook" },
     { name = "sphinx", version = "8.1.3", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.11'" },
@@ -688,6 +689,9 @@ docs = [
     { name = "sphinx", version = "9.1.0", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.12'" },
     { name = "sphinx-rtd-theme" },
 ]
+lint = [
+    { name = "prek" },
+]
 
 [package.metadata]
 requires-dist = [
@@ -705,6 +709,7 @@ dev = [
     { name = "nbsphinx", specifier = ">=0.9.0" },
     { name = "numpy", specifier = ">=1.24.0" },
     { name = "pillow", specifier = ">=9.0.0" },
+    { name = "prek", specifier = ">=0.5.4" },
     { name = "pytest", specifier = ">=9.0" },
     { name = "pytest-notebook", specifier = ">=0.10" },
     { name = "sphinx", specifier = ">=7.0.0" },
@@ -721,6 +726,7 @@ docs = [
     { name = "sphinx", specifier = ">=7.0.0" },
     { name = "sphinx-rtd-theme", specifier = ">=1.3.0" },
 ]
+lint = [{ name = "prek", specifier = ">=0.5.4" }]
 
 [[package]]
 name = "fonttools"
@@ -2778,6 +2784,22 @@ wheels = [
     { url = "https://files.pythonhosted.org/packages/54/20/4d324d65cc6d9205fabedc306948156824eb9f0ee1633355a8f7ec5c66bf/pluggy-1.6.0-py3-none-any.whl", hash = "sha256:e920276dd6813095e9377c0bc5566d94c932c33b27a3e3945d8389c374dd4746", size = 20538, upload-time = "2025-05-15T12:30:06.134Z" },
 ]
 
+[[package]]
+name = "prek"
+version = "0.5.4"
+source = { registry = "https://pypi.org/simple" }
+sdist = { url = "https://files.pythonhosted.org/packages/30/f1/8f3143530bf43fa82f8873f4e536f318b6143dd5c09171aba0d465e21f04/prek-0.5.4.tar.gz", hash = "sha256:7f5bc061880141e362ef422e52188d0090794c0622e442e3157686d1e342780a", size = 565182, upload-time = "2026-09-28T04:54:43.453Z" }
+wheels = [
+    { url = "https://files.pythonhosted.org/packages/78/a7/f3a4a252c0846c5846dd86118ee0c1fddb7381dc35e3d44c2a5f3cea1757/prek-0.5.4-py3-none-macosx_10_12_x86_64.whl", hash = "sha256:550ecaae0601f9405526c866a57cc6882ed1e9d8696ac1117b8703bb37d087eb", size = 6011113, upload-time = "2026-09-28T04:54:28.107Z" },
+    { url = "https://files.pythonhosted.org/packages/1b/e3/26638cf090b620e63e60a639b8e7d7eb38b9fbea74a0827a182fc922cf8e/prek-0.5.4-py3-none-macosx_11_0_arm64.whl", hash = "sha256:a7ef67cc23e23c6e9c6ca0b27a7b12c19f16774d36b1c88ed035f2639995718d", size = 5534338, upload-time = "2026-09-28T04:54:30.214Z" },
+    { url = "https://files.pythonhosted.org/packages/a0/08/64b668eb088204650b96183cea739b56e6b179130c4398a8a83cbe9b125a/prek-0.5.4-py3-none-manylinux_2_17_aarch64.manylinux2014_aarch64.musllinux_1_1_aarch64.whl", hash = "sha256:59f2206850f643bafe46b28b4431525f8dcc5b7a763f271179bff10b8157870d", size = 5840937, upload-time = "2026-09-28T04:54:32.226Z" },
+    { url = "https://files.pythonhosted.org/packages/1f/ee/a0f847194eb6a2b67921b4c94c4fbf3f33a785affb190c6a2097344dbdd6/prek-0.5.4-py3-none-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:7d5b0cab1219a7bc9cc90785c7436f563765440616543ede7dfa3ba5c332cede", size = 6234354, upload-time = "2026-09-28T04:54:34.176Z" },
+    { url = "https://files.pythonhosted.org/packages/32/9a/14186d90ed29b46650452068d6bd1ac233ca0ea41a3b35081c37f6b77b1c/prek-0.5.4-py3-none-manylinux_2_28_aarch64.whl", hash = "sha256:8fc22db8132dab9aed4e152bf6a58a7f791f4c91c74305d65d96c48b8ec0bc62", size = 5845649, upload-time = "2026-09-28T04:54:36.252Z" },
+    { url = "https://files.pythonhosted.org/packages/e0/88/ff5bbbada4098f59740ab767f5c2f82d4110b90ea7db309a9bcf21f3dcb9/prek-0.5.4-py3-none-musllinux_1_1_x86_64.whl", hash = "sha256:19d4f3986dcd1dfa15f9417e32a75fb2ba084156dbd7ace161f0764281b8064b", size = 6351516, upload-time = "2026-09-28T04:54:38.539Z" },
+    { url = "https://files.pythonhosted.org/packages/c8/1a/0cc410232ef42750b3bf34babd6def4a3943ec98259d9fecda3b26fcd63f/prek-0.5.4-py3-none-win_amd64.whl", hash = "sha256:3100233ba802eff6a4b09bc045c8f50d599bb9d810fa234823bc1642938b224f", size = 5755956, upload-time = "2026-09-28T04:54:40.227Z" },
+    { url = "https://files.pythonhosted.org/packages/98/78/3fd53cc54468db37462885d3139eae40bb469d51beae4575c986b190ab6e/prek-0.5.4-py3-none-win_arm64.whl", hash = "sha256:9fc4ad2d647f70a8a3a9900849e1dd89472cf5a71d830f083b62cf87a0406e27", size = 5502478, upload-time = "2026-09-28T04:54:41.973Z" },
+]
+
 [[package]]
 name = "prometheus-client"
 version = "0.24.1"

fastquat-1.0b2/.github/workflows/ci.yml

@@ -1,67 +0,0 @@
-name: CI
-
-on:
-  push:
-    branches:
-    - main
-  pull_request:
-    branches:
-    - main
-
-concurrency:
-  group: ${{ github.workflow }}-${{ github.ref }}
-  cancel-in-progress: true
-
-jobs:
-
-  pre-commit:
-    name: Pre-commit
-    runs-on: ubuntu-latest
-    steps:
-    - uses: actions/checkout@v4
-    - name: Set up Python
-      uses: actions/setup-python@v5
-      with:
-        python-version: '3.x'
-    - name: Install dependencies
-      run: |
-        python -m pip install --upgrade pip
-        python -m pip install pre-commit
-    - name: Run pre-commit
-      run: pre-commit run --all-files --show-diff-on-failure --color=always
-
-  test:
-    name: Run tests on Python ${{ matrix.python-version }}
-    runs-on: ubuntu-latest
-    strategy:
-      matrix:
-        python-version: ['3.10', '3.11', '3.12', '3.13', '3.14']
-
-    steps:
-      - name: Checkout code
-        uses: actions/checkout@v4
-        with:
-          fetch-depth: 0
-
-      - name: Set up Python ${{ matrix.python-version }}
-        uses: actions/setup-python@v5
-        with:
-          python-version: ${{ matrix.python-version }}
-
-      - name: Install uv
-        uses: astral-sh/setup-uv@v6
-        with:
-          enable-cache: true
-
-      - name: Build wheel
-        run: uv build --wheel
-
-      - name: Install project's wheel and dependencies
-        run: |
-          uv sync --locked --no-install-project
-          uv pip install dist/fastquat-*.whl
-
-      - name: Run tests
-        run: |
-          source .venv/bin/activate
-          pytest -v