from functools import lru_cache
import math
import numpy as np
import tensorflow as tf
from . import _lib
from ._backend import BACKEND
from .slm import SLM, DMD
from typing import Tuple
__all__ = ["SLMSimulation", "DMDSimulation"]
[docs]
class SLMSimulation(object):
"""The class implements a Simulation plan for the SLM.
Parameters
----------
slm: pySLM2.slm.SLM
padding_x: int
Number of pixels padded in the x direction. (see the note for details)
padding_y: int
Number of pixels padded in the y direction. (see the note for details)
.. Note::
The padded Fourier plane is defined as following::
+----------------------+ ^
| | | padding_y * pixel_size
| | |
| +------------+ | v
| | | |
| | DMD | |
| | | |
| +------------+ |
| |
| padded FP |
+----------------------+
<---->
padding_x * pixel_size
"""
def __init__(self, slm: SLM, padding_x: int=0, padding_y: int=0):
if not isinstance(padding_x, int):
raise TypeError("padding_x must be an integer.")
if not isinstance(padding_y, int):
raise TypeError("padding_y must be an integer.")
if isinstance(slm, SLM):
self._slm = slm
else:
raise TypeError("dmd must be a SLM object.")
self._padding_x = padding_x
self._padding_y = padding_y
self._input_field = None
self._output_field = None
self._image_plane_field = None
[docs]
def clear(self):
"""Clear the simulation result."""
self._input_field = None
self._output_field = None
self._image_plane_field = None
@property
def fourier_plane_pixel_area(self) -> float:
"""float"""
return self._slm._pixel_size ** 2
@property
def image_plane_pixel_area(self) -> float:
return self._slm.scaling_factor ** 2 / (self.Nx * self.Ny) / self._slm.pixel_size ** 2
@property
def padding_x(self) -> int:
return self._padding_x
@property
def padding_y(self) -> int:
return self._padding_y
@property
def Nx(self) -> int:
return self._slm.Nx + 2 * self._padding_x
@property
def Ny(self) -> int:
return self._slm.Ny + 2 * self._padding_y
@lru_cache()
def _image_plane_padded_grid(self):
kx_atom, ky_atom = tf.constant(np.fft.fftfreq(self.Nx, self._slm.pixel_size),
dtype=BACKEND.dtype), \
tf.constant(-np.fft.fftfreq(self.Ny, self._slm.pixel_size),
dtype=BACKEND.dtype)
kx_atom, ky_atom = tf.signal.fftshift(kx_atom), tf.signal.fftshift(ky_atom)
x_atom = kx_atom * self._slm.scaling_factor
y_atom = ky_atom * self._slm.scaling_factor
return tf.meshgrid(x_atom, y_atom)
@property
def image_plane_padded_grid(self) -> Tuple[np.ndarray, np.ndarray]:
x, y = self._image_plane_padded_grid()
return np.array(x), np.array(y)
@lru_cache()
def _fourier_plane_padded_grid(self):
pix_j = tf.range(-self._padding_x, self._slm.Nx + self._padding_x, dtype=BACKEND.dtype)
pix_i = tf.range(-self._padding_y, self._slm.Ny + self._padding_y, dtype=BACKEND.dtype)
pix_ii, pix_jj = tf.meshgrid(pix_i, pix_j, indexing="ij")
return self._slm._convert_pixel_index_to_dmd_coordinate(pix_ii, pix_jj)
@property
def fourier_plane_padded_grid(self) -> Tuple[np.ndarray, np.ndarray]:
x, y = self._fourier_plane_padded_grid()
return np.array(x), np.array(y)
[docs]
def propagate_to_image(self, input_profile):
"""
Parameters
----------
input_profile
"""
input_profile = self._slm.profile_to_tensor(input_profile, complex=True)
self._input_field = tf.pad(input_profile,
[[self._padding_y, self._padding_y], [self._padding_x, self._padding_x]])
output_profile = input_profile * self._slm._state_tensor()
self._output_field = tf.pad(output_profile,
[[self._padding_y, self._padding_y], [self._padding_x, self._padding_x]])
_image_plane_field_unormalized = tf.signal.fftshift(
_lib._inverse_fourier_transform(tf.signal.ifftshift(self._output_field)))
# TODO proper explanation
self._image_plane_field = _image_plane_field_unormalized * (self.Nx * self.Ny * self.fourier_plane_pixel_area / self.scaling_factor)
@tf.function
def _field_to_intensity(self, field_tensor):
return tf.math.real(field_tensor)**2 + tf.math.imag(field_tensor)**2
@property
def _input_intensity(self):
return None if self._input_field is None else self._field_to_intensity(self._input_field)
@property
def _output_intensity(self):
return None if self._output_field is None else self._field_to_intensity(self._output_field)
@property
def _image_plane_intensity(self):
return None if self._image_plane_field is None else self._field_to_intensity(self._image_plane_field)
def _pack_tensor_to_array(self, tensor):
return None if tensor is None else np.array(tensor)
@property
def input_field(self) -> np.ndarray:
return self._pack_tensor_to_array(self._input_field)
@property
def output_field(self) -> np.ndarray:
return self._pack_tensor_to_array(self._output_field)
@property
def image_plane_field(self) -> np.ndarray:
"np.ndarray"
return self._pack_tensor_to_array(self._image_plane_field)
@property
def input_intensity(self) -> np.ndarray:
"np.ndarray"
return self._pack_tensor_to_array(self._input_intensity)
@property
def output_intensity(self) -> np.ndarray:
"np.ndarray"
return self._pack_tensor_to_array(self._output_intensity)
@property
def image_plane_intensity(self) -> np.ndarray:
"np.ndarray"
return self._pack_tensor_to_array(self._image_plane_intensity)
[docs]
def get_output_power(self) -> float:
"""Calculates and returns the total power output from the SLM.
Returns
-------
power: float
The total power output from the SLM.
"""
return np.sum(self.output_intensity) * self.fourier_plane_pixel_area
[docs]
def get_image_plane_power(self) -> float:
"""Calculates and returns the total power at the images plane.
Returns
-------
power: float
The total power at images plane.
"""
return np.sum(self.image_plane_intensity) * self.image_plane_pixel_area
@property
def scaling_factor(self) -> float:
return self._slm.scaling_factor
[docs]
class DMDSimulation(SLMSimulation):
"""The class implements a Simulation plan for the DMD.
Parameters
----------
dmd: pySLM2.slm.DMD
padding_x: int
Number of pixels padded in the x direction. (see pySLM2.simulation.SLMSimulation for details)
padding_y: int
Number of pixels padded in the y direction. (see pySLM2.simulation.SLMSimulation for details)
"""
def __init__(self, dmd, padding_x=0, padding_y=0):
if not isinstance(dmd, DMD):
raise TypeError("dmd must be a DMD object.")
super(DMDSimulation, self).__init__(dmd, padding_x=padding_x, padding_y=padding_y)
@property
def first_order_origin(self):
assert isinstance(self._slm, DMD)
return self._slm.first_order_origin
[docs]
def block_zeroth_order(self, r=None):
"""On the DMDSimulation.image_plane array, set the centre values (0th order diffraction) to zero."""
assert isinstance(self._slm, DMD)
if self._output_field is None:
raise TypeError("Run propagate_to_image first to initialise images plane light field")
if r is None:
r = self.scaling_factor / self._slm._p.value() / 2
x, y = self._image_plane_padded_grid()
r2 = x ** 2 + y ** 2
mask = r2 < r ** 2
self._image_plane_field = tf.where(mask, tf.zeros_like(self._image_plane_field, dtype=BACKEND.dtype_complex),
self._image_plane_field)