diffbio 0.1.7__tar.gz → 0.1.8__tar.gz

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.5
 Name: diffbio
-Version: 0.1.7
+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,7 +45,7 @@ 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.8

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

@@ -39,7 +39,7 @@ 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
@@ -71,7 +71,7 @@ license = {file = "LICENSE"}
 name = "diffbio"
 readme = "README.md"
 requires-python = ">=3.12,<3.14"
-version = "0.1.7"
+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.7 → 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.7 → 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.