Merge pull request #7 from catheybl/main

Virac Parser

virtaccl/EPICS_Server/ca_server.py

@@ -1,15 +1,29 @@
+import sys
 from threading import Thread
 from datetime import datetime
 from time import sleep
 
 from math import floor
-from typing import Any
+from typing import Any, Dict
 
 from pcaspy import Driver
 from pcaspy.cas import epicsTimeStamp
 from pcaspy import SimpleServer
 
 from virtaccl.server import Server
+from virtaccl.virtual_accelerator import VA_Parser
+
+
+def add_epics_arguments(va_parser: VA_Parser) -> VA_Parser:
+    # Number (in seconds) that determine some delay parameter in the server. Not exactly sure how it works, so use at
+    # your own risk.
+    va_parser.add_server_argument('--ca_proc', default=0.1, type=float,
+                                  help='Number (in seconds) that determine some delay parameter in the server. Not '
+                                       'exactly sure how it works, so use at your own risk.')
+
+    va_parser.add_server_argument('--print_pvs', dest='print_pvs', action='store_true',
+                                  help="Will print all server PVs. Will NOT run the virtual accelerator.")
+    return va_parser
 
 
 def to_epics_timestamp(t: datetime):
@@ -39,16 +53,24 @@ def setParam(self, reason, value, timestamp=None):
 
 
 class EPICS_Server(Server):
-    def __init__(self, prefix='', process_delay=0.1):
+    def __init__(self, prefix='', process_delay=0.1, print_pvs=False):
         super().__init__()
         self.prefix = prefix
         self.driver = None
         self.process_delay = process_delay
+        self.print_pvs = print_pvs
 
     def _CA_events(self, server):
         while True:
             server.process(self.process_delay)
 
+    def add_parameters(self, new_parameters: Dict[str, Dict[str, Any]]):
+        super().add_parameters(new_parameters)
+        if self.print_pvs:
+            for key in self.get_parameter_keys():
+                print(key)
+            sys.exit()
+
     def set_parameter(self, reason: str, value: Any, timestamp: datetime = None):
         if timestamp is not None:
             timestamp = epics_now(timestamp)

virtaccl/PyORBIT_Model/pyorbit_va_arguments.py

@@ -0,0 +1,26 @@
+from virtaccl.virtual_accelerator import VA_Parser
+
+
+def add_pyorbit_arguments(va_parser: VA_Parser) -> VA_Parser:
+    # Lattice xml input file and the sequences desired from that file.
+    va_parser.add_model_argument('--lattice', type=str, help='Pathname of lattice file.')
+    va_parser.add_model_argument("--start", default="MEBT", type=str,
+                                 help='Desired sequence of the lattice to start the model with.')
+    va_parser.add_model_argument("end", nargs='?', type=str,
+                                 help='Desired sequence of the lattice to end the model with.')
+    va_parser.add_model_argument('--space_charge', const=0.01, nargs='?', type=float,
+                                 help="Adds Uniform Ellipse Space Charge nodes to the lattice. The minimum distance "
+                                      "in meters between nodes can be specified; the default is 0.01m if no minimum "
+                                      "is given. If the argument is not used, no space charge nodes are added.")
+
+    # Desired initial bunch file and the desired number of particles from that file.
+    va_parser.add_model_argument('--bunch', type=str, help='Pathname of input bunch file.')
+    va_parser.add_model_argument('--particle_number', default=1000, type=int,
+                                 help='Number of particles to use.')
+    va_parser.add_model_argument('--beam_current', default=38.0, type=float,
+                                 help='Initial beam current in mA.')
+    va_parser.add_model_argument('--save_bunch', const='end_bunch.dat', nargs='?', type=str,
+                                 help="Saves the bunch at the end of the lattice after each track in the given "
+                                      "location. If no location is given, the bunch is saved as 'end_bunch.dat' in "
+                                      "the working directory.")
+    return va_parser

virtaccl/site/BTF/btf_virtual_accelerator.py

@@ -12,6 +12,7 @@
 from orbit.py_linac.lattice_modifications import Add_quad_apertures_to_lattice, Add_rfgap_apertures_to_lattice
 
 from virtaccl.PyORBIT_Model.pyorbit_child_nodes import BPMclass, FCclass, BCMclass
+from virtaccl.PyORBIT_Model.pyorbit_va_arguments import add_pyorbit_arguments
 from virtaccl.site.BTF.orbit_model.btf_lattice_factory import PyORBIT_Lattice_Factory
 
 from orbit.core.bunch import Bunch
@@ -25,49 +26,42 @@
 from virtaccl.site.BTF.orbit_model.btf_child_nodes import BTF_Screenclass, BTF_Slitclass
 
 from virtaccl.PyORBIT_Model.pyorbit_lattice_controller import OrbitModel
-from virtaccl.EPICS_Server.ca_server import EPICS_Server
+from virtaccl.EPICS_Server.ca_server import EPICS_Server, add_epics_arguments
 
-from virtaccl.virtual_accelerator import va_parser, virtual_accelerator
+from virtaccl.virtual_accelerator import virtual_accelerator, VA_Parser
 
 
 def main():
     loc = Path(__file__).parent
-    parser, va_version = va_parser()
-    parser.description = 'Run the SNS Linac PyORBIT virtual accelerator server. Version ' + va_version
+    va_parser = VA_Parser()
+    va_parser.set_description('Run the BTF PyORBIT virtual accelerator server.')
+
+    va_parser = add_pyorbit_arguments(va_parser)
+    # Set the defaults for the PyORBIT model.
+    va_parser.change_argument_default('--lattice', loc / 'orbit_model/btf_lattice_straight.xml')
+    va_parser.change_argument_default('--start', 'MEBT1')
+    va_parser.change_argument_default('end', 'MEBT2')
+    va_parser.change_argument_default('--bunch', loc / 'orbit_model/parmteq_bunch_RFQ_output_1.00e+05.dat')
+    va_parser.change_argument_default('--beam_current', 50.0)
+
+    va_parser = add_epics_arguments(va_parser)
+    va_parser.add_server_argument('--print_settings', action='store_true',
+                                  help="Will only print setting PVs. Will NOT run the virtual accelerator.")
 
     # Json file that contains a dictionary connecting EPICS name of devices with their associated element model names.
-    parser.add_argument('--file', '-f', default=loc / 'btf_config.json', type=str,
-                        help='Pathname of config json file.')
-
-    # Lattice xml input file and the sequences desired from that file.
-    parser.add_argument('--lattice', default=loc / 'orbit_model/btf_lattice_straight.xml', type=str,
-                        help='Pathname of lattice file')
-    parser.add_argument("--start", default="MEBT1", type=str,
-                        help='Desired sequence of the lattice to start the model with (default=MEBT1).')
-    parser.add_argument("end", nargs='?', default="MEBT2", type=str,
-                        help='Desired sequence of the lattice to end the model with (default=MEBT2).')
-
-    # Desired initial bunch file and the desired number of particles from that file.
-    parser.add_argument('--bunch', default=loc / 'orbit_model/parmteq_bunch_RFQ_output_1.00e+05.dat', type=str,
-                        help='Pathname of input bunch file.')
-    parser.add_argument('--particle_number', default=10000, type=int,
-                        help='Number of particles to use (default=1000).')
-    parser.add_argument('--beam_current', default=50.0, type=float,
-                        help='Initial beam current in mA. (default=30.0).')
-    parser.add_argument('--save_bunch', const='end_bunch.dat', nargs='?', type=str,
-                        help="Saves the bunch at the end of the lattice after each track in the given location. "
-                             "If no location is given, the bunch is saved as 'end_bunch.dat' in the working directory. "
-                             "(Default is that the bunch is not saved.)")
+    va_parser.add_argument('--config_file', '-f', default=loc / 'btf_config.json', type=str,
+                           help='Pathname of config json file.')
 
     # Json file that contains a dictionary connecting EPICS name of devices with their phase offset.
-    parser.add_argument('--phase_offset', default=None, type=str,
-                        help='Pathname of phase offset file.')
+    va_parser.add_argument('--phase_offset', default=None, type=str,
+                           help='Pathname of phase offset file.')
 
-    args = parser.parse_args()
+    va_parser = va_parser.initialize_arguments()
+    args = va_parser.parse_args()
     debug = args.debug
     save_bunch = args.save_bunch
 
-    config_file = Path(args.file)
+    config_file = Path(args.config_file)
     with open(config_file, "r") as json_file:
         devices_dict = json.load(json_file)
 
@@ -300,9 +294,15 @@ def main():
             slit_device = BTF_Actuator(name, ele_name, speed=speed, limit=limit)
             beam_line.add_device(slit_device)
 
-    server = EPICS_Server()
+    if args.print_settings:
+        for key in beam_line.get_setting_keys():
+            print(key)
+        sys.exit()
+
+    delay = args.ca_proc
+    server = EPICS_Server(process_delay=delay, print_pvs=args.print_pvs)
 
-    virtual_accelerator(model, beam_line, server, parser)
+    virtual_accelerator(model, beam_line, server, va_parser)
 
 
 if __name__ == '__main__':

virtaccl/site/SNS_IDmp/IDmp_maker.py

@@ -16,7 +16,7 @@
 from virtaccl.PyORBIT_Model.bunch_generator import BunchGenerator
 
 
-def get_IDMP_lattice_and_bunch(particle_number=1000, x_off=0, xp_off=0, y_off=0, yp_off=0):
+def get_IDMP_lattice_and_bunch(particle_number=1000, x_off=0, xp_off=0, y_off=0, yp_off=0, debug: bool = False):
     # Field strength and length of the quadrupoles
     quad_field = 0.5
     dch_field = 0.01
@@ -100,14 +100,15 @@ def get_IDMP_lattice_and_bunch(particle_number=1000, x_off=0, xp_off=0, y_off=0,
     # Dump
 
     # Define the sequence and add the list of nodes to the sequence.
-    fodo = Sequence('FODO')
-    fodo.setNodes(list_of_nodes)
+    idmp = Sequence('IDmp')
+    idmp.setNodes(list_of_nodes)
 
     # Define the lattice, add the list of nodes to the lattice, and initialize the lattice.
     my_lattice = LinacAccLattice('My Lattice')
     my_lattice.setNodes(list_of_nodes)
     my_lattice.initialize()
-    print("Total length=", my_lattice.getLength())
+    if debug:
+        print("Total length=", my_lattice.getLength())
 
     # -----TWISS Parameters at the entrance of MEBT ---------------
     # transverse emittances are unnormalized and in pi*mm*mrad
@@ -116,8 +117,9 @@ def get_IDMP_lattice_and_bunch(particle_number=1000, x_off=0, xp_off=0, y_off=0,
     mass = 0.939294  # in [GeV]
     gamma = (mass + e_kin_ini) / mass
     beta = math.sqrt(gamma * gamma - 1.0) / gamma
-    print("relat. gamma=", gamma)
-    print("relat.  beta=", beta)
+    if debug:
+        print("relat. gamma=", gamma)
+        print("relat.  beta=", beta)
     frequency = 402.5e6
     v_light = 2.99792458e8  # in [m/sec]
 
@@ -144,16 +146,18 @@ def get_IDMP_lattice_and_bunch(particle_number=1000, x_off=0, xp_off=0, y_off=0,
     emittZ = emittZ * gamma ** 3 * beta ** 2 * mass
     betaZ = betaZ / (gamma ** 3 * beta ** 2 * mass)
 
-    print(" ========= PyORBIT Twiss ===========")
-    print(" aplha beta emitt[mm*mrad] X= %6.4f %6.4f %6.4f " % (alphaX, betaX, emittX * 1.0e6))
-    print(" aplha beta emitt[mm*mrad] Y= %6.4f %6.4f %6.4f " % (alphaY, betaY, emittY * 1.0e6))
-    print(" aplha beta emitt[mm*MeV] Z= %6.4f %6.4f %6.4f " % (alphaZ, betaZ, emittZ * 1.0e6))
+    if debug:
+        print(" ========= PyORBIT Twiss ===========")
+        print(" aplha beta emitt[mm*mrad] X= %6.4f %6.4f %6.4f " % (alphaX, betaX, emittX * 1.0e6))
+        print(" aplha beta emitt[mm*mrad] Y= %6.4f %6.4f %6.4f " % (alphaY, betaY, emittY * 1.0e6))
+        print(" aplha beta emitt[mm*MeV] Z= %6.4f %6.4f %6.4f " % (alphaZ, betaZ, emittZ * 1.0e6))
 
     twissX = TwissContainer(alphaX, betaX, emittX)
     twissY = TwissContainer(alphaY, betaY, emittY)
     twissZ = TwissContainer(alphaZ, betaZ, emittZ)
 
-    print("Start Bunch Generation.")
+    if debug:
+        print("Start Bunch Generation.")
     bunch_gen = BunchGenerator(twissX, twissY, twissZ)
 
     # set the initial kinetic energy in GeV

virtaccl/site/SNS_IDmp/IDmp_virtual_accelerator.py

@@ -1,63 +1,55 @@
-# Channel access server used to generate fake PV signals analogous to accelerator components.
-# The main body of the script instantiates PVs from a file passed by command line argument.
 import json
+import sys
 from pathlib import Path
 
 from virtaccl.PyORBIT_Model.pyorbit_child_nodes import BPMclass, WSclass, ScreenClass
+from virtaccl.PyORBIT_Model.pyorbit_va_arguments import add_pyorbit_arguments
 from virtaccl.site.SNS_Linac.virtual_devices import (Quadrupole, Corrector, Quadrupole_Power_Supply,
                                                      Corrector_Power_Supply, WireScanner, BPM, P_BPM, Screen)
 from virtaccl.site.SNS_Linac.virtual_devices_SNS import SNS_Dummy_BCM, SNS_Dummy_ICS
 
 from virtaccl.PyORBIT_Model.pyorbit_lattice_controller import OrbitModel
 from virtaccl.beam_line import BeamLine
-from virtaccl.EPICS_Server.ca_server import EPICS_Server
-from virtaccl.virtual_accelerator import va_parser, virtual_accelerator
+from virtaccl.EPICS_Server.ca_server import EPICS_Server, add_epics_arguments
+from virtaccl.virtual_accelerator import virtual_accelerator, VA_Parser
 
 from virtaccl.site.SNS_IDmp.IDmp_maker import get_IDMP_lattice_and_bunch
 
 
 def main():
     loc = Path(__file__).parent
-    parser, va_version = va_parser()
-    parser.description = 'Run the SNS Injection Dump PyORBIT virtual accelerator server. Version ' + va_version
+    va_parser = VA_Parser()
+    va_parser.set_description('Run the SNS Injection Dump PyORBIT virtual accelerator server.')
+
+    va_parser = add_pyorbit_arguments(va_parser)
+    # Set the defaults for the PyORBIT model.
+    va_parser.remove_argument('--lattice')
+    va_parser.remove_argument('--start')
+    va_parser.remove_argument('end')
+    va_parser.remove_argument('--bunch')
+
+    va_parser = add_epics_arguments(va_parser)
+    va_parser.add_server_argument('--print_settings', action='store_true',
+                                  help="Will only print setting PVs. Will NOT run the virtual accelerator.")
 
     # Json file that contains a dictionary connecting EPICS name of devices with their associated element model names.
-    parser.add_argument('--file', '-f', default=loc / 'va_config.json', type=str,
-                        help='Pathname of config json file.')
-
-    # Lattice xml input file and the sequences desired from that file.
-    parser.add_argument('--lattice', default=loc / 'orbit_model/sns_linac.xml', type=str,
-                        help='Pathname of lattice file')
-    parser.add_argument("--start", default="MEBT", type=str,
-                        help='Desired sequence of the lattice to start the model with (default=MEBT).')
-    parser.add_argument("end", nargs='?', default="HEBT1", type=str,
-                        help='Desired sequence of the lattice to end the model with (default=HEBT1).')
-
-    # Desired initial bunch file and the desired number of particles from that file.
-    parser.add_argument('--bunch', default=loc / 'orbit_model/MEBT_in.dat', type=str,
-                        help='Pathname of input bunch file.')
-    parser.add_argument('--particle_number', default=1000, type=int,
-                        help='Number of particles to use (default=1000).')
-    parser.add_argument('--beam_current', default=38.0, type=float,
-                        help='Initial beam current in mA. (default=38.0).')
-    parser.add_argument('--save_bunch', const='end_bunch.dat', nargs='?', type=str,
-                        help="Saves the bunch at the end of the lattice after each track in the given location. "
-                             "If no location is given, the bunch is saved as 'end_bunch.dat' in the working directory. "
-                             "(Default is that the bunch is not saved.)")
+    va_parser.add_argument('--config_file', '-f', default=loc / 'va_config.json', type=str,
+                           help='Pathname of config json file.')
 
     # Json file that contains a dictionary connecting EPICS name of devices with their phase offset.
-    parser.add_argument('--phase_offset', default=None, type=str,
-                        help='Pathname of phase offset file.')
+    va_parser.add_argument('--phase_offset', default=None, type=str,
+                           help='Pathname of phase offset file.')
 
-    args = parser.parse_args()
+    va_parser = va_parser.initialize_arguments()
+    args = va_parser.parse_args()
     debug = args.debug
 
-    config_file = Path(args.file)
+    config_file = Path(args.config_file)
     with open(config_file, "r") as json_file:
         devices_dict = json.load(json_file)
 
     part_num = args.particle_number
-    lattice, bunch = get_IDMP_lattice_and_bunch(part_num, x_off=2, xp_off=0.3)
+    lattice, bunch = get_IDMP_lattice_and_bunch(part_num, x_off=2, xp_off=0.3, debug=debug)
     model = OrbitModel(input_bunch=bunch, debug=debug)
     model.define_custom_node(BPMclass.node_type, BPMclass.parameter_list, diagnostic=True)
     model.define_custom_node(WSclass.node_type, WSclass.parameter_list, diagnostic=True)
@@ -133,9 +125,15 @@ def main():
     dummy_device = SNS_Dummy_ICS("ICS_Tim")
     beam_line.add_device(dummy_device)
 
-    server = EPICS_Server()
+    if args.print_settings:
+        for key in beam_line.get_setting_keys():
+            print(key)
+        sys.exit()
+
+    delay = args.ca_proc
+    server = EPICS_Server(process_delay=delay, print_pvs=args.print_pvs)
 
-    virtual_accelerator(model, beam_line, server, parser)
+    virtual_accelerator(model, beam_line, server, va_parser)
 
     print('Exiting. Thank you for using our virtual accelerator!')