-
Notifications
You must be signed in to change notification settings - Fork 2
/
Copy pathUTILITY_quickstart.py
1153 lines (857 loc) · 41.8 KB
/
UTILITY_quickstart.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
433
434
435
436
437
438
439
440
441
442
443
444
445
446
447
448
449
450
451
452
453
454
455
456
457
458
459
460
461
462
463
464
465
466
467
468
469
470
471
472
473
474
475
476
477
478
479
480
481
482
483
484
485
486
487
488
489
490
491
492
493
494
495
496
497
498
499
500
501
502
503
504
505
506
507
508
509
510
511
512
513
514
515
516
517
518
519
520
521
522
523
524
525
526
527
528
529
530
531
532
533
534
535
536
537
538
539
540
541
542
543
544
545
546
547
548
549
550
551
552
553
554
555
556
557
558
559
560
561
562
563
564
565
566
567
568
569
570
571
572
573
574
575
576
577
578
579
580
581
582
583
584
585
586
587
588
589
590
591
592
593
594
595
596
597
598
599
600
601
602
603
604
605
606
607
608
609
610
611
612
613
614
615
616
617
618
619
620
621
622
623
624
625
626
627
628
629
630
631
632
633
634
635
636
637
638
639
640
641
642
643
644
645
646
647
648
649
650
651
652
653
654
655
656
657
658
659
660
661
662
663
664
665
666
667
668
669
670
671
672
673
674
675
676
677
678
679
680
681
682
683
684
685
686
687
688
689
690
691
692
693
694
695
696
697
698
699
700
701
702
703
704
705
706
707
708
709
710
711
712
713
714
715
716
717
718
719
720
721
722
723
724
725
726
727
728
729
730
731
732
733
734
735
736
737
738
739
740
741
742
743
744
745
746
747
748
749
750
751
752
753
754
755
756
757
758
759
760
761
762
763
764
765
766
767
768
769
770
771
772
773
774
775
776
777
778
779
780
781
782
783
784
785
786
787
788
789
790
791
792
793
794
795
796
797
798
799
800
801
802
803
804
805
806
807
808
809
810
811
812
813
814
815
816
817
818
819
820
821
822
823
824
825
826
827
828
829
830
831
832
833
834
835
836
837
838
839
840
841
842
843
844
845
846
847
848
849
850
851
852
853
854
855
856
857
858
859
860
861
862
863
864
865
866
867
868
869
870
871
872
873
874
875
876
877
878
879
880
881
882
883
884
885
886
887
888
889
890
891
892
893
894
895
896
897
898
899
900
901
902
903
904
905
906
907
908
909
910
911
912
913
914
915
916
917
918
919
920
921
922
923
924
925
926
927
928
929
930
931
932
933
934
935
936
937
938
939
940
941
942
943
944
945
946
947
948
949
950
951
952
953
954
955
956
957
958
959
960
961
962
963
964
965
966
967
968
969
970
971
972
973
974
975
976
977
978
979
980
981
982
983
984
985
986
987
988
989
990
991
992
993
994
995
996
997
998
999
1000
from pytao import Tao
import numpy as np
import matplotlib.pyplot as plt
import matplotlib.cm as cm
import matplotlib.patches as mpatches
import re
import io
from os import path,environ
import pandas as pd
import random
from IPython.display import display, clear_output, update_display
import bayes_opt
import shutil
from scipy.special import jn as besselj
import scipy
from scipy.optimize import curve_fit
from pmd_beamphysics import ParticleGroup
#from pmd_beamphysics.statistics import resample_particles
import pmd_beamphysics.statistics
from UTILITY_plotMod import plotMod, slicePlotMod, floorplanPlot
from UTILITY_linacPhaseAndAmplitude import getLinacMatchStrings, setLinacPhase, setLinacGradientAuto
from UTILITY_modifyAndSaveInputBeam import modifyAndSaveInputBeam
from UTILITY_setLattice import setLattice, getBendkG, getQuadkG, getSextkG, setBendkG, setQuadkG, setSextkG, setXOffset, setYOffset, setKickerkG
from UTILITY_impact import runImpact
from UTILITY_OpenPMDtoBmad import OpenPMD_to_Bmad
from UTILITY_finalFocusSolver import finalFocusSolver
import os
import yaml
filePathGlobal = None
def initializeTao(
filePath = None,
lastTrackedElement = "end",
csrTF = False,
inputBeamFilePathSuffix = None,
numMacroParticles = None,
loadDefaultLatticeTF = True,
defaultsFile = None,
runImpactTF = False,
scratchPath = None,
randomizeFileNames = False,
**kwargs
):
#######################################################################
#Set file path
#######################################################################
global filePathGlobal
if not filePath:
filePath = os.getcwd()
if not scratchPath:
scratchPath = filePath
os.environ['FACET2_LATTICE'] = filePath
filePathGlobal = filePath
print('Environment set to: ', environ['FACET2_LATTICE'])
#######################################################################
#Launch and configure Tao
#######################################################################
tao=Tao('-init {:s}/bmad/models/f2_elec/tao.init -noplot'.format(environ['FACET2_LATTICE']))
tao.cmd("set beam add_saved_at = DTOTR, XTCAVF, M2EX, PR10571, PR10711, CN2069") #The beam is saved at all MARKER elements already; this list just supplements
tao.cmd(f'set beam_init track_end = {lastTrackedElement}') #See track_start and track_end values with `show beam`
print(f"Tracking to {lastTrackedElement}")
tao.cmd(f'call {filePath}/bmad/models/f2_elec/scripts/Activate_CSR.tao')
if csrTF:
tao.cmd('csron')
print("CSR on")
else:
tao.cmd('csroff')
print("CSR off")
if loadDefaultLatticeTF:
print("Overwriting lattice with setLattice() defaults")
setLattice(tao, verbose = True, defaultsFile = defaultsFile) #Set lattice to my latest default config
else:
print("Not using setLattice(). Golden lattice")
#######################################################################
#Import or generate input beam file
#######################################################################
if randomizeFileNames:
#True-random path for this particular instance
randomPath = str(int.from_bytes(os.urandom(8), "big"))
activeFilePath = f'{scratchPath}/beams/activeBeamFile_{randomPath}.h5'
patchFilePath = f'{scratchPath}/beams/patchBeamFile_{randomPath}.h5'
else:
activeFilePath = f'{scratchPath}/beams/activeBeamFile.h5'
patchFilePath = f'{scratchPath}/beams/patchBeamFile.h5'
# Create 'beams' folder if it doesn't exist
os.makedirs(f"{scratchPath}/beams", exist_ok=True)
if runImpactTF:
if not numMacroParticles:
print("Define numMacroParticles to run Impact")
return
runImpact(
filePath = filePath,
numMacroParticles = numMacroParticles,
**kwargs
)
inputBeamFilePath = f'{filePath}/beams/ImpactBeam.h5'
else:
if inputBeamFilePathSuffix:
inputBeamFilePath = f'{filePath}{inputBeamFilePathSuffix}'
else: #If tracking wasn't requested and a beamfile wasn't specified just grab a random beam... assume the user only wants to do single-particle sims
print("WARNING! No beam file is specified!")
#inputBeamFilePath = f'{filePath}/beams/activeBeamFile.h5'
inputBeamFilePath = f'{filePath}/beams/L0AFEND_facet2-lattice.h5'
if numMacroParticles:
print(f"Number of macro particles = {numMacroParticles}")
else:
print(f"Number of macro particles defined by input file")
#Create the beam
modifyAndSaveInputBeam(
inputBeamFilePath,
numMacroParticles = (None if runImpactTF else numMacroParticles),
outputBeamFilePath = activeFilePath
)
tao.cmd(f'set beam_init position_file={activeFilePath}')
tao.cmd('reinit beam')
print(f"Beam created, written to {activeFilePath}, and reinit to tao")
#Save things into the tao object
tao.inputBeamFilePath = inputBeamFilePath
tao.activeFilePath = activeFilePath
tao.patchFilePath = patchFilePath
#tao.activeBeam = activeBeam
return tao
# def reinitActiveBeam(tao):
# #Take the beam stored in the tao object (tao.activeBeam), save it to a file, load and reinit tao with that file
# (tao.activeBeam).write(tao.activeFilePath)
# tao.cmd(f'set beam_init position_file={tao.activeFilePath}')
# tao.cmd('reinit beam')
# #os.remove(tao.activeFilePath)
# def reinitPatchBeam(tao, P):
# #Take the provided beam, save it to a file, load and reinit tao with that file
# P.write(tao.patchFilePath)
# tao.cmd(f'set beam_init position_file={tao.patchFilePath}')
# tao.cmd('reinit beam')
# #os.remove(tao.patchFilePath)
def trackBeam(
tao,
trackStart = "L0AFEND",
trackEnd = "end",
laserHeater = False,
centerBC14 = False,
assertBC14Energy = False,
centerBC20 = False,
assertBC20Energy = False,
allCollimatorRules = None,
centerMFFF = False,
verbose = False,
**kwargs,
):
"""
Tracks the beam in activeBeamFile.h5 through the lattice presently in tao from trackStart to trackEnd
Some special options are available but disabled by default
* Centering
* At some selected treaty points, remove net offsets to transverse position and angle
* Assert energy
* Centering must be enabled. Can either set True (for default energy at that point) or the desired energy in eV. This is sort of a virtual energy feedback
* Laser heater
* Refer to addLHmodulation(). Need to pass additional options to trackBeam() as **kwargs!
* BC20 collimators
* Refer to collimateBeam(). Collimator positions passed as allCollimatorRules
"""
global filePathGlobal
tao.cmd(f'set beam_init position_file={tao.activeFilePath}')
tao.cmd('reinit beam')
if verbose: print(f"Loaded {tao.activeFilePath}")
tao.cmd(f'set beam_init track_start = {trackStart}')
tao.cmd(f'set beam_init track_end = {trackEnd}')
if verbose: print(f"Set track_start = {trackStart}, track_end = {trackEnd}")
#Adding S-location checks so center* commands won't trigger unnecessarily
trackStartS = tao.ele_param(trackStart,"ele.s")['ele_s']
trackEndS = tao.ele_param(trackEnd,"ele.s")['ele_s']
laserHeaterS = tao.ele_param("HTRUNDF","ele.s")['ele_s']
BC14BEGS = tao.ele_param("BEGBC14_1","ele.s")['ele_s']
BC20BEGS = tao.ele_param("BEGBC20","ele.s")['ele_s']
BC20COLLS = tao.ele_param("CN2069","ele.s")['ele_s']
MFFFS = tao.ele_param("MFFF","ele.s")['ele_s']
if laserHeater and trackStartS < laserHeaterS < trackEndS:
#Will track from start to HTRUNDF, get the beam, modify it, export it, import it, update track_start and track_end
tao.cmd(f'set beam_init track_end = HTRUNDF')
if verbose: print(f"Set track_end = HTRUNDF")
if verbose: print(f"Tracking!")
trackBeamHelper(tao)
P = getBeamAtElement(tao, "HTRUNDF", tToZ = False)
PAfterLHmodulation, deltagamma, t = addLHmodulation(P, **kwargs,);
writeBeam(PAfterLHmodulation, tao.patchFilePath)
if verbose: print(f"Beam with LH modulation written to {tao.patchFilePath}")
tao.cmd(f'set beam_init position_file={tao.patchFilePath}')
tao.cmd('reinit beam')
if verbose: print(f"Loaded {tao.patchFilePath}")
tao.cmd(f'set beam_init track_start = HTRUNDF')
tao.cmd(f'set beam_init track_end = {trackEnd}')
if verbose: print(f"Set track_start = HTRUNDF, track_end = {trackEnd}")
if centerBC14 and trackStartS < BC14BEGS < trackEndS:
tao.cmd(f'set beam_init track_end = BEGBC14_1')
if verbose: print(f"Set track_end = BEGBC14_1")
if verbose: print(f"Tracking!")
trackBeamHelper(tao)
P = getBeamAtElement(tao, "BEGBC14_1", tToZ = False)
if assertBC14Energy:
if type(assertBC14Energy) is bool:
assertBC14Energy = 4.5e9
if verbose: print(f"""Also setting BC14 energy = {1e-9 * assertBC14Energy} GeV, from {1e-9 * P["mean_energy"]} GeV""")
PMod = centerBeam(P, assertEnergy = assertBC14Energy)
else:
PMod = centerBeam(P)
writeBeam(PMod, tao.patchFilePath)
if verbose: print(f"Beam centered at BEGBC14 written to {tao.patchFilePath}")
tao.cmd(f'set beam_init position_file={tao.patchFilePath}')
tao.cmd('reinit beam')
if verbose: print(f"Loaded {tao.patchFilePath}")
tao.cmd(f'set beam_init track_start = BEGBC14_1')
tao.cmd(f'set beam_init track_end = {trackEnd}')
if verbose: print(f"Set track_start = BEGBC14_1, track_end = {trackEnd}")
if centerBC20 and trackStartS < BC20BEGS < trackEndS:
tao.cmd(f'set beam_init track_end = BEGBC20')
if verbose: print(f"Set track_end = BEGBC20")
if verbose: print(f"Tracking!")
trackBeamHelper(tao)
P = getBeamAtElement(tao, "BEGBC20", tToZ = False)
if assertBC20Energy:
if type(assertBC20Energy) is bool:
assertBC20Energy = 10e9
if verbose: print(f"""Also setting BC20 energy = {1e-9 * assertBC20Energy} GeV, from {1e-9 * P["mean_energy"]} GeV""")
PMod = centerBeam(P, assertEnergy = assertBC20Energy)
else:
PMod = centerBeam(P)
writeBeam(PMod, tao.patchFilePath)
if verbose: print(f"Beam centered at BEGBC20 written to {tao.patchFilePath}")
tao.cmd(f'set beam_init position_file={tao.patchFilePath}')
tao.cmd('reinit beam')
if verbose: print(f"Loaded {tao.patchFilePath}")
tao.cmd(f'set beam_init track_start = BEGBC20')
tao.cmd(f'set beam_init track_end = {trackEnd}')
if verbose: print(f"Set track_start = BEGBC20, track_end = {trackEnd}")
if allCollimatorRules and trackStartS < BC20COLLS < trackEndS:
tao.cmd(f'set beam_init track_end = CN2069')
if verbose: print(f"Set track_end = CN2069")
if verbose: print(f"Tracking!")
trackBeamHelper(tao)
P = getBeamAtElement(tao, "CN2069", tToZ = False)
PMod = collimateBeam(P, allCollimatorRules)
writeBeam(PMod, tao.patchFilePath)
if verbose: print(f"Collimated beam written to {tao.patchFilePath}. Rules: {allCollimatorRules}")
tao.cmd(f'set beam_init position_file={tao.patchFilePath}')
tao.cmd('reinit beam')
if verbose: print(f"Loaded {tao.patchFilePath}")
tao.cmd(f'set beam_init track_start = CN2069')
tao.cmd(f'set beam_init track_end = {trackEnd}')
if verbose: print(f"Set track_start = CN2069, track_end = {trackEnd}")
if centerMFFF and trackStartS < MFFFS < trackEndS:
tao.cmd(f'set beam_init track_end = MFFF')
if verbose: print(f"Set track_end = MFFF")
if verbose: print(f"Tracking!")
trackBeamHelper(tao)
P = getBeamAtElement(tao, "MFFF", tToZ = False)
PMod = centerBeam(P)
writeBeam(PMod, tao.patchFilePath)
if verbose: print(f"Beam centered at MFFF written to {tao.patchFilePath}")
tao.cmd(f'set beam_init position_file={tao.patchFilePath}')
tao.cmd('reinit beam')
if verbose: print(f"Loaded {tao.patchFilePath}")
tao.cmd(f'set beam_init track_start = MFFF')
tao.cmd(f'set beam_init track_end = {trackEnd}')
if verbose: print(f"Set track_start = MFFF, track_end = {trackEnd}")
if verbose: print(f"Tracking!")
trackBeamHelper(tao)
if verbose: print(f"trackBeam() exiting")
#For backwards compatibility, return to activeBeamFile. Might be unnecessary
# tao.cmd(f'set beam_init position_file={filePathGlobal}/beams/activeBeamFile.h5')
# tao.cmd('reinit beam')
def trackBeamLEGACY(tao):
#This is the pre-2024-08-23 version of trackBeam(), retained for debugging purposes. Can be deleted
tao.cmd('set global track_type = beam') #set "track_type = single" to return to single particle
tao.cmd('set global track_type = single') #return to single to prevent accidental long re-evaluation
def trackBeamHelper(tao):
"""Wrap some of the tao commands with a try/except. This way if tracking doesn't work, we failsafe to track_type = single"""
try:
tao.cmd('set global track_type = beam') #set "track_type = single" to return to single particle
except:
print("Beam tracking failed. Resetting track_type = single")
tao.cmd('set global track_type = single') #return to single to prevent accidental long re-evaluation
raise #Rethrow the error. Adding this line causes trackBeam() to see the error and also fail instead of potentially entering a weird state
tao.cmd('set global track_type = single') #return to single to prevent accidental long re-evaluation
return
def getBeamAtElement(tao, eleString, tToZ = True):
P = ParticleGroup(data=tao.bunch_data(eleString))
P = P[P.status == 1]
#Naive implementation for "typical" beams. ParticleGroup has .drift_to_z but I couldn't get it to work...
if tToZ:
P.z = -299792458 * P["delta_t"]
#P.t = 0 * P.t #I haven't decided the best practice for this yet. Technically the beam is not self-consistent without t being set to zero but not doing so is convenient for backwards compatibility
return P
def nudgeMacroparticleWeights(
PInput,
trailingBunchFraction = None,
trailingBunchType = None
):
"""
This is NOT a robust function. Don't trust it to do what you want
Presently splits based on z and a user-specified charge ratio. Lots of things can go wrong if you aren't careful!
Borrowing stuff from 2024-03-29_nudgeMacroparticleWeights.ipynb
"""
P = PInput.copy()
zVals = (P["delta_z"]).copy()
zVals = np.sort(zVals)
splitZ = zVals[int(trailingBunchFraction * len(zVals))]
startingWeight = P.weight[0]
startingWeight
witnessWeight = 0.999*startingWeight
driverWeight = 1.001*startingWeight
if trailingBunchType == "witness":
trailingBunchWeight = witnessWeight
leadingBunchWeight = driverWeight
if trailingBunchType == "driver":
trailingBunchWeight = driverWeight
leadingBunchWeight = witnessWeight
newWeightArr = np.full(np.size(P.weight), -1.1)
for i in range(np.size(newWeightArr)):
if P["delta_z"][i] < splitZ:
newWeightArr[i] = trailingBunchWeight
else:
newWeightArr[i] = leadingBunchWeight
P.weight = newWeightArr
return P
def getDriverAndWitness(P):
#See, e.g. "2024-07-01 Nudge Macroparticle Weights.ipynb" for details
weights = np.sort(np.unique(P.weight))
if len(weights) != 2:
print("WARNING! Expected drive/witness structure not found")
return
PWitness = P[P.weight == weights[0]]
PDrive = P[P.weight == weights[1]]
return PDrive, PWitness
def writeBeam(P, fileName):
""" Writes the beam as an h5 with E. Cropp's timeOffset fix """
P.write(fileName)
OpenPMD_to_Bmad(fileName)
def makeBeamActiveBeamFile(P, tao = None):
#Weird structure on this function for backwards compatiblity. If no tao object is provided, put the beam in the nominal location. Otherwise, put it in the right place
global filePathGlobal
if tao:
writeBeam(P, tao.activeFilePath)
else:
print(f"WARNING! No tao object provided. Writing beam to {filePathGlobal}/beams/activeBeamFile.h5... hope that's what you wanted")
writeBeam(P, f"{filePathGlobal}/beams/activeBeamFile.h5")
def smallestInterval(nums, percentage=0.9):
"""Give the smallest interval containing a desired percentage of provided points"""
nums = nums.copy() # Create a copy to avoid modifying the original list
nums.sort()
n = len(nums)
k = int(n * percentage)
min_range = float('inf')
interval = (None, None)
for i in range(n - k + 1):
current_range = nums[i + k - 1] - nums[i]
if current_range < min_range:
min_range = current_range
interval = (nums[i], nums[i + k - 1])
return interval[1]-interval[0]
def smallestIntervalImpliedSigma(nums, percentage=0.9):
interval = smallestInterval(nums, percentage)
intervalToSigmaFactor = scipy.special.erfinv(percentage) * (2 * np.sqrt(2))
return interval/intervalToSigmaFactor
#See "Discussion of alternative emittance and spot size calculations.ipynb"
#See also "2024-07-01 RMS vs FWHM at PENT.ipynb"
def smallestIntervalImpliedEmittanceModelFunction(z, sigmax, sigmaxp, rho):
return np.sqrt(sigmax**2 + 2 * z * rho * sigmax * sigmaxp + z**2 * sigmaxp**2)
def smallestIntervalImpliedEmittance(P, plane = "x", percentage = 0.9, verbose = False):
#zValues = np.arange(-20, 20, 0.1)
zValues = np.array([-131.072, -65.536, -32.768, -16.384, -8.192, -4.096, -2.048, -1.024,
-0.512, -0.256, -0.128, -0.064, -0.032, -0.016, -0.008, -0.004,
-0.002, -0.001, 0.0, 0.001, 0.002, 0.004, 0.008, 0.016, 0.032,
0.064, 0.128, 0.256, 0.512, 1.024, 2.048, 4.096, 8.192,
16.384, 32.768, 65.536, 131.072]) #(-Reverse[#])~Join~{0}~Join~# &@PowerRange[0.001, 200, 2]
if plane == "x":
sigmaXResults = [ smallestIntervalImpliedSigma(P.x + z * P.xp, percentage = percentage) for z in zValues]
sigmaXResultsExact = [ np.std(P.x + z * P.xp) for z in zValues]
elif plane == "y":
sigmaXResults = [ smallestIntervalImpliedSigma(P.y + z * P.yp, percentage = percentage) for z in zValues]
sigmaXResultsExact = [ np.std(P.y + z * P.yp) for z in zValues]
else:
return
#sigmaXResults = [ smallestIntervalImpliedSigma(P.x + z * P.xp, percentage = percentage) for z in zValues]
#sigmaXResultsExact = [ np.std(P.x + z * P.xp) for z in zValues]
# Fit the model to the data
popt, pcov = curve_fit(smallestIntervalImpliedEmittanceModelFunction, zValues, sigmaXResults, p0=[1, 1, 0])
# Extract optimal parameters
sigmax_opt, sigmaxp_opt, rho_opt = popt
emit_opt = np.sqrt( sigmax_opt**2 * sigmaxp_opt**2 - (rho_opt * sigmax_opt * sigmaxp_opt)**2 )
if verbose:
print(f"""True sigma_x, sigma_xp, rho: {P.std("x")}, {P.std("xp")}, {P["cov_x__xp"] / (P.std("x") * P.std("xp"))}""")
print(f"Optimizer parameters: sigma_x = {sigmax_opt}, sigma_xp = {sigmaxp_opt}, rho = {rho_opt}")
plt.scatter(zValues, sigmaXResultsExact, label='True rms')
plt.scatter(zValues, sigmaXResults, label='Inferred rms')
plt.plot(zValues, smallestIntervalImpliedEmittanceModelFunction(zValues, *popt), label='Fitted function', color='red')
plt.xlabel('Drift [m]')
plt.ylabel('Sigma_x [m]')
plt.legend()
plt.show()
print(f"""Actual emittance: \t {P["norm_emit_x"]}""")
print(f"""Fit emittance: \t\t {emit_opt * P["mean_gamma"]}""")
return emit_opt * P["mean_gamma"]
def emittance(P, plane = "x", fraction = 0.9):
"""Just a wrapper for the OpenPMD functions which let you specify the fraction used in the emittance calculation"""
return P.twiss(plane = plane, fraction = fraction)[f"norm_emit_{plane}"]
def displayMatrix(matrix):
display(pd.DataFrame(matrix).style.hide(axis="index").hide(axis="columns"))
def getMatrix(tao, start, end, order = 1, startOffset = 0, endOffset = 0, print = False):
"""Return zero or first order transport matrix from start to end, with offsets. Optionally print in a human readable format"""
startS = tao.ele_head(start)["s"] + startOffset
endS = tao.ele_head(end)["s"] + endOffset
#This disgusting workaround is required because PyTao implementation of taylor_map doesn't support -s mode
if order == 0:
allStrings = tao.show(f"taylor_map -order 0 -s {startS} {endS}")
transportMatrix = np.array([float(str) for str in allStrings[0].split()])
elif order == 1:
allStrings = tao.show(f"taylor_map -s {startS} {endS}")
splitStrings = [str.split() for str in allStrings]
transportMatrix = np.array([ [float(str) for str in row[0:6]] for row in splitStrings[2:] ])
else:
print("Invalid matrix order requested")
return
if print:
displayMatrix(transportMatrix)
return transportMatrix
def getMatrixLEGACY(tao, start, end, order = 1, print = False):
"""Return zero or first order transport matrix from start to end. Optionally print in a human readable format"""
if order == 0:
transportMatrix = (tao.matrix(start, end))["vec0"]
elif order == 1:
transportMatrix = (tao.matrix(start, end))["mat6"]
else:
print("Invalid matrix order requested")
return
if print:
#display(pd.DataFrame(transportMatrix).style.hide(axis="index").hide(axis="columns"))
displayMatrix(transportMatrix)
return transportMatrix
def setLatticeAndGetMatrix(tao, start, end, startOffset = 0, endOffset = 0, defaultSettings = None, overrideSettings = {}):
"""
Load lattice based on default and override settings. Then provide the transfer matrix between two elements, optionally with offsets
If no defaultSettings are specified, the golden lattice from defaults.yml will be loaded
"""
if not defaultSettings:
defaultSettings = loadConfig("setLattice_configs/defaults.yml")
setLattice(tao, **( defaultSettings | overrideSettings ) )
return getMatrix(tao, start, end, startOffset = startOffset, endOffset = endOffset)
def addLHmodulation(
inputBeam,
#Elaser,
showplots=False,
laserHeater_laserEnergy = 0.5e-3,
laserHeater_sigma_t = (2 / 2.355) * 1e-12,
laserHeater_offset = -0.5
):
""" From C. Emma, 2024-08-23 """
# Hardcode FACET-II laser and undulator parameters
# Laser parameters
Elaser = laserHeater_laserEnergy
lambda_laser = 760e-9
sigmar_laser = 200e-6
sigmat_laser = laserHeater_sigma_t # (2 / 2.355) * 1e-12
Plaser = Elaser / np.sqrt(2 * np.pi * sigmat_laser**2)
offset = laserHeater_offset #-0.5 # laser to e-beam offset if you want you can add it
# Undulator parameters
K = 1.1699
lambdaw = 0.054
Nwig = 9
# Electron beam
outputBeam = inputBeam.copy()
x = inputBeam.x - np.mean(inputBeam.x)
y = inputBeam.y - np.mean(inputBeam.y)
gamma = inputBeam.gamma
gamma0 = np.mean(inputBeam.gamma)
t = inputBeam.t-np.mean(inputBeam.t);
# Calculated parameters
lambda_r = lambdaw / 2 / gamma0**2 * (1 + K**2 / 2) # Assumes planar undulator
omega = 2 * np.pi * 299792458 / lambda_r # Resonant frequency
JJ = besselj(0, K**2 / (4 + 2 * K**2)) - besselj(1, K**2 / (4 + 2 * K**2))
totalLaserEnergy = np.sqrt(2 * np.pi * sigmat_laser**2) * Plaser
# Laser is assumed Gaussian with peak power Plaser
# This formula from eq. 8 Huang PRSTAB 074401 2004
mod_amplitude = np.sqrt(Plaser / 8.7e9) * K * lambdaw * Nwig / gamma0 / sigmar_laser * JJ
#print(mod_amplitude / np.sqrt(Plaser))
# offset = 1.0 # temporal offset between laser and e-beam in units of laser wavelengths
# Calculate induced modulation deltagamma
deltagamma = mod_amplitude * np.exp(-0.25 * (x**2 + y**2) / sigmar_laser**2) * \
np.sin(omega * t + offset * 2 * np.pi) * \
np.exp(-0.5 * ((t - offset * sigmat_laser) / sigmat_laser)**2)
outputBeam.gamma = inputBeam.gamma + deltagamma
return outputBeam, deltagamma, t
def centerBeam(
P,
centerType = "median",
assertEnergy = None
):
"""
Shifts x, y, xp, and yp of a beam to zero
centerType is either "median" or "mean"
"""
PMod = P
if centerType == "median":
PMod.x = P.x - np.median(P.x)
PMod.y = P.y - np.median(P.y)
PMod.px = P.px - np.median(P.px)
PMod.py = P.py - np.median(P.py)
if assertEnergy:
PMod.pz = P.pz * assertEnergy / np.median(P.pz)
return PMod
if centerType == "mean":
PMod.x = P.x - np.mean(P.x)
PMod.y = P.y - np.mean(P.y)
PMod.px = P.px - np.mean(P.px)
PMod.py = P.py - np.mean(P.py)
if assertEnergy:
PMod.pz = P.pz * assertEnergy / np.mean(P.pz)
return PMod
return
def calcBMAG(b0, a0, b, a):
#From Lucretia
#For a bit more detail, see "BMAG from Lucretia.nb"
#Not validated!!!
# function [B,Bpsi]=bmag(b0,a0,b,a)
# %
# % [B,Bpsi]=bmag(b0,a0,b,a);
# %
# % Compute BMAG and its phase from Twiss parameters
# %
# % INPUTs:
# %
# % b0 = matched beta
# % a0 = matched alpha
# % b = mismatched beta
# % a = mismatched alpha
# %
# % OUTPUTs:
# %
# % B = mismatch amplitude
# % Bpsi = mismatch phase (deg)
g0 = (1 + a0 ** 2) / b0
g = (1 + a ** 2) / b
B = (b0 * g - 2.0 * a0 * a + g0 * b) / 2
return B
def collimateBeam(
P,
allCollimatorRules = None
):
"""
allCollimatorRules is a list of lists. Each sublist should have exactly two elements for the lower and upper x position of a collimator
Arbitrarily many collimators can be defined this way; therefore it works for notch and/or jaw collimators
"""
PMod = P.copy()
for collimatorRange in allCollimatorRules:
print(collimatorRange)
all_indices = np.arange(len(PMod.x))
killedIndices = np.where(np.logical_and(PMod.x > collimatorRange[0], PMod.x < collimatorRange[1]))[0]
survivingIndices = np.setdiff1d(all_indices, killedIndices)
# OpenPMD checks the length so I can't just remove the "killed" particles
# Also, for compatibility, I don't want to change either the weight or status of the killed particles
filtered_data = {
"x": PMod.x[survivingIndices],
"y": PMod.y[survivingIndices],
"z": PMod.z[survivingIndices],
"px": PMod.px[survivingIndices],
"py": PMod.py[survivingIndices],
"pz": PMod.pz[survivingIndices],
"t": PMod.t[survivingIndices],
"status": PMod.status[survivingIndices],
"weight": PMod.weight[survivingIndices],
"species": PMod.species
}
# Create a new ParticleGroup instance with the filtered data
PMod = ParticleGroup(data=filtered_data)
print(f"New particle count: {len(PMod.x)}")
print(f"{len(PMod.x)}")
return PMod
def sortIndices(lst):
#Returns the indices of the sorted elements, e.g. [1, 3, 5, 2, 4] --> [0, 3, 1, 4, 2]
return [i for i, _ in sorted(enumerate(lst), key=lambda x: x[1])]
def sliceBeam(
P,
sortKey = None,
numBeamlets = None
):
"""Sort a beam by sortKey, then slice it into numBeamlets of equal count"""
sortedIndices = sortIndices(P[sortKey])
subsetIndices = np.array_split(sortedIndices, numBeamlets)
resultBeamlets = []
for activeSubsetIndices in subsetIndices:
PMod = P.copy()
# OpenPMD checks the length so I can't just remove the "killed" particles
# Also, for compatibility, I don't want to change either the weight or status of the killed particles
filtered_data = {
"x": PMod.x[activeSubsetIndices],
"y": PMod.y[activeSubsetIndices],
"z": PMod.z[activeSubsetIndices],
"px": PMod.px[activeSubsetIndices],
"py": PMod.py[activeSubsetIndices],
"pz": PMod.pz[activeSubsetIndices],
"t": PMod.t[activeSubsetIndices],
"status": PMod.status[activeSubsetIndices],
"weight": PMod.weight[activeSubsetIndices],
"species": PMod.species
}
# Create a new ParticleGroup instance with the filtered data
PMod = ParticleGroup(data=filtered_data)
#print(f"New particle count: {len(PMod.x)}")
#print(f"{len(PMod.x)}")
resultBeamlets.append(PMod)
return resultBeamlets
def getSingleBeamSlice(
P,
sortKey = None,
minVal = None,
maxVal = None
):
"""Return a beamlet of particles which satisfy the inequality """
# Get indices where sortKey is within the given range
mask = (P[sortKey] >= minVal) & (P[sortKey] <= maxVal)
if not np.any(mask):
raise ValueError("No particles found in the specified range.")
# Filter data based on the mask
filtered_data = {
"x": P.x[mask],
"y": P.y[mask],
"z": P.z[mask],
"px": P.px[mask],
"py": P.py[mask],
"pz": P.pz[mask],
"t": P.t[mask],
"status": P.status[mask],
"weight": P.weight[mask],
"species": P.species
}
PMod = ParticleGroup(data=filtered_data)
return PMod
def loadConfig(file, loaded_files=None):
"""Code to load nested config files... ChatGPT is the author, beware!"""
if loaded_files is None:
loaded_files = set()
if file in loaded_files:
return {} # Avoid circular imports
loaded_files.add(file)
with open(file, 'r') as f:
data = yaml.safe_load(f) or {}
# Handle includes
includes = data.pop('include', [])
merged_data = {}
for include_file in includes:
merged_data.update(loadConfig(include_file, loaded_files))
merged_data.update(data) # Later settings override earlier ones
return merged_data
def getBeamSpecs(P, targetTwiss = None):
"""
Returns a collection of convenient beam parameters as a dictionary
Will automatically detect and add extra measurements for two-bunch beams
targetTwiss can either be in the form [betaX, alphaX, betaY, alphaY] or
for a very limited number of treaty point elements, can instead provide the element name. This will use the golden lattice targetTwiss
Presently defined: "PR10571", "BEGBC20", "MFFF", "PENT"
"""
savedData = {}
#A silly trick; set() will give back unique values
bunchCount = len(set(P.weight))
if bunchCount == 1:
PDrive = P.copy()
beamsToEvaluate = ["PDrive"]
elif bunchCount == 2:
PDrive, PWitness = getDriverAndWitness(P)
beamsToEvaluate = ["PDrive", "PWitness"]
else:
print("bunchCount doesn't make sense. Aborting")
return
if targetTwiss:
if isinstance(targetTwiss, str):
if targetTwiss == "PR10571":
#PR10571 lucretia live model lattice 2024-10-16
targetBetaX = 5.7
targetBetaY = 2.6
targetAlphaX = -2.1
targetAlphaY = 0.0
if targetTwiss == "BEGBC20":
#BEGBC20 lucretia live model lattice 2024-10-16
targetBetaX = 11.5
targetBetaY = 27.3
targetAlphaX = 0.7
targetAlphaY = 1.2
if targetTwiss == "MFFF":
#MFFF lucretia live model lattice 2024-10-16
targetBetaX = 11.6
targetAlphaX = -0.64
targetBetaY = 25.2
targetAlphaY = -1.6
elif targetTwiss == "PENT":
#PENT lucretia live model lattice 2024-10-16
targetBetaX = 0.5
targetAlphaX = 0.0
targetBetaY = 0.5
targetAlphaY = 0.0
else:
#print("Not a valid treaty point. Aborting")
#return
#Invalid treaty point; setting to None to avoid BMAG evaluation
targetTwiss = None
else:
targetBetaX, targetAlphaX, targetBetaY, targetAlphaY = targetTwiss
for PActiveStr in beamsToEvaluate:
PActive = locals()[PActiveStr]
# for val in ["mean_x", "mean_y", "sigma_x", "sigma_y", "mean_xp", "mean_yp"]:
# savedData[f"{PActiveStr}_{val}"] = PActive[val]
savedData[f"{PActiveStr}_median_x"] = np.median(PActive.x)
savedData[f"{PActiveStr}_median_y"] = np.median(PActive.y)
savedData[f"{PActiveStr}_median_xp"] = np.median(PActive.xp)
savedData[f"{PActiveStr}_median_yp"] = np.median(PActive.yp)
savedData[f"{PActiveStr}_median_energy"] = np.median(PActive.energy)
savedData[f"{PActiveStr}_sigmaSI90_x"] = smallestIntervalImpliedSigma(PActive.x, percentage = 0.90)
savedData[f"{PActiveStr}_sigmaSI90_y"] = smallestIntervalImpliedSigma(PActive.y, percentage = 0.90)
savedData[f"{PActiveStr}_sigmaSI90_z"] = smallestIntervalImpliedSigma(PActive.t * 3e8, percentage=0.9)
savedData[f"{PActiveStr}_sigmaSI90_xp"] = smallestIntervalImpliedSigma(PActive.xp, percentage = 0.90)
savedData[f"{PActiveStr}_sigmaSI90_yp"] = smallestIntervalImpliedSigma(PActive.yp, percentage = 0.90)
savedData[f"{PActiveStr}_sigmaSI90_energy"] = smallestIntervalImpliedSigma(PActive.energy, percentage = 0.90)
savedData[f"{PActiveStr}_emitSI90_x"] = smallestIntervalImpliedEmittance(PActive, plane = "x", percentage = 0.90)
savedData[f"{PActiveStr}_emitSI90_y"] = smallestIntervalImpliedEmittance(PActive, plane = "y", percentage = 0.90)
savedData[f"{PActiveStr}_norm_emit_x"] = (PActive.twiss(plane = "x", fraction = 0.9))["norm_emit_x"]
savedData[f"{PActiveStr}_norm_emit_y"] = (PActive.twiss(plane = "y", fraction = 0.9))["norm_emit_y"]
if bunchCount == 2:
savedData[f"{PActiveStr}_zCentroid"] = np.median(PActive.t * 3e8)
savedData[f"{PActiveStr}_charge_nC"] = PActive.charge * 1e9
PActiveTwiss = PActive.twiss(plane = "x", fraction = 0.9) | PActive.twiss(plane = "y", fraction = 0.9)
if targetTwiss:
savedData[f"{PActiveStr}_BMAG_x"] = calcBMAG(targetBetaX, targetAlphaX, PActiveTwiss["beta_x"], PActiveTwiss["alpha_x"])
savedData[f"{PActiveStr}_BMAG_y"] = calcBMAG(targetBetaY, targetAlphaY, PActiveTwiss["beta_y"], PActiveTwiss["alpha_y"])
# Get BMAGs by energy slice
slicedBeamlets = sliceBeam( PActive , sortKey = "pz", numBeamlets = 5 )
slicedTwiss = [ ( beamlet.twiss(plane = "x", fraction = 0.9) | beamlet.twiss(plane = "y", fraction = 0.9) ) for beamlet in slicedBeamlets ]
savedData[f"{PActiveStr}_sliced_BMAG_x"] = [ calcBMAG(targetBetaX, targetAlphaX, beamletTwiss["beta_x"], beamletTwiss["alpha_x"]) for beamletTwiss in slicedTwiss ]
savedData[f"{PActiveStr}_sliced_BMAG_y"] = [ calcBMAG(targetBetaY, targetAlphaY, beamletTwiss["beta_y"], beamletTwiss["alpha_y"]) for beamletTwiss in slicedTwiss ]
if bunchCount == 2:
savedData["bunchSpacing"] = savedData["PWitness_zCentroid"] - savedData["PDrive_zCentroid"]
savedData["transverseCentroidOffset"] = np.sqrt(
(savedData["PDrive_median_x"] - savedData["PWitness_median_x"])**2 +
(savedData["PDrive_median_y"] - savedData["PWitness_median_y"])**2
)
#savedData["lostChargeFraction"] = 1 - (P.charge / PInit.charge)
return savedData
#Here's a version that would work if the axes are coupled.... they really, really, really shouldn't ever be though
def launchTwissCorrectionObjective(params, tao, evalElement, targetBetaX, targetAlphaX, targetBetaY, targetAlphaY):