Repository: atcollab/at (language: python)

Issue:
Momentum offset for 6d orbit
This PR is proposal to answer #528 and other similar problems related to 6D optics/orbit calculations with input momentum offset.

The `frequency_control` decorator is moved to `orbit.py` and applied on `find_orbit6` allowing to input `dp`, `dct` or `df ` and shift the ring frequency internally, This is done n a copy of `ring` and does not modify the input lattice.

Drawback: the `dp`, `dct` or `df ` offset are applied w.r.t to the **NOMINAL** lattice and therefore overwrite existing frequency shifts, this is safer as it prevents applying these offsets several times.
The help was modified to explain this behavior.

The behavior of `find_orbit6` with `dp` argument is now identical to the one of `find_orbit4` with the same input arguments, which is convenient for codes having to handle lattices without or without radiations.

I know this has been a controversial question and we may decide to discard this proposal, however the present situation raising errors all the time is not satisfactory and needs to be improved. Any comments or ideas welcome!


Discussion / hints:
Hello @swhite2401

Moving the decorator to `orbit.py` is a good thing. This will allow to decorate almost anything. But decorating `find_orbit6` does not solve everything: we still need to decorate the high-level function (those involving matrix computations), because tracking around the closed orbit to get the matrix needs the modified ring with the right RF frequency. So we'll be piling up multiple checks. This can be avoided by keeping a non-decorated `find_orbit6` for internal use and a decorated one for user access.

Or alternatively keep find_orbit6 non-decorated (with a warning or error, as you like), but modify `find_orbit` to handle dp, dct and df. We just need to ensure that internally we use the specific functions rather that the general one (anyway it would not be wrong, just adding useless controls).

Then we would decorate `linopt6`, not `get_tune` (it uses `linopt6`), but `get_chrom`, to have the right central frequency…

Or go for the brute force solution: decorate everything. The penalty does not look so tremendous. Much less than the one for using a 6d lattice when it's not useful. We are just leaving out the pedagogic aspect: one should not use a 6d lattice to compute the chromaticity…

The changes that have been made have fixed the chromaticity function.

```
import numpy as np
import at
import matplotlib.pyplot as plt

def func(dp, fit):
    val = np.zeros_like(dp)
    for o in np.arange(len(fit)):
        val += fit[o] * dp**o * np.math.factorial(o)
    return val

ring = at.load_mat('/machfs/carver/Repo/EBSlattices/AT/S28F.mat', key='LOW_EMIT_RING_INJ')

(fitx, fity), dpa, qz = at.chromaticity(ring, dpm = 0.02, npoints=31, order=3)
```

![image](https://user-images.githubusercontent.com/16048082/206707695-0d7f66ed-9cb8-4a29-9c92-4ac850320acd.png)

 > The changes that have been made have fixed the chromaticity function.

No, they don't: `get_chrom` does not crash any more, but `dp` is still ignored for 6d lattices:

![Figure_2](https://user-images.githubusercontent.com/20038121/206895872-bfe04c9e-1eb1-4b3e-987b-cffd2854f4c1.png)

(Example with the `hmba.mat` test lattice)
So we have to choose between 3 possibilities:
1. Decorate everything (`find_orbit` **and** higher level functions) which may lead to  >= 3 times the 6d check. Example `linopt6` -> `find_m66` -> `find_orbit6`,
2. leave things as they are and fix `get_chrom`,
3. decorate `find_orbit6` but keep a non-decorated version for internal use, thus avoiding duplicated tests. Note that this is equivalent to point 2, with some renaming and  adding a decorated `find_orbit` for users.

Any preference ?
Hello, @lfarv is correct , I have overlooked `findm66`.
For what concerns the fix, it is complicated am I cannot convince myself on the best option.

However, if we are allowing `dp, dct, df ` arguments for 6d functions (such as `find_orbit6` or `linopt6`) we have to make them useful otherwise it makes no sense. This is why I would be tempted to say that all 6d functions allowing for these arguments should be decorated which is option 1... it is true that this leads to many useless checks but I don't think this would impact performance.


> This is why I would be tempted to say that all 6d functions allowing for these arguments should be decorated which is option 1

Not exactly: in all 3 options, the idea is to decorate all high level 6d functions (`linopt6`...). Options 2 and 3 simply try to avoid duplicate checks by keeping private non-decorated low-level functions for internal use. If it's limited to keep an undecorated `find_orbit6`, that's rather easy, but for longer call chains like `get_something()` -> `linopt6()` -> `find_m66()` -> `find_orbit6()`, that may get tricky.

So I finally agree that option 1 is the easiest solution… 
So now, if we continue with this branch, we **must** decorate:
- `find_m66` (otherwise it's wrong in 6d with dp specified),
- `get_chrom` to have dp taken into account (see above),
- `get_tune`, which may call `line_pass`, so needs the right RF frequency. Decorating the nested `get_centroid` would work, but would modify the lattice twice, once for `linopt6` and then for `get_centroid`

It looks useless for `linopt6`, `get_optics` and derivatives (all they do is processing the output of `find_m66`, they don't care about the modified ring), 
`get_chrom` and `get_w` acitivated in `linopt6` will give inconsistent results because they shift the frequency around `ring.rf_frequency`, also `find_orbit6` would need to be decorated I believe.

Do we expect others? In `radiation.py` for example?

I can see 2 options:
- we pursue with this branch as discussed
- we remove the arguments that cannot be used and decorate with check_6d(True)

Since the second option is what we add in the past and was not popular I would suggest to continue with this branch.
Is it possible to decorate function in matlab?
> also find_orbit6 would need to be decorated I believe.

Yes, you are right!

> Is it possible to decorate function in matlab?

The decorator syntax does not exist, but there is a wrapper which does the equivalent: see `wrapper6d.m`. But the problem is worse in Matlab, because checking 6d needs a scan of the whole lattice (no lattice object to store attributes). So avoiding unnecessary checks is important, and we'll probably have to keep private raw versions… Or finally store the radiation state in the RingParam element.
> Do we expect others? In radiation.py for example?

Everything is correctly decorated with `check_6d(True)` except `get_radiation_integrals`. `get_radiation_integrals` does not care of 6d true or false, but we already noticed that off-momentum computation misses the quadrupole contribution, so removing the `dp` input makes sense.

@lfarv, can I leave this to you? (I can do the python if needed but the matlab I am not comfortable with...)

Fix patch:
diff --git a/atmat/atphysics/LinearOptics/atlinopt6.m b/atmat/atphysics/LinearOptics/atlinopt6.m
index 407330098..689089b90 100644
--- a/atmat/atphysics/LinearOptics/atlinopt6.m
+++ b/atmat/atphysics/LinearOptics/atlinopt6.m
@@ -80,9 +80,9 @@
 %
 %  See also atlinopt2 atlinopt4 tunechrom
 
-[ringdata,elemdata]=wrapper6d(ring,@xlinopt6,varargin{:});
+[ringdata,elemdata]=frequency_control(@xlinopt6,ring,varargin{:});
 
-    function [ringdata,elemdata] = xlinopt6(ring, is6d, varargin)
+    function [ringdata,elemdata] = xlinopt6(ring, varargin)
         clight = PhysConstant.speed_of_light_in_vacuum.value;   % m/s
         [get_chrom, varargs]=getflag(varargin, 'get_chrom');
         [get_w, varargs]=getflag(varargs, 'get_w');
@@ -91,6 +91,7 @@
         [DPStep,~]=getoption(varargs,'DPStep');
         [cavargs,varargs]=getoption(varargs,{'cavpts'});
         [refpts,varargs]=getargs(varargs,1,'check',@(arg) isnumeric(arg) || islogical(arg));
+        is_6d=getoption(varargs,'is_6d',[]); % Always set by frequency_control, keep in varargs
 
         if isempty(twiss_in)        % Circular machine
             [orbs,orbitin]=findorbit(ring,refpts,dpargs{:},varargs{:});
@@ -120,15 +121,15 @@
             )';
         ringdata=struct('tune',tunes,'damping_time',damping_times);
 
-        if is6d                     % 6D processing
+        if is_6d                     % 6D processing
             [alpha,beta,disp]=cellfun(@output6,ri,'UniformOutput',false);
             if get_w || get_chrom
                 frf=get_rf_frequency(ring);
                 DFStep=-DPStep*mcf(atradoff(ring))*frf;
                 rgup=atsetcavity(ring,'Frequency',frf+0.5*DFStep,cavargs{:});
                 rgdn=atsetcavity(ring,'Frequency',frf-0.5*DFStep,cavargs{:});
-                [~,o1P]=findorbit(rgup,'guess',orbitin,varargs{:});
-                [~,o1M]=findorbit(rgdn,'guess',orbitin,varargs{:});
+                [~,o1P]=findorbit6(rgup,'guess',orbitin,varargs{:});
+                [~,o1M]=findorbit6(rgdn,'guess',orbitin,varargs{:});
                 if get_w
                     [ringdata.chromaticity,w]=chrom_w(rgup,rgdn,o1P,o1M,refpts);
                     [elemdata.W]=deal(w{:});
@@ -180,7 +181,7 @@
 
         function [vals,ms,phis,rmats,as]=build_1turn_map(ring,refpts,orbit,varargin)
             % Build the initial distribution at entrance of the transfer line
-            if is6d
+            if is_6d
                 [mt,ms]=findm66(ring,refpts,'orbit',orbit,varargs{:});
             else
                 [mt,ms]=findm44(ring,NaN,refpts,'orbit',orbit,varargs{:});
diff --git a/atmat/atphysics/LinearOptics/findm44.m b/atmat/atphysics/LinearOptics/findm44.m
index 56175ba10..ad5abbc56 100644
--- a/atmat/atphysics/LinearOptics/findm44.m
+++ b/atmat/atphysics/LinearOptics/findm44.m
@@ -78,6 +78,7 @@
 varargs=getdparg(varargs);
 [dp,varargs]=getoption(varargs,'dp',0.0);
 [dpargs,varargs]=getoption(varargs,{'dct','df'});
+[~,varargs]=getoption(varargs,'is_6d',[]); % Consume the is_6d option
 [refpts,orbitin,varargs]=getargs(varargs,[],orbitin,'check',@(x) ~(ischar(x) || isstring(x))); %#ok<ASGLU>
 
 if islogical(refpts)
diff --git a/atmat/atphysics/LinearOptics/findm66.m b/atmat/atphysics/LinearOptics/findm66.m
index cd3377f83..a0e187de0 100644
--- a/atmat/atphysics/LinearOptics/findm66.m
+++ b/atmat/atphysics/LinearOptics/findm66.m
@@ -1,4 +1,4 @@
-function [M66, varargout] = findm66(LATTICE, varargin)
+function varargout = findm66(LATTICE, varargin)
 %FINDM66 numerically finds the 6x6 transfer matrix of an accelerator lattice
 %  by differentiation of LINEPASS near the closed orbit
 %  FINDM66 uses FINDORBIT6 to search for the closed orbit in 6-d
@@ -7,6 +7,18 @@
 % M66 = FINDM66(RING) finds the full one-turn 6-by-6
 %    matrix at the entrance of the first element
 %
+%[...]=FINDM66(RING,...,'dp',DP) Specify the momentum deviation when
+%   radiation is OFF (default: 0)
+%
+%[...]=FINDM66(RING,...,'dct',DCT) Specify the path lengthening when
+%   radiation is OFF (default: 0)
+%
+%[...]=FINDM66(RING,...,'df',DF) Specify the RF frequency deviation
+%   radiation is OFF (default: 0)
+%
+%[...]=FINDM66(RING,...,'orbit',ORBIT) Specify the orbit at the entrance
+%   of the ring, if known.
+%
 % [M66,T] = FINDM66(RING,REFPTS) in addition to M finds
 %    6-by-6 transfer matrixes  between entrances of
 %    the first element and each element indexed by REFPTS.
@@ -37,52 +49,54 @@
 if ~iscell(LATTICE)
     error('First argument must be a cell array');
 end
-NE = length(LATTICE);
-[XYStep,varargs]=getoption(varargin,'XYStep');	% Step size for numerical differentiation	%1.e-8
-[DPStep,varargs]=getoption(varargs,'DPStep');	% Step size for numerical differentiation	%1.e-6
-[orbitin,varargs]=getoption(varargs,'orbit',[]);
-[refpts,orbitin,varargs]=getargs(varargs,[],orbitin,'check',@(x) ~(ischar(x) || isstring(x))); %#ok<ASGLU>
+[varargout{1:nargout}] = frequency_control(@xfindm66,LATTICE,varargin{:});
 
-if islogical(refpts)
-    refpts(end+1:NE+1)=false;
-elseif isnumeric(refpts)
-    refpts=setelems(false(1,NE+1),refpts);
-else
-    error('REFPTS must be numeric or logical');
-end
+    function [M66, varargout] = xfindm66(LATTICE, varargin)
+        NE = length(LATTICE);
+        [XYStep,varargs]=getoption(varargin,'XYStep');	% Step size for numerical differentiation	%1.e-8
+        [DPStep,varargs]=getoption(varargs,'DPStep');	% Step size for numerical differentiation	%1.e-6
+        [orbitin,varargs]=getoption(varargs,'orbit',[]);
+        [dpargs,varargs]=getoption(varargs,{'dp','dct','df'});
+        [is_6d,varargs]=getoption(varargs,'is_6d',[]); % Always set by frequency_control
+        [refpts,orbitin,varargs]=getargs(varargs,[],orbitin,'check',@(x) ~(ischar(x) || isstring(x))); %#ok<ASGLU>
 
-if isempty(orbitin)
-    if check_radiation(LATTICE)
-        orbitin = findorbit6(LATTICE,'XYStep',XYStep,'DPStep',DPStep);
-    else
-        [~, orbitin] = findorbit4(LATTICE,0.0,'XYStep',XYStep);
-    end
-end
+        if islogical(refpts)
+            refpts(end+1:NE+1)=false;
+        elseif isnumeric(refpts)
+            refpts=setelems(false(1,NE+1),refpts);
+        else
+            error('REFPTS must be numeric or logical');
+        end
 
-refs=setelems(refpts,NE+1);
-reqs=refpts(refs);
+        if isempty(orbitin)
+            [~, orbitin] = findorbit(LATTICE,'XYStep',XYStep,'DPStep',DPStep,'is_6d',is_6d,dpargs{:});
+        end
 
-% Build a diagonal matrix of initial conditions
-%scaling=2*XYStep*[1 0.1 1 0.1 1 1];
-scaling=XYStep*[1 1 1 1 0 0] + DPStep*[0 0 0 0 1 1];
-D6 = 0.5*diag(scaling);
-% Add to the orbit_in. First 12 columns for derivative
-% 13-th column is for closed orbit
-RIN = orbitin + [D6 -D6 zeros(6,1)];
-ROUT = linepass(LATTICE,RIN,refs);
-TMAT3 = reshape(ROUT,6,13,[]);
-M66 = (TMAT3(:,1:6,end)-TMAT3(:,7:12,end))./scaling;
+        refs=setelems(refpts,NE+1);
+        reqs=refpts(refs);
 
-if nargout >= 2 % Calculate matrices at all REFPTS.
-    varargout{1} = (TMAT3(:,1:6,reqs)-TMAT3(:,7:12,reqs))./scaling;
-    % Return closed orbit if requested
-    if nargout >= 3
-        varargout{2}=squeeze(TMAT3(:,13,reqs));
-    end
-end
+        % Build a diagonal matrix of initial conditions
+        %scaling=2*XYStep*[1 0.1 1 0.1 1 1];
+        scaling=XYStep*[1 1 1 1 0 0] + DPStep*[0 0 0 0 1 1];
+        D6 = 0.5*diag(scaling);
+        % Add to the orbit_in. First 12 columns for derivative
+        % 13-th column is for closed orbit
+        RIN = orbitin + [D6 -D6 zeros(6,1)];
+        ROUT = linepass(LATTICE,RIN,refs);
+        TMAT3 = reshape(ROUT,6,13,[]);
+        M66 = (TMAT3(:,1:6,end)-TMAT3(:,7:12,end))./scaling;
 
-    function mask=setelems(mask,idx)
-        mask(idx)=true;
-    end
+        if nargout >= 2 % Calculate matrices at all REFPTS.
+            varargout{1} = (TMAT3(:,1:6,reqs)-TMAT3(:,7:12,reqs))./scaling;
+            % Return closed orbit if requested
+            if nargout >= 3
+                varargout{2}=squeeze(TMAT3(:,13,reqs));
+            end
+        end
+
+        function mask=setelems(mask,idx)
+            mask(idx)=true;
+        end
 
+    end
 end
diff --git a/atmat/atphysics/Orbit/findorbit.m b/atmat/atphysics/Orbit/findorbit.m
index 70d06949b..5c77ecaaa 100644
--- a/atmat/atphysics/Orbit/findorbit.m
+++ b/atmat/atphysics/Orbit/findorbit.m
@@ -30,39 +30,41 @@
 %
 % See also FINDORBIT4, FINDSYNCORBIT, FINDORBIT6.
 
-[orbitin,varargs]=getoption(varargin,'orbit',[]);
-[refpts,varargs]=getargs(varargs,[],'check',@(arg) isnumeric(arg) || islogical(arg));
-[dp,varargs]=getoption(varargs,'dp',NaN);
-[dct,varargs]=getoption(varargs,'dct',NaN);
-[df,varargs]=getoption(varargs,'df',NaN);
-if isempty(orbitin)
-    if check_6d(ring)    % Radiation ON: 6D orbit
-        if isfinite(dp) || isfinite(dct) || isfinite(df)
-            warning('AT:linopt','In 6D, "dp", "dct" and "df" are ignored');
+[orbs,orbitin]=frequency_control(@xfindorbit,ring,varargin{:});
+
+    function [orbs,orb0]=xfindorbit(ring,varargin)
+        [orb0,varargs]=getoption(varargin,'orbit',[]);
+        [refpts,varargs]=getargs(varargs,[],'check',@(arg) isnumeric(arg) || islogical(arg));
+        [dp,varargs]=getoption(varargs,'dp',NaN);
+        [dct,varargs]=getoption(varargs,'dct',NaN);
+        [df,varargs]=getoption(varargs,'df',NaN);
+        [is_6d,varargs]=getoption(varargs,'is_6d',[]); % Always set by frequency_control
+        if isempty(orb0)
+            if is_6d   % Radiation ON: 6D orbit
+                orb0=xorbit_6(ring,varargs{:});
+            elseif isfinite(df)
+                [cell_l,cell_frev,cell_h]=atGetRingProperties(ring,'cell_length',...
+                    'cell_harmnumber','cell_revolution_frequency');
+                dct=-cell_l*df/(cell_frev*cell_h+df);
+                orb0=xorbit_ct(ring,dct,varargs{:});
+            elseif isfinite(dct)
+                orb0=xorbit_ct(ring,dct,varargs{:});
+            else
+                orb0=xorbit_dp(ring,dp,varargs{:});
+            end
+            args={'KeepLattice'};
+        else
+            args={};
         end
-        orbitin=xorbit_6(ring,varargs{:});
-    elseif isfinite(df)
-        [cell_l,cell_frev,cell_h]=atGetRingProperties(ring,'cell_length',...
-            'cell_harmnumber','cell_revolution_frequency');
-        dct=-cell_l*df/(cell_frev*cell_h+df);
-        orbitin=xorbit_ct(ring,dct,varargs{:});
-    elseif isfinite(dct)
-        orbitin=xorbit_ct(ring,dct,varargs{:});
-    else
-        orbitin=xorbit_dp(ring,dp,varargs{:});
-    end
-    args={'KeepLattice'};
-else
-    args={};
-end
 
-if islogical(refpts)
-    refpts=find(refpts);
-end
-if isempty(refpts)
-    % return only the fixed point at the entrance of RING{1}
-    orbs=orbitin;
-else
-    orbs=linepass(ring,orbitin,refpts,args{:});
-end
+        if islogical(refpts)
+            refpts=find(refpts);
+        end
+        if isempty(refpts)
+            % return only the fixed point at the entrance of RING{1}
+            orbs=orb0;
+        else
+            orbs=linepass(ring,orb0,refpts,args{:});
+        end
+    end
 end
diff --git a/atmat/atphysics/Orbit/findorbit4.m b/atmat/atphysics/Orbit/findorbit4.m
index ad2cbc67b..135481fde 100644
--- a/atmat/atphysics/Orbit/findorbit4.m
+++ b/atmat/atphysics/Orbit/findorbit4.m
@@ -61,6 +61,7 @@
 [dct,varargs]=getoption(varargs,'dct',NaN);
 [df,varargs]=getoption(varargs,'df',NaN);
 [refpts,varargs]=getargs(varargs,[],'check',@(arg) isnumeric(arg) || islogical(arg));
+[~,varargs]=getoption(varargs,'is_6d',[]); % Consume the is_6d option
 if isempty(orbitin)
     if isfinite(df)
         [cell_l,cell_frev,cell_h]=atGetRingProperties(ring,'cell_length',...
diff --git a/atmat/atphysics/Orbit/findorbit6.m b/atmat/atphysics/Orbit/findorbit6.m
index 19c1ece23..575887c55 100644
--- a/atmat/atphysics/Orbit/findorbit6.m
+++ b/atmat/atphysics/Orbit/findorbit6.m
@@ -38,6 +38,17 @@
 %   from the range 1 to length(RING)+1.
 %   See further explanation of REFPTS in the 'help' for FINDSPOS
 %
+% FINDORBIT6(...,'dp',DP)
+%   Specify the off-momentum. The RF frequency will be adjusted to get the
+%   desired value
+%
+% FINDORBIT6(...,'dct',DCT)
+%   Specify the path lengthening. The RF frequency will be adjusted to get
+%   the desired value
+%
+% FINDORBIT6(...,'df',DF)
+%   Specify the RF frequency deviation
+%
 % FINDORBIT6(RING,REFPTS,GUESS)
 % FINDORBIT6(...,'guess',GUESS)     The search for the fixed point
 %	starts from initial condition GUESS. Otherwise the search starts from
@@ -56,22 +67,27 @@
 if ~iscell(ring)
     error('First argument must be a cell array');
 end
-[orbitin,varargs]=getoption(varargin,'orbit',[]);
-[refpts,varargs]=getargs(varargs,[],'check',@(arg) isnumeric(arg) || islogical(arg));
-if isempty(orbitin)
-    orbitin=xorbit_6(ring,varargs{:});
-    args={'KeepLattice'};
-else
-    args={};
-end
+[orb6,orbitin] = frequency_control(@xfindorbit6,ring,varargin{:});
 
-if islogical(refpts)
-    refpts=find(refpts);
-end
-if isempty(refpts)
-    % return only the fixed point at the entrance of RING{1}
-    orb6=orbitin;
-else
-    orb6 = linepass(ring,orbitin,refpts,args{:});
-end
+    function[orb6,orbitin] = xfindorbit6(ring,varargin)
+        [orbitin,varargs]=getoption(varargin,'orbit',[]);
+        [refpts,varargs]=getargs(varargs,[],'check',@(arg) isnumeric(arg) || islogical(arg));
+        [~,varargs]=getoption(varargs,'is_6d',[]); %% Consume the is_6d option
+        if isempty(orbitin)
+            orbitin=xorbit_6(ring,varargs{:});
+            args={'KeepLattice'};
+        else
+            args={};
+        end
+
+        if islogical(refpts)
+            refpts=find(refpts);
+        end
+        if isempty(refpts)
+            % return only the fixed point at the entrance of RING{1}
+            orb6=orbitin;
+        else
+            orb6 = linepass(ring,orbitin,refpts,args{:});
+        end
+    end
 end
diff --git a/atmat/atphysics/Orbit/findsyncorbit.m b/atmat/atphysics/Orbit/findsyncorbit.m
index f4686f410..aa374ce33 100644
--- a/atmat/atphysics/Orbit/findsyncorbit.m
+++ b/atmat/atphysics/Orbit/findsyncorbit.m
@@ -60,6 +60,7 @@
 [dp,varargs]=getoption(varargs,'dp',NaN);
 [df,varargs]=getoption(varargs,'df',NaN);
 [refpts,varargs]=getargs(varargs,[],'check',@(arg) isnumeric(arg) || islogical(arg));
+[~,varargs]=getoption(varargs,'is_6d',[]); % Consume the is_6d option
 if isempty(orbitin)
     if isfinite(df)
         [cell_l,cell_frev,cell_h]=atGetRingProperties(ring,'cell_length',...
diff --git a/atmat/atphysics/TuneAndChromaticity/tunechrom.m b/atmat/atphysics/TuneAndChromaticity/tunechrom.m
index e2bf5e516..5a8d267cd 100644
--- a/atmat/atphysics/TuneAndChromaticity/tunechrom.m
+++ b/atmat/atphysics/TuneAndChromaticity/tunechrom.m
@@ -40,15 +40,16 @@
 if cpl1 || ~cpl2
     warning('AT:ObsoleteParameter','The "coupled" flag is ignored: coupling is always assumed');
 end
-[varargout{1:nargout}]=wrapper6d(ring,@xtunechrom,allargs{:});
+[varargout{1:nargout}]=frequency_control(@xtunechrom,ring,allargs{:});
 
-    function [tune, chrom] = xtunechrom(ring,is6d,varargin)
+    function [tune, chrom] = xtunechrom(ring,varargin)
         [oldchrom,varargs]=getflag(varargin,'chrom');
         [chrom,varargs]=getflag(varargs,'get_chrom');
         [dpargs,varargs]=getoption(varargs,{'orbit','dp','dct','df'});
         [cavargs,varargs]=getoption(varargs,{'cavpts'});
         [DPStep,~]=getoption(varargs,'DPStep');
-        if is6d
+        is_6d=getoption(varargs,'is_6d',[]); % Always set by frequency_control, keep in varargs
+        if is_6d
             tunefunc=@tune6;
         else
             tunefunc=@tune4;
@@ -59,13 +60,13 @@
         tune=tunefunc(ring,'orbit',orbitin,varargs{:});
 
         if chrom || oldchrom || nargout == 2
-            if is6d
+            if is_6d
                 frf=get_rf_frequency(ring);
                 DFStep=-DPStep*mcf(atradoff(ring))*frf;
                 rgup=atsetcavity(ring,'Frequency',frf+0.5*DFStep,cavargs{:});
                 rgdn=atsetcavity(ring,'Frequency',frf-0.5*DFStep,cavargs{:});
-                [~,o1P]=findorbit(rgup,'guess',orbitin,varargs{:});
-                [~,o1M]=findorbit(rgdn,'guess',orbitin,varargs{:});
+                [~,o1P]=findorbit6(rgup,'guess',orbitin,varargs{:});
+                [~,o1M]=findorbit6(rgdn,'guess',orbitin,varargs{:});
                 deltap=o1P(5)-o1M(5);
             else
                 dp=orbitin(5);
diff --git a/atmat/atutils/frequency_control.m b/atmat/atutils/frequency_control.m
new file mode 100644
index 000000000..da498a1cf
--- /dev/null
+++ b/atmat/atutils/frequency_control.m
@@ -0,0 +1,33 @@
+function varargout = frequency_control(func,ring,varargin)
+%FREQUENCY_CONTROL  Private. Handle off-momentum for 6D lattice
+%
+%   VARARGOUT=FREQUENCY_CONTROL(FUNC,RING,VARARGIN)
+%
+%FUNC   Wrapped function, called as FUNC(RING,VARARGIN{:},'is6d',IS6D)
+
+[warningdp6d,varargs]=getoption(varargin,'WarningDp6D');
+[is_6d, varargs]= getoption(varargs, 'is_6d', []);
+if isempty(is_6d), is_6d=atGetRingProperties(ring,'is_6d'); end
+if is_6d
+    [dpargs,varargs]=getoption(varargs,{'dp','dct','df'});
+    if length(dpargs) > 2
+        error('AT:OffMomentum', ['Off-momentum specification: ',...
+            'only one keyword in "dp", "dct" or "df" may be specified.'])
+    end
+    if ~isempty(dpargs)
+        if warningdp6d
+            warning('AT:Dp6D','\n%s\n%s\n%s',...
+                'Specifying "dp, dct or df" for a 6D lattice creates a copy with a modified RF frequency.',...
+                'For a better efficiency, handle the RF frequency beforehand,',...
+                'or to avoid this warning, use "setoption(''WarningDp6D'',false)"');
+        end
+        [cavargs,~]=getoption(varargs,{'cavpts'});
+        ring2=atsetcavity(ring,'Frequency','nominal',dpargs{:},cavargs{:});
+        [varargout{1:nargout}]=func(ring2,varargs{:},'is_6d',is_6d);
+    else
+        [varargout{1:nargout}]=func(ring,varargs{:},'is_6d',is_6d);
+    end
+else
+    [varargout{1:nargout}]=func(ring,varargs{:},'is_6d',is_6d);
+end
+end
\ No newline at end of file


Test patch:
diff --git a/atmat/attests/pytests.m b/atmat/attests/pytests.m
index 52460d388..647a2dd66 100644
--- a/atmat/attests/pytests.m
+++ b/atmat/attests/pytests.m
@@ -7,23 +7,23 @@
             "machine_data/spear3.m"];
     end
     
-    properties
-        ring4
-        ring6
-    end
-    
     properties(TestParameter)
         dp = {0., -0.01, 0.01};
-        dct = {0., -0.00005, 0.00005};
         lat = struct("hmba", "hmba","dba","dba","spear3","spear3");
         rad = struct("radoff","ring4","radon","ring6");
         lat2 = struct("hmba", "hmba","spear3","spear3");
     end
+    
+    properties
+        ring4
+        ring6
+    end
 
     methods(TestClassSetup)
         function load_lattice(testCase)
             % Shared setup for the entire test class
             t=warning('off','AT:atradon:NOCavity');
+            setoption('WarningDp6D',false);
             for fpath=testCase.mlist
                 [~,fname,~]=fileparts(fpath);
                 [testCase.ring4.(fname),testCase.ring6.(fname)]=mload(fpath);
@@ -35,8 +35,8 @@ function load_lattice(testCase)
                 pr=atwritepy(mr,'keep_all',true);
                 ring4.m=mr;
                 ring4.p=pr;
-                ring6.m=atradon(mr);
-                ring6.p=pr.radiation_on(pyargs('copy',true));
+                ring6.m=atenable_6d(mr);
+                ring6.p=pr.enable_6d(pyargs('copy',true));
             end
         end
     end
@@ -45,8 +45,18 @@ function load_lattice(testCase)
         % Setup for each test
     end
 
+    methods(Static)
+        function [dct,df]=dctdf(r4, dpp)
+            [frf,l]=atGetRingProperties(r4,'rf_frequency','cell_length');
+            [~,o0]=findorbit4(r4,dp=dpp);
+            o1=ringpass(r4, o0);
+            dct=o1(6);
+            df=-frf*dct/(l+dct);
+        end
+    end
+
     methods(Test, TestTags="GitHub")
-        % These tests may rub in GitHub actions
+        % These tests may run in GitHub actions
 
         function lattice_pass(testCase,lat,rad)
             % Test ob basic tracking
@@ -71,8 +81,9 @@ function orbit4(testCase,lat,dp)
             testCase.verifyEqual(morbit4,porbit4,AbsTol=1.E-15);
         end
 
-        function syncorbit(testCase,lat,dct)
+        function syncorbit(testCase,lat,dp)
             lattice=testCase.ring4.(lat);
+            [dct,~]=testCase.dctdf(lattice.m, dp);
             % python
             a=cell(lattice.p.find_sync_orbit(dct));
             [psyncorb,~]=deal(a{:});
@@ -82,15 +93,15 @@ function syncorbit(testCase,lat,dct)
             testCase.verifyEqual(msyncorb,psyncorb,AbsTol=1.E-15);
         end
 
-        function orbit6(testCase,lat2)
+        function orbit6(testCase,lat2,dp)
             lattice=testCase.ring6.(lat2);
             % python
-            a=cell(lattice.p.find_orbit6());
+            a=cell(lattice.p.find_orbit6(pyargs(dp=dp)));
             [porbit6,~]=deal(a{:});
             porbit6=double(porbit6)';
             % Matlab
-            [~,morbit6]=findorbit6(lattice.m);
-            testCase.verifyEqual(morbit6,porbit6,AbsTol=5.E-15 );
+            [~,morbit6]=findorbit6(lattice.m,dp=dp);
+            testCase.verifyEqual(morbit6,porbit6,AbsTol=2.E-12);
         end
 
         function m44(testCase,lat2,dp)
@@ -114,6 +125,34 @@ function m66(testCase,lat2)
             pm66=double(pm66);
             testCase.verifyEqual(mm66,pm66,AbsTol=5.E-9);
         end
+
+        function offmomdp(testCase, dp)
+            % Checks that off-momentum is correctly taken into account
+            % for 6D lattices
+            [ring1, elem1]=atlinopt6(testCase.ring4.hmba.m,'get_chrom', dp=dp);
+            [ring2, elem2]=atlinopt6(testCase.ring6.hmba.m,'get_chrom', dp=dp);
+            testCase.verifyEqual(ring1.tune, ring2.tune(1:2), AbsTol=1e-6);
+            testCase.verifyEqual(ring1.chromaticity, ring2.chromaticity(1:2), AbsTol=5.e-5);
+            testCase.verifyEqual(elem1.beta, elem2.beta, AbsTol=1.e-4)
+        end
+
+        function offmomdct(testCase, dp)
+            [dct,~]=testCase.dctdf(testCase.ring4.hmba.m, dp);
+            [ring1, elem1]=atlinopt6(testCase.ring4.hmba.m,'get_chrom', dct=dct);
+            [ring2, elem2]=atlinopt6(testCase.ring6.hmba.m,'get_chrom', dct=dct);
+            testCase.verifyEqual(ring1.tune, ring2.tune(1:2), AbsTol=1e-6);
+            testCase.verifyEqual(ring1.chromaticity, ring2.chromaticity(1:2), AbsTol=5.e-5);
+            testCase.verifyEqual(elem1.beta, elem2.beta, AbsTol=1.e-4)
+        end
+
+        function offmomdf(testCase, dp)
+            [~,df]=testCase.dctdf(testCase.ring4.hmba.m, dp);
+            [ring1, elem1]=atlinopt6(testCase.ring4.hmba.m,'get_chrom', df=df);
+            [ring2, elem2]=atlinopt6(testCase.ring6.hmba.m,'get_chrom', df=df);
+            testCase.verifyEqual(ring1.tune, ring2.tune(1:2), AbsTol=1e-6);
+            testCase.verifyEqual(ring1.chromaticity, ring2.chromaticity(1:2), AbsTol=5.e-5);
+            testCase.verifyEqual(elem1.beta, elem2.beta, AbsTol=1.e-4)
+        end
     end
 
     methods(Test)