msdt 1.0.0__py3-none-any.whl

msdt/__init__.py

@@ -0,0 +1,112 @@
+"""msdt - multi-medium DT50 persistence prediction library.
+
+Combines the four medium-optimal heteroscedastic models
+(soil / sediment / water_ph7 / water_ph9). Given a DreaMS mass-spectrometry
+embedding (1024-d) it returns:
+  - predictive interval (mu +/- 1.96*sigma_latent, log10-days and days)
+  - persistence probabilities p(P) / p(vP) (120 d / 180 d thresholds)
+  - P / vP decisions made with the best threshold of the chosen scenario
+    (``prec`` = precision-oriented 0.7, ``rec`` = recall-oriented 0.7)
+
+Usage:
+    import msdt
+    out = msdt.predict("soil", "prec", dreams_1024d, models_dir="D:/models")
+    outs = msdt.predict_batch("water_ph7", "rec", dreams_Nx1024, models_dir="D:/models")
+"""
+from __future__ import annotations
+
+from typing import Any, Dict, List, Optional, Union
+
+import numpy as np
+
+from .engine import (
+    MEDIUM_CFG,
+    MEDIUMS,
+    SCENARIOS,
+    get_engine,
+    validate_scenario,
+    _validate_vector,
+)
+
+__all__ = ["predict", "predict_batch", "MEDIUMS", "SCENARIOS"]
+__version__ = "1.0.0"
+
+
+def _single_result(medium: str, scenario: str, r: Dict[str, Any]) -> Dict[str, Any]:
+    return {
+        "medium": medium,
+        "scenario": scenario,
+        "mu_log10": float(r["mu_log10"][0]),
+        "sigma_latent_log10": float(r["sigma_latent_log10"][0]),
+        "interval_log10": [float(r["interval_log10"][0, 0]), float(r["interval_log10"][0, 1])],
+        "interval_days": [float(r["interval_days"][0, 0]), float(r["interval_days"][0, 1])],
+        "p_P": float(r["p_P"][0]),
+        "p_vP": float(r["p_vP"][0]),
+        "threshold_P": float(r["threshold_P"]),
+        "threshold_vP": float(r["threshold_vP"]),
+        "is_P": bool(r["is_P"][0]),
+        "is_vP": bool(r["is_vP"][0]),
+    }
+
+
+def predict(
+    medium: str,
+    scenario: str,
+    dreams: Union[np.ndarray, List[float]],
+    models_dir: Optional[str] = None,
+) -> Dict[str, Any]:
+    """Predict for a single compound.
+
+    Args:
+        medium: "soil" | "sediment" | "water_ph7" | "water_ph9"
+        scenario: "prec" (precision >= 0.7, max recall) | "rec" (recall >= 0.7, max precision)
+        dreams: DreaMS embedding, shape (1024,)
+        models_dir: optional directory containing {medium}/model_final.pth;
+            overrides MSDT_MODELS_DIR
+
+    Returns:
+        dict with predictive interval, p(P)/p(vP), thresholds and P/vP decisions
+    """
+    medium = str(medium).lower().strip()
+    scenario = validate_scenario(scenario)
+    x = _validate_vector(dreams, batch=False)
+    r = get_engine(medium, models_dir).run(x, scenario)
+    return _single_result(medium, scenario, r)
+
+
+def predict_batch(
+    medium: str,
+    scenario: str,
+    dreams: Union[np.ndarray, List[List[float]]],
+    models_dir: Optional[str] = None,
+) -> Dict[str, Any]:
+    """Predict for a batch of compounds.
+
+    Args:
+        medium: "soil" | "sediment" | "water_ph7" | "water_ph9"
+        scenario: "prec" | "rec"
+        dreams: DreaMS embedding matrix, shape (N, 1024)
+        models_dir: optional directory containing {medium}/model_final.pth
+
+    Returns:
+        dict; scalar fields become arrays of length N (interval_* are (N, 2))
+    """
+    medium = str(medium).lower().strip()
+    scenario = validate_scenario(scenario)
+    x = _validate_vector(dreams, batch=True)
+    r = get_engine(medium, models_dir).run(x, scenario)
+    return {
+        "medium": medium,
+        "scenario": scenario,
+        "n": int(len(x)),
+        "mu_log10": r["mu_log10"],
+        "sigma_latent_log10": r["sigma_latent_log10"],
+        "interval_log10": r["interval_log10"],
+        "interval_days": r["interval_days"],
+        "p_P": r["p_P"],
+        "p_vP": r["p_vP"],
+        "threshold_P": float(r["threshold_P"]),
+        "threshold_vP": float(r["threshold_vP"]),
+        "is_P": r["is_P"],
+        "is_vP": r["is_vP"],
+    }