From 789dfe6960a519c15c07cf53c35ada037e9d8576 Mon Sep 17 00:00:00 2001 From: noakaplan675 Date: Tue, 1 Jul 2025 10:32:15 -0400 Subject: [PATCH 1/3] added 3D attractor plotting --- src/orc/utils/visualization.py | 105 +++++++++++++++++++++++++++++++++ 1 file changed, 105 insertions(+) diff --git a/src/orc/utils/visualization.py b/src/orc/utils/visualization.py index 78cb567..a83907c 100644 --- a/src/orc/utils/visualization.py +++ b/src/orc/utils/visualization.py @@ -179,3 +179,108 @@ def imshow_1D_spatiotemp(U, plt.show() # TODO: add plot_attrator function to visualize 2D/3D attractors in state space + + + +def plot_in_3D_state_space(U_lst, + time_series_labels=None, + line_formats=["-", "r--"], + state_var_names=None, + figsize = (20,8), + title = None, + **plot_kwargs): + """ **STEP 1 - clear, one-line summary:** Plot time series data to visualize 3D attractors in state space. + + Parameters + ---------- + **STEP 2 - + : + ** + U_lst : 2D array or list of 2D arrays + If a 2D array, shape should be (Nt, 3) with 3 state + variables, where Nt is the number of time points. If a list of 2D arrays, + each array should have shape (Nt, 3) and represent different time series. + time_series_labels : list of strings, optional + List of strings containing the labels for each time series to be shown in + a legend. If None, no labels will be shown. + line_formats : list of strings, optional + List of strings containing the line formats for each time series. If None, + default line format will be used. + state_var_names : list of strings, optional + List of strings containing the names of the state variables. If None, + no axis labels will be shown. + figsize : tuple, optional + Size of the figure to be created. Default is (20, 8). + title : string, optional + Title of the plot. If None, no title is shown. + plot_kwargs : dict, optional + Additional arguments to pass to the plot function. + """ + + # **STEP 3** + # Input validation + if not isinstance(U_lst, list): + if not isinstance(U_lst, jnp.ndarray | np.ndarray) or U_lst.ndim != 2: + raise TypeError("U_lst must be a 2D JAX or NumPy array or a list of \ + 2D JAX/NumPy arrays.") + U_lst = [U_lst] + else: + if not all( + isinstance(U, jnp.ndarray | np.ndarray) and U.ndim == 2 for U in U_lst): + raise TypeError("All elements in U_lst must be 2D JAX or NumPy arrays.") + if not all(U.shape == U_lst[0].shape for U in U_lst): + raise ValueError("All arrays in U_lst must have the same shape.") + + Nu = U_lst[0].shape[1] + + if time_series_labels is not None: + if not isinstance(time_series_labels, list): + raise TypeError("time_series_labels must be a list of strings.") + if len(time_series_labels) != len(U_lst): + raise ValueError(f"Length of time_series_labels ({len(time_series_labels)})\ + must match the number of time series ({len(U_lst)}).") + + if line_formats is not None: + if not isinstance(line_formats, list): + raise TypeError("line_formats must be a list of strings.") + if len(line_formats) != len(U_lst): + raise ValueError(f"Length of line_formats ({len(line_formats)}) must \ + match the number of time series ({len(U_lst)}).") + + if state_var_names is not None: + if not isinstance(state_var_names, list): + raise TypeError("state_var_names must be a list of strings.") + if len(state_var_names) != Nu: + raise ValueError(f"Length of state_var_names ({len(state_var_names)}) \ + must match the number of state variables ({Nu}).") + + # **STEP 4** + # defaults + plot_kwargs.setdefault('linewidth', 2) + + # **STEP 5** + # handle optional inputs + if time_series_labels is None: + time_series_labels = [None for _ in range(len(U_lst))] + if line_formats is None: + line_formats = ['-' for _ in range(len(U_lst))] + + # **STEP 5 - CODE:** + # plot + fig, axs = plt.subplots(subplot_kw={"projection": "3d"}) + # Ensure axs is always iterable, even if Nu=1 + if Nu == 1: + axs = [axs] + for j, Y in enumerate(U_lst): + axs.plot(Y[:, 0], Y[:, 1], Y[:, 2], line_formats[j], + label=time_series_labels[j], + **plot_kwargs) + if state_var_names is not None: + axs.set(xlabel=state_var_names[0]) + axs.set(ylabel=state_var_names[1]) + axs.set(zlabel=state_var_names[2]) + if time_series_labels[0] is not None: + axs.legend(loc='upper right') + if title is not None: + axs.set_title(title, fontsize=14) + plt.show() \ No newline at end of file From a92c32fbb7592dce02c0bef763464df57a174a78 Mon Sep 17 00:00:00 2001 From: noakaplan675 Date: Sat, 5 Jul 2025 04:35:08 -0400 Subject: [PATCH 2/3] added 3D plot of attractors - after passing tests --- src/orc/utils/visualization.py | 31 +++++++++++-------------------- 1 file changed, 11 insertions(+), 20 deletions(-) diff --git a/src/orc/utils/visualization.py b/src/orc/utils/visualization.py index a83907c..221a36c 100644 --- a/src/orc/utils/visualization.py +++ b/src/orc/utils/visualization.py @@ -178,24 +178,17 @@ def imshow_1D_spatiotemp(U, plt.colorbar(pad = 0.01, label = r'$u$') plt.show() -# TODO: add plot_attrator function to visualize 2D/3D attractors in state space - - - def plot_in_3D_state_space(U_lst, time_series_labels=None, - line_formats=["-", "r--"], + line_formats=None, state_var_names=None, figsize = (20,8), title = None, **plot_kwargs): - """ **STEP 1 - clear, one-line summary:** Plot time series data to visualize 3D attractors in state space. + """Plot time series data to visualize 3D attractors in state space. Parameters ---------- - **STEP 2 - - : - ** U_lst : 2D array or list of 2D arrays If a 2D array, shape should be (Nt, 3) with 3 state variables, where Nt is the number of time points. If a list of 2D arrays, @@ -216,8 +209,6 @@ def plot_in_3D_state_space(U_lst, plot_kwargs : dict, optional Additional arguments to pass to the plot function. """ - - # **STEP 3** # Input validation if not isinstance(U_lst, list): if not isinstance(U_lst, jnp.ndarray | np.ndarray) or U_lst.ndim != 2: @@ -230,7 +221,7 @@ def plot_in_3D_state_space(U_lst, raise TypeError("All elements in U_lst must be 2D JAX or NumPy arrays.") if not all(U.shape == U_lst[0].shape for U in U_lst): raise ValueError("All arrays in U_lst must have the same shape.") - + Nu = U_lst[0].shape[1] if time_series_labels is not None: @@ -254,27 +245,23 @@ def plot_in_3D_state_space(U_lst, raise ValueError(f"Length of state_var_names ({len(state_var_names)}) \ must match the number of state variables ({Nu}).") - # **STEP 4** # defaults plot_kwargs.setdefault('linewidth', 2) - # **STEP 5** # handle optional inputs if time_series_labels is None: time_series_labels = [None for _ in range(len(U_lst))] if line_formats is None: line_formats = ['-' for _ in range(len(U_lst))] - # **STEP 5 - CODE:** # plot - fig, axs = plt.subplots(subplot_kw={"projection": "3d"}) + fig, axs = plt.subplots(subplot_kw={"projection": "3d"}, figsize = figsize) # Ensure axs is always iterable, even if Nu=1 if Nu == 1: axs = [axs] for j, Y in enumerate(U_lst): - axs.plot(Y[:, 0], Y[:, 1], Y[:, 2], line_formats[j], - label=time_series_labels[j], - **plot_kwargs) + axs.plot(Y[:, 0], Y[:, 1], Y[:, 2], line_formats[j], + label=time_series_labels[j], **plot_kwargs) if state_var_names is not None: axs.set(xlabel=state_var_names[0]) axs.set(ylabel=state_var_names[1]) @@ -283,4 +270,8 @@ def plot_in_3D_state_space(U_lst, axs.legend(loc='upper right') if title is not None: axs.set_title(title, fontsize=14) - plt.show() \ No newline at end of file + plt.show() + + + +# TODO: add plot_attrator function to visualize 2D attractors in state space \ No newline at end of file From 4b5fa3161fcde3b74d9f6c07a7c56c5cead085c4 Mon Sep 17 00:00:00 2001 From: noakaplan675 Date: Tue, 8 Jul 2025 15:07:40 -0400 Subject: [PATCH 3/3] added 3D plot of attractors - after passing tests (final?) --- src/orc/utils/visualization.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/orc/utils/visualization.py b/src/orc/utils/visualization.py index 221a36c..8a4742c 100644 --- a/src/orc/utils/visualization.py +++ b/src/orc/utils/visualization.py @@ -274,4 +274,4 @@ def plot_in_3D_state_space(U_lst, -# TODO: add plot_attrator function to visualize 2D attractors in state space \ No newline at end of file +# TODO: add plot_attrator function to visualize 2D attractors in state space