timewise 0.5.4__py3-none-any.whl → 1.0.0a1__py3-none-any.whl

timewise/__init__.py

@@ -1,5 +1 @@
-from timewise.wise_data_by_visit import WiseDataByVisit
-from timewise.wise_bigdata_desy_cluster import WISEDataDESYCluster
-from timewise.parent_sample_base import ParentSampleBase
-
-__version__ = "0.5.4"
+__version__ = "1.0.0a1"

timewise/backend/__init__.py

@@ -0,0 +1,6 @@
+from typing import Union
+
+from .base import Backend
+from .filesystem import FileSystemBackend
+
+BackendType = Union[FileSystemBackend]

timewise/backend/base.py

@@ -0,0 +1,36 @@
+import abc
+from typing import Any
+from pydantic import BaseModel
+from astropy.table import Table
+from ..types import TaskID
+
+
+class Backend(abc.ABC, BaseModel):
+    type: str
+    base_path: Any
+    """
+    Abstract persistence backend for jobs, results, and markers.
+    Works with generic TaskIDs so it can be reused across Downloader/Processor.
+    """
+
+    # --- metadata ---
+    @abc.abstractmethod
+    def meta_exists(self, task: TaskID) -> bool: ...
+    @abc.abstractmethod
+    def save_meta(self, task: TaskID, meta: dict[str, Any]) -> None: ...
+    @abc.abstractmethod
+    def load_meta(self, task: TaskID) -> dict[str, Any] | None: ...
+
+    # --- Markers ---
+    @abc.abstractmethod
+    def mark_done(self, task: TaskID) -> None: ...
+    @abc.abstractmethod
+    def is_done(self, task: TaskID) -> bool: ...
+
+    # --- Data ---
+    @abc.abstractmethod
+    def save_data(self, task: TaskID, content: Table) -> None: ...
+    @abc.abstractmethod
+    def load_data(self, task: TaskID) -> Table: ...
+    @abc.abstractmethod
+    def data_exists(self, task: TaskID) -> bool: ...

timewise/backend/filesystem.py

@@ -0,0 +1,80 @@
+import json
+import logging
+from pathlib import Path
+from typing import Any, Literal
+from astropy.table import Table
+
+from .base import Backend
+from ..types import TaskID
+
+
+logger = logging.getLogger(__name__)
+
+
+class FileSystemBackend(Backend):
+    type: Literal["filesystem"] = "filesystem"
+    base_path: Path
+
+    # ----------------------------
+    # Helpers for paths
+    # ----------------------------
+    def _meta_path(self, task: TaskID) -> Path:
+        return self.base_path / f"{task}.meta.json"
+
+    def _marker_path(self, task: TaskID) -> Path:
+        return self.base_path / f"{task}.ok"
+
+    def _data_path(self, task: TaskID) -> Path:
+        return self.base_path / f"{task}.fits"
+
+    # ----------------------------
+    # Metadata
+    # ----------------------------
+    def save_meta(self, task: TaskID, meta: dict[str, Any]) -> None:
+        path = self._meta_path(task)
+        tmp = path.with_suffix(".tmp")
+        tmp.parent.mkdir(parents=True, exist_ok=True)
+        logger.debug(f"writing {path}")
+        tmp.write_text(json.dumps(meta, indent=2))
+        tmp.replace(path)
+
+    def load_meta(self, task: TaskID) -> dict[str, Any] | None:
+        path = self._meta_path(task)
+        if not path.exists():
+            return None
+        return json.loads(path.read_text())
+
+    def meta_exists(self, task: TaskID) -> bool:
+        return self._meta_path(task).exists()
+
+    # ----------------------------
+    # Markers
+    # ----------------------------
+    def mark_done(self, task: TaskID) -> None:
+        mp = self._marker_path(task)
+        mp.parent.mkdir(parents=True, exist_ok=True)
+        logger.debug(f"writing {mp}")
+        mp.write_text("done")
+
+    def is_done(self, task: TaskID) -> bool:
+        return self._marker_path(task).exists()
+
+    # ----------------------------
+    # Data
+    # ----------------------------
+    def save_data(self, task: TaskID, content: Table) -> None:
+        path = self._data_path(task)
+        tmp = path.with_suffix(".tmp")
+        tmp.parent.mkdir(parents=True, exist_ok=True)
+        logger.debug(f"writing {path}")
+        content.write(tmp, format="fits")
+        tmp.replace(path)
+
+    def load_data(self, task: TaskID) -> Table:
+        path = self._data_path(task)
+        if not path.exists():
+            raise FileNotFoundError(path)
+        return Table.read(path, format="fits")
+
+    def data_exists(self, task: TaskID) -> bool:
+        return self._data_path(task).exists()

timewise/chunking.py

@@ -0,0 +1,50 @@
+from typing import Iterator
+from pathlib import Path
+import numpy as np
+from numpy import typing as npt
+import pandas as pd
+import logging
+
+logger = logging.getLogger(__name__)
+
+
+class Chunk:
+    def __init__(
+        self, chunk_id: int, indices: npt.ArrayLike, row_indices: npt.ArrayLike
+    ):
+        self.chunk_id = chunk_id
+        self.indices = indices
+        self.row_numbers = row_indices
+
+
+class Chunker:
+    def __init__(self, input_csv: Path, chunk_size: int):
+        self.input_csv = input_csv
+        self.chunk_size = chunk_size
+        self._n_rows = self._count_rows()
+        logger.debug(f"found {self._n_rows} rows in {self.input_csv}")
+
+    def _count_rows(self) -> int:
+        chunk = 1024 * 1024  # Process 1 MB at a time.
+        f = np.memmap(self.input_csv)
+        num_newlines = sum(
+            np.sum(f[i : i + chunk] == ord("\n")) for i in range(0, len(f), chunk)
+        )
+        del f
+        return num_newlines - 1  # one header row
+
+    def __len__(self) -> int:
+        return int(np.ceil(self._n_rows / self.chunk_size))
+
+    def __iter__(self) -> Iterator[Chunk]:
+        for chunk_id in range(len(self)):
+            yield self.get_chunk(chunk_id)
+
+    def get_chunk(self, chunk_id: int) -> Chunk:
+        if chunk_id >= len(self):
+            raise IndexError(f"Invalid chunk_id {chunk_id}")
+        start = chunk_id * self.chunk_size
+        stop = min(start + self.chunk_size, self._n_rows)
+        indices = pd.read_csv(self.input_csv, skiprows=start, nrows=stop - start).index
+        logger.debug(f"chunk {chunk_id}: from {start} to {stop}")
+        return Chunk(chunk_id, indices, np.arange(start=start, stop=stop))

timewise/cli.py

@@ -1,18 +1,124 @@
 import logging
-from argparse import ArgumentParser
+from typing import Annotated, Literal, List
+from pathlib import Path
 
-from timewise.general import main_logger
-from timewise.config_loader import TimewiseConfigLoader
+import typer
 
+from .config import TimewiseConfig
+from .plot.diagnostic import make_plot
 
-logger = logging.getLogger(__name__)
+from rich.logging import RichHandler
 
 
-def timewise_cli():
-    parser = ArgumentParser()
-    parser.add_argument("config", type=str, help="Path to timewise config file")
-    parser.add_argument("-l", "--logging-level", default="INFO", type=str)
-    cfg = vars(parser.parse_args())
+app = typer.Typer(help="Timewsie CLI")
 
-    main_logger.setLevel(cfg.pop("logging_level"))
-    TimewiseConfigLoader.run_yaml(cfg["config"])
+config_path_type = Annotated[
+    Path, typer.Argument(help="Pipeline config file (YAML/JSON)")
+]
+ampel_config_path_type = Annotated[Path, typer.Argument(help="AMPEL config YAML")]
+
+
+# --- Global callback (runs before every command) ---
+@app.callback()
+def main(
+    ctx: typer.Context,
+    log_level: str = typer.Option(
+        "INFO",
+        "--log-level",
+        "-l",
+        help="Logging level (DEBUG, INFO, WARNING, ERROR, CRITICAL)",
+        case_sensitive=False,
+    ),
+):
+    """Global options for all Timewise commands."""
+    # Normalize log level
+    level = getattr(logging, log_level.upper(), None)
+    if not isinstance(level, int):
+        raise typer.BadParameter(f"Invalid log level: {log_level}")
+
+    # Rich logging
+    logging.basicConfig(
+        handlers=[RichHandler(rich_tracebacks=True, markup=True)],
+    )
+    logging.getLogger("timewise").setLevel(level)
+
+    # Store log level in context for subcommands
+    ctx.obj = {"log_level": level}
+
+
+@app.command(help="Download WISE photometry from IRSA")
+def download(
+    config_path: config_path_type,
+):
+    TimewiseConfig.from_yaml(config_path).download.build_downloader().run()
+
+
+@app.command(help="Prepares the AMPEL job file so AMPEL can be run manually")
+def prepare_ampel(
+    config_path: config_path_type,
+):
+    cfg = TimewiseConfig.from_yaml(config_path)
+    ampel_interface = cfg.build_ampel_interface()
+    p = ampel_interface.prepare(config_path)
+    typer.echo(f"AMPEL job file: {p}")
+
+
+@app.command(help="Processes the lightcurves using AMPEL")
+def process(
+    config_path: config_path_type,
+    ampel_config_path: ampel_config_path_type,
+):
+    cfg = TimewiseConfig.from_yaml(config_path)
+    ampel_interface = cfg.build_ampel_interface()
+    ampel_interface.run(config_path, ampel_config_path)
+
+
+@app.command(help="Write stacked lightcurves to disk")
+def export(
+    config_path: config_path_type,
+    output_directory: Annotated[Path, typer.Argument(help="output directory")],
+    indices: Annotated[
+        list[int] | None,
+        typer.Option(
+            "-i", "--indices", help="Indices to export, defaults to all indices"
+        ),
+    ] = None,
+):
+    TimewiseConfig.from_yaml(config_path).build_ampel_interface().export_many(
+        output_directory, indices
+    )
+
+
+@app.command(help="Run download, process and export")
+def run_chain(
+    config_path: config_path_type,
+    ampel_config_path: ampel_config_path_type,
+    output_directory: Annotated[Path, typer.Argument(help="output directory")],
+    indices: Annotated[
+        list[int] | None,
+        typer.Option(
+            "-i", "--indices", help="Indices to export, defaults to all indices"
+        ),
+    ] = None,
+):
+    download(config_path)
+    process(config_path, ampel_config_path)
+    export(config_path, output_directory, indices)
+
+
+@app.command(help="Make diagnostic plots")
+def plot(
+    config_path: config_path_type,
+    indices: Annotated[
+        List[int],
+        typer.Argument(help="Identifiers of the objects for which to create plots"),
+    ],
+    output_directory: Annotated[Path, typer.Argument(help="Output directory")],
+    cutout: Annotated[
+        Literal["sdss", "panstarrs"],
+        typer.Option("-c", "--cutout", help="Which survey to use for cutouts"),
+    ] = "panstarrs",
+):
+    make_plot(
+        config_path, indices=indices, cutout=cutout, output_directory=output_directory
+    )

timewise/config.py

@@ -0,0 +1,34 @@
+from pathlib import Path
+import yaml
+
+import numpy as np
+from pydantic import BaseModel, model_validator
+
+from .io import DownloadConfig
+from .process import AmpelConfig, AmpelInterface
+
+
+class TimewiseConfig(BaseModel):
+    download: DownloadConfig
+    ampel: AmpelConfig
+
+    @classmethod
+    def from_yaml(cls, path: str | Path):
+        path = Path(path)
+        assert path.exists(), f"{path} not found!"
+        with path.open("r") as f:
+            config_dict = yaml.safe_load(f)
+        return cls.model_validate(config_dict)
+
+    @model_validator(mode="after")
+    def validate_query_original_id_key(self) -> "TimewiseConfig":
+        unique_keys = np.unique([q.original_id_key for q in self.download.queries])
+        assert len(unique_keys) == 1, (
+            "Can not use different 'original_id_key' in queries!"
+        )
+        return self
+
+    def build_ampel_interface(self) -> AmpelInterface:
+        return self.ampel.build_interface(
+            self.download.queries[0].original_id_key, self.download.input_csv
+        )

timewise/io/__init__.py

@@ -0,0 +1 @@
+from .config import DownloadConfig

timewise/io/config.py

@@ -0,0 +1,64 @@
+from pathlib import Path
+from typing import List
+import pandas as pd
+from pydantic import BaseModel, Field, model_validator
+
+from .download import Downloader
+from ..query import QueryType
+from ..backend import BackendType
+from ..types import TYPE_MAP
+
+
+class DownloadConfig(BaseModel):
+    input_csv: Path
+    chunk_size: int = 500_000
+    max_concurrent_jobs: int = 4
+    poll_interval: float = 10.0
+    queries: List[QueryType] = Field(..., description="One or more queries per chunk")
+    backend: BackendType = Field(..., discriminator="type")
+
+    service_url: str = "https://irsa.ipac.caltech.edu/TAP"
+
+    @model_validator(mode="after")
+    def validate_input_csv_columns(self) -> "DownloadConfig":
+        """Ensure that the input CSV contains all columns required by queries."""
+        # only validate if the CSV actually exists
+        if not self.input_csv.exists():
+            raise ValueError(f"CSV file does not exist: {self.input_csv}")
+
+        # read just the header and first 10 lines
+        input_table = pd.read_csv(self.input_csv, nrows=10)
+
+        missing_columns = set()
+        wrong_dtype = set()
+        for qc in self.queries:
+            for col, dtype in qc.input_columns.items():
+                if col not in input_table.columns:
+                    missing_columns.add(col)
+                else:
+                    try:
+                        input_table[col].astype(TYPE_MAP[dtype])
+                    except Exception:
+                        wrong_dtype.add(col)
+
+        msg = f"CSV file {self.input_csv}: "
+        if missing_columns:
+            raise KeyError(msg + f"Missing required columns: {sorted(missing_columns)}")
+        if wrong_dtype:
+            raise TypeError(
+                msg
+                + f"Columns not convertable to right data type: {sorted(wrong_dtype)}"
+            )
+
+        return self
+
+    def build_downloader(self) -> Downloader:
+        return Downloader(
+            service_url=self.service_url,
+            input_csv=self.input_csv,
+            chunk_size=self.chunk_size,
+            backend=self.backend,
+            queries=self.queries,
+            max_concurrent_jobs=self.max_concurrent_jobs,
+            poll_interval=self.poll_interval,
+        )