Improvement of RF frequency control

atmat/atphysics/LongitudinalDynamics/atsetcavity.m

@@ -92,9 +92,12 @@
             if isfinite(dct)
                 frev=frev * lcell/(lcell+dct);
             elseif isfinite(dp)
-                [~,ringrad]=check_radiation(ring,false,'force');
-                etac=1/gamma0^2 - mcf(ringrad);
-                frev=frev + frev*etac*dp;
+                % Find the path lengthening for dp
+                [~,rnorad]=check_radiation(ring,false,'force');
+                [~,orbitin]=findorbit4(rnorad,dp);
+                orbitout=ringpass(rnorad,orbitin);
+                dct=orbitout(6);
+                frev=frev * lcell/(lcell+dct);
             end
             frequency = frev * props_harmnumber(harmnumber,props);
         else

pyat/at/acceptance/acceptance.py

@@ -1,21 +1,24 @@
 import numpy
 from .boundary import GridMode
+# noinspection PyProtectedMember
 from .boundary import boundary_search
 import multiprocessing
 from ..lattice import Lattice
+from ..physics import frequency_control
 
 
 __all__ = ['get_acceptance', 'get_1d_acceptance', 'get_horizontal_acceptance',
            'get_vertical_acceptance', 'get_momentum_acceptance']
 
 
+@frequency_control
 def get_acceptance(ring, planes, npoints, amplitudes, nturns=1024,
                    refpts=None, dp=None, offset=None, bounds=None,
                    grid_mode=GridMode.RADIAL, use_mp=False, verbose=True,
                    start_method=None, divider=2, shift_zero=1.0e-9):
     """
     Computes the acceptance at repfts observation points
-    Grid Coordiantes ordering is as follows: CARTESIAN: (x,y), RADIAL/RECURSIVE
+    Grid Coordinates ordering is as follows: CARTESIAN: (x,y), RADIAL/RECURSIVE
     (r, theta). Scalar inputs can be used for 1D grid.
     The grid can be changed using grid_mode input:
     at.GridMode.CARTESIAN: (x,y) grid
@@ -192,7 +195,7 @@ def get_1d_acceptance(ring, plane, resolution, amplitude, nturns=1024, dp=None,
                              nturns=nturns, dp=dp, refpts=refpts,
                              grid_mode=grid_mode, use_mp=use_mp,
                              verbose=verbose, start_method=start_method,
-                             divider=2, shift_zero=0.0)
+                             divider=divider, shift_zero=0.0)
     return numpy.squeeze(b), s, g
 
 

pyat/at/acceptance/boundary.py

@@ -4,12 +4,10 @@
 grid definitions
 """
 
-import at
 from at.lattice import AtError
 from at.tracking import lattice_pass, patpass
 from enum import Enum
 import numpy
-import warnings
 from scipy.ndimage import binary_dilation, binary_opening
 from collections import namedtuple
 import time
@@ -79,20 +77,14 @@ def set_ring_orbit(ring, dp, obspt, orbit):
     Returns a ring starting at obspt and initial
     closed orbit
     """
-    if ring.radiation:
-        newring = ring.set_rf_frequency(dp=dp, copy=True)
-        dp = None
-    else:
-        newring = ring
-
     if obspt is not None:
         assert numpy.isscalar(obspt), 'Scalar value needed for obspt'
-        newring = newring.rotate(obspt)
+        ring = ring.rotate(obspt)
 
     if orbit is None:
-        orbit, _ = newring.find_orbit(dp=dp)
+        orbit = ring.find_orbit(dp=dp)[0]
 
-    return orbit, newring
+    return orbit, ring
 
 
 def grid_configuration(planes, npoints, amplitudes, grid_mode, bounds=None,

pyat/at/physics/revolution.py

@@ -4,9 +4,40 @@
 from ..tracking import lattice_pass
 from .orbit import find_orbit4
 import numpy
+import functools
 
-__all__ = ['get_mcf', 'get_slip_factor', 'get_revolution_frequency',
-           'set_rf_frequency']
+__all__ = ['frequency_control', 'get_mcf', 'get_slip_factor',
+           'get_revolution_frequency', 'set_rf_frequency']
+
+
+def frequency_control(func):
+    """Function to be used as decorator for ``func(ring, *args, **kwargs)``
+
+    If ``ring.radiation`` is ``True`` and ``dp``, ``dct`` or ``df`` is
+    specified, make a copy of ``ring`` with a modified RF frequency, remove
+    ``dp``, ``dct`` or ``df`` from ``kwargs`` and call ``func`` with the
+    modified ``ring``.
+
+    If ``ring.radiation`` is ``False``, ``func`` is called unchanged.
+
+    Examples::
+
+        @frequency_control
+        def func(ring, *args, dp=None, dct=None, **kwargs):
+            pass
+    """
+    @functools.wraps(func)
+    def wrapper(ring, *args, **kwargs):
+        if ring.radiation:
+            momargs = {}
+            for key in ['dp', 'dct', 'df']:
+                v = kwargs.pop(key, None)
+                if v is not None:
+                    momargs[key] = v
+            if len(momargs) > 0:
+                ring = set_rf_frequency(ring, **momargs, copy=True)
+        return func(ring, *args, **kwargs)
+    return wrapper
 
 
 @check_radiation(False)
@@ -50,7 +81,7 @@ def get_slip_factor(ring, **kwargs):
     return etac
 
 
-def get_revolution_frequency(ring, dp=None, dct=None, **kwargs):
+def get_revolution_frequency(ring, dp=None, dct=None):
     """Compute the revolution frequency of the full ring [Hz]
 
     PARAMETERS
@@ -59,18 +90,17 @@ def get_revolution_frequency(ring, dp=None, dct=None, **kwargs):
     KEYWORDS
         dp=0.0          momentum deviation.
         dct=0.0         Path length deviation
-        keep_lattice    Assume no lattice change since the previous tracking.
-                        Defaults to False
-        dp_step=1.0E-6  momentum deviation used for differentiation
     """
     lcell = ring.get_s_pos(len(ring))[0]
     frev = ring.beta * clight / lcell / ring.periodicity
     if dct is not None:
         frev *= lcell / (lcell + dct)
     elif dp is not None:
+        # Find the path lengthening for dp
         rnorad = ring.radiation_off(copy=True) if ring.radiation else ring
-        etac = get_slip_factor(rnorad, **kwargs)
-        frev += frev * etac * dp
+        orbit = lattice_pass(rnorad, rnorad.find_orbit4(dp=dp)[0])
+        dct = numpy.squeeze(orbit)[5]
+        frev *= lcell / (lcell + dct)
     return frev