# Mesh Alignment

> NumKong for Python

### Many-to-Many Shape Distances

Point clouds are packed once each, centered on their centroids and rotated onto their principal axes.
The packed sets then score every row cloud against every column cloud, and the clouds may differ in size.
Chamfer and Hausdorff compare nearest partners under the 8 sign flips of the principal axes.
Wasserstein-Procrustes solves for the exact transport plan and the best rotation or reflection together.
On equal sizes it is at most the Kabsch RMSD of a known correspondence, and equals it when that correspondence is an optimal plan.
Gromov-Wasserstein compares the clouds' own distance matrices.

```python
import numpy as np
import numkong as nk

molecules = [np.random.randn(atoms, 3) for atoms in (20, 31, 60)]
rows = nk.shapes_pack(molecules)  # float64 input packs as float32
columns = nk.shapes_pack(molecules[1:])
assert rows.shape == (3, 111) and columns.shape == (2, 91)  # clouds and total points

rotations = np.empty((3, 2, 3, 3), dtype=np.float32)
distances = nk.wassersteinprocrustes_packed(rows, columns, rotations=rotations)
assert np.asarray(distances).shape == (3, 2)
assert np.asarray(distances)[1, 0] < 1e-6, "a molecule against itself"

chamfers = nk.chamfers_packed(rows[:1], columns)  # slices are views over the same packs
assert np.asarray(chamfers).dtype == np.float64
assert np.allclose(np.asarray(rows.unpack(1)), molecules[1], atol=1e-5)  # the inverse of packing

grown = nk.shapes_pack_normalized([molecules[0] * 3.7 + 1.0])  # divided by its RMS radius
originals = nk.shapes_pack_normalized(molecules)
assert np.asarray(nk.wassersteinprocrustes_packed(originals, grown))[0, 0] < 1e-5
```

Each rotation takes column cloud `b` onto row cloud `a` through the centroids: `a ≈ (b - b.mean(0)) @ R.T + a.mean(0)`.
The optional `rotations=` output is a writable C-contiguous `float32` buffer of shape `[rows, columns, 3, 3]`, and the functions always return the distances.
`nk.shapes_pack_normalized` also divides every cloud by its RMS distance from the centroid, so every kind ignores size as well as pose.
Between normalized packs, the rotation also scales by `scale_a / scale_b`, and each kind compares the packed coordinates, so mixing absolute and normalized packs in one call is the caller's choice.
The pack layout is private, and the `centroids`, `frames`, `scales`, `spreads` and `max_radii` properties report every cloud's pose and bounds.
`unpack(index)` returns one cloud as a `[points, 3]` tensor in the pack dtype and its original pose and scale, within one storage rounding.
`nk.ShapesPackedMatrix.pack_size(points, dtype)` returns the bytes one pack takes.
Chamfer and Hausdorff take `limit=`: distances up to it are exact, and a pair past it may stop early and report a lower bound that also lies past it.
Chamfer, Wasserstein-Procrustes and Gromov-Wasserstein allocate their workspace inside the call, or take a writable `workspace=` of `nk.chamfers_workspace_size(a, b, threads=threads)` bytes, with its `wassersteinprocrustes` and `gromovwassersteins` siblings.
`nk.rmsds_packed(a, b)` scores clouds of equally many points, point i matching point i under the best rotation.
Every kind takes `out=`, a writable `[rows, columns]` buffer of its output dtype.
Every kind takes `threads=`, splitting the row clouds with a workspace per worker, and both packing functions take it to split the clouds.

Packs need no GPU kernel: the host gathers them, and `copy` places the table and every pack in one block of unified memory on the GPU the mask names.
Every shape function then runs that mask in one launch on `stream=`, ignoring `threads=`, so it needs `out=` and, where the workspace size is not zero, a single window of `workspace=`.

```python
gpu = nk.cuda_capabilities_enabled(0)
packs = nk.shapes_pack(clouds, dtype="f16", threads=8).copy(capabilities=gpu)
distances = nk.from_dlpack(torch.empty((n, n), device="cuda"))
nk.chamfers_packed(packs, packs, out=distances, limit=2.5, stream=stream)
nk.synchronize(gpu, stream=stream)
```

A slice of copied packs in steps would rebuild the table on the host, so it raises `BufferError` on the GPU until it is copied itself.
Metal runs `float16` packs only, and writes only into blocks of its own allocator, so a plain NumPy `out=` raises `BufferError`:

```python
gpu = nk.metal_capabilities_enabled(0)
packs = nk.shapes_pack(clouds, dtype="f16").copy(gpu)
distances = nk.zeros((n, n), dtype="float32").copy(capabilities=gpu)
nk.chamfers_packed(packs, packs, out=distances, limit=2.5)
nk.synchronize(gpu)
```

`import numkong`

`from numkong import ShapesPackedMatrix`

## ShapesPackedMatrix

Point clouds packed by shapes_pack() or shapes_pack_normalized(), centered and rotated onto principal axes.

```python
class ShapesPackedMatrix
```

Indexing and slicing return views sharing the packs, which chamfers_packed() and its siblings
score as rows or columns. The pack layout is private: the properties report every cloud's pose
and bounds, and unpack() returns any cloud in its original pose.

## pack_size

Return the bytes one packed cloud of the given points and dtype takes.

```python
@classmethod
def pack_size(
    points: int,
    /,
    dtype: Literal['f16', 'float16', 'f32', 'float32', 'f64', 'float64'] | type[float16],
) -> int
```

## unpack

Return one cloud as [points,3] in the pack dtype and its original pose and scale, within one rounding.

```python
def unpack(index: int, /) -> Tensor
```

## copy

Gather the table and packs into one block of the unified memory of `capabilities`, for GPU shape functions.

```python
def copy(
    capabilities: Capability | int,
    *,
    stream: int | None = None,
) -> ShapesPackedMatrix
```

## __len__

Return the number of packed clouds.

```python
def __len__() -> int
```

## __getitem__

Return a view over one cloud or a slice of them, sharing the packs.

```python
def __getitem__(key: int | slice, /) -> ShapesPackedMatrix
```

## __repr__

Return a string representation.

```python
def __repr__() -> str
```

## clouds

Number of packed clouds.

```python
clouds: int
```

## points

Total points across the packed clouds.

```python
points: int
```

## shape

Dimensions (clouds, points) of the packed set.

```python
shape: tuple[int, int]
```

## centroids

Mean input point of every cloud, [clouds,3], float32 for float16 packs and float64 for float32 ones.

```python
centroids: Tensor
```

## frames

Unit quaternions (x, y, z, w) of Q rotating every centered cloud onto its principal axes, [clouds,4].

```python
frames: Tensor
```

## spreads

Descending eigenvalues of every packed cloud's second moments, [clouds,3].

```python
spreads: Tensor
```

## scales

What every rotated cloud was divided by, its RMS distance from the centroid or 1, [clouds].

```python
scales: Tensor
```

## max_radii

Largest packed point norm of every cloud, [clouds].

```python
max_radii: Tensor
```

## dtype

Data type of the packed points.

```python
dtype: Literal['float16', 'float32']
```

## nbytes

Size of the packs in bytes.

```python
nbytes: int
```

## shapes_pack

Center every cloud and rotate it onto its principal axes, float16 input as float16, else float32.

```python
def shapes_pack(
    points: _BufferType | list[_BufferType] | tuple[_BufferType, ...],
    /,
    dtype: Literal['f16', 'float16', 'f32', 'float32', 'f64', 'float64'] | type[float16] | None = None,
    *,
    threads: int = 1,
    capabilities: Capability | int | None = ...,
    stream: int | None = ...,
) -> ShapesPackedMatrix
```

- `capabilities`: Capabilities to run within: every CPU one by default, or for packed operands the mask that packed them.
- `stream`: A GPU stream pointer, left None on the CPU.

## shapes_pack_normalized

Pack every cloud like shapes_pack() and divide it by its RMS radius, so every kind ignores its size.

```python
def shapes_pack_normalized(
    points: _BufferType | list[_BufferType] | tuple[_BufferType, ...],
    /,
    dtype: Literal['f16', 'float16', 'f32', 'float32', 'f64', 'float64'] | type[float16] | None = None,
    *,
    threads: int = 1,
    capabilities: Capability | int | None = ...,
    stream: int | None = ...,
) -> ShapesPackedMatrix
```

- `capabilities`: Capabilities to run within: every CPU one by default, or for packed operands the mask that packed them.
- `stream`: A GPU stream pointer, left None on the CPU.

## rmsds_packed

Compute superposed RMSDs of clouds with equally many points, point i matching point i.

```python
def rmsds_packed(
    a: ShapesPackedMatrix,
    b: ShapesPackedMatrix,
    /,
    *,
    out: _BufferType | None = None,
    threads: int = 1,
    capabilities: Capability | int | None = ...,
    stream: int | None = ...,
) -> Tensor
```

- `capabilities`: Capabilities to run within: every CPU one by default, or for packed operands the mask that packed them.
- `stream`: A GPU stream pointer, left None on the CPU.

## hausdorffs_packed

Compute Hausdorff distances, the worst nearest-partner distance either way, best of 8 axis flips.

```python
def hausdorffs_packed(
    a: ShapesPackedMatrix,
    b: ShapesPackedMatrix,
    /,
    *,
    rotations: _BufferType | None = None,
    limit: float | None = None,
    out: _BufferType | None = None,
    threads: int = 1,
    capabilities: Capability | int | None = ...,
    stream: int | None = ...,
) -> Tensor
```

- `capabilities`: Capabilities to run within: every CPU one by default, or for packed operands the mask that packed them.
- `stream`: A GPU stream pointer, left None on the CPU.

## chamfers_packed

Compute Chamfer distances, the mean nearest-partner distance both ways, halved, best of 8 axis flips.

```python
def chamfers_packed(
    a: ShapesPackedMatrix,
    b: ShapesPackedMatrix,
    /,
    *,
    rotations: _BufferType | None = None,
    limit: float | None = None,
    workspace: _BufferType | None = None,
    out: _BufferType | None = None,
    threads: int = 1,
    capabilities: Capability | int | None = ...,
    stream: int | None = ...,
) -> Tensor
```

- `capabilities`: Capabilities to run within: every CPU one by default, or for packed operands the mask that packed them.
- `stream`: A GPU stream pointer, left None on the CPU.

## wassersteinprocrustes_packed

Compute Wasserstein-Procrustes distances, the exact transport RMSD after the best rotation or reflection.

```python
def wassersteinprocrustes_packed(
    a: ShapesPackedMatrix,
    b: ShapesPackedMatrix,
    /,
    *,
    rotations: _BufferType | None = None,
    workspace: _BufferType | None = None,
    out: _BufferType | None = None,
    threads: int = 1,
    capabilities: Capability | int | None = ...,
    stream: int | None = ...,
) -> Tensor
```

- `capabilities`: Capabilities to run within: every CPU one by default, or for packed operands the mask that packed them.
- `stream`: A GPU stream pointer, left None on the CPU.

## gromovwassersteins_packed

Compute Gromov-Wasserstein distances between the clouds' own distance matrices.

```python
def gromovwassersteins_packed(
    a: ShapesPackedMatrix,
    b: ShapesPackedMatrix,
    /,
    *,
    rotations: _BufferType | None = None,
    workspace: _BufferType | None = None,
    out: _BufferType | None = None,
    threads: int = 1,
    capabilities: Capability | int | None = ...,
    stream: int | None = ...,
) -> Tensor
```

- `capabilities`: Capabilities to run within: every CPU one by default, or for packed operands the mask that packed them.
- `stream`: A GPU stream pointer, left None on the CPU.

## chamfers_workspace_size

Return the bytes of the `workspace` that chamfers_packed() needs, any alignment included.

```python
def chamfers_workspace_size(
    a: ShapesPackedMatrix,
    b: ShapesPackedMatrix,
    /,
    *,
    threads: int = 1,
    capabilities: Capability | int | None = ...,
    stream: int | None = ...,
) -> int
```

- `capabilities`: Capabilities to run within: every CPU one by default, or for packed operands the mask that packed them.
- `stream`: A GPU stream pointer, left None on the CPU.

## wassersteinprocrustes_workspace_size

Return the bytes of the `workspace` that wassersteinprocrustes_packed() needs.

```python
def wassersteinprocrustes_workspace_size(
    a: ShapesPackedMatrix,
    b: ShapesPackedMatrix,
    /,
    *,
    threads: int = 1,
    capabilities: Capability | int | None = ...,
    stream: int | None = ...,
) -> int
```

- `capabilities`: Capabilities to run within: every CPU one by default, or for packed operands the mask that packed them.
- `stream`: A GPU stream pointer, left None on the CPU.

## gromovwassersteins_workspace_size

Return the bytes of the `workspace` that gromovwassersteins_packed() needs.

```python
def gromovwassersteins_workspace_size(
    a: ShapesPackedMatrix,
    b: ShapesPackedMatrix,
    /,
    *,
    threads: int = 1,
    capabilities: Capability | int | None = ...,
    stream: int | None = ...,
) -> int
```

- `capabilities`: Capabilities to run within: every CPU one by default, or for packed operands the mask that packed them.
- `stream`: A GPU stream pointer, left None on the CPU.