cuperiod 1.0.0__py3-none-any.whl

cuperiod/__init__.py

@@ -0,0 +1,85 @@
+"""cuPeriod: optimized, GPU-accelerated periodograms for astronomy.
+
+Quick start
+-----------
+>>> import cuperiod as cup                                    # doctest: +SKIP
+>>> pg = cup.periodogram((time, mag, mag_err), "GLS")         # doctest: +SKIP
+>>> pg.best_period()                                          # doctest: +SKIP
+>>> for peak in pg.best_periods(10):                          # doctest: +SKIP
+...     print(peak.period, peak.power)
+
+The recommended import alias is ``cup``. The two entry points are
+:func:`periodogram` (a single light curve) and :func:`batch_periodograms` (many).
+"""
+
+from __future__ import annotations
+
+# Importing the methods package registers the built-in methods (GLS, BLS, ...).
+from cuperiod import methods as methods
+from cuperiod.api import best_periods, periodogram, to_input
+from cuperiod.batch import BatchSummary, batch_periodograms
+from cuperiod.core.columns import ColumnMap, Domain
+from cuperiod.core.config import (
+    BatchSettings,
+    BLSSettings,
+    CESettings,
+    GLSSettings,
+    MHAOVSettings,
+    PDMSettings,
+    StringLengthSettings,
+    TLSSettings,
+)
+from cuperiod.core.device import GpuInfo, free_gpu_memory, gpu_info, suggest_gpu_workers
+from cuperiod.core.errors import (
+    BackendUnavailableError,
+    ColumnResolutionError,
+    CuPeriodError,
+    InsufficientDataError,
+    UnknownMethodError,
+)
+from cuperiod.core.grid import GridSpec, log_period_grid, uniform_frequency_grid
+from cuperiod.core.lightcurve import LightCurve, MultiBandLightCurve
+from cuperiod.core.result import MultiResult, Peak, Periodogram
+from cuperiod.methods.base import MethodInfo, get_method, list_methods, method_names
+
+__version__ = "1.0.0"
+
+__all__ = [
+    "BLSSettings",
+    "BackendUnavailableError",
+    "BatchSettings",
+    "BatchSummary",
+    "CESettings",
+    "ColumnMap",
+    "ColumnResolutionError",
+    "CuPeriodError",
+    "Domain",
+    "GLSSettings",
+    "GpuInfo",
+    "GridSpec",
+    "InsufficientDataError",
+    "LightCurve",
+    "MHAOVSettings",
+    "MethodInfo",
+    "MultiBandLightCurve",
+    "MultiResult",
+    "PDMSettings",
+    "Peak",
+    "Periodogram",
+    "StringLengthSettings",
+    "TLSSettings",
+    "UnknownMethodError",
+    "__version__",
+    "batch_periodograms",
+    "best_periods",
+    "free_gpu_memory",
+    "get_method",
+    "gpu_info",
+    "list_methods",
+    "log_period_grid",
+    "method_names",
+    "periodogram",
+    "suggest_gpu_workers",
+    "to_input",
+    "uniform_frequency_grid",
+]

cuperiod/api.py

@@ -0,0 +1,187 @@
+"""The single-light-curve API: :func:`periodogram`.
+
+This is the one entry point most users need. It accepts a light curve in any convenient
+form (a :class:`LightCurve`/:class:`MultiBandLightCurve`, a ``(t, y[, dy])`` tuple, a
+pandas/astropy/pyarrow table, or a dict of arrays), one or several method names, and an
+optional settings/grid/backend, and returns a :class:`Periodogram` (single method) or a
+:class:`MultiResult` (several methods).
+"""
+
+from __future__ import annotations
+
+from collections.abc import Mapping, Sequence
+from typing import Any
+
+from pydantic_settings import BaseSettings
+
+from cuperiod.core.columns import ColumnMap, Domain
+from cuperiod.core.grid import GridSpec
+from cuperiod.core.lightcurve import LightCurve, MultiBandLightCurve
+from cuperiod.core.result import MultiResult, Periodogram
+from cuperiod.methods.base import PeriodogramMethod, get_method
+
+#: A single method's settings, a ``{method: settings}`` mapping, or ``None`` (defaults).
+SettingsInput = BaseSettings | Mapping[str, BaseSettings] | None
+
+
+def _looks_like_arrays(obj: Sequence[Any]) -> bool:
+    """Whether ``obj`` is a ``(t, y)`` / ``(t, y, dy)`` tuple of array-likes."""
+    if not isinstance(obj, (tuple, list)) or not (2 <= len(obj) <= 3):
+        return False
+    first = obj[0]
+    return hasattr(first, "__len__") and not isinstance(first, (str, bytes))
+
+
+def to_input(
+    data: Any,
+    *,
+    columns: ColumnMap | None = None,
+    domain: Domain | None = None,
+) -> LightCurve | MultiBandLightCurve:
+    """Coerce a user input into a :class:`LightCurve` or :class:`MultiBandLightCurve`.
+
+    Parameters
+    ----------
+    data : various
+        A :class:`LightCurve`/:class:`MultiBandLightCurve`, a ``(t, y[, dy])`` tuple, a
+        ``{band: LightCurve}`` mapping, a pandas/astropy/pyarrow table, or a
+        ``{column: array}`` mapping.
+    columns : ColumnMap, optional
+        Column overrides for table inputs.
+    domain : Domain, optional
+        Brightness-domain override.
+
+    Returns
+    -------
+    LightCurve or MultiBandLightCurve
+    """
+    if isinstance(data, (LightCurve, MultiBandLightCurve)):
+        if domain is not None and isinstance(data, LightCurve):
+            return data.in_domain(domain)
+        return data
+    if _looks_like_arrays(data):
+        t, v = data[0], data[1]
+        e = data[2] if len(data) == 3 else None
+        return LightCurve.from_arrays(t, v, e, domain=domain or Domain.MAGNITUDE)
+    if isinstance(data, Mapping):
+        values = list(data.values())
+        if values and all(isinstance(v, LightCurve) for v in values):
+            return MultiBandLightCurve.from_light_curves(data)
+        # dict of arrays -> a table
+        names, getter = _mapping_table(data)
+        return LightCurve._from_table(names, getter, columns, domain, None)
+    # pandas / astropy / pyarrow table
+    return LightCurve.from_dataframe(data, columns=columns, domain=domain)
+
+
+def _mapping_table(data: Mapping[str, Any]) -> tuple[list[str], Any]:
+    from cuperiod.core.lightcurve import _adapt_table
+
+    return _adapt_table(data)
+
+
+def _select_settings(
+    method: PeriodogramMethod, settings: SettingsInput
+) -> BaseSettings:
+    """Pick the settings object that applies to ``method``."""
+    if settings is None:
+        return method.coerce_settings(None)
+    if isinstance(settings, BaseSettings):
+        if isinstance(settings, method.settings_cls):
+            return settings
+        return method.coerce_settings(None)
+    # mapping of {method_name: settings}
+    for key, value in settings.items():
+        if key.upper() == method.name:
+            return method.coerce_settings(value)
+    return method.coerce_settings(None)
+
+
+def _multiband_grid(
+    method: PeriodogramMethod, mblc: MultiBandLightCurve, settings: BaseSettings
+) -> GridSpec:
+    """Build a method's default grid from a multi-band stacked baseline."""
+    time, value, error, _ = mblc.finite().stacked()
+    synthetic = LightCurve.from_arrays(time, value, error)
+    return method.default_grid(synthetic, settings)
+
+
+def _run_one(
+    method: PeriodogramMethod,
+    lc: LightCurve | MultiBandLightCurve,
+    settings: BaseSettings,
+    backend: str,
+    grid: GridSpec | None,
+) -> Periodogram:
+    resolved_backend = method.resolve_backend(backend)
+    if isinstance(lc, MultiBandLightCurve):
+        if not method.supports_multiband:
+            raise ValueError(f"{method.name} does not support multi-band input")
+        g = grid or _multiband_grid(method, lc, settings)
+        return method.multiband_power(g, lc, settings, resolved_backend)
+    single = lc.in_domain(method.natural_domain) if method.natural_domain else lc
+    g = grid or method.default_grid(single, settings)
+    return method.power(g, single, settings, resolved_backend, engine=None)
+
+
+def periodogram(
+    data: Any,
+    method: str | Sequence[str] = "GLS",
+    *,
+    backend: str = "auto",
+    grid: GridSpec | None = None,
+    settings: SettingsInput = None,
+    columns: ColumnMap | None = None,
+    domain: Domain | None = None,
+) -> Periodogram | MultiResult:
+    """Compute one or more periodograms for a single light curve.
+
+    Parameters
+    ----------
+    data : various
+        The light curve (see :func:`to_input` for accepted forms).
+    method : str or sequence of str, default "GLS"
+        One method name, or several. Case-insensitive (``"gls"`` == ``"GLS"``).
+    backend : str, default "auto"
+        ``"auto"`` (GPU when available, else CPU), ``"cpu"``, ``"gpu"``, or a concrete
+        backend name.
+    grid : GridSpec, optional
+        Custom trial grid; defaults to the method's grid for this light curve.
+    settings : settings model or mapping, optional
+        A single method's settings, or a ``{method: settings}`` mapping.
+    columns : ColumnMap, optional
+        Column overrides for table inputs.
+    domain : Domain, optional
+        Brightness-domain override.
+
+    Returns
+    -------
+    Periodogram or MultiResult
+        A :class:`Periodogram` if ``method`` is a single string, else a
+        :class:`MultiResult` keyed by method name.
+
+    Examples
+    --------
+    >>> pg = periodogram((t, mag, err), "GLS")          # doctest: +SKIP
+    >>> pg.best_period()                                # doctest: +SKIP
+    >>> res = periodogram(df, ["GLS", "BLS"])           # doctest: +SKIP
+    """
+    lc = to_input(data, columns=columns, domain=domain)
+    names = [method] if isinstance(method, str) else list(method)
+    results: dict[str, Periodogram] = {}
+    for name in names:
+        m = get_method(name)
+        results[m.name] = _run_one(m, lc, _select_settings(m, settings), backend, grid)
+    if isinstance(method, str):
+        return results[get_method(method).name]
+    return MultiResult(results)
+
+
+def best_periods(
+    result: Periodogram | MultiResult, n: int = 10, **kwargs: Any
+) -> Any:
+    """Convenience wrapper around ``result.best_periods(n)``."""
+    return result.best_periods(n, **kwargs)
+
+
+__all__ = ["best_periods", "periodogram", "to_input"]

cuperiod/batch/__init__.py

@@ -0,0 +1,7 @@
+"""Batch processing: scale periodograms over many light curves."""
+
+from __future__ import annotations
+
+from cuperiod.batch.runner import BatchSummary, batch_periodograms
+
+__all__ = ["BatchSummary", "batch_periodograms"]

cuperiod/batch/io.py

@@ -0,0 +1,232 @@
+"""Batch input resolution and result sinks.
+
+Inputs can be an iterable of light curves, a glob pattern, a directory of light-curve
+files, or a ``(DataFrame, group_column)`` pair; :func:`resolve_inputs` normalizes any of
+them to ``(key, source)`` items, where ``source`` is loaded lazily in the worker so file
+paths (not big arrays) cross the process boundary. Results are flattened to one row per
+light curve and written as Parquet (default) or CSV.
+"""
+
+from __future__ import annotations
+
+import glob as _glob
+from collections.abc import Iterable, Mapping, Sequence
+from pathlib import Path
+from typing import Any
+
+from cuperiod.core.columns import ColumnMap, Domain
+from cuperiod.core.lightcurve import LightCurve, MultiBandLightCurve
+from cuperiod.core.result import Periodogram
+
+#: File extensions treated as light-curve tables when scanning a directory.
+LIGHTCURVE_SUFFIXES = (
+    ".parquet",
+    ".pq",
+    ".csv",
+    ".ecsv",
+    ".fits",
+    ".fit",
+    ".dat",
+    ".txt",
+    ".tsv",
+)
+
+Source = LightCurve | MultiBandLightCurve | Path
+InputItem = tuple[str, Source]
+
+
+def _is_dataframe(obj: Any) -> bool:
+    return hasattr(obj, "groupby") and hasattr(obj, "columns")
+
+
+def resolve_inputs(
+    inputs: Any,
+    *,
+    columns: ColumnMap | None = None,
+    domain: Domain | None = None,
+    band_column: str | None = None,
+) -> list[InputItem]:
+    """Normalize a batch input into ``(key, source)`` items.
+
+    Parameters
+    ----------
+    inputs : various
+        Iterable of light curves / ``(key, lc)`` pairs / paths, a glob string, a
+        directory path, or a ``(DataFrame, group_column)`` tuple.
+    columns, domain : optional
+        Forwarded to file/table loading.
+    band_column : str, optional
+        If given with a ``(DataFrame, group_column)`` input, each group is loaded as a
+        :class:`MultiBandLightCurve` split on this column.
+
+    Returns
+    -------
+    list of (str, Source)
+        Stable keys paired with a light curve or a path to load lazily.
+    """
+    # (DataFrame, group_column)
+    if isinstance(inputs, tuple) and len(inputs) == 2 and _is_dataframe(inputs[0]):
+        df, group_col = inputs
+        items: list[InputItem] = []
+        for key, sub in df.groupby(group_col, sort=False):
+            if band_column is not None:
+                lc: Source = MultiBandLightCurve.from_dataframe(
+                    sub, band_column=band_column, columns=columns, domain=domain
+                )
+            else:
+                lc = LightCurve.from_dataframe(sub, columns=columns, domain=domain)
+            items.append((str(key), lc))
+        return items
+
+    # Glob / directory / single file given as a string or Path.
+    if isinstance(inputs, (str, Path)):
+        return _resolve_path_like(inputs)
+
+    # An iterable of light curves / (key, lc) pairs / paths.
+    if isinstance(inputs, Iterable):
+        return _resolve_iterable(inputs, columns=columns, domain=domain)
+
+    raise TypeError(f"unsupported batch input type: {type(inputs)!r}")
+
+
+def _resolve_path_like(inputs: str | Path) -> list[InputItem]:
+    text = str(inputs)
+    if any(ch in text for ch in "*?[") or "**" in text:
+        paths = sorted(Path(p) for p in _glob.glob(text, recursive=True))
+        return [(p.stem, p) for p in paths]
+    path = Path(inputs)
+    if path.is_dir():
+        paths = sorted(
+            p for p in path.iterdir() if p.suffix.lower() in LIGHTCURVE_SUFFIXES
+        )
+        return [(p.stem, p) for p in paths]
+    if path.is_file():
+        return [(path.stem, path)]
+    raise FileNotFoundError(f"no such file, directory, or glob match: {inputs!r}")
+
+
+def _resolve_iterable(
+    inputs: Iterable[Any], *, columns: ColumnMap | None, domain: Domain | None
+) -> list[InputItem]:
+    items: list[InputItem] = []
+    for i, entry in enumerate(inputs):
+        if (
+            isinstance(entry, tuple)
+            and len(entry) == 2
+            and isinstance(entry[0], (str, int))
+            and isinstance(entry[1], (LightCurve, MultiBandLightCurve, str, Path))
+        ):
+            key, src = entry
+            src = Path(src) if isinstance(src, str) else src
+            items.append((str(key), src))
+        elif isinstance(entry, (LightCurve, MultiBandLightCurve)):
+            key = str(entry.meta.get("id", i)) if entry.meta else str(i)
+            items.append((key, entry))
+        elif isinstance(entry, (str, Path)):
+            items.append((Path(entry).stem, Path(entry)))
+        else:
+            raise TypeError(f"unsupported batch entry at index {i}: {type(entry)!r}")
+    return items
+
+
+def load_source(
+    source: Source,
+    *,
+    columns: ColumnMap | None = None,
+    domain: Domain | None = None,
+) -> LightCurve | MultiBandLightCurve:
+    """Materialize a source: return a light curve unchanged, or load it from a path."""
+    if isinstance(source, (LightCurve, MultiBandLightCurve)):
+        return source
+    return LightCurve.from_file(source, columns=columns, domain=domain)
+
+
+def periodogram_to_row(
+    key: str, pg: Periodogram, *, n_best: int, store_raw: bool
+) -> dict[str, Any]:
+    """Flatten a :class:`Periodogram` into one result row (peaks + optional raw)."""
+    row: dict[str, Any] = {
+        "key": key,
+        "method": pg.method,
+        "backend": pg.backend,
+        "n_samples": pg.n_samples,
+        "baseline": pg.baseline,
+        "best_period": pg.best_period() if pg.size else float("nan"),
+    }
+    peaks = pg.best_periods(n_best, alias_diverse=(pg.method == "BLS"))
+    for rank in range(1, n_best + 1):
+        peak = peaks[rank - 1] if rank <= len(peaks) else None
+        row[f"period_{rank}"] = peak.period if peak else float("nan")
+        row[f"power_{rank}"] = peak.power if peak else float("nan")
+        if peak is not None:
+            for ekey, eval_ in peak.extra.items():
+                row[f"{ekey}_{rank}"] = eval_
+    if store_raw:
+        freq, power = pg.downsample()
+        row["pgram_frequency"] = freq.tolist()
+        row["pgram_power"] = power.tolist()
+    return row
+
+
+def write_rows(rows: Sequence[Mapping[str, Any]], path: Path) -> None:
+    """Write result rows to ``path`` as Parquet (``.parquet``) or CSV (``.csv``)."""
+    path = Path(path)
+    path.parent.mkdir(parents=True, exist_ok=True)
+    if path.suffix.lower() in {".parquet", ".pq"}:
+        _write_parquet(rows, path)
+    elif path.suffix.lower() == ".csv":
+        _write_csv(rows, path)
+    else:
+        raise ValueError(
+            f"unsupported sink format {path.suffix!r}; use .parquet or .csv"
+        )
+
+
+def _union_columns(rows: Sequence[Mapping[str, Any]]) -> list[str]:
+    seen: dict[str, None] = {}
+    for row in rows:
+        for key in row:
+            seen.setdefault(key, None)
+    return list(seen)
+
+
+def _write_parquet(rows: Sequence[Mapping[str, Any]], path: Path) -> None:
+    import pyarrow as pa
+    import pyarrow.parquet as pq
+
+    columns = _union_columns(rows)
+    data = {col: [row.get(col) for row in rows] for col in columns}
+    table = pa.table(data)
+    tmp = path.with_suffix(path.suffix + ".tmp")
+    pq.write_table(table, tmp)
+    tmp.replace(path)
+
+
+def _write_csv(rows: Sequence[Mapping[str, Any]], path: Path) -> None:
+    import csv
+
+    columns = _union_columns(rows)
+    list_cols = [
+        col for col in columns if any(isinstance(row.get(col), list) for row in rows)
+    ]
+    if list_cols:
+        raise ValueError(
+            f"CSV sink cannot store array-valued columns {list_cols} (e.g. from "
+            "store_raw); use a .parquet sink."
+        )
+    tmp = path.with_suffix(path.suffix + ".tmp")
+    with tmp.open("w", newline="", encoding="utf-8") as fh:
+        writer = csv.DictWriter(fh, fieldnames=columns, extrasaction="ignore")
+        writer.writeheader()
+        writer.writerows(rows)
+    tmp.replace(path)
+
+
+__all__ = [
+    "InputItem",
+    "Source",
+    "load_source",
+    "periodogram_to_row",
+    "resolve_inputs",
+    "write_rows",
+]