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()
def _load_file(path):
try:
return yaml.load_file(path or '')
except FileNotFoundError:
return ''