From b3de7f3d4de89e44d2ebb2d90577f117a8bb0dfa Mon Sep 17 00:00:00 2001 From: Joseph Murtagh Date: Fri, 4 Apr 2025 16:54:00 +0100 Subject: [PATCH 1/9] add some visualisation functionality --- src/layup/visualise.py | 372 +++++++++++++++++++++++++++++++++++++++++ 1 file changed, 372 insertions(+) create mode 100644 src/layup/visualise.py diff --git a/src/layup/visualise.py b/src/layup/visualise.py new file mode 100644 index 00000000..90ed63bb --- /dev/null +++ b/src/layup/visualise.py @@ -0,0 +1,372 @@ +import logging +import os +from pathlib import Path +from typing import Literal + +import numpy as np +import matplotlib.pyplot as plt +from matplotlib.collections import LineCollection + +from layup.utilities.file_io.CSVReader import CSVDataReader +from layup.utilities.file_io.HDF5Reader import HDF5DataReader + +logger = logging.getLogger(__name__) + +# Required column names for each orbit format +REQUIRED_COLUMN_NAMES = { + "BCART": ["ObjID", "FORMAT", "x", "y", "z", "xdot", "ydot", "zdot", "epochMJD_TDB"], + "BCOM": ["ObjID", "FORMAT", "q", "e", "inc", "node", "argPeri", "t_p_MJD_TDB", "epochMJD_TDB"], + "BKEP": ["ObjID", "FORMAT", "a", "e", "inc", "node", "argPeri", "ma", "epochMJD_TDB"], + "CART": ["ObjID", "FORMAT", "x", "y", "z", "xdot", "ydot", "zdot", "epochMJD_TDB"], + "COM": ["ObjID", "FORMAT", "q", "e", "inc", "node", "argPeri", "t_p_MJD_TDB", "epochMJD_TDB"], + "KEP": ["ObjID", "FORMAT", "a", "e", "inc", "node", "argPeri", "ma", "epochMJD_TDB"], +} + + +def construct_ellipse( + a: float, + e: float, + i: float, + omega: float, + Omega: float, + M: float, +): + """ + Construct an ellipse of position vectors using the formalism defined + in chapter 1.2 of Morbidelli 2011, "Modern Celestial Mechanics" + + Paramters + ---------- + a : float + The semimajor axis of the orbit (in au) + + e : float + The eccentricty of the orbit + + i : float + The inclination of the orbit (in degrees) + + omega : float + The argument of perihelion of the orbit (in degrees) + + Omega : float + The longitude of the ascending node of the orbit (in degrees) + + M : float + The mean anomaly of the orbit (in degrees) + + Returns + -------- + r : numpy array + The rotated barycentric position vectors of the orbit ellipse + """ + # we want an eccentrict anomaly array for the entire ellipse, but + # to avoid all orbits starting at perihelion in the plot we will do + # a quick numerical estimate of the eccentric anomaly from the input + # mean anomaly via fixed point iteration and then wrap the linspace + # around this value from 0 to 2pi + E_init = M # initial guesses using mean anomaly (shouldn't be too far off) + + for tries in range(100): + E_new = M + e * np.sin(E_init) + if np.abs(E_new - E_init) < 1e-8: + break + E_init = E_new + if tries == 99: + print('didnt converge') + raise ValueError("Did not converge for all M values") + + N = 2000 + start = np.linspace(E_new, 360, N//2, endpoint=True) + end = np.linspace(0, E_new, N - len(start)) + E = np.concatenate((start, end)) * np.pi/180 + + # or, if you don't care, define eccentric anomaly for all ellipses + # E = np.linspace(0, 2*np.pi, 10000) + + # define position vectors in orthogonal reference frame with origin + # at ellipse main focus with q1 oriented towards periapsis (see chapter + # 1.2 of Morbidelli 2011, "Modern Celestial Mechanics" for details) + q1 = a * (np.cos(E) - e) + q2 = a * np.sqrt(1. - e**2) * np.sin(E) + + if type(q1) == float: + q = np.array([q1, q2, 0.]) + else: + q = np.array([q1, q2, np.zeros_like(q1)]) + + # define rotation matrix to map orthogonal reference frame to + # barycentric reference frame + c0 = np.cos(np.pi * Omega/180) + s0 = np.sin(np.pi * Omega/180) + co = np.cos(np.pi * omega/180) + so = np.sin(np.pi * omega/180) + ci = np.cos(np.pi * i/180) + si = np.sin(np.pi * i/180) + + R = np.array([ + [c0 * co - s0 * so * ci, -c0 * so - s0 * co * ci, s0 * si], + [s0 * co + c0 * so * ci, -s0 * so + c0 * co * ci, -c0 * si], + [si * so, si * co, ci ] + ]) + + # apply rotation matrix + r = np.einsum('ij,j...', R, q) + + return r + +def plot_ellipse( + orb_array, + plot_type, + output +): + """ + Handle the plotting of the orbit ellipses depending on whether the user wants + matplot flat 2D plots, or 3D interactive plots with plotly + + Parameters + ----------- + orb_array : numpy structured array + The orbits as read in by data readers + plot_type : str + The plotting engine to use. Can be either 'matplot' or 'plotly' currently + output : str + The output file stem. + """ + + if plot_type == 'matplot': + import matplotlib.pyplot as plt + from matplotlib.collections import LineCollection + + plt.rcParams.update( + {'axes.linewidth': 1.875, + 'grid.linewidth': 1.5, + 'lines.linewidth': 2.25, + 'lines.markersize': 9.0, + 'patch.linewidth': 1.5, + 'xtick.major.width': 1.875, + 'ytick.major.width': 1.875, + 'xtick.minor.width': 1.5, + 'ytick.minor.width': 1.5, + 'xtick.major.size': 9.0, + 'ytick.major.size': 9.0, + 'xtick.minor.size': 6.0, + 'ytick.minor.size': 6.0, + 'font.size': 18.0, + 'axes.labelsize': 18.0, + 'axes.titlesize': 18.0, + 'xtick.labelsize': 16.5, + 'ytick.labelsize': 16.5, + 'legend.fontsize': 16.5, + 'legend.title_fontsize': 18.0} + ) + + # TODO: make a y-z plot as well as an x-y plot and have them side-by-side? + fig, axs = plt.subplots(figsize=(15,9)) + + axs.tick_params( + labelbottom=True, + labeltop=True, + labelleft=True, + labelright=True + ) + + fig.patch.set_facecolor('k') + axs.set_facecolor('k') + + axs.spines['bottom'].set_color('white') + axs.spines['top'].set_color('white') + axs.spines['right'].set_color('white') + axs.spines['left'].set_color('white') + + axs.tick_params(axis='x', colors='white') + axs.tick_params(axis='y', colors='white') + + axs.xaxis.label.set_color('white') + axs.yaxis.label.set_color('white') + + axs.set_xlabel('x [AU]') + axs.set_ylabel('y [AU]') + + for obj in orb_array: + posvec = construct_ellipse(obj[2], obj[3], obj[4], obj[6], obj[5], obj[7]) + + # doing some matplot magic to make the orbits fade in opacity further away from the final position in array + points = np.array([posvec[:,0], posvec[:,1]]).T.reshape(-1, 1, 2) + segments = np.concatenate([points[:-1], points[1:]], axis=1) + colors = [(1,1,1,a) for a in np.power(np.linspace(0, 1, len(segments)), 3)] + + lc = LineCollection(segments, colors=colors, linewidth=1) + axs.add_collection(lc) + + # TODO: read in semimajor values and determine max, use that to figure out plot extents? + x_extent = 50 + y_extent = 50 + + axs.set( + xlim=(-x_extent,x_extent), + ylim=(-y_extent,y_extent), + aspect="equal" if x_extent == y_extent else "auto" + ) + + # TODO: redo file saving and make it go to user defined path + plt.savefig(output, dpi=300) + plt.show() + plt.close() + + + elif plot_type == 'plotly': + import plotly.graph_objects as go + + fig = go.Figure() + + fig.update_layout(plot_bgcolor='rgb(0,0,0)', + paper_bgcolor='rgb(0,0,0)', + scene = dict( + xaxis = dict( + backgroundcolor="rgba(0, 0, 0,0)", + gridcolor="rgb(106, 110, 117)", + title_font=dict(color="rgb(255,255,255)"), + tickfont=dict(color="rgb(255,255,255)"), + showbackground=True, + zerolinecolor="white", + ), + yaxis = dict( + backgroundcolor="rgba(0, 0, 0,0)", + gridcolor="rgb(106, 110, 117)", + title_font=dict(color="rgb(255,255,255)"), + tickfont=dict(color="rgb(255,255,255)"), + showbackground=True, + zerolinecolor="white" + ), + zaxis = dict( + backgroundcolor="rgba(0, 0, 0,0)", + gridcolor="rgb(106, 110, 117)", + title_font=dict(color="rgb(255,255,255)"), + tickfont=dict(color="rgb(255,255,255)"), + showbackground=True, + zerolinecolor="white", + ), + ), + ) + fig.update_yaxes(title_font_color='white', color='white') + fig.update_xaxes(title_font_color='white', color='white') + + + xs = [] + ys = [] + zs = [] + names = [] + + for obj in orb_array: + posvec = construct_ellipse(obj[2], obj[3], obj[4], obj[6], obj[5], obj[7]) + xs.append(posvec[-1,0]) + ys.append(posvec[-1,1]) + zs.append(posvec[-1,2]) + names.append(obj[0]) + + normx = (posvec[:,0] - posvec[:,0].min()) / (posvec[:,0].max() - posvec[:,0].min()) + opacity = np.log1p(normx * 9) / np.log1p(9) + + # TODO: fix hover boxes to reflect objID properly + fig.add_trace(go.Scatter3d( + x=posvec[:,0], y=posvec[:,1], z=posvec[:,2], + mode="lines+markers", + marker=dict(color=opacity, size=5), + line=dict(color='white'), + customdata=[obj[0]], + hovertemplate="%{customdata[0]}
(x: %{x:.2f}, y: %{y:.2f}, z: %{z:.2f})" + )) + + # TODO: add object a/e/i as hover tooltip info? + fig.add_trace(go.Scatter3d( + x=xs, y=ys, z=zs, + mode="markers", + marker=dict(color='rgba(255, 0, 0, 1)', size=5), + customdata=names, # Pass planet name as custom data + hovertemplate="%{customdata}
(x: %{x:.2f}, y: %{y:.2f}, z: %{z:.2f})" # Define hover format + )) + + fig.update_layout(showlegend=False) + fig.show() + + # TODO: add in file saving? or can user do it from plotly output + + + +def visualise_cli( + input: str, + output_file_stem: str, + plot_type: Literal["matplot", "plotly"] = "matplot", + input_format: Literal["BCART", "BCOM", "BKEP", "CART", "COM", "KEP"] = "BKEP", + file_format: Literal["csv", "hdf5"] = "csv", + num_orbs: int = 1_000, +): + """ + Read in an orbits file and plot a user determined random number of orbits. + + Parameters + ----------- + inputs : str + The path to the input orbits file + + output : str + The output file stem name. + + plot_type : str (default=matplot) + The backend to use for plotting orbits. Must be one of: 'matplot' or 'plotly' + + input_format : str (default=BKEP) + The orbit format of the input orbits. Must be one of: 'BCART', 'BCOM', 'BKEP', 'CART', 'COM', or 'KEP' + + file_format : str (default=csv) + The format of the input orbits file. Must be one of: 'csv' or 'hdf5' + + num_orbs : int, optional (default=1_000) + The number of random orbits to plot + """ + input_file = Path(input) + + output_file = Path(f"{output_file_stem}.pdf") + + output_directory = output_file.parent.resolve() + if not output_directory.exists(): + logging.error(f"Output directory {output_directory} does not exist") + + # TODO: currently assumes BKEP input. need to convert if not in BKEP + required_columns = REQUIRED_COLUMN_NAMES[input_format] + + if file_format == "hdf5": + reader = HDF5DataReader( + input_file, + form_column_name="FORMAT", + required_column_names=required_columns, + ) + else: + reader = CSVDataReader( + input_file, + format_column_name="FORMAT", + required_column_names=required_columns, + ) + + random_orbs = reader.read_rows(block_start=0, block_size=1000) + random_orbs = random_orbs[np.random.choice(random_orbs.size, size=num_orbs, replace=False)] + + plot_ellipse(random_orbs, plot_type, output_file) + + + +def main(): + + # TODO: add in actual layup argument stuff to this, instead of hardcoded values/paths + visualise_cli( + input='/Users/josephmurtagh/Downloads/cent_orbs.csv', + output_file_stem='testplot', + plot_type='matplot', + file_format='csv', + num_orbs=100, + ) + +if __name__ == "__main__": + main() + From 2a1f98750ebebe05da23a04a9d6d07620a251911 Mon Sep 17 00:00:00 2001 From: Joseph Murtagh Date: Fri, 2 May 2025 14:19:09 -0700 Subject: [PATCH 2/9] update CLI plotting functionality --- src/layup/matplot2d.mplstyle | 20 ++ src/layup/visualize.py | 547 +++++++++++++++++++++++++++++++++ src/layup_cmdline/visualize.py | 59 +++- 3 files changed, 624 insertions(+), 2 deletions(-) create mode 100644 src/layup/matplot2d.mplstyle create mode 100644 src/layup/visualize.py diff --git a/src/layup/matplot2d.mplstyle b/src/layup/matplot2d.mplstyle new file mode 100644 index 00000000..c4debece --- /dev/null +++ b/src/layup/matplot2d.mplstyle @@ -0,0 +1,20 @@ +axes.linewidth : 1.875 +grid.linewidth : 1.5 +lines.linewidth : 2.25 +lines.markersize : 9.0 +patch.linewidth : 1.5 +xtick.major.width : 1.875 +ytick.major.width : 1.875 +xtick.minor.width : 1.5 +ytick.minor.width : 1.5 +xtick.major.size : 9.0 +ytick.major.size : 9.0 +xtick.minor.size : 6.0 +ytick.minor.size : 6.0 +font.size : 18.0 +axes.labelsize : 18.0 +xtick.labelsize : 16.5 +ytick.labelsize : 16.5 +legend.fontsize : 16.5 +legend.title_fontsize : 18.0 +axes.titlesize : 32 \ No newline at end of file diff --git a/src/layup/visualize.py b/src/layup/visualize.py new file mode 100644 index 00000000..e0e0f751 --- /dev/null +++ b/src/layup/visualize.py @@ -0,0 +1,547 @@ +import logging +import os +from pathlib import Path +from typing import Literal + +import numpy as np + +from layup.utilities.file_io.CSVReader import CSVDataReader +from layup.utilities.file_io.HDF5Reader import HDF5DataReader + +logger = logging.getLogger(__name__) + +# Required column names for each orbit format +REQUIRED_COLUMN_NAMES = { + "BCART": ["ObjID", "FORMAT", "x", "y", "z", "xdot", "ydot", "zdot", "epochMJD_TDB"], + "BCOM": ["ObjID", "FORMAT", "q", "e", "inc", "node", "argPeri", "t_p_MJD_TDB", "epochMJD_TDB"], + "BKEP": ["ObjID", "FORMAT", "a", "e", "inc", "node", "argPeri", "ma", "epochMJD_TDB"], + "CART": ["ObjID", "FORMAT", "x", "y", "z", "xdot", "ydot", "zdot", "epochMJD_TDB"], + "COM": ["ObjID", "FORMAT", "q", "e", "inc", "node", "argPeri", "t_p_MJD_TDB", "epochMJD_TDB"], + "KEP": ["ObjID", "FORMAT", "a", "e", "inc", "node", "argPeri", "ma", "epochMJD_TDB"], +} + + +def construct_ellipse( + a: float, + e: float, + i: float, + omega: float, + Omega: float, + M: float, +): + """ + Construct an ellipse of position vectors using the formalism defined + in chapter 1.2 of Morbidelli 2011, "Modern Celestial Mechanics" + + Paramters + ---------- + a : float + The semimajor axis of the orbit (in au) + + e : float + The eccentricty of the orbit + + i : float + The inclination of the orbit (in degrees) + + omega : float + The argument of perihelion of the orbit (in degrees) + + Omega : float + The longitude of the ascending node of the orbit (in degrees) + + M : float + The mean anomaly of the orbit (in degrees) + + Returns + -------- + r : numpy array + The rotated barycentric position vectors of the orbit ellipse + """ + # we want an eccentric anomaly array for the entire ellipse, but + # to avoid all orbits starting at perihelion in the plot we will do + # a quick numerical estimate of the eccentric anomaly from the input + # mean anomaly via fixed point iteration and then wrap the linspace + # around this value from 0 to 2pi + E_init = M # initial guesses using mean anomaly (shouldn't be too far off) + + for tries in range(100): + E_new = M + e * np.sin(E_init) + if np.abs(E_new - E_init) < 1e-8: + break + E_init = E_new + if tries == 99: + print('didnt converge') + raise ValueError("Did not converge for all M values") + + N = 100 + start = np.linspace(E_new, 360, N//2, endpoint=True) + end = np.linspace(0, E_new, N - len(start)) + E = np.concatenate((start, end)) * np.pi/180 + + # or, if you don't care, define eccentric anomaly for all ellipses + # E = np.linspace(0, 2*np.pi, 10000) + + # define position vectors in orthogonal reference frame with origin + # at ellipse main focus with q1 oriented towards periapsis (see chapter + # 1.2 of Morbidelli 2011, "Modern Celestial Mechanics" for details) + q1 = a * (np.cos(E) - e) + q2 = a * np.sqrt(1. - e**2) * np.sin(E) + + if type(q1) == float: + q = np.array([q1, q2, 0.]) + else: + q = np.array([q1, q2, np.zeros_like(q1)]) + + # define rotation matrix to map orthogonal reference frame to + # barycentric reference frame + c0 = np.cos(np.pi * Omega/180) + s0 = np.sin(np.pi * Omega/180) + co = np.cos(np.pi * omega/180) + so = np.sin(np.pi * omega/180) + ci = np.cos(np.pi * i/180) + si = np.sin(np.pi * i/180) + + R = np.array([ + [c0 * co - s0 * so * ci, -c0 * so - s0 * co * ci, s0 * si], + [s0 * co + c0 * so * ci, -s0 * so + c0 * co * ci, -c0 * si], + [si * so, si * co, ci ] + ]) + + # apply rotation matrix + r = np.einsum('ij,j...', R, q) + + return r + +def matplot_2D( + orb_array, + planets, + no_planets, + no_sun, + output, + fade +): + ''' + Create a 2D orbit distribution plot using matplot as a backend. + + Parameters + ------- + orb_array : structured numpy array + Numpy array containing orbital element information + + planets : list + List of desired planet orbits to be plotted + + no_planets : bool + Flag to turn off planet orbits + + no_sun : bool + Flag to turn of sun plotting + + output : str + Output file stem + + fade : bool + Flag to turn on object orbit fading + ''' + import matplotlib.pyplot as plt + from matplotlib.collections import LineCollection + + plt.style.use('./src/layup/matplot2d.mplstyle') + + fig, axs = plt.subplots(1, 2, layout='constrained', figsize=(20,9)) + fig.patch.set_facecolor('k') + + for ax in axs: + ax.tick_params( + labelbottom=True, + labeltop=True, + labelleft=True, + labelright=True, + colors='white' + ) + + ax.set_facecolor('k') + + for spine in ax.spines.values(): + spine.set_color('white') + + ax.xaxis.label.set_color('white') + ax.yaxis.label.set_color('white') + + + if np.max(orb_array['a'] < 2): + maxQ = np.max(orb_array['a'] * (1 + orb_array['e'])) + 0.2 + ax.set( + xlim=(-maxQ, maxQ), + ylim=(-maxQ, maxQ), + aspect="equal" + ) + elif np.max(orb_array['a'] <= 5): + maxQ = np.max(orb_array['a'] * (1 + orb_array['e'])) + 1 + ax.set( + xlim=(-maxQ, maxQ), + ylim=(-maxQ, maxQ), + aspect="equal" + ) + else: + maxQ = np.max(orb_array['a'] * (1 + orb_array['e'])) + 5 + ax.set( + xlim=(-maxQ, maxQ), + ylim=(-maxQ, maxQ), + aspect="equal" + ) + + # add objects + for obj in orb_array: + posvec = construct_ellipse(obj[2], obj[3], obj[4], obj[6], obj[5], obj[7]) + + # doing some matplot magic to make the orbits fade in opacity further away from the final position in array + points = np.array([posvec[:,0], posvec[:,1]]).T.reshape(-1, 1, 2) + segments = np.concatenate([points[:-1], points[1:]], axis=1) + if fade: + colors = [(14/255,84/255,118/255,a) for a in np.power(np.linspace(0, 1, len(segments)), 5)] + lc = LineCollection(segments, colors=colors, linewidth=1) + else: + lc = LineCollection(segments, colors=(14/225, 84/225, 118/225, 0.6), linewidth=1) + axs[0].add_collection(lc) + + points = np.array([posvec[:,0], posvec[:,2]]).T.reshape(-1, 1, 2) + segments = np.concatenate([points[:-1], points[1:]], axis=1) + if fade: + colors = [(14/255,84/255,118/255,a) for a in np.power(np.linspace(0, 1, len(segments)), 5)] + lc = LineCollection(segments, colors=colors, linewidth=1, zorder=10) + else: + lc = LineCollection(segments, colors=(14/225, 84/225, 118/225, 0.6), linewidth=1) + axs[1].add_collection(lc) + + # add the sun + if not no_sun: + axs[0].scatter(0, 0, s=100, color='xkcd:sunflower yellow') + + # add planets + if not no_planets: + planets_dic = { + 'Me':[0.387, 7.0, '#E7E8EC'], + 'V':[0.723,3.4, '#E39E1C'], + 'E':[1., 0, '#6B93D6'], + 'Ma':[1.524, 1.9, '#C1440E'], + 'J':[5.204, 1.3, '#D8CA9D'], + 'S':[9.582, 2.5, '#EAD6B8'], + 'U':[19.218, 0.8, '#D1E7E7'], + 'N':[30.11, 1.8, '#85ADDB'] + } + + theta = np.linspace(0, 2 * np.pi, 200) + dictfilt = lambda x, y: dict([ (i,x[i]) for i in x if i in set(y) ]) + + planets_dic = dictfilt(planets_dic, planets) + for _, (k, v) in enumerate(planets_dic.items()): + axs[0].add_patch(plt.Circle((0,0), v[0], color=v[2], fill=False, alpha=0.9, linestyle='dotted')) + # axs[0].text(v[0]+1, 0, k[0], color=v[2]) + + axs[1].plot(v[0]*np.cos(theta), v[0]*np.cos(theta)*np.sin(np.radians(v[1])), color=v[2], linestyle='dotted') + + + axs[0].set_xlabel('x [AU]') + axs[0].set_ylabel('y [AU]') + axs[1].set_xlabel('x [AU]') + axs[1].set_ylabel('z [AU]') + + axs[0].set_title("Bird's Eye", fontdict={"color":'white'}) + axs[1].set_title("Edge-on", fontdict={"color":'white'}) + + plt.savefig(output, dpi=300) + plt.show() + plt.close() + + +def matplot_3D( + orb_array, + planets, + no_planets, + no_sun, + output, + fade +): + ''' + Create a 3D orbit distribution plot using matplot as a backend. + + Parameters + ------- + orb_array : structured numpy array + Numpy array containing orbital element information + + planets : list + List of desired planet orbits to be plotted + + no_planets : bool + Flag to turn off planet orbits + + no_sun : bool + Flag to turn of sun plotting + + output : str + Output file stem + + fade : bool + Flag to turn on object orbit fading + ''' + import matplotlib.pyplot as plt + from mpl_toolkits.mplot3d.art3d import Line3DCollection + + fig = plt.figure(figsize=(15,9)) + fig.subplots_adjust(left=0, right=1, bottom=0, top=1) + axs = fig.add_subplot(111, projection='3d') + + fig.patch.set_facecolor('black') + axs.set_facecolor('black') + axs.xaxis.set_pane_color((.0, .0, .0, 1.0)) # RGBA: white + axs.yaxis.set_pane_color((.0, .0, .0, 1.0)) + axs.zaxis.set_pane_color((.0, .0, .0, 1.0)) + + axs.xaxis._axinfo['grid']['color'] = (0.1, 0.1, 0.1, 1) # grey + axs.yaxis._axinfo['grid']['color'] = (0.1, 0.1, 0.1, 1) # grey + axs.zaxis._axinfo['grid']['color'] = (0.1, 0.1, 0.1, 1) # grey + + axs.xaxis.label.set_color('white') + axs.yaxis.label.set_color('white') + axs.zaxis.label.set_color('white') + + axs.tick_params(axis='x',colors='white', length=0) + axs.tick_params(axis='y',colors='white', length=0) + axs.tick_params(axis='z',colors='white', length=0) + + for axis in [axs.xaxis, axs.yaxis, axs.zaxis]: + axis._axinfo['tick']['inward_factor'] = 0 + axis._axinfo['tick']['outward_factor'] = 0 + axis._axinfo['tick']['size'] = 0 + + axs.set_xlabel('x [au]') + axs.set_ylabel('y [au]') + axs.set_zlabel('z [au]') + + for obj in orb_array: + posvec = construct_ellipse(obj[2], obj[3], obj[4], obj[6], obj[5], obj[7]) + + points = np.array([posvec[:,0], posvec[:,1], posvec[:,2]]).T.reshape(-1, 1, 3) + segments = np.concatenate([points[:-1], points[1:]], axis=1) + if fade: + alpha = np.power(np.linspace(0, 1, len(segments)), 3) + colors = np.column_stack((np.ones(len(segments)), # R + np.ones(len(segments)), # G + np.ones(len(segments)), # B + alpha)) + lc = Line3DCollection(segments, colors=colors, linewidth=1) + else: + lc = Line3DCollection(segments, colors=(14/225, 84/225, 118/225, 0.6), linewidth=1) + axs.add_collection3d(lc) + + # add sun + if not no_sun: + axs.scatter(0, 0, s=100, color='xkcd:sunflower yellow') + + # add planets + if not no_planets: + planets_dic = { + 'Me':[0.387, 7.0, '#E7E8EC'], + 'V':[0.723,3.4, '#E39E1C'], + 'E':[1., 0, '#6B93D6'], + 'Ma':[1.524, 1.9, '#C1440E'], + 'J':[5.204, 1.3, '#D8CA9D'], + 'S':[9.582, 2.5, '#EAD6B8'], + 'U':[19.218, 0.8, '#D1E7E7'], + 'N':[30.11, 1.8, '#85ADDB'] + } + + theta = np.linspace(0, 2 * np.pi, 200) + dictfilt = lambda x, y: dict([ (i,x[i]) for i in x if i in set(y) ]) + + planets_dic = dictfilt(planets_dic, planets) + for _, (k, v) in enumerate(planets_dic.items()): + axs.plot(v[0]*np.cos(theta), v[0]*np.sin(theta)*np.cos(np.radians(v[1])), v[0]*np.cos(theta)*np.sin(np.radians(v[1])), color=v[2], alpha=0.9, linestyle='dotted') + axs.text(v[0], 0, 0, k[0], color=v[2]) + + + plt.savefig(output, dpi=300, bbox_inches='tight') + plt.show() + + +def plotly_3D( + orb_array, + planets, + no_planets, + no_sun, + output, + fade +): + ''' + Create a 3D orbit distribution plot using plotly as a backend. + + Parameters + ------- + orb_array : structured numpy array + Numpy array containing orbital element information + + planets : list + List of desired planet orbits to be plotted + + no_planets : bool + Flag to turn off planet orbits + + no_sun : bool + Flag to turn of sun plotting + + output : str + Output file stem + + fade : bool + Flag to turn on object orbit fading + ''' + import plotly.graph_objects as go + + fig = go.Figure() + + fig.update_layout( + plot_bgcolor='rgb(0,0,0)', + paper_bgcolor='rgb(0,0,0)', + scene = dict( + xaxis = dict( + backgroundcolor="rgba(0, 0, 0,0)", + gridcolor="rgb(106, 110, 117)", + title_font=dict(color="rgb(255,255,255)"), + tickfont=dict(color="rgb(255,255,255)"), + showbackground=True, + zerolinecolor="white", + ), + yaxis = dict( + backgroundcolor="rgba(0, 0, 0,0)", + gridcolor="rgb(106, 110, 117)", + title_font=dict(color="rgb(255,255,255)"), + tickfont=dict(color="rgb(255,255,255)"), + showbackground=True, + zerolinecolor="white" + ), + zaxis = dict( + backgroundcolor="rgba(0, 0, 0,0)", + gridcolor="rgb(106, 110, 117)", + title_font=dict(color="rgb(255,255,255)"), + tickfont=dict(color="rgb(255,255,255)"), + showbackground=True, + zerolinecolor="white", + ), + ), + ) + + fig.update_yaxes(title_font_color='white', color='white') + fig.update_xaxes(title_font_color='white', color='white') + + if not no_sun: + fig.add_trace(go.scatter(x=0, y=0, s=100, color='xkcd:sunflower yellow')) + + for obj in orb_array: + posvec = construct_ellipse(obj[2], obj[3], obj[4], obj[6], obj[5], obj[7]) + + normx = (posvec[:,0] - posvec[:,0].min()) / (posvec[:,0].max() - posvec[:,0].min()) + opacity = np.log1p(normx * 9) / np.log1p(9) + + fig.add_trace(go.Scatter( + x=posvec[:,0], y=posvec[:,1], + mode="lines", + line=dict(color=opacity, colorscale='gray', cmin=0, cmax=1,), + customdata=[obj[0]], + hovertemplate="%{customdata[0]}
(x: %{x:.2f}, y: %{y:.2f}, z: %{z:.2f})" + )) + + fig.update_layout(showlegend=False) + fig.show() + + + +def visualize_cli( + input: str, + output_file_stem: str, + planets: list, + file_format: Literal["csv", "hdf5"] = "csv", + input_format: Literal["BCART", "BCOM", "BKEP", "CART", "COM", "KEP"] = "BKEP", + backend: Literal["matplot", "plotly"] = "matplot", + dimensions: Literal["2D", "3D"] = "2D", + num_orbs: int = 1_000, + no_planets: bool = False, + no_sun: bool = False, + fade: bool = False +): + """ + Read in an orbits file and plot a user determined random number of orbits. + + Parameters + ----------- + inputs : str + The path to the input orbits file + + output : str + The output file stem name. + + planets : list + List of planets to be plotted from ['Me', 'V', 'E', 'Ma', J', 'S', 'U', 'N] + + file_format : str (default=csv) + The format of the input orbits file. Must be one of: 'csv' or 'hdf5' + + input_format : str (default=BKEP) + The orbit format of the input orbits. Must be one of: 'BCART', 'BCOM', 'BKEP', 'CART', 'COM', or 'KEP' + + backend : str (default=matplot) + The backend to use for plotting orbits. Must be one of: 'matplot' or 'plotly' + + dimensions : str (default=2D) + Number of dimensions for plotting purposes. Must be one of: '2D' or '3D' (TODO: add 4D plotting ;)) + + num_orbs : int (default=1_000) + The number of random orbits to plot + + no_planets : bool (default=False) + Flag to turn off the planet orbits + + no_sun : bool (default=False) + Flag to turn off the Sun + + fade : bool (default=False) + Flag to turn off object orbit fading + """ + input_file = Path(input) + + output_file = Path(f"{output_file_stem}") + + output_directory = output_file.parent.resolve() + if not output_directory.exists(): + logging.error(f"Output directory {output_directory} does not exist") + + # TODO: currently assumes BKEP input. need to convert if not in BKEP + required_columns = REQUIRED_COLUMN_NAMES[input_format] + + if file_format == "hdf5": + reader = HDF5DataReader( + input_file, + form_column_name="FORMAT", + required_column_names=required_columns, + ) + else: + reader = CSVDataReader( + input_file, + format_column_name="FORMAT", + required_column_names=required_columns, + ) + + random_orbs = reader.read_rows(block_start=0, block_size=1000) + random_orbs = random_orbs[np.random.choice(random_orbs.size, size=num_orbs, replace=False)] + + if backend == 'matplot': + if dimensions == '2D': + matplot_2D(random_orbs, planets, no_planets, no_sun, output_file_stem, fade) + else: + matplot_3D(random_orbs, planets, no_planets, no_sun, output_file_stem, fade) + elif backend == 'plotly': + if dimensions == '3D': + plotly_3D(random_orbs) + diff --git a/src/layup_cmdline/visualize.py b/src/layup_cmdline/visualize.py index 658e2747..5a057ef7 100644 --- a/src/layup_cmdline/visualize.py +++ b/src/layup_cmdline/visualize.py @@ -33,12 +33,18 @@ def main(): default="csv", required=False, ) + optional.add_argument( + "--orbit_type", + help="orbit reference frame of orbit from [COM, BCOM, KEP, BKEP, CART, BCART]", + dest="orbit_type", + type=str, + ) optional.add_argument( "-n", "--num", help="random number of orbits to take from input file", dest="n", - type=str, + type=int, default=1000, required=False, ) @@ -69,6 +75,38 @@ def main(): default="output", required=False, ) + optional.add_argument( + "--fade", + help="fade out the orbits of input objects", + dest="fade", + action='store_true' + ) + optional.add_argument( + "--planets", + "-p", + help="choose which planets to overplot", + dest="planets", + nargs='+', + required=False + ) + optional.add_argument( + "--no_planets", + help="overplot the planets. default is True", + dest="no_planets", + action="store_true" + ) + optional.add_argument( + "--no_sun", + help="overplot the sun. default is True", + dest="no_sun", + action="store_true" + ) + optional.add_argument( + "-f", + "--force", + action="store_true", + help="Overwrite output file", + ) args = parser.parse_args() @@ -76,7 +114,8 @@ def main(): def execute(args): - + from layup.visualize import visualize_cli + from layup.utilities.cli_utilities import warn_or_remove_file from layup.utilities.file_access_utils import find_file_or_exit, find_directory_or_exit # from layup.visualize import visualize_cli @@ -95,12 +134,28 @@ def execute(args): else: sys.exit("ERROR: File format must be 'csv' or 'hdf5'") + # check for overwriting output file + warn_or_remove_file(str(output_file), args.force, logger) + if args.d.upper() not in ["2D", "3D"]: sys.exit("ERROR: -d --dimensions must be '2D' or '3D'") if args.b not in ["matplot", "plotly"]: sys.exit("ERROR: -b --backend must be 'matplot' or 'plotly'") + visualize_cli( + input=args.input, + output_file_stem=args.o, + planets=args.planets, + input_format=args.orbit_type, + backend=args.b, + dimensions=args.d, + num_orbs=args.n, + no_planets=args.no_planets, + no_sun=args.no_sun, + fade=args.fade + ) + if __name__ == "__main__": main() From 03841d617eacde56283fd898478642ee495ae6b2 Mon Sep 17 00:00:00 2001 From: Joseph Murtagh Date: Fri, 2 May 2025 14:24:55 -0700 Subject: [PATCH 3/9] lint the files --- src/layup/visualize.py | 408 ++++++++++++++++----------------- src/layup_cmdline/visualize.py | 24 +- 2 files changed, 200 insertions(+), 232 deletions(-) diff --git a/src/layup/visualize.py b/src/layup/visualize.py index e0e0f751..625e2d98 100644 --- a/src/layup/visualize.py +++ b/src/layup/visualize.py @@ -22,12 +22,12 @@ def construct_ellipse( - a: float, - e: float, - i: float, - omega: float, - Omega: float, - M: float, + a: float, + e: float, + i: float, + omega: float, + Omega: float, + M: float, ): """ Construct an ellipse of position vectors using the formalism defined @@ -37,16 +37,16 @@ def construct_ellipse( ---------- a : float The semimajor axis of the orbit (in au) - + e : float The eccentricty of the orbit - + i : float The inclination of the orbit (in degrees) omega : float The argument of perihelion of the orbit (in degrees) - + Omega : float The longitude of the ascending node of the orbit (in degrees) @@ -61,9 +61,9 @@ def construct_ellipse( # we want an eccentric anomaly array for the entire ellipse, but # to avoid all orbits starting at perihelion in the plot we will do # a quick numerical estimate of the eccentric anomaly from the input - # mean anomaly via fixed point iteration and then wrap the linspace + # mean anomaly via fixed point iteration and then wrap the linspace # around this value from 0 to 2pi - E_init = M # initial guesses using mean anomaly (shouldn't be too far off) + E_init = M # initial guesses using mean anomaly (shouldn't be too far off) for tries in range(100): E_new = M + e * np.sin(E_init) @@ -71,63 +71,59 @@ def construct_ellipse( break E_init = E_new if tries == 99: - print('didnt converge') + print("didnt converge") raise ValueError("Did not converge for all M values") N = 100 - start = np.linspace(E_new, 360, N//2, endpoint=True) + start = np.linspace(E_new, 360, N // 2, endpoint=True) end = np.linspace(0, E_new, N - len(start)) - E = np.concatenate((start, end)) * np.pi/180 + E = np.concatenate((start, end)) * np.pi / 180 # or, if you don't care, define eccentric anomaly for all ellipses # E = np.linspace(0, 2*np.pi, 10000) - # define position vectors in orthogonal reference frame with origin + # define position vectors in orthogonal reference frame with origin # at ellipse main focus with q1 oriented towards periapsis (see chapter # 1.2 of Morbidelli 2011, "Modern Celestial Mechanics" for details) q1 = a * (np.cos(E) - e) - q2 = a * np.sqrt(1. - e**2) * np.sin(E) + q2 = a * np.sqrt(1.0 - e**2) * np.sin(E) if type(q1) == float: - q = np.array([q1, q2, 0.]) + q = np.array([q1, q2, 0.0]) else: q = np.array([q1, q2, np.zeros_like(q1)]) - # define rotation matrix to map orthogonal reference frame to + # define rotation matrix to map orthogonal reference frame to # barycentric reference frame - c0 = np.cos(np.pi * Omega/180) - s0 = np.sin(np.pi * Omega/180) - co = np.cos(np.pi * omega/180) - so = np.sin(np.pi * omega/180) - ci = np.cos(np.pi * i/180) - si = np.sin(np.pi * i/180) - - R = np.array([ - [c0 * co - s0 * so * ci, -c0 * so - s0 * co * ci, s0 * si], - [s0 * co + c0 * so * ci, -s0 * so + c0 * co * ci, -c0 * si], - [si * so, si * co, ci ] - ]) + c0 = np.cos(np.pi * Omega / 180) + s0 = np.sin(np.pi * Omega / 180) + co = np.cos(np.pi * omega / 180) + so = np.sin(np.pi * omega / 180) + ci = np.cos(np.pi * i / 180) + si = np.sin(np.pi * i / 180) + + R = np.array( + [ + [c0 * co - s0 * so * ci, -c0 * so - s0 * co * ci, s0 * si], + [s0 * co + c0 * so * ci, -s0 * so + c0 * co * ci, -c0 * si], + [si * so, si * co, ci], + ] + ) # apply rotation matrix - r = np.einsum('ij,j...', R, q) + r = np.einsum("ij,j...", R, q) return r -def matplot_2D( - orb_array, - planets, - no_planets, - no_sun, - output, - fade -): - ''' + +def matplot_2D(orb_array, planets, no_planets, no_sun, output, fade): + """ Create a 2D orbit distribution plot using matplot as a backend. Parameters ------- orb_array : structured numpy array - Numpy array containing orbital element information + Numpy array containing orbital element information planets : list List of desired planet orbits to be plotted @@ -143,134 +139,116 @@ def matplot_2D( fade : bool Flag to turn on object orbit fading - ''' + """ import matplotlib.pyplot as plt from matplotlib.collections import LineCollection - plt.style.use('./src/layup/matplot2d.mplstyle') + plt.style.use("./src/layup/matplot2d.mplstyle") - fig, axs = plt.subplots(1, 2, layout='constrained', figsize=(20,9)) - fig.patch.set_facecolor('k') + fig, axs = plt.subplots(1, 2, layout="constrained", figsize=(20, 9)) + fig.patch.set_facecolor("k") for ax in axs: - ax.tick_params( - labelbottom=True, - labeltop=True, - labelleft=True, - labelright=True, - colors='white' - ) + ax.tick_params(labelbottom=True, labeltop=True, labelleft=True, labelright=True, colors="white") - ax.set_facecolor('k') + ax.set_facecolor("k") for spine in ax.spines.values(): - spine.set_color('white') - - ax.xaxis.label.set_color('white') - ax.yaxis.label.set_color('white') - - - if np.max(orb_array['a'] < 2): - maxQ = np.max(orb_array['a'] * (1 + orb_array['e'])) + 0.2 - ax.set( - xlim=(-maxQ, maxQ), - ylim=(-maxQ, maxQ), - aspect="equal" - ) - elif np.max(orb_array['a'] <= 5): - maxQ = np.max(orb_array['a'] * (1 + orb_array['e'])) + 1 - ax.set( - xlim=(-maxQ, maxQ), - ylim=(-maxQ, maxQ), - aspect="equal" - ) + spine.set_color("white") + + ax.xaxis.label.set_color("white") + ax.yaxis.label.set_color("white") + + if np.max(orb_array["a"] < 2): + maxQ = np.max(orb_array["a"] * (1 + orb_array["e"])) + 0.2 + ax.set(xlim=(-maxQ, maxQ), ylim=(-maxQ, maxQ), aspect="equal") + elif np.max(orb_array["a"] <= 5): + maxQ = np.max(orb_array["a"] * (1 + orb_array["e"])) + 1 + ax.set(xlim=(-maxQ, maxQ), ylim=(-maxQ, maxQ), aspect="equal") else: - maxQ = np.max(orb_array['a'] * (1 + orb_array['e'])) + 5 - ax.set( - xlim=(-maxQ, maxQ), - ylim=(-maxQ, maxQ), - aspect="equal" - ) + maxQ = np.max(orb_array["a"] * (1 + orb_array["e"])) + 5 + ax.set(xlim=(-maxQ, maxQ), ylim=(-maxQ, maxQ), aspect="equal") # add objects for obj in orb_array: posvec = construct_ellipse(obj[2], obj[3], obj[4], obj[6], obj[5], obj[7]) # doing some matplot magic to make the orbits fade in opacity further away from the final position in array - points = np.array([posvec[:,0], posvec[:,1]]).T.reshape(-1, 1, 2) + points = np.array([posvec[:, 0], posvec[:, 1]]).T.reshape(-1, 1, 2) segments = np.concatenate([points[:-1], points[1:]], axis=1) if fade: - colors = [(14/255,84/255,118/255,a) for a in np.power(np.linspace(0, 1, len(segments)), 5)] + colors = [ + (14 / 255, 84 / 255, 118 / 255, a) for a in np.power(np.linspace(0, 1, len(segments)), 5) + ] lc = LineCollection(segments, colors=colors, linewidth=1) else: - lc = LineCollection(segments, colors=(14/225, 84/225, 118/225, 0.6), linewidth=1) + lc = LineCollection(segments, colors=(14 / 225, 84 / 225, 118 / 225, 0.6), linewidth=1) axs[0].add_collection(lc) - points = np.array([posvec[:,0], posvec[:,2]]).T.reshape(-1, 1, 2) + points = np.array([posvec[:, 0], posvec[:, 2]]).T.reshape(-1, 1, 2) segments = np.concatenate([points[:-1], points[1:]], axis=1) if fade: - colors = [(14/255,84/255,118/255,a) for a in np.power(np.linspace(0, 1, len(segments)), 5)] + colors = [ + (14 / 255, 84 / 255, 118 / 255, a) for a in np.power(np.linspace(0, 1, len(segments)), 5) + ] lc = LineCollection(segments, colors=colors, linewidth=1, zorder=10) else: - lc = LineCollection(segments, colors=(14/225, 84/225, 118/225, 0.6), linewidth=1) + lc = LineCollection(segments, colors=(14 / 225, 84 / 225, 118 / 225, 0.6), linewidth=1) axs[1].add_collection(lc) - + # add the sun if not no_sun: - axs[0].scatter(0, 0, s=100, color='xkcd:sunflower yellow') + axs[0].scatter(0, 0, s=100, color="xkcd:sunflower yellow") # add planets if not no_planets: planets_dic = { - 'Me':[0.387, 7.0, '#E7E8EC'], - 'V':[0.723,3.4, '#E39E1C'], - 'E':[1., 0, '#6B93D6'], - 'Ma':[1.524, 1.9, '#C1440E'], - 'J':[5.204, 1.3, '#D8CA9D'], - 'S':[9.582, 2.5, '#EAD6B8'], - 'U':[19.218, 0.8, '#D1E7E7'], - 'N':[30.11, 1.8, '#85ADDB'] + "Me": [0.387, 7.0, "#E7E8EC"], + "V": [0.723, 3.4, "#E39E1C"], + "E": [1.0, 0, "#6B93D6"], + "Ma": [1.524, 1.9, "#C1440E"], + "J": [5.204, 1.3, "#D8CA9D"], + "S": [9.582, 2.5, "#EAD6B8"], + "U": [19.218, 0.8, "#D1E7E7"], + "N": [30.11, 1.8, "#85ADDB"], } theta = np.linspace(0, 2 * np.pi, 200) - dictfilt = lambda x, y: dict([ (i,x[i]) for i in x if i in set(y) ]) + dictfilt = lambda x, y: dict([(i, x[i]) for i in x if i in set(y)]) planets_dic = dictfilt(planets_dic, planets) for _, (k, v) in enumerate(planets_dic.items()): - axs[0].add_patch(plt.Circle((0,0), v[0], color=v[2], fill=False, alpha=0.9, linestyle='dotted')) + axs[0].add_patch(plt.Circle((0, 0), v[0], color=v[2], fill=False, alpha=0.9, linestyle="dotted")) # axs[0].text(v[0]+1, 0, k[0], color=v[2]) - axs[1].plot(v[0]*np.cos(theta), v[0]*np.cos(theta)*np.sin(np.radians(v[1])), color=v[2], linestyle='dotted') - + axs[1].plot( + v[0] * np.cos(theta), + v[0] * np.cos(theta) * np.sin(np.radians(v[1])), + color=v[2], + linestyle="dotted", + ) - axs[0].set_xlabel('x [AU]') - axs[0].set_ylabel('y [AU]') - axs[1].set_xlabel('x [AU]') - axs[1].set_ylabel('z [AU]') + axs[0].set_xlabel("x [AU]") + axs[0].set_ylabel("y [AU]") + axs[1].set_xlabel("x [AU]") + axs[1].set_ylabel("z [AU]") - axs[0].set_title("Bird's Eye", fontdict={"color":'white'}) - axs[1].set_title("Edge-on", fontdict={"color":'white'}) + axs[0].set_title("Bird's Eye", fontdict={"color": "white"}) + axs[1].set_title("Edge-on", fontdict={"color": "white"}) plt.savefig(output, dpi=300) plt.show() plt.close() -def matplot_3D( - orb_array, - planets, - no_planets, - no_sun, - output, - fade -): - ''' +def matplot_3D(orb_array, planets, no_planets, no_sun, output, fade): + """ Create a 3D orbit distribution plot using matplot as a backend. Parameters ------- orb_array : structured numpy array - Numpy array containing orbital element information + Numpy array containing orbital element information planets : list List of desired planet orbits to be plotted @@ -286,102 +264,100 @@ def matplot_3D( fade : bool Flag to turn on object orbit fading - ''' + """ import matplotlib.pyplot as plt from mpl_toolkits.mplot3d.art3d import Line3DCollection - fig = plt.figure(figsize=(15,9)) + fig = plt.figure(figsize=(15, 9)) fig.subplots_adjust(left=0, right=1, bottom=0, top=1) - axs = fig.add_subplot(111, projection='3d') + axs = fig.add_subplot(111, projection="3d") - fig.patch.set_facecolor('black') - axs.set_facecolor('black') - axs.xaxis.set_pane_color((.0, .0, .0, 1.0)) # RGBA: white - axs.yaxis.set_pane_color((.0, .0, .0, 1.0)) - axs.zaxis.set_pane_color((.0, .0, .0, 1.0)) + fig.patch.set_facecolor("black") + axs.set_facecolor("black") + axs.xaxis.set_pane_color((0.0, 0.0, 0.0, 1.0)) # RGBA: white + axs.yaxis.set_pane_color((0.0, 0.0, 0.0, 1.0)) + axs.zaxis.set_pane_color((0.0, 0.0, 0.0, 1.0)) - axs.xaxis._axinfo['grid']['color'] = (0.1, 0.1, 0.1, 1) # grey - axs.yaxis._axinfo['grid']['color'] = (0.1, 0.1, 0.1, 1) # grey - axs.zaxis._axinfo['grid']['color'] = (0.1, 0.1, 0.1, 1) # grey + axs.xaxis._axinfo["grid"]["color"] = (0.1, 0.1, 0.1, 1) # grey + axs.yaxis._axinfo["grid"]["color"] = (0.1, 0.1, 0.1, 1) # grey + axs.zaxis._axinfo["grid"]["color"] = (0.1, 0.1, 0.1, 1) # grey - axs.xaxis.label.set_color('white') - axs.yaxis.label.set_color('white') - axs.zaxis.label.set_color('white') + axs.xaxis.label.set_color("white") + axs.yaxis.label.set_color("white") + axs.zaxis.label.set_color("white") - axs.tick_params(axis='x',colors='white', length=0) - axs.tick_params(axis='y',colors='white', length=0) - axs.tick_params(axis='z',colors='white', length=0) + axs.tick_params(axis="x", colors="white", length=0) + axs.tick_params(axis="y", colors="white", length=0) + axs.tick_params(axis="z", colors="white", length=0) for axis in [axs.xaxis, axs.yaxis, axs.zaxis]: - axis._axinfo['tick']['inward_factor'] = 0 - axis._axinfo['tick']['outward_factor'] = 0 - axis._axinfo['tick']['size'] = 0 + axis._axinfo["tick"]["inward_factor"] = 0 + axis._axinfo["tick"]["outward_factor"] = 0 + axis._axinfo["tick"]["size"] = 0 - axs.set_xlabel('x [au]') - axs.set_ylabel('y [au]') - axs.set_zlabel('z [au]') + axs.set_xlabel("x [au]") + axs.set_ylabel("y [au]") + axs.set_zlabel("z [au]") for obj in orb_array: posvec = construct_ellipse(obj[2], obj[3], obj[4], obj[6], obj[5], obj[7]) - points = np.array([posvec[:,0], posvec[:,1], posvec[:,2]]).T.reshape(-1, 1, 3) + points = np.array([posvec[:, 0], posvec[:, 1], posvec[:, 2]]).T.reshape(-1, 1, 3) segments = np.concatenate([points[:-1], points[1:]], axis=1) if fade: alpha = np.power(np.linspace(0, 1, len(segments)), 3) - colors = np.column_stack((np.ones(len(segments)), # R - np.ones(len(segments)), # G - np.ones(len(segments)), # B - alpha)) - lc = Line3DCollection(segments, colors=colors, linewidth=1) + colors = np.column_stack( + (np.ones(len(segments)), np.ones(len(segments)), np.ones(len(segments)), alpha) # R # G # B + ) + lc = Line3DCollection(segments, colors=colors, linewidth=1) else: - lc = Line3DCollection(segments, colors=(14/225, 84/225, 118/225, 0.6), linewidth=1) - axs.add_collection3d(lc) + lc = Line3DCollection(segments, colors=(14 / 225, 84 / 225, 118 / 225, 0.6), linewidth=1) + axs.add_collection3d(lc) # add sun if not no_sun: - axs.scatter(0, 0, s=100, color='xkcd:sunflower yellow') - + axs.scatter(0, 0, s=100, color="xkcd:sunflower yellow") + # add planets if not no_planets: planets_dic = { - 'Me':[0.387, 7.0, '#E7E8EC'], - 'V':[0.723,3.4, '#E39E1C'], - 'E':[1., 0, '#6B93D6'], - 'Ma':[1.524, 1.9, '#C1440E'], - 'J':[5.204, 1.3, '#D8CA9D'], - 'S':[9.582, 2.5, '#EAD6B8'], - 'U':[19.218, 0.8, '#D1E7E7'], - 'N':[30.11, 1.8, '#85ADDB'] + "Me": [0.387, 7.0, "#E7E8EC"], + "V": [0.723, 3.4, "#E39E1C"], + "E": [1.0, 0, "#6B93D6"], + "Ma": [1.524, 1.9, "#C1440E"], + "J": [5.204, 1.3, "#D8CA9D"], + "S": [9.582, 2.5, "#EAD6B8"], + "U": [19.218, 0.8, "#D1E7E7"], + "N": [30.11, 1.8, "#85ADDB"], } theta = np.linspace(0, 2 * np.pi, 200) - dictfilt = lambda x, y: dict([ (i,x[i]) for i in x if i in set(y) ]) + dictfilt = lambda x, y: dict([(i, x[i]) for i in x if i in set(y)]) planets_dic = dictfilt(planets_dic, planets) for _, (k, v) in enumerate(planets_dic.items()): - axs.plot(v[0]*np.cos(theta), v[0]*np.sin(theta)*np.cos(np.radians(v[1])), v[0]*np.cos(theta)*np.sin(np.radians(v[1])), color=v[2], alpha=0.9, linestyle='dotted') + axs.plot( + v[0] * np.cos(theta), + v[0] * np.sin(theta) * np.cos(np.radians(v[1])), + v[0] * np.cos(theta) * np.sin(np.radians(v[1])), + color=v[2], + alpha=0.9, + linestyle="dotted", + ) axs.text(v[0], 0, 0, k[0], color=v[2]) - - plt.savefig(output, dpi=300, bbox_inches='tight') + plt.savefig(output, dpi=300, bbox_inches="tight") plt.show() -def plotly_3D( - orb_array, - planets, - no_planets, - no_sun, - output, - fade -): - ''' +def plotly_3D(orb_array, planets, no_planets, no_sun, output, fade): + """ Create a 3D orbit distribution plot using plotly as a backend. Parameters ------- orb_array : structured numpy array - Numpy array containing orbital element information + Numpy array containing orbital element information planets : list List of desired planet orbits to be plotted @@ -397,79 +373,86 @@ def plotly_3D( fade : bool Flag to turn on object orbit fading - ''' + """ import plotly.graph_objects as go fig = go.Figure() fig.update_layout( - plot_bgcolor='rgb(0,0,0)', - paper_bgcolor='rgb(0,0,0)', - scene = dict( - xaxis = dict( + plot_bgcolor="rgb(0,0,0)", + paper_bgcolor="rgb(0,0,0)", + scene=dict( + xaxis=dict( backgroundcolor="rgba(0, 0, 0,0)", gridcolor="rgb(106, 110, 117)", title_font=dict(color="rgb(255,255,255)"), tickfont=dict(color="rgb(255,255,255)"), showbackground=True, zerolinecolor="white", - ), - yaxis = dict( + ), + yaxis=dict( backgroundcolor="rgba(0, 0, 0,0)", gridcolor="rgb(106, 110, 117)", title_font=dict(color="rgb(255,255,255)"), tickfont=dict(color="rgb(255,255,255)"), showbackground=True, - zerolinecolor="white" - ), - zaxis = dict( + zerolinecolor="white", + ), + zaxis=dict( backgroundcolor="rgba(0, 0, 0,0)", gridcolor="rgb(106, 110, 117)", title_font=dict(color="rgb(255,255,255)"), tickfont=dict(color="rgb(255,255,255)"), - showbackground=True, + showbackground=True, zerolinecolor="white", - ), ), - ) - - fig.update_yaxes(title_font_color='white', color='white') - fig.update_xaxes(title_font_color='white', color='white') + ), + ) + + fig.update_yaxes(title_font_color="white", color="white") + fig.update_xaxes(title_font_color="white", color="white") if not no_sun: - fig.add_trace(go.scatter(x=0, y=0, s=100, color='xkcd:sunflower yellow')) + fig.add_trace(go.scatter(x=0, y=0, s=100, color="xkcd:sunflower yellow")) for obj in orb_array: posvec = construct_ellipse(obj[2], obj[3], obj[4], obj[6], obj[5], obj[7]) - normx = (posvec[:,0] - posvec[:,0].min()) / (posvec[:,0].max() - posvec[:,0].min()) - opacity = np.log1p(normx * 9) / np.log1p(9) + normx = (posvec[:, 0] - posvec[:, 0].min()) / (posvec[:, 0].max() - posvec[:, 0].min()) + opacity = np.log1p(normx * 9) / np.log1p(9) - fig.add_trace(go.Scatter( - x=posvec[:,0], y=posvec[:,1], + fig.add_trace( + go.Scatter( + x=posvec[:, 0], + y=posvec[:, 1], mode="lines", - line=dict(color=opacity, colorscale='gray', cmin=0, cmax=1,), + line=dict( + color=opacity, + colorscale="gray", + cmin=0, + cmax=1, + ), customdata=[obj[0]], - hovertemplate="%{customdata[0]}
(x: %{x:.2f}, y: %{y:.2f}, z: %{z:.2f})" - )) + hovertemplate="%{customdata[0]}
(x: %{x:.2f}, y: %{y:.2f}, z: %{z:.2f})", + ) + ) fig.update_layout(showlegend=False) fig.show() - def visualize_cli( - input: str, - output_file_stem: str, - planets: list, - file_format: Literal["csv", "hdf5"] = "csv", - input_format: Literal["BCART", "BCOM", "BKEP", "CART", "COM", "KEP"] = "BKEP", - backend: Literal["matplot", "plotly"] = "matplot", - dimensions: Literal["2D", "3D"] = "2D", - num_orbs: int = 1_000, - no_planets: bool = False, - no_sun: bool = False, - fade: bool = False + input: str, + output_file_stem: str, + planets: list, + file_format: Literal["csv", "hdf5"] = "csv", + input_format: Literal["BCART", "BCOM", "BKEP", "CART", "COM", "KEP"] = "BKEP", + backend: Literal["matplot", "plotly"] = "matplot", + dimensions: Literal["2D", "3D"] = "2D", + num_orbs: int = 1_000, + no_planets: bool = False, + no_sun: bool = False, + fade: bool = False, ): """ Read in an orbits file and plot a user determined random number of orbits. @@ -478,7 +461,7 @@ def visualize_cli( ----------- inputs : str The path to the input orbits file - + output : str The output file stem name. @@ -490,7 +473,7 @@ def visualize_cli( input_format : str (default=BKEP) The orbit format of the input orbits. Must be one of: 'BCART', 'BCOM', 'BKEP', 'CART', 'COM', or 'KEP' - + backend : str (default=matplot) The backend to use for plotting orbits. Must be one of: 'matplot' or 'plotly' @@ -536,12 +519,11 @@ def visualize_cli( random_orbs = reader.read_rows(block_start=0, block_size=1000) random_orbs = random_orbs[np.random.choice(random_orbs.size, size=num_orbs, replace=False)] - if backend == 'matplot': - if dimensions == '2D': + if backend == "matplot": + if dimensions == "2D": matplot_2D(random_orbs, planets, no_planets, no_sun, output_file_stem, fade) else: matplot_3D(random_orbs, planets, no_planets, no_sun, output_file_stem, fade) - elif backend == 'plotly': - if dimensions == '3D': + elif backend == "plotly": + if dimensions == "3D": plotly_3D(random_orbs) - diff --git a/src/layup_cmdline/visualize.py b/src/layup_cmdline/visualize.py index 5a057ef7..0601fce5 100644 --- a/src/layup_cmdline/visualize.py +++ b/src/layup_cmdline/visualize.py @@ -76,30 +76,16 @@ def main(): required=False, ) optional.add_argument( - "--fade", - help="fade out the orbits of input objects", - dest="fade", - action='store_true' + "--fade", help="fade out the orbits of input objects", dest="fade", action="store_true" ) optional.add_argument( - "--planets", - "-p", - help="choose which planets to overplot", - dest="planets", - nargs='+', - required=False + "--planets", "-p", help="choose which planets to overplot", dest="planets", nargs="+", required=False ) optional.add_argument( - "--no_planets", - help="overplot the planets. default is True", - dest="no_planets", - action="store_true" + "--no_planets", help="overplot the planets. default is True", dest="no_planets", action="store_true" ) optional.add_argument( - "--no_sun", - help="overplot the sun. default is True", - dest="no_sun", - action="store_true" + "--no_sun", help="overplot the sun. default is True", dest="no_sun", action="store_true" ) optional.add_argument( "-f", @@ -153,7 +139,7 @@ def execute(args): num_orbs=args.n, no_planets=args.no_planets, no_sun=args.no_sun, - fade=args.fade + fade=args.fade, ) From 7793c19c622c4d03647f7cf2cc9015d8cfd29264 Mon Sep 17 00:00:00 2001 From: Joseph Murtagh Date: Fri, 2 May 2025 14:30:06 -0700 Subject: [PATCH 4/9] lint file --- src/layup_cmdline/visualize.py | 7 ++++++- 1 file changed, 6 insertions(+), 1 deletion(-) diff --git a/src/layup_cmdline/visualize.py b/src/layup_cmdline/visualize.py index 0601fce5..6a35bcfe 100644 --- a/src/layup_cmdline/visualize.py +++ b/src/layup_cmdline/visualize.py @@ -79,7 +79,12 @@ def main(): "--fade", help="fade out the orbits of input objects", dest="fade", action="store_true" ) optional.add_argument( - "--planets", "-p", help="choose which planets to overplot", dest="planets", nargs="+", required=False + "--planets", + "-p", + help="choose which planets to overplot. must be from ['Me', 'V', 'E', 'Ma', 'J', 'S', 'U', 'N']", + dest="planets", + nargs="+", + required=False, ) optional.add_argument( "--no_planets", help="overplot the planets. default is True", dest="no_planets", action="store_true" From ab33cff7b06ca204590f3acd38c4f8831530b90b Mon Sep 17 00:00:00 2001 From: Joseph Murtagh Date: Fri, 2 May 2025 14:40:46 -0700 Subject: [PATCH 5/9] remove old file --- src/layup/visualise.py | 372 ----------------------------------------- 1 file changed, 372 deletions(-) delete mode 100644 src/layup/visualise.py diff --git a/src/layup/visualise.py b/src/layup/visualise.py deleted file mode 100644 index 90ed63bb..00000000 --- a/src/layup/visualise.py +++ /dev/null @@ -1,372 +0,0 @@ -import logging -import os -from pathlib import Path -from typing import Literal - -import numpy as np -import matplotlib.pyplot as plt -from matplotlib.collections import LineCollection - -from layup.utilities.file_io.CSVReader import CSVDataReader -from layup.utilities.file_io.HDF5Reader import HDF5DataReader - -logger = logging.getLogger(__name__) - -# Required column names for each orbit format -REQUIRED_COLUMN_NAMES = { - "BCART": ["ObjID", "FORMAT", "x", "y", "z", "xdot", "ydot", "zdot", "epochMJD_TDB"], - "BCOM": ["ObjID", "FORMAT", "q", "e", "inc", "node", "argPeri", "t_p_MJD_TDB", "epochMJD_TDB"], - "BKEP": ["ObjID", "FORMAT", "a", "e", "inc", "node", "argPeri", "ma", "epochMJD_TDB"], - "CART": ["ObjID", "FORMAT", "x", "y", "z", "xdot", "ydot", "zdot", "epochMJD_TDB"], - "COM": ["ObjID", "FORMAT", "q", "e", "inc", "node", "argPeri", "t_p_MJD_TDB", "epochMJD_TDB"], - "KEP": ["ObjID", "FORMAT", "a", "e", "inc", "node", "argPeri", "ma", "epochMJD_TDB"], -} - - -def construct_ellipse( - a: float, - e: float, - i: float, - omega: float, - Omega: float, - M: float, -): - """ - Construct an ellipse of position vectors using the formalism defined - in chapter 1.2 of Morbidelli 2011, "Modern Celestial Mechanics" - - Paramters - ---------- - a : float - The semimajor axis of the orbit (in au) - - e : float - The eccentricty of the orbit - - i : float - The inclination of the orbit (in degrees) - - omega : float - The argument of perihelion of the orbit (in degrees) - - Omega : float - The longitude of the ascending node of the orbit (in degrees) - - M : float - The mean anomaly of the orbit (in degrees) - - Returns - -------- - r : numpy array - The rotated barycentric position vectors of the orbit ellipse - """ - # we want an eccentrict anomaly array for the entire ellipse, but - # to avoid all orbits starting at perihelion in the plot we will do - # a quick numerical estimate of the eccentric anomaly from the input - # mean anomaly via fixed point iteration and then wrap the linspace - # around this value from 0 to 2pi - E_init = M # initial guesses using mean anomaly (shouldn't be too far off) - - for tries in range(100): - E_new = M + e * np.sin(E_init) - if np.abs(E_new - E_init) < 1e-8: - break - E_init = E_new - if tries == 99: - print('didnt converge') - raise ValueError("Did not converge for all M values") - - N = 2000 - start = np.linspace(E_new, 360, N//2, endpoint=True) - end = np.linspace(0, E_new, N - len(start)) - E = np.concatenate((start, end)) * np.pi/180 - - # or, if you don't care, define eccentric anomaly for all ellipses - # E = np.linspace(0, 2*np.pi, 10000) - - # define position vectors in orthogonal reference frame with origin - # at ellipse main focus with q1 oriented towards periapsis (see chapter - # 1.2 of Morbidelli 2011, "Modern Celestial Mechanics" for details) - q1 = a * (np.cos(E) - e) - q2 = a * np.sqrt(1. - e**2) * np.sin(E) - - if type(q1) == float: - q = np.array([q1, q2, 0.]) - else: - q = np.array([q1, q2, np.zeros_like(q1)]) - - # define rotation matrix to map orthogonal reference frame to - # barycentric reference frame - c0 = np.cos(np.pi * Omega/180) - s0 = np.sin(np.pi * Omega/180) - co = np.cos(np.pi * omega/180) - so = np.sin(np.pi * omega/180) - ci = np.cos(np.pi * i/180) - si = np.sin(np.pi * i/180) - - R = np.array([ - [c0 * co - s0 * so * ci, -c0 * so - s0 * co * ci, s0 * si], - [s0 * co + c0 * so * ci, -s0 * so + c0 * co * ci, -c0 * si], - [si * so, si * co, ci ] - ]) - - # apply rotation matrix - r = np.einsum('ij,j...', R, q) - - return r - -def plot_ellipse( - orb_array, - plot_type, - output -): - """ - Handle the plotting of the orbit ellipses depending on whether the user wants - matplot flat 2D plots, or 3D interactive plots with plotly - - Parameters - ----------- - orb_array : numpy structured array - The orbits as read in by data readers - plot_type : str - The plotting engine to use. Can be either 'matplot' or 'plotly' currently - output : str - The output file stem. - """ - - if plot_type == 'matplot': - import matplotlib.pyplot as plt - from matplotlib.collections import LineCollection - - plt.rcParams.update( - {'axes.linewidth': 1.875, - 'grid.linewidth': 1.5, - 'lines.linewidth': 2.25, - 'lines.markersize': 9.0, - 'patch.linewidth': 1.5, - 'xtick.major.width': 1.875, - 'ytick.major.width': 1.875, - 'xtick.minor.width': 1.5, - 'ytick.minor.width': 1.5, - 'xtick.major.size': 9.0, - 'ytick.major.size': 9.0, - 'xtick.minor.size': 6.0, - 'ytick.minor.size': 6.0, - 'font.size': 18.0, - 'axes.labelsize': 18.0, - 'axes.titlesize': 18.0, - 'xtick.labelsize': 16.5, - 'ytick.labelsize': 16.5, - 'legend.fontsize': 16.5, - 'legend.title_fontsize': 18.0} - ) - - # TODO: make a y-z plot as well as an x-y plot and have them side-by-side? - fig, axs = plt.subplots(figsize=(15,9)) - - axs.tick_params( - labelbottom=True, - labeltop=True, - labelleft=True, - labelright=True - ) - - fig.patch.set_facecolor('k') - axs.set_facecolor('k') - - axs.spines['bottom'].set_color('white') - axs.spines['top'].set_color('white') - axs.spines['right'].set_color('white') - axs.spines['left'].set_color('white') - - axs.tick_params(axis='x', colors='white') - axs.tick_params(axis='y', colors='white') - - axs.xaxis.label.set_color('white') - axs.yaxis.label.set_color('white') - - axs.set_xlabel('x [AU]') - axs.set_ylabel('y [AU]') - - for obj in orb_array: - posvec = construct_ellipse(obj[2], obj[3], obj[4], obj[6], obj[5], obj[7]) - - # doing some matplot magic to make the orbits fade in opacity further away from the final position in array - points = np.array([posvec[:,0], posvec[:,1]]).T.reshape(-1, 1, 2) - segments = np.concatenate([points[:-1], points[1:]], axis=1) - colors = [(1,1,1,a) for a in np.power(np.linspace(0, 1, len(segments)), 3)] - - lc = LineCollection(segments, colors=colors, linewidth=1) - axs.add_collection(lc) - - # TODO: read in semimajor values and determine max, use that to figure out plot extents? - x_extent = 50 - y_extent = 50 - - axs.set( - xlim=(-x_extent,x_extent), - ylim=(-y_extent,y_extent), - aspect="equal" if x_extent == y_extent else "auto" - ) - - # TODO: redo file saving and make it go to user defined path - plt.savefig(output, dpi=300) - plt.show() - plt.close() - - - elif plot_type == 'plotly': - import plotly.graph_objects as go - - fig = go.Figure() - - fig.update_layout(plot_bgcolor='rgb(0,0,0)', - paper_bgcolor='rgb(0,0,0)', - scene = dict( - xaxis = dict( - backgroundcolor="rgba(0, 0, 0,0)", - gridcolor="rgb(106, 110, 117)", - title_font=dict(color="rgb(255,255,255)"), - tickfont=dict(color="rgb(255,255,255)"), - showbackground=True, - zerolinecolor="white", - ), - yaxis = dict( - backgroundcolor="rgba(0, 0, 0,0)", - gridcolor="rgb(106, 110, 117)", - title_font=dict(color="rgb(255,255,255)"), - tickfont=dict(color="rgb(255,255,255)"), - showbackground=True, - zerolinecolor="white" - ), - zaxis = dict( - backgroundcolor="rgba(0, 0, 0,0)", - gridcolor="rgb(106, 110, 117)", - title_font=dict(color="rgb(255,255,255)"), - tickfont=dict(color="rgb(255,255,255)"), - showbackground=True, - zerolinecolor="white", - ), - ), - ) - fig.update_yaxes(title_font_color='white', color='white') - fig.update_xaxes(title_font_color='white', color='white') - - - xs = [] - ys = [] - zs = [] - names = [] - - for obj in orb_array: - posvec = construct_ellipse(obj[2], obj[3], obj[4], obj[6], obj[5], obj[7]) - xs.append(posvec[-1,0]) - ys.append(posvec[-1,1]) - zs.append(posvec[-1,2]) - names.append(obj[0]) - - normx = (posvec[:,0] - posvec[:,0].min()) / (posvec[:,0].max() - posvec[:,0].min()) - opacity = np.log1p(normx * 9) / np.log1p(9) - - # TODO: fix hover boxes to reflect objID properly - fig.add_trace(go.Scatter3d( - x=posvec[:,0], y=posvec[:,1], z=posvec[:,2], - mode="lines+markers", - marker=dict(color=opacity, size=5), - line=dict(color='white'), - customdata=[obj[0]], - hovertemplate="%{customdata[0]}
(x: %{x:.2f}, y: %{y:.2f}, z: %{z:.2f})" - )) - - # TODO: add object a/e/i as hover tooltip info? - fig.add_trace(go.Scatter3d( - x=xs, y=ys, z=zs, - mode="markers", - marker=dict(color='rgba(255, 0, 0, 1)', size=5), - customdata=names, # Pass planet name as custom data - hovertemplate="%{customdata}
(x: %{x:.2f}, y: %{y:.2f}, z: %{z:.2f})" # Define hover format - )) - - fig.update_layout(showlegend=False) - fig.show() - - # TODO: add in file saving? or can user do it from plotly output - - - -def visualise_cli( - input: str, - output_file_stem: str, - plot_type: Literal["matplot", "plotly"] = "matplot", - input_format: Literal["BCART", "BCOM", "BKEP", "CART", "COM", "KEP"] = "BKEP", - file_format: Literal["csv", "hdf5"] = "csv", - num_orbs: int = 1_000, -): - """ - Read in an orbits file and plot a user determined random number of orbits. - - Parameters - ----------- - inputs : str - The path to the input orbits file - - output : str - The output file stem name. - - plot_type : str (default=matplot) - The backend to use for plotting orbits. Must be one of: 'matplot' or 'plotly' - - input_format : str (default=BKEP) - The orbit format of the input orbits. Must be one of: 'BCART', 'BCOM', 'BKEP', 'CART', 'COM', or 'KEP' - - file_format : str (default=csv) - The format of the input orbits file. Must be one of: 'csv' or 'hdf5' - - num_orbs : int, optional (default=1_000) - The number of random orbits to plot - """ - input_file = Path(input) - - output_file = Path(f"{output_file_stem}.pdf") - - output_directory = output_file.parent.resolve() - if not output_directory.exists(): - logging.error(f"Output directory {output_directory} does not exist") - - # TODO: currently assumes BKEP input. need to convert if not in BKEP - required_columns = REQUIRED_COLUMN_NAMES[input_format] - - if file_format == "hdf5": - reader = HDF5DataReader( - input_file, - form_column_name="FORMAT", - required_column_names=required_columns, - ) - else: - reader = CSVDataReader( - input_file, - format_column_name="FORMAT", - required_column_names=required_columns, - ) - - random_orbs = reader.read_rows(block_start=0, block_size=1000) - random_orbs = random_orbs[np.random.choice(random_orbs.size, size=num_orbs, replace=False)] - - plot_ellipse(random_orbs, plot_type, output_file) - - - -def main(): - - # TODO: add in actual layup argument stuff to this, instead of hardcoded values/paths - visualise_cli( - input='/Users/josephmurtagh/Downloads/cent_orbs.csv', - output_file_stem='testplot', - plot_type='matplot', - file_format='csv', - num_orbs=100, - ) - -if __name__ == "__main__": - main() - From 1e70a38b3830d4c3a837158cf26472e3085dc735 Mon Sep 17 00:00:00 2001 From: Joseph Murtagh Date: Tue, 20 May 2025 16:46:23 -0400 Subject: [PATCH 6/9] remove .mplstyle file, add same funcitonality as function decorator --- src/layup/matplot2d.mplstyle | 20 -------------------- src/layup/visualize.py | 26 +++++++++++++++++++++++--- 2 files changed, 23 insertions(+), 23 deletions(-) delete mode 100644 src/layup/matplot2d.mplstyle diff --git a/src/layup/matplot2d.mplstyle b/src/layup/matplot2d.mplstyle deleted file mode 100644 index c4debece..00000000 --- a/src/layup/matplot2d.mplstyle +++ /dev/null @@ -1,20 +0,0 @@ -axes.linewidth : 1.875 -grid.linewidth : 1.5 -lines.linewidth : 2.25 -lines.markersize : 9.0 -patch.linewidth : 1.5 -xtick.major.width : 1.875 -ytick.major.width : 1.875 -xtick.minor.width : 1.5 -ytick.minor.width : 1.5 -xtick.major.size : 9.0 -ytick.major.size : 9.0 -xtick.minor.size : 6.0 -ytick.minor.size : 6.0 -font.size : 18.0 -axes.labelsize : 18.0 -xtick.labelsize : 16.5 -ytick.labelsize : 16.5 -legend.fontsize : 16.5 -legend.title_fontsize : 18.0 -axes.titlesize : 32 \ No newline at end of file diff --git a/src/layup/visualize.py b/src/layup/visualize.py index 625e2d98..2d2f999d 100644 --- a/src/layup/visualize.py +++ b/src/layup/visualize.py @@ -4,6 +4,7 @@ from typing import Literal import numpy as np +import matplotlib as mpl from layup.utilities.file_io.CSVReader import CSVDataReader from layup.utilities.file_io.HDF5Reader import HDF5DataReader @@ -115,7 +116,28 @@ def construct_ellipse( return r - +@mpl.rc_context({ + 'axes.linewidth' : 1.875, + 'grid.linewidth' : 1.5, + 'lines.linewidth' : 2.25, + 'lines.markersize' : 9.0, + 'patch.linewidth' : 1.5, + 'xtick.major.width' : 1.875, + 'ytick.major.width' : 1.875, + 'xtick.minor.width' : 1.5, + 'ytick.minor.width' : 1.5, + 'xtick.major.size' : 9.0, + 'ytick.major.size' : 9.0, + 'xtick.minor.size' : 6.0, + 'ytick.minor.size' : 6.0, + 'font.size' : 18.0, + 'axes.labelsize' : 18.0, + 'xtick.labelsize' : 16.5, + 'ytick.labelsize' : 16.5, + 'legend.fontsize' : 16.5, + 'legend.title_fontsize' : 18.0, + 'axes.titlesize' : 32 +}) def matplot_2D(orb_array, planets, no_planets, no_sun, output, fade): """ Create a 2D orbit distribution plot using matplot as a backend. @@ -143,8 +165,6 @@ def matplot_2D(orb_array, planets, no_planets, no_sun, output, fade): import matplotlib.pyplot as plt from matplotlib.collections import LineCollection - plt.style.use("./src/layup/matplot2d.mplstyle") - fig, axs = plt.subplots(1, 2, layout="constrained", figsize=(20, 9)) fig.patch.set_facecolor("k") From 836f0b7d968803c6728c32c6b5bf9c9907353978 Mon Sep 17 00:00:00 2001 From: Joseph Murtagh Date: Tue, 20 May 2025 16:52:32 -0400 Subject: [PATCH 7/9] linting --- src/layup/visualize.py | 47 ++++++++++++++++++++++-------------------- 1 file changed, 25 insertions(+), 22 deletions(-) diff --git a/src/layup/visualize.py b/src/layup/visualize.py index 2d2f999d..d2b42de2 100644 --- a/src/layup/visualize.py +++ b/src/layup/visualize.py @@ -116,28 +116,31 @@ def construct_ellipse( return r -@mpl.rc_context({ - 'axes.linewidth' : 1.875, - 'grid.linewidth' : 1.5, - 'lines.linewidth' : 2.25, - 'lines.markersize' : 9.0, - 'patch.linewidth' : 1.5, - 'xtick.major.width' : 1.875, - 'ytick.major.width' : 1.875, - 'xtick.minor.width' : 1.5, - 'ytick.minor.width' : 1.5, - 'xtick.major.size' : 9.0, - 'ytick.major.size' : 9.0, - 'xtick.minor.size' : 6.0, - 'ytick.minor.size' : 6.0, - 'font.size' : 18.0, - 'axes.labelsize' : 18.0, - 'xtick.labelsize' : 16.5, - 'ytick.labelsize' : 16.5, - 'legend.fontsize' : 16.5, - 'legend.title_fontsize' : 18.0, - 'axes.titlesize' : 32 -}) + +@mpl.rc_context( + { + "axes.linewidth": 1.875, + "grid.linewidth": 1.5, + "lines.linewidth": 2.25, + "lines.markersize": 9.0, + "patch.linewidth": 1.5, + "xtick.major.width": 1.875, + "ytick.major.width": 1.875, + "xtick.minor.width": 1.5, + "ytick.minor.width": 1.5, + "xtick.major.size": 9.0, + "ytick.major.size": 9.0, + "xtick.minor.size": 6.0, + "ytick.minor.size": 6.0, + "font.size": 18.0, + "axes.labelsize": 18.0, + "xtick.labelsize": 16.5, + "ytick.labelsize": 16.5, + "legend.fontsize": 16.5, + "legend.title_fontsize": 18.0, + "axes.titlesize": 32, + } +) def matplot_2D(orb_array, planets, no_planets, no_sun, output, fade): """ Create a 2D orbit distribution plot using matplot as a backend. From ee1b810cae4244e53f0e3b2fa01e09f7c69b42f5 Mon Sep 17 00:00:00 2001 From: Joseph Murtagh Date: Thu, 23 Oct 2025 15:48:38 +0100 Subject: [PATCH 8/9] added pedro suggested changes, added some user safeguarding for planet and sun plotting --- src/layup/visualize.py | 443 ++++++++++++++++++--------------- src/layup_cmdline/visualize.py | 12 +- 2 files changed, 250 insertions(+), 205 deletions(-) diff --git a/src/layup/visualize.py b/src/layup/visualize.py index d2b42de2..2cc9a8ee 100644 --- a/src/layup/visualize.py +++ b/src/layup/visualize.py @@ -1,5 +1,4 @@ import logging -import os from pathlib import Path from typing import Literal @@ -21,6 +20,83 @@ "KEP": ["ObjID", "FORMAT", "a", "e", "inc", "node", "argPeri", "ma", "epochMJD_TDB"], } +DEFAULT_PLOT_RC = { + "axes.linewidth": 1.875, + "grid.linewidth": 1.5, + "lines.linewidth": 2.25, + "lines.markersize": 9.0, + "patch.linewidth": 1.5, + "xtick.major.width": 1.875, + "ytick.major.width": 1.875, + "xtick.minor.width": 1.5, + "ytick.minor.width": 1.5, + "xtick.major.size": 9.0, + "ytick.major.size": 9.0, + "xtick.minor.size": 6.0, + "ytick.minor.size": 6.0, + "font.size": 18.0, + "axes.labelsize": 18.0, + "xtick.labelsize": 16.5, + "ytick.labelsize": 16.5, + "legend.fontsize": 16.5, + "legend.title_fontsize": 18.0, + "axes.titlesize": 32, +} + + +def get_default_fig( + kind: Literal["2D", "3D"] = "2D" +): + import matplotlib.pyplot as plt + + if kind == "2D": + fig, axs = plt.subplots(1,2, layout="constrained", figsize=(20,9)) + fig.patch.set_facecolor("k") + + for ax in axs: + ax.tick_params(labelbottom=True, labeltop=True, labelleft=True, labelright=True, colors="white") + ax.set_facecolor("k") + for spine in ax.spines.values(): + spine.set_color("white") + ax.xaxis.label.set_color("white") + ax.yaxis.label.set_color("white") + + axs[0].set_xlabel("x [AU]") + axs[0].set_ylabel("y [AU]") + axs[1].set_xlabel("x [AU]") + axs[1].set_ylabel("z [AU]") + + axs[0].set_title("Bird's Eye", fontdict={"color": "white"}) + axs[1].set_title("Edge-on", fontdict={"color": "white"}) + + return fig, axs + + elif kind == "3D": + fig = plt.figure(figsize=(15,9)) + fig.subplots_adjust(left=0, right=1, bottom=0, top=1) + ax = fig.add_subplot(111, projection="3d") + + fig.patch.set_facecolor("black") + ax.set_facecolor("black") + + ax.tick_params(axis="x", colors="white", length=0) + ax.tick_params(axis="y", colors="white", length=0) + ax.tick_params(axis="z", colors="white", length=0) + + ax.set_xlabel("x [au]") + ax.set_ylabel("y [au]") + ax.set_xlabel("z [au]") + + for axis in [ax.xaxis, ax.yaxis, ax.zaxis]: + axis.set_pane_color((0.0, 0.0, 0.0, 1.0)) + axis._axinfo["grid"]["color"] = (0.1, 0.1, 0.1, 1.0) + axis.label.set_color("white") + axis._axinfo["tick"]["inward_factor"] = 0 + axis._axinfo["tick"]["outward_factor" \ + ""] = 0 + axis._axinfo["tick"]["size"] = 0 + return fig, ax + def construct_ellipse( a: float, @@ -65,6 +141,7 @@ def construct_ellipse( # mean anomaly via fixed point iteration and then wrap the linspace # around this value from 0 to 2pi E_init = M # initial guesses using mean anomaly (shouldn't be too far off) + assert np.all(e < 1.), "e must be < 1 (bound elliptical orbits)" for tries in range(100): E_new = M + e * np.sin(E_init) @@ -85,9 +162,15 @@ def construct_ellipse( # define position vectors in orthogonal reference frame with origin # at ellipse main focus with q1 oriented towards periapsis (see chapter - # 1.2 of Morbidelli 2011, "Modern Celestial Mechanics" for details) - q1 = a * (np.cos(E) - e) - q2 = a * np.sqrt(1.0 - e**2) * np.sin(E) + # 1.2 of Morbidelli 2011, "Modern Celestial Mechanics" for details, or + # §15 p38 eq. 15.12 of Landau and Lifshitz, "Mechanics" for the case + # of a hyperbolic orbit) + if e < 1.: + q1 = a * (np.cos(E) - e) + q2 = a * np.sqrt(1.0 - e**2) * np.sin(E) + elif e >= 1.: + q1 = a * (e - np.cosh(E)) + q2 = a * np.sqrt(e**2 - 1.0) * np.sinh(E) if type(q1) == float: q = np.array([q1, q2, 0.0]) @@ -117,31 +200,7 @@ def construct_ellipse( return r -@mpl.rc_context( - { - "axes.linewidth": 1.875, - "grid.linewidth": 1.5, - "lines.linewidth": 2.25, - "lines.markersize": 9.0, - "patch.linewidth": 1.5, - "xtick.major.width": 1.875, - "ytick.major.width": 1.875, - "xtick.minor.width": 1.5, - "ytick.minor.width": 1.5, - "xtick.major.size": 9.0, - "ytick.major.size": 9.0, - "xtick.minor.size": 6.0, - "ytick.minor.size": 6.0, - "font.size": 18.0, - "axes.labelsize": 18.0, - "xtick.labelsize": 16.5, - "ytick.labelsize": 16.5, - "legend.fontsize": 16.5, - "legend.title_fontsize": 18.0, - "axes.titlesize": 32, - } -) -def matplot_2D(orb_array, planets, no_planets, no_sun, output, fade): +def matplot_2D(orb_array, planets, plot_planets, plot_sun, output, fade, fig=None, axs=None): """ Create a 2D orbit distribution plot using matplot as a backend. @@ -153,35 +212,33 @@ def matplot_2D(orb_array, planets, no_planets, no_sun, output, fade): planets : list List of desired planet orbits to be plotted - no_planets : bool - Flag to turn off planet orbits + plot_planets : bool + Flag to turn on planet orbits - no_sun : bool - Flag to turn of sun plotting + plot_sun : bool + Flag to turn on sun plotting output : str Output file stem fade : bool Flag to turn on object orbit fading + + fig : matplot figure object (optional, default=None) + User created matplot figure object for custom plotting + + axs : matplot axis object (optional, default=None) + User created corresponding matplot axis object for custom plotting """ import matplotlib.pyplot as plt from matplotlib.collections import LineCollection - fig, axs = plt.subplots(1, 2, layout="constrained", figsize=(20, 9)) - fig.patch.set_facecolor("k") + created_fig = False + if fig is None or axs is None: + fig, axs = get_default_fig("2D") + created_fig = True for ax in axs: - ax.tick_params(labelbottom=True, labeltop=True, labelleft=True, labelright=True, colors="white") - - ax.set_facecolor("k") - - for spine in ax.spines.values(): - spine.set_color("white") - - ax.xaxis.label.set_color("white") - ax.yaxis.label.set_color("white") - if np.max(orb_array["a"] < 2): maxQ = np.max(orb_array["a"] * (1 + orb_array["e"])) + 0.2 ax.set(xlim=(-maxQ, maxQ), ylim=(-maxQ, maxQ), aspect="equal") @@ -220,11 +277,11 @@ def matplot_2D(orb_array, planets, no_planets, no_sun, output, fade): axs[1].add_collection(lc) # add the sun - if not no_sun: + if plot_sun: axs[0].scatter(0, 0, s=100, color="xkcd:sunflower yellow") # add planets - if not no_planets: + if plot_planets and planets: planets_dic = { "Me": [0.387, 7.0, "#E7E8EC"], "V": [0.723, 3.4, "#E39E1C"], @@ -241,30 +298,26 @@ def matplot_2D(orb_array, planets, no_planets, no_sun, output, fade): planets_dic = dictfilt(planets_dic, planets) for _, (k, v) in enumerate(planets_dic.items()): - axs[0].add_patch(plt.Circle((0, 0), v[0], color=v[2], fill=False, alpha=0.9, linestyle="dotted")) - # axs[0].text(v[0]+1, 0, k[0], color=v[2]) + axs[0].add_patch(plt.Circle((0, 0), v[0], color=v[2], fill=False, alpha=0.9, linestyle="dotted", zorder=100)) + axs[0].text(v[0]+1, 0, k[0], color=v[2]) axs[1].plot( v[0] * np.cos(theta), v[0] * np.cos(theta) * np.sin(np.radians(v[1])), color=v[2], linestyle="dotted", + zorder=100 ) - axs[0].set_xlabel("x [AU]") - axs[0].set_ylabel("y [AU]") - axs[1].set_xlabel("x [AU]") - axs[1].set_ylabel("z [AU]") - - axs[0].set_title("Bird's Eye", fontdict={"color": "white"}) - axs[1].set_title("Edge-on", fontdict={"color": "white"}) - - plt.savefig(output, dpi=300) - plt.show() - plt.close() + if created_fig: + fig.savefig(output, dpi=300) + plt.show() + plt.close(fig) + else: + return fig, axs -def matplot_3D(orb_array, planets, no_planets, no_sun, output, fade): +def matplot_3D(orb_array, planets, plot_planets, plot_sun, output, fade, fig=None, axs=None): """ Create a 3D orbit distribution plot using matplot as a backend. @@ -276,51 +329,31 @@ def matplot_3D(orb_array, planets, no_planets, no_sun, output, fade): planets : list List of desired planet orbits to be plotted - no_planets : bool - Flag to turn off planet orbits + plot_planets : bool + Flag to turn on planet orbits - no_sun : bool - Flag to turn of sun plotting + plot_sun : bool + Flag to turn on sun plotting output : str Output file stem fade : bool Flag to turn on object orbit fading + + fig : matplot figure object (optional, default=None) + User created matplot figure object for custom plotting + + axs : matplot axis object (optional, default=None) + User created corresponding matplot axis object for custom plotting """ import matplotlib.pyplot as plt from mpl_toolkits.mplot3d.art3d import Line3DCollection - fig = plt.figure(figsize=(15, 9)) - fig.subplots_adjust(left=0, right=1, bottom=0, top=1) - axs = fig.add_subplot(111, projection="3d") - - fig.patch.set_facecolor("black") - axs.set_facecolor("black") - axs.xaxis.set_pane_color((0.0, 0.0, 0.0, 1.0)) # RGBA: white - axs.yaxis.set_pane_color((0.0, 0.0, 0.0, 1.0)) - axs.zaxis.set_pane_color((0.0, 0.0, 0.0, 1.0)) - - axs.xaxis._axinfo["grid"]["color"] = (0.1, 0.1, 0.1, 1) # grey - axs.yaxis._axinfo["grid"]["color"] = (0.1, 0.1, 0.1, 1) # grey - axs.zaxis._axinfo["grid"]["color"] = (0.1, 0.1, 0.1, 1) # grey - - axs.xaxis.label.set_color("white") - axs.yaxis.label.set_color("white") - axs.zaxis.label.set_color("white") - - axs.tick_params(axis="x", colors="white", length=0) - axs.tick_params(axis="y", colors="white", length=0) - axs.tick_params(axis="z", colors="white", length=0) - - for axis in [axs.xaxis, axs.yaxis, axs.zaxis]: - axis._axinfo["tick"]["inward_factor"] = 0 - axis._axinfo["tick"]["outward_factor"] = 0 - axis._axinfo["tick"]["size"] = 0 - - axs.set_xlabel("x [au]") - axs.set_ylabel("y [au]") - axs.set_zlabel("z [au]") + created_fig = False + if fig is None or axs is None: + fig, axs = get_default_fig("3D") + created_fig = True for obj in orb_array: posvec = construct_ellipse(obj[2], obj[3], obj[4], obj[6], obj[5], obj[7]) @@ -338,11 +371,11 @@ def matplot_3D(orb_array, planets, no_planets, no_sun, output, fade): axs.add_collection3d(lc) # add sun - if not no_sun: + if plot_sun: axs.scatter(0, 0, s=100, color="xkcd:sunflower yellow") # add planets - if not no_planets: + if plot_planets and planets: planets_dic = { "Me": [0.387, 7.0, "#E7E8EC"], "V": [0.723, 3.4, "#E39E1C"], @@ -366,102 +399,107 @@ def matplot_3D(orb_array, planets, no_planets, no_sun, output, fade): color=v[2], alpha=0.9, linestyle="dotted", + zorder=10000 ) axs.text(v[0], 0, 0, k[0], color=v[2]) - plt.savefig(output, dpi=300, bbox_inches="tight") - plt.show() - - -def plotly_3D(orb_array, planets, no_planets, no_sun, output, fade): - """ - Create a 3D orbit distribution plot using plotly as a backend. - - Parameters - ------- - orb_array : structured numpy array - Numpy array containing orbital element information - - planets : list - List of desired planet orbits to be plotted - - no_planets : bool - Flag to turn off planet orbits - - no_sun : bool - Flag to turn of sun plotting - - output : str - Output file stem - - fade : bool - Flag to turn on object orbit fading - """ - import plotly.graph_objects as go - - fig = go.Figure() - - fig.update_layout( - plot_bgcolor="rgb(0,0,0)", - paper_bgcolor="rgb(0,0,0)", - scene=dict( - xaxis=dict( - backgroundcolor="rgba(0, 0, 0,0)", - gridcolor="rgb(106, 110, 117)", - title_font=dict(color="rgb(255,255,255)"), - tickfont=dict(color="rgb(255,255,255)"), - showbackground=True, - zerolinecolor="white", - ), - yaxis=dict( - backgroundcolor="rgba(0, 0, 0,0)", - gridcolor="rgb(106, 110, 117)", - title_font=dict(color="rgb(255,255,255)"), - tickfont=dict(color="rgb(255,255,255)"), - showbackground=True, - zerolinecolor="white", - ), - zaxis=dict( - backgroundcolor="rgba(0, 0, 0,0)", - gridcolor="rgb(106, 110, 117)", - title_font=dict(color="rgb(255,255,255)"), - tickfont=dict(color="rgb(255,255,255)"), - showbackground=True, - zerolinecolor="white", - ), - ), - ) - - fig.update_yaxes(title_font_color="white", color="white") - fig.update_xaxes(title_font_color="white", color="white") - - if not no_sun: - fig.add_trace(go.scatter(x=0, y=0, s=100, color="xkcd:sunflower yellow")) - - for obj in orb_array: - posvec = construct_ellipse(obj[2], obj[3], obj[4], obj[6], obj[5], obj[7]) - - normx = (posvec[:, 0] - posvec[:, 0].min()) / (posvec[:, 0].max() - posvec[:, 0].min()) - opacity = np.log1p(normx * 9) / np.log1p(9) - - fig.add_trace( - go.Scatter( - x=posvec[:, 0], - y=posvec[:, 1], - mode="lines", - line=dict( - color=opacity, - colorscale="gray", - cmin=0, - cmax=1, - ), - customdata=[obj[0]], - hovertemplate="%{customdata[0]}
(x: %{x:.2f}, y: %{y:.2f}, z: %{z:.2f})", - ) - ) - - fig.update_layout(showlegend=False) - fig.show() + if created_fig: + fig.savefig(output, dpi=300, bbox_inches="tight") + plt.show() + plt.close(fig) + else: + return fig, axs + +# TODO: develop plotly plotting functionality better +# def plotly_3D(orb_array, planets, no_planets, sun, output, fade): +# """ +# Create a 3D orbit distribution plot using plotly as a backend. + +# Parameters +# ------- +# orb_array : structured numpy array +# Numpy array containing orbital element information + +# planets : list +# List of desired planet orbits to be plotted + +# no_planets : bool +# Flag to turn off planet orbits + +# sun : bool +# Flag to turn on sun plotting + +# output : str +# Output file stem + +# fade : bool +# Flag to turn on object orbit fading +# """ +# import plotly.graph_objects as go + +# fig = go.Figure() + +# fig.update_layout( +# plot_bgcolor="rgb(0,0,0)", +# paper_bgcolor="rgb(0,0,0)", +# scene=dict( +# xaxis=dict( +# backgroundcolor="rgba(0, 0, 0,0)", +# gridcolor="rgb(106, 110, 117)", +# title_font=dict(color="rgb(255,255,255)"), +# tickfont=dict(color="rgb(255,255,255)"), +# showbackground=True, +# zerolinecolor="white", +# ), +# yaxis=dict( +# backgroundcolor="rgba(0, 0, 0,0)", +# gridcolor="rgb(106, 110, 117)", +# title_font=dict(color="rgb(255,255,255)"), +# tickfont=dict(color="rgb(255,255,255)"), +# showbackground=True, +# zerolinecolor="white", +# ), +# zaxis=dict( +# backgroundcolor="rgba(0, 0, 0,0)", +# gridcolor="rgb(106, 110, 117)", +# title_font=dict(color="rgb(255,255,255)"), +# tickfont=dict(color="rgb(255,255,255)"), +# showbackground=True, +# zerolinecolor="white", +# ), +# ), +# ) + +# fig.update_yaxes(title_font_color="white", color="white") +# fig.update_xaxes(title_font_color="white", color="white") + +# if sun: +# fig.add_trace(go.scatter(x=0, y=0, s=100, color="xkcd:sunflower yellow")) + +# for obj in orb_array: +# posvec = construct_ellipse(obj[2], obj[3], obj[4], obj[6], obj[5], obj[7]) + +# normx = (posvec[:, 0] - posvec[:, 0].min()) / (posvec[:, 0].max() - posvec[:, 0].min()) +# opacity = np.log1p(normx * 9) / np.log1p(9) + +# fig.add_trace( +# go.Scatter( +# x=posvec[:, 0], +# y=posvec[:, 1], +# mode="lines", +# line=dict( +# color=opacity, +# colorscale="gray", +# cmin=0, +# cmax=1, +# ), +# customdata=[obj[0]], +# hovertemplate="%{customdata[0]}
(x: %{x:.2f}, y: %{y:.2f}, z: %{z:.2f})", +# ) +# ) + +# fig.update_layout(showlegend=False) +# fig.show() def visualize_cli( @@ -473,8 +511,8 @@ def visualize_cli( backend: Literal["matplot", "plotly"] = "matplot", dimensions: Literal["2D", "3D"] = "2D", num_orbs: int = 1_000, - no_planets: bool = False, - no_sun: bool = False, + plot_planets: bool = True, + plot_sun: bool = True, fade: bool = False, ): """ @@ -506,11 +544,11 @@ def visualize_cli( num_orbs : int (default=1_000) The number of random orbits to plot - no_planets : bool (default=False) - Flag to turn off the planet orbits + plot_planets : bool (default=True) + Flag to turn on the planet orbits - no_sun : bool (default=False) - Flag to turn off the Sun + plot_sun : bool (default=True) + Flag to turn on the Sun fade : bool (default=False) Flag to turn off object orbit fading @@ -540,13 +578,16 @@ def visualize_cli( ) random_orbs = reader.read_rows(block_start=0, block_size=1000) + if num_orbs > random_orbs.size: + logger.warning(f"Requested {num_orbs} orbits, but only {random_orbs.size} available from input. Capping at {random_orbs.size} instead.") + num_orbs = random_orbs.size random_orbs = random_orbs[np.random.choice(random_orbs.size, size=num_orbs, replace=False)] if backend == "matplot": if dimensions == "2D": - matplot_2D(random_orbs, planets, no_planets, no_sun, output_file_stem, fade) + matplot_2D(random_orbs, planets, plot_planets, plot_sun, output_file_stem, fade) else: - matplot_3D(random_orbs, planets, no_planets, no_sun, output_file_stem, fade) - elif backend == "plotly": - if dimensions == "3D": - plotly_3D(random_orbs) + matplot_3D(random_orbs, planets, plot_planets, plot_sun, output_file_stem, fade) + # elif backend == "plotly": + # if dimensions == "3D": + # plotly_3D(random_orbs) diff --git a/src/layup_cmdline/visualize.py b/src/layup_cmdline/visualize.py index 6a35bcfe..75d3e967 100644 --- a/src/layup_cmdline/visualize.py +++ b/src/layup_cmdline/visualize.py @@ -87,10 +87,10 @@ def main(): required=False, ) optional.add_argument( - "--no_planets", help="overplot the planets. default is True", dest="no_planets", action="store_true" + "--plot_planets", help="overplot the planets. default is True", dest="plot_planets", action="store_true" ) optional.add_argument( - "--no_sun", help="overplot the sun. default is True", dest="no_sun", action="store_true" + "--plot_sun", help="overplot the sun. default is True", dest="plot_sun", action="store_true" ) optional.add_argument( "-f", @@ -134,6 +134,10 @@ def execute(args): if args.b not in ["matplot", "plotly"]: sys.exit("ERROR: -b --backend must be 'matplot' or 'plotly'") + if args.plot_planets and not args.planets: + logger.warning("WARNING: --plot-planets given without --planets . defaulting to all planets") + args.planets = ["Me", "V", "E", "Ma", "J", "S", "U", "N"] + visualize_cli( input=args.input, output_file_stem=args.o, @@ -142,8 +146,8 @@ def execute(args): backend=args.b, dimensions=args.d, num_orbs=args.n, - no_planets=args.no_planets, - no_sun=args.no_sun, + plot_planets=args.plot_planets, + plot_sun=args.plot_sun, fade=args.fade, ) From ef98cb19cc1a523b92e382a4930cf64e5aef9da5 Mon Sep 17 00:00:00 2001 From: Joseph Murtagh Date: Thu, 23 Oct 2025 16:00:13 +0100 Subject: [PATCH 9/9] fixed linting --- src/layup/visualize.py | 36 ++++++++++++++++++---------------- src/layup_cmdline/visualize.py | 5 ++++- 2 files changed, 23 insertions(+), 18 deletions(-) diff --git a/src/layup/visualize.py b/src/layup/visualize.py index 2cc9a8ee..bd179c66 100644 --- a/src/layup/visualize.py +++ b/src/layup/visualize.py @@ -44,13 +44,11 @@ } -def get_default_fig( - kind: Literal["2D", "3D"] = "2D" -): +def get_default_fig(kind: Literal["2D", "3D"] = "2D"): import matplotlib.pyplot as plt if kind == "2D": - fig, axs = plt.subplots(1,2, layout="constrained", figsize=(20,9)) + fig, axs = plt.subplots(1, 2, layout="constrained", figsize=(20, 9)) fig.patch.set_facecolor("k") for ax in axs: @@ -70,9 +68,9 @@ def get_default_fig( axs[1].set_title("Edge-on", fontdict={"color": "white"}) return fig, axs - + elif kind == "3D": - fig = plt.figure(figsize=(15,9)) + fig = plt.figure(figsize=(15, 9)) fig.subplots_adjust(left=0, right=1, bottom=0, top=1) ax = fig.add_subplot(111, projection="3d") @@ -92,8 +90,7 @@ def get_default_fig( axis._axinfo["grid"]["color"] = (0.1, 0.1, 0.1, 1.0) axis.label.set_color("white") axis._axinfo["tick"]["inward_factor"] = 0 - axis._axinfo["tick"]["outward_factor" \ - ""] = 0 + axis._axinfo["tick"]["outward_factor" ""] = 0 axis._axinfo["tick"]["size"] = 0 return fig, ax @@ -141,7 +138,7 @@ def construct_ellipse( # mean anomaly via fixed point iteration and then wrap the linspace # around this value from 0 to 2pi E_init = M # initial guesses using mean anomaly (shouldn't be too far off) - assert np.all(e < 1.), "e must be < 1 (bound elliptical orbits)" + assert np.all(e < 1.0), "e must be < 1 (bound elliptical orbits)" for tries in range(100): E_new = M + e * np.sin(E_init) @@ -162,13 +159,13 @@ def construct_ellipse( # define position vectors in orthogonal reference frame with origin # at ellipse main focus with q1 oriented towards periapsis (see chapter - # 1.2 of Morbidelli 2011, "Modern Celestial Mechanics" for details, or + # 1.2 of Morbidelli 2011, "Modern Celestial Mechanics" for details, or # §15 p38 eq. 15.12 of Landau and Lifshitz, "Mechanics" for the case # of a hyperbolic orbit) - if e < 1.: + if e < 1.0: q1 = a * (np.cos(E) - e) q2 = a * np.sqrt(1.0 - e**2) * np.sin(E) - elif e >= 1.: + elif e >= 1.0: q1 = a * (e - np.cosh(E)) q2 = a * np.sqrt(e**2 - 1.0) * np.sinh(E) @@ -298,15 +295,17 @@ def matplot_2D(orb_array, planets, plot_planets, plot_sun, output, fade, fig=Non planets_dic = dictfilt(planets_dic, planets) for _, (k, v) in enumerate(planets_dic.items()): - axs[0].add_patch(plt.Circle((0, 0), v[0], color=v[2], fill=False, alpha=0.9, linestyle="dotted", zorder=100)) - axs[0].text(v[0]+1, 0, k[0], color=v[2]) + axs[0].add_patch( + plt.Circle((0, 0), v[0], color=v[2], fill=False, alpha=0.9, linestyle="dotted", zorder=100) + ) + axs[0].text(v[0] + 1, 0, k[0], color=v[2]) axs[1].plot( v[0] * np.cos(theta), v[0] * np.cos(theta) * np.sin(np.radians(v[1])), color=v[2], linestyle="dotted", - zorder=100 + zorder=100, ) if created_fig: @@ -399,7 +398,7 @@ def matplot_3D(orb_array, planets, plot_planets, plot_sun, output, fade, fig=Non color=v[2], alpha=0.9, linestyle="dotted", - zorder=10000 + zorder=10000, ) axs.text(v[0], 0, 0, k[0], color=v[2]) @@ -410,6 +409,7 @@ def matplot_3D(orb_array, planets, plot_planets, plot_sun, output, fade, fig=Non else: return fig, axs + # TODO: develop plotly plotting functionality better # def plotly_3D(orb_array, planets, no_planets, sun, output, fade): # """ @@ -579,7 +579,9 @@ def visualize_cli( random_orbs = reader.read_rows(block_start=0, block_size=1000) if num_orbs > random_orbs.size: - logger.warning(f"Requested {num_orbs} orbits, but only {random_orbs.size} available from input. Capping at {random_orbs.size} instead.") + logger.warning( + f"Requested {num_orbs} orbits, but only {random_orbs.size} available from input. Capping at {random_orbs.size} instead." + ) num_orbs = random_orbs.size random_orbs = random_orbs[np.random.choice(random_orbs.size, size=num_orbs, replace=False)] diff --git a/src/layup_cmdline/visualize.py b/src/layup_cmdline/visualize.py index 75d3e967..fa20010d 100644 --- a/src/layup_cmdline/visualize.py +++ b/src/layup_cmdline/visualize.py @@ -87,7 +87,10 @@ def main(): required=False, ) optional.add_argument( - "--plot_planets", help="overplot the planets. default is True", dest="plot_planets", action="store_true" + "--plot_planets", + help="overplot the planets. default is True", + dest="plot_planets", + action="store_true", ) optional.add_argument( "--plot_sun", help="overplot the sun. default is True", dest="plot_sun", action="store_true"