Source code for pytac.load_csv

"""Module to load the elements of the machine from multiple csv files stored in the same directory."""
import os
import csv
from pytac import lattice, element, device, units, utils
import collections


[docs]def get_div_rigidity(energy): rigidity = utils.rigidity(energy) def div_rigidity(input): return input / rigidity return div_rigidity
[docs]def get_mult_rigidity(energy): rigidity = utils.rigidity(energy) def mult_rigidity(input): return input * rigidity return mult_rigidity
[docs]def load_unitconv(directory, mode, lattice): """Load the unit conversion objects from a file. Args: directory(string): The directory where the data is stored. mode(string): The name of the mode that is used. lattice(Lattice): The lattice object that will be used. """ data = collections.defaultdict(list) uc = {} with open(os.path.join(directory, mode, 'uc_poly_data.csv')) as poly: csv_reader = csv.DictReader(poly) for item in csv_reader: data[(int(item['uc_id']))].append((int(item['coeff']), float(item['val']))) for d in data: u = units.PolyUnitConv([x[1] for x in reversed(sorted(data[d]))]) uc[d] = u data.clear() with open(os.path.join(directory, mode, 'uc_pchip_data.csv')) as pchip: csv_reader = csv.DictReader(pchip) for item in csv_reader: data[(int(item['uc_id']))].append((float(item['eng']), float(item['phy']))) for d in data: eng = [x[0] for x in sorted(data[d])] phy = [x[1] for x in sorted(data[d])] u = units.PchipUnitConv(eng, phy) uc[d] = u with open(os.path.join(directory, mode, 'unitconv.csv')) as unitconv: csv_reader = csv.DictReader(unitconv) for item in csv_reader: element = lattice[int(item['el_id']) - 1] if 'QUAD' in element.families or 'SEXT' in element.families: uc[int(item['uc_id'])].f1 = get_div_rigidity(lattice.get_energy()) uc[int(item['uc_id'])].f2 = get_mult_rigidity(lattice.get_energy()) element._uc[item['field']] = uc[int(item['uc_id'])]
[docs]def load(mode, control_system=None, directory=None): """Load the elements of a lattice from a directory. Args: mode(string): The name of the mode to be loaded. control_system(ControlSystem): The control system to be used. If none is provided an EpicsControlSystem will be created. directory(string): Directory where to load the files from. If no directory is given the data directory at the root of the repository is used. Returns: Lattice: The lattice containing all elements. """ if control_system is None: # Don't import epics unless we need it to avoid unnecessary # installation of cothread from pytac import epics control_system = epics.EpicsControlSystem() if directory is None: directory = os.path.join(os.path.dirname(os.path.abspath(__file__)), 'data') lat = lattice.Lattice(mode, control_system, 3000) with open(os.path.join(directory, mode, 'elements.csv')) as elements: csv_reader = csv.DictReader(elements) for item in csv_reader: cell = int(item['cell']) if item['cell'] else None e = element.Element(item['name'], float(item['length']), item['type'], cell) e.add_to_family(item['type']) lat.add_element(e) with open(os.path.join(directory, mode, 'devices.csv')) as devices: csv_reader = csv.DictReader(devices) for item in csv_reader: name = item['name'] enable_pv = item['enable_pv'] enable_value = item['enable_value'] get_pv = item['get_pv'] set_pv = item['set_pv'] pve = True if enable_pv and enable_value: pve = device.PvEnabler(enable_pv, enable_value, control_system) d = device.Device(name, control_system, pve, get_pv, set_pv) lat[int(item['id']) - 1].add_device(item['field'], d, control_system) with open(os.path.join(directory, mode, 'families.csv')) as families: csv_reader = csv.DictReader(families) for item in csv_reader: lat[int(item['id']) - 1].add_to_family(item['family']) if os.path.exists(os.path.join(directory, mode, 'unitconv.csv')): load_unitconv(directory, mode, lat) return lat