astro-otter 0.0.1__py3-none-any.whl → 0.1.0__py3-none-any.whl

otter/io/transient.py

@@ -3,10 +3,13 @@ Class for a transient,
 basically just inherits the dict properties with some overwriting
 """
 
+from __future__ import annotations
 import warnings
 from copy import deepcopy
 import re
 from collections.abc import MutableMapping
+from typing_extensions import Self
+import logging
 
 import numpy as np
 import pandas as pd
@@ -15,9 +18,6 @@ import astropy.units as u
 from astropy.time import Time
 from astropy.coordinates import SkyCoord
 
-from synphot.units import VEGAMAG, convert_flux
-from synphot.spectrum import SourceSpectrum
-
 from ..exceptions import (
     FailedQueryError,
     IOError,
@@ -25,10 +25,12 @@ from ..exceptions import (
     TransientMergeError,
 )
 from ..util import XRAY_AREAS
+from .host import Host
 
 warnings.simplefilter("once", RuntimeWarning)
 warnings.simplefilter("once", UserWarning)
 np.seterr(divide="ignore")
+logger = logging.getLogger(__name__)
 
 
 class Transient(MutableMapping):
@@ -37,7 +39,9 @@ class Transient(MutableMapping):
         Overwrite the dictionary init
 
         Args:
-            d [dict]: A transient dictionary
+            d (dict): A transient dictionary
+            name (str): The default name of the transient, default is None and it will
+                        be inferred from the input dictionary.
         """
         self.data = d
 
@@ -68,7 +72,7 @@ class Transient(MutableMapping):
         """
 
         if isinstance(keys, (list, tuple)):
-            return Transient({key: self[key] for key in keys})
+            return Transient({key: (self[key] if key in self else []) for key in keys})
         elif isinstance(keys, str) and "/" in keys:  # this is for a path
             s = "']['".join(keys.split("/"))
             s = "['" + s
@@ -85,6 +89,10 @@ class Transient(MutableMapping):
             return self.data[keys]
 
     def __setitem__(self, key, value):
+        """
+        Override set item to work with the '/' syntax
+        """
+
         if isinstance(key, str) and "/" in key:  # this is for a path
             s = "']['".join(key.split("/"))
             s = "['" + s
@@ -109,7 +117,7 @@ class Transient(MutableMapping):
 
     def __repr__(self, html=False):
         if not html:
-            return str(self.data)
+            return f"Transient(\n\tName: {self.default_name},\n\tKeys: {self.keys()}\n)"
         else:
             html = ""
 
@@ -188,6 +196,19 @@ class Transient(MutableMapping):
                 + " You can set strict_merge=False to override the check"
             )
 
+        # create set of the allowed keywords
+        allowed_keywords = {
+            "name",
+            "date_reference",
+            "coordinate",
+            "distance",
+            "filter_alias",
+            "schema_version",
+            "photometry",
+            "classification",
+            "host",
+        }
+
         # create a blank dictionary since we don't want to overwrite this object
         out = {}
 
@@ -222,24 +243,8 @@ class Transient(MutableMapping):
                 continue
 
             # There are some special keys that we are expecting
-            if key == "name":
-                self._merge_names(other, out)
-            elif key == "coordinate":
-                self._merge_coords(other, out)
-            elif key == "date_reference":
-                self._merge_date(other, out)
-            elif key == "distance":
-                self._merge_distance(other, out)
-            elif key == "filter_alias":
-                self._merge_filter_alias(other, out)
-            elif key == "schema_version":
-                self._merge_schema_version(other, out)
-            elif key == "photometry":
-                self._merge_photometry(other, out)
-            elif key == "spectra":
-                self._merge_spectra(other, out)
-            elif key == "classification":
-                self._merge_class(other, out)
+            if key in allowed_keywords:
+                Transient._merge_arbitrary(key, self, other, out)
             else:
                 # this is an unexpected key!
                 if strict_merge:
@@ -260,14 +265,17 @@ class Transient(MutableMapping):
         # now return out as a Transient Object
         return Transient(out)
 
-    def get_meta(self, keys=None):
+    def get_meta(self, keys=None) -> Self:
         """
         Get the metadata (no photometry or spectra)
 
         This essentially just wraps on __getitem__ but with some checks
 
         Args:
-            keys [list[str]] : list of keys
+            keys (list[str]) : list of keys to get the metadata for from the transient
+
+        Returns:
+            A Transient object of just the meta data
         """
         if keys is None:
             keys = list(self.keys())
@@ -296,9 +304,16 @@ class Transient(MutableMapping):
 
         return self[keys]
 
-    def get_skycoord(self, coord_format="icrs"):
+    def get_skycoord(self, coord_format="icrs") -> SkyCoord:
         """
         Convert the coordinates to an astropy SkyCoord
+
+        Args:
+            coord_format (str): Astropy coordinate format to convert the SkyCoord to
+                                defaults to icrs.
+
+        Returns:
+            Astropy.coordinates.SkyCoord of the default coordinate for the transient
         """
 
         # now we can generate the SkyCoord
@@ -309,9 +324,12 @@ class Transient(MutableMapping):
 
         return coord
 
-    def get_discovery_date(self):
+    def get_discovery_date(self) -> Time:
         """
-        Get the default discovery date
+        Get the default discovery date for this Transient
+
+        Returns:
+            astropy.time.Time of the default discovery date
         """
         key = "date_reference"
         date = self._get_default(key, filt='df["date_type"] == "discovery"')
@@ -320,11 +338,14 @@ class Transient(MutableMapping):
         else:
             f = "mjd"
 
-        return Time(date["value"], format=f)
+        return Time(str(date["value"]).strip(), format=f)
 
-    def get_redshift(self):
+    def get_redshift(self) -> float:
         """
-        Get the default redshift
+        Get the default redshift of this Transient
+
+        Returns:
+            Float value of the default redshift
         """
         f = "df['distance_type']=='redshift'"
         default = self._get_default("distance", filt=f)
@@ -333,7 +354,73 @@ class Transient(MutableMapping):
         else:
             return default["value"]
 
-    def _get_default(self, key, filt=""):
+    def get_classification(self) -> tuple(str, float, list):
+        """
+        Get the default classification of this Transient.
+        This normally corresponds to the highest confidence classification that we have
+        stored for the transient.
+
+        Returns:
+            The default object class as a string, the confidence level in that class,
+            and a list of the bibcodes corresponding to that classification. Or, None
+            if there is no classification.
+        """
+        default = self._get_default("classification")
+        if default is None:
+            return default
+        return default.object_class, default.confidence, default.reference
+
+    def get_host(self, max_hosts=3, **kwargs) -> list[Host]:
+        """
+        Gets the default host information of this Transient. This returns an otter.Host
+        object. If no host is known in OTTER, it uses astro-ghost to find the best
+        match.
+
+        Args:
+            max_hosts [int] : The maximum number of hosts to return
+            **kwargs : keyword arguments to be passed to getGHOST
+
+        Returns:
+            A list of otter.Host objects. This is useful becuase the Host objects have
+            useful methods for querying public catalogs for data of the host.
+        """
+        # first try to get the host information from our local database
+        if "host" in self:
+            host = [
+                Host(transient_name=self.default_name, **dict(h)) for h in self["host"]
+            ]
+
+        # then try astro-ghost
+        else:
+            logger.warn(
+                "No host known, trying to find it with astro-ghost. \
+                See https://uiucsnastro-ghost.readthedocs.io/en/latest/index.html"
+            )
+
+            # this import has to be here otherwise the code breaks
+            from astro_ghost.ghostHelperFunctions import getTransientHosts, getGHOST
+
+            getGHOST(real=False, verbose=1)
+            res = getTransientHosts(
+                [self.default_name], [self.get_skycoord()], verbose=False
+            )
+
+            host = [
+                Host(
+                    host_ra=row["raStack"],
+                    host_dec=row["decStack"],
+                    host_ra_units="deg",
+                    host_dec_units="deg",
+                    host_name=row["objName"],
+                    transient_name=self.default_name,
+                    reference=["astro-ghost"],
+                )
+                for i, row in res.iterrows()
+            ]
+
+        return host
+
+    def _get_default(self, key, filt=None):
         """
         Get the default of key
 
@@ -346,7 +433,8 @@ class Transient(MutableMapping):
             raise KeyError(f"This transient does not have {key} associated with it!")
 
         df = pd.DataFrame(self[key])
-        df = df[eval(filt)]  # apply the filters
+        if filt is not None:
+            df = df[eval(filt)]  # apply the filters
 
         if "default" in df:
             # first try to get the default
@@ -390,11 +478,37 @@ class Transient(MutableMapping):
         wave_unit: u.Unit = "nm",
         by: str = "raw",
         obs_type: str = None,
-    ):
+    ) -> pd.DataFrame:
         """
         Ensure the photometry associated with this transient is all in the same
         units/system/etc
-        """
+
+        Args:
+            flux_unit (astropy.unit.Unit): The astropy unit or string representation of
+                                           an astropy unit to convert and return the
+                                           flux as.
+            date_unit (str): Valid astropy date format string.
+            freq_unit (astropy.unit.Unit): The astropy unit or string representation of
+                                           an astropy unit to convert and return the
+                                           frequency as.
+            wave_unit (astropy.unit.Unit): The astropy unit or string representation of
+                                           an astropy unit to convert and return the
+                                           wavelength as.
+            by (str): Either 'raw' or 'value'. 'raw' is the default and is highly
+                      recommended! If 'value' is used it may skip some photometry.
+                      See the schema definition to understand this keyword completely
+                      before using it.
+            obs_type (str): "radio", "xray", or "uvoir". If provided, it only returns
+                            data taken within that range of wavelengths/frequencies.
+                            Default is None which will return all of the data.
+
+        Returns:
+            A pandas DataFrame of the cleaned up photometry in the requested units
+        """
+        # these imports need to be here for some reason
+        # otherwise the code breaks
+        from synphot.units import VEGAMAG, convert_flux
+        from synphot.spectrum import SourceSpectrum
 
         # check inputs
         if by not in {"value", "raw"}:
@@ -429,6 +543,9 @@ class Transient(MutableMapping):
             else:
                 by = "value"
 
+        # skip rows where 'by' is nan
+        df = df[df[by].notna()]
+
         # drop irrelevant obs_types before continuing
         if obs_type is not None:
             valid_obs_types = {"radio", "uvoir", "xray"}
@@ -474,8 +591,9 @@ class Transient(MutableMapping):
                 )
 
             unit = unit[0]
+            isvegamag = "vega" in unit.lower()
             try:
-                if "vega" in unit.lower():
+                if isvegamag:
                     astropy_units = VEGAMAG
                 else:
                     astropy_units = u.Unit(unit)
@@ -506,25 +624,24 @@ class Transient(MutableMapping):
                 astropy_units
             )  # assume error and values have the same unit
 
-            # get the effective wavelength
+            # get and save the effective wavelength
             if "freq_eff" in data and not np.isnan(data["freq_eff"].iloc[0]):
-                freq_units = data["freq_units"]
-                if len(np.unique(freq_units)) > 1:
-                    raise OtterLimitationError(
-                        "Can not convert different units to the same unit!"
-                    )
-
-                freq_eff = np.array(data["freq_eff"]) * u.Unit(freq_units.iloc[0])
-                wave_eff = freq_eff.to(u.AA, equivalencies=u.spectral())
+                zz = zip(data["freq_eff"], data["freq_units"])
+                freq_eff = u.Quantity([vv * u.Unit(uu) for vv, uu in zz], freq_unit)
+                wave_eff = freq_eff.to(wave_unit, equivalencies=u.spectral())
 
             elif "wave_eff" in data and not np.isnan(data["wave_eff"].iloc[0]):
-                wave_units = data["wave_units"]
-                if len(np.unique(wave_units)) > 1:
-                    raise OtterLimitationError(
-                        "Can not convert different units to the same unit!"
-                    )
+                zz = zip(data["wave_eff"], data["wave_units"])
+                wave_eff = u.Quantity([vv * u.Unit(uu) for vv, uu in zz], wave_unit)
+                freq_eff = wave_eff.to(freq_unit, equivalencies=u.spectral())
+
+            else:
+                raise ValueError("No known frequency or wavelength, please fix!")
 
-                wave_eff = np.array(data["wave_eff"]) * u.Unit(wave_units.iloc[0])
+            data["converted_wave"] = wave_eff.value
+            data["converted_wave_unit"] = wave_unit
+            data["converted_freq"] = freq_eff.value
+            data["converted_freq_unit"] = freq_unit
 
             # convert using synphot
             # stuff has to be done slightly differently for xray than for the others
@@ -544,39 +661,48 @@ class Transient(MutableMapping):
 
                 # we also need to make this wave_min and wave_max
                 # instead of just the effective wavelength like for radio and uvoir
-                wave_eff = np.array(
-                    list(zip(data["wave_min"], data["wave_max"]))
-                ) * u.Unit(wave_units.iloc[0])
+                zz = zip(data["wave_min"], data["wave_max"], data["wave_units"])
+                wave_eff = u.Quantity(
+                    [np.array([m, M]) * u.Unit(uu) for m, M, uu in zz],
+                    u.Unit(wave_unit),
+                )
 
             else:
                 area = None
 
-            # we unfortunately have to loop over the points here because
-            # syncphot does not work with a 2D array of min max wavelengths
-            # for converting counts to other flux units. It also can't convert
-            # vega mags with a wavelength array because it then interprets that as the
-            # wavelengths corresponding to the SourceSpectrum.from_vega()
-            flux, flux_err = [], []
-            for wave, xray_point, xray_point_err in zip(wave_eff, q, q_err):
-                f_val = convert_flux(
-                    wave,
-                    xray_point,
-                    u.Unit(flux_unit),
-                    vegaspec=SourceSpectrum.from_vega(),
-                    area=area,
-                )
-                f_err = convert_flux(
-                    wave,
-                    xray_point_err,
-                    u.Unit(flux_unit),
-                    vegaspec=SourceSpectrum.from_vega(),
-                    area=area,
-                )
+            if obstype == "xray" or isvegamag:
+                # we unfortunately have to loop over the points here because
+                # syncphot does not work with a 2D array of min max wavelengths
+                # for converting counts to other flux units. It also can't convert
+                # vega mags with a wavelength array because it interprets that as the
+                # wavelengths corresponding to the SourceSpectrum.from_vega()
+
+                flux, flux_err = [], []
+                for wave, xray_point, xray_point_err in zip(wave_eff, q, q_err):
+                    f_val = convert_flux(
+                        wave,
+                        xray_point,
+                        u.Unit(flux_unit),
+                        vegaspec=SourceSpectrum.from_vega(),
+                        area=area,
+                    )
+                    f_err = convert_flux(
+                        wave,
+                        xray_point_err,
+                        u.Unit(flux_unit),
+                        vegaspec=SourceSpectrum.from_vega(),
+                        area=area,
+                    )
+
+                    # then we take the average of the minimum and maximum values
+                    # computed by syncphot
+                    flux.append(np.mean(f_val).value)
+                    flux_err.append(np.mean(f_err).value)
 
-                # then we take the average of the minimum and maximum values
-                # computed by syncphot
-                flux.append(np.mean(f_val).value)
-                flux_err.append(np.mean(f_err).value)
+            else:
+                # this will be faster and cover most cases
+                flux = convert_flux(wave_eff, q, u.Unit(flux_unit))
+                flux_err = convert_flux(wave_eff, q_err, u.Unit(flux_unit))
 
             flux = np.array(flux) * u.Unit(flux_unit)
             flux_err = np.array(flux_err) * u.Unit(flux_unit)
@@ -590,7 +716,7 @@ class Transient(MutableMapping):
         outdata = pd.concat(outdata)
 
         # copy over the flux units
-        outdata["converted_flux_unit"] = [flux_unit] * len(outdata)
+        outdata["converted_flux_unit"] = flux_unit
 
         # make sure all the datetimes are in the same format here too!!
         times = [
@@ -598,27 +724,7 @@ class Transient(MutableMapping):
             for d, f in zip(outdata.date, outdata.date_format.str.lower())
         ]
         outdata["converted_date"] = times
-        outdata["converted_date_unit"] = [date_unit] * len(outdata)
-
-        # same with frequencies and wavelengths
-        freqs = []
-        waves = []
-
-        for _, row in df.iterrows():
-            if "freq_eff" in row and not np.isnan(row["freq_eff"]):
-                val = row["freq_eff"] * u.Unit(row["freq_units"])
-            elif "wave_eff" in df and not np.isnan(row["wave_eff"]):
-                val = row["wave_eff"] * u.Unit(row["wave_units"])
-            else:
-                raise ValueError("No known frequency or wavelength, please fix!")
-
-            freqs.append(val.to(freq_unit, equivalencies=u.spectral()).value)
-            waves.append(val.to(wave_unit, equivalencies=u.spectral()).value)
-
-        outdata["converted_freq"] = freqs
-        outdata["converted_wave"] = waves
-        outdata["converted_wave_unit"] = [wave_unit] * len(outdata)
-        outdata["converted_freq_unit"] = [freq_unit] * len(outdata)
+        outdata["converted_date_unit"] = date_unit
 
         return outdata
 
@@ -707,16 +813,6 @@ class Transient(MutableMapping):
         bothlines = [{"value": k, "reference": t1map[k] + t2map[k]} for k in inboth]
         out[key]["alias"] = line2 + line1 + bothlines
 
-    def _merge_coords(t1, t2, out):  # noqa: N805
-        """
-        Merge the coordinates subdictionaries for t1 and t2 and put it in out
-
-        Use pandas to drop any duplicates
-        """
-        key = "coordinate"
-
-        Transient._merge_arbitrary(key, t1, t2, out)
-
     def _merge_filter_alias(t1, t2, out):  # noqa: N805
         """
         Combine the filter alias lists across the transient objects
@@ -748,8 +844,15 @@ class Transient(MutableMapping):
         key = "photometry"
 
         out[key] = deepcopy(t1[key])
-        refs = np.array([d["reference"] for d in out[key]])
+        refs = []  # np.array([d["reference"] for d in out[key]])
         # merge_dups = lambda val: np.sum(val) if np.any(val.isna()) else val.iloc[0]
+        for val in out[key]:
+            if isinstance(val, list):
+                refs += val
+            elif isinstance(val, np.ndarray):
+                refs += list(val)
+            else:
+                refs.append(val)
 
         for val in t2[key]:
             # first check if t2's reference is in out
@@ -774,12 +877,6 @@ class Transient(MutableMapping):
 
                 out[key][i1] = newdict  # replace the dictionary at i1 with the new dict
 
-    def _merge_spectra(t1, t2, out):  # noqa: N805
-        """
-        Combine spectra sources
-        """
-        pass
-
     def _merge_class(t1, t2, out):  # noqa: N805
         """
         Combine the classification attribute
@@ -815,22 +912,6 @@ class Transient(MutableMapping):
             else:
                 item["default"] = False
 
-    def _merge_date(t1, t2, out):  # noqa: N805
-        """
-        Combine epoch data across two transients and write it to "out"
-        """
-        key = "date_reference"
-
-        Transient._merge_arbitrary(key, t1, t2, out)
-
-    def _merge_distance(t1, t2, out):  # noqa: N805
-        """
-        Combine distance information for these two transients
-        """
-        key = "distance"
-
-        Transient._merge_arbitrary(key, t1, t2, out)
-
     @staticmethod
     def _merge_arbitrary(key, t1, t2, out):
         """
@@ -840,24 +921,36 @@ class Transient(MutableMapping):
         a NxM pandas dataframe!
         """
 
-        df1 = pd.DataFrame(t1[key])
-        df2 = pd.DataFrame(t2[key])
+        if key == "name":
+            t1._merge_names(t2, out)
+        elif key == "filter_alias":
+            t1._merge_filter_alias(t2, out)
+        elif key == "schema_version":
+            t1._merge_schema_version(t2, out)
+        elif key == "photometry":
+            t1._merge_photometry(t2, out)
+        elif key == "classification":
+            t1._merge_class(t2, out)
+        else:
+            # this is where we can standardize some of the merging
+            df1 = pd.DataFrame(t1[key])
+            df2 = pd.DataFrame(t2[key])
 
-        merged_with_dups = pd.concat([df1, df2]).reset_index(drop=True)
+            merged_with_dups = pd.concat([df1, df2]).reset_index(drop=True)
 
-        # have to get the indexes to drop using a string rep of the df
-        # this is cause we have lists in some cells
-        to_drop = merged_with_dups.astype(str).drop_duplicates().index
+            # have to get the indexes to drop using a string rep of the df
+            # this is cause we have lists in some cells
+            to_drop = merged_with_dups.astype(str).drop_duplicates().index
 
-        merged = merged_with_dups.iloc[to_drop].reset_index(drop=True)
+            merged = merged_with_dups.iloc[to_drop].reset_index(drop=True)
 
-        outdict = merged.to_dict(orient="records")
+            outdict = merged.to_dict(orient="records")
 
-        outdict_cleaned = Transient._remove_nans(
-            outdict
-        )  # clear out the nans from pandas conversion
+            outdict_cleaned = Transient._remove_nans(
+                outdict
+            )  # clear out the nans from pandas conversion
 
-        out[key] = outdict_cleaned
+            out[key] = outdict_cleaned
 
     @staticmethod
     def _remove_nans(d):

otter/plotter/otter_plotter.py

@@ -6,19 +6,21 @@ Currently supported backends are:
 - plotly
 """
 
+from __future__ import annotations
 import importlib
 
 
 class OtterPlotter:
-    def __init__(self, backend):
-        """
-        Handles the backend for the "plotter" module
+    """
+    Handles the backend for the "plotter" module
 
-        Args:
-            backend [string]: a string of the module name to import and use
-                              as the backend. Currently supported are "matplotlib",
-                              "matplotlib.pyplot", "plotly", and "plotly.graph_objects"
-        """
+    Args:
+        backend (string): a string of the module name to import and use
+                          as the backend. Currently supported are "matplotlib",
+                          "matplotlib.pyplot", "plotly", and "plotly.graph_objects"
+    """
+
+    def __init__(self, backend):
         if backend == "matplotlib.pyplot":
             self.backend = backend
         elif backend == "pyplot.graph_objects":