Allow toggling of additional simulator behaviour

src/virtac/__main__.py

@@ -1,7 +1,7 @@
-import argparse
 import logging
 import os
 import socket
+from argparse import ArgumentError, ArgumentParser
 from pathlib import Path
 from warnings import warn
 
@@ -21,7 +21,7 @@
 
 def parse_arguments():
     """Parse command line arguments sent to virtac"""
-    parser = argparse.ArgumentParser()
+    parser = ArgumentParser()
     parser.add_argument(
         "ring_mode",
         nargs="?",
@@ -35,6 +35,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",
@@ -57,6 +79,29 @@ def parse_arguments():
     return parser.parse_args()
 
 
+def check_sim_params(args):
+    """Check that we have a valid combination of simulation parameters."""
+    if args.disable_radiation:
+        if args.linopt_function == "linopt6":
+            raise ArgumentError(
+                None,
+                f"Cannot disable radiation when using linopt function: "
+                f"{args.linopt_function}",
+            )
+        if not args.disable_emittance:
+            raise ArgumentError(
+                None,
+                "You cannot calculate emittance with radiation disabled",
+            )
+    else:
+        if args.linopt_function == "linopt2" or args.linopt_function == "linopt4":
+            raise ArgumentError(
+                None,
+                "You must disable radiation to use linopt function: "
+                f"{args.linopt_function}",
+            )
+
+
 def configure_ca():
     """Setup channel access settings for our CA server and for accessing PVs
     from other IOCs. We will be creating a python softioc IOC which automatically
@@ -110,6 +155,7 @@ def main():
     logging.basicConfig(level=log_level, format=LOG_FORMAT)
 
     configure_ca()
+    check_sim_params(args)
 
     # Determine the ring mode
     if args.ring_mode is not None:
@@ -138,7 +184,10 @@ def main():
         DATADIR / ring_mode / "feedback.csv",
         DATADIR / ring_mode / "mirrored.csv",
         DATADIR / ring_mode / "tunefb.csv",
+        args.linopt_function,
         args.disable_emittance,
+        args.disable_chromaticity,
+        args.disable_radiation,
         args.disable_tfb,
     )
 

src/virtac/virtac_server.py

@@ -60,7 +60,10 @@ def __init__(
         feedback_csv: str | None = None,
         mirror_csv: str | None = None,
         tune_csv: str | None = None,
+        linopt_function: str = "linopt6",
         disable_emittance: bool = False,
+        disable_chromaticity: bool = False,
+        disable_radiation: bool = False,
         disable_tunefb: bool = False,
     ):
         """
@@ -79,11 +82,20 @@ def __init__(
             disable_emittance (bool): Whether emittance should be disabled.
             disable_tunefb (bool): Whether tune feedback should be disabled.
         """
+        self._linopt_function: str = linopt_function
         self._disable_emittance: bool = disable_emittance
+        self._disable_chromaticity: bool = disable_chromaticity
+        self._disable_radiation: bool = disable_radiation
+
         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,
+            self._linopt_function,
+            self._disable_emittance,
+            self._disable_chromaticity,
+            self._disable_radiation,
+            self.update_pvs,
         )
         self.lattice.set_default_data_source(pytac.SIM)
         # Holding dictionary for all PVs
@@ -537,19 +549,32 @@ def print_virtac_stats(self, verbosity: int = 0):
             "\t Tune feedbacks is "
             f"{('disabled' if self._disable_tunefb else 'enabled')}"
         )
+        print(f"\t Linear optics function is {self._linopt_function}")
         print(
             "\t Emittance calculations are "
             f"{('disabled' if self._disable_emittance else 'enabled')}"
         )
+        print(
+            "\t Chromaticity calculations are "
+            f"{('disabled' if self._disable_chromaticity else 'enabled')}"
+        )
+        print(
+            "\t Radiation calculations are "
+            f"{('disabled' if self._disable_chromaticity 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)}")