Plot trajectory in ric frame

python/nyx_space/analysis/__init__.py

@@ -0,0 +1,23 @@
+"""
+Nyx, blazing fast astrodynamics
+Copyright (C) 2023 Christopher Rabotin <christopher.rabotin@gmail.com>
+
+This program is free software: you can redistribute it and/or modify
+it under the terms of the GNU Affero General Public License as published
+by the Free Software Foundation, either version 3 of the License, or
+(at your option) any later version.
+
+This program is distributed in the hope that it will be useful,
+but WITHOUT ANY WARRANTY; without even the implied warranty of
+MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+GNU Affero General Public License for more details.
+
+You should have received a copy of the GNU Affero General Public License
+along with this program.  If not, see <https://www.gnu.org/licenses/>.
+"""
+
+from .traj import diff_traj_parquet
+
+__all__ = [
+    "diff_traj_parquet",
+]

python/nyx_space/analysis/traj.py

@@ -0,0 +1,87 @@
+"""
+Nyx, blazing fast astrodynamics
+Copyright (C) 2023 Christopher Rabotin <christopher.rabotin@gmail.com>
+
+This program is free software: you can redistribute it and/or modify
+it under the terms of the GNU Affero General Public License as published
+by the Free Software Foundation, either version 3 of the License, or
+(at your option) any later version.
+
+This program is distributed in the hope that it will be useful,
+but WITHOUT ANY WARRANTY; without even the implied warranty of
+MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+GNU Affero General Public License for more details.
+
+You should have received a copy of the GNU Affero General Public License
+along with this program.  If not, see <https://www.gnu.org/licenses/>.
+"""
+
+import pandas as pd
+from nyx_space.time import Unit, TimeSeries
+from nyx_space.mission_design import TrajectoryLoader
+from nyx_space.monte_carlo import StateParameter
+
+
+def diff_traj_parquet(path1: str, path2: str, step=Unit.Minute * 1) -> pd.DataFrame:
+    """
+    Load both trajectory files from parquet, sample them at the provided step, and return an error data frame that can be plotted with the typical tools.
+    """
+
+    oe_list_str = [
+        "sma",
+        "ecc",
+        "inc",
+        "raan",
+        "aop",
+        "ta",
+        "tlong",
+        "aol",
+        "x",
+        "y",
+        "z",
+        "vx",
+        "vy",
+        "vz",
+    ]
+
+    oe_list = [StateParameter(name) for name in oe_list_str]
+
+    def loader(path: str) -> TrajectoryLoader:
+        # Try to be somewhat clever
+        if path.lower().endswith(".oem"):
+            # Load as OEM and build the parquet file
+            return TrajectoryLoader(
+                str(path), "oem", path[:-3] + "parquet"
+            ).to_orbit_traj()
+        else:
+            return TrajectoryLoader(str(path), "parquet").to_orbit_traj()
+
+    traj1 = loader(path1)
+    traj2 = loader(path2)
+
+    err_data = {"Epoch:Gregorian UTC": []}
+    for oe in oe_list:
+        err_data[str(oe)] = []
+
+    # Sample to build the data.
+
+    for epoch in TimeSeries(
+        traj1.first().epoch, traj1.last().epoch, step, inclusive=True
+    ):
+        try:
+            state1 = traj1.at(epoch)
+            state2 = traj2.at(epoch)
+        except:
+            # Missing data at either of these, let's ignore
+            pass
+        else:
+            err_data["Epoch:Gregorian UTC"] += [str(epoch)]
+            for oe in oe_list:
+                val1 = state1.value_of(oe)
+                val2 = state2.value_of(oe)
+                err_data[str(oe)] += [val1 - val2]
+
+    # Build the data frame
+    err_df = pd.DataFrame(err_data, columns=list(err_data.keys()))
+
+    return err_df

python/nyx_space/plots/__init__.py

@@ -18,13 +18,14 @@
 
 from .gauss_markov import plot_gauss_markov
 from .od import plot_covar, plot_estimates, plot_measurements
-from .traj import plot_traj, plot_ground_track
+from .traj import plot_traj, plot_ground_track, plot_traj_errors
 
 __all__ = [
     "plot_gauss_markov",
     "plot_covar",
     "plot_estimates",
     "plot_traj",
+    "plot_traj_errors",
     "plot_ground_track",
     "plot_measurements",
 ]

python/nyx_space/plots/traj.py

@@ -27,7 +27,6 @@
     colors,
     finalize_plot,
     body_color,
-    nyx_tpl,
 )
 
 
@@ -205,10 +204,9 @@ def plot_orbit_elements(
     Args:
         dfs (pandas.DataFrame): The data frame containing the trajectory (or a list thereof)
         title (str): The title of the plot
+        names (List[str]): The names of each data frame
         html_out (str, optional): The path to save the HTML to. Defaults to None.
         copyright (str, optional): The copyright to display on the plot. Defaults to None.
-        fig (plotly.graph_objects.Figure, optional): The figure to add the trajectory to. Defaults to None.
-        center (str, optional): The name of the center object, e.g. `Luna` (Moon). Defaults to "Earth".
         show (bool, optional): Whether to show the plot. Defaults to True. If set to false, the figure will be returned.
     """
 
@@ -246,15 +244,21 @@ def plot_orbit_elements(
     row_i = 0
     col_i = 0
 
+    color_values = list(colors.values())
+
     for col in columns:
         for k, df in enumerate(dfs):
             try:
                 name = f"{names[k]} {col}"
             except IndexError:
                 name = col
 
+            # Build the color for this data frame
+            color = color_values[k % len(color_values)]
+            color = f"rgb({int(color[0])}, {int(color[1])}, {int(color[2])})"
+
             fig.add_trace(
-                go.Scatter(x=df["Epoch"], y=df[col], name=name),
+                go.Scatter(x=df["Epoch"], y=df[col], name=name, marker_color=color),
                 row=row_i + 1,
                 col=col_i + 1,
             )
@@ -273,3 +277,117 @@ def plot_orbit_elements(
         fig.show()
     else:
         return fig
+
+
+def plot_traj_errors(
+    dfs,
+    title,
+    names=[],
+    ric=True,
+    vertical=False,
+    velocity=False,
+    html_out=None,
+    copyright=None,
+    show=True,
+):
+    """
+    Plot the trajectory error data in the Radial, In-track, Cross-track frame.
+
+    Args:
+        dfs (List[pandas.DataFrame]): The list of trajectory data frames
+        title (str): The title of the plot
+        names (List[str]): The name of each trajectory
+        ric (bool): Set to False to plot Cartesian errors instead of RIC.
+        vertical (bool): Set True to plot components vertically instead of stacked.
+        velocity (bool): Set True to also plot velocity errors instead of the position errors.
+        html_out (str, optional): Path to save HTML plot to.
+        copyright (str, optional): Copyright notice.
+        show (bool, optional): Whether to show the plot.
+
+    Returns:
+        plotly.graph_objects.Figure: The Figure instance
+    """
+
+    if not isinstance(dfs, list):
+        dfs = [dfs]
+
+    for df in dfs:
+        pd_ok_epochs = []
+        for epoch in df["Epoch:Gregorian UTC"]:
+            epoch = epoch.replace("UTC", "").strip()
+            if "." not in epoch:
+                epoch += ".0"
+            pd_ok_epochs += [epoch]
+        df["Epoch"] = pd.to_datetime(pd_ok_epochs)
+
+    if not isinstance(names, list):
+        names = [names]
+
+    if velocity:
+        if ric:
+            columns = [
+                "delta_vx_ric (km/s)",
+                "delta_vy_ric (km/s)",
+                "delta_vz_ric (km/s)",
+            ]
+        else:
+            columns = [
+                "vx (km/s)",
+                "vy (km/s)",
+                "vz (km/s)",
+            ]
+    else:
+        if ric:
+            columns = [
+                "delta_x_ric (km)",
+                "delta_y_ric (km)",
+                "delta_z_ric (km)",
+            ]
+        else:
+            columns = [
+                "x (km)",
+                "y (km)",
+                "z (km)",
+            ]
+
+    if vertical:
+        fig = make_subplots(rows=1, cols=3, subplot_titles=columns)
+    else:
+        fig = make_subplots(rows=3, cols=1, subplot_titles=columns)
+
+    row_i = 0
+    col_i = 0
+
+    for col in columns:
+        for k, df in enumerate(dfs):
+            try:
+                name = f"{names[k]} {col}"
+            except IndexError:
+                name = col
+
+            try:
+                fig.add_trace(
+                    go.Scatter(x=df["Epoch"], y=df[col], name=name),
+                    row=row_i + 1,
+                    col=col_i + 1,
+                )
+            except KeyError:
+                raise KeyError(
+                    f"Rebuild the trajectory and export the RIC frame: missing `{col}`"
+                )
+        if vertical:
+            col_i += 1
+        else:
+            row_i += 1
+
+    finalize_plot(fig, title, copyright=copyright)
+
+    if html_out:
+        with open(html_out, "w") as f:
+            f.write(fig.to_html())
+        print(f"Saved HTML to {html_out}")
+
+    if show:
+        fig.show()
+    else:
+        return fig