test both ol and new matching

docs/p/index.rst

@@ -44,9 +44,9 @@ Sub-packages
    :recursive:
 
    at.lattice
+   at.latticetools
    at.tracking
    at.physics
-   at.latticetools
    at.load
    at.matching
    at.acceptance

pyat/test/test_matching.py

@@ -7,103 +7,94 @@
 @pytest.fixture(scope='session')
 def test_ring(hmba_lattice):
     ring = hmba_lattice.deepcopy()
-    ring.periodicity = 1
-    sf = ring.get_uint32_index('SF*')
-    ring[sf[1]:sf[1] + 1] = ring[sf[1]].divide([0.5, 0.5])
-    ring[sf[0]:sf[0] + 1] = ring[sf[0]].divide([0.5, 0.5])
+    sf = at.get_refpts(ring, 'SF*')
+    sf1 = ring[sf[0]].divide([0.5, 0.5])
+    sf2 = ring[sf[1]].divide([0.5, 0.5])
+    ring.pop(sf[1])
+    ring.insert(sf[1], sf2[1])
+    ring.insert(sf[1], sf2[0])
+    ring.pop(sf[0])
+    ring.insert(sf[0], sf1[1])
+    ring.insert(sf[0], sf1[0])
     return ring
 
 
 def test_linopt_matching(test_ring):
-
-    def varqp(refs):
-        name = refs
-#       bnds = [0., 5.] if refs[1] == 'F' else [-5., 0]
-#       return at.latticetools.ElementVariable(ik, 'K', name=name, bounds=bnds)
-        return at.latticetools.ElementVariable(refs, 'K', name=name)
-
     # Define the location of constraint
-    sf = test_ring.get_uint32_index('SF*')[1::2]
-    center = test_ring.get_uint32_index('CellCenter')
+    sf = at.get_refpts(test_ring, 'SF*')[1::2]
+    center = at.get_refpts(test_ring, 'CellCenter')
+
     # Define the variables
     names = ['QF1*', 'QD2*', 'QD3*', 'QF4*', 'QD5*', 'QF6*', 'QF8*']
-    variables = at.VariableList([varqp(refs) for refs in names])
+    bounds = [[0, 5], [-5, 0], [-5, 0], [0, 5], [-5, 0], [0, 5], [0, 5]]
+    variables = [at.ElementVariable(at.get_refpts(test_ring, nm), 'PolynomB',
+                                    index=1, name=nm)
+                 for nm, bnd in zip(names, bounds)]
 
     # Define an evaluation function for the phase advance (both planes)
     # noinspection PyUnusedLocal
-    def phase_advance(ring, elemdata):
-        mu = elemdata.mu
-        return ((mu[-1] - mu[0]) / 2.0 / np.pi) % 1.0
+    def mu_diff(lindata, tune, chrom):
+        delta_mu = (lindata[1].mu - lindata[0].mu) / (2 * np.pi)
+        return delta_mu % 1.0
 
     # Define the constraints
-    constraints = at.ObservableList(test_ring)
-    constraints.append(at.GlobalOpticsObservable('tune', target=[0.38, 0.85]))
+    cst1 = at.LinoptConstraints(test_ring, method=at.linopt2)
+    cst1.add('tunes', [0.38, 0.85], name='tunes')
     # Start of the ring
-    constraints.append(
-        at.LocalOpticsObservable(0, 'alpha', target=[0., 0.]))
-    constraints.append(
-        at.LocalOpticsObservable(0, 'dispersion', plane='px', target=0.))
+    cst1.add('alpha', [0.0, 0.0], refpts=0, name='alpha_0')
+    cst1.add('dispersion', 0.0, refpts=0, index=1, name="eta'_0")
     # Center of the cell
-    constraints.append(
-        at.LocalOpticsObservable(center, 'alpha', target=[0., 0.]))
-    constraints.append(
-        at.LocalOpticsObservable(center, 'beta', plane='v', target=4.67))
-    constraints.append(
-        at.LocalOpticsObservable(center, 'dispersion', plane='px', target=0.))
+    cst1.add('beta', 4.69999257, refpts=center, index=1, name='betay_c')
+    cst1.add('alpha', [0.0, 0.0], refpts=center, name='alpha_c')
+    cst1.add('dispersion', 0.0, refpts=center, index=1, name="eta'_c")
     # Focusing sextupoles
-    constraints.append(
-        at.LocalOpticsObservable(sf, 'alpha', plane='v', target=[0.68, -0.68]))
-    constraints.append(
-        at.LocalOpticsObservable(sf, 'beta', plane='v', target=[5.4, 5.4]))
-    constraints.append(
-        at.LocalOpticsObservable(sf, 'dispersion', plane='x', target=0.0882))
+    cst1.add('beta', [5.40000677, 5.39998124], refpts=sf, index=1,
+             name='betay_sf')
+    cst1.add('alpha', [0.68000392, -0.67999686], refpts=sf, index=1,
+             name='alphay_sf')
+    cst1.add('dispersion', [0.08820002, 0.08819999], refpts=sf, index=0,
+             name="eta_sf")
     # Phase advance
-    constraints.append(at.LocalOpticsObservable(sf, phase_advance,
-                                                target=[0.49721338, 0.48228011],
-                                                summary=True))
+    cst1.add(mu_diff, [0.49721338, 0.48228011], refpts=sf, name='delta phi')
+
     # Perform the fit
-    newring = at.match(test_ring, variables, constraints, method='lm')
+    newring = at.match(test_ring, variables, (cst1,), method='lm', verbose=1)
 
     # check the residuals
-    constraints.evaluate(newring)
-    assert_close(constraints.sum_residuals, 0.0, rtol=0.0, atol=5e-8)
+    assert_close(cst1.evaluate(newring), 0, rtol=0.0, atol=1e-4)
 
 
 def test_envelope_matching(test_ring):
-
-    def varqp(refs):
-        name = refs
-        return at.latticetools.ElementVariable(refs, 'K', name=name)
-
     # Define the variables
     test_ring = test_ring.radiation_on(copy=True)
     names = ['QF1*', 'QD2*']
-    variables = at.VariableList([varqp(refs) for refs in names])
-    variables.append(at.ElementVariable(0, 'Voltage', name='voltage'))
+    variables = [at.ElementVariable(at.get_refpts(test_ring, nm), 'PolynomB',
+                                    index=1, name=nm) for nm in names]
+    variables += [at.ElementVariable([0], 'Voltage', name='voltage')]
+    # [0] instead of 0 because of bug in utils.py
 
     # Define the constraints
-    constraints = at.ObservableList(test_ring)
-    constraints.append(
-        at.EmittanceObservable('tunes6', target=[0.38, 0.85, 1.1e-4],
-                               weight=[1., 1., 0.0001]))
+    lopcst = at.EnvelopeConstraints(test_ring)
+    lopcst.add('tunes', [0.38, 0.85, 1.1e-4])
 
     # Perform the fit
-    newring = at.match(test_ring, variables, constraints)
+    newring = at.match(test_ring, variables, (lopcst,), verbose=1)
 
     # check the residuals
-    constraints.evaluate(newring)
-    assert_close(constraints.sum_residuals, 0, rtol=0.0, atol=1.e-10)
+    residual = lopcst.evaluate(newring.radiation_on(copy=True))
+    assert_close(residual, 0, rtol=0.0, atol=6e-9)
 
     # Define the constraints
-    constraints = at.ObservableList(test_ring)
-    constraints.append(
-        at.GlobalOpticsObservable('tune', plane=slice(2), target=[0.38, 0.85]))
-    constraints.append(
-        at.EmittanceObservable('f_s', target=1300.0, weight=1000.0))
+    lincst = at.LinoptConstraints(test_ring)
+    lincst.add('tunes', [0.38, 0.85], name='tunes')
+    lopcst = at.EnvelopeConstraints(test_ring)
+    lopcst.add('tunes', 1.1e-4, index=2, name='sync. tune')
 
     # Perform the fit
-    newring = at.match(test_ring, variables, constraints)
+    newring = at.match(test_ring, variables, (lincst, lopcst), verbose=1)
 
     # check the residuals
-    constraints.evaluate(newring)
-    assert_close(constraints.sum_residuals, 0, rtol=0.0, atol=1.e-10)
+    linresidual = lincst.evaluate(newring)
+    lopresidual = lopcst.evaluate(newring.radiation_on(copy=True))
+    assert_close(linresidual, 0, rtol=0.0, atol=6e-9)
+    assert_close(lopresidual, 0, rtol=0.0, atol=3e-8)

pyat/test/test_new_matching.py

@@ -0,0 +1,112 @@
+import at
+from at.latticetools import ElementVariable, VariableList
+from at.latticetools import GlobalOpticsObservable, LocalOpticsObservable
+from at.latticetools import EmittanceObservable, ObservableList, match
+from numpy.testing import assert_allclose as assert_close
+import pytest
+import numpy as np
+
+
+@pytest.fixture(scope='session')
+def test_ring(hmba_lattice):
+    ring = hmba_lattice.deepcopy()
+    ring.periodicity = 1
+    sf = ring.get_uint32_index('SF*')
+    ring[sf[1]:sf[1] + 1] = ring[sf[1]].divide([0.5, 0.5])
+    ring[sf[0]:sf[0] + 1] = ring[sf[0]].divide([0.5, 0.5])
+    return ring
+
+
+def test_linopt_matching(test_ring):
+
+    def varqp(refs):
+        name = refs
+#       bnds = [0., 5.] if refs[1] == 'F' else [-5., 0]
+#       return ElementVariable(ik, 'K', name=name, bounds=bnds)
+        return ElementVariable(refs, 'K', name=name)
+
+    # Define the location of constraint
+    sf = test_ring.get_uint32_index('SF*')[1::2]
+    center = test_ring.get_uint32_index('CellCenter')
+    # Define the variables
+    names = ['QF1*', 'QD2*', 'QD3*', 'QF4*', 'QD5*', 'QF6*', 'QF8*']
+    variables = VariableList([varqp(refs) for refs in names])
+
+    # Define an evaluation function for the phase advance (both planes)
+    # noinspection PyUnusedLocal
+    def phase_advance(ring, elemdata):
+        mu = elemdata.mu
+        return ((mu[-1] - mu[0]) / 2.0 / np.pi) % 1.0
+
+    # Define the constraints
+    constraints = ObservableList(test_ring)
+    constraints.append(GlobalOpticsObservable('tune', target=[0.38, 0.85]))
+    # Start of the ring
+    constraints.append(
+        LocalOpticsObservable(0, 'alpha', target=[0., 0.]))
+    constraints.append(
+        LocalOpticsObservable(0, 'dispersion', plane='px', target=0.))
+    # Center of the cell
+    constraints.append(
+        LocalOpticsObservable(center, 'alpha', target=[0., 0.]))
+    constraints.append(
+        LocalOpticsObservable(center, 'beta', plane='v', target=4.67))
+    constraints.append(
+        LocalOpticsObservable(center, 'dispersion', plane='px', target=0.))
+    # Focusing sextupoles
+    constraints.append(
+        LocalOpticsObservable(sf, 'alpha', plane='v', target=[0.68, -0.68]))
+    constraints.append(
+        LocalOpticsObservable(sf, 'beta', plane='v', target=[5.4, 5.4]))
+    constraints.append(
+        LocalOpticsObservable(sf, 'dispersion', plane='x', target=0.0882))
+    # Phase advance
+    constraints.append(LocalOpticsObservable(sf, phase_advance,
+                                                target=[0.49721338, 0.48228011],
+                                                summary=True))
+    # Perform the fit
+    newring = match(test_ring, variables, constraints, method='lm')
+
+    # check the residuals
+    constraints.evaluate(newring)
+    assert_close(constraints.sum_residuals, 0.0, rtol=0.0, atol=5e-8)
+
+
+def test_envelope_matching(test_ring):
+
+    def varqp(refs):
+        name = refs
+        return ElementVariable(refs, 'K', name=name)
+
+    # Define the variables
+    test_ring = test_ring.radiation_on(copy=True)
+    names = ['QF1*', 'QD2*']
+    variables = VariableList([varqp(refs) for refs in names])
+    variables.append(ElementVariable(0, 'Voltage', name='voltage'))
+
+    # Define the constraints
+    constraints = ObservableList(test_ring)
+    constraints.append(
+        EmittanceObservable('tunes6', target=[0.38, 0.85, 1.1e-4],
+                               weight=[1., 1., 0.0001]))
+
+    # Perform the fit
+    newring = match(test_ring, variables, constraints)
+
+    # check the residuals
+    constraints.evaluate(newring)
+    assert_close(constraints.sum_residuals, 0, rtol=0.0, atol=1.e-10)
+
+    # Define the constraints
+    constraints = ObservableList(test_ring)
+    constraints.append(
+        GlobalOpticsObservable('tune', plane=slice(2), target=[0.38, 0.85]))
+    constraints.append(
+        EmittanceObservable('f_s', target=1300.0, weight=1000.0))
+
+    # Perform the fit
+    newring = match(test_ring, variables, constraints)
+
+    # check the residuals
+    constraints.evaluate(newring)
+    assert_close(constraints.sum_residuals, 0, rtol=0.0, atol=1.e-10)