from ...lib.utils.math import normalized_linear_interpolator
import json
import numpy as np
import math
D2R = np.pi / 180
R2D = 180 / np.pi
# conversion from THz to A^-2
K = 1.9958584
THZ2MEV = 4.1356675
[docs]
def x2qe(x, axes, offset):
"""Transform coordinates from 2D plane to 4D Q-E space."""
return np.dot(x, axes) + offset
[docs]
def qe2x(qe, axes, offset):
"""Transform coordinates from 4D Q-E space to 2D plane."""
axes = np.atleast_2d(axes)
return np.linalg.solve(axes @ axes.T, np.dot(qe - offset, axes.T).T).T
[docs]
class Sample:
"""Class representing a sample."""
def __init__(self, name: str, lattice, angles) -> None:
self.name = name
self.lattice = np.asarray(lattice)
self.angles = np.asarray(angles)
def __repr__(self) -> str:
return f"name: {self.name}, lattice: {self.lattice}, angles: {self.angles}"
[docs]
class Orientation:
"""Class representing an orientation of a sample."""
def __init__(self, orient1, orient2, psi0: float) -> None:
self.orient1 = np.asarray(orient1)
self.orient2 = np.asarray(orient2)
self.psi0 = psi0
def __repr__(self) -> str:
return f"orient1: {self.orient1}, orient2: {self.orient2}, psi0: {self.psi0}"
[docs]
class InstrumentConfiguration:
"""Class representing an instrument configuration."""
def __init__(self, scan_mode: str, scan_constant: float, senses,
coupled=False, psi360=True) -> None:
assert scan_constant > 0
self.scan_mode = scan_mode
self.scan_constant = scan_constant
self.senses = np.asarray(senses)
self.coupled = coupled
self.psi360 = psi360
def __repr__(self) -> str:
return f"scan mode: {self.scan_mode}, scan_constant: {self.scan_constant}\n" \
f"senses: {self.senses}, coupled: {self.coupled}, psi360: {self.psi360}"
[docs]
class Instrument:
"""Class representing an instrument."""
def __init__(self, angle_velocities, d_mono: float, d_ana: float) -> None:
self.angle_velocities = np.asarray(angle_velocities)
self.d_mono = d_mono
self.d_ana = d_ana
def __repr__(self) -> str:
return f"angle velocities: {self.angle_velocities}, \n" \
f"d_mono: {self.d_mono}, d_ana: {self.d_ana}"
[docs]
class AngleMap:
"""Class for mapping Q-E points to corresponding angles of instrument axes."""
def __init__(self, sample: Sample, orientation: Orientation,
instrument_config: InstrumentConfiguration,
instrument: Instrument) -> None:
self._sample = sample
self._orientation = orientation
self._instrument_config = instrument_config
self._instrument = instrument
self._lattice = self._sample.lattice
self._angles = self._sample.angles
self._orient1 = self._orientation.orient1
self._orient2 = self._orientation.orient2
self._psi0 = self._orientation.psi0
self._scan_mode = self._instrument_config.scan_mode
self._scan_constant = self._instrument_config.scan_constant
self._senses = self._instrument_config.senses
self._coupled = self._instrument_config.coupled
self._psi360 = self._instrument_config.psi360
self._d_mono = self._instrument.d_mono
self._d_ana = self._instrument.d_ana
# calculate reciprocal lattice
self._lattice_rec, self._angles_rec = self._lattice_rec_f()
# matrix for rotation about z lab system with psi
self._matrix = self._matrix_crystal2lab()
# matrix for rotation about sgx, sgy for psi = 0; not changed so far
self._matrix_cardan = np.identity(3)
def __repr__(self) -> str:
return f'direct lattice: {self._lattice}\n' \
f'direct angles: {self._angles}\n' \
f'plane vectors: {self._orient1} {self._orient2}\n' \
f'psi0: {self._psi0 * R2D}\n' \
f'recip. lattice: {self._lattice_rec}\n' \
f'recip. angles: {[x * R2D for x in self._angles_rec]}\n' \
f'zone axis: {self._cal_zone()}\n' \
f'cardan matrix: {self._matrix_cardan}\n' \
f'hkl2qcart matrix: {self._matrix}'
def __call__(self, qe):
qhkl = qe[:3]
e = qe[3] / THZ2MEV # convert from meV to THz
# compute phi, psi, alpha
qlab = self._hkl2qlab(qhkl[0], qhkl[1], qhkl[2])
Y = self._cal_Y(self._orient1, self._orient2, qhkl)
if self._scan_mode in ['CKI', 'DIFF']:
ki = self._scan_constant
kf = self._cal_kf(e, ki)
phi = self._cal_phi(qlab, ki, kf, self._senses[1])
alpha = self._cal_alpha1(qlab, e, ki, self._senses[1])
psi = self._cal_psi(Y, alpha)
elif self._scan_mode == 'CKF':
kf = self._scan_constant
ki = self._cal_ki1(e, kf)
phi = self._cal_phi(qlab, ki, kf, self._senses[1])
alpha = self._cal_alpha1(qlab, e, ki, self._senses[1])
psi = self._cal_psi(Y, alpha)
elif self._scan_mode == 'CPSI':
psi = self._scan_constant
alpha = self._cal_alpha2(Y, psi)
ki = self._cal_ki3(qlab, e, alpha)
sphi = 1 if alpha > 0 else -1
kf = self._cal_kf(e, ki)
phi = self._cal_phi(qlab, ki, kf, sphi)
elif self._scan_mode == 'CPHI':
phi = self._scan_constant
sphi = np.sign(phi)
ki = self._cal_ki2(qlab, e, phi)
kf = self._cal_kf(e, ki)
alpha = self._cal_alpha1(qlab, e, ki, sphi)
psi = self._cal_psi(Y, alpha)
psi -= self._psi0
if self._coupled:
psi -= phi
while psi < -180:
psi += 360
while psi > 180:
psi -= 360
if self._psi360 and psi < 0:
psi += 360
# transform ki and kf to angles
ki_angle = self._senses[0] * 2 * self._cal_k_angle(self._d_mono, order=1, k=ki)
kf_angle = self._senses[2] * 2 * self._cal_k_angle(self._d_ana, order=1, k=kf)
return [ki_angle, phi, psi, kf_angle]
def is_valid(self, qe, *, with_angles=False):
try:
angles = self(qe)
return True if with_angles is False else (True, angles)
except:
return False if with_angles is False else (False, None)
def _lattice_rec_f(self):
astar = np.zeros(3)
alphastar = np.zeros(3)
V = self._cal_volume_real()
co = np.cos(self._angles * D2R)
si = np.sin(self._angles * D2R)
astar[0] = self._lattice[1] * self._lattice[2] * si[0] / V
astar[1] = self._lattice[2] * self._lattice[0] * si[1] / V
astar[2] = self._lattice[0] * self._lattice[1] * si[2] / V
alphastar[0] = np.arccos((co[1] * co[2] - co[0]) / (si[1] * si[2]))
alphastar[1] = np.arccos((co[0] * co[2] - co[1]) / (si[0] * si[2]))
alphastar[2] = np.arccos((co[0] * co[1] - co[2]) / (si[0] * si[1]))
return [astar, alphastar]
def _cal_volume_real(self):
co = np.cos(self._angles * D2R)
mul = self._lattice[0] * self._lattice[1] * self._lattice[2]
V = mul * np.sqrt(1 - np.dot(co, co) + 2 * co[0] * co[1] * co[2])
return V
def _matrix_crystal2lab(self):
B = np.zeros((3, 3))
U = np.zeros((3, 3))
B[0, 0] = self._lattice_rec[0] * 2 * np.pi
B[0, 1] = self._lattice_rec[1] * np.cos(self._angles_rec[2]) * 2 * np.pi
B[0, 2] = self._lattice_rec[2] * np.cos(self._angles_rec[1]) * 2 * np.pi
B[1, 0] = 0
B[1, 1] = self._lattice_rec[1] * np.sin(self._angles_rec[2]) * 2 * np.pi
B[1, 2] = -self._lattice_rec[2] * np.sin(self._angles_rec[1]) * \
np.cos(self._angles[0] * D2R) * 2 * np.pi
B[2, 0] = 0
B[2, 1] = 0
B[2, 2] = 2 * np.pi / self._lattice[2]
for i in range(3):
for j in range(3):
if -0.000001 < B[i, j] < 0.000001:
B[i, j] = 0.0
vec1 = np.dot(B, self._orient1)
vec2 = np.dot(B, self._orient2)
vec3 = np.cross(vec1, vec2)
vec2 = np.cross(vec3, vec1)
# xXX use self.coordinatesystem instead of hardcoded 1?
U = self._cal_Umatrix(vec1, vec2, vec3, 1)
return np.dot(U, B)
def _cal_Umatrix(self, vec1, vec2, vec3, direction):
U = np.zeros((3, 3))
U[2] = vec3 / np.linalg.norm(vec3)
if direction == 1:
U[0] = vec1 / np.linalg.norm(vec1)
U[1] = vec2 / np.linalg.norm(vec2)
elif direction == 2:
U[0] = vec2 / np.linalg.norm(vec2)
U[1] = vec1 / np.linalg.norm(vec1)
elif direction == -1:
U[0] = -vec1 / np.linalg.norm(vec1)
U[1] = vec2 / np.linalg.norm(vec2)
elif direction == -2:
U[0] = -vec2 / np.linalg.norm(vec2)
U[1] = vec1 / np.linalg.norm(vec1)
return U
def _hkl2qcart(self, h, k, l):
"""Return the cartesian coordinates of given (h,k,l) Miller indices."""
return np.dot(self._matrix, np.array([h, k, l], float))
def _hkl2qlab(self, h, k, l):
"""Transform a vector given in real lattice with (h,k,l) Miller indices
in Qlab with coordinates in system:
x in beam direction, z direction upwards, y making a right handed system.
"""
hklcart = self._hkl2qcart(h, k, l)
result = np.dot(hklcart, self._matrix_cardan)
if abs(result[2]) > 0.001:
raise Exception('out of plane vector; check your scattering plane')
return result
def _cal_Y(self, r1, r2, hkl):
"""Calculate angle between 1st orientation reflection and Q vector."""
# this function doesn't use r1 or r2...
crit = 0.000001
hkl = self._hkl2qlab(hkl[0], hkl[1], hkl[2])
qs = hkl[0]**2 + hkl[1]**2
qabs = np.sqrt(qs)
Y = hkl[1] / qabs
a1 = np.arcsin(Y)
if -crit < hkl[0] < crit:
Y = np.pi / 2
if hkl[1] < crit:
Y = -Y
elif hkl[0] < -crit:
Y = -a1
if hkl[1] > crit:
Y += np.pi
else:
Y -= np.pi
elif hkl[0] > crit:
Y = a1
Y *= R2D
return Y
def _cal_kf(self, e, ki):
"""Calculate the outgoing wavevector for given energy transfer and
incoming wavevector.
"""
kf = ki**2 - K * e
if kf > 0.000001:
kf = np.sqrt(kf)
return kf
else:
raise Exception('energy transfer of %s THz not possible '
'with k_i = %s' % (e, ki))
def _cal_phi(self, q, ki, kf, sense):
"""Return the sample scattering angle."""
qabs = np.linalg.norm(q)
temp = (ki**2 + kf**2 - qabs**2) / (2.0 * ki * kf)
if -1 <= temp <= 1:
phi = np.arctan2(np.sqrt(1 - temp**2), temp) * sense * R2D
return phi
else:
raise Exception('scattering triangle not closed when '
'calculating phi angle')
def _cal_alpha1(self, qlab, e, ki, sense):
"""Calculate the angle alpha (ki, Q) for given Qlab vector, energy
transfer, incoming wavevector and scattering sense (sample).
"""
qabs = np.linalg.norm(qlab)
temp = (qabs**2 + K * e) / (2 * qabs * ki)
if -1 <= temp < 1:
alpha = np.arctan2(np.sqrt(1 - temp**2), temp) * sense * R2D
return alpha
else:
raise Exception('energy transfer of %s THz not possible;'
' scattering triangle not closed' % e)
def _cal_alpha2(self, y, psi):
"""Calculate the angle alpha (ki, Q) for given angle Y (ki, r1)
and rotation angle (sample).
"""
alpha = y - psi
if alpha < -180:
alpha += 360
elif alpha > 180:
alpha -= 360
return alpha
def _cal_psi(self, y, alpha):
"""Calculate the rotation angle sample for given angle Y (ki, r1) and
angle alpha (ki, Q).
"""
psi = y - alpha
if psi < -180:
psi += 360
if psi > 180:
psi -= 360
return psi
def _cal_ki1(self, e, kf):
"""Calculate the incoming wavevector for given energy transfer and
outgoing wavevector.
"""
ki = kf**2 + K * e
if ki > 0.000001:
ki = np.sqrt(ki)
return ki
else:
raise Exception('energy transfer of %s THz not possible '
'with k_f = %s' % (e, kf))
def _cal_ki2(self, qlab, e, phi):
"""Calculate the incoming wavevector for given Qlab vector, energy
transfer and sample scattering angle.
"""
ki = 0
phi *= D2R
qabs = np.linalg.norm(qlab)
a1 = (qabs / np.sin(phi))**2
a2 = K * e
a3 = a2 / np.sin(phi)
a3 = a1**2 - a3**2
if a3 > 0:
ki = a1 + a2 + np.cos(phi) * np.sqrt(a3)
ki /= 2
if ki > 0.000001:
ki = np.sqrt(ki)
return ki
else:
raise Exception('energy transfer of %s THz not possible '
'with phi = %s' % (e, phi))
def _cal_ki3(self, qlab, e, alpha):
"""Calculate the incoming wavevector for given Qlab vector, energy
transfer and angle alpha(ki, Q).
"""
alpha *= D2R
qabs = np.linalg.norm(qlab)
ki = (qabs**2 + K * e) / (2.0 * qabs * np.cos(alpha))
if ki > 0.000001:
return ki
else:
raise Exception('energy transfer of %s THz not possible;'
' scattering triangle not closed' % e)
def _cal_zone(self):
return np.cross(self._orient1, self._orient2)
def _cal_k_angle(self, dvalue, order, k):
return math.degrees(math.asin(math.pi * order / (k * dvalue)))
[docs]
class AngleMapInducedMetric:
"""Class representing a metric on 2D plane induced by an angle map."""
def __init__(self, angle_map: AngleMap, axes, offset, instrument: Instrument) -> None:
self.angle_map = angle_map
self.axes = axes
self.offset = offset
self.instrument = instrument
def __call__(self, x1, x2) -> float:
qe_1 = x2qe(x1, self.axes, self.offset)
qe_2 = x2qe(x2, self.axes, self.offset)
angles_1 = np.asarray(self.angle_map(qe_1))
valid_2, angles_2 = self.angle_map.is_valid(qe_2, with_angles=True)
return max(abs((angles_1 - angles_2) / self.instrument.angle_velocities)) \
if valid_2 else np.Inf
[docs]
class Scenario:
"""Class representing an experimental scenario."""
def __init__(self, name: str, sample: Sample, orientation: Orientation,
axes, offset, limits,
instrument_config: InstrumentConfiguration,
instrument: Instrument, intensity_threshold: float,
counting_time: float, intensity_function) -> None:
self.name = name
self.sample = sample
self.orientation = orientation
self.axes = np.asarray(axes)
self.offset = np.asarray(offset)
self.limits = np.asarray(limits)
self.instrument_config = instrument_config
self.instrument = instrument
self.intensity_threshold = intensity_threshold
self.counting_time = counting_time
self.intensity_function = intensity_function
self.angle_map = AngleMap(sample=self.sample, orientation=self.orientation,
instrument_config=self.instrument_config,
instrument=self.instrument)
self.angle_map_metric = AngleMapInducedMetric(angle_map=self.angle_map,
axes=self.axes, offset=self.offset,
instrument=self.instrument)
@staticmethod
def load(filename: str):
with open(filename, 'r') as file:
json_test_case = json.load(file)
name = json_test_case['name']
json_sample = json_test_case['sample']
sample = Sample(name=json_sample['name'],
lattice=json_sample['lattice'],
angles=json_sample['angles'])
json_orientation = json_test_case['orientation']
orientation = Orientation(orient1=json_orientation['orient1'],
orient2=json_orientation['orient2'],
psi0=json_orientation['psi0'])
axes = json_test_case['axes']
offset = json_test_case['offset']
limits = json_test_case['limits']
json_instrument_config = json_test_case['instrument_configuration']
instrument_config = InstrumentConfiguration(scan_mode=json_instrument_config['scan_mode'],
scan_constant=json_instrument_config['scan_constant'],
senses=json_instrument_config['senses'])
json_instrument = json_test_case['instrument']
instrument = Instrument(angle_velocities=json_instrument['angle_velocities'],
d_mono=json_instrument['d_mono'],
d_ana=json_instrument['d_ana'])
intensity_threshold = json_test_case['intensity_threshold']
counting_time = json_test_case['counting_time']
json_intensity_function = json_test_case['intensity_function']
locations = json_intensity_function['locations']
intensities = json_intensity_function['intensities']
intensity_function = normalized_linear_interpolator(locations, intensities)
scenario = Scenario(name=name, sample=sample, orientation=orientation,
axes=axes, offset=offset, limits=limits,
instrument_config=instrument_config,
instrument=instrument, intensity_threshold=intensity_threshold,
counting_time=counting_time,
intensity_function=intensity_function)
return scenario
def prepare_locs_plot(locs, *, axes, offset=None):
"""Merge the coordinates in `locs` with `axes` and `offset` according to TAS conventions."""
if offset is None:
offset = 4 * [0]
assert np.shape(locs)[1] == len(axes)
assert np.shape(axes)[1] == len(offset) == 4
locs_plot = np.array(locs) # copy `locs`
axes = np.asarray(axes)
offset = np.asarray(offset)
if _is_inelastic_scan(axes):
if _is_classical_inelastic_scan(axes):
axis_q = axes[0, :3]
offset_q = offset[:3]
# Q axis
if _q_axis_compatible_with_offset(axis_q, offset_q):
locs_plot[:, 0] += _compatibility_factor_q_axis_offset(axis_q, offset_q)
# energy axis
locs_plot[:, 1] += offset[3]
else:
axis_q1 = axes[0, :3]
axis_q2 = axes[1, :3]
offset_q = offset[:3]
if _q_axis_compatible_with_offset(axis_q1, offset_q):
locs_plot[:, 0] += _compatibility_factor_q_axis_offset(axis_q1, offset_q)
if _q_axis_compatible_with_offset(axis_q2, offset_q):
locs_plot[:, 1] += _compatibility_factor_q_axis_offset(axis_q2, offset_q)
return locs_plot
def _q_axis_compatible_with_offset(axis, offset):
"""Check whether Q `axis` is compatible with `offset`.
Compatibility is equal to whether the vector of non-zero entries in `axis`
and the vector of correponding entries in `offset` are linearly dependent."""
return np.linalg.matrix_rank([axis[axis != 0], offset[axis != 0]]) == 1
def _compatibility_factor_q_axis_offset(axis, offset):
"""Compute coordinate of the point `offset` along the `axis`."""
return offset[axis != 0][0] / axis[axis != 0][0]
def _is_inelastic_scan(axes):
"""Check whether `axes` represents a scan with zero energy component."""
return np.any(axes[:, 3] == 1)
def _is_classical_inelastic_scan(axes):
"""Check whether `axes` represents a scan with one plain energy axis."""
return len([axis for axis in axes if _is_plain_energy_axis(axis)]) == 1
def _is_plain_energy_axis(axis):
"""Check whether only the energy component of `axis` is non-zero."""
return np.all(axis[:3] == 0) and axis[3] != 0
def axes_plot_labels(axes, *, offset=None, units=None):
"""Compute plotting labels for `axes`."""
if offset is None:
offset = len(axes[0]) * [0]
axes = np.asarray(axes)
offset = np.array(offset) # copy `offset`
assert np.shape(axes)[1] == len(offset) == 4
if _is_inelastic_scan(axes):
if _is_classical_inelastic_scan(axes):
axis_q = axes[0, :3]
offset_q = offset[:3]
if units is None:
units = ["r.l.u.", "meV"]
if _q_axis_compatible_with_offset(axis_q, offset_q):
offset_q[axis_q != 0] = 0
x1_label = _format_plot_label(_merge_axis_offset(axis_q, offset_q), unit=units[0])
x2_label = f"Energy transfer [{units[1]}]"
else:
x1_label = _format_plot_label(_merge_axis_offset(axes[0], offset))
x2_label = _format_plot_label(_merge_axis_offset(axes[1], offset))
else:
axis_q1 = axes[0, :3]
axis_q2 = axes[1, :3]
offset_q1 = np.array(offset[:3])
offset_q2 = np.array(offset[:3])
if units is None:
units = "r.l.u."
if _q_axis_compatible_with_offset(axis_q1, offset_q1):
offset_q1[axis_q1 != 0] = 0
if _q_axis_compatible_with_offset(axis_q2, offset_q2):
offset_q2[axis_q2 != 0] = 0
# ignore offset components belonging to the other axis
offset_q1[axis_q2 != 0] = 0
offset_q2[axis_q1 != 0] = 0
x1_label = _format_plot_label(_merge_axis_offset(axis_q1, offset_q1), unit=units)
x2_label = _format_plot_label(_merge_axis_offset(axis_q2, offset_q2), unit=units)
return x1_label, x2_label
_hkle_names = ['h', 'k', 'l', 'e']
def _merge_axis_offset(axis, offset):
"""Merge `axis` and `offset` for plotting label."""
def _merge_axis_offset_i(axis_i, offset_i, hkl_name):
# make integers if possible
if axis_i == int(axis_i):
axis_i = int(axis_i)
if offset_i == int(offset_i):
offset_i = int(offset_i)
s = ""
if axis_i != 0:
if offset_i != 0:
s += f"{offset_i}"
if axis_i > 0:
s += '+'
if abs(axis_i) != 1:
s += f"{axis_i}"
elif axis_i == -1:
s += '-'
s += hkl_name
else:
s += f"{offset_i}"
return s
i_hkl_name = np.argmax(axis != 0)
return [_merge_axis_offset_i(axis[i], offset[i], _hkle_names[i_hkl_name if i < 3 else 3]) for i in range(len(axis))]
def _format_plot_label(comps, unit=None):
return f"$({','.join(comps)})$" + (f" [{unit}]" if unit is not None else "")
