|
7 | 7 | ################################################################################ |
8 | 8 | import os |
9 | 9 | import subprocess |
| 10 | +import ast |
10 | 11 | import numpy as np |
11 | 12 | import tfs |
12 | 13 | import xtrack as xt |
@@ -203,6 +204,7 @@ def twiss_sad( |
203 | 204 | radiation: bool = False, |
204 | 205 | rad_compensation: bool = False, |
205 | 206 | rad_taper: bool = False, |
| 207 | + delta0: float = 0.0, |
206 | 208 | additional_commands: str = ''): |
207 | 209 | """ |
208 | 210 | Generate a SAD command to compute the twiss parameters of a lattice. |
@@ -257,7 +259,7 @@ def twiss_sad( |
257 | 259 | USE {line_name}; |
258 | 260 |
|
259 | 261 | {additional_commands}; |
260 | | -
|
| 262 | +DP0 = {delta0}; |
261 | 263 | {RF_FLAG} |
262 | 264 | {RAD_FLAG} |
263 | 265 | {RADCOMP_FLAG} |
@@ -418,6 +420,89 @@ def twiss_sad( |
418 | 420 | ######################################## |
419 | 421 | return tw_sad |
420 | 422 |
|
| 423 | +################################################################################ |
| 424 | +# Calculate Transfer Matrix |
| 425 | +################################################################################ |
| 426 | +def sad_transfer_matrix( |
| 427 | + lattice_filename: str, |
| 428 | + line_name: str, |
| 429 | + start_element: str | None = None, |
| 430 | + end_element: str | None = None): |
| 431 | + |
| 432 | + ######################################## |
| 433 | + # Ensure both or neither of start_element and end_element are provided |
| 434 | + ######################################## |
| 435 | + if start_element is not None and end_element is None: |
| 436 | + raise ValueError("If start_element is provided, end_element must also be provided") |
| 437 | + if start_element is None and end_element is not None: |
| 438 | + raise ValueError("If end_element is provided, start_element must also be provided") |
| 439 | + |
| 440 | + ######################################## |
| 441 | + # Generate the Transfer Matrix command |
| 442 | + ######################################## |
| 443 | + if start_element is not None and end_element is not None: |
| 444 | + SAD_COMMAND = f""" |
| 445 | +FFS; |
| 446 | +
|
| 447 | +GetMAIN["./{lattice_filename}"]; |
| 448 | +USE {line_name}; |
| 449 | +
|
| 450 | +CALC4D; |
| 451 | +CALC; |
| 452 | +TM = TransferMatrix["{start_element}", "{end_element}"]; |
| 453 | +Print[TM]; |
| 454 | +abort; |
| 455 | +""" |
| 456 | + else: |
| 457 | + SAD_COMMAND = f""" |
| 458 | +FFS; |
| 459 | +
|
| 460 | +GetMAIN["./{lattice_filename}"]; |
| 461 | +USE {line_name}; |
| 462 | +
|
| 463 | +CALC4D; |
| 464 | +CALC; |
| 465 | +TM = TransferMatrix[1, -1]; |
| 466 | +Print[TM]; |
| 467 | +abort; |
| 468 | +""" |
| 469 | + |
| 470 | + ######################################## |
| 471 | + # Write and execute the SAD command |
| 472 | + ######################################## |
| 473 | + with open("temporary_sad_transfer_matrix.sad", "w") as f: |
| 474 | + f.write(SAD_COMMAND) |
| 475 | + |
| 476 | + try: |
| 477 | + result = subprocess.run( |
| 478 | + ["sad", "temporary_sad_transfer_matrix.sad"], |
| 479 | + capture_output = True, |
| 480 | + text = True, |
| 481 | + timeout = 30) |
| 482 | + os.remove("temporary_sad_transfer_matrix.sad") |
| 483 | + except subprocess.TimeoutExpired: |
| 484 | + print("Timed out at 30s") |
| 485 | + |
| 486 | + os.remove("temporary_sad_transfer_matrix.sad") |
| 487 | + raise subprocess.TimeoutExpired(cmd = ["sad", "temporary_sad_transfer_matrix.sad"], timeout = 30) |
| 488 | + |
| 489 | + # Take the raw text output |
| 490 | + raw = result.stdout |
| 491 | + |
| 492 | + # Isolate the matrix |
| 493 | + start = raw.find("{{") |
| 494 | + end = raw.rfind("}}") |
| 495 | + if start == -1 or end == -1: |
| 496 | + raise ValueError("Matrix not found in subprocess output") |
| 497 | + matrix_str = raw[start:end + 2] |
| 498 | + |
| 499 | + # Convert to numpy array |
| 500 | + cleaned = matrix_str.replace("}", "]").replace("{", "[") |
| 501 | + matrix_list = ast.literal_eval(cleaned) |
| 502 | + R_matrix = np.array(matrix_list, dtype = float) |
| 503 | + |
| 504 | + return R_matrix |
| 505 | + |
421 | 506 | ################################################################################ |
422 | 507 | # Rebuild SAD lattice (post GEO for solenoid) |
423 | 508 | ################################################################################ |
|
0 commit comments