From 695697064e86fb566164c75d3d98c990b072fc15 Mon Sep 17 00:00:00 2001 From: Axel Huebl Date: Tue, 24 Sep 2024 11:58:54 -0700 Subject: [PATCH] Examples: Update Programmable Signature For cycle argument --- examples/pytorch_surrogate_model/run_ml_surrogate_15_stage.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/examples/pytorch_surrogate_model/run_ml_surrogate_15_stage.py b/examples/pytorch_surrogate_model/run_ml_surrogate_15_stage.py index 94c205c80..a0ac965b2 100644 --- a/examples/pytorch_surrogate_model/run_ml_surrogate_15_stage.py +++ b/examples/pytorch_surrogate_model/run_ml_surrogate_15_stage.py @@ -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 @@ -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}"]