# nanomanifold

Fast, batched and differentiable SO(3)/SE(3) transforms for any backend (NumPy, PyTorch, JAX, ...)

Works directly on arrays, defined as:

- **SO(3)**: unit quaternions `[w, x, y, z]` for 3D rotations, shape `(..., 4)`
- **SE(3)**: concatenated `[quat, translation]`, shape `(..., 7)`

```python
import numpy as np
from nanomanifold import SO3, SE3

# Rotations stored as quaternion arrays [w,x,y,z]
q = SO3.from_axis_angle(np.array([0, 0, 1]), np.pi/4)  # 45° around Z
points = np.array([[1, 0, 0], [0, 1, 0]])
rotated = SO3.rotate_points(q, points)

# Rigid transforms stored as 7D arrays [quat, translation]
T = SE3.from_rt(q, np.array([1, 0, 0]))  # rotation + translation
transformed = SE3.transform_points(T, points)
```

## Installation

```bash
pip install nanomanifold
```

## Quick Start

### Rotations (SO3)

```python
from nanomanifold import SO3

# Create rotations
q1 = SO3.from_axis_angle([1, 0, 0], np.pi/2)    # 90° around X
q2 = SO3.from_euler([0, 0, np.pi/4])            # 45° around Z
q3 = SO3.from_matrix(rotation_matrix)

# Compose and interpolate
q_combined = SO3.multiply(q1, q2)
q_halfway = SO3.slerp(q1, q2, t=0.5)

# Apply to points
points = np.array([[1, 0, 0], [0, 1, 0]])
rotated = SO3.rotate_points(q_combined, points)
```

### Rigid Transforms (SE3)

```python
from nanomanifold import SE3

# Create transforms
T1 = SE3.from_rt(q1, [1, 2, 3])               # rotation + translation
T2 = SE3.from_matrix(transformation_matrix)

# Compose and interpolate
T_combined = SE3.multiply(T1, T2)
T_inverse = SE3.inverse(T_combined)
T_halfway = SE3.slerp(T1, T2, t=0.5)

# Apply to points
transformed = SE3.transform_points(T_combined, points)
```

## API Reference

All functions are available via `nanomanifold.SO3` and `nanomanifold.SE3`. Shapes follow the
Array API convention and accept arbitrarily batched inputs.

### SO3 (3D Rotations)

| Function                              | Signature                                 |
| ------------------------------------- | ----------------------------------------- |
| `canonicalize(q)`                     | `(...,4) -> (...,4)`                      |
| `to_axis_angle(q)`                    | `(...,4) -> (...,3)`                      |
| `from_axis_angle(axis_angle)`         | `(...,3) -> (...,4)`                      |
| `to_euler(q, convention="ZYX")`       | `(...,4) -> (...,3)`                      |
| `from_euler(euler, convention="ZYX")` | `(...,3) -> (...,4)`                      |
| `to_matrix(q)`                        | `(...,4) -> (...,3,3)`                    |
| `from_matrix(R)`                      | `(...,3,3) -> (...,4)`                    |
| `from_quat_xyzw(quat)`                | `(...,4) -> (...,4)`                      |
| `to_quat_xyzw(quat)`                  | `(...,4) -> (...,4)`                      |
| `to_6d(q)`                            | `(...,4) -> (...,6)`                      |
| `from_6d(d6)`                         | `(...,6) -> (...,4)`                      |
| `multiply(q1, q2)`                    | `(...,4), (...,4) -> (...,4)`             |
| `inverse(q)`                          | `(...,4) -> (...,4)`                      |
| `rotate_points(q, points)`            | `(...,4), (...,N,3) -> (...,N,3)`         |
| `slerp(q1, q2, t)`                    | `(...,4), (...,4), (...,N) -> (...,N,4)`  |
| `distance(q1, q2)`                    | `(...,4), (...,4) -> (...)`               |
| `log(q)`                              | `(...,4) -> (...,3)`                      |
| `exp(tangent)`                        | `(...,3) -> (...,4)`                      |
| `hat(w)`                              | `(...,3) -> (...,3,3)`                    |
| `vee(W)`                              | `(...,3,3) -> (...,3)`                    |
| `weighted_mean(quats, weights)`       | `sequence of (...,4), (...,N) -> (...,4)` |
| `mean(quats)`                         | `sequence of (...,4) -> (...,4)`          |
| `random(*shape)`                      | `(...,4)`                                 |

### SE3 (Rigid Transforms)

| Function                              | Signature                                 |
| ------------------------------------- | ----------------------------------------- |
| `canonicalize(se3)`                   | `(...,7) -> (...,7)`                      |
| `from_rt(quat, translation)`          | `(...,4), (...,3) -> (...,7)`             |
| `to_rt(se3)`                          | `(...,7) -> (quat, translation)`          |
| `from_matrix(T)`                      | `(...,4,4) -> (...,7)`                    |
| `to_matrix(se3)`                      | `(...,7) -> (...,4,4)`                    |
| `multiply(se3_1, se3_2)`              | `(...,7), (...,7) -> (...,7)`             |
| `inverse(se3)`                        | `(...,7) -> (...,7)`                      |
| `transform_points(se3, points)`       | `(...,7), (...,N,3) -> (...,N,3)`         |
| `slerp(se3_1, se3_2, t)`             | `(...,7), (...,7), (...,N) -> (...,N,7)`  |
| `log(se3)`                            | `(...,7) -> (...,6)`                      |
| `exp(tangent)`                        | `(...,6) -> (...,7)`                      |
| `hat(v)`                              | `(...,6) -> (...,4,4)`                    |
| `vee(M)`                              | `(...,4,4) -> (...,6)`                    |
| `weighted_mean(transforms, weights)`  | `sequence of (...,7), (...,N) -> (...,7)` |
| `mean(transforms)`                    | `sequence of (...,7) -> (...,7)`          |
| `random(*shape)`                      | `(...,7)`                                 |

## Pairwise Conversions (`SO3.conversions`)

Convert directly between any two rotation representations without going through
quaternions manually. All 30 pairwise functions follow the naming pattern
`from_{source}_to_{target}`.

Representations: `axis_angle`, `euler`, `matrix`, `quat_wxyz`, `quat_xyzw`, `sixd`.

| Function                                                        | Signature                   |
| --------------------------------------------------------------- | --------------------------- |
| `SO3.conversions.from_axis_angle_to_matrix(aa)`                 | `(...,3) -> (...,3,3)`      |
| `SO3.conversions.from_axis_angle_to_euler(aa, convention)`      | `(...,3) -> (...,3)`        |
| `SO3.conversions.from_axis_angle_to_quat_wxyz(aa)`              | `(...,3) -> (...,4)`        |
| `SO3.conversions.from_axis_angle_to_quat_xyzw(aa)`              | `(...,3) -> (...,4)`        |
| `SO3.conversions.from_axis_angle_to_sixd(aa)`                   | `(...,3) -> (...,6)`        |
| `SO3.conversions.from_euler_to_axis_angle(e, convention)`       | `(...,3) -> (...,3)`        |
| `SO3.conversions.from_euler_to_matrix(e, convention)`           | `(...,3) -> (...,3,3)`      |
| `SO3.conversions.from_euler_to_quat_wxyz(e, convention)`        | `(...,3) -> (...,4)`        |
| `SO3.conversions.from_euler_to_quat_xyzw(e, convention)`        | `(...,3) -> (...,4)`        |
| `SO3.conversions.from_euler_to_sixd(e, convention)`             | `(...,3) -> (...,6)`        |
| `SO3.conversions.from_matrix_to_axis_angle(R)`                  | `(...,3,3) -> (...,3)`      |
| `SO3.conversions.from_matrix_to_euler(R, convention)`           | `(...,3,3) -> (...,3)`      |
| `SO3.conversions.from_matrix_to_quat_wxyz(R)`                   | `(...,3,3) -> (...,4)`      |
| `SO3.conversions.from_matrix_to_quat_xyzw(R)`                   | `(...,3,3) -> (...,4)`      |
| `SO3.conversions.from_matrix_to_sixd(R)`                        | `(...,3,3) -> (...,6)`      |
| `SO3.conversions.from_quat_wxyz_to_axis_angle(q)`               | `(...,4) -> (...,3)`        |
| `SO3.conversions.from_quat_wxyz_to_euler(q, convention)`        | `(...,4) -> (...,3)`        |
| `SO3.conversions.from_quat_wxyz_to_matrix(q)`                   | `(...,4) -> (...,3,3)`      |
| `SO3.conversions.from_quat_wxyz_to_quat_xyzw(q)`               | `(...,4) -> (...,4)`        |
| `SO3.conversions.from_quat_wxyz_to_sixd(q)`                     | `(...,4) -> (...,6)`        |
| `SO3.conversions.from_quat_xyzw_to_axis_angle(q)`               | `(...,4) -> (...,3)`        |
| `SO3.conversions.from_quat_xyzw_to_euler(q, convention)`        | `(...,4) -> (...,3)`        |
| `SO3.conversions.from_quat_xyzw_to_matrix(q)`                   | `(...,4) -> (...,3,3)`      |
| `SO3.conversions.from_quat_xyzw_to_quat_wxyz(q)`               | `(...,4) -> (...,4)`        |
| `SO3.conversions.from_quat_xyzw_to_sixd(q)`                     | `(...,4) -> (...,6)`        |
| `SO3.conversions.from_sixd_to_axis_angle(d6)`                   | `(...,6) -> (...,3)`        |
| `SO3.conversions.from_sixd_to_euler(d6, convention)`            | `(...,6) -> (...,3)`        |
| `SO3.conversions.from_sixd_to_matrix(d6)`                       | `(...,6) -> (...,3,3)`      |
| `SO3.conversions.from_sixd_to_quat_wxyz(d6)`                    | `(...,6) -> (...,4)`        |
| `SO3.conversions.from_sixd_to_quat_xyzw(d6)`                    | `(...,6) -> (...,4)`        |

## Backend-Explicit Mode

By default, nanomanifold auto-detects the array backend via `array_api_compat`. Every function also
accepts an optional `xp` keyword argument to specify the backend explicitly. This is required for
`torch.compile(fullgraph=True)`, since Dynamo cannot trace the dynamic dispatch:

```python
import torch
from nanomanifold import SO3, SE3

@torch.compile(fullgraph=True)
def forward(q1, q2, T1, T2):
    q_mid = SO3.slerp(q1, q2, torch.tensor([0.5]), xp=torch)
    T_mid = SE3.slerp(T1, T2, torch.tensor([0.5]), xp=torch)
    return q_mid, T_mid
```
