From 69e311cb85cfb058e9f2523a75bf8b4b113154a8 Mon Sep 17 00:00:00 2001 From: Alexander Steppke Date: Mon, 24 Nov 2025 15:05:43 +0100 Subject: [PATCH] Added RSM, in a first iteration --- src/cristallina/rsm.py | 901 +++++++++++++++++++++++++++++++++++++++++ 1 file changed, 901 insertions(+) create mode 100644 src/cristallina/rsm.py diff --git a/src/cristallina/rsm.py b/src/cristallina/rsm.py new file mode 100644 index 0000000..eb6d0a1 --- /dev/null +++ b/src/cristallina/rsm.py @@ -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