Skip to content

Commit

Permalink
Examples: Update Programmable Signature
Browse files Browse the repository at this point in the history
For cycle argument
  • Loading branch information
ax3l committed Sep 24, 2024
1 parent a3549e1 commit 6956970
Showing 1 changed file with 2 additions and 2 deletions.
4 changes: 2 additions & 2 deletions examples/pytorch_surrogate_model/run_ml_surrogate_15_stage.py
Original file line number Diff line number Diff line change
Expand Up @@ -169,7 +169,7 @@ def __init__(self, stage_i, surrogate_model, surrogate_length, stage_start):
self.push = self.surrogate_push
self.ds = surrogate_length

def surrogate_push(self, pc, step):
def surrogate_push(self, pc, step, cycle):
ref_part = pc.ref_particle()
ref_z_i = ref_part.z
ref_z_i_LPA = ref_z_i - self.stage_start
Expand Down Expand Up @@ -304,7 +304,7 @@ def __init__(self, sim, stage_i, lattice_index, x_or_y):
self.x_or_y = x_or_y
self.push = self.set_lens

def set_lens(self, pc, step):
def set_lens(self, pc, step, cycle):
# get envelope parameters
rbc = pc.reduced_beam_characteristics()
alpha = rbc[f"alpha_{self.x_or_y}"]
Expand Down

0 comments on commit 6956970

Please sign in to comment.