Also scaling families in tune and chromaticity fits (#783)

add scaling

pyat/at/matching/globalfit.py

@@ -4,14 +4,28 @@
 """
 from typing import Optional
 import numpy
-from ..lattice import Lattice, Refpts, set_value_refpts
+from ..lattice import Lattice, Refpts
+from ..lattice import get_value_refpts, set_value_refpts
 from ..lattice import AtWarning, AtError
 from ..physics import get_tune, get_chrom
 
 
 __all__ = ['fit_tune', 'fit_chrom']
 
 
+def _set_magnets(ring, refpts, attname, delta, index=None,
+                 increment=True, regex=False, scaling=False):
+    if scaling:
+        val0 = get_value_refpts(ring, refpts, attname, index=index,
+                                regex=regex)
+        set_value_refpts(ring, refpts, attname, val0*(1+delta/numpy.mean(val0)),
+                         index=index, regex=regex)                        
+    else:
+        set_value_refpts(ring, refpts, attname, delta, index=index,
+                         increment=True, regex=regex)
+
+
+
 def _get_tune(ring: Lattice, dp: float, **kwargs):
     get_integer = kwargs.pop('fit_integer', False)
     return get_tune(ring, dp=dp, get_integer=get_integer)[0:2]
@@ -26,29 +40,29 @@ def _fit_tune_chrom(ring: Lattice, index: int, func,
                     tol: Optional[float] = 1.0e-12,
                     dp: Optional[float] = 0, niter: Optional[int] = 3,
                     delta: Optional[float] = None,
-                    regex=False, **kwargs):
+                    regex=False, scaling=False, **kwargs):
 
     def _get_resp(ring: Lattice, index: int, func, refpts, attname,
-                  delta, dp, regex=False, **kwargs):
-        set_value_refpts(ring, refpts, attname, delta, index=index,
-                         increment=True, regex=regex)
+                  delta, dp, regex=False, scaling=False, **kwargs):
+        _set_magnets(ring, refpts, attname, delta, index=index,
+                     increment=True, regex=regex, scaling=scaling)
         datap = func(ring, dp, **kwargs)
-        set_value_refpts(ring, refpts, attname, -2*delta, index=index,
-                         increment=True, regex=regex)
+        _set_magnets(ring, refpts, attname, -2*delta, index=index,
+                     increment=True, regex=regex, scaling=scaling)
         datan = func(ring, dp, **kwargs)
-        set_value_refpts(ring, refpts, attname, delta, index=index,
-                         increment=True, regex=regex)
+        _set_magnets(ring, refpts, attname, delta, index=index,
+                     increment=True, regex=regex, scaling=scaling)
         data = numpy.subtract(datap, datan)/(2*delta)
         return data
 
     def _fit(ring, index, func, refpts1, refpts2, newval, J,
              dp: Optional[float] = 0, regex=False, **kwargs):
         val = func(ring, dp, **kwargs)
         dk = numpy.linalg.solve(J, numpy.subtract(newval, val))
-        set_value_refpts(ring, refpts1, 'PolynomB', dk[0], index=index,
-                         increment=True, regex=regex)
-        set_value_refpts(ring, refpts2, 'PolynomB', dk[1], index=index,
-                         increment=True, regex=regex)
+        _set_magnets(ring, refpts1, 'PolynomB', dk[0], index=index,
+                     increment=True, regex=regex, scaling=scaling)
+        _set_magnets(ring, refpts2, 'PolynomB', dk[1], index=index,
+                     increment=True, regex=regex, scaling=scaling)
         val = func(ring, dp, **kwargs)
         sumsq = numpy.sum(numpy.square(numpy.subtract(val, newval)))
         return sumsq
@@ -96,6 +110,8 @@ def fit_tune(ring: Lattice, refpts1: Refpts, refpts2: Refpts, newval,
         regex:      Using regular expressions for refpt string matching;
                     Default: False
         KStep: gradient variation applied to magnets. Default 1e-5
+        scaling: Scales the families instead of incrementing.
+                 Default False
 
     Typical usage:
     at.fit_tune(ring, refpts1, refpts2, [0.1,0.25])
@@ -127,6 +143,8 @@ def fit_chrom(ring: Lattice, refpts1: Refpts, refpts2: Refpts, newval,
         regex:      Using regular expressions for refpt string matching;
                     Default: False
         HStep: gradient variation applied to magnets. Default 1e-4
+        scaling: Scales the families instead of incrementing them.
+                 Default False
 
     Typical usage:
     at.fit_chrom(ring, refpts1, refpts2, [10,5])