diff --git a/pyat/at/matching/globalfit.py b/pyat/at/matching/globalfit.py index 8ec1f1bf7..3154b47a3 100644 --- a/pyat/at/matching/globalfit.py +++ b/pyat/at/matching/globalfit.py @@ -4,7 +4,8 @@ """ 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 @@ -12,6 +13,19 @@ __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,18 +40,18 @@ 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 @@ -45,10 +59,10 @@ 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])