Fix NameError: move import math to top of model setup cell

When residual_library_encoder=False, the if-block containing import math
was skipped, causing NameError on math.exp in the library bias print.

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>

Author: vitkl

docs/notebooks/model_comparisons/bone_marrow_multimodal_tutorial_early_stopping.ipynb

@@ -908,88 +908,7 @@
     "tags": []
    },
    "outputs": [],
-   "source": [
-    "model = regularizedvi.RegularizedMultimodalVI(\n",
-    "    mdata,\n",
-    "    # Per-modality architecture sizes\n",
-    "    n_hidden={\"rna\": 512, \"atac\": 256},\n",
-    "    n_latent={\"rna\": 128, \"atac\": 64},\n",
-    "    n_layers=1,\n",
-    "    # Z combination: concatenation preserves modality-specific signal\n",
-    "    latent_mode=\"concatenation\",\n",
-    "    # Per-modality flags\n",
-    "    additive_background_modalities=[\"rna\"],  # ambient correction for RNA only\n",
-    "    feature_scaling_modalities=[\"rna\", \"atac\"] if use_feature_scaling else [],  # per-covariate scaling\n",
-    "    # Dispersion\n",
-    "    dispersion=\"gene-batch\",\n",
-    "    regularise_dispersion=True,\n",
-    "    regularise_dispersion_prior=regularise_dispersion_prior,\n",
-    "    dispersion_hyper_prior_alpha=dispersion_hyper_prior_alpha,\n",
-    "    dispersion_hyper_prior_beta=dispersion_hyper_prior_beta,\n",
-    "    # Background prior\n",
-    "    additive_bg_prior_alpha=additive_bg_prior_alpha,\n",
-    "    additive_bg_prior_beta=additive_bg_prior_beta,\n",
-    "    regularise_background=regularise_background,\n",
-    "    # Library centering\n",
-    "    library_log_means_centering_sensitivity=library_log_means_centering_sensitivity,\n",
-    "    library_log_vars_weight=library_log_vars_weight,\n",
-    "    # Learnable modality scaling\n",
-    "    learnable_modality_scaling=learnable_modality_scaling,\n",
-    "    modality_scale_prior_concentration=modality_scale_prior_concentration,\n",
-    "    # Pearson correlation metrics\n",
-    "    compute_pearson=compute_pearson,\n",
-    "    # Decoder regularisation and initialisation\n",
-    "    decoder_weight_l2=decoder_weight_l2,\n",
-    "    decoder_cov_weight_l2=decoder_cov_weight_l2,\n",
-    "    init_decoder_bias=init_decoder_bias,\n",
-    "    bg_init_gene_fraction=bg_init_gene_fraction,\n",
-    "    decoder_bias_multiplier=_decoder_bias_multiplier,\n",
-    "    # Residual library encoder\n",
-    "    residual_library_encoder=residual_library_encoder,\n",
-    "    library_obs_w_prior_rate=library_obs_w_prior_rate,\n",
-    ")\n",
-    "\n",
-    "print(model)\n",
-    "print(f\"\\nTotal latent dim: {model.module.total_latent_dim}\")\n",
-    "print(f\"  RNA Z dim: {model.module.n_latent_dict['rna']}\")\n",
-    "print(f\"  ATAC Z dim: {model.module.n_latent_dict['atac']}\")\n",
-    "print(f\"\\nRegion factors modalities: {model.module.feature_scaling_modalities}\")\n",
-    "for name in model.module.feature_scaling:\n",
-    "    print(f\"  {name} feature scaling shape: {model.module.feature_scaling[name].shape}\")\n",
-    "print(f\"\\nRegularise background: {regularise_background}\")\n",
-    "if regularise_background:\n",
-    "    print(f\"  Prior: Gamma({additive_bg_prior_alpha}, {additive_bg_prior_beta})\")\n",
-    "    print(f\"  Prior mean: {additive_bg_prior_alpha / additive_bg_prior_beta:.4f}\")\n",
-    "print(f\"Library centering sensitivity: {library_log_means_centering_sensitivity}\")\n",
-    "print(f\"Library log vars weight: {library_log_vars_weight}\")\n",
-    "print(f\"Decoder weight L2: {decoder_weight_l2}\")\n",
-    "print(f\"Decoder cov weight L2: {decoder_cov_weight_l2}\")\n",
-    "print(f\"Init decoder bias: {init_decoder_bias}\")\n",
-    "print(f\"BG init gene fraction: {bg_init_gene_fraction}\")\n",
-    "print(f\"Decoder bias multiplier: {_decoder_bias_multiplier}\")\n",
-    "print(f\"Dispersion prior: Gamma({dispersion_hyper_prior_alpha:.1f}, {dispersion_hyper_prior_beta:.2f})\")\n",
-    "print(f\"Residual library encoder: {residual_library_encoder}\")\n",
-    "print(f\"Library obs w prior rate: {library_obs_w_prior_rate}\")\n",
-    "print(f\"Learnable modality scaling: {learnable_modality_scaling}\")\n",
-    "if learnable_modality_scaling:\n",
-    "    for name in model.module.modality_scale_init:\n",
-    "        print(f\"  {name} init: {model.module.modality_scale_init[name]:.3f}\")\n",
-    "print(f\"Compute Pearson: {compute_pearson}\")\n",
-    "\n",
-    "# Library encoder bias inspection\n",
-    "if hasattr(model.module, \"library_obs_w_mu\"):\n",
-    "    import math\n",
-    "    from torch.distributions import LogNormal\n",
-    "    import torch\n",
-    "\n",
-    "    for name in model.module.modality_names:\n",
-    "        w_mu = model.module.library_obs_w_mu[name]\n",
-    "        w_sigma = torch.exp(model.module.library_obs_w_log_sigma[name])\n",
-    "        print(f\"  library_obs_w_{name}: E[w]={LogNormal(w_mu, w_sigma).mean.item():.3f}\")\n",
-    "for name in model.module.modality_names:\n",
-    "    bias_val = model.module.l_encoders[name].mean_encoder.bias.item()\n",
-    "    print(f\"  l_encoder[{name}].bias = {bias_val:.4f} (exp = {math.exp(bias_val):.4f})\")"
-   ]
+   "source": "import math\n\nmodel = regularizedvi.RegularizedMultimodalVI(\n    mdata,\n    # Per-modality architecture sizes\n    n_hidden={\"rna\": 512, \"atac\": 256},\n    n_latent={\"rna\": 128, \"atac\": 64},\n    n_layers=1,\n    # Z combination: concatenation preserves modality-specific signal\n    latent_mode=\"concatenation\",\n    # Per-modality flags\n    additive_background_modalities=[\"rna\"],  # ambient correction for RNA only\n    feature_scaling_modalities=[\"rna\", \"atac\"] if use_feature_scaling else [],  # per-covariate scaling\n    # Dispersion\n    dispersion=\"gene-batch\",\n    regularise_dispersion=True,\n    regularise_dispersion_prior=regularise_dispersion_prior,\n    dispersion_hyper_prior_alpha=dispersion_hyper_prior_alpha,\n    dispersion_hyper_prior_beta=dispersion_hyper_prior_beta,\n    # Background prior\n    additive_bg_prior_alpha=additive_bg_prior_alpha,\n    additive_bg_prior_beta=additive_bg_prior_beta,\n    regularise_background=regularise_background,\n    # Library centering\n    library_log_means_centering_sensitivity=library_log_means_centering_sensitivity,\n    library_log_vars_weight=library_log_vars_weight,\n    # Learnable modality scaling\n    learnable_modality_scaling=learnable_modality_scaling,\n    modality_scale_prior_concentration=modality_scale_prior_concentration,\n    # Pearson correlation metrics\n    compute_pearson=compute_pearson,\n    # Decoder regularisation and initialisation\n    decoder_weight_l2=decoder_weight_l2,\n    decoder_cov_weight_l2=decoder_cov_weight_l2,\n    init_decoder_bias=init_decoder_bias,\n    bg_init_gene_fraction=bg_init_gene_fraction,\n    decoder_bias_multiplier=_decoder_bias_multiplier,\n    # Residual library encoder\n    residual_library_encoder=residual_library_encoder,\n    library_obs_w_prior_rate=library_obs_w_prior_rate,\n)\n\nprint(model)\nprint(f\"\\nTotal latent dim: {model.module.total_latent_dim}\")\nprint(f\"  RNA Z dim: {model.module.n_latent_dict['rna']}\")\nprint(f\"  ATAC Z dim: {model.module.n_latent_dict['atac']}\")\nprint(f\"\\nRegion factors modalities: {model.module.feature_scaling_modalities}\")\nfor name in model.module.feature_scaling:\n    print(f\"  {name} feature scaling shape: {model.module.feature_scaling[name].shape}\")\nprint(f\"\\nRegularise background: {regularise_background}\")\nif regularise_background:\n    print(f\"  Prior: Gamma({additive_bg_prior_alpha}, {additive_bg_prior_beta})\")\n    print(f\"  Prior mean: {additive_bg_prior_alpha / additive_bg_prior_beta:.4f}\")\nprint(f\"Library centering sensitivity: {library_log_means_centering_sensitivity}\")\nprint(f\"Library log vars weight: {library_log_vars_weight}\")\nprint(f\"Decoder weight L2: {decoder_weight_l2}\")\nprint(f\"Decoder cov weight L2: {decoder_cov_weight_l2}\")\nprint(f\"Init decoder bias: {init_decoder_bias}\")\nprint(f\"BG init gene fraction: {bg_init_gene_fraction}\")\nprint(f\"Decoder bias multiplier: {_decoder_bias_multiplier}\")\nprint(f\"Dispersion prior: Gamma({dispersion_hyper_prior_alpha:.1f}, {dispersion_hyper_prior_beta:.2f})\")\nprint(f\"Residual library encoder: {residual_library_encoder}\")\nprint(f\"Library obs w prior rate: {library_obs_w_prior_rate}\")\nprint(f\"Learnable modality scaling: {learnable_modality_scaling}\")\nif learnable_modality_scaling:\n    for name in model.module.modality_scale_init:\n        print(f\"  {name} init: {model.module.modality_scale_init[name]:.3f}\")\nprint(f\"Compute Pearson: {compute_pearson}\")\n\n# Library encoder bias inspection\nif hasattr(model.module, \"library_obs_w_mu\"):\n    from torch.distributions import LogNormal\n    import torch\n\n    for name in model.module.modality_names:\n        w_mu = model.module.library_obs_w_mu[name]\n        w_sigma = torch.exp(model.module.library_obs_w_log_sigma[name])\n        print(f\"  library_obs_w_{name}: E[w]={LogNormal(w_mu, w_sigma).mean.item():.3f}\")\nfor name in model.module.modality_names:\n    bias_val = model.module.l_encoders[name].mean_encoder.bias.item()\n    print(f\"  l_encoder[{name}].bias = {bias_val:.4f} (exp = {math.exp(bias_val):.4f})\")"
   },
   {
    "cell_type": "markdown",