diffbio 0.1.3__tar.gz → 0.1.4__tar.gz

{diffbio-0.1.3 → diffbio-0.1.4}/PKG-INFO

@@ -1,7 +1,7 @@
 Metadata-Version: 2.5
 Name: diffbio
-Version: 0.1.3
-Summary: DiffBio: End-to-end differentiable bioinformatics pipelines built on Datarax, Artifex, Opifex, and Calibrax
+Version: 0.1.4
+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
 Project-URL: Source, https://github.com/avitai/DiffBio
@@ -74,6 +74,7 @@ Requires-Dist: griffe>=1.7.3; extra == 'all'
 Requires-Dist: import-linter>=2.5; extra == 'all'
 Requires-Dist: interrogate>=1.7.0; extra == 'all'
 Requires-Dist: ipykernel>=6.29.5; extra == 'all'
+Requires-Dist: jax-metal>=0.1.0; (sys_platform == 'darwin' and platform_machine == 'arm64') and extra == 'all'
 Requires-Dist: jax[cuda12]>=0.11.1; extra == 'all'
 Requires-Dist: lineax>=0.0.8; extra == 'all'
 Requires-Dist: matplotlib>=3.7; extra == 'all'
@@ -101,6 +102,7 @@ Requires-Dist: pytest-timeout>=2.1; extra == 'all'
 Requires-Dist: pytest-xdist>=3.6; extra == 'all'
 Requires-Dist: pytest>=8.3.5; extra == 'all'
 Requires-Dist: python-dotenv>=1; extra == 'all'
+Requires-Dist: pyyaml>=6; extra == 'all'
 Requires-Dist: radon>=6.0.1; extra == 'all'
 Requires-Dist: rdkit>=2025.9.3; extra == 'all'
 Requires-Dist: ruff>=0.1.5; extra == 'all'
@@ -185,6 +187,8 @@ Requires-Dist: pymdown-extensions>=10.14.3; extra == 'docs'
 Provides-Extra: genomics
 Requires-Dist: pyfaidx>=0.8.0; extra == 'genomics'
 Requires-Dist: pysam>=0.22.0; extra == 'genomics'
+Provides-Extra: metal
+Requires-Dist: jax-metal>=0.1.0; (sys_platform == 'darwin' and platform_machine == 'arm64') and extra == 'metal'
 Provides-Extra: soft-ops-advanced
 Requires-Dist: lineax>=0.0.8; extra == 'soft-ops-advanced'
 Requires-Dist: optimistix>=0.0.9; extra == 'soft-ops-advanced'
@@ -203,6 +207,7 @@ Requires-Dist: pytest-randomly>=3.16.0; extra == 'test'
 Requires-Dist: pytest-timeout>=2.1; extra == 'test'
 Requires-Dist: pytest-xdist>=3.6; extra == 'test'
 Requires-Dist: pytest>=8.3.5; extra == 'test'
+Requires-Dist: pyyaml>=6; extra == 'test'
 Provides-Extra: torch-io
 Requires-Dist: torch>=1.13.0; extra == 'torch-io'
 Description-Content-Type: text/markdown

{diffbio-0.1.3 → diffbio-0.1.4}/benchmarks/README.md

@@ -90,13 +90,13 @@ Download datasets before running:
 
 ```bash
 # Single-cell (immune_human, 2GB)
-mkdir -p /media/mahdi/ssd23/Data/scib
-wget -O /media/mahdi/ssd23/Data/scib/Immune_ALL_human.h5ad \
+mkdir -p /mnt/ssd2/Data/scib
+wget -O /mnt/ssd2/Data/scib/Immune_ALL_human.h5ad \
     "https://ndownloader.figshare.com/files/25717328"
 
 # Trajectory (pancreas, 51MB)
-mkdir -p /media/mahdi/ssd23/Data/scvelo
-wget -O /media/mahdi/ssd23/Data/scvelo/endocrinogenesis_day15.h5ad \
+mkdir -p /mnt/ssd2/Data/scvelo
+wget -O /mnt/ssd2/Data/scvelo/endocrinogenesis_day15.h5ad \
     "https://github.com/theislab/scvelo_notebooks/raw/master/data/Pancreas/endocrinogenesis_day15.h5ad"
 ```
 

{diffbio-0.1.3 → diffbio-0.1.4}/pyproject.toml

@@ -64,16 +64,16 @@ dependencies = [
   "h5py>=3.7",
   "jax-md>=0.2.27"
 ]
-description = "DiffBio: End-to-end differentiable bioinformatics pipelines built on Datarax, Artifex, Opifex, and Calibrax"
+description = "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"
 keywords = ["jax", "flax", "bioinformatics", "differentiable", "variant-calling", "alignment", "machine-learning"]
 license = {file = "LICENSE"}
 name = "diffbio"
 readme = "README.md"
 requires-python = ">=3.12,<3.14"
-version = "0.1.3"
+version = "0.1.4"
 
 [project.optional-dependencies]
-all = ["diffbio[benchmark,chem,cuda12,dev,docs,genomics,soft-ops-advanced,soft-ops-ot,test,torch-io]"]
+all = ["diffbio[benchmark,chem,cuda12,dev,docs,genomics,metal,soft-ops-advanced,soft-ops-ot,test,torch-io]"]
 benchmark = [
   # benchmarks/crossmodality/*.py call datasets.load_dataset. It used to arrive
   # transitively through the sibling packages; they now declare their heavy
@@ -131,6 +131,8 @@ docs = [
   "pymdown-extensions>=10.14.3"
 ]
 genomics = ["pysam>=0.22.0", "pyfaidx>=0.8.0"]
+# Apple Silicon GPU backend; setup.sh selects it on arm64 macOS.
+metal = ["jax-metal>=0.1.0; sys_platform == 'darwin' and platform_machine == 'arm64'"]
 soft-ops-advanced = ["optimistix>=0.0.9", "lineax>=0.0.8"]
 soft-ops-ot = ["diffbio[soft-ops-advanced]", "ott-jax>=0.5.0"]
 test = [
@@ -143,7 +145,9 @@ test = [
   "pytest-env>=1.0.1",
   "pytest-randomly>=3.16.0",
   "pytest-timeout>=2.1",
-  "pytest-xdist>=3.6"
+  "pytest-xdist>=3.6",
+  # tests/test_ci_shards.py reads the CI workflow.
+  "pyyaml>=6"
 ]
 torch-io = ["torch>=1.13.0"]
 
@@ -341,7 +345,7 @@ save_data = true
 warmup = true
 warmup_iterations = 3
 
-[tool.pytest-env]
+[tool.pytest_env]
 JAX_ENABLE_X64 = "0"
 JAX_SKIP_CUDA_CONSTRAINTS_CHECK = "1"
 LD_LIBRARY_PATH = "/usr/local/cuda/lib64"

{diffbio-0.1.3 → diffbio-0.1.4}/src/diffbio/core/soft_ops/__init__.py

@@ -81,6 +81,9 @@ from diffbio.core.soft_ops.sorting import (
     top_k_mask,
 )
 
+# --- Range-aware normalization ---
+from diffbio.core.soft_ops.normalization import temperature_softmax
+
 # --- Quantile ---
 from diffbio.core.soft_ops.quantile import (
     argmedian,
@@ -126,6 +129,7 @@ from diffbio.core.soft_ops.straight_through import (
 )
 
 __all__ = [
+    "temperature_softmax",
     # Types
     "SoftBool",
     "SoftIndex",

diffbio-0.1.4/src/diffbio/core/soft_ops/normalization.py

@@ -0,0 +1,174 @@
+"""Temperature softmax with jointly evaluated derivative coefficients.
+
+Ordinary division AD forms ``T**-2`` before multiplying the softmax tail.
+The product can be finite even when that reciprocal overflows or the tail
+underflows. Evaluate probability products, score gaps and reciprocal powers
+together in log space. A recursive coefficient JVP preserves mixed derivatives
+at zero gaps; differentiating an outer zero selection would lose them.
+
+This requires representable coefficient sums, not merely a representable final
+contraction. Work grows with donor count and derivative order: the first JVP
+has quadratic donor work. It is intended for small donor panels.
+"""
+
+from functools import partial
+
+import jax
+import jax.numpy as jnp
+from jax import Array
+
+from diffbio.core.soft_ops._utils import canonicalize_axis
+
+
+type _Specification = tuple[tuple[int, ...], tuple[tuple[int, int], ...], int]
+type _Inputs = tuple[Array, Array, Array]
+
+
+def _logits(scores: Array, temperature: Array, covered: Array) -> Array:
+    """Shift over covered donors before division, guarding excluded operands."""
+    maximum = jnp.max(jnp.where(covered, scores, -jnp.inf), axis=0, keepdims=True)
+    shifted = jnp.where(covered, scores - maximum, 0.0)
+    return jnp.where(covered, shifted / temperature, -jnp.inf)
+
+
+def _coefficient_value(
+    scores: Array, temperature: Array, covered: Array, specification: _Specification
+) -> Array:
+    """Evaluate a probability/gap monomial divided by a temperature power."""
+    indices, gaps, power = specification
+    log_weights = jax.nn.log_softmax(_logits(scores, temperature, covered), axis=0)
+    magnitude = sum(log_weights[index] for index in indices) - power * jnp.log(temperature)
+    sign = jnp.ones_like(magnitude)
+    nonzero = jnp.ones_like(magnitude, dtype=jnp.bool_)
+    for left, right in gaps:
+        gap = scores[left] - scores[right]
+        nonzero = nonzero & (gap != 0)
+        magnitude = magnitude + jnp.log(jnp.where(gap != 0, jnp.abs(gap), 1.0))
+        sign = sign * jnp.sign(gap)
+    # Guard the exponential too: a zero gap must not create 0 * inf.
+    value = sign * jnp.exp(jnp.where(nonzero, magnitude, 0.0))
+    return jnp.where(nonzero, value, 0.0)
+
+
+_coefficient = jax.custom_jvp(_coefficient_value, nondiff_argnums=(3,))
+
+
+def _coefficient_jvp(
+    specification: _Specification, primals: _Inputs, tangents: _Inputs
+) -> tuple[Array, Array]:
+    """Differentiate complete coefficients, including zero-gap extensions."""
+    scores, _, _ = primals
+    score_dot, temperature_dot, _ = tangents
+    indices, gaps, power = specification
+    value = _coefficient(*primals, specification)
+    derivative = jnp.zeros_like(value)
+    thermal = jnp.zeros_like(value)
+    for index in indices:
+        for other in range(scores.shape[0]):
+            if other != index:
+                extended = (*indices, other)
+                coefficient = _coefficient(*primals, (extended, gaps, power + 1))
+                derivative = derivative + coefficient * (score_dot[index] - score_dot[other])
+                thermal = thermal + _coefficient(
+                    *primals, (extended, (*gaps, (other, index)), power + 2)
+                )
+    for position, (left, right) in enumerate(gaps):
+        remaining = gaps[:position] + gaps[position + 1 :]
+        derivative = derivative + _coefficient(*primals, (indices, remaining, power)) * (
+            score_dot[left] - score_dot[right]
+        )
+    if power:
+        thermal = thermal - power * _coefficient(*primals, (indices, gaps, power + 1))
+    return value, derivative + thermal * temperature_dot
+
+
+_coefficient.defjvp(_coefficient_jvp)
+
+
+@jax.custom_jvp
+def _weights(scores: Array, temperature: Array, covered: Array) -> Array:
+    """Keep the native shifted forward softmax arithmetic."""
+    return jax.nn.softmax(_logits(scores, temperature, covered), axis=0)
+
+
+def _weights_jvp(primals: _Inputs, tangents: _Inputs) -> tuple[Array, Array]:
+    """Contract pairwise sensitivities without rounding either probability first."""
+    scores, _, _ = primals
+    score_dot, temperature_dot, _ = tangents
+    value = _weights(*primals)
+    rows = []
+    for index in range(scores.shape[0]):
+        derivative = jnp.zeros_like(value[index])
+        for other in range(scores.shape[0]):
+            if other != index:
+                probability = (index, other)
+                spatial = _coefficient(*primals, (probability, (), 1))
+                thermal = _coefficient(*primals, (probability, ((other, index),), 2))
+                derivative = derivative + spatial * (score_dot[index] - score_dot[other])
+                derivative = derivative + thermal * temperature_dot
+        rows.append(derivative)
+    return value, jnp.stack(rows)
+
+
+_weights.defjvp(_weights_jvp)
+
+
+@partial(jax.jit, static_argnames=("axis",))
+def temperature_softmax(
+    scores: Array,
+    temperature: float | Array = 1.0,
+    *,
+    axis: int = -1,
+    where: Array | None = None,
+) -> Array:
+    """Normalize small score axes with range-aware temperature derivatives.
+
+    Equivalent to masked ``jax.nn.softmax(scores / temperature, axis=axis)``
+    for valid inputs. Shift before division and evaluate complete derivative
+    coefficients in signed log space, preserving sensitivities even when a
+    probability rounds to zero or one. Supports JVP, VJP and higher derivatives.
+
+    This is an opt-in operation for small axes, not the default sorting kernel:
+    first derivatives require quadratic axis work and higher orders cost more.
+    Coefficients, score differences and their necessary sums must be representable;
+    a representable final contraction alone does not guarantee a finite derivative.
+
+    Args:
+        scores: Real floating array. Included scores and their pairwise differences
+            must be finite. Excluded scores are ignored, including NaN/infinity.
+        temperature: Positive finite real scalar. Callers own runtime validation;
+            nonpositive/nonfinite values are outside this numerical contract.
+        axis: Static reduction axis. Must be nonempty.
+        where: Boolean mask broadcastable to the score shape. Empty slices return
+            zero weights and zero derivatives, matching native masked softmax.
+
+    Returns:
+        Weights with the score shape and standard floating dtype promotion.
+
+    Raises:
+        TypeError: Scores are not floating, temperature is complex, or mask is not boolean.
+        ValueError: Temperature is not scalar, the axis is invalid/empty, or mask cannot broadcast.
+    """
+    scores = jnp.asarray(scores)
+    temperature = jnp.asarray(temperature)
+    if not jnp.issubdtype(scores.dtype, jnp.floating):
+        raise TypeError("scores must have a real floating dtype")
+    if jnp.issubdtype(temperature.dtype, jnp.complexfloating):
+        raise TypeError("temperature must be real")
+    if temperature.ndim != 0:
+        raise ValueError("temperature must be scalar")
+    temperature = temperature.astype(jnp.result_type(scores, temperature))
+    axis = canonicalize_axis(axis, scores.ndim)
+    if scores.shape[axis] == 0:
+        raise ValueError("the normalization axis must be nonempty")
+    covered = jnp.ones_like(scores, dtype=jnp.bool_) if where is None else jnp.asarray(where)
+    if covered.dtype != jnp.bool_:
+        raise TypeError("where must have boolean dtype")
+    covered = jnp.broadcast_to(covered, scores.shape)
+    scores, covered = jnp.moveaxis(scores, axis, 0), jnp.moveaxis(covered, axis, 0)
+    scores = jnp.where(covered, scores, 0.0)
+    populated = jnp.any(covered, axis=0, keepdims=True)
+    first = (jnp.arange(scores.shape[0]) == 0).reshape((-1,) + (1,) * (scores.ndim - 1))
+    safe_coverage = covered | (first & ~populated)
+    value = _weights(scores, temperature, safe_coverage)
+    return jnp.moveaxis(jnp.where(populated, value, 0.0), 0, axis)

{diffbio-0.1.3 → diffbio-0.1.4}/src/diffbio/sources/archive_ii.py

@@ -32,7 +32,7 @@ from flax import nnx
 
 logger = logging.getLogger(__name__)
 
-_DEFAULT_DATA_DIR = "/media/mahdi/ssd23/Works/RNAFoldAssess/tutorial/processed_data"
+_DEFAULT_DATA_DIR = "/mnt/ssd2/Works/RNAFoldAssess/tutorial/processed_data"
 _STRUCTURE_FILENAME = "example_data_structure.csv"
 
 

{diffbio-0.1.3 → diffbio-0.1.4}/src/diffbio/sources/balifam.py

@@ -35,7 +35,7 @@ from flax import nnx
 
 logger = logging.getLogger(__name__)
 
-_DEFAULT_DATA_DIR = "/media/mahdi/ssd23/Works/balifam"
+_DEFAULT_DATA_DIR = "/mnt/ssd2/Works/balifam"
 
 
 def _parse_fasta(path: Path) -> list[tuple[str, str]]:

{diffbio-0.1.3 → diffbio-0.1.4}/src/diffbio/sources/bengrn_ground_truth.py

@@ -3,7 +3,7 @@
 Loads mESC expression data and ChIP+Perturb ground truth edges from
 the benGRN repository (Stone & Sroy gold standards).
 
-Data source: /media/mahdi/ssd23/Works/benGRN/data/GroundTruth/
+Data source: /mnt/ssd2/Works/benGRN/data/GroundTruth/
 
 References:
     - benGRN: https://github.com/your-org/benGRN
@@ -28,7 +28,7 @@ from datarax.core.data_source import DataSourceModule
 
 logger = logging.getLogger(__name__)
 
-_BASE_DIR = Path("/media/mahdi/ssd23/Works/benGRN/data/GroundTruth/stone_and_sroy")
+_BASE_DIR = Path("/mnt/ssd2/Works/benGRN/data/GroundTruth/stone_and_sroy")
 
 
 @dataclass(frozen=True, kw_only=True)

{diffbio-0.1.3 → diffbio-0.1.4}/src/diffbio/sources/encode_peaks.py

@@ -33,7 +33,7 @@ from flax import nnx
 
 logger = logging.getLogger(__name__)
 
-_DEFAULT_DATA_PATH = "/media/mahdi/ssd23/Data/encode/CTCF_K562_narrowPeak.bed.gz"
+_DEFAULT_DATA_PATH = "/mnt/ssd2/Data/encode/CTCF_K562_narrowPeak.bed.gz"
 
 
 @dataclass(frozen=True, kw_only=True)

{diffbio-0.1.3 → diffbio-0.1.4}/src/diffbio/sources/immune_human.py

@@ -48,7 +48,7 @@ class ImmuneHumanConfig(StructuralConfig):
         embedding_key: Key in obsm for precomputed embeddings.
     """
 
-    data_dir: str = "/media/mahdi/ssd23/Data/scib"
+    data_dir: str = "/mnt/ssd2/Data/scib"
     subsample: int | None = None
     batch_key: str = "batch"
     label_key: str = "final_annotation"

{diffbio-0.1.3 → diffbio-0.1.4}/src/diffbio/sources/pancreas.py

@@ -42,7 +42,7 @@ class PancreasConfig(StructuralConfig):
         cluster_key: Column in obs for cell type labels.
     """
 
-    data_dir: str = "/media/mahdi/ssd23/Data/scvelo"
+    data_dir: str = "/mnt/ssd2/Data/scvelo"
     subsample: int | None = None
     cluster_key: str = "clusters"
 

{diffbio-0.1.3 → diffbio-0.1.4}/src/diffbio/sources/seqfish.py

@@ -46,7 +46,7 @@ class SeqFISHConfig(StructuralConfig):
         spatial_key: Key in obsm for spatial coordinates.
     """
 
-    data_dir: str = "/media/mahdi/ssd23/Data/spatial"
+    data_dir: str = "/mnt/ssd2/Data/spatial"
     subsample: int | None = None
     label_key: str = "celltype_mapped_refined"
     spatial_key: str = "spatial"