Merge pull request #317 from fnalacceleratormodeling/316-easy-way-to-…

…see-nonlinear-maps

new script print_map_o2.py to print the map of a lattice to 3rd order…

examples/nonlinear_maps/CMakeLists.txt

@@ -0,0 +1 @@
+copy_file(print_map_o3.py nonlinear_maps)

examples/nonlinear_maps/print_map_o3.py

@@ -0,0 +1,61 @@
+#!/usr/bin/env python
+import sys
+import numpy as np
+import synergia as syn
+
+# print the third order map for a lattice
+
+# the coordinate names
+cnm = ["x", "xp", "y", "yp", "cdt", "dpop" ]
+
+# given a lattice, print it's 3rd order maps
+
+def print_o3(lattice, ofile=sys.stdout):
+
+    # calculate one turn map at given order (3 here)
+    mapping = syn.simulation.Lattice_simulator.get_one_turn_map_o3(lattice)
+
+    #  iterate through the components and fields
+    for comp in range(6):
+        print(f'\nCoordinate {cnm[comp]} = ', file=ofile)
+        trigon = mapping.component(comp)
+
+        for pwr in range(trigon.power()+1):
+           for idx in range(trigon.count(pwr)): 
+               idx_to_exp = "syn.foundation.Trigon_index_to_exp_o{}(idx)".format(pwr)
+               coeff = trigon.get_term(pwr, idx)
+               if abs(coeff) < 1.0e-15:
+                   continue
+
+               print(f'{coeff} ', end='', file=ofile)
+
+               # get the exponents of the polynomial term
+               poly_exps = eval(idx_to_exp)
+
+               for c,p in enumerate(poly_exps):
+                   if p != 0:
+                       print(f'{cnm[c]}^{p} ', end='', file=ofile)
+
+               print(file=ofile)
+
+usage_txt = """
+print_map_o3.py <lattice-file> <seq-name>
+
+prints the 3rd order maps for the sequence <seq-name> within the
+lattice file given by <lattice-file>. There must be a BEAM statement
+within the file to specify the particle mass and energy.
+"""
+
+def main(argv):
+    if len(argv) < 3:
+        print(usage_txt)
+        sys.exit(10)
+
+    reader = syn.lattice.MadX_reader()
+    lattice = reader.get_lattice(argv[2], argv[1])
+    syn.simulation.Lattice_simulator.tune_circular_lattice(lattice)
+
+    print_o3(lattice)
+
+if __name__ == "__main__":
+    main(sys.argv)

src/analysis_tools/print_map_o3.py

@@ -0,0 +1,61 @@
+#!/usr/bin/env python
+import sys
+import numpy as np
+import synergia as syn
+
+# print the third order map for a lattice
+
+# the coordinate names
+cnm = ["x", "xp", "y", "yp", "cdt", "dpop" ]
+
+# given a lattice, print it's 3rd order maps
+
+def print_o3(lattice, ofile=sys.stdout):
+
+    # calculate one turn map at given order (3 here)
+    mapping = syn.simulation.Lattice_simulator.get_one_turn_map_o3(lattice)
+
+    #  iterate through the components and fields
+    for comp in range(6):
+        print(f'\nCoordinate {cnm[comp]} = ', file=ofile)
+        trigon = mapping.component(comp)
+
+        for pwr in range(trigon.power()+1):
+           for idx in range(trigon.count(pwr)): 
+               idx_to_exp = "syn.foundation.Trigon_index_to_exp_o{}(idx)".format(pwr)
+               coeff = trigon.get_term(pwr, idx)
+               if abs(coeff) < 1.0e-15:
+                   continue
+
+               print(f'{coeff} ', end='', file=ofile)
+
+               # get the exponents of the polynomial term
+               poly_exps = eval(idx_to_exp)
+
+               for c,p in enumerate(poly_exps):
+                   if p != 0:
+                       print(f'{cnm[c]}^{p} ', end='', file=ofile)
+
+               print(file=ofile)
+
+usage_txt = """
+print_map_o3.py <lattice-file> <seq-name>
+
+prints the 3rd order maps for the sequence <seq-name> within the
+lattice file given by <lattice-file>. There must be a BEAM statement
+within the file to specify the particle mass and energy.
+"""
+
+def main(argv):
+    if len(argv) < 3:
+        print(usage_txt)
+        sys.exit(10)
+
+    reader = syn.lattice.MadX_reader()
+    lattice = reader.get_lattice(argv[2], argv[1])
+    syn.simulation.Lattice_simulator.tune_circular_lattice(lattice)
+
+    print_o3(lattice)
+
+if __name__ == "__main__":
+    main(sys.argv)

src/synergia/tools/CMakeLists.txt

@@ -2,8 +2,10 @@ file(COPY ${CMAKE_CURRENT_SOURCE_DIR}/__init__.py
      DESTINATION ${CMAKE_CURRENT_BINARY_DIR})
 file(COPY ${CMAKE_CURRENT_SOURCE_DIR}/three_bump.py
      DESTINATION ${CMAKE_CURRENT_BINARY_DIR})
+file(COPY ${CMAKE_CURRENT_SOURCE_DIR}/print_map_o3.py
+     DESTINATION ${CMAKE_CURRENT_BINARY_DIR})
 
-install(FILES __init__.py three_bump.py
+install(FILES __init__.py three_bump.py print_map_o3.py
         DESTINATION ${PYTHON_INSTALL_DIR}/synergia/tools)
 
 add_subdirectory(tests)

src/synergia/tools/__init__.py

@@ -1 +1,2 @@
 from .three_bump import *
+from .print_map_o3 import *

src/synergia/tools/print_map_o3.py

@@ -0,0 +1,61 @@
+#!/usr/bin/env python
+import sys
+import numpy as np
+import synergia as syn
+
+# print the third order map for a lattice
+
+# the coordinate names
+cnm = ["x", "xp", "y", "yp", "cdt", "dpop" ]
+
+# given a lattice, print it's 3rd order maps
+
+def print_o3(lattice, ofile=sys.stdout):
+
+    # calculate one turn map at given order (3 here)
+    mapping = syn.simulation.Lattice_simulator.get_one_turn_map_o3(lattice)
+
+    #  iterate through the components and fields
+    for comp in range(6):
+        print(f'\nCoordinate {cnm[comp]} = ', file=ofile)
+        trigon = mapping.component(comp)
+
+        for pwr in range(trigon.power()+1):
+           for idx in range(trigon.count(pwr)): 
+               idx_to_exp = "syn.foundation.Trigon_index_to_exp_o{}(idx)".format(pwr)
+               coeff = trigon.get_term(pwr, idx)
+               if abs(coeff) < 1.0e-15:
+                   continue
+
+               print(f'{coeff} ', end='', file=ofile)
+
+               # get the exponents of the polynomial term
+               poly_exps = eval(idx_to_exp)
+
+               for c,p in enumerate(poly_exps):
+                   if p != 0:
+                       print(f'{cnm[c]}^{p} ', end='', file=ofile)
+
+               print(file=ofile)
+
+usage_txt = """
+print_map_o3.py <lattice-file> <seq-name>
+
+prints the 3rd order maps for the sequence <seq-name> within the
+lattice file given by <lattice-file>. There must be a BEAM statement
+within the file to specify the particle mass and energy.
+"""
+
+def main(argv):
+    if len(argv) < 3:
+        print(usage_txt)
+        sys.exit(10)
+
+    reader = syn.lattice.MadX_reader()
+    lattice = reader.get_lattice(argv[2], argv[1])
+    syn.simulation.Lattice_simulator.tune_circular_lattice(lattice)
+
+    print_o3(lattice)
+
+if __name__ == "__main__":
+    main(sys.argv)

src/synergia/tools/three_bump.py

@@ -53,17 +53,17 @@ def _construct(self):
         self.bump_lattice = synergia.lattice.Lattice("bump", element_adaptor_map)
         self.bump_lattice.set_reference_particle(self.lattice.get_reference_particle())
 
-        print "length of lattice_elements: ", len(self.lattice_elements)
+        print("length of lattice_elements: ", len(self.lattice_elements))
         elem_names = [e.get_name() for e in self.lattice_elements]
-        print "elem_names[0:9]: ", elem_names[0:9]
+        print("elem_names[0:9]: ", elem_names[0:9])
         try:
             start_idx = elem_names.index(self.start_name)
         except:
-            raise RuntimeError, "Three_bump: start_name: %s not found"%self.start_name
+            raise RuntimeError("Three_bump: start_name: %s not found"%self.start_name)
         try:
             end_idx = elem_names.index(self.end_name)
         except:
-            raise RuntimeError, "Three_bump: end_name: %s not found"%self.end_name
+            raise RuntimeError("Three_bump: end_name: %s not found"%self.end_name)
 
         if start_idx < end_idx:
             for elem in self.lattice_elements[start_idx:end_idx+1]:
@@ -80,9 +80,9 @@ def _construct(self):
         #    elem.set_string_attribute("extractor_type", "chef_propagate")
 
         if self.verbose > 2:
-            print "bump lattice:"
-            print self.bump_lattice.as_string()
-            print self.bump_idx
+            print("bump lattice:")
+            print(self.bump_lattice.as_string())
+            print(self.bump_idx)
 
         bump_elements = self.bump_lattice.get_elements()
         bump_enames = [e.get_name() for e in bump_elements]
@@ -96,7 +96,7 @@ def _construct(self):
             try:
                 hc_idx = bump_enames.index(self.hcorr_names[i])
             except:
-                raise RuntimeError, "Three_bump: hcorr_name: %s not found"%self.hcorr_names[i]
+                raise RuntimeError("Three_bump: hcorr_name: %s not found"%self.hcorr_names[i])
 
             self.hcorr_elements.append(bump_elements[hc_idx])
             self.hcorr_idx.append(self.bump_idx[hc_idx])
@@ -108,15 +108,15 @@ def _construct(self):
             try:
                 vc_idx = bump_enames.index(self.vcorr_names[i])
             except:
-                raise RuntimeError, "Three_bump: vcorr_name: %s not found"%self.vcorr_names[i]
+                raise RuntimeError("Three_bump: vcorr_name: %s not found"%self.vcorr_names[i])
 
             self.vcorr_elements.append(bump_elements[vc_idx])
             self.vcorr_idx.append(self.bump_idx[vc_idx])
 
         try:
             target_idx = bump_enames.index(self.target_name)
         except:
-            raise RuntimeError, "Three_bump: target_name: %s not found"%self.target_name
+            raise RuntimeError("Three_bump: target_name: %s not found"%self.target_name)
 
         self.target_elem = bump_elements[target_idx]
         self.target_elem.set_string_attribute("force_diagnostics", "true")
@@ -127,14 +127,14 @@ def _construct(self):
     # print out information about the bump settings
 
     def information(self):
-        print "bump_lattice: ", len(self.bump_lattice.get_elements()), " elements, length: ", self.bump_lattice.get_length()
-        print "horizontal correctors: "
+        print("bump_lattice: ", len(self.bump_lattice.get_elements()), " elements, length: ", self.bump_lattice.get_length())
+        print("horizontal correctors: ")
         for i in range(3):
-            print "\t%s, kick = %.16g"%(self.hcorr_elements[i].get_name(), self.hcorr_elements[i].get_double_attribute("kick"))
-        print "vertical correctors: "
+            print("\t%s, kick = %.16g"%(self.hcorr_elements[i].get_name(), self.hcorr_elements[i].get_double_attribute("kick")))
+        print("vertical correctors: ")
         for i in range(3):
-            print "\t%s, kick = %.16g"%(self.vcorr_elements[i].get_name(), self.vcorr_elements[i].get_double_attribute("kick"))
-        print "target element name: ", self.target_name
+            print("\t%s, kick = %.16g"%(self.vcorr_elements[i].get_name(), self.vcorr_elements[i].get_double_attribute("kick")))
+        print("target element name: ", self.target_name)
 
 
     ##################################################
@@ -145,10 +145,10 @@ def information(self):
 
     def set_corrector_elements(self, hcorr_values, vcorr_values):
         if len(self.hcorr_elements) != len(hcorr_values):
-            raise RuntimeError, "set_corrector_elements: len(hcorr_elements) != len(hcorr_values)"
+            raise RuntimeError("set_corrector_elements: len(hcorr_elements) != len(hcorr_values)")
 
         if len(self.vcorr_elements) != len(vcorr_values):
-            raise RuntimeError, "set_corrector_elements: len(vcorr_elements) != len(vcorr_values)"
+            raise RuntimeError("set_corrector_elements: len(vcorr_elements) != len(vcorr_values)")
 
         for i in range(len(hcorr_values)):
             self.hcorr_elements[i].set_double_attribute("kick", hcorr_values[i])
@@ -172,8 +172,8 @@ def propagate_zero(self):
         #stepper = synergia.simulation.Independent_stepper(self.bump_lattice, 1, 1)
         stepper = synergia.simulation.Independent_stepper_elements(self.bump_lattice, 1, 1)
         if self.verbose > 5:
-            print "bump chef beamline"
-            print synergia.lattice.chef_beamline_as_string(stepper.get_lattice_simulator().get_chef_lattice().get_sliced_beamline())
+            print("bump chef beamline")
+            print(synergia.lattice.chef_beamline_as_string(stepper.get_lattice_simulator().get_chef_lattice().get_sliced_beamline()))
         # 3 particles is the minimum so that the diagnostics don't crash
         bunch = synergia.bunch.Bunch(refpart, 3, 1.0e10, comm)
         bunch.get_local_particles()[:,0:6] = 0.0
@@ -226,7 +226,7 @@ def set_bump(self, desired_position, coords=(0, 2)):
         solver.set(tmp)
 
         if self.verbose:
-            print "  bump solver residuals:"
+            print("  bump solver residuals:")
             #print "  %5s %9s %9s %9s %9s  %9s  %9s" %("iter", "hc1", "hc2", "hc3", "vc1", "vc2", "vc3")
 
         for iter in range(100):
@@ -235,14 +235,14 @@ def set_bump(self, desired_position, coords=(0, 2)):
             x = solver.getx()
             f = solver.getf()
             if self.verbose:
-                print "  %5d % .7g % .7g % .7g % .7g  % .7g  % .7g" %(iter, f[0], f[1], f[2], f[3], f[4], f[5])
+                print("  %5d % .7g % .7g % .7g % .7g  % .7g  % .7g" %(iter, f[0], f[1], f[2], f[3], f[4], f[5]))
 
             status = multiroots.test_residual(f, 1.0e-13)
             if status == pygslerrno.GSL_SUCCESS:
-                if self.verbose: print "Converged!!"
+                if self.verbose: print("Converged!!")
                 break
         else:
-            raise ValueError, "too many iterations"
+            raise ValueError("too many iterations")
 
         # set the correctors in the original lattice
         for i in range(3):
@@ -257,15 +257,15 @@ def set_bump(self, desired_position, coords=(0, 2)):
 #  just a little tester for the class
 if __name__ == "__main__":
     lattice = synergia.lattice.MadX_reader().get_lattice("model", "tests/lattices/foborodobo128.madx")
-    print "read lattice: ", len(lattice.get_elements()), " elements, length = ", lattice.get_length()
+    print("read lattice: ", len(lattice.get_elements()), " elements, length = ", lattice.get_length())
 
     hcorr_names = ('hc1', 'hc2', 'hc3')
     vcorr_names = ('vc1', 'vc2', 'vc3')
     three_bump = Three_bump(lattice, 'm1', 'm2', hcorr_names, vcorr_names, 'm3', (0,2), True)
 
     bump_settings = three_bump.set_bump((0.001, -0.0005))
 
-    print "bump_settings: ", bump_settings[0], bump_settings[1], bump_settings[2], bump_settings[3], bump_settings[4], bump_settings[5]
+    print("bump_settings: ", bump_settings[0], bump_settings[1], bump_settings[2], bump_settings[3], bump_settings[4], bump_settings[5])
     three_bump.information()
 
     # propagate the whole lattice now
@@ -280,4 +280,4 @@ def set_bump(self, desired_position, coords=(0, 2)):
     bunch_simulator.add_per_step(synergia.bunch.Diagnostics_bulk_track("step_tracks.h5", 1))
     propagator = synergia.simulation.Propagator(stepper)
     propagator.propagate(bunch_simulator, 1, 1, 1)
-    print "final coordinates: ", np.array2string(bunch.get_local_particles()[0, 0:6])
+    print("final coordinates: ", np.array2string(bunch.get_local_particles()[0, 0:6]))