diffbio 0.1.6__tar.gz → 0.1.8__tar.gz

{diffbio-0.1.6 → diffbio-0.1.8}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.5
 Name: diffbio
-Version: 0.1.6
+Version: 0.1.8
 Summary: End-to-end differentiable bioinformatics for JAX/Flax NNX: alignment, mapping, assembly, variant calling, RNA-seq, single-cell, epigenomics, CRISPR, metabolomics, multi-omics, protein and RNA structure, molecular dynamics and drug-discovery operators composed into trainable pipelines on datarax, artifex, opifex and calibrax
 Project-URL: Bug Tracker, https://github.com/avitai/DiffBio/issues
 Project-URL: Documentation, https://diffbio.readthedocs.io
@@ -45,10 +45,10 @@ Classifier: Topic :: Software Development :: Libraries
 Classifier: Topic :: Software Development :: Libraries :: Python Modules
 Requires-Python: <3.14,>=3.12
 Requires-Dist: anndata>=0.9.1
-Requires-Dist: avitai-artifex>=0.1.9
+Requires-Dist: avitai-artifex>=0.1.10
 Requires-Dist: beartype>=0.14.1
 Requires-Dist: biopython>=1.81
-Requires-Dist: calibrax>=0.1.6
+Requires-Dist: calibrax>=0.1.8
 Requires-Dist: chex>=0.1.7
 Requires-Dist: datarax>=0.1.11
 Requires-Dist: flax>=0.12.9
@@ -352,11 +352,11 @@ result, _, _ = pipeline.apply(batch_data, {}, None)
 ### Training a Pipeline
 
 ```python
-from diffbio.utils import (
+from calibrax.metrics.functional import softmax_cross_entropy
 from substrax.optim import OptimizerConfig
+from diffbio.utils import (
     Trainer,
     TrainingConfig,
-    cross_entropy_loss,
     create_synthetic_training_data,
     data_iterator,
 )
@@ -384,10 +384,9 @@ trainer = Trainer(pipeline, config)
 
 # Define loss function
 def loss_fn(predictions, targets):
-    return cross_entropy_loss(
+    return softmax_cross_entropy(
         predictions["logits"],
         targets["labels"],
-        num_classes=3,
     )
 
 # Train

{diffbio-0.1.6 → diffbio-0.1.8}/README.md

@@ -138,11 +138,11 @@ result, _, _ = pipeline.apply(batch_data, {}, None)
 ### Training a Pipeline
 
 ```python
-from diffbio.utils import (
+from calibrax.metrics.functional import softmax_cross_entropy
 from substrax.optim import OptimizerConfig
+from diffbio.utils import (
     Trainer,
     TrainingConfig,
-    cross_entropy_loss,
     create_synthetic_training_data,
     data_iterator,
 )
@@ -170,10 +170,9 @@ trainer = Trainer(pipeline, config)
 
 # Define loss function
 def loss_fn(predictions, targets):
-    return cross_entropy_loss(
+    return softmax_cross_entropy(
         predictions["logits"],
         targets["labels"],
-        num_classes=3,
     )
 
 # Train

{diffbio-0.1.6 → diffbio-0.1.8}/pyproject.toml

@@ -39,11 +39,11 @@ dependencies = [
   # Datarax - operator, source, and pipeline contracts
   "datarax>=0.1.11",
   # Artifex - modeling, modality, and transformer substrate
-  "avitai-artifex>=0.1.9",
+  "avitai-artifex>=0.1.10",
   # Opifex - scientific ML, operator learning, and advanced optimization
   "opifex>=0.2.7",
   # Calibrax - benchmarking, comparison, profiling, and regression control
-  "calibrax>=0.1.6",
+  "calibrax>=0.1.8",
   # Substrax - device identity, meshes, runtime settings and checkpoint store shared by the
   # siblings; 0.1.6 adds substrax.runtime, substrax.artifacts and substrax.testing.
   "substrax>=0.1.9",
@@ -71,7 +71,7 @@ license = {file = "LICENSE"}
 name = "diffbio"
 readme = "README.md"
 requires-python = ">=3.12,<3.14"
-version = "0.1.6"
+version = "0.1.8"
 
 [project.optional-dependencies]
 all = ["diffbio[benchmark,chem,cuda12,dev,docs,genomics,metal,soft-ops-advanced,soft-ops-ot,test,torch-io]"]

{diffbio-0.1.6 → diffbio-0.1.8}/src/diffbio/losses/self_supervised_losses.py

@@ -11,10 +11,11 @@ term, which the ``core.losses`` ledger flags as the one genuinely new B6 piece.
 from __future__ import annotations
 
 import jax
+from calibrax.metrics.functional import mae, mse
 from flax import nnx
 
-_EPSILON = 1.0e-8
-_VALID_LOSS_TYPES = ("mse", "mae")
+_LOSSES = {"mse": mse, "mae": mae}
+_VALID_LOSS_TYPES = tuple(_LOSSES)
 
 
 def masked_value_loss(
@@ -23,28 +24,28 @@ def masked_value_loss(
     mask: jax.Array,
     *,
     loss_type: str = "mse",
-    epsilon: float = _EPSILON,
 ) -> jax.Array:
     """Reconstruction loss averaged over masked positions only.
 
+    The reduction is calibrax's masked mean: positions where ``mask`` is nonzero are
+    averaged, the rest contribute nothing, and an all-zero mask gives ``0`` rather than
+    ``0 / 0``.
+
     Args:
         predictions: Predicted values, any shape broadcastable with ``mask``.
         targets: Ground-truth values, same shape as ``predictions``.
         mask: ``1`` at positions to reconstruct (the masked-out values), ``0`` elsewhere.
         loss_type: ``"mse"`` (squared error) or ``"mae"`` (absolute error).
-        epsilon: Stabilizer so an all-zero mask returns ``0`` rather than ``0 / 0``.
 
     Returns:
-        The mask-weighted mean reconstruction error at the masked positions.
+        The mean reconstruction error at the masked positions.
 
     Raises:
         ValueError: If ``loss_type`` is not one of ``("mse", "mae")``.
     """
-    if loss_type not in _VALID_LOSS_TYPES:
+    if loss_type not in _LOSSES:
         raise ValueError(f"loss_type must be one of {_VALID_LOSS_TYPES}, got {loss_type!r}")
-    residual = predictions - targets
-    error = residual**2 if loss_type == "mse" else abs(residual)
-    return (mask * error).sum() / (mask.sum() + epsilon)
+    return _LOSSES[loss_type](predictions, targets, mask=mask)
 
 
 class MaskedValueLoss(nnx.Module):

{diffbio-0.1.6 → diffbio-0.1.8}/src/diffbio/operators/epigenomics/contextual.py

@@ -9,6 +9,7 @@ import jax
 import jax.numpy as jnp
 import optax
 from artifex.generative_models.core.layers import TransformerEncoder
+from calibrax.metrics.functional import mse
 from datarax.core.config import OperatorConfig
 from datarax.core.operator import OperatorModule
 from flax import nnx
@@ -250,8 +251,7 @@ def compute_chromatin_guidance_loss(
     )
     predicted_contacts = jax.nn.sigmoid(similarity)
     pair_mask = sequence_mask[:, :, None] * sequence_mask[:, None, :]
-    squared_error = jnp.square(predicted_contacts - chromatin_contacts) * pair_mask
-    return squared_error.sum() / jnp.maximum(pair_mask.sum(), 1.0)
+    return mse(predicted_contacts, chromatin_contacts, mask=pair_mask)
 
 
 def compute_contextual_epigenomics_loss(

{diffbio-0.1.6 → diffbio-0.1.8}/src/diffbio/operators/multiomics/multiomics_vae.py

@@ -23,7 +23,6 @@ from typing import Any
 import jax
 import jax.numpy as jnp
 from artifex.generative_models.core.base import MLP
-from artifex.generative_models.core.losses.base import reduce_loss
 from artifex.generative_models.core.losses.divergence import gaussian_kl_divergence
 from datarax.core.config import OperatorConfig
 from datarax.core.operator import require_key
@@ -304,8 +303,7 @@ class DifferentiableMultiOmicsVAE(LossBalancingMixin, EncoderDecoderOperator):
         for i in range(n_modalities):
             counts = data[self._input_key(i)]
             per_sample = jnp.sum((counts - reconstructions[i]) ** 2, axis=-1)
-            mean_recon = reduce_loss(per_sample, reduction="mean")
-            total_recon = total_recon + weights[i] * mean_recon
+            total_recon = total_recon + weights[i] * jnp.mean(per_sample)
 
         # KL divergence (batch_sum: sum over latent, mean over batch)
         kl = gaussian_kl_divergence(mu_joint, logvar_joint, reduction="batch_sum")

{diffbio-0.1.6 → diffbio-0.1.8}/src/diffbio/operators/singlecell/enhanced_batch_correction.py

@@ -29,8 +29,8 @@ from artifex.generative_models.core.losses.adversarial import (
     wasserstein_discriminator_loss,
     wasserstein_generator_loss,
 )
-from artifex.generative_models.core.losses.base import reduce_loss
 from artifex.generative_models.core.losses.divergence import maximum_mean_discrepancy
+from calibrax.metrics.functional import mse
 from datarax.core.config import OperatorConfig
 from datarax.core.operator import OperatorModule
 from flax import nnx
@@ -294,10 +294,7 @@ class DifferentiableMMDBatchCorrection(LossBalancingMixin, OperatorModule):
         reconstructed = self._decode(latent)
 
         # Losses
-        reconstruction_loss = reduce_loss(
-            (reconstructed - expression) ** 2,
-            reduction="mean",
-        )
+        reconstruction_loss = mse(reconstructed, expression)
         mmd_loss = self._compute_pairwise_mmd(latent, batch_labels)
 
         result = {
@@ -474,10 +471,7 @@ class DifferentiableWGANBatchCorrection(LossBalancingMixin, OperatorModule):
         disc_scores = self._discriminate(latent_reversed)
 
         # Reconstruction loss
-        reconstruction_loss = reduce_loss(
-            (reconstructed - expression) ** 2,
-            reduction="mean",
-        )
+        reconstruction_loss = mse(reconstructed, expression)
 
         # Identify "real" (batch 0) and "fake" (batch != 0) for WGAN framing.
         # The discriminator tries to distinguish batch 0 from the rest.

{diffbio-0.1.6 → diffbio-0.1.8}/src/diffbio/pipelines/joint_training.py

@@ -36,7 +36,7 @@ from substrax.optim import OptimizerConfig, create_optimizer
 
 from diffbio.losses.singlecell_losses import gene_weight_sparsity_loss
 from diffbio.pipelines.joint_preprocessing import JointPreprocessingPipeline
-from diffbio.utils.training import cross_entropy_loss
+from calibrax.metrics.functional import softmax_cross_entropy
 
 _NUM_LOSSES = 2
 
@@ -94,7 +94,7 @@ def _classification_loss(
 ) -> jnp.ndarray:
     """Cross-entropy of the pipeline's predicted logits against the labels."""
     output, _, _ = pipeline.apply({"counts": counts}, {}, None)
-    return cross_entropy_loss(output["logits"], labels, num_classes=n_classes)
+    return softmax_cross_entropy(output["logits"], labels)
 
 
 def _global_norm(grads: nnx.State) -> jnp.ndarray:

{diffbio-0.1.6 → diffbio-0.1.8}/src/diffbio/pipelines/minibatch_training.py

@@ -27,7 +27,7 @@ import numpy as np
 from flax import nnx
 from substrax.optim import create_optimizer, OptimizerConfig
 
-from diffbio.utils.training import cross_entropy_loss
+from calibrax.metrics.functional import softmax_cross_entropy
 
 
 # ``features`` may be a single array or an arbitrary pytree of per-sample arrays (each
@@ -140,9 +140,7 @@ def train_minibatch(
     optimizer = create_optimizer(model, config.optimizer)
 
     def loss_fn(module: nnx.Module, batch_features: Any, batch_labels: jnp.ndarray) -> jnp.ndarray:
-        loss = cross_entropy_loss(
-            forward_fn(module, batch_features), batch_labels, num_classes=n_classes
-        )
+        loss = softmax_cross_entropy(forward_fn(module, batch_features), batch_labels)
         if aux_loss_fn is not None:
             loss = loss + aux_loss_fn(module)
         return loss

{diffbio-0.1.6 → diffbio-0.1.8}/src/diffbio/utils/__init__.py

@@ -20,7 +20,6 @@ from diffbio.utils.quality import apply_quality_filter
 from diffbio.utils.training import (
     default_training_optimizer,
     create_synthetic_training_data,
-    cross_entropy_loss,
     data_iterator,
     Trainer,
     TrainingConfig,
@@ -42,7 +41,6 @@ __all__ = [
     "TrainingState",
     "default_training_optimizer",
     "create_synthetic_training_data",
-    "cross_entropy_loss",
     "data_iterator",
     # Quality utilities
     "apply_quality_filter",

{diffbio-0.1.6 → diffbio-0.1.8}/src/diffbio/utils/training.py

@@ -12,6 +12,7 @@ from typing import Any, Callable
 import jax
 import jax.numpy as jnp
 from datarax.core.operator import OperatorModule
+from calibrax.metrics.functional import softmax_cross_entropy  # noqa: F401 - the docstring example
 from flax import nnx
 from jaxtyping import Array, Float
 from substrax.optim import OptimizerConfig, create_optimizer
@@ -62,26 +63,6 @@ class TrainingState:
             self.loss_history = []
 
 
-def cross_entropy_loss(
-    logits: Float[Array, "... num_classes"],
-    labels: Float[Array, "..."],
-    num_classes: int = 3,
-) -> Float[Array, ""]:
-    """Compute cross-entropy loss for variant classification.
-
-    Args:
-        logits: Raw model predictions
-        labels: Integer class labels
-        num_classes: Number of classes
-
-    Returns:
-        Scalar loss value
-    """
-    one_hot_labels = jax.nn.one_hot(labels.astype(jnp.int32), num_classes)
-    log_probs = jax.nn.log_softmax(logits, axis=-1)
-    return -jnp.mean(jnp.sum(one_hot_labels * log_probs, axis=-1))
-
-
 class Trainer:
     """Training loop for DiffBio pipelines using Flax NNX patterns.
 
@@ -96,10 +77,7 @@ class Trainer:
         trainer = Trainer(pipeline, TrainingConfig(learning_rate=1e-3))
         # Define loss function
         def loss_fn(predictions, targets):
-            return cross_entropy_loss(
-                predictions["logits"],
-                targets["labels"],
-            )
+            return softmax_cross_entropy(predictions["logits"], targets["labels"])
         # Train
         trainer.train(data_iterator_fn, loss_fn)
         trained_pipeline = trainer.pipeline