jaxmodels

A simple API for global epistasis modeling.

class multidms.jaxmodels.Data(x_wt: Int[Array, 'n_mutations'], X: Int[Array, 'n_variants n_mutations'], functional_scores: Float[Array, 'n_variants'], pre_count_wt: Int[Array, ''] | None = None, post_count_wt: Int[Array, ''] | None = None, pre_counts: Int[Array, ' n_variants'] | None = None, post_counts: Int[Array, ' n_variants'] | None = None)

Bases: Module

Data for a DMS experiment.

x_wt: Int[Array, 'n_mutations']

Binary encoding of the wildtype sequence.

X: Int[Array, 'n_variants n_mutations']

Variant encoding matrix (sparse format).

functional_scores: Float[Array, 'n_variants']

Functional scores for each variant.

pre_count_wt: Int[Array, ''] | None = None

Wildtype pre-selection count (optional).

post_count_wt: Int[Array, ''] | None = None

Wildtype post-selection count (optional).

pre_counts: Int[Array, 'n_variants'] | None = None

Pre-selection counts for each variant (optional).

post_counts: Int[Array, 'n_variants'] | None = None

Post-selection counts for each variant (optional).

static from_multidms(multidms_data: Data, condition: str) → Self

Create data from a multidms data object.

Parameters:

  • multidms_data – The data to use. Note the WT must be the first variant in each condition.

  • condition – The condition to extract data for.

Returns:

Data object with count data if available in the source.

class multidms.jaxmodels.Latent(β0: Any, β: Any)

Bases: Module

Model a latent phenotype.

β0: Float[Array, '']

Intercept.

β: Float[Array, 'n_mutations']

Mutation effects.

static from_params(β0: Float, β: Float[Array, 'n_mutations']) → Self

Create a latent model from explicit parameters.

Parameters:

  • β0 – Intercept value.

  • β – Mutation effects array.

Returns:

Latent model with specified parameters.

static zeros(n_mutations: int, β0: Float = 0.0) → Self

Create a zero-initialized latent model with optional intercept.

Parameters:

  • n_mutations – Number of mutations.

  • β0 – Intercept value (default: 0.0).

Returns:

Latent model with β set to zeros and specified β0.

static warmstart(data: Data, l2reg: float = 0.0) → Self

Warmstart the latent model.

Parameters:

  • data – Data to initialize the model for.

  • l2reg – L2 regularization strength for warmstart.

Returns:

Latent model initialized with warmstart parameters.

class multidms.jaxmodels.GlobalEpistasis

Bases: Module, ABC

Global epistasis model.

class multidms.jaxmodels.Identity

Bases: GlobalEpistasis

Identity function.

class multidms.jaxmodels.Sigmoid

Bases: GlobalEpistasis

Sigmoid function.

class multidms.jaxmodels.OutputActivation

Bases: Module, ABC

Activation applied to the model’s predicted functional score.

class multidms.jaxmodels.IdentityOutput

Bases: OutputActivation

Pass-through output activation (no floor). The default.

class multidms.jaxmodels.Softplus(lower_bound: float = -3.5, hinge_scale: float = 0.1)

Bases: OutputActivation

Smooth differentiable lower floor on the predicted functional score.

Reproduces the manuscript-era biophysical.py::softplus_activation:

\[t(y) = \lambda \cdot \log\!\left(1 + e^{(y - l)/\lambda}\right) + l\]

where \(l\) is lower_bound and \(\lambda\) is hinge_scale. For \(y \gg l\) it is approximately the identity; for \(y \ll l\) it flattens toward \(l\) (from above; \(t(y) > l\) for all finite \(y\)). The floor is soft — at \(y = l\), \(t = l + \lambda\log 2\), not exactly \(l\).

lower_bound: float = -3.5
hinge_scale: float = 0.1
class multidms.jaxmodels.Model(φ: dict[str, Latent], α: Float[Array, ''] | dict[str, Float[Array, '']], logθ: dict[str, Float[Array, '']], reference_condition: str, global_epistasis: GlobalEpistasis = Identity(), output_activation: OutputActivation = IdentityOutput())

Bases: Module

Model DMS data.

φ: dict[str, multidms.jaxmodels.Latent]

Latent models for each condition.

α: Float[Array, ''] | dict[str, jaxtyping.Float[Array, '']]

Fitness-functional score scaling factor.

logθ: dict[str, jaxtyping.Float[Array, '']]

Overdispersion parameter for each condition.

reference_condition: str

The condition to use as a reference.

global_epistasis: GlobalEpistasis = Identity()
output_activation: OutputActivation = IdentityOutput()
predict_score(data_sets: dict[str, multidms.jaxmodels.Data]) → dict[str, jaxtyping.Float[Array, 'n_variants']]

Predict functional scores, interpreted as \(\log_e\) enrichment wrt WT.

Parameters:

data_sets – Data sets for each condition.

predict_post_count(data_sets: dict[str, multidms.jaxmodels.Data]) → dict[str, jaxtyping.Float[Array, 'n_variants']]

Predict post-selection counts.

Parameters:

data_sets – Data sets for each condition.

multidms.jaxmodels.count_loss(model: Model, data_sets: dict[str, multidms.jaxmodels.Data]) → dict[str, jaxtyping.Float[Array, '']]

Count-based loss.

Parameters:

  • model – Model to evaluate.

  • data_sets – Data sets for each condition.

Returns:

Loss for each condition.

multidms.jaxmodels.functional_score_loss(model: Model, data_sets: dict[str, multidms.jaxmodels.Data], δ: Float = 1.0) → dict[str, jaxtyping.Float[Array, '']]

Huber loss on functional scores.

Returns mean Huber loss per variant for each condition, so that conditions contribute equally to the total objective regardless of variant count.

Parameters:

  • model – Model to evaluate.

  • data_sets – Data sets for each condition.

  • δ – Huber loss parameter.

Returns:

Mean loss for each condition.

multidms.jaxmodels.fit(data_sets: dict[str, multidms.jaxmodels.Data], reference_condition: str, l2reg: ~jaxtyping.Float = 0.0, fusionreg: ~jaxtyping.Float = 0.0, beta0_ridge: ~jaxtyping.Float = 0.0, scale_fusion_by_n: bool = False, block_iters: int = 10, block_tol: ~jaxtyping.Float = 1e-06, ge_kwargs: dict[str, typing.Any] = {}, cal_kwargs: dict[str, typing.Any] = {}, global_epistasis: ~multidms.jaxmodels.GlobalEpistasis = Identity(), output_activation: ~multidms.jaxmodels.OutputActivation = IdentityOutput(), loss_fn: ~typing.Callable[[~multidms.jaxmodels.Model, dict[str, multidms.jaxmodels.Data]], dict[str, jaxtyping.Float[Array, '']]] = <function functional_score_loss>, loss_kwargs: dict[str, typing.Any] = {'δ': 1.0}, warmstart: bool = True, recompute_scale: bool = True, beta0_init: dict[str, jaxtyping.Float] | None = None, beta_init: dict[str, jaxtyping.Float[Array, 'n_mutations']] | None = None, alpha_init: ~jaxtyping.Float | dict[str, jaxtyping.Float] | None = None, share_alpha: bool = True, beta_clip_range: tuple[jaxtyping.Float, jaxtyping.Float] | None = None, verbose: bool = True) → tuple[multidms.jaxmodels.Model, pandas.DataFrame]

Fit a model to data.

Parameters:

  • data_sets – Data to fit to. Each key is a condition.

  • reference_condition – The condition to use as a reference.

  • l2reg – L2 regularization strength for mutation effects.

  • fusionreg – Fusion (shift lasso) regularization strength.

  • beta0_ridge – Ridge penalty for β0 differences from reference condition.

  • scale_fusion_by_n – If True, weight each condition’s fusion penalty by n_ref / n_d, reducing shrinkage for data-poor conditions.

  • block_iters – Number iterations for block coordinate descent.

  • block_tol – Tolerance on objective function for block coordinate descent.

  • ge_kwargs – Keyword arguments for the global epistasis model optimizer.

  • cal_kwargs – Keyword arguments for the experimental calibration parameter optimizer.

  • global_epistasis – Global epistasis model.

  • output_activation – Output activation applied to predicted functional scores. Defaults to IdentityOutput (no floor).

  • loss_fn – Loss function.

  • loss_kwargs – Keyword arguments for the loss function.

  • warmstart – Whether to use Ridge regression warmstart (default: True). If True, performs Ridge regression to initialize parameters. The warmstart values will be overridden by any explicit values provided in beta0_init or beta_init.

  • recompute_scale – If True (default), recompute the objective scale normalizer at the start of every outer sweep (current behavior). If False, compute scale once after warmstart and hold it constant, so the lasso threshold (fusionreg / scale) is stationary and objective_error measures a true relative change in the scaled objective across sweeps. See issue #246.

  • beta0_init – Initial β0 (intercept) values for each condition. If None, uses zeros (or warmstart values if warmstart=True). If dict provided, uses those values for specified conditions.

  • beta_init – Initial β (mutation effects) values for each condition. If None, uses zeros (or warmstart values if warmstart=True). If dict provided, uses those values for specified conditions.

  • alpha_init – Initial α (fitness-functional score scaling) value. Float applies to all conditions; dict maps condition names to per-condition values. If None, uses 1.0.

  • share_alpha – If True (default), optimize a single shared α across all conditions. If False, each condition gets its own α.

  • beta_clip_range – Optional tuple of (min, max) values for clipping β parameters. If None, no clipping is applied. Example: (-10.0, 10.0). This constrains mutation effect parameters during optimization to prevent extreme values.

  • verbose – Whether to print progress information during fitting (default: True). If False, suppresses all print output.

Returns:

Tuple of (fitted model, convergence trajectory DataFrame).

The DataFrame has one row per outer iteration with columns:

  • iteration, objective_total_trajectory, objective_error_trajectory, loss_trajectory, loss_per_variant_trajectory

  • Per-condition loss: loss_{condition} (mean Huber loss per variant for that condition) and loss_per_variant_{condition} (identical to loss_{condition} since the loss is already per-variant). Per-condition losses sum to loss_trajectory. loss_per_variant_trajectory is loss_trajectory divided by the number of conditions (average per-condition mean loss).

  • Block-level diagnostics for each optimization block (calibration_error, calibration_stepsize, calibration_iter_num, beta0_error, etc.)

  • Per-condition parameters: alpha_{condition}, theta_{condition}, beta0_{condition}, sparsity_{condition} (non-reference only)