blastwave 1.0.0__py3-none-any.whl

blastwave/query/alert.py

@@ -0,0 +1,252 @@
+"""
+Base Class for an Alert Survey Client
+"""
+
+from abc import ABC, abstractmethod
+from typing import Callable
+
+import numpy as np
+import pandas as pd
+from astropy.coordinates import SkyCoord
+from babamul.models import LsstAlert, ZtfAlert
+
+from blastwave.models import BOOMQuery, Source
+from blastwave.query.boom import BoomClient
+from blastwave.utils.photometry import (
+    deduplicate_lsst_photometry,
+    deduplicate_ztf_photometry,
+)
+
+parse_map = {
+    "ZTF_alerts": deduplicate_ztf_photometry,
+    "LSST_alerts": deduplicate_lsst_photometry,
+}
+
+GenerateFunc = Callable[[dict, pd.DataFrame], Source]
+generator_mapping: dict[str, GenerateFunc] = {
+    "ZTF_alerts": Source.from_ztf,
+    "LSST_alerts": Source.from_lsst,
+}
+
+
+class AlertClient(BoomClient, ABC):
+    """
+    Base class for Alert Survey Client
+    """
+
+    @property
+    @abstractmethod
+    def base_catalog(self) -> str:
+        """The BOOM catalog name for this survey."""
+
+    @property
+    @abstractmethod
+    def base_model(self) -> type[LsstAlert] | type[ZtfAlert]:
+        """The Babamul model for this survey."""
+
+    @property
+    def catalog(self) -> str:
+        """
+        The BOOM catalog name for this survey.
+        """
+        return f"{self.base_catalog}_alerts"
+
+    @property
+    def aux_table(self) -> str:
+        """
+        The BOOM aux catalog name for this survey.
+
+        :return: Name of the BOOM aux catalog
+        """
+        return f"{self.catalog}_aux"
+
+    def get_latest_alert(
+        self,
+        object_id: str | int,
+        catalog: str | None = None,
+        projection: dict | None = None,
+    ) -> dict:
+        """
+        Get latest alert for a given object.
+
+        :param object_id: Object ID
+        :param catalog: Catalog name (e.g. ZTF_alerts or LSST_alerts)
+        :param projection: Mongodb projection for fields to return
+        :return: Nested dictionary of the latest alert for the given object ID
+        """
+        catalog = catalog or self.catalog
+        query = BOOMQuery(
+            catalog_name=catalog,
+            filter={"objectId": {"$eq": str(object_id)}},
+            limit=1,
+            projection=projection,
+            sort={"candidate.jd": -1},
+        )
+        return self.query(query)[0]
+
+    def get_aux_data(
+        self,
+        object_id: str | int,
+        catalog: str | None = None,
+        projection: dict | None = None,
+    ) -> dict:
+        """
+        Get aux data for a given object.
+
+        :param object_id: Object ID
+        :param catalog: Parent catalog name (e.g. ZTF_alerts or LSST_alerts)
+        :param projection: Projection for fields to return
+        :return: Dictionary of the aux data for the given object ID
+        """
+
+        aux_table = f"{catalog}_aux" if catalog is not None else self.aux_table
+        query = BOOMQuery(
+            catalog_name=aux_table,
+            filter={"_id": {"$eq": str(object_id)}},
+            projection=projection,
+        )
+        return self.query(query)[0]
+
+    def get_full_data(
+        self,
+        object_id: str | int,
+        catalog: str | None = None,
+        alert_projection: dict | None = None,
+        aux_projection: dict | None = None,
+    ) -> dict:
+        """
+        Get full data for a given object, combining the latest alert and aux data.
+
+        :param object_id: Object ID
+        :param catalog: Parent catalog name (e.g. ZTF_alerts or LSST_alerts)
+        :param alert_projection: Projection for fields to return for alert
+        :param aux_projection: Projection for fields to return for aux data
+        :return: Dictionary of the aux data for the given object ID
+        """
+        latest = self.get_latest_alert(
+            object_id, catalog=catalog, projection=alert_projection
+        )
+        aux_data = self.get_aux_data(
+            object_id, catalog=catalog, projection=aux_projection
+        )
+        aux_data.update(latest)
+        aux_data = self.update_crossmatches(aux_data)
+        return aux_data
+
+    def update_crossmatches(self, aux_data: dict) -> dict:
+        """
+        Update the cross-matches in the aux data to include the latest alert data.
+
+        :param aux_data: Dictionary of aux data
+        :return: Updated aux data with cross-matches
+        """
+
+        src_position = SkyCoord(
+            aux_data["candidate"]["ra"], aux_data["candidate"]["dec"], unit="deg"
+        )
+
+        for cat in ["LSPSC", "NED"]:
+            matches = self.cone_search(
+                src_position.ra.deg, src_position.dec.deg, catalog=cat, limit=1
+            )
+
+            # Should be zero or 1 for now
+            for match in matches:
+                pos = match.pop("coordinates")["radec_geojson"]
+                sep = float(
+                    src_position.separation(
+                        SkyCoord(
+                            pos["coordinates"][0] + 180.0,
+                            pos["coordinates"][1],
+                            unit="deg",
+                        )
+                    ).arcsec
+                )
+                match["distance_arcsec"] = sep
+
+            aux_data["cross_matches"][cat] = matches
+        return aux_data
+
+    def get_match_photometry(self, catalog: str, aliases: dict) -> pd.DataFrame:
+        """
+        Get photometry for a given survey.
+
+        :param catalog: Catalog name
+        :param aliases: Dictionary of aliases
+        :return: DataFrame of photometry for the survey
+        """
+        parse_f = parse_map[catalog]
+        survey = catalog.split("_")[0]
+        try:
+            match_name = aliases[survey][0]
+            match = self.get_full_data(match_name, catalog=catalog)
+            photometry = parse_f(match)
+            new_df = pd.DataFrame(photometry)
+            new_df["survey"] = survey
+        except (KeyError, IndexError):
+            new_df = pd.DataFrame()
+        return new_df
+
+    def get_all_match_photometry(self, aliases: dict) -> pd.DataFrame:
+        """
+        Get photometry for all cross-matched surveys.
+
+        :param aliases: Aliases
+        :return: DataFrame of photometry for all cross-matched surveys
+        """
+        all_match_photometry = []
+        for survey in parse_map:
+            all_match_photometry.append(self.get_match_photometry(survey, aliases))
+        return pd.concat(all_match_photometry, ignore_index=True)
+
+    def get_all_photometry(
+        self, full_alert: dict, catalog: str | None = None
+    ) -> pd.DataFrame:
+        """
+        Get photometry for the base survey and all cross-matched surveys.
+
+        :param full_alert: Full alert from get_full_data
+        :param catalog: Catalog name
+        :return: DataFrame of photometry for all cross-matched surveys
+        """
+        catalog = self.resolve_catalog(catalog)
+        parse_f = parse_map[catalog]
+        photometry_df = parse_f(full_alert)
+        photometry_df["survey"] = catalog.split("_")[0]
+        match_photometry = self.get_all_match_photometry(full_alert["aliases"])
+
+        if len(match_photometry) > 0:
+            photometry_df = pd.concat(
+                [photometry_df, match_photometry], ignore_index=True
+            )
+
+        photometry_df = (
+            photometry_df.sort_values(by="jd").reset_index(drop=True)
+        ).replace({None: np.nan})
+        return photometry_df
+
+    def get_source(
+        self,
+        object_id: str | int,
+        catalog: str | None = None,
+        alert_projection: dict | None = None,
+        aux_projection: dict | None = None,
+    ) -> Source:
+        """
+        Get the full source data for a given object ID.
+
+        :param object_id: Object ID
+        :param catalog: Catalog name
+        :param alert_projection: Projection for fields to return for alert
+        :param aux_projection: Projection for fields to return for aux data
+        :return: Dictionary of source data
+        """
+        catalog = self.resolve_catalog(catalog)
+        full_data = self.get_full_data(
+            object_id,
+            catalog=catalog,
+            alert_projection=alert_projection,
+            aux_projection=aux_projection,
+        )
+        photometry_df = self.get_all_photometry(full_data, catalog=catalog)
+        return generator_mapping[catalog](full_data, photometry_df)

blastwave/query/boom.py

@@ -0,0 +1,339 @@
+"""
+Client for querying BOOM
+"""
+
+import logging
+import os
+from typing import Callable, Mapping, Optional
+from urllib.parse import urljoin
+
+import numpy as np
+import requests
+from dotenv import load_dotenv
+from urllib3.util import Retry
+
+from blastwave.errors import BOOMCredentialsError
+from blastwave.models.query import BOOMQuery, CatalogQuery, FilterQuery
+from blastwave.query.timeout import DEFAULT_TIMEOUT, TimeoutHTTPAdapter
+
+logger = logging.getLogger(__name__)
+
+BASE_URL = "https://api.kaboom.caltech.edu"
+
+
+class BoomClient:
+    """
+    Basic Boom client class for executing functions
+
+    :param base_url: Base URL to query (default: https://api.kaboom.caltech.edu)
+    """
+
+    def __init__(
+        self,
+        base_url: str = BASE_URL,
+    ):
+        self.base_url = base_url
+        self._session: None | requests.Session = None
+        self._session_headers: None | dict = None
+        self._catalogs: None | list[str] = None
+
+    @property
+    def catalog(self) -> str | None:
+        """Functions as a default catalog name for queries."""
+        return None
+
+    @staticmethod
+    def set_up_session() -> requests.Session:
+        """
+        Set up a session for sending requests to BOOM.
+
+        :return: None
+        """
+        # session to talk to SkyPortal
+        session = requests.Session()
+
+        retries = Retry(
+            total=5,
+            backoff_factor=2,
+            status_forcelist=[405, 429, 500, 502, 503, 504],
+            allowed_methods=["HEAD", "GET", "PUT", "POST", "PATCH"],
+        )
+        adapter = TimeoutHTTPAdapter(max_retries=retries)
+        session.mount("https://", adapter)
+        session.mount("http://", adapter)
+        return session
+
+    def get_session_headers(self) -> dict:
+        """
+        Get session headers
+
+        :return: Session headers
+        """
+        if self._session_headers is None:
+            self._session_headers = {
+                "Content-Type": "application/json",
+                "Authorization": f"Bearer {self._get_boom_token()}",
+            }
+        assert self._session_headers is not None
+        return self._session_headers
+
+    def get_session(self) -> requests.Session:
+        """
+        Wrapper for getting the session.
+        If the session is not set up, it will be set up.
+
+        :return: Session
+        """
+        if self._session is None:
+            self._session = self.set_up_session()
+        assert self._session is not None
+        return self._session
+
+    def _get_boom_token(self) -> str:
+        """
+        Get Boom token from environment variable.
+
+        :return: Boom token
+        """
+        load_dotenv()
+
+        user = os.getenv("BOOM_API_USER")
+
+        if user is None:
+            err = (
+                "No Boom API user specified. "
+                "Add this to .env or set BOOM_API_USER environment variable."
+            )
+            logger.error(err)
+            raise BOOMCredentialsError(err)
+
+        password = os.getenv("BOOM_API_PASSWORD")
+
+        if password is None:
+            err = (
+                "No Boom API password specified. "
+                "Add this to .env or set BOOM_API_PASSWORD environment variable."
+            )
+            logger.error(err)
+            raise BOOMCredentialsError(err)
+
+        boom_token = self.get_fresh_token(user, password)
+        return boom_token
+
+    def api(
+        self, method: str, endpoint: str, data: Optional[Mapping] = None
+    ) -> requests.Response:
+        """
+        Make an API call to a BOOM instance
+
+        headers = {'Authorization': f'token {self.token}'}
+        response = requests.request(method, endpoint, json_dict=data, headers=headers)
+
+        :param method: HTTP method
+        :param endpoint: API endpoint
+        :param data: JSON data to send
+        :return: response from API call
+        """
+        method = method.lower()
+
+        session = self.get_session()
+        headers = self.get_session_headers()
+
+        methods: dict[str, Callable[..., requests.Response]] = {
+            "head": session.head,
+            "get": session.get,
+            "post": session.post,
+            "put": session.put,
+            "patch": session.patch,
+            "delete": session.delete,
+        }
+
+        if endpoint is None:
+            raise ValueError("Endpoint not specified")
+        if method not in ["head", "get", "post", "put", "patch", "delete"]:
+            raise ValueError(f"Unsupported method: {method}")
+
+        url = urljoin(self.base_url, endpoint)
+
+        if method == "get":
+            response = methods[method](
+                url,
+                params=data,
+                headers=headers,
+            )
+        else:
+            response = methods[method](
+                url,
+                json=data,
+                headers=headers,
+            )
+
+        return response
+
+    def query(self, query: BOOMQuery | dict) -> list[dict]:
+        """
+        Make a find query call
+
+        :param query: Query to execute
+        :return: List of results
+        """
+        boom_query = BOOMQuery.model_validate(query)
+        res = self.api(
+            "post", "queries/find", data=boom_query.model_dump(exclude_none=True)
+        )
+        res.raise_for_status()
+        return res.json()["data"]
+
+    def count(self, query: CatalogQuery | FilterQuery | BOOMQuery | dict) -> int:
+        """
+        Make a count call
+
+        :param query: Query to count
+        :return: Number of records
+        """
+        boom_query = FilterQuery.model_validate(query)
+        res = self.api("post", "queries/count", data=boom_query.model_dump())
+        res.raise_for_status()
+        return res.json()["data"]
+
+    def get_fresh_token(self, username: str, password: str) -> str:
+        """
+        Get a fresh token from the Boom API.
+
+        :param username: Username for Boom API
+        :param password: Password for Boom API
+        :return: New token from Boom API
+        """
+        response = requests.post(
+            urljoin(self.base_url, "auth"),
+            headers={"Content-Type": "application/x-www-form-urlencoded"},
+            data={"username": username, "password": password},
+            timeout=DEFAULT_TIMEOUT,
+        )
+
+        if response.status_code != 200:
+            logger.error(f"Failed to get new token: {response.text}")
+            raise BOOMCredentialsError(f"Failed to get new token: {response.text}")
+
+        return response.json()["access_token"]
+
+    def ping(self) -> requests.Response:
+        """
+        Make a ping call
+
+        :return: Response from ping call
+        """
+        return self.api("get", "")
+
+    def get_catalogs(self) -> list[str]:
+        """
+        Get a list of catalogs available in BOOM
+
+        :return: List of catalogs
+        """
+        if self._catalogs is None:
+            res = self.api("get", "catalogs")
+            res.raise_for_status()
+            self._catalogs = [x["name"] for x in res.json()["data"]]
+
+        assert self._catalogs is not None
+        return self._catalogs
+
+    def resolve_catalog(self, catalog: str | None = None) -> str:
+        """
+        Resolve a catalog name
+
+        :param catalog: Catalog name
+        :return: Catalog name to use in query
+        """
+        name = catalog or self.catalog
+        if name is None:
+            raise ValueError("No catalog specified")
+        return name
+
+    def get_entry_count(self, catalog: str | None = None) -> int:
+        """
+        Get the number of entries in a catalog
+
+        :param catalog: Catalog name
+        :return: Number of entries in catalog
+        """
+        query = CatalogQuery(catalog_name=self.resolve_catalog(catalog))
+        res = self.api(
+            "post",
+            "queries/estimated_count",
+            data=query.model_dump(exclude_none=True),
+        )
+        res.raise_for_status()
+        return res.json()["data"]
+
+    def get_catalog_indexes(self, catalog: str | None = None) -> int:
+        """
+        Get the indexes columns for a catalog
+
+        :param catalog: Catalog name
+        :return: Example json
+        """
+        res = self.api("get", f"/catalogs/{self.resolve_catalog(catalog)}/indexes")
+        res.raise_for_status()
+        return res.json()["data"]
+
+    def get_sample_data(self, catalog: str | None = None) -> dict:
+        """
+        Get a sample data entry from a catalog
+
+        :param catalog: Catalog name
+        :return: Example json
+        """
+        res = self.api("get", f"/catalogs/{self.resolve_catalog(catalog)}/sample")
+        res.raise_for_status()
+        return res.json()["data"][0]
+
+    @staticmethod
+    def get_near_sphere_dist(radius_arcsec: float) -> float:
+        """
+        Convert a radius in arcseconds to a radius
+        in meters (on Earth surface) for use in a $nearSphere query
+
+        :param radius_arcsec: Radius in arcseconds
+        :return: Radius in meters
+        """
+        return ((radius_arcsec / 3600.0) * np.pi / 180.0) * 6371008.8
+
+    def cone_search(
+        self,
+        ra: float,
+        dec: float,
+        radius_arcsec: float = 1.0,
+        catalog: str | None = None,
+        limit: int | None = 1,
+    ) -> list[dict]:
+        """
+        Function to do a cone search
+
+        :param ra: Ra (decimal degrees)
+        :param dec: Declination (decimal degrees)
+        :param radius_arcsec: Search radius (arcsec)
+        :param catalog: Catalog name
+        :param limit: Number of results to return
+        :return: List of search results (possibly of length 0 if there are no results)
+        """
+
+        res = self.query(
+            BOOMQuery(
+                catalog_name=self.resolve_catalog(catalog),
+                filter={
+                    "coordinates.radec_geojson": {
+                        "$nearSphere": {
+                            "$geometry": {
+                                "type": "Point",
+                                "coordinates": [ra - 180.0, dec],
+                            },
+                            "$maxDistance": self.get_near_sphere_dist(radius_arcsec),
+                        }
+                    }
+                },
+                limit=limit,
+            )
+        )
+        return res

blastwave/query/lsst.py

@@ -0,0 +1,16 @@
+"""
+LSST Client
+"""
+
+from babamul.models import LsstAlert
+
+from blastwave.query.alert import AlertClient
+
+
+class LSSTClient(AlertClient):
+    """
+    LSST Client
+    """
+
+    base_catalog = "LSST"
+    base_model = LsstAlert

blastwave/query/timeout.py

@@ -0,0 +1,32 @@
+"""
+Timeout HTTP Adapter
+"""
+
+from requests.adapters import HTTPAdapter
+
+DEFAULT_TIMEOUT = 60
+
+
+class TimeoutHTTPAdapter(HTTPAdapter):
+    """
+    HTTP adapter that sets a default timeout for all requests.
+    """
+
+    def __init__(self, *args, **kwargs):
+        self.timeout = DEFAULT_TIMEOUT
+        if "timeout" in kwargs:
+            self.timeout = kwargs["timeout"]
+            del kwargs["timeout"]
+        super().__init__(*args, **kwargs)
+
+    def send(self, request, *args, **kwargs):
+        """
+        Send a request with a default timeout.
+
+        :param request: request to send
+        :param args: args to pass to super().send
+        :param kwargs: kwargs to pass to super().send
+        :return: response from request
+        """
+        kwargs.setdefault("timeout", self.timeout)
+        return super().send(request, *args, **kwargs)

blastwave/query/ztf.py

@@ -0,0 +1,16 @@
+"""
+ZTF Client
+"""
+
+from babamul.models import ZtfAlert
+
+from blastwave.query.alert import AlertClient
+
+
+class ZTFClient(AlertClient):
+    """
+    ZTF Client
+    """
+
+    base_catalog = "ZTF"
+    base_model = ZtfAlert

blastwave/utils/__init__.py

@@ -0,0 +1,12 @@
+"""
+Util classes for blastwave
+"""
+
+from blastwave.utils.cache import (
+    get_crossmatch_path,
+    get_data_dir,
+    get_photometry_path,
+    get_source_path,
+)
+from blastwave.utils.combine import combine_sources, load_consolidated_sources
+from blastwave.utils.plot import plot_lightcurve