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

timewise/process/stacking.py

@@ -0,0 +1,322 @@
+import logging
+from typing import cast, Dict, Any
+
+from scipy import stats
+import numpy as np
+from numpy import typing as npt
+import pandas as pd
+
+from ..util.visits import get_visit_map
+from timewise.process import keys
+
+
+logger = logging.getLogger(__name__)
+
+
+# zero points come from https://wise2.ipac.caltech.edu/docs/release/allsky/expsup/sec4_4h.html#conv2flux
+# published in Jarret et al. (2011): https://ui.adsabs.harvard.edu/abs/2011ApJ...735..112J/abstract
+MAGNITUDE_ZEROPOINTS: Dict[str, float] = {"w1": 20.752, "w2": 19.596}
+# in Jy
+FLUX_ZEROPOINTS = {"w1": 309.54, "w2": 171.787}
+
+
+def calculate_epochs(
+    f: pd.Series,
+    e: pd.Series,
+    visit_mask: npt.NDArray[np.int64],
+    counts: npt.NDArray[np.int64],
+    remove_outliers: bool,
+    outlier_threshold: float,
+    outlier_quantile: float,
+    outlier_mask: npt.NDArray[np.bool_] | None = None,
+) -> tuple[
+    npt.NDArray[np.float64],
+    npt.NDArray[np.float64],
+    npt.NDArray[np.bool_],
+    npt.NDArray[np.bool_],
+    npt.NDArray[np.bool_],
+    npt.NDArray[np.int64],
+]:
+    """
+    Calculates the visits within a raw lightcurve.
+
+    :param f: the fluxes
+    :type f: np.array
+    :param e: the flux errors
+    :type e: np.array
+    :param visit_mask: the visit mask
+    :type visit_mask: np.array
+    :param counts: the counts
+    :type counts: np.array
+    :param remove_outliers: whether to remove outliers
+    :type remove_outliers: bool
+    :param outlier_mask: the outlier mask
+    :type outlier_mask: np.array
+    :return: the epoch
+    :rtype: float
+    """
+
+    if len(f) == 0:
+        return (
+            np.array([]),
+            np.array([]),
+            np.array([]),
+            np.array([]),
+            np.array([]),
+            np.array([]),
+        )
+
+    u_lims = pd.isna(e)
+    nan_mask = pd.isna(f)
+
+    # ---------------------   remove outliers in the bins   ---------------------- #
+
+    # if we do not want to clean outliers just set the threshold to infinity
+    _outlier_threshold = np.inf if not remove_outliers else outlier_threshold
+
+    # set up empty masks
+    outlier_mask = cast(
+        npt.NDArray[np.bool_],
+        (np.array([False] * len(f)) if outlier_mask is None else outlier_mask),
+    )
+    median = np.full_like(counts, np.nan, dtype=float)
+    u = np.full_like(counts, np.nan, dtype=float)
+    use_mask = np.full_like(counts, False, dtype=bool)
+    n_points = counts
+
+    # set up dummy values for number of remaining outliers
+    n_remaining_outlier = np.inf
+
+    # ---------------------   flag upper limits   ---------------------- #
+    bin_n_ulims: npt.NDArray[np.int64] = np.bincount(
+        visit_mask, weights=u_lims, minlength=len(counts)
+    )
+    bin_ulim_bool = cast(npt.NDArray[np.bool_], (counts - bin_n_ulims) == 0)
+    use_mask_ul = ~u_lims | (u_lims & bin_ulim_bool[visit_mask])
+
+    n_loops = 0
+
+    # recalculate uncertainty and median as long as no outliers left
+    while n_remaining_outlier > 0:
+        # make a mask of values to use
+        use_mask = ~outlier_mask & use_mask_ul & ~nan_mask  # type: ignore[operator]
+        n_points = np.bincount(visit_mask, weights=use_mask)
+        zero_points_mask = cast(npt.NDArray[np.bool_], n_points == 0)
+
+        # -------------------------   calculate median   ------------------------- #
+        median = np.zeros_like(counts, dtype=float)
+        visits_at_least_one_point = np.unique(visit_mask[~zero_points_mask[visit_mask]])
+        visits_zero_points = np.unique(visit_mask[zero_points_mask[visit_mask]])
+        median[visits_at_least_one_point] = np.array(
+            [
+                np.median(f[(visit_mask == i) & use_mask])
+                for i in visits_at_least_one_point
+            ]
+        )
+        median[visits_zero_points] = np.nan
+
+        # median is NaN for visits with 0 detections, (i.e. detections in one band and not the other)
+        # if median is NaN for other visits raise Error
+        if np.any(np.isnan(median[n_points > 0])):
+            nan_indices = np.where(np.isnan(median))[0]
+            msg = ""
+            for inan_index in nan_indices:
+                nanf = f[visit_mask == inan_index]
+                msg += f"median is nan for {inan_index}th bin\n{nanf}\n\n"
+            raise ValueError(msg)
+
+        # ---------------------   calculate uncertainty   ---------------------- #
+        mean_deviation = np.bincount(
+            visit_mask[use_mask],
+            weights=(f[use_mask] - median[visit_mask[use_mask]]) ** 2,
+            minlength=len(counts),
+        )
+        one_points_mask = n_points <= 1
+        # calculate standard deviation
+        std = np.zeros_like(counts, dtype=float)
+        std[~one_points_mask] = (
+            np.sqrt(mean_deviation[~one_points_mask])
+            / (n_points[~one_points_mask] - 1)
+            * stats.t.interval(0.68, df=n_points[~one_points_mask] - 1)[1]
+            # for visits with small number of detections we have to correct according to the t distribution
+        )
+        std[one_points_mask] = -np.inf
+
+        # calculate the propagated errors of the single exposure measurements
+        single_exp_measurement_errors = np.sqrt(
+            np.bincount(
+                visit_mask[use_mask],
+                weights=e[use_mask] ** 2,
+                minlength=len(counts),
+            )
+        )
+        e_meas = np.zeros_like(std, dtype=float)
+        e_meas[~zero_points_mask] = (
+            single_exp_measurement_errors[n_points > 0] / n_points[n_points > 0]
+        )
+        e_meas[zero_points_mask] = np.nan
+        # take the maximum value of the measured single exposure errors and the standard deviation
+        u = np.maximum(std, e_meas)
+
+        # Estimate the spread of the flux.
+        # To be robust against outliers, do that with quantiles instead of std
+        qs = np.zeros_like(counts, dtype=float)
+        qs[one_points_mask] = 1e-10
+        visits_at_least_two_point = np.unique(visit_mask[~one_points_mask[visit_mask]])
+        qs[visits_at_least_two_point] = np.array(
+            [
+                np.quantile(
+                    abs(f[(visit_mask == i) & use_mask] - median[i]),
+                    outlier_quantile,
+                    method="interpolated_inverted_cdf",
+                )
+                for i in visits_at_least_two_point
+            ]
+        )
+
+        # ---------------------   remove outliers in the bins   ---------------------- #
+        remaining_outliers = (
+            abs(median[visit_mask] - f) > _outlier_threshold * qs[visit_mask]
+        ) & ~outlier_mask
+        outlier_mask |= remaining_outliers
+        n_remaining_outlier = sum(remaining_outliers) if remove_outliers else 0
+        # setting remaining_outliers to 0 will exit the while loop
+
+        n_loops += 1
+
+        if n_loops > 20:
+            raise Exception(f"{n_loops}!")
+
+    return median, u, bin_ulim_bool, outlier_mask, use_mask, n_points
+
+
+def stack_visits(
+    lightcurve: pd.DataFrame,
+    outlier_threshold: float,
+    outlier_quantile: float,
+    clean_outliers: bool = True,
+):
+    """
+    Combine the data by visits of the satellite of one region in the sky.
+    The visits typically consist of some tens of observations. The individual visits are separated by about
+    six months.
+    The mean flux for one visit is calculated by the weighted mean of the data.
+    The error on that mean is calculated by the root-mean-squared and corrected by the t-value.
+    Outliers per visit are identified if they are more than 100 times the rms away from the mean. These outliers
+    are removed from the calculation of the mean and the error if self.clean_outliers_when_stacking is True.
+
+    :param lightcurve: the raw lightcurve
+    :type lightcurve: pandas.DataFrame
+    :return: the stacked lightcurve
+    :rtype: pandas.DataFrame
+    """
+
+    # -------------------------   create visit mask   -------------------------- #
+    visit_map = get_visit_map(lightcurve.mjd)
+    counts = np.bincount(visit_map)
+
+    stacked_data: Dict[str, Any] = dict()
+
+    # -------------------------   calculate mean mjd   -------------------------- #
+    stacked_data["mean_mjd"] = np.bincount(visit_map, weights=lightcurve.mjd) / counts
+
+    # -------------------------   loop through bands   -------------------------- #
+    for b in ["w1", "w2"]:
+        # loop through magnitude and flux and save the respective datapoints
+
+        outlier_masks: Dict[str, Any] = dict()
+        use_masks = dict()
+        bin_ulim_bools = dict()
+
+        for lum_ext in [keys.FLUX_EXT, keys.MAG_EXT]:
+            f = lightcurve[f"{b}{lum_ext}"]
+            e = lightcurve[f"{b}{keys.ERROR_EXT}{lum_ext}"]
+
+            # we will flag outliers based on the flux only
+            remove_outliers = lum_ext == keys.FLUX_EXT and clean_outliers
+            outlier_mask = outlier_masks.get(keys.FLUX_EXT, None)
+
+            mean, u, bin_ulim_bool, outlier_mask, use_mask, n_points = calculate_epochs(
+                f,
+                e,
+                visit_map,
+                counts,
+                remove_outliers=remove_outliers,
+                outlier_mask=outlier_mask,
+                outlier_quantile=outlier_quantile,
+                outlier_threshold=outlier_threshold,
+            )
+            n_outliers = np.sum(outlier_mask)
+
+            if n_outliers > 0:
+                logger.debug(
+                    f"removed {n_outliers} outliers by brightness for {b} {lum_ext}"
+                )
+
+            stacked_data[f"{b}{keys.MEAN}{lum_ext}"] = mean
+            stacked_data[f"{b}{lum_ext}{keys.RMS}"] = u
+            stacked_data[f"{b}{lum_ext}{keys.UPPER_LIMIT}"] = bin_ulim_bool
+            stacked_data[f"{b}{lum_ext}{keys.NPOINTS}"] = n_points
+
+            outlier_masks[lum_ext] = outlier_mask
+            use_masks[lum_ext] = use_mask
+            bin_ulim_bools[lum_ext] = bin_ulim_bool
+
+        # -------  calculate the zeropoints per exposure ------- #
+        # this might look wrong since we use the flux mask on the magnitudes but it s right
+        # for each flux measurement we need the corresponding magnitude to get the zeropoint
+        mags = lightcurve[f"{b}{keys.MAG_EXT}"]
+        inst_fluxes = lightcurve[f"{b}{keys.FLUX_EXT}"]
+        pos_m = inst_fluxes > 0  # select only positive fluxes, i.e. detections
+        zp_mask = pos_m & use_masks[keys.FLUX_EXT]
+
+        # calculate zero points
+        zps = np.zeros_like(inst_fluxes)
+        zps[zp_mask] = mags[zp_mask] + 2.5 * np.log10(inst_fluxes[zp_mask])
+        # find visits with no zeropoints
+        n_valid_zps = np.bincount(visit_map, weights=zp_mask)
+        at_least_one_valid_zp = n_valid_zps > 0
+        # calculate the median zeropoint for each visit
+        zps_median = np.zeros_like(n_valid_zps, dtype=float)
+        zps_median[n_valid_zps > 0] = np.array(
+            [
+                np.median(zps[(visit_map == i) & zp_mask])
+                for i in np.unique(visit_map[at_least_one_valid_zp[visit_map]])
+            ]
+        )
+        # if there are only non-detections then fall back to default zeropoint
+        zps_median[n_valid_zps == 0] = MAGNITUDE_ZEROPOINTS[b]
+        # if the visit only has upper limits then use the fall-back zeropoint
+        zps_median[bin_ulim_bools[keys.FLUX_EXT]] = MAGNITUDE_ZEROPOINTS[b]
+
+        # ---------------   calculate flux density from instrument flux   ---------------- #
+        # get the instrument flux [digital numbers], i.e. source count
+        inst_fluxes_e = lightcurve[f"{b}{keys.ERROR_EXT}{keys.FLUX_EXT}"]
+
+        # calculate the proportionality constant between flux density and source count
+        mag_zp = FLUX_ZEROPOINTS[b] * 1e3  # in mJy
+        flux_dens_const = mag_zp * 10 ** (-zps_median / 2.5)
+
+        # calculate flux densities from instrument counts
+        flux_densities = inst_fluxes * flux_dens_const[visit_map]
+        flux_densities_e = inst_fluxes_e * flux_dens_const[visit_map]
+
+        # bin flux densities
+        mean_fd, u_fd, ul_fd, outlier_mask_fd, use_mask_fd, n_points_fd = (
+            calculate_epochs(
+                flux_densities,
+                flux_densities_e,
+                visit_map,
+                counts,
+                remove_outliers=False,
+                outlier_mask=outlier_masks[keys.FLUX_EXT],
+                outlier_threshold=outlier_threshold,
+                outlier_quantile=outlier_quantile,
+            )
+        )
+        stacked_data[f"{b}{keys.MEAN}{keys.FLUX_DENSITY_EXT}"] = mean_fd
+        stacked_data[f"{b}{keys.FLUX_DENSITY_EXT}{keys.RMS}"] = u_fd
+        stacked_data[f"{b}{keys.FLUX_DENSITY_EXT}{keys.UPPER_LIMIT}"] = ul_fd
+        stacked_data[f"{b}{keys.FLUX_DENSITY_EXT}{keys.NPOINTS}"] = n_points_fd
+
+    return pd.DataFrame(stacked_data)

timewise/process/template.yml

@@ -0,0 +1,49 @@
+channel:
+- access:
+  - ZTF_PUB
+  name: wise
+  policy: []
+  version: 0
+mongo:
+  prefix: MONGODB_NAME
+  reset: true
+name: timewise
+task:
+- config:
+    compiler_opts: TiCompilerOptions
+    directives:
+    - channel: wise
+      ingest:
+        mux:
+          combine:
+          - state_t2:
+            - unit: T2StackVisits
+            unit: T1HDBSCAN
+            config:
+              original_id_key: ORIGINAL_ID_KEY
+              input_mongo_db_name: INPUT_MONGODB_NAME
+              plot: true
+          unit: TiMongoMuxer
+    iter_max: 1000000
+    shaper: TiDataPointShaper
+    supplier:
+      config:
+        dpid: hash
+        loader:
+          config:
+            timewise_config_file: TIMEWISE_CONFIG_PATH
+            stock_id_column_name: ORIGINAL_ID_KEY
+          unit: TimewiseFileLoader
+      unit: TimewiseAlertSupplier
+  multiplier: 1
+  title: t0
+  template:
+    live:
+      - resolve_run_time_aliases
+      - hash_t2_config
+  unit: AlertConsumer
+- config:
+    log_profile: default
+  multiplier: 1
+  title: t2
+  unit: T2Worker

timewise/query/__init__.py

@@ -0,0 +1,6 @@
+from pydantic import Field
+from typing import Union, Annotated, TypeAlias
+from .positional import PositionalQuery
+
+# Discriminated union of all query types
+QueryType: TypeAlias = Annotated[Union[PositionalQuery], Field(discriminator="type")]

timewise/query/base.py

@@ -0,0 +1,45 @@
+import abc
+from typing import ClassVar, List
+from pydantic import BaseModel
+from hashlib import sha256
+
+from ..tables import TableType
+
+
+class Query(abc.ABC, BaseModel):
+    type: str
+    upload_name: ClassVar[str] = "mine"
+
+    original_id_key: str = "orig_id"
+    constraints: List[str] = [
+        "nb < 2",
+        "na < 1",
+        "cc_flags like '00%'",
+        "qi_fact >= 1",
+        "saa_sep >= 5",
+        "moon_masked like '00%'",
+    ]
+    columns: List[str]
+    table: TableType
+
+    @property
+    @abc.abstractmethod
+    def input_columns(self) -> dict[str, str]: ...
+
+    @abc.abstractmethod
+    def build(self) -> str: ...
+
+    @property
+    def adql(self) -> str:
+        """ADQL string computed once per instance."""
+        return self.build()
+
+    @property
+    def hash(self) -> str:
+        return (
+            self.type
+            + "_"
+            + self.table.name
+            + "_"
+            + sha256(self.adql.encode()).hexdigest()
+        )

timewise/query/positional.py

@@ -0,0 +1,40 @@
+import logging
+from typing import Literal, Dict
+
+from .base import Query
+
+logger = logging.getLogger(__name__)
+
+
+class PositionalQuery(Query):
+    type: Literal["positional"] = "positional"
+    radius_arcsec: float
+
+    @property
+    def input_columns(self) -> Dict[str, str]:
+        return {"ra": "float", "dec": "float", self.original_id_key: "int"}
+
+    def build(self) -> str:
+        logger.debug(f"constructing positional query for {self.table.name}")
+
+        q = "SELECT \n\t"
+        for k in self.columns:
+            q += f"{self.table.name}.{k}, "
+        q += f"\n\tmine.{self.original_id_key} \n"
+        q += f"FROM\n\tTAP_UPLOAD.{self.upload_name} AS mine \n"
+        q += f"RIGHT JOIN\n\t{self.table.name} \n"
+        q += "WHERE \n"
+        q += (
+            f"\tCONTAINS(POINT('J2000',{self.table.name}.{self.table.ra_column},{self.table.name}.{self.table.dec_column}),"
+            f"CIRCLE('J2000',mine.ra,mine.dec,{self.radius_arcsec / 3600:.18f}))=1 "
+        )
+
+        if len(self.constraints) > 0:
+            q += " AND (\n"
+            for c in self.constraints:
+                q += f"\t{self.table.name}.{c} AND \n"
+            q = q.strip(" AND \n")
+            q += "\t)"
+
+        logger.debug(f"\n{q}")
+        return q

timewise/tables/__init__.py

@@ -0,0 +1,10 @@
+from pydantic import Field
+from typing import Union, Annotated
+
+from .allwise_p3as_mep import allwise_p3as_mep
+from .neowiser_p1bs_psd import neowiser_p1bs_psd
+
+
+TableType = Annotated[
+    Union[allwise_p3as_mep, neowiser_p1bs_psd], Field(discriminator="name")
+]

timewise/tables/allwise_p3as_mep.py

@@ -0,0 +1,22 @@
+from typing import Literal, ClassVar, Type, Dict
+from .base import TableConfig
+
+
+class allwise_p3as_mep(TableConfig):
+    name: Literal["allwise_p3as_mep"] = "allwise_p3as_mep"
+    columns_dtypes: ClassVar[Dict[str, Type]] = {
+        "ra": float,
+        "dec": float,
+        "mjd": float,
+        "cntr_mf": str,
+        "w1mpro_ep": float,
+        "w1sigmpro_ep": float,
+        "w2mpro_ep": float,
+        "w2sigmpro_ep": float,
+        "w1flux_ep": float,
+        "w1sigflux_ep": float,
+        "w2flux_ep": float,
+        "w2sigflux_ep": float,
+    }
+    ra_column: ClassVar[str] = "ra"
+    dec_column: ClassVar[str] = "dec"

timewise/tables/base.py

@@ -0,0 +1,9 @@
+from typing import ClassVar, Dict, Type
+from pydantic import BaseModel
+
+
+class TableConfig(BaseModel):
+    name: str
+    columns_dtypes: ClassVar[Dict[str, Type]]
+    ra_column: ClassVar[str]
+    dec_column: ClassVar[str]

timewise/tables/neowiser_p1bs_psd.py

@@ -0,0 +1,22 @@
+from typing import Literal, ClassVar, Dict, Type
+from .base import TableConfig
+
+
+class neowiser_p1bs_psd(TableConfig):
+    name: Literal["neowiser_p1bs_psd"] = "neowiser_p1bs_psd"
+    columns_dtypes: ClassVar[Dict[str, Type]] = {
+        "ra": float,
+        "dec": float,
+        "mjd": float,
+        "allwise_cntr": str,
+        "w1mpro": float,
+        "w1sigmpro": float,
+        "w2mpro": float,
+        "w2sigmpro": float,
+        "w1flux": float,
+        "w1sigflux": float,
+        "w2flux": float,
+        "w2sigflux": float,
+    }
+    ra_column: ClassVar[str] = "ra"
+    dec_column: ClassVar[str] = "dec"

timewise/types.py

@@ -0,0 +1,30 @@
+from typing import TypedDict, NamedTuple
+
+from timewise.query import QueryType
+
+
+class TAPJobMeta(TypedDict):
+    url: str
+    status: str
+    submitted: str
+    last_checked: str
+    input_length: int
+    query: str
+    query_config: QueryType | dict
+    completed_at: str
+
+
+class TaskID(NamedTuple):
+    """
+    Generic identifier for a unit of work.
+    Can be extended by Downloader/Processor as needed.
+    """
+
+    namespace: str  # e.g. "downloader", "processor"
+    key: str  # unique string, e.g. "chunk_0001_q0" or "mask_2025-01-01"
+
+    def __str__(self):
+        return f"{self.namespace}_{self.key}"
+
+
+TYPE_MAP = {"float": float, "str": str, "int": int}

timewise/util/backoff.py

@@ -0,0 +1,12 @@
+import logging
+
+
+logger = logging.getLogger(__name__)
+
+
+def backoff_hndlr(details):
+    logger.info(
+        "Backing off {wait:0.1f} seconds after {tries} tries "
+        "calling function {target} with args {args} and kwargs "
+        "{kwargs}".format(**details)
+    )

timewise/util/csv_utils.py

@@ -0,0 +1,12 @@
+import numpy as np
+from pathlib import Path
+
+
+def get_n_rows(path: str | Path):
+    chunk = 1024 * 1024  # Process 1 MB at a time.
+    f = np.memmap(path)
+    num_newlines = sum(
+        np.sum(f[i : i + chunk] == ord("\n")) for i in range(0, len(f), chunk)
+    )
+    del f
+    return num_newlines

timewise/util/error_threading.py

@@ -0,0 +1,70 @@
+import sys
+from queue import Queue
+from threading import Thread, Event
+from typing import Any, Optional, Callable, Mapping
+
+
+class ErrorQueue(Queue):
+    """Queue subclass whose join() re-raises exceptions from worker threads."""
+
+    def __init__(self, stop_event: Event, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        self.error_queue: Queue = Queue()
+        self.stop_event = stop_event
+
+    def report_error(self, exc_info):
+        """Called by workers to push an exception into the error queue."""
+        self.error_queue.put(exc_info)
+        # Also decrement unfinished_tasks, so join() won't block forever
+        with self.all_tasks_done:
+            self.unfinished_tasks = max(0, self.unfinished_tasks - 1)
+            self.all_tasks_done.notify_all()
+
+    def join(self):
+        """Wait until all tasks are done, or raise if a worker failed."""
+        with self.all_tasks_done:
+            while self.unfinished_tasks:
+                if not self.error_queue.empty():
+                    exc_info = self.error_queue.get()
+                    self.stop_event.set()
+                    raise exc_info[1].with_traceback(exc_info[2])
+                self.all_tasks_done.wait()
+
+    def raise_errors(self):
+        """
+        Raise the first worker exception, if any.
+        """
+        if not self.error_queue.empty():
+            exc_info = self.error_queue.get()
+            raise exc_info[1].with_traceback(exc_info[2])
+
+
+class ExceptionSafeThread(Thread):
+    """Thread subclass that reports uncaught exceptions to the ErrorQueue."""
+
+    def __init__(
+        self,
+        error_queue: Any,
+        *,
+        group: Optional[Any] = None,
+        target: Optional[Callable[..., Any]] = None,
+        name: Optional[str] = None,
+        args: tuple = (),
+        kwargs: Optional[Mapping[str, Any]] = None,
+        daemon: Optional[bool] = None,
+    ):
+        super().__init__(
+            group=group,
+            target=target,
+            name=name,
+            args=args,
+            kwargs=kwargs,
+            daemon=daemon,
+        )
+        self.error_queue = error_queue
+
+    def run(self):
+        try:
+            super().run()
+        except Exception:
+            self.error_queue.report_error(sys.exc_info())