Modify envelope tracker

examples/Envelope/run.sh

@@ -2,15 +2,28 @@
 
 set -x 
 
-python test_env_2d_fodo.py
-python test_env_2d_fodo_speed.py
-python test_env_3d_drift.py
+python test_env_2d_fodo.py --sc 0
+python test_env_2d_fodo.py --sc 1
+python test_env_2d_fodo.py --sc 1 --offset-x 0.001
+python test_env_2d_fodo.py --sc 1 --tilt 45.0
+python test_env_2d_fodo_speed.py --sc 0
+python test_env_2d_fodo_speed.py --sc 1
+python test_env_3d_drift.py --sc 0
+python test_env_3d_drift.py --sc 1
+python test_env_3d_drift.py --sc 1 --rms-y 0.002 --tilt-z 45.0
+python test_env_3d_drift.py --sc 1 --rms-z 0.002 --tilt-x 45.0
 
 cd sns_linac
-python test_sns_linac.py
+python test_sns_linac.py --sc 0
+python test_sns_linac.py --sc 1 --dist kv
+python test_sns_linac.py --sc 1 --dist waterbag
+python test_sns_linac.py --sc 1 --dist gauss
 cd ..
 
 cd sns_ring
-python test_sns_ring.py
-python test_sns_ring_speed.py
+python test_sns_ring.py --sc 0
+python test_sns_ring.py --sc 1
+python test_sns_ring.py --sc 1 --tilt 45.0
+python test_sns_ring_speed.py --sc 0
+python test_sns_ring_speed.py --sc 1
 cd ..

examples/Envelope/sns_linac/test_sns_linac.py

@@ -11,8 +11,10 @@
 import argparse
 import math
 import os
+import pathlib
 import random
 import sys
+import time
 
 import numpy as np
 import matplotlib.pyplot as plt
@@ -30,7 +32,6 @@
 from orbit.bunch_generators import KVDist3D
 from orbit.bunch_utils import collect_bunch
 from orbit.envelope import Envelope
-from orbit.envelope import EnvelopeTracker
 from orbit.lattice import AccLattice
 from orbit.lattice import AccNode
 from orbit.lattice import AccActionsContainer
@@ -66,16 +67,47 @@ def parse_args() -> argparse.Namespace:
     return parser.parse_args()
 
 
-def main(args: argparse.Namespace) -> None:
+def make_lattice(args: argparse.Namespace) -> LinacAccLattice:
+    seq_names = [
+        "MEBT",
+        "DTL1",
+        "DTL2",
+        "DTL3",
+        "DTL4",
+        "DTL5",
+        "DTL6",
+        "CCL1",
+        "CCL2",
+        "CCL3",
+        "CCL4",
+        "SCLMed",
+        "SCLHigh",
+        "HEBT1",
+        "HEBT2",
+    ]
+    if args.seq_stop:
+        index = seq_names.index(args.seq_stop) + 1
+        seq_names = seq_names[:index]
 
-    output_dir = "outputs"
-    os.makedirs(output_dir, exist_ok=True)
+    sns_linac_factory = SNS_LinacLatticeFactory()
+    sns_linac_factory.setMaxDriftLength(args.sc_path_length_min)
+    lattice = sns_linac_factory.getLinacAccLattice(seq_names, "inputs/sns_linac.xml")
 
-    random.seed(23)
+    for node in lattice.getNodes():
+        try:
+            node.setUsageFringeFieldIN(False)
+            node.setUsageFringeFieldOUT(False)
+        except:
+            pass
 
-    # Bunch
-    # --------------------------------------------------------------------------------
+    rf_gaps = lattice.getRF_Gaps()
+    for rf_gap in rf_gaps:
+        rf_gap.setCppGapModel(MatrixRfGap())
 
+    return lattice
+
+
+def make_bunch(args: argparse.Namespace) -> Bunch:
     kin_energy = 0.0025  # [GeV]
     mass = mass_proton + 2.0 * mass_electron
     frequency = 402.5e06
@@ -110,72 +142,34 @@ def main(args: argparse.Namespace) -> None:
 
     for _ in range(args.nparts):
         bunch.addParticle(*dist.getCoordinates())
+    return bunch
 
-    # Lattice
-    # --------------------------------------------------------------------------------
 
-    seq_names = [
-        "MEBT",
-        "DTL1",
-        "DTL2",
-        "DTL3",
-        "DTL4",
-        "DTL5",
-        "DTL6",
-        "CCL1",
-        "CCL2",
-        "CCL3",
-        "CCL4",
-        "SCLMed",
-        "SCLHigh",
-        "HEBT1",
-        "HEBT2",
-    ]
-    if args.seq_stop:
-        index = seq_names.index(args.seq_stop) + 1
-        seq_names = seq_names[:index]
-
-    sns_linac_factory = SNS_LinacLatticeFactory()
-    sns_linac_factory.setMaxDriftLength(args.sc_path_length_min)
-    lattice = sns_linac_factory.getLinacAccLattice(seq_names, "inputs/sns_linac.xml")
-
-    for node in lattice.getNodes():
-        try:
-            node.setUsageFringeFieldIN(False)
-            node.setUsageFringeFieldOUT(False)
-        except:
-            pass
-
-    rf_gaps = lattice.getRF_Gaps()
-    for rf_gap in rf_gaps:
-        rf_gap.setCppGapModel(MatrixRfGap())
+def main(args: argparse.Namespace) -> None:
 
-    for index, node in enumerate(lattice.getNodes()):
-        print(index, type(node), node.getName())
+    path = pathlib.Path(__file__)
+    output_dir = os.path.join("outputs", path.stem, time.strftime("%Y%m%d_%H%M%S"))
+    os.makedirs(output_dir, exist_ok=True)
 
-    lattice.trackDesignBunch(bunch)
+    random.seed(23)
 
     # Track envelope
-    # --------------------------------------------------------------------------------
-
-    twiss_calc = BunchTwissAnalysis()
-    twiss_calc.analyzeBunch(bunch)
-
-    cov_matrix = np.zeros((6, 6))
-    for i in range(6):
-        for j in range(6):
-            cov_matrix[i, j] = cov_matrix[j, i] = twiss_calc.getCorrelation(i, j)
+    bunch = make_bunch(args)
+    envelope = Envelope(bunch=bunch)
 
-    envelope = Envelope(bunch=bunch, cov_matrix=cov_matrix, intensity=intensity)
-
-    tracker = EnvelopeTracker(lattice, sc=("3d" if args.sc else None))
+    lattice = make_lattice(args)
+    lattice.trackDesignBunch(bunch)
 
     histories = {}
-    histories["envelope"] = tracker.track_history(envelope)
+    histories["envelope"] = lattice.trackEnvelope(
+        envelope,
+        history=True,
+        sc=("3d" if args.sc else None)
+    )
 
     # Track bunch
-    # --------------------------------------------------------------------------------
-
+    bunch = make_bunch(args)
+    lattice = make_lattice(args)
     lattice.trackDesignBunch(bunch)
 
     if args.sc:

examples/Envelope/sns_ring/test_sns_ring.py

@@ -5,8 +5,8 @@
 import math
 import os
 import pathlib
-import time
 import sys
+import time
 
 import numpy as np
 import matplotlib.pyplot as plt
@@ -16,12 +16,12 @@
 from orbit.core.spacecharge import SpaceChargeCalc2p5D
 from orbit.bunch_utils import collect_bunch
 from orbit.envelope import Envelope
-from orbit.envelope import EnvelopeTracker
 from orbit.core.spacecharge import SpaceChargeCalc2p5D
 from orbit.space_charge.sc2p5d import setSC2p5DAccNodes
 from orbit.teapot import TEAPOT_Ring
 from orbit.teapot import TEAPOT_MATRIX_Lattice
 from orbit.teapot import teapot
+from orbit.teapot import BendTEAPOT
 from orbit.utils.consts import mass_proton
 
 sys.path.append("..")
@@ -65,7 +65,7 @@ def parse_args() -> argparse.Namespace:
 
 def main(args: argparse.Namespace) -> None:
     path = pathlib.Path(__file__)
-    output_dir = os.path.join("outputs", path.stem)
+    output_dir = os.path.join("outputs", path.stem, time.strftime("%Y%m%d_%H%M%S"))
     os.makedirs(output_dir, exist_ok=True)
 
     # Lattice
@@ -122,7 +122,7 @@ def main(args: argparse.Namespace) -> None:
 
     if args.tilt:
         rot_matrix = np.identity(6)
-        rot_matrix[:4, :4] = build_rotation_matrix_xy(angle=(args.tilt * math.pi))
+        rot_matrix[:4, :4] = build_rotation_matrix_xy(angle=np.radians(args.tilt))
         cov_matrix = np.linalg.multi_dot([rot_matrix, cov_matrix, rot_matrix.T])
 
     if args.mismatch_x or args.mismatch_y:
@@ -157,12 +157,11 @@ def main(args: argparse.Namespace) -> None:
     print("TRACK ENVELOPE")
 
     envelope = Envelope(
-        bunch=bunch,
+        sync_part=sync_part,
         cov_matrix=cov_matrix_init,
         centroid=centroid_init,
         intensity=args.intensity,
     )
-    tracker = EnvelopeTracker(lattice, sc=("2d" if args.sc else None))
 
     history_keys = [
         "rms_x",
@@ -178,7 +177,7 @@ def main(args: argparse.Namespace) -> None:
 
     for turn in range(args.turns + 1):
         if turn > 0:
-            tracker.track_ring(envelope)
+            lattice.trackEnvelopeRing(envelope, sc=("2d" if args.sc else None))
 
         cov_matrix = envelope.cov_matrix
         centroid = envelope.centroid

examples/Envelope/sns_ring/test_sns_ring_speed.py

@@ -12,16 +12,17 @@
 from orbit.core.bunch import Bunch
 from orbit.core.spacecharge import SpaceChargeCalc2p5D
 from orbit.envelope import Envelope
-from orbit.envelope import EnvelopeTracker
 from orbit.core.spacecharge import SpaceChargeCalc2p5D
 from orbit.space_charge.sc2p5d import setSC2p5DAccNodes
 from orbit.teapot import TEAPOT_Ring
 from orbit.teapot import TEAPOT_MATRIX_Lattice
+from orbit.teapot import BendTEAPOT
 from orbit.utils.consts import mass_proton
 
 sys.path.append("..")
 from utils import gen_dist
 
+
 parser = argparse.ArgumentParser()
 parser.add_argument("--bunch-length", type=float, default=120.0)
 parser.add_argument("--kin-energy", type=float, default=1.300)
@@ -45,6 +46,11 @@
     except:
         pass
 
+for node in lattice.getNodes():
+    if type(node) is BendTEAPOT:
+        node.setParam("ea1", 0.0)
+        node.setParam("ea2", 0.0)
+
 for node in lattice.getNodes():
     max_length = 1.0
     if node.getLength() > max_length:
@@ -80,19 +86,19 @@
 print("ENVELOPE")
 
 envelope = Envelope(
-    bunch=bunch,
+    sync_part=sync_part,
     cov_matrix=cov_matrix_init,
     intensity=args.intensity,
 )
-tracker = EnvelopeTracker(lattice, sc=("2d" if args.sc else None))
+envelope_sc = "2d" if args.sc else None
 
 start_time = time.time()
 
 profiler = cProfile.Profile()
 profiler.enable()
 
 for turn in trange(args.turns):
-    tracker.track_ring(envelope)
+    lattice.trackEnvelopeRing(envelope, sc=envelope_sc)
 
 time_per_turn = (time.time() - start_time) / args.turns
 

examples/Envelope/test_env_2d_fodo.py

@@ -5,6 +5,7 @@
 import math
 import os
 import pathlib
+import time
 
 import numpy as np
 import matplotlib.pyplot as plt
@@ -14,7 +15,6 @@
 from orbit.core.spacecharge import SpaceChargeCalc2p5D
 from orbit.bunch_utils import collect_bunch
 from orbit.envelope import Envelope
-from orbit.envelope import EnvelopeTracker
 from orbit.lattice import AccLattice
 from orbit.lattice import AccNode
 from orbit.core.spacecharge import SpaceChargeCalc2p5D
@@ -59,12 +59,8 @@ def parse_args() -> argparse.Namespace:
 
 
 def main(args: argparse.Namespace) -> None:
-
-    # Setup
-    # ------------------------------------------------------------------------------
-
     path = pathlib.Path(__file__)
-    output_dir = os.path.join("outputs", path.stem)
+    output_dir = os.path.join("outputs", path.stem, time.strftime("%Y%m%d_%H%M%S"))
     os.makedirs(output_dir, exist_ok=True)
 
     # Create lattice
@@ -120,7 +116,7 @@ def main(args: argparse.Namespace) -> None:
     # Tilt
     if args.tilt:
         rot_matrix = np.identity(6)
-        rot_matrix[:4, :4] = build_rotation_matrix_xy(angle=(args.tilt * math.pi))
+        rot_matrix[:4, :4] = build_rotation_matrix_xy(angle=np.radians(args.tilt))
         cov_matrix = np.linalg.multi_dot([rot_matrix, cov_matrix, rot_matrix.T])
 
     # Mismatch
@@ -135,7 +131,7 @@ def main(args: argparse.Namespace) -> None:
 
     # Create envelope
     envelope = Envelope(
-        bunch=bunch,
+        sync_part=sync_part,
         cov_matrix=cov_matrix_init,
         centroid=centroid_init,
         intensity=args.intensity,
@@ -146,12 +142,12 @@ def main(args: argparse.Namespace) -> None:
 
     print("TRACK ENVELOPE")
 
-    tracker = EnvelopeTracker(lattice, sc=("2d" if args.sc else None))
+    envelope_sc = "2d" if args.sc else None
 
     history = {"xrms": [], "yrms": [], "xavg": [], "yavg": []}
     for turn in range(args.turns):
         if turn > 0:
-            tracker.track_ring(envelope)
+            lattice.trackEnvelopeRing(envelope, sc=envelope_sc)
 
         cov_matrix = envelope.cov_matrix
         centroid = envelope.centroid