Atavedata (#678)

* atavedata modified to avoid large errors when the lattice has errors(polynom(2) and displacement errors in sextupoles, rotated quadrupoles). Also dipole edge focusing effects have been included.

* dipersion averages also included, don't use wiedemann formulas!

* function name corrected

* small bug corrected

* debuging average optics comparison between matlab and python

* implementing Farvacque comment on repeated ring selection

* PEP8

* adding roll and closed orbit effect on quad strengths

* eps() instead of 1e-12

* numpy eps instead of 0.0

atmat/atphysics/atavedata.m

@@ -13,6 +13,7 @@
 %[LINDATA,AVEBETA,AVEMU,AVEDISP,TUNES,CHROMS]=ATAVEDATA(RING,DPP,REFPTS,ORBITIN)
 %    does not search for closed orbit. instead ORBITIN is used
 
+
 lr=length(ring)+1;
 if islogical(refpts)
     refs=[refpts(:);false(lr-length(refpts),1)];
@@ -28,6 +29,8 @@
 avebeta=cat(1,lindata.beta); %refpts
 avemu=cat(1,lindata.mu); %refpts
 avedisp=cat(2,lindata.Dispersion)'; %refpts
+ClosedOrbit=cat(2,lindata.ClosedOrbit)'; %refpts
+
 
 if any(long)
     initial=[long(needed(1:end-1));false]; %needed
@@ -40,44 +43,75 @@
     alpha0=cat(1,lind(initial).alpha); %long
     mu0=avemu(lg,:); %long
     disp0=avedisp(lg,:); %long
+    ClosedOrbit0=ClosedOrbit(lg,:); %long
 
-    beta1=cat(1,lind(final).beta); %long
-    alpha1=cat(1,lind(final).alpha); %long
     mu1=cat(1,lind(final).mu); %long
     disp1=cat(2,lind(final).Dispersion)'; %long
+    ClosedOrbit1=cat(2,lind(final).ClosedOrbit)'; %long
 
     L2=[L L]; %long
-    avebeta(lg,:)=betadrift(beta0,beta1,alpha0,L2);
+    avebeta(lg,:)=betadrift(beta0,alpha0,L2);
     avemu(lg,:)=0.5*(mu0+mu1);
     avedisp(lg,[1 3])=(disp1(:,[1 3])+disp0(:,[1 3]))*0.5;
-
-    foc=atgetcells(ring(long),'PolynomB',@(el,polb) length(polb)>=2 && polb(2)~=0); %long
+    avedisp(lg,[2 4])=(disp1(:,[1 3])-disp0(:,[1 3]))./L2;
+    foc=atgetcells(ring(long),'PolynomB',@(el,polb) length(polb)>=2 && polb(2)~=0||polb(3)~=0); %long
     if any(foc)
-        qp=false(size(long));
-        qp(long)=foc;
-        K=zeros(size(L)); %long
-        K(foc)=atgetfieldvalues(ring(qp),'PolynomB',{2});
+        qp=false(size(lg));
+        qp(lg)=foc;
+
+        %Extract element parameters
+        reng_selection=ring(refpts(qp));
+        L=atgetfieldvalues(reng_selection,'Length','Default',0);
+        q=eps()+atgetfieldvalues(reng_selection,'PolynomB',{2},'Default',0);
+        m=atgetfieldvalues(reng_selection,'PolynomB',{3},'Default',0);
+        R11=atgetfieldvalues(reng_selection,'R2',{1,1},'Default',1);
+        dx=(atgetfieldvalues(reng_selection,'T2',{1},'Default',0)-atgetfieldvalues(reng_selection,'T1',{1},'Default',0))/2;
+        ba=atgetfieldvalues(reng_selection,'BendingAngle','Default',0);
+        irho=ba./L;
+        e1=atgetfieldvalues(reng_selection,'EntranceAngle','Default',0);
+        Fint=atgetfieldvalues(reng_selection,'Fint','Default',0);
+        gap=atgetfieldvalues(reng_selection,'gap','Default',0);
+
+        %Hard edge model on dipoles
+        d_csi=ba.*gap.*Fint.*(1+sin(e1).^2)./cos(e1)./L;
+        Cp=[ba.*tan(e1)./L -ba.*tan(e1-d_csi)./L];
+        for ii=1:2
+            alpha0(foc,ii)=alpha0(foc,ii)-beta0(foc,ii).*Cp(:,ii);
+            disp0(foc,2+(ii-1)*2)=disp0(foc,2+(ii-1)*2)-disp0(foc,1+(ii-1)*2).*Cp(:,ii);
+        end
+
+        % Additional components
+        dx0=(ClosedOrbit0(foc,1)+ClosedOrbit1(foc,1))/2;
+        K=q.*R11+2*m.*(dx0-dx);
         K2=[K -K]; %long
-        sel=false(size(avebeta,1)); %refpts
-        sel(lg)=foc;
-        avebeta(sel,:)=betafoc(beta1(foc,:),alpha0(foc,:),alpha1(foc,:),K2(foc,:),L2(foc,:));
-        avedisp(sel,[1 3])=dispfoc(disp0(foc,[2 4]),disp1(foc,[2 4]),K2(foc,:),L2(foc,:));
+        K2(:,1)=K2(:,1)+irho.^2;
+        irho2=[irho 0*irho];
+
+        % Apply formulas
+        avebeta(qp,:)=betafoc(beta0(foc,:),alpha0(foc,:),K2,L2(foc,:));
+        avedisp(qp,:)=dispfoc(disp0(foc,:),irho2,K2,L2(foc,:));
     end
-    avedisp(lg,[2 4])=(disp1(:,[1 3])-disp0(:,[1 3]))./L2;
+
 end
 
-    function avebeta=betadrift(beta0,beta1,alpha0,L)
+    function avebeta=betadrift(beta0,alpha0,L)
         gamma0=(alpha0.*alpha0+1)./beta0;
-        avebeta=0.5*(beta0+beta1)-gamma0.*L.*L/6;
+        avebeta=beta0-alpha0.*L+gamma0.*L.*L/3;
     end
 
-    function avebeta=betafoc(beta1,alpha0,alpha1,K,L)
-        gamma1=(alpha1.*alpha1+1)./beta1;
-        avebeta=0.5*((gamma1+K.*beta1).*L+alpha1-alpha0)./K./L;
+    function avebeta=betafoc(beta0,alpha0,K,L)
+        gamma0=(alpha0.*alpha0+1)./beta0;
+        avebeta=((beta0+gamma0./K).*L+(beta0-gamma0./K).*sin(2.0*sqrt(K).*L)./sqrt(K)/2.0+...
+            (cos(2.0*sqrt(K).*L)-1.0).*alpha0./K)./L/2.0;
     end
 
-    function avedisp=dispfoc(dispp0,dispp1,K,L)
-        avedisp=(dispp0-dispp1)./K./L;
+    function avedisp=dispfoc(disp0,ir,K,L)
+        avedisp=disp0;
+        avedisp(:,[1 3])=(disp0(:,[1 3]).*(sin(sqrt(K).*L)./sqrt(K))+...
+            disp0(:,[2 4]).*(1-cos(sqrt(K).*L))./K)./L+...
+            ir.*(L-sin(sqrt(K).*L)./sqrt(K))./K./L;
+        avedisp(:,[2 4])=(disp0(:,[2 3]).*(sin(sqrt(K).*L)./sqrt(K))-...
+            disp0(:,[1 3]).*(1-cos(sqrt(K).*L)))./L+...
+            ir.*(L-cos(sqrt(K).*L))./K./L;
     end
-
 end

pyat/at/physics/linear.py

@@ -1070,7 +1070,7 @@ def linopt(ring: Lattice, dp: float = 0.0, refpts: Refpts = None,
 
 # noinspection PyPep8Naming
 @check_6d(False)
-def avlinopt(ring: Lattice, dp: float = None, refpts: Refpts = None, **kwargs):
+def avlinopt(ring, dp, refpts, **kwargs):
     r"""Linear analysis of a lattice with average values
 
     :py:func:`avlinopt` returns average beta, mu, dispersion over the lattice
@@ -1124,10 +1124,10 @@ def avlinopt(ring: Lattice, dp: float = None, refpts: Refpts = None, **kwargs):
     Returns:
         elemdata:   Linear optics at the points refered to by *refpts*,
           if refpts is :py:obj:`None` an empty lindata structure is returned.
-        avebeta:    Average beta functions [:math:`\hat{\beta_x},\hat{\beta_y}`]
-          at *refpts*
-        avemu:      Average phase advances [:math:`\hat{\mu_x},\hat{\mu_y}`]
-          at *refpts*
+        avebeta:    Average beta functions
+          [:math:`\hat{\beta_x},\hat{\beta_y}`] at *refpts*
+        avemu:      Average phase advances
+          [:math:`\hat{\mu_x},\hat{\mu_y}`] at *refpts*
         avedisp:    Average dispersion [:math:`\hat{\eta_x}, \hat{\eta'_x},
           \hat{\eta_y}, \hat{\eta'_y}`] at *refpts*
         avespos:    Average s position at *refpts*
@@ -1141,31 +1141,116 @@ def avlinopt(ring: Lattice, dp: float = None, refpts: Refpts = None, **kwargs):
     def get_strength(elem):
         try:
             k = elem.PolynomB[1]
+        except (AttributeError, IndexError):
+            k = numpy.finfo(numpy.float64).eps
+        return k
+
+    def get_sext_strength(elem):
+        try:
+            k = elem.PolynomB[2]
+        except (AttributeError, IndexError):
+            k = 0.0
+        return k
+
+    def get_roll(elem):
+        try:
+            k = elem.R2[0][0]
+        except (AttributeError, IndexError):
+            k = 1.0
+        return k
+
+    def get_dx(elem):
+        try:
+            k = (elem.T2[0]-elem.T1[0])/2.0
+        except (AttributeError, IndexError):
+            k = 0.0
+        return k
+
+    def get_bendingangle(elem):
+        try:
+            k = elem.BendingAngle
         except (AttributeError, IndexError):
             k = 0.0
         return k
 
-    def betadrift(beta0, beta1, alpha0, lg):
-        gamma0 = (alpha0 * alpha0 + 1) / beta0
-        return 0.5 * (beta0 + beta1) - gamma0 * lg * lg / 6
+    def get_e1(elem):
+        try:
+            k = elem.EntranceAngle
+        except (AttributeError, IndexError):
+            k = 0.0
+        return k
 
-    def betafoc(beta1, alpha0, alpha1, k2, lg):
-        gamma1 = (alpha1 * alpha1 + 1) / beta1
-        return 0.5 * ((gamma1 + k2 * beta1) * lg + alpha1 - alpha0) / k2 / lg
+    def get_fint(elem):
+        try:
+            k = elem.FringeInt1
+        except (AttributeError, IndexError):
+            k = 0.0
+        return k
 
-    def dispfoc(dispp0, dispp1, k2, lg):
-        return (dispp0 - dispp1) / k2 / lg
+    def get_gap(elem):
+        try:
+            k = elem.FullGap
+        except (AttributeError, IndexError):
+            k = 0.0
+        return k
 
+    def sini(x, L):
+        r = x.copy()
+        r[x > 0] = numpy.sin(numpy.sqrt(x[x > 0])*L[x > 0]
+                             )/numpy.sqrt(x[x > 0])
+        r[x < 0] = numpy.sinh(numpy.sqrt(-x[x < 0])*L[x < 0]
+                              )/numpy.sqrt(-x[x < 0])
+        return r
+
+    def cosi(x, L):
+        r = x.copy()
+        r[x > 0] = numpy.cos(numpy.sqrt(x[x > 0])*L[x > 0])
+        r[x < 0] = numpy.cosh(numpy.sqrt(-x[x < 0])*L[x < 0])
+        return r
+
+    def betadrift(beta0, alpha0, lg):
+        gamma0 = (alpha0 * alpha0 + 1.0) / beta0
+        return beta0-alpha0*lg+gamma0*lg*lg/3
+
+    def betafoc(beta0, alpha0, kk, lg):
+        gamma0 = (alpha0 * alpha0 + 1.0) / beta0
+        return ((beta0+gamma0/kk)*lg +
+                (beta0-gamma0/kk)*sini(kk, 2.0*lg)/2.0 +
+                (cosi(kk, 2.0*lg)-1.0)*alpha0/kk)/lg/2.0
+
+    def dispfoc(disp0, ir, k2, lg):
+        avedisp = disp0.copy()
+        avedisp[:, 0::2] = (disp0[:, 0::2]*(sini(k2, lg)) +
+                            disp0[:, 1::2]*(1.0-cosi(k2, lg))/k2 +
+                            ir*(lg-sini(k2, lg))/k2)/lg
+        avedisp[:, 1::2] = (disp0[:, 0::2]*(sini(k2, lg)) -
+                            disp0[:, 0::2]*(1.0-cosi(k2, lg)) +
+                            ir*(lg-cosi(k2, lg))/k2)/lg
+        return avedisp
+
+    # selected list
     boolrefs = get_bool_index(ring, refpts)
-    length = numpy.array([el.Length for el in ring[boolrefs]])
-    strength = numpy.array([get_strength(el) for el in ring[boolrefs]])
+    ring_selected = ring[refpts]
+    L = numpy.array([el.Length for el in ring_selected])
+    K = numpy.array([get_strength(el) for el in ring_selected])
+    sext_strength = numpy.array([get_sext_strength(el)
+                                 for el in ring_selected])
+    roll = numpy.array([get_roll(el) for el in ring_selected])
+    ba = numpy.array([get_bendingangle(el) for el in ring_selected])
+    e1 = numpy.array([get_e1(el) for el in ring_selected])
+    Fint = numpy.array([get_fint(el) for el in ring_selected])
+    gap = numpy.array([get_gap(el) for el in ring_selected])
+    dx = numpy.array([get_dx(el) for el in ring_selected])
+    irho = ba.copy()
+    d_csi = ba.copy()
+
+    # whole ring list
     longelem = get_bool_index(ring, None)
-    longelem[boolrefs] = (length != 0)
+    longelem[boolrefs] = (L != 0)
 
     shorti_refpts = (~longelem) & boolrefs
     longi_refpts = longelem & boolrefs
     longf_refpts = numpy.roll(longi_refpts, 1)
-
     all_refs = shorti_refpts | longi_refpts | longf_refpts
     _, bd, d_all = linopt4(ring, refpts=all_refs, dp=dp,
                            get_chrom=True, **kwargs)
@@ -1175,32 +1260,42 @@ def dispfoc(dispp0, dispp1, k2, lg):
     avemu = lindata.mu.copy()
     avedisp = lindata.dispersion.copy()
     aves = lindata.s_pos.copy()
+    ClosedOrbit = lindata.closed_orbit.copy()
 
     di = d_all[longi_refpts[all_refs]]
     df = d_all[longf_refpts[all_refs]]
 
-    long = (length != 0.0)
-    kfoc = (strength != 0.0)
-    foc = long & kfoc
-    nofoc = long & (~kfoc)
-    K2 = numpy.stack((strength[foc], -strength[foc]), axis=1)
-    fff = foc[long]
-    length = length.reshape((-1, 1))
-
-    avemu[long] = 0.5 * (di.mu + df.mu)
-    aves[long] = 0.5 * (df.s_pos + di.s_pos)
-    avebeta[nofoc] = \
-        betadrift(di.beta[~fff], df.beta[~fff], di.alpha[~fff], length[nofoc])
-    avebeta[foc] = \
-        betafoc(df.beta[fff], di.alpha[fff], df.alpha[fff], K2, length[foc])
-    avedisp[numpy.ix_(long, [1, 3])] = \
-        (df.dispersion[:, [0, 2]] - di.dispersion[:, [0, 2]]) / length[long]
+    b_long = (L != 0.0)
+    b_foc = (K != 0.0)
+    b_foc_long = b_long & b_foc
+    b_drf = b_long & (~b_foc)
+
+    dx0 = ClosedOrbit[:, 0]
+    dx0[b_long] = (di.closed_orbit[:, 0]+df.closed_orbit[:, 0])/2
+    K = K*roll + 2*sext_strength*(dx0-dx)
+
+    K2 = numpy.stack((K[b_foc_long], -K[b_foc_long]), axis=1)
+    fff = b_foc_long[b_long]
+    irho[b_long] = ba[b_long]/L[b_long]
+    d_csi[b_long] = ba[b_long]*(gap[b_long]*Fint[b_long] *
+                                (1.0 + numpy.sin(e1[b_long])**2) /
+                                numpy.cos(e1[b_long])/L[b_long])
+    L2 = numpy.stack((L[b_foc_long], L[b_foc_long]), axis=1)
+    irho = irho.reshape((-1, 1))
+    L = L.reshape((-1, 1))
+
+    avemu[b_long] = 0.5 * (di.mu + df.mu)
+    aves[b_long] = 0.5 * (df.s_pos + di.s_pos)
+    avebeta[b_drf] = betadrift(di.beta[~fff], di.alpha[~fff], L[b_drf])
+    avebeta[b_foc_long] = betafoc(di.beta[fff], di.alpha[fff], K2, L2)
+    avedisp[numpy.ix_(b_long, [1, 3])] = (df.dispersion[:, [0, 2]] -
+                                          di.dispersion[:, [0, 2]]) / L[b_long]
     idx = numpy.ix_(~fff, [0, 2])
-    avedisp[numpy.ix_(nofoc, [0, 2])] = (di.dispersion[idx] +
+    avedisp[numpy.ix_(b_drf, [0, 2])] = (di.dispersion[idx] +
                                          df.dispersion[idx]) * 0.5
-    idx = numpy.ix_(fff, [1, 3])
-    avedisp[numpy.ix_(foc, [0, 2])] = \
-        dispfoc(di.dispersion[idx], df.dispersion[idx], K2, length[foc])
+    avedisp[b_foc_long, :] = dispfoc(di.dispersion[fff, :],
+                                     irho[b_foc_long], K2, L2)
+
     return lindata, avebeta, avemu, avedisp, aves, bd.tune, bd.chromaticity
 
 
@@ -1238,7 +1333,7 @@ def get_tune(ring: Lattice, *, method: str = 'linopt',
         fmin (float):           Lower tune bound. Default: 0
         fmax (float):           Upper tune bound. Default: 1
         hann (bool):            Turn on Hanning window.
-          Default: :py:obj:`False`. Work only for ``'fft'`` 
+          Default: :py:obj:`False`. Work only for ``'fft'``
         get_integer(bool):   Turn on integer tune (slower)
 
     Returns:
@@ -1259,7 +1354,7 @@ def gen_centroid(ring, ampl, nturns, remove_dc, ld):
         return numpy.conjugate(p2.T.view(dtype=complex).T)
     get_integer = kwargs.pop('get_integer', False)
     if get_integer:
-        assert method == 'linopt',\
+        assert method == 'linopt', \
            'Integer tune only accessible with method=linopt'
     if method == 'linopt':
         if get_integer: