""" Generate SIM patterns using lattice periodicity vectors Va and Vb, and duplicating roi_size single unit cell. See the supplemental material of https://doi.org/10.1038/nmeth.1734 for more discussion of similar approaches. """ from typing import Union, Optional from collections.abc import Sequence from warnings import warn from time import perf_counter import datetime import json from pathlib import Path import numpy as np from PIL import Image from scipy.fft import fftshift, fftfreq from scipy.signal import fftconvolve from scipy.signal.windows import hann import tifffile import matplotlib.pyplot as plt from matplotlib.figure import Figure from matplotlib.colors import PowerNorm from matplotlib.patches import Circle from mcsim.analysis.sim_reconstruction import get_peak_value from mcsim.analysis.fft import conj_transpose_fft, ft2 from mcsim.analysis.simulate_dmd import xy2uvector, blaze_envelope, _dlp_1stgen_axis from localize_psf.affine import xform_sinusoid_params, xform_mat, params2xform try: import cupy as cp except ImportError: cp = None if cp: array = Union[np.ndarray, cp.ndarray] else: array = np.ndarray def get_sim_pattern(dmd_size: Sequence[int, int], vec_a: array, vec_b: array, nphases: int, phase_index: int, geometry: str = "orthogonal") -> tuple[array, array]: """ Convenience function for generating SIM patterns from the tile_patterns() function. This function can run either on the CPU or the GPU. If vec_a is a NumPy array, it will run on the CPU. If vec_a is a CuPy array it will run on the GPU. :param dmd_size: [nx, ny] :param vec_a: first lattice vector, [dxa, dya] :param vec_b: second lattice vector, [dxb, dyb] :param nphases: number of phase shifts required. This effects the filling of the pattern :param phase_index: integer in range(nphases) :param geometry: "orthogonal" for pixel addressing geometries similar to the DLP6500, or "diamond" for pixel geometries like the DLP4500. :return pattern, cell: 'pattern' is an array giving the desired pattern and 'cell' is an array giving a single unit cell of the pattern """ xp = cp if cp and isinstance(vec_a, cp.ndarray) else np vec_a = xp.asarray(vec_a) vec_b = xp.asarray(vec_b) if not _verify_int(vec_b / nphases): raise ValueError("At least one component of vec_b was not divisible by nphases") cell, x_cell, y_cell = get_sim_unit_cell(vec_a, vec_b, nphases, phase_index) nx, ny = dmd_size # get indices of pixels jj, ii = np.meshgrid(xp.arange(nx), xp.arange(ny)) # convert indices to geometry if geometry == "orthogonal": x = jj y = ii elif geometry == "diamond": p = jj * xp.sqrt(2) + xp.mod(ii, 2) / xp.sqrt(2) m = ii / xp.sqrt(2) x = (p + m) / xp.sqrt(2) y = (p - m) / xp.sqrt(2) else: raise ValueError(f"Geometry `{geometry:s}` is not supported.") pts = xp.stack((x, y), axis=-1) pts_cell, na, nb = reduce2cell(pts, vec_a, vec_b) pts_cell = xp.round(pts_cell, 12).astype(int) # # generate pattern by choosing coordinate in unit cell xinds = pts_cell[..., 0] - x_cell[0] yinds = pts_cell[..., 1] - y_cell[0] pattern = cell[(yinds, xinds)].astype(bool) return pattern, cell # tool for manipulating unit cells def tile_pattern(dmd_size: Sequence[int, int], vec_a: array, vec_b: array, start_coord: Sequence[int, int], cell: array, x_cell: array, y_cell: array, do_cell_reduction: bool = True) -> array: """ Generate SIM patterns using lattice vectors vec_a = [dxa, dya] and vec_b = [dxb, 0], and duplicating roi_size single unit cell. See the supplemental material of https://doi.org/10.1038/nmeth.1734 for more information. Note: we interpret the pattern params(x, y) = M[i_y, i_x], where M is the matrix representing the pattern. Matlab will display the matrix with i_y = 0 on top, so the pattern we really want is the matrix flipped along the first dimension. :param dmd_size: [nx, ny] :param vec_a: [dxa, dya] :param vec_b: [dxb, dyb] :param start_coord: [x, y]. DMD coordinate where the origin of the unit cell will be situated. These coordinates are relative to the image corner :param cell: :param x_cell: :param y_cell: :param do_cell_reduction: whether to call get_minimal_cell() before tiling :return pattern: """ xp = cp if cp and isinstance(vec_a, cp.ndarray) else np # todo: much slower than the old function because looping and doing pixel assignment instead of concatenating vec_a = xp.array(vec_a, copy=True) vec_b = xp.array(vec_b, copy=True) start_coord = xp.array(start_coord, copy=True) nx, ny = dmd_size if do_cell_reduction: # this will typically make the vectors shorter and more orthogonal, so tiling is easier cell, x_cell, y_cell, vec_a, vec_b = get_minimal_cell(cell, x_cell, y_cell, vec_a, vec_b) pattern = xp.zeros((ny, nx)) * xp.nan dy, dx = cell.shape # find maximum integer multiples of the periodicity vectors that we need # n * vec_a + m * vec_b + start_coord = corners # [[dxa, dxb], [dya, dyb]] * [[n], [m]] = [[cx], [cy]] - [[sx], [sy]] sx, sy = start_coord mat = xp.linalg.inv(xp.array([[vec_a[0], vec_b[0]], [vec_a[1], vec_b[1]]])) n1, m1 = mat.dot(xp.array([[pattern.shape[1] - sx], [pattern.shape[0] - sy]], dtype=float)) n2, m2 = mat.dot(xp.array([[0. - sx], [pattern.shape[0] - sy]], dtype=float)) n3, m3 = mat.dot(xp.array([[pattern.shape[1] - sx], [0. - sy]], dtype=float)) n4, m4 = mat.dot(xp.array([[0. - sx], [0. - sy]], dtype=float)) na_min = int(xp.floor(xp.array([n1, n2, n3, n4]).min())) na_max = int(xp.ceil(xp.array([n1, n2, n3, n4]).max())) nb_min = int(xp.floor(xp.array([m1, m2, m3, m4]).min())) nb_max = int(xp.ceil(xp.array([m1, m2, m3, m4]).max())) niterations = (na_max - na_min) * (nb_max - nb_min) if niterations > 1e3: # if number of iterations is large, reduce number of tilings required by doubling unit cell # todo: could probably find the optimal number of doublings/tilings. This is important to get this to be fast # todo: in a general case # todo: right now pretty fast for 'reasonable' patterns, but still seems to be slow for some sets of patterns. na_max_doublings = int(xp.floor(xp.log2((na_max - na_min)))) na_doublings = int(xp.array([int(xp.round(na_max_doublings / 2)), 1]).max()) nb_max_doublings = int(xp.floor(xp.log2((nb_max - nb_min)))) nb_doublings = int(xp.array([int(xp.round(nb_max_doublings / 2)), 1]).max()) large_pattern, xpatt, ypatt = double_cell(cell, x_cell, y_cell, vec_a, vec_b, na=na_doublings, nb=nb_doublings) # finish by tiling pattern = tile_pattern(dmd_size, 2 ** na_doublings * vec_a, 2 ** nb_doublings * vec_b, start_coord, large_pattern, xpatt, ypatt, do_cell_reduction=False) else: # for smaller iteration number, tile directly for n in range(na_min, na_max + 1): for m in range(nb_min, nb_max + 1): # account for fact the origin of the cell may not be at the lower left corner. # (0, 0) position of the cell should be at vec_a * n + vec_b * m + start_coord xzero, yzero = vec_a * n + vec_b * m + start_coord xstart = int(xzero + xp.min(x_cell)) ystart = int(yzero + xp.min(y_cell)) xend = xstart + int(dx) yend = ystart + int(dy) if xend < 0 or yend < 0 or xstart > pattern.shape[1] or ystart > pattern.shape[0]: continue if xstart < 0: xstart_cell = -xstart xstart = 0 else: xstart_cell = 0 if xend > pattern.shape[1]: xend = pattern.shape[1] xend_cell = xstart_cell + (xend - xstart) if ystart < 0: ystart_cell = -ystart ystart = 0 else: ystart_cell = 0 if yend > pattern.shape[0]: yend = pattern.shape[0] yend_cell = ystart_cell + (yend - ystart) pattern[ystart:yend, xstart:xend] = xp.nansum( np.concatenate((pattern[ystart:yend, xstart:xend, None], cell[ystart_cell:yend_cell, xstart_cell:xend_cell, None]), axis=2), axis=2) assert not xp.any(xp.isnan(pattern)) pattern = xp.asarray(pattern, dtype=bool) return pattern def double_cell(cell: array, x: array, y: array, vec_a: array, vec_b: array, na: int = 1, nb: int = 0) -> tuple[np.ndarray, np.ndarray, np.ndarray]: """ Create new unit cell by doubling the original one by a factor of na along vec_a and nb along vec_b :param cell: initial cell :param x: x-coordinates of cell :param y: y-coordinates of cell :param vec_a: periodicity vector a :param vec_b: periodicity vector b :param na: number of times to double unit cell along vec_a :param nb: number of times to double cell along vec_b :return big_cell, xs, ys: doubled cell and x- and y-coordinates of double cell """ xp = cp if cp and isinstance(cell, cp.ndarray) else np x = xp.asarray(x) y = xp.asarray(y) vec_a = xp.array(vec_a, copy=True) vec_b = xp.array(vec_b, copy=True) if not (na == 1 and nb == 0): big_cell = cell xs = x ys = y for ii in range(na): big_cell, xs, ys = double_cell(big_cell, xs, ys, 2**ii * vec_a, vec_b, na=1, nb=0) for jj in range(nb): big_cell, xs, ys = double_cell(big_cell, xs, ys, 2**jj * vec_b, 2**na * vec_a, na=1, nb=0) else: dyc, dxc = cell.shape v1 = xp.array([0, 0]) v2 = 2*vec_a v3 = vec_b v4 = 2*vec_a + vec_b xs = xp.arange(int(xp.stack([v1[0], v2[0], v3[0], v4[0]]).min()), int(xp.stack([v1[0], v2[0], v3[0], v4[0]]).max()) + 1) ys = xp.arange(int(xp.stack([v1[1], v2[1], v3[1], v4[1]]).min()), int(xp.stack([v1[1], v2[1], v3[1], v4[1]]).max()) + 1) big_cell = xp.zeros((len(ys), len(xs))) * xp.nan for n in [0, 1]: xzero, yzero = vec_a * n # todo: should round or ceil/floor? istart_x = int(xzero - xp.min(xs) + xp.min(x)) istart_y = int(yzero - xp.min(ys) + xp.min(y)) big_cell[istart_y:istart_y+dyc, istart_x:istart_x+dxc][np.logical_not(np.isnan(cell))] = \ cell[np.logical_not(np.isnan(cell))] return big_cell, xs, ys def get_sim_unit_cell(vec_a: array, vec_b: array, nphases: int, phase_index: int = 0) -> tuple[array, array, array]: """ Get unit cell, which can be repeated to form SIM pattern. :param vec_a: first lattice vector :param vec_b: second lattice vector :param nphases: number of phase shifts. Required to determine the on and off pixels in cell. :param phase_index: :return cell, x, y: square array representing cell. Ones and zeroes give on and off points, and nans are points that are not part of the unit cell, but are necessary to pad the array to make it squares """ xp = cp if cp is not None and isinstance(vec_a, cp.ndarray) else np vec_a = xp.asarray(vec_a) vec_b = xp.asarray(vec_b) # ensure both vec_b components are divisible by nphases if not _verify_int(vec_b / nphases): raise ValueError("At least one component of vec_b was not divisible by nphases") # get full unit cell cell, x_cell, y_cell = get_unit_cell(vec_a, vec_b) # get reduced unit cell from vec_a, vec_b/nphases. If we set all of these positions to 1, # then we get perfect tiling. vec_b_sub = vec_b // nphases cell_sub, x_cell_sub, y_cell_sub = get_unit_cell(vec_a, vec_b_sub) to_use = xp.logical_not(np.isnan(cell_sub)) # get coordinates in main cell xx_cs, yy_cs = np.meshgrid(x_cell_sub, y_cell_sub) xx_cs += vec_b_sub[0] * phase_index yy_cs += vec_b_sub[1] * phase_index # ensure using coordinates in unit cell pts_cell, _, _ = reduce2cell(xp.stack((xx_cs, yy_cs), axis=-1), vec_a, vec_b) pts_cell = xp.round(pts_cell, 12).astype(int) # convert to indices and set pattern xinds = (pts_cell[..., 0] - x_cell[0])[to_use] yinds = (pts_cell[..., 1] - y_cell[0])[to_use] cell[(yinds, xinds)] = 1 with np.errstate(invalid='ignore'): if xp.nansum(cell) != xp.sum(cell >= 0) / nphases: raise ValueError("Cell does not have appropriate number of 'on' pixels") return cell, x_cell, y_cell def get_unit_cell(vec_a: array, vec_b: array) -> tuple[array, array, array]: """ Generate a mask which represents one unit cell of a pattern for given vectors. This mask is a square array with NaNs at positions outside the unit cell, and zeros at points in the cell. The unit cell is the area enclosed by [0, vec_a, vec_b, vec_a + vec_b]. For pixels, we say that an entire pixel is within the cell if its center is. For a pixel with center exactly on one of the edges of the cell, we say it is inside if it lies on the lines from [0, vec_b] or [0, vec_a] and outside if its lies on the lines from [vec_a, vec_a + vec_b] or [vec_b, vec_a + vec_b]. This choice avoids including pixels twice. :param vec_a: [dxa, dya] :param vec_b: [dxb, dyb] :return cell, x, y: cell values are 0 for valid points and NaN for invalid points. x- and y- coordinates """ xp = cp if cp and isinstance(vec_a, cp.ndarray) else np # test that vec_a and vec_b components are integers if not _verify_int(vec_a) or not _verify_int(vec_b): raise ValueError(f"At least one component of lattice vectors={vec_a}, {vec_b} cannot " f"be interpreted as an integer") # copy vector data, so don't affect inputs vec_a = xp.array(vec_a, copy=True, dtype=int) vec_b = xp.array(vec_b, copy=True, dtype=int) # check vectors are linearly independent if (vec_a[0] * vec_b[1] - vec_a[1] * vec_b[0]) == 0: raise ValueError("vec_a and vec_b are linearly dependent.") # square array containing unit cell, with points not in unit cell nans dy = xp.abs(vec_a[1]) + xp.abs(vec_b[1]) dx = xp.abs(vec_a[0]) + xp.abs(vec_b[0]) # x-coordinates massaged so that origin is at x=0 x = xp.arange(dx) if vec_a[0] < 0 and vec_b[0] >= 0: x = x + vec_a[0] + 1 elif vec_a[0] >= 0 and vec_b[0] < 0: x = x + vec_b[0] + 1 elif vec_a[0] < 0 and vec_b[0] < 0: x = x + vec_a[0] + vec_b[0] + 1 # y-coordinates massaged so that origin is at y=0 y = xp.arange(dy) if vec_a[1] < 0 and vec_b[1] >= 0: y = y + vec_a[1] + 1 elif vec_a[1] >= 0 and vec_b[1] < 0: y = y + vec_b[1] + 1 elif vec_a[1] < 0 and vec_b[1] < 0: y = y + vec_a[1] + vec_b[1] + 1 xx, yy = xp.meshgrid(x, y) # generate cell pts = xp.stack((xx, yy), axis=-1) cell = xp.array(test_in_cell(pts, vec_a, vec_b), dtype=float) cell[cell == False] = np.nan cell[cell == True] = 0 # check unit cell has correct volume cell_volume = xp.abs(vec_a[0] * vec_b[1] - vec_a[1] * vec_b[0]) assert xp.nansum(xp.logical_not(xp.isnan(cell))) == cell_volume return cell, x, y def test_in_cell(points: array, va: array, vb: array) -> array: """ Test if points (x, y) are in the unit cell for a given pair of unit vectors. We suppose the unit cell is the region enclosed by 0, va, vb, and va + vb. Point on the boundary are considered inside if they are on the lines 0 -> va or 0 ->vb, and outside if they are on the lines va -> va+vb or vb -> va + vb :param points: array of size n0 x n1 x ... x nm x 2, where points[..., 0] are the x-coordinates and points[..., 1] are the y-coordinates :param va: first lattice vector :param vb: second lattice vector :return in_cell: boolean array of size n0 x n1 x ... x nm indicating if coordinate is in the unit cell or not """ xp = cp if cp and isinstance(va, cp.ndarray) else np va = xp.array(va, copy=True).ravel() vb = xp.array(vb, copy=True).ravel() x = points[..., 0] y = points[..., 1] def line(x, p1, p2): return ((p2[1] - p1[1]) * x + p1[1] * p2[0] - p1[0] * p2[1]) / (p2[0] - p1[0]) precision = 10 # strategy: consider parellel lines from line1 = [0,0] -> va and line2 = vb -> va + vb # if point is on opposite sides of line1 and line2, or exactly on line1 then it is inside the cell # if it is one of the same sides of line1 and line2, or exactly on line2, it is outside if va[0] != 0: gthan_a1 = xp.round(line(x, [0, 0], va), precision) > xp.round(y, precision) eq_a1 = xp.round(line(x, [0, 0], va), precision) == xp.round(y, precision) gthan_a2 = xp.round(line(x, vb, va + vb), precision) > xp.round(y, precision) eq_a2 = xp.round(line(x, vb, va + vb), precision) == xp.round(y, precision) else: # if x-component of va = 0. Then x-component of vb cannot be zero, else linearly dependent gthan_a1 = xp.round(x, precision) > 0 eq_a1 = xp.round(x, precision) == 0 gthan_a2 = xp.round(x, precision) > xp.round(vb[0], precision) eq_a2 = xp.round(x, precision) == xp.round(vb[0], precision) # same strategy for vb if vb[0] != 0: gthan_b1 = xp.round(line(x, [0, 0], vb), precision) > xp.round(y, precision) eq_b1 = xp.round(line(x, [0, 0], vb), precision) == xp.round(y, precision) gthan_b2 = xp.round(line(x, va, va + vb), precision) > xp.round(y, precision) eq_b2 = xp.round(line(x, va, va + vb), precision) == xp.round(y, precision) else: # if x-component of vb = 0. Then x-component of va cannot be zero, else linearly dependent gthan_b1 = xp.round(x, precision) > 0 eq_b1 = xp.round(x, precision) == 0 gthan_b2 = xp.round(x, precision) > xp.round(va[0], precision) eq_b2 = xp.round(x, precision) == xp.round(va[0], precision) in_cell_a = xp.logical_and(np.logical_or(gthan_a1 != gthan_a2, eq_a1), xp.logical_not(eq_a2)) in_cell_b = xp.logical_and(np.logical_or(gthan_b1 != gthan_b2, eq_b1), xp.logical_not(eq_b2)) in_cell = xp.logical_and(in_cell_a, in_cell_b) return in_cell def reduce2cell(point: array, va: array, vb: array) -> tuple[array, array, array]: """ Given a vector, reduce it to coordinates within the unit cell. To ensure all reduced points actually lie in the unit cell, run >>> assert np.all(test_in_cell(point, va, vb)[0], va, vb) :param point: of size n0 x n1 x ... x nm x 2 :param va: first lattice vector, of size 2 x 1 :param vb: second lattice vector, of size 2 x 1 :return point_red, na_out, nb_out: """ xp = cp if cp is not None and isinstance(point, cp.ndarray) else np if not _verify_int(va) or not _verify_int(vb): raise ValueError(f"At least one component of lattice vectors={va}, {vb} cannot " f"be interpreted as an integer") point = xp.asarray(point) va = xp.asarray(va) vb = xp.asarray(vb) ra, rb = get_reciprocal_vects(va, vb) # need to round to avoid problems with machine precision na_out = xp.floor(xp.round(xp.dot(point, ra), 10)).astype(int) nb_out = xp.floor(xp.round(xp.dot(point, rb), 10)).astype(int) point_red = point - (na_out * va + nb_out * vb) return point_red, na_out, nb_out def convert_cell(cell1: array, x1: array, y1: array, va1: array, vb1: array, va2: array, vb2: array) -> tuple[array, array, array]: """ Given a unit cell described by vectors va1 and vb2, convert to equivalent description in terms of va2, vb2 :param cell1: initial cell :param x1: initial coordinates :param y1: :param va1: initial lattice vectors :param vb1: :param va2: new lattice vectors :param vb2: :return cell2, x2, y2: """ xp = cp if cp and isinstance(cell1, cp.ndarray) else np x1 = xp.asarray(x1) y1 = xp.asarray(y1) va1 = xp.asarray(va1) vb1 = xp.asarray(vb1) va2 = xp.asarray(va2) vb2 = xp.asarray(vb2) # todo: add check that va1/vb1 and va2/vb2 describe same lattice for vec in [va1, vb1, va2, vb2]: if not _verify_int(vec): raise ValueError(f"At least one component of lattice {vec} cannot " f"be interpreted as an integer") cell2, x2, y2 = get_unit_cell(va2, vb2) xx2, yy2 = xp.meshgrid(x2, y2) p1, _, _ = reduce2cell(xp.stack((xx2.ravel(), yy2.ravel()), axis=1), va1, vb1) xind = p1[:, 0] - x1.min() yind = p1[:, 1] - y1.min() cell2 += cell1[(yind, xind)].reshape(xx2.shape) # preserver NaNs return cell2, x2, y2 def get_minimal_cell(cell: array, x: array, y: array, va: array, vb: array) -> tuple[array, array, array, array, array]: """ Convert to cell using the smallest lattice vectors :param cell: :param x: :param y: :param va: first lattice vector :param vb: second lattice vector :return cell_m, x_m, y_m, va_m, vb_m: """ va_m, vb_m = reduce_basis(va, vb) cell_m, x_m, y_m = convert_cell(cell, x, y, va, vb, va_m, vb_m) return cell_m, x_m, y_m, va_m, vb_m def show_cell(v1: np.ndarray, v2: np.ndarray, cell: np.ndarray, x: np.ndarray, y: np.ndarray, aspect_equal: bool = True, ax=None, **kwargs): """ Plot unit cell and periodicity vectors :param v1: first lattice vector :param v2: second lattice vector :param cell: array representing values of pattern in unit cell :param x: x-coordinates of cell :param y: y-coordinates of cell :param aspect_equal: :param ax: axes on which to plot cell. If not provided, generate a new figure :return fig, ax: handle to resulting figure and axes """ if ax is None: fig = plt.figure(**kwargs) ax = fig.add_subplot(1, 1, 1) else: fig = ax.get_figure() ax.set_title("Unit cell") # plot cell if aspect_equal: aspect = "equal" else: aspect = "auto" ax.imshow(np.abs(cell), origin='lower', extent=[x[0] - 0.5, x[-1] + 0.5, y[0] - 0.5, y[-1] + 0.5], aspect=aspect, interpolation="none") # plot lattice vectors if v1 is not None and v2 is not None: v1 = np.asarray(v1).ravel() v2 = np.asarray(v2).ravel() ax.plot([0, v1[0]], [0, v1[1]], 'r', label=f"$v_1$=({v1[0]:d}, {v1[1]:d})") ax.plot([0, v2[0]], [0, v2[1]], 'g', label=f"$v_2$=({v2[0]:d}, {v2[1]:d})") ax.plot([v2[0], v2[0] + v1[0]], [v2[1], v2[1] + v1[1]], 'r') ax.plot([v1[0], v2[0] + v1[0]], [v1[1], v2[1] + v1[1]], 'g') ax.legend() ax.set_xlabel("x") ax.set_ylabel("y") return fig, ax # determine parameters of SIM patterns def get_reciprocal_vects(vec_a: array, vec_b: array) -> tuple[array, array]: """ Compute the reciprocal vectors for (real-space) lattice vectors vec_a and vec_b. If we call the lattice vectors a_i and the reciprocal vectors b_j, then these should be defined such that dot(a_i, b_j) = delta_{ij}, i.e. exp[ i 2*pi*ai * bj] = 1. This means the b_j are frequency like (i.e. equivalent of Hz). Note that they are instead sometimes defined dot(a_i, b_j) = 2*pi * delta_{ij}, which would make them angular-frequency like (i.e. in radians). :param vec_a: first lattice vector, [vx, vy] :param vec_b: second lattice vector :return rv1, rv2: frequency-like reciprocal lattice vectors """ xp = cp if cp and isinstance(vec_a, cp.ndarray) else np vec_a = xp.asarray(vec_a) vec_b = xp.asarray(vec_b) # check this directly, as sometimes due to numerical precision np.linalg.inv() will not throw error err_msg = "vec_a and vec_b are linearly dependent, so their reciprocal vectors could not be computed." if (vec_a[0] * vec_b[1] - vec_a[1] * vec_b[0]) == 0: raise ValueError(err_msg) a_mat = xp.stack((vec_a, vec_b), axis=0) try: inv_a = xp.linalg.inv(a_mat) except xp.linalg.LinAlgError: raise ValueError(err_msg) rv1 = inv_a[:, 0][:, None] rv2 = inv_a[:, 1][:, None] # todo: return 1D vectors instead of 2D vectors # rv1 = inv_a[:, 0] # rv2 = inv_a[:, 1] return rv1, rv2 def get_sim_angle(vec_a: array, vec_b: array) -> float: """ Get angle of SIM pattern in radians :param vec_a: first lattice vector, [vx, vy] :param vec_b: second lattice vector :return angle: angle in radians """ xp = cp if cp and isinstance(vec_a, cp.ndarray) else np recp_va, recp_vb = get_reciprocal_vects(vec_a, vec_b) angle = xp.angle(recp_vb[0, 0] + 1j * recp_vb[1, 0]) return float(xp.mod(angle, 2*np.pi)) def get_sim_period(vec_a: array, vec_b: array) -> float: """ Get period of SIM pattern constructed from periodicity vectors. The period is the distance between parallel lines pointing in the direction of vec_a passing through the points 0 and vec_b_temp respectively. We construct this by taking the projection of vec_b along the perpendicular to vec_a. NOTE: to say this another way, the period is given by the reciprocal lattice vector orthogonal to vec_a. :param vec_a: first lattice vector, [vx, vy] :param vec_b: second lattice vector :return period: pattern period """ xp = cp if cp and isinstance(vec_a, cp.ndarray) else np vec_b = xp.asarray(vec_b) uvec_perp_a = xp.array([vec_a[1], -vec_a[0]]) / np.sqrt(vec_a[0]**2 + vec_a[1]**2) period = xp.abs(uvec_perp_a.dot(vec_b)) return float(period) def get_sim_frqs(vec_a: array, vec_b: array) -> tuple[float, float]: """ Get spatial frequency of SIM pattern constructed from periodicity vectors. :param vec_a: first lattice vector, [vx, vy] :param vec_b: second lattice vector :return fx, fy: """ recp_va, recp_vb = get_reciprocal_vects(vec_a, vec_b) fx = recp_vb[0, 0] fy = recp_vb[1, 0] return fx, fy def get_sim_phase(vec_a: array, vec_b: array, nphases: int, phase_index: int, pattern_size: Sequence[int], use_fft_origin: bool = True): """ Get phase of dominant frequency component in the SIM pattern. .. math:: P(x, y) = 0.5 \\left(1 + \\cos(2 \\pi f_x x + 2\\pi f_y y + \\phi \\right) :param vec_a: first lattice vector, [vx, vy] :param vec_b: second lattive vector :param nphases: number of equal phase shifts for SIM pattern :param phase_index: 0, ..., nphases-1 :param pattern_size: [nx, ny] :param use_fft_origin: origin to use for computing the phase. If True, will assume the coordinates are the same as used in an FFT (i.e. before performing an ifftshift, with the 0 near the center). If False, will suppose the origin is at pattern[0, 0]. :return phase: phase of the SIM pattern at the dominant frequency component (which is recp_vec_b) """ xp = cp if cp and isinstance(vec_a, cp.ndarray) else np vec_b = xp.asarray(vec_b) cell, xs, ys = get_sim_unit_cell(vec_a, vec_b, nphases) fourier_component, _ = get_pattern_fourier_component(cell, xs, ys, vec_a, vec_b, 0, 1, nphases, phase_index, use_fft_origin=use_fft_origin, dmd_size=pattern_size) phase = xp.angle(fourier_component) return xp.mod(phase, 2*np.pi) def ldftfreq(vec_a: np.ndarray, vec_b: np.ndarray) -> tuple[array, array, array]: """ Get the Fourier frequencies for the lattice fourier transform (LDFT). See ldft2() for more details. This function is an analog of fftfreq(). :param vec_a: first lattice vector :param vec_b: second lattice vector :return fvecs, f1, f2: fvecs are the frequencies. Return NaNs for values outside the unit cell. f1 and f2 are integers which define the frequencies as multiplies of the reciprocal lattice vectors b1 and b2. fvecs = f1 * b1 + f2 * b2. """ xp = cp if cp and isinstance(vec_a, cp.ndarray) else np b1, b2 = get_reciprocal_vects(vec_a, vec_b) det = _convert_int(xp.linalg.det(xp.stack((vec_a, vec_b), axis=1))) b1_int_approx = b1.ravel() * det b2_int_approx = b2.ravel() * det b1_int = _convert_int(b1_int_approx) b2_int = _convert_int(b2_int_approx) # todo: to match DFT, want this to produce one more negative frequency than positive frequency # todo: don't think this works that way at the moment? bcell, f1, f2 = get_unit_cell(b1_int, b2_int) fvecs = xp.expand_dims(f1, axis=(0, 2)) * xp.expand_dims(b1.ravel(), axis=(0, 1)) + \ xp.expand_dims(f2, axis=(1, 2)) * xp.expand_dims(b2.ravel(), axis=(0, 1)) fvecs[xp.isnan(bcell), :] = np.nan return fvecs, f1, f2 def ldft2(unit_cell: array, x: array, y: array, vec_a: array, vec_b: array) -> array: """ Compute the 2D lattice DFT (LDFT) of a given pattern defined on a unit cell. Here the indices of the pattern are given by the pixel locations in the unit cell. The frequencies are also defined on a unit cell, but in this case generated by the unit vectors b1 * det(A) and b2 * det(A), where det(A) is the determinant of a matrix with rows or columns given by the periodicity vectors a1, a2. The unit cell and coordinates can be obtained from get_sim_unit_cell(). Alternatively, the skeleton of the unit cell can be obtained from get_unit_cell() and then the values on the unit cell may be chosen by the user. Use ldftfreq() to get the corresponding frequencies. :math: `LDFT[g](f_1, f_2) = \\sum_{nx, ny} \\exp[-2 \\pi i [f_1 b_1 (n_x, n_y) + f_2 b_2 (n_x, n_y)]] g(n_x, n_y) Instead of (nx, ny), one can also work with coefficients d1, d2 for the vectors a1 and a2. i.e. (nx, ny) = d1(nx, ny) * a1 + d2(nx, ny) * a2 :param unit_cell: array giving values of the unit cell at each coordinate, or NaN's if the position is not in the unit cell :param x: x-coordinates of the unit cell. Same shape as unit_cell :param y: y-coordinates of the unit cell. :param vec_a: first lattice vector :param vec_b: second lattice vector :return ldft: points which are not in the frequency unit cell are returned as NaNs """ # todo: what exactly is the relationship between this and get_pattern_fourier_component()? # todo: presumably this computes all Fourier components rather than just one # todo: what is relationship between this and get_efield_fourier_components()? # todo: should combine these two functions ... xp = cp if cp and isinstance(vec_a, cp.ndarray) else np fvecs, f1, f2 = ldftfreq(vec_a, vec_b) xx, yy = np.meshgrid(x, y) # todo: can I rewrite to do this with an fft on a larger rectangle? ldft = xp.nansum(unit_cell * xp.exp(-1j*2*np.pi * (fvecs[..., 0][..., None, None] * xx + fvecs[..., 1][..., None, None] * yy)), axis=(-1, -2)) ldft[xp.isnan(fvecs[..., 0])] = xp.nan + 1j * xp.nan return ldft def get_pattern_fourier_component(unit_cell: array, x: array, y: array, vec_a: array, vec_b: array, na: int, nb: int, nphases: int = 3, phase_index: int = 0, use_fft_origin: bool = True, dmd_size: Optional[Sequence[int]] = None) -> tuple[np.ndarray, np.ndarray]: """ Get fourier component at f = n * recp_vec_a + m * recp_vec_b. :math: `\\mathcal{F}(f) = \\sum_r f(r) * \\exp(-1j 2 \\pi f r)` :param np.array unit_cell: unit cell, as produced by get_sim_unit_cell() :param x: x-coordinates of unit cell :param y: y-coordinates of unit cell :param vec_a: first lattice vector :param vec_b: second lattice vector :param na: integer multiples of recp_vec_a :param nb: integer multiples of recp_vec_b :param nphases: only relevant for calculating phase :param phase_index: only relevant for calculating phase :param use_fft_origin: Specifies where the origin of the array is, which affects the phase. if not using the fft_origin, using the corner :param dmd_size: [nx, ny], only required if origin is "fft" :return fcomponent, frq_vector: fourier component of pattern at frq_vector = recp_vec_a * n + recp_vec_b * m """ # todo: vectorize in na/nb # todo: change this to accept frequency instead of index? Then can use this as helper function for ldft2() # todo: accept any pattern, not just SIM pattern xp = cp if cp and isinstance(unit_cell, cp.ndarray) else np x = xp.asarray(x) y = xp.asarray(y) vec_a = xp.asarray(vec_a) vec_b = xp.asarray(vec_b) recp_vect_a, recp_vect_b = get_reciprocal_vects(vec_a, vec_b) frq_vector = na * recp_vect_a + nb * recp_vect_b # fourier component is integral over unit cell xxs, yys = xp.meshgrid(x, y) fcomponent = xp.nansum(unit_cell * xp.exp(-1j*2*np.pi * (frq_vector[0] * xxs + frq_vector[1] * yys))) # correct phase for start coord of pattern "on" mirrors # todo: better to include this shift in definition fo unit_cell, and make # get_sim_unit_cell() accept a phase argument start_coord = xp.array(vec_b) / nphases * phase_index phase = xp.angle(fcomponent) - 2 * np.pi * start_coord.dot(frq_vector) if use_fft_origin: if dmd_size is None: raise ValueError("dmd_size was None, but must be specified when use_fft_origin is True") # now correct for DMD size # todo: relies on assumption that the cell zero coordinate is placed at DMD[0, 0] nx, ny = dmd_size x_pattern = np.arange(nx) - (nx // 2) y_pattern = np.arange(ny) - (ny // 2) # center coordinate in the edge coordinate system center_coord = xp.array([-x_pattern[0], -y_pattern[0]]) phase = phase + 2 * np.pi * center_coord.dot(frq_vector) fcomponent = xp.abs(fcomponent) * xp.exp(1j * phase) return fcomponent[0], frq_vector def get_efield_fourier_components(unit_cell: np.ndarray, x: np.ndarray, y: np.ndarray, vec_a: np.ndarray, vec_b: np.ndarray, nphases: int, phase_index: int, dmd_size: Sequence[int], nmax: int = 20, use_fft_origin: bool = True, ctf: Optional[np.ndarray] = None): """ Generate many Fourier components of pattern :param unit_cell: :param x: :param y: :param vec_a: :param vec_b: :param nphases: :param phase_index: :param dmd_size: :param nmax: :param use_fft_origin: :param ctf: coherent transfer function to apply :return efield, ns, ms, vecs: evaluated at the frequencyes vecs = ns * recp_va + ms * recp_vb """ warn("get_efield_fourier_components() is deprecated in favor of ldft2()") if ctf is None: def ctf(fx, fy): return 1 rva, rvb = get_reciprocal_vects(vec_a, vec_b) # first, get electric field fourier components ns = np.arange(-nmax, nmax + 1) ms = np.arange(-nmax, nmax + 1) ninds = 2 * nmax + 1 vecs = np.zeros((ninds, ninds, 2)) efield_fc = np.zeros((ninds, ninds), dtype=complex) # calculate half of values, as can get other half with E(-f) = E^*(f) for ii in range(nmax, len(ns)): for jj in range(len(ms)): # maximum pattern size is f = 0.5 1/mirrors, after this Fourier transform repeats information v = rva * ns[ii] + rvb * ms[jj] # if np.linalg.norm(v) > 1: # if np.linalg.norm(v) > 0.5: if np.abs(v[0]) > 0.5 or np.abs(v[1]) > 0.5: efield_fc[ii, jj] = 0 vecs[ii, jj] = v[:, 0] else: efield_fc[ii, jj], v = get_pattern_fourier_component(unit_cell, x, y, vec_a, vec_b, ns[ii], ms[jj], nphases, phase_index, use_fft_origin=use_fft_origin, dmd_size=dmd_size) vecs[ii, jj] = v[:, 0] # E(-f) = E^*(f) efield_fc[:nmax] = np.flip(efield_fc[nmax + 1:], axis=(0, 1)).conj() vecs[:nmax] = -np.flip(vecs[nmax + 1:], axis=(0, 1)) # apply OTF efield_fc = efield_fc * ctf(vecs[:, :, 0], vecs[:, :, 1]) # divide by volume of unit cell (i.e. maximum possible Fourier component) with np.errstate(invalid='ignore'): efield_fc = efield_fc / np.nansum(unit_cell >= 0) return efield_fc, ns, ms, vecs def get_int_fc(efield_ft: np.ndarray) -> np.ndarray: """ Generate intensity fourier components from efield fourier components :param efield_ft: electric field Fourier components nvec1 x nvec2 array, where efield_fc[ii, jj] is the electric field at frequencies f = ii * v1 + jj * v2. :return intensity_ft: intensity Fourier components at the same frequencies, f = ii * v1 + jj * v2 """ # todo: move to beam propagation tools? # todo: optionally include effect of polarization as in otf_tools.get_int_fc_polarized()? # I(f) = autocorrelation[E(f)] = convolution[E(f), E^*(-f)] intensity_ft = fftconvolve(efield_ft, conj_transpose_fft(efield_ft, axes=(0, 1)), mode='same') return intensity_ft # other fourier component function def get_intensity_fourier_components(unit_cell: np.ndarray, x: np.ndarray, y: np.ndarray, vec_a: np.ndarray, vec_b: np.ndarray, fmax: float, nphases: int, phase_index: int, dmd_size: Sequence[int, int], nmax: int = 20, use_fft_origin: bool = True, include_blaze_correction: bool = True, dmd_params: Optional[dict] = None) -> \ tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray, np.ndarray]: """ Utility function for computing many electric field and intensity components of the Fourier pattern, including the effect of the Blaze angle and system numerical aperture # todo: deprecate this in favor of get_int_fc() and get_efield_fourier_components() # todo: instead of setting nmax, just generate all e-field components that do not get blocked # todo: debating moving this function to simulate_dmd.py instead Given an electric field in fourier space E(k), the intensity I(k) = \\sum_q E(q) E^*(q-k). For a pattern where P(r)^2 = P(r), these must be equal, giving P(k) = \\sum_q P(q) P(q-k). But the relevant quantity after passing through the microscope is P(k) * bandlimit(k), where bandlimit(k) = 1 for k <= fmax, and 0 otherwise. Then the intensity pattern should be \\sum_q P(q) P(q-k) * bandlimit(q) * bandlimit(q-k) :param unit_cell: unit cell :param x: x-coordinates of unit cell :param y: y-coordinates of unit cell :param vec_a: :param vec_b: :param fmax: maximum pass frequency for electric field in 1/mirrors. i.e. fmax = NA/lambda without the factor of 2 that appears for the intensity. Note that fmax <= 1, which is the maximum frequency supported by the DMD. :param nphases: :param phase_index: :param dmd_size: [nx, ny] :param nmax: :param use_fft_origin: origin used to compute pattern phases "fft" or "" :param include_blaze_correction: if True, include blaze corrections :param dmd_params: dictionary {'wavelength', 'dx', 'dy', 'wx', 'wy', 'theta_ins': [tx_in, ty_in], 'theta_outs': [tx_out, ty_out]} :return intensity_fc: fourier components of intensity (band limited) :return efield_fc: fourier components of efield (band limited) :return ns: vec = ns * recp_vec_a + ms * recp_vec_b :return ms: vec = ns * recp_vec_a + ms * recp_vec_b :return vecs: ns * recp_vec_a + ms * recp_vec_b """ warn("get_intensity_fourier_components() is deprecated") if dmd_params is None and include_blaze_correction is True: raise ValueError("dmd_params must be supplied as include_blaze_correction is True") if dmd_params is not None: wavelength = dmd_params["wavelength"] gamma = dmd_params["gamma"] dx = dmd_params["dx"] dy = dmd_params["dy"] wx = dmd_params["wx"] wy = dmd_params["wy"] tin_x, tin_y = dmd_params['theta_ins'] tout_x, tout_y = dmd_params['theta_outs'] # get minimal lattice vectors # todo: use minimal lattice vectors to do the computation # va_m, vb_m = reduce_basis(vec_a, vec_b) # cell_m, x_m, y_m = convert_cell(unit_cell, x, y, vec_a, vec_b, va_m, vb_m) # todo: compute nmax # first, get electric field fourier components ns = np.arange(-nmax, nmax + 1) ms = np.arange(-nmax, nmax + 1) vecs = np.zeros((len(ns), len(ms), 2)) efield_fc = np.zeros((len(ns), len(ms)), dtype=complex) blaze_env = np.zeros(efield_fc.shape) # todo: calculating at f and -f is redundant for ii in range(len(ns)): for jj in range(len(ms)): efield_fc[ii, jj], v = get_pattern_fourier_component(unit_cell, x, y, vec_a, vec_b, ns[ii], ms[jj], nphases, phase_index, use_fft_origin=use_fft_origin, dmd_size=dmd_size) vecs[ii, jj] = v[:, 0] if include_blaze_correction: # wavelength * frq = theta in Fraunhofer approximation uvec_in = xy2uvector(tin_x, tin_y, True) uvec_out = xy2uvector(tout_x + wavelength * vecs[ii, jj][0] / dx, tout_y + wavelength * vecs[ii, jj][1] / dy, False) bma = uvec_out - uvec_in blaze_env[ii, jj] = blaze_envelope(wavelength, gamma, wx, wy, bma, _dlp_1stgen_axis) efield_fc[ii, jj] = efield_fc[ii, jj] * blaze_env[ii, jj] # divide by volume of unit cell (i.e. maximum possible Fourier component) with np.errstate(invalid='ignore'): efield_fc = efield_fc / np.nansum(unit_cell >= 0) # band limit frqs = np.linalg.norm(vecs, axis=2) # enforce maximum allowable frequency from DMD efield_fc = efield_fc * (frqs <= 0.5) # enforce maximum allowable frequency from imaging system efield_fc = efield_fc * (frqs <= fmax) # I(f) = autocorrelation[E(f)] = convolution[E(f), E^*(-f)] # note: the flip operation only for taking f-> -f only works assuming that array size is odd, with f=0 at the center intensity_fc = fftconvolve(efield_fc, np.flip(efield_fc, axis=(0, 1)).conj(), mode='same') # enforce maximum allowable frequency (should only be machine precision errors) intensity_fc = intensity_fc * (frqs <= 1) intensity_fc = intensity_fc * (frqs <= 2*fmax) return intensity_fc, efield_fc, ns, ms, vecs def get_intensity_fourier_components_xform(pattern: np.ndarray, affine_xform: np.ndarray, roi: Sequence[int, int, int, int], vec_a: np.ndarray, vec_b: np.ndarray, fmax: float, nmax: int = 20, cam_size: Sequence[int, int] = (2048, 2048), include_blaze_correction: bool = True, dmd_params: Optional[dict] = None): """ Utility function for computing many electric field and intensity components of the Fourier pattern, including the effect of the Blaze angle and system numerical aperture. To correct for ROI effects, extract from affine transformed pattern # todo: instead of setting nmax, just generate all e-field components that do not get blocked # todo: debating moving this function to simulate_dmd.py instead Given an electric field in fourier space E(k), the intensity I(k) =\\sum_q E(q) E^*(q-k). For a pattern where P(r)^2 = P(r), these must be equal, giving P(k) = \\sum_q P(q) P(q-k). But the relevant quantity after passing through the microscope is P(k) * bandlimit(k), where bandlimit(k) = 1 for k <= fmax, and 0 otherwise. Then the intensity pattern should be \\sum_q P(q) P(q-k) * bandlimit(q) * bandlimit(q-k) :param pattern: :param affine_xform: :param roi: :param vec_a: :param vec_b: :param fmax: :param nmax: :param cam_size: (ny, nx) :param include_blaze_correction: if True, include blaze corrections :param dmd_params: dictionary {'wavelength', 'dx', 'dy', 'wx', 'wy', 'theta_ins': [tx_in, ty_in], 'theta_outs': [tx_out, ty_out]} :return: """ if dmd_params is None and include_blaze_correction is True: raise ValueError("dmd_params must be supplied as include_blaze_correction is True") if dmd_params is not None: wavelength = dmd_params["wavelength"] gamma = dmd_params["gamma"] dx = dmd_params["dx"] dy = dmd_params["dy"] wx = dmd_params["wx"] wy = dmd_params["wy"] tin_x, tin_y = dmd_params['theta_ins'] tout_x, tout_y = dmd_params['theta_outs'] recp_va, recp_vb = get_reciprocal_vects(vec_a, vec_b) # todo: generate roi directly instead of cropping # todo: should be roi[0], not roi[1]? xform_roi = params2xform([1, 0, -roi[2], 1, 0, -roi[1]]).dot(affine_xform) nx_roi = roi[3] - roi[2] ny_roi = roi[1] - roi[0] xxi_roi, yyi_roi = np.meshgrid(range(nx_roi), range(ny_roi)) # todo: need to test this after swapping order of coordinates for xform_mat swap_xy = np.array([[0, 1, 0], [1, 0, 0], [0, 0, 1]]) xform_roi_yx = swap_xy.dot(xform_roi.dot(swap_xy)) pattern_xformed = xform_mat(pattern, xform_roi_yx, (yyi_roi, xxi_roi), mode="linear") pattern_xformed_ft = ft2(pattern_xformed) fxs = fftshift(fftfreq(pattern_xformed.shape[1], 1)) fys = fftshift(fftfreq(pattern_xformed.shape[0], 1)) # first, get electric field fourier components ns = np.arange(-nmax, nmax + 1) ms = np.arange(-nmax, nmax + 1) vecs = np.zeros((len(ns), len(ms), 2)) vecs_xformed = np.zeros(vecs.shape) efield_fc_xformed = np.zeros((len(ns), len(ms)), dtype=complex) blaze_env = np.zeros(efield_fc_xformed.shape) # todo: calculating @ freq and -freq is redundant for ii in range(len(ns)): for jj in range(len(ms)): vecs[ii, jj] = ns[ii] * recp_va[:, 0] + ms[jj] * recp_vb[:, 0] vecs_xformed[ii, jj, 0], vecs_xformed[ii, jj, 1], _ = xform_sinusoid_params(vecs[ii, jj, 0], vecs[ii, jj, 1], 0, affine_xform) try: efield_fc_xformed[ii, jj] = get_peak_value(pattern_xformed_ft, fxs, fys, vecs_xformed[ii, jj], peak_pixel_size=2) except: # todo: what exception is this supposed to catch? efield_fc_xformed[ii, jj] = 0 if include_blaze_correction: # wavelength * frq = theta in Fraunhofer approximation uvec_in = xy2uvector(tin_x, tin_y, True) uvec_out = xy2uvector(tout_x + wavelength * vecs[ii, jj][0] / dx, tout_y + wavelength * vecs[ii, jj][1] / dy, False) # amb = uvec_in - uvec_out bma = uvec_out - uvec_in blaze_env[ii, jj] = blaze_envelope(wavelength, gamma, wx, wy, bma, _dlp_1stgen_axis) efield_fc_xformed[ii, jj] = efield_fc_xformed[ii, jj] * blaze_env[ii, jj] # divide by DC component efield_fc_xformed = efield_fc_xformed / np.max(np.abs(efield_fc_xformed)) # hack to get to agree with nphases = 3 # todo: why is this here? efield_fc_xformed = efield_fc_xformed / 3 # band limit frqs = np.linalg.norm(vecs, axis=2) efield_fc_xformed = efield_fc_xformed * (frqs <= 0.5) efield_fc_xformed = efield_fc_xformed * (frqs <= fmax) # intensity fourier components from autocorrelation intensity_fc_xformed = fftconvolve(efield_fc_xformed, conj_transpose_fft(efield_fc_xformed, axes=(0, 1)), mode='same') intensity_fc_xformed = intensity_fc_xformed * (frqs <= 1) intensity_fc_xformed = intensity_fc_xformed * (frqs <= 2*fmax) return intensity_fc_xformed, efield_fc_xformed, ns, ms, vecs, vecs_xformed def show_fourier_components(vec_a: np.ndarray, vec_b: np.ndarray, fmax: float, int_fc: np.ndarray, efield_fc: np.ndarray, ns: np.ndarray, ms: np.ndarray, vecs: np.ndarray, plot_lims: Sequence[float, float] = (1e-4, 1), gamma: float = 0.1, figsize: Sequence[float, float] = (20., 10.), **kwargs) -> Figure: """ Display strength of fourier components for a given pattern. Display function for data generated with `get_bandlimited_fourier_components()`. See that function for more information about parameters. Additional keyword arguments are passed through to plt.figure() :param vec_a: :param vec_b: :param fmax: maximum frequency for electric field :param int_fc: :param efield_fc: :param ns: :param ms: :param vecs: :param plot_lims: limits in plots :param gamma: gamma to use in power law normalization of plots :param figsize: :return figh: handle to figure """ recp_va, recp_vb = get_reciprocal_vects(vec_a, vec_b) recp_va_reduced, recp_vb_reduced = reduce_recp_basis(vec_a, vec_b) # norm to use when plotting ft_norm = PowerNorm(vmin=plot_lims[0], vmax=plot_lims[1], gamma=gamma) # ################################ # plot results # ################################ figh = plt.figure(figsize=figsize, **kwargs) grid = figh.add_gridspec(2, 6, wspace=0.4) figh.suptitle(f"Pattern fourier weights versus position and reciprocal lattice vector\n" f" va=({vec_a[0]:d}, {vec_a[1]:d});" f" vb=({vec_b[0]:d}, {vec_b[1]})," f" max efield frq=1/{1/fmax:.2f} 1/mirrors") marker_size = 2 # ################################ # electric fields # ################################ # fourier components scatter plot ax = figh.add_subplot(grid[0, :2]) ax.set_facecolor((0., 0., 0.)) ax.axis('equal') im = ax.scatter(vecs[:, :, 0].ravel(), vecs[:, :, 1].ravel(), s=marker_size, c=np.abs(efield_fc).ravel(), norm=ft_norm) ax.scatter([recp_va[0], recp_vb[0]], [recp_va[1], recp_vb[1]], edgecolor='r', facecolor='none') ax.scatter([recp_va_reduced[0], recp_vb_reduced[0]], [recp_va_reduced[1], recp_vb_reduced[1]], edgecolor="m", facecolor="none") ax.add_artist(Circle((0, 0), radius=fmax, color='r', fill=0, ls='-')) ax.set_xlim([-fmax, fmax]) ax.set_ylim([-fmax, fmax]) cb = plt.colorbar(im) ax.set_xlabel('$f_x$ (1/mirror)') ax.set_ylabel('$f_y$ (1/mirror)') cb.set_label('|FT(f)|') ax.set_title('efield versus freq') # fourier components image ax = figh.add_subplot(grid[0, 2:4]) im = ax.imshow(np.abs(efield_fc), extent=[ns[0] - 0.5, ns[-1] + 0.5, ms[-1] + 0.5, ms[0] - 0.5], norm=ft_norm) ax.set_xlabel("$n_1 v_1$ ($n_1$)") ax.set_ylabel("$n_2 v_2$ ($n_2$)") cb = plt.colorbar(im) cb.set_label('|FT(f)|') ax.set_title('efield versus recp vect') # ################################ # intensity # ################################ ax = figh.add_subplot(grid[1, :2]) ax.set_facecolor((0., 0., 0.)) ax.axis('equal') im = ax.scatter(vecs[:, :, 0].ravel(), vecs[:, :, 1].ravel(), s=marker_size, c=np.abs(int_fc).ravel(), norm=ft_norm) ax.scatter([recp_va[0], recp_vb[0]], [recp_va[1], recp_vb[1]], edgecolor='r', facecolor='none') ax.scatter([recp_va_reduced[0], recp_vb_reduced[0]], [recp_va_reduced[1], recp_vb_reduced[1]], edgecolor="m", facecolor="none") ax.add_artist(Circle((0, 0), radius=(2*fmax), color='r', fill=0, ls='-')) ax.add_artist(Circle((0, 0), radius=fmax, color='r', fill=0, ls='-')) cb = plt.colorbar(im) ax.set_xlim([-2*fmax, 2*fmax]) ax.set_ylim([-2*fmax, 2*fmax]) ax.set_xlabel('$f_x$ (1/mirror)') ax.set_ylabel('$f_y$ (1/mirror)') cb.set_label('|FT(f)|') ax.set_title('intensity versus freq') # intensity image ax = figh.add_subplot(grid[1, 2:4]) im = ax.imshow(np.abs(int_fc), extent=[ns[0] - 0.5, ns[-1] + 0.5, ms[-1] + 0.5, ms[0] - 0.5], norm=ft_norm) cb = plt.colorbar(im) ax.set_xlabel('recp vec as') ax.set_ylabel('recp vec bs') cb.set_label('|FT(f)|') ax.set_title('intensity versus recp vect') # ################################ # 1D plots # ################################ # efield and intensity 1D ax = figh.add_subplot(grid[0, 4:]) ax.set_title("|FT| vs frq") ax.set_xlabel("$f$ (1/mirrors)") # only plot one of +/- f, and only plot if above certain threshold vec_mag = np.linalg.norm(vecs, axis=-1) nmax1 = _convert_int(0.5 * (int_fc.shape[0] - 1)) nmax2 = _convert_int(0.5 * (int_fc.shape[1] - 1)) to_use = np.ones(int_fc.shape, dtype=int) xx, yy = np.meshgrid(range(-nmax2, nmax2 + 1), range(-nmax1, nmax1 + 1)) to_use[xx > yy] = 0 to_use[np.logical_and(xx == yy, yy < 0)] = 0 to_plot_int = np.logical_and(to_use, np.abs(int_fc) >= plot_lims[0]) to_plot_e = np.logical_and(to_plot_int, vec_mag <= fmax) ylim = [plot_lims[0], plot_lims[1] * 1.2] ax.plot([fmax, fmax], ylim, 'k') ax.plot([2*fmax, 2*fmax], ylim, 'k') ax.plot(vec_mag[to_plot_int], np.abs(int_fc[to_plot_int]), '.', label="I") ax.plot(vec_mag[to_plot_e], np.abs(efield_fc[to_plot_e]), 'x', label="E") ax.set_yscale("log") ax.set_ylim(ylim) ax.set_xlim([-0.1 * 2 * fmax, 2.2 * fmax]) ax.legend() # E/I phases 1D ax = figh.add_subplot(grid[1, 4:]) ax.set_title("Fourier component phase vs frq") ax.set_xlabel("Frequency (1/mirrors)") ylim = [-np.pi - 0.2, np.pi + 0.2] ax.plot([fmax, fmax], ylim, 'k') ax.plot([2 * fmax, 2 * fmax], ylim, 'k') ax.plot(vec_mag[to_plot_int], np.angle(int_fc[to_plot_int]), '.', label="I") ax.plot(vec_mag[to_plot_e], np.angle(efield_fc[to_plot_e]), 'x', label="E") ax.set_ylim(ylim) ax.set_xlim([-0.1 * 2 * fmax, 2.2 * fmax]) ax.legend() return figh # Lagrange-Gauss basis reduction def reduce_basis(va: array, vb: array) -> tuple[array, array]: """ Find the "smallest" set of basis vectors using Lagrange-Gauss basis reduction. :param va: first lattice vector :param vb: second lattice vector :return va, vb: """ xp = cp if cp and isinstance(va, cp.ndarray) else np va = xp.array(va, copy=True) va = va.reshape([2, ]) vb = xp.array(vb, copy=True) vb = vb.reshape([2, ]) Ba = xp.linalg.norm(va)**2 mu = xp.vdot(va, vb) / Ba vb = vb - xp.round(mu) * va Bb = xp.linalg.norm(vb)**2 swapped = -1 while Bb < Ba: va, vb = vb, va swapped *= -1 Ba = Bb mu = xp.inner(va, vb) / Ba vb = vb - xp.round(mu) * va Bb = xp.linalg.norm(vb) ** 2 if swapped == 1: va, vb = vb, va # ensure integers va_int = _convert_int(va) vb_int = _convert_int(vb) return va_int, vb_int def reduce_recp_basis(va: array, vb: array) -> tuple[array, array]: """ Compute the shortest pair of reciprocal basis vectors. These vectors may not be dual to the lattice vectors in the sense that vi * rsj = delta_{ij}, but they do form a basis for the reciprocal lattice vectors. :param va: first lattice vector :param vb: second lattice vector :return rsa, rsb: reduced reciprocal vectors """ va, vb = reduce_basis(va, vb) return get_reciprocal_vects(va, vb) def get_closest_lattice_vec(point: np.ndarray, va: np.ndarray, vb: np.ndarray) -> tuple[np.ndarray, int, int]: """ Find the closest lattice vector to point :param point: :param va: first lattice vector :param vb: second lattice vector :return vec, na_min, nb_min: """ point = np.array(point, copy=True) point = point.reshape([2, ]) # get reduced lattice basis vectors var, vbr = reduce_basis(va, vb) # get reduced reciprocal vectors rva, rvb = get_reciprocal_vects(var, vbr) frac_a = np.vdot(point, rva) nas = [int(np.ceil(frac_a)), int(np.floor(frac_a))] frac_b = np.vdot(point, rvb) nbs = [int(np.ceil(frac_b)), int(np.floor(frac_b))] # possible choices diff = np.inf for na in nas: for nb in nbs: v_diff = point - na * var - nb * vbr diff_current = np.linalg.norm(v_diff) if diff_current < diff: nar_min = na nbr_min = nb diff = diff_current vec = na*var + nb*vbr # convert back to initial basis lattice vectors # get reciprocal vectors ra, rb = get_reciprocal_vects(va, vb) # and how they are related to initial lattice vectors var_ints = np.array([np.vdot(var, ra), np.vdot(var, rb)]) vbr_ints = np.array([np.vdot(vbr, ra), np.vdot(vbr, rb)]) na_min = _convert_int(nar_min * var_ints[0] + nbr_min * vbr_ints[0]) nb_min = _convert_int(nar_min * var_ints[1] + nbr_min * vbr_ints[1]) return vec, na_min, nb_min def get_closest_recip_vec(recp_point: np.ndarray, va: np.ndarray, vb: np.ndarray) -> tuple[np.ndarray, int, int]: """ Find the closest reciprocal lattive vector, f = na * rva + nb * rvb, to a given point in reciprocal space, recp_point. :param recp_point: :param va: first lattice vector :param vb: second lattice vector :return vec, na_min, nb_min: vec = na_min * rva + nb_min * rvb """ recp_point = np.array(recp_point, copy=True) recp_point = recp_point.reshape([2, ]) va = np.array(va, copy=True) va = va.reshape([2, ]) vb = np.array(vb, copy=True) vb = vb.reshape([2, ]) det = va[0] * vb[1] - va[1] * vb[0] rva, rvb = get_reciprocal_vects(va, vb) # use get_closest_lattice_vec() function after scaling rva, rvb to have integer components vec, na_min, nb_min = get_closest_lattice_vec(recp_point * det, rva * det, rvb * det) vec = vec / det return vec, na_min, nb_min # working with grayscale patterns def binarize(pattern_gray: array, mode: str = "floyd-steinberg") -> array: """ Binarize a gray scale pattern :param pattern_gray: gray scale pattern, with values in the range [0, 1] :param mode: "floyd-steinberg" to specify the Floyd-Steinberg error diffusion algorithm, "jjn" to use the error diffusion algorithm of Jarvis, Judis, and Ninke https:doi.org/10.1016/S0146-664X(76)80003-2, "random" to use a random dither, or "round" to round to the nearest value :return pattern_binary: binary approximation of pattern_gray """ xp = cp if cp and isinstance(pattern_gray, cp.ndarray) else np pattern_gray = xp.array(pattern_gray, copy=True) if xp.any(pattern_gray) > 1 or xp.any(pattern_gray) < 0: raise ValueError("pattern values must be in [0, 1]") ny, nx = pattern_gray.shape if mode == "floyd-steinberg": # error diffusion Kernel = # 1/16 * [[_ # 7], [3, 5, 1]] pattern_bin = xp.zeros(pattern_gray.shape, dtype=bool) for ii in range(ny): for jj in range(nx): pattern_bin[ii, jj] = xp.round(pattern_gray[ii, jj]) err = pattern_gray[ii, jj] - pattern_bin[ii, jj] if jj < (nx - 1): pattern_gray[ii, jj+1] += err * 7/16 if ii < (ny - 1): if jj > 0: pattern_gray[ii + 1, jj - 1] += err * 3/16 pattern_gray[ii + 1, jj] += err * 5/16 if jj < (ny - 1): pattern_gray[ii + 1, jj + 1] += err * 1/16 elif mode == "jjn": # error diffusion Kernel = # 1/48 * [[_, _, #, 7, 5], [3, 5, 7, 5, 3], [1, 3, 5, 3, 1]] pattern_bin = xp.zeros(pattern_gray.shape, dtype=bool) for ii in range(ny): for jj in range(nx): pattern_bin[ii, jj] = xp.round(pattern_gray[ii, jj]) err = pattern_gray[ii, jj] - pattern_bin[ii, jj] if jj < (nx - 1): pattern_gray[ii, jj + 1] += err * 7/48 if jj < (nx - 2): pattern_gray[ii, jj + 2] += err * 5/48 if ii < (ny - 1): if jj > 1: pattern_gray[ii + 1, jj - 2] += err * 3/48 if jj > 0: pattern_gray[ii + 1, jj - 1] += err * 5/48 pattern_gray[ii + 1, jj] += err * 7/48 if jj < (ny - 1): pattern_gray[ii + 1, jj + 1] += err * 5/48 if jj < (ny - 2): pattern_gray[ii + 1, jj + 2] += err * 3/48 if ii < (ny - 2): if jj > 1: pattern_gray[ii + 2, jj - 2] += err * 1/48 if jj > 0: pattern_gray[ii + 2, jj - 1] += err * 3/48 pattern_gray[ii + 2, jj] += err * 5/48 if jj < (ny - 1): pattern_gray[ii + 2, jj + 1] += err * 3/48 if jj < (ny - 2): pattern_gray[ii + 2, jj + 2] += err * 1/48 elif mode == "random": pattern_bin = xp.random.binomial(1, pattern_gray).astype(bool) elif mode == "round": pattern_bin = xp.round(pattern_gray).astype(bool) else: raise ValueError(f"mode must be 'floyd-steinberg', 'random', or 'round' but was '{mode:s}'") return pattern_bin # utility functions def min_angle_diff(angle1: array, angle2: array, mode: str = 'normal') -> array: """ Find minimum magnitude of angular difference between two angles. :param angle1: in radians :param angle2: in radians :param mode: "normal" or "half" :return angle_diff: """ xp = cp if cp and isinstance(angle1, cp.ndarray) else np angle2 = xp.asarray(angle2) # take difference modulo 2pi, which gives positive distance angle_diff = xp.asarray(xp.mod(angle1 - angle2, 2*np.pi)) # still want smallest magnitude difference (negative or positive). If larger than pi, can express as smaller # magnitude negative distance ind_greater_pi = angle_diff > np.pi angle_diff[ind_greater_pi] = angle_diff[ind_greater_pi] - 2 * np.pi if mode == 'normal': pass elif mode == 'half': # compute differences allowing theta and theta + pi to be equivalent angle_diff_pi = min_angle_diff(angle1, angle2 + np.pi, mode='normal') to_switch = xp.abs(angle_diff_pi) < xp.abs(angle_diff) angle_diff[to_switch] = angle_diff_pi[to_switch] else: raise ValueError(f"'mode' must be 'normal' or 'half', but was '{mode:s}'") return angle_diff # generate single pattern def find_closest_pattern(period: float, angle: float, nphases: int = 1, avec_max_size: int = 40, bvec_max_size: int = 40): """ Find pattern vectors for pattern with an approximate period and angle that also satisfies the perfect phase shift condition :param period: pattern period in pixels :param angle: pattern angle :param nphases: :param avec_max_size: maximum size of first lattice vector :param bvec_max_size: maximum size of second lattice vector :return avec, bvec, period_real, angle_real: """ angles_proposed, bvecs_proposed = find_allowed_angles(period, nphases, bvec_max_size, restrict_to_coordinate_axes=False) ia = np.argmin(np.abs(angle - angles_proposed)) a = angles_proposed[ia] bvec = bvecs_proposed[ia] # approximate a-vector x, y, seq = find_rational_approx_angle(a, avec_max_size) avec = np.array([x, y]) period_real = get_sim_period(avec, bvec) angle_real = get_sim_angle(avec, bvec) return avec, bvec, period_real, angle_real # tools for finding nearest SIM pattern set def find_closest_multicolor_set(period: float, nangles: int, nphases: int, wavelengths: Optional[Sequence[float]] = None, bvec_max_size: int = 40, avec_max_size: int = 40, atol: float = np.pi/180, ptol_relative: float = 0.1, angle_sep_tol: float = 5*np.pi/180, max_solutions_to_search: int = 20, pitch: float = 7560., minimize_leakage: bool = True) -> tuple[np.ndarray, np.ndarray]: """ Generate set of SIM patterns for multiple colors with period close to specified value and maximizing distance between angles. The patterns are determined such that the diffracted orders will pass through the same positions in the Fourier plane of the imaging sytem. i.e. the fractional resolution increase in SIM should be the same for all the colors. :param period: pattern period in pixels. If using multiple colors, specify this for the shortest wavelength :param nangles: number of angles :param nphases: number of phases :param wavelengths: list of wavelengths in consistent units. If set to None, then will assume only one wavelength. :param bvec_max_size: maximum allowed size of second lattice vector, in mirrors :param avec_max_size: maximum allowed size of first lattice vector, in mirrors :param atol: maximum allowed deviation between angles for different colors. :param ptol_relative: maximum tolerance for period deviations, as a fraction of the period :param angle_sep_tol: maximum deviation between adjacent pattern angles from the desired value which would lead to equally spaced patterns. :param max_solutions_to_search: maximum number of angle combinations to search for furthest distance to leakage peaks :param pitch: DMD micromirror spacing in the same units as wavelength :param minimize_leakage: whether to do leakage minimization :return vec_as, vec_bs: """ # todo: still problems with even number phase shifts if wavelengths is None: wavelengths = [1] # factor to multiply the period by for each wavelength factors = np.sort(wavelengths / np.min(wavelengths)) periods = period * factors # get allowed angles in range [0, pi] for all wavelengths angles_all = [] bvs_all = [] for p in periods: a, b = find_allowed_angles(p, nphases, bvec_max_size, restrict_to_coordinate_axes=False) angles_all.append(a) bvs_all.append(b) # only keep angles that are very similar # todo: could in principle keep increasing bvec_max_size until have enough angles to work with angles_kept = [[] for _ in wavelengths] bvs_kept = [[] for _ in wavelengths] # todo: could check difference accounting for e.g. 0, pi being same. Probably these edge cases not very important. # todo: could think about dynamically changing size of max_bvecs and max_avecs until have an appropriate size. for a in angles_all[0]: keep = True for angs in angles_all: if np.min(np.abs(a - angs)) > atol: keep = False break if keep: for ii, angs in enumerate(angles_all): ind = np.argmin(np.abs(a - angs)) angles_kept[ii].append(angs[ind]) bvs_kept[ii].append(bvs_all[ii][ind]) # do typical minimization using one set of these angles. i.e. "one-color" minimization # todo: Now want to do minimization over bvec size, so want to include all colors. angles = np.asarray(angles_kept[0]) expected_angle_sep = np.pi / nangles min_sep = expected_angle_sep - angle_sep_tol max_sep = expected_angle_sep + angle_sep_tol angle_inds = np.asarray(range(len(angles))) # for each angle, find the allowed successor angles successor_inds = [angle_inds[np.logical_and(angles > a + min_sep, angles < a + max_sep)] for a in angles] # list of lists, with each sublist giving possible set of angles # can grow these sets by looking only at the last angle and its possible successors sets_inds = [[ind] for ind in angle_inds] for ii in range(1, nangles): sets_inds_new = [] for set_current in sets_inds: for successor_ind in successor_inds[set_current[-1]]: sets_inds_new.append(set_current + [successor_ind]) sets_inds = sets_inds_new # arrays of angles and indices sets_inds = np.array(sets_inds) angle_sets = angles[sets_inds] # get rid of anywhere separation between n-1 and 0th is too large too_big = np.abs(min_angle_diff(angle_sets[:, 0], angle_sets[:, nangles - 1], mode='half') - expected_angle_sep) > angle_sep_tol # cost on bvector norms bvs_norms = np.array([np.linalg.norm(bv) for bv in bvs_kept[0]]) cost = np.sum(bvs_norms[sets_inds] / nphases, axis=1) cost[too_big] = np.nan # sort choices by cost # isort = np.flip(np.argsort(cost.ravel())) isort = np.argsort(cost.ravel()) sets_inds_sort = sets_inds[isort] csort = cost.ravel()[isort] # isort = isort[np.logical_not(np.isnan(csort))] sets_inds_sort = sets_inds_sort[np.logical_not(np.isnan(csort))] # csort = csort[np.logical_not(np.isnan(csort))] if not minimize_leakage: # take closest solution sopt = sets_inds_sort[0] vec_bs = [[bvs_wvl[s] for s in sopt] for bvs_wvl in bvs_kept] angles_opt = [np.asarray([angs_wvl[s] for s in sopt]) for angs_wvl in angles_kept] vec_as = [[find_rational_approx_angle(a, avec_max_size)[-1][-1] for a in a_wavlen] for a_wavlen in angles_opt] min_leakage_angle = np.nan else: # loop over so many possible solutions and check which has most leeway wrt leakage orders # todo: should ensure that all the angle sets looped over are close enough to the optimum min_leakage_angle = 0 for sopt in sets_inds_sort[:max_solutions_to_search]: # list of lists. List holds lists of and b-vectors for each wavelength angles_opt = [np.asarray([angs_wvl[s] for s in sopt]) for angs_wvl in angles_kept] vec_bs_proposed = [[bvs_wvl[s] for s in sopt] for bvs_wvl in bvs_kept] # find vec_as satisfying approximate angle vec_as_proposed = [[] for _ in wavelengths] min_leakage_dist_wvlen = [[] for _ in wavelengths] for ii in range(len(wavelengths)): vec_as_accepted = [[] for _ in angles_opt[ii]] for jj, (a, vb) in enumerate(zip(angles_opt[ii], vec_bs_proposed[ii])): xsh, ysh, vec_a_seq = find_rational_approx_angle(a, avec_max_size) vec_as_accepted[jj] = [va for va in vec_a_seq if (va[0] * vb[1] - va[1] * vb[0]) != 0 and min_angle_diff(get_sim_angle(va, vb), a, mode='half') < atol and np.abs((get_sim_period(va, vb) - periods[ii]) / periods[ii]) < ptol_relative] # ####################################### # find set of vec_as with maximum distance to nearest leakage orders # ####################################### vava = np.meshgrid(*[range(len(v)) for v in vec_as_accepted], indexing='ij') min_dists = np.zeros(vava[0].shape) for kk in range(vava[0].size): ind_prop = np.unravel_index(kk, vava[0].shape) vec_as_curr = [vec_as_accepted[ll][vava[ll][ind_prop]] for ll in range(nangles)] # min_dists[ind_prop], _, _ = find_nearest_leakage_peaks(vec_as_curr, vec_bs_proposed[ii], nphases) min_dists[ind_prop], _, _ = find_nearest_leakage_peaks(vec_as_curr, vec_bs_proposed[ii], nphases, wavelength=wavelengths[ii], pitch=pitch) ind_min = np.argmax(min_dists) sub_min = np.unravel_index(ind_min, min_dists.shape) # multiply by wavelength factor to account for the fact the scale of the Fourier plane # changes with wavelength # min_leakage_dist_wvlen[ii] = factors[ii] * min_dists[sub_min] min_leakage_dist_wvlen[ii] = min_dists[sub_min] vec_as_proposed[ii] = [vec_as_accepted[ll][vava[ll][sub_min]] for ll in range(nangles)] # accept new set of angles if closest leakage order is further than what we already have proposed_min_leakage_dist = np.min(min_leakage_dist_wvlen) if proposed_min_leakage_dist > min_leakage_angle: min_leakage_angle = proposed_min_leakage_dist vec_as = vec_as_proposed vec_bs = vec_bs_proposed return np.array(vec_as), np.array(vec_bs) def find_allowed_angles(period: float, nphases: int, nmax: int, restrict_to_coordinate_axes: bool = False): """ Given a DMD pattern with fixed period of absolute value p, get allowed pattern angles in the range [0, pi] for which the pattern allows perfect phase shifting for nphases. :math: `p = dxb \\cos \\theta + dyb \\sin \\theta` For theta in [0, pi] we can take x=cos(theta), and sin(theta) = sqrt(1-x^2). We get a quadratic equation in x, x^2 * (dxb**2/dyb**2 + 1) - x * (2*P*dxb/dyb**2) + (P**2/dxb**2 - 1) = 0 :param period: period in pixels :param nphases: number of phase shifts :param nmax: :param restrict_to_coordinate_axes: deprecated...used to allow running old behavior when adding functionality :return angles, vbs: """ # allowed vector components if restrict_to_coordinate_axes: ns = np.arange(nphases, nmax, nphases) dxb = np.concatenate((ns, np.zeros(ns.shape))) dyb = np.concatenate((np.zeros(ns.shape), ns)) else: # with two vector components, can no longer restrict all to be positive dxs = np.arange(nphases, nmax, nphases, dtype=float) dxs = np.concatenate((np.flip(-dxs), np.array([0]), dxs), axis=0) dys = np.arange(0, nmax, nphases, dtype=float) dxb, dyb = np.meshgrid(dxs, dys) # exclude vb = [0, 0] dxb, dyb = dxb[dxb**2 + dyb**2 > 0], dyb[dxb**2 + dyb**2 > 0] # (1) P = dxb * cos(theta) + dyb * sin(theta) # (2) P = dxb * x + dyb * sqrt(1-x**2) # squaring both sides gives # (3) (P - dxb * x)**2 = dyb**2 * (1 - x**2) # A*x**2 + B*x + C = 0 # two solutions, expect one in [0, pi/2] and one in [pi/2, pi]. # BUT it is possible these are not both solutions to the original equation. This can happen if the portion in # paranetheses on the LHS of (3) is negative A = dxb**2 + dyb**2 B = - 2 * period * dxb C = period**2 - dyb**2 with np.errstate(invalid='ignore'): # get solutions to the squared problem x1 = 0.5 * (-B + np.sqrt(B**2 - 4 * A * C)) / A # only keep ones that also satisfy the base problem x1[np.abs(dxb * x1 + dyb * np.sqrt(1 - x1**2) - period) > 1e-7] = np.nan x2 = 0.5 * (-B - np.sqrt(B**2 - 4 * A * C)) / A x2[np.abs(dxb * x2 + dyb * np.sqrt(1 - x2 ** 2) - period) > 1e-7] = np.nan # also negative period solutions. Only change here is B -> -B x3 = 0.5 * (B + np.sqrt(B**2 - 4 * A * C)) / A x3[np.abs(dxb * x3 + dyb * np.sqrt(1 - x3 ** 2) + period) > 1e-7] = np.nan x4 = 0.5 * (B - np.sqrt(B**2 - 4 * A * C)) / A x4[np.abs(dxb * x4 + dyb * np.sqrt(1 - x4 ** 2) + period) > 1e-7] = np.nan # get final angles and vectors angles = np.concatenate((np.arccos(x1).ravel(), np.arccos(x2).ravel(), np.arccos(x3).ravel(), np.arccos(x4).ravel())) # exclude nans vbs = 4 * [[int(dx), int(dy)] for dx, dy in zip(dxb.ravel(), dyb.ravel())] # test # dxbt, dybt = zip(*vbs) # ps = np.cos(angles) * np.asarray(dxbt) + np.sin(angles) * np.asarray(dybt) vbs = [v for v, a in zip(vbs, angles) if not np.isnan(a)] angles = angles[np.logical_not(np.isnan(angles))] # sort lists by size of angles isort = np.argsort(angles) angles = angles[isort] vbs = [vbs[ii] for ii in isort] return angles, vbs def find_rational_approx_angle(angle: float, nmax: int): """ Find the closest allowed vector compatible with the pixel array :param angle: desired angle in radians :param nmax: maximum size of the x- and y-components of the lattice vector in pixels. :return xshift, yshift, vecs: """ # todo: how to simplify these cases # first convert angle to [0, pi/2], so can do rational approximation for positive fraction angle_2p = np.mod(angle, 2*np.pi) if angle_2p <= np.pi/2: angle_pos = angle_2p case = 1 elif np.pi/2 < angle_2p <= np.pi: angle_pos = np.pi - angle_2p case = 2 elif np.pi < angle_2p <= 3*np.pi/2: angle_pos = angle_2p - np.pi case = 3 elif 3*np.pi/2 < angle_2p <= 2*np.pi: angle_pos = 2*np.pi - angle_2p case = 4 else: raise ValueError('disallowed angle') slope = np.tan(angle_pos) slope_inverted = False if slope > 1: slope = 1 / slope slope_inverted = True # use Farey sequence and binary search. See e.g. # https://www.johndcook.com/blog/2010/10/20/best-rational-approximation/ fr_lb = [0, 1] fr_ub = [1, 1] approximate_seq = [] while True: mediant_num = fr_lb[0] + fr_ub[0] mediant_denom = fr_lb[1] + fr_ub[1] if mediant_denom >= nmax: break mediant = mediant_num / mediant_denom if mediant == slope: fr_ub = [mediant_num, mediant_denom] fr_lb = [mediant_num, mediant_denom] approximate_seq.append(fr_ub) if mediant > slope: fr_ub = [mediant_num, mediant_denom] # compare new bound to last best estimate if approximate_seq == []: approximate_seq.append(fr_ub) else: current_est = mediant_num/mediant_denom best_est = approximate_seq[-1][0] / approximate_seq[-1][1] if np.abs(current_est - slope) <= np.abs(best_est - slope): approximate_seq.append(fr_ub) else: fr_lb = [mediant_num, mediant_denom] # compare new bound to last best estimate if approximate_seq == []: approximate_seq.append(fr_lb) else: current_est = mediant_num / mediant_denom best_est = approximate_seq[-1][0] / approximate_seq[-1][1] if np.abs(current_est - slope) <= np.abs(best_est - slope): approximate_seq.append(fr_lb) if slope_inverted: slope = 1 / slope approximate_seq = [np.flip(s) for s in approximate_seq] # todo: don't really understand why each sign needs to be so. Thought I had it figured out, but found I had to # change all except case 1 # tan(theta) = (-dxa) / dya if case == 1: vecs = [[-s[0], s[1]] for s in approximate_seq] elif case == 2: vecs = [[-s[0], -s[1]] for s in approximate_seq] elif case == 3: vecs = [[s[0], -s[1]] for s in approximate_seq] elif case == 4: vecs = [[s[0], s[1]] for s in approximate_seq] xshift = vecs[-1][0] yshift = vecs[-1][1] return xshift, yshift, vecs def find_allowed_periods(angle: float, nphases: int, nmax: int) -> tuple[array, array, array]: """ Given a DMD pattern with fixed angle, get allowed pattern periods which allow perfect phase shifting for nphases :param angle: :param nphases: :param nmax: :return periods, ls, is_xlike: """ # Recall that for vec_a = [dxa, dya] and vec_b = [dxb, 0], and dxb = l * nphases for perfect phase shifting # period = dxb * dya/|vec_a| = dxb * cos(theta) # theta = angle(vec_a_perp) = arctan(-dxa / dya) # P = np.cos(theta) * l*nphases # on the other hand, if vec_b = [0, dyb] # period = dyb * -dxa/|vec_a = dyb * sin(theta) # P = np.sin(theta) * l*nphases ls = np.arange(1, int(np.floor(nmax / nphases))) p1 = np.cos(angle) * ls * nphases p2 = np.sin(angle) * ls * nphases # store data about angles is_xlike = np.concatenate((np.ones(p1.size), np.zeros(p2.size)), axis=0) ls_all = np.concatenate((ls, ls), axis=0) periods = np.concatenate((p1, p2), axis=0) # sort lists by size of angles combined_list = list(zip(periods, ls_all, is_xlike)) combined_list.sort(key=lambda v: v[0]) periods, ls_all, is_xlike = zip(*combined_list) return np.asarray(periods), np.asarray(ls_all), np.asarray(is_xlike) def find_nearest_leakage_peaks(vec_as: np.ndarray, vec_bs: np.ndarray, nphases: int = 3, minimum_relative_peak_size: float = 1e-3, wavelength: float = 1., pitch: float = 7560.): """ Find minimum distance between main pattern frequency and leakage frequencies from other patterns in the set :param vec_as: list of first lattice vectors :param vec_bs: list of second lattice vectors :param nphases: :param minimum_relative_peak_size: peaks smaller than this size (compared with the maximum peak, i.e. the DC peak) will not be included. :param wavelength: can be provided so that distance will be appropriately scaled for different wavelengths :param pitch: :return min_angle_all, min_angle_leakage_peaks, leakage_order_pattern_index: """ # find frequencies nangles = len(vec_as) # frqs = [get_sim_frqs(va, vb) for va, vb in zip(vec_as, vec_bs)] cells, xs, ys = zip(*[get_sim_unit_cell(va, vb, nphases) for va, vb in zip(vec_as, vec_bs)]) xxs, yys = zip(*[np.meshgrid(x, y) for x, y in zip(xs, ys)]) recp_vects = [get_reciprocal_vects(va, vb) for va, vb in zip(vec_as, vec_bs)] min_dists = np.ones((nangles, nangles)) * np.inf for ii in range(nangles): for jj in range(nangles): if ii == jj: # allow nearby combinations except for ns, ms = np.meshgrid([-1, 0, 1], [-1, 0, 1, 2]) ns = ns.ravel() ms = ms.ravel() ns, ms = ns[np.logical_and(ns != 0, ms != 1)], ms[np.logical_and(ns != 0, ms != 1)] else: # want to find combinations of reciprocal vectors of one pattern that are closest to # those of another pattern # n * r1 + m * r2 ~ recp_vec_b, with n,m integers # first, solve for n, m real numbers mat = np.linalg.inv(np.concatenate((recp_vects[jj][0], recp_vects[jj][1]), axis=1)) n, m = mat.dot(recp_vects[ii][1]) # now check distances for nearby reciprocal vectors. Expand our search in case some of the nearby peaks # have very little weight ns, ms = np.meshgrid(list(range(int(np.floor(n)) - 2, int(np.ceil(n)) + 3)), list(range(int(np.floor(m)) - 2, int(np.ceil(m)) + 3))) ns = ns.ravel() ms = ms.ravel() for n, m in zip(ns, ms): vec = n * recp_vects[jj][0] + m * recp_vects[jj][1] # peak weight is the Fourier transform over the unit cell (divided by the DC component) weight = np.abs(np.nansum(cells[jj] * np.exp(1j * 2 * np.pi * (vec[0] * xxs[jj] + vec[1] * yys[jj]))) / np.nansum(cells[jj])) # if weight is too small, don't count distance if weight < minimum_relative_peak_size: continue # if weight is not too small, then set distance if it is smaller than what we already have dist = np.linalg.norm(vec - recp_vects[ii][1]) * wavelength / pitch if dist < min_dists[ii, jj]: min_dists[ii, jj] = dist # minimum distance for each pattern min_angle_leakage_peaks = np.nanmin(min_dists, axis=1) leakage_order_pattern_index = np.nanargmin(min_dists, axis=1) # minimum distance over patterns min_angle_all = np.min(min_angle_leakage_peaks) return min_angle_all, min_angle_leakage_peaks, leakage_order_pattern_index # functions for obtaining and exporting results def vects2pattern_data(dmd_size: Sequence[int, int], vec_as: np.ndarray, vec_bs: np.ndarray, nphases: int = 3, wavelength: Optional[float] = None, invert: bool = False, pitch: float = 7560, generate_patterns: bool = True, geometry: str = "orthogonal") -> tuple[np.ndarray, dict]: """ Generate pattern and useful data (angles, phases, frequencies, reciprocal vectors, ...) from the lattice vectors for a given pattern set. :param dmd_size: [nx, ny] :param vec_as: NumPy array, size nangles x nphases x 2 :param vec_bs: :param nphases: :param wavelength: wavelength in nm :param invert: whether pattern is "inverted", i.e. if the roll of "OFF" and "ON" should be flipped :param pitch: pixel pitch :param generate_patterns: :param geometry: see get_sim_pattern() :return patterns, data: """ vec_as = np.array(vec_as, copy=True) vec_bs = np.array(vec_bs, copy=True) # extract dmd size nx, ny = dmd_size if wavelength is None: wavelength = 1 nangles, _ = vec_as.shape patterns = np.zeros((nangles, nphases, ny, nx)) phases = np.zeros((nangles, nphases)) angles = np.zeros(nangles) periods = np.zeros(nangles) frqs = np.zeros((nangles, 2)) recp_vects_a = np.zeros((nangles, 2)) recp_vects_b = np.zeros((nangles, 2)) # loop over wavelengths min_leakage_angle, _, _ = find_nearest_leakage_peaks(vec_as, vec_bs, nphases, minimum_relative_peak_size=1e-3, wavelength=wavelength, pitch=pitch) # loop over angles and find the closest available patterns for ii in range(nangles): ra, rb = get_reciprocal_vects(vec_as[ii], vec_bs[ii]) recp_vects_a[ii] = ra.ravel() recp_vects_b[ii] = rb.ravel() periods[ii] = get_sim_period(vec_as[ii], vec_bs[ii]) angles[ii] = get_sim_angle(vec_as[ii], vec_bs[ii]) frqs[ii] = get_sim_frqs(vec_as[ii], vec_bs[ii]) for jj in range(nphases): phases[ii, jj] = get_sim_phase(vec_as[ii], vec_bs[ii], nphases, jj, pattern_size=dmd_size, use_fft_origin=True) if generate_patterns: patterns[ii, jj], c = get_sim_pattern([nx, ny], vec_as[ii], vec_bs[ii], nphases, jj, geometry=geometry) if invert: patterns = 1 - patterns data = {'vec_as': vec_as, 'vec_bs': vec_bs, 'frqs': frqs, 'angles': angles, 'periods': periods, 'phases': phases, 'nx': int(dmd_size[0]), 'ny': int(dmd_size[1]), 'recp_vects_a': recp_vects_a, 'recp_vects_b': recp_vects_b, 'min_leakage_angle': float(min_leakage_angle), 'dmd_pitch': float(pitch), 'wavelength': float(wavelength)} return patterns, data def plot_sim_pattern_sets(patterns: np.ndarray, vas: np.ndarray, vbs: np.ndarray, wavelength: Optional[float] = None, pitch: float = 7560., figsize: Sequence[float, float] = (16., 12.), **kwargs) -> Figure: """ Plot all angles/phases in pattern set, as well as their Fourier transforms. Additional keyword arguments will be passed through to plt.figure() :param patterns: :param vas: :param vbs: :param wavelength: :param pitch: :param figsize: :return figh: handle to resulting figure """ nangles, nphases, ny, nx = patterns.shape _, data = vects2pattern_data([nx, ny], vas, vbs, nphases=nphases, wavelength=wavelength, generate_patterns=False, pitch=pitch) periods = data["periods"] phases = data["phases"] angles = data["angles"] frqs = data["frqs"] min_leakage_angle = data["min_leakage_angle"] # display summary of patterns figh = plt.figure(figsize=figsize, **kwargs) grid = figh.add_gridspec(nrows=nphases + 1, ncols=nangles) if wavelength is not None: figh.suptitle(f"sim pattern diagnostic, wavelength = {wavelength:.0f}nm," f" min leakage angle={min_leakage_angle * 180/np.pi:.3f}deg") else: figh.suptitle(f"sim pattern diagnostic, min leakage angle = {min_leakage_angle:.3f}") # ############################### # real space patterns # ############################### for ii in range(nangles): for jj in range(nphases): ax = figh.add_subplot(grid[jj, ii]) # cut_size = int(np.max(np.abs(vas)) * np.ceil(period / np.max(np.abs(vbs)))) cut_size = int(np.max([np.max(np.abs(vas)), np.max(np.abs(vbs))])) ax.imshow(patterns[ii, jj, :cut_size, :cut_size], cmap="bone") ax.set_ylabel(f"phase={phases[ii, jj] * 180 / np.pi:.2f}deg") if jj == 0: ax.set_title(f"angle={angles[ii] * 180 / np.pi:.2f}deg," f" p={periods[ii]:.2f}\n" f"a = [{vas[ii, 0]:d}, {vas[ii, 1]:d}]," f" b=[{vbs[ii, 0]:d}, {vbs[ii, 1]:d}]") else: ax.set_xticklabels([]) ax.set_yticklabels([]) # ############################### # Fourier transforms # ############################### fx = fftshift(fftfreq(nx, 1)) dfx = fx[1] - fx[0] fy = fftshift(fftfreq(ny, 1)) dfy = fy[1] - fy[0] df_min = np.min([fx[1] - fx[0], fy[1] - fy[0]]) extent = [fx[0] - 0.5 * dfx, fx[-1] + 0.5 * dfx, fy[-1] + 0.5 * dfy, fy[0] - 0.5 * dfy] for ii in range(nangles): ax = figh.add_subplot(grid[nphases, ii]) # 2D window from broadcasting apodization = np.expand_dims(hann(nx), axis=0) * np.expand_dims(hann(ny), axis=1) ft = ft2(patterns[ii, 0] * apodization) ax.imshow(np.abs(ft) / np.abs(ft).max(), norm=PowerNorm(gamma=0.1), extent=extent, cmap="bone") # dominant frequencies of underlying patterns for rr in range(nangles): if rr == ii: color = 'r' else: color = 'm' ax.add_artist(Circle((frqs[rr, 0], frqs[rr, 1]), radius=5 * df_min, color=color, fill=0, ls='-')) ax.add_artist(Circle((-frqs[rr, 0], -frqs[rr, 1]), radius=5 * df_min, color=color, fill=0, ls='-')) ax.set_ylabel('ft') return figh def export_pattern_set(dmd_size: Sequence[int, int], vec_as: np.ndarray, vec_bs: np.ndarray, nphases: int = 3, invert: bool = False, pitch: float = 7560., wavelength: float = 1., save_dir: Union[str, Path] = 'sim_patterns', plot_results: bool = False, geometry: str = "orthogonal") -> tuple[np.ndarray, dict, Figure]: """ Export a single set of SIM patterns, i.e. single wavelength, single period :param dmd_size: [nx, ny] :param np.array vec_as: nangles x nphases x 2 :param np.array vec_bs: :param nphases: :param invert: :param pitch: :param wavelength: :param save_dir: :param plot_results: :param geometry: see get_sim_pattern() :return patterns, data, figh: """ save_dir = Path(save_dir) save_dir.mkdir(exist_ok=True) tstamp = datetime.datetime.now().strftime('%Y_%m_%d_%H;%M;%S') patterns, data = vects2pattern_data(dmd_size, vec_as, vec_bs, nphases=nphases, wavelength=None, invert=invert, pitch=pitch, geometry=geometry) periods = data["periods"] angles = data["angles"] phases = data["phases"] data["tstamp"] = tstamp nangles, _, ny, nx = patterns.shape for k, v in data.items(): if isinstance(v, np.ndarray): data[k] = v.tolist() fpath = save_dir / (f"sim_patterns_period={np.mean(periods):.2f}_" f"nangles={nangles:d}.json") with open(fpath, 'w') as f: json.dump(data, f, indent="\t") # save patterns to separate PNG files for ii in range(nangles): for jj in range(nphases): ind = ii * nphases + jj # save file # need to convert so not float to save as PNG im = Image.fromarray(patterns[ii, jj].astype('bool')) im.save(save_dir / f"{ind:02d}_period={periods[ii]:.2f}_" f"angle={angles[ii] * 180/np.pi:.1f}deg_" f"phase={phases[ii, jj]:.2f}.png") # save patterns in tif stack fpath = save_dir / (f"sim_patterns_period={np.mean(periods):.2f}_" f"nangles={nangles:d}_" f"nphases={nphases:d}.tif") tifffile.imwrite(fpath, tifffile.transpose_axes(patterns.astype(np.uint8).reshape((nangles * nphases, ny, nx)), "CYX", asaxes="TZQCYXS"), imagej=True) if plot_results: figh = plot_sim_pattern_sets(patterns, vec_as, vec_bs, wavelength) figh.savefig(save_dir / f"period={np.mean(periods):.2f}_pattern_summary.png") else: figh = None return patterns, data, figh # main function for generating SIM patterns at several frequencies and wavelengths def export_all_pattern_sets(dmd_size: Sequence[int, int], periods: Sequence[float], nangles: int = 3, nphases: int = 3, wavelengths: Optional[Sequence[float]] = None, invert: Sequence[bool] = False, pitch: float = 7560., save_dir: Union[str, Path] = 'sim_patterns', plot_results: bool = True, geometry: str = "orthogonal", **kwargs): """ Generate SIM pattern sets and save results. Additional keyword arguments will be passed through to find_closest_multicolor_set(). Use them to set the angle/period tolerances and search range for that function. :param dmd_size: [nx, ny] :param periods: list of approximate periods :param nangles: number of angles :param nphases: number of phases :param wavelengths: list of wavelengths in nanometers. If set to None, will assume only one wavelength. :param invert: :param pitch: :param save_dir: directory to save results :param plot_results: :param geometry: see get_sim_pattern() :return data_all: [[dict_period0_wlen0, dict_period0, wlen1, ...], [dict_period1, wlen0, ...]] list of lists of dictionary objects. First level sublists are data for different periods, second sublevel is data for different wavelengths. """ if wavelengths is None: wavelengths = [1] if not isinstance(wavelengths, list): wavelengths = [wavelengths] nwavelengths = len(wavelengths) if not isinstance(periods, list): periods = [periods] save_dir = Path(save_dir) save_dir.mkdir(exist_ok=True) data_all = [] # generate sets for period in periods: data_period = [] # directory to save results sub_dir = f"period={period:.1f}_nangles={nangles:d}" pattern_save_dir = Path(save_dir, sub_dir) if not pattern_save_dir.exists(): pattern_save_dir.mkdir() vec_as, vec_bs = find_closest_multicolor_set(period, nangles, nphases, wavelengths=wavelengths, **kwargs) # loop over wavelengths for kk in range(nwavelengths): if nwavelengths == 1: wavlen_savedir = pattern_save_dir else: wavlen_savedir = pattern_save_dir / f"wavelength={wavelengths[kk]:d}nm" patterns, data, figh = export_pattern_set(dmd_size, vec_as[kk], vec_bs[kk], nphases=nphases, wavelength=wavelengths[kk], invert=invert[kk], save_dir=wavlen_savedir, pitch=pitch, plot_results=plot_results, geometry=geometry) data_period.append(data) data_all.append(data_period) return data_all # export calibration patterns def aberration_map_pattern(dmd_size: Sequence[int, int], vec_a: np.ndarray, vec_b: np.ndarray, nphases: int, centers, radius: int = 20, phase_indices: int = 0): """ Generate patterns to calibrate DMD aberrations using the approach of https://doi.org/10.1364/OE.24.013881 Each pattern contains two small patches of lattice. If we measure the interference of the beams diffracted from the two patches, we can extract the surface profile of the DMD. :param dmd_size: (nx, ny) :param vec_a: :param vec_b: :param nphases: number of phase shifts allowed :param centers: :param radius: radius, must be an integer :param phase_indices: :return: """ if not isinstance(radius, int): raise ValueError("radius must be an integer") centers = np.array(centers, dtype=int) if centers.ndim == 1: centers = np.expand_dims(centers, axis=0) phase_indices = np.atleast_1d(phase_indices) if len(phase_indices) == 1 and centers.shape[0] > 1: phase_indices = np.ones(centers.shape[0]) * phase_indices[0] # get patches pattern_patches = [] for ii in range(nphases): pattern_patch, _ = get_sim_pattern([2 * radius + 1, 2 * radius + 1], vec_a, vec_b, nphases, ii) pattern_patches.append(pattern_patch) pattern_patches = np.asarray(pattern_patches) xx, yy = np.meshgrid(range(pattern_patches.shape[2]), range(pattern_patches.shape[1])) xx = xx - xx.mean() yy = yy - yy.mean() pattern_patches[:, np.sqrt(xx**2 + yy**2) > radius] = 0 # get pattern nx, ny = dmd_size pattern = np.zeros((ny, nx)) for ii in range(len(centers)): pattern[centers[ii, 1] - radius: centers[ii, 1] + radius + 1, centers[ii, 0] - radius: centers[ii, 0] + radius + 1] = pattern_patches[phase_indices[ii]] return pattern def checkerboard(dmd_size: Sequence[int, int], n_on: int, n_off: Optional[int] = None): """ Create checkerboard pattern :param dmd_size: [nx, ny] :param n_on: :param n_off: :return np.array pattern: """ # default is use same number of off and on pixels if n_off is None: n_off = n_on nx, ny = dmd_size n_cell = n_on + n_off cell = np.zeros((n_cell, n_cell)) cell[:n_on, :n_on] = 1 nx_tiles = int(np.ceil(nx / cell.shape[1])) ny_tiles = int(np.ceil(ny / cell.shape[0])) mask = np.tile(cell, [ny_tiles, nx_tiles]) mask = mask[0:ny, 0:nx] return mask def export_calibration_patterns(dmd_size: Sequence[int, int], save_dir: Union[str, Path] = '', circle_radii: Sequence[int] = (1, 2, 3, 4, 5, 10, 25, 50, 100, 200, 300)): """ Produce calibration patterns for the DMD, which are all on, all off, center-circles of several sizes, and checkerboard patterns of several sizes :param dmd_size: [nx, ny] :param save_dir: :param circle_radii: :return: """ save_dir = Path(save_dir) save_dir.mkdir(exist_ok=True) nx = dmd_size[0] ny = dmd_size[1] # all mirrors on all_on = np.ones((ny, nx)) im = Image.fromarray(all_on.astype('bool')) im.save(save_dir / "on.png") # all mirror off all_off = np.zeros((ny, nx)) im = Image.fromarray(all_off.astype('bool')) im.save(save_dir / "off.png") # circles of different radii centered in the middle of DMD xx, yy = np.meshgrid(range(nx), range(ny)) xc = (nx - 1) / 2 yc = (ny - 1) / 2 rr = np.sqrt((xx - xc)**2 + (yy - yc)**2) for r in circle_radii: mask = np.zeros((ny, nx)) mask[rr <= r] = 1 im = Image.fromarray(mask.astype('bool')) im.save(save_dir / f"circle_on_r={r:d}.png") im = Image.fromarray((1 - mask).astype('bool')) im.save(save_dir / f"circle_off_r={r:d}.png") # checkerboard patterns with different spacing periods = np.concatenate((np.arange(2, 12, 1), np.arange(12, 30, 2), np.arange(30, 200, 10))) for p in periods: on_pix = int(np.ceil(p / 2)) mask = checkerboard(dmd_size, on_pix) im = Image.fromarray(mask.astype('bool')) im.save(save_dir / f"checkerboard_period={p:d}.png") # patterns with variable spacing periods = range(2, 20, 2) for ii, p in enumerate(periods): cell = np.zeros((p, nx)) on_pix = int(np.ceil(p / 2)) cell[:on_pix, :] = 1 cell = np.tile(cell, [4, 1]) if ii == 0: mask = cell else: mask = np.concatenate((mask, cell), axis=0) mask = mask[:, :mask.shape[0]] mask_block = np.concatenate((mask, np.rot90(mask)), axis=1) mask_block2 = np.concatenate((np.rot90(mask), mask), axis=1) mask_superblock = np.concatenate((mask_block, mask_block2)) ny_reps = int(np.ceil(ny / mask_superblock.shape[0])) nx_reps = int(np.ceil(nx / mask_superblock.shape[1])) mask = np.tile(mask_superblock, [ny_reps, nx_reps]) mask = mask[0:ny, 0:nx] im = Image.fromarray(mask.astype('bool')) im.save(save_dir / f"variable_pattern_periods={periods[0]:d}_to_{periods[-1]:d}.png") # pattern with three corners corner_size = 300 corner_pattern = np.zeros((ny, nx), dtype=bool) corner_pattern[:corner_size, :corner_size] = 1 corner_pattern[:corner_size, -corner_size:] = 1 corner_pattern[-corner_size:, :corner_size] = 1 im = Image.fromarray(corner_pattern) im.save(f"three_corners_{corner_size:d}.png") def export_otf_test_set(dmd_size: Sequence[int], pmin: float = 4.5, pmax: float = 50, nperiods: int = 20, nangles: int = 12, nphases: int = 3, avec_max_size: int = 40, bvec_max_size: int = 40, phase_index: int = 0, save_dir: Optional[Union[str, Path]] = None, verbose: bool = True) -> tuple[np.ndarray, dict]: """ Generate many SIM-like patterns at a variety of different periods and angles :param dmd_size: [nx, ny] :param pmin: minimum period (mirrors) :param pmax: maximum period (mirrors) :param nperiods: number of periods :param nangles: number of angles, equally space in [0, pi) :param nphases: 1/the filling fraction of the lattice :param avec_max_size: constrain maximum size of lattice vectors. Larger lattice vectors typically allow more accurate approximation of patterns with a given period and angle, but introduce additional stray diffraction orders at other frequencies. :param bvec_max_size: :param phase_index: :param save_dir: If provided, save patterns as PNG and Tiff files and pattern data as json :param verbose: print information as patterns are generated :return patterns, data: """ if save_dir is not None: save_dir = Path(save_dir) save_dir.mkdir(parents=True, exist_ok=False) nx, ny = dmd_size # equally spaced values in frequency space frqs = np.linspace(1 / pmax, 1 / pmin, nperiods) periods = np.flip(1/frqs) # angles equally distributed in [0, pi) angles = np.arange(nangles) * np.pi / nangles # arrays to store pattern information patterns = np.zeros((nperiods, nangles, ny, nx), dtype=bool) real_angles = np.zeros((nperiods, nangles)) real_frqs = np.zeros((nperiods, nangles, 2)) real_periods = np.zeros((nperiods, nangles)) real_phases = np.zeros((nperiods, nangles)) vec_as = np.zeros((nperiods, nangles, 2), dtype=int) vec_bs = np.zeros_like(vec_as) # find nearest patterns for each request angle and period tstart = perf_counter() for ii, p in enumerate(periods): for jj, a in enumerate(angles): tstart_pattern = perf_counter() vec_as[ii, jj], vec_bs[ii, jj], real_periods[ii, jj], real_angles[ii, jj] = \ find_closest_pattern(p, a, nphases=nphases, avec_max_size=avec_max_size, bvec_max_size=bvec_max_size) patterns[ii, jj], _ = get_sim_pattern(dmd_size, vec_as[ii, jj], vec_bs[ii, jj], nphases, phase_index) real_phases[ii, jj] = get_sim_phase(vec_as[ii, jj], vec_bs[ii, jj], nphases, phase_index, dmd_size, use_fft_origin=True) if verbose: tnow = perf_counter() print(f"generated pattern {ii * len(angles) + jj + 1:d}/{len(periods) * len(angles):d}" f" in {tnow - tstart_pattern:.2f}s," f"elapsed time {tnow - tstart:.2f}s", end="\r") if verbose: print("") pattern_on = np.ones((ny, nx), dtype=np.uint8) pattern_off = np.zeros((ny, nx), dtype=np.uint8) # convert arrays to lists to make json compatible data = {'tstamp': datetime.datetime.now().strftime('%Y_%m_%d_%H;%M;%S'), 'vec_as': vec_as.tolist(), 'vec_bs': vec_bs.tolist(), 'angles': real_angles.tolist(), 'periods': real_periods.tolist(), 'frequencies': real_frqs.tolist(), 'phases': real_phases.tolist(), 'nphases': nphases, 'phase_index': phase_index, 'units': 'um', 'notes': 'total number of patterns should be nphases*nangles + 2.' ' The last two patterns are all ON and all OFF respectively.'} # export results if save_dir is not None: # save pattern info with open(save_dir / "pattern_data.json", "w") as f: json.dump(data, f, indent="\t") # save patterns as set of pngs for ii in range(nperiods): for jj in range(nangles): ind = ii * nangles + jj # need to convert so not float to save as PNG im = Image.fromarray(patterns[ii, jj].astype('bool')) im.save(save_dir / f"{ind:03d}_" f"pattern_period={real_periods[ii, jj]:.3f}_" f"angle={real_angles[ii, jj] * 180/np.pi:.2f}deg.png") # save all on im = Image.fromarray(pattern_on.astype('bool')) im.save(save_dir / f"{nperiods * nangles:03d}_pattern_all_on.png") # save all off im = Image.fromarray(pattern_off.astype('bool')) im.save(save_dir / f"{nperiods * nangles + 1:03d}_pattern_all_off.png") # save patterns as tif patterns_reshaped = np.reshape(patterns, [patterns.shape[0] * patterns.shape[1], patterns.shape[2], patterns.shape[3]]) patterns_reshaped = np.concatenate((patterns_reshaped, np.expand_dims(pattern_on, axis=0), np.expand_dims(pattern_off, axis=0)), axis=0) tifffile.imwrite(save_dir / "patterns.tif", patterns_reshaped) return patterns, data def _verify_int(m: Union[array, float]) -> bool: """ :param m: :return: """ xp = cp if cp and isinstance(m, cp.ndarray) else np m = xp.asarray(m) mint = xp.round(m).astype(int) return xp.all(xp.abs(m - mint) < 1e-10) def _convert_int(m: array): if not _verify_int(m): raise ValueError("array could not be converted to integers") xp = cp if cp and isinstance(m, cp.ndarray) else np return xp.round(xp.asarray(m)).astype(int)