common.py

import torch
import torch.nn as nn
import numpy as np
import PIL.Image as Image
import matplotlib.pyplot as plt
import torchvision.transforms as transforms
from torch import optim
from torch.autograd import Variable
import torch.nn.functional as F


def unwrap(x):
    y = x % (2 * np.pi)
    return torch.where(y > np.pi, 2 * np.pi - y, y)


def fft2dc(x):
    return np.fft.fftshift(np.fft.fft2(x))


def ifft2dc(x):
    return np.fft.ifft2(np.fft.fftshift(x))


def Phase_unwrapping(in_):
    f = np.zeros((1000, 1000))
    for ii in range(1000):
        for jj in range(1000):
            x = ii - 1000 / 2
            y = jj - 1000 / 2
            f[ii, jj] = x**2 + y**2
    a = ifft2dc(fft2dc(np.cos(in_) * ifft2dc(fft2dc(np.sin(in_)) * f)) / (f + 0.000001))
    b = ifft2dc(fft2dc(np.sin(in_) * ifft2dc(fft2dc(np.cos(in_)) * f)) / (f + 0.000001))
    out = np.real(a - b)
    return out


def propagator(Nx, Ny, z, wavelength, deltaX, deltaY):
    k = 1 / wavelength
    x = np.expand_dims(
        np.arange(np.ceil(-Nx / 2), np.ceil(Nx / 2), 1) * (1 / (Nx * deltaX)), axis=0
    )
    y = np.expand_dims(
        np.arange(np.ceil(-Ny / 2), np.ceil(Ny / 2), 1) * (1 / (Ny * deltaY)), axis=1
    )
    y_new = np.repeat(y, Nx, axis=1)
    x_new = np.repeat(x, Ny, axis=0)
    kp = np.sqrt(y_new**2 + x_new**2)
    term = k**2 - kp**2
    term = np.maximum(term, 0)
    phase = np.exp(1j * 2 * np.pi * z * np.sqrt(term))
    return phase