xsuite / xplt

Plotting library for Xsuite and other accelerator physics codes
https://xsuite.github.io/xplt/
Apache License 2.0
10 stars 4 forks source link

Thick elements and labels from list in FloorPlot #4

Closed giadarol closed 8 months ago

giadarol commented 8 months ago

This PR adds:

Here a small example (using a lattice available in the xtrack test data):

from cpymad.madx import Madx
import xtrack as xt

test_data_folder = '../../test_data/'
mad = Madx()

mad.call(test_data_folder + 'pimms/PIMM_orig.seq')
mad.call(test_data_folder + 'pimms/betatron.str')
mad.beam(particle='proton', gamma=1.21315778) # 200 MeV
mad.use('pimms')
seq = mad.sequence.pimms
def_expr = True

line = xt.Line.from_madx_sequence(seq, deferred_expressions=def_expr)
line.particle_ref = xt.Particles(gamma0=seq.beam.gamma,
                                 mass0=seq.beam.mass * 1e9,
                                 q0=seq.beam.charge)

sv = line.survey()

tt = line.get_table()
ttsext = tt.rows[tt.element_type == 'Sextupole']

import xplt
import matplotlib.pyplot as plt
plt.close('all')
xplt.FloorPlot(sv, line, labels=ttsext.name)

plt.show()
image

PS: I would like to show these features for the Xsuite tutorial at the Slow Extraction Workshop. What do you think?


My contribution follows "inbound=outbound" licensing as defined by the GitHub Terms of Service.

giadarol commented 8 months ago

I just saw the test not passing because of the import of xtrack. That can be avoided. I'll update the PR

giadarol commented 8 months ago

Just updated. I also commented out the alterations to the matplotlib rcParams, as this changes the aspect of all other plots within the same session, which I find a bit invasive :-)

eltos commented 8 months ago

I just saw the test not passing because of the import of xtrack. That can be avoided. I'll update the PR

Hi @giadarol and many thanks for adding support for thick elements here!
I made a few more adjustments to handle line==None properly in case of having a quick look at MAD-X survey (without the nice wedges of course). Does this work for you?

eltos commented 8 months ago

Just updated. I also commented out the alterations to the matplotlib rcParams, as this changes the aspect of all other plots within the same session, which I find a bit invasive :-)

I agree that this is invasive, however instead of uncommenting that, it should go into a matplotlib style sheet, so it can be called by the user manually (I would still like to use these "good defaults" in my Jupyter notebooks with a single line):

xplt.apply_style()
eltos commented 8 months ago

PS: I would like to show these features for the Xsuite tutorial at the Slow Extraction Workshop. What do you think?

It would be an honour :)

giadarol commented 8 months ago

I just saw the test not passing because of the import of xtrack. That can be avoided. I'll update the PR

Hi @giadarol and many thanks for adding support for thick elements here! I made a few more adjustments to handle line==None properly in case of having a quick look at MAD-X survey (without the nice wedges of course). Does this work for you?

yes!

giadarol commented 8 months ago

PS: I would like to show these features for the Xsuite tutorial at the Slow Extraction Workshop. What do you think?

It would be an honour :)

Thank you! Would you manage to make also a PyPI release by then? So that people will be able to run the example notebooks that I will show after a simple pip install of the packages

eltos commented 8 months ago

Yes, I can make a release in the next days.

eltos commented 8 months ago

@giadarol I pushed a new version, CI is running and it should soon appear on PyPI. However, I noticed that there seems to be an issue in XZ-projection (projection="XZ"): for thick elements, the wedge is slightly displaced. In the default ZX-projection it works though.

xt.Multipole xt.Bend
grafik grafik