diff --git a/docs/api/metrics.md b/docs/api/metrics.md
index 3f74eee..32b71b8 100644
--- a/docs/api/metrics.md
+++ b/docs/api/metrics.md
@@ -1,153 +1,9 @@
-# Metrics API
+# Metrics
 
-The `tcri.metrics` module (imported as `tcri.tl`) provides functions for calculating information-theoretic metrics on paired single-cell RNA and TCR sequencing data.
+Information-theoretic metrics over the clone–phenotype joint distribution:
+entropies, mutual information, flux, and related summaries. Exposed as
+``tcri.tl``.
 
-## clonotypic_entropy
-
-```python
-def clonotypic_entropy(adata, covariate, phenotype, temperature=1.0):
-    """
-    Calculate the clonotypic entropy for each value of the covariate.
-    
-    Parameters
-    ----------
-    adata : AnnData
-        AnnData object with model results registered
-    covariate : str
-        Name of the covariate in adata.obs
-    phenotype : str
-        Name of the phenotype in adata.obs
-    temperature : float, default=1.0
-        Temperature parameter for softening/sharpening distributions
-        
-    Returns
-    -------
-    dict
-        Dictionary mapping covariate values to entropy values
-    """
-    pass
-```
-
-## phenotypic_entropy
-
-```python
-def phenotypic_entropy(adata, covariate, clonotype, temperature=1.0):
-    """
-    Calculate the phenotypic entropy for each value of the covariate.
-    
-    Parameters
-    ----------
-    adata : AnnData
-        AnnData object with model results registered
-    covariate : str
-        Name of the covariate in adata.obs
-    clonotype : str
-        Name of the clonotype field in adata.obs
-    temperature : float, default=1.0
-        Temperature parameter for softening/sharpening distributions
-        
-    Returns
-    -------
-    dict
-        Dictionary mapping covariate values to entropy values
-    """
-    pass
-```
-
-## mutual_information
-
-```python
-def mutual_information(adata, covariate, temperature=1.0, weighted=False):
-    """
-    Calculate the mutual information between phenotypes and TCR clonotypes
-    for each value of the covariate.
-    
-    Parameters
-    ----------
-    adata : AnnData
-        AnnData object with model results registered
-    covariate : str
-        Name of the covariate in adata.obs
-    temperature : float, default=1.0
-        Temperature parameter for softening/sharpening distributions
-    weighted : bool, default=False
-        Whether to weight by clone size
-        
-    Returns
-    -------
-    dict
-        Dictionary mapping covariate values to mutual information values
-    """
-    pass
-```
-
-## clonality
-
-```python
-def clonality(adata):
-    """
-    Calculate clonality metrics for the data.
-    
-    Parameters
-    ----------
-    adata : AnnData
-        AnnData object with TCR information
-        
-    Returns
-    -------
-    dict
-        Dictionary containing clonality metrics
-    """
-    pass
-```
-
-## flux
-
-```python
-def flux(adata, from_this, to_that, clones=None, temperature=1.0):
-    """
-    Calculate phenotypic flux between two covariate values.
-    
-    Parameters
-    ----------
-    adata : AnnData
-        AnnData object with model results registered
-    from_this : str
-        Starting covariate value
-    to_that : str
-        Ending covariate value
-    clones : list, optional
-        List of clone IDs to include
-    temperature : float, default=1.0
-        Temperature parameter for softening/sharpening distributions
-        
-    Returns
-    -------
-    numpy.ndarray
-        Flux matrix of shape (n_phenotypes, n_phenotypes)
-    """
-    pass
-```
-
-## Usage Examples
-
-```python
-import tcri
-import scanpy as sc
-
-# Load data
-adata = sc.read_h5ad("your_data.h5ad")
-
-# Initialize and train model
-model = tcri.TCRIModel(adata)
-model.train()
-tcri.pp.register_model(adata, model)
-
-# Calculate metrics
-mi = tcri.tl.mutual_information(adata, "timepoint")
-entropy = tcri.tl.clonotypic_entropy(adata, "timepoint", "phenotype")
-clonality = tcri.tl.clonality(adata)
-
-# Calculate flux between timepoints
-flux_matrix = tcri.tl.flux(adata, from_this="T1", to_that="T2")
+```{eval-rst}
+.. automodule:: tcri.metrics._metrics
 ```
diff --git a/docs/api/model.md b/docs/api/model.md
index 0002f2f..e630e92 100644
--- a/docs/api/model.md
+++ b/docs/api/model.md
@@ -1,170 +1,14 @@
-# Model API
+# Model
 
-## TCRIModel
+The deep-learning model that jointly embeds gene expression and clonotype
+information and learns the clone–phenotype distribution.
 
-The `TCRIModel` class implements a hierarchical Bayesian model for analyzing TCR and gene expression data.
-
-```python
-class TCRIModel:
-    """
-    TCRi Model for joint analysis of gene expression and TCR data.
-    
-    This model implements a hierarchical Bayesian framework that learns
-    a joint representation of gene expression and TCR sequences.
-    
-    Parameters
-    ----------
-    adata : AnnData
-        AnnData object containing gene expression and TCR information
-    n_latent : int, default=10
-        Dimension of the latent space
-    n_hidden : int, default=128
-        Number of hidden units in the neural networks
-    global_scale : float, default=10.0
-        Scale parameter for the global prior
-    local_scale : float, default=5.0
-        Scale parameter for the local prior
-    prior_temperature : float, default=1.0
-        Temperature for sharpening the clone-phenotype prior distributions
-    guide_temperature : float, default=1.0
-        Temperature for sharpening learned parameters in the guide and get_p_ct()
-    use_enumeration : bool, default=False
-        Whether to use enumeration for discrete variables
-    device : str, optional
-        Device to use for computation ("cpu" or "cuda")
-    """
-    
-    def train(self, max_epochs=50, batch_size=128, lr=1e-3, 
-    margin_scale=0.0, margin_value=2.0, adaptive_margin=False, 
-    reconstruction_loss_scale=1e-2, n_steps_kl_warmup=1000):
-        """
-        Train the model.
-        
-        Parameters
-        ----------
-        max_epochs : int, default=50
-            Maximum number of epochs to train for
-        batch_size : int, default=128
-            Batch size for training
-        lr : float, default=1e-3
-            Learning rate
-        margin_scale : float, default=0.0
-            Scale for the margin loss
-        margin_value : float, default=2.0
-            Value for the margin
-        adaptive_margin : bool, default=False
-            Whether to use adaptive margin
-        reconstruction_loss_scale : float, default=1e-2
-            Scale for the reconstruction loss
-        n_steps_kl_warmup : int, default=1000
-            Number of steps for KL warmup
-        """
-        pass
-    
-    def get_latent_representation(self, adata=None, batch_size=256):
-        """
-        Get the latent representation for the data.
-        
-        Parameters
-        ----------
-        adata : AnnData, optional
-            AnnData object to get latent representation for.
-            If None, uses the training data.
-        batch_size : int, default=256
-            Batch size for inference
-            
-        Returns
-        -------
-        ndarray
-            Latent representation of shape (n_cells, n_latent)
-        """
-        pass
-    
-    def get_phenotype_probabilities(self, adata=None, batch_size=256):
-        """
-        Get phenotype probabilities for the data.
-        
-        Parameters
-        ----------
-        adata : AnnData, optional
-            AnnData object to get probabilities for.
-            If None, uses the training data.
-        batch_size : int, default=256
-            Batch size for inference
-            
-        Returns
-        -------
-        ndarray
-            Phenotype probabilities of shape (n_cells, n_phenotypes)
-        """
-        pass
-    
-    def save(self, path):
-        """
-        Save the model to a file.
-        
-        Parameters
-        ----------
-        path : str
-            Path to save the model to
-        """
-        pass
-    
-    @classmethod
-    def load(cls, path, adata=None):
-        """
-        Load a model from a file.
-        
-        Parameters
-        ----------
-        path : str
-            Path to load the model from
-        adata : AnnData, optional
-            AnnData object to use with the model
-            
-        Returns
-        -------
-        TCRIModel
-            Loaded model
-        """
-        pass
+```{eval-rst}
+.. autoclass:: tcri.model._model.TCRIModel
 ```
 
-## Usage Example
-
-```python
-import tcri
-import scanpy as sc
-
-# Load data
-adata = sc.read_h5ad("your_data.h5ad")
-
-# Initialize model
-model = tcri.TCRIModel(
-    adata,
-    n_latent=10,
-    n_hidden=128,
-    global_scale=10.0,
-    local_scale=5.0
-)
-
-# Train model
-model.train(
-    max_epochs=50,
-    batch_size=128,
-    lr=1e-3,
-    reconstruction_loss_scale=1e-2
-)
-
-# Get latent representations
-latent_z = model.get_latent_representation(adata)
-
-# Get phenotype probabilities
-probs = model.get_phenotype_probabilities(adata)
-
-# Save model
-model.save("your_model.pkl")
-
-# Load model
-loaded_model = tcri.TCRIModel.load("your_model.pkl", adata)
+```{note}
+Model persistence is handled by the session helpers in the
+[Utilities](utils.md) API — ``save_tcri_session`` and ``load_tcri_session`` —
+not by methods on the model object.
 ```
diff --git a/docs/api/plotting.md b/docs/api/plotting.md
index d195117..d19035c 100644
--- a/docs/api/plotting.md
+++ b/docs/api/plotting.md
@@ -1,180 +1,8 @@
-# Plotting API
+# Plotting
 
-The `tcri.plotting` module (imported as `tcri.pl`) provides functions for visualizing the results of TCRi analyses.
+Visualization functions for clonal dynamics, phenotype probabilities, and
+information-theoretic summaries. Exposed as ``tcri.pl``.
 
-## polar_plot
-
-```python
-def polar_plot(adata, phenotypes=None, statistic="distribution", method="joint_distribution", splitby=None, color_dict=None, temperature=1.0):
-    """
-    Create a polar plot of phenotype distributions or statistics.
-    
-    Parameters
-    ----------
-    adata : AnnData
-        AnnData object with model results registered
-    phenotypes : list, optional
-        List of phenotype names to include
-    statistic : str, default="distribution"
-        Statistic to plot, one of "distribution", "entropy", "mi"
-    method : str, default="joint_distribution"
-        Method to compute distributions
-    splitby : str, optional
-        Variable to split the plot by
-    color_dict : dict, optional
-        Dictionary mapping phenotypes to colors
-    temperature : float, default=1.0
-        Temperature parameter for distributions
-        
-    Returns
-    -------
-    matplotlib.figure.Figure
-        The polar plot figure
-    """
-    pass
-```
-
-## probability_ternary
-
-```python
-def probability_ternary(adata, phenotype_names, splitby=None, conditions=None, top_n=None):
-    """
-    Create a ternary plot of phenotype probabilities.
-    
-    Parameters
-    ----------
-    adata : AnnData
-        AnnData object with model results registered
-    phenotype_names : list
-        List of exactly 3 phenotype names to plot
-    splitby : str, optional
-        Variable to split the plot by
-    conditions : list, optional
-        List of condition values to include
-    top_n : int, optional
-        Number of top clones to highlight
-        
-    Returns
-    -------
-    matplotlib.figure.Figure
-        The ternary plot figure
-    """
-    pass
-```
-
-## mutual_information
-
-```python
-def mutual_information(adata, splitby=None, temperature=1.0, n_samples=0, 
-                     normalized=True, palette=None, save=None, 
-                     legend_fontsize=6, bbox_to_anchor=(1.15,1.), 
-                     figsize=(8,4), rotation=90, weighted=True, 
-                     return_plot=True):
-    """
-    Plot mutual information between phenotypes and TCR clonotypes.
-    
-    Parameters
-    ----------
-    adata : AnnData
-        AnnData object with model results registered
-    splitby : str, optional
-        Variable to split the plot by
-    temperature : float, default=1.0
-        Temperature parameter for distributions
-    n_samples : int, default=0
-        Number of samples to draw
-    normalized : bool, default=True
-        Whether to normalize the mutual information
-    palette : dict, optional
-        Color palette for the plot
-    save : str, optional
-        Path to save the figure
-    legend_fontsize : int, default=6
-        Font size for the legend
-    bbox_to_anchor : tuple, default=(1.15,1.)
-        Position for the legend
-    figsize : tuple, default=(8,4)
-        Figure size
-    rotation : int, default=90
-        Rotation for x-axis labels
-    weighted : bool, default=True
-        Whether to weight by clone size
-    return_plot : bool, default=True
-        Whether to return the plot
-        
-    Returns
-    -------
-    matplotlib.figure.Figure, optional
-        The mutual information plot
-    """
-    pass
-```
-
-## clonality
-
-```python
-def clonality(adata, groupby=None, splitby=None, s=10, order=None, 
-             figsize=(12,5), palette=None):
-    """
-    Plot clonality metrics.
-    
-    Parameters
-    ----------
-    adata : AnnData
-        AnnData object with TCR information
-    groupby : str, optional
-        Variable to group by
-    splitby : str, optional
-        Variable to split by
-    s : int, default=10
-        Point size
-    order : list, optional
-        Order of categories
-    figsize : tuple, default=(12,5)
-        Figure size
-    palette : dict, optional
-        Color palette
-        
-    Returns
-    -------
-    matplotlib.figure.Figure
-        The clonality plot
-    """
-    pass
-```
-
-## Usage Examples
-
-```python
-import tcri
-import scanpy as sc
-
-# Load data
-adata = sc.read_h5ad("your_data.h5ad")
-
-# Initialize and train model
-model = tcri.TCRIModel(adata)
-model.train()
-tcri.pp.register_model(adata, model)
-
-# Create plots
-tcri.pl.polar_plot(adata, statistic="distribution")
-
-tcri.pl.probability_ternary(
-    adata,
-    ["Phenotype1", "Phenotype2", "Phenotype3"],
-    splitby="condition"
-)
-
-tcri.pl.mutual_information(
-    adata,
-    splitby="timepoint",
-    figsize=(10,6)
-)
-
-tcri.pl.clonality(
-    adata,
-    groupby="condition",
-    splitby="timepoint"
-)
+```{eval-rst}
+.. automodule:: tcri.plotting._plotting
 ```
diff --git a/docs/api/preprocessing.md b/docs/api/preprocessing.md
index 39f1cfc..907915d 100644
--- a/docs/api/preprocessing.md
+++ b/docs/api/preprocessing.md
@@ -1,123 +1,8 @@
-# Preprocessing API
+# Preprocessing
 
-The `tcri.preprocessing` module (imported as `tcri.pp`) provides functions for preprocessing data and preparing it for analysis with TCRi.
+Model registration and the joint-distribution machinery, plus clone/phenotype
+bookkeeping helpers. Exposed as ``tcri.pp``.
 
-## register_model
-
-```python
-def register_model(adata, model, phenotype_prob_slot="X_tcri_phenotypes", phenotype_assignment_obs="tcri_phenotype", latent_slot="X_tcri", batch_size=256):
-    """
-    Register model outputs in the AnnData object.
-    
-    Parameters
-    ----------
-    adata : AnnData
-        AnnData object to register model outputs in
-    model : TCRIModel
-        Trained TCRIModel
-    phenotype_prob_slot : str, default="X_tcri_phenotypes"
-        Key in adata.obsm where phenotype probabilities will be stored
-    phenotype_assignment_obs : str, default="tcri_phenotype"
-        Key in adata.obs where phenotype assignments will be stored
-    latent_slot : str, default="X_tcri"
-        Key in adata.obsm where latent representations will be stored
-    batch_size : int, default=256
-        Batch size for model inference
-        
-    Returns
-    -------
-    AnnData
-        Updated AnnData object with model outputs registered
-    """
-    pass
-```
-
-## joint_distribution
-
-```python
-def joint_distribution(adata, covariate_label, temperature=1.0, n_samples=0, clones=None, weighted=False):
-    """
-    Compute joint distribution of phenotypes and TCR clonotypes conditioned on a covariate.
-    
-    Parameters
-    ----------
-    adata : AnnData
-        AnnData object with model results registered
-    covariate_label : str
-        Label of the covariate to condition on (must be in adata.obs)
-    temperature : float, default=1.0
-        Temperature parameter for softening/sharpening distributions
-    n_samples : int, default=0
-        Number of samples to draw from the posterior
-    clones : list, optional
-        List of clone IDs to include in the analysis
-    weighted : bool, default=False
-        Whether to weight the distribution by clone size
-        
-    Returns
-    -------
-    dict
-        Dictionary mapping covariate values to joint distribution matrices
-    """
-    pass
-```
-
-## global_joint_distribution
-
-```python
-def global_joint_distribution(adata, temperature=1.0, n_samples=0):
-    """
-    Compute global joint distribution of phenotypes and TCR clonotypes.
-    
-    Parameters
-    ----------
-    adata : AnnData
-        AnnData object with model results registered
-    temperature : float, default=1.0
-        Temperature parameter for softening/sharpening distributions
-    n_samples : int, default=0
-        Number of samples to draw from the posterior
-        
-    Returns
-    -------
-    numpy.ndarray
-        Joint distribution matrix of shape (n_clones, n_phenotypes)
-    """
-    pass
-```
-
-## Usage Examples
-
-```python
-import tcri
-import scanpy as sc
-
-# Load data
-adata = sc.read_h5ad("your_data.h5ad")
-
-# Initialize and train model
-model = tcri.TCRIModel(adata)
-model.train()
-
-# Register model outputs in AnnData
-tcri.pp.register_model(
-    adata,
-    model,
-    phenotype_prob_slot="X_tcri_phenotypes",
-    phenotype_assignment_obs="tcri_phenotype",
-    latent_slot="X_tcri"
-)
-
-# Compute joint distributions
-joint_dist = tcri.pp.joint_distribution(
-    adata,
-    covariate_label="timepoint",  # Replace with your covariate
-    temperature=1.0
-)
-
-# Compute global joint distribution
-global_dist = tcri.pp.global_joint_distribution(
-    adata,
-    temperature=1.0
-)
+```{eval-rst}
+.. automodule:: tcri.preprocessing._preprocessing
 ```
diff --git a/docs/api/utils.md b/docs/api/utils.md
index fbdb2d5..78905fd 100644
--- a/docs/api/utils.md
+++ b/docs/api/utils.md
@@ -1,91 +1,8 @@
-# Utils API
+# Utilities
 
-The `tcri.utils` module (imported as `tcri.ut`) provides utility functions for working with TCRi data.
+Session save/load, safe AnnData I/O, and assorted helpers. Exposed as
+``tcri.ut``.
 
-## get_clones
-
-```python
-def get_clones(adata, clone_key="clone_id"):
-    """
-    Get unique clone IDs from the AnnData object.
-    
-    Parameters
-    ----------
-    adata : AnnData
-        AnnData object with TCR information
-    clone_key : str, default="clone_id"
-        Key in adata.obs containing clone IDs
-        
-    Returns
-    -------
-    list
-        List of unique clone IDs
-    """
-    pass
-```
-
-## get_phenotypes
-
-```python
-def get_phenotypes(adata, phenotype_key="phenotype"):
-    """
-    Get unique phenotype labels from the AnnData object.
-    
-    Parameters
-    ----------
-    adata : AnnData
-        AnnData object with phenotype information
-    phenotype_key : str, default="phenotype"
-        Key in adata.obs containing phenotype labels
-        
-    Returns
-    -------
-    list
-        List of unique phenotype labels
-    """
-    pass
-```
-
-## normalize_distribution
-
-```python
-def normalize_distribution(distribution):
-    """
-    Normalize a distribution to sum to 1.
-    
-    Parameters
-    ----------
-    distribution : numpy.ndarray
-        Distribution to normalize
-        
-    Returns
-    -------
-    numpy.ndarray
-        Normalized distribution
-    """
-    pass
-```
-
-## Usage Examples
-
-```python
-import tcri
-import scanpy as sc
-
-# Load data
-adata = sc.read_h5ad("your_data.h5ad")
-
-# Get unique clones
-clones = tcri.ut.get_clones(adata, clone_key="clone_id")
-print(f"Found {len(clones)} unique clones")
-
-# Get unique phenotypes
-phenotypes = tcri.ut.get_phenotypes(adata, phenotype_key="phenotype")
-print(f"Found {len(phenotypes)} unique phenotypes: {phenotypes}")
-
-# Normalize a distribution
-import numpy as np
-dist = np.random.rand(10)
-normalized_dist = tcri.ut.normalize_distribution(dist)
-print(f"Sum of normalized distribution: {normalized_dist.sum():.6f}")
+```{eval-rst}
+.. automodule:: tcri.utils._utils
 ```
diff --git a/docs/conf.py b/docs/conf.py
index 7e0b8de..b7c355b 100644
--- a/docs/conf.py
+++ b/docs/conf.py
@@ -64,7 +64,6 @@
 
 html_theme_options = {
     'logo_only': False,
-    'display_version': True,
     'prev_next_buttons_location': 'bottom',
     'style_external_links': False,
     'style_nav_header_background': '#2980B9',
@@ -81,6 +80,19 @@
 # -- Options for autodoc extension -------------------------------------------
 autodoc_member_order = 'bysource'
 autodoc_typehints = 'description'
+autodoc_default_options = {
+    'members': True,
+    'undoc-members': True,
+    'show-inheritance': True,
+}
+# Show clean object names (e.g. ``joint_distribution``) rather than the full
+# private module path in signatures and headings.
+add_module_names = False
+# Generate stub pages for any autosummary directives.
+autosummary_generate = True
+# Don't fail the whole build if an optional/heavy import is unavailable at
+# doc-build time; autodoc will note the missing object instead.
+autodoc_mock_imports = []
 
 # -- Options for napoleon extension ------------------------------------------
 napoleon_google_docstring = True
diff --git a/tcri/plotting/_plotting.py b/tcri/plotting/_plotting.py
index 8793623..adb38a6 100644
--- a/tcri/plotting/_plotting.py
+++ b/tcri/plotting/_plotting.py
@@ -864,12 +864,12 @@ def ridge_delta_entropy(
     """
     Ridge plot of Δ-entropy posteriors per phenotype.
 
-    For each phenotype the *first two* groups in `order_group`
-    are compared and annotated:
+    For each phenotype the *first two* groups in ``order_group``
+    are compared and annotated::
 
-        ┌──────────┐
-    CR  │          │ NR
-        ★ p-value / stars
+            ┌──────────┐
+        CR  │          │ NR
+            ★ p-value / stars
 
     Bracket anchors are placed at the group means.
     """