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

timewise/plot/lightcurve.py

@@ -0,0 +1,112 @@
+from typing import Dict
+import warnings
+import matplotlib.pyplot as plt
+import pandas as pd
+import numpy as np
+
+from timewise.process import keys
+
+
+BAND_PLOT_COLORS = {"w1": "r", "w2": "b"}
+
+
+def plot_lightcurve(
+    lum_key: str,
+    stacked_lightcurve: pd.DataFrame | None = None,
+    raw_lightcurve: pd.DataFrame | None = None,
+    ax: plt.Axes | None = None,
+    colors: Dict[str, str] | None = None,
+    add_to_label: str = "",
+    **kwargs,
+) -> tuple[plt.Figure, plt.Axes]:
+    assert (stacked_lightcurve is not None) or (raw_lightcurve is not None)
+
+    if not colors:
+        colors = BAND_PLOT_COLORS
+
+    if not ax:
+        fig, ax = plt.subplots(**kwargs)
+    else:
+        fig = plt.gcf()
+
+    for b in ["w1", "w2"]:
+        try:
+            if isinstance(stacked_lightcurve, pd.DataFrame):
+                ul_mask_stacked = np.array(
+                    stacked_lightcurve[f"{b}{lum_key}{keys.UPPER_LIMIT}"]
+                ).astype(bool)
+                ax.errorbar(
+                    stacked_lightcurve[keys.MEAN + "_mjd"][~ul_mask_stacked],
+                    stacked_lightcurve[f"{b}{keys.MEAN}{lum_key}"][~ul_mask_stacked],
+                    yerr=stacked_lightcurve[f"{b}{lum_key}{keys.RMS}"][
+                        ~ul_mask_stacked
+                    ],
+                    label=f"{b}{add_to_label} stacked",
+                    ls="",
+                    marker="s",
+                    c=colors[b],
+                    markersize=4,
+                    markeredgecolor="k",
+                    ecolor="k",
+                    capsize=2,
+                )
+                ax.scatter(
+                    stacked_lightcurve[keys.MEAN + "_mjd"][ul_mask_stacked],
+                    stacked_lightcurve[f"{b}{keys.MEAN}{lum_key}"][ul_mask_stacked],
+                    marker="v",
+                    c=colors[b],
+                    alpha=0.7,
+                    s=2,
+                )
+
+            if isinstance(raw_lightcurve, pd.DataFrame):
+                m = ~raw_lightcurve[f"{b}{lum_key}"].isna()
+                ul_mask_raw = raw_lightcurve[f"{b}{keys.ERROR_EXT}{lum_key}"].isna()
+
+                tot_m = m & ~ul_mask_raw
+                if np.any(tot_m):
+                    ax.errorbar(
+                        raw_lightcurve.mjd[tot_m],
+                        raw_lightcurve[f"{b}{lum_key}"][tot_m],
+                        yerr=raw_lightcurve[f"{b}{keys.ERROR_EXT}{lum_key}"][tot_m],
+                        label=f"{b}{add_to_label}",
+                        ls="",
+                        marker="o",
+                        c=colors[b],
+                        markersize=4,
+                        alpha=0.3,
+                    )
+
+                single_ul_m = m & ul_mask_raw
+                if np.any(single_ul_m):
+                    label = (
+                        f"{b}{add_to_label} upper limits" if not np.any(tot_m) else ""
+                    )
+                    ax.scatter(
+                        raw_lightcurve.mjd[single_ul_m],
+                        raw_lightcurve[f"{b}{lum_key}"][single_ul_m],
+                        marker="d",
+                        c=colors[b],
+                        alpha=0.3,
+                        s=1,
+                        label=label,
+                    )
+
+        except KeyError as e:
+            raise KeyError(f"Could not find brightness key {e}!")
+
+    if lum_key == keys.MAG_EXT:
+        ylim = ax.get_ylim()
+        ax.set_ylim(max(ylim), min(ylim))
+
+    ax.set_xlabel("MJD")
+    ax.set_ylabel(lum_key)
+    with warnings.catch_warnings():
+        warnings.filterwarnings(
+            "ignore",
+            message="No artists with labels found to put in legend",
+            category=UserWarning,
+        )
+        ax.legend()
+
+    return fig, ax

timewise/plot/panstarrs.py

@@ -0,0 +1,260 @@
+import logging
+import hashlib
+from pathlib import Path
+from io import BytesIO
+from PIL import Image
+
+import requests
+import matplotlib.pyplot as plt
+import numpy
+from astropy.table import Table
+import astropy.config.paths
+
+
+logger = logging.getLogger(__name__)
+
+
+class PanSTARRSQueryError(Exception):
+    pass
+
+
+def load_cache_or_download(url):
+    logger.debug(f"loading or downloading {url}")
+    h = hashlib.md5(url.encode()).hexdigest()
+    cache_dir = Path(astropy.config.paths.get_cache_dir())
+    cache_file = cache_dir / (h + ".cache")
+    logger.debug(f"cache file is {cache_file}")
+    if not cache_file.is_file():
+        logger.debug(f"downloading {url}")
+        r = requests.get(url)
+        with open(cache_file, "wb") as f:
+            f.write(r.content)
+        return r.content
+    else:
+        logger.debug(f"loading {cache_file}")
+        with open(cache_file, "rb") as f:
+            return f.read()
+
+
+def annotate_not_available(ax):
+    xlim = ax.get_xlim()
+    ylim = ax.get_ylim()
+    x = sum(xlim) / 2
+    y = sum(ylim) / 2
+    logger.debug(f"annotate_not_available at {x}, {y}")
+    ax.annotate(
+        "Outside\nPanSTARRS\nFootprint",
+        (x, y),
+        color="red",
+        ha="center",
+        va="center",
+        fontsize=10,
+    )
+
+
+def getimages(ra, dec, filters="grizy"):
+    """Query ps1filenames.py service to get a list of images
+
+    ra, dec = position in degrees
+    size = image size in pixels (0.25 arcsec/pixel)
+    filters = string with filters to include
+    Returns a table with the results
+    """
+
+    service = "https://ps1images.stsci.edu/cgi-bin/ps1filenames.py"
+    url = f"{service}?ra={ra}&dec={dec}&filters={filters}"
+    content = load_cache_or_download(url)
+    table = Table.read(content.decode(), format="ascii")
+    return table
+
+
+def geturl(
+    ra, dec, size=240, output_size=None, filters="grizy", format="jpg", color=False
+):
+    """Get URL for images in the table
+
+    ra, dec = position in degrees
+    size = extracted image size in pixels (0.25 arcsec/pixel)
+    output_size = output (display) image size in pixels (default = size).
+                  output_size has no effect for fits format images.
+    filters = string with filters to include
+    format = data format (options are "jpg", "png" or "fits")
+    color = if True, creates a color image (only for jpg or png format).
+            Default is return a list of URLs for single-filter grayscale images.
+    Returns a string with the URL
+    """
+
+    if color and format == "fits":
+        raise ValueError("color images are available only for jpg or png formats")
+    if format not in ("jpg", "png", "fits"):
+        raise ValueError("format must be one of jpg, png, fits")
+    table = getimages(ra, dec, filters=filters)
+    if len(table) == 0:
+        raise PanSTARRSQueryError("No images available")
+    url = (
+        f"https://ps1images.stsci.edu/cgi-bin/fitscut.cgi?"
+        f"ra={ra}&dec={dec}&size={size}&format={format}"
+    )
+    if output_size:
+        url = url + "&output_size={}".format(output_size)
+    # sort filters from red to blue
+    flist = ["yzirg".find(x) for x in table["filter"]]
+    table = table[numpy.argsort(flist)]
+    if color:
+        if len(table) > 3:
+            # pick 3 filters
+            table = table[[0, len(table) // 2, len(table) - 1]]
+        for i, param in enumerate(["red", "green", "blue"]):
+            url = url + "&{}={}".format(param, table["filename"][i])
+    else:
+        urlbase = url + "&red="
+        url = []
+        for filename in table["filename"]:
+            url.append(urlbase + filename)
+    return url
+
+
+def getcolorim(ra, dec, size=240, output_size=None, filters="grizy", format="jpg"):
+    """Get color image at a sky position
+
+    ra, dec = position in degrees
+    size = extracted image size in pixels (0.25 arcsec/pixel)
+    output_size = output (display) image size in pixels (default = size).
+                  output_size has no effect for fits format images.
+    filters = string with filters to include
+    format = data format (options are "jpg", "png")
+    Returns the image
+    """
+
+    if format not in ("jpg", "png"):
+        raise ValueError("format must be jpg or png")
+    url = geturl(
+        ra,
+        dec,
+        size=size,
+        filters=filters,
+        output_size=output_size,
+        format=format,
+        color=True,
+    )
+    content = load_cache_or_download(url)
+    im = Image.open(BytesIO(content))
+    return im
+
+
+def getgrayim(ra, dec, size=240, output_size=None, filter="g", format="jpg"):
+    """Get grayscale image at a sky position
+
+    ra, dec = position in degrees
+    size = extracted image size in pixels (0.25 arcsec/pixel)
+    output_size = output (display) image size in pixels (default = size).
+                  output_size has no effect for fits format images.
+    filter = string with filter to extract (one of grizy)
+    format = data format (options are "jpg", "png")
+    Returns the image
+    """
+
+    if format not in ("jpg", "png"):
+        raise ValueError("format must be jpg or png")
+    if filter not in list("grizy"):
+        raise ValueError("filter must be one of grizy")
+    url = geturl(
+        ra, dec, size=size, filters=filter, output_size=output_size, format=format
+    )
+    content = load_cache_or_download(url[0])
+    im = Image.open(BytesIO(content))
+    return im
+
+
+def plot_panstarrs_cutout(
+    ra,
+    dec,
+    arcsec,
+    interactive=False,
+    fn=None,
+    title=None,
+    save=False,
+    ax=False,
+    plot_color_image=False,
+    height=2.5,
+):
+    arcsec_per_px = 0.25
+    ang_px = int(arcsec / arcsec_per_px)
+
+    imshow_kwargs = {
+        "origin": "upper",
+        "extent": ([arcsec / 2, -arcsec / 2, -arcsec / 2, arcsec / 2]),
+    }
+    scatter_args = [0, 0]
+    scatter_kwargs = {"marker": "x", "color": "red"}
+
+    if not plot_color_image:
+        filters = "grizy"
+        if not ax:
+            fig, axss = plt.subplots(
+                2,
+                len(filters),
+                sharex="all",
+                sharey="all",
+                gridspec_kw={"wspace": 0, "hspace": 0, "height_ratios": [1, 8]},
+                figsize=(height * 5, height),
+            )
+        else:
+            fig = plt.gcf()
+            axss = ax
+
+        for j, fil in enumerate(list(filters)):
+            axs = axss[1]
+            try:
+                im = getgrayim(ra, dec, size=ang_px, filter=fil)
+                axs[j].imshow(im, cmap="gray", **imshow_kwargs)
+            except PanSTARRSQueryError:
+                axs[j].set_xlim(-arcsec / 2, arcsec / 2)
+                axs[j].set_ylim(-arcsec / 2, arcsec / 2)
+                annotate_not_available(axs[j])
+
+            axs[j].scatter(*scatter_args, **scatter_kwargs)
+            axs[j].set_title(fil)
+            axss[0][j].axis("off")
+
+    else:
+        logger.debug("plotting color image")
+        if not ax:
+            fig, axss = plt.subplots(figsize=(height, height))
+        else:
+            fig = plt.gcf()
+            axss = ax
+
+        try:
+            im = getcolorim(ra, dec, size=ang_px)
+            axss.imshow(im, **imshow_kwargs)
+        except PanSTARRSQueryError:
+            axss.set_xlim(-arcsec / 2, arcsec / 2)
+            axss.set_ylim(-arcsec / 2, arcsec / 2)
+            annotate_not_available(axss)
+        axss.scatter(*scatter_args, **scatter_kwargs)
+
+    _this_title = title if title else f"{ra}_{dec}"
+    si = "-" if dec > 0 else "+"
+    ylabel = f"Dec {si} {abs(dec):.2f} [arcsec]"
+    xlabel = f"RA - {ra:.2f} [arcsec]"
+    try:
+        axss.set_title(_this_title)
+        axss.set_xlabel(xlabel)
+        axss.set_ylabel(ylabel)
+        axss.grid(ls=":", alpha=0.5)
+    except AttributeError:  # in this case axss is an array
+        fig.supylabel(ylabel)
+        fig.supxlabel(xlabel)
+        fig.suptitle(_this_title)
+        for a in axss.flatten():
+            a.grid(ls=":", alpha=0.5)
+
+    if save:
+        logger.info(f"saving under {fn}")
+        fig.savefig(fn)
+
+    if interactive:
+        return fig, axss
+
+    plt.close()

timewise/plot/sdss.py

@@ -0,0 +1,109 @@
+import requests
+import os
+import getpass
+import logging
+import matplotlib.pyplot as plt
+import backoff
+
+from ..util.backoff import backoff_hndlr
+
+
+logger = logging.getLogger(__name__)
+
+
+def get_sdss_credentials():
+    if not os.environ.get("SDSS_USERID"):
+        os.environ["SDSS_USERID"] = input("Enter SDSS user ID:")
+    if not os.environ.get("SDSS_USERPW"):
+        os.environ["SDSS_USERPW"] = getpass.getpass("Enter SDSS password:")
+    return os.environ["SDSS_USERID"], os.environ["SDSS_USERPW"]
+
+
+def login_to_sciserver():
+    try:
+        from SciServer import Authentication
+    except ModuleNotFoundError:
+        raise ModuleNotFoundError(
+            "Please install SciServer (https://github.com/sciserver/SciScript-Python) "
+            "if you want to see SDSS cutouts!"
+        )
+
+    uid, pw = get_sdss_credentials()
+    logger.debug(f"logging in to SciServer with username {uid}")
+    Authentication.login(uid, pw)
+
+
+@backoff.on_exception(
+    backoff.expo, requests.RequestException, max_tries=50, on_backoff=backoff_hndlr
+)
+def get_cutout(*args, **kwargs):
+    login_to_sciserver()
+    from SciServer import SkyServer
+
+    return SkyServer.getJpegImgCutout(*args, **kwargs)
+
+
+def plot_sdss_cutout(
+    ra,
+    dec,
+    arcsec=20,
+    arcsec_per_px=0.1,
+    interactive=False,
+    fn=None,
+    title=None,
+    save=False,
+    ax=False,
+    height=2.5,
+):
+    ang_px = int(arcsec / arcsec_per_px)
+
+    if not ax:
+        fig, ax = plt.subplots(figsize=(height, height))
+    else:
+        fig = plt.gcf()
+
+    try:
+        im = get_cutout(ra, dec, scale=arcsec_per_px, height=ang_px, width=ang_px)
+        ax.imshow(
+            im,
+            origin="upper",
+            extent=(
+                (
+                    arcsec / 2,
+                    -arcsec / 2,
+                    -arcsec / 2,
+                    arcsec / 2,
+                )
+            ),
+            cmap="gray",
+        )
+        ax.scatter(0, 0, marker="x", color="red")
+
+    except Exception as e:
+        if "outside the SDSS footprint" in str(e):
+            xlim = ax.get_xlim()
+            ylim = ax.get_ylim()
+            x = sum(xlim) / 2
+            y = sum(ylim) / 2
+            ax.annotate(
+                "Outside SDSS Footprint",
+                (x, y),
+                color="red",
+                ha="center",
+                va="center",
+                fontsize=20,
+            )
+        else:
+            raise Exception(e)
+
+    if title:
+        ax.set_title(title)
+
+    if save:
+        logger.debug(f"saving under {fn}")
+        fig.savefig(fn)
+
+    if interactive:
+        return fig, ax
+
+    plt.close()

timewise/process/__init__.py

@@ -0,0 +1,2 @@
+from .config import AmpelConfig
+from .interface import AmpelInterface

timewise/process/config.py

@@ -0,0 +1,30 @@
+from pathlib import Path
+import logging
+
+from pydantic import BaseModel
+
+from .interface import AmpelInterface
+
+
+logger = logging.getLogger(__name__)
+DEFAULT_TEMPLATE_PATH = Path(__file__).parent / "template.yml"
+
+
+class AmpelConfig(BaseModel):
+    mongo_db_name: str
+    template_path: Path = DEFAULT_TEMPLATE_PATH
+    uri: str = "localhost:27017"
+
+    @property
+    def input_mongo_db_name(self) -> str:
+        return self.mongo_db_name + "_input"
+
+    def build_interface(self, original_id_key: str, input_csv: Path) -> AmpelInterface:
+        return AmpelInterface(
+            mongo_db_name=self.mongo_db_name,
+            orig_id_key=original_id_key,
+            input_csv=input_csv,
+            input_mongo_db_name=self.input_mongo_db_name,
+            template_path=self.template_path,
+            uri=self.uri,
+        )

timewise/process/interface.py

@@ -0,0 +1,143 @@
+from pathlib import Path
+import logging
+from typing import Iterable, List, cast
+
+import numpy as np
+from numpy import typing as npt
+import pandas as pd
+from pymongo import MongoClient, ASCENDING
+from pymongo.collection import Collection
+from pymongo.database import Database
+from ampel.cli.JobCommand import JobCommand
+from ampel.types import DataPointId, StockId
+
+
+logger = logging.getLogger(__name__)
+
+
+class AmpelInterface:
+    def __init__(
+        self,
+        mongo_db_name: str,
+        orig_id_key: str,
+        input_csv: Path,
+        input_mongo_db_name: str,
+        template_path: str | Path,
+        uri: str,
+    ):
+        self.mongo_db_name = mongo_db_name
+        self.orig_id_key = orig_id_key
+        self.input_csv = input_csv
+        self.input_mongo_db_name = input_mongo_db_name
+        self.template_path = Path(template_path)
+        self.uri = uri
+
+    def import_input(self):
+        # if collection already exists, assume import was already done
+        if "input" in self.client[self.input_mongo_db_name].list_collection_names():
+            logger.debug(
+                f"'input' collection already exists in '{self.input_mongo_db_name}'."
+            )
+            return
+
+        logger.debug(f"importing {self.input_csv} into {self.input_mongo_db_name}")
+        col = self.client[self.input_mongo_db_name]["input"]
+
+        # create an index from stock id
+        col.create_index([(self.orig_id_key, ASCENDING)], unique=True)
+        col.insert_many(pd.read_csv(self.input_csv).to_dict(orient="records"))
+
+    def make_ampel_job_file(self, cfg_path: Path) -> Path:
+        logger.debug(f"loading ampel job template from {self.template_path}")
+        with self.template_path.open("r") as f:
+            template = f.read()
+
+        ampel_job = (
+            template.replace("TIMEWISE_CONFIG_PATH", str(cfg_path))
+            .replace("ORIGINAL_ID_KEY", self.orig_id_key)
+            .replace("INPUT_MONGODB_NAME", self.input_mongo_db_name)
+            .replace("MONGODB_NAME", self.mongo_db_name)
+        )
+
+        ampel_job_path = cfg_path.parent / f"{cfg_path.stem}_ampel_job.yml"
+        logger.info(f"writing ampel job to {ampel_job_path}")
+        with ampel_job_path.open("w") as f:
+            f.write(ampel_job)
+
+        return ampel_job_path
+
+    def prepare(self, cfg_path: Path) -> Path:
+        self.import_input()
+        return self.make_ampel_job_file(cfg_path)
+
+    def run(self, timewise_cfg_path: Path, ampel_config_path: Path):
+        ampel_job_path = self.prepare(timewise_cfg_path)
+        cmd = JobCommand()
+        parser = cmd.get_parser()
+        args = vars(
+            parser.parse_args(
+                ["--schema", str(ampel_job_path), "--config", str(ampel_config_path)]
+            )
+        )
+        logger.debug(args)
+        cmd.run(args, unknown_args=())
+
+    @property
+    def client(self) -> MongoClient:
+        return MongoClient(self.uri)
+
+    @property
+    def db(self) -> Database:
+        return self.client[self.mongo_db_name]
+
+    @property
+    def t0(self) -> Collection:
+        return self.db["t0"]
+
+    @property
+    def t1(self) -> Collection:
+        return self.db["t1"]
+
+    @property
+    def t2(self) -> Collection:
+        return self.db["t2"]
+
+    def extract_stacked_lightcurve(self, stock_id: StockId) -> pd.DataFrame:
+        records = []
+        for i, ic in enumerate(
+            self.t2.find({"stock": stock_id, "unit": "T2StackVisits"})
+        ):
+            stock_id_str = str(stock_id)
+            assert i == 0, f"More than one stacked lightcurve found for {stock_id_str}!"
+            assert len(ic["body"]) == 1, (
+                f"None or more than one stacking result found for {stock_id_str}!"
+            )
+            records = ic["body"][0]
+        return pd.DataFrame(records)
+
+    def extract_datapoints(self, stock_id: StockId) -> pd.DataFrame:
+        records = []
+        index = []
+        for ic in self.t0.find({"stock": stock_id}):
+            records.append(ic["body"])
+            index.append(ic["id"])
+        return pd.DataFrame(records, index=index)
+
+    def extract_selected_datapoint_ids(self, stock_id: StockId) -> List[DataPointId]:
+        d = self.t1.find_one({"stock": stock_id})
+        if d is None:
+            return []
+        return d["dps"]
+
+    def export_stacked_lightcurve(self, stock_id: StockId, filename: Path):
+        logger.debug(f"Exporting stacked lightcurve for {str(stock_id)} to {filename}")
+        self.extract_stacked_lightcurve(stock_id).to_csv(filename)
+
+    def export_many(
+        self, directory: Path, stock_ids: Iterable[StockId] | StockId | None = None
+    ):
+        if stock_ids is None:
+            stock_ids = pd.read_csv(self.input_csv)[self.orig_id_key]
+        directory.mkdir(exist_ok=True, parents=True)
+        for s in np.atleast_1d(cast(npt.ArrayLike, stock_ids)):
+            self.export_stacked_lightcurve(s.item(), directory / f"{s}.csv")

timewise/process/keys.py

@@ -0,0 +1,10 @@
+MEAN = "mean"
+MEDIAN = "median"
+RMS = "rms"
+UPPER_LIMIT = "ul"
+NPOINTS = "npoints"
+ZEROPOINT_EXT = "zeropoint"
+FLUX_EXT = "flux"
+FLUX_DENSITY_EXT = "fluxdensity"
+MAG_EXT = "mpro"
+ERROR_EXT = "sig"