Skip to content

API Documentation

Here lies the official top-level API for interacting with jax-unirep.

Calculating Representations

jax_unirep.get_reps

jax_unirep.get_reps(seqs, model=None, mlstm_size=1900)

Get reps of proteins.

This function generates representations of protein sequences using the mLSTM model from the UniRep paper.

Each element of the output 3-tuple is a np.array of shape (n_input_sequences, mlstm_size):

  • h_avg: Average hidden state of the mLSTM over the whole sequence.
  • h_final: Final hidden state of the mLSTM
  • c_final: Final cell state of the mLSTM

You should not use this function if you want to do further JAX-based computations on the output vectors! In that case, call the MLSTM directly, so that the JAX arrays it returns can be passed into the next step instead of being converted to np.arrays.

All three published model sizes are supported. The model knows its own depth and width, so nothing needs to be declared about its architecture: the 1900 model has one mLSTM cell and the 256 and 64 models have four.

:param seqs: A list of sequences as strings, or a single string. :param model: The MLSTM to featurize with, as returned by load_model() or fit(). When given, its own width is used and mlstm_size is ignored. :param mlstm_size: Which set of pre-trained weights to load when model is None. One of 1900, 256 or 64. :returns: A 3-tuple of np.arrays containing the reps, in the order h_avg, h_final, and c_final. Each np.array has shape (n_sequences, mlstm_size).

jax_unirep.fusion_reps

jax_unirep.fusion_reps(seqs, model=None, mlstm_size=1900)

Get "UniRep Fusion" representations of proteins.

Alley et al. 2019 define UniRep Fusion as the concatenation of all three representations -- average hidden, final hidden and final cell state -- into a single vector, and use it for the supervised stability and quantitative function prediction tasks. For the 1900 model that is 5700 dimensions.

This is a one-line composition of get_reps, and exists because the concatenation is a named quantity from the paper rather than an obvious thing to guess:

h_avg, h_final, c_final = get_reps(seqs)
fusion = np.hstack([h_avg, h_final, c_final])

If you are fine-tuning rather than featurizing, do not reach for this. Concatenate inside your own equinox.Module instead, so the gradient reaches the mLSTM -- see the "End-to-end differentiable models" section of the docs.

:param seqs: A list of sequences as strings, or a single string. :param model: The MLSTM to featurize with, as returned by load_model() or fit(). When given, its own width is used and mlstm_size is ignored. :param mlstm_size: Which set of pre-trained weights to load when model is None. One of 1900, 256 or 64. :returns: An np.array of shape (n_sequences, 3 * mlstm_size), the components in the order h_avg, h_final, c_final.

Evotuning

jax_unirep.fit

jax_unirep.fit(sequences, n_epochs, model=None, batch_method='length', batch_size=25, step_size=0.0001, holdout_seqs=None, proj_name='temp', epochs_per_print=1)

Return an mLSTM fitted to predict the next letter in each AA sequence.

The training loop is as follows, depending on the batching strategy:

Length batching:

  • At each iteration, of all sequence lengths present in sequences, one length gets chosen at random.
  • Next, batch_size number of sequences of the chosen length get selected at random.
  • If there are less sequences of a given length than batch_size, all sequences of that length get chosen.
  • Those sequences then get passed through the model. No padding of sequences occurs.

To get batching of sequences by length done, we call on batch_sequences from our utils.py module, which returns a list of sub-lists, in which each sub-list contains the indices in the original list of sequences that are of a particular length.

Random batching:

  • Before training, all sequences get padded to be the same length as the longest sequence in sequences.
  • Then, at each iteration, we randomly sample batch_size sequences and pass them through the model.

The training loop does not adhere to the common notion of epochs, where all sequences would be seen by the model exactly once per epoch. Instead sequences always get sampled at random, and one epoch approximately consists of round(len(sequences) / batch_size) weight updates. Asymptotically, this should be approximately equivalent to doing epoch passes over the dataset.

You can optionally dump weights and print losses every epochs_per_print epochs to monitor training progress. For ergonomics, training/holdout set losses are estimated on a batch size the same as batch_size, rather than calculated exactly on the entire set. Dumped weights are written in the same layout load_model reads, so training can be resumed from load_model(folderpath=...).

Parameters

  • sequences: List of sequences to evotune on.
  • n_epochs: The number of iterations to evotune on.
  • model: The MLSTM to tune. Defaults to the pre-trained mLSTM1900 from the paper. Pass MLSTM(n_cells=..., output_dim=..., key=...) to start from randomly initialized weights of any size, or load_model(folderpath=...) to resume from dumped weights.
  • batch_method: One of "length" or "random". Defaults to "length", which groups sequences of identical length and pads nothing. "random" pads every sequence to the longest in the whole dataset, which on a realistic length distribution wastes about half the compute and feeds gap characters through the recurrent state. Prefer "length" unless the sequences are already near-uniform.
  • batch_size: If random batching is used, number of sequences per batch. As a rule of thumb, batch size of 50 consumes about 5GB of GPU RAM.
  • step_size: The learning rate.
  • holdout_seqs: Holdout set, an optional input.
  • proj_name: The directory path for weights to be output to.
  • epochs_per_print: Number of epochs to progress before printing and dumping of weights. Must be greater than or equal to 1.

Returns

The tuned MLSTM.

jax_unirep.evotune

jax_unirep.evotune(sequences, model=None, n_trials=20, n_epochs_config=None, learning_rate_config=None, n_splits=5, out_dom_seqs=None)

Evolutionarily tune the model to a set of sequences.

Evotuning is described in the original UniRep and eUniRep papers. This reimplementation of evotune provides a nicer API that automatically handles multiple sequences of variable lengths.

Evotuning always needs a starter model. By default, the pre-trained weights from the Nature Methods paper are used. However, other pre-trained weights are legitimate.

We first use optuna to figure out how many epochs to fit before overfitting happens. To save on computation time, the number of trials run defaults to 20, but can be configured.

If you want to start from randomly initialized weights of any size:

from jax.random import PRNGKey
from jax_unirep.evotuning import evotune
from jax_unirep.models import MLSTM

model = MLSTM(n_cells=4, output_dim=256, key=PRNGKey(0))
study, tuned_model = evotune(sequences, model=model)

or from previously dumped weights:

from jax_unirep.models import load_model

model = load_model(folderpath="path/to/weights/folder")

The model states its own architecture, so nothing needs to be told what size it is.

This function is intended as an automagic way of identifying the best model and training routine hyperparameters. If you want more control over how fitting happens, please use the fit() function directly. There is an example in the examples/ directory that shows how to use it.

Parameters

  • sequences: Sequences to evotune against.
  • model: The MLSTM to tune. Defaults to the pre-trained mLSTM1900 from the paper.
  • n_trials: The number of trials Optuna should attempt.
  • n_epochs_config: A dictionary of kwargs to trial.suggest_float, which are: name, low, high, step. This controls how many epochs to have Optuna test. See source code for default configuration, at the definition of n_epochs_kwargs.
  • learning_rate_config: A dictionary of kwargs to trial.suggest_float, which are: name, low, high. This controls the learning rate of the model. See source code for default configuration, at the definition of learning_rate_kwargs.
  • n_splits: The number of folds of cross-validation to do.
  • out_dom_seqs: Out-domain holdout set of sequences, to check for loss on to prevent overfitting.

Returns

  • study: The optuna study object, containing information about all evotuning trials.
  • tuned_model: The final, optimized MLSTM.

Sampling

jax_unirep.sample_one_chain

jax_unirep.sample_one_chain(starter_sequence, n_steps, scoring_func, is_accepted_kwargs={}, trust_radius=7, propose_kwargs={})

Return one chain of MCMC samples of new sequences.

Given a starter_sequence, this function will sample one chain of protein sequences, scored using a user-provided scoring_func.

Design choices made here include the following.

Firstly, we record all sequences that were sampled, and not just the accepted ones. This behaviour differs from other MCMC samplers that record only the accepted values. We do this just in case sequences that are still "good" (but not better than current) are rejected. The effect here is that we get a cluster of sequences that are one-apart from newly accepted sequences.

Secondly, we check the Hamming distance between the newly proposed sequences and the original. This corresponds to the "trust radius" specified in the jax-unirep paper. If the hamming distance > trust radius, we reject the sequence outright.

A dictionary containing the following key-value pairs are returned:

  • "sequences": All proposed sequences.
  • "scores": All scores from the scoring function.
  • "accept": Whether the sequence was accepted as the new 'current sequence' on which new sequences are proposed.

This can be turned into a pandas DataFrame.

Parameters

  • starter_sequence: The starting sequence.
  • n_steps: Number of steps for the MC chain to walk.
  • scoring_func: Scoring function for a new sequence. It should only accept a string sequence.
  • is_accepted_kwargs: Dictionary of kwargs to pass into is_accepted function. See is_accepted docstring for more details.
  • trust_radius: Maximum allowed number of mutations away from starter sequence.
  • propose_kwargs: Dictionary of kwargs to pass into propose function. See propose docstring for more details.
  • verbose: Whether or not to print iteration number and associated sequence + score. Defaults to False

Returns

A dictionary with sequences, accept and score as keys.