From da71ec5424998f6d8716d018c57ac4618076088e Mon Sep 17 00:00:00 2001
From: Claudio Scafuri <claudio.scafuri@elettra.eu>
Date: Fri, 12 Jul 2024 16:19:39 +0200
Subject: [PATCH] improvements, new tests

---
 boosterubung.py      | 15 ++++++++++++---
 dummy_test.py        | 40 ++++++++++++++++++++++++++++++++++++++++
 elettra2test.py      |  6 ++++--
 lattgetinfo.py       |  6 ++++--
 list_ele_elettra2.py | 16 ++++++++++------
 list_lattice.py      | 37 ++++++++++++++++++++++---------------
 spectrm.py           | 40 ++++++++++++++++++++++++++++++++++++++++
 tracke2.py           | 33 +++++++++++++++++++++++++++++++++
 tracke2b.py          | 30 ++++++++++++++++++++++++++++++
 9 files changed, 195 insertions(+), 28 deletions(-)
 create mode 100644 dummy_test.py
 create mode 100755 spectrm.py
 create mode 100755 tracke2.py
 create mode 100755 tracke2b.py

diff --git a/boosterubung.py b/boosterubung.py
index 6ba15e7..7a94283 100755
--- a/boosterubung.py
+++ b/boosterubung.py
@@ -11,6 +11,9 @@ refpts = range(len(ring) + 1)
 
 #--- lionop needs radiation off
 ring.radiation_off()
+QD=ring[1] #get firs quad for testting embbeed quad
+QD.PolynomB[0]=0.001 # H kick
+QD.PolynomA[0]=0.00  # V kick
 elemdata0, tune, chrom, elemdata=at.linopt(ring,0.0001,list(refpts))
 tuneH=tune[0]
 tuneV=tune[1]
@@ -39,7 +42,7 @@ e_eV=ring.energy
 
 
 mp.figure()
-mp.subplot(211)
+mp.subplot(311)
 mp.plot(s_pos,beta_x)
 mp.plot(s_pos,beta_y)
 mp.plot(s_pos,dispersion_x)
@@ -48,11 +51,17 @@ mp.ylabel('m')
 mp.title('beta_x  beta_y disp_x' )
 
 
-mp.subplot(212)
-mp.plot(s_pos,closed_orbit_x)
+mp.subplot(312)
+mp.plot(s_pos, closed_orbit_x)
 mp.xlabel('s [m]')
 mp.ylabel('m')
 mp.title('orb_x' )
+
+mp.subplot(313)
+mp.plot(s_pos, closed_orbit_y)
+mp.xlabel('s [m]')
+mp.ylabel('m')
+mp.title('orb_y' )
 mp.show()
 print('tune H:',tuneH,'tune V:',tuneV)
 print('E:', e_eV/1e9,' GeV')
diff --git a/dummy_test.py b/dummy_test.py
new file mode 100644
index 0000000..c09db77
--- /dev/null
+++ b/dummy_test.py
@@ -0,0 +1,40 @@
+#!/usr/bin/python3
+import at
+# test handling of embedded corrector
+
+import at
+
+import matplotlib as mp
+
+D1 = at.Drift('DR01', 3.0)
+QF = at.Quadrupole('QF', 1.0, 0.2)
+D2 = D1
+D2.FamName = 'DR02'
+D2.Length = 2
+QD = at.deepcopy(QF)
+QD.FamName = 'QD'
+QD.K = -0.4
+C = at.Corrector('CHV',0.1,[0 , 0])
+M = at.Marker('Marker')
+
+fodocell = [QF, D1, QD, C, D2, M, D2, C, QF]
+cells = [fodocell, fodocell, fodocell, fodocell, fodocell,
+           fodocell, fodocell, fodocell, fodocell, fodocell]
+ring = at.Lattice(cells,energy=1e9)
+
+#ring[3] = at.Corrector('C',0.1,[-0.00002, 0.00001])
+refpts = range(len(ring) + 1)
+ring.radiation_off()
+elemdata0, tune, chrom, elemdata=at.linopt4(ring,0.0001,list(refpts))
+s_pos = elemdata['s_pos']
+closed_orbit_x=elemdata['closed_orbit'][:,0]
+closed_orbit_xp=elemdata['closed_orbit'][:,1]
+closed_orbit_y=elemdata['closed_orbit'][:,2]
+closed_orbit_yp=elemdata['closed_orbit'][:,3]
+
+mp.figure()
+mp.plot(s_pos,closed_orbit_x,closed_orbit_y)
+mp.xlabel('s [m]')
+mp.ylabel('m')
+mp.title('orb_x' )
+mp.show()
diff --git a/elettra2test.py b/elettra2test.py
index cbd0cf2..00b3c82 100755
--- a/elettra2test.py
+++ b/elettra2test.py
@@ -2,7 +2,7 @@
 import os
 import at
 import matplotlib.pylab as mp
-lattfile = "../../machine/lattice/elettra2/srElettra2_9_4_beta.mat"
+lattfile = "~/src/gitlab/dt/machine/lattice/elettra2/srElettra2_9_4_beta.mat"
 #lattfile = "~/esrf_visit/last.mat"
 lattfile = os.path.expanduser(lattfile)
 ring=at.load_lattice(lattfile, energy=2.4e9)
@@ -56,7 +56,7 @@ mp.title('orb_x' )
 mp.show()
 print('tune H:',tuneH,'tune V:',tuneV)
 print('chrom H:',chromH,' chromV:', chromV)
-print('E:', e_eV/1e9,' GeV')
+prinpt('E:', e_eV/1e9,' GeV')
 
 # find index in ring strucutre a specific quadrupole
 
@@ -73,3 +73,5 @@ ringparams = ring.radiation_parameters()
 ringparams = ring.envelope_parameters(params=ringparams)
 print(ringparams)
 
+
+
diff --git a/lattgetinfo.py b/lattgetinfo.py
index 2467c0b..71a6e75 100755
--- a/lattgetinfo.py
+++ b/lattgetinfo.py
@@ -3,10 +3,12 @@
 import os
 import LatticeInfo
 homedir=os.path.expanduser('~')
-latticetfile = homedir + "/src/gitlab/dt/machine/lattice/elettra2/srElettra2_9_4_beta.mat"
+#lattname = "srElettra2_9_4_beta.mat"
+lattname = "srElettra2_high_betax_long_straights_ring.m"
+latticefile = homedir + "/src/gitlab/dt/machine/lattice/elettra2/" + lattname
 
 
-linfo = LatticeInfo.LatticeInfo(latticetfile, 2.4e9)
+linfo = LatticeInfo.LatticeInfo(latticefile, 2.4e9)
 print(linfo.field_names())
 linfo.printall()
 """
diff --git a/list_ele_elettra2.py b/list_ele_elettra2.py
index 9c7aae8..4696af3 100755
--- a/list_ele_elettra2.py
+++ b/list_ele_elettra2.py
@@ -2,10 +2,13 @@
 # extract index of various  elements and print them
 import at
 import re
+import os
 #import matplotlib.pylab as mp
-
-#ring=at.load_lattice("../../machine/lattice/elettra2/srElettra2_9_4_beta.mat",energy=2.4e9)
-ring=at.load_lattice("../../machine/lattice/elettra2/srElettra2_low_beta.mat",energy=2.4e9)
+homedir=os.path.expanduser('~')
+#lattname = "srElettra2_9_4_beta.mat"
+lattname = "srElettra2_high_betax_long_straights_ring.m"
+latticefile = homedir + "/src/gitlab/dt/machine/lattice/elettra2/" + lattname
+ring=at.load_lattice(latticefile, energy=2.4e9)
 refpts = range(len(ring) + 1)
 
 
@@ -17,7 +20,7 @@ for i in iquads:
     q=ring[i]
     print(q.FamName,q.Length,q.K,q.PolynomA[1],q.PolynomB[1])
     
-#find index of reverse bends Q24R quads
+#find index of reverse bends Q24AB quads
 print("Name, Length,K,PolynomA,PolynomB")
 irb = ring.uint32_refpts("QAB*")
 for i in irb:
@@ -77,7 +80,8 @@ for i in ib64:
 
 # get bpms and their s position
 
-ibpms = ring.uint32_refpts("bpm")
+ibpms = ring.uint32_refpts(at.Monitor)
+print("num bpm:", len(ibpms))
 print("ind. ,name,  s[m]")
 for i in ibpms:
 
@@ -94,7 +98,7 @@ for i in ichv:
     print(i,chv.FamName,sp[0])
 
 #get RF cavities
-selector = "RF"
+selector = "CAV"
 irfcavs = ring.uint32_refpts(selector)
 for i in irfcavs:
     rfcav=ring[i]
diff --git a/list_lattice.py b/list_lattice.py
index 9d2f2f9..acaf083 100755
--- a/list_lattice.py
+++ b/list_lattice.py
@@ -4,7 +4,7 @@
 still missing: get TYPE of magnets etc... which must be got form another source
 """
 import at
-
+import os
 
 def lattlister(ring,geo,selector):
     idlist = ring.uint32_refpts(selector)
@@ -27,9 +27,13 @@ def lattlister(ring,geo,selector):
         print(i,'\t',ename,'\t',length,'\t',spos,'\t',iscorr,'\t',isskew,'\t',x,'\t',y,'\t',angle)
         #return i,ename,length,spos,iscorr,isskew
 
-ring=at.load_lattice("../../machine/lattice/elettra2/srElettra2_9_4_beta.mat",energy=2.4e9)
+homedir=os.path.expanduser('~')
+#lattname = "srElettra2_9_4_beta.mat"
+lattname = "srElettra2_high_betax_long_straights_ring.m"
+latticefile = homedir + "/src/gitlab/dt/machine/lattice/elettra2/" + lattname
+ring=at.load_lattice(latticefile, energy=2.4e9)
 """
-up to now the low_beta lattice file can be ignred. It will have the same structure of rElettra2_9_4_beta but different
+up to now the low_beta lattice file can be ignored. It will have the same structure of rElettra2_9_4_beta but different
 magnet strenghts
 ring=at.load_lattice("../../machine/lattice/elettra2/srElettra2_low_beta.mat",energy=2.4e9)
 """
@@ -38,35 +42,38 @@ geo,radius=ring.get_geometry((0,0,0),True)
 print('i','\t','ename','\t','length','\t','spos','\t','iscorr','\t','isskew','\t','x','\t','y','\t','angle')
 
 
-# find index of quadrupoles in ring
+# find index and print of quadrupoles in ring
 lattlister(ring,geo,at.Quadrupole)
 
     
-#find index of reverse bends Q24R quads
-lattlister(ring,geo,"QAB*")
+#find index and print of reverse bends Q24R quads
+#lattlister(ring,geo,"QAB*")
     
-# find index of sextupoles in ring
+# find index and print of sextupoles in ring
 lattlister(ring,geo,at.Sextupole)
 
-#find index of octupoles in ring
+#find index and print of octupoles in ring
 lattlister(ring,geo,at.Octupole)
     
-#find index of b64 dipoles in ring
+#find index and print of b64 dipoles in ring
 #this is triky..... We have to set name to indexes ... or modify something in the lattice.
-lattlister(ring,geo,"b1")
+#lattlister(ring,geo,"b1")
     
-#find index of b80 dipoles in ring
+#find index and print of b80 dipoles in ring
 #this is enven trikier! triky..... We have to set name to indexes and elemets are a composition of pieces
 
 # TO DO
 
 # get bpms and their s position
-lattlister(ring,geo,"bpm*")
+lattlister(ring,geo,"BPM*")
+
+# get corectors  C
+#lattlister(ring,geo,"C*")
 
-# get corectors  
-lattlister(ring,geo,"chv*")
+#get correctors c
+#lattlister(ring,geo,"c*")
 
 #get RF cavities
-lattlister(ring,geo,"RF")
+lattlister(ring,geo,"CAV")
 
 
diff --git a/spectrm.py b/spectrm.py
new file mode 100755
index 0000000..7de1743
--- /dev/null
+++ b/spectrm.py
@@ -0,0 +1,40 @@
+#!/usr/bin/python3
+
+import matplotlib.pyplot as plt
+import at as pa
+# Define the particle beam
+particle_type = pa.ParticleType.proton
+particle_energy = 10e6  # MeV
+beam_sigma = 0.5e-3  # m
+
+# Create the particle beam
+proton_beam = pa.ParticleBeam(particle_type, particle_energy, 1, beam_sigma)
+# Define the magnetic elements
+dipole_length = 0.1  # m
+dipole_strength = 0.1  # Tm
+quadrupole_length = 0.05  # m
+quadrupole_strength = 1  # T/m
+
+# Create the dipole magnet
+dipole = pa.Dipole(dipole_length, dipole_strength)
+
+# Create the quadrupole magnet
+quadrupole = pa.Quadrupole(quadrupole_length, quadrupole_strength)
+# Define the number of particles and the tracking steps
+n_particles = 100
+n_steps = 10
+
+# Track the particles through the magnetic elements
+tracker = pa.Tracker(dipole, quadrupole)
+tracked_particles = tracker.track(proton_beam, n_steps)
+
+
+# Calculate the particle energies
+particle_energies = tracked_particles.x[-1] * proton_beam.energy / beam_sigma
+
+# Create a histogram of the particle energies
+plt.hist(particle_energies, bins=50, density=True)
+plt.xlabel("Particle Energy (MeV)")
+plt.ylabel("Probability Density")
+plt.show()
+
diff --git a/tracke2.py b/tracke2.py
new file mode 100755
index 0000000..2b0a386
--- /dev/null
+++ b/tracke2.py
@@ -0,0 +1,33 @@
+#!/usr/bin/python3
+import os
+import at
+import at.tracking
+import matplotlib.pylab as plt
+import numpy as np
+lattfile = "~/src/gitlab/dt/machine/lattice/elettra2/srElettra2_9_4_beta.mat"
+#lattfile = "~/esrf_visit/last.mat"
+lattfile = os.path.expanduser(lattfile)
+ring=at.load_lattice(lattfile, energy=2.4e9)
+ring.radiation_on()
+ring.rf_voltage=0
+bpms = ring.uint32_refpts("BPM*")
+print (len(bpms))
+nturns=200
+Z01 = np.array([.001, 0, 0, 0, 0, 0])
+Z02 = np.array([.00, 0.001, 0, 0, 0, 0])
+Z03 = np.array([.001, 0, 0.001, 0.0, 0.0, 0])
+Z1=at.lattice_pass(ring,Z01,nturns,refpts=bpms)
+Z2=at.lattice_pass(ring,Z02,nturns,refpts=bpms)
+print(Z03)
+Z3=at.lattice_pass(ring,Z03,nturns,refpts=bpms)
+print(Z03)
+plt.figure()
+#plt.plot(Z1[0, 0, 0, :], Z1[1, 0, 0, :],'.')
+#plt.plot(Z2[0, 0, 0, :], Z2[1, 0, 0, :],'.')
+plt.plot(Z3[0, 0, 0, :], Z3[1, 0, 0, :],'.')
+
+plt.figure()
+#plt.plot(Z3[1, 0, :, :],Z3[3, 0, :, :],'.')
+plt.plot(Z3[3, 0, :, :])
+plt.show()
+
diff --git a/tracke2b.py b/tracke2b.py
new file mode 100755
index 0000000..bd7bb0e
--- /dev/null
+++ b/tracke2b.py
@@ -0,0 +1,30 @@
+#!/usr/bin/python3
+import os
+import at
+import at.tracking
+import matplotlib.pylab as plt
+import numpy as np
+lattfile = "~/src/gitlab/dt/machine/lattice/elettra2/srElettra2_9_4_beta.mat"
+#lattfile = "~/esrf_visit/last.mat"
+lattfile = os.path.expanduser(lattfile)
+ring=at.load_lattice(lattfile, energy=2.4e9)
+ring.radiation_off()
+ring.rf_voltage=0
+
+nturns=20000
+Z01 = np.array([.001, 0, 0, 0, 0, 0])
+Z02 = np.array([.00, 0.001, 0, 0, 0, 0])
+Z03 = np.array([.001, 0, 0.001, 0.0, 0.0, 0])
+Z1=at.lattice_pass(ring,Z01,nturns)
+Z2=at.lattice_pass(ring,Z02,nturns)
+Z3=at.lattice_pass(ring,Z03,nturns)
+plt.figure()
+#plt.plot(Z1[0, 0, 0, :], Z1[1, 0, 0, :],'.')
+#plt.plot(Z2[0, 0, 0, :], Z2[1, 0, 0, :],'.')
+plt.plot(Z3[0, 0, 0, :], Z3[1, 0, 0, :],'.')
+
+plt.figure()
+plt.plot(Z3[0, 0, 0, :],Z3[3, 0, 0, :],'.')
+
+plt.show()
+
-- 
GitLab