Skip to content

Commit a81ee59

Browse files
JPTS2JPTS2
authored
Merge pull request #14 from JPTS2/0.1.12
0.1.12 Add computation of chromatic functions and second order dispersions for SAD twiss
2 parents cbd3262 + 019a476 commit a81ee59

File tree

4 files changed

+287
-8
lines changed

4 files changed

+287
-8
lines changed

‎sad2xs/sad_helpers/__init__.py‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -10,7 +10,7 @@
1010
# SAD Helper Functions
1111
################################################################################
1212
from .rebuild_lattice import rebuild_sad_lattice
13-
from .twiss import twiss_sad
13+
from .twiss import twiss_sad, compute_chromatic_functions, compute_second_order_dispersions
1414
from .emit import emit_sad
1515
from .track import track_sad
1616
from .survey import survey_sad

‎sad2xs/sad_helpers/track.py‎

Lines changed: 4 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -288,7 +288,7 @@ def track_sad(
288288
!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!
289289
! Save to file
290290
!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!
291-
Put[beam, "temp_sad_track.dat"];
291+
Put[beam[[2]], "temp_sad_track.dat"];
292292
293293
!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!
294294
! Close process
@@ -306,7 +306,7 @@ def track_sad(
306306
!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!
307307
! Save to file
308308
!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!
309-
Put[beam, "temp_sad_track.dat"];
309+
Put[beam[[2]], "temp_sad_track.dat"];
310310
311311
!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!
312312
! Close process
@@ -384,7 +384,7 @@ def track_sad(
384384
########################################
385385
print("Loading output file")
386386
with open("temp_sad_track.dat", "r", encoding = "utf-8") as f:
387-
output = f.read()
387+
raw_output = f.read()
388388

389389
########################################
390390
# Remove temporary data
@@ -398,7 +398,7 @@ def track_sad(
398398
print("Processing outputs")
399399

400400
# Fix Mathematica"s ".00123" → "0.00123"
401-
output = re.sub(r"(?<![\d])\.(\d+)", r"0.\1", output)
401+
output = re.sub(r"(?<![\d])\.(\d+)", r"0.\1", raw_output)
402402

403403
# Extract all floats in a single pass
404404
output = np.fromstring(

‎sad2xs/sad_helpers/twiss.py‎

Lines changed: 280 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -137,7 +137,7 @@ def generate_twiss_print_function():
137137
return twiss_command
138138

139139
################################################################################
140-
# Closed Ring 4D Twiss Function
140+
# Twiss
141141
################################################################################
142142
def twiss_sad(
143143
lattice_filepath: str,
@@ -364,3 +364,282 @@ def twiss_sad(
364364
# Return the TwissTable
365365
########################################
366366
return tw_sad
367+
368+
################################################################################
369+
# Add second order dispersions to twiss
370+
################################################################################
371+
def compute_second_order_dispersions(
372+
lattice_filepath: str,
373+
line_name: str,
374+
sad_twiss: xt.TwissTable | None = None,
375+
reverse_element_order: bool = False,
376+
reverse_bend_direction: bool = False,
377+
closed: bool = True,
378+
calc6d: bool = False,
379+
rfsw: bool = True,
380+
rad: bool = False,
381+
radcod: bool = False,
382+
radtaper: bool = False,
383+
delta0: float = 0.0,
384+
ddelta: float = 1E-4,
385+
additional_commands: str = "",
386+
wall_time: int = 60,
387+
sad_path: str = "sad"):
388+
"""
389+
Compute the second order dispersions and add them to the provided twiss table.
390+
391+
With thanks to G. Broggi for the method.
392+
"""
393+
394+
########################################
395+
# Compute twiss at delta0
396+
########################################
397+
tw_on = twiss_sad(
398+
lattice_filepath = lattice_filepath,
399+
line_name = line_name,
400+
reverse_element_order = reverse_element_order,
401+
reverse_bend_direction = reverse_bend_direction,
402+
closed = closed,
403+
calc6d = calc6d,
404+
rfsw = rfsw,
405+
rad = rad,
406+
radcod = radcod,
407+
radtaper = radtaper,
408+
delta0 = delta0,
409+
additional_commands = additional_commands,
410+
wall_time = wall_time,
411+
sad_path = sad_path)
412+
413+
# If a reference twiss is provided, check consistency
414+
if sad_twiss is not None:
415+
assert np.allclose(tw_on.s, sad_twiss.s), \
416+
"Twiss s positions do not match!"
417+
assert np.allclose(tw_on.x, sad_twiss.x), \
418+
"Twiss x positions do not match!"
419+
assert np.allclose(tw_on.px, sad_twiss.px), \
420+
"Twiss px positions do not match!"
421+
assert np.allclose(tw_on.y, sad_twiss.y), \
422+
"Twiss y positions do not match!"
423+
assert np.allclose(tw_on.py, sad_twiss.py), \
424+
"Twiss py positions do not match!"
425+
assert np.allclose(tw_on.zeta, sad_twiss.zeta), \
426+
"Twiss zeta positions do not match!"
427+
assert np.allclose(tw_on.delta, sad_twiss.delta), \
428+
"Twiss delta positions do not match!"
429+
430+
########################################
431+
# Compute off momentum twiss at delta0 + ddelta
432+
########################################
433+
tw_off = twiss_sad(
434+
lattice_filepath = lattice_filepath,
435+
line_name = line_name,
436+
reverse_element_order = reverse_element_order,
437+
reverse_bend_direction = reverse_bend_direction,
438+
closed = closed,
439+
calc6d = calc6d,
440+
rfsw = rfsw,
441+
rad = rad,
442+
radcod = radcod,
443+
radtaper = radtaper,
444+
delta0 = delta0 + ddelta,
445+
additional_commands = additional_commands,
446+
wall_time = wall_time,
447+
sad_path = sad_path)
448+
449+
########################################
450+
# Obtain required data
451+
########################################
452+
x_on = tw_on.x
453+
x_off = tw_off.x
454+
dx_on = tw_on.dx
455+
456+
px_on = tw_on.px
457+
px_off = tw_off.px
458+
dpx_on = tw_on.dpx
459+
460+
y_on = tw_on.y
461+
y_off = tw_off.y
462+
dy_on = tw_on.dy
463+
464+
py_on = tw_on.py
465+
py_off = tw_off.py
466+
dpy_on = tw_on.dpy
467+
468+
########################################
469+
# Compute second order dispersions
470+
########################################
471+
ddx = 2 * (x_off - x_on - dx_on * ddelta) / (ddelta**2)
472+
ddpx = 2 * (px_off - px_on - dpx_on * ddelta) / (ddelta**2)
473+
ddy = 2 * (y_off - y_on - dy_on * ddelta) / (ddelta**2)
474+
ddpy = 2 * (py_off - py_on - dpy_on * ddelta) / (ddelta**2)
475+
476+
########################################
477+
# Add these values to the twiss
478+
########################################
479+
if sad_twiss is None:
480+
sad_twiss = tw_on
481+
482+
sad_twiss["ddx"] = ddx
483+
sad_twiss["ddpx"] = ddpx
484+
sad_twiss["ddy"] = ddy
485+
sad_twiss["ddpy"] = ddpy
486+
487+
########################################
488+
# Return the twiss
489+
########################################
490+
return sad_twiss
491+
492+
################################################################################
493+
# Add chromatic functions
494+
################################################################################
495+
def compute_chromatic_functions(
496+
lattice_filepath: str,
497+
line_name: str,
498+
sad_twiss: xt.TwissTable | None = None,
499+
reverse_element_order: bool = False,
500+
reverse_bend_direction: bool = False,
501+
closed: bool = True,
502+
calc6d: bool = False,
503+
rfsw: bool = True,
504+
rad: bool = False,
505+
radcod: bool = False,
506+
radtaper: bool = False,
507+
delta0: float = 0.0,
508+
ddelta: float = 1E-4,
509+
additional_commands: str = "",
510+
wall_time: int = 60,
511+
sad_path: str = "sad"):
512+
"""
513+
Compute the chromatic functions and add them to the provided twiss table.
514+
515+
With thanks to G. Broggi for the method.
516+
"""
517+
518+
########################################
519+
# Compute twiss at delta0
520+
########################################
521+
tw_on = twiss_sad(
522+
lattice_filepath = lattice_filepath,
523+
line_name = line_name,
524+
reverse_element_order = reverse_element_order,
525+
reverse_bend_direction = reverse_bend_direction,
526+
closed = closed,
527+
calc6d = calc6d,
528+
rfsw = rfsw,
529+
rad = rad,
530+
radcod = radcod,
531+
radtaper = radtaper,
532+
delta0 = delta0,
533+
additional_commands = additional_commands,
534+
wall_time = wall_time,
535+
sad_path = sad_path)
536+
537+
# If a reference twiss is provided, check consistency
538+
if sad_twiss is not None:
539+
assert np.allclose(tw_on.s, sad_twiss.s), \
540+
"Twiss s positions do not match!"
541+
assert np.allclose(tw_on.x, sad_twiss.x), \
542+
"Twiss x positions do not match!"
543+
assert np.allclose(tw_on.px, sad_twiss.px), \
544+
"Twiss px positions do not match!"
545+
assert np.allclose(tw_on.y, sad_twiss.y), \
546+
"Twiss y positions do not match!"
547+
assert np.allclose(tw_on.py, sad_twiss.py), \
548+
"Twiss py positions do not match!"
549+
assert np.allclose(tw_on.zeta, sad_twiss.zeta), \
550+
"Twiss zeta positions do not match!"
551+
assert np.allclose(tw_on.delta, sad_twiss.delta), \
552+
"Twiss delta positions do not match!"
553+
554+
########################################
555+
# Compute off momentum twiss at delta0 + ddelta
556+
########################################
557+
tw_plus = twiss_sad(
558+
lattice_filepath = lattice_filepath,
559+
line_name = line_name,
560+
reverse_element_order = reverse_element_order,
561+
reverse_bend_direction = reverse_bend_direction,
562+
closed = closed,
563+
calc6d = calc6d,
564+
rfsw = rfsw,
565+
rad = rad,
566+
radcod = radcod,
567+
radtaper = radtaper,
568+
delta0 = delta0 + ddelta,
569+
additional_commands = additional_commands,
570+
wall_time = wall_time,
571+
sad_path = sad_path)
572+
573+
########################################
574+
# Compute off momentum twiss at delta0 - ddelta
575+
########################################
576+
tw_minus = twiss_sad(
577+
lattice_filepath = lattice_filepath,
578+
line_name = line_name,
579+
reverse_element_order = reverse_element_order,
580+
reverse_bend_direction = reverse_bend_direction,
581+
closed = closed,
582+
calc6d = calc6d,
583+
rfsw = rfsw,
584+
rad = rad,
585+
radcod = radcod,
586+
radtaper = radtaper,
587+
delta0 = delta0 - ddelta,
588+
additional_commands = additional_commands,
589+
wall_time = wall_time,
590+
sad_path = sad_path)
591+
592+
########################################
593+
# Obtain required data
594+
########################################
595+
betx_on = tw_on.betx
596+
betx_plus = tw_plus.betx
597+
betx_minus = tw_minus.betx
598+
599+
bety_on = tw_on.bety
600+
bety_plus = tw_plus.bety
601+
bety_minus = tw_minus.bety
602+
603+
alfx_on = tw_on.alfx
604+
alfx_plus = tw_plus.alfx
605+
alfx_minus = tw_minus.alfx
606+
607+
alfy_on = tw_on.alfy
608+
alfy_plus = tw_plus.alfy
609+
alfy_minus = tw_minus.alfy
610+
611+
########################################
612+
# Compute chromatic functions
613+
########################################
614+
# See MAD8 physics manual section 6.3
615+
dbetx = (betx_plus - betx_minus) / (2 * ddelta)
616+
dbety = (bety_plus - bety_minus) / (2 * ddelta)
617+
dalfx = (alfx_plus - alfx_minus) / (2 * ddelta)
618+
dalfy = (alfy_plus - alfy_minus) / (2 * ddelta)
619+
620+
bx_chrom = dbetx / betx_on
621+
by_chrom = dbety / bety_on
622+
623+
ax_chrom = dalfx - dbetx * alfx_on / betx_on
624+
ay_chrom = dalfy - dbety * alfy_on / bety_on
625+
626+
wx_chrom = np.sqrt(ax_chrom**2 + bx_chrom**2)
627+
wy_chrom = np.sqrt(ay_chrom**2 + by_chrom**2)
628+
629+
########################################
630+
# Add these values to the twiss
631+
########################################
632+
if sad_twiss is None:
633+
sad_twiss = tw_on
634+
635+
sad_twiss["bx_chrom"] = bx_chrom
636+
sad_twiss["by_chrom"] = by_chrom
637+
sad_twiss["ax_chrom"] = ax_chrom
638+
sad_twiss["ay_chrom"] = ay_chrom
639+
sad_twiss["wx_chrom"] = wx_chrom
640+
sad_twiss["wy_chrom"] = wy_chrom
641+
642+
########################################
643+
# Return the twiss
644+
########################################
645+
return sad_twiss

‎setup.py‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -22,8 +22,8 @@
2222
################################################################################
2323
setup(
2424
name = "sad2xs",
25-
version = "0.1.11",
26-
date = "20-11-2025",
25+
version = "0.1.12",
26+
date = "21-11-2025",
2727
description = "Conversion of SAD lattices to Xtrack format",
2828
long_description = description,
2929
long_description_content_type = "text/markdown",

0 commit comments

Comments
 (0)