pyaging 0.5.2__tar.gz → 0.5.3__tar.gz

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.5
 Name: pyaging
-Version: 0.5.2
+Version: 0.5.3
 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.2`, to pin downloads; it defaults to `main`. Package-only releases do not create new HF tags. Version 0.5.3 uses the existing model artifacts.
+
+### 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.3}/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.2`, to pin downloads; it defaults to `main`. Package-only releases do not create new HF tags. Version 0.5.3 uses the existing model artifacts.
+
+### 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.3}/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",
 ]
 

{pyaging-0.5.2 → pyaging-0.5.3}/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.3"

{pyaging-0.5.2 → pyaging-0.5.3}/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.3}/src/pyaging/data/feature_ranges.json

@@ -80,8 +80,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.3}/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.3}/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 that sets up pyaging's legacy text progress output."""
 
     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.3}/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.2 → pyaging-0.5.3}/src/pyaging/predict/__init__.py

@@ -1,5 +1,6 @@
 # pyaging/predict/__init__.py
 
+from ._cache import ClockCache
 from ._pred import predict_age
 from ._pred_utils import (
     add_pred_ages_and_clock_metadata_adata,
@@ -21,6 +22,7 @@ from ._transforms import (
 )
 
 __all__ = [
+    "ClockCache",
     "add_pred_ages_and_clock_metadata_adata",
     "binarize",
     "check_feature_ranges",

pyaging-0.5.3/src/pyaging/predict/_cache.py

@@ -0,0 +1,56 @@
+"""Explicit, bounded reuse of prepared models across prediction calls."""
+
+from collections import OrderedDict
+from numbers import Integral
+
+import torch
+
+from ..utils._hf import get_data_revision
+
+
+class ClockCache:
+    """Keep a bounded number of prepared clocks for repeated datasets.
+
+    Pass the same instance to ``predict_age(..., clock_cache=cache)`` to reuse
+    loaded models, including their device transfer and evaluation setup. Keys
+    include the lowercase clock name, device, and ``PYAGING_DATA_REVISION``.
+    Input matrices and cohort preprocessing are never cached.
+
+    ``maxsize`` is a positive integer, defaulting to 2. A cache miss evicts the
+    least recently used clock *before* loading when the cache is full, so a
+    failed replacement load may leave one fewer cached model. Models remain
+    on their selected device until eviction, ``clear()``, or cache disposal.
+
+    The cache represents a snapshot: call ``clear()`` to pick up changed files
+    at a mutable Hugging Face revision such as ``main``. No global cache is
+    created. An instance is intended for sequential calls, not shared threads.
+    """
+
+    def __init__(self, maxsize: int = 2):
+        if isinstance(maxsize, bool) or not isinstance(maxsize, Integral) or maxsize <= 0:
+            raise ValueError("maxsize must be a positive integer.")
+        self._maxsize = int(maxsize)
+        self._models = OrderedDict()
+
+    @property
+    def maxsize(self) -> int:
+        """Maximum number of prepared models held by this cache."""
+        return self._maxsize
+
+    def clear(self) -> None:
+        """Release all model references held by this cache."""
+        self._models.clear()
+
+    def _get_or_load(self, clock_name, device, load):
+        device = torch.device(device)
+        if device.type == "cuda" and device.index is None:
+            device = torch.device("cuda", torch.cuda.current_device())
+        key = (clock_name.lower(), str(device), get_data_revision())
+        if key in self._models:
+            self._models.move_to_end(key)
+            return self._models[key]
+        if len(self._models) >= self.maxsize:
+            self._models.popitem(last=False)
+        model = load()
+        self._models[key] = model
+        return model

{pyaging-0.5.2 → pyaging-0.5.3}/src/pyaging/predict/_pred.py

@@ -1,10 +1,13 @@
 import gc
+from functools import partial
 
 import anndata
 import torch
 
 from ..logger._live import ClockRunDisplay, DisplayLogger, display_enabled, quiet_hf_bars
+from ._cache import ClockCache
 from ._pred_utils import (
+    _validate_batch_size,
     add_pred_ages_and_clock_metadata_adata,
     apply_cohort_transform,
     build_cohort_feature_matrix,
@@ -23,6 +26,8 @@ def predict_age(
     batch_size: int = 1024,
     clean: bool = True,
     verbose: bool = True,
+    *,
+    clock_cache: ClockCache | None = None,
 ) -> None:
     """
     Predicts biological age using specified aging clocks.
@@ -56,6 +61,11 @@ def predict_age(
         notebooks and terminals, a plain summary when output is captured,
         and fully silent when False. Defaults to True.
 
+    clock_cache: ClockCache or None
+        Optional bounded cache of prepared models to reuse across datasets.
+        Keep the same cache instance for repeated calls. Clear it to refresh
+        weights at mutable revisions such as ``main``. Defaults to no reuse.
+
     Returns
     -------
     None
@@ -87,6 +97,9 @@ def predict_age(
       the samples to centre against. Without it the cohort centres on every sample. Predictions are
       differences against that reference, not absolute values.
 
+    Pasta, Reg, and PastaMouse fill missing values with a median computed over the whole cohort
+    before ranking each sample. Their inference batch size does not change that reference.
+
     The function automatically handles the transfer of data and models to the appropriate compute
     device (CPU or GPU) based on system configuration.
 
@@ -97,6 +110,10 @@ def predict_age(
     >>> adata.obs["horvath2013"]  # Access predicted ages by clock name
 
     """
+    _validate_batch_size(batch_size)
+    if adata.n_obs == 0:
+        raise ValueError("Prediction requires at least one sample.")
+
     # Ensure clock_names is a list with lowercase names
     if isinstance(clock_names, str):
         clock_names = [clock_names]
@@ -118,7 +135,12 @@ def predict_age(
             pipeline_logger = DisplayLogger(lambda m, name=clock_name: display.warn(name, m))
 
             # Load and prepare the clock
-            model = load_clock(clock_name, device, dir, pipeline_logger)
+            if clock_cache is None:
+                model = load_clock(clock_name, device, dir, pipeline_logger)
+            else:
+                model = clock_cache._get_or_load(
+                    clock_name, device, partial(load_clock, clock_name, device, dir, pipeline_logger)
+                )
 
             # Clocks saved before either attribute existed lack both.
             transform_name = getattr(model, "cohort_transform", None)
@@ -171,9 +193,12 @@ def predict_age(
             if clean:
                 del adata.obsm[f"X_{clock_name}"]
 
-            # Flush memory
-            gc.collect()
-            torch.cuda.empty_cache()
+            # Release references before collecting memory and loading the next
+            # clock, so two large models need not coexist on the GPU.
+            del model, predicted_ages_tensor
+            if clock_cache is None:
+                gc.collect()
+                torch.cuda.empty_cache()
 
             display.finish_clock(clock_name)
         display.finish(n_samples=adata.n_obs)