6164332637

Committed 12 Sep 2023 08:07PM UTC coverage: 35.385% (-12.6%) from 47.98%

Build # 6164332637

Build Type

push

github

Committed by

web-flow

Commit Message

Merge pull request #114 from moeyensj/jm/2.0-remove-old-tests

Fix CI for v2.0 changes

Run Details

1719 of 4858 relevant lines covered (35.38%)

0.35 hits per line

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

15.16

/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

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.")
        return self._cartesian

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

moeyensj / thor / 6164332637

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