diff --git a/phases/full_polarization.py b/phases/full_polarization.py index 48e1b96..f0d8ba8 100644 --- a/phases/full_polarization.py +++ b/phases/full_polarization.py @@ -131,24 +131,42 @@ class UndRadial(UndShiftRadialBase): class UndShiftQuadrants(Adjustable): - def __init__(self, ID, accuracy=0.001): + + def __init__(self, ID, which, accuracy=0.001): ID += "SHIFT" - super().__init__(ID) - names = ["TL", "BL", "TR", "BR"] + + which = which.upper() + + opposites = { + "TL": "BR", + "BL": "TR" + } + + inverted_opposites = {v: k for k, v in opposites.items()} + opposites.update(inverted_opposites) + all_names = opposites.keys() + opposite = opposites[which] + + self.names = names = [which, opposite] + for n in all_names: + if n not in names: + names.append(n) + pvnames = [f"{ID}-{n}" for n in names] adjs = {n.lower(): PVAdjustable(pvn + "-SET", pvn, accuracy=accuracy) for n, pvn in zip(names, pvnames)} self.adjs = SimpleDevice(ID, **adjs) + self.which_adj = which_adj = adjs[which.lower()] + super().__init__(F"{ID}-{which}", units=which_adj.units) + self.pv_on = PV(ID + "-ON") self.pv_go = PV(ID + "-GO") def set_target_value(self, value): - tl = value - br = -tl - bl = tr = bl = 0 - vals = (tl, bl, tr, br) - tasks = [a.set_target_value(v) for a, v in zip(self.adjs, vals)] + names = (n.lower() for n in self.names) + vals = (value, -value, 0, 0) + tasks = [adjs[n].set_target_value(v) for n, v in zip(names, vals)] sleep(0.3) self.pv_on.put(1) @@ -160,7 +178,7 @@ class UndShiftQuadrants(Adjustable): def get_current_value(self): - return self.adjs.tl.get_current_value() + return self.which_adj.get_current_value() def is_moving(self): return any(a.is_moving() for a in self.adjs)