diff --git a/src/virtac/__main__.py b/src/virtac/__main__.py index 0a30ddb..2eb825e 100644 --- a/src/virtac/__main__.py +++ b/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 typing import cast 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() -> None: 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() -> None: 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, ) diff --git a/src/virtac/virtac_server.py b/src/virtac/virtac_server.py index 9915e0a..0807b9d 100644 --- a/src/virtac/virtac_server.py +++ b/src/virtac/virtac_server.py @@ -57,12 +57,15 @@ class VirtacServer: def __init__( self, ring_mode: str, - limits_csv: Path | None = None, - bba_csv: Path | None = None, - feedback_csv: Path | None = None, - mirror_csv: Path | None = None, - tune_csv: Path | None = None, + limits_csv: str, + bba_csv: str | None = None, + 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, ) -> None: """ @@ -81,11 +84,20 @@ def __init__( disable_emittance: Whether emittance should be disabled. disable_tunefb: 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 @@ -515,19 +527,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._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)}")