Append coordinate reference frame to coordinate dimension columns whe…

…n serializing to and from a dataframe

adam_core/coordinates/cartesian.py

@@ -1,5 +1,5 @@
 import logging
-from typing import TYPE_CHECKING, Literal, Optional
+from typing import TYPE_CHECKING, Optional
 
 import numpy as np
 import pandas as pd
@@ -341,24 +341,20 @@ def to_dataframe(
         )
 
     @classmethod
-    def from_dataframe(
-        cls, df: pd.DataFrame, frame: Literal["ecliptic", "equatorial"]
-    ) -> "CartesianCoordinates":
+    def from_dataframe(cls, df: pd.DataFrame) -> "CartesianCoordinates":
         """
         Create coordinates from a pandas DataFrame.
 
         Parameters
         ----------
         df : `~pandas.Dataframe`
             DataFrame containing coordinates.
-        frame : {"ecliptic", "equatorial"}
-            Frame in which coordinates are defined.
 
         Returns
         -------
         coords : `~adam_core.coordinates.cartesian.CartesianCoordinates`
             Cartesian coordinates.
         """
         return coords_from_dataframe(
-            cls, df, coord_names=["x", "y", "z", "vx", "vy", "vz"], frame=frame
+            cls, df, coord_names=["x", "y", "z", "vx", "vy", "vz"]
         )

adam_core/coordinates/cometary.py

@@ -1,4 +1,4 @@
-from typing import TYPE_CHECKING, Literal, Optional
+from typing import TYPE_CHECKING, Optional
 
 import numpy as np
 import pandas as pd
@@ -356,24 +356,20 @@ def to_dataframe(
         )
 
     @classmethod
-    def from_dataframe(
-        cls, df: pd.DataFrame, frame: Literal["ecliptic", "equatorial"]
-    ) -> "CometaryCoordinates":
+    def from_dataframe(cls, df: pd.DataFrame) -> "CometaryCoordinates":
         """
         Create coordinates from a pandas DataFrame.
 
         Parameters
         ----------
         df : `~pandas.Dataframe`
             DataFrame containing coordinates.
-        frame : {"ecliptic", "equatorial"}
-            Frame in which coordinates are defined.
 
         Returns
         -------
         coords : `~adam_core.coordinates.cometary.CometaryCoordinates`
             Cometary coordinates.
         """
         return coords_from_dataframe(
-            cls, df, coord_names=["q", "e", "i", "raan", "ap", "tp"], frame=frame
+            cls, df, coord_names=["q", "e", "i", "raan", "ap", "tp"]
         )

adam_core/coordinates/io.py

@@ -1,4 +1,4 @@
-from typing import TYPE_CHECKING, List, Literal, Optional, Type, Union
+from typing import TYPE_CHECKING, List, Optional, Type, Union
 
 import numpy as np
 import pandas as pd
@@ -35,7 +35,9 @@ def coords_to_dataframe(
     coords : {CartesianCoordinates, CometaryCoordinates, KeplerianCoordinates, SphericalCoordinates}
         Coordinates to store.
     coord_names : list of str
-        Names of the coordinates to store.
+        Names of the coordinates to store. The coordinate reference frame will be appended to the
+        coordinate names, e.g., "x" will become "x_ec" for ecliptic coordinates and "x_eq" for
+        equatorial coordinates.
     sigmas : bool, optional
         If None, will check if any sigmas are defined (via covariance matrix) and add them to the dataframe.
         If True, include 1-sigma uncertainties in the DataFrame regardless. If False, do not include 1-sigma
@@ -54,6 +56,13 @@ def coords_to_dataframe(
     # Gather times and put them into a dataframe
     df_times = coords.time.to_dataframe(flatten=True)
 
+    if coords.frame == "ecliptic":
+        coord_names = [f"{coord}_ec" for coord in coord_names]
+    elif coords.frame == "equatorial":
+        coord_names = [f"{coord}_eq" for coord in coord_names]
+    else:
+        raise ValueError(f"Frame {coords.frame} not recognized.")
+
     df_coords = pd.DataFrame(
         data=coords.values,
         columns=coord_names,
@@ -92,7 +101,6 @@ def coords_from_dataframe(
     cls: "Type[CoordinateType]",
     df: pd.DataFrame,
     coord_names: List[str],
-    frame: Literal["ecliptic", "equatorial"],
 ) -> "CoordinateType":
     """
     Return coordinates from a pandas DataFrame that was generated with
@@ -104,9 +112,9 @@ def coords_from_dataframe(
     df : `~pandas.Dataframe`
         DataFrame containing coordinates and covariances.
     coord_names : list of str
-        Names of the coordinates dimensions.
-    frame : {"ecliptic", "equatorial"}
-        Frame in which the coordinates are defined.
+        Names of the coordinate dimensions. The coordinate reference frame will be appended to the
+        coordinate names, e.g., "x" will become "x_ec" for ecliptic coordinates and "x_eq" for
+        equatorial coordinates.
 
     Returns
     -------
@@ -123,7 +131,26 @@ def coords_from_dataframe(
         df[["origin.code"]].rename(columns={"origin.code": "code"})
     )
     covariances = CoordinateCovariances.from_dataframe(df, coord_names=coord_names)
-    coords = {col: df[col].values for col in coord_names}
+
+    coord_names_ec = [f"{coord}_ec" for coord in coord_names]
+    coord_names_eq = [f"{coord}_eq" for coord in coord_names]
+    if all(col in df.columns for col in coord_names_ec):
+        coords = {
+            col: df[col_frame].values
+            for col, col_frame in zip(coord_names, coord_names_ec)
+        }
+        frame = "ecliptic"
+    elif all(col in df.columns for col in coord_names_eq):
+        coords = {
+            col: df[col_frame].values
+            for col, col_frame in zip(coord_names, coord_names_eq)
+        }
+        frame = "equatorial"
+    else:
+        raise ValueError(
+            "DataFrame does not contain the expected columns for the coordinate values:\n"
+            f"Expected: {coord_names_ec} or {coord_names_eq}\n"
+        )
 
     return cls.from_kwargs(
         **coords, time=times, origin=origin, frame=frame, covariance=covariances

adam_core/coordinates/keplerian.py

@@ -1,4 +1,4 @@
-from typing import TYPE_CHECKING, Literal, Optional
+from typing import TYPE_CHECKING, Optional
 
 import numpy as np
 import pandas as pd
@@ -332,24 +332,20 @@ def to_dataframe(
         )
 
     @classmethod
-    def from_dataframe(
-        cls, df: pd.DataFrame, frame: Literal["ecliptic", "equatorial"]
-    ) -> "KeplerianCoordinates":
+    def from_dataframe(cls, df: pd.DataFrame) -> "KeplerianCoordinates":
         """
         Create coordinates from a pandas DataFrame.
 
         Parameters
         ----------
         df : `~pandas.Dataframe`
             DataFrame containing coordinates.
-        frame : {"ecliptic", "equatorial"}
-            Frame in which coordinates are defined.
 
         Returns
         -------
         coords : `~adam_core.coordinates.keplerian.KeplerianCoordinates`
             Keplerian coordinates.
         """
         return coords_from_dataframe(
-            cls, df, coord_names=["a", "e", "i", "raan", "ap", "M"], frame=frame
+            cls, df, coord_names=["a", "e", "i", "raan", "ap", "M"]
         )

adam_core/coordinates/spherical.py

@@ -1,4 +1,4 @@
-from typing import TYPE_CHECKING, Literal, Optional
+from typing import TYPE_CHECKING, Optional
 
 import numpy as np
 import pandas as pd
@@ -280,27 +280,20 @@ def to_dataframe(
         )
 
     @classmethod
-    def from_dataframe(
-        cls, df: pd.DataFrame, frame: Literal["ecliptic", "equatorial"]
-    ) -> "SphericalCoordinates":
+    def from_dataframe(cls, df: pd.DataFrame) -> "SphericalCoordinates":
         """
         Create coordinates from a pandas DataFrame.
 
         Parameters
         ----------
         df : `~pandas.Dataframe`
             DataFrame containing coordinates.
-        frame : {"ecliptic", "equatorial"}
-            Frame in which coordinates are defined.
 
         Returns
         -------
         coords : `~adam_core.coordinates.spherical.SphericalCoordinates`
             Spherical coordinates.
         """
         return coords_from_dataframe(
-            cls,
-            df,
-            coord_names=["rho", "lon", "lat", "vrho", "vlon", "vlat"],
-            frame=frame,
+            cls, df, coord_names=["rho", "lon", "lat", "vrho", "vlon", "vlat"]
         )

adam_core/coordinates/tests/test_io.py

@@ -0,0 +1,111 @@
+import numpy as np
+from astropy.time import Time
+
+from ..cartesian import CartesianCoordinates
+from ..covariances import CoordinateCovariances
+from ..origin import Origin
+from ..times import Times
+
+coords_ec = CartesianCoordinates.from_kwargs(
+    time=Times.from_astropy(
+        Time([59000.0, 59001.0, 59002.0], format="mjd", scale="tdb")
+    ),
+    x=[1, 2, 3],
+    y=[4, 5, 6],
+    z=[7, 8, 9],
+    vx=[0.1, 0.2, 0.3],
+    vy=[0.4, 0.5, 0.6],
+    vz=[0.7, 0.8, 0.9],
+    covariance=CoordinateCovariances.from_sigmas(
+        np.array(
+            [
+                [0.01, 0.01, 0.01, 0.01, 0.01, 0.01],
+                [0.02, 0.02, 0.02, 0.02, 0.02, 0.02],
+                [0.03, 0.03, 0.03, 0.03, 0.03, 0.03],
+            ]
+        ),
+    ),
+    frame="ecliptic",
+    origin=Origin.from_kwargs(code=["SUN" for i in range(3)]),
+)
+
+coords_eq = CartesianCoordinates.from_kwargs(
+    time=coords_ec.time,
+    x=coords_ec.x,
+    y=coords_ec.y,
+    z=coords_ec.z,
+    vx=coords_ec.vx,
+    vy=coords_ec.vy,
+    vz=coords_ec.vz,
+    covariance=coords_ec.covariance,
+    frame="equatorial",
+    origin=coords_ec.origin,
+)
+
+
+def test_coords_to_dataframe():
+    # Cartesian coordinates defined in the ecliptic frame
+    df = coords_ec.to_dataframe()
+    np.testing.assert_equal(df["x_ec"].values, coords_ec.x.to_numpy())
+    np.testing.assert_equal(df["y_ec"].values, coords_ec.y.to_numpy())
+    np.testing.assert_equal(df["z_ec"].values, coords_ec.z.to_numpy())
+    np.testing.assert_equal(df["vx_ec"].values, coords_ec.vx.to_numpy())
+    np.testing.assert_equal(df["vy_ec"].values, coords_ec.vy.to_numpy())
+    np.testing.assert_equal(df["vz_ec"].values, coords_ec.vz.to_numpy())
+    np.testing.assert_equal(
+        df["cov_x_x"].values, coords_ec.covariances.sigmas[:, 0] ** 2
+    )
+    np.testing.assert_equal(
+        df["cov_y_y"].values, coords_ec.covariances.sigmas[:, 1] ** 2
+    )
+    np.testing.assert_equal(
+        df["cov_z_z"].values, coords_ec.covariances.sigmas[:, 2] ** 2
+    )
+    np.testing.assert_equal(
+        df["cov_vx_vx"].values, coords_ec.covariances.sigmas[:, 3] ** 2
+    )
+    np.testing.assert_equal(
+        df["cov_vy_vy"].values, coords_ec.covariances.sigmas[:, 4] ** 2
+    )
+    np.testing.assert_equal(
+        df["cov_vz_vz"].values, coords_ec.covariances.sigmas[:, 5] ** 2
+    )
+
+    # Cartesian coordinates defined in the equatorial frame
+    df = coords_eq.to_dataframe()
+    np.testing.assert_equal(df["x_eq"].values, coords_eq.x.to_numpy())
+    np.testing.assert_equal(df["y_eq"].values, coords_eq.y.to_numpy())
+    np.testing.assert_equal(df["z_eq"].values, coords_eq.z.to_numpy())
+    np.testing.assert_equal(df["vx_eq"].values, coords_eq.vx.to_numpy())
+    np.testing.assert_equal(df["vy_eq"].values, coords_eq.vy.to_numpy())
+    np.testing.assert_equal(df["vz_eq"].values, coords_eq.vz.to_numpy())
+    np.testing.assert_equal(
+        df["cov_x_x"].values, coords_eq.covariances.sigmas[:, 0] ** 2
+    )
+    np.testing.assert_equal(
+        df["cov_y_y"].values, coords_eq.covariances.sigmas[:, 1] ** 2
+    )
+    np.testing.assert_equal(
+        df["cov_z_z"].values, coords_eq.covariances.sigmas[:, 2] ** 2
+    )
+    np.testing.assert_equal(
+        df["cov_vx_vx"].values, coords_eq.covariances.sigmas[:, 3] ** 2
+    )
+    np.testing.assert_equal(
+        df["cov_vy_vy"].values, coords_eq.covariances.sigmas[:, 4] ** 2
+    )
+    np.testing.assert_equal(
+        df["cov_vz_vz"].values, coords_eq.covariances.sigmas[:, 5] ** 2
+    )
+
+
+def test_coords_from_dataframe():
+    # Cartesian coordinates defined in the ecliptic frame
+    df = coords_ec.to_dataframe()
+    coords_ec2 = CartesianCoordinates.from_dataframe(df)
+    assert coords_ec2.frame == coords_ec.frame
+
+    # Cartesian coordinates defined in the equatorial frame
+    df = coords_eq.to_dataframe()
+    coords_eq2 = CartesianCoordinates.from_dataframe(df)
+    assert coords_eq2.frame == coords_eq.frame

adam_core/observers/observers.py

@@ -150,18 +150,17 @@ def from_dataframe(cls, df: pd.DataFrame) -> Self:
             columns={
                 "obs_jd1_tdb": "jd1_tdb",
                 "obs_jd2_tdb": "jd2_tdb",
-                "obs_x": "x",
-                "obs_y": "y",
-                "obs_z": "z",
-                "obs_vx": "vx",
-                "obs_vy": "vy",
-                "obs_vz": "vz",
+                "obs_x_ec": "x_ec",
+                "obs_y_ec": "y_ec",
+                "obs_z_ec": "z_ec",
+                "obs_vx_ec": "vx_ec",
+                "obs_vy_ec": "vy_ec",
+                "obs_vz_ec": "vz_ec",
                 "obs_origin.code": "origin.code",
             }
         )
         coordinates = CartesianCoordinates.from_dataframe(
             df_renamed,
-            frame="ecliptic",
         )
         return cls.from_kwargs(
             code=df["obs_code"].values,

adam_core/observers/tests/test_state.py

@@ -24,7 +24,7 @@ def test_get_observer_state_X05():
             index_col=False,
             float_precision="round_trip",
         )
-        states_expected = CartesianCoordinates.from_dataframe(states_df, "ecliptic")
+        states_expected = CartesianCoordinates.from_dataframe(states_df)
 
         if origin == "sun":
             origin_code = OriginCodes.SUN
@@ -64,7 +64,7 @@ def test_get_observer_state_I41():
             index_col=False,
             float_precision="round_trip",
         )
-        states_expected = CartesianCoordinates.from_dataframe(states_df, "ecliptic")
+        states_expected = CartesianCoordinates.from_dataframe(states_df)
 
         if origin == "sun":
             origin_code = OriginCodes.SUN
@@ -104,7 +104,7 @@ def test_get_observer_state_W84():
             index_col=False,
             float_precision="round_trip",
         )
-        states_expected = CartesianCoordinates.from_dataframe(states_df, "ecliptic")
+        states_expected = CartesianCoordinates.from_dataframe(states_df)
 
         if origin == "sun":
             origin_code = OriginCodes.SUN
@@ -144,7 +144,7 @@ def test_get_observer_state_000():
             index_col=False,
             float_precision="round_trip",
         )
-        states_expected = CartesianCoordinates.from_dataframe(states_df, "ecliptic")
+        states_expected = CartesianCoordinates.from_dataframe(states_df)
 
         if origin == "sun":
             origin_code = OriginCodes.SUN
@@ -184,7 +184,7 @@ def test_get_observer_state_500():
             index_col=False,
             float_precision="round_trip",
         )
-        states_expected = CartesianCoordinates.from_dataframe(states_df, "ecliptic")
+        states_expected = CartesianCoordinates.from_dataframe(states_df)
 
         if origin == "sun":
             origin_code = OriginCodes.SUN