mhcflurry.affinity package

Helper modules for class I affinity prediction.

Submodules

mhcflurry.affinity.calibration_sizing module

Calibration sizing and cache helpers for class I affinity predictors.

mhcflurry.affinity.calibration_sizing.peptide_sequences_fingerprint(sequences)

Order-and-content-sensitive SHA-256 of a peptide list.

Length-prefixes each peptide so e.g. ["AB", "C"] and ["A", "BC"] cannot collide by concatenation. Used as the cache key for the fast calibration peptide-stage cache; collisions there silently reuse the wrong tensors and produce wrong PercentRankTransforms.

class mhcflurry.affinity.calibration_sizing.CalibrationFastCache

Bases: object

Per-predictor-instance state for calibrate_percentile_ranks_fast.

Holds the device-resident peptide-stage tensors and the motif-summary helper state that survive across calibrate tasks within one worker. Centralizing both fields here makes the cache lifecycle visible — it used to live behind dynamic setattr of private names.

stage_signature

cached_stages

motif_signature

motif_state

clear()

mhcflurry.affinity.calibration_sizing.peptide_encoding_feature_dim(model)

Return the raw encoded peptide feature width for model if known.

mhcflurry.affinity.calibration_sizing.resolve_peptide_feature_dim(model, peptide_feature_dim, peak_bytes)

Resolve the width of cached peptide features used during calibration.

peptide_feature_dim is the actual width of forward_peptide_stage(...) when the caller has probed it. Without a probe, fall back to shape-derived estimates and then to the generic peak memory estimate.

class mhcflurry.affinity.calibration_sizing.CalibrationSizingEnv(free_memory_fraction, reserve_fraction, reserve_min_bytes, fixed_safety_multiplier)

Bases: tuple

Create new instance of CalibrationSizingEnv(free_memory_fraction, reserve_fraction, reserve_min_bytes, fixed_safety_multiplier)

fixed_safety_multiplier

Alias for field number 3

free_memory_fraction

Alias for field number 0

reserve_fraction

Alias for field number 1

reserve_min_bytes

Alias for field number 2

mhcflurry.affinity.calibration_sizing.read_calibration_sizing_env(free_memory_fraction, safety_multiplier)

Resolve the MHCFLURRY_CALIBRATE_AUTO_* overrides for the sizer.

free_memory_fraction and safety_multiplier are the caller defaults used when the matching env var is unset. Invalid or out-of-range values raise ValueError (via env_float) rather than being silently ignored. Returns a CalibrationSizingEnv.

mhcflurry.affinity.calibration_sizing.cuda_device_memory_bytes(device)

Return best-effort (free_bytes, total_bytes) for a CUDA device.

free comes from the shared free_device_memory_bytes probe. total uses torch.cuda.get_device_properties when available and otherwise falls back to free – tests and nonstandard CUDA wrappers may not expose device properties, and the explicit reserve still keeps the budget bounded.

mhcflurry.affinity.calibration_sizing.cuda_calibration_total_rows(model, device, n_peptides, num_workers_per_gpu, env, num_cached_networks, peptide_feature_dim, num_sub_networks, cuda_overhead_bytes)

Size the cartesian forward row budget on a CUDA device.

Carves a per-worker VRAM budget out of free memory, subtracts the fixed overhead (the persistent peptide-feature cache plus safety-padded CUDA/runtime scratch and small state), and converts the remaining forward budget into a row count. env is a CalibrationSizingEnv from read_calibration_sizing_env; see auto_size_calibration_batches for the full peak model. Returns the row budget for choose_calibration_batch_shape.

mhcflurry.affinity.calibration_sizing.auto_size_calibration_batches(model, device, n_peptides, n_alleles, num_workers_per_gpu=1, free_memory_fraction=None, num_cached_networks=1, peptide_feature_dim=None, num_sub_networks=None, cuda_overhead_bytes=2147483648, safety_multiplier=1.3, fixed_peptide_batch=None, fixed_allele_batch=None)

Split the auto-sized batch budget between peptide and allele axes.

fixed_peptide_batch / fixed_allele_batch pin one axis to a user-supplied value (mixed pin/auto mode). The pinned axis is held constant and only the other axis is sized against the VRAM budget, so the auto axis shrinks to keep allele_batch × peptide_batch within the per-worker budget instead of being sized as if the pinned axis were also auto (which would underestimate peak VRAM).

Models the per-worker VRAM peak as:

peak = cuda_overhead
     + cache_bytes                # peptide-stage cache,
                                  # ``num_cached_networks ×
                                  # n_peptides × peptide_feature_dim × 4``
     + cartesian_intermediate     # transient forward
                                  # ``a_size × p_batch ×
                                  # peak_bytes_per_row``
     + small_state                # log-IC50 acc, ic50_unique,
                                  # motif state (~1 GB)

This includes explicit free-memory headroom plus a safety factor on CUDA/runtime scratch allocations to absorb fragmentation that mem_get_info can’t see.

peak_bytes_per_row is calibrated for the cartesian fast path. For a merged ensemble it uses one sub-network’s hidden activation peak plus the small retained per-sub-network outputs, matching MergedClass1NeuralNetwork.forward_cartesian_from_peptide_stage which runs each sub-network to completion before starting the next. The cartesian-intermediate term scales with a_size × p_batch, which is exactly the rate the budget carves out — preserving the existing total_rows = forward_budget // peak_bytes math.

Returns the chosen (peptide_batch, allele_batch).

mhcflurry.affinity.calibration_sizing.estimate_calibration_peak_bytes_per_row(model)

Estimate cartesian calibration forward peak bytes per row.

The generic prediction estimator is deliberately conservative for arbitrary merged forwards. Calibration has a more specific execution shape: MergedClass1NeuralNetwork.forward_cartesian_from_peptide_stage evaluates sub-networks serially and retains only their final outputs before combining them. Hidden-layer peak memory is therefore the max sub-network peak, not the sum of every sub-network peak.

mhcflurry.affinity.calibration_sizing.choose_calibration_batch_shape(total_rows, n_peptides, n_alleles, min_peptide_batch, max_allele_batch=256, fixed_peptide_batch=None, fixed_allele_batch=None)

Choose (peptide_batch, allele_batch) under a row budget.

Minimize the number of cartesian forward chunks rather than filling one axis greedily. This matters for the production shape (tens of alleles × tens of thousands of peptides), where many equivalent row budgets can differ by 20-40% in Python-level loop count and kernel launches.

fixed_peptide_batch / fixed_allele_batch pin one axis (mixed pin/auto mode). The pinned axis is held at the user value and the other axis is sized as total_rows // pinned so the product stays within the budget; the pinned axis is not capped by max_allele_batch (the user asked for it explicitly).

mhcflurry.affinity.calibration_sizing.calibration_stage_cache_signature(encoded_peptides, networks, device)

Return the key for a reusable peptide-stage calibration cache.

Keyed on the peptide-set fingerprint, the network object identities (id), and the device – NOT on weight content. Adding, removing, or replacing ensemble models changes the networks list (new objects -> new ids) and so invalidates the cache, but mutating an existing network’s weights in place (e.g. a re-fit) does not. A weight-content fingerprint is deliberately not used here: borrow_cached_network serves architecturally-identical networks from a single process-wide MODELS_CACHE module, so the underlying torch parameter storage is shared across ensemble members and is not a reliable per-network weight signal. Callers that mutate weights in place between fast-calibrate calls on the same predictor instance must therefore call clear_calibration_fast_cache() first (see calibrate_percentile_ranks_fast).

mhcflurry.affinity.calibration_sizing.probe_peptide_feature_dim(net_obj, encoded_peptides, device)

Run a 1-row forward through forward_peptide_stage to record the actual width of the peptide feature tensor.

The auto-sizer’s cache estimate depends on this and the encoding-shape heuristic under-counts when the model configures peptide_dense_layer_sizes or LC layers.

Returns None on probe failure so callers can fall back to the heuristic.

mhcflurry.affinity.calibration_sizing.calibration_fast_cache(self)

Return (creating if needed) the per-instance fast-calibrate cache.

See CalibrationFastCache. Lazy so that predictors loaded for prediction never pay the allocation cost.

mhcflurry.affinity.calibration_sizing.clear_calibration_fast_cache(self)

Drop any cached fast-calibrate state on this predictor.

Long-lived workers that finish calibrate but stay alive for prediction can reclaim the (potentially many GB) of device- resident peptide-stage tensors via this hook.

mhcflurry.affinity.model_selection module

Model-selection helpers for class I affinity predictors.

mhcflurry.affinity.model_selection.model_select(predictor, predictor_class, score_function, alleles=None, min_models=1, max_models=10000)

Perform model selection using a user-specified scoring function.

This works only with allele-specific models, not pan-allele models.

Model selection is done using a “step up” variable selection procedure, in which models are repeatedly added to an ensemble until the score stops improving.

Parameters:

predictorClass1AffinityPredictor

Predictor containing the candidate models.

predictor_classtype

Predictor class used for candidate and selected ensembles.

score_functionClass1AffinityPredictor -> float function

Scoring function

alleleslist of string, optional

If not specified, model selection is performed for all alleles.

min_modelsint, optional

Min models to select per allele

max_modelsint, optional

Max models to select per allele

Returns:

Class1AffinityPredictorpredictor containing the selected models

mhcflurry.affinity.persistence module

Persistence helpers for class I affinity predictors.

mhcflurry.affinity.persistence.save_predictor(predictor, models_dir, model_names_to_write=None, write_metadata=True)

Serialize the predictor to a directory on disk. If the directory does not exist it will be created.

The serialization format consists of a file called “manifest.csv” with the configurations of each Class1NeuralNetwork, along with per-network files giving the model weights. If there are pan-allele predictors in the ensemble, the pseudosequences are also stored in the directory. There is also a small file “info.txt” with basic metadata: the save time, user and host (not the original training time).

Parameters:

predictorClass1AffinityPredictor

Predictor to serialize.

models_dirstring

Path to directory. It will be created if it doesn’t exist.

model_names_to_writelist of string, optional

Only write the weights for the specified models. Useful for incremental updates during training. Passing an explicit empty list writes no model artifacts; this is used by calibration-only updates that replace percentile calibration (CSV or JSON) without touching the manifest, weights, model provenance, allele sequences, or optimization metadata. Explicit metadata_dataframes are still written when write_metadata is true.

write_metadataboolean, optional

Whether to write optional metadata

mhcflurry.affinity.persistence.load_predictor(predictor_class, models_dir=None, max_models=None, optimization_level=None, optimization_level_default=None)

Deserialize a predictor from a directory on disk.

Parameters:

predictor_classtype

Predictor class to instantiate.

models_dirstring

Path to directory. If unspecified the default downloaded models are used.

max_modelsint, optional

Maximum number of Class1NeuralNetwork instances to load

optimization_levelint or None

If >0, model optimization will be attempted. When None, use optimization_level_default supplied by the public loader.

optimization_level_defaultint or None

Fallback supplied by the public loader when no explicit level is set. That loader reads MHCFLURRY_OPTIMIZATION_LEVEL with a default of 1.

Returns:

Class1AffinityPredictor instance