def main(args=None): opts = docopt(__doc__, args) input_file = opts['<ERRORS>'] output_file = opts['--output'] if not output_file: output_file = os.path.splitext(input_file)[0] + '-fit.yml' BPM_ERR = float(opts['--stderr']) record_files = ( opts['<MEASURED>'] or '../data/orm/2018-10-20-orm_measurements/M8-E108-F1-I9-G1/*.yml') model_path = opts['--model'] or '../hit_models/hht3' ana = Analysis.app(model_path, record_files) ana.model.update_twiss_args(dict(x=0, y=0, px=0, py=0)) if opts['--absolute']: ana.mode = 'abs' elif opts['--hybrid']: ana.mode = 'hybrid' else: ana.mode = 'orm' ana.measured.stderr = (ana.measured.stderr**2 + BPM_ERR**2)**0.5 if opts['-i']: init_errors = parse_errors(yaml.load_file(opts['-i'])) ana.apply_errors(init_errors.keys(), init_errors.values()) if opts['--backtrack']: from_monitors = ['t3dg2g', 't3dg1g', 't3df1'] ana.ensure_monitors_available(from_monitors) ana.setup_backtracking(ana.extrapolate(from_monitors, to='#e')) ana.init() errors = parse_errors(yaml.load_file(input_file)) options = dict( algorithm='lstsq', mode='xy', iterations=8, delta=1e-3, bounds=Bounds(-0.001, 0.001), fourier=False, rcond=1e-6, ) result = ana.fit(errors, save_to=output_file, **options) print(result)
def main(args=None): opts = docopt(__doc__, args) model_file = opts['--model'] or '../hit_models/hht3' record_files = ( opts['RECORDS'] or ['../data/orm/2018-10-20-orm_measurements/M8-E108-F1-I9-G1']) if len(record_files) == 1 and os.path.isdir(record_files[0]): default_prefix = record_files[0] + '_' record_files = record_files[0] + '/*.yml' else: default_prefix = 'orbits' if opts['--init']: g_spec = yaml.load_file(opts['--init']) else: g_spec = {} g_errors = list(map(parse_error, g_spec.keys())) g_values = list(g_spec.values()) if opts['--error']: specs = yaml.load_file(opts['--error']) default_prefix += '/' + os.path.splitext( os.path.basename(opts['--error']))[0] else: specs = [{}] if isinstance(specs, dict): specs = [specs] prefix = (opts['--output'] or default_prefix) + '/' os.makedirs(prefix, exist_ok=True) if opts['--init']: copyfile(opts['--init'], f'{prefix}spec_init.yml') if opts['--error']: copyfile(opts['--error'], f'{prefix}spec_error.yml') ana = Analysis.app(model_file, record_files) model = ana.model strengths = ana.measured.strengths model.madx.eoption(add=True) model.update_globals(strengths.items()) for i, spec in enumerate(specs): errors = list(map(parse_error, spec.keys())) + g_errors values = list(spec.values()) + g_values errname = '_' + repr(errors[0]) if len(spec) == 1 else '' output_orbits(ana, f'{prefix}/model_{i}{errname}/', errors, values)
def _load_params(data, name, path): """Load parameter dict from file if necessary.""" vals = data.get(name, {}) if isinstance(vals, str): data[name] = yaml.load_file(os.path.join(path, vals)) if len(data[name]) == 1 and name in data[name]: data[name] = data[name][name]
def load_file(cls, filename, madx=None, *, undo_stack=None, interpolate=0, **madx_kwargs): """Load model from .madx or .yml file and pass additional arguments to the Model constructor.""" madx = madx or Madx(**madx_kwargs) madx.option(echo=False) filename = os.path.abspath(filename) path, name = os.path.split(filename) ext = os.path.splitext(name)[1].lower() if ext in ('.yml', '.yaml'): data = yaml.load_file(filename) path = os.path.join(path, data.get('path', '.')) _load_params(data, 'beam', path) _load_params(data, 'twiss', path) for fname in data.get('init-files', []): _call(madx, path, fname) else: _call(madx, path, filename) seqname = _guess_main_sequence(madx) data = _get_seq_model(madx, seqname) data['init-files'] = [filename] return cls(madx, data, undo_stack=undo_stack, filename=filename, interpolate=interpolate)
def main(args=None): opts = docopt(__doc__, args) prefix = opts['--output'] + '/' os.makedirs(prefix, exist_ok=True) model = opts['--model'] record_files = ( opts['RECORDS'] or '../data/orm/2018-10-20-orm_measurements/M8-E108-F1-I9-G1/*.yml') ana = Analysis.app(model, record_files) if opts['--base-orbit']: ana.init() base_orm = ana.model_orbits else: base_orm = None for filename in opts['--errors']: import_errors(ana.model, yaml.load_file(filename)) ana.init({}) if opts['--backtrack']: ana._init_twiss = ana.extrapolate(['t3dg2g', 't3dg1g', 't3df1'], to='#s') ana.info() ana.plot_monitors(base_orm=base_orm, save_to=f'{prefix}mon_response') ana.plot_steerers(base_orm=base_orm, save_to=f'{prefix}mon_response') ana.plot_orbit(save_to=f'{prefix}orbit')
def plot_single_optic_measurements(): """Create plots for the 'single optic' measurements (i.e. combinations of .knobs.str + .bpms.yml exports).""" # Currently, we have only gantry measurements of this type: model_name = 'hht3' bpm_suffix = '.bpms.yml' str_suffix = '.knobs.str' bpm_files = [ fname for folder in single_optic_measurements_folders for fname in glob(DATA_PREFIX + folder + '*' + bpm_suffix) ] for bpm_file in bpm_files: str_file = strip_suffix(bpm_file, bpm_suffix) + str_suffix img_base = strip_suffix(bpm_file, bpm_suffix) print("Processing", bpm_file) model = Model.load_file(f'../hit_models/{model_name}.cpymad.yml', stdout=False) model.interpolate = 400 model.call(str_file) # dict with {<monitor>: {'x': <X>, 'y': <Y>, ...}} measured = load_file(bpm_file)['monitor'] measured = sort_dict(measured, key=model.elements.index) bpm_names = sorted(measured, key=model.elements.index) bpm_data = np.array([[ model.elements[bpm_name].position, measured[bpm_name]['x'], measured[bpm_name]['y'] ] for bpm_name in bpm_names]) model_hebt = interpolate(model, measured, HEBT_BPMS) model_isoc = interpolate(model, measured, ISOC_BPMS) curves = [ (bpm_data.T, args('o', label='measured', color='C2')), (model_hebt, args('-', label='model', color='C0')), (model_isoc, args('-', label='model backtracked', color='C1', ls='dashed')), ] fig = plot_twissfigure(model, curves, ylim=(-25, +25)) savefig(fig, img_base + '-fit-25') fig = plot_twissfigure(model, curves, ylim=(-40, +40)) savefig(fig, img_base + '-fit-40') fig = plot_twissfigure(model, curves, ylim=(-200, +200)) savefig(fig, img_base + '-fit-200')
def import_params(filename, data_key=None): _, ext = os.path.splitext(filename.lower()) if ext in ('.yml', '.yaml'): data = yaml.load_file(filename) if data_key: data = data[data_key] elif ext == '.str': data = read_str_file(filename) if data is None: raise ValueError( "SyntaxError in {!r}. Not the simplest of .str files?" .format(filename)) else: raise ValueError( "Unknown file format for import: {!r}".format(filename)) return data
def import_params(filename, data_key=None): _, ext = os.path.splitext(filename.lower()) if ext in ('.yml', '.yaml'): data = yaml.load_file(filename) if data_key: data = data[data_key] elif ext == '.str': data = read_str_file(filename) if data is None: raise ValueError( "SyntaxError in {!r}. Not the simplest of .str files?".format( filename)) else: raise ValueError( "Unknown file format for import: {!r}".format(filename)) return data
def load_record_file(filename): """ Load a YAML file in the format produced by one of the madgui dialogs: - Online control -> ORM measurement - Online control -> Beam diagnostic -> Offsets -> Calibrate """ data = yaml.load_file(filename) if 'base_optics' in data: # exported by `Online -> Beam diagnostic -> Offsets -> Calibrate`… data = _convert_orm_export(data) strengths = data['model'] records = {(monitor, tuple(record['optics'].items())): [shot[monitor][:2] for shot in record['shots']] for record in data['records'] for monitor in data['monitors']} return strengths, records
def load_sd_values(self, filename): from madgui.util.yaml import load_file self.sd_file = filename = os.path.abspath(filename) data = load_file(filename) cols = { 'envx': 'widthx', 'envy': 'widthy', 'x': 'posx', 'y': 'posy', } self.auto_sd.set(False) self._lib.set_sd_values({ cols[param] + '_' + elem: value for elem, values in data['monitor'].items() for param, value in values.items() })
def plot_orms(): print("Plotting ORMs for each session") from orm_util import Analysis for seq, folder in orm_measurements_folders: record_files = DATA_PREFIX + folder + '*.yml' name = seq[-2:].upper() + (' (pre gantry)' if seq == 'hht3' else '') print(record_files) ana = Analysis.session(f'../hit_models/{seq}', record_files) ana.mode = 'hybrid' hebt_bpms = [m for m in ana.monitors if m in HEBT_BPMS] gant_bpms = [ m for m in ana.monitors if m not in HEBT_BPMS and m not in ('t3dg1g', 't3df1') ] isoc_bpms = [m for m in ana.monitors if m in ISOC_BPMS] ana.measured.stderr = (ana.measured.stderr**2 + BPM_ERR**2)**0.5 ana.model.update_twiss_args(dict(x=0, y=0, px=0, py=0)) ana.init() base_orbits = ana.model_orbits ana.info([ana.monitors.index(m) for m in hebt_bpms]) ana.info([ana.monitors.index(m) for m in gant_bpms]) if hebt_bpms: fig = ana.plot_orm(hebt_bpms) suptitle(fig, f"Orbit response: {name}") savefig(fig, DATA_PREFIX + folder[:-1] + '-orm-init0-horiz') if gant_bpms: fig = ana.plot_orm(gant_bpms) suptitle(fig, "Orbit response: T3 (gantry)") savefig(fig, DATA_PREFIX + folder[:-1] + '-orm-init0-gantry') fit_orms = [] for errfile in fit_files: errors = parse_errors(load_file(errfile)) errors = filter_errors(errors, ana.model) ana.apply_errors(errors.keys(), errors.values()) ana.init() fit_orms.append(ana.model_orbits) ana.model_orbits = base_orbits for ifit, fitted in enumerate(fit_orms): if hebt_bpms: fig = ana.plot_orm(hebt_bpms, fitted) suptitle(fig, f"Orbit response: {name}") savefig( fig, DATA_PREFIX + folder[:-1] + f'-orm-f{ifit}-init0-horiz') if gant_bpms: fig = ana.plot_orm(gant_bpms, fitted) suptitle(fig, "Orbit response: T3 (gantry)") savefig( fig, DATA_PREFIX + folder[:-1] + f'-orm-f{ifit}-init0-gantry') ana = Analysis.session(f'../hit_models/{seq}', record_files) ana.mode = 'hybrid' ana.measured.stderr = (ana.measured.stderr**2 + BPM_ERR**2)**0.5 ana.model.update_twiss_args(dict(x=0, y=0, px=0, py=0)) ana.init() # As it turns out, the ORMs obtained by using different initial # conditions are very similar to those assuming x=y=px=py=0. We # therefore do currently repeat the fitted error plots here. ana.apply_errors([], []) if len(hebt_bpms) >= 2: ana.backtrack(hebt_bpms) if hebt_bpms: fig = ana.plot_orm(hebt_bpms) suptitle(fig, f"Orbit response: {name}") savefig(fig, DATA_PREFIX + folder[:-1] + '-orm-backfit-horiz') if gant_bpms: fig = ana.plot_orm(gant_bpms) suptitle(fig, "Orbit response: T3 (gantry)") savefig(fig, DATA_PREFIX + folder[:-1] + '-orm-backfit-gantry') if len(isoc_bpms) >= 2: ana.backtrack(isoc_bpms) if hebt_bpms: fig = ana.plot_orm(hebt_bpms) suptitle(fig, f"Orbit response: {name}") savefig(fig, DATA_PREFIX + folder[:-1] + '-orm-backisoc-horiz') if gant_bpms: fig = ana.plot_orm(gant_bpms) suptitle(fig, "Orbit response: T3 (gantry)") savefig(fig, DATA_PREFIX + folder[:-1] + '-orm-backisoc-gantry')
def _load_file(path): try: return yaml.load_file(path or '') except FileNotFoundError: return ''
def main(model_file, spec_file, record_file): """ Usage: analysis MODEL PARAMS RECORDS MODEL must be the path of the model/sequence file to initialize MAD-X. PARAMS is a YAML file with arguments for the "measurement" procedure, it must contain at least a list of `monitors` and `optics` and should contain keys for errors to be inserted: `knobs`, `ealign`, `efcomp` RECORDS is the name of the YAML output file where """ init_app([], gui=False) config = load_config(isolated=True) with ExitStack() as stack: setup_args = yaml.load_file(spec_file)['procedure'] session = stack.enter_context(Session(config)) session.control._settings.update({ 'shot_interval': 0.001, 'jitter': setup_args.get('jitter', True), 'auto_params': False, 'auto_sd': True, }) session.load_model(model_file, stdout=False, command_log=lambda text: print("X:>", text)) session.control.set_backend('hit_acs.plugin:TestACS') session.control.connect() session.control.write_all() corrector = Corrector(session) corrector.setup({ 'monitors': setup_args['monitors'], 'optics': setup_args['optics'], }) # FIXME: this is not yet compatible with general parameter errors. In # order to fix this, the hit_acs test backend will have to use an # independent model! model = session.model() import_errors(model, setup_args['errors']) model.twiss.invalidate() corrector.set_optics_delta(setup_args.get('optics_deltas', {}), setup_args.get('default_delta', 1e-4)) corrector.open_export(record_file) widget = mock.Mock() procbot = ProcBot(widget, corrector) num_mons = len(setup_args['monitors']) num_optics = len(setup_args['optics']) + 1 if setup_args.get('jitter', True): num_ignore = setup_args.get('num_ignore', 1) num_shots = setup_args.get('num_shots', 5) else: num_ignore = 0 num_shots = 1 procbot.start(num_ignore, num_shots, gui=False) total_steps = num_mons * (num_optics + 1) * (num_ignore + num_shots) i = 0 while procbot.running and i < 2 * total_steps: procbot._feed(None, None) time.sleep(0.010) i += 1 assert not procbot.running
def load_from(self, filename): data = yaml.load_file(filename)['orm'] self.saved_orms.push([ORM_Entry(*entry) for entry in data])
def load_from(self, filename): offsets = yaml.load_file(filename)['offsets'] self._offsets.clear() self._offsets.update(offsets) self.update()