diff --git a/README.md b/README.md index 9539fdc..7d101ef 100644 --- a/README.md +++ b/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: diff --git a/data/README.md b/demo/README.md similarity index 100% rename from data/README.md rename to demo/README.md diff --git a/data/cr6261_h1.csv b/demo/cr6261_h1.csv similarity index 100% rename from data/cr6261_h1.csv rename to demo/cr6261_h1.csv diff --git a/lock_gp/train.py b/demo/train.py similarity index 84% rename from lock_gp/train.py rename to demo/train.py index 525da9f..2426f4c 100644 --- a/lock_gp/train.py +++ b/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,13 +74,18 @@ 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) @@ -88,19 +93,18 @@ def main() -> None: 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) diff --git a/demo/util.py b/demo/util.py new file mode 100644 index 0000000..ed4b2b0 --- /dev/null +++ b/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) diff --git a/lock_gp/gp_wrapper.py b/lock_gp/gp_wrapper.py index 893e201..3cfe9c5 100644 --- a/lock_gp/gp_wrapper.py +++ b/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), ) diff --git a/lock_gp/lock_kernel.py b/lock_gp/lock_kernel.py index 45e444d..44ec17a 100644 --- a/lock_gp/lock_kernel.py +++ b/lock_gp/lock_kernel.py @@ -1,5 +1,6 @@ from __future__ import annotations +from collections.abc import Sequence from functools import partial import torch @@ -7,8 +8,16 @@ 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): @@ -16,16 +25,19 @@ class LinearGlobalLOCKKernel(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( @@ -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), diff --git a/lock_gp/tanimoto_kernel.py b/lock_gp/tanimoto_kernel.py index 08433a6..1a489e8 100644 --- a/lock_gp/tanimoto_kernel.py +++ b/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) diff --git a/lock_gp/utils.py b/lock_gp/utils.py index 81ac40b..03a0b64 100644 --- a/lock_gp/utils.py +++ b/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, diff --git a/pyproject.toml b/pyproject.toml index 147adda..53ca319 100644 --- a/pyproject.toml +++ b/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" diff --git a/tests/test_alphabet_validation.py b/tests/test_alphabet_validation.py new file mode 100644 index 0000000..ab7ef42 --- /dev/null +++ b/tests/test_alphabet_validation.py @@ -0,0 +1,43 @@ +from __future__ import annotations + +import pytest +import torch + +from lock_gp.blosum50 import get_blosum50_matrix +from lock_gp.lock_kernel import LinearGlobalLOCKKernel, NonlinearLocalLOCKKernel, build_lock_kernel +from lock_gp.tanimoto_kernel import TanimotoKernel + + +def test_tanimoto_kernel_rejects_mismatched_alphabet_order() -> None: + alphabet, _ = get_blosum50_matrix() + + with pytest.raises(ValueError, match="BLOSUM50 alphabet order"): + TanimotoKernel(num_positions=1, alphabet=list(reversed(alphabet))) + + +def test_lock_kernel_rejects_mismatched_alphabet_order() -> None: + alphabet, _ = get_blosum50_matrix() + + with pytest.raises(ValueError, match="BLOSUM50 alphabet order"): + build_lock_kernel(num_positions=1, alphabet=list(reversed(alphabet))) + + +def test_tanimoto_kernel_forward_rejects_mismatched_alphabet_size() -> None: + alphabet, _ = get_blosum50_matrix() + kernel = TanimotoKernel(num_positions=1, alphabet=alphabet) + x = torch.zeros(2, 1, len(alphabet) - 1, dtype=torch.float64) + + with pytest.raises(ValueError, match="alphabet size"): + kernel.forward(x, x) + + +@pytest.mark.parametrize("kernel_cls", [LinearGlobalLOCKKernel, NonlinearLocalLOCKKernel]) +def test_lock_kernel_forward_rejects_mismatched_alphabet_size( + kernel_cls: type[LinearGlobalLOCKKernel | NonlinearLocalLOCKKernel], +) -> None: + alphabet, _ = get_blosum50_matrix() + kernel = kernel_cls(num_positions=1, alphabet=alphabet) + x = torch.zeros(2, 1, len(alphabet) - 1, dtype=torch.float64) + + with pytest.raises(ValueError, match="alphabet size"): + kernel.forward(x, x) diff --git a/tests/test_e2e.py b/tests/test_e2e.py index a817a00..b2e8701 100644 --- a/tests/test_e2e.py +++ b/tests/test_e2e.py @@ -21,7 +21,7 @@ def test_train_end_to_end() -> None: result = subprocess.run( - [sys.executable, "-m", "lock_gp.train", "--num-training", "256"], + [sys.executable, "-m", "demo.train", "--data", "demo/cr6261_h1.csv", "--train-size", "256"], capture_output=True, text=True, cwd=REPO_ROOT,