hologradpy.torch_functions ========================== .. py:module:: hologradpy.torch_functions .. autoapi-nested-parse:: Module containing PyTorch-specific functions to perform conjugate gradient minimisation. Classes ------- .. autoapisummary:: hologradpy.torch_functions.ASM hologradpy.torch_functions.VirtualSlm hologradpy.torch_functions.PhaseRetrieval Functions --------- .. autoapisummary:: hologradpy.torch_functions.check_device hologradpy.torch_functions.gpu_to_numpy hologradpy.torch_functions.fft hologradpy.torch_functions.ifft hologradpy.torch_functions.asm hologradpy.torch_functions.rms hologradpy.torch_functions.eff hologradpy.torch_functions.loss_fn_fid hologradpy.torch_functions.loss_fn_amp hologradpy.torch_functions.camera_calibration hologradpy.torch_functions.camera_feedback Module Contents --------------- .. py:function:: check_device(verbose=None) Check if GPU is available. :param bool verbose: Verbose output? :return: 'cuda' if GPU available, otherwise 'cpu'. .. py:function:: gpu_to_numpy(gpu_tensor) .. py:function:: fft(e_in, shift=True, norm=None) Performs the FFT. :param e_in: Input electric field. :param bool shift: Perform FFT shift? :param norm: Normalisation of FFT. :return: FFT of ``e_in``. .. py:function:: ifft(e_in, shift=True, norm=None) Performs the IFFT. :param e_in: Input electric field. :param bool shift: Perform IFFT shift? :param norm: Normalisation of IFFT. :return: IFFT of ``e_in``. .. py:function:: asm(e_in, e_lens, theta1, theta2=None, shift=True) Performs the angular spectrum method (ASM) twice, from the SLM to the Fourier lens and from the Fourier lens to the camera. :param e_in: Electric field at the SLM. :param e_lens: Electric field of the lens (phase and aperture). :param theta1: Propagation phase from SLM to lens. :param theta2: Propagation phase from lens to camera. If not provided it is assumed ``theta2 = theta1``. :param bool shift: Perform FFT shift? :return: Electric field at the camera. .. py:class:: ASM(slm_disp_obj, pms_obj, pix_res, npix_tot, pd1, pd2, extent_lens, xf, shift=False, precision='single', device='cuda') This class models the propagation of light from the SLM to the Fourier lens and from the Fourier lens to the image plane using the angular spectrum method. The lens is modelled as a doublet. The ASM wavefront correction is also calculated in this class. .. py:attribute:: theta1 .. py:attribute:: theta2 .. py:attribute:: shift :value: False .. py:attribute:: e_lens .. py:attribute:: phi_corr_native .. py:attribute:: phi_corr .. py:method:: forward(e_in) This function performs the simulation. :param e_in: Electric field at the SLM plane. :return: Electric field at the image plane. .. py:class:: VirtualSlm(slm_disp_obj, pms_obj, phi, npix_pad, npix=None, e_slm=None, kernel_ct=None, pix_res=None, propagation_type='fft', extent_lens=None, pd1=None, pd2=None, xf=None, device='cpu', slm_mask=None, precision=None, fft_shift=True) Bases: :py:obj:`torch.nn.Module` This class models pixel crosstalk on the SLM and the propagation of light from the SLM to the camera. .. py:attribute:: slm_disp_obj .. py:attribute:: pms_obj .. py:attribute:: precision :value: None .. py:attribute:: dtype_r .. py:attribute:: dtype_c .. py:attribute:: fft_shift :value: True .. py:attribute:: device :value: 'cpu' .. py:attribute:: npix_pad .. py:attribute:: pix_res :value: None .. py:attribute:: npix_full .. py:attribute:: propagation_type :value: 'fft' .. py:attribute:: npix :value: None .. py:attribute:: e_slm .. py:attribute:: slm_mask .. py:attribute:: pad .. py:attribute:: counter :value: 0 .. py:attribute:: phi_disp .. py:method:: set_phi(new_phi) Set SLM phase from numpy array. :param ndarray new_phi: SLM phase [rad]. .. py:method:: forward() Model the SLM and simulate the propagation of light from the SLM plane to the image plane. This method is used by gradient-based optimizers. :return: Electric field in the image plane. .. py:function:: rms(signal, i_target, i_out, frac) Calculate normalised root-mean-squared error between two images inside a region of interest. Only pixels which are brighter than ``frac * np.max(i_target_norm)`` are taken into account, where ``i_target_norm`` is the normalised target intensity pattern. :param signal: Binary mask containing region of interest (signal region). :param i_target: Target intensity pattern. :param i_out: Intensity pattern of light potential. :param frac: Threshold as explained above. :return: Normalised rms error. .. py:function:: eff(signal, i_out) Calculates the predicted efficiency of a light potential by dividing the pixel sum in the signal region by the pixel sum in the entire pattern. :param signal: Binary mask containing the signal region. :param i_out: Intensity pattern of the light potential. :return: Predicted efficiency. .. py:function:: loss_fn_fid(e_out, i_tar, phi_tar, signal) Phase and amplitude cost function from https://doi.org/10.1364/OE.25.011692. :param e_out: Electric field at the image plane. :param i_tar: Target intensity pattern. :param phi_tar: Target phase pattern. :param signal: Binary mask containing signal region. :return: Cost. .. py:function:: loss_fn_amp(e_out, i_tar, signal) Amplitude-only cost function from https://doi.org/10.1364/OE.22.026548. :param e_out: Electric field at the image plane. :param i_tar: Target intensity pattern. :param signal: Binary mask containing signal region. :return: Cost. .. py:class:: PhaseRetrieval(slm_obj, n_iter=10, i_tar=None, phi_tar=None, signal_region=None, save=False, n_save=10, loss_type='amp', optim_type='cg') This function calculates the SLM phase pattern for a given target light potential in the image plane using conjugate gradient minimisation or stochastic gradient descent (Adam). .. py:attribute:: slm_obj .. py:attribute:: signal :value: None .. py:attribute:: i_tar :value: None .. py:attribute:: fft_shift .. py:attribute:: phi_tar :value: None .. py:attribute:: phi_tar_t .. py:attribute:: signal_t .. py:attribute:: i_tar_t .. py:attribute:: save :value: False .. py:attribute:: n_save :value: 10 .. py:attribute:: n_iter :value: 10 .. py:attribute:: loss :value: 0 .. py:attribute:: closure_counter :value: 0 .. py:attribute:: callback_counter :value: 0 .. py:attribute:: phi :value: [] .. py:attribute:: eta_pred :value: [] .. py:attribute:: eff_pred :value: [] .. py:attribute:: loss_type :value: 'amp' .. py:attribute:: optim_type :value: 'cg' .. py:attribute:: optimizer :value: None .. py:method:: set_target(target) Sets the target light potential. :param target: Target light potential. .. py:method:: set_optimizer() Sets the optimisation algorithm based on ``self.optim_type``. .. py:method:: loss_fn(e_out) Defines the loss function based on ``self.loss_type``. :param e_out: Electric field at the image plane. :return: Loss value. .. py:method:: callback(x) This function is called after every iteration of the optimisation. It saves intermediate SLM phase patterns and the electric field in the image plane if ``save=True``. The progress of the optimisation is printed after every iteration. .. py:method:: retrieve_phase() Performs phase retrieval algorithm. :return: Optimised SLM phase(s), (RMS error and efficiency if ``save=True``) .. py:function:: camera_calibration(slm_obj, slm_disp_obj, cam_obj, pms_obj, save=None, exp_time=None, checkerboard_rows=None, checkerboard_columns=None, checkerboard_square_size=None, linear_phase=None) This function performs the camera calibration to obtain the coordinate transform between the camera image and the computational image plane. To do this, an SLM phase pattern is calculated for a checkerboard-shaped target light potential using CG minimisation and displayed on the SLM. The corners of the checkerboard in the resulting camera image are detected and fitted to the corners of the checkerboard in the computational image plane using an affine transformation. :param slm_obj: Virtual SLM object created by :py:class:`VirtualSlm`. :param slm_disp_obj: Object created by a subclass of :py:class:`hardware.SlmBase`. :param cam_obj: Object created by a subclass of :py:class:`hardware.CameraBase`. :param pms_obj: Object created by a subclass of :py:class:`hardware.ParamsBase`. :param bool save: Save data? :param exp_time: Exposure time. :param checkerboard_rows: :param checkerboard_columns: :param checkerboard_square_size: :param linear_phase: :return: .. py:function:: camera_feedback(phase_retrieval_obj, slm_disp_obj, cam_obj, tf, itf, iter_fb=1, iter_cg=None, detect_vortices=False, threshold_vtx=0.2, n_save=10, n_avg=10, exp_time=1000, fb_blur=0, alpha=None, convergence=False, iter_convergence=None, path=None) This function implements a camera feedback algorithm to reduce experimental errors in the light potentials (see ``_). Before applying any camera feedback, optical vortices in the light potential are detected using the patterns.detect_vortices() function and removed if required. After vortices are removed, the optimised phase pattern is displayed on the SLM and a camera image, M, is recorded. To create the target light potential for the next feedback iteration, T[..., i + 1], a discrepancy, D, between the camera image and the original target light potential, T[..., 0], is calculated and added to the previous light potential, T[..., i]. At the end of each feedback iteration, the root-mean-squared error (RMSE) and the peak signal-to-noise ratio (PSNR) of the camera image are calculated and saved. To find the experimental convergence of the CG minimisation, intermediate SLM phase patterns are saved and displayed on the SLM. A camera image is taken for each pattern and the RMSE is calculated. :param phase_retrieval_obj: Instance of the class :py:class:`PhaseRetrieval`. :param slm_disp_obj: Object created by a subclass of :py:class:`hardware.SlmBase`. :param cam_obj: Object created by a subclass of :py:class:`hardware.CameraBase`. :param tf: Affine transform matrix. :param itf: Inverse affine transform matrix. :param iter_fb: Number of feedback iterations. :param iter_cg: Number of conjugate gradient iterations per feedback iteration. :param bool detect_vortices: Detect vortices? :param threshold_vtx: See ``patterns.detect_vortices()`` :param n_save: Save data for every ``n_save`` th CG iteration. :param n_avg: Number of camera frames to capture and average per feedback iteration. :param exp_time: Exposure time. :param fb_blur: Width of blurring kernel for camera image [px]. :param alpha: Feedback gain parameter for each feedback iteration. :param bool convergence: Save intermediate phase patterns during CG minimisation? :param iter_convergence: During which feedback iterations to save intermediate phase patterns. :param path: Save path. :return: See code.