import time
import numpy as np
import os
from rich import box
from rich.console import Console
from rich.table import Table
from typeguard import typechecked
from bec_lib import bec_logger
logger = bec_logger.logger
[docs]
class OMNYAlignmentError(Exception):
pass
[docs]
class OMNYAlignmentMixin:
default_correction_file = "correction_omny_202204.txt"
default_correction_file_x = "correction_omny_202204_x.txt"
[docs]
def reset_correction(self, use_default_correction=True):
"""
Reset the correction to the default values.
If use_default_correction is False, the correction will be set to empty values.
Otherwise the default values will be loaded.
Args:
use_default_correction (bool, optional): If set to true, a call reset the correction to the default values. Defaults to True.
"""
self.corr_pos_x = []
self.corr_angle_x = []
self.corr_pos_y = []
self.corr_angle_y = []
self.corr_pos_y_2 = []
self.corr_angle_y_2 = []
if use_default_correction:
try:
self.read_additional_correction_x(self.default_correction_file_x)
logger.info(f"Applying default x correction from {self.default_correction_file_x}")
except FileNotFoundError:
logger.warning(
f"Could not find default correction file {self.default_correction_file_x}."
)
logger.warning("Not applying any correction.")
try:
self.read_additional_correction_y(self.default_correction_file)
logger.info(f"Applying default y correction from {self.default_correction_file}")
except FileNotFoundError:
logger.warning(
f"Could not find default correction file {self.default_correction_file}."
)
logger.warning("Not applying any correction.")
def reset_tomo_alignment_fit(self):
self.client.delete_global_var("tomo_alignment_fit")
[docs]
def read_alignment_offset(
self,
dir_path=os.path.expanduser("~/Data10/specES1/internal/"),
setup="omny",
use_vertical_default_values=True,
):
"""
Read the alignment offset from the given directory and set the global parameter
tomo_alignment_fit.
Args:
dir_path (str, optional): The directory to read the alignment offset from. Defaults to os.path.expanduser("~/Data10/specES1/internal/").
"""
tomo_alignment_fit = np.zeros((2, 5))
with open(os.path.join(dir_path, "ptychotomoalign_Ax.txt"), "r") as file:
tomo_alignment_fit[0][0] = file.readline()
with open(os.path.join(dir_path, "ptychotomoalign_Bx.txt"), "r") as file:
tomo_alignment_fit[0][1] = file.readline()
with open(os.path.join(dir_path, "ptychotomoalign_Cx.txt"), "r") as file:
tomo_alignment_fit[0][2] = file.readline()
with open(os.path.join(dir_path, "ptychotomoalign_Ay.txt"), "r") as file:
tomo_alignment_fit[1][0] = file.readline()
with open(os.path.join(dir_path, "ptychotomoalign_By.txt"), "r") as file:
tomo_alignment_fit[1][1] = file.readline()
with open(os.path.join(dir_path, "ptychotomoalign_Cy.txt"), "r") as file:
tomo_alignment_fit[1][2] = file.readline()
with open(os.path.join(dir_path, "ptychotomoalign_Ay3.txt"), "r") as file:
tomo_alignment_fit[1][3] = file.readline()
with open(os.path.join(dir_path, "ptychotomoalign_Cy3.txt"), "r") as file:
tomo_alignment_fit[1][4] = file.readline()
print("New alignment parameters loaded:")
print(
f"X Amplitude {tomo_alignment_fit[0][0]}, "
f"X Phase {tomo_alignment_fit[0][1]}, "
f"X Offset {tomo_alignment_fit[0][2]}, "
f"Y Amplitude {tomo_alignment_fit[1][0]}, "
f"Y Phase {tomo_alignment_fit[1][1]}, "
f"Y Offset {tomo_alignment_fit[1][2]}, "
f"Y 3rd Order Amplitude {tomo_alignment_fit[1][3]}, "
f"Y 3rd Order Phase {tomo_alignment_fit[1][4]} ."
)
if use_vertical_default_values:
print(
f"Using default values for vertical alignment for setup {setup}. Optional: use_vertical_default_values=False"
)
if setup == "flomni":
tomo_alignment_fit[1][0] = 0
tomo_alignment_fit[1][1] = 0
tomo_alignment_fit[1][2] = 0
tomo_alignment_fit[1][3] = 0
tomo_alignment_fit[1][4] = 0
elif setup == "omny":
tomo_alignment_fit[1][0] = 2.588628
tomo_alignment_fit[1][1] = -2.385422
tomo_alignment_fit[1][2] = 0
tomo_alignment_fit[1][3] = 1.010583
tomo_alignment_fit[1][4] = -1.359157
print("Follwing parameters will be used:")
print(
f"X Amplitude {tomo_alignment_fit[0][0]}, "
f"X Phase {tomo_alignment_fit[0][1]}, "
f"X Offset {tomo_alignment_fit[0][2]}, "
f"Y Amplitude {tomo_alignment_fit[1][0]}, "
f"Y Phase {tomo_alignment_fit[1][1]}, "
f"Y Offset {tomo_alignment_fit[1][2]}, "
f"Y 3rd Order Amplitude {tomo_alignment_fit[1][3]}, "
f"Y 3rd Order Phase {tomo_alignment_fit[1][4]} ."
)
self.client.set_global_var("tomo_alignment_fit", tomo_alignment_fit.tolist())
# x amp, phase, offset, y amp, phase, offset, 3rd order amp, 3rd order phase
# 0 0 0 1 0 2 1 0 1 1 1 2 1 3 1 4
[docs]
def get_alignment_offset(self, angle: float):
"""
Compute the alignment offset for the given angle.
Args:
angle (float): The angle to compute the alignment offset for.
Returns:
tuple: The alignment offset in x, y and z direction.
"""
tomo_alignment_fit = self.client.get_global_var("tomo_alignment_fit")
if tomo_alignment_fit is None:
print("Not applying any alignment offsets. No tomo alignment fit data available.\n")
return (0, 0, 0)
# x amp, phase, offset, y amp, phase, offset
# 0 0 0 1 0 2 1 0 1 1 1 2
correction_x = (
tomo_alignment_fit[0][0] * np.sin(np.radians(angle) + tomo_alignment_fit[0][1])
+ tomo_alignment_fit[0][2]
)
correction_y = (
tomo_alignment_fit[1][0] * np.sin(np.radians(angle) + tomo_alignment_fit[1][1])
+ tomo_alignment_fit[1][2]
+ tomo_alignment_fit[1][3] * np.sin(3 * np.radians(angle) + tomo_alignment_fit[1][4])
)
correction_z = tomo_alignment_fit[0][0] * np.sin(
np.radians(angle + 90) + tomo_alignment_fit[0][1]
)
print(
f"Alignment offset x {correction_x}, y {correction_y}, z {correction_z} for angle"
f" {angle}\n"
)
return (correction_x, correction_y, correction_z)
def _read_correction_file(self, correction_file: str):
with open(correction_file, "r") as f:
num_elements = f.readline()
int_num_elements = int(num_elements.split(" ")[2])
corr_pos = []
corr_angle = []
for j in range(int_num_elements * 2):
line = f.readline()
value = line.split(" ")[2]
name = line.split(" ")[0].split("[")[0]
if name == "corr_pos":
corr_pos.append(float(value))
elif name == "corr_angle":
corr_angle.append(float(value))
print(
f"Loading default mirror correction from file {correction_file} containing {int_num_elements} elements."
)
# print(corr_pos)
return corr_pos, corr_angle
def read_additional_correction_x(self, correction_file: str):
self.corr_pos_x, self.corr_angle_x = self._read_correction_file(correction_file)
def read_additional_correction_y(self, correction_file: str):
self.corr_pos_y, self.corr_angle_y = self._read_correction_file(correction_file)
def read_additional_correction_y_2(self, correction_file: str):
self.corr_pos_y_2, self.corr_angle_y_2 = self._read_correction_file(correction_file)
def compute_additional_correction_x(self, angle):
return self._compute_additional_correction(angle, iteration="x1")
def compute_additional_correction_y(self, angle):
return self._compute_additional_correction(angle, iteration="y1")
def compute_additional_correction_y_2(self, angle):
return self._compute_additional_correction(angle, iteration="y2")
def _compute_additional_correction(self, angle, iteration="y1"):
if iteration == "x1":
corr_pos = self.corr_pos_x
corr_angle = self.corr_angle_x
elif iteration == "y1":
corr_pos = self.corr_pos_y
corr_angle = self.corr_angle_y
elif iteration == "y2":
corr_pos = self.corr_pos_y_2
corr_angle = self.corr_angle_y_2
if not corr_pos:
print(f"Not applying any additional correction {iteration}. No data available.\n")
return 0
# find index of closest angle
for j, _ in enumerate(corr_pos):
newangledelta = np.fabs(corr_angle[j] - angle)
if j == 0:
angledelta = newangledelta
additional_correction_shift = corr_pos[j]
continue
if newangledelta < angledelta:
additional_correction_shift = corr_pos[j]
angledelta = newangledelta
if additional_correction_shift == 0 and angle > corr_angle[-1]:
additional_correction_shift = corr_pos[-1]
print(f"Additional correction shift {iteration}: {additional_correction_shift}")
return additional_correction_shift