atcollab / at

Accelerator Toolbox
Apache License 2.0
49 stars 32 forks source link

Select elements with wildcards in PyAt #100

Closed abdomit closed 5 years ago

abdomit commented 5 years ago

Hello,

First I'd like to thanks all the people working on PyAt. It is really great for Matlab-allergic people (like me). I am a PyAt user since a month, and it's working great.

But one thing I couldn't find in PyAt is a way to easily get elements from the lattice by their FamName, using a string supporting wildcards. This is something possible in Matlab-At, and it is really convenient (specially when the lattice you use is only made of Multipole elements).

For instance I would like to get all focusing quadrupole in my ring by just typing ring['QD*']. Maybe there is already a way to do it in PyAt, but I couldn't find it.

I did a quick implementation of this here. It is just 5 lines long, doesn't support Python2 and don't address the __setitem__ method. If there is no other way of doing this and you think it is worth it, I could do a pull request with it, trying to make it cleaner in the meantime.

Thanks, Benoit

lfarv commented 5 years ago

Hello @abdomit,

I thought of that, but I figured out that you can get the same thing either as an iterator with:

qd_iter=(elem for elem in ring if fnmatch.fnmatch(elem.FamName, 'QD*'))

for a single usage, or as a list with:

qd_list=[elem for elem in ring if fnmatch.fnmatch(elem.FamName, 'QD*')]

if you want to use the list several times… And as elements are indeed references, modifying the element in qd_list does modify it in ring. Modifying lattices using only simple commands give something like:

qd_strength = [elem.PolynomB[1] for elem in ring if fnmatch.fnmatch(elem.FamName, 'QD*')]

or in the other direction:

qd_iter=(elem for elem in ring if fnmatch.fnmatch(elem.FamName, 'QD*'))
for elem, strength in zip(qd_iter, qd_strength):
    elem.PolynomB[1] = strength

I have 2 comments on indexing directly the lattice:

  1. It creates a new list, which you can avoid using iterators if you do not need to keep the list
  2. It is restricted to a match on FamName, while you can imagine any other selection criteria

However, if there is a need for it, why not!

P.S. selecting focusing quads looking for 'QD*' might give surprising results!

abdomit commented 5 years ago

Hello Laurent,

I have 2 comments on indexing directly the lattice:

  1. It creates a new list, which you can avoid using iterators if you do not need to keep the list

Ok, but an iterator is not very user-friendy. You cannot address the nth element with

qd_list[n]

Moreover this list will not take much place in memory because it is basically just a list of references.

  1. It is restricted to a match on FamName, while you can imagine any other selection criteria

We could also have to possibility to give an Element class, for instance asking all the BPMs with

ring[at.Monitor]

Then concerning the possibility to select elements using other criteria, that would be interesting, but I think it should more look like this:

qd_list[qd_list.PolynomB[1] > 1.0]

However, if there is a need for it, why not!

In my opinion it is very very useful! I would be interested to know if other at-users feel the same.

P.S. selecting focusing quads looking for 'QD*' might give surprising results!

Now you know why I am not in the beam dynamics group!

T-Nicholls commented 5 years ago

If I may weigh in, I had a similar need, to get elements by class, and created a local utility function for the task when I first started using AT. As for implementation, supporting things like 'QD*' and element instances/types as valid arguments to __getitem__ makes me uncomfortable for a number of reasons. In Pytac we have a method on the lattice that gets elements by family; so might I suggest something similar for class and family name in this case.

abdomit commented 5 years ago

Done in #103

abdomit commented 5 years ago

Sorry to reopen this issue. I realized (a bit too late) that something was missing.

Now we have a get_elements function, but it would be even better to have in addition to this a corresponding method for Lattice object (as suggested by @lfarv: https://github.com/atcollab/at/pull/102#issuecomment-478246951).

lfarv commented 5 years ago

Ok, I take care of this in #99, almost ready!

abdomit commented 5 years ago

get_elements is now a method of Lattice object (#99).

Thanks @lfarv!