Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Add option to plot error #148

Merged
merged 1 commit into from
Aug 21, 2023
Merged

Add option to plot error #148

Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
79 changes: 56 additions & 23 deletions pykoop/koopman_pipeline.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -2691,6 +2691,7 @@ def plot_predicted_trajectory(
relift_state: bool = True,
plot_lifted: bool = False,
plot_input: bool = False,
plot_error: bool = False,
episode_feature: Optional[bool] = None,
plot_ground_truth: bool = True,
episode_style: Optional[str] = None,
Expand Down Expand Up @@ -2719,6 +2720,9 @@ def plot_predicted_trajectory(
plot_input : bool
If true, plot the input as well as the state. If false, plot
only the original state (default).
plot_error : bool
If true, plot the prediction error instead of the state. If false,
plot the predicted state and ground truth (default).
episode_feature : Optional[bool]
True if first feature indicates which episode a timestep is from.
If ``None``, ``self.episode_feature_`` is used.
Expand Down Expand Up @@ -2797,33 +2801,62 @@ def plot_predicted_trajectory(
# Plot results
for row in range(n_row):
for ep in range(n_eps):
if episode_style == 'overlay':
line_pred = ax[row, 0].plot(
eps[ep][1][:, row],
label=f'Ep. {int(eps[ep][0])} prediction',
**plot_args,
)
if eps_gt is not None and row < n_states:
ax[row, 0].plot(
eps_gt[ep][1][:, row],
label=f'Ep. {int(eps[ep][0])} ground truth',
linestyle='--',
color=line_pred[0].get_color(),
if plot_error:
if episode_style == 'overlay':
if eps_gt is not None and row < n_states:
ax[row, 0].plot(
eps_gt[ep][1][:, row] - eps[ep][1][:, row],
label=f'Ep. {int(eps[ep][0])} error',
**plot_args,
)
else:
ax[row, 0].plot(
eps[ep][1][:, row],
label=f'Ep. {int(eps[ep][0])} input',
**plot_args,
)
else:
if eps_gt is not None and row < n_states:
ax[row, ep].plot(
eps_gt[ep][1][:, row] - eps[ep][1][:, row],
label=f'Prediction error',
linestyle='--',
**plot_args,
)
else:
ax[row, ep].plot(
eps[ep][1][:, row],
label=f'Input',
**plot_args,
)
else:
if episode_style == 'overlay':
line_pred = ax[row, 0].plot(
eps[ep][1][:, row],
label=f'Ep. {int(eps[ep][0])} prediction',
**plot_args,
)
else:
line_pred = ax[row, ep].plot(
eps[ep][1][:, row],
label=f'Prediction',
**plot_args,
)
if eps_gt is not None and row < n_states:
ax[row, ep].plot(
eps_gt[ep][1][:, row],
label=f'Ground truth',
linestyle='--',
if eps_gt is not None and row < n_states:
ax[row, 0].plot(
eps_gt[ep][1][:, row],
label=f'Ep. {int(eps[ep][0])} ground truth',
linestyle='--',
color=line_pred[0].get_color(),
**plot_args,
)
else:
line_pred = ax[row, ep].plot(
eps[ep][1][:, row],
label=f'Prediction',
**plot_args,
)
if eps_gt is not None and row < n_states:
ax[row, ep].plot(
eps_gt[ep][1][:, row],
label=f'Ground truth',
linestyle='--',
**plot_args,
)
# Set y labels
if plot_lifted:
names = self.get_feature_names_out(
Expand Down