Change Legend Labels for plot_trajectory() function

[tfs563]

Hi,

I'm using the plot_trajectory() function available to pyAT, which can be found at this link to model an energy spectrometer, and I wanted to get rid of the z labels that you get when this function plots with the default settings. I also want to change the x labels to illustrate momentum(say p, p_low, and p_high) instead of the x position.

The current plot I get is the following for reference

Is there a way to do this with this plot function?

[TeresiaOlsson]

Hi,

I also needed to make changes to a plot once. The easiest way I found was to get the figure from the axes and then I could do any modifications by looking up how to make modifications to an existing figure in matplotlib.

plot_trajectory should return the axes (like this https://atcollab.github.io/at/p/api/at.plot.generic.html) and then you can use the following function to get the figure:

https://matplotlib.org/3.2.2/api/_as_gen/matplotlib.axes.Axes.get_figure.html

[tfs563]

Hi,

I tried to use the function recommended to get the figure, but I get an error message of "AttributeError: 'tuple' object has no attribute 'figure'". So, if I am reading this correct, then the plot_trajectory() function I am using doesn't provide the axes needed for me to modify them?

[lcarver]

Hello,
In a generic sense, you can do something like this.

st is a starting ring index for the plot
en is the ending ring index for the plot

ax = at.plot.plot_synopt(ring[st:en])
ax.yaxis.set_label_position("right")
ax.yaxis.tick_right()
ax.yaxis.set_visible(False)
ax2 = ax.twinx()
ax2.plot(sPos[bpm_inds]-sPos[st], 1e6*orball[bpm_inds,2], color='k', marker='x', linestyle='None', label='Fitted Orb at BPMs')
ax2.yaxis.set_label_position("left")
ax2.yaxis.tick_left()
ax2.xaxis.set_visible(True)
ax2.set_zorder(3)
ax2.set_xlabel('s [m]')
ax2.set_ylabel('y-orbit [um]')
ax2.legend()
plt.show()

this allows you to plot whatever you want with the AT synoptic at the bottom.

[TeresiaOlsson]

I did it like this for plot_beta. I think the same should be possible for plot_trajectory since both are calling the same plotting function. In this way I could modify only the parts I wanted of the existing figure.

def plot_optics(ring):

# Plot the optics functions
[left,right,syn] = ring.plot_beta()

# Get the figure to be able to make changes
fig = left.get_figure()

# Change width of the function
fig.set_figwidth(11)

# Set axes ranges
left.set_ylim(0,30)
right.set_ylim(0,1)
align.yaxes(left, 0, right, 0, 0.15)

# Set location of the legend
legend = left.get_legend()
legend.set_loc('upper right')

# Change legend names
legend.get_texts()[2].set_text('$\\eta_x$')

# Change label name
right.set_ylabel('$\\eta_x$')

# Save the plot to file
plt.savefig('optics.png')

[tfs563]

Hello, In a generic sense, you can do something like this.

st is a starting ring index for the plot en is the ending ring index for the plot

ax = at.plot.plot_synopt(ring[st:en])
ax.yaxis.set_label_position("right")
ax.yaxis.tick_right()
ax.yaxis.set_visible(False)
ax2 = ax.twinx()
ax2.plot(sPos[bpm_inds]-sPos[st], 1e6*orball[bpm_inds,2], color='k', marker='x', linestyle='None', label='Fitted Orb at BPMs')
ax2.yaxis.set_label_position("left")
ax2.yaxis.tick_left()
ax2.xaxis.set_visible(True)
ax2.set_zorder(3)
ax2.set_xlabel('s [m]')
ax2.set_ylabel('y-orbit [um]')
ax2.legend()
plt.show()

this allows you to plot whatever you want with the AT synoptic at the bottom.

I gave this a try, but it tells me that sPos is not defined. I also recognized that the line where sPos is located has data not relevant to what I am trying to plot, and so I have to replace it with my own array of data. However, I don't have those arrays because I am trying to use the trajectory plot from plot_trajectory() and switch the legend labels. However, when I use any command relevant to get the legend from this plot (from here: https://matplotlib.org/stable/api/_as_gen/matplotlib.pyplot.legend.html or here: https://matplotlib.org/3.2.2/api/_as_gen/matplotlib.axes.Axes.get_figure.html), it tells me there is no attribute legend or attribute figure and so I don't seem able to alter it all.

[tfs563]

I did it like this for plot_beta. I think the same should be possible for plot_trajectory since both are calling the same plotting function. In this way I could modify only the parts I wanted of the existing figure.

def plot_optics(ring):

# Plot the optics functions
[left,right,syn] = ring.plot_beta()

# Get the figure to be able to make changes
fig = left.get_figure()

# Change width of the function
fig.set_figwidth(11)

# Set axes ranges
left.set_ylim(0,30)
right.set_ylim(0,1)
align.yaxes(left, 0, right, 0, 0.15)

# Set location of the legend
legend = left.get_legend()
legend.set_loc('upper right')

# Change legend names
legend.get_texts()[2].set_text('$\\eta_x$')

# Change label name
right.set_ylabel('$\\eta_x$')

# Save the plot to file
plt.savefig('optics.png')

I gave this a try, but still ran into issues. To see if this type of answer would fix it, I even used the plot_beta function with my lattice to see if it was maybe the plot_trajectory function that had issues, and I still ran into cases of the plot having no attributes such as the 'Legend' object has no attribute 'set_loc' and the set_texts commands altering nothing to the plot_beta function. Moving over to the plot_trajectory function still gave me issues, such as

I'm really starting to think the plot_trajectory function is impossible to change, or something is wrong with how things are setup on my end since I cannot alter anything to the plot functions.

[TeresiaOlsson]

Hmm... Maybe it could be something with the module versions.

Can you perhaps upload an example of your script and then I can test it in my environment to see if it works for me or I also get the same errors?

[tfs563]

Hmm... Maybe it could be something with the module versions.

Can you perhaps upload an example of your script and then I can test it in my environment to see if it works for me or I also get the same errors?

Sure, here is effectively the relevant parts of my script (other things are just cosmetic changes for altering plots and wouldn't influence AT)

import at
import at.plot as atplt
import numpy as np
import math
import matplotlib.pyplot as plt
import matplotlib.axes as axes
from math import pi
from scipy import constants
from scipy import integrate
from decimal import Decimal
from uncertainties import ufloat
from uncertainties import unumpy
from uncertainties.umath import *
import pandas as pd
import glob
import cmd
plt.rcParams["figure.figsize"] = [9.0, 6.0]

# Lattice elements 
Dr1 = at.Drift('Dr', 0.5)
Bend1 = at.Dipole('Bend', 0.9666, 0.6789) 
Dr2 = at.Drift('Dr', 1)
Monitor1 = at.Monitor('Monitor')

# Lattice cell
ECS_cell_simple = at.Lattice([Dr1, Bend1, Dr2, Monitor1], name='ECS_Dipole_cell', energy=250E6) 


# Generate a set of particles with different momentum to track
x_01 = 0.001, 0.001, 0.001
xp_01 = 0.000 
y_01 = 0.001, 0.001, 0.001 
yp_01 = 0 
dp_01 = 0, -0.0001, 0.0001 
ct_01 = 0 

N=3
zeros = [0]*N

# Generate the distribution of particles 
r_in_list = np.array([x_01, zeros, zeros, zeros, dp_01, zeros])

#Plot particle trajectories
ECS_E_spectrometer = ECS_cell_simple.plot_trajectory(r_in_list, legend=True)

# Plot the trajectory function with alterations
[left,right,syn] = ECS_cell_simple.plot_trajectory(r_in_list)

# Get the figure to be able to make changes
fig = left.get_figure()

# Change width of the function
fig.set_figwidth(11)

# Set axes ranges
left.set_ylim(0,30)
right.set_ylim(0,1)
align.yaxes(left, 0, right, 0, 0.15)

# Set location of the legend
legend = left.get_legend()
legend.set_loc('upper right')

# Change legend names
legend.get_texts()[2].set_text('$\\eta_x$')

# Change label name
right.set_ylabel('$\\eta_x$')

This is the basic script both without the solution you provided and with it included.

[TeresiaOlsson]

It worked for me. I modified your script a bit to plot the exact same figure twice and then make modifications to the second figure. The only difference I saw compared to plot_beta is that plot_trajectory has no right axes so that one will be None and you can't make any modifications on it.

Otherwise the only thing I could think off that might be the problem is the version of matplotlib. I have version 3.8.0.

This is what I did:

import at
import at.plot as atplt
import matplotlib.pyplot as plt
import numpy as np
# import math
# import matplotlib.pyplot as plt
# import matplotlib.axes as axes
# from math import pi
# from scipy import constants
# from scipy import integrate
# from decimal import Decimal
# from uncertainties import ufloat
# from uncertainties import unumpy
# from uncertainties.umath import *
# import pandas as pd
# import glob
# import cmd
# plt.rcParams["figure.figsize"] = [9.0, 6.0]

# Lattice elements 
Dr1 = at.Drift('Dr', 0.5)
Bend1 = at.Dipole('Bend', 0.9666, 0.6789) 
Dr2 = at.Drift('Dr', 1)
Monitor1 = at.Monitor('Monitor')

# Lattice cell
ECS_cell_simple = at.Lattice([Dr1, Bend1, Dr2, Monitor1], name='ECS_Dipole_cell', energy=250E6) 

# Generate a set of particles with different momentum to track
x_01 = 0.001, 0.001, 0.001
xp_01 = 0.000 
y_01 = 0.001, 0.001, 0.001 
yp_01 = 0 
dp_01 = 0, -0.0001, 0.0001 
ct_01 = 0 

N=3
zeros = [0]*N

# Generate the distribution of particles 
r_in_list = np.array([x_01, zeros, zeros, zeros, dp_01, zeros])

#Plot particle trajectories
ECS_E_spectrometer = ECS_cell_simple.plot_trajectory(r_in_list.copy(), legend=True)

# Plot the trajectory function with alterations
[left,right,syn] = ECS_cell_simple.plot_trajectory(r_in_list.copy(), legend=True)

# Get the figure to be able to make changes
fig = left.get_figure()

# Change width of the function
fig.set_figwidth(11)

# Set axes ranges
left.set_ylim(0,2e-3)
#right.set_ylim(0,1) # There is not right y axis in this plot so right is None
#align.yaxes(left, 0, right, 0, 0.15)

# Set location of the legend
legend = left.get_legend()
legend.set_loc('upper left')

# Change legend names
legend.get_texts()[1].set_text('test')
legend.get_texts()[2].set_text('test')

# Change label name
#right.set_ylabel('$\\eta_x$')

Without modifications:

With modifications:

[tfs563]

It worked for me. I modified your script a bit to plot the exact same figure twice and then make modifications to the second figure. The only difference I saw compared to plot_beta is that plot_trajectory has no right axes so that one will be None and you can't make any modifications on it.

Otherwise the only thing I could think off that might be the problem is the version of matplotlib. I have version 3.8.0.

This is what I did:

import at
import at.plot as atplt
import matplotlib.pyplot as plt
import numpy as np
# import math
# import matplotlib.pyplot as plt
# import matplotlib.axes as axes
# from math import pi
# from scipy import constants
# from scipy import integrate
# from decimal import Decimal
# from uncertainties import ufloat
# from uncertainties import unumpy
# from uncertainties.umath import *
# import pandas as pd
# import glob
# import cmd
# plt.rcParams["figure.figsize"] = [9.0, 6.0]

# Lattice elements 
Dr1 = at.Drift('Dr', 0.5)
Bend1 = at.Dipole('Bend', 0.9666, 0.6789) 
Dr2 = at.Drift('Dr', 1)
Monitor1 = at.Monitor('Monitor')

# Lattice cell
ECS_cell_simple = at.Lattice([Dr1, Bend1, Dr2, Monitor1], name='ECS_Dipole_cell', energy=250E6) 

# Generate a set of particles with different momentum to track
x_01 = 0.001, 0.001, 0.001
xp_01 = 0.000 
y_01 = 0.001, 0.001, 0.001 
yp_01 = 0 
dp_01 = 0, -0.0001, 0.0001 
ct_01 = 0 

N=3
zeros = [0]*N

# Generate the distribution of particles 
r_in_list = np.array([x_01, zeros, zeros, zeros, dp_01, zeros])

#Plot particle trajectories
ECS_E_spectrometer = ECS_cell_simple.plot_trajectory(r_in_list.copy(), legend=True)

# Plot the trajectory function with alterations
[left,right,syn] = ECS_cell_simple.plot_trajectory(r_in_list.copy(), legend=True)

# Get the figure to be able to make changes
fig = left.get_figure()

# Change width of the function
fig.set_figwidth(11)

# Set axes ranges
left.set_ylim(0,2e-3)
#right.set_ylim(0,1) # There is not right y axis in this plot so right is None
#align.yaxes(left, 0, right, 0, 0.15)

# Set location of the legend
legend = left.get_legend()
legend.set_loc('upper left')

# Change legend names
legend.get_texts()[1].set_text('test')
legend.get_texts()[2].set_text('test')

# Change label name
#right.set_ylabel('$\\eta_x$')

Without modifications:

With modifications:

Both good and bad news when I tried this again. The good news is that I did have an older version of matplotlib, and after I updated it to the latest version (3.8.4) I no longer got any attribute errors - or any errors - when running the script. The bad news is the plot still refuses to update. I get the exact same plot for both, which tells me it's still plotting the second plot_trajectory function, but the alterations aren't being applied to it for some reason.

I printed out the [left,right,syn] = ECS_cell_simple.plot_trajectory(r_in_list.copy(), legend=True) command both altogether and separately to see if that would give insight into what's wrong. I used the following

# Plot the trajectory function with alterations
[left,right,syn] = ECS_cell_simple.plot_trajectory(r_in_list.copy(), legend=True)
print([left,right,syn])
print(left)
print(right)
print(syn)

and got these as my outputs

Do they match what you get if you did the same thing?

[TeresiaOlsson]

Are you perhaps running it in a jupyter notebook? I discovered today that it didn't work for me when doing that. Then the plot didn't update. I had to load the following

from IPython.display import display

and then force the cell to update the plot with

display(fig)

after making the changes.

[tfs563]

Are you perhaps running it in a jupyter notebook? I discovered today that it didn't work for me when doing that. Then the plot didn't update. I had to load the following

from IPython.display import display

and then force the cell to update the plot with

display(fig)

after making the changes.

Yup, I was using jupyter notebook to run the scripts. After loading the same package and using the display(fig) line in my code I was also able to get the plots changed.

Thank you so much for the help!