Notebook demonstrating bulge+disk inference (#74)

* notebook demonstration * notebook updates * move notebooks * demonstrate inference with bulge + disk notebook * trying to fix efficiency, need to explore more options * comment
LSSTDESC · Jan 14, 2025 · 070f885 · 070f885
1 parent c72f59e
commit 070f885
Show file tree

Hide file tree

Showing 4 changed files with 1,339 additions and 2 deletions.
diff --git a/notebooks/bulge_disk_exp1.ipynb b/notebooks/bulge_disk_exp1.ipynb
diff --git a/notebooks/exp4_check1.ipynb → notebooks/old/exp4_check1.ipynb b/notebooks/exp4_check1.ipynb → notebooks/old/exp4_check1.ipynb
diff --git a/notebooks/shape-noise-cancellation-draft1.ipynb b/notebooks/shape-noise-cancellation-draft1.ipynb
@@ -365,7 +365,7 @@
    ],
    "source": [
     "from math import ceil\n",
-    "n_batches = 10\n",
+    "n_batches = 10 # rule of thumb: ~100\n",
     "batch_size = ceil(n_gals / n_batches)\n",
     "\n",
     "# g_pos_list = [] \n",
@@ -563,7 +563,7 @@
   "kernelspec": {
    "display_name": "bpd_gpu3",
    "language": "python",
-   "name": "bpd_gpu3"
+   "name": "python3"
   },
   "language_info": {
    "codemirror_mode": {

diff --git a/notebooks/spergel-speed1.ipynb b/notebooks/spergel-speed1.ipynb
@@ -0,0 +1,306 @@
+{
+ "cells": [
+  {
+   "cell_type": "code",
+   "execution_count": 1,
+   "id": "bc3945e8-3e07-416d-88d3-c3b35421afad",
+   "metadata": {
+    "tags": []
+   },
+   "outputs": [],
+   "source": [
+    "import os\n",
+    "os.environ[\"CUDA_VISIBLE_DEVICES\"] = \"0\"\n",
+    "os.environ[\"JAX_ENABLE_X64\"] = \"True\"\n",
+    "\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 45,
+   "id": "d7831e80-3e63-4f59-b180-3d984e45fa85",
+   "metadata": {
+    "tags": []
+   },
+   "outputs": [],
+   "source": [
+    "from functools import partial\n",
+    "\n",
+    "import jax\n",
+    "import jax.numpy as jnp\n",
+    "import jax_galsim as xgalsim\n",
+    "import matplotlib.pyplot as plt \n",
+    "\n",
+    "from jax import random\n",
+    "from jax import vmap, grad, jit\n",
+    "\n",
+    "import numpy as np\n",
+    "\n",
+    "import galsim\n",
+    "\n",
+    "import matplotlib.pyplot as plt\n",
+    "\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 6,
+   "id": "6e850712-6982-410d-a35d-ce77cfa93e5a",
+   "metadata": {
+    "tags": []
+   },
+   "outputs": [],
+   "source": [
+    "@jax.jit\n",
+    "def render_bd(\n",
+    "    lf, scale_radius, q, beta, x, y, *, psf_hlr=0.7, slen=53, fft_size=256, pixel_scale=0.2\n",
+    "):\n",
+    "    gsparams = xgalsim.GSParams(minimum_fft_size=fft_size, maximum_fft_size=fft_size)\n",
+    "\n",
+    "    bulge = xgalsim.Spergel(nu=-0.6, flux=10**lf, scale_radius=scale_radius).shear(\n",
+    "        q=q,\n",
+    "        beta=beta * xgalsim.radians,\n",
+    "    )\n",
+    "\n",
+    "    psf = xgalsim.Gaussian(flux=1.0, half_light_radius=0.7)\n",
+    "    gal_conv = xgalsim.Convolve([bulge, psf]).withGSParams(gsparams)\n",
+    "    galaxy_image = gal_conv.drawImage(nx=slen, ny=slen, scale=pixel_scale, offset=(x,y)).array\n",
+    "    return galaxy_image"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 39,
+   "id": "6c228ea9-e513-4aa9-a318-f4d852fbc46d",
+   "metadata": {
+    "tags": []
+   },
+   "outputs": [],
+   "source": [
+    "def render_bd_galsim(\n",
+    "    lf, scale_radius, q, beta, x, y, *, psf_hlr=0.7, slen=53, fft_size=256, pixel_scale=0.2\n",
+    "):\n",
+    "\n",
+    "    bulge = galsim.Spergel(nu=-0.6, flux=10**lf, scale_radius=scale_radius).shear(\n",
+    "        q=q,\n",
+    "        beta=beta * galsim.radians,\n",
+    "    )\n",
+    "\n",
+    "    psf = galsim.Gaussian(flux=1.0, half_light_radius=0.7)\n",
+    "    gal_conv = galsim.Convolve([bulge, psf])\n",
+    "    galaxy_image = gal_conv.drawImage(nx=slen, ny=slen, scale=pixel_scale, offset=(x,y)).array\n",
+    "    return galaxy_image"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 35,
+   "id": "e54f7103-90fe-467d-87da-13dcbec17b7c",
+   "metadata": {
+    "tags": []
+   },
+   "outputs": [
+    {
+     "data": {
+      "text/plain": [
+       "66"
+      ]
+     },
+     "execution_count": 35,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "# image size? \n",
+    "bulge = galsim.Spergel(nu=-0.6, flux=10**5, scale_radius=0.7).shear(\n",
+    "    q=0.2,\n",
+    "    beta=np.pi/2 * galsim.radians,\n",
+    ")\n",
+    "\n",
+    "psf = galsim.Gaussian(flux=1.0, half_light_radius=0.7)\n",
+    "gal_conv = galsim.Convolve([bulge, psf])\n",
+    "gal_conv.getGoodImageSize(0.2)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 36,
+   "id": "b7bf387e-86d4-426b-875c-87aa95706215",
+   "metadata": {
+    "tags": []
+   },
+   "outputs": [],
+   "source": [
+    "image  = render_bd(5.0, 0.7, 0.2, jnp.pi / 2, 0.0, 0.0) # compile"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 37,
+   "id": "83a10db6-cdeb-4645-82b9-05bf04249c30",
+   "metadata": {
+    "tags": []
+   },
+   "outputs": [
+    {
+     "data": {
+      "text/plain": [
+       "<matplotlib.image.AxesImage at 0x7fa48c5e8320>"
+      ]
+     },
+     "execution_count": 37,
+     "metadata": {},
+     "output_type": "execute_result"
+    },
+    {
+     "data": {
+      "image/png": "iVBORw0KGgoAAAANSUhEUgAAAaAAAAGfCAYAAAAZGgYhAAAAOnRFWHRTb2Z0d2FyZQBNYXRwbG90bGliIHZlcnNpb24zLjEwLjAsIGh0dHBzOi8vbWF0cGxvdGxpYi5vcmcvlHJYcgAAAAlwSFlzAAAPYQAAD2EBqD+naQAAJAVJREFUeJzt3X9sleX9//HXfdqeUwR6Kr9aOijD+AN/BIyVH/2qm4NOwscQHP2DGZIxR2Z0hQi4bLJM0GRLiSaibIBmY5AlY0yWoMF9hjNVanSUQZWIuvUDho0ucIpuH05LoaenPdf3Dz+e7Yxe19bDaa/T9vlI7gTOde77XOc6bV/nPve77wbGGCMAAAZZyPcEAAAjEwEEAPCCAAIAeEEAAQC8IIAAAF4QQAAALwggAIAXBBAAwAsCCADgBQEEAPCicKAOvHXrVj399NOKxWKaNWuWfvSjH2nOnDn/dr9UKqUzZ85o7NixCoJgoKYHABggxhh1dHSooqJCoZDjPMcMgD179phwOGx+9rOfmQ8++MB885vfNKWlpaatre3f7tva2moksbGxsbEN8a21tdX58z4wJvfNSOfOnavZs2frxz/+saRPz2qmTp2q1atX67HHHnPuG4/HVVpaqjv1XypUUa6nBvzDYJ9hB5Z3giY1uPOg/zAGWI+Sekv/rfPnzysajVrvl/OP4Lq7u9Xc3Kz169enbwuFQqqpqdGhQ4cuu38ikVAikUj/v6Oj4/8mVqTCgADCAMqXANIgB5AIIAyw//sS+3eXUXJehPDJJ5+ot7dXZWVlGbeXlZUpFotddv/6+npFo9H0NnXq1FxPCQCQh7xXwa1fv17xeDy9tba2+p4SAGAQ5PwjuAkTJqigoEBtbW0Zt7e1tam8vPyy+0ciEUUikVxPAwCQ53J+BhQOh1VVVaWGhob0balUSg0NDaqurs71wwFuQWDfsj5mKLttsI73b+c/AGsCZGFAfg9o3bp1WrFihW6//XbNmTNHzz77rDo7O/XAAw8MxMMBAIagAQmgZcuW6eOPP9aGDRsUi8V066236sCBA5cVJgAARq4B+T2gK9He3q5oNKq7tYQybFy5gfhY6Uo+/sqlgfj9ofz6cYAhqsckdVAvKx6Pq6SkxHq/PPlOAgCMNAQQAMALAggA4MWAdcMGhrR8uc7j4prjYPeXA7IwBL7LAADDEQEEAPCCAAIAeEEAAQC8IIAAAF4QQAAALyjDxtCXR12cg1Bu52JSg9wax7WWtOlBjnEGBADwggACAHhBAAEAvCCAAABeEEAAAC8IIACAF5RhY+TKsuN1rkuts30sZ4k2nbIxBHAGBADwggACAHhBAAEAvCCAAABeEEAAAC8IIACAF5RhA33IutQ6y9JuK0fJdNYl2kCe4AwIAOAFAQQA8IIAAgB4QQABALwggAAAXhBAAAAvKMPG0BDkvix6QEqtc90pO5VdV2t3iXaWnbJdr4Gh7Bv9xxkQAMALAggA4AUBBADwggACAHhBAAEAvCCAAABeUIYN9CXLUusg23JxC+N6i5htOTWQJzgDAgB4QQABALwggAAAXhBAAAAvCCAAgBcEEADAC8qwMfRl2/F6IEqtQ7l9Txek7OXU2ZZoByHHMSntxiDiDAgA4AUBBADwggACAHhBAAEAvCCAAABeEEAAAC8ow8bINRCl1jnuhu16rKxLtHuvYD5ADnEGBADwggACAHhBAAEAvCCAAABeEEAAAC8IIACAF/0OoDfffFOLFy9WRUWFgiDQSy+9lDFujNGGDRs0efJkjRo1SjU1NTpx4kSu5ovhLAjsW9bHDNk3l1DIvhUUWLcgx5vrsZxzHIg1cR5zAF47DHv9/orr7OzUrFmztHXr1j7Hn3rqKW3ZskXPP/+8Dh8+rNGjR2vhwoXq6uq64skCAIaPfv8i6qJFi7Ro0aI+x4wxevbZZ/X9739fS5YskST9/Oc/V1lZmV566SV99atfvbLZAgCGjZxeAzp16pRisZhqamrSt0WjUc2dO1eHDh3qc59EIqH29vaMDQAw/OU0gGKxmCSprKws4/aysrL02L+qr69XNBpNb1OnTs3llAAAecp7Fdz69esVj8fTW2trq+8pAQAGQU4DqLy8XJLU1taWcXtbW1t67F9FIhGVlJRkbACA4S+n3bCnT5+u8vJyNTQ06NZbb5Uktbe36/Dhw3r44Ydz+VAYaRwlwoGjq7XzkK4S4YIC+34Frm7Yuf1QITCOjteu/Vydsh17utbSpBzPzTFPwKbfAXThwgWdPHky/f9Tp07p2LFjGjdunCorK7VmzRr94Ac/0HXXXafp06fr8ccfV0VFhe67775czhsAMMT1O4COHj2qL33pS+n/r1u3TpK0YsUK7dq1S9/5znfU2dmpBx98UOfPn9edd96pAwcOqLi4OHezBgAMeYExxnUmP+ja29sVjUZ1t5aoMCjyPR0MJtdHYtl+BOfaz/VRWpH9a28wP4JzfbRleh0feyWT2e3neryU40eF6yO4/PoRg0HQY5I6qJcVj8ed1/W9V8EBAEYmAggA4AUBBADwIqdl2EDecV0fcpVaZ1mindV1LOf1E/vxAsc1GeOao+uaTK99CMg1zoAAAF4QQAAALwggAIAXBBAAwAsCCADgBQEEAPCCMmwMfa52O66y6JCrFY/jWyOUbadsy1yMfR9n2xxXpypHqbVrTYyzlRA12sgtzoAAAF4QQAAALwggAIAXBBAAwAsCCADgBQEEAPCCMmwMCc6/epptx2tXybSr1DrsqH92lXbb5unoah0UOP5Cabf9oYICe8m0SWXXKTtwdOY2VGgjC5wBAQC8IIAAAF4QQAAALwggAIAXBBAAwAsCCADgBWXYGFyu7tTOTsyuQ2bZ8brQ/uXvLLV2dMo2hY4SZ2s3bEfpc4+9vtnxrGWMo4u2o8O2s1O24/Hcr51jLo7njuGPMyAAgBcEEADACwIIAOAFAQQA8IIAAgB4QQABALygDBtDg6vM11Fq7ep4HTjKsJ2l1pGwfcyxn/XtnqNKWcke65CrDDvodbSn7rEf0/Q61tlZak07bPQfZ0AAAC8IIACAFwQQAMALAggA4AUBBADwggACAHhBGTaGhpCri7ZjzFlqbe947Sy1HmUfS4Xt3bCNpSQ8cHSnDnVn2SHccUxXGbar7Nv5GlCFjSxwBgQA8IIAAgB4QQABALwggAAAXhBAAAAvCCAAgBeUYSNvBI4y38BRah0U2EufXR2vTcRRhu0ote4Z4yjDjtjnkirs+zmEeox1n1DCXt/s/OZ1dMMOkkn7fo61dB3TVaJtKNGGBWdAAAAvCCAAgBcEEADACwIIAOAFAQQA8IIAAgB4QRk2hoaQ471SkePL2NXVujhiHXOVWiej9vLt5FX2eZqCvkuVg157GXbRxezeIxb22LthB66O164SbVcXbSALnAEBALwggAAAXhBAAAAvCCAAgBcEEADAi34FUH19vWbPnq2xY8dq0qRJuu+++9TS0pJxn66uLtXV1Wn8+PEaM2aMamtr1dbWltNJAwCGvn6VYTc2Nqqurk6zZ89WT0+Pvve97+mee+7Rhx9+qNGjR0uS1q5dq9/85jfau3evotGoVq1apaVLl+rtt98ekCeAPOToXK3A8Z7HNebq0uzseG0vp+4day/D7i6179c1zj6X7jGOrtCWaQaO6ubwBUdZt6tDuKO0O+h2lFp3JexjrtfA+bra5yLZy8VlXPthOOhXAB04cCDj/7t27dKkSZPU3NysL3zhC4rH49qxY4d2796t+fPnS5J27typG2+8UU1NTZo3b17uZg4AGNKu6BpQPB6XJI0bN06S1NzcrGQyqZqamvR9ZsyYocrKSh06dOhKHgoAMMxk3QkhlUppzZo1uuOOO3TLLbdIkmKxmMLhsEpLSzPuW1ZWplgs1udxEomEEol/nPa3t7dnOyUAwBCS9RlQXV2d3n//fe3Zs+eKJlBfX69oNJrepk6dekXHAwAMDVkF0KpVq/TKK6/ojTfe0JQpU9K3l5eXq7u7W+fPn8+4f1tbm8rLy/s81vr16xWPx9Nba2trNlMCAAwx/QogY4xWrVqlffv26fXXX9f06dMzxquqqlRUVKSGhob0bS0tLTp9+rSqq6v7PGYkElFJSUnGBgAY/vp1Daiurk67d+/Wyy+/rLFjx6av60SjUY0aNUrRaFQrV67UunXrNG7cOJWUlGj16tWqrq6mAg7/XshRVlxgf68UFNm7U6dGObpal9j3uzTeXnJ8aZJ9nt1Re+lwyvJwIVcD6rijrDuwzzGUtD+3gov2NQldtO+nAnuJtnG8duq1D2Fk61cAbd++XZJ09913Z9y+c+dOff3rX5ckbd68WaFQSLW1tUokElq4cKG2bduWk8kCAIaPfgWQ+Q9+May4uFhbt27V1q1bs54UAGD4oxccAMALAggA4AUBBADwggACAHiRdSseIBuBq9Ta1UXb1Yk5bC8d7h1tLznujtq//Lsm2OdycbK9g7MZ320dK4z03fY6mbDPo+cq+/yDlP39Y9El+zGLOhxl2B2uMmz7a+B87RyvuaFEe0TjDAgA4AUBBADwggACAHhBAAEAvCCAAABeEEAAAC8ow0b+CDneDxXav1RNxF463HOVfb9E1F4e3DXR3vewoOKidezask+sYxOKL/R5+yddY6z7nIxMsI51dY+2joXb7c8t8nf7mhQ61tL1GjhfO8CCrxoAgBcEEADACwIIAOAFAQQA8IIAAgB4QQABALygDBv5w9FROXCU+abCjm7So+0dnLtL7I+XHJe0jt1c/rF1rGbiH61jnw/3XaL95257qXUouNE69kF7xDrW/bGjNN2xJsaxliHHa2Bc3bABC86AAABeEEAAAC8IIACAFwQQAMALAggA4AUBBADwgjJsDK7A8Z7H1VG5wFE6XOgo0S6ylwf3FNsfrmhst3XshrFt1rH/d9UJ69g1hX0fs6Lwf637tI4dZx37n7GTrGM9xfYybNeauNbS9Ro4XzvXa65exxiGO86AAABeEEAAAC8IIACAFwQQAMALAggA4AVVcBj6XI1KC1wVX/ZDFoV7rGNXF120jk0MJaxjEwrG9Hl7R+pCVo/lmmPK8dxca+KsZgNyjK82AIAXBBAAwAsCCADgBQEEAPCCAAIAeEEAAQC8oAwbQ18qZR0K9RrrWGCvYlay2/6t8b/Jq6xjH6ci1rGxvZ393sf1WK45Fjqem2tNXGsJ5BpnQAAALwggAIAXBBAAwAsCCADgBQEEAPCCAAIAeEEZNgaXcZT5ukqAe3utQ0GPoww7aS85LuyyP9yljrB1rKWjzDr2++LrrGNnwp/0efufuydk9VhJxxyLHc/NtSautXS9Bs7XzvWaY0TjDAgA4AUBBADwggACAHhBAAEAvCCAAABeEEAAAC8ow0b+MPbyYOMo8w267a2fCzvtpcPh9gLrWNHf7d8a/xObaB1LmcA6NqH4Qp+3f9I1xrrPyTZ7ibZrjuF2R/m5Y01ca+l6DVyvHWDDGRAAwAsCCADgBQEEAPCCAAIAeEEAAQC8IIAAAF70qwx7+/bt2r59u/785z9Lkm6++WZt2LBBixYtkiR1dXXp0Ucf1Z49e5RIJLRw4UJt27ZNZWX2jr5AmqvMt8deHhwkktaxwov2/SLxIutY8cf292ad4ausY39KlNvnEul7Lj0J+7dh8Dd7x+vRH9tLviNx+1q61sS1lsbxGjhfO8CiX2dAU6ZM0aZNm9Tc3KyjR49q/vz5WrJkiT744ANJ0tq1a7V//37t3btXjY2NOnPmjJYuXTogEwcADG39OgNavHhxxv9/+MMfavv27WpqatKUKVO0Y8cO7d69W/Pnz5ck7dy5UzfeeKOampo0b9683M0aADDkZX0NqLe3V3v27FFnZ6eqq6vV3NysZDKpmpqa9H1mzJihyspKHTp0yHqcRCKh9vb2jA0AMPz1O4COHz+uMWPGKBKJ6KGHHtK+fft00003KRaLKRwOq7S0NOP+ZWVlisVi1uPV19crGo2mt6lTp/b7SQAAhp5+B9ANN9ygY8eO6fDhw3r44Ye1YsUKffjhh1lPYP369YrH4+mttbU162MBAIaOfjcjDYfDuvbaayVJVVVVOnLkiJ577jktW7ZM3d3dOn/+fMZZUFtbm8rL7ZVBkUhEkUik/zMHAAxpV9wNO5VKKZFIqKqqSkVFRWpoaFBtba0kqaWlRadPn1Z1dfUVTxTDg0k5uiY7OioHvfYOzuq2lw4XdHZbx8KuMuxR9jETsn9w0H3R/mYqVdT3WJF9+grH7aXWxZ/Y1ysct5dMu9bEtZZyvAbG2cmcTtnoW78CaP369Vq0aJEqKyvV0dGh3bt36+DBg3r11VcVjUa1cuVKrVu3TuPGjVNJSYlWr16t6upqKuAAAJfpVwCdO3dOX/va13T27FlFo1HNnDlTr776qr785S9LkjZv3qxQKKTa2tqMX0QFAOBfBcZ17uxBe3u7otGo7tYSFQb2j0GQxwL7x0YK7B9fBUX290OhUcX2/cbY/6BbanyJdayrfLR1rLPc/rV3aZL9+XVH7d9OKcshQ1l+BDfqnP2xRsfsBy2OdVrHQn+z/xqEudD3H9STpNSlLvt+SUcHBcMfuRuOekxSB/Wy4vG4Skrs34P0ggMAeEEAAQC8IIAAAF5ccRk2kDOOcl3T67hWkLRf7wgu2UuOi9rtnaZHFTmuVZkC61hhp/2ajbF8twWOSyThC/Y1Kf67vSy6qD27NTGOtXS+BpRaIwucAQEAvCCAAABeEEAAAC8IIACAFwQQAMALAggA4AVl2Mg9ZwsVV+sVRwsfVzfsHnsdc5CwlxwXdCSsY+EC+1wCY2/TU3jJ/p7OWI4Z9NrXq+iifb2K4o4u4I7n5loT41hL52vgbKlDux30jTMgAIAXBBAAwAsCCADgBQEEAPCCAAIAeEEAAQC8oAwbQ0PK1Q3bUTrsKDl2/QXWwgvZvTcr6LJ3yk4V9l2GHeqxlyKHEvbS58ILjufWZS/Ddq2Jcy1drwGQBc6AAABeEEAAAC8IIACAFwQQAMALAggA4AUBBADwgjJs5A2TcnRGdnVNzrpTtr2btArs5dSub5pUt30/U9D3+72g117eHOq2P7fgkqMM2/HcXB2vjWMtjeM1cL52gAVnQAAALwggAIAXBBAAwAsCCADgBQEEAPCCAAIAeEEZNoaGbEu0HSXHStpLlYNEdu/NQj2ObynbIR1NpgNHd+rA2dXaUWLuWhPXWlJqjRzjDAgA4AUBBADwggACAHhBAAEAvCCAAABeEEAAAC8ow8bQYBy1yin7mHF0mnZ2yk7au1oH9iO6O3MHlj0dpc9Bj+N4jhJtd8fr7NbS+RoAWeAMCADgBQEEAPCCAAIAeEEAAQC8IIAAAF4QQAAALyjDxuBydVt2tYWWvSzauMqYXWXFrlLlwP7ezFWGHbhKnEOWPV1dpl0l5t1Zdrx2HdP5+ji4SrSzPSaGPc6AAABeEEAAAC8IIACAFwQQAMALAggA4AUBBADwgjJsDAnGUaocBI4yX0d3ahNylFoXOPbrtj9cUOB4T5dFN2x352rHHF37uTp2O9bZ9RoA2eAMCADgBQEEAPCCAAIAeEEAAQC8IIAAAF4QQAAAL64ogDZt2qQgCLRmzZr0bV1dXaqrq9P48eM1ZswY1dbWqq2t7UrnCdiZlHUzxlg3pVL2Ldlj31K91s309Di2Xstm38f1WO452p+bc00cawnkWtYBdOTIEb3wwguaOXNmxu1r167V/v37tXfvXjU2NurMmTNaunTpFU8UADC8ZBVAFy5c0PLly/WTn/xEV199dfr2eDyuHTt26JlnntH8+fNVVVWlnTt36ve//72amppyNmkAwNCXVQDV1dXp3nvvVU1NTcbtzc3NSiaTGbfPmDFDlZWVOnToUJ/HSiQSam9vz9gAAMNfv1vx7NmzR++8846OHDly2VgsFlM4HFZpaWnG7WVlZYrFYn0er76+Xk8++WR/pwEAGOL6dQbU2tqqRx55RL/4xS9UXFyckwmsX79e8Xg8vbW2tubkuACA/NavAGpubta5c+d02223qbCwUIWFhWpsbNSWLVtUWFiosrIydXd36/z58xn7tbW1qby8vM9jRiIRlZSUZGwAgOGvXx/BLViwQMePH8+47YEHHtCMGTP03e9+V1OnTlVRUZEaGhpUW1srSWppadHp06dVXV2du1kD/ylXB+dsO2W7ukkHrvd0WZQyO8qfjaOLdrYdr4HB1K8AGjt2rG655ZaM20aPHq3x48enb1+5cqXWrVuncePGqaSkRKtXr1Z1dbXmzZuXu1kDAIa8nP89oM2bNysUCqm2tlaJREILFy7Utm3bcv0wAIAhLjDO8/jB197ermg0qru1RIVBke/pYDDZ/mCb5PxoKwhluZ/rj8cV2b/23H90LsfdrVwfwbn+6Fwymd1+rsdzfXTn6pSQXz9iMAh6TFIH9bLi8bjzuj694AAAXhBAAAAvCCAAgBc5L0IABoTz2oTrOo/jkI5rE65Sa9cVjSDI7fWObEuts720m/V1HiALnAEBALwggAAAXhBAAAAvCCAAgBcEEADACwIIAOAFZdjIH67SYVebHucxXaXDjhrtVHYlxybbeVoP6FiTLOc4IOXUtNtBFjgDAgB4QQABALwggAAAXhBAAAAvCCAAgBcEEADAC8qwMXI5Oj8b1x89dZU/h3L8ns7xWM6O166u1kCe4AwIAOAFAQQA8IIAAgB4QQABALwggAAAXhBAAAAvKMPG0Ofo7mxS9vdYQchRTu3YL+sS7SxkXWrtXJPs9gNyjTMgAIAXBBAAwAsCCADgBQEEAPCCAAIAeEEAAQC8oAwb6IurHDnLEu2sUDKNYYwzIACAFwQQAMALAggA4AUBBADwggACAHhBAAEAvKAMG0ODqyt0EDj2y7ZTdnbHdJVoZyXLUusB6Xjteg2ALHAGBADwggACAHhBAAEAvCCAAABeEEAAAC8IIACAF5RhA31wlTFnXaKdY85Sa2AI4AwIAOAFAQQA8IIAAgB4QQABALwggAAAXhBAAAAvKMPGyOUqmQ7s782yLtHOQtal1oNYDg5kizMgAIAXBBAAwAsCCADgBQEEAPCCAAIAeEEAAQC86FcAPfHEEwqCIGObMWNGeryrq0t1dXUaP368xowZo9raWrW1teV80kAGY+zbYE8lZXK6Df4TyJ+1xPDX7zOgm2++WWfPnk1vb731Vnps7dq12r9/v/bu3avGxkadOXNGS5cuzemEAQDDQ79/EbWwsFDl5eWX3R6Px7Vjxw7t3r1b8+fPlyTt3LlTN954o5qamjRv3rwrny0AYNjo9xnQiRMnVFFRoWuuuUbLly/X6dOnJUnNzc1KJpOqqalJ33fGjBmqrKzUoUOHrMdLJBJqb2/P2AAAw1+/Amju3LnatWuXDhw4oO3bt+vUqVO666671NHRoVgspnA4rNLS0ox9ysrKFIvFrMesr69XNBpNb1OnTs3qiQAAhpZ+fQS3aNGi9L9nzpypuXPnatq0aXrxxRc1atSorCawfv16rVu3Lv3/9vZ2QggARoArKsMuLS3V9ddfr5MnT6q8vFzd3d06f/58xn3a2tr6vGb0mUgkopKSkowNADD8XVEAXbhwQR999JEmT56sqqoqFRUVqaGhIT3e0tKi06dPq7q6+oonCgwqk7Jv+WIozBFw6NdHcN/+9re1ePFiTZs2TWfOnNHGjRtVUFCg+++/X9FoVCtXrtS6des0btw4lZSUaPXq1aqurqYCDgBwmX4F0F//+lfdf//9+tvf/qaJEyfqzjvvVFNTkyZOnChJ2rx5s0KhkGpra5VIJLRw4UJt27ZtQCYOABjaAmPy61ec29vbFY1GdbeWqDAo8j0dDHVBbv9A3KfHzJMOVgPxUVt+/TjAENVjkjqolxWPx53X9fPkOwkAMNIQQAAALwggAIAX/e4FBwwprmsa2V4fyvbai+3a0WCXTXOdB3mCMyAAgBcEEADACwIIAOAFAQQA8IIAAgB4kXdVcJ81ZuhRUqJYBwNqALokOFEFh5GhR0lJ//h5bpN3AdTR0SFJekv/7XkmGPYG++cwP/cxwnR0dCgajVrH864XXCqV0pkzZzR27FgFQZD+A3Wtra38raB/wrpcjjXpG+vSN9alb7lYF2OMOjo6VFFRoVDIfqUn786AQqGQpkyZctnt/LG6vrEul2NN+sa69I116duVrovrzOczFCEAALwggAAAXuR9AEUiEW3cuFGRSMT3VPIK63I51qRvrEvfWJe+Dea65F0RAgBgZMj7MyAAwPBEAAEAvCCAAABeEEAAAC/yOoC2bt2qz3/+8youLtbcuXP1hz/8wfeUBtWbb76pxYsXq6KiQkEQ6KWXXsoYN8Zow4YNmjx5skaNGqWamhqdOHHCz2QHUX19vWbPnq2xY8dq0qRJuu+++9TS0pJxn66uLtXV1Wn8+PEaM2aMamtr1dbW5mnGA2/79u2aOXNm+pcHq6ur9dvf/jY9PtLWw2bTpk0KgkBr1qxJ3zYS1+aJJ55QEAQZ24wZM9Ljg7UmeRtAv/rVr7Ru3Tpt3LhR77zzjmbNmqWFCxfq3Llzvqc2aDo7OzVr1ixt3bq1z/GnnnpKW7Zs0fPPP6/Dhw9r9OjRWrhwobq6ugZ5poOrsbFRdXV1ampq0muvvaZkMql77rlHnZ2d6fusXbtW+/fv1969e9XY2KgzZ85o6dKlHmc9sKZMmaJNmzapublZR48e1fz587VkyRJ98MEHkkbeevTlyJEjeuGFFzRz5syM20fq2tx88806e/ZsenvrrbfSY4O2JiZPzZkzx9TV1aX/39vbayoqKkx9fb3HWfkjyezbty/9/1QqZcrLy83TTz+dvu38+fMmEomYX/7ylx5m6M+5c+eMJNPY2GiM+XQdioqKzN69e9P3+eMf/2gkmUOHDvma5qC7+uqrzU9/+lPWwxjT0dFhrrvuOvPaa6+ZL37xi+aRRx4xxozcr5WNGzeaWbNm9Tk2mGuSl2dA3d3dam5uVk1NTfq2UCikmpoaHTp0yOPM8sepU6cUi8Uy1igajWru3Lkjbo3i8bgkady4cZKk5uZmJZPJjLWZMWOGKisrR8Ta9Pb2as+ePers7FR1dfWIXw9Jqqur07333puxBtLI/lo5ceKEKioqdM0112j58uU6ffq0pMFdk7xrRipJn3zyiXp7e1VWVpZxe1lZmf70pz95mlV+icViktTnGn02NhKkUimtWbNGd9xxh2655RZJn65NOBxWaWlpxn2H+9ocP35c1dXV6urq0pgxY7Rv3z7ddNNNOnbs2Ihcj8/s2bNH77zzjo4cOXLZ2Ej9Wpk7d6527dqlG264QWfPntWTTz6pu+66S++///6grkleBhDwn6qrq9P777+f8fn1SHXDDTfo2LFjisfj+vWvf60VK1aosbHR97S8am1t1SOPPKLXXntNxcXFvqeTNxYtWpT+98yZMzV37lxNmzZNL774okaNGjVo88jLj+AmTJiggoKCy6ou2traVF5e7mlW+eWzdRjJa7Rq1Sq98soreuONNzL+hEd5ebm6u7t1/vz5jPsP97UJh8O69tprVVVVpfr6es2aNUvPPffciF0P6dOPk86dO6fbbrtNhYWFKiwsVGNjo7Zs2aLCwkKVlZWN2LX5Z6Wlpbr++ut18uTJQf16ycsACofDqqqqUkNDQ/q2VCqlhoYGVVdXe5xZ/pg+fbrKy8sz1qi9vV2HDx8e9mtkjNGqVau0b98+vf7665o+fXrGeFVVlYqKijLWpqWlRadPnx72a/PPUqmUEonEiF6PBQsW6Pjx4zp27Fh6u/3227V8+fL0v0fq2vyzCxcu6KOPPtLkyZMH9+slpyUNObRnzx4TiUTMrl27zIcffmgefPBBU1paamKxmO+pDZqOjg7z7rvvmnfffddIMs8884x59913zV/+8hdjjDGbNm0ypaWl5uWXXzbvvfeeWbJkiZk+fbq5dOmS55kPrIcffthEo1Fz8OBBc/bs2fR28eLF9H0eeughU1lZaV5//XVz9OhRU11dbaqrqz3OemA99thjprGx0Zw6dcq899575rHHHjNBEJjf/e53xpiRtx4u/1wFZ8zIXJtHH33UHDx40Jw6dcq8/fbbpqamxkyYMMGcO3fOGDN4a5K3AWSMMT/60Y9MZWWlCYfDZs6cOaapqcn3lAbVG2+8YSRdtq1YscIY82kp9uOPP27KyspMJBIxCxYsMC0tLX4nPQj6WhNJZufOnen7XLp0yXzrW98yV199tbnqqqvMV77yFXP27Fl/kx5g3/jGN8y0adNMOBw2EydONAsWLEiHjzEjbz1c/jWARuLaLFu2zEyePNmEw2Hzuc99zixbtsycPHkyPT5Ya8KfYwAAeJGX14AAAMMfAQQA8IIAAgB4QQABALwggAAAXhBAAAAvCCAAgBcEEADACwIIAOAFAQQA8IIAAgB4QQABALz4/8TXk4rDmIZ2AAAAAElFTkSuQmCC",
+      "text/plain": [
+       "<Figure size 640x480 with 1 Axes>"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    }
+   ],
+   "source": [
+    "plt.imshow(image)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 17,
+   "id": "10c8f524-dd45-4114-a756-93c6e5991eb6",
+   "metadata": {
+    "tags": []
+   },
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "210 μs ± 2.42 μs per loop (mean ± std. dev. of 7 runs, 1,000 loops each)\n"
+     ]
+    }
+   ],
+   "source": [
+    "# timing\n",
+    "%timeit _ = render_bd(5.0, 1.0, 0.2, jnp.pi / 2, 0.0, 0.0)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 41,
+   "id": "b3b3ca3c-dd74-4325-ab2c-48868cd44329",
+   "metadata": {
+    "tags": []
+   },
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "413 μs ± 18.3 μs per loop (mean ± std. dev. of 7 runs, 1,000 loops each)\n"
+     ]
+    }
+   ],
+   "source": [
+    "%timeit _ = render_bd_galsim(5.0, 1.0, 0.2, jnp.pi / 2, 0.0, 0.0)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 49,
+   "id": "b6d75ca1-6c98-47a5-a206-10cfb15eb574",
+   "metadata": {
+    "tags": []
+   },
+   "outputs": [],
+   "source": [
+    "from bpd.draw import draw_gaussian\n",
+    "\n",
+    "\n",
+    "draw_gaussian_jitted = jax.jit(partial(draw_gaussian, slen=53, fft_size=252))\n",
+    "_ = draw_gaussian_jitted(f=1e6, hlr=1.0, e1=0.2, e2=0.2, g1=0.02, g2=0.0,x=0,y=0)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 51,
+   "id": "938e232a-c52a-48d6-8c76-ab62f58a041c",
+   "metadata": {
+    "tags": []
+   },
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "241 μs ± 1.38 μs per loop (mean ± std. dev. of 7 runs, 1,000 loops each)\n"
+     ]
+    }
+   ],
+   "source": [
+    "# compare with Gaussian\n",
+    "%timeit _ = draw_gaussian_jitted(f=1e6, hlr=1.0, e1=0.2, e2=0.2, g1=0.02, g2=0.0,x=0,y=0)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 53,
+   "id": "cc054d34-5ef1-4b17-bdb1-d0c69b6bdded",
+   "metadata": {
+    "tags": []
+   },
+   "outputs": [
+    {
+     "data": {
+      "text/plain": [
+       "106"
+      ]
+     },
+     "execution_count": 53,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "# check good size of Gaussian to see how much it matters\n",
+    "gal = galsim.Gaussian(flux=1e5, half_light_radius=2.0)\n",
+    "psf = galsim.Gaussian(flux=1.0, half_light_radius=0.7)\n",
+    "gal_conv = galsim.Convolve([gal, psf])\n",
+    "gal_conv.getGoodImageSize(0.2)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "32b9d712-c53b-42a1-aa2e-ad8f059ad1ef",
+   "metadata": {},
+   "outputs": [],
+   "source": []
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "bpd_gpu3",
+   "language": "python",
+   "name": "bpd_gpu3"
+  },
+  "language_info": {
+   "codemirror_mode": {
+    "name": "ipython",
+    "version": 3
+   },
+   "file_extension": ".py",
+   "mimetype": "text/x-python",
+   "name": "python",
+   "nbconvert_exporter": "python",
+   "pygments_lexer": "ipython3",
+   "version": "3.12.8"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 5
+}