Working with tables

Several methods Xsuite return information in the form of Table objects. Examples of such methods are xtrack.Line.get_table(), xtrack.Line.twiss(), xtrack.Line.survey(), xtrack.Line.vars.get_table(), xtrack.Environment.vars.get_table().

See Table class in the Reference guide for the base table API, and Survey for the API of line.survey().

Table objects offer several capabilities to access data, for example by selecting rows and columns in various ways. This is illustrated in the following example:

import numpy as np
import xtrack as xt

pi = np.pi
lbend = 3

# Build a simple ring
env = xt.Environment()
line = env.new_line(components=[
    env.new('mqf.1', xt.Quadrupole, length=0.3, k1=0.1),
    env.new('d1.1',  xt.Drift, length=1),
    env.new('mb1.1', xt.Bend, length=lbend, angle=pi / 2),
    env.new('d2.1',  xt.Drift, length=1),

    env.new('mqd.1', xt.Quadrupole, length=0.3, k1=-0.7),
    env.new('d3.1',  xt.Drift, length=1),
    env.new('mb2.1', xt.Bend, length=lbend, angle=pi / 2),
    env.new('d4.1',  xt.Drift, length=1),

    env.new('mqf.2', xt.Quadrupole, length=0.3, k1=0.1),
    env.new('d1.2',  xt.Drift, length=1),
    env.new('mb1.2', xt.Bend, length=lbend, angle=pi / 2),
    env.new('d2.2',  xt.Drift, length=1),

    env.new('mqd.2', xt.Quadrupole, length=0.3, k1=-0.7),
    env.new('d3.2',  xt.Drift, length=1),
    env.new('mb2.2', xt.Bend, length=lbend, angle=pi / 2),
    env.new('d4.2',  xt.Drift, length=1),
])
line.set_particle_ref('proton', p0c=1.2e9)

# Get the twiss table of the line with the strengths
tab = line.twiss4d(strengths=True)

# The table can be printed
tab.show()
# prints:
#
# name       element_type             s          betx          bety ...
# mqf.1      Quadrupole               0       1.27738       4.79104
# d1.1       Drift                  0.3       1.27738       4.79104
# mb1.1      Bend                   1.3       2.26749       5.20838
# d2.1       Drift                  4.3       2.88391       8.99176
# mqd.1      Quadrupole             5.3       1.72481       11.0967
# d3.1       Drift                  5.6       1.72481       11.0967
# mb2.1      Bend                   6.6       2.88391       8.99176
# d4.1       Drift                  9.6       2.26749       5.20838
# mqf.2      Quadrupole            10.6       1.27738       4.79104
# d1.2       Drift                 10.9       1.27738       4.79104
# mb1.2      Bend                  11.9       2.26749       5.20838
# d2.2       Drift                 14.9       2.88391       8.99176
# mqd.2      Quadrupole            15.9       1.72481       11.0967
# d3.2       Drift                 16.2       1.72481       11.0967
# mb2.2      Bend                  17.2       2.88391       8.99176
# d4.2       Drift                 20.2       2.26749       5.20838
# _end_point                       21.2       1.27738       4.79104

# Single values can be accessed using the column name and the row name. For example:
tab['s', 'mb2.1'] # is 6.6

# Entire columns can be accessed using the column name. For example:
tab['s'] # is [0.0, 0.3, 1.3, 4.3, 5.3, 5.6, 6.6, 9.6, 10.6, 10.9, 11.9, ...

# The `.cols` attribute can be used to access multiple columns (the output is 
# a Table object). For example:
tab.cols['betx bety alfx alfy']
# returns:
#
# TwissTable: 17 rows, 5 cols
# name                betx          bety          alfx          alfy
# mqf.1            1.27738       4.79104     0.0997349      0.103198
# d1.1             1.27738       4.79104    -0.0997349     -0.103198
# mb1.1            2.26749       5.20838     -0.890376     -0.314144
# d2.1             2.88391       8.99176      0.890376     -0.946981
# mqd.1            1.72481       11.0967      0.268732      -1.15793
# d3.1             1.72481       11.0967     -0.268732       1.15793
# mb2.1            2.88391       8.99176     -0.890376      0.946981
# d4.1             2.26749       5.20838      0.890376      0.314144
# mqf.2            1.27738       4.79104     0.0997349      0.103198
# d1.2             1.27738       4.79104    -0.0997349     -0.103198
# mb1.2            2.26749       5.20838     -0.890376     -0.314144
# d2.2             2.88391       8.99176      0.890376     -0.946981
# mqd.2            1.72481       11.0967      0.268732      -1.15793
# d3.2             1.72481       11.0967     -0.268732       1.15793
# mb2.2            2.88391       8.99176     -0.890376      0.946981
# d4.2             2.26749       5.20838      0.890376      0.314144
# _end_point       1.27738       4.79104     0.0997349      0.103198

# Simple expressions can be used in the cols
tab.cols['betx dx dx/sqrt(betx)']
# returns:
#
# TwissTable: 17 rows, 4 cols
# name                betx            dx dx/sqrt(betx)
# mqf.1            1.27738       2.27721       2.01486
# d1.1             1.27738       2.27721       2.01486
# mb1.1            2.26749       2.24303       1.48958
# d2.1             2.88391       1.84457       1.08619
# mqd.1            1.72481       1.67012       1.27168
# d3.1             1.72481       1.67012       1.27168
# mb2.1            2.88391       1.84457       1.08619
# d4.1             2.26749       2.24303       1.48958
# mqf.2            1.27738       2.27721       2.01486
# d1.2             1.27738       2.27721       2.01486
# mb1.2            2.26749       2.24303       1.48958
# d2.2             2.88391       1.84457       1.08619
# mqd.2            1.72481       1.67012       1.27168
# d3.2             1.72481       1.67012       1.27168
# mb2.2            2.88391       1.84457       1.08619
# d4.2             2.26749       2.24303       1.48958
# _end_point       1.27738       2.27721       2.01486

# The .rows attribute can be used access selected rows
tab.rows[['mqf.1', 'mqd.1']]
# returns:
#
# TwissTable: 2 rows, 98 cols
# TwissTable: 2 rows, 98 cols
# name  element_type             s          betx          bety ...
# mqf.1 Quadrupole               0       1.27738       4.79104
# mqd.1 Quadrupole             5.3       1.72481       11.0967

# Regular expressions can be used to select rows
tab.rows['mb.*']
# returns:
#
# TwissTable: 4 rows, 98 cols
# name  element_type             s          betx          bety ...
# mb1.1 Bend                   1.3       2.26749       5.20838
# mb2.1 Bend                   6.6       2.88391       8.99176
# mb1.2 Bend                  11.9       2.26749       5.20838
# mb2.2 Bend                  17.2       2.88391       8.99176

# Elements can be selected by type using the match search (applicable to any column)
tab.rows.match(element_type='Quadrupole')
# returns:
#
# TwissTable: 4 rows, 98 cols
# name  element_type             s          betx          bety ...
# mqf.1 Quadrupole               0       1.27738       4.79104
# mqd.1 Quadrupole             5.3       1.72481       11.0967
# mqf.2 Quadrupole            10.6       1.27738       4.79104
# mqd.2 Quadrupole            15.9       1.72481       11.0967

# rows.match supports regular expressions
tab.rows.match(element_type='Quad.*|Be.*')
# returns:
# TwissTable: 8 rows, 98 cols
# name  element_type             s          betx          bety ...
# mqf.1 Quadrupole               0       1.27738       4.79104
# mb1.1 Bend                   1.3       2.26749       5.20838
# mqd.1 Quadrupole             5.3       1.72481       11.0967
# mb2.1 Bend                   6.6       2.88391       8.99176
# mqf.2 Quadrupole            10.6       1.27738       4.79104
# mb1.2 Bend                  11.9       2.26749       5.20838
# mqd.2 Quadrupole            15.9       1.72481       11.0967
# mb2.2 Bend                  17.2       2.88391       8.99176

# rows.match_not can be used to select rows not matching a given condition.
# For example to select all elements that are not drifts:
tab.rows.match_not(element_type='Drift')
#
# TwissTable: 9 rows, 98 cols
# name  element_type             s          betx          bety ...
# mqd.1 Quadrupole             5.3       1.72481       11.0967
# d3.1  Drift                  5.6       1.72481       11.0967
# mb2.1 Bend                   6.6       2.88391       8.99176
# d4.1  Drift                  9.6       2.26749       5.20838
# mqf.2 Quadrupole            10.6       1.27738       4.79104
# d1.2  Drift                 10.9       1.27738       4.79104
# mb1.2 Bend                  11.9       2.26749       5.20838
# d2.2  Drift                 14.9       2.88391       8.99176
# mqd.2 Quadrupole            15.9       1.72481       11.0967

# A section of the table can be selected using names
tab.rows['mqd.1':'mqd.2']
# returns:
#
# TwissTable: 9 rows, 98 cols
# name  element_type             s          betx          bety ...
# mqd.1 Quadrupole             5.3       1.72481       11.0967
# d3.1  Drift                  5.6       1.72481       11.0967
# mb2.1 Bend                   6.6       2.88391       8.99176
# d4.1  Drift                  9.6       2.26749       5.20838
# mqf.2 Quadrupole            10.6       1.27738       4.79104
# d1.2  Drift                 10.9       1.27738       4.79104
# mb1.2 Bend                  11.9       2.26749       5.20838
# d2.2  Drift                 14.9       2.88391       8.99176
# mqd.2 Quadrupole            15.9       1.72481       11.0967

# A section of the table can be selected using any given column. For example,
# to select all elements between s=3 and s=7:
tab.rows[3.0:7.0:'s']
# returns:
#
# TwissTable: 4 rows, 98 cols
# name  element_type             s          betx          bety ...
# d2.1  Drift                  4.3       2.88391       8.99176
# mqd.1 Quadrupole             5.3       1.72481       11.0967
# d3.1  Drift                  5.6       1.72481       11.0967
# mb2.1 Bend                   6.6       2.88391       8.99176

# A section of the table can be selected using indexes relative one element
# (e.g. to get from three elements upstream of 'mqd.1' to two elements
# downstream of 'mb2.1')
tab.rows['mqd.1<<3':'mb2.1>>2']
# returns:
#
# TwissTable: 8 rows, 98 cols
# name  element_type             s          betx          bety ...
# d1.1  Drift                  0.3       1.27738       4.79104
# mb1.1 Bend                   1.3       2.26749       5.20838
# d2.1  Drift                  4.3       2.88391       8.99176
# mqd.1 Quadrupole             5.3       1.72481       11.0967
# d3.1  Drift                  5.6       1.72481       11.0967
# mb2.1 Bend                   6.6       2.88391       8.99176
# d4.1  Drift                  9.6       2.26749       5.20838
# mqf.2 Quadrupole            10.6       1.27738       4.79104

# As rows and column selectors return Table objects they can be chained for example
# in the following we select rows in the range 'd1.1'-'d2.2', which are not of type Drift,
# and that match the regular expression 'mb.*' and we select the columns `betx` and `bety`:
tab.rows['d1.1':'d2.2'].rows.match_not(element_type='Drift').rows.match(name='mb.*').cols['betx bety']
# returns:
#
# TwissTable: 3 rows, 3 cols
# name           betx          bety
# mb1.1       2.26749       5.20838
# mb2.1       2.88391       8.99176
# mb1.2       2.26749       5.20838

# Complete source: xtrack/examples/toy_ring/004b_work_with_table.py