pyaging 0.5.2__tar.gz → 0.5.4__tar.gz

{pyaging-0.5.2 → pyaging-0.5.4}/.gitignore

@@ -10,6 +10,10 @@ __pycache__/
 # Tool caches
 .pytest_cache/
 .ruff_cache/
+.coverage
+.coverage.*
+coverage.xml
+htmlcov/
 
 # Packaging artifacts
 build/
@@ -35,7 +39,7 @@ example_data/
 
 # Generated clock weights and aggregate metadata (published to Hugging Face)
 clocks/weights/
-clocks/metadata/all_clock_metadata.pt
+clocks/metadata/*.pt
 clocks/blood_chemistry_example.pkl
 clocks/blood_chemistry_example.csv
 

{pyaging-0.5.2 → pyaging-0.5.4}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.5
 Name: pyaging
-Version: 0.5.2
+Version: 0.5.4
 Summary: A Python-based compendium of GPU-optimized aging clocks.
 Project-URL: homepage, https://github.com/lucascamillomd/pyaging
 Project-URL: repository, https://github.com/lucascamillomd/pyaging
@@ -19,11 +19,13 @@ Classifier: Topic :: Scientific/Engineering :: Bio-Informatics
 Classifier: Typing :: Typed
 Requires-Python: >=3.11
 Requires-Dist: anndata>=0.10.3
+Requires-Dist: httpx<1,>=0.23
 Requires-Dist: huggingface-hub<2,>=1.3
 Requires-Dist: numpy>=2.0.0
 Requires-Dist: pandas>=2.2.0
 Requires-Dist: rich>=13.0
 Requires-Dist: scikit-learn>=1.3.2
+Requires-Dist: scipy>=1.11
 Requires-Dist: torch>=2.2.0
 Provides-Extra: histone
 Requires-Dist: pybigwig>=0.3.23; extra == 'histone'
@@ -79,7 +81,21 @@ pya.pred.predict_age(adata, ["Horvath2013", "AltumAge", "DunedinPACE"])
 adata.obs.head()
 ```
 
-Clock weights are downloaded on demand from per-clock repositories under the [`pyaging` Hugging Face organization](https://huggingface.co/pyaging) (example data comes from [`lucascamillomd/pyaging-data`](https://huggingface.co/lucascamillomd/pyaging-data)). Set the `PYAGING_DATA_REVISION` environment variable to a release tag (e.g. `v0.3.1`) to pin downloads to a specific data revision for reproducibility; it defaults to `main`, the live data release.
+Clock weights are downloaded on demand from per-clock repositories under the [`pyaging` Hugging Face organization](https://huggingface.co/pyaging) (example data comes from [`lucascamillomd/pyaging-data`](https://huggingface.co/lucascamillomd/pyaging-data)). Set `PYAGING_DATA_REVISION` to an existing data release tag, such as `v0.5.4`, to pin downloads; it defaults to `main`. Version 0.5.4 adds CellDRIFT, MiAge, PCBrainAge, IntrinClock370, PhenoAge partitions, standalone PCGrimAge proxies, and 90 proteomic OrganAge outputs. See the [input and output guide](https://pyaging.readthedocs.io/en/latest/clocks_0_5_4.html). Model weights retain the original authors’ licensing terms.
+
+### Run the same clocks across datasets
+
+Reuse loaded models with a bounded cache:
+
+```python
+clocks = ["Horvath2013", "AltumAge"]
+cache = pya.pred.ClockCache(maxsize=len(clocks))
+for adata in datasets:  # Each dataset is a separate AnnData object.
+    pya.pred.predict_age(adata, clocks, clock_cache=cache)
+cache.clear()
+```
+
+The cache avoids repeated download checks, deserialization, and device transfers. It holds at most `maxsize` models and keeps separate entries for each device and data revision. It is optional; ordinary calls retain no model cache. Clear it to refresh weights from a moving `main` revision. See the [prediction guide](https://pyaging.readthedocs.io/en/latest/pyaging.predict.html) for memory and cohort guidance.
 
 ## ❓ Can't find an aging clock?
 

{pyaging-0.5.2 → pyaging-0.5.4}/README.md

@@ -48,7 +48,21 @@ pya.pred.predict_age(adata, ["Horvath2013", "AltumAge", "DunedinPACE"])
 adata.obs.head()
 ```
 
-Clock weights are downloaded on demand from per-clock repositories under the [`pyaging` Hugging Face organization](https://huggingface.co/pyaging) (example data comes from [`lucascamillomd/pyaging-data`](https://huggingface.co/lucascamillomd/pyaging-data)). Set the `PYAGING_DATA_REVISION` environment variable to a release tag (e.g. `v0.3.1`) to pin downloads to a specific data revision for reproducibility; it defaults to `main`, the live data release.
+Clock weights are downloaded on demand from per-clock repositories under the [`pyaging` Hugging Face organization](https://huggingface.co/pyaging) (example data comes from [`lucascamillomd/pyaging-data`](https://huggingface.co/lucascamillomd/pyaging-data)). Set `PYAGING_DATA_REVISION` to an existing data release tag, such as `v0.5.4`, to pin downloads; it defaults to `main`. Version 0.5.4 adds CellDRIFT, MiAge, PCBrainAge, IntrinClock370, PhenoAge partitions, standalone PCGrimAge proxies, and 90 proteomic OrganAge outputs. See the [input and output guide](https://pyaging.readthedocs.io/en/latest/clocks_0_5_4.html). Model weights retain the original authors’ licensing terms.
+
+### Run the same clocks across datasets
+
+Reuse loaded models with a bounded cache:
+
+```python
+clocks = ["Horvath2013", "AltumAge"]
+cache = pya.pred.ClockCache(maxsize=len(clocks))
+for adata in datasets:  # Each dataset is a separate AnnData object.
+    pya.pred.predict_age(adata, clocks, clock_cache=cache)
+cache.clear()
+```
+
+The cache avoids repeated download checks, deserialization, and device transfers. It holds at most `maxsize` models and keeps separate entries for each device and data revision. It is optional; ordinary calls retain no model cache. Clear it to refresh weights from a moving `main` revision. See the [prediction guide](https://pyaging.readthedocs.io/en/latest/pyaging.predict.html) for memory and cohort guidance.
 
 ## ❓ Can't find an aging clock?
 

{pyaging-0.5.2 → pyaging-0.5.4}/pyproject.toml

@@ -34,8 +34,10 @@ dependencies = [
     "pandas>=2.2.0",
     "torch>=2.2.0",
     "numpy>=2.0.0",
+    "scipy>=1.11",
     "anndata>=0.10.3",
     "huggingface-hub>=1.3,<2",
+    "httpx>=0.23,<1",
     "rich>=13.0",
 ]
 
@@ -100,7 +102,7 @@ only-include = ["src/pyaging", "README.md", "LICENSE", "pyproject.toml"]
 
 [tool.ruff]
 line-length = 120
-extend-exclude = ["*.ipynb", "dev", "docs/superpowers"]
+extend-exclude = ["*.ipynb"]
 
 [tool.ruff.lint]
 select = [

{pyaging-0.5.2 → pyaging-0.5.4}/src/pyaging/__init__.py

@@ -4,4 +4,4 @@ from . import data, logger, models, utils
 from . import predict as pred
 from . import preprocess as pp
 
-__version__ = "0.5.2"
+__version__ = "0.5.4"

{pyaging-0.5.2 → pyaging-0.5.4}/src/pyaging/data/_data.py

@@ -2,6 +2,7 @@ import shutil
 from pathlib import Path
 
 from ..logger._live import SimpleStep, display_enabled, live_step, quiet_hf_bars
+from ..utils._files import atomic_output_path
 from ..utils._hf import download_hf_file
 
 _EXAMPLE_DATA_FILENAMES = {
@@ -63,13 +64,15 @@ def download_example_data(data_type: str, dir: str = "pyaging_data", verbose: bo
     enabled = display_enabled(verbose)
     filename = _EXAMPLE_DATA_FILENAMES[data_type]
     destination = Path(dir) / filename
+    if destination.is_dir():
+        raise IsADirectoryError(destination)
     if destination.exists():
         SimpleStep(filename, enabled=enabled).done(f"example data already at {destination}")
         return str(destination)
 
     with quiet_hf_bars(verbose), live_step(f"downloading {filename}", verbose) as (step, pipeline_logger):
         cache_path = download_hf_file(filename, dir, pipeline_logger, indent_level=1)
-        destination.parent.mkdir(parents=True, exist_ok=True)
-        shutil.copy(cache_path, destination)
+        with atomic_output_path(destination) as temporary_path:
+            shutil.copy(cache_path, temporary_path)
         step.done(f"example data at {destination}")
     return str(destination)

{pyaging-0.5.2 → pyaging-0.5.4}/src/pyaging/data/feature_ranges.json

@@ -2,6 +2,7 @@
   "schema_version": 1,
   "modality_defaults": {
     "DNA methylation": {"unit": "beta value", "low": 0.0, "high": 1.0},
+    "proteomics": {"unit": "NPX", "low": null, "high": null},
     "transcriptomics": {"unit": null, "low": 0.0, "high": null},
     "transcriptomics (relative)": {"unit": null, "low": null, "high": null},
     "histone modification": {"unit": null, "low": 0.0, "high": null},
@@ -80,8 +81,8 @@
     "grimage2timp1": {"unit": "pg/mL", "low": 1000.0, "high": 5000000.0},
     "hdl_cholesterol": {"unit": "mmol/L", "low": 0.1, "high": 5.0},
     "health_compared_to_one_year_ago": {"unit": "NHANES change code (1 = better, 2 = worse, 3 = same, 7 = refused, 9 = don't know)", "low": 1.0, "high": 9.0},
-    "healthcare_use_index": {"unit": "NHANES visit-count code (0 = none, 8 = 16 or more)", "low": 0.0, "high": 8.0},
-    "healthcare_visits_past_year": {"unit": "NHANES visit-count code (0 = none, 8 = 16 or more, 77 = refused, 99 = don't know)", "low": 0.0, "high": 99.0},
+    "healthcare_use_index": {"unit": "NHANES 1999-2000 HUQ050 visit-count code (0 = none, 5 = 13 or more)", "low": 0.0, "high": 5.0},
+    "healthcare_visits_past_year": {"unit": "NHANES 1999-2000 HUQ050 visit-count code (0 = none, 1 = 1, 2 = 2 to 3, 3 = 4 to 9, 4 = 10 to 12, 5 = 13 or more, 77 = refused, 99 = don't know); collapse the 0 to 8 HUQ051 scale of later cycles onto it first", "low": 0.0, "high": 99.0},
     "hematocrit": {"unit": "%", "low": 5.0, "high": 75.0},
     "hemoglobin": {"unit": "g/dL", "low": 2.0, "high": 25.0},
     "hemoglobin_a1c": {"unit": "%", "low": 2.0, "high": 20.0},

{pyaging-0.5.2 → pyaging-0.5.4}/src/pyaging/logger/_live.py

@@ -21,6 +21,7 @@ import contextlib
 import sys
 import threading
 import time
+from contextvars import ContextVar
 
 from rich.console import Console, Group
 
@@ -52,6 +53,7 @@ DOT = "⏺" if sys.platform == "darwin" else "●"
 ELBOW = "⎿"
 
 _console = Console()
+_hf_bars_disabled = ContextVar("pyaging_hf_bars_disabled", default=False)
 
 
 class Sparkle:
@@ -107,20 +109,17 @@ def display_enabled(verbose) -> bool:
 
 @contextlib.contextmanager
 def quiet_hf_bars(verbose=True):
-    """Suppress the Hub's own download bars wherever the display renders its
-    own progress, and always at verbose=False. Non-interactive verbose runs
-    keep the bars: they are the only live progress signal there."""
-    from huggingface_hub.utils import are_progress_bars_disabled, disable_progress_bars, enable_progress_bars
+    """Suppress bars for this run without changing global Hub settings.
 
+    Non-interactive verbose runs keep the bars: they are the only live
+    progress signal there. Nested quiet contexts remain quiet.
+    """
     animated = _console.is_jupyter or _console.is_interactive
-    were_enabled = (not verbose or animated) and not are_progress_bars_disabled()
-    if were_enabled:
-        disable_progress_bars()
+    token = _hf_bars_disabled.set(_hf_bars_disabled.get() or not verbose or animated)
     try:
         yield
     finally:
-        if were_enabled:
-            enable_progress_bars()
+        _hf_bars_disabled.reset(token)
 
 
 def _bar(completed: float, total: float | None, width: int = 24, pulse: bool = False):
@@ -143,7 +142,7 @@ def live_step(label: str, verbose):
     The step renders animated when interactive, prints final lines when
     piped, and is fully silent at verbose=False - callers keep one code path.
     """
-    with SimpleStep(label, enabled=display_enabled(verbose)) as step:
+    with quiet_hf_bars(verbose), SimpleStep(label, enabled=display_enabled(verbose)) as step:
         yield step, DisplayLogger(step.warn)
 
 

{pyaging-0.5.2 → pyaging-0.5.4}/src/pyaging/logger/_logger.py

@@ -4,6 +4,7 @@ import functools
 import logging
 import sys
 import time
+from collections.abc import Sized
 from contextlib import contextmanager
 
 
@@ -45,7 +46,7 @@ def format_logging_message(msg, logging_level, indent_level=1, indent_space_num=
 
 
 class Logger:
-    """Dynamo-specific logger that sets up logging for the package."""
+    """Logger for pyaging's text output and file logging."""
 
     FORMAT = "%(message)s"
 
@@ -69,13 +70,6 @@ class Logger:
         self.logger.propagate = False
         self.log_time()
 
-        # Other global initialization
-        silence_logger("anndata")
-        silence_logger("h5py")
-        silence_logger("numba")
-        silence_logger("pysam")
-        silence_logger("pystan")
-
         if level is not None:
             self.logger.setLevel(level)
         else:
@@ -112,8 +106,10 @@ class Logger:
         """
         previous = self.namespace
         self.namespace = namespace
-        yield
-        self.namespace = previous
+        try:
+            yield
+        finally:
+            self.namespace = previous
 
     def namespace_message(self, message):
         """Add namespace information at the beginning of the logging message.
@@ -296,16 +292,17 @@ class Logger:
         ts :
             total size
         """
-        if self.report_hook_percent_state is None:
-            self.report_hook_percent_state = 0
-
-        if ts == -1:
+        if ts <= 0:
+            self.report_hook_percent_state = None
             return
 
-        cur_percent = rs * bn / ts
+        if bn == 0 or self.report_hook_percent_state is None:
+            self.report_hook_percent_state = 0
+
+        cur_percent = min(rs * bn / ts, 1)
 
         if cur_percent - self.report_hook_percent_state > 0.01:
-            self.report_progress(count=rs * bn, total=ts, indent_level=self.indent_level)
+            self.report_progress(percent=cur_percent * 100, indent_level=self.indent_level)
             self.report_hook_percent_state = cur_percent
         if rs * bn >= ts:
             self.report_hook_percent_state = None
@@ -318,8 +315,8 @@ class LoggerManager:
     CRITICAL = logging.CRITICAL
     EXCEPTION = logging.ERROR
 
-    main_logger = Logger("dynamo")
-    temp_timer_logger = Logger("dynamo-temp-timer-logger")
+    main_logger = Logger("pyaging")
+    temp_timer_logger = Logger("pyaging.temp-timer")
 
     @staticmethod
     def get_main_logger():
@@ -337,23 +334,24 @@ class LoggerManager:
     def progress_logger(generator, logger=None, progress_name="", indent_level=1):
         if logger is None:
             logger = LoggerManager.get_temp_timer_logger()
-        iterator = iter(generator)
         logger.log_time()
-        i = 0
+        total = len(generator) if isinstance(generator, Sized) else None
         prev_progress_percent = 0
-        while i < len(generator):
-            i += 1
-            new_progress_percent = i / len(generator) * 100
+        for i, item in enumerate(generator, start=1):
+            if total is None or total == 0:
+                yield item
+                continue
+            new_progress_percent = i / total * 100
             # report every `interval` percent
             if new_progress_percent - prev_progress_percent > 1 or new_progress_percent >= 100:
                 logger.report_progress(
                     count=i,
-                    total=len(generator),
+                    total=total,
                     progress_name=progress_name,
                     indent_level=indent_level,
                 )
                 prev_progress_percent = new_progress_percent
-            yield next(iterator)
+            yield item
 
 
 def main_info(message, indent_level=1):
@@ -435,4 +433,4 @@ def main_info_verbose_timeit(msg):
 
 
 def main_set_level(level):
-    set_logger_level("dynamo", level)
+    LoggerManager.main_logger.setLevel(level)

{pyaging-0.5.2 → pyaging-0.5.4}/src/pyaging/models/__init__.py

@@ -1,5 +1,6 @@
 # pyaging/models/__init__.py
 
+from ._author_extensions import DNAmPhenoAgeNonPRC, DNAmPhenoAgePRC, IntrinClock370
 from ._base_models import (
     AltumAgeNeuralNetwork,
     LinearModel,
@@ -173,8 +174,36 @@ from ._models import (
     epiTOC2,
     stemTOC,
 )
+from ._organage import OrganAge
+from ._pc_extensions import (
+    PCBrainAge,
+    PCGrimAgeADM,
+    PCGrimAgeB2M,
+    PCGrimAgeCystatinC,
+    PCGrimAgeGDF15,
+    PCGrimAgeLeptin,
+    PCGrimAgePackYrs,
+    PCGrimAgePAI1,
+    PCGrimAgeTIMP1,
+)
+from ._replication import CellDRIFT, MiAge
 
 __all__ = [
+    "CellDRIFT",
+    "DNAmPhenoAgeNonPRC",
+    "DNAmPhenoAgePRC",
+    "IntrinClock370",
+    "MiAge",
+    "OrganAge",
+    "PCBrainAge",
+    "PCGrimAgeADM",
+    "PCGrimAgeB2M",
+    "PCGrimAgeCystatinC",
+    "PCGrimAgeGDF15",
+    "PCGrimAgeLeptin",
+    "PCGrimAgePackYrs",
+    "PCGrimAgePAI1",
+    "PCGrimAgeTIMP1",
     "ABEC",
     "ADBahadoSingh",
     "AltumAge",

pyaging-0.5.4/src/pyaging/models/_author_extensions.py

@@ -0,0 +1,28 @@
+"""Author-defined IntrinClock and PhenoAge variants."""
+
+import torch
+
+from ._base_models import pyagingModel
+from ._models import IntrinClock
+
+
+class IntrinClock370(IntrinClock):
+    """The original cv.glmnet model at lambda.1se, as used by the author demo."""
+
+
+class DNAmPhenoAgePRC(pyagingModel):
+    """Polycomb-related PhenoAge contribution, with no age intercept.
+
+    Matches the author's ``calcPRCPhenoAge(..., imputation = FALSE)``:
+    missing CpGs and supplied NA beta values make zero contribution.
+    """
+
+    def preprocess(self, x):
+        return torch.where(torch.isnan(x), torch.zeros_like(x), x)
+
+    def postprocess(self, x):
+        return x
+
+
+class DNAmPhenoAgeNonPRC(DNAmPhenoAgePRC):
+    """Complementary non-polycomb PhenoAge contribution, without intercept."""

{pyaging-0.5.2 → pyaging-0.5.4}/src/pyaging/models/_models.py

@@ -42,6 +42,24 @@ LINAGE2_COMORBIDITY_ITEMS = (
 )
 
 
+def _global_median(x):
+    """Keep the rank clocks' lower median and all-missing zero fallback."""
+    median = torch.nanmedian(x)
+    return torch.where(torch.isnan(median), torch.zeros_like(median), median)
+
+
+def _average_ranks(values):
+    """Assign average one-based ranks without scalar tensor comparisons."""
+    sorted_values, order = torch.sort(values)
+    _, counts = torch.unique_consecutive(sorted_values, return_counts=True)
+    ends = counts.cumsum(0)
+    starts = ends - counts
+    averages = (starts + ends + 1).to(values.dtype) / 2
+    ranks = torch.empty_like(values)
+    ranks[order] = torch.repeat_interleave(averages, counts, output_size=values.numel())
+    return ranks
+
+
 class AltumAge(pyagingModel):
     def __init__(self):
         super().__init__()
@@ -1138,6 +1156,7 @@ class PCGrimAge(pyagingModel):
         self.features_PCTIMP1 = None
 
     def forward(self, x):
+        x = self.preprocess(x)
         CpGs = x[:, :-2]
         Female = x[:, -2].unsqueeze(1)
         Age = x[:, -1].unsqueeze(1)
@@ -1278,28 +1297,21 @@ class Pasta(pyagingModel):
         """
         Assign average ranks (1-based) per vector, handling ties.
         """
-        sorted_vals, sorted_idx = torch.sort(values)
-        ranks = torch.empty_like(sorted_vals, dtype=values.dtype)
+        return _average_ranks(values)
 
-        n = values.numel()
-        start = 0
-        while start < n:
-            end = start + 1
-            while end < n and sorted_vals[end] == sorted_vals[start]:
-                end += 1
-            avg_rank = (start + end - 1) / 2.0 + 1.0
-            ranks[sorted_idx[start:end]] = avg_rank
-            start = end
+    def prepare_cohort_context(self, x):
+        """Compute the missing-value fill once across the complete cohort."""
+        return _global_median(x)
 
-        return ranks
+    def predict_with_cohort_context(self, x, context):
+        """Score a batch with its cohort's median without storing run state."""
+        return self.postprocess(self.base_model(self.preprocess(x, context)))
 
-    def preprocess(self, x):
+    def preprocess(self, x, median=None):
         """
         Fill missing values with the global median then rank-normalize per sample.
         """
-        median = torch.nanmedian(x)
-        if torch.isnan(median):
-            median = torch.tensor(0.0, device=x.device, dtype=x.dtype)
+        median = _global_median(x) if median is None else median.to(device=x.device, dtype=x.dtype)
         x = torch.where(torch.isnan(x), median, x)
 
         ranked = torch.empty_like(x, dtype=x.dtype)
@@ -1349,21 +1361,33 @@ class PastaMouse(Pasta):
         else:
             self.reference_values = [self.full_reference_values[i] for i in self.mouse_feature_indices]
 
-    def _expand_with_reference(self, x):
-        """
-        Reconstruct the full 8113-length input expected by the base model by
-        inserting reference values for human-only genes.
-        """
+    def _full_reference_tensor(self, x):
         if self.base_model_features is None or self.mouse_feature_indices is None:
             raise ValueError("PastaMouse must be configured with set_mouse_features before inference.")
 
         if self.full_reference_values is None:
-            ref_full = torch.zeros(len(self.base_model_features), device=x.device, dtype=x.dtype)
-        elif isinstance(self.full_reference_values, torch.Tensor):
-            ref_full = self.full_reference_values.to(device=x.device, dtype=x.dtype)
-        else:
-            ref_full = torch.tensor(self.full_reference_values, device=x.device, dtype=x.dtype)
+            return torch.zeros(len(self.base_model_features), device=x.device, dtype=x.dtype)
+        return torch.as_tensor(self.full_reference_values, device=x.device, dtype=x.dtype)
+
+    def prepare_cohort_context(self, x):
+        """Include human reference values without expanding the full cohort."""
+        reference = self._full_reference_tensor(x)
+        inserted = torch.ones(reference.shape, dtype=torch.bool, device=x.device)
+        inserted[self.mouse_feature_indices] = False
+        # Missing references do not affect nanmedian, which is the common case.
+        reference = reference[inserted & ~torch.isnan(reference)]
+        values = torch.cat([x.reshape(-1), reference.repeat(x.size(0))]) if reference.numel() else x
+        return _global_median(values)
+
+    def predict_with_cohort_context(self, x, context):
+        return super().predict_with_cohort_context(self._expand_with_reference(x), context)
 
+    def _expand_with_reference(self, x):
+        """
+        Reconstruct the full 8113-length input expected by the base model by
+        inserting reference values for human-only genes.
+        """
+        ref_full = self._full_reference_tensor(x)
         full_x = ref_full.unsqueeze(0).repeat(x.size(0), 1)
         full_x[:, self.mouse_feature_indices] = x
         return full_x
@@ -1386,28 +1410,21 @@ class Reg(pyagingModel):
         """
         Assign average ranks (1-based) per vector, handling ties.
         """
-        sorted_vals, sorted_idx = torch.sort(values)
-        ranks = torch.empty_like(sorted_vals, dtype=values.dtype)
+        return _average_ranks(values)
 
-        n = values.numel()
-        start = 0
-        while start < n:
-            end = start + 1
-            while end < n and sorted_vals[end] == sorted_vals[start]:
-                end += 1
-            avg_rank = (start + end - 1) / 2.0 + 1.0
-            ranks[sorted_idx[start:end]] = avg_rank
-            start = end
+    def prepare_cohort_context(self, x):
+        """Compute the missing-value fill once across the complete cohort."""
+        return _global_median(x)
 
-        return ranks
+    def predict_with_cohort_context(self, x, context):
+        """Score a batch with its cohort's median without storing run state."""
+        return self.postprocess(self.base_model(self.preprocess(x, context)))
 
-    def preprocess(self, x):
+    def preprocess(self, x, median=None):
         """
         Fill missing values with the global median then rank-normalize per sample.
         """
-        median = torch.nanmedian(x)
-        if torch.isnan(median):
-            median = torch.tensor(0.0, device=x.device, dtype=x.dtype)
+        median = _global_median(x) if median is None else median.to(device=x.device, dtype=x.dtype)
         x = torch.where(torch.isnan(x), median, x)
 
         ranked = torch.empty_like(x, dtype=x.dtype)
@@ -1708,11 +1725,11 @@ class PhenoAge(pyagingModel):
         Applies a convertion from a CDF of the mortality score from a Gompertz
         distribution to phenotypic age.
         """
-        # lambda
-        lambda_ = torch.tensor(0.0192, device=x.device, dtype=x.dtype)
-        mortality_score = 1 - torch.exp(-torch.exp(x) * (torch.exp(120 * lambda_) - 1) / lambda_)
-        age = 141.50225 + torch.log(-0.00553 * torch.log(1 - mortality_score)) / 0.090165
-        return age
+        # log(1 - mortality) equals the negative Gompertz hazard exactly.
+        # Cancel the nested exp/log terms before finite precision rounds the
+        # mortality to zero or one, preserving the published constants.
+        log_scale = math.log(0.00553 * math.expm1(120 * 0.0192) / 0.0192)
+        return 141.50225 + (x + log_scale) / 0.090165
 
 
 class PhenoAgeSaoPaulo(pyagingModel):
@@ -2273,7 +2290,7 @@ class stemTOC(pyagingModel):
         quantiles = []
         for row in x:
             filtered_row = row[row != -1]
-            quantile_95 = torch.quantile(filtered_row, 0.95) if len(filtered_row) > 0 else torch.tensor(float("nan"))
+            quantile_95 = torch.quantile(filtered_row, 0.95) if len(filtered_row) > 0 else x.new_tensor(float("nan"))
             quantiles.append(quantile_95)
         return torch.vstack(quantiles)
 
@@ -2290,7 +2307,7 @@ class epiTOC1(pyagingModel):
         means = []
         for row in x:
             filtered_row = row[row != -1]
-            mean = torch.mean(filtered_row) if len(filtered_row) > 0 else torch.tensor(float("nan"))
+            mean = torch.mean(filtered_row) if len(filtered_row) > 0 else x.new_tensor(float("nan"))
             means.append(mean)
         return torch.vstack(means)
 

pyaging-0.5.4/src/pyaging/models/_organage.py

@@ -0,0 +1,33 @@
+"""Goeminne et al. plasma-protein OrganAge models.
+
+Original author code and coefficients:
+https://github.com/ludgergoeminne/organAging
+DOI: 10.1016/j.cmet.2024.10.005
+
+The released coefficients are subject to the authors' Academic / Non-Commercial
+License. Each converted artifact carries that license and its pinned provenance.
+"""
+
+from ._base_models import pyagingModel
+
+
+class OrganAge(pyagingModel):
+    """Score Olink NPX values with a released OrganAge coefficient vector.
+
+    Protein names are the authors' case-sensitive symbols, including hyphens.
+    NPX values are already normalized; no cohort scaling is applied here. The
+    authors recommend fold one for new data. Chronological models include an
+    intercept and return years; mortality models have no intercept and return
+    relative natural-log mortality hazards.
+
+    The authors' example omits absent protein columns and propagates NaN values
+    in present columns. Zero-filled absent features preserve that omission in
+    pyaging, while present NaNs remain NaNs. Any imputation performed before
+    scoring is the caller's separate data-preparation decision.
+    """
+
+    def preprocess(self, x):
+        return x
+
+    def postprocess(self, x):
+        return x