Demo training layout

README.md

@@ -17,13 +17,19 @@ This repo was tested with **PyTorch 2.6**.
 
 ## Data
 
-- `data/cr6261_h1.csv` with columns `sequence` and `fitness`.
+- `demo/cr6261_h1.csv` with columns `sequence` and `fitness`.
 - Alphabet is fixed to the canonical 20 amino acids plus gap token (`-`).
 
 ## How to run
 
 ```bash
-uv run python -m lock_gp.train --num-training 256
+uv run python -m demo.train --train-size 256
+```
+
+To train on custom data, pass a CSV with `sequence` and `fitness` columns:
+
+```bash
+uv run python -m demo.train --data path/to/data.csv --train-size 256
 ```
 
 The script:

lock_gp/train.py → demo/train.py

@@ -8,9 +8,9 @@
 import torch
 from scipy.stats import pearsonr, spearmanr
 
-from .blosum50 import get_blosum50_matrix
-from .utils import encode_one_hot
-from .gp_wrapper import GPWrapper, LinearGP, LockGP, TanimotoGP
+from demo.util import encode_one_hot
+from lock_gp.blosum50 import get_blosum50_matrix
+from lock_gp.gp_wrapper import GPWrapper, LinearGP, LockGP, TanimotoGP
 
 
 def train_test_split(
@@ -74,33 +74,37 @@ def evaluate(y_true: torch.Tensor, y_pred_mean: torch.Tensor, y_pred_var: torch.
 def main() -> None:
     """Fit Linear, LOCK, and Tanimoto GPs to CR6261-H1 dataset."""
     parser = argparse.ArgumentParser(description="LOCK GP Demo")
-    parser.add_argument("--num-training", type=int, default=256, help="Number of training points.")
+    parser.add_argument(
+        "--data",
+        type=Path,
+        default=Path(__file__).parent / "cr6261_h1.csv",
+        help="Path to a CSV with sequence and fitness columns.",
+    )
+    parser.add_argument("--train-size", type=int, default=256, help="Number of training points.")
     args = parser.parse_args()
 
     device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 
-    data_path = Path("data") / "cr6261_h1.csv"
-    df = pd.read_csv(data_path)
+    df = pd.read_csv(args.data)
     sequences = df["sequence"].astype(str).tolist()
     y_np = df["fitness"].to_numpy(dtype=np.float64)
 
     alphabet, _ = get_blosum50_matrix()
     x = encode_one_hot(sequences, alphabet, dtype=torch.float64)
     y = torch.tensor(y_np, dtype=torch.float64)
 
-    x_train, y_train, x_test, y_test = train_test_split(x, y, num_train=args.num_training, device=device)
+    x_train, y_train, x_test, y_test = train_test_split(x, y, num_train=args.train_size, device=device)
 
     print(f"Device: {device}     Train size: {len(y_train)}     Test size: {len(y_test)}")
 
-    gp_configs: list[tuple[str, type[GPWrapper]]] = [
-        ("Linear GP", LinearGP),
-        ("Tanimoto GP", TanimotoGP),
-        ("LOCK GP", LockGP),
+    gp_models: list[tuple[str, GPWrapper]] = [
+        ("Linear GP", LinearGP()),
+        ("Tanimoto GP", TanimotoGP(alphabet=alphabet)),
+        ("LOCK GP", LockGP(alphabet=alphabet)),
     ]
 
     results: dict[str, dict[str, float]] = {}
-    for name, gp_cls in gp_configs:
-        gp = gp_cls()
+    for name, gp in gp_models:
         gp.fit(x_train, y_train)
         mean, var = gp.predict(x_test)
         results[name] = evaluate(y_test, mean, var)

demo/util.py

@@ -0,0 +1,37 @@
+from __future__ import annotations
+
+from collections.abc import Sequence
+
+import torch
+import torch.nn.functional as F
+
+
+def encode_one_hot(
+    sequences: list[str],
+    alphabet: Sequence[str],
+    dtype: torch.dtype = torch.float64,
+) -> torch.Tensor:
+    """
+    Convert sequences to one-hot tensor of shape (N, L, A).
+
+    Args:
+        sequences: List of equal-length sequences.
+        alphabet: Ordered alphabet of valid tokens.
+        dtype: Floating-point dtype of the returned one-hot tensor. Defaults
+            to ``torch.float64``.
+
+    Raises:
+        ValueError: if any sequence contains a token outside the alphabet or
+            if the sequences are not all the same length.
+    """
+    seq_length = len(sequences[0])
+    for seq in sequences:
+        if len(seq) != seq_length:
+            raise ValueError("All sequences must have the same length.")
+    alphabet_index = {tok: i for i, tok in enumerate(alphabet)}
+    unknown = {tok for seq in sequences for tok in seq if tok not in alphabet_index}
+    if unknown:
+        raise ValueError(f"Unknown token(s) {sorted(unknown)} encountered.")
+
+    indices = torch.tensor([[alphabet_index[tok] for tok in seq] for seq in sequences], dtype=torch.int64)
+    return F.one_hot(indices, num_classes=len(alphabet)).to(dtype=dtype)

lock_gp/gp_wrapper.py

@@ -1,6 +1,6 @@
 from __future__ import annotations
 
-from collections.abc import Callable
+from collections.abc import Callable, Sequence
 
 import gpytorch
 import torch
@@ -11,9 +11,9 @@
 from gpytorch.kernels import LinearKernel, ScaleKernel
 from gpytorch.priors import GammaPrior
 
-from .exact_gp import ExactGPModel
-from .lock_kernel import build_lock_kernel
-from .tanimoto_kernel import TanimotoKernel
+from lock_gp.exact_gp import ExactGPModel
+from lock_gp.lock_kernel import build_lock_kernel
+from lock_gp.tanimoto_kernel import TanimotoKernel
 import logging
 import gc
 
@@ -138,19 +138,19 @@ def __init__(self) -> None:
 class LockGP(GPWrapper):
     """GP with the composite LOCK kernel."""
 
-    def __init__(self) -> None:
+    def __init__(self, alphabet: Sequence[str]) -> None:
         super().__init__(
-            kernel_factory=lambda x: build_lock_kernel(num_positions=x.shape[1])
+            kernel_factory=lambda x: build_lock_kernel(num_positions=x.shape[1], alphabet=alphabet)
         )
 
 
 class TanimotoGP(GPWrapper):
     """GP with a scaled Tanimoto kernel using BLOSUM50 encoding."""
 
-    def __init__(self) -> None:
+    def __init__(self, alphabet: Sequence[str]) -> None:
         super().__init__(
             kernel_factory=lambda x: ScaleKernel(
-                TanimotoKernel(num_positions=x.shape[1]),
+                TanimotoKernel(num_positions=x.shape[1], alphabet=alphabet),
                 outputscale_prior=GammaPrior(2.0, 2.0),
                 outputscale_constraint=GreaterThan(1e-4),
             )

lock_gp/lock_kernel.py

@@ -1,31 +1,43 @@
 from __future__ import annotations
 
+from collections.abc import Sequence
 from functools import partial
 
 import torch
 from gpytorch.constraints import GreaterThan
 from gpytorch.kernels import AdditiveKernel, Kernel, ProductKernel, ScaleKernel
 from gpytorch.priors import GammaPrior, NormalPrior
 
-from .blosum50 import get_blosum50_matrix
-from .utils import _gpytorch_default_setting_closure, reshape_inputs
+from lock_gp.blosum50 import get_blosum50_matrix
+from lock_gp.utils import _gpytorch_default_setting_closure, reshape_inputs
+
+
+def _validate_alphabet_size(x1: torch.Tensor, x2: torch.Tensor, expected_size: int) -> None:
+    if x1.shape[2] != expected_size or x2.shape[2] != expected_size:
+        raise ValueError(
+            f"Expected one-hot inputs to use alphabet size {expected_size}; "
+            f"got {x1.shape[2]} and {x2.shape[2]}."
+        )
 
 
 class LinearGlobalLOCKKernel(Kernel):
     """Linear LOCK kernel with a single global learnable exponent."""
 
     has_lengthscale = False
 
-    def __init__(self, num_positions: int, **kwargs) -> None:
+    def __init__(self, num_positions: int, alphabet: Sequence[str], **kwargs) -> None:
         """
         Initialize the kernel.
 
         Args:
             num_positions: Number of positions in the aligned sequence.
+            alphabet: Alphabet used to one-hot encode inputs.
             **kwargs: Additional arguments passed to parent Kernel class.
         """
         super().__init__(**kwargs)
-        _, blosum = get_blosum50_matrix(normalize=True)
+        blosum_alphabet, blosum = get_blosum50_matrix(normalize=True)
+        if list(alphabet) != blosum_alphabet:
+            raise ValueError("Input alphabet must match the BLOSUM50 alphabet order.")
         self.register_buffer("blosum", blosum)
         self.num_positions = num_positions
         self.register_parameter(
@@ -54,6 +66,7 @@ def forward(self, x1: torch.Tensor, x2: torch.Tensor, diag: bool = False, **kwar
         """
         x1 = reshape_inputs(x1, self.num_positions)
         x2 = reshape_inputs(x2, self.num_positions)
+        _validate_alphabet_size(x1, x2, expected_size=self.blosum.shape[0])
         if diag:
             return self.num_positions * torch.ones(x1.shape[0], dtype=x1.dtype, device=x1.device)
         global_exponent = self.log_global_exponent.exp().to(dtype=x1.dtype, device=x1.device)
@@ -67,16 +80,19 @@ class NonlinearLocalLOCKKernel(Kernel):
 
     has_lengthscale = False
 
-    def __init__(self, num_positions: int, **kwargs) -> None:
+    def __init__(self, num_positions: int, alphabet: Sequence[str], **kwargs) -> None:
         """
         Initialize the kernel.
 
         Args:
             num_positions: Number of positions in the aligned sequence.
+            alphabet: Alphabet used to one-hot encode inputs.
             **kwargs: Additional arguments passed to parent Kernel class.
         """
         super().__init__(**kwargs)
-        _, blosum = get_blosum50_matrix(normalize=True)
+        blosum_alphabet, blosum = get_blosum50_matrix(normalize=True)
+        if list(alphabet) != blosum_alphabet:
+            raise ValueError("Input alphabet must match the BLOSUM50 alphabet order.")
         self.register_buffer("blosum", blosum)
         self.num_positions = num_positions
         self.register_parameter(
@@ -115,6 +131,7 @@ def forward(self, x1: torch.Tensor, x2: torch.Tensor, diag: bool = False, **kwar
         """
         x1 = reshape_inputs(x1, self.num_positions)
         x2 = reshape_inputs(x2, self.num_positions)
+        _validate_alphabet_size(x1, x2, expected_size=self.blosum.shape[0])
         if diag:
             return torch.ones(x1.shape[0], dtype=x1.dtype, device=x1.device)
         global_exponent = self.log_global_exponent.exp().to(dtype=x1.dtype, device=x1.device)
@@ -124,7 +141,7 @@ def forward(self, x1: torch.Tensor, x2: torch.Tensor, diag: bool = False, **kwar
         return torch.einsum("bla,BlA,laA->bB", x1, x2, blosum).exp()
 
 
-def build_lock_kernel(num_positions: int) -> Kernel:
+def build_lock_kernel(num_positions: int, alphabet: Sequence[str]) -> Kernel:
     """
     Build the full LOCK kernel.
 
@@ -135,13 +152,14 @@ def build_lock_kernel(num_positions: int) -> Kernel:
 
     Args:
         num_positions: Number of positions in the aligned sequence.
+        alphabet: Alphabet used to one-hot encode inputs.
 
     Returns:
         Composite LOCK kernel.
     """
-    linear1 = LinearGlobalLOCKKernel(num_positions=num_positions)
-    linear2 = LinearGlobalLOCKKernel(num_positions=num_positions)
-    nonlinear = NonlinearLocalLOCKKernel(num_positions=num_positions)
+    linear1 = LinearGlobalLOCKKernel(num_positions=num_positions, alphabet=alphabet)
+    linear2 = LinearGlobalLOCKKernel(num_positions=num_positions, alphabet=alphabet)
+    nonlinear = NonlinearLocalLOCKKernel(num_positions=num_positions, alphabet=alphabet)
     return AdditiveKernel(
         ScaleKernel(
             ProductKernel(nonlinear, linear1),

lock_gp/tanimoto_kernel.py

@@ -1,10 +1,20 @@
 from __future__ import annotations
 
+from collections.abc import Sequence
+
 import torch
 from gpytorch.kernels import Kernel
 
-from .blosum50 import get_blosum50_matrix
-from .utils import reshape_inputs
+from lock_gp.blosum50 import get_blosum50_matrix
+from lock_gp.utils import reshape_inputs
+
+
+def _validate_alphabet_size(x1: torch.Tensor, x2: torch.Tensor, expected_size: int) -> None:
+    if x1.shape[2] != expected_size or x2.shape[2] != expected_size:
+        raise ValueError(
+            f"Expected one-hot inputs to use alphabet size {expected_size}; "
+            f"got {x1.shape[2]} and {x2.shape[2]}."
+        )
 
 
 class TanimotoKernel(Kernel):
@@ -16,15 +26,18 @@ class TanimotoKernel(Kernel):
 
     has_lengthscale = False
 
-    def __init__(self, num_positions: int) -> None:
+    def __init__(self, num_positions: int, alphabet: Sequence[str]) -> None:
         """
         Initialize the kernel.
 
         Args:
             num_positions: Number of variable positions in the aligned sequence.
+            alphabet: Alphabet used to one-hot encode inputs.
         """
         super().__init__()
-        _, blosum = get_blosum50_matrix(normalize=False)
+        blosum_alphabet, blosum = get_blosum50_matrix(normalize=False)
+        if list(alphabet) != blosum_alphabet:
+            raise ValueError("Input alphabet must match the BLOSUM50 alphabet order.")
         eigenvalues, eigenvectors = torch.linalg.eigh(blosum)
         encoding = eigenvectors @ torch.diag(torch.sqrt(torch.clamp(eigenvalues, min=0.0)))
         self.register_buffer("encoding", encoding)
@@ -45,6 +58,7 @@ def forward(self, x1: torch.Tensor, x2: torch.Tensor, diag: bool = False, **kwar
         """
         x1 = reshape_inputs(x1, self.num_positions)
         x2 = reshape_inputs(x2, self.num_positions)
+        _validate_alphabet_size(x1, x2, expected_size=self.encoding.shape[0])
         if diag:
             return torch.ones(x1.shape[0], dtype=x1.dtype, device=x1.device)
 

lock_gp/utils.py

@@ -1,10 +1,7 @@
 from __future__ import annotations
 
-from collections.abc import Sequence
-
 import gpytorch
 import torch
-import torch.nn.functional as F
 
 
 def reshape_inputs(x: torch.Tensor, num_positions: int) -> torch.Tensor:
@@ -32,37 +29,6 @@ def reshape_inputs(x: torch.Tensor, num_positions: int) -> torch.Tensor:
     return x.view(x.shape[0], num_positions, alphabet_size)
 
 
-def encode_one_hot(
-    sequences: list[str],
-    alphabet: Sequence[str],
-    dtype: torch.dtype = torch.float64,
-) -> torch.Tensor:
-    """
-    Convert sequences to one-hot tensor of shape (N, L, A).
-
-    Args:
-        sequences: List of equal-length sequences.
-        alphabet: Ordered alphabet of valid tokens.
-        dtype: Floating-point dtype of the returned one-hot tensor. Defaults
-            to ``torch.float64``.
-
-    Raises:
-        ValueError: if any sequence contains a token outside the alphabet or
-            if the sequences are not all the same length.
-    """
-    seq_length = len(sequences[0])
-    for seq in sequences:
-        if len(seq) != seq_length:
-            raise ValueError("All sequences must have the same length.")
-    alphabet_index = {tok: i for i, tok in enumerate(alphabet)}
-    unknown = {tok for seq in sequences for tok in seq if tok not in alphabet_index}
-    if unknown:
-        raise ValueError(f"Unknown token(s) {sorted(unknown)} encountered.")
-
-    indices = torch.tensor([[alphabet_index[tok] for tok in seq] for seq in sequences], dtype=torch.int64)
-    return F.one_hot(indices, num_classes=len(alphabet)).to(dtype=dtype)
-
-
 def _gpytorch_default_setting_closure(
     module: gpytorch.Module,
     value: torch.Tensor,

pyproject.toml

@@ -30,6 +30,12 @@ include = ["lock_gp"]
 [tool.ruff]
 line-length = 120
 
+[tool.ruff.lint]
+extend-select = ["TID252"]
+
+[tool.ruff.lint.flake8-tidy-imports]
+ban-relative-imports = "all"
+
 [[tool.uv.index]]
 name = "pypi"
 url = "https://pypi.org/simple"