Hla 853 toggleable simulation parameters

pyproject.toml

@@ -15,10 +15,10 @@ description = "Diamond virtual accelerator"
 dependencies = [
     "numpy",
     "scipy",
-    "pytac==0.6.0",
+    "pytac>=1.1.0",
     "cothread",
     "softioc",
-    "atip>=0.2.0",
+    "atip>=0.3.0",
     "epicscorelibs",
 ]
 

src/virtac/__main__.py

@@ -1,11 +1,12 @@
-import argparse
 import logging
 import os
 import socket
+from argparse import ArgumentParser
 from pathlib import Path
 from typing import cast
 from warnings import warn
 
+from atip.simulator import SimParams
 from cothread.catools import ca_nothing, caget
 from softioc import builder, softioc
 
@@ -21,7 +22,7 @@
 
 def parse_arguments():
     """Parse command line arguments sent to virtac"""
-    parser = argparse.ArgumentParser()
+    parser = ArgumentParser()
     parser.add_argument(
         "ring_mode",
         nargs="?",
@@ -35,6 +36,28 @@ def parse_arguments():
         action="store_true",
         default=False,
     )
+    parser.add_argument(
+        "-c",
+        "--disable-chromaticity",
+        help="Disable chromaticity calculations",
+        action="store_true",
+        default=False,
+    )
+    parser.add_argument(
+        "-r",
+        "--disable-radiation",
+        help="Disable radiation calculations in the simulation",
+        action="store_true",
+        default=False,
+    )
+    parser.add_argument(
+        "-l",
+        "--linopt-function",
+        help="Which pyAT linear optics function to use: linopt2, linopt4, linopt6. "
+        "Default is linopt6",
+        default="linopt6",
+        type=str,
+    )
     parser.add_argument(
         "-t",
         "--disable-tfb",
@@ -111,6 +134,13 @@ def main() -> None:
 
     configure_ca()
 
+    sim_params = SimParams(
+        args.linopt_function,
+        not args.disable_emittance,
+        not args.disable_chromaticity,
+        not args.disable_radiation,
+    )
+
     # Determine the ring mode
     if args.ring_mode is not None:
         ring_mode = args.ring_mode
@@ -138,7 +168,7 @@ def main() -> None:
         DATADIR / ring_mode / "feedback.csv",
         DATADIR / ring_mode / "mirrored.csv",
         DATADIR / ring_mode / "tunefb.csv",
-        args.disable_emittance,
+        sim_params,
         args.disable_tfb,
     )
 

src/virtac/virtac_server.py

@@ -62,7 +62,7 @@ def __init__(
         feedback_csv: Path | None = None,
         mirror_csv: Path | None = None,
         tune_csv: Path | None = None,
-        disable_emittance: bool = False,
+        sim_params: atip.simulator.SimParams | None = None,
         disable_tunefb: bool = False,
     ) -> None:
         """
@@ -78,19 +78,27 @@ def __init__(
                 mirror records, for more information see create_csv.py.
             tune_csv: The filepath to the .csv file from which to load the tune
                 feedback records, for more information see create_csv.py.
-            disable_emittance: Whether emittance should be disabled.
+            sim_params (SimParams | None): An optional dataclass containing the pyAT
+                simulation parameters to use.
             disable_tunefb: Whether tune feedback should be disabled.
         """
-        self._disable_emittance: bool = disable_emittance
+
+        if sim_params is None:
+            sim_params = atip.simulator.SimParams()
+        self._sim_params: atip.simulator.SimParams = sim_params
         self._disable_tunefb: bool = disable_tunefb
         self._pv_monitoring: bool = True
+
         self.lattice: pytac.lattice.EpicsLattice = atip.utils.loader(
-            ring_mode, self.update_pvs, self._disable_emittance
+            ring_mode,
+            sim_params,
+            self.update_pvs,
         )
         self.lattice.set_default_data_source(pytac.SIM)
+
         # Holding dictionary for all PVs
         self._pv_dict: dict[str, BasePV] = {}
-        # Dictionary for the PVs which should be automatically updated when the
+        # Dictionary for the PVs which need to be automatically updated when the
         # simulation data is recalculated
         self._readback_pvs_dict: dict[str, ReadSimPV] = {}
 
@@ -262,7 +270,7 @@ def _create_lattice_pvs(self, limits_dict: LimitsDictType) -> None:
         """
         lat_field_dict = cast(dict[str, list[str]], self.lattice.get_fields())
         lat_field_set = set(lat_field_dict[pytac.LIVE]) & set(lat_field_dict[pytac.SIM])
-        if self._disable_emittance:
+        if not self._sim_params.emittance:
             lat_field_set -= {"emittance_x", "emittance_y"}
         for field in lat_field_set:
             # Ignore basic devices as they do not have PVs.
@@ -313,7 +321,7 @@ def _create_feedback_records(self, feedback_csv: Path) -> None:
 
         # We can choose to not calculate emittance as it is not always required,
         # which decreases computation time.
-        if not self._disable_emittance:
+        if self._sim_params.emittance:
             name = "SR-DI-EMIT-01:STATUS"
             record_data = RecordData(RecordTypes.MBBI, zrvl="0", zrst="Successful")
             emit_status_pv = BasePV(name, record_data)
@@ -515,19 +523,32 @@ def print_virtac_stats(self, verbosity: int = 0) -> None:
             "\t Tune feedbacks is "
             f"{('disabled' if self._disable_tunefb else 'enabled')}"
         )
+        print(f"\t Linear optics function is {self._sim_params.linopt}")
         print(
             "\t Emittance calculations are "
-            f"{('disabled' if self._disable_emittance else 'enabled')}"
+            f"{('disabled' if not self._sim_params.emittance else 'enabled')}"
+        )
+        print(
+            "\t Chromaticity calculations are "
+            f"{('disabled' if not self._sim_params.chromaticity else 'enabled')}"
+        )
+        print(
+            "\t Radiation calculations are "
+            f"{('disabled' if not self._sim_params.radiation else 'enabled')}"
         )
         print(
             f"\t PV monitoring is {('enabled' if self._pv_monitoring else 'disabled')}"
         )
 
-        print(f"\t Total pvs: {len(self._pv_dict)}")
+        print(f"\t Total pvs: {len(self._pv_dict)}, consisting of:")
         for pv_type, count in pv_type_count.items():
             print(f"\t\t {pv_type.__name__} pvs: {count}")
+        print(
+            "\t Number of PVs to update after simulation recalculation: "
+            f"{len(self._readback_pvs_dict)}"
+        )
 
         if verbosity >= 1:
-            print("\tAvailable PVs")
+            print("\t Available PVs")
             for pv in self._pv_dict.values():
-                print(f"\t\t{pv.name}, {type(pv)}")
+                print(f"\t\t {pv.name}, {type(pv)}")