Miscellaneous changes to pyat

pyat/at/lattice/lattice_object.py

@@ -5,7 +5,7 @@
 from scipy.constants import physical_constants as cst
 from warnings import warn
 from at.lattice import elements, get_s_pos, checktype, uint32_refpts
-from at.lattice import AtWarning, AtError
+from at.lattice import AtWarning, AtError, get_ring_energy
 from at.physics import find_orbit4, find_orbit6, find_sync_orbit, find_m44
 from at.physics import find_m66, linopt, ohmi_envelope, get_mcf
 
@@ -53,7 +53,6 @@ def full_scan(elems, keep_all=False, **kwargs):
             """
             params = []
             valid_elems = []
-            energies = []
             thetas = []
             attributes = dict(_radiation=False)
 
@@ -68,8 +67,6 @@ def full_scan(elems, keep_all=False, **kwargs):
                     warn(AtWarning('item {0} ({1}) is not an AT element: '
                     'ignored'.format(idx, elem)))
                     continue
-                if hasattr(elem, 'Energy'):
-                    energies.append(elem.Energy)
                 if isinstance(elem, elements.Dipole):
                     thetas.append(elem.BendingAngle)
                 if elem.PassMethod.endswith(
@@ -90,13 +87,7 @@ def full_scan(elems, keep_all=False, **kwargs):
                 attributes['name'] = ''
                 if 'energy' not in kwargs:
                     # Guess energy from the Energy attribute of the elems
-                    if not energies:
-                        raise AtError('Lattice energy is not defined')
-                    energy = max(energies)
-                    if min(energies) < energy:
-                        warn(AtWarning('Inconsistent energy values, '
-                                       '"energy" set to {0}'.format(energy)))
-                    attributes['energy'] = energy
+                    attributes['energy'] = get_ring_energy(elems)
                 if 'periodicity' not in kwargs:
                     # Guess periodicity from the bending angles of the lattice
                     try:
@@ -155,8 +146,7 @@ def __getitem__(self, key):
             return super(Lattice, self).__getitem__(key)
         except TypeError:
             key = uint32_refpts(key, len(self))
-            items = [super(Lattice, self).__getitem__(i) for i in key]
-            return items[0] if len(items) == 1 else items
+            return [super(Lattice, self).__getitem__(i) for i in key]
 
     def __setitem__(self, key, values):
         try:

pyat/at/lattice/utils.py

@@ -16,6 +16,8 @@
 """
 import numpy
 import itertools
+from warnings import warn
+from at.lattice import elements
 
 
 class AtError(Exception):
@@ -185,6 +187,41 @@ def get_s_pos(ring, refpts=None):
     return s_pos[refpts]
 
 
+def get_ring_energy(ring):
+    """Establish the energy of the ring from the Energy attribute of the
+    elements. Energies of RingParam elements are most prioritised, if none are
+    found then the energies from RFCavity elements will be used, if none are
+    found then the energies from all elements will be used. An error will be
+    raised if no elements have a 'Energy' attribute or if inconsistent values
+    for energy are found.
+
+    Args:
+        ring: sequence of elements of which you wish to establish the energy.
+    """
+    rp_energies = []
+    rf_energies = []
+    energies = []
+    for elem in ring:
+        if hasattr(elem, 'Energy'):
+            energies.append(elem.Energy)
+            if isinstance(elem, elements.RingParam):
+                rp_energies.append(elem.Energy)
+            elif isinstance(elem, elements.RFCavity):
+                rf_energies.append(elem.Energy)
+    if not energies:
+        raise AtError('Lattice energy is not defined.')
+    elif rp_energies:
+        energy = max(rp_energies)
+    elif rf_energies:
+        energy = max(rf_energies)
+    else:
+        energy = max(energies)
+    if len(set(energies)) > 1:
+        warn(AtWarning('Inconsistent energy values in ring, {0} has been '
+                       'used.'.format(energy)))
+    return energy
+
+
 def tilt_elem(elem, rots):
     """
     set a new tilt angle to an element. The rotation matrices are stored in the R1 and R2 attributes

pyat/at/load/utils.py

@@ -148,12 +148,10 @@ def sanitise_class(index, class_name, kwargs):
             AttributeError: if the PassMethod and Class are incompatible.
         """
         def err_message(message, *args):
-            location = '' if index is None else 'Error in element {0}: '.format(
-                index)
-            return ''.join(
-                (location,
-                 'PassMethod {0} is not compatible with '.format(pass_method),
-                 message.format(*args), '\n{0}'.format(kwargs)))
+            location = ': ' if index is None else ' {0}: '.format(index)
+            return ''.join(('Error in element', location,
+                            'PassMethod {0} '.format(pass_method),
+                             message.format(*args), '\n{0}'.format(kwargs)))
 
         pass_method = kwargs.get('PassMethod')
         if pass_method is not None:
@@ -164,19 +162,22 @@ def err_message(message, *args):
                                                       '../../integrators',
                                                       pass_method + extension))
             if not os.path.exists(file_path):
-                raise AttributeError('PassMethod {0} does not exist.'
-                                     '\n{1}'.format(pass_method, kwargs))
+                raise AttributeError(err_message("does not exist."))
             elif (pass_method == 'IdentityPass') and (length != 0.0):
-                raise AttributeError(err_message("length {0}.", length))
+                raise AttributeError(err_message("is not compatible with "
+                                                 "length {0}.", length))
             elif pass_to_class is not None:
                 if pass_to_class != class_name:
-                    raise AttributeError(err_message("Class {0}.", class_name))
+                    raise AttributeError(err_message("is not compatible with "
+                                                     "Class {0}.", class_name))
             elif class_name in ['Marker', 'Monitor', 'RingParam']:
                 if pass_method != 'IdentityPass':
-                    raise AttributeError(err_message("Class {0}.", class_name))
+                    raise AttributeError(err_message("is not compatible with "
+                                                     "Class {0}.", class_name))
             elif class_name == 'Drift':
                 if pass_method != 'DriftPass':
-                    raise AttributeError(err_message("Class {0}.", class_name))
+                    raise AttributeError(err_message("is not compatible with "
+                                                     "Class {0}.", class_name))
 
     class_name = find_class_name(elem_dict, quiet=quiet)
     if check:

pyat/at/physics/linear.py

@@ -146,7 +146,7 @@ def get_twiss(ring, dp=0.0, refpts=None, get_chrom=False, orbit=None,
                       l_down['closed_orbit'])[:, :4] / ddp
         disp0 = (d0_up['closed_orbit'] - d0_down['closed_orbit'])[:4] / ddp
     else:
-        chrom = None
+        chrom = numpy.array([numpy.NaN, numpy.NaN])
         dispersion = numpy.NaN
         disp0 = numpy.NaN
 
@@ -314,7 +314,7 @@ def analyze(r44):
                       l_down['closed_orbit'])[:, :4] / ddp
         disp0 = (d0_up['closed_orbit'] - d0_down['closed_orbit'])[:4] / ddp
     else:
-        chrom = None
+        chrom = numpy.array([numpy.NaN, numpy.NaN])
         dispersion = numpy.NaN
         disp0 = numpy.NaN
 

pyat/at/physics/radiation.py

@@ -3,9 +3,8 @@
 """
 import numpy
 from scipy.linalg import inv, det, solve_sylvester
-from warnings import warn
 import at
-from at.lattice import uint32_refpts, RFCavity, AtError, AtWarning
+from at.lattice import uint32_refpts, get_ring_energy
 from at.tracking import lattice_pass
 from at.physics import find_orbit6, find_m66, find_elem_m66, get_tunes_damp
 # noinspection PyUnresolvedReferences
@@ -120,21 +119,7 @@ def propag(m, cumb, orbit6):
         if isinstance(ring, at.lattice.Lattice):
             energy = ring.energy
         else:
-            rf_energies = []
-            energies = []
-            for elem in ring:
-                if hasattr(elem, 'Energy'):
-                    energies.append(elem.Energy)
-                if isinstance(elem, RFCavity):
-                    rf_energies.append(elem.Energy)
-            if len(energies) is 0:
-                raise AtError('Lattice energy is not defined')
-            else:
-                energy = max(energies if len(rf_energies) is 0
-                             else rf_energies)
-                if len(set(energies)) > 1:
-                    warn(AtWarning('Inconsistent energy values in ring, {0} '
-                                   'has been used'.format(energy)))
+            energy = get_ring_energy(ring)
 
     if orbit is None:
         orbit, _ = find_orbit6(ring, keep_lattice=keep_lattice)
@@ -171,8 +156,11 @@ def propag(m, cumb, orbit6):
     data0 = numpy.rec.fromarrays(
         (rr, rr4, mring, orbit, emitxy, emitxyz),
         dtype=ENVELOPE_DTYPE)
-    data = numpy.rec.fromrecords(
-        list(map(propag, ms, bbcum[uint32refs], orbs[uint32refs, :])),
-        dtype=ENVELOPE_DTYPE)
+    if uint32refs.shape == (0,):
+        data = numpy.recarray((0,), dtype=ENVELOPE_DTYPE)
+    else:
+        data = numpy.rec.fromrecords(
+            list(map(propag, ms, bbcum[uint32refs], orbs[uint32refs, :])),
+            dtype=ENVELOPE_DTYPE)
 
     return data0, r66data, data

pyat/requirements.txt

@@ -1,4 +1,4 @@
 pytest>=2.9
-numpy>=1.13
+numpy>=1.10
 scipy>=0.16
 pytest-lazy-fixture

pyat/test/test_lattice_object.py

@@ -82,7 +82,7 @@ def test_lattice_creation_warnings_and_errors():
 def test_lattice_string_ordering():
     l = Lattice([elements.Drift('D0', 1.0, attr1=numpy.array(0))], name='lat',
                 energy=5, periodicity=1, attr2=3)
-    # Dictionary ordering is only in Python >= 3.6
+    # Default dictionary ordering is only in Python >= 3.6
     if sys.version_info < (3, 6):
         assert l.__str__().startswith("Lattice(<1 elements>, ")
         assert l.__str__().endswith(", attr2=3)")