def main():
filenames = glob.glob(
'../data/orm/2019-01-20_quadscan/M8-E108-F1-I9-G1/*/*_X.yml')
raw_data = [yaml.safe_load(read_file(filename)) for filename in filenames]
m = Madx(stdout=False)
m.verbose()
m.call('../hit_models/hht3/sequence.madx', chdir=True)
m.command.beam()
m.use('hht3')
#m.call('../hit_models/hht3/strengths0.madx', chdir=True)
for data in raw_data:
mean = mean_values(data)
err = stddevs(data)
kl = np.array([
sum(optics.values())
for optics, group in itertools.groupby(data['records'],
key=lambda r: r['optics'])
])
mon = data['monitors'][0]
knob, = data['optics'][0].keys()
quad = [
elem.name for elem in m.sequence.hht3.expanded_elements
if elem.base_name == 'quadrupole'
and knob in m.expr_vars(elem.defs.k1)
][0]
m.globals.update(data['base_optics'])
for i, x in enumerate('xy'):
#plt.subplot(2, 1, i+1)
plt.title(f"pos{x}_{mon}({knob})")
plt.ylabel("{x} [m]")
plt.xlabel(f"{knob} [$m^{{-1}}$]")
plt.errorbar(kl, mean[:, 0, i], err[:, 0, i], label=x)
for pos in np.linspace(-0.002, 0.002, 5):
modelled = np.array([
track(m, optics, range=f'{quad}/{mon}', **{x: pos})[i]
for optics, group in itertools.groupby(
data['records'], key=lambda r: r['optics'])
])
i0 = np.argmin(kl)
modelled += mean[i0, 0, i] - modelled[i0]
plt.plot(kl, modelled, label=f'${x}_0$={pos}')
plt.legend()
plt.show()
