Save memory by removing figure information after plotting

Author: kadrlica

03_Image_Display_and_Manipulation.ipynb

@@ -72,7 +72,7 @@
     "%matplotlib inline\n",
     "import matplotlib.pyplot as plt      # imports matplotlib.pyplot as plt\n",
     "import warnings                      # imports the warnings library\n",
-    "\n",
+    "import gc                            # imports python's garbage collector\n",
     "from astropy.wcs import WCS          # imports astropy's World Coordinate System function WCS"
    ]
   },
@@ -95,6 +95,30 @@
     "warnings.simplefilter(\"ignore\", category=UserWarning)"
    ]
   },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "Matplotlib stores the data array associated with an image that is plotted. Since the LSST images are large (~4k x 4k pixels), this can eventually lead to a memory overflow, which will cause the notebook kernel to die. To mitigate this issue, we define a function to clean up after we plot them."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "def remove_figure(fig):\n",
+    "    \"\"\"Remove a figure to reduce memory footprint. \"\"\"\n",
+    "    # get the axes and clear their images\n",
+    "    for ax in fig.get_axes():\n",
+    "        for im in ax.get_images():\n",
+    "            im.remove()\n",
+    "    fig.clf()      # clear the figure\n",
+    "    plt.close(fig) # close the figure\n",
+    "    gc.collect()   # call the garbage collector"
+   ]
+  },
   {
    "cell_type": "markdown",
    "metadata": {},
@@ -234,15 +258,17 @@
    "outputs": [],
    "source": [
     "# create a matplotlib.pyplot figure\n",
-    "plt.figure()\n",
+    "fig,ax = plt.subplots()\n",
     "# get an alias to the lsst.afw.display.Display() method\n",
-    "display = afwDisplay.Display()\n",
+    "display = afwDisplay.Display(frame=fig)\n",
     "# set the image stretch algorithm and range\n",
     "display.scale('asinh', 'zscale')\n",
     "# load the image into the display\n",
     "display.mtv(calexp.image)\n",
     "# show the corresponding pyplot figure\n",
-    "plt.show()"
+    "plt.show()\n",
+    "# clean up memory\n",
+    "remove_figure(fig)"
    ]
   },
   {
@@ -260,17 +286,20 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "plt.figure()\n",
+    "fig = plt.figure()\n",
     "# Set the figure's projection to be the WCS of the calexp\n",
     "plt.subplot(projection=WCS(calexp.getWcs().getFitsMetadata()))\n",
     "# Display the calexp image data array using the gray colormap (cmap)\n",
     "#  and use approximately the same min and max scale values as above\n",
-    "plt.imshow(calexp.image.array, cmap='gray', vmin=-200.0, vmax=400, origin='lower')\n",
+    "im = plt.imshow(calexp.image.array, cmap='gray', vmin=-200.0, vmax=400, origin='lower')\n",
     "# Add solid white grid lines\n",
     "plt.grid(color='white', ls='solid')\n",
     "# Label axes\n",
     "plt.xlabel('Right Ascension')\n",
-    "plt.ylabel('Declination')"
+    "plt.ylabel('Declination')\n",
+    "plt.show()\n",
+    "# Clean up memory\n",
+    "remove_figure(fig)"
    ]
   },
   {
@@ -298,7 +327,9 @@
     "display2.mtv(calexp.mask)\n",
     "\n",
     "plt.tight_layout()\n",
-    "plt.show()"
+    "plt.show()\n",
+    "\n",
+    "remove_figure(fig)"
    ]
   },
   {
@@ -314,11 +345,12 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "fig = plt.figure()\n",
-    "display = afwDisplay.Display()\n",
+    "fig,ax = plt.subplots()\n",
+    "display = afwDisplay.Display(frame=fig)\n",
     "display.scale('linear', 'zscale')\n",
     "display.mtv(calexp.maskedImage)\n",
-    "plt.show()"
+    "plt.show()\n",
+    "remove_figure(fig)"
    ]
   },
   {
@@ -355,9 +387,9 @@
    "outputs": [],
    "source": [
     "# create a matplotlib.pyplot figure\n",
-    "plt.figure()\n",
+    "fig,ax = plt.subplots()\n",
     "# get an alias to the lsst.afw.display.Display() method\n",
-    "afw_display = afwDisplay.Display()\n",
+    "afw_display = afwDisplay.Display(frame=fig)\n",
     "# set the image stretch algorithm and range\n",
     "afw_display.scale('asinh', 'zscale')\n",
     "# set the transparency of the mask plane (1 = opaque)\n",
@@ -367,7 +399,9 @@
     "# load the image and mask plane into the display\n",
     "afw_display.mtv(calexp)\n",
     "# show the corresponding pyplot figure\n",
-    "plt.show()"
+    "plt.show()\n",
+    "# clean up memory\n",
+    "remove_figure(fig)"
    ]
   },
   {
@@ -544,14 +578,16 @@
    "outputs": [],
    "source": [
     "# create a matplotlib.pyplot figure\n",
-    "plt.figure()\n",
+    "fig,ax = plt.subplots()\n",
     "# get an alias to the lsst.afw.display.Display() method\n",
-    "display = afwDisplay.Display()\n",
+    "display = afwDisplay.Display(frame=fig)\n",
     "# set the image stretch algorithm and range\n",
     "display.scale('asinh', 'zscale')\n",
     "# load the image into the display\n",
     "display.mtv(coadd_calexp.image)\n",
-    "plt.show()"
+    "plt.show()\n",
+    "# clean up memory\n",
+    "remove_figure(fig)"
    ]
   },
   {
@@ -616,11 +652,13 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "fig = plt.figure()\n",
-    "afw_display = afwDisplay.Display(1)\n",
+    "fig,ax = plt.subplots()\n",
+    "afw_display = afwDisplay.Display(frame=fig)\n",
     "afw_display.scale('asinh', 'zscale')\n",
     "afw_display.mtv(cutout_image.image)\n",
-    "plt.gca().axis('off')"
+    "ax.axis('off')\n",
+    "plt.show()\n",
+    "remove_figure(fig)"
    ]
   },
   {
@@ -766,7 +804,10 @@
     "\n",
     "ax[0].set_axis_off()\n",
     "ax[1].set_axis_off()\n",
-    "plt.show()"
+    "plt.show()\n",
+    "\n",
+    "# clean up memory\n",
+    "remove_figure(fig)"
    ]
   },
   {
@@ -789,13 +830,6 @@
     "* [afw.display GitHub website](https://github.com/RobertLuptonTheGood/afw/tree/master/python/lsst/afw/display)  \n",
     "* [Getting Started on Image Display (pipelines.lsst.io)](https://pipelines.lsst.io/getting-started/display.html)"
    ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {},
-   "outputs": [],
-   "source": []
   }
  ],
  "metadata": {