dnakit 0.1.0.dev0__py3-none-any.whl

dnakit/evaluation/distribution.py

@@ -0,0 +1,287 @@
+"""Transparent descriptive distribution comparison between two DNA collections."""
+
+from __future__ import annotations
+
+import math
+from collections import Counter
+
+from dnakit.core import DNARecord
+from dnakit.descriptors import exact_repeat_fraction, gc_at_content
+
+from ._shared import EvaluationInput, materialize_input, mean, record_for, report, require_nonempty
+from .config import DistributionEvaluationConfig
+from .results import EvaluationReport
+
+
+def _ecdf_distance(left: tuple[float, ...], right: tuple[float, ...]) -> float:
+    """Return the exact two-sample Kolmogorov-Smirnov D statistic."""
+
+    if not left or not right:
+        return 0.0
+    support = sorted(set(left) | set(right))
+    left_sorted = sorted(left)
+    right_sorted = sorted(right)
+    left_index = 0
+    right_index = 0
+    distance = 0.0
+    for value in support:
+        while left_index < len(left_sorted) and left_sorted[left_index] <= value:
+            left_index += 1
+        while right_index < len(right_sorted) and right_sorted[right_index] <= value:
+            right_index += 1
+        distance = max(
+            distance,
+            abs(left_index / len(left_sorted) - right_index / len(right_sorted)),
+        )
+    return distance
+
+
+def _total_variation(left: Counter[str], right: Counter[str]) -> float:
+    left_total = sum(left.values())
+    right_total = sum(right.values())
+    if left_total == 0 and right_total == 0:
+        return 0.0
+    support = set(left) | set(right)
+    return 0.5 * math.fsum(
+        abs(
+            (left[key] / left_total if left_total else 0.0)
+            - (right[key] / right_total if right_total else 0.0)
+        )
+        for key in support
+    )
+
+
+def _kmer_counts(records: tuple[DNARecord, ...], k: int, canonical: bool) -> Counter[str]:
+    complement = str.maketrans("ACGT", "TGCA")
+    counts: Counter[str] = Counter()
+    for record in records:
+        for part in record.sequence.parts:
+            if not isinstance(part, str):
+                continue
+            for start in range(max(0, len(part) - k + 1)):
+                word = part[start : start + k]
+                if set(word) - set("ACGT"):
+                    continue
+                if canonical:
+                    word = min(word, word.translate(complement)[::-1])
+                counts[word] += 1
+    return counts
+
+
+def _kmer_observation_count(
+    records: tuple[DNARecord, ...],
+    k: int,
+    *,
+    initial_count: int,
+    maximum: int,
+) -> int:
+    """Count eligible k-mers without allocating a frequency table.
+
+    The scan stops at the first observation beyond ``maximum`` so an
+    over-budget input cannot build either collection's ``Counter`` first.
+    """
+
+    observation_count = initial_count
+    canonical_symbols = frozenset("ACGT")
+    for record in records:
+        for part in record.sequence.parts:
+            if not isinstance(part, str):
+                continue
+            for start in range(max(0, len(part) - k + 1)):
+                if set(part[start : start + k]) - canonical_symbols:
+                    continue
+                observation_count += 1
+                if observation_count > maximum:
+                    from dnakit.exceptions import ConfigurationError
+
+                    raise ConfigurationError(
+                        "Distribution comparison exceeds max_kmer_observations.",
+                        code="DISTRIBUTION_KMER_LIMIT",
+                        context={
+                            "observation_count": observation_count,
+                            "max_kmer_observations": maximum,
+                        },
+                    )
+    return observation_count
+
+
+def _motif_rates(records: tuple[DNARecord, ...], motifs: tuple[str, ...]) -> dict[str, float]:
+    denominator = sum(record.sequence.symbol_length for record in records)
+    return {
+        motif: (
+            sum(
+                part.count(motif)
+                for record in records
+                for part in record.sequence.parts
+                if isinstance(part, str)
+            )
+            / denominator
+            if denominator
+            else 0.0
+        )
+        for motif in motifs
+    }
+
+
+def _feature_distances(
+    left: tuple[DNARecord, ...],
+    right: tuple[DNARecord, ...],
+    config: DistributionEvaluationConfig,
+) -> tuple[dict[str, float], dict[str, object]]:
+    distances: dict[str, float] = {}
+    details: dict[str, object] = {}
+    if "length" in config.features:
+        left_values = tuple(float(record.sequence.symbol_length) for record in left)
+        right_values = tuple(float(record.sequence.symbol_length) for record in right)
+        distances["length"] = _ecdf_distance(left_values, right_values)
+        details["length"] = {
+            "distance": distances["length"],
+            "method": "two-sample empirical-CDF supremum distance (KS D; no p-value)",
+            "left_mean": mean(left_values),
+            "right_mean": mean(right_values),
+        }
+    if "gc" in config.features:
+        left_values = tuple(
+            result.gc_fraction
+            for record in left
+            if (result := gc_at_content(record, ambiguity_policy="ignore")).gc_fraction is not None
+        )
+        right_values = tuple(
+            result.gc_fraction
+            for record in right
+            if (result := gc_at_content(record, ambiguity_policy="ignore")).gc_fraction is not None
+        )
+        distances["gc"] = _ecdf_distance(left_values, right_values)
+        details["gc"] = {
+            "distance": distances["gc"],
+            "method": "two-sample empirical-CDF supremum distance (KS D; no p-value)",
+            "left_mean": mean(left_values),
+            "right_mean": mean(right_values),
+            "undefined_gc_omitted": (len(left) - len(left_values))
+            + (len(right) - len(right_values)),
+        }
+    if "kmer" in config.features:
+        left_observations = _kmer_observation_count(
+            left,
+            config.k,
+            initial_count=0,
+            maximum=config.max_kmer_observations,
+        )
+        observation_count = _kmer_observation_count(
+            right,
+            config.k,
+            initial_count=left_observations,
+            maximum=config.max_kmer_observations,
+        )
+        left_counts = _kmer_counts(left, config.k, config.canonical)
+        right_counts = _kmer_counts(right, config.k, config.canonical)
+        distances["kmer"] = _total_variation(left_counts, right_counts)
+        details["kmer"] = {
+            "distance": distances["kmer"],
+            "method": "total variation over pooled exact k-mer frequency vectors",
+            "left_observations": left_observations,
+            "right_observations": observation_count - left_observations,
+        }
+    if "motif" in config.features:
+        left_rates = _motif_rates(left, config.motifs)
+        right_rates = _motif_rates(right, config.motifs)
+        per_motif = {motif: abs(left_rates[motif] - right_rates[motif]) for motif in config.motifs}
+        distances["motif"] = math.fsum(per_motif.values()) / len(per_motif)
+        details["motif"] = {
+            "distance": distances["motif"],
+            "method": "mean absolute difference of non-overlapping motif counts per symbol",
+            "left_rates": left_rates,
+            "right_rates": right_rates,
+            "absolute_differences": per_motif,
+        }
+    if "repeat" in config.features:
+        left_values = tuple(
+            exact_repeat_fraction(
+                record,
+                ambiguity_policy="ignore",
+                max_comparisons=config.max_repeat_comparisons_per_sequence,
+            ).repeat_fraction
+            for record in left
+        )
+        right_values = tuple(
+            exact_repeat_fraction(
+                record,
+                ambiguity_policy="ignore",
+                max_comparisons=config.max_repeat_comparisons_per_sequence,
+            ).repeat_fraction
+            for record in right
+        )
+        distances["repeat"] = _ecdf_distance(left_values, right_values)
+        details["repeat"] = {
+            "distance": distances["repeat"],
+            "method": "two-sample empirical-CDF supremum distance of exact-repeat union fraction",
+            "left_mean": mean(left_values),
+            "right_mean": mean(right_values),
+        }
+    return distances, details
+
+
+def evaluate_distribution_similarity(
+    left: EvaluationInput,
+    right: EvaluationInput,
+    *,
+    config: DistributionEvaluationConfig | None = None,
+) -> EvaluationReport:
+    """Compare length, GC, k-mer, motif, and repeat distributions without p-values."""
+
+    resolved = DistributionEvaluationConfig() if config is None else config
+    if not isinstance(resolved, DistributionEvaluationConfig):
+        raise TypeError("config must be DistributionEvaluationConfig or None.")
+    left_items = materialize_input(left, limits=resolved.limits)
+    right_items = materialize_input(right, limits=resolved.limits)
+    require_nonempty(left_items, "distribution similarity")
+    require_nonempty(right_items, "distribution similarity")
+    if (
+        sum(item.sequence.symbol_length for item in left_items + right_items)
+        > resolved.limits.max_total_symbols
+    ):
+        from dnakit.exceptions import ConfigurationError
+
+        raise ConfigurationError(
+            "Combined distribution inputs exceed max_total_symbols.",
+            code="EVALUATION_SYMBOL_LIMIT",
+        )
+    left_records = tuple(record_for(item) for item in left_items)
+    right_records = tuple(record_for(item) for item in right_items)
+    distances, details = _feature_distances(left_records, right_records, resolved)
+    aggregate_distance = math.fsum(distances.values()) / len(distances)
+    return report(
+        name="distribution_similarity",
+        method="mean-of-transparent-feature-distances",
+        version="eval-distribution-v1",
+        parameters={
+            "features": resolved.features,
+            "k": resolved.k,
+            "canonical": resolved.canonical,
+            "motifs": resolved.motifs,
+            "aggregation": "1 - arithmetic mean of feature distances",
+            "inference": "descriptive distances only; no hypothesis-test p-values",
+            "motif_counting": "str.count non-overlapping",
+            "gap_policy": "analyze text parts separately; no k-mer or motif crosses a Gap",
+            "topology_policy": "linearized at stored origin; no circular-origin wrap",
+            "ambiguity_policy": (
+                "length retains IUPAC symbols; GC, k-mer, motif, and repeat calculations "
+                "exclude affected noncanonical observations"
+            ),
+            "max_kmer_observations": resolved.max_kmer_observations,
+            "max_repeat_comparisons_per_sequence": resolved.max_repeat_comparisons_per_sequence,
+            "limits": resolved.limits,
+        },
+        metrics={
+            "score": 1.0 - aggregate_distance,
+            "distribution_similarity": 1.0 - aggregate_distance,
+            "aggregate_distance": aggregate_distance,
+            "feature_distances": distances,
+            "feature_details": details,
+            "left_count": len(left_records),
+            "right_count": len(right_records),
+        },
+    )
+
+
+__all__ = ["evaluate_distribution_similarity"]

dnakit/evaluation/frechet.py

@@ -0,0 +1,223 @@
+"""Fréchet distance between two DNA representation distributions."""
+
+from __future__ import annotations
+
+import math
+from typing import Any
+
+from dnakit.exceptions import (
+    BackendExecutionError,
+    BackendUnavailableError,
+    ConfigurationError,
+)
+from dnakit.representations import RepresentationBackend, extract_representations
+
+from ._shared import EvaluationInput, materialize_input, record_for, report, require_nonempty
+from .config import FrechetDistanceConfig
+from .results import EvaluationReport
+
+
+def _require_numpy() -> Any:
+    try:
+        import numpy as np
+    except ImportError as exc:
+        raise BackendUnavailableError(
+            "Fréchet DNA distance requires NumPy and a representation backend.",
+            code="MISSING_NEURAL_DEPENDENCY",
+            hint='Install the neural extra with: python -m pip install "dnakit[neural]"',
+        ) from exc
+    return np
+
+
+def _l2_normalize(matrix: Any, np: Any) -> Any:
+    norms = np.linalg.norm(matrix, axis=1, keepdims=True)
+    zero_rows = np.flatnonzero(norms[:, 0] == 0)
+    if len(zero_rows):
+        raise ConfigurationError(
+            "Cannot L2-normalize zero-valued model representations.",
+            code="ZERO_REPRESENTATION_VECTOR",
+            context={"row_indices": tuple(int(value) for value in zero_rows[:20])},
+        )
+    return matrix / norms
+
+
+def _frechet_components(left: Any, right: Any, np: Any) -> dict[str, float]:
+    """Compute Gaussian Fréchet terms through an exact sample-space identity.
+
+    If ``A`` and ``B`` are centered sample matrices scaled by ``sqrt(n - 1)``,
+    the covariance cross term equals the nuclear norm of ``A @ B.T``.  This
+    avoids materializing dense feature-by-feature covariance matrices.
+    """
+
+    left_mean = np.mean(left, axis=0, dtype=np.float64)
+    right_mean = np.mean(right, axis=0, dtype=np.float64)
+    mean_delta = left_mean - right_mean
+    mean_component = float(np.dot(mean_delta, mean_delta))
+
+    left_scaled = (left - left_mean) / math.sqrt(int(left.shape[0]) - 1)
+    right_scaled = (right - right_mean) / math.sqrt(int(right.shape[0]) - 1)
+    left_covariance_trace = float(np.sum(left_scaled * left_scaled, dtype=np.float64))
+    right_covariance_trace = float(np.sum(right_scaled * right_scaled, dtype=np.float64))
+    try:
+        cross_gram = left_scaled @ right_scaled.T
+        singular_values = np.linalg.svd(cross_gram, compute_uv=False)
+    except (ValueError, np.linalg.LinAlgError) as exc:
+        raise BackendExecutionError(
+            "Could not compute the Fréchet covariance cross term.",
+            code="FRECHET_DECOMPOSITION_FAILED",
+            context={"error_type": type(exc).__name__},
+        ) from exc
+    cross_nuclear_norm = float(np.sum(singular_values, dtype=np.float64))
+    raw_covariance_component = (
+        left_covariance_trace + right_covariance_trace - 2.0 * cross_nuclear_norm
+    )
+    tolerance = 1e-10 * max(
+        1.0,
+        left_covariance_trace + right_covariance_trace,
+        cross_nuclear_norm,
+    )
+    if raw_covariance_component < -tolerance:
+        raise BackendExecutionError(
+            "The Fréchet covariance term is numerically inconsistent.",
+            code="INVALID_FRECHET_COVARIANCE_TERM",
+            context={
+                "covariance_component": raw_covariance_component,
+                "tolerance": tolerance,
+            },
+        )
+    covariance_component = max(0.0, raw_covariance_component)
+    distance = mean_component + covariance_component
+    if not math.isfinite(distance):
+        raise BackendExecutionError(
+            "The Fréchet distance is not finite.",
+            code="INVALID_FRECHET_DISTANCE",
+        )
+    return {
+        "frechet_distance": distance,
+        "mean_component": mean_component,
+        "covariance_component": covariance_component,
+        "left_covariance_trace": left_covariance_trace,
+        "right_covariance_trace": right_covariance_trace,
+        "cross_covariance_nuclear_norm": cross_nuclear_norm,
+    }
+
+
+def evaluate_frechet_distance(
+    left: EvaluationInput,
+    right: EvaluationInput,
+    *,
+    config: FrechetDistanceConfig | None = None,
+    backend: RepresentationBackend | None = None,
+) -> EvaluationReport:
+    """Compare two DNA collections with a Gaussian Fréchet representation distance.
+
+    The default representation model is LucaOne, matching ``DATA-027``.  A
+    standard LucaOne backend still requires explicit ``allow_remote_code=True``
+    after the caller reviews its checkpoint code.  Passing ``backend`` enables
+    controlled custom encoders and tests without a checkpoint download.
+    """
+
+    resolved = FrechetDistanceConfig() if config is None else config
+    if not isinstance(resolved, FrechetDistanceConfig):
+        raise ConfigurationError(
+            "config must be FrechetDistanceConfig or None.",
+            code="INVALID_FRECHET_DISTANCE_CONFIG",
+        )
+    left_items = materialize_input(left, limits=resolved.limits)
+    right_items = materialize_input(right, limits=resolved.limits)
+    require_nonempty(left_items, "Fréchet DNA distance")
+    require_nonempty(right_items, "Fréchet DNA distance")
+    if len(left_items) < 2 or len(right_items) < 2:
+        raise ConfigurationError(
+            "Fréchet DNA distance requires at least two records in each collection.",
+            code="FRECHET_MINIMUM_SAMPLE_SIZE",
+            context={"left_count": len(left_items), "right_count": len(right_items)},
+        )
+    total_symbols = sum(item.sequence.symbol_length for item in left_items + right_items)
+    if total_symbols > resolved.limits.max_total_symbols:
+        raise ConfigurationError(
+            "Combined Fréchet inputs exceed max_total_symbols.",
+            code="EVALUATION_SYMBOL_LIMIT",
+            context={
+                "total_symbols": total_symbols,
+                "max_total_symbols": resolved.limits.max_total_symbols,
+            },
+        )
+    cross_gram_elements = len(left_items) * len(right_items)
+    if cross_gram_elements > resolved.max_cross_gram_elements:
+        raise ConfigurationError(
+            "Fréchet DNA distance exceeds max_cross_gram_elements.",
+            code="FRECHET_CROSS_GRAM_LIMIT",
+            context={
+                "cross_gram_elements": cross_gram_elements,
+                "max_cross_gram_elements": resolved.max_cross_gram_elements,
+            },
+        )
+
+    records = tuple(record_for(item) for item in left_items + right_items)
+    representation_result = extract_representations(
+        records,
+        config=resolved.representation,
+        backend=backend,
+    )
+    np = _require_numpy()
+    matrix = np.asarray(representation_result.representations, dtype=np.float64)
+    if resolved.normalize:
+        matrix = _l2_normalize(matrix, np)
+    left_matrix = matrix[: len(left_items)]
+    right_matrix = matrix[len(left_items) :]
+    components = _frechet_components(left_matrix, right_matrix, np)
+
+    return report(
+        name="frechet_dna_distance",
+        method="gaussian-frechet-distance-over-dna-model-representations",
+        version="eval-frechet-dna-v1",
+        parameters={
+            "formula": (
+                "||mu_left-mu_right||_2^2 + trace(Sigma_left + Sigma_right "
+                "- 2*(Sigma_left^(1/2)*Sigma_right*Sigma_left^(1/2))^(1/2))"
+            ),
+            "encoder": representation_result.model_name,
+            "checkpoint_path": representation_result.checkpoint_path,
+            "pooling": representation_result.pooling,
+            "representation_config": {
+                "model": representation_result.model_name,
+                "checkpoint_path": representation_result.checkpoint_path,
+                "model_source_path": resolved.representation.model_source_path,
+                "pooling": representation_result.pooling,
+                "ambiguity_policy": representation_result.ambiguity_policy,
+                "device": resolved.representation.device,
+                "dtype": resolved.representation.dtype,
+                "batch_size": resolved.representation.batch_size,
+                "max_length": resolved.representation.max_length,
+                "allow_remote_code": resolved.representation.allow_remote_code,
+            },
+            "normalize": resolved.normalize,
+            "sample_covariance_ddof": 1,
+            "numerical_method": "exact sample-space cross-Gram nuclear norm",
+            "lower_is_better": True,
+            "max_cross_gram_elements": resolved.max_cross_gram_elements,
+            "limits": resolved.limits,
+            "inference": (
+                "FCD-analogous descriptive distance; not ChemNet FCD and not an "
+                "experimentally validated DNA quality score"
+            ),
+        },
+        metrics={
+            **components,
+            "left_count": len(left_items),
+            "right_count": len(right_items),
+            "embedding_dimension": representation_result.embedding_dimension,
+            "left_covariance_rank_upper_bound": min(
+                len(left_items) - 1,
+                representation_result.embedding_dimension,
+            ),
+            "right_covariance_rank_upper_bound": min(
+                len(right_items) - 1,
+                representation_result.embedding_dimension,
+            ),
+        },
+    )
+
+
+__all__ = ["evaluate_frechet_distance"]