Save progress

Author: austin-hoover

py/orbit/envelope/danilov_env_solver_lattice_modifications.py

@@ -5,8 +5,8 @@
 from ..lattice import AccNode
 from ..lattice import AccNodeBunchTracker
 
-from .danilov_envelope_solver_nodes import DanilovEnvelopeSolverNode20
-from .danilov_envelope_solver_nodes import DanilovEnvelopeSolverNode22
+from .danilov_env_solver_nodes import Danilov20EnvSolverNode
+from .danilov_env_solver_nodes import Danilov22EnvSolverNode
 
 
 class Parent:
@@ -30,9 +30,9 @@ def add_danilov_envolope_solver_nodes(
     lattice: AccLattice,
     path_length_max: float,
     path_length_min: float,
-    solver_node_constructor: DanilovEnvelopeSolverNode20 | DanilovEnvelopeSolverNode22,
+    solver_node_constructor: Danilov20EnvSolverNode | Danilov22EnvSolverNode,
     solver_node_constructor_kwargs: dict,
-) -> list[DanilovEnvelopeSolverNode20 | DanilovEnvelopeSolverNode22]:
+) -> list[Danilov20EnvSolverNode | Danilov22EnvSolverNode]:
 
     nodes = lattice.getNodes()
     if not nodes:
@@ -102,7 +102,7 @@ def add_danilov_envelope_solver_nodes_20(
         lattice=lattice,
         path_length_max=path_length_max,
         path_length_min=path_length_min,
-        solver_node_constructor=DanilovEnvelopeSolverNode20,
+        solver_node_constructor=Danilov20EnvSolverNode,
         solver_node_constructor_kwargs=kwargs,
     )
     for solver_node in solver_nodes:

py/orbit/envelope/danilov_env_solver_nodes.py

@@ -1,33 +1,36 @@
+from typing import Any
+
 from ..lattice import AccActionsContainer
 from ..lattice import AccLattice
 from ..lattice import AccNode
 from ..lattice import AccNodeBunchTracker
 from ..utils import orbitFinalize
 
-from orbit.core.envelope import Danilov20EnvelopeSolver
-from orbit.core.envelope import Danilov22EnvelopeSolver
+from orbit.core.envelope import EnvSolverDanilov20
+from orbit.core.envelope import EnvSolverDanilov22
 
 
-class DanilovEnvelopeSolverNode(AccNodeBunchTracker):
+class EnvSolverNode(AccNodeBunchTracker):
     def __init__(
         self,
-        solver: DanilovEnvelopeSolver20 | DanilovEnvelopeSolver22,
-        name: str = None,
+        solver: Any,
+        name: str = "",
         kick_length: float = 0.0,
-        perveance: float = 0.0,
     ) -> None:
         super().__init__(name=name)
-        self.setType("DanilovEnvSolver")
+        self.setType("EnvSolver")
         self.setLength(0.0)
 
-        self.solver = solver
+        self.solver = None
         self.kick_length = kick_length
-        self.perveance = perveance
         self.active = True
 
-    def set_active(self, setting: bool) -> None:
+    def setActive(self, setting: bool) -> None:
         self.active = setting
 
+    def setKickLength(self, kick_length: float) -> None:
+        self.kick_length = kick_length
+
     def isRFGap(self) -> bool:
         # In case this node is used in linac tracking
         return False
@@ -42,28 +45,24 @@ def track(self, params_dict: dict) -> None:
         bunch = params_dict["bunch"]
         self.solver.trackBunch(bunch, self.kick_length)
 
-    def set_perveance(self, perveance: float) -> None:
-        self.solver.setPerveance(perveance)
 
-    def set_kick_length(self, kick_length: float) -> None:
-        self.kick_length = kick_length
+class Danilov22EnvSolverNode(EnvSolverNode):
+    def __init__(self, perveance: float, **kwargs) -> None:
+        super().__init__(**kwargs)
+        self.solver = EnvSolverDanilov22(perveance)
 
+    def setPerveance(self, perveance: float) -> None:
+        self.solver.setPerveance(perveance)
 
-class DanilovEnvelopeSolverNode20(DanilovEnvelopeSolverNode):
-    def __init__(self, eps_x: float, eps_y: float, **kwargs) -> None:
-        super().__init__(solver=None, **kwargs)
-        self.eps_x = eps_x
-        self.eps_y = eps_y
-        self.solver = DanilovEnvelopeSolver20(self.perveance, self.eps_x, self.eps_y)
 
-    def set_emittances(self, eps_x: float, eps_y: float) -> None:
-        self.eps_x = eps_x
-        self.eps_y = eps_y
-        self.solver.setEmittanceX(eps_x)
-        self.solver.setEmittanceY(eps_y)
+class Danilov20EnvSolverNode(EnvSolverNode):
+    def __init__(self, eps_x: float, eps_y: float, perveance: float, **kwargs) -> None:
+        super().__init__(**kwargs)
+        self.solver = EnvSolverDanilov20(perveance, eps_x, eps_y)
 
+    def setPerveance(self, perveance: float) -> None:
+        self.solver.setPerveance(perveance)
 
-class DanilovEnvelopeSolverNode22(DanilovEnvelopeSolverNode):
-    def __init__(self, **kwargs) -> None:
-        super().__init__(solver=None, **kwargs)
-        self.solver = DanilovEnvelopeSolver22(self.perveance)
+    def setEmittances(self, eps_x: float, eps_y: float) -> None:
+        self.solver.setEmittanceX(eps_x)
+        self.solver.setEmittanceY(eps_y)