nwgrad

High-performance C++ sequence alignment library with gradient computation for substitution matrix optimization — exposed to Python via nanobind.

nwgrad implements Needleman-Wunsch (global) and Smith-Waterman (local) alignment with both linear and affine gap penalties. Its primary novel feature is computing the gradient of the alignment score with respect to substitution matrix entries, enabling substitution matrix optimization in gradient-based ML pipelines.

Features

• Global (NW) and local (SW) alignment — shared template core, selected at compile time
• Linear and affine gap models — single gap penalty or gap-open + gap-extend
• Hard subgradient — substitution-pair counts along the optimal traceback
• Soft (differentiable) gradient — forward-backward in log-space; log-partition function and expected substitution counts
• Multithreaded batch processing — lock-free work dispatch, per-thread gradient accumulation
• Guide-banded DP — cheap re-alignment under a new matrix around a cached alignment path
• Double precision throughout — no float truncation in the hot path
• Zero-copy Python interface — nanobind buffer protocol; no unnecessary array copies

Installation

pip install nwgrad

Or from source:

git clone https://github.com/michalsta/nwgrad
cd nwgrad
pip install .

Requires Python ≥ 3.8 and a C++20 compiler (GCC 11+ or Clang 13+).

Quick Start

import numpy as np
import nwgrad

# Predefined matrices ship as ready-to-use SubstMatrix objects (they carry
# their own alphabet, so you never have to specify the symbol order yourself):
from nwgrad.matrices import BLOSUM62        # SubstMatrix, NCBI 23-symbol alphabet
print(BLOSUM62.alphabet)                    # 'ARNDCQEGHILKMFPSTWYVBZX'

# Or build one from your own (N, N) array.  A SubstMatrix always pairs a matrix
# with the alphabet that indexes its rows/columns:
dna_mat = nwgrad.SubstMatrix(np.eye(4) * 2 - 1, alphabet="ACGT")

Single pair — SeqPair

SeqPair takes an AlignParams object that bundles the substitution matrix with gap penalties. Pass a SubstMatrix directly — its alphabet travels with it, so amino acids map to matrix cells unambiguously:

params = nwgrad.AlignParams(BLOSUM62,
                            gap_open_a=11.0, gap_extend_a=1.0,
                            gap_open_b=11.0, gap_extend_b=1.0)

sp = nwgrad.SeqPair(
    "PLEASANTLY", "MEANLY", params,
    gap_model="affine",   # "linear" | "affine"
    mode="global",        # "global" | "local"
    grad_mode="hard",     # "hard" | "soft" | "none"
)

# Align, then read the score and a printable alignment:
score, grad = sp.score_and_grad()   # one call: allocates DP, aligns, computes grad
print(score)
print(sp.formatted())               # PLEASANTLY / match line / -MEAN---LY
print(sp.aligned())                 # ('PLEASANTLY', '-MEAN---LY')

# `grad` is an AlignParams; inspect it as a dict:
g = grad.to_dict()
print(g["matrix"].shape, g["gap_open_a"], g["gap_extend_b"])

align_full() (the DP) and compute_grad() (reading the gradient off the cached DP tables) are exposed separately so you can re-align many times and only pay for the gradient when you need it; score_and_grad() above just fuses the common case. The hard gradient is a subgradient read from the Viterbi path; the soft gradient (grad_mode="soft") is the exact gradient of the log-partition function via forward–backward.

After updating the matrix, realign_banded() re-scores the pair cheaply around the existing alignment path:

new_params = nwgrad.AlignParams(new_subst_matrix,
                                gap_open_a=11.0, gap_extend_a=1.0,
                                gap_open_b=11.0, gap_extend_b=1.0)
sp.set_params(new_params)
sp.realign_banded(bandwidth=20)
sp.compute_grad()

Batch — SeqPairBatch

from nwgrad.matrices import BLOSUM62

params = nwgrad.AlignParams(BLOSUM62,
                            gap_open_a=11.0, gap_extend_a=1.0,
                            gap_open_b=11.0, gap_extend_b=1.0)
pairs = [
    nwgrad.SeqPair(a, b, params,
                   gap_model="affine", mode="global", grad_mode="soft")
    for a, b in zip(seqs_a, seqs_b)
]

batch = nwgrad.SeqPairBatch(n_threads=8)
for sp in pairs:
    batch.add(sp)

# Align all pairs in parallel; score and grad are cached on each SeqPair
total_log_z = batch.score_and_grad()

# Sum gradients across all pairs → AlignParams
grad = batch.compute_grad()

# Update matrix and re-run with banded DP around the existing paths
new_params = nwgrad.AlignParams(new_subst_matrix,
                                gap_open_a=11.0, gap_extend_a=1.0,
                                gap_open_b=11.0, gap_extend_b=1.0)
batch.set_params(new_params)
total_log_z = batch.score_and_grad(bandwidth=20)
grad = batch.compute_grad()

SeqPair objects are constructed once and reused across optimization iterations. Only the AlignParams changes; the sequence data and, when applicable, the alignment path are preserved.

API Reference

Alphabet

nwgrad.Alphabet.get(symbols: str) -> Alphabet

The character↔index mapping that a matrix, its gradients, and its sequences all share. Alphabets are interned: Alphabet.get("ACGT") is nwgrad.DNA, so two matrices are compatible exactly when their alphabets are the same object.

Named alphabets:

Extensions append at the end, so the canonical 20 keep indices 0–19 and a 20×20 matrix in that order embeds as the top-left block of any extended one.

The matrices in nwgrad.matrices use a different ordering. BLOSUM, PAM and VTML are all in NCBI column order (NCBI_PROTEIN), and NUC44 is in IUPAC order (IUPAC_DNA). These are not extensions of PROTEIN / DNA — BLOSUM62 does not embed in PROTEIN_X. Combining a matrix over one alphabet with a gradient over another raises, rather than silently misreading the columns.

An Alphabet governs legality and ordering only — scoring is entirely the matrix's job, including for X and N.

For anything not listed, build your own: nwgrad.Alphabet.get("ACGTRYSWKM").

SubstMatrix

nwgrad.SubstMatrix(matrix: np.ndarray, alphabet: str = "ACDEFGHIKLMNPQRSTVWY")

Constructs a substitution matrix from an (N, N) float64 numpy array where N = len(alphabet). The default alphabet is the 20 canonical amino acids (ACDEFGHIKLMNPQRSTVWY). Any square matrix with a matching alphabet is accepted — including DNA ("ACGT"), extended amino acids, or any other symbol set. The alphabet is given as a string; if you are holding an Alphabet object, pass its .symbols.

Stored as a dense N × N block indexed by alphabet position — 16 doubles for DNA. Sequences are validated and encoded to indices once, at the boundary, so the DP inner loop is a single O(1) lookup and never touches a character. Asymmetric matrices are fully supported.

Characters outside the alphabet raise ValueError; they are not silently scored as zero. Case is significant — 'd' is not 'D', so soft-masked FASTA must be upper-cased by the caller.

Predefined matrices are available from nwgrad.matrices as ready-to-use SubstMatrix objects (each already carrying the correct alphabet): BLOSUM45/50/62/80/90, PAM30/70/250, VTML40/80/160/200, and NUC44 (DNA).

from nwgrad.matrices import BLOSUM62
BLOSUM62.alphabet            # 'ARNDCQEGHILKMFPSTWYVBZX'
BLOSUM62.to_matrix()         # (23, 23) float64 array

AlignParams

nwgrad.AlignParams(matrix, gap_open_a=0.0, gap_extend_a=0.0,
                           gap_open_b=0.0, gap_extend_b=0.0)
# or, from a raw array + explicit alphabet:
nwgrad.AlignParams(array, alphabet="ACDEFGHIKLMNPQRSTVWY", gap_open_a=0.0, ...)

Bundles a substitution matrix with (possibly asymmetric) gap penalties. The first form takes a SubstMatrix directly and is preferred — the alphabet travels with the matrix, so there is no risk of mismatching the symbol order. Suffix _a penalties apply to gaps in sequence A, _b to gaps in sequence B.

AlignParams is also the gradient type returned by SeqPair.grad / compute_grad(); to_dict() is the easy way to inspect a computed gradient.

Sign convention. Every field of a gradient — the matrix entries and the four gap fields — is the derivative of the score with respect to that field. The score subtracts gap penalties, so a gradient's gap fields come out non-positive while its matrix fields come out non-negative. That is what lets one update rule move the whole struct in the ascent direction:

params = params + lr * grad    # element-wise over matrix and all four gap fields

SeqPair

nwgrad.SeqPair(
    seq_a, seq_b, params,   # params: AlignParams
    gap_model="affine",     # "linear" | "affine"
    mode="global",          # "global" | "local"
    grad_mode="hard",       # "hard" | "soft" | "none"
)

Persistent sequence-pair object. Sequences and alignment mode are fixed at construction; the alignment parameters can be swapped cheaply via set_params(). The pair holds its AlignParams alive automatically, so you don't need to keep a separate reference to it.

Methods:

Properties:

Gradient modes:

SeqPairBatch

nwgrad.SeqPairBatch(n_threads=0)

n_threads=0 (default) uses hardware_concurrency. add() keeps each SeqPair (and, transitively, its AlignParams) alive for the lifetime of the batch.

Methods:

The batch is also a sequence: len(batch) is the number of pairs and batch[i] returns the i-th SeqPair (negative indices allowed), so per-pair results can be read back without keeping a separate list.

Properties:

Typical optimization loop:

from nwgrad.matrices import BLOSUM62

# Learnable matrix starts from BLOSUM62; keep its alphabet to stay consistent.
alphabet = BLOSUM62.alphabet
mat_array = BLOSUM62.to_matrix()

# Build pairs once
params = nwgrad.AlignParams(nwgrad.SubstMatrix(mat_array, alphabet),
                            gap_open_a=11.0, gap_extend_a=1.0,
                            gap_open_b=11.0, gap_extend_b=1.0)
batch = nwgrad.SeqPairBatch(n_threads=8)
for a, b in zip(seqs_a, seqs_b):
    batch.add(nwgrad.SeqPair(a, b, params,
                             gap_model="affine", mode="global", grad_mode="soft"))

for step in range(n_steps):
    total_log_z = batch.score_and_grad(bandwidth=bw if step > 0 else 0)
    grad = batch.compute_grad()
    mat_array += lr * grad.matrix.to_matrix()   # both in `alphabet` order
    params = nwgrad.AlignParams(nwgrad.SubstMatrix(mat_array, alphabet),
                                gap_open_a=11.0, gap_extend_a=1.0,
                                gap_open_b=11.0, gap_extend_b=1.0)
    batch.set_params(params)

BatchAligner

Stateless batch alignment: constructs a fresh DP buffer per call and does not preserve alignment paths across calls. Simpler API when you do not need to reuse paths for banded re-alignment.

nwgrad.BatchAligner(
    params,               # AlignParams — matrix and gap penalties together
    band=0,               # 0 = full DP; >0 = banded half-width
    gap_model="affine",   # "linear" | "affine"
    mode="global",        # "global" | "local"
    grad_mode="hard",     # "hard" | "soft" | "none"
    n_threads=1,
)

.align(sequences_a, sequences_b, aligned_a=[], aligned_b=[]) -> BatchResult

Aligns each pair (sequences_a[i], sequences_b[i]). aligned_a / aligned_b, if given, are gapped alignment strings (one per pair) used as banding guides.

BatchResult

Single-pair convenience functions

Stateless functions that create and destroy their DP tables on every call. Useful for one-off alignments; prefer SeqPair / SeqPairBatch for loops over many pairs or iterative optimization.

All twelve share the same signature — the gap penalties live in the AlignParams, not in the argument list, so the linear and affine variants differ only in which of its gap fields they read:

f(seq_a, seq_b, params, band=0, aligned_a="", aligned_b="")

band > 0 (or a non-empty aligned_a / aligned_b guide pair) runs a banded DP instead of the full one. The sequences are plain str and are validated and encoded against params's alphabet on the way in.

Score only — return float:

Hard gradient — return (score: float, grad: AlignParams):

Soft gradient — return (log_z: float, grad: AlignParams):

(Note the naming is irregular: the affine score functions suffix _affine, while the affine gradient functions infix it.)

score, grad = nwgrad.nw_affine_grad("PLEASANTLY", "MEANLY", params)

Gap penalty conventions

For linear gap model, a gap of length k costs gap_extend * k.

For affine gap model, a gap of length k costs gap_open + gap_extend * k.

This matches BioPython's convention where open is charged once per gap regardless of length.

Performance notes

• Full O(m×n) DP tables are retained for gradient computation — no Hirschberg-style memory optimisation.
• SeqPairBatch.score_and_grad() allocates one DpBuffer per thread (sized to the largest sequence pair in the batch) and reuses it across all assigned pairs. Pair-owned DP tables are never allocated, keeping peak memory at n_threads × max(m×n) rather than N × max(m×n).
• Work is distributed via a shared std::atomic counter — no per-task mutex, no work-stealing queue.
• Gradient accumulation is per-thread; a single mutex is taken once at join to merge partial gradients.

C++ header-only library

The C++ implementation is header-only (src/nwgrad/cpp/nwgrad/). To use it directly from C++:

python -m nwgrad --include
# prints: /path/to/nwgrad/cpp

Add that path to your include path and #include "nwgrad/aligner.hpp" etc.

Building from source

pip install scikit-build-core nanobind
pip install -e ".[dev]"
pytest tests/Python/

C++ unit tests (requires CMake). The extension target does find_package(nanobind), so even a tests-only configure needs nanobind discoverable:

cmake -S . -B build -DCMAKE_BUILD_TYPE=Debug \
      -Dnanobind_DIR="$(python -m nanobind --cmake_dir)"
cmake --build build
ctest --test-dir build

-DNWGRAD_SANITIZE=ON builds the C++ tests with AddressSanitizer and UndefinedBehaviorSanitizer (-fno-sanitize-recover=all, so the first violation fails the run rather than printing and carrying on). CI runs this under both GCC and Clang.

License

MIT — see LICENSE.