model¶

Defines Model objects for global epistasis modeling using JAX.

class multidms.model.Model(data: Data, loss_type: Literal['functional_score_loss', 'count_loss'] = 'functional_score_loss', ge_type: Literal['Identity', 'Sigmoid'] = 'Sigmoid', l2reg: float = 0.0, fusionreg: float = 0.0, beta0_ridge: float = 0.0, scale_fusion_by_n: bool = False, output_floor: float | None = None, output_floor_hinge: float = 0.1)¶

Bases: object

Model for global epistasis analysis of DMS experiments.

This class wraps the jaxmodels backend to provide a user-friendly interface for fitting global epistasis models to deep mutational scanning data.

Parameters:

data (multidms.Data) – Preprocessed DMS data object containing variants and functional scores.
loss_type ({'functional_score_loss', 'count_loss'}) – Type of loss function to use. ‘functional_score_loss’ for standard functional score fitting, ‘count_loss’ for count-based enrichment models.
ge_type ({'Identity', 'Sigmoid'}) – Global epistasis model type. ‘Identity’ for no global epistasis (linear), ‘Sigmoid’ for sigmoidal transformation.
l2reg (float) – L2 regularization strength for mutation effects (default: 0.0).
fusionreg (float) – Fusion regularization strength for shift parameters (default: 0.0).
beta0_ridge (float) – Ridge penalty for β0 offsets from reference condition (default: 0.0).

Example

>>> import pandas as pd
>>> from multidms import Data, Model
>>> df = pd.DataFrame({
...     'condition': ['a', 'a', 'b', 'b'],
...     'aa_substitutions': ['', 'M1A', '', 'M1A'],
...     'func_score': [0.0, 1.2, 0.1, 1.5]
... })
>>> data = Data(df, reference='a')  
>>> model = Model(data, ge_type='Sigmoid', l2reg=0.01)
>>> model  
Model(ge_type='Sigmoid', loss_type='functional_score_loss')

property data: Data¶: The Data object used for model fitting.

property params¶: Model parameters (available after fit).

property converged: bool¶

Whether the model fitting converged.

Convergence is determined by whether the objective error (relative change in the objective function) at the last iteration was below the tolerance used during fitting.

property convergence_trajectory_df: DataFrame¶

Convergence trajectory with diagnostics over iterations.

Returns:: One row per iteration with columns for overall objective (iteration, objective_total_trajectory, objective_error_trajectory, loss_trajectory, loss_per_variant_trajectory), per-condition loss (loss_{cond}, loss_per_variant_{cond}), block-level diagnostics (e.g. calibration_error, beta0_stepsize), shared alpha (alpha), per-condition parameters (beta0_{cond}, sparsity_{cond}), and theta_{cond} when count data is present.
Return type:: pd.DataFrame

fit(warmstart: bool = True, recompute_scale: bool = True, maxiter: int = 10, tol: float = 1e-06, beta0_init: dict = None, beta_init: dict = None, alpha_init: float | dict = None, share_alpha: bool = True, beta_clip_range: tuple = None, ge_kwargs: dict = None, cal_kwargs: dict = None, loss_kwargs: dict = None, verbose: bool = True)¶

Fit the model to data.

Parameters:

warmstart (bool) – Whether to use Ridge regression for parameter initialization (default: True).
recompute_scale (bool) – If True (default), recompute the objective normalizer each outer sweep (current behavior). If False, compute it once after warmstart and hold it constant — the fixed-scale convergence fix (#246).
maxiter (int) – Maximum number of optimization iterations (default: 10).
tol (float) – Convergence tolerance on objective function (default: 1e-6).
beta0_init (dict, optional) – Initial β0 values per condition.
beta_init (dict, optional) – Initial β values per condition.
alpha_init (float or dict, optional) – Initial α scaling value. Float applies to all conditions; dict maps condition names to per-condition values.
share_alpha (bool) – If True (default), optimize a single shared α. If False, each condition gets its own α.
beta_clip_range (tuple, optional) – Tuple of (min, max) values for clipping β parameters during optimization. Example: (-10.0, 10.0). If None, no clipping is applied.
ge_kwargs (dict, optional) – Keyword arguments for global epistasis optimizer (e.g., tol, maxiter, maxls).
cal_kwargs (dict, optional) – Keyword arguments for calibration (α) optimizer (e.g., tol, maxiter, maxls).
loss_kwargs (dict, optional) – Keyword arguments for the loss function (e.g., δ for Huber loss).
verbose (bool) – Whether to print progress information during fitting (default: True).

Returns:

Returns self for method chaining.

Return type:

self

get_mutations_df(phenotype_as_effect: bool = True, times_seen_threshold: int = 0) → DataFrame¶

Extract mutation-level parameters and predicted functional scores.

Parameters:

phenotype_as_effect (bool) – If True, report mutation effects. If False, report raw latent phenotypes.
times_seen_threshold (int) – Minimum number of times a mutation must be seen in ALL conditions to be included. Default is 0 (no filtering).

Returns:

DataFrame with mutations as rows (index) and columns:

beta_{condition} for each condition
shift_{condition} for each non-reference condition
predicted_func_score_{condition} for each condition

Shift parameters represent the difference in beta values between each condition and the reference condition. Predicted functional scores are the model’s predictions for each single mutation on its condition-specific wild-type background.

Return type:

pd.DataFrame

Example

For a model with conditions [‘a’, ‘b’] where ‘a’ is reference, the returned columns are: beta_a, beta_b, shift_b, predicted_func_score_a, predicted_func_score_b. One row per mutation.

get_variants_df(phenotype_as_effect: bool = True) → DataFrame¶

Extract variant-level predictions.

Parameters:

phenotype_as_effect (bool) – If True, report effects. If False, report raw latent phenotypes.

Returns:

Variant-level predictions merged with original data. Includes columns:

predicted_func_score: model-predicted functional score, i.e. α * (g(φ(X)) - g(φ(x_wt)))
predicted_latent: latent phenotype φ(X)
predicted_fitness: predicted fitness in g(φ) space, i.e. predicted_func_score / α + g(φ(x_wt))
measured_fitness: measured fitness in g(φ) space, i.e. func_score / α + g(φ(x_wt))

Return type:

pd.DataFrame

property training_loss: dict¶

Per-condition and total loss on training data.

Returns:: Dictionary mapping condition names and "total" to their training loss values.
Return type:: dict[str, float]
Raises:: ValueError – If model has not been fitted.

eval_loss(df)¶

Evaluate the model’s loss on an arbitrary DataFrame.

Parameters:: df (pd.DataFrame) – DataFrame with columns ‘condition’, ‘aa_substitutions’, ‘func_score’.
Returns:: Per-condition losses and "total" loss.
Return type:: dict[str, float]
Raises:: ValueError – If model is not fitted, required columns are missing, conditions are invalid, or substitutions contain unseen mutations.

add_phenotypes_to_df(df: DataFrame, substitutions_col: str = 'aa_substitutions', condition_col: str = 'condition', predicted_phenotype_col: str = 'predicted_func_score', overwrite_cols: bool = False) → DataFrame¶

Add model predictions to a DataFrame of variants.

Parameters:

df (pd.DataFrame) – DataFrame with columns specified by condition_col and substitutions_col. Additional columns will be preserved in output.
substitutions_col (str) – Column in df giving variants as substitution strings. Default is ‘aa_substitutions’.
condition_col (str) – Column in df giving the condition for each variant. Values must exist in the model’s conditions. Default is ‘condition’.
predicted_phenotype_col (str) – Name of column to add containing predicted functional scores. Default is ‘predicted_func_score’.
overwrite_cols (bool) – If the specified predicted phenotype column already exists in df, overwrite it? If False, raise an error.

Returns:

A copy of df with predictions added. Always includes:

predicted_func_score (or custom name): predicted functional score
predicted_latent: latent phenotype φ(X)
predicted_fitness: predicted fitness in g(φ) space

If func_score column is present in df, also includes:

measured_fitness: measured fitness in g(φ) space

Return type:

pd.DataFrame

Raises:

ValueError – If model is not fitted, required columns are missing, indices are not unique, conditions are invalid, or substitutions contain mutations not seen during training.

Example

>>> import pandas as pd
>>> from multidms import Data, Model
>>> df_train = pd.DataFrame({
...     'condition': ['a', 'a', 'b', 'b'],
...     'aa_substitutions': ['', 'M1A', '', 'M1A'],
...     'func_score': [0.0, 1.2, 0.1, 1.5]
... })
>>> data = Data(df_train, reference='a')  
>>> model = Model(data, ge_type='Identity', l2reg=0.01)
>>> _ = model.fit(maxiter=5, warmstart=False, verbose=False)
>>> df_new = pd.DataFrame({
...     'condition': ['a', 'b'],
...     'aa_substitutions': ['M1A', 'M1A']
... })
>>> result = model.add_phenotypes_to_df(df_new)
>>> 'predicted_func_score' in result.columns
True
>>> 'predicted_latent' in result.columns
True
>>> 'predicted_fitness' in result.columns
True
>>> len(result)
2

get_ge_landscape_df(n_curve_points: int = 200, space: str = 'fitness') → tuple¶

Get data for plotting the global epistasis landscape.

Returns a tuple of (variants_df, curve_df). variants_df contains per-variant latent phenotype and fitness columns from get_variants_df(), plus a wildtype_latent column for reference-line plotting. The curve DataFrame depends on space.

Parameters:

n_curve_points (int) – Number of points for the curve grid.
space (str) – Which landscape to build. "fitness" (default) returns the shared global-epistasis curve g(φ). "func_score" returns one predicted-functional-score curve per condition, α · (g(φ) − g(φ_wt)).

Returns:

(variants_df, curve_df) where:

variants_df has all columns from get_variants_df() plus wildtype_latent.
For space="fitness": curve_df has columns predicted_latent and ge_curve_value.
For space="func_score": curve_df is long-form with columns condition, predicted_latent, and func_score_curve_value (n_conditions × n_curve_points rows).

Return type:

tuple[pd.DataFrame, pd.DataFrame]

get_ge_params_df() → DataFrame¶

Per-condition parameters that place the GE landscape.

The wildtype latent phenotype decomposes exactly into the condition’s intercept plus the summed effect of its bundle mutations (those separating the condition’s wildtype sequence from the reference wildtype sequence):

φ_wt(d) = β0(d) + Σ_{m ∈ bundle(d)} β(d)[m]

The reference condition has an empty bundle, so its bundle_sum is zero and wildtype_latent equals beta0.

Returns:

One row per condition, in Data.conditions order, with columns:

condition : str Condition name.
alpha : float The fitness-to-functional-score scale α for this condition. Constant down the column when the model was fit with share_alpha=True (the default). Note α scales the curve and does not affect wildtype placement.
beta0 : float The condition’s latent intercept β0.
bundle_sum : float Summed β over the condition’s bundle mutations.
wildtype_latent : float beta0 + bundle_sum; equals Model.wildtype_latent.
n_bundle_mutations : int Number of bundle mutations for the condition.

Return type:

pandas.DataFrame

Raises:

ValueError – If the model has not been fitted.

plot_ge_landscape(n_curve_points=200, space='fitness', **kwargs)¶

Plot the global epistasis landscape.

Convenience wrapper that calls get_ge_landscape_df() and delegates to multidms.plot.ge_landscape().

Parameters:

n_curve_points (int) – Number of points for the curve grid. Default 200.
space (str) – "fitness" (default) plots the shared g(φ) curve; "func_score" plots one predicted-functional-score curve per condition. Passed to both get_ge_landscape_df() and multidms.plot.ge_landscape().
**kwargs – Passed to multidms.plot.ge_landscape(). The parameter annotation is on by default here: unless params_df is given explicitly or annotate_params=False, this method supplies get_ge_params_df() so the chart is annotated with α, β0 and the bundle sum.

Returns:

Interactive Altair chart.

Return type:

alt.LayerChart

get_ge_curve(grid_min: float = -5.0, grid_max: float = 5.0, n_points: int = 200) → DataFrame¶

Evaluate the global epistasis function over a latent phenotype grid.

Parameters:

grid_min (float) – Minimum latent phenotype value for the grid.
grid_max (float) – Maximum latent phenotype value for the grid.
n_points (int) – Number of points in the grid.

Returns:

DataFrame with columns ‘latent’ and ‘observed’.

Return type:

pd.DataFrame

property wildtype_latent: dict¶

Wildtype latent phenotype for each condition.

Returns:: Dictionary mapping condition names to the wildtype’s latent phenotype value in that condition.
Return type:: dict[str, float]

model¶

multidms

Navigation

Related Topics