Allow toggling of additional simulator behaviour

Allow selecting linopt function, disabling chromaticity,
disabling radiation

Author: ptsOSL

src/atip/load_sim.py

@@ -37,7 +37,15 @@ def load_from_filepath(
     return load(pytac_lattice, at_lattice, callback, disable_emittance)
 
 
-def load(pytac_lattice, at_lattice, callback=None, disable_emittance=False):
+def load(
+    pytac_lattice,
+    at_lattice,
+    linopt_function="linopt6",
+    disable_emittance=False,
+    disable_chromaticity=False,
+    disable_radiation=False,
+    callback=None,
+):
     """Load simulator data sources onto the lattice and its elements.
 
     Args:
@@ -58,7 +66,14 @@ def load(pytac_lattice, at_lattice, callback=None, disable_emittance=False):
             f"(AT:{len(at_lattice)} Pytac:{len(pytac_lattice)})."
         )
     # Initialise an instance of the ATSimulator Object.
-    atsim = ATSimulator(at_lattice, callback, disable_emittance)
+    atsim = ATSimulator(
+        at_lattice,
+        linopt_function,
+        disable_emittance,
+        disable_chromaticity,
+        disable_radiation,
+        callback,
+    )
     # Set the simulator data source on the Pytac lattice.
     pytac_lattice.set_data_source(ATLatticeDataSource(atsim), pytac.SIM)
     # Load the sim onto each element.

src/atip/simulator.py

@@ -24,7 +24,10 @@ class LatticeData:
 def calculate_optics(
     at_lattice: at.lattice_object.Lattice,
     refpts: ArrayLike,
+    linopt_function: str = "linopt6",
     disable_emittance: bool = False,
+    disable_chromaticity: bool = False,
+    disable_radiation: bool = False,
 ) -> LatticeData:
     """Perform the physics calculations on the lattice.
 
@@ -43,21 +46,61 @@ def calculate_optics(
         LatticeData: The calculated lattice data.
     """
     logging.debug("Starting physics calculations.")
-
-    orbit0, _ = at_lattice.find_orbit6()
-    logging.debug("Completed orbit calculation.")
-
-    _, beamdata, twiss = at_lattice.linopt6(
-        refpts=refpts, get_chrom=True, orbit=orbit0, keep_lattice=True
+    logging.debug(
+        f"Using simulation params: {linopt_function}, disable_emittance="
+        f"{disable_emittance}, disable_chromaticity={disable_chromaticity}, "
+        f"disable_radiation={disable_radiation}"
     )
+    if linopt_function == "linopt6":
+        orbit0, _ = at_lattice.find_orbit6()
+        logging.debug("Completed orbit calculation.")
+
+        _, beamdata, twiss = at_lattice.linopt6(
+            refpts=refpts,
+            get_chrom=not disable_chromaticity,
+            orbit=orbit0,
+            keep_lattice=True,
+        )
+    elif linopt_function == "linopt4":
+        orbit0, _ = at_lattice.find_orbit4()
+        logging.debug("Completed orbit calculation.")
+
+        _, beamdata, twiss = at_lattice.linopt6(
+            refpts=refpts,
+            get_chrom=not disable_chromaticity,
+            orbit=orbit0,
+            keep_lattice=True,
+        )
+    elif linopt_function == "linopt2":
+        orbit0, _ = at_lattice.find_orbit()
+        logging.debug("Completed orbit calculation.")
+
+        _, beamdata, twiss = at_lattice.linopt2(
+            refpts=refpts,
+            get_chrom=not disable_chromaticity,
+            orbit=orbit0,
+            keep_lattice=True,
+        )
+    else:
+        logging.error(
+            f"Error. Invalid linopt function selected: {linopt_function}. Simulation "
+            "not computed."
+        )
+
     logging.debug("Completed linear optics calculation.")
 
     if not disable_emittance:
         emitdata = at_lattice.ohmi_envelope(orbit=orbit0, keep_lattice=True)
         logging.debug("Completed emittance calculation")
     else:
         emitdata = ()
-    radint = at_lattice.get_radiation_integrals(twiss=twiss)
+
+    if not disable_radiation:
+        radint = at_lattice.get_radiation_integrals(twiss=twiss)
+        logging.debug("Completed radiation calculation")
+    else:
+        radint = ()
+
     logging.debug("All calculation complete.")
     return LatticeData(twiss, beamdata.tune, beamdata.chromaticity, emitdata, radint)
 
@@ -98,7 +141,15 @@ class ATSimulator:
                                                     physics data upon a change.
     """
 
-    def __init__(self, at_lattice, callback=None, disable_emittance=False):
+    def __init__(
+        self,
+        at_lattice,
+        linopt_function="linopt6",
+        disable_emittance=False,
+        disable_chromaticity=False,
+        disable_radiation=False,
+        callback=None,
+    ):
         """
         .. Note:: To avoid errors, the physics data must be initially
            calculated here, during creation, otherwise it could be accidentally
@@ -122,12 +173,22 @@ def __init__(self, at_lattice, callback=None, disable_emittance=False):
             )
         self._at_lat = at_lattice
         self._rp = numpy.ones(len(at_lattice) + 1, dtype=bool)
+        self._linopt_function = linopt_function
         self._disable_emittance = disable_emittance
-        self._at_lat.radiation_on()
+        self._disable_chromaticity = disable_chromaticity
+        self._disable_radiation = disable_radiation
+
+        if not self._disable_radiation:
+            self._at_lat.radiation_on()
 
         # Initial phys data calculation.
         self._lattice_data = calculate_optics(
-            self._at_lat, self._rp, self._disable_emittance
+            self._at_lat,
+            self._rp,
+            self._linopt_function,
+            self._disable_emittance,
+            self._disable_chromaticity,
+            self._disable_radiation,
         )
 
         # Threading stuff initialisation.
@@ -196,7 +257,12 @@ def _recalculate_phys_data(self, callback):
             if bool(self._paused) is False:
                 try:
                     self._lattice_data = calculate_optics(
-                        self._at_lat, self._rp, self._disable_emittance
+                        self._at_lat,
+                        self._rp,
+                        self._linopt_function,
+                        self._disable_emittance,
+                        self._disable_chromaticity,
+                        self._disable_radiation,
                     )
                 except Exception as e:
                     warn(at.AtWarning(e), stacklevel=1)

src/atip/utils.py

@@ -29,7 +29,14 @@ def load_at_lattice(mode="I04", **kwargs):
     return at_lattice
 
 
-def loader(mode="I04", callback=None, disable_emittance=False):
+def loader(
+    mode="I04",
+    linopt_function="linopt6",
+    disable_emittance=False,
+    disable_chromaticity=False,
+    disable_radiation=False,
+    callback=None,
+):
     """Load a unified lattice of the specifed mode.
 
     .. Note:: A unified lattice is a Pytac lattice where the corresponding AT
@@ -52,7 +59,15 @@ def loader(mode="I04", callback=None, disable_emittance=False):
         periodicity=1,
         energy=pytac_lattice.get_value("energy", units=pytac.PHYS),
     )
-    lattice = atip.load_sim.load(pytac_lattice, at_lattice, callback, disable_emittance)
+    lattice = atip.load_sim.load(
+        pytac_lattice,
+        at_lattice,
+        linopt_function,
+        disable_emittance,
+        disable_chromaticity,
+        disable_radiation,
+        callback,
+    )
     return lattice