def update_histo(self):
if len(self.position) > 0:
M_position = compton.transform(self.M, np.array(self.position))
B, _ = np.histogramdd(M_position,
self.bin_edges,
weights=np.array(self.edep))
self.hist += B
self.position = []
self.edep = []
def calc_vectorfield(conf):
M = compton.build_transformation(conf['Slab']['Transformation'], mm, deg)
M = inv(M)
grid = np.array(conf['Slab']['Volume']) * mm
dx = np.array(conf['VectorField']['GridResolution']) * mm
result = []
with h5py.File(conf['VectorField']['Input'], 'r') as fin:
for group_name, gin in tqdm.tqdm(fin.items(),
bar_format=fmt,
desc='Reading trajectories'):
x = gin['x'][:] * meter
Mx = np.array([compton.transform(M, q) for q in x])
grid_mask = compton.in_volume(grid, Mx)
grid_mask += shift(grid_mask, -10) + shift(grid_mask, 1)
Mx = Mx[grid_mask].copy()
delta = np.diff(Mx, axis=0)
n = delta / norm(delta, axis=1, keepdims=True)
result.append(np.array((Mx[:-1], n)))
x, n = np.hstack(result)
mirror_x = x * np.array((1, -1, 1))
bar = tqdm.tqdm(total=4,
bar_format=fmt,
desc='Creating smooth vectorfield')
delaunay = spatial.Delaunay(np.concatenate((x, mirror_x)))
bar.update(1)
nearest_interp = NearestNDInterpolator(x, n)
linear_interp = LinearNDInterpolator(x, n)
bar.update(1)
epsilon = 0.1 * dx
dims = [np.arange(grid[q, 0], grid[q, 1] + epsilon, dx) for q in [0, 1, 2]]
# Use linear interpolation within convex hull, else use nearest-neighbor.
X, Y, Z = np.meshgrid(*dims, indexing='ij')
n_XYZ = linear_interp((X, np.abs(Y), Z))
bar.update(1)
n_XYZ_nearest = nearest_interp((X, np.abs(Y), Z))
bar.update(1)
bar.close()
mask = np.isnan(n_XYZ)
n_XYZ[mask] = n_XYZ_nearest[mask]
n_XYZ[Y < 0] *= np.array((1, -1, 1))
if 'DiagnosticOutput' in conf['VectorField']:
with h5py.File(conf['VectorField']['DiagnosticOutput'], 'w') as fout:
fout['nx'] = n_XYZ[:, :, :, 0]
fout['ny'] = n_XYZ[:, :, :, 1]
fout['nz'] = n_XYZ[:, :, :, 2]
interpolator = RegularGridInterpolator(dims,
n_XYZ,
bounds_error=False,
fill_value=None)
return interpolator, delaunay
def plot_sextupole_magnet_profile(ax, conf):
from numpy.linalg import inv
M = compton.build_transformation(conf['Transformation'], mm, deg)
Minv = inv(M)
c0 = conf['c0'] * mm
b1 = conf['b1'] * mm
a0 = c0 * (3 * (b1 / c0)**2 - 1)**(1.0 / 3)
from scipy.interpolate import interp1d
y = np.arange(1, 90 * mm, 1.0 * mm)
z = np.sqrt((y**3 + a0**3) / (3 * y))
y_z = interp1d(z, y)
zscint = np.arange(95, 300, 1.0 * mm)
pos = np.array((np.zeros_like(zscint), np.zeros_like(zscint), zscint))
tpos = np.array([compton.transform(M, x) for x in pos.T]).T
tpos[1] = y_z(tpos[2])
pos = np.array([compton.transform(Minv, x) for x in tpos.T]).T
ax.plot(pos[2] / mm, pos[1] / mm, color='k', linewidth=0.4)
ax.plot(pos[2] / mm, -pos[1] / mm, color='k', linewidth=0.4)
def project_histogram(filename, conf, indices):
fin = h5py.File(filename, 'r')
edep = fin['edep'][:] * keV
pos = fin['position'][:] * mm
num_events = fin['edep'].attrs['num_events']
M = compton.build_transformation(conf, mm, deg)
tpos = np.array([compton.transform(M, x) for x in pos])
# from pyevtk.hl import pointsToVTK
# pointsToVTK("tpos", *np.ascontiguousarray(tpos.T), data=None)
H, edges = np.histogramdd(tpos,
tuple(x.vals for x in indices),
weights=(edep / num_events))
return H, edges
def ProcessHits(self, step, history):
event_id = g4.gEventManager.GetConstCurrentEvent().GetEventID()
if event_id != self.curr_event:
self.curr_event = event_id
self.tally = 0
x = G4ThreeVector_to_list(step.GetPreStepPoint().GetPosition())
Mx = compton.transform(self.M, np.array(x))
if self.tally < self.threshold:
if compton.in_volume(self.vol, np.array((Mx, )))[0]:
self.tally += step.GetTotalEnergyDeposit()
if self.tally >= self.threshold:
self.num_flagged += 1
if self.num_flagged <= self.limit_count:
g4.gApplyUICommand('/event/keepCurrentEvent')
def plot_sextupole_magnet_profile(ax, conf):
from numpy.linalg import inv
M = compton.build_transformation(conf['Transformation'], mm, deg)
Minv = inv(M)
c0 = conf['c0'] * mm
b1 = conf['b1'] * mm
xoff = conf['xoff'] * mm
alpha0 = conf['alpha0'] * deg
a0 = c0 * (3 * (b1 / c0)**2 - 1)**(1.0 / 3)
y = np.arange(1 * mm, 200 * mm, 1 * mm)
z = np.sqrt((y**3 + a0**3) / (3 * y))
x = xoff + np.tan(alpha0) * z
pos = np.array((x, y, z))
tpos = np.array([compton.transform(M, x) for x in pos.T]).T
_, yscint, zscint = tpos
ax.plot(zscint / mm, yscint / mm, color='#cccccc', linewidth=0.4, zorder=1)
ax.plot(zscint / mm,
-yscint / mm,
color='#cccccc',
linewidth=0.4,
zorder=1)
def main():
common.setup_plot()
energy = []
position = []
theta0 = 28.0 * deg
M = compton.build_transformation(
[['TranslateX', 'RotateY', 'TranslateZ'], [-40.1, -28.0, -30.0]], mm,
deg)
with h5py.File('trajectories/trajectories.h5', 'r') as fin:
for gin in fin.values():
E0 = get_energy(gin)
x0 = gin['x'][0] * meter
x1 = gin['x'][-1] * meter
if x0[1] != 0.0:
continue
x1 = compton.transform(M, x1)
energy.append(E0)
position.append(x1[2])
energy = np.array(energy)
position = np.array(position)
args = np.argsort(energy)
energy = energy[args]
position = position[args]
mod = lmfit.Model(fit_func)
params = mod.make_params(c0=0.0, c1=0.0, c2=0.0)
result = mod.fit(data=energy / MeV, x=position, params=params)
print(result.fit_report())
v = result.params.valuesdict()
print(v['c0'])
print(v['c1'])
print(v['c2'])
x_fit = np.linspace(position[0], position[-1], 200)
fig = plot.figure(figsize=(244.0 / 72, 120.0 / 72))
ax = fig.add_subplot(1, 1, 1)
ax.semilogy(x_fit / mm, result.eval(x=x_fit), linewidth=0.6)
ax.semilogy(position / mm,
energy / MeV,
marker='.',
ls='',
markersize=0.001,
color='k')
ax.text(0.05,
0.8,
r'$E(z_s)/{\rm MeV} = \exp (c_0 + c_1 z_s + c_2 z_s^2)$',
fontsize=7.0,
transform=ax.transAxes)
text = r'$c_0 = {:.3}$'.format(v['c0']) + '\n'
text += r'$c_1 = {:.3}'.format(num2tex(v['c1'] * mm))
text += r'\;{\rm mm}^{-1}$' + '\n'
text += r'$c_2 = {:.3}'.format(num2tex(v['c2'] * mm**2))
text += r'\;{\rm mm}^{-2}$' + '\n'
print(text)
ax.text(0.5, 0.2, text, transform=ax.transAxes)
ax.set_xlabel(r'$z_s$ (mm)', labelpad=-1.0)
ax.set_ylabel(r"Energy (MeV)", labelpad=2.0)
# ax.set_xlim(0, 250)
ax.set_ylim(0.1, 20)
ax.xaxis.set_minor_locator(matplotlib.ticker.AutoMinorLocator())
filename = 'out/energy-scale.pdf'
os.makedirs(os.path.dirname(filename), exist_ok=True)
plot.savefig(filename, transparent=True)
def calc_lattice(conf, interpolator, delaunay):
bar = tqdm.tqdm(total=2, bar_format=fmt, desc='Creating pore lattice')
pore_conf = conf['Pores']
a = pore_conf['a'] * mm
delta_z = pore_conf['delta_z'] * mm
z0 = pore_conf['z0'] * mm
lz = pore_conf['lz'] * mm
ly = pore_conf['ly'] * mm
magnet_buffer = pore_conf['MagnetBuffer'] * mm
max_length = pore_conf['MaxLength'] * mm
volume = np.array(conf['Slab']['Volume']) * mm
M = compton.build_transformation(conf['Slab']['Transformation'], mm, deg)
zy_lattice = gen_lattice_points(a, (0, lz), (-0.5 * ly, 0.5 * ly))
zy_lattice += np.array((delta_z, 0))
# zy --> xyz
tri_lattice = np.hstack((volume[0, 0] * np.ones(
(len(zy_lattice), 1)), np.flip(zy_lattice, axis=1)))
tri_lattice = tri_lattice[tri_lattice[:, 2] > z0]
tri_lattice = tri_lattice[delaunay.find_simplex(tri_lattice) >= 0]
bar.update(1)
M_tri_lattice = np.array([compton.transform(M, q) for q in tri_lattice])
# discard lattice points within a given distance from magnet pole
tri_lattice = tri_lattice[calculate_pole_mask(conf, M_tri_lattice,
magnet_buffer)]
bar.update(1)
bar.close()
hex_lattice = {}
N = 6
phi = (30 * deg) + np.linspace(0, 2 * np.pi, N, endpoint=False)
hex_coords = (a / np.sqrt(3)) * np.array(
(np.zeros_like(phi), np.sin(phi), np.cos(phi))).T
# integrate hexagonal lattice through vectorfield
for i, p in tqdm.tqdm(list(enumerate(tri_lattice)),
bar_format=fmt,
desc='Integrating lattice trajectories'):
for q in hex_coords:
coord = p + q
key = (round(coord[1], 3), round(coord[2], 3))
if key in hex_lattice:
continue
if delaunay.find_simplex(coord) == -1:
# hex lattice point is outside trajectory hull
hex_lattice[key] = None
continue
trajectory = calc_trajectory(interpolator, coord, max_length)
if trajectory == None:
# trajectory didn't cross x=0 plane
hex_lattice[key] = None
continue
hex_lattice[key] = trajectory
# assemble result
result = []
for p in tri_lattice:
valid_hex = True
trajectories = []
for q in hex_coords:
coord = p + q
key = (round(coord[1], 3), round(coord[2], 3))
if hex_lattice[key] == None:
valid_hex = False
break
trajectories.append(hex_lattice[key])
if valid_hex:
result.append(trajectories)
return result
def main():
args = get_args()
conf = args.conf
# create safe interpreter for evaluation of configuration expressions
aeval = asteval.Interpreter(use_numpy=True)
for q in common.units.__all__:
aeval.symtable[q] = common.units.__dict__[q]
pconf = conf['Projection']
M = compton.build_transformation(pconf['Transformation'], mm, deg)
prefilter = np.array(pconf['Prefilter']) * mm
postfilter = np.array(pconf['Postfilter']) * mm
energy = []
position = []
x = []
E0 = []
with h5py.File(conf['Files']['Input'], 'r') as fin:
for gin in fin.values():
x.append((gin['x'][0] * meter,
compton.transform(M, gin['x'][-1] * meter)))
E0.append(get_energy(gin))
x = np.array(x)
E0 = np.array(E0)
prefilter_mask = compton.in_volume(prefilter, x[:, 0, :])
x_pre = x[prefilter_mask, :, :]
E0_pre = E0[prefilter_mask]
postfilter_mask = compton.in_volume(postfilter, x_pre[:, 1, :])
x_post = x_pre[postfilter_mask, :, :]
E0_post = E0_pre[postfilter_mask]
energy = E0_post.copy()
position = x_post[:, 1, 2].copy()
args = np.argsort(energy)
energy = energy[args]
position = position[args]
mod = lmfit.Model(fit_func)
params = mod.make_params(c0=0.0, c1=0.0, c2=0.0, c3=0.0)
result = mod.fit(data=np.log(energy / MeV), x=position, params=params)
v = result.params.valuesdict()
x_fit = np.linspace(position[0], position[-1], 200)
common.setup_plot()
fig = plot.figure(figsize=np.array(conf['Plot']['FigSize']) / 72)
ax = fig.add_subplot(1, 1, 1)
axes = [get_axis(aeval, *conf['Plot'][x]) for x in ['XAxis', 'YAxis']]
ax.semilogy(x_fit / axes[0].unit,
np.exp(result.eval(x=x_fit)),
linewidth=0.6)
ax.semilogy(position / axes[0].unit,
energy / axes[1].unit,
marker='.',
ls='',
markersize=2.0,
markeredgewidth=0,
color='k')
aeval.symtable['fitval'] = v
aeval.symtable['num2tex'] = num2tex
plot_annotation(ax, aeval, conf['Plot'])
ax.set_xlabel(axes[0].label, labelpad=-1.0)
ax.set_ylabel(axes[1].label, labelpad=2.0)
ax.set_xlim(*axes[0].xlim)
ax.set_ylim(*axes[1].xlim)
ax.xaxis.set_minor_locator(matplotlib.ticker.AutoMinorLocator())
filename = conf['Files']['PlotOutput']
path = os.path.dirname(filename)
if path != '':
os.makedirs(path, exist_ok=True)
plot.savefig(filename, transparent=True)
if 'CalcOutput' in conf['Files']:
filename = conf['Files']['CalcOutput']
path = os.path.dirname(filename)
if path != '':
os.makedirs(path, exist_ok=True)
calc_output = {
'EnergyScaleCoefficients': {
'c0': float(v['c0']),
'c1': float(v['c1'] * mm),
'c2': float(v['c2'] * mm**2),
'c3': float(v['c3'] * mm**3)
}
}
with open(filename, 'w') as fout:
toml.dump(calc_output, fout)
