quat-numpy 0.2.0__py3-none-any.whl

quat/__init__.py

@@ -0,0 +1,102 @@
+# =========================================================================
+#  LLM-generated code.  See README.md for full disclosure.
+# =========================================================================
+
+"""Quaternion Algebra Library — quat package.
+
+Provides:
+  Quaternion  — single quaternion value
+  QuatVector  — 1-d collection of quaternions
+  QuatMatrix  — 2-d quaternion matrix
+  QuatTensor  — 3-d quaternion tensor
+
+  quat()      — convenience constructor
+  dict_to_quat_matrix, dict_to_quat_tensor, labels_to_quat_vector
+
+Algebra primitives (from quat.algebra):
+  _hamilton, _CONJ, _REAL_LEFT
+
+Basis constants (from quat.core):
+  _I, _J, _K, _ZERO, _R, _ONE_Q
+
+Utilities (from quat.utils):
+  to_ndarray, from_ndarray, from_components, broadcast_quat_shapes, stack_quat,
+  isnan, isinf, isfinite, isclose
+
+Serialization (from quat.serialization):
+  to_json, from_json, to_bytes, from_bytes,
+  to_scipy_rotation, from_scipy_rotation
+
+Optimized (from quat.optimized):
+  hamilton_einsum, quat_matmul, conjugate_batch,
+  norm_squared_batch, normalize_batch
+
+Linear algebra (from quat.linalg):
+  svd, rank, condition_number, pseudo_inverse,
+  trace, det, norm, solve
+
+Random (from quat.random):
+  random_quat, random_unit_quat, random_quat_vector,
+  random_quat_matrix, random_quat_tensor
+
+Interpolation (from quat.interpolate):
+  slerp, slerp_vector, squad
+
+Signal processing (from quat.signal):
+  qfft, iqfft, qfft2, iqfft2,
+  qconv, qconv2,
+  lowpass, highpass, bandpass, bandstop
+"""
+
+__version__ = "0.2.0"
+from quat.algebra import _hamilton, _CONJ, _REAL_LEFT
+from quat.core import Quaternion, quat, _I, _J, _K, _ZERO, _R, _ONE_Q
+from quat.collections import (
+    QuatVector, QuatMatrix, QuatTensor,
+    dict_to_quat_matrix, dict_to_quat_tensor, labels_to_quat_vector,
+)
+from quat.utils import to_ndarray, from_ndarray, from_components, broadcast_quat_shapes, stack_quat
+from quat.utils import isnan, isinf, isfinite, isclose
+from quat.serialization import (
+    to_json, from_json, to_bytes, from_bytes,
+    to_scipy_rotation, from_scipy_rotation,
+)
+from quat.optimized import (
+    hamilton_einsum, quat_matmul,
+    conjugate_batch, norm_squared_batch, normalize_batch,
+)
+from quat.linalg import (
+    svd, rank, condition_number, pseudo_inverse,
+    trace, det, norm, solve,
+)
+from quat.random import (
+    random_quat, random_unit_quat, random_quat_vector,
+    random_quat_matrix, random_quat_tensor,
+)
+from quat.interpolate import slerp, slerp_vector, squad
+from quat.signal import (
+    qfft, iqfft, qfft2, iqfft2,
+    qconv, qconv2,
+    lowpass, highpass, bandpass, bandstop,
+)
+
+__all__ = [
+    'Quaternion', 'QuatVector', 'QuatMatrix', 'QuatTensor',
+    'quat', 'dict_to_quat_matrix', 'dict_to_quat_tensor', 'labels_to_quat_vector',
+    '_hamilton', '_CONJ', '_REAL_LEFT', '_I', '_J', '_K', '_ZERO', '_R', '_ONE_Q',
+    'to_ndarray', 'from_ndarray', 'from_components', 'broadcast_quat_shapes', 'stack_quat',
+    'isnan', 'isinf', 'isfinite', 'isclose',
+    'to_json', 'from_json', 'to_bytes', 'from_bytes',
+    'to_scipy_rotation', 'from_scipy_rotation',
+    'hamilton_einsum', 'quat_matmul',
+    'conjugate_batch', 'norm_squared_batch', 'normalize_batch',
+    'svd', 'rank', 'condition_number', 'pseudo_inverse',
+    'trace', 'det', 'norm', 'solve',
+    'random_quat', 'random_unit_quat', 'random_quat_vector',
+    'random_quat_matrix', 'random_quat_tensor',
+    'slerp', 'slerp_vector', 'squad',
+    'qfft', 'iqfft', 'qfft2', 'iqfft2',
+    'qconv', 'qconv2',
+    'lowpass', 'highpass', 'bandpass', 'bandstop',
+    '__version__',
+]

quat/_arrayops.py

@@ -0,0 +1,56 @@
+"""NumPy interop helpers for quaternion collection types.
+
+Provides shared implementations of ``data``, ``to_array``, ``to_numpy``,
+``__array__``, and ``__array_ufunc__`` so QuatVector, QuatMatrix, and
+QuatTensor share a single code path for these operations.
+"""
+from __future__ import annotations
+import numpy as np
+
+
+def _data_copy(data: np.ndarray) -> np.ndarray:
+    """Return a defensive copy of the underlying quaternion data."""
+    return data.copy()
+
+
+def _to_numpy(data: np.ndarray, copy: bool = True,
+              dtype: np.dtype | None = None) -> np.ndarray:
+    """Export quaternion data to an ndarray, optionally with dtype conversion.
+
+    When ``copy=False`` and ``dtype=None``, returns the internal array
+    without copying — caller must not mutate it.
+    """
+    if copy is False and dtype is None:
+        return data
+    return np.array(data, dtype=dtype, copy=copy)
+
+
+def _to_array(data: np.ndarray) -> np.ndarray:
+    """Return a copy of the quaternion data as an ndarray."""
+    return data.copy()
+
+
+def _dispatch_collection_ufunc(
+    self, ufunc, method: str, *inputs, **kwargs
+) -> object:
+    """Shared ``__array_ufunc__`` dispatch for QuatVector/QuatMatrix/QuatTensor.
+
+    The *self* parameter is unused in the body (dispatch works via *inputs*)
+    but is kept for signature compatibility with NumPy's protocol.
+
+    Returns ``NotImplemented`` for unsupported ufuncs or reduction methods.
+    """
+    if method != '__call__' or kwargs.get('out') is not None:
+        return NotImplemented
+    a, b = (inputs[0], inputs[1]) if len(inputs) == 2 else (inputs[0], None)
+    if ufunc is np.add:
+        return a + b
+    if ufunc is np.subtract:
+        return a - b
+    if ufunc is np.multiply:
+        return a * b
+    if ufunc is np.true_divide or ufunc is np.floor_divide:
+        return a / b
+    if ufunc is np.negative:
+        return -a
+    return NotImplemented

quat/_checks.py

@@ -0,0 +1,30 @@
+"""Validation helpers for quaternion collection types.
+
+Each function accepts a ``(..., 4)`` ndarray of quaternion component data
+and returns a per-element boolean ndarray (or scalar for Quaternion's 1-D
+case).  Used by QuatVector, QuatMatrix, and QuatTensor to avoid duplicating
+the same numpy calls across three classes.
+"""
+from __future__ import annotations
+import numpy as np
+
+
+def _vec_isnan(data: np.ndarray) -> np.ndarray:
+    """True where any quaternion component is NaN along the last axis."""
+    return np.any(np.isnan(data), axis=-1)
+
+
+def _vec_isinf(data: np.ndarray) -> np.ndarray:
+    """True where any quaternion component is infinite along the last axis."""
+    return np.any(np.isinf(data), axis=-1)
+
+
+def _vec_isfinite(data: np.ndarray) -> np.ndarray:
+    """True where all quaternion components are finite along the last axis."""
+    return np.all(np.isfinite(data), axis=-1)
+
+
+def _vec_isclose(data: np.ndarray, other_data: np.ndarray,
+                 rtol: float, atol: float) -> np.ndarray:
+    """Element-wise closeness test — all four components must be close along the last axis."""
+    return np.isclose(data, other_data, rtol=rtol, atol=atol).all(axis=-1)

quat/_serialize.py

@@ -0,0 +1,81 @@
+"""Shared serialization primitives for quaternion types.
+
+Provides JSON and binary serialization helpers used by Quaternion,
+QuatVector, QuatMatrix, and QuatTensor.  Each type delegates its
+``to_json`` / ``from_json`` / ``to_bytes`` / ``from_bytes`` methods
+to these functions, passing only its type-specific label or id.
+"""
+from __future__ import annotations
+import json as _json
+import struct as _struct
+import numpy as np
+
+_TYPE_IDS: dict[int, str] = {0: "Quaternion", 1: "QuatVector", 2: "QuatMatrix", 3: "QuatTensor"}
+"""Mapping from binary type_id to type name string."""
+
+_TYPE_NAMES: dict[str, int] = {v: k for k, v in _TYPE_IDS.items()}
+"""Reverse mapping from type name to binary type_id."""
+
+
+def _serialize_to_json(type_name: str, data: np.ndarray) -> str:
+    """Serialize quaternion data to a JSON string.
+
+    Args:
+        type_name: one of ``"Quaternion"``, ``"QuatVector"``, etc.
+        data: ndarray of shape ``(..., 4)``.
+
+    Returns:
+        JSON string ``{"type": "<type_name>", "data": [[...], ...]}``.
+    """
+    return _json.dumps({"type": type_name, "data": data.tolist()})
+
+
+def _deserialize_from_json(s: str, cls_map: dict) -> object:
+    """Deserialize a JSON string back to a quaternion type instance.
+
+    Args:
+        s: JSON string produced by ``_serialize_to_json``.
+        cls_map: dict mapping type_name → constructor callable
+            (e.g. ``{"QuatVector": QuatVector}``).
+
+    Returns:
+        Instance of the type indicated by the JSON ``"type"`` field.
+    """
+    d = _json.loads(s)
+    return cls_map[d["type"]](np.array(d["data"], dtype=float))
+
+
+def _serialize_bytes_shaped(type_id: int, data: np.ndarray) -> bytes:
+    """Serialize quaternion data to compact binary form.
+
+    Format: ``<i`` type_id + ``<i`` ndim + int32[ndim] shape + float64[...].
+
+    Args:
+        type_id: binary type identifier (1=QuatVector, 2=QuatMatrix, 3=QuatTensor).
+        data: ndarray of shape ``(..., 4)``.
+
+    Returns:
+        Packed bytes.
+    """
+    data64 = data.astype(np.float64)
+    shape = np.array(data64.shape, dtype=np.int32)
+    return _struct.pack('<ii', type_id, len(shape)) + shape.tobytes() + data64.tobytes()
+
+
+def _deserialize_bytes_shaped(b: bytes, cls_map: dict) -> object:
+    """Deserialize binary bytes back to a quaternion collection type.
+
+    Args:
+        b: bytes produced by ``_serialize_bytes_shaped``.
+        cls_map: dict mapping type_id → constructor callable.
+
+    Returns:
+        Instance of the type indicated by the stored *type_id*.
+    """
+    type_id, ndim = _struct.unpack_from('<ii', b, 0)
+    offset = 8
+    shape = np.frombuffer(b[offset:offset + ndim * 4], dtype=np.int32)
+    offset += ndim * 4
+    size = int(np.prod(shape))
+    data = np.frombuffer(b[offset:offset + size * 8], dtype=np.float64).reshape(shape)
+    return cls_map[type_id](data)

quat/algebra.py

@@ -0,0 +1,90 @@
+# =========================================================================
+#  LLM-generated code.  See README.md for full disclosure.
+# =========================================================================
+
+"""Low-level quaternion algebra — constants, Hamilton product, real-matrix tensor."""
+from __future__ import annotations
+import numpy as np
+import numpy.typing as npt
+
+_CONJ = np.array([1., -1., -1., -1.])
+"""Conjugate mask: ``_CONJ * q`` negates the three imaginary components."""
+
+_REAL_LEFT = np.zeros((4, 4, 4))
+"""Left-regular real-representation tensor.
+
+``L(q)[r, c] = Σ_k _REAL_LEFT[r, c, k] * q[k]`` yields the 4×4 real matrix
+satisfying ``L(q) @ vec(x) = vec(q * x)``.
+"""
+_REAL_LEFT[0, 0, 0] = 1;   _REAL_LEFT[0, 1, 1] = -1;  _REAL_LEFT[0, 2, 2] = -1;  _REAL_LEFT[0, 3, 3] = -1
+_REAL_LEFT[1, 0, 1] = 1;   _REAL_LEFT[1, 1, 0] = 1;   _REAL_LEFT[1, 2, 3] = -1;  _REAL_LEFT[1, 3, 2] = 1
+_REAL_LEFT[2, 0, 2] = 1;   _REAL_LEFT[2, 1, 3] = 1;   _REAL_LEFT[2, 2, 0] = 1;   _REAL_LEFT[2, 3, 1] = -1
+_REAL_LEFT[3, 0, 3] = 1;   _REAL_LEFT[3, 1, 2] = -1;  _REAL_LEFT[3, 2, 1] = 1;   _REAL_LEFT[3, 3, 0] = 1
+
+_HAMILTON_TENSOR = np.zeros((4, 4, 4))
+_HAMILTON_TENSOR[0, 0, 0] = 1;    _HAMILTON_TENSOR[0, 1, 1] = -1;   _HAMILTON_TENSOR[0, 2, 2] = -1;   _HAMILTON_TENSOR[0, 3, 3] = -1
+_HAMILTON_TENSOR[1, 0, 1] = 1;    _HAMILTON_TENSOR[1, 1, 0] = 1;    _HAMILTON_TENSOR[1, 2, 3] = 1;    _HAMILTON_TENSOR[1, 3, 2] = -1
+_HAMILTON_TENSOR[2, 0, 2] = 1;    _HAMILTON_TENSOR[2, 1, 3] = -1;   _HAMILTON_TENSOR[2, 2, 0] = 1;    _HAMILTON_TENSOR[2, 3, 1] = 1
+_HAMILTON_TENSOR[3, 0, 3] = 1;    _HAMILTON_TENSOR[3, 1, 2] = 1;    _HAMILTON_TENSOR[3, 2, 1] = -1;   _HAMILTON_TENSOR[3, 3, 0] = 1
+
+
+_SMALL_THRESHOLD = 500
+"""Total-element threshold below which component-wise arithmetic is fastest."""
+
+_LARGE_THRESHOLD = 5000
+"""Total-element threshold above which full einsum optimisation pays off."""
+
+
+def _hamilton(p: npt.NDArray, q: npt.NDArray) -> npt.NDArray:
+    """Vectorized Hamilton (quaternion) product.
+
+    Dispatches to the optimal kernel based on data size:
+      - small (<=500 elements): component-wise arithmetic
+      - medium (500-5000): einsum without contraction-path optimisation
+      - large (>5000): einsum with full contraction-path optimisation
+
+    Supports arbitrary leading-dimension broadcasting.
+
+    Example:
+        >>> from quat.algebra import _hamilton
+        >>> import numpy as np
+        >>> p = np.array([1., 0., 0., 0.])   # real unit
+        >>> q = np.array([[1., 2., 3., 4.]]) # batch of one
+        >>> _hamilton(p, q)
+        array([[1., 2., 3., 4.]])
+        >>> i = np.array([0., 1., 0., 0.])
+        >>> j = np.array([0., 0., 1., 0.])
+        >>> _hamilton(i, j)   # i*j = k
+        array([0., 0., 0., 1.])
+    """
+    total_elements = p.size + q.size
+    if total_elements <= _SMALL_THRESHOLD:
+        return _hamilton_component(p, q)
+    if total_elements <= _LARGE_THRESHOLD:
+        return _hamilton_einsum_noopt(p, q)
+    return _hamilton_einsum(p, q)
+
+
+def _hamilton_component(p: npt.NDArray, q: npt.NDArray) -> npt.NDArray:
+    """Component-wise Hamilton product — fastest for small batches."""
+    a1, b1, c1, d1 = p[..., 0], p[..., 1], p[..., 2], p[..., 3]
+    a2, b2, c2, d2 = q[..., 0], q[..., 1], q[..., 2], q[..., 3]
+    shp = np.broadcast_shapes(p.shape[:-1], q.shape[:-1]) + (4,)
+    out = np.empty(shp)
+    out[..., 0] = a1*a2 - b1*b2 - c1*c2 - d1*d2
+    out[..., 1] = a1*b2 + b1*a2 + c1*d2 - d1*c2
+    out[..., 2] = a1*c2 - b1*d2 + c1*a2 + d1*b2
+    out[..., 3] = a1*d2 + b1*c2 - c1*b2 + d1*a2
+    return out
+
+
+def _hamilton_einsum_noopt(p: npt.NDArray, q: npt.NDArray) -> npt.NDArray:
+    """Einsum without contraction-path optimisation — faster for medium batches
+    where the optimisation overhead outweighs its benefit."""
+    return np.einsum('rck,...c,...k->...r', _HAMILTON_TENSOR, p, q, optimize=False)
+
+
+def _hamilton_einsum(p: npt.NDArray, q: npt.NDArray) -> npt.NDArray:
+    """Einsum with full contraction-path optimisation — best throughput for large
+    batches where the optimisation cost is amortised."""
+    return np.einsum('rck,...c,...k->...r', _HAMILTON_TENSOR, p, q, optimize=True)