Merge branch 'master' into modify_beamloading

Author: Lee Carver

pyat/at.c

@@ -342,7 +342,7 @@ static PyObject *at_atpass(PyObject *self, PyObject *args, PyObject *kwargs) {
                              "bunch_spos", "bunch_currents", "start_elem", "end_elem", NULL};
     static double lattice_length = 0.0;
     static int last_turn = 0;
-    static int valid = 0;
+    static PyObject *lattice_id = NULL;
 
     PyObject *lattice;
     PyObject *particle;
@@ -434,11 +434,12 @@ static PyObject *at_atpass(PyObject *self, PyObject *args, PyObject *kwargs) {
     }
     #endif /*_OPENMP*/
 
-    if (!(keep_lattice && valid)) {
+    if (!(keep_lattice && lattice_id == lattice)) {
         PyObject **element;
         double *elem_length;
         track_function *integrator;
         PyObject **pyintegrator;
+
         /* Release the stored elements */
         for (elem_index=0; elem_index < num_elements; elem_index++) {
             free(elemdata_list[elem_index]);
@@ -499,7 +500,7 @@ static PyObject *at_atpass(PyObject *self, PyObject *args, PyObject *kwargs) {
             *elem_length++ = length;
             Py_INCREF(el);                          /* Keep a reference to each element in case of reuse */
         }
-        valid = 0;
+        lattice_id = NULL;
     }
 
     /* Check for partial turn tracking */
@@ -641,7 +642,7 @@ static PyObject *at_atpass(PyObject *self, PyObject *args, PyObject *kwargs) {
         param.nturn++;
         start_elem = 0;
     }
-    valid = 1;      /* Tracking successful: the lattice can be reused */
+    lattice_id = lattice;     /* Tracking successful: the lattice can be reused */
     last_turn = param.nturn;  /* Store turn number in a static variable */
 
     #ifdef _OPENMP

pyat/at/lattice/elements/basic_elements.py

@@ -556,24 +556,24 @@ def __init__(
         if taux == 0.0:
             dampx = 1
         else:
-            dampx = np.exp(-1 / taux)
+            dampx = np.exp(-2 / taux)
 
         assert tauy >= 0.0, "tauy must be greater than or equal to 0"
         if tauy == 0.0:
             dampy = 1
         else:
-            dampy = np.exp(-1 / tauy)
+            dampy = np.exp(-2 / tauy)
 
         assert tauz >= 0.0, "tauz must be greater than or equal to 0"
         if tauz == 0.0:
             dampz = 1
         else:
-            dampz = np.exp(-1 / tauz)
+            dampz = np.exp(-2 / tauz)
 
         kwargs.setdefault("PassMethod", self.default_pass[True])
         kwargs.setdefault("U0", U0)
         kwargs.setdefault(
-            "damp_mat_diag", np.array([dampx, dampx, dampy, dampy, dampz, dampz])
+            "damp_mat_diag", np.array([1, dampx, 1, dampy, dampz, 1])
         )
 
         super().__init__(family_name, **kwargs)

pyat/at/tracking/track.py

@@ -1,14 +1,14 @@
 from __future__ import annotations
 
 __all__ = [
-    "lattice_track",
+    "MPMode",
     "element_track",
-    "internal_lpass",
+    "gpu_core_count",
+    "gpu_info",
     "internal_epass",
+    "internal_lpass",
     "internal_plpass",
-    "gpu_info",
-    "gpu_core_count",
-    "MPMode",
+    "lattice_track",
 ]
 
 import multiprocessing
@@ -31,7 +31,7 @@
 
 class MPMode(Enum):
     """
-    Multi Processing mode
+    Multi Processing mode.
     """
 
     CPU = 1  #: CPU multiprocessing
@@ -55,14 +55,14 @@ class MPMode(Enum):
 
 
 def _atpass_fork(seed, rank, rin, **kwargs):
-    """Single forked job"""
+    """Single forked job."""
     reset_rng(rank=rank, seed=seed)
     result = _atpass(_globring, rin, **kwargs)
     return rin, result
 
 
 def _atpass_spawn(ring, seed, rank, rin, **kwargs):
-    """Single spawned job"""
+    """Single spawned job."""
     reset_rng(rank=rank, seed=seed)
     result = _atpass(ring, rin, **kwargs)
     return rin, result
@@ -99,23 +99,23 @@ def _lattice_pass(
     r_in,
     nturns: int = 1,
     refpts: Refpts = End,
-    no_varelem=True,
+    no_varelem: bool = True,
     seed: int | None = None,
     use_gpu: bool = False,
     **kwargs,
 ):
     kwargs["reuse"] = kwargs.pop("keep_lattice", False)
     if no_varelem:
         lattice = disable_varelem(lattice)
-    else:
-        if sum(variable_refs(lattice)) > 0:
-            kwargs["reuse"] = False
+    elif sum(variable_refs(lattice)) > 0:
+        kwargs["reuse"] = False
     refs = get_uint32_index(lattice, refpts)
     if seed is not None:
         reset_rng(seed=seed)
     if use_gpu:
         if not (iscuda() or isopencl()):
-            raise AtError("No GPU support enabled")
+            msg = "No GPU support enabled"
+            raise AtError(msg)
         else:
             return _gpupass(lattice, r_in, nturns, refpts=refs, **kwargs)
     else:
@@ -129,11 +129,10 @@ def _plattice_pass(
     nturns: int = 1,
     refpts: Refpts = End,
     seed: int | None = None,
-    pool_size: int = None,
-    start_method: str = None,
+    pool_size: int | None = None,
+    start_method: str | None = None,
     **kwargs,
 ):
-    verbose = kwargs.pop("verbose", False)
     refpts = get_uint32_index(lattice, refpts)
     any_collective = has_collective(lattice)
     kwargs["reuse"] = kwargs.pop("keep_lattice", False)
@@ -208,7 +207,8 @@ def lattice_track(
           continue the sequence
         keep_lattice (bool):    Use elements persisted from a previous
           call. If :py:obj:`True`, assume that the lattice has not changed
-          since the previous call.
+          since the previous call. Ignored if *lattice* differs from the cached one.
+          Default: :py:obj:`False`
         keep_counter (bool):    Keep the turn number from the previous
           call.
         turn (int):             Starting turn number. Ignored if
@@ -379,7 +379,7 @@ def element_track(element: Element, r_in, in_place: bool = False, **kwargs):
     """
     :py:func:`element_track` tracks particles through one element of a
     calling the element-specific tracking function specified in the
-    Element's *PassMethod* field
+    Element's *PassMethod* field.
 
     Usage:
       >>> element_track(element, r_in)