Cristallina/cristallina_analysis_package
0
1
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases 3 Wiki Activity

Added RSM, in a first iteration

Browse Source
This commit is contained in:
steppke_a
2025-11-24 15:05:43 +01:00
parent e95b5a1670
commit 69e311cb85
1 changed files with 901 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files

901
src/cristallina/rsm.py Normal file
View File

@@ -0,0 +1,901 @@
import re
from copy import copy
from collections import defaultdict, deque
from pathlib import Path
from typing import Optional
import logging
import time
import numpy as np
from matplotlib import pyplot as plt
from sfdata import SFDataFiles, sfdatafile, SFScanInfo
from joblib import Memory
from . import utils
from . import channels, analysis
from .utils import ROI
logger = logging.getLogger(__name__)
def setup_cachedirs(pgroup=None, cachedir=None):
"""
Sets the path to a persistent cache directory either from the given p-group (e.g. "p20841")
or an explicitly given directory.
If heuristics fail we use "/tmp" as a non-persistent alternative.
"""
if cachedir is not None:
# explicit directory given, use this choice
memory = Memory(cachedir, verbose=0, compress=2)
return memory
try:
if pgroup is None:
pgroup_no = utils.heuristic_extract_pgroup()
else:
parts = re.split(r"(\d.*)", pgroup) # ['p', '2343', '']
pgroup_no = parts[-2]
candidates = [f"/sf/cristallina/data/p{pgroup_no}/work",
f"/sf/cristallina/data/p{pgroup_no}/res",]
for cache_parent_dir in candidates:
if Path(cache_parent_dir).exists():
cachedir = Path(cache_parent_dir) / "cachedir"
break
except KeyError as e:
print(e)
cachedir = "/das/work/units/cristallina/p19739/cachedir"
try:
memory = Memory(cachedir, verbose=0, compress=2)
except PermissionError as e:
cachedir = "/tmp"
memory = Memory(cachedir, verbose=0, compress=2)
return memory
memory = setup_cachedirs()
@memory.cache() # we ignore batch_size for caching purposes
def RSM(run_number, sam, det, roi, offset=0, lower_cutoff_threshold=10, detector=channels.JF):
det = copy(det)
sam=copy(sam)
roi=copy(roi)
scan = utils.scan_info(run_number, use_uscan=True)
for i in range(2):
try:
channel_Exray='SAROP31-ODCC110:MOT_ENY.RBV'
energy=np.array(scan[0][channel_Exray])[1][0]
# detector 2 theta
channel_detangle='SARES31-GPS:ROT2THETA-BS'
tth=np.mean(np.array(scan[0][channel_detangle])[1])
break
except KeyError:
print("Waiting")
time.sleep(60)
readbacks=scan.readbacks
stacks, IO_norms = analysis.calculate_JF_stacks(scan,
lower_cutoff_threshold=lower_cutoff_threshold,
exclude_steps=None,
recompute=False,
detector=detector)
roi.values = np.array(stacks)[:, roi.rows, roi.cols]
roi.values = np.where(roi.values < lower_cutoff_threshold, 0, roi.values)
roi.values = roi.values/IO_norms[:,None, None]
roi.intensities = np.sum(roi.values[:], axis=(1, 2))
roi.intensities_normalized = roi.intensities/(roi.width*roi.height)
roi.intensities_I0 = roi.intensities/IO_norms
roicoord=[roi.left, roi.right, roi.bottom, roi.top]
coors_h_all=[]
coors_k_all=[]
coors_l_all=[]
inten_all=[]
det.move_angles(gamma=tth)
for i in range(len(readbacks)):
#
sam.rotate_sample(alpha=readbacks[i]+offset)
exp = Experiment(sam, det, energy)
# change the coordinates into h,k,l
hkl_temp=exp.ROI_hkl(roicoord)
if len(inten_all)==0:
coors_h_all=np.reshape(hkl_temp[0],-1)
coors_k_all=np.reshape(hkl_temp[1],-1)
coors_l_all=np.reshape(hkl_temp[2],-1)
inten_all=np.reshape(roi.values[i],-1)
else:
coors_h_all=np.concatenate((coors_h_all,np.reshape(hkl_temp[0],-1)))
coors_k_all=np.concatenate((coors_k_all,np.reshape(hkl_temp[1],-1)))
coors_l_all=np.concatenate((coors_l_all,np.reshape(hkl_temp[2],-1)))
inten_all=np.concatenate((inten_all,np.reshape(roi.values[i],-1)))
return coors_h_all, coors_k_all, coors_l_all, inten_all
@memory.cache()
def RSM_interpolate(coors_h_all, coors_k_all, coors_l_all, inten_all, hbound, kbound, gridsize):
lmin, lmax = coors_l_all.min(),coors_l_all.max()
l_grid=np.arange(lmin,lmax+gridsize,gridsize)
k_grid=np.arange(kbound[0], kbound[1]+gridsize, gridsize)
h_grid=np.arange(hbound[0], hbound[1]+gridsize, gridsize)
[H,K,L]=np.meshgrid(h_grid,k_grid,l_grid)
# interpolation
grided_all=griddata(np.transpose(np.array([coors_h_all,coors_k_all,coors_l_all])),inten_all,np.transpose(np.array([H,K,L])),method='nearest')
return grided_all
@memory.cache()
def RSM_full(run_number, sam, det, roi, hbound, kbound, gridsize, offset=0, lower_cutoff_threshold=10, detector=channels.JF):
det = copy(det)
sam=copy(sam)
roi=copy(roi)
coors_h_all, coors_k_all, coors_l_all, inten_all = RSM(run_number, sam, det, roi, offset=offset, lower_cutoff_threshold=lower_cutoff_threshold, detector=detector)
return RSM_interpolate(coors_h_all, coors_k_all, coors_l_all, inten_all, hbound, kbound, gridsize)
# based on henry's codes
#numpy imports
import numpy as np
import uncertainties.unumpy as unumpy
#visualization
import matplotlib.pyplot as plt
from matplotlib import animation, rc, gridspec
#scipy imports
from scipy.spatial.transform import Rotation as Rot
from scipy.optimize import brentq, curve_fit, minimize
from scipy.interpolate import griddata
from tqdm import tqdm
# data handling
import re
import concurrent.futures
import sys
import os
#constants
hbar_eV = 6.582119569e-16
c_mps = 2.99792458e8
#for ease of use later
xhat = np.array((1,0,0))
yhat = np.array((0,1,0))
zhat = np.array((0,0,1))
#Helper functions
def real_space_vectors(sample):
#this function takes alpha, beta, gamma and unit cell sizes and returns unit cell vectors a,b,c
#equations were copied from the internet.
#It works correctly for alp=bet=90deg (as in tas2). I didnt check for the other angles.
a=sample["a"]*xhat
b=sample["b"]*np.array([np.cos(sample["gam"]*np.pi/180),np.sin(sample["gam"]*np.pi/180),0])
c1=np.cos(sample["bet"]*np.pi/180)
c2=(np.cos(sample["alp"]*np.pi/180)-np.cos(sample["bet"]*np.pi/180)*np.cos(sample["gam"]*np.pi/180))/np.sin(sample["gam"]*np.pi/180)
c3=np.sqrt(1-(np.cos(sample["bet"]*np.pi/180))**2-c2**2)
c=sample["c"]*np.array([c1,c2,c3])
return a, b, c
def reciprocal_space_vectors(a, b, c):
#Takes real space basis, a, b and c and returns reciprocal space basis arec, brec, crec.
denom = a @ np.cross(b, c)
arec = 2*np.pi*np.cross(b, c)/denom
brec = 2*np.pi*np.cross(c, a)/denom
crec = 2*np.pi*np.cross(a, b)/denom
return arec, brec, crec
def rotate(vector, axis, angle, degrees = True):
#rotates vector around axis by angle degrees.
axis = axis/np.sqrt(axis @ axis) #normalize just in case
if degrees:
p = np.pi/180 #conversion constant
else:
p = 1 #no conversion needed
r = Rot.from_rotvec(axis*angle*p)
return r.apply(vector)
def E_to_k(energy):
"""
Converts energy to wavevector magnitude
Input:
energy : float, energy of incoming radiation in eV
Output:
wavevector : wavevector magnitude in 1/nm
"""
return 1e-9*energy/(hbar_eV*c_mps)
class Detector():
'''
Description:
Holds Information about the detector, specifically for a Bernina Experiment. Detector position and tilt can be given as well as pixel size and offset between
detector panels. Contains useful functions that return the real space position of each pixel on the detector and visualization tools that show the detector in
real space.
Inputs:
Location:
gamma : float, Angle in degrees
delta : float, Angle in degrees
R_m : float, Distance from sample in meters
Location:
center_pixel : tuple, Center pixel indices
px_size_um : float, Side length of a pixel in um
Orientation:
psi : float, Angle representing the rotation of detector around detector normal
beta : float, Angle representing rotation around x-pixel axis
xi : float, Angle representing rotation around y-pixel axis
'''
### Setup ###
def __init__(self, R_m, px_size_um, center_pixel, detector_size_px = (1030,514), gamma = 0, delta = 0, psi = 0, beta = 0, xi = 0):
self.px_size_um = px_size_um
self.center_pixel = center_pixel
self.detector_size_px = detector_size_px
self.phi = psi
self.beta = beta
self.xi = xi
self.R_m = R_m
self.pixel_positions_real = np.array(()) #for plotting later
self.move_angles(gamma, delta)
self.get_cartesian_coordinates()
def move_angles(self, gamma = None, delta = None, phi = None, beta = None, xi = None):
'''
Rotates the detector to gamma and delta. Redefines r_m and delta_hat and gamma_hat
Input:
Location:
gamma : float, Rotation angle in azimuthal direction
delta : float, Rotation angle in polar direction
Orientation:
phi : float, Angle representing the rotation of detector around detector normal
beta : float, Angle representing rotation around x-pixel axis
xi : float, Angle representing rotation around y-pixel axis
'''
if gamma != None:
self.gamma = gamma
if delta != None:
self.delta = delta
if phi != None:
self.phi = phi
if beta != None:
self.beta = beta
if xi != None:
self.xi = xi
self.get_cartesian_coordinates()
gamma_rad = np.pi*self.gamma/180
delta_rad = np.pi*self.delta/180
#change delta_hat, gamma_hat perpendicular to R
delta_hat_perp = np.array((np.sin(delta_rad)*np.cos(gamma_rad), np.sin(delta_rad)*np.sin(gamma_rad),np.cos(delta_rad)))
gamma_hat_perp = np.array((np.sin(gamma_rad), -np.cos(gamma_rad),0))
#Rotate delta_hat, gamma_hat by phi
delta_hat_p = rotate(delta_hat_perp, self.cart_vec, self.phi, degrees = True)
gamma_hat_p = rotate(gamma_hat_perp, self.cart_vec, self.phi, degrees = True)
#rotate gamma hat prime by beta around delta hat
self.gamma_hat = rotate(gamma_hat_p, delta_hat_p, self.beta, degrees = True)
#rotate delta hat prime by beta around gamma hat
self.delta_hat = rotate(delta_hat_p, self.gamma_hat, self.xi, degrees = True)
def move_to_cartesian(self,vec):
'''
Moves the detector to be in the path of a cartesian vector.
Note, off angles dont work here
Input:
vec : 1D np.ndarray of size 3, location to move the detector in cartesian coordinates
'''
projxy = (vec @ xhat)*xhat + (vec @ yhat)*yhat #project k_out onto xy plane for gamma
projxy = projxy/np.sqrt(projxy@projxy) #normalize
projz = (vec@zhat)*zhat
projz = projz/np.sqrt(projz@projz)
vec_normalize = vec/np.sqrt(vec@vec)
self.delta = np.sign(zhat @ vec) * np.arcsin(projz@vec_normalize) * 180/np.pi
self.gamma = -np.sign(yhat @ vec) * np.arccos(-projxy@xhat) * 180/np.pi
self.move_angles(self.gamma, self.delta)
### Computational Tools ###
def get_cartesian_coordinates(self):
"""
Returns cartesian coordinates of the detector as a vector.
"""
gamma_rad = np.pi*self.gamma/180
delta_rad = np.pi*self.delta/180
self.cart_vec = self.R_m*(-np.cos(gamma_rad)*np.cos(delta_rad)*xhat-np.sin(gamma_rad)*np.cos(delta_rad)*yhat+np.sin(delta_rad)*zhat)
def pixel_cartesian_coordinate(self, i_y, i_x):
"""
Find the cartesian coordinate of a pixel with index i_y and i_x
Inputs:
i_y : int, defined above
i_x : int, defined above
"""
i_x = self.center_pixel[0] - i_x
i_y = self.center_pixel[1] - i_y
pixel_position_relative = 1e-6*self.px_size_um*(np.array((i_x*self.gamma_hat[0], i_x*self.gamma_hat[1], i_x*self.gamma_hat[2])) + np.array((i_y*self.delta_hat[0], i_y*self.delta_hat[1], i_y*self.delta_hat[2])))
pixel_position_real = self.cart_vec + pixel_position_relative
return pixel_position_real
def whole_detector_cartesian_coordinates(self):
"""
Calculates the real space coordinates for each pixel of the detector in its defined postion as a numpy array.
"""
i_x = self.center_pixel[0] - np.arange(self.detector_size_px[0])
i_y = self.center_pixel[1] - np.arange(self.detector_size_px[1])
ii_x, ii_y = np.meshgrid(i_x, i_y)
#probably a better way to do this part
sum_one = 1e-6*self.px_size_um*np.array((ii_x*self.gamma_hat[0], ii_x*self.gamma_hat[1], ii_x*self.gamma_hat[2]))
sum_two = 1e-6*self.px_size_um*np.array((ii_y*self.delta_hat[0], ii_y*self.delta_hat[1], ii_y*self.delta_hat[2]))
pixel_positions_relative = sum_one + sum_two
pixel_positions_real = self.cart_vec + np.transpose(pixel_positions_relative, [1,2,0])
self.pixel_positions_real = np.transpose(pixel_positions_real,[2,0,1])
def ROI_cartesian_coordinates(self, ROI):
"""
Calculates the real space coordinates for each pixel of an ROI of adetector in its defined postion as a numpy array.
Input:
ROI : tuple of size 4, in [xmin, xmax, ymin, ymax]
Output:
pixel_positions_ROI : 3D np.ndarray of shape (3x(2*y_half)x(2*x_half)) holding the real space position of each pixel in ROI.
"""
self.whole_detector_cartesian_coordinates()
x_min, x_max, y_min, y_max = ROI
pixel_positions_ROI = self.pixel_positions_real[:,y_min:y_max, x_min:x_max]
return pixel_positions_ROI
class Sample():
'''
Description:
Stores information about sample geometry and orientation and contains useful computational tools as well as visualization of real/reciprocal space lattice
basis vectors.
Inputs:
Required:
sample_dict : dictionary, containing lattice constants and angles.
Example, {"a":0.336, "b":0.336, "c": 0.5853, "alp":90, "bet":90, "gam":120}
sample_name : string, Name of sample
Example, "1T-TaS2"
sample_normal : 1D np.ndarray of size 3, vector defining the sample normal in real space vector basis.
Not Required:
alpha : float, Angle in degrees, angle according to experiment, roughly correct.
ov : float Angle in degrees, angle according to experiment, relatively correct between runs.
ov0: float, Angle in degrees, how much the sample is offset to the experimental angles in ov
chi0: float, Angle in degrees, how much the sample is offset to the experimental angles in chi
alpha0: float, Angle in degrees, how much the sample is offset to the experimental angles in alpha
'''
### SETUP ###
def __init__(self, sample_dict, sample_name, surface_normal = np.array((0,0,1)), ov0 = 0, chi0 = 0, alpha0 = 0, alpha = 0, ov = 0, chi = 0):
self.sample_dict = sample_dict
self.sample_name = sample_name
self.surface_normal = surface_normal
self.alpha0 = alpha0
self.ov0 = ov0
self.chi0 = chi0
self.define_off_angles(alpha0, ov0, chi0)
self.initialize_lattice_vectors()
self.alpha = alpha
self.ov = ov
self.chi = chi
self.alpha_axis = np.array([0,0,1])
self.ov_axis = rotate(yhat, self.alpha_axis, self.alpha, degrees=True)
self.rotate_sample(alpha, ov)
def define_off_angles(self, alpha0 = 0, ov0 = 0, chi0 = 0):
"""
Initialize sample with off angles
"""
self.alpha0 = alpha0
self.ov0 = ov0
self.chi0 = chi0
self.alpha0_axis = zhat
self.ov0_axis = rotate(yhat, self.alpha0_axis, self.alpha0, degrees=True)
chi0_axis = rotate(xhat, self.alpha0_axis, self.alpha0, degrees=True)
self.chi0_axis = rotate(chi0_axis, self.ov0_axis, self.ov0, degrees=True)
self.initialize_lattice_vectors()
def change_lattice_constants(self, a = None, b = None, c = None):
if a != None:
self.sample_dict["a"] = a
if b != None:
self.sample_dict["b"] = b
if c != None:
self.sample_dict["c"] = c
self.initialize_lattice_vectors()
self.rotate_sample(self.alpha, self.ov)
def initialize_lattice_vectors(self):
"""
Set up the real space vectors and reciprocal space vectors
"""
self.a_unrotated, self.b_unrotated, self.c_unrotated = real_space_vectors(self.sample_dict) #basis vectors unrotated
self.arec0_unrotated, self.brec0_unrotated, self.crec0_unrotated = reciprocal_space_vectors(self.a_unrotated, self.b_unrotated, self.c_unrotated)
self.surface_normal_init = self.a_unrotated * self.surface_normal[0] + self.b_unrotated * self.surface_normal[1] + self.c_unrotated * self.surface_normal[2]
self.surface_normal_init = self.surface_normal_init/np.linalg.norm(self.surface_normal_init)
#get rotated basis vectors
self.a0 = self.vector_setup(self.a_unrotated)
self.b0 = self.vector_setup(self.b_unrotated)
self.c0 = self.vector_setup(self.c_unrotated)
self.surface_normal_init = self.vector_setup(self.surface_normal_init)
#compute reciprocal space vectors
self.arec0, self.brec0, self.crec0 = reciprocal_space_vectors(self.a0, self.b0, self.c0)
def vector_setup(self, vector):
"""
Input
vector: np.ndarray of size 3
Output
vector: np.ndarray of size 3
Helper function for initialize_lattice_vectors, rotates each basis vector to the correct position.
"""
#first deal with surface normal
surface_normal_vec = self.a_unrotated * self.surface_normal[0] + self.b_unrotated * self.surface_normal[1] + self.c_unrotated * self.surface_normal[2]
surface_normal_vec = surface_normal_vec/np.linalg.norm(surface_normal_vec)
angle = np.arccos(surface_normal_vec @ zhat)*180/np.pi
axis = np.cross(surface_normal_vec, zhat)/np.linalg.norm(np.cross(surface_normal_vec, zhat))
vector = rotate(vector, axis, angle, degrees = True)
#rotate to beamline coordinates (horizontal)
vector = rotate(vector, xhat, 90, degrees = True)
vector = rotate(vector, yhat, 90, degrees = True)
#rotate to account for misalignment
vector = rotate(vector, self.alpha0_axis, self.alpha0, degrees = True)
vector = rotate(vector, self.ov0_axis, self.ov0, degrees = True)
vector = rotate(vector, self.chi0_axis, self.chi0, degrees = True)
return vector
### ROTATE SAMPLE ###
def rotate_sample_vector(self, vector):
"""
Input
vector: np.ndarray of size 3
Output
vector: np.ndarray of size 3
Description
Rotates a sample vector by self.alpha and self.ov
"""
vector = rotate(vector, self.alpha_axis, self.alpha)
vector = rotate(vector, self.ov_axis, self.ov)
vector = rotate(vector, self.chi_axis, self.chi)
return vector
def rotate_sample(self, alpha = None, ov = None, chi = None):
'''
Rotates the sample by alpha and then by ov.
'''
if alpha != None:
self.alpha = alpha
if ov != None:
self.ov = ov
if chi!= None:
self.chi = chi
self.ov_axis = rotate(yhat, self.alpha_axis, self.alpha, degrees=True) #update ovrotate axis for rotation
chi_axis = rotate(xhat, self.alpha_axis, self.alpha, degrees=True)
self.chi_axis = rotate(chi_axis, self.ov_axis, self.ov, degrees = True)
self.arec = self.rotate_sample_vector(self.arec0)
self.brec = self.rotate_sample_vector(self.brec0)
self.crec = self.rotate_sample_vector(self.crec0)
self.a = self.rotate_sample_vector(self.a0)
self.b = self.rotate_sample_vector(self.b0)
self.c = self.rotate_sample_vector(self.c0)
self.surface_normal_updated = self.rotate_sample_vector(self.surface_normal_init)
### Computational Tools ###
def peak_q(self, peak_indices):
'''
Calculated the coordinates for a specified peak (h, k, l) in reciprocal space.
Inputs:
peak_indices : tuple/array of size 3, (h, k, l) indices.
Example: (0, 1, 4)
Outputs:
q : 1D np.ndarray of size 3, cartesian momentum space vector of given peak indices
'''
q = peak_indices[0]*self.arec + peak_indices[1]*self.brec + peak_indices[2]*self.crec
return q
def get_basis_change_matrix(self, return_inverse = False, form = "rotated"):
"""
Calculates matrix to go from cartesian basis to hkl basis
Inputs:
return_inverse : boolean, whether to return inverse of default matrix.
Options:
form : string, whether to return rotated basis or unrotated basis. Either "rotated" or "unrotated"
"""
if form == "rotated":
A = np.transpose(np.array((self.arec, self.brec, self.crec)))
if form == "unrotated":
a = self.arec0_unrotated
b = self.brec0_unrotated
c = self.crec0_unrotated
A = np.transpose(np.array((a, b, c)))
if return_inverse:
Ainv = np.linalg.inv(A)
return Ainv
else:
return A
def convert_hkl_to_Q(self, hkl):
return hkl[0]*self.arec0_unrotated+hkl[1]*self.brec0_unrotated+hkl[2]*self.crec0_unrotated
### Visualization Tools ###
def quiver_plot(self, space = "Reciprocal", figax = None):
"""
Plots real and reciprocal space crystal basis vectors in 3D. Good for debugging.
Inputs:
Not Required:
space : string, either "Reciprocal" or "Real". By default, space = "Reciprocal"
Returns:
fig : matplotlib figure
ax : matplotlib axis
"""
if figax == None:
fig = plt.figure(figsize = (8,8))
ax = fig.add_subplot(projection='3d')
else:
fig = figax[0]
ax = figax[1]
if space in ["Real", "All"]:
ax.quiver(0, 0, 0, *self.a,color = "red", length = 0.05, normalize = True, label = "Real Space Basis")
ax.quiver(0, 0, 0, *self.b,color = "red", length = 0.05, normalize = True)
ax.quiver(0, 0, 0, *self.c,color = "red", length = 0.05, normalize = True)
# ax.quiver(0, 0, 0, *self.surface_normal_updated,color = "green", length = 0.05, normalize = True, label ="Surface Normal")
if space in ['Reciprocal', "All"]:
ax.quiver(0, 0, 0, *self.arec,color = "blue", length = np.linalg.norm(self.arec), normalize = True, label = "Reciprocal Space Basis")
ax.quiver(0, 0, 0, *self.brec,color = "blue", length = np.linalg.norm(self.brec), normalize = True)
ax.quiver(0, 0, 0, *self.crec,color = "blue", length = np.linalg.norm(self.crec), normalize = True)
ax.quiver(0, 0, 0, *self.surface_normal_updated,color = "green", length = 35, normalize = True, label ="Surface Normal")
ax.set_xlim((-50, 50))
ax.set_ylim((-50, 50))
ax.set_zlim((-50, 50))
ax.set_aspect('auto')
ax.set_xlabel("x")
ax.set_ylabel("y")
ax.set_zlabel("z")
ax.legend()
return fig, ax
class Experiment():
'''
Description:
Given a sample, a detector, and the energy of an incoming beam, this class contains important tools to calculate scattering conditions to
find a specific peak or functions to calculate the momentum transfer of each pixel for use in the Reconstruction class.
Input:
Sample : Sample object for the experiment
Detector : Detector object for the experiment
energy : float, Energy of the incoming beam
'''
### Setup ###
def __init__(self, Sample, Detector, energy):
self.energy = energy
self.Sample = Sample
self.Detector = Detector
self.set_energy(energy) #incoming beam
self.pixel_hkl = np.array(())
### Changes in Experimental Setup ###
def make_angle_dict(self):
"""
Creates a dictionary containing alpha, ov, delta, and gamma for the experiment in its current state.
Output:
angle_dict : dictionary, holds angle info of experiment
"""
angle_dict = {}
angle_dict["alpha"] = self.Sample.alpha
angle_dict["ov"] = self.Sample.ov
angle_dict["delta"] = self.Detector.delta
angle_dict["gamma"] = self.Detector.gamma
return angle_dict
def move_from_string(self, string, value):
if string in ["energy", "energies"]:
self.set_energy(value)
elif string in ["alp", "alphas", "alpha"]:
self.Sample.rotate_stample(alpha = value)
elif string in ["ov", "phi"]:
self.Sample.rotate_stample(ov = value)
elif string in ["delta", "delt", "deltas"]:
self.Detector.move_angles(delta = value)
elif string in ["gamma", "gam", "gammas"]:
self.Detector.move_angles(gamma = value)
else:
print("String not recognized, choose from ['gamma', 'delta', 'energy', 'alpha', 'phi']")
def move_from_dict(self, angle_dict):
"""
Takes in a dictionary containing angles. Moves the sample and detector accordingly
Input:
angle_dict : dictionary, holds angle info of experiment
"""
alpha = angle_dict["alpha"]
ov = angle_dict["ov"]
self.Sample.rotate_sample(ov = ov, alpha = alpha)
delta = angle_dict["delta"]
gamma = angle_dict["gamma"]
self.Detector.move_angles(delta = delta, gamma = gamma)
def set_energy(self, energy):
self.energy = energy
self.k_in = E_to_k(self.energy)*np.array([-1,0,0]) #incoming beam
### Computational Tools ###
def whole_detector_hkl(self):
"""
Calculate momentum transfer in HKL basis for the entire detector.
Output:
pixel_hkl : 3d np.ndarray of shape (3,y_size,x_size)
"""
#real space coordinates
self.Detector.whole_detector_cartesian_coordinates()
#calculate momentum transfer
pixel_k_out = np.linalg.norm(self.k_in)*self.Detector.pixel_positions_real/np.linalg.norm(self.Detector.pixel_positions_real, axis = 0)
self.pixel_Q = -np.transpose(self.k_in - np.transpose(pixel_k_out,[1,2,0]),[2,0,1])
#solve for hkl of each pixel
Ainv = self.Sample.get_basis_change_matrix(return_inverse = True)
self.pixel_hkl = np.tensordot(Ainv, self.pixel_Q, axes = ([1],[0]))
return self.pixel_hkl
def ROI_hkl(self, ROI):
"""
Calculate momentum transfer in HKL basis for ROI of the detector.
Input:
ROI : tuple of size 4, in [xmin, xmax, ymin, ymax]
Output:
pixel_hkl : 3d np.ndarray of shape (3, 2*y_half, 2*x_half)
"""
real_space_coords = self.Detector.ROI_cartesian_coordinates(ROI)
pixel_k_out = np.linalg.norm(self.k_in)*real_space_coords/np.linalg.norm(real_space_coords, axis = 0)
pixel_Q = -np.transpose(self.k_in - np.transpose(pixel_k_out,[1,2,0]),[2,0,1])
#solve for hkl of each pixel
Ainv = self.Sample.get_basis_change_matrix(return_inverse = True)
pixel_hkl = np.tensordot(Ainv, pixel_Q, axes = ([1],[0]))
return pixel_hkl
def ROI_Q(self, ROI):
"""
Calculate momentum transfer for ROI of the detector
Input:
ROI : tuple of size 4, in [xmin, xmax, ymin, ymax]
Output:
pixel_hkl : 3d np.ndarray of shape (3, 2*y_half, 2*x_half)
"""
real_space_coords = self.Detector.ROI_cartesian_coordinates(ROI)
pixel_k_out = np.linalg.norm(self.k_in)*real_space_coords/np.linalg.norm(real_space_coords, axis = 0)
pixel_Q = -np.transpose(self.k_in - np.transpose(pixel_k_out,[1,2,0]),[2,0,1])
return pixel_Q
def ROI_Q_angle(self, ROI, theta=0):
"""
Calculate momentum transfer for ROI of the detector, assuming there is a coordinate system which rotates with theta
Input:
ROI : tuple of size 4, in [xmin, xmax, ymin, ymax]
theta : sample rotation angle in degree
Output:
pixel_hkl : 3d np.ndarray of shape (3, 2*y_half, 2*x_half)
"""
real_space_coords = self.Detector.ROI_cartesian_coordinates(ROI)
pixel_k_out = np.linalg.norm(self.k_in)*real_space_coords/np.linalg.norm(real_space_coords, axis = 0)
pixel_Q = -np.transpose(self.k_in - np.transpose(pixel_k_out,[1,2,0]),[2,0,1])
theta_rad=theta/180*np.pi
pixel_Q_afterrotate=np.array(pixel_Q)
pixel_Q_afterrotate[0]=pixel_Q[0]*np.cos(theta_rad)+pixel_Q[1]*np.sin(theta_rad)
pixel_Q_afterrotate[1]=-pixel_Q[0]*np.sin(theta_rad)+pixel_Q[1]*np.cos(theta_rad)
return pixel_Q_afterrotate
def point_to_hkl(self, ind_y, ind_x):
"""
Calculate momentum transfer in HKL basis for a single pixel of the detector given ind_y and ind_x.
Inputs:
ind_y : int, defined above
ind_x : int, defined above
"""
# Maps a single pixel to hkl coordinate.
pixel_coords_real = self.Detector.pixel_cartesian_coordinate(ind_y, ind_x)
k_out = np.linalg.norm(self.k_in)*pixel_coords_real/np.linalg.norm(pixel_coords_real)
Q = -(self.k_in - k_out)
Ainv = self.Sample.get_basis_change_matrix(return_inverse = True)
return Ainv @ Q
def visualize_scattering_xyz(self):
"""
Makes a plot showing the scattering vectors in reciprocal space in the xyz basis.
"""
fig, ax = self.Sample.quiver_plot(space = "Reciprocal")
try:
detector_position = self.Detector.cart_vec
except:
detector_position = self.Detector.center_px_position
ax.quiver(np.linalg.norm(self.k_in),0,0,*(self.k_in), length = np.linalg.norm(self.k_in), normalize = True, color = "purple", label = "Incoming beam (K_in)")
k_out_plotting = np.linalg.norm(self.k_in)*detector_position/np.linalg.norm(detector_position)
ax.quiver(0,0,0,*(k_out_plotting), length = np.linalg.norm(k_out_plotting), normalize = True, color = "black", label = "Detector Location (K_out)")
Q_plotting = self.k_in - k_out_plotting
ax.quiver(0,0,0,*(-Q_plotting),length = np.linalg.norm(Q_plotting), normalize = True, color = "orange", label = "Momentum Transfer (Q)")
self.whole_detector_hkl()
Q_pix = self.pixel_Q[:,::50,::50].reshape(3,-1)
ax.scatter(Q_pix[0,:], Q_pix[1,:], Q_pix[2,:], marker = ".", color = "brown")
center = self.pixel_Q[:,self.Detector.center_pixel[1], self.Detector.center_pixel[0]]
ax.scatter(center[0], center[1], center[2], color = "red")
ax.legend()
return fig, ax