From d8fe838fb92f638375d770da8bc161fb7a048e42 Mon Sep 17 00:00:00 2001 From: Viljar Femoen Date: Fri, 19 Sep 2025 16:33:13 +0200 Subject: [PATCH 1/6] Simplify simulation: no dependencies, single file, simple scope --- src/instamatic/camera/camera.py | 2 +- src/instamatic/config/camera/simulate.yaml | 14 + src/instamatic/simulation/camera.py | 351 ++++++++++++++++++--- 3 files changed, 316 insertions(+), 51 deletions(-) diff --git a/src/instamatic/camera/camera.py b/src/instamatic/camera/camera.py index fa806b07..902292ea 100644 --- a/src/instamatic/camera/camera.py +++ b/src/instamatic/camera/camera.py @@ -18,7 +18,7 @@ def get_cam(interface: str = None): simulate = config.settings.simulate if simulate or interface == 'simulate': - from instamatic.camera.camera_simu import CameraSimu as cam + from instamatic.simulation.camera import CameraSimulation as cam elif interface == 'simulateDLL': from instamatic.camera.camera_gatan import CameraDLL as cam elif interface in ('orius', 'gatan'): diff --git a/src/instamatic/config/camera/simulate.yaml b/src/instamatic/config/camera/simulate.yaml index 02e091f8..6b37e890 100644 --- a/src/instamatic/config/camera/simulate.yaml +++ b/src/instamatic/config/camera/simulate.yaml @@ -19,3 +19,17 @@ possible_binsizes: [1] stretch_amplitude: 2.43 stretch_azimuth: 83.37 dead_time: 0.0 + +# Simulation parameters +n_crystals: 1000000 +grid_diameter: 3.05 # mm +min_crystal_radius: 100 # nm +max_crystal_radius: 1000 # nm +unit_cell_a: 10 # Å +unit_cell_b: 15 # Å +unit_cell_c: 25 # Å +unit_cell_alpha: 90 # Degrees +unit_cell_beta: 90 # Degrees +unit_cell_gamma: 120 # Degrees +max_excitation_error: 0.01 +spot_radius: 3 # pixels diff --git a/src/instamatic/simulation/camera.py b/src/instamatic/simulation/camera.py index 29ccd2f8..cc54ac48 100644 --- a/src/instamatic/simulation/camera.py +++ b/src/instamatic/simulation/camera.py @@ -1,29 +1,164 @@ from __future__ import annotations -from typing import Tuple +from typing import NamedTuple, Tuple -from numpy import ndarray +import numpy as np +from instamatic._typing import float_deg from instamatic.camera.camera_base import CameraBase -from instamatic.simulation.stage import Stage +from instamatic.config import microscope as microscope_config + + +class IlluminatedArea(NamedTuple): + """Area of stage that would show on an image. + + x, y: coordinates relative to optical axis + r: radius. + """ + + x: float + y: float + r: float + + +def _get_reciprocal_unit_cell( + a: float, b: float, c: float, alpha: float, beta: float, gamma: float +) -> np.ndarray: + # Copied from diffpy.structure + ca = np.cos(np.deg2rad(alpha)) + cb = np.cos(np.deg2rad(beta)) + cg = np.cos(np.deg2rad(gamma)) + sa = np.sin(np.deg2rad(alpha)) + sb = np.sin(np.deg2rad(beta)) + cgr = (ca * cb - cg) / (sa * sb) + sgr = np.sqrt(1.0 - cgr * cgr) + Vunit = np.sqrt(1.0 + 2.0 * ca * cb * cg - ca * ca - cb * cb - cg * cg) + ar = sa / (a * Vunit) + base = np.array( + [[1.0 / ar, -cgr / sgr / ar, cb * a], [0.0, b * sa, b * ca], [0.0, 0.0, c]], + ) + recbase = np.linalg.inv(base) + return recbase + + +def _euler_angle_z_to_matrix(theta: float_deg) -> np.ndarray: + c = np.cos(np.deg2rad(theta)) + s = np.sin(np.deg2rad(theta)) + return np.asarray( + [ + [c, -s, 0, 0], + [s, c, 0, 0], + [0, 0, 1, 0], + [0, 0, 0, 1], + ] + ) + + +def _euler_angle_x_to_matrix(theta: float_deg) -> np.ndarray: + c = np.cos(np.deg2rad(theta)) + s = np.sin(np.deg2rad(theta)) + return np.asarray( + [ + [1, 0, 0, 0], + [0, c, -s, 0], + [0, s, c, 0], + [0, 0, 0, 1], + ] + ) + + +def _euler_angle_y_to_matrix(theta: float_deg) -> np.ndarray: + c = np.cos(np.deg2rad(theta)) + s = np.sin(np.deg2rad(theta)) + return np.asarray( + [ + [c, 0, s, 0], + [0, 1, 0, 0], + [-s, 0, c, 0], + [0, 0, 0, 1], + ] + ) class CameraSimulation(CameraBase): streamable = True def __init__(self, name: str = 'simulate'): + """Simple simulation of the TEM. MUST START IN IMAGE MODE! + + Consisting of randomly dispersed crystals. + Each crystal is perfectly flat, and perfectly round. + All crystals have the same unit cell. + + The simulation assumes a 200kV convergent beam, always. + Magnification and/or camera length is read from the TEM, which the simulation connects to after a few seconds. + This is used in combination with the detector pixel size to scale the image/diffraction pattern. + + Image mode: + Crystals have different "thickness", yielding different gray-values proportional to thickness. + This gray-value never changes, regardless of tilt, orientation ect. + + Diffraction mode: + All crystals visible in image mode contribute to the diffraction pattern. + Only reflection positions are accurately simulated. + Intensities are computed by a single Gaussian, rather than proper structure factor computing. + Also scaled with excitation error + """ super().__init__(name) self.ready = False - # TODO put parameters into config - self.stage = Stage() + # Real-space setup + grid_radius_nm = self.grid_diameter * 1e6 / 2 + self.crystal_x = np.random.uniform(-grid_radius_nm, grid_radius_nm, self.n_crystals) + self.crystal_y = np.random.uniform(-grid_radius_nm, grid_radius_nm, self.n_crystals) + self.crystal_r = np.random.uniform( + self.min_crystal_radius, self.max_crystal_radius, self.n_crystals + ) + self.crystal_t = np.random.uniform(0.4, 1.0, self.n_crystals) + self.crystal_euler_angle_phi_1 = np.random.uniform(0, 360, self.n_crystals) + self.crystal_euler_angle_psi = np.random.uniform(0, 180, self.n_crystals) + self.crystal_euler_angle_phi_2 = np.random.uniform(0, 360, self.n_crystals) + + # Reciprocal-space setup + reciprocal_unit_cell = _get_reciprocal_unit_cell( + self.unit_cell_a, + self.unit_cell_b, + self.unit_cell_c, + self.unit_cell_alpha, + self.unit_cell_beta, + self.unit_cell_gamma, + ) + d_min = 0.8 # Å + max_h = round(self.unit_cell_a / d_min) + max_k = round(self.unit_cell_b / d_min) + max_l = round(self.unit_cell_c / d_min) + hkl = np.asarray( + [ + [h, k, l] + for h in range(-max_h, max_h + 1) # noqa: E741 + for k in range(-max_k, max_k + 1) # noqa: E741 + for l in range(-max_l, max_l + 1) # noqa: E741 + ] + ).T + # Filter out the edges that are too high res + k_vec = reciprocal_unit_cell @ hkl + k = np.linalg.norm(k_vec, axis=0) + self.reflections = k_vec[:, k < 1 / d_min] + # Calculate intensity of reflections + k = k[k < 1 / d_min] + F = 5 * np.exp(-1.5 * k**2) # "Structure factor" scaling with a Gaussian + self.reflection_intensity = F**2 + + # Keep track of these, not accessible when outside corresponding mode self.mag = None + self.camera_length = None def establish_connection(self): pass def actually_establish_connection(self): + # Hack to establish connection with running TEM if self.ready: return import time @@ -34,9 +169,10 @@ def actually_establish_connection(self): ctrl = get_instance() self.tem = ctrl.tem - ctrl.stage.set(z=0, a=0, b=0) - print(self.tem.getStagePosition()) - print(self.stage.samples[0].x, self.stage.samples[0].y) + if self.tem.getFunctionMode() == 'diff': + raise RuntimeError('Simulation cannot start in diffraction mode') + + self.mag = self.tem.getMagnification() self.ready = True @@ -44,7 +180,113 @@ def release_connection(self): self.tem = None self.ready = False - def get_image(self, exposure: float = None, binsize: int = None, **kwargs) -> ndarray: + def _stage_pose(self) -> np.ndarray: + """Compute the 4x4 affine transform of the stage. + + Multiply with a position on the stage (x, y) (including z=0 and + the final 1) to get a position relative to the optical axis. + """ + x, y, z, a, b = self.tem.getStagePosition() + + # Using 'ZXY' intrinsic euler angles + Z = _euler_angle_z_to_matrix(self.camera_rotation_vs_stage_xy) + X = _euler_angle_x_to_matrix(a) + Y = _euler_angle_y_to_matrix(b) + + T = np.asarray( + [ + [1, 0, 0, -x], + [0, 1, 0, -y], + [0, 0, 1, -z], + [0, 0, 0, 1], + ] + ) + + return Z @ X @ Y @ T + + def get_realspace_image( + self, xx: np.ndarray, yy: np.ndarray, crystal_indices: np.ndarray + ) -> np.ndarray: + # Only crystals and void, no mesh or anything + out = np.zeros_like(xx) + for x, y, r, t in zip( + self.crystal_x[crystal_indices], + self.crystal_y[crystal_indices], + self.crystal_r[crystal_indices], + self.crystal_t[crystal_indices], + ): + # thickness multiplied by mask of where the crystal is + mask = ((xx - x) ** 2 + (yy - y) ** 2) < r**2 + out += t * mask.astype(float) + return out + + def get_diffraction_image( + self, shape: Tuple[int, int], crystal_indices: np.ndarray + ) -> np.ndarray: + out = np.zeros(shape) + pose = self._stage_pose() + + # Handle camera length + wavelength = microscope_config.wavelength + detector_half_width = self.physical_pixelsize * self.dimensions[0] // 2 + max_theta = np.arctan(detector_half_width / (self.camera_length)) + max_recip_len = 2 * np.sin(max_theta) / wavelength + for phi1, psi, phi2, t in zip( + self.crystal_euler_angle_phi_1[crystal_indices], + self.crystal_euler_angle_psi[crystal_indices], + self.crystal_euler_angle_phi_2[crystal_indices], + self.crystal_t[crystal_indices], + ): + # Crystal orientation, Bunge convention for euler angles + X1 = _euler_angle_x_to_matrix(phi1)[:-1, :-1] + Z = _euler_angle_z_to_matrix(psi)[:-1, :-1] + X2 = _euler_angle_x_to_matrix(phi2)[:-1, :-1] + R = pose[:-1, :-1] @ X1 @ Z @ X2 + beam_direction = R @ np.array([0, 0, -1]) + + # Find intersect with Ewald's sphere, by computing excitation error + # Instead of rotating all vectors, rotate Ewald's sphere (i.e. the beam) to find intersections. + # Then only rotate the intersecting vectors. + # Using notation from https://en.wikipedia.org/wiki/Line%E2%80%93sphere_intersection + r = 1 / wavelength + u = beam_direction + c = r * u + o = self.reflections + + diff = o.T - c + dot = np.dot(u, diff.T) + nabla = dot**2 - np.sum(diff**2, axis=1) + r**2 + # We know the relrods (assuming infinite length) are all going to intersect twice + # Therefore, no need to look at the sign of nabla + # We also know only the smaller root is important, i.e. -sqrt(nabla) + sqrt_nabla = np.sqrt(nabla) + d = -dot - sqrt_nabla + + intersection = np.abs(d) < self.max_excitation_error + vecs = self.reflections[:, intersection] + intensities = self.reflection_intensity[intersection] + # Linear excitation error scaling + intensities *= 1 - np.abs(d[intersection]) / self.max_excitation_error + + # Project vectors onto camera, ignoring curvature + projected_vecs = R.T @ vecs + + # Prepare image + for intensity, (xv, yv, _) in zip(intensities, projected_vecs.T): + x = xv / max_recip_len * shape[0] / 2 + shape[0] / 2 + y = yv / max_recip_len * shape[1] / 2 + shape[1] / 2 + if not 0 <= x < shape[0]: + continue + if not 0 <= y < shape[1]: + continue + min_x = round(max(0, x - self.spot_radius)) + max_x = round(min(shape[1], x + self.spot_radius)) + min_y = round(max(0, y - self.spot_radius)) + max_y = round(min(shape[0], y + self.spot_radius)) + out[min_y:max_y, min_x:max_x] = intensity + return out + + def get_image(self, exposure: float = None, binsize: int = None, **kwargs) -> np.ndarray: self.actually_establish_connection() if exposure is None: @@ -52,54 +294,63 @@ def get_image(self, exposure: float = None, binsize: int = None, **kwargs) -> nd if binsize is None: binsize = self.default_binsize - # TODO this has inconsistent units. Assume m, deg - pos = self.tem.getStagePosition() - if pos is not None and len(pos) == 5: - x, y, z, alpha, beta = pos - self.stage.set_position(x=x, y=y, z=z, alpha_tilt=alpha, beta_tilt=beta) - mode = self.tem.getFunctionMode() - - # Get real-space extent + mag = self.tem.getMagnification() if mode == 'diff': - # TODO this has inconsistent units. Assume mm - self.camera_length = self.tem.getMagnification() + self.camera_length = mag else: - mag = self.tem.getMagnification() - if isinstance(mag, (float, int)): - self.mag = mag - else: - print(mag, type(mag)) - if self.mag is None: - raise ValueError('Must start in image mode') - - # TODO consider beam shift, tilt ect. - x_min, x_max, y_min, y_max = self._mag_to_ranges(self.mag) - x_min += self.stage.x - x_max += self.stage.x - y_min += self.stage.y - y_max += self.stage.y - - # TODO I mean properly considering them, this has no regard for units ect - bx, by = self.tem.getBeamShift() - x_min += bx - x_max += bx - y_min += by - y_max += by + self.mag = mag shape_x, shape_y = self.get_camera_dimensions() shape = (shape_x // binsize, shape_y // binsize) + # Compute "illuminated area" (really, what part of the stage is visible on the image) + real_pixel_size = self.physical_pixelsize * 1e6 # nm + stage_pixel_size = real_pixel_size / self.mag + r = stage_pixel_size * (self.dimensions[0] ** 2 + self.dimensions[1] ** 2) ** 0.5 / 2 + + # Should gun shift also be considered? + x, y = self.tem.getBeamShift() + # x, y, r is now center and radius of beam, compared to optical axis + + # Compute pixel coordinates on stage that are illuminated + R = self._stage_pose() + rp = r / 2**0.5 # Convert from full radius to square inside the circle + x_beam_on_stage = np.linspace(x - rp, x + rp, shape[0]) + y_beam_on_stage = np.linspace(y - rp, y + rp, shape[1]) + # meshgrid of x, y, z and 1 (for the affine transform) + z = self.tem.getStagePosition()[2] + p_beam_on_stage = np.array( + [p.flatten() for p in np.meshgrid(x_beam_on_stage, y_beam_on_stage, [z], [1])] + ) + p_beam_in_column = np.linalg.inv(R) @ p_beam_on_stage + xx = p_beam_in_column[0, :].reshape(shape) + yy = p_beam_in_column[1, :].reshape(shape) + # xx and yy are now 2d coordinate arrays of each pixel on the stage + + # Find which crystals are illuminated + p_crystals = np.vstack( + ( + self.crystal_x, + self.crystal_y, + np.zeros(self.n_crystals), + np.ones(self.n_crystals), + ) + ) + x_beam_center_on_stage, y_beam_center_on_stage, _, _ = np.linalg.inv(R) @ np.array( + [x, y, 0, 1] + ) + c_ind = ( + (p_crystals[0, :] - x_beam_center_on_stage) ** 2 + + (p_crystals[1, :] - y_beam_center_on_stage) ** 2 + ) < (self.crystal_r + r) ** 2 + # c_ind is now a boolean array of which crystals contribute to the image/diffraction pattern + if mode == 'diff': - return self.stage.get_diffraction_pattern( - shape=shape, x_min=x_min, x_max=x_max, y_min=y_min, y_max=y_max + return self.get_diffraction_image( + shape, + c_ind, ) else: - return self.stage.get_image( - shape=shape, x_min=x_min, x_max=x_max, y_min=y_min, y_max=y_max - ) - - def _mag_to_ranges(self, mag: float) -> Tuple[float, float, float, float]: - # assume 50x = 2mm full size - half_width = 50 * 1e6 / mag # 2mm/2 in nm is 1e6 - return -half_width, half_width, -half_width, half_width + img = (self.get_realspace_image(xx, yy, c_ind) * 0xFFFF).astype(int) + return img From 69ab9cd5005953849d4fd7ac8896d1206c25cfe3 Mon Sep 17 00:00:00 2001 From: Viljar Femoen Date: Fri, 19 Sep 2025 16:33:32 +0200 Subject: [PATCH 2/6] Delete superfluous files --- src/instamatic/simulation/crystal.py | 255 ---------------------- src/instamatic/simulation/grid.py | 130 ------------ src/instamatic/simulation/sample.py | 81 ------- src/instamatic/simulation/stage.py | 293 -------------------------- src/instamatic/simulation/warnings.py | 5 - tests/test_simulation/__init__.py | 0 tests/test_simulation/test_crystal.py | 102 --------- tests/test_simulation/test_grid.py | 50 ----- tests/test_simulation/test_sample.py | 42 ---- tests/test_simulation/test_stage.py | 38 ---- 10 files changed, 996 deletions(-) delete mode 100644 src/instamatic/simulation/crystal.py delete mode 100644 src/instamatic/simulation/grid.py delete mode 100644 src/instamatic/simulation/sample.py delete mode 100644 src/instamatic/simulation/stage.py delete mode 100644 src/instamatic/simulation/warnings.py delete mode 100644 tests/test_simulation/__init__.py delete mode 100644 tests/test_simulation/test_crystal.py delete mode 100644 tests/test_simulation/test_grid.py delete mode 100644 tests/test_simulation/test_sample.py delete mode 100644 tests/test_simulation/test_stage.py diff --git a/src/instamatic/simulation/crystal.py b/src/instamatic/simulation/crystal.py deleted file mode 100644 index 13875be3..00000000 --- a/src/instamatic/simulation/crystal.py +++ /dev/null @@ -1,255 +0,0 @@ -from __future__ import annotations - -from typing import Type, TypeVar - -import numpy as np -from diffpy import structure as diffpy - -Crystal_T = TypeVar('Crystal_T', bound='Crystal') - - -class Crystal: - def __init__( - self, a: float, b: float, c: float, alpha: float, beta: float, gamma: float - ) -> None: - """Simulate a primitive crystal given the unit cell. No additional - symmetry is imposed. - - Standard orientation as defined in diffpy. - - Parameters - ---------- - a : float - Unit cell length a, in Å - b : float - Unit cell length b, in Å - c : float - Unit cell length c, in Å - alpha : float - Angle between b and c, in degrees - beta : float - Angle between a and c, in degrees - gamma : float - Angle between a and b, in degrees - """ - self.a = a - self.b = b - self.c = c - self.alpha = alpha - self.beta = beta - self.gamma = gamma - - self.lattice = diffpy.Lattice(self.a, self.b, self.c, self.alpha, self.beta, self.gamma) - self.structure = diffpy.Structure( - atoms=[diffpy.Atom(xyz=[0, 0, 0])], - lattice=self.lattice, - ) - - @property - def a_vec(self) -> np.ndarray: - return self.lattice.cartesian((1, 0, 0)) - - @property - def b_vec(self) -> np.ndarray: - return self.lattice.cartesian((0, 1, 0)) - - @property - def c_vec(self) -> np.ndarray: - return self.lattice.cartesian((0, 0, 1)) - - @property - def a_star_vec(self) -> np.ndarray: - return self.lattice.reciprocal().cartesian((1, 0, 0)) - - @property - def b_star_vec(self) -> np.ndarray: - return self.lattice.reciprocal().cartesian((0, 1, 0)) - - @property - def c_star_vec(self) -> np.ndarray: - return self.lattice.reciprocal().cartesian((0, 0, 1)) - - @classmethod - def default(cls: Type[Crystal_T]) -> Crystal_T: - return cls(1, 2, 3, 90, 100, 110) - - def real_space_lattice(self, d_max: float) -> np.ndarray: - """Get the real space lattice as a (n, 3) shape array. - - Parameters - ---------- - d_max: float - The maximum d-spacing - - Returns - ------- - np.ndarray - Shape (n, 3), lattice points - """ - max_h = int(d_max // self.a) - max_k = int(d_max // self.b) - max_l = int(d_max // self.c) - hkls = np.array( - [ - (h, k, l) - for h in range(-max_h, max_h + 1) # noqa: E741 - for k in range(-max_k, max_k + 1) # noqa: E741 - for l in range(-max_l, max_l + 1) # noqa: E741 - ] - ) - vecs = self.lattice.cartesian(hkls) - return vecs - - def reciprocal_space_lattice(self, d_min: float) -> np.ndarray: - """Get the reciprocal space lattice as a (n, 3) shape array for input - n. - - Parameters - ---------- - d_min: float - Minimum d-spacing included - - Returns - ------- - np.ndarray - Shape (n, 3), lattice points - """ - max_h = int(d_min // self.lattice.ar) - max_k = int(d_min // self.lattice.br) - max_l = int(d_min // self.lattice.cr) - hkls = np.array( - [ - (h, k, l) - for h in range(-max_h, max_h + 1) # noqa: E741 - for k in range(-max_k, max_k + 1) # noqa: E741 - for l in range(-max_l, max_l + 1) # noqa: E741 - ] - ) - vecs = self.lattice.reciprocal().cartesian(hkls) - return vecs - - def diffraction_pattern_mask( - self, - shape: tuple[int, int], - d_min: float, - rotation_matrix: np.ndarray, - wavelength: float, - excitation_error: float, - ) -> np.ndarray: - """Get a diffraction pattern with a given shape, up to a given - resolution, in a given orientation and wavelength. - - Parameters - ---------- - shape : tuple[int, int] - Output shape - d_min : float - Minimum d-spacing, in Å - rotation_matrix : np.ndarray - Orientation - wavelength : float - Wavelength of incident beam, in Å - excitation_error : float - Excitation error used for intensity calculation, in reciprocal Å - - Returns - ------- - np.ndarray - Diffraction pattern - """ - # TODO calibration - out = np.zeros(shape, dtype=bool) - - # TODO this depends on convergence angle - spot_radius = 3 # pixels - - vecs = self.reciprocal_space_lattice(d_min) - d = np.sum(vecs**2, axis=1) - vecs = vecs[d < d_min**2] - - k = 2 * np.pi / wavelength - k_vec = rotation_matrix @ np.array([0, 0, -k]) - - # Find intersect with Ewald's sphere - q_squared = np.sum((vecs - k_vec) ** 2, axis=1) - vecs = vecs[ - (q_squared > (k - excitation_error) ** 2) - & (q_squared < (k + excitation_error) ** 2) - ] - - # Project onto screen - vecs_xy = (rotation_matrix.T @ vecs.T).T[:, :-1] # ignoring curvature - - # Make image - for vec in vecs_xy: - x = int(vec[0] * d_min * shape[1] / 2) + shape[1] // 2 - y = int(vec[1] * d_min * shape[0] / 2) + shape[0] // 2 - min_x = max(0, x - spot_radius) - max_x = min(shape[1], x + spot_radius) - min_y = max(0, y - spot_radius) - max_y = min(shape[0], y + spot_radius) - out[min_y:max_y, min_x:max_x] = 1 - return out - - def __str__(self) -> str: - return f'{self.__class__.__name__}(a = {self.a}, b = {self.b}, c = {self.c}, alpha = {self.alpha}, beta = {self.beta}, gamma = {self.gamma})' - - -class CubicCrystal(Crystal): - def __init__(self, a: float) -> None: - super().__init__(a, a, a, 90, 90, 90) - - @classmethod - def default(cls: Type[Crystal_T]) -> Crystal_T: - return cls(1) - - -class HexagonalCrystal(Crystal): - def __init__(self, a: float, c: float) -> None: - super().__init__(a, a, c, 90, 90, 120) - - @classmethod - def default(cls: Type[Crystal_T]) -> Crystal_T: - return cls(1, 2) - - -class TrigonalCrystal(Crystal): - def __init__(self, a: float, alpha: float) -> None: - super().__init__(a, a, a, alpha, alpha, alpha) - - @classmethod - def default(cls: Type[Crystal_T]) -> Crystal_T: - return cls(1, 100) - - -class TetragonalCrystal(Crystal): - def __init__(self, a: float, c: float) -> None: - super().__init__(a, a, c, 90, 90, 90) - - @classmethod - def default(cls: Type[Crystal_T]) -> Crystal_T: - return cls(1, 2) - - -class OrthorhombicCrystal(Crystal): - def __init__(self, a: float, b: float, c: float) -> None: - super().__init__(a, b, c, 90, 90, 90) - - @classmethod - def default(cls: Type[Crystal_T]) -> Crystal_T: - return cls(1, 2, 3) - - -class MonoclinicCrystal(Crystal): - def __init__(self, a: float, b: float, c: float, beta: float) -> None: - super().__init__(a, b, c, 90, beta, 90) - - @classmethod - def default(cls: Type[Crystal_T]) -> Crystal_T: - return cls(1, 2, 3, 100) - - -class TriclinicCrystal(Crystal): - @classmethod - def default(cls: Type[Crystal_T]) -> Crystal_T: - return cls(1, 2, 3, 90, 100, 110) diff --git a/src/instamatic/simulation/grid.py b/src/instamatic/simulation/grid.py deleted file mode 100644 index 9d093e63..00000000 --- a/src/instamatic/simulation/grid.py +++ /dev/null @@ -1,130 +0,0 @@ -from __future__ import annotations - -import warnings - -import numpy as np - -from instamatic.simulation.warnings import NotImplementedWarning - -# TODO carbon lace - - -class Grid: - def __init__( - self, - diameter: float = 3.05, - mesh: int = 200, - pitch: float = 125, - hole_width: float = 90, - bar_width: float = 35, - rim_width: float = 0.225, - ): - """TEM grid. - - Parameters - ---------- - diameter : float, optional - [mm] Total diameter, by default 3.05 - mesh : int, optional - [lines/inch] Hole density, by default 200 - pitch : float, optional - [µm], by default 125 - hole_width : float, optional - [µm], by default 90 - bar_width : float, optional - [µm], by default 35 - rim_width : float, optional - [mm], by default 0.225 - """ - # TODO make mesh set the pitch, bar width and pitch set the hole width ect. - self.diameter_mm = diameter - self.radius_nm = 1e6 * diameter / 2 - self.mesh = mesh - self.pitch_um = pitch - self.hole_width_um = hole_width - self.bar_width_um = bar_width - self.rim_width_mm = rim_width - self.grid_width_um = self.bar_width_um + self.hole_width_um - self.grid_width_nm = 1e3 * self.grid_width_um - - def get_rim_filter(self, x: np.ndarray, y: np.ndarray) -> np.ndarray: - r_nm = 1e6 * (self.diameter_mm / 2 - self.rim_width_mm) - return x**2 + y**2 >= r_nm**2 - - def get_hole_filter(self, x: np.ndarray, y: np.ndarray) -> np.ndarray: - x = np.remainder(np.abs(x), self.grid_width_nm) - y = np.remainder(np.abs(y), self.grid_width_nm) - - # Assume no grid in center, i.e. middle of the bar width - half_bar_width_nm = 1e3 * self.bar_width_um / 2 - return ( - (x < half_bar_width_nm) - | (x > (self.grid_width_nm - half_bar_width_nm)) - | (y < half_bar_width_nm) - | (y > (self.grid_width_nm - half_bar_width_nm)) - ) - - def get_center_mark(self, x: np.ndarray, y: np.ndarray) -> np.ndarray: - # TODO - warnings.warn('Center mark is not implemented yet', NotImplementedWarning) - return np.zeros(x.shape, dtype=bool) - - def array_from_coords(self, x: np.ndarray, y: np.ndarray) -> np.ndarray: - """Get mask array for given coordinate arrays (output from - np.meshgrid). (x, y) = (0, 0) is in the center of the grid. - - Parameters - ---------- - x : np.ndarray - x-coordinates - y : np.ndarray - y-coordinates - - Returns - ------- - np.ndarray - Mask array, False where the grid is blocking - """ - rim_filter = self.get_rim_filter(x, y) - grid_filter = self.get_hole_filter(x, y) - - # TODO proper logic for this, - # as the mark includes a hole in the center which will be overridden by the grid filter - center_mark = self.get_center_mark(x, y) - - return rim_filter | grid_filter | center_mark - - def array( - self, - shape: tuple[int, int], - x_min: float, - x_max: float, - y_min: float, - y_max: float, - ) -> np.ndarray: - """Get mask array for given ranges. (x, y) = (0, 0) is in the center of - the grid. - - Parameters - ---------- - shape : tuple[int, int] - Output shape - x_min : float - [nm] Lower bound for x (left) - x_max : float - [nm] Upper bound for x (right) - y_min : float - [nm] Lower bound for y (bottom) - y_max : float - [nm] Upper bound for y (top) - - Returns - ------- - np.ndarray - Mask array, False where the grid is blocking - """ - x, y = np.meshgrid( - np.linspace(x_min, x_max, shape[1]), - np.linspace(y_min, y_max, shape[0]), - ) - return self.array_from_coords(x, y) diff --git a/src/instamatic/simulation/sample.py b/src/instamatic/simulation/sample.py deleted file mode 100644 index f4786f58..00000000 --- a/src/instamatic/simulation/sample.py +++ /dev/null @@ -1,81 +0,0 @@ -from __future__ import annotations - -from dataclasses import dataclass - -import numpy as np - - -@dataclass -class Sample: - x: float - y: float - r: float - thickness: float # between 0 and 1 - euler_angle_phi_1: float - euler_angle_psi: float - euler_angle_phi_2: float - crystal_index: int = 0 # used for lookup in a list of crystals - - def __post_init__(self): - cp1 = np.cos(self.euler_angle_phi_1) - cp = np.cos(self.euler_angle_psi) - cp2 = np.cos(self.euler_angle_phi_2) - sp1 = np.sin(self.euler_angle_phi_1) - sp = np.sin(self.euler_angle_psi) - sp2 = np.sin(self.euler_angle_phi_2) - r1 = np.array([[cp1, sp1, 0], [-sp1, cp1, 0], [0, 0, 1]]) - r2 = np.array([[1, 0, 0], [0, cp, sp], [0, -sp, cp]]) - r3 = np.array([[cp2, sp2, 0], [-sp2, cp2, 0], [0, 0, 1]]) - self.rotation_matrix = r1 @ r2 @ r3 - - def pixel_contains_crystal(self, x: np.ndarray, y: np.ndarray) -> np.ndarray: - """Given arrays of x- and y- coordinates in the lab frame, calculate - whether the crystal overlaps with these positions. - - Parameters - ---------- - x : np.ndarray - x coordinates - y : np.ndarray - y coordinates - - Returns - ------- - np.ndarray - Same shape as inputs, dtype bool - """ - return (x - self.x) ** 2 + (y - self.y) ** 2 < self.r**2 - - def range_might_contain_crystal( - self, - x_min: float, - x_max: float, - y_min: float, - y_max: float, - ) -> bool: - """Simple estimate of whether a range contains the crystal. This check - is fast but inaccurate. False positives are possible, false negatives - are impossible. - - Parameters - ---------- - x_min : float - Lower bound for x - x_max : float - Upper bound for x - y_min : float - Lower bound for y - y_max : float - Upper bound for y - - Returns - ------- - bool - True if range contains crystal - """ - # TODO ensure the docstring is true, regarding false negatives. - # TODO improve estimate? - # TODO handle this correctly when stage is rotated... - in_x = x_min <= self.x + self.r and self.x - self.r <= x_max - in_y = y_min <= self.y + self.r and self.y - self.r <= y_max - return in_x and in_y diff --git a/src/instamatic/simulation/stage.py b/src/instamatic/simulation/stage.py deleted file mode 100644 index a5b57036..00000000 --- a/src/instamatic/simulation/stage.py +++ /dev/null @@ -1,293 +0,0 @@ -from __future__ import annotations - -import warnings - -import numpy as np -from scipy.spatial.transform import Rotation - -from instamatic.simulation.crystal import Crystal -from instamatic.simulation.grid import Grid -from instamatic.simulation.sample import Sample -from instamatic.simulation.warnings import NotImplementedWarning - - -class Stage: - def __init__( - self, - num_crystals: int = 100_000, - min_crystal_size: float = 100, - max_crystal_size: float = 1000, - random_seed: int = 100, - ) -> None: - """Handle many samples on a grid. - - Parameters - ---------- - num_crystals : int, optional - Number of crystals to disperse on the grid, by default 100_000 - min_crystal_size : float, optional - Minimum radius of the crystals, in nm, by default 100 - max_crystal_size : float, optional - Maximum radius of the crystals, in nm, by default 1000 - random_seed : int, optional - Seed for random number generation, by default 100 - """ - # TODO make this settable - self.x = 0 - self.y = 0 - self.z = 0 - self.alpha_tilt = 0 - self.beta_tilt = 0 - self.in_plane_rotation = 0 # TODO change this with focus/magnification - self.rotation_matrix = np.eye(3) - self.origin = np.array([0, 0, 0]) - - # TODO parameters - self.grid = Grid() - - self.rng = np.random.Generator(np.random.PCG64(random_seed)) - - # TODO parameters - # TODO multiple phases - # TODO amorphous phase - self.crystal = Crystal(*self.rng.uniform(5, 25, 3), *self.rng.uniform(80, 110, 3)) - - self.samples = [ - Sample( - x=self.rng.uniform(-self.grid.radius_nm, self.grid.radius_nm), - y=self.rng.uniform(-self.grid.radius_nm, self.grid.radius_nm), - r=self.rng.uniform(min_crystal_size, max_crystal_size), - thickness=self.rng.uniform(0, 1), - euler_angle_phi_1=self.rng.uniform(0, 2 * np.pi), - euler_angle_psi=self.rng.uniform(0, np.pi), - euler_angle_phi_2=self.rng.uniform(0, 2 * np.pi), - ) - for _ in range(num_crystals) - ] - - def set_position( - self, - x: float = None, - y: float = None, - z: float = None, - alpha_tilt: float = None, - beta_tilt: float = None, - ): - if x is not None: - self.x = x - if y is not None: - self.y = y - if z is not None: - self.z = z - if alpha_tilt is not None: - warnings.warn( - 'Tilting is not fully implemented yet', - NotImplementedWarning, - stacklevel=2, - ) - self.alpha_tilt = alpha_tilt - if beta_tilt is not None: - warnings.warn( - 'Tilting is not fully implemented yet', - NotImplementedWarning, - stacklevel=2, - ) - self.beta_tilt = beta_tilt - - # TODO define orientation. Is this matrix multiplied with lab coordinates to get sample coordinates? - self.rotation_matrix = Rotation.from_euler( - 'ZXY', - [self.in_plane_rotation, self.alpha_tilt, self.beta_tilt], - degrees=True, - ).as_matrix() - - def image_extent_to_sample_coordinates( - self, - shape: tuple[int, int], - x_min: float, - x_max: float, - y_min: float, - y_max: float, - ) -> tuple[np.ndarray, np.ndarray]: - """Get arrays of grid positions with a given shape and extent in lab - coordinates. - - Parameters - ---------- - shape : tuple[int, int] - Output shape - x_min : float - Lower bound of x - x_max : float - Upper bound of x - y_min : float - Lower bound of y - y_max : float - Upper bound of y - - Returns - ------- - tuple[np.ndarray, np.ndarray] - x, y. 2D arrays of floats - """ - if self.alpha_tilt != 0 or self.beta_tilt != 0: - warnings.warn( - 'Tilting is not fully implemented yet', NotImplementedWarning, stacklevel=2 - ) - # https://en.wikipedia.org/wiki/Line%E2%80%93plane_intersection - n = self.rotation_matrix @ np.array([0, 0, 1]) - p0 = self.origin - l = np.array([0, 0, 1]) # noqa: E741 - l0 = np.array( - [ - p.flatten() - for p in np.meshgrid( - np.linspace(x_min, x_max, shape[1]), - np.linspace(y_min, y_max, shape[0]), - [0], - ) - ] - ) - - p = l0 + np.array([0, 0, 1])[:, np.newaxis] * np.dot(-l0.T + p0, n) / np.dot(l, n) - - x, y, z = self.rotation_matrix.T @ p - x = x.reshape(shape) - y = y.reshape(shape) - return x, y - - def get_image( - self, - shape: tuple[int, int], - x_min: float, - x_max: float, - y_min: float, - y_max: float, - ) -> np.ndarray: - """Get image array for given ranges. (x, y) = (0, 0) is in the center - of the grid. - - Parameters - ---------- - shape : tuple[int, int] - Output shape - x_min : float - [nm] Lower bound for x (left) - x_max : float - [nm] Upper bound for x (right) - y_min : float - [nm] Lower bound for y (bottom) - y_max : float - [nm] Upper bound for y (top) - - Returns - ------- - np.ndarray - Image - """ - x, y = self.image_extent_to_sample_coordinates( - shape=shape, x_min=x_min, x_max=x_max, y_min=y_min, y_max=y_max - ) - - grid_mask = self.grid.array_from_coords(x, y) - - sample_data = np.ones(shape, dtype=int) * 1000 - for ind, sample in enumerate(self.samples): - if not sample.range_might_contain_crystal( - x_min=x_min, x_max=x_max, y_min=y_min, y_max=y_max - ): - continue - # TODO better logic here - sample_data[sample.pixel_contains_crystal(x, y)] = 1000 * (1 - sample.thickness) - - sample_data[grid_mask] = 0 - - return sample_data - - def get_diffraction_pattern( - self, - shape: tuple[int, int], - x_min: float, - x_max: float, - y_min: float, - y_max: float, - camera_length: float = 150, - ) -> np.ndarray: - """Get diffraction pattern array for given ranges. (x, y) = (0, 0) is - in the center of the grid. - - Parameters - ---------- - shape : tuple[int, int] - Output shape - x_min : float - [nm] Lower bound for x (left) - x_max : float - [nm] Upper bound for x (right) - y_min : float - [nm] Lower bound for y (bottom) - y_max : float - [nm] Upper bound for y (top) - camera_length : float - [cm] Camera length, for calibration - - Returns - ------- - np.ndarray - diffraction pattern - """ - x, y = self.image_extent_to_sample_coordinates( - shape=shape, x_min=x_min, x_max=x_max, y_min=y_min, y_max=y_max - ) - d_min = 1.0 - - grid_mask = self.grid.array_from_coords(x, y) - - if np.all(grid_mask): - # no transmission - return np.zeros(shape, dtype=int) - - reflections = np.zeros(shape, dtype=bool) - - # Direct beam - reflections[ - shape[0] // 2 - 4 : shape[0] // 2 + 4, shape[1] // 2 - 4 : shape[1] // 2 + 4 - ] = 1 - - for sample in self.samples: - if not sample.range_might_contain_crystal( - x_min=x_min, x_max=x_max, y_min=y_min, y_max=y_max - ): - continue - pos = sample.pixel_contains_crystal(x, y) - if np.all(grid_mask[pos]): - # Crystal is completely on the grid - continue - - reflections |= self.crystal.diffraction_pattern_mask( - shape, - d_min=d_min, - rotation_matrix=self.rotation_matrix @ sample.rotation_matrix, - wavelength=0.02, - excitation_error=0.01, - ) - # TODO diffraction shift - - # TODO noise - - # Simple scaling - # TODO improve, proper form factors maybe - # TODO camera length - kx, ky = np.meshgrid( - np.linspace(-1 / d_min, 1 / d_min, shape[1]), - np.linspace(-1 / d_min, 1 / d_min, shape[0]), - ) - k_squared = kx**2 + ky**2 - scale = 1 / (3 * k_squared + 1) - - scale[~reflections] = 0 - - # Convert to int array - scale = (scale * 0x8000).astype(int) - - return scale diff --git a/src/instamatic/simulation/warnings.py b/src/instamatic/simulation/warnings.py deleted file mode 100644 index dfca684b..00000000 --- a/src/instamatic/simulation/warnings.py +++ /dev/null @@ -1,5 +0,0 @@ -from __future__ import annotations - - -class NotImplementedWarning(UserWarning): - pass diff --git a/tests/test_simulation/__init__.py b/tests/test_simulation/__init__.py deleted file mode 100644 index e69de29b..00000000 diff --git a/tests/test_simulation/test_crystal.py b/tests/test_simulation/test_crystal.py deleted file mode 100644 index 8a1ad4ea..00000000 --- a/tests/test_simulation/test_crystal.py +++ /dev/null @@ -1,102 +0,0 @@ -from __future__ import annotations - -from typing import Type - -import pytest - -from instamatic.simulation.crystal import ( - Crystal, - CubicCrystal, - HexagonalCrystal, - MonoclinicCrystal, - OrthorhombicCrystal, - TetragonalCrystal, - TriclinicCrystal, - TrigonalCrystal, -) - - -def test_crystal_init(): - Crystal(1, 1, 1, 1, 1, 1) - - -@pytest.mark.parametrize( - 'crystal', - [ - Crystal, - CubicCrystal, - HexagonalCrystal, - TrigonalCrystal, - TetragonalCrystal, - OrthorhombicCrystal, - MonoclinicCrystal, - TriclinicCrystal, - ], -) -def test_crystal_default(crystal: Type[Crystal]): - c = crystal.default() - assert isinstance(c, Crystal) - - -def test_get_lattice_cubic(): - c = CubicCrystal.default() - lat = c.real_space_lattice(1) - assert pytest.approx(lat) == [ - (-1, -1, -1), - (-1, -1, 0), - (-1, -1, 1), - (-1, 0, -1), - (-1, 0, 0), - (-1, 0, 1), - (-1, 1, -1), - (-1, 1, 0), - (-1, 1, 1), - (0, -1, -1), - (0, -1, 0), - (0, -1, 1), - (0, 0, -1), - (0, 0, 0), - (0, 0, 1), - (0, 1, -1), - (0, 1, 0), - (0, 1, 1), - (1, -1, -1), - (1, -1, 0), - (1, -1, 1), - (1, 0, -1), - (1, 0, 0), - (1, 0, 1), - (1, 1, -1), - (1, 1, 0), - (1, 1, 1), - ] - lat = c.reciprocal_space_lattice(1) - assert pytest.approx(lat) == [ - (-1, -1, -1), - (-1, -1, 0), - (-1, -1, 1), - (-1, 0, -1), - (-1, 0, 0), - (-1, 0, 1), - (-1, 1, -1), - (-1, 1, 0), - (-1, 1, 1), - (0, -1, -1), - (0, -1, 0), - (0, -1, 1), - (0, 0, -1), - (0, 0, 0), - (0, 0, 1), - (0, 1, -1), - (0, 1, 0), - (0, 1, 1), - (1, -1, -1), - (1, -1, 0), - (1, -1, 1), - (1, 0, -1), - (1, 0, 0), - (1, 0, 1), - (1, 1, -1), - (1, 1, 0), - (1, 1, 1), - ] diff --git a/tests/test_simulation/test_grid.py b/tests/test_simulation/test_grid.py deleted file mode 100644 index f1116517..00000000 --- a/tests/test_simulation/test_grid.py +++ /dev/null @@ -1,50 +0,0 @@ -from __future__ import annotations - -import numpy as np -import pytest - -from instamatic.simulation.grid import Grid - - -def test_init(): - Grid() - - -def test_get_array(): - g = Grid( - hole_width=9, - bar_width=1, - ) - arr = g.array( - shape=(40, 40), - x_min=g.grid_width_nm, - x_max=3 * g.grid_width_nm, - y_min=-4 * g.grid_width_nm, - y_max=-2 * g.grid_width_nm, - ) - - # Grid - assert np.all(arr[:, 0]) - assert np.all(arr[:, -1]) - assert np.all(arr[0, :]) - assert np.all(arr[-1, :]) - assert np.all(arr[:, 19]) - assert np.all(arr[:, 20]) - assert np.all(arr[19, :]) - assert np.all(arr[20, :]) - - # Holes - assert np.sum(arr[1:19, 1:19]) == 0 - assert np.sum(arr[1:19, 21:-1]) == 0 - assert np.sum(arr[21:-1, 1:19]) == 0 - assert np.sum(arr[21:-1, 21:-1]) == 0 - - -@pytest.mark.xfail(reason='TODO') -def test_get_array_including_center(): - assert False, 'TODO' - - -@pytest.mark.xfail(reason='TODO') -def test_get_array_including_rim(): - assert False, 'TODO' diff --git a/tests/test_simulation/test_sample.py b/tests/test_simulation/test_sample.py deleted file mode 100644 index f01022dd..00000000 --- a/tests/test_simulation/test_sample.py +++ /dev/null @@ -1,42 +0,0 @@ -from __future__ import annotations - -import pytest - -from instamatic.simulation.sample import Sample - - -def test_init(): - s = Sample(0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0) - assert isinstance(s, Sample) - - -def test_range_might_contain_crystal(): - s = Sample(0, 0, 1, 0, 0, 0, 0) - assert s.range_might_contain_crystal(-1, 1, -1, 1) - assert not s.range_might_contain_crystal(9, 10, 9, 10) - assert not s.range_might_contain_crystal(-10, -9, -10, -9) - assert not s.range_might_contain_crystal(-10, -9, -10, 10) - assert not s.range_might_contain_crystal(-10, 10, -10, -9) - assert s.range_might_contain_crystal(0, 1, 0, 1) - # TODO expand - - -@pytest.mark.xfail(reason='TODO') -def test_pixel_contains_crystal(): - assert False, 'TODO' - - -def test_range_might_contain_crystal_false_positive(): - s = Sample(0, 0, 1, 0, 0, 0, 0) - x_min = 0.9 - x_max = 1 - y_min = 0.9 - y_max = 1 - assert s.range_might_contain_crystal(x_min, x_max, y_min, y_max) - assert not s.pixel_contains_crystal(x_min, y_min) - # TODO expand - - -@pytest.mark.xfail(reason='Need to figure out how this can be done') -def test_range_might_contain_crystal_false_negative(): - assert False, 'TODO' diff --git a/tests/test_simulation/test_stage.py b/tests/test_simulation/test_stage.py deleted file mode 100644 index bd6e4c00..00000000 --- a/tests/test_simulation/test_stage.py +++ /dev/null @@ -1,38 +0,0 @@ -from __future__ import annotations - -import pytest - -from instamatic.simulation.stage import Stage -from instamatic.simulation.warnings import NotImplementedWarning - - -def test_init_default(): - s = Stage() - assert isinstance(s, Stage) - - -def test_set_position(): - s = Stage() - s.set_position(x=10) - assert s.x == 10 - s.set_position(z=10) - assert s.x == 10 - assert s.z == 10 - with pytest.warns(NotImplementedWarning): - s.set_position(alpha_tilt=1) - with pytest.warns(NotImplementedWarning): - s.set_position(beta_tilt=1) - with pytest.warns(NotImplementedWarning): - s.set_position(x=1, y=1, z=1, alpha_tilt=1, beta_tilt=1) - - -@pytest.mark.xfail(reason='TODO') -def test_tilt(): - # Somehow check that the projected coordinates using stage.image_extent_to_sample_coordinates are correct - assert False, 'TODO' - - -@pytest.mark.xfail(reason='TODO') -def test_image_rotation(): - # Image rotates with focus ect. - assert False, 'TODO' From 8b96d89378a19c918a7b0208fb18a26f21ac2be4 Mon Sep 17 00:00:00 2001 From: Viljar Femoen Date: Fri, 19 Sep 2025 17:04:14 +0200 Subject: [PATCH 3/6] Move implementation to previous camera simulation --- src/instamatic/camera/camera.py | 2 +- src/instamatic/camera/camera_simu.py | 408 +++++++++++++++++++++----- src/instamatic/simulation/__init__.py | 0 src/instamatic/simulation/camera.py | 356 ---------------------- 4 files changed, 339 insertions(+), 427 deletions(-) delete mode 100644 src/instamatic/simulation/__init__.py delete mode 100644 src/instamatic/simulation/camera.py diff --git a/src/instamatic/camera/camera.py b/src/instamatic/camera/camera.py index 902292ea..fa806b07 100644 --- a/src/instamatic/camera/camera.py +++ b/src/instamatic/camera/camera.py @@ -18,7 +18,7 @@ def get_cam(interface: str = None): simulate = config.settings.simulate if simulate or interface == 'simulate': - from instamatic.simulation.camera import CameraSimulation as cam + from instamatic.camera.camera_simu import CameraSimu as cam elif interface == 'simulateDLL': from instamatic.camera.camera_gatan import CameraDLL as cam elif interface in ('orius', 'gatan'): diff --git a/src/instamatic/camera/camera_simu.py b/src/instamatic/camera/camera_simu.py index 8e0bc9f2..96e84515 100644 --- a/src/instamatic/camera/camera_simu.py +++ b/src/instamatic/camera/camera_simu.py @@ -3,27 +3,162 @@ import atexit import logging import time -from typing import Generator, List, Optional, Tuple +from typing import Tuple import numpy as np -from instamatic import config +from instamatic._typing import float_deg from instamatic.camera.camera_base import CameraBase +from instamatic.config import microscope as microscope_config logger = logging.getLogger(__name__) -class CameraSimu(CameraBase): - """Simple class that simulates the camera interface and mocks the method - calls.""" +def _get_reciprocal_unit_cell( + a: float, b: float, c: float, alpha: float, beta: float, gamma: float +) -> np.ndarray: + # Copied from diffpy.structure + ca = np.cos(np.deg2rad(alpha)) + cb = np.cos(np.deg2rad(beta)) + cg = np.cos(np.deg2rad(gamma)) + sa = np.sin(np.deg2rad(alpha)) + sb = np.sin(np.deg2rad(beta)) + cgr = (ca * cb - cg) / (sa * sb) + sgr = np.sqrt(1.0 - cgr * cgr) + Vunit = np.sqrt(1.0 + 2.0 * ca * cb * cg - ca * ca - cb * cb - cg * cg) + ar = sa / (a * Vunit) + base = np.array( + [[1.0 / ar, -cgr / sgr / ar, cb * a], [0.0, b * sa, b * ca], [0.0, 0.0, c]], + ) + recbase = np.linalg.inv(base) + return recbase + + +def _euler_angle_z_to_matrix(theta: float_deg) -> np.ndarray: + c = np.cos(np.deg2rad(theta)) + s = np.sin(np.deg2rad(theta)) + return np.asarray( + [ + [c, -s, 0, 0], + [s, c, 0, 0], + [0, 0, 1, 0], + [0, 0, 0, 1], + ] + ) + + +def _euler_angle_x_to_matrix(theta: float_deg) -> np.ndarray: + c = np.cos(np.deg2rad(theta)) + s = np.sin(np.deg2rad(theta)) + return np.asarray( + [ + [1, 0, 0, 0], + [0, c, -s, 0], + [0, s, c, 0], + [0, 0, 0, 1], + ] + ) + + +def _euler_angle_y_to_matrix(theta: float_deg) -> np.ndarray: + c = np.cos(np.deg2rad(theta)) + s = np.sin(np.deg2rad(theta)) + return np.asarray( + [ + [c, 0, s, 0], + [0, 1, 0, 0], + [-s, 0, c, 0], + [0, 0, 0, 1], + ] + ) + +class CameraSimu(CameraBase): streamable = True - def __init__(self, name='simulate'): - """Initialize camera module.""" + def __init__(self, name: str = 'simulate'): + """Simple simulation of the TEM. MUST START IN IMAGE MODE! + + Consisting of randomly dispersed crystals. + Each crystal is perfectly flat, and perfectly round. + All crystals have the same unit cell. + + The simulation assumes a 200kV convergent beam, always. + Magnification and/or camera length is read from the TEM, which the simulation connects to after a few seconds. + This is used in combination with the detector pixel size to scale the image/diffraction pattern. + + Image mode: + Crystals have different "thickness", yielding different gray-values proportional to thickness. + This gray-value never changes, regardless of tilt, orientation ect. + + Diffraction mode: + All crystals visible in image mode contribute to the diffraction pattern. + Only reflection positions are accurately simulated. + Intensities are computed by a single Gaussian, rather than proper structure factor computing. + Also scaled with excitation error + """ + # Defaults + self.n_crystals = 1000000 + self.grid_diameter = 3.05 # mm + self.min_crystal_radius = 100 # nm + self.max_crystal_radius = 1000 # nm + self.unit_cell_a = 10 # Å + self.unit_cell_b = 15 # Å + self.unit_cell_c = 25 # Å + self.unit_cell_alpha = 90 # Degrees + self.unit_cell_beta = 90 # Degrees + self.unit_cell_gamma = 120 # Degrees + self.max_excitation_error = 0.01 + self.spot_radius = 3 # pixels super().__init__(name) - self.establish_connection() + self.ready = False + + # Real-space setup + grid_radius_nm = self.grid_diameter * 1e6 / 2 + self.crystal_x = np.random.uniform(-grid_radius_nm, grid_radius_nm, self.n_crystals) + self.crystal_y = np.random.uniform(-grid_radius_nm, grid_radius_nm, self.n_crystals) + self.crystal_r = np.random.uniform( + self.min_crystal_radius, self.max_crystal_radius, self.n_crystals + ) + self.crystal_t = np.random.uniform(0.4, 1.0, self.n_crystals) + self.crystal_euler_angle_phi_1 = np.random.uniform(0, 360, self.n_crystals) + self.crystal_euler_angle_psi = np.random.uniform(0, 180, self.n_crystals) + self.crystal_euler_angle_phi_2 = np.random.uniform(0, 360, self.n_crystals) + + # Reciprocal-space setup + reciprocal_unit_cell = _get_reciprocal_unit_cell( + self.unit_cell_a, + self.unit_cell_b, + self.unit_cell_c, + self.unit_cell_alpha, + self.unit_cell_beta, + self.unit_cell_gamma, + ) + d_min = 0.8 # Å + max_h = round(self.unit_cell_a / d_min) + max_k = round(self.unit_cell_b / d_min) + max_l = round(self.unit_cell_c / d_min) + hkl = np.asarray( + [ + [h, k, l] + for h in range(-max_h, max_h + 1) # noqa: E741 + for k in range(-max_k, max_k + 1) # noqa: E741 + for l in range(-max_l, max_l + 1) # noqa: E741 + ] + ).T + # Filter out the edges that are too high res + k_vec = reciprocal_unit_cell @ hkl + k = np.linalg.norm(k_vec, axis=0) + self.reflections = k_vec[:, k < 1 / d_min] + # Calculate intensity of reflections + k = k[k < 1 / d_min] + F = 5 * np.exp(-1.5 * k**2) # "Structure factor" scaling with a Gaussian + self.reflection_intensity = F**2 + + # Keep track of these, not accessible when outside corresponding mode + self.mag = None + self.camera_length = None msg = f'Camera {self.get_name()} initialized' logger.info(msg) @@ -36,63 +171,208 @@ def __init__(self, name='simulate'): self._autoincrement = True self._start_record_time = -1 - def get_image(self, exposure=None, binsize=None, **kwargs) -> np.ndarray: - """Image acquisition routine. If the exposure and binsize are not - given, the default values are read from the config file. + def establish_connection(self): + pass + + def actually_establish_connection(self): + # Hack to establish connection with running TEM + if self.ready: + return + import time - Parameters - ---------- - exposure : float - Exposure time in seconds. - binsize : int - Which binning to use. + time.sleep(2) + from instamatic.controller import get_instance + + ctrl = get_instance() + self.tem = ctrl.tem + + if self.tem.getFunctionMode() == 'diff': + # raise RuntimeError('Simulation cannot start in diffraction mode') + # Assume something reasonable instead + self.mag = 25000 + else: + self.mag = self.tem.getMagnification() - Returns - ------- - arr : np.ndarray + self.ready = True + + def release_connection(self): + self.tem = None + self.ready = False + + def _stage_pose(self) -> np.ndarray: + """Compute the 4x4 affine transform of the stage. + + Multiply with a position on the stage (x, y) (including z=0 and + the final 1) to get a position relative to the optical axis. """ + x, y, z, a, b = self.tem.getStagePosition() + + # Using 'ZXY' intrinsic euler angles + Z = _euler_angle_z_to_matrix(self.camera_rotation_vs_stage_xy) + X = _euler_angle_x_to_matrix(a) + Y = _euler_angle_y_to_matrix(b) + + T = np.asarray( + [ + [1, 0, 0, -x], + [0, 1, 0, -y], + [0, 0, 1, -z], + [0, 0, 0, 1], + ] + ) + + return Z @ X @ Y @ T + + def get_realspace_image( + self, xx: np.ndarray, yy: np.ndarray, crystal_indices: np.ndarray + ) -> np.ndarray: + # Only crystals and void, no mesh or anything + out = np.zeros_like(xx) + for x, y, r, t in zip( + self.crystal_x[crystal_indices], + self.crystal_y[crystal_indices], + self.crystal_r[crystal_indices], + self.crystal_t[crystal_indices], + ): + # thickness multiplied by mask of where the crystal is + mask = ((xx - x) ** 2 + (yy - y) ** 2) < r**2 + out += t * mask.astype(float) + return out + + def get_diffraction_image( + self, shape: Tuple[int, int], crystal_indices: np.ndarray + ) -> np.ndarray: + out = np.zeros(shape) + pose = self._stage_pose() + + # Handle camera length + wavelength = microscope_config.wavelength + detector_half_width = self.physical_pixelsize * self.dimensions[0] // 2 + max_theta = np.arctan(detector_half_width / (self.camera_length)) + max_recip_len = 2 * np.sin(max_theta) / wavelength + for phi1, psi, phi2, t in zip( + self.crystal_euler_angle_phi_1[crystal_indices], + self.crystal_euler_angle_psi[crystal_indices], + self.crystal_euler_angle_phi_2[crystal_indices], + self.crystal_t[crystal_indices], + ): + # Crystal orientation, Bunge convention for euler angles + X1 = _euler_angle_x_to_matrix(phi1)[:-1, :-1] + Z = _euler_angle_z_to_matrix(psi)[:-1, :-1] + X2 = _euler_angle_x_to_matrix(phi2)[:-1, :-1] + R = pose[:-1, :-1] @ X1 @ Z @ X2 + beam_direction = R @ np.array([0, 0, -1]) + + # Find intersect with Ewald's sphere, by computing excitation error + # Instead of rotating all vectors, rotate Ewald's sphere (i.e. the beam) to find intersections. + # Then only rotate the intersecting vectors. + # Using notation from https://en.wikipedia.org/wiki/Line%E2%80%93sphere_intersection + r = 1 / wavelength + u = beam_direction + c = r * u + o = self.reflections + + diff = o.T - c + dot = np.dot(u, diff.T) + nabla = dot**2 - np.sum(diff**2, axis=1) + r**2 + # We know the relrods (assuming infinite length) are all going to intersect twice + # Therefore, no need to look at the sign of nabla + # We also know only the smaller root is important, i.e. -sqrt(nabla) + sqrt_nabla = np.sqrt(nabla) + d = -dot - sqrt_nabla + + intersection = np.abs(d) < self.max_excitation_error + vecs = self.reflections[:, intersection] + intensities = self.reflection_intensity[intersection] + # Linear excitation error scaling + intensities *= 1 - np.abs(d[intersection]) / self.max_excitation_error + + # Project vectors onto camera, ignoring curvature + projected_vecs = R.T @ vecs + + # Prepare image + for intensity, (xv, yv, _) in zip(intensities, projected_vecs.T): + x = xv / max_recip_len * shape[0] / 2 + shape[0] / 2 + y = yv / max_recip_len * shape[1] / 2 + shape[1] / 2 + if not 0 <= x < shape[0]: + continue + if not 0 <= y < shape[1]: + continue + min_x = round(max(0, x - self.spot_radius)) + max_x = round(min(shape[1], x + self.spot_radius)) + min_y = round(max(0, y - self.spot_radius)) + max_y = round(min(shape[0], y + self.spot_radius)) + out[min_y:max_y, min_x:max_x] = intensity + return out + + def get_image(self, exposure: float = None, binsize: int = None, **kwargs) -> np.ndarray: + self.actually_establish_connection() + if exposure is None: exposure = self.default_exposure - if not binsize: + if binsize is None: binsize = self.default_binsize - dim_x, dim_y = self.get_camera_dimensions() - - dim_x = int(dim_x / binsize) - dim_y = int(dim_y / binsize) - - time.sleep(exposure) - - arr = np.random.randint(256, size=(dim_x, dim_y)) - - return arr - - def get_movie( - self, - n_frames: int, - *, - exposure: Optional[float] = None, - binsize: Optional[int] = None, - **kwargs, - ) -> Generator[np.ndarray, None, None]: - """Movie acquisition routine. If the exposure and binsize are not - given, the default values are read from the config file. - - Parameters - ---------- - n_frames : int - Number of frames to collect - exposure : float - Exposure time in seconds. - binsize : int - Which binning to use. - - Yields - ------- - image : np.ndarray - """ - for _ in range(n_frames): - yield self.get_image(exposure=exposure, binsize=binsize, **kwargs) + mode = self.tem.getFunctionMode() + mag = self.tem.getMagnification() + if mode == 'diff': + self.camera_length = mag + else: + self.mag = mag + + shape_x, shape_y = self.get_camera_dimensions() + shape = (shape_x // binsize, shape_y // binsize) + + # Compute "illuminated area" (really, what part of the stage is visible on the image) + real_pixel_size = self.physical_pixelsize * 1e6 # nm + stage_pixel_size = real_pixel_size / self.mag + r = stage_pixel_size * (self.dimensions[0] ** 2 + self.dimensions[1] ** 2) ** 0.5 / 2 + + # Should gun shift also be considered? + x, y = self.tem.getBeamShift() + # x, y, r is now center and radius of beam, compared to optical axis + + # Compute pixel coordinates on stage that are illuminated + R = self._stage_pose() + rp = r / 2**0.5 # Convert from full radius to square inside the circle + x_beam_on_stage = np.linspace(x - rp, x + rp, shape[0]) + y_beam_on_stage = np.linspace(y - rp, y + rp, shape[1]) + # meshgrid of x, y, z and 1 (for the affine transform) + z = self.tem.getStagePosition()[2] + p_beam_on_stage = np.array( + [p.flatten() for p in np.meshgrid(x_beam_on_stage, y_beam_on_stage, [z], [1])] + ) + p_beam_in_column = np.linalg.inv(R) @ p_beam_on_stage + xx = p_beam_in_column[0, :].reshape(shape) + yy = p_beam_in_column[1, :].reshape(shape) + # xx and yy are now 2d coordinate arrays of each pixel on the stage + + # Find which crystals are illuminated + p_crystals = np.vstack( + ( + self.crystal_x, + self.crystal_y, + np.zeros(self.n_crystals), + np.ones(self.n_crystals), + ) + ) + x_beam_center_on_stage, y_beam_center_on_stage, _, _ = np.linalg.inv(R) @ np.array( + [x, y, 0, 1] + ) + c_ind = ( + (p_crystals[0, :] - x_beam_center_on_stage) ** 2 + + (p_crystals[1, :] - y_beam_center_on_stage) ** 2 + ) < (self.crystal_r + r) ** 2 + # c_ind is now a boolean array of which crystals contribute to the image/diffraction pattern + + if mode == 'diff': + return self.get_diffraction_image( + shape, + c_ind, + ) + else: + img = (self.get_realspace_image(xx, yy, c_ind) * 0xFFFF).astype(int) + return img def acquire_image(self) -> int: """For TVIPS compatibility.""" @@ -112,18 +392,6 @@ def get_image_dimensions(self) -> Tuple[int, int]: return dim_x, dim_y - def establish_connection(self) -> None: - """Establish connection to the camera.""" - res = 1 - if res != 1: - raise RuntimeError(f'Could not establish camera connection to {self.name}') - - def release_connection(self) -> None: - """Release the connection to the camera.""" - name = self.get_name() - msg = f"Connection to camera '{name}' released" - logger.info(msg) - # Mimic EMMENU API def get_emmenu_version(self) -> str: diff --git a/src/instamatic/simulation/__init__.py b/src/instamatic/simulation/__init__.py deleted file mode 100644 index e69de29b..00000000 diff --git a/src/instamatic/simulation/camera.py b/src/instamatic/simulation/camera.py deleted file mode 100644 index cc54ac48..00000000 --- a/src/instamatic/simulation/camera.py +++ /dev/null @@ -1,356 +0,0 @@ -from __future__ import annotations - -from typing import NamedTuple, Tuple - -import numpy as np - -from instamatic._typing import float_deg -from instamatic.camera.camera_base import CameraBase -from instamatic.config import microscope as microscope_config - - -class IlluminatedArea(NamedTuple): - """Area of stage that would show on an image. - - x, y: coordinates relative to optical axis - r: radius. - """ - - x: float - y: float - r: float - - -def _get_reciprocal_unit_cell( - a: float, b: float, c: float, alpha: float, beta: float, gamma: float -) -> np.ndarray: - # Copied from diffpy.structure - ca = np.cos(np.deg2rad(alpha)) - cb = np.cos(np.deg2rad(beta)) - cg = np.cos(np.deg2rad(gamma)) - sa = np.sin(np.deg2rad(alpha)) - sb = np.sin(np.deg2rad(beta)) - cgr = (ca * cb - cg) / (sa * sb) - sgr = np.sqrt(1.0 - cgr * cgr) - Vunit = np.sqrt(1.0 + 2.0 * ca * cb * cg - ca * ca - cb * cb - cg * cg) - ar = sa / (a * Vunit) - base = np.array( - [[1.0 / ar, -cgr / sgr / ar, cb * a], [0.0, b * sa, b * ca], [0.0, 0.0, c]], - ) - recbase = np.linalg.inv(base) - return recbase - - -def _euler_angle_z_to_matrix(theta: float_deg) -> np.ndarray: - c = np.cos(np.deg2rad(theta)) - s = np.sin(np.deg2rad(theta)) - return np.asarray( - [ - [c, -s, 0, 0], - [s, c, 0, 0], - [0, 0, 1, 0], - [0, 0, 0, 1], - ] - ) - - -def _euler_angle_x_to_matrix(theta: float_deg) -> np.ndarray: - c = np.cos(np.deg2rad(theta)) - s = np.sin(np.deg2rad(theta)) - return np.asarray( - [ - [1, 0, 0, 0], - [0, c, -s, 0], - [0, s, c, 0], - [0, 0, 0, 1], - ] - ) - - -def _euler_angle_y_to_matrix(theta: float_deg) -> np.ndarray: - c = np.cos(np.deg2rad(theta)) - s = np.sin(np.deg2rad(theta)) - return np.asarray( - [ - [c, 0, s, 0], - [0, 1, 0, 0], - [-s, 0, c, 0], - [0, 0, 0, 1], - ] - ) - - -class CameraSimulation(CameraBase): - streamable = True - - def __init__(self, name: str = 'simulate'): - """Simple simulation of the TEM. MUST START IN IMAGE MODE! - - Consisting of randomly dispersed crystals. - Each crystal is perfectly flat, and perfectly round. - All crystals have the same unit cell. - - The simulation assumes a 200kV convergent beam, always. - Magnification and/or camera length is read from the TEM, which the simulation connects to after a few seconds. - This is used in combination with the detector pixel size to scale the image/diffraction pattern. - - Image mode: - Crystals have different "thickness", yielding different gray-values proportional to thickness. - This gray-value never changes, regardless of tilt, orientation ect. - - Diffraction mode: - All crystals visible in image mode contribute to the diffraction pattern. - Only reflection positions are accurately simulated. - Intensities are computed by a single Gaussian, rather than proper structure factor computing. - Also scaled with excitation error - """ - super().__init__(name) - - self.ready = False - - # Real-space setup - grid_radius_nm = self.grid_diameter * 1e6 / 2 - self.crystal_x = np.random.uniform(-grid_radius_nm, grid_radius_nm, self.n_crystals) - self.crystal_y = np.random.uniform(-grid_radius_nm, grid_radius_nm, self.n_crystals) - self.crystal_r = np.random.uniform( - self.min_crystal_radius, self.max_crystal_radius, self.n_crystals - ) - self.crystal_t = np.random.uniform(0.4, 1.0, self.n_crystals) - self.crystal_euler_angle_phi_1 = np.random.uniform(0, 360, self.n_crystals) - self.crystal_euler_angle_psi = np.random.uniform(0, 180, self.n_crystals) - self.crystal_euler_angle_phi_2 = np.random.uniform(0, 360, self.n_crystals) - - # Reciprocal-space setup - reciprocal_unit_cell = _get_reciprocal_unit_cell( - self.unit_cell_a, - self.unit_cell_b, - self.unit_cell_c, - self.unit_cell_alpha, - self.unit_cell_beta, - self.unit_cell_gamma, - ) - d_min = 0.8 # Å - max_h = round(self.unit_cell_a / d_min) - max_k = round(self.unit_cell_b / d_min) - max_l = round(self.unit_cell_c / d_min) - hkl = np.asarray( - [ - [h, k, l] - for h in range(-max_h, max_h + 1) # noqa: E741 - for k in range(-max_k, max_k + 1) # noqa: E741 - for l in range(-max_l, max_l + 1) # noqa: E741 - ] - ).T - # Filter out the edges that are too high res - k_vec = reciprocal_unit_cell @ hkl - k = np.linalg.norm(k_vec, axis=0) - self.reflections = k_vec[:, k < 1 / d_min] - # Calculate intensity of reflections - k = k[k < 1 / d_min] - F = 5 * np.exp(-1.5 * k**2) # "Structure factor" scaling with a Gaussian - self.reflection_intensity = F**2 - - # Keep track of these, not accessible when outside corresponding mode - self.mag = None - self.camera_length = None - - def establish_connection(self): - pass - - def actually_establish_connection(self): - # Hack to establish connection with running TEM - if self.ready: - return - import time - - time.sleep(2) - from instamatic.controller import get_instance - - ctrl = get_instance() - self.tem = ctrl.tem - - if self.tem.getFunctionMode() == 'diff': - raise RuntimeError('Simulation cannot start in diffraction mode') - - self.mag = self.tem.getMagnification() - - self.ready = True - - def release_connection(self): - self.tem = None - self.ready = False - - def _stage_pose(self) -> np.ndarray: - """Compute the 4x4 affine transform of the stage. - - Multiply with a position on the stage (x, y) (including z=0 and - the final 1) to get a position relative to the optical axis. - """ - x, y, z, a, b = self.tem.getStagePosition() - - # Using 'ZXY' intrinsic euler angles - Z = _euler_angle_z_to_matrix(self.camera_rotation_vs_stage_xy) - X = _euler_angle_x_to_matrix(a) - Y = _euler_angle_y_to_matrix(b) - - T = np.asarray( - [ - [1, 0, 0, -x], - [0, 1, 0, -y], - [0, 0, 1, -z], - [0, 0, 0, 1], - ] - ) - - return Z @ X @ Y @ T - - def get_realspace_image( - self, xx: np.ndarray, yy: np.ndarray, crystal_indices: np.ndarray - ) -> np.ndarray: - # Only crystals and void, no mesh or anything - out = np.zeros_like(xx) - for x, y, r, t in zip( - self.crystal_x[crystal_indices], - self.crystal_y[crystal_indices], - self.crystal_r[crystal_indices], - self.crystal_t[crystal_indices], - ): - # thickness multiplied by mask of where the crystal is - mask = ((xx - x) ** 2 + (yy - y) ** 2) < r**2 - out += t * mask.astype(float) - return out - - def get_diffraction_image( - self, shape: Tuple[int, int], crystal_indices: np.ndarray - ) -> np.ndarray: - out = np.zeros(shape) - pose = self._stage_pose() - - # Handle camera length - wavelength = microscope_config.wavelength - detector_half_width = self.physical_pixelsize * self.dimensions[0] // 2 - max_theta = np.arctan(detector_half_width / (self.camera_length)) - max_recip_len = 2 * np.sin(max_theta) / wavelength - for phi1, psi, phi2, t in zip( - self.crystal_euler_angle_phi_1[crystal_indices], - self.crystal_euler_angle_psi[crystal_indices], - self.crystal_euler_angle_phi_2[crystal_indices], - self.crystal_t[crystal_indices], - ): - # Crystal orientation, Bunge convention for euler angles - X1 = _euler_angle_x_to_matrix(phi1)[:-1, :-1] - Z = _euler_angle_z_to_matrix(psi)[:-1, :-1] - X2 = _euler_angle_x_to_matrix(phi2)[:-1, :-1] - R = pose[:-1, :-1] @ X1 @ Z @ X2 - beam_direction = R @ np.array([0, 0, -1]) - - # Find intersect with Ewald's sphere, by computing excitation error - # Instead of rotating all vectors, rotate Ewald's sphere (i.e. the beam) to find intersections. - # Then only rotate the intersecting vectors. - # Using notation from https://en.wikipedia.org/wiki/Line%E2%80%93sphere_intersection - r = 1 / wavelength - u = beam_direction - c = r * u - o = self.reflections - - diff = o.T - c - dot = np.dot(u, diff.T) - nabla = dot**2 - np.sum(diff**2, axis=1) + r**2 - # We know the relrods (assuming infinite length) are all going to intersect twice - # Therefore, no need to look at the sign of nabla - # We also know only the smaller root is important, i.e. -sqrt(nabla) - sqrt_nabla = np.sqrt(nabla) - d = -dot - sqrt_nabla - - intersection = np.abs(d) < self.max_excitation_error - vecs = self.reflections[:, intersection] - intensities = self.reflection_intensity[intersection] - # Linear excitation error scaling - intensities *= 1 - np.abs(d[intersection]) / self.max_excitation_error - - # Project vectors onto camera, ignoring curvature - projected_vecs = R.T @ vecs - - # Prepare image - for intensity, (xv, yv, _) in zip(intensities, projected_vecs.T): - x = xv / max_recip_len * shape[0] / 2 + shape[0] / 2 - y = yv / max_recip_len * shape[1] / 2 + shape[1] / 2 - if not 0 <= x < shape[0]: - continue - if not 0 <= y < shape[1]: - continue - min_x = round(max(0, x - self.spot_radius)) - max_x = round(min(shape[1], x + self.spot_radius)) - min_y = round(max(0, y - self.spot_radius)) - max_y = round(min(shape[0], y + self.spot_radius)) - out[min_y:max_y, min_x:max_x] = intensity - return out - - def get_image(self, exposure: float = None, binsize: int = None, **kwargs) -> np.ndarray: - self.actually_establish_connection() - - if exposure is None: - exposure = self.default_exposure - if binsize is None: - binsize = self.default_binsize - - mode = self.tem.getFunctionMode() - mag = self.tem.getMagnification() - if mode == 'diff': - self.camera_length = mag - else: - self.mag = mag - - shape_x, shape_y = self.get_camera_dimensions() - shape = (shape_x // binsize, shape_y // binsize) - - # Compute "illuminated area" (really, what part of the stage is visible on the image) - real_pixel_size = self.physical_pixelsize * 1e6 # nm - stage_pixel_size = real_pixel_size / self.mag - r = stage_pixel_size * (self.dimensions[0] ** 2 + self.dimensions[1] ** 2) ** 0.5 / 2 - - # Should gun shift also be considered? - x, y = self.tem.getBeamShift() - # x, y, r is now center and radius of beam, compared to optical axis - - # Compute pixel coordinates on stage that are illuminated - R = self._stage_pose() - rp = r / 2**0.5 # Convert from full radius to square inside the circle - x_beam_on_stage = np.linspace(x - rp, x + rp, shape[0]) - y_beam_on_stage = np.linspace(y - rp, y + rp, shape[1]) - # meshgrid of x, y, z and 1 (for the affine transform) - z = self.tem.getStagePosition()[2] - p_beam_on_stage = np.array( - [p.flatten() for p in np.meshgrid(x_beam_on_stage, y_beam_on_stage, [z], [1])] - ) - p_beam_in_column = np.linalg.inv(R) @ p_beam_on_stage - xx = p_beam_in_column[0, :].reshape(shape) - yy = p_beam_in_column[1, :].reshape(shape) - # xx and yy are now 2d coordinate arrays of each pixel on the stage - - # Find which crystals are illuminated - p_crystals = np.vstack( - ( - self.crystal_x, - self.crystal_y, - np.zeros(self.n_crystals), - np.ones(self.n_crystals), - ) - ) - x_beam_center_on_stage, y_beam_center_on_stage, _, _ = np.linalg.inv(R) @ np.array( - [x, y, 0, 1] - ) - c_ind = ( - (p_crystals[0, :] - x_beam_center_on_stage) ** 2 - + (p_crystals[1, :] - y_beam_center_on_stage) ** 2 - ) < (self.crystal_r + r) ** 2 - # c_ind is now a boolean array of which crystals contribute to the image/diffraction pattern - - if mode == 'diff': - return self.get_diffraction_image( - shape, - c_ind, - ) - else: - img = (self.get_realspace_image(xx, yy, c_ind) * 0xFFFF).astype(int) - return img From 4e57ebbfa7d3f41962801334d86299dfe331e2a1 Mon Sep 17 00:00:00 2001 From: Viljar Femoen Date: Fri, 19 Sep 2025 17:04:32 +0200 Subject: [PATCH 4/6] Fix failing test (more accurate sim broke it) --- tests/test_ctrl.py | 1 + 1 file changed, 1 insertion(+) diff --git a/tests/test_ctrl.py b/tests/test_ctrl.py index 7771f683..7ef94a3c 100644 --- a/tests/test_ctrl.py +++ b/tests/test_ctrl.py @@ -229,6 +229,7 @@ def test_screen(ctrl): def test_align_to(ctrl): reference = ctrl.get_raw_image() pos = ctrl.stage.xy + ctrl.stage.x += 10 shift = ctrl.align_to(reference, apply=True) From f369af9eb193ee7e1efaa690dca0616d5293895d Mon Sep 17 00:00:00 2001 From: Viljar Femoen Date: Fri, 19 Sep 2025 17:04:48 +0200 Subject: [PATCH 5/6] Remove diffpy.structure from requirements --- pyproject.toml | 1 - requirements.txt | 1 - 2 files changed, 2 deletions(-) diff --git a/pyproject.toml b/pyproject.toml index 9485f68b..df763278 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -58,7 +58,6 @@ dependencies = [ "typing_extensions >= 4.0.0", "virtualbox >= 2.0.0", "pyserialem >= 0.3.2", - "diffpy.structure", ] [project.urls] diff --git a/requirements.txt b/requirements.txt index d0d58ae1..068b696c 100644 --- a/requirements.txt +++ b/requirements.txt @@ -17,4 +17,3 @@ tqdm >= 4.41.1 typing_extensions >= 4.0.0 virtualbox >= 2.0.0 pyserialem >= 0.3.2 -diffpy.structure From 215de3d5c76f854d63d58ca1820dc0a89b5039bf Mon Sep 17 00:00:00 2001 From: Viljar Femoen Date: Mon, 22 Sep 2025 11:04:37 +0200 Subject: [PATCH 6/6] Ensure one crystal at (0, 0), and add cross at (0, 0) --- src/instamatic/camera/camera_simu.py | 30 +++++++++++++++++++++++----- 1 file changed, 25 insertions(+), 5 deletions(-) diff --git a/src/instamatic/camera/camera_simu.py b/src/instamatic/camera/camera_simu.py index 96e84515..ba9d559d 100644 --- a/src/instamatic/camera/camera_simu.py +++ b/src/instamatic/camera/camera_simu.py @@ -80,6 +80,7 @@ def __init__(self, name: str = 'simulate'): """Simple simulation of the TEM. MUST START IN IMAGE MODE! Consisting of randomly dispersed crystals. + One crystal is guaranteeed to be at (0, 0) on the stage Each crystal is perfectly flat, and perfectly round. All crystals have the same unit cell. @@ -90,6 +91,7 @@ def __init__(self, name: str = 'simulate'): Image mode: Crystals have different "thickness", yielding different gray-values proportional to thickness. This gray-value never changes, regardless of tilt, orientation ect. + A cross centered at (0, 0) is added for convenience Diffraction mode: All crystals visible in image mode contribute to the diffraction pattern. @@ -108,7 +110,7 @@ def __init__(self, name: str = 'simulate'): self.unit_cell_alpha = 90 # Degrees self.unit_cell_beta = 90 # Degrees self.unit_cell_gamma = 120 # Degrees - self.max_excitation_error = 0.01 + self.max_excitation_error = 0.005 self.spot_radius = 3 # pixels super().__init__(name) @@ -126,6 +128,11 @@ def __init__(self, name: str = 'simulate'): self.crystal_euler_angle_psi = np.random.uniform(0, 180, self.n_crystals) self.crystal_euler_angle_phi_2 = np.random.uniform(0, 360, self.n_crystals) + # Ensure one crystal is always at (0, 0) and 1µm radius + self.crystal_x[0] = 0 + self.crystal_y[0] = 0 + self.crystal_r[0] = 1000 + # Reciprocal-space setup reciprocal_unit_cell = _get_reciprocal_unit_cell( self.unit_cell_a, @@ -193,6 +200,11 @@ def actually_establish_connection(self): else: self.mag = self.tem.getMagnification() + # If the TEM is a simulated one (i.e. random beam shift), reset beam shift + # Otherwise, reseting the stage will not get you to (0, 0) + if self.tem.name == 'simulate': + self.tem.setBeamShift(0, 0) + self.ready = True def release_connection(self): @@ -237,7 +249,14 @@ def get_realspace_image( # thickness multiplied by mask of where the crystal is mask = ((xx - x) ** 2 + (yy - y) ** 2) < r**2 out += t * mask.astype(float) - return out + # Invert and scale + out = 0xF000 * (1 - out) + # Add some noise + out *= np.random.uniform(0.9, 1.1, out.shape) + # Add cross at (0, 0) + width = 50 # nm + out[(np.abs(xx) < width) | (np.abs(yy) < width)] = 0 + return out.astype(int) def get_diffraction_image( self, shape: Tuple[int, int], crystal_indices: np.ndarray @@ -303,7 +322,9 @@ def get_diffraction_image( min_y = round(max(0, y - self.spot_radius)) max_y = round(min(shape[0], y + self.spot_radius)) out[min_y:max_y, min_x:max_x] = intensity - return out + # Scale. Direct beam intensity is 25 + out = out * 0x8000 / 25 + return out.astype(int) def get_image(self, exposure: float = None, binsize: int = None, **kwargs) -> np.ndarray: self.actually_establish_connection() @@ -371,8 +392,7 @@ def get_image(self, exposure: float = None, binsize: int = None, **kwargs) -> np c_ind, ) else: - img = (self.get_realspace_image(xx, yy, c_ind) * 0xFFFF).astype(int) - return img + return self.get_realspace_image(xx, yy, c_ind) def acquire_image(self) -> int: """For TVIPS compatibility."""