Protein ML in Python: Build an Amino Acid Composition Baseline

The nanoBERT and VHHBERT comparison suggested testing a simple baseline before fine-tuning an encoder. Here is a concrete starting point: count amino acids, add sequence length, and fit a small regression model.

This tutorial builds the complete path from sequences to held-out predictions. It runs on a CPU, accepts a CSV file, and saves the split so a later embedding experiment can use the same records. The built-in data and targets are synthetic; the example demonstrates the workflow and makes no claim about real protein properties.

1. Decide what the baseline can measure

For each sequence, compute the fraction of each of the 20 canonical amino acids, then append its length. The result has 21 features. For example, AAGC gives an alanine fraction of 2 / 4 = 0.5 and a length of 4.

This representation discards residue order. AAGC and CGAA produce identical features. It is therefore a deliberately limited reference for checking whether a more expensive representation adds predictive value on your task.

We compare two predictors on the same test rows:

Predictor	Input	Purpose
Training median	No sequence information	Establish an error level without sequence features
Composition + Ridge	20 fractions and length	Test whether these simple features reduce error

The median dummy regressor always predicts the training target median. Ridge regression fits a linear predictor with an L2 penalty. Here alpha=1.0 is a fixed demonstration setting, not a tuned choice.

2. Install and prepare the input

In a macOS or Linux terminal:

mkdir composition-demo
cd composition-demo
python3 -m venv .venv
source .venv/bin/activate
python -m pip install 'scikit-learn==1.9.1'

The example was tested with Python 3.14.7 and scikit-learn 1.9.1. No GPU or pretrained checkpoint is needed.

For your own data, provide these CSV columns:

Column	Meaning
`id`	Unique record identifier
`sequence`	Nonempty sequence using the 20 canonical amino acids
`group`	A pre-established family, lineage, or other evaluation group
`target`	One finite numeric measurement, with a consistent meaning and unit

The script normalizes case and trims surrounding whitespace. It rejects ambiguous residues such as X, gaps, and internal spaces. Decide how those records should be handled before running it; silently removing residues changes the features.

3. Run a complete baseline

Save the following as composition_baseline.py, or download the script.

import argparse
import csv
import random
from pathlib import Path

import numpy as np
import sklearn
from sklearn.dummy import DummyRegressor
from sklearn.linear_model import Ridge
from sklearn.metrics import mean_absolute_error
from sklearn.model_selection import GroupShuffleSplit
from sklearn.pipeline import make_pipeline
from sklearn.preprocessing import StandardScaler

AA = "ACDEFGHIKLMNPQRSTVWY"


def demo_rows():
    rng = random.Random(17)
    rows = []
    for family in range(60):
        length = rng.randrange(60, 101)
        weights = [rng.uniform(1, 15) if a == "A" else 1 for a in AA]
        parent = "".join(rng.choices(AA, weights=weights, k=length))
        for variant in range(3):
            seq = parent[:-3] + "".join(rng.choices(AA, k=3))
            # Artificial target designed to depend on composition and length.
            target = 4 * seq.count("A") / len(seq) + 0.01 * len(seq)
            target += rng.gauss(0, 0.03)
            rows.append(dict(id=f"f{family}_v{variant}", sequence=seq,
                             group=f"f{family}", target=target))
    return rows


def prepare(rows):
    if not rows:
        raise ValueError("No data rows")
    ids, sequences, groups, targets = [], [], [], []
    for row in rows:
        ident, group = str(row["id"]).strip(), str(row["group"]).strip()
        seq = str(row["sequence"]).strip().upper()
        target = float(row["target"])
        if not ident or not group or not seq or set(seq) - set(AA):
            raise ValueError("Each row needs id, group, and a canonical sequence")
        if not np.isfinite(target):
            raise ValueError("Targets must be finite numbers")
        ids.append(ident)
        sequences.append(seq)
        groups.append(group)
        targets.append(target)
    if len(set(ids)) != len(ids):
        raise ValueError("IDs must be unique")
    if len(set(groups)) < 4:
        raise ValueError("Provide at least four groups for this demo workflow")
    X = np.array([[s.count(a) / len(s) for a in AA] + [len(s)]
                  for s in sequences], dtype=float)
    return X, np.array(targets), np.array(groups), ids, sequences


def main():
    parser = argparse.ArgumentParser(description="Grouped composition baseline")
    parser.add_argument("--csv", type=Path, help="id,sequence,group,target columns")
    parser.add_argument("--seed", type=int, default=42)
    parser.add_argument("--split-out", type=Path, default=Path("split.csv"))
    args = parser.parse_args()
    if args.csv:
        with args.csv.open(newline="", encoding="utf-8-sig") as handle:
            reader = csv.DictReader(handle)
            required = {"id", "sequence", "group", "target"}
            if not required <= set(reader.fieldnames or []):
                raise ValueError("Missing CSV columns: id,sequence,group,target")
            rows = list(reader)
    else:
        rows = demo_rows()
    X, y, groups, ids, sequences = prepare(rows)
    splitter = GroupShuffleSplit(n_splits=1, test_size=0.25,
                                 random_state=args.seed)
    train, test = next(splitter.split(X, y, groups))
    overlap = set(groups[train]) & set(groups[test])
    if overlap:
        raise ValueError("Group overlap")
    if {sequences[i] for i in train} & {sequences[i] for i in test}:
        raise ValueError("Identical sequences cross the split; revise groups")
    # Exclusive creation protects an existing split record from replacement.
    with args.split_out.open("x", newline="", encoding="utf-8") as handle:
        writer = csv.writer(handle)
        writer.writerow(["id", "group", "split"])
        for label, indices in [("train", train), ("test", test)]:
            writer.writerows((ids[i], groups[i], label) for i in indices)
    print(f"scikit-learn={sklearn.__version__}; synthetic={args.csv is None}")
    print(f"rows: train={len(train)} test={len(test)}; features={X.shape[1]}")
    print(f"groups: train={len(set(groups[train]))} "
          f"test={len(set(groups[test]))}; overlap={len(overlap)}")
    models = {
        "median": DummyRegressor(strategy="median"),
        "composition_ridge": make_pipeline(StandardScaler(), Ridge(alpha=1.0)),
    }
    for name, model in models.items():
        model.fit(X[train], y[train])
        error = mean_absolute_error(y[test], model.predict(X[test]))
        print(f"{name}: MAE={error:.4f}")
    print(f"saved split: {args.split_out}")


if __name__ == "__main__":
    main()

The default generator creates 60 artificial families with three variants each. Its target is explicitly constructed from alanine fraction and length plus noise. That makes it useful for checking the implementation, but it gives the composition model an advantage by design.

Run the help command and the demonstration:

python composition_baseline.py --help
python composition_baseline.py

Observed output in the tested environment:

scikit-learn=1.9.1; synthetic=True
rows: train=135 test=45; features=21
groups: train=45 test=15; overlap=0
median: MAE=0.3388
composition_ridge: MAE=0.0275
saved split: split.csv

The script creates split.csv in the current directory. It refuses to overwrite an existing file; choose a new name when repeating a run:

python composition_baseline.py --split-out split-repeat.csv
python composition_baseline.py --csv measurements.csv --split-out split-real.csv

The second command requires your own complete dataset. Four groups are the script’s minimum input check, not a recommendation for a reliable study.

4. Read the split and the error correctly

GroupShuffleSplit assigns whole groups to one side of a split. Its test_size=0.25 reserves a quarter of the groups, rounded up, so the fraction of rows can differ when groups have unequal sizes.

The code checks both group overlap and exact sequence overlap. It does not discover related sequences or verify that the supplied group labels are biologically meaningful. Establish those groups before fitting the models. For a prediction task involving new antigens, a lineage-only split is insufficient; see the antibody evaluation guide.

The scaler is fitted inside a pipeline using training rows only. Test rows use those learned scaling parameters. This follows scikit-learn’s guidance on preprocessing leakage.

mean_absolute_error averages the absolute differences between predictions and targets. Lower is better, and the result has the target’s units. Here those units are artificial. The printed score weights rows equally; large families therefore contribute more than small ones.

5. Reuse this baseline for an embedding comparison

Keep split.csv, the input dataset, and the environment record:

python -m pip freeze > requirements-tested.txt

For the next experiment, join embeddings to records by id, and load the saved train and test assignments. Check that every expected ID appears exactly once. Do not rely on array row order or regenerate the split after reordering the CSV.

Use the same target, held-out rows, metric, and tuning budget for composition and embeddings. If you tune regularization, do so with grouped validation inside the training portion. Keep the final test targets out of that selection process. A single split is a starting check; a serious comparison needs enough independent groups and an uncertainty analysis matched to the sampling design.

Troubleshooting

Error or observation	Action
`Missing CSV columns`	Use the exact header `id,sequence,group,target`.
Canonical sequence error	Inspect empty sequences, gaps, spaces, and noncanonical residues. Apply an explicit data policy.
`Identical sequences cross the split`	Revisit grouping of duplicate records before training.
`FileExistsError`	Pass a fresh `--split-out` path to preserve the previous assignment.
Ridge does not beat the median	Check target units, sample size, grouping, and whether composition contains useful signal. Do not change the split just to improve the score.

Next experiment

Run this baseline on a consistently measured property, then compare frozen protein embeddings using the saved split. The useful result is how much error the representation removes beyond the median and composition references on the evaluation you actually care about.