Update python

Author: roflmaostc

SAS_pytorch.ipynb

@@ -66,6 +66,10 @@
    "outputs": [],
    "source": [
     "def zero_pad(arr, device=device):\n",
+    "    '''\n",
+    "    Pad arr with zeros to double the size. First dim is assumed to be batch dim which\n",
+    "    won't be changed\n",
+    "    '''\n",
     "    out_arr = torch.zeros(arr.shape[0], arr.shape[1] * 2, arr.shape[2] * 2, device=device, dtype=arr.dtype)\n",
     "    \n",
     "    as1 = (arr.shape[1] + 1) // 2\n",
@@ -74,7 +78,10 @@
     "    return out_arr\n",
     "\n",
     "def zero_unpad(arr, original_shape):\n",
-    "        \n",
+    "    '''\n",
+    "    Strip off padding of arr with zeros to halve the size. First dim is assumed to be batch dim which\n",
+    "    won't be changed\n",
+    "    '''\n",
     "    as1 = (original_shape[1] + 1) // 2\n",
     "    as2 = (original_shape[2] + 1) // 2\n",
     "    return arr[:, as1:as1 + original_shape[1], as2:as2 + original_shape[2]]\n"
@@ -86,7 +93,22 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "def scalable_angular_spectrum(psi, z, lbd, L, device=device):\n",
+    "def scalable_angular_spectrum(psi, z, lbd, L, device=device, skip_final_phase=True):\n",
+    "    '''\n",
+    "    Returns the complex electrical field psi propagated with the Scalable Angular Spectrum Method.\n",
+    "    \n",
+    "        Parameters:\n",
+    "            psi (torch.tensor): the quadratically shaped input field, with leading batch dimension\n",
+    "            z (number): propagation distance\n",
+    "            lbd (number): wavelength\n",
+    "            L (number): physical sidelength of the input field\n",
+    "            skip_final_phase=True: Skip final multiplication of phase factor. For M>2 undersampled,\n",
+    "            \n",
+    "        Returns:\n",
+    "            psi_final (torch.tensor): Propagated field\n",
+    "            Q (number): Output field size, corresponds to magnificiation * L\n",
+    "    \n",
+    "    '''\n",
     "    N = psi.shape[-1]\n",
     "    z_limit = (- 4 * L * np.sqrt(8*L**2 / N**2 + lbd**2) * np.sqrt(L**2 * 1 / (8 * L**2 + N**2 * lbd**2))\\\n",
     "               / (lbd * (-1+2 * np.sqrt(2) * np.sqrt(L**2 * 1 / (8 * L**2 + N**2 * lbd**2)))))\n",
@@ -137,14 +159,22 @@
     "    dq = lbd * z / L_new\n",
     "    Q = dq * N * pad_factor\n",
     "    \n",
+    "    q_y = torch.fft.ifftshift(torch.tensor(np.linspace(-Q/2, Q/2, N_new, endpoint=False),\n",
+    "                               device=device).reshape(1, 1, N_new), dim=(-1))\n",
+    "    q_x = q_y.reshape(1, N_new, 1)\n",
     "    \n",
-    "    # skip final H_2 phase\n",
     "    H_1 = torch.exp(1j * k / (2 * z) * (x**2 + y**2))\n",
-    "    psi_p_final = torch.fft.fftshift(torch.fft.fft2(H_1 * psi_precomp), dim=(-1,-2))\n",
     "\n",
+    "    if skip_final_phase:\n",
+    "        psi_p_final = torch.fft.fftshift(torch.fft.fft2(H_1 * psi_precomp), dim=(-1,-2))\n",
+    "    else:\n",
+    "        H_2 = np.exp(1j * k * z) * torch.exp(1j * k / (2 * z) * (q_x**2 + q_y**2))\n",
+    "        psi_p_final = torch.fft.fftshift(H_2 * torch.fft.fft2(H_1 * psi_precomp), dim=(-1,-2))\n",
+    "    \n",
+    "    \n",
     "    psi_final = zero_unpad(psi_p_final, psi.shape)\n",
     "    \n",
-    "    return psi_final, Q \n"
+    "    return psi_final, Q / 2\n"
    ]
   },
   {
@@ -178,7 +208,7 @@
     {
      "data": {
       "text/plain": [
-       "<matplotlib.image.AxesImage at 0x7f13aae403d0>"
+       "<matplotlib.image.AxesImage at 0x7fbe06e3f3a0>"
       ]
      },
      "execution_count": 7,
@@ -203,14 +233,12 @@
   {
    "cell_type": "code",
    "execution_count": 8,
-   "metadata": {
-    "scrolled": false
-   },
+   "metadata": {},
    "outputs": [
     {
      "data": {
       "text/plain": [
-       "<matplotlib.image.AxesImage at 0x7f13c43bdbe0>"
+       "<matplotlib.image.AxesImage at 0x7fbe1c2431c0>"
       ]
      },
      "execution_count": 8,
@@ -241,7 +269,7 @@
      "name": "stdout",
      "output_type": "stream",
      "text": [
-      "5.52 ms ± 148 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)\n"
+      "5.6 ms ± 292 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)\n"
      ]
     }
    ],