4708907600

Committed 15 Apr 2023 05:15PM UTC coverage: 48.357%. First build

Build Type

github

Committed by GitHub

Commit Message

Merge 6dfe698e9 into cc697a037

Pull Request Pull Request #98: Update workflows, add Docker compose recipe, and add linting with pre-commit

Run Details

956 of 956 new or added lines in 82 files covered. (100.0%)

3091 of 6392 relevant lines covered (48.36%)

0.48 hits per line

Source File
Press 'n' to go to next uncovered line, 'b' for previous

42.17

/thor/orbits/orbits.py

import ast
import logging
from typing import Any, Tuple

import numpy as np
import pandas as pd
from astropy import units as u
from astropy.time import Time

from ..utils import _checkTime, getHorizonsVectors
from .kepler import convertOrbitalElements

CARTESIAN_COLS = ["x", "y", "z", "vx", "vy", "vz"]
CARTESIAN_UNITS = [u.au, u.au, u.au, u.au / u.d, u.au / u.d, u.au / u.d]
KEPLERIAN_COLS = ["a", "e", "i", "omega", "w", "M"]
KEPLERIAN_UNITS = [u.au, u.dimensionless_unscaled, u.deg, u.deg, u.deg, u.deg]

logger = logging.getLogger(__name__)

__all__ = ["Orbits"]


def _convertOrbitUnits(orbits, orbit_units, desired_units):
    orbits_ = orbits.copy()
    for i, (unit_desired, unit_given) in enumerate(zip(desired_units, orbit_units)):
        if unit_desired != unit_given:
            orbits_[:, i] = orbits_[:, i] * unit_given.to(unit_desired)

    return orbits_


def _convertCovarianceUnits(
    covariances,
    orbit_units,
    desired_units,
):
    orbit_units_ = np.array([orbit_units])
    covariance_units = np.dot(orbit_units_.T, orbit_units_)

    desired_units_ = np.array([desired_units])
    desired_units_ = np.dot(desired_units_.T, desired_units_)

    covariances_ = np.stack(covariances)
    for i in range(6):
        for j in range(6):
            unit_desired = desired_units_[i, j]
            unit_given = covariance_units[i, j]
            if unit_desired != unit_given:
                covariances_[:, i, j] = covariances_[:, i, j] * unit_given.to(
                    unit_desired
                )

    covariances_ = np.split(covariances_, covariances_.shape[0], axis=0)
    covariances_ = [cov[0] for cov in covariances_]
    return covariances_


class Orbits:
    """
    Orbits Class



    """

    def __init__(
        self,
        orbits,
        epochs,
        orbit_type="cartesian",
        orbit_units=CARTESIAN_UNITS,
        covariance=None,
        ids=None,
        H=None,
        G=None,
        orbit_class=None,
        additional_data=None,
    ):
        """
        Class to store orbits and a variety of other useful attributes and
        data required to:
            propagate orbits
            generate ephemerides
            transform orbital elements

        """

        # Make sure that the given epoch(s) are an astropy time object
        if len(epochs) > 0:
            _checkTime(epochs, "epoch")

        # Make sure that each orbit has a an epoch
        assert len(epochs) == orbits.shape[0]
        self._epochs = epochs
        self.num_orbits = orbits.shape[0]

        # Make sure the passed orbits are one of the the supported types
        self._cartesian = None
        self._keplerian = None

        if orbit_type == "cartesian":
            if orbit_units != CARTESIAN_UNITS:
                orbits_ = _convertOrbitUnits(orbits, orbit_units, CARTESIAN_UNITS)
            else:
                orbits_ = orbits.copy()

            self._cartesian = orbits_

        elif orbit_type == "keplerian":
            if orbit_units != KEPLERIAN_UNITS:
                orbits = _convertOrbitUnits(orbits, orbit_units, KEPLERIAN_UNITS)
            else:
                orbits_ = orbits.copy()

            self._keplerian = orbits_

        else:
            err = "orbit_type has to be one of {'cartesian', 'keplerian'}."
            raise ValueError(err)

        self.orbit_type = orbit_type

        # If object IDs have been passed make sure each orbit has one
        if ids is not None:
            assert len(ids) == self.num_orbits
            self.ids = np.asarray(ids, dtype="<U60")
        else:
            self.ids = np.array(
                ["{}".format(i) for i in range(self.num_orbits)], dtype="<U60"
            )

        # If H magnitudes have been passed make sure each orbit has one
        if H is not None:
            assert len(H) == self.num_orbits
            self.H = np.asarray(H)
        else:
            self.H = None

        # If the slope parameter G has been passed make sure each orbit has one
        if G is not None:
            assert len(G) == self.num_orbits
            self.G = np.asarray(G)
        else:
            self.G = None

        # If covariances matrixes have been passed make sure each orbit has one
        self._cartesian_covariance = None
        self._keplerian_covariance = None
        if covariance is not None:
            assert len(covariance) == self.num_orbits
            if orbit_type == "cartesian":
                if orbit_units != CARTESIAN_UNITS:
                    covariance_ = _convertCovarianceUnits(
                        covariance,
                        orbit_units,
                        CARTESIAN_UNITS,
                    )
                else:
                    covariance_ = covariance

                self._cartesian_covariance = covariance_

            elif orbit_type == "keplerian":
                if orbit_units != KEPLERIAN_UNITS:
                    covariance_ = _convertCovarianceUnits(
                        covariance,
                        orbit_units,
                        KEPLERIAN_UNITS,
                    )
                else:
                    covariance_ = covariance

                self._keplerian_covariance = covariance_

            else:
                pass

        if orbit_class is not None:
            assert len(orbit_class) == self.num_orbits
            self.orbit_class = np.asarray(orbit_class)
        else:
            self.orbit_class = None

        if additional_data is not None:
            if isinstance(additional_data, pd.DataFrame):
                assert len(additional_data) == self.num_orbits
                self.additional_data = additional_data
            else:
                raise TypeError("additional_data must be a pandas DataFrame")
        else:
            self.additional_data = None

        return

    def __repr__(self):
        rep = "Orbits: {}\n"
        return rep.format(self.num_orbits)

    def __len__(self):
        return self.num_orbits

    def __getitem__(self, i):

        # Convert integer index to a slice so that array
        # based properties are not inappropriately reshaped
        if isinstance(i, int):
            i = slice(i, i + 1)

        # Thrown an index error when the index is out of
        # range
        if i.start >= self.num_orbits:
            raise IndexError

        # Extract the relevant epochs
        _epochs = self.__dict__["_epochs"][i]

        # Extract the relevant orbital elements
        args = ()
        covariance = None
        if self.__dict__["_cartesian"] is not None and len(args) == 0:
            _cartesian = self.__dict__["_cartesian"][i]
            args = (_cartesian, _epochs)
            orbit_type = "cartesian"
            if self._cartesian_covariance is not None:
                covariance = self._cartesian_covariance[i]
        else:
            _cartesian = None

        if self.__dict__["_keplerian"] is not None and len(args) == 0:
            _keplerian = self.__dict__["_keplerian"][i]
            args = (_keplerian, _epochs)
            orbit_type = "_keplerian"
            if self._keplerian_covariance is not None:
                covariance = self._keplerian_covariance[i]
        else:
            _keplerian = None

        kwargs = {
            "orbit_type": orbit_type,
            "covariance": covariance,
        }
        for kwarg in ["ids", "H", "G", "orbit_class", "additional_data"]:
            if isinstance(self.__dict__[kwarg], np.ndarray):
                kwargs[kwarg] = self.__dict__[kwarg][i]
            elif isinstance(self.__dict__[kwarg], pd.DataFrame):
                kwargs[kwarg] = self.__dict__[kwarg].iloc[i]
                kwargs[kwarg].reset_index(inplace=True, drop=True)
            else:
                kwargs[kwarg] = self.__dict__[kwarg]

        orbit = Orbits(*args, **kwargs)
        orbit._cartesian = _cartesian
        orbit._keplerian = _keplerian
        orbit._cartesian_covariance = self._cartesian_covariance
        orbit._keplerian_covariance = self._keplerian_covariance
        return orbit

    def __eq__(self, other):

        eq = True
        while eq:
            if len(self) != len(other):
                return False

            if not np.all(
                np.isclose(self.cartesian, other.cartesian, rtol=1e-15, atol=1e-15)
            ):
                return False

            if not np.all(
                np.isclose(
                    self.epochs.tdb.mjd, other.epochs.tdb.mjd, rtol=1e-15, atol=1e-15
                )
            ):
                return False

            if not np.all(self.ids == other.ids):
                return False

            break

        return eq

    @property
    def epochs(self):
        return self._epochs

    @property
    def cartesian(self):
        if not isinstance(self._cartesian, np.ndarray):
            logger.debug(
                "Cartesian elements are not defined. Converting Keplerian elements to Cartesian."
            )
            self._cartesian = convertOrbitalElements(
                self._keplerian,
                "keplerian",
                "cartesian",
            )
        return self._cartesian

    @property
    def keplerian(self):
        if not isinstance(self._keplerian, np.ndarray):
            logger.debug(
                "Keplerian elements are not defined. Converting Cartesian elements to Keplerian."
            )
            self._keplerian = convertOrbitalElements(
                self._cartesian,
                "cartesian",
                "keplerian",
            )
        return self._keplerian

    @property
    def cartesian_covariance(self):
        return self._cartesian_covariance

    @property
    def keplerian_covariance(self):
        return self._keplerian_covariance

    def split(self, chunk_size: int) -> list:
        """
        Split orbits into new orbit classes of size chunk_size.

        Parameters
        ----------
        chunk_size : int
            Size of each chunk of orbits.

        Returns
        -------
        list of Orbits
        """
        objs = []
        for chunk in range(0, self.num_orbits, chunk_size):
            objs.append(self[chunk : chunk + chunk_size])
        return objs

    def assignOrbitClasses(self):

        a = self.keplerian[:, 0]
        e = self.keplerian[:, 1]
        i = self.keplerian[:, 2]
        q = a * (1 - e)
        p = Q = a * (1 + e)

        classes = np.array(["AST" for i in range(len(self.keplerian))])

        classes_dict = {
            "AMO": np.where((a > 1.0) & (q > 1.017) & (q < 1.3)),
            "MBA": np.where((a > 1.0) & (q < 1.017)),
            "ATE": np.where((a < 1.0) & (Q > 0.983)),
            "CEN": np.where((a > 5.5) & (a < 30.3)),
            "IEO": np.where((Q < 0.983))[0],
            "IMB": np.where((a < 2.0) & (q > 1.666))[0],
            "MBA": np.where((a > 2.0) & (a < 3.2) & (q > 1.666)),
            "MCA": np.where((a < 3.2) & (q > 1.3) & (q < 1.666)),
            "OMB": np.where((a > 3.2) & (a < 4.6)),
            "TJN": np.where((a > 4.6) & (a < 5.5) & (e < 0.3)),
            "TNO": np.where((a > 30.1)),
            "PAA": np.where((e == 1)),
            "HYA": np.where((e > 1)),
        }
        for c, v in classes_dict.items():
            classes[v] = c

        self.orbit_class = classes
        return

    @staticmethod
    def fromHorizons(obj_ids, t0):
        """
        Query Horizons for state vectors for each object ID at time t0.
        This is a convenience function and should not be used to query for state
        vectors for many objects or at many times.

        Parameters
        ----------
        obj_ids : `~numpy.ndarray` (N)
            Object IDs / designations recognizable by HORIZONS.
        t0 : `~astropy.core.time.Time` (1)
            Astropy time object at which to gather state vectors.

        Return
        ------
        `~thor.orbits.orbits.Orbits`
            THOR Orbits class
        """

        if len(t0) != 1:
            err = "t0 should be a single time."
            raise ValueError(err)

        horizons_vectors = getHorizonsVectors(
            obj_ids, t0, location="@sun", id_type="smallbody", aberrations="geometric"
        )

        orbits = Orbits(
            horizons_vectors[["x", "y", "z", "vx", "vy", "vz"]].values,
            t0 + np.zeros(len(obj_ids)),
            orbit_type="cartesian",
            orbit_units=CARTESIAN_UNITS,
            ids=horizons_vectors["targetname"].values,
            H=horizons_vectors["H"].values,
            G=horizons_vectors["G"].values,
        )
        return orbits

    @staticmethod
    def fromMPCOrbitCatalog(mpcorb):

        cols = ["a_au", "e", "i_deg", "ascNode_deg", "argPeri_deg", "meanAnom_deg"]
        additional_cols = mpcorb.columns[
            ~mpcorb.columns.isin(cols + ["mjd_tt", "designation", "H_mag", "G"])
        ]
        orbits = mpcorb[cols].values
        epochs = Time(mpcorb["mjd_tt"].values, scale="tt", format="mjd")
        args = (orbits, epochs)

        kwargs = {
            "orbit_type": "keplerian",
            "ids": mpcorb["designation"].values,
            "H": mpcorb["H_mag"].values,
            "G": mpcorb["G"].values,
            "orbit_units": KEPLERIAN_UNITS,
            "additional_data": mpcorb[additional_cols],
        }
        return Orbits(*args, **kwargs)

    def to_df(
        self,
        include_units: bool = True,
        include_keplerian: bool = False,
        include_cartesian: bool = True,
    ) -> pd.DataFrame:
        """
        Convert orbits into a pandas dataframe.

        Returns
        -------
        dataframe : `~pandas.DataFrame`
            Dataframe containing orbits, epochs and IDs. If H, G, and covariances are defined then
            those are also added to the dataframe.
        """
        if self.num_orbits > 0:
            data = {"orbit_id": self.ids, "mjd_tdb": self.epochs.tdb.mjd}
        else:
            data = {
                "orbit_id": [],
                "mjd_tdb": [],
            }

        if include_units:

            units_index = ["--", "mjd [TDB]"]
            data["epoch"] = data["mjd_tdb"]
            data.pop("mjd_tdb")

        if include_cartesian:
            for i in range(6):
                data[CARTESIAN_COLS[i]] = self.cartesian[:, i]
            if include_units:
                orbit_units_str = []
                for unit in CARTESIAN_UNITS:
                    orbit_units_str.append(str(unit).lower())

                units_index += orbit_units_str

        if include_keplerian:
            for i in range(6):
                data[KEPLERIAN_COLS[i]] = self.keplerian[:, i]

            if include_units:
                orbit_units_str = []
                for unit in KEPLERIAN_UNITS:
                    if unit == u.dimensionless_unscaled:
                        orbit_units_str.append("--")
                    else:
                        orbit_units_str.append(str(unit).lower())
                units_index += orbit_units_str

        if self._cartesian_covariance is not None and include_cartesian:
            data["covariance"] = self.cartesian_covariance
            if include_units:
                units_index.append("--")

        if self._keplerian_covariance is not None and include_keplerian:
            data["covariance"] = self.keplerian_covariance
            if include_units:
                units_index.append("--")

        if self.H is not None:
            data["H"] = self.H
            if include_units:
                units_index.append("mag")

        if self.G is not None:
            data["G"] = self.G
            if include_units:
                units_index.append("--")

        if self.orbit_class is not None:
            data["orbit_class"] = self.orbit_class
            if include_units:
                units_index.append("--")

        dataframe = pd.DataFrame(data)
        if self.additional_data is not None:
            dataframe = dataframe.join(self.additional_data)
            if include_units:
                for col in self.additional_data.columns:
                    units_index.append("--")

        if include_units:
            dataframe.columns = pd.MultiIndex.from_arrays(
                [dataframe.columns, np.array(units_index)]
            )

        return dataframe

    @staticmethod
    def from_df(dataframe: pd.DataFrame):
        """
        Read orbits from a dataframe. Required columns are
        the epoch at which the orbits are defined ('mjd_tdb') and the 6 dimensional state.
        If the states are cartesian then the expected columns are ('x', 'y', 'z', 'vx', 'vy', 'vz'),
        if the states are keplerian then the expected columns are ("a", "e", "i", "raan", "argperi", "M").

        Parameters
        ----------
        dataframe : `~pandas.DataFrame`
            Dataframe containing either Cartesian or Keplerian orbits.

        Returns
        -------
        `~thor.orbits.orbits.Orbits`
            THOR Orbit class with the orbits read from the input dataframe.
        """
        # Extract the epochs from the given dataframe
        dataframe_ = dataframe.copy()

        if isinstance(dataframe.columns, pd.MultiIndex):
            dataframe_.columns = dataframe_.columns.droplevel(level=1)
            dataframe_.rename(columns={"epoch": "mjd_tdb"}, inplace=True)

        if len(dataframe_) > 0:
            epochs = Time(dataframe_["mjd_tdb"].values, format="mjd", scale="tdb")
        else:
            epochs = np.array([])

        # If the dataframe's index is not sorted and increasing, reset it
        if not np.all(dataframe_.index.values == np.arange(0, len(dataframe_))):
            dataframe_.reset_index(inplace=True, drop=True)

        columns_required = ["orbit_id", "mjd_tdb"]
        cartesian = None
        keplerian = None
        if np.all(pd.Series(CARTESIAN_COLS).isin(dataframe_.columns)):
            columns_required += CARTESIAN_COLS
            cartesian = dataframe_[CARTESIAN_COLS].values
            args = (cartesian, epochs)
            orbit_type = "cartesian"
        elif np.all(pd.Series(KEPLERIAN_COLS).isin(dataframe_.columns)):
            columns_required += KEPLERIAN_COLS
            keplerian = dataframe_[KEPLERIAN_COLS].values
            args = (keplerian, epochs)
            orbit_type = "keplerian"

        kwargs = {"ids": dataframe_["orbit_id"].values, "orbit_type": orbit_type}

        columns_optional = ["covariance", "H", "G", "orbit_class"]
        for col in columns_optional:
            if col in dataframe_.columns:
                kwargs[col] = dataframe_[col].values

        columns = columns_required + columns_optional
        if len(dataframe_.columns[~dataframe_.columns.isin(columns)]) > 0:
            kwargs["additional_data"] = dataframe_[
                dataframe_.columns[~dataframe_.columns.isin(columns)]
            ]

        orbits = Orbits(*args, **kwargs)
        orbits._keplerian = keplerian
        orbits._cartesian = cartesian
        # orbits._keplerian_covariance = keplerian_covariance
        # orbits._cartesian_covariance = cartesian_covariance
        return orbits

    def to_csv(
        self, file: str, include_cartesian: bool = True, include_keplerian: bool = False
    ):
        """
        Save orbits to a csv. Orbits are always saved with the
        units of each quantity to avoid ambiguity and confusion.

        Parameters
        ----------
        file : str
            Name or path of file including extension to which to save
            orbits.
        include_cartesian : bool
            Save Cartesian elements.
        include_keplerian : bool
            Save Keplerian elements.
        """
        df = self.to_df(
            include_units=True,
            include_cartesian=include_cartesian,
            include_keplerian=include_keplerian,
        )
        df.to_csv(file, index=False, float_format="%.15e", encoding="utf-8")
        return

    @staticmethod
    def from_csv(file: str):
        """
        Read orbits from a csv.

        Parameters
        ----------
        file : str
            Name or path of file including extension to which to read
            orbits.
        """
        df = pd.read_csv(
            file,
            index_col=None,
            header=[0, 1],
            converters={
                "cartesian_covariance": lambda x: np.array(
                    ast.literal_eval(",".join(x.replace("[ ", "[").split()))
                ),
                "keplerian_covariance": lambda x: np.array(
                    ast.literal_eval(",".join(x.replace("[ ", "[").split()))
                ),
            },
            dtype={
                "orbit_id": str,
                "test_orbit_id": str,
            },
            float_precision="round_trip",
        )
        df.rename(columns={"epoch": "mjd_tdb"}, inplace=True)
        return Orbits.from_df(df)

    def to_hdf(
        self, file: str, include_cartesian: bool = True, include_keplerian: bool = False
    ):
        """
        Save orbits to a HDF5 file. Orbits are always saved with the
        units of each quantity to avoid ambiguity and confusion.

        Parameters
        ----------
        file : str
            Name or path of file including extension to which to save
            orbits.
        include_cartesian : bool
            Save Cartesian elements.
        include_keplerian : bool
            Save Keplerian elements.
        """
        df = self.to_df(
            include_units=True,
            include_cartesian=include_cartesian,
            include_keplerian=include_keplerian,
        )
        df.to_hdf(file, key="data")
        return

    @staticmethod
    def from_hdf(file: str):
        """
        Read orbits from a HDF5 file.

        Parameters
        ----------
        file : str
            Name or path of file including extension to which to read
            orbits.
        """
        df = pd.read_hdf(file, key="data")
        return Orbits.from_df(df)

1	import ast	1✔
2	import logging	1✔
3	from typing import Any, Tuple	1✔
4
5	import numpy as np	1✔
6	import pandas as pd	1✔
7	from astropy import units as u	1✔
8	from astropy.time import Time	1✔
9
10	from ..utils import _checkTime, getHorizonsVectors	1✔
11	from .kepler import convertOrbitalElements	1✔
12
13	CARTESIAN_COLS = ["x", "y", "z", "vx", "vy", "vz"]	1✔
14	CARTESIAN_UNITS = [u.au, u.au, u.au, u.au / u.d, u.au / u.d, u.au / u.d]	1✔
15	KEPLERIAN_COLS = ["a", "e", "i", "omega", "w", "M"]	1✔
16	KEPLERIAN_UNITS = [u.au, u.dimensionless_unscaled, u.deg, u.deg, u.deg, u.deg]	1✔
17
18	logger = logging.getLogger(__name__)	1✔
19
20	__all__ = ["Orbits"]	1✔
21
22
23	def _convertOrbitUnits(orbits, orbit_units, desired_units):	1✔
24	orbits_ = orbits.copy()	×
25	for i, (unit_desired, unit_given) in enumerate(zip(desired_units, orbit_units)):	×
26	if unit_desired != unit_given:	×
27	orbits_[:, i] = orbits_[:, i] * unit_given.to(unit_desired)	×
28
29	return orbits_	×
30
31
32	def _convertCovarianceUnits(	1✔
33	covariances,
34	orbit_units,
35	desired_units,
36	):
37	orbit_units_ = np.array([orbit_units])	×
38	covariance_units = np.dot(orbit_units_.T, orbit_units_)	×
39
40	desired_units_ = np.array([desired_units])	×
41	desired_units_ = np.dot(desired_units_.T, desired_units_)	×
42
43	covariances_ = np.stack(covariances)	×
44	for i in range(6):	×
45	for j in range(6):	×
46	unit_desired = desired_units_[i, j]	×
47	unit_given = covariance_units[i, j]	×
48	if unit_desired != unit_given:	×
49	covariances_[:, i, j] = covariances_[:, i, j] * unit_given.to(	×
50	unit_desired
51	)
52
53	covariances_ = np.split(covariances_, covariances_.shape[0], axis=0)	×
54	covariances_ = [cov[0] for cov in covariances_]	×
55	return covariances_	×
56
57
58	class Orbits:	1✔
59	"""
60	Orbits Class
61
62
63
64	"""
65
66	def __init__(	1✔
67	self,
68	orbits,
69	epochs,
70	orbit_type="cartesian",
71	orbit_units=CARTESIAN_UNITS,
72	covariance=None,
73	ids=None,
74	H=None,
75	G=None,
76	orbit_class=None,
77	additional_data=None,
78	):
79	"""
80	Class to store orbits and a variety of other useful attributes and
81	data required to:
82	propagate orbits
83	generate ephemerides
84	transform orbital elements
85
86	"""
87
88	# Make sure that the given epoch(s) are an astropy time object
89	if len(epochs) > 0:	1✔
90	_checkTime(epochs, "epoch")	1✔
91
92	# Make sure that each orbit has a an epoch
93	assert len(epochs) == orbits.shape[0]	1✔
94	self._epochs = epochs	1✔
95	self.num_orbits = orbits.shape[0]	1✔
96
97	# Make sure the passed orbits are one of the the supported types
98	self._cartesian = None	1✔
99	self._keplerian = None	1✔
100
101	if orbit_type == "cartesian":	1✔
102	if orbit_units != CARTESIAN_UNITS:	1✔
103	orbits_ = _convertOrbitUnits(orbits, orbit_units, CARTESIAN_UNITS)	×
104	else:
105	orbits_ = orbits.copy()	1✔
106
107	self._cartesian = orbits_	1✔
108
109	elif orbit_type == "keplerian":	×
110	if orbit_units != KEPLERIAN_UNITS:	×
111	orbits = _convertOrbitUnits(orbits, orbit_units, KEPLERIAN_UNITS)	×
112	else:
113	orbits_ = orbits.copy()	×
114
115	self._keplerian = orbits_	×
116
117	else:
118	err = "orbit_type has to be one of {'cartesian', 'keplerian'}."	×
119	raise ValueError(err)	×
120
121	self.orbit_type = orbit_type	1✔
122
123	# If object IDs have been passed make sure each orbit has one
124	if ids is not None:	1✔
125	assert len(ids) == self.num_orbits	1✔
126	self.ids = np.asarray(ids, dtype="<U60")	1✔
127	else:
128	self.ids = np.array(	1✔
129	["{}".format(i) for i in range(self.num_orbits)], dtype="<U60"
130	)
131
132	# If H magnitudes have been passed make sure each orbit has one
133	if H is not None:	1✔
134	assert len(H) == self.num_orbits	×
135	self.H = np.asarray(H)	×
136	else:
137	self.H = None	1✔
138
139	# If the slope parameter G has been passed make sure each orbit has one
140	if G is not None:	1✔
141	assert len(G) == self.num_orbits	×
142	self.G = np.asarray(G)	×
143	else:
144	self.G = None	1✔
145
146	# If covariances matrixes have been passed make sure each orbit has one
147	self._cartesian_covariance = None	1✔
148	self._keplerian_covariance = None	1✔
149	if covariance is not None:	1✔
150	assert len(covariance) == self.num_orbits	×
151	if orbit_type == "cartesian":	×
152	if orbit_units != CARTESIAN_UNITS:	×
153	covariance_ = _convertCovarianceUnits(	×
154	covariance,
155	orbit_units,
156	CARTESIAN_UNITS,
157	)
158	else:
159	covariance_ = covariance	×
160
161	self._cartesian_covariance = covariance_	×
162
163	elif orbit_type == "keplerian":	×
164	if orbit_units != KEPLERIAN_UNITS:	×
165	covariance_ = _convertCovarianceUnits(	×
166	covariance,
167	orbit_units,
168	KEPLERIAN_UNITS,
169	)
170	else:
171	covariance_ = covariance	×
172
173	self._keplerian_covariance = covariance_	×
174
175	else:
176	pass	×
177
178	if orbit_class is not None:	1✔
179	assert len(orbit_class) == self.num_orbits	×
180	self.orbit_class = np.asarray(orbit_class)	×
181	else:
182	self.orbit_class = None	1✔
183
184	if additional_data is not None:	1✔
185	if isinstance(additional_data, pd.DataFrame):	×
186	assert len(additional_data) == self.num_orbits	×
187	self.additional_data = additional_data	×
188	else:
189	raise TypeError("additional_data must be a pandas DataFrame")	×
190	else:
191	self.additional_data = None	1✔
192
193	return	1✔
194
195	def __repr__(self):	1✔
196	rep = "Orbits: {}\n"	×
197	return rep.format(self.num_orbits)	×
198
199	def __len__(self):	1✔
200	return self.num_orbits	1✔
201
202	def __getitem__(self, i):	1✔
203
204	# Convert integer index to a slice so that array
205	# based properties are not inappropriately reshaped
206	if isinstance(i, int):	1✔
207	i = slice(i, i + 1)	×
208
209	# Thrown an index error when the index is out of
210	# range
211	if i.start >= self.num_orbits:	1✔
212	raise IndexError	×
213
214	# Extract the relevant epochs
215	_epochs = self.__dict__["_epochs"][i]	1✔
216
217	# Extract the relevant orbital elements
218	args = ()	1✔
219	covariance = None	1✔
220	if self.__dict__["_cartesian"] is not None and len(args) == 0:	1✔
221	_cartesian = self.__dict__["_cartesian"][i]	1✔
222	args = (_cartesian, _epochs)	1✔
223	orbit_type = "cartesian"	1✔
224	if self._cartesian_covariance is not None:	1✔
225	covariance = self._cartesian_covariance[i]	×
226	else:
227	_cartesian = None	×
228
229	if self.__dict__["_keplerian"] is not None and len(args) == 0:	1✔
230	_keplerian = self.__dict__["_keplerian"][i]	×
231	args = (_keplerian, _epochs)	×
232	orbit_type = "_keplerian"	×
233	if self._keplerian_covariance is not None:	×
234	covariance = self._keplerian_covariance[i]	×
235	else:
236	_keplerian = None	1✔
237
238	kwargs = {	1✔
239	"orbit_type": orbit_type,
240	"covariance": covariance,
241	}
242	for kwarg in ["ids", "H", "G", "orbit_class", "additional_data"]:	1✔
243	if isinstance(self.__dict__[kwarg], np.ndarray):	1✔
244	kwargs[kwarg] = self.__dict__[kwarg][i]	1✔
245	elif isinstance(self.__dict__[kwarg], pd.DataFrame):	1✔
246	kwargs[kwarg] = self.__dict__[kwarg].iloc[i]	×
247	kwargs[kwarg].reset_index(inplace=True, drop=True)	×
248	else:
249	kwargs[kwarg] = self.__dict__[kwarg]	1✔
250
251	orbit = Orbits(args, *kwargs)	1✔
252	orbit._cartesian = _cartesian	1✔
253	orbit._keplerian = _keplerian	1✔
254	orbit._cartesian_covariance = self._cartesian_covariance	1✔
255	orbit._keplerian_covariance = self._keplerian_covariance	1✔
256	return orbit	1✔
257
258	def __eq__(self, other):	1✔
259
260	eq = True	×
261	while eq:	×
262	if len(self) != len(other):	×
263	return False	×
264
265	if not np.all(	×
266	np.isclose(self.cartesian, other.cartesian, rtol=1e-15, atol=1e-15)
267	):
268	return False	×
269
270	if not np.all(	×
271	np.isclose(
272	self.epochs.tdb.mjd, other.epochs.tdb.mjd, rtol=1e-15, atol=1e-15
273	)
274	):
275	return False	×
276
277	if not np.all(self.ids == other.ids):	×
278	return False	×
279
280	break	×
281
282	return eq	×
283
284	@property	1✔
285	def epochs(self):	1✔
286	return self._epochs	1✔
287
288	@property	1✔
289	def cartesian(self):	1✔
290	if not isinstance(self._cartesian, np.ndarray):	1✔
291	logger.debug(	×
292	"Cartesian elements are not defined. Converting Keplerian elements to Cartesian."
293	)
294	self._cartesian = convertOrbitalElements(	×
295	self._keplerian,
296	"keplerian",
297	"cartesian",
298	)
299	return self._cartesian	1✔
300
301	@property	1✔
302	def keplerian(self):	1✔
303	if not isinstance(self._keplerian, np.ndarray):	×
304	logger.debug(	×
305	"Keplerian elements are not defined. Converting Cartesian elements to Keplerian."
306	)
307	self._keplerian = convertOrbitalElements(	×
308	self._cartesian,
309	"cartesian",
310	"keplerian",
311	)
312	return self._keplerian	×
313
314	@property	1✔
315	def cartesian_covariance(self):	1✔
316	return self._cartesian_covariance	×
317
318	@property	1✔
319	def keplerian_covariance(self):	1✔
320	return self._keplerian_covariance	×
321
322	def split(self, chunk_size: int) -> list:	1✔
323	"""
324	Split orbits into new orbit classes of size chunk_size.
325
326	Parameters
327	----------
328	chunk_size : int
329	Size of each chunk of orbits.
330
331	Returns
332	-------
333	list of Orbits
334	"""
335	objs = []	×
336	for chunk in range(0, self.num_orbits, chunk_size):	×
337	objs.append(self[chunk : chunk + chunk_size])	×
338	return objs	×
339
340	def assignOrbitClasses(self):	1✔
341
342	a = self.keplerian[:, 0]	×
343	e = self.keplerian[:, 1]	×
344	i = self.keplerian[:, 2]	×
345	q = a * (1 - e)	×
346	p = Q = a * (1 + e)	×
347
348	classes = np.array(["AST" for i in range(len(self.keplerian))])	×
349
350	classes_dict = {	×
351	"AMO": np.where((a > 1.0) & (q > 1.017) & (q < 1.3)),
352	"MBA": np.where((a > 1.0) & (q < 1.017)),
353	"ATE": np.where((a < 1.0) & (Q > 0.983)),
354	"CEN": np.where((a > 5.5) & (a < 30.3)),
355	"IEO": np.where((Q < 0.983))[0],
356	"IMB": np.where((a < 2.0) & (q > 1.666))[0],
357	"MBA": np.where((a > 2.0) & (a < 3.2) & (q > 1.666)),
358	"MCA": np.where((a < 3.2) & (q > 1.3) & (q < 1.666)),
359	"OMB": np.where((a > 3.2) & (a < 4.6)),
360	"TJN": np.where((a > 4.6) & (a < 5.5) & (e < 0.3)),
361	"TNO": np.where((a > 30.1)),
362	"PAA": np.where((e == 1)),
363	"HYA": np.where((e > 1)),
364	}
365	for c, v in classes_dict.items():	×
366	classes[v] = c	×
367
368	self.orbit_class = classes	×
369	return	×
370
371	@staticmethod	1✔
372	def fromHorizons(obj_ids, t0):	1✔
373	"""
374	Query Horizons for state vectors for each object ID at time t0.
375	This is a convenience function and should not be used to query for state
376	vectors for many objects or at many times.
377
378	Parameters
379	----------
380	obj_ids : `~numpy.ndarray` (N)
381	Object IDs / designations recognizable by HORIZONS.
382	t0 : `~astropy.core.time.Time` (1)
383	Astropy time object at which to gather state vectors.
384
385	Return
386	------
387	`~thor.orbits.orbits.Orbits`
388	THOR Orbits class
389	"""
390
391	if len(t0) != 1:	×
392	err = "t0 should be a single time."	×
393	raise ValueError(err)	×
394
395	horizons_vectors = getHorizonsVectors(	×
396	obj_ids, t0, location="@sun", id_type="smallbody", aberrations="geometric"
397	)
398
399	orbits = Orbits(	×
400	horizons_vectors[["x", "y", "z", "vx", "vy", "vz"]].values,
401	t0 + np.zeros(len(obj_ids)),
402	orbit_type="cartesian",
403	orbit_units=CARTESIAN_UNITS,
404	ids=horizons_vectors["targetname"].values,
405	H=horizons_vectors["H"].values,
406	G=horizons_vectors["G"].values,
407	)
408	return orbits	×
409
410	@staticmethod	1✔
411	def fromMPCOrbitCatalog(mpcorb):	1✔
412
413	cols = ["a_au", "e", "i_deg", "ascNode_deg", "argPeri_deg", "meanAnom_deg"]	×
414	additional_cols = mpcorb.columns[	×
415	~mpcorb.columns.isin(cols + ["mjd_tt", "designation", "H_mag", "G"])
416	]
417	orbits = mpcorb[cols].values	×
418	epochs = Time(mpcorb["mjd_tt"].values, scale="tt", format="mjd")	×
419	args = (orbits, epochs)	×
420
421	kwargs = {	×
422	"orbit_type": "keplerian",
423	"ids": mpcorb["designation"].values,
424	"H": mpcorb["H_mag"].values,
425	"G": mpcorb["G"].values,
426	"orbit_units": KEPLERIAN_UNITS,
427	"additional_data": mpcorb[additional_cols],
428	}
429	return Orbits(args, *kwargs)	×
430
431	def to_df(	1✔
432	self,
433	include_units: bool = True,
434	include_keplerian: bool = False,
435	include_cartesian: bool = True,
436	) -> pd.DataFrame:
437	"""
438	Convert orbits into a pandas dataframe.
439
440	Returns
441	-------
442	dataframe : `~pandas.DataFrame`
443	Dataframe containing orbits, epochs and IDs. If H, G, and covariances are defined then
444	those are also added to the dataframe.
445	"""
446	if self.num_orbits > 0:	×
447	data = {"orbit_id": self.ids, "mjd_tdb": self.epochs.tdb.mjd}	×
448	else:
449	data = {	×
450	"orbit_id": [],
451	"mjd_tdb": [],
452	}
453
454	if include_units:	×
455
456	units_index = ["--", "mjd [TDB]"]	×
457	data["epoch"] = data["mjd_tdb"]	×
458	data.pop("mjd_tdb")	×
459
460	if include_cartesian:	×
461	for i in range(6):	×
462	data[CARTESIAN_COLS[i]] = self.cartesian[:, i]	×
463	if include_units:	×
464	orbit_units_str = []	×
465	for unit in CARTESIAN_UNITS:	×
466	orbit_units_str.append(str(unit).lower())	×
467
468	units_index += orbit_units_str	×
469
470	if include_keplerian:	×
471	for i in range(6):	×
472	data[KEPLERIAN_COLS[i]] = self.keplerian[:, i]	×
473
474	if include_units:	×
475	orbit_units_str = []	×
476	for unit in KEPLERIAN_UNITS:	×
477	if unit == u.dimensionless_unscaled:	×
478	orbit_units_str.append("--")	×
479	else:
480	orbit_units_str.append(str(unit).lower())	×
481	units_index += orbit_units_str	×
482
483	if self._cartesian_covariance is not None and include_cartesian:	×
484	data["covariance"] = self.cartesian_covariance	×
485	if include_units:	×
486	units_index.append("--")	×
487
488	if self._keplerian_covariance is not None and include_keplerian:	×
489	data["covariance"] = self.keplerian_covariance	×
490	if include_units:	×
491	units_index.append("--")	×
492
493	if self.H is not None:	×
494	data["H"] = self.H	×
495	if include_units:	×
496	units_index.append("mag")	×
497
498	if self.G is not None:	×
499	data["G"] = self.G	×
500	if include_units:	×
501	units_index.append("--")	×
502
503	if self.orbit_class is not None:	×
504	data["orbit_class"] = self.orbit_class	×
505	if include_units:	×
506	units_index.append("--")	×
507
508	dataframe = pd.DataFrame(data)	×
509	if self.additional_data is not None:	×
510	dataframe = dataframe.join(self.additional_data)	×
511	if include_units:	×
512	for col in self.additional_data.columns:	×
513	units_index.append("--")	×
514
515	if include_units:	×
516	dataframe.columns = pd.MultiIndex.from_arrays(	×
517	[dataframe.columns, np.array(units_index)]
518	)
519
520	return dataframe	×
521
522	@staticmethod	1✔
523	def from_df(dataframe: pd.DataFrame):	1✔
524	"""
525	Read orbits from a dataframe. Required columns are
526	the epoch at which the orbits are defined ('mjd_tdb') and the 6 dimensional state.
527	If the states are cartesian then the expected columns are ('x', 'y', 'z', 'vx', 'vy', 'vz'),
528	if the states are keplerian then the expected columns are ("a", "e", "i", "raan", "argperi", "M").
529
530	Parameters
531	----------
532	dataframe : `~pandas.DataFrame`
533	Dataframe containing either Cartesian or Keplerian orbits.
534
535	Returns
536	-------
537	`~thor.orbits.orbits.Orbits`
538	THOR Orbit class with the orbits read from the input dataframe.
539	"""
540	# Extract the epochs from the given dataframe
541	dataframe_ = dataframe.copy()	1✔
542
543	if isinstance(dataframe.columns, pd.MultiIndex):	1✔
544	dataframe_.columns = dataframe_.columns.droplevel(level=1)	1✔
545	dataframe_.rename(columns={"epoch": "mjd_tdb"}, inplace=True)	1✔
546
547	if len(dataframe_) > 0:	1✔
548	epochs = Time(dataframe_["mjd_tdb"].values, format="mjd", scale="tdb")	1✔
549	else:
550	epochs = np.array([])	×
551
552	# If the dataframe's index is not sorted and increasing, reset it
553	if not np.all(dataframe_.index.values == np.arange(0, len(dataframe_))):	1✔
554	dataframe_.reset_index(inplace=True, drop=True)	×
555
556	columns_required = ["orbit_id", "mjd_tdb"]	1✔
557	cartesian = None	1✔
558	keplerian = None	1✔
559	if np.all(pd.Series(CARTESIAN_COLS).isin(dataframe_.columns)):	1✔
560	columns_required += CARTESIAN_COLS	1✔
561	cartesian = dataframe_[CARTESIAN_COLS].values	1✔
562	args = (cartesian, epochs)	1✔
563	orbit_type = "cartesian"	1✔
564	elif np.all(pd.Series(KEPLERIAN_COLS).isin(dataframe_.columns)):	×
565	columns_required += KEPLERIAN_COLS	×
566	keplerian = dataframe_[KEPLERIAN_COLS].values	×
567	args = (keplerian, epochs)	×
568	orbit_type = "keplerian"	×
569
570	kwargs = {"ids": dataframe_["orbit_id"].values, "orbit_type": orbit_type}	1✔
571
572	columns_optional = ["covariance", "H", "G", "orbit_class"]	1✔
573	for col in columns_optional:	1✔
574	if col in dataframe_.columns:	1✔
575	kwargs[col] = dataframe_[col].values	×
576
577	columns = columns_required + columns_optional	1✔
578	if len(dataframe_.columns[~dataframe_.columns.isin(columns)]) > 0:	1✔
579	kwargs["additional_data"] = dataframe_[	×
580	dataframe_.columns[~dataframe_.columns.isin(columns)]
581	]
582
583	orbits = Orbits(args, *kwargs)	1✔
584	orbits._keplerian = keplerian	1✔
585	orbits._cartesian = cartesian	1✔
586	# orbits._keplerian_covariance = keplerian_covariance
587	# orbits._cartesian_covariance = cartesian_covariance
588	return orbits	1✔
589
590	def to_csv(	1✔
591	self, file: str, include_cartesian: bool = True, include_keplerian: bool = False
592	):
593	"""
594	Save orbits to a csv. Orbits are always saved with the
595	units of each quantity to avoid ambiguity and confusion.
596
597	Parameters
598	----------
599	file : str
600	Name or path of file including extension to which to save
601	orbits.
602	include_cartesian : bool
603	Save Cartesian elements.
604	include_keplerian : bool
605	Save Keplerian elements.
606	"""
607	df = self.to_df(	×
608	include_units=True,
609	include_cartesian=include_cartesian,
610	include_keplerian=include_keplerian,
611	)
612	df.to_csv(file, index=False, float_format="%.15e", encoding="utf-8")	×
613	return	×
614
615	@staticmethod	1✔
616	def from_csv(file: str):	1✔
617	"""
618	Read orbits from a csv.
619
620	Parameters
621	----------
622	file : str
623	Name or path of file including extension to which to read
624	orbits.
625	"""
626	df = pd.read_csv(	1✔
627	file,
628	index_col=None,
629	header=[0, 1],
630	converters={
631	"cartesian_covariance": lambda x: np.array(
632	ast.literal_eval(",".join(x.replace("[ ", "[").split()))
633	),
634	"keplerian_covariance": lambda x: np.array(
635	ast.literal_eval(",".join(x.replace("[ ", "[").split()))
636	),
637	},
638	dtype={
639	"orbit_id": str,
640	"test_orbit_id": str,
641	},
642	float_precision="round_trip",
643	)
644	df.rename(columns={"epoch": "mjd_tdb"}, inplace=True)	1✔
645	return Orbits.from_df(df)	1✔
646
647	def to_hdf(	1✔
648	self, file: str, include_cartesian: bool = True, include_keplerian: bool = False
649	):
650	"""
651	Save orbits to a HDF5 file. Orbits are always saved with the
652	units of each quantity to avoid ambiguity and confusion.
653
654	Parameters
655	----------
656	file : str
657	Name or path of file including extension to which to save
658	orbits.
659	include_cartesian : bool
660	Save Cartesian elements.
661	include_keplerian : bool
662	Save Keplerian elements.
663	"""
664	df = self.to_df(	×
665	include_units=True,
666	include_cartesian=include_cartesian,
667	include_keplerian=include_keplerian,
668	)
669	df.to_hdf(file, key="data")	×
670	return	×
671
672	@staticmethod	1✔
673	def from_hdf(file: str):	1✔
674	"""
675	Read orbits from a HDF5 file.
676
677	Parameters
678	----------
679	file : str
680	Name or path of file including extension to which to read
681	orbits.
682	"""
683	df = pd.read_hdf(file, key="data")	×
684	return Orbits.from_df(df)	×

moeyensj / thor / 4708907600

Source File Press 'n' to go to next uncovered line, 'b' for previous

Source File
Press 'n' to go to next uncovered line, 'b' for previous