Python API¶

Training and prediction¶

xiRTNET(params, input_dim, embedding_dim) validates the network configuration. For recurrent input, input_dim is the padded residue-sequence width; build_model(alphabet=...) supplies the vocabulary size including padding index zero. The recurrent trainable embedding width is configured by embedding.length, not inferred from the constructor's compatibility argument.

import yaml

from xirt.xirtnet import xiRTNET

with open("xirt_params.yaml") as stream:
    params = yaml.safe_load(stream)

network = xiRTNET(params, input_dim=30, embedding_dim=25)
network.build_model(siamese=True, alphabet=25)
network.compile()

# Supply your encoded arrays and targets in consistent sample/task order.
# x1, x2: integer arrays (samples, 30), zero-padded.
# targets: one target array per network.tasks entry.
# train/validation partitions must be selected before fitting preprocessing.
history = network.model.fit(
    [x1, x2], targets,
    validation_data=([validation_x1, validation_x2], validation_targets),
    epochs=10, batch_size=128,
)
predictions = network.model.predict([test_x1, test_x2])
metrics = network.model.evaluate([test_x1, test_x2], test_targets)

The arrays in this snippet are placeholders; the training notebook includes data preparation and a runnable example. A non-Siamese model accepts one array. Fraction tasks precede continuous tasks, with names sorted within each group; use network.tasks rather than guessing head order.

network.model is a torch.nn.Module with xiRT's compatibility fit/predict/evaluate methods. network.compile() configures its optimizer(s). fit returns an object with .history holding epoch metric lists. These adapters are not Keras objects and do not accept every Keras extension. Direct forward is for tensor-based PyTorch usage.

Network and checkpoints¶

xirt.predictor.Predictor ¶

Load a self-contained RT or FAIMS checkpoint for label-free prediction and XAI.

Raw RT inputs follow the checkpoint's saved columns/preprocessing. Tutorial checkpoints embed their explicit tokenizer metadata. FDR, scores and target labels are not needed for inference. Frozen ESM encoders use the ESM API.

encode ¶

encode(frame)

Encode raw peptides without sorting/dropping rows or fitting preprocessing.

explain ¶

explain(frame, *, task=None, steps=64)

Return residue/numeric attributions mapped back to the input peptide order.

load `classmethod` ¶

load(location, *, device='cpu')

Load RT and FAIMS checkpoints through one public entry point.

predict ¶

predict(frame)

Return task-prediction columns in original target units and input row order.

xirt.predictor.PeptideInputs `dataclass` ¶

Encoded branches and their row/token mapping for prediction and explanation.

xirt.predictor.prediction_frame ¶

prediction_frame(
    model, predictions, *, index=None, original_units=True
)

Format native task outputs consistently for API, CLI and legacy ModelData.

xirt.xirtnet.xiRTNET ¶

Build, train and transfer xiRT PyTorch models using legacy YAML sections.

Example

network = xiRTNET(params, input_dim=30, embedding_dim=25) # doctest: +SKIP network.build_model(siamese=True, alphabet=25) # doctest: +SKIP network.compile() # doctest: +SKIP

adjust_model ¶

adjust_model() -> None

Replace changed task outputs while retaining compatible encoder and head weights.

build_model ¶

build_model(
    siamese: bool = False, alphabet: int | None = None
) -> None

Create a new model with shared branches and deterministic initialization.

compile ¶

compile() -> None

Initialize optimizers for the current model.

fine_tune ¶

fine_tune(
    location: str | Path, strategy: str, **options: Any
) -> None

Load a base checkpoint and configure native full, frozen, LoKr or LoRA training.

get_callbacks ¶

get_callbacks(suffix: str = '') -> dict[str, Any]

Return explicit training options and create the callback output directory.

get_param_overview ¶

get_param_overview() -> None

Print total, trainable and frozen parameter counts.

load_model ¶

load_model(location: str | Path) -> None

Reconstruct a saved model on the configured device.

Example

network.load_model("retention.pt") # doctest: +SKIP predicted_rt = network.model.predict(features) # doctest: +SKIP

load_weights ¶

load_weights(location: str | Path) -> None

Load weights into an existing architecture.

print_layers ¶

print_layers() -> None

Print the actual PyTorch module tree.

xirt.xirtnet.RetentionModel ¶

Bases: Module

Shared peptide encoder and task-specific heads, with NumPy/Pandas adapters.

Parameters:

Name	Type	Description	Default
`params`	`dict[str, Any]`	Validated architecture mapping.	required
`input_dim`	`int`	Padded sequence length or embedding width.	required
`alphabet`	`int`	Vocabulary size including padding index zero.	required
`siamese`	`bool`	Whether two peptides share the encoder.	required

Example

model = RetentionModel(params, 30, 25, True) # doctest: +SKIP outputs = model(torch.ones(2, 30, dtype=torch.long), ... torch.ones(2, 30, dtype=torch.long)) # doctest: +SKIP

device `property` ¶

device: device

Return the device holding model parameters.

optimizer `property` ¶

optimizer: Optimizer

Return the primary optimizer; Muon also has an AdamW companion.

compile ¶

compile() -> None

Initialize fresh optimizers; call after replacing heads or trainability.

configure_fine_tuning ¶

configure_fine_tuning(options: dict[str, Any]) -> None

Set trainability after loading base weights; recompile before training.

count_params ¶

count_params() -> int

Return total parameter count, excluding running-statistics buffers.

encode ¶

encode(x: Tensor) -> Tensor

Encode pre-padded tokens, ignoring padding, or accept frozen embeddings.

evaluate ¶

evaluate(
    x: Any, y: Any, batch_size: int = 512, verbose: int = 0
) -> list[float]

Return total loss, task losses and configured metrics, in metrics_names order.

explain ¶

explain(
    x: Any,
    *,
    task: str | None = None,
    index: Any = None,
    tokens: Any = None,
    steps: int = 64,
) -> Any

Explain a continuous prediction at residue level using integrated gradients.

Embeddings are integrated from zero with sequence lengths held fixed; numeric features are integrated from their fitted training means. Returned contributions and completeness errors use the target's units.

fit ¶

fit(
    x: Any,
    y: Any,
    validation_data: tuple[Any, Any] | None = None,
    validation_split: float = 0.0,
    epochs: int = 1,
    batch_size: int = 128,
    verbose: int = 1,
    callbacks: dict[str, Any] | None = None,
) -> SimpleNamespace

Train shuffled minibatches and restore the best validation checkpoint.

Explicit validation data is preferred; validation_split reserves the tail before shuffling. No sample is discarded, including singleton batches. Returns an object with history, a mapping of epoch metric lists.

Example

history = model.fit((first, second), [observed_rt], ... validation_data=(val_features, val_targets), ... epochs=10, batch_size=64) # doctest: +SKIP

forward ¶

forward(
    first: Tensor,
    second: Tensor | None = None,
    *,
    numeric: Tensor | None = None,
) -> dict[str, Tensor]

Return task logits; probabilistic activation is applied only for prediction.

load `classmethod` ¶

load(
    location: str | Path, *, device: str = "cpu"
) -> RetentionModel

Reconstruct a native checkpoint through the public tensor-model API.

load_weights ¶

load_weights(location: str | Path) -> None

Load strictly compatible 2.0 tensor weights; legacy HDF5 is not accepted.

predict ¶

predict(
    x: Any, batch_size: int = 512, verbose: int = 0
) -> Any

Predict in bounded batches; return an array for one task, otherwise a list.

Example

predicted_rt = model.predict((first, second)) # doctest: +SKIP

predict_frame ¶

predict_frame(
    x: Any, *, index: Any = None, batch_size: int = 512
) -> pd.DataFrame

Return named predictions in original units, preserving an optional input index.

save ¶

save(location: str | Path) -> None

Save a versioned tensor checkpoint with reconstruction and vocabulary metadata.

Example

model.save("retention.pt") # doctest: +SKIP

target_transform ¶

target_transform(task: str) -> tuple[float, float, str]

Return multiplier, offset and units mapping a scalar head to reported values.

train ¶

train(mode: bool = True) -> RetentionModel

Keep frozen encoder statistics unchanged while training heads/adapters.

xirt.xirtnet.read_checkpoint ¶

read_checkpoint(location: str | Path) -> dict[str, Any]

Load safe tensors/primitive metadata and reject legacy checkpoint formats.

Example

data = read_checkpoint("xirt_model_01.pt") # doctest: +SKIP data["format_version"] # doctest: +SKIP 2

Configuration and training utilities¶

xirt.config.NetworkConfig ¶

Bases: BaseModel

validate_architecture ¶

validate_architecture() -> NetworkConfig

xirt.config.FineTuningConfig ¶

Bases: BaseModel

Train all recurrent weights, only heads, or GRU adapters plus heads.

validate_factor ¶

validate_factor() -> FineTuningConfig

Accept positive factor bounds or balanced factorization.

xirt.training.make_optimizers ¶

make_optimizers(
    model: Module, options: dict[str, Any]
) -> list[torch.optim.Optimizer]

Build optimizers, assigning each trainable parameter exactly once.

Muon receives hidden two-dimensional weights only; embeddings, prediction layers and vector parameters use AdamW. Lion is an optional dependency.

Example

opts = make_optimizers(nn.Linear(4, 1), {"optimizer": "adamw", ... "learningrate": 0.001}) type(opts[0]).name 'AdamW'

xirt.training.seed_everything ¶

seed_everything(
    seed: int, deterministic: bool = True
) -> None

Seed Python, NumPy and Torch; optionally require deterministic kernels.

Example

seed_everything(42)

xirt.training.task_loss ¶

task_loss(
    logits: Tensor,
    target: Tensor,
    name: str,
    gamma: float = 2.0,
    alpha: float | None = None,
) -> Tensor

Compute a scalar task loss from raw logits and shape-matched targets.

binary_focal supports ordinal multi-label targets; categorical focal supports one-hot targets. Gamma zero reproduces ordinary cross entropy. Alpha, when provided, is a binary positive-class weight only.

Example

task_loss(torch.zeros(2, 1), torch.ones(2, 1), "mse").item() 1.0

Frozen features and differentiable ESM-C¶

The optional ESM SDK is imported when needed, not to run a core recurrent model. Consult the ESM guide for licensing, modification limitations, local weights, and cache semantics.

xirt.embeddings.ESMCEncoder ¶

Frozen, deduplicated mean-residue ESM-C embeddings with optional disk cache.

Injected models must implement forward(sequence_tokens=...), returning an object with .embeddings, and expose the public tokenizer (including get_vocab) and embedding size. Every model's actual initial state_dict and tokenizer vocabulary are hashed, including downloaded pretrained weights. CPU copies are bounded to one tensor at a time, not the entire backbone. Treat the model as immutable after construction.

Loading this example requires optional ESM dependencies and local model weights:

encoder = ESMCEncoder(cache_dir="embedding-cache") # doctest: +SKIP encoder.encode(["ACDE", "W"]).shape # doctest: +SKIP (2, 960)

init ¶

__init__(
    model_name: str = "esmc_300m",
    device: str | device = "cpu",
    cache_dir: str | Path | None = None,
    batch_size: int = 8,
    model: Module | None = None,
) -> None

Load or freeze an injected local model and initialize its cache identity.

encode ¶

encode(sequences: Sequence[str]) -> np.ndarray

Return (number of sequences, embedding_dim) float32 mean embeddings.

encoder.encode(["ACDE", "W"]).dtype # doctest: +SKIP dtype('float32')

xirt.embeddings.tokenize_batch ¶

tokenize_batch(
    sequences: Sequence[str],
    model: Module,
    device: str | device | None = None,
    *,
    max_length: int = MAX_RESIDUES,
) -> tuple[torch.Tensor, torch.Tensor]

Tokenize with public ESM tokenizer; return padded IDs and residue-only mask.

The boolean mask excludes BOS, EOS and PAD, regardless of their token IDs. Length checks happen before tokenization; no sequence is silently truncated.

Given a locally loaded ESM-compatible model:

tokens, residues = tokenize_batch(["ACDE", "W"], model) # doctest: +SKIP residues.sum(dim=1).tolist() # doctest: +SKIP [4, 1]

xirt.embeddings.ESMCRetentionModel ¶

Bases: Module

Differentiable local ESM-C with optional LoRA/LoKr and a retention head.

The entire pretrained backbone is frozen before transformer Linear adapters are injected. Save full state_dict plus constructor configuration and target names; reconstruct the same architecture before load_state_dict.

Loading this example requires optional ESM dependencies and local model weights:

model = ESMCRetentionModel(rank=2, factor=2, alpha=2) # doctest: +SKIP tokens, residues = model.tokenize(["ACDE", "W"]) # doctest: +SKIP model(tokens, residues).square().mean().backward() # doctest: +SKIP

init ¶

__init__(
    model_name: str = "esmc_300m",
    device: str | device = "cpu",
    *,
    model: Module | None = None,
    output_dim: int = 1,
    rank: int = 4,
    factor: int = 4,
    alpha: float = 4.0,
    target_names: Sequence[str] | None = None,
    adapter_kind: str = "lokr",
) -> None

Freeze the backbone, inject selected transformer adapters and create a head.

encode ¶

encode(
    token_ids: Tensor, residue_mask: Tensor
) -> torch.Tensor

Return differentiable residue-only mean embeddings for shared Siamese heads.

forward ¶

forward(
    token_ids: Tensor, residue_mask: Tensor
) -> torch.Tensor

Predict retention values through differentiable residue pooling and the head.

tokenize ¶

tokenize(
    sequences: Sequence[str],
) -> tuple[torch.Tensor, torch.Tensor]

Prepare IDs and residue masks on the backbone's current device.

tokens, residues = model.tokenize(["ACDE"]) # doctest: +SKIP

train ¶

train(mode: bool = True) -> ESMCRetentionModel

Set head training mode while preserving deterministic backbone evaluation.

Recurrent adapters and linear LoKr¶

xirt.adapters.LoKrLinear ¶

Bases: _LoKrUpdate

Frozen Linear plus a trainable, initially zero Kronecker-product update.

layer = LoKrLinear(nn.Linear(12, 8), rank=2, factor=2, alpha=2) layer(torch.ones(3, 12)).shape torch.Size([3, 8])

bias `property` ¶

bias: Tensor | None

Return the original frozen bias, or None for a bias-free base.

weight `property` ¶

weight: Tensor

Return the differentiable base-plus-adapter effective weight.

init ¶

__init__(
    base: Linear,
    rank: int = 4,
    factor: int = 4,
    alpha: float = 4.0,
) -> None

Freeze the base and preserve the established LoKr state_dict layout.

forward ¶

forward(inputs: Tensor) -> torch.Tensor

Apply the frozen linear mapping plus the trainable Kronecker update.

xirt.adapters.apply_lokr ¶

apply_lokr(
    model: Module,
    target_names: Sequence[str] | None = None,
    *,
    rank: int = 4,
    factor: int = 4,
    alpha: float = 4.0,
) -> tuple[str, ...]

Replace selected Linear descendants in-place, returning their exact paths.

None selects all Linear descendants. Explicit targets must be exact named_modules paths. Only replaced base layers are frozen; callers can freeze the whole backbone first. Repeated calls skip already wrapped descendants.

heads = nn.Sequential(nn.Linear(12, 8), nn.ReLU()) apply_lokr(heads, rank=2, factor=2, alpha=2) ('0',)

xirt.adapters.apply_lora_gru ¶

apply_lora_gru(
    model: Module, *, rank: int = 4, alpha: float = 4.0
) -> tuple[str, ...]

Parametrize every GRU input/hidden gate matrix with a zero-initialized LoRA update.

xirt.adapters.apply_lokr_gru ¶

apply_lokr_gru(
    model: Module,
    *,
    rank: int = 4,
    factor: int = 4,
    alpha: float = 4.0,
) -> tuple[str, ...]

Adapt input/hidden GRU gate matrices, preserving layers and directions.

Biases and other modules are not frozen here: freeze a pretrained model first when performing parameter-efficient fine-tuning. Reconstruct the same parametrizations before loading an adapted checkpoint.

gru = nn.GRU(4, 8, batch_first=True) apply_lokr_gru(gru, rank=2) ('weight_ih_l0', 'weight_hh_l0')

Local experiment tracking¶

xirt.tracking.TrainingRun ¶

Record configuration, data fingerprints, epoch metrics and model artifacts.

Parameters:

Name	Type	Description	Default
`options`	`dict[str, Any]`	Tracking configuration with `enabled`, `uri`, `experiment` and optional `run_name`. Only local stores/localhost are accepted.	required
`params`	`dict[str, Any]`	Architecture and training configuration.	required
`inputs`	`tuple[Tensor, ...]`	Training features, used only for SHA256 fingerprints (not uploaded).	required
`targets`	`tuple[Tensor, ...]`	Optional training labels, fingerprinted without uploading them.	`()`

Example

with TrainingRun({"enabled": False}, {}, ()) as run: ... run.log_metrics({"loss": 1.0}, step=0)

enter ¶

__enter__() -> TrainingRun

Create an independent run when tracking is enabled.

exit ¶

__exit__(
    exc_type: Any, exc_value: Any, traceback: Any
) -> None

Mark failed training runs as failed; never swallow training/logging errors.

log_artifact ¶

log_artifact(path: str | Path) -> None

Log a safe xiRT checkpoint or reproducibility artifact.

log_metrics ¶

log_metrics(metrics: dict[str, float], step: int) -> None

Log scalar metrics at their epoch index, without changing global MLflow state.

FAIMS compensation voltage¶

These APIs predict compensation voltage in volts, not CCS. See the FAIMS guide for inputs, grouping, units and coverage limitations.

xirt.faims.FAIMSPredictor ¶

Train and persist a PTM-aware xiRT predictor of observed FAIMS CV in volts.

Fit requires explicitly disjoint training and validation peptide pairs. All preprocessing is fitted on training data only. Inference needs no run, target, FDR or TT columns; unseen PTMs are counted and encoded as <unk>.

encode ¶

encode(frame: DataFrame) -> PeptideInputs

Encode only the checkpoint's required inputs, retaining original peptide order mapping.

explain ¶

explain(frame: DataFrame, *, steps: int = 64) -> Any

Explain CV at residue level, including charge/m/z contributions in volts.

fit ¶

fit(
    train: DataFrame,
    validation: DataFrame,
    *,
    verbose: int = 0,
) -> FAIMSPredictor

Fit vocabulary, target scaling and network without test-set access.

load `classmethod` ¶

load(
    location: str | Path, *, device: str = "cpu"
) -> FAIMSPredictor

Reconstruct a FAIMS checkpoint, rejecting RT-only checkpoints.

predict ¶

predict(frame: DataFrame) -> pd.DataFrame

Predict volts in input row order and expose unseen-token coverage.

save ¶

save(location: str | Path) -> None

Save native tensor weights, vocabulary, scaling and CV semantics.

xirt.faims.prepare_faims_observations ¶

prepare_faims_observations(
    raw: DataFrame,
    source: str,
    *,
    maximum_fdr: float = 0.01,
    cv_column: str | None = None,
) -> tuple[pd.DataFrame, dict[str, Any]]

Filter TT CSMs and aggregate repeated observations to mean observed CV.

Each precursor is a source, modified peptide pair, link sites and charge. Counts and observed CV spread retain information about label ambiguity. This target is not an intensity-calibrated optimum or calibrated CCS.

xirt.faims.split_precursors ¶

split_precursors(
    frame: DataFrame,
    *,
    seed: int = 42,
    fractions: tuple[float, float, float] = (
        0.75,
        0.125,
        0.125,
    ),
) -> pd.DataFrame

Assign deterministic train/validation/test partitions by unmodified pair.

xirt.faims.default_faims_params ¶

default_faims_params(
    *,
    seed: int = 42,
    precursor_features: bool = True,
    features: str = "charge",
) -> dict[str, Any]

Return a shared CV GRU with sequence, charge, or charge+m/z context.

precursor retains the original charge/mass/m/z architecture for old checkpoints; the focused comparison uses sequence, charge and charge-mz.

xirt.faims.cv_metrics ¶

cv_metrics(observed: Any, predicted: Any) -> dict[str, Any]

Compute voltage errors; R² is explicitly undefined for constant/tiny sets.

xirt.faims.precursor_descriptors ¶

precursor_descriptors(frame: DataFrame) -> pd.DataFrame

Return symmetric, target-free precursor and sequence composition features.

Explanations and shared QC¶

See prediction explanations for baselines and limitations.

xirt.explain.Explanation `dataclass` ¶

Signed contributions and their completeness check in the target's units.

residues includes separate terminus rows when the tokenizer uses them. numeric records values, training-mean baselines and contributions. summary preserves sample order; its completeness error is prediction minus baseline minus the sum of all token and numeric contributions.

The embedding baseline is zero at unchanged sequence lengths. It is not an empty peptide, a physical counterfactual or evidence of causality.

xirt.explain.integrated_gradients ¶

integrated_gradients(
    model: Any,
    inputs: tuple[Tensor, ...],
    *,
    task: str | None = None,
    index: Any = None,
    tokens: list[list[list[str]]] | None = None,
    steps: int = 64,
) -> Explanation

Integrate a scalar head along an embedding/precursor path using public PyTorch APIs.

Gauss-Legendre integration holds token IDs, peptide lengths and padding masks fixed. No derivatives of integer token IDs are taken. Numeric features move jointly from their fitted training means to their values. Model mode, hooks and existing parameter gradients are preserved.

xirt.qc.prediction_metrics ¶

prediction_metrics(observed, predicted)

Return regression errors and correlations, with undefined scores as None.

Values must be aligned, finite and in the same units. Predictive R² is distinct from Pearson r squared; a constant prediction has undefined r.

xirt.qc.prediction_qc ¶

prediction_qc(
    observed, predicted, *, name="Retention time", unit=""
)

Create the same annotated prediction/residual QC figure for RT and FAIMS.

xirt.qc.plot_predictions ¶

plot_predictions(
    observed,
    predicted,
    *,
    name="Retention time",
    unit="",
    ax=None,
    color=None,
)

Plot observed/predicted values with identity line and annotated correlations.

xirt.qc.plot_learning_curve ¶

plot_learning_curve(history, *, scale=1.0, unit='')

Plot training and validation RMSE from scalar standardized MSE histories.

xirt.qc.plot_residue_attributions ¶

plot_residue_attributions(explanation, *, sample=0)

Plot signed residue and precursor contributions in the prediction's units.

Positive contributions raise the prediction relative to the explanation baseline; they are model attributions, not causal chemical effects.

Python API¶

Training and prediction¶

Network and checkpoints¶

xirt.predictor.Predictor ¶

encode ¶

explain ¶

load classmethod ¶

predict ¶

xirt.predictor.PeptideInputs dataclass ¶

xirt.predictor.prediction_frame ¶

xirt.xirtnet.xiRTNET ¶

adjust_model ¶

build_model ¶

compile ¶

fine_tune ¶

get_callbacks ¶

get_param_overview ¶

load_model ¶

load_weights ¶

print_layers ¶

xirt.xirtnet.RetentionModel ¶

device property ¶

optimizer property ¶

compile ¶

configure_fine_tuning ¶

count_params ¶

encode ¶

evaluate ¶

explain ¶

fit ¶

forward ¶

load classmethod ¶

load_weights ¶

predict ¶

predict_frame ¶

save ¶

target_transform ¶

train ¶

xirt.xirtnet.read_checkpoint ¶

Configuration and training utilities¶

xirt.config.NetworkConfig ¶

validate_architecture ¶

xirt.config.FineTuningConfig ¶

validate_factor ¶

xirt.training.make_optimizers ¶

xirt.training.seed_everything ¶

xirt.training.task_loss ¶

Frozen features and differentiable ESM-C¶

xirt.embeddings.ESMCEncoder ¶

__init__ ¶

encode ¶

xirt.embeddings.tokenize_batch ¶

xirt.embeddings.ESMCRetentionModel ¶

__init__ ¶

encode ¶

forward ¶

tokenize ¶

train ¶

Recurrent adapters and linear LoKr¶

xirt.adapters.LoKrLinear ¶

bias property ¶

weight property ¶

__init__ ¶

forward ¶

xirt.adapters.apply_lokr ¶

xirt.adapters.apply_lora_gru ¶

xirt.adapters.apply_lokr_gru ¶

Local experiment tracking¶

xirt.tracking.TrainingRun ¶

__enter__ ¶

__exit__ ¶

log_artifact ¶

log_metrics ¶

FAIMS compensation voltage¶

xirt.faims.FAIMSPredictor ¶

encode ¶

explain ¶

fit ¶

load classmethod ¶

predict ¶

load `classmethod` ¶

xirt.predictor.PeptideInputs `dataclass` ¶

device `property` ¶

optimizer `property` ¶

load `classmethod` ¶

init ¶

init ¶

bias `property` ¶

weight `property` ¶

init ¶

enter ¶

exit ¶

load `classmethod` ¶

xirt.explain.Explanation `dataclass` ¶