at.beam improvement (#537)

* add x-y and x-delta coupling

* bug fix

pyat/at/tracking/particles.py

@@ -2,6 +2,7 @@
 Functions relating to particle generation
 """
 import numpy as np
+from numpy.linalg import cholesky, LinAlgError
 from warnings import warn
 from ..physics import Orbit, ohmi_envelope
 from ..lattice import AtError, AtWarning, Lattice, random
@@ -63,47 +64,48 @@ def ignore(mess):
         if any(ok):
             warn(AtWarning('{0}: {1} ignored'.format(mess, ', '.join(ok))))
 
-    def r_4d(ax, ay, bx, by):
+    def r_s():
+        rs = np.zeros((2, 2))
+        if any(('espread' in kwargs, 'blength' in kwargs)):
+            blength, espread = require(
+                'Both "blength" and "espread" are required',
+                'blength', 'espread')
+            ems = blength * espread
+            if ems > 0.0:
+                rs = np.array([[espread / blength, 0], [0, blength / espread]])
+        else:
+            ems = None
+        return ems, rs
+
+    def r_4d(ax, ay, bx, by, rs):
         r6 = np.zeros((3, 6, 6))
         r6[0, :2, :2] = np.array([[bx, -ax], [-ax, (1+ax**2)/bx]])
         r6[1, 2:4, 2:4] = np.array([[by, -ay], [-ay, (1+ay**2)/by]])
+        r6[2, 4:, 4:] = rs
         return r6
 
-    def expand(rmat):
+    def r_expand(rmat, rs):
         if rmat.shape[0] < 3:
             r6 = np.zeros((3, 6, 6))
             r6[:2, :4, :4] = rmat
+            r6[2, 4:, 4:] = rs
         else:
             r6 = rmat
         return r6
 
-    def sigma_from_rmat(rmat, emit):
+    def sigma_from_rmat(rmat, emit, embl):
 
         if emit is None:
             emx, emy = require('Emittances must be provided', 'emitx', 'emity')
-            ems = None
+            ems = embl
         else:
             emx = kwargs.pop('emitx', emit[0])
             emy = kwargs.pop('emity', emit[1])
-            ems = emit[2]
-            if any(('espread' in kwargs, 'blength' in kwargs)):
-                blength, espread = require(
-                    'Both "blength" and "espread" are required',
-                    'blength', 'espread')
-                ems = blength * espread
+            ems = emit[2] if embl is None else embl
 
         sigm = emx * rmat[0, :, :] + emy * rmat[1, :, :]
         if ems is None:
-            if any(('espread' in kwargs, 'blength' in kwargs)):
-                blength, espread = require(
-                    'Both "blength" and "espread" are required',
-                    'blength', 'espread')
-                rmat[2, 4:6, 4:6] = np.array([[espread / blength, 0],
-                                              [0, blength / espread]])
-                ems = blength * espread
-                sigm += ems * rmat[2, :, :]
-            else:
-                warn(AtWarning('Monochromatic beam: no energy spread'))
+            warn(AtWarning('Monochromatic beam: no energy spread'))
         else:
             sigm += ems * rmat[2, :, :]
         return sigm
@@ -123,36 +125,38 @@ def sigma_from_rmat(rmat, emit):
     elif 'twiss_in' in kwargs:
         twin = kwargs.pop('twiss_in')
         message = "'twiss_in' is provided"
-#       orbit = kwargs.pop('orbit', twin.getattr('closed_orbit', np.zeros(6)))
+        em56, r56 = r_s()
         try:
             Rin = twin['R']
         except (ValueError, KeyError):  # record arrays throw ValueError !
             Rm = r_4d(twin['alpha'][0], twin['beta'][0],
-                      twin['alpha'][1], twin['beta'][1])
+                      twin['alpha'][1], twin['beta'][1], r56)
         else:
-            Rm = expand(Rin)
-        sigmat = sigma_from_rmat(Rm, None)
+            Rm = r_expand(Rin, r56)
+        sigmat = sigma_from_rmat(Rm, None, em56)
 
     elif 'ring' in kwargs:
         ring = kwargs.pop('ring')
         message = "'ring' is provided"
+        em56, r56 = r_s()
         el0, _, _ = ring.linopt6(orbit=kwargs.pop('orbit', None))
         orbit = el0.closed_orbit
-        Rm = expand(el0['R'])
+        Rm = r_expand(el0['R'], r56)
         if ring.is_6d:
             _, beam0, _ = ohmi_envelope(ring, orbit=orbit)
-            sigmat = sigma_from_rmat(Rm, beam0.mode_emittances)
+            sigmat = sigma_from_rmat(Rm, beam0.mode_emittances, em56)
         else:
-            sigmat = sigma_from_rmat(Rm, None)
+            sigmat = sigma_from_rmat(Rm, None, em56)
 
     else:
         message = "Uncoupled beam"
         alphax, alphay, betax, betay = require(
             'Neither "twiss_in" nor "ring" is provided',
             'alphax', 'alphay', 'betax', 'betay')
         # orbit = kwargs.pop('orbit', np.zeros(6))
-        Rm = r_4d(alphax, alphay, betax, betay)
-        sigmat = sigma_from_rmat(Rm, None)
+        em56, r56 = r_s()
+        Rm = r_4d(alphax, alphay, betax, betay, r56)
+        sigmat = sigma_from_rmat(Rm, None, em56)
 
     ignore(message)
     return sigmat
@@ -175,19 +179,39 @@ def beam(nparts: int, sigma, orbit: Orbit = None):
     """
     def _get_single_plane(slc):
         try:
-            return np.linalg.cholesky(sigma[slc, slc])
-        except np.linalg.LinAlgError:
+            return cholesky(sigma[slc, slc])
+        except LinAlgError:
             return np.zeros((2, 2))
 
     v = random.thread.normal(size=(sigma.shape[1], nparts))
 
     try:
-        lmat = np.linalg.cholesky(sigma)
-    except np.linalg.LinAlgError:
+        # Try full 6x6 matrix
+        lmat = cholesky(sigma)
+        print("h, v, delta")
+    except LinAlgError:
         lmat = np.zeros((6, 6))
-        lmat[:2, :2] = _get_single_plane(slice(2))
-        lmat[2:4, 2:4] = _get_single_plane(slice(2, 4))
-        lmat[4:, 4:] = _get_single_plane(slice(4, 6))
+        try:
+            # Try x-y 4x4 matrix
+            sel = range(4)
+            idx = np.ix_(sel, sel)
+            lmat[idx] = cholesky(sigma[idx])
+            lmat[4:, 4:] = _get_single_plane(slice(4, 6))
+            print("h, v")
+        except LinAlgError:
+            try:
+                # Try x-delta 4x4 matrix
+                sel = [0, 1, 4, 5]
+                idx = np.ix_(sel, sel)
+                lmat[idx] = cholesky(sigma[idx])
+                lmat[2:4, 2:4] = _get_single_plane(slice(2, 4))
+                print("h, delta")
+            except LinAlgError:
+                # Uncoupled
+                lmat[:2, :2] = _get_single_plane(slice(2))
+                lmat[2:4, 2:4] = _get_single_plane(slice(2, 4))
+                lmat[4:, 4:] = _get_single_plane(slice(4, 6))
+                print("uncoupled")
 
     particle_dist = np.asfortranarray(lmat @ v)