diffbio 0.1.0__py3-none-any.whl

diffbio/operators/metabolomics/spectral_similarity.py

@@ -0,0 +1,315 @@
+"""Differentiable spectral similarity operator for metabolomics (MS2DeepScore-style).
+
+This module implements a Siamese neural network for predicting molecular
+structural similarity from tandem mass spectra (MS/MS), based on the
+MS2DeepScore architecture.
+
+The approach uses a shared base network to generate spectral embeddings,
+then computes cosine similarity between embeddings to predict structural
+similarity (Tanimoto scores).
+
+Architecture based on:
+    Huber et al. (2021). "MS2DeepScore: a novel deep learning similarity
+    measure to compare tandem mass spectra." Journal of Cheminformatics.
+
+Key features:
+
+- Siamese architecture with shared weights for spectrum encoding
+- 200-dimensional spectral embeddings
+- Cosine similarity for structure prediction
+- Monte-Carlo dropout for uncertainty estimation
+"""
+
+import logging
+from dataclasses import dataclass
+from typing import Any
+
+import jax.numpy as jnp
+from artifex.generative_models.core.base import MLP
+from datarax.core.config import OperatorConfig
+from datarax.core.operator import OperatorModule
+from flax import nnx
+
+logger = logging.getLogger(__name__)
+
+
+@dataclass(frozen=True)
+class SpectralSimilarityConfig(OperatorConfig):
+    """Configuration for DifferentiableSpectralSimilarity.
+
+    Attributes:
+        n_bins: Number of m/z bins for spectrum discretization.
+            Default 1000 (10-1000 m/z at 1 m/z resolution).
+            Original MS2DeepScore uses 10000 bins at 0.1 m/z resolution.
+        embedding_dim: Dimension of spectral embeddings. Default 200.
+        hidden_dims: Tuple of hidden layer dimensions. Default (512, 256).
+            Original MS2DeepScore uses (500, 500).
+        dropout_rate: Dropout rate for regularization. Default 0.2.
+        min_mz: Minimum m/z value for binning. Default 0.0.
+        max_mz: Maximum m/z value for binning. Default 1000.0.
+        use_batch_norm: Whether to use batch normalization. Default True.
+    """
+
+    n_bins: int = 1000
+    embedding_dim: int = 200
+    hidden_dims: tuple[int, ...] = (512, 256)
+    dropout_rate: float = 0.2
+    min_mz: float = 0.0
+    max_mz: float = 1000.0
+    use_batch_norm: bool = True
+
+
+class DifferentiableSpectralSimilarity(OperatorModule):
+    """Siamese neural network for spectral similarity prediction.
+
+    This operator implements the MS2DeepScore architecture for predicting
+    molecular structural similarity from tandem mass spectra. The network
+    uses a shared encoder to generate spectral embeddings, then computes
+    cosine similarity between pairs of embeddings.
+
+    The operator supports two modes of operation:
+    1. Single spectrum input: Generates embeddings for spectra
+    2. Paired spectra input: Computes similarity between spectrum pairs
+
+    Architecture:
+        Input (n_bins) -> shared encoder MLP -> Embedding (embedding_dim)
+
+    Attributes:
+        config: SpectralSimilarityConfig with hyperparameters.
+        backbone: Shared Artifex encoder backbone producing spectral embeddings.
+
+    Example:
+        ```python
+        config = SpectralSimilarityConfig(n_bins=1000, embedding_dim=200)
+        operator = DifferentiableSpectralSimilarity(config, rngs=nnx.Rngs(42))
+        # Get embeddings for spectra
+        spectra = jax.random.uniform(jax.random.PRNGKey(0), (10, 1000))
+        result, _, _ = operator.apply({"spectra": spectra}, {}, None)
+        embeddings = result["embeddings"]  # (10, 200)
+        # Compute pairwise similarity
+        spectra_a = jax.random.uniform(jax.random.PRNGKey(0), (5, 1000))
+        spectra_b = jax.random.uniform(jax.random.PRNGKey(1), (5, 1000))
+        result, _, _ = operator.apply(
+            {"spectra_a": spectra_a, "spectra_b": spectra_b}, {}, None
+        )
+        similarity = result["similarity_scores"]  # (5,) in [-1, 1]
+        ```
+    """
+
+    def __init__(self, config: SpectralSimilarityConfig, *, rngs: nnx.Rngs) -> None:
+        """Initialize the spectral similarity operator.
+
+        Args:
+            config: Configuration with network hyperparameters.
+            rngs: Flax NNX random number generators.
+        """
+        super().__init__(config, rngs=rngs)
+        self.config = config
+
+        self.backbone = MLP(
+            hidden_dims=[*config.hidden_dims, config.embedding_dim],
+            in_features=config.n_bins,
+            activation="relu",
+            dropout_rate=config.dropout_rate,
+            output_activation=None,
+            use_batch_norm=config.use_batch_norm,
+            rngs=rngs,
+        )
+
+    def encode(self, spectra: jnp.ndarray) -> jnp.ndarray:
+        """Encode binned spectra into embeddings.
+
+        BatchNorm and Dropout respect the model's train/eval mode:
+        - Call model.train() before training to enable dropout and update batch stats
+        - Call model.eval() before inference to disable dropout and use running stats
+
+        Args:
+            spectra: Binned spectra with shape (n_spectra, n_bins).
+
+        Returns:
+            Embeddings with shape (n_spectra, embedding_dim).
+        """
+        backbone_output = self.backbone(spectra)
+        if isinstance(backbone_output, tuple):
+            raise TypeError("Spectral similarity backbone must return a single tensor output.")
+        return backbone_output
+
+    def cosine_similarity(
+        self, embeddings_a: jnp.ndarray, embeddings_b: jnp.ndarray
+    ) -> jnp.ndarray:
+        """Compute cosine similarity between embedding pairs.
+
+        Args:
+            embeddings_a: First set of embeddings (n, embedding_dim).
+            embeddings_b: Second set of embeddings (n, embedding_dim).
+
+        Returns:
+            Cosine similarity scores with shape (n,).
+        """
+        # Normalize embeddings
+        norm_a = jnp.linalg.norm(embeddings_a, axis=-1, keepdims=True)
+        norm_b = jnp.linalg.norm(embeddings_b, axis=-1, keepdims=True)
+
+        # Avoid division by zero
+        norm_a = jnp.maximum(norm_a, 1e-8)
+        norm_b = jnp.maximum(norm_b, 1e-8)
+
+        embeddings_a_normalized = embeddings_a / norm_a
+        embeddings_b_normalized = embeddings_b / norm_b
+
+        # Compute cosine similarity
+        similarity = jnp.sum(embeddings_a_normalized * embeddings_b_normalized, axis=-1)
+
+        return similarity
+
+    def apply(
+        self,
+        data: dict[str, Any],
+        state: dict[str, Any],
+        metadata: dict[str, Any] | None,
+        random_params: Any = None,
+        stats: dict[str, Any] | None = None,
+    ) -> tuple[dict[str, Any], dict[str, Any], dict[str, Any] | None]:
+        """Apply the spectral similarity operator.
+
+        The operator supports two input modes:
+
+        1. Single spectra mode (embedding generation):
+           Input: {"spectra": (n, n_bins)}
+           Output: {"embeddings": (n, embedding_dim)}
+
+        2. Paired spectra mode (similarity computation):
+           Input: {"spectra_a": (n, n_bins), "spectra_b": (n, n_bins)}
+           Output: {"similarity_scores": (n,), "embeddings_a": ..., "embeddings_b": ...}
+
+        Args:
+            data: Input data dictionary with spectra.
+            state: Per-element state (passed through).
+            metadata: Optional metadata (passed through).
+            random_params: Random parameters (unused).
+            stats: Optional statistics (unused).
+
+        Returns:
+            Tuple of (output_data, state, metadata).
+        """
+        if "spectra" in data:
+            # Single spectra mode: compute embeddings
+            spectra = data["spectra"]
+            embeddings = self.encode(spectra)
+
+            output = {**data, "embeddings": embeddings}
+
+        elif "spectra_a" in data and "spectra_b" in data:
+            # Paired spectra mode: compute similarity
+            spectra_a = data["spectra_a"]
+            spectra_b = data["spectra_b"]
+
+            embeddings_a = self.encode(spectra_a)
+            embeddings_b = self.encode(spectra_b)
+
+            similarity_scores = self.cosine_similarity(embeddings_a, embeddings_b)
+
+            output = {
+                **data,
+                "embeddings_a": embeddings_a,
+                "embeddings_b": embeddings_b,
+                "similarity_scores": similarity_scores,
+            }
+
+        else:
+            raise ValueError(
+                "Input must contain either 'spectra' or both 'spectra_a' and 'spectra_b'"
+            )
+
+        return output, state, metadata
+
+
+def bin_spectrum(
+    mz_values: jnp.ndarray,
+    intensities: jnp.ndarray,
+    n_bins: int = 1000,
+    min_mz: float = 0.0,
+    max_mz: float = 1000.0,
+    normalize: bool = True,
+) -> jnp.ndarray:
+    """Bin a mass spectrum into fixed-width m/z bins.
+
+    This function discretizes a continuous mass spectrum (m/z, intensity pairs)
+    into a fixed-size vector suitable for neural network input.
+
+    Args:
+        mz_values: Array of m/z values with shape (n_peaks,).
+        intensities: Array of intensity values with shape (n_peaks,).
+        n_bins: Number of bins to use. Default 1000.
+        min_mz: Minimum m/z value for binning. Default 0.0.
+        max_mz: Maximum m/z value for binning. Default 1000.0.
+        normalize: Whether to normalize intensities to max=1.0. Default True.
+
+    Returns:
+        Binned spectrum with shape (n_bins,).
+
+    Example:
+        ```python
+        mz = jnp.array([100.0, 200.0, 300.0])
+        intensity = jnp.array([0.5, 1.0, 0.3])
+        binned = bin_spectrum(mz, intensity, n_bins=100)
+        binned.shape
+        ```
+        (100,)
+    """
+    # Compute bin edges
+    bin_width = (max_mz - min_mz) / n_bins
+
+    # Compute bin indices for each m/z value
+    bin_indices = ((mz_values - min_mz) / bin_width).astype(jnp.int32)
+
+    # Clip to valid range
+    bin_indices = jnp.clip(bin_indices, 0, n_bins - 1)
+
+    # Create empty binned spectrum
+    binned = jnp.zeros(n_bins)
+
+    # Accumulate intensities in bins (use segment_sum for differentiability)
+    binned = binned.at[bin_indices].add(intensities)
+
+    # Normalize if requested
+    if normalize:
+        max_intensity = jnp.maximum(jnp.max(binned), 1e-8)
+        binned = binned / max_intensity
+
+    return binned
+
+
+def create_spectral_similarity(
+    n_bins: int = 1000,
+    embedding_dim: int = 200,
+    hidden_dims: tuple[int, ...] = (512, 256),
+    dropout_rate: float = 0.2,
+    seed: int = 42,
+) -> DifferentiableSpectralSimilarity:
+    """Factory function to create a spectral similarity operator.
+
+    Args:
+        n_bins: Number of m/z bins. Default 1000.
+        embedding_dim: Embedding dimension. Default 200.
+        hidden_dims: Hidden layer dimensions. Default (512, 256).
+        dropout_rate: Dropout rate. Default 0.2.
+        seed: Random seed. Default 42.
+
+    Returns:
+        Configured DifferentiableSpectralSimilarity operator.
+
+    Example:
+        ```python
+        operator = create_spectral_similarity(n_bins=500, embedding_dim=128)
+        spectra = jax.random.uniform(jax.random.PRNGKey(0), (10, 500))
+        result, _, _ = operator.apply({"spectra": spectra}, {}, None)
+        ```
+    """
+    config = SpectralSimilarityConfig(
+        n_bins=n_bins,
+        embedding_dim=embedding_dim,
+        hidden_dims=hidden_dims,
+        dropout_rate=dropout_rate,
+    )
+    return DifferentiableSpectralSimilarity(config, rngs=nnx.Rngs(seed))

diffbio/operators/molecular_dynamics/__init__.py

@@ -0,0 +1,51 @@
+"""Molecular dynamics operators for DiffBio.
+
+This module provides differentiable operators for molecular dynamics simulations,
+wrapping JAX-MD functionality for seamless integration with DiffBio pipelines.
+
+Operators:
+    ForceFieldOperator: Compute energies and forces from particle positions
+    MDIntegratorOperator: Time integration for MD simulations
+
+Factory Functions:
+    create_force_field: Create force field operator with specified potential
+    create_integrator: Create integrator operator with specified type
+    create_lennard_jones_operator: Create LJ force field operator (convenience)
+    create_verlet_integrator: Create velocity Verlet integrator (convenience)
+
+Enums:
+    PotentialType: Enumeration of supported potential types
+
+References:
+    Schoenholz & Cubuk (2020). JAX, M.D.: A Framework for Differentiable Physics.
+    NeurIPS 2020.
+"""
+
+from diffbio.operators.molecular_dynamics.force_field import (
+    ForceFieldConfig,
+    ForceFieldOperator,
+    create_force_field,
+    create_lennard_jones_operator,
+)
+from diffbio.operators.molecular_dynamics.integrator import (
+    MDIntegratorConfig,
+    MDIntegratorOperator,
+    create_integrator,
+    create_verlet_integrator,
+)
+from diffbio.operators.molecular_dynamics.primitives import PotentialType
+
+__all__ = [
+    # Enums
+    "PotentialType",
+    # Force field
+    "ForceFieldConfig",
+    "ForceFieldOperator",
+    "create_force_field",
+    "create_lennard_jones_operator",
+    # Integrator
+    "MDIntegratorConfig",
+    "MDIntegratorOperator",
+    "create_integrator",
+    "create_verlet_integrator",
+]

diffbio/operators/molecular_dynamics/force_field.py

@@ -0,0 +1,265 @@
+"""Force field operators wrapping JAX-MD.
+
+This module provides differentiable force field operators that compute
+molecular energies and forces using JAX-MD's efficient implementations.
+"""
+
+import logging
+from dataclasses import dataclass
+from typing import Any
+
+import jax
+from datarax.core.config import OperatorConfig
+from datarax.core.operator import OperatorModule
+from flax import nnx
+from jaxtyping import Array
+
+from diffbio.operators.molecular_dynamics.primitives import (
+    PotentialType,
+    create_displacement_fn,
+    create_energy_fn,
+    create_force_fn,
+)
+
+logger = logging.getLogger(__name__)
+
+
+@dataclass(frozen=True)
+class ForceFieldConfig(OperatorConfig):
+    """Configuration for force field operator.
+
+    Attributes:
+        potential_type: Type of potential ("lennard_jones", "morse", "soft_sphere").
+        sigma: Length scale parameter (particle diameter).
+        epsilon: Energy scale parameter (well depth).
+        cutoff: Cutoff distance for interactions (in units of sigma). None for no cutoff.
+        box_size: Size of periodic box. None for non-periodic.
+        alpha: Morse potential width parameter (only for morse).
+        reference_positions: Optional reference positions for computing
+            geometric losses (chamfer/EMD) via artifex. If provided,
+            ``chamfer_distance`` and ``earth_mover_distance`` are added
+            to the output dict. Shape must match particle positions.
+        geometric_loss_weight: Weight for geometric loss terms.
+    """
+
+    potential_type: str = "lennard_jones"
+    sigma: float = 1.0
+    epsilon: float = 1.0
+    cutoff: float | None = 2.5
+    box_size: float | None = None
+    alpha: float = 5.0  # Morse potential parameter
+    geometric_loss_weight: float = 0.0
+
+
+class ForceFieldOperator(OperatorModule):
+    """Differentiable force field operator using JAX-MD.
+
+    Computes potential energy and forces for a system of particles using
+    classical pairwise potentials. Forces are computed automatically via
+    JAX's automatic differentiation.
+
+    Supported potentials:
+        - lennard_jones: Standard 12-6 LJ potential
+        - morse: Morse potential for bonded interactions
+        - soft_sphere: Soft repulsive potential
+
+    Example:
+        ```python
+        config = ForceFieldConfig(potential_type="lennard_jones", box_size=10.0)
+        operator = ForceFieldOperator(config, rngs=nnx.Rngs(42))
+        data = {"positions": positions}  # (n_particles, dim)
+        result, state, meta = operator.apply(data, {}, None)
+        energy = result["energy"]  # scalar
+        forces = result["forces"]  # (n_particles, dim)
+        ```
+    """
+
+    def __init__(self, config: ForceFieldConfig, *, rngs: nnx.Rngs | None = None):
+        """Initialize force field operator.
+
+        Args:
+            config: Force field configuration.
+            rngs: Flax NNX random number generators.
+        """
+        super().__init__(config, rngs=rngs)
+        self.config: ForceFieldConfig = config
+
+        # Pre-create energy and force functions (efficiency: only created once)
+        self._displacement_fn, _ = create_displacement_fn(config.box_size)
+        self._energy_fn = create_energy_fn(
+            self._displacement_fn,
+            potential_type=config.potential_type,
+            sigma=config.sigma,
+            epsilon=config.epsilon,
+            cutoff=config.cutoff,
+            alpha=config.alpha,
+        )
+        self._force_fn = create_force_fn(self._energy_fn)
+
+    def apply(
+        self,
+        data: dict[str, Any],
+        state: dict[str, Any],
+        metadata: dict[str, Any] | None,
+        random_params: Any = None,
+        stats: dict[str, Any] | None = None,
+    ) -> tuple[dict[str, Any], dict[str, Any], dict[str, Any] | None]:
+        """Compute energy and forces for particle positions.
+
+        Args:
+            data: Input data containing:
+                - positions: Particle positions (n_particles, dim) or
+                             (batch, n_particles, dim)
+            state: Per-element state (passed through).
+            metadata: Optional metadata.
+            random_params: Unused random parameters.
+            stats: Optional statistics dictionary.
+
+        Returns:
+            Tuple of:
+                - data with added "energy" and "forces" keys
+                - unchanged state
+                - unchanged metadata
+        """
+        positions = data["positions"]
+
+        # Handle batched input
+        if positions.ndim == 3:
+            # Batched: (batch, n_particles, dim)
+            batch_apply = jax.vmap(self._compute_single)
+            energy_vals, forces = batch_apply(positions)
+        else:
+            # Single: (n_particles, dim)
+            energy_vals, forces = self._compute_single(positions)
+
+        result = {
+            **data,
+            "energy": energy_vals,
+            "forces": forces,
+        }
+
+        # Optionally compute geometric losses against reference positions
+        if self.config.geometric_loss_weight > 0 and "reference_positions" in data:
+            result.update(
+                _compute_geometric_losses(
+                    positions, data["reference_positions"], self.config.geometric_loss_weight
+                )
+            )
+
+        return result, state, metadata
+
+    def _compute_single(self, positions: Array) -> tuple[Array, Array]:
+        """Compute energy and forces for a single configuration.
+
+        Args:
+            positions: Particle positions (n_particles, dim).
+
+        Returns:
+            Tuple of (energy, forces).
+        """
+        # Use pre-created functions from __init__
+        total_energy = self._energy_fn(positions)
+        forces = self._force_fn(positions)
+
+        return total_energy, forces
+
+
+def _compute_geometric_losses(
+    predicted: Array,
+    reference: Array,
+    weight: float,
+) -> dict[str, Array]:
+    """Compute artifex geometric losses between predicted and reference positions.
+
+    Args:
+        predicted: Predicted positions (batch, n_particles, dim) or (n_particles, dim).
+        reference: Reference positions with same shape.
+        weight: Scaling weight for the losses.
+
+    Returns:
+        Dict with ``chamfer_distance`` and ``earth_mover_distance`` keys.
+    """
+    from artifex.generative_models.core.losses.geometric import (  # noqa: PLC0415
+        chamfer_distance,
+        earth_mover_distance,
+    )
+
+    # Ensure batch dimension
+    if predicted.ndim == 2:
+        predicted = predicted[None, ...]
+        reference = reference[None, ...]
+
+    cd = chamfer_distance(predicted, reference) * weight
+    emd = earth_mover_distance(predicted, reference) * weight
+    return {"chamfer_distance": cd, "earth_mover_distance": emd}
+
+
+def create_force_field(
+    potential_type: str | PotentialType = PotentialType.LENNARD_JONES,
+    sigma: float = 1.0,
+    epsilon: float = 1.0,
+    cutoff: float | None = 2.5,
+    box_size: float | None = None,
+    alpha: float = 5.0,
+    seed: int = 42,
+) -> ForceFieldOperator:
+    """Create a force field operator with specified potential.
+
+    Args:
+        potential_type: Type of potential ("lennard_jones", "morse", "soft_sphere")
+            or PotentialType enum.
+        sigma: Particle diameter (length scale).
+        epsilon: Well depth (energy scale).
+        cutoff: Cutoff distance in units of sigma. None for no cutoff.
+        box_size: Periodic box size. None for non-periodic.
+        alpha: Morse potential width parameter.
+        seed: Random seed for initialization.
+
+    Returns:
+        Configured ForceFieldOperator.
+    """
+    # Convert enum to string if needed
+    if isinstance(potential_type, PotentialType):
+        potential_type = potential_type.value
+
+    config = ForceFieldConfig(
+        potential_type=potential_type,
+        sigma=sigma,
+        epsilon=epsilon,
+        cutoff=cutoff,
+        box_size=box_size,
+        alpha=alpha,
+    )
+    return ForceFieldOperator(config, rngs=nnx.Rngs(seed))
+
+
+def create_lennard_jones_operator(
+    sigma: float = 1.0,
+    epsilon: float = 1.0,
+    cutoff: float | None = 2.5,
+    box_size: float | None = None,
+    seed: int = 42,
+) -> ForceFieldOperator:
+    """Create a Lennard-Jones force field operator.
+
+    This is a convenience function for creating a force field operator
+    with Lennard-Jones potential.
+
+    Args:
+        sigma: Particle diameter (length scale).
+        epsilon: Well depth (energy scale).
+        cutoff: Cutoff distance in units of sigma. None for no cutoff.
+        box_size: Periodic box size. None for non-periodic.
+        seed: Random seed for initialization.
+
+    Returns:
+        Configured ForceFieldOperator.
+    """
+    return create_force_field(
+        potential_type=PotentialType.LENNARD_JONES,
+        sigma=sigma,
+        epsilon=epsilon,
+        cutoff=cutoff,
+        box_size=box_size,
+        seed=seed,
+    )