Sequence alignment operators

.gitignore

@@ -7,9 +7,11 @@
 .coverage
 .hypothesis/
 .idea/
+.ipynb_checkpoints/
 .pytest_cache/
 .ruff_cache/
 __pycache__/
 build/
 dist/
+notebooks/
 venv/

.pre-commit-config.yaml

@@ -3,11 +3,11 @@ repos:
       - id: "check-toml"
       - id: "check-yaml"
     repo: "https://github.com/pre-commit/pre-commit-hooks"
-    rev: "v4.5.0"
+    rev: "v4.6.0"
   - hooks:
       - args:
         - "--fix"
         id: "ruff"
       - id: "ruff-format"
     repo: "https://github.com/astral-sh/ruff-pre-commit"
-    rev: "v0.3.5"
+    rev: "v0.3.7"

docs/index.md

@@ -18,6 +18,7 @@
 ::: beignet.invert_quaternion
 ::: beignet.invert_rotation_matrix
 ::: beignet.invert_rotation_vector
+::: beignet.needleman_wunsch
 ::: beignet.quaternion_identity
 ::: beignet.quaternion_magnitude
 ::: beignet.quaternion_mean
@@ -41,4 +42,5 @@
 ::: beignet.rotation_vector_to_euler_angle
 ::: beignet.rotation_vector_to_quaternion
 ::: beignet.rotation_vector_to_rotation_matrix
+::: beignet.smith_waterman
 ::: beignet.translation_identity

pyproject.toml

@@ -7,8 +7,13 @@ requires = [
 ]
 
 [project]
-authors = [{ email = "[email protected]", name = "Allen Goodman" }]
-dependencies = ["torch"]
+authors = [
+    { email = "[email protected]", name = "Allen Goodman" },
+]
+dependencies = [
+    "pooch",
+    "torch",
+]
 dynamic = ["version"]
 license = { file = "LICENSE" }
 name = "beignet"
@@ -27,6 +32,11 @@ test = [
 ]
 
 [tool.ruff]
+exclude = [
+    "./src/beignet/constants/_substitution_matrices.py",
+]
+
+[tool.ruff.lint]
 select = [
     "B",  # FLAKE8-BUGBEAR
     "E",  # PYCODESTYLE ERRORS

src/beignet/__init__.py

@@ -28,6 +28,7 @@
 from ._invert_quaternion import invert_quaternion
 from ._invert_rotation_matrix import invert_rotation_matrix
 from ._invert_rotation_vector import invert_rotation_vector
+from ._needleman_wunsch import needleman_wunsch
 from ._quaternion_identity import quaternion_identity
 from ._quaternion_magnitude import quaternion_magnitude
 from ._quaternion_mean import quaternion_mean
@@ -65,6 +66,7 @@
 from ._rotation_vector_to_rotation_matrix import (
     rotation_vector_to_rotation_matrix,
 )
+from ._smith_waterman import smith_waterman
 from ._translation_identity import translation_identity
 
 __all__ = [
@@ -86,9 +88,11 @@
     "invert_quaternion",
     "invert_rotation_matrix",
     "invert_rotation_vector",
+    "needleman_wunsch",
     "quaternion_identity",
     "quaternion_magnitude",
     "quaternion_mean",
+    "quaternion_slerp",
     "quaternion_to_euler_angle",
     "quaternion_to_rotation_matrix",
     "quaternion_to_rotation_vector",
@@ -108,6 +112,6 @@
     "rotation_vector_to_euler_angle",
     "rotation_vector_to_quaternion",
     "rotation_vector_to_rotation_matrix",
-    "quaternion_slerp",
+    "smith_waterman",
     "translation_identity",
 ]

src/beignet/_needleman_wunsch.py

@@ -0,0 +1,194 @@
+import math
+import operator
+
+import torch
+import torch.func
+import torch.nn.functional
+from torch import Tensor
+
+
+def needleman_wunsch(
+    input: Tensor,
+    lengths: Tensor,
+    gap_penalty: float = 0.0,
+    temperature: float = 1.0,
+):
+    def fn(x: Tensor, shape: Tensor) -> Tensor:
+        padded = torch.nn.functional.pad(x, [1, 0, 1, 0])
+
+        i = torch.add(
+            torch.subtract(
+                torch.arange(
+                    padded.size(1),
+                )[None, :],
+                torch.flip(
+                    torch.arange(
+                        padded.size(0),
+                    ),
+                    dims=[0],
+                )[:, None],
+            ),
+            operator.sub(
+                padded.size(0),
+                1,
+            ),
+        )
+
+        j = torch.floor_divide(
+            torch.add(
+                torch.flip(
+                    torch.arange(
+                        padded.size(0),
+                    ),
+                    dims=[0],
+                )[:, None],
+                torch.arange(
+                    padded.size(1),
+                )[None, :],
+            ),
+            2,
+        )
+
+        m = operator.sub(
+            operator.add(
+                padded.size(0),
+                padded.size(1),
+            ),
+            1,
+        )
+
+        n = operator.floordiv(
+            operator.add(
+                padded.size(0),
+                padded.size(1),
+            ),
+            2,
+        )
+
+        y = torch.zeros([m, n], dtype=padded.dtype)
+
+        initialization = torch.zeros(
+            [
+                padded.size(0),
+                padded.size(1),
+            ],
+            dtype=padded.dtype,
+        )
+
+        initialization[:, 0] = torch.multiply(
+            torch.arange(padded.size(0)),
+            gap_penalty,
+        )
+
+        initialization[0, :] = torch.multiply(
+            torch.arange(padded.size(1)),
+            gap_penalty,
+        )
+
+        initialization = y.index_put(
+            [i, j],
+            initialization,
+        )
+
+        previous = torch.zeros(n)
+
+        previous_previous = torch.zeros(n)
+
+        traceback = torch.zeros([m, n])
+
+        mask = y.index_put(
+            [i, j],
+            torch.nn.functional.pad(
+                torch.multiply(
+                    torch.less(
+                        torch.arange(x.size(0)),
+                        shape[0],
+                    )[:, None],
+                    torch.less(
+                        torch.arange(x.size(1)),
+                        shape[1],
+                    )[None, :],
+                ),
+                [1, 0, 1, 0],
+            ).to(x.dtype),
+        )
+
+        striped_indexes = torch.fmod(
+            torch.add(
+                torch.arange(m),
+                math.fmod(
+                    padded.size(0),
+                    2,
+                ),
+            ),
+            2,
+        )
+
+        padded = y.index_put([i, j], padded)
+
+        for index in range(m):
+            # TRACEBACK:
+            traceback[index] = torch.add(
+                # APPLY MASK:
+                torch.multiply(
+                    # SMOOTH:
+                    torch.multiply(
+                        # APPLY SOFTMAX:
+                        torch.logsumexp(
+                            torch.divide(
+                                torch.stack(
+                                    [
+                                        # ALIGN:
+                                        torch.add(
+                                            previous_previous,
+                                            padded[index],
+                                        ),
+                                        previous + gap_penalty,
+                                        # CHANGE DIRECTION:
+                                        torch.add(
+                                            torch.add(
+                                                # INSERT:
+                                                torch.multiply(
+                                                    torch.nn.functional.pad(
+                                                        previous[:-1],
+                                                        [1, 0],
+                                                    ),
+                                                    striped_indexes[index],
+                                                ),
+                                                # DELETE:
+                                                torch.multiply(
+                                                    torch.nn.functional.pad(
+                                                        previous[+1:],
+                                                        [0, 1],
+                                                    ),
+                                                    operator.sub(
+                                                        1,
+                                                        striped_indexes[index],
+                                                    ),
+                                                ),
+                                            ),
+                                            gap_penalty,
+                                        ),
+                                    ],
+                                ),
+                                temperature,
+                            ),
+                            dim=0,
+                        ),
+                        temperature,
+                    ),
+                    mask[index],
+                ),
+                initialization[index],
+            )
+
+            previous_previous, previous = previous, traceback[index]
+
+        return traceback[i, j][shape[0], shape[1]]
+
+    output = torch.empty_like(input)
+
+    for index in range(input.shape[0]):
+        output[index] = torch.func.grad(fn)(input[index], lengths[index])
+
+    return output