Add in-memory cache for interpolation matrices (#23)

* Add in-memory cache for interpolation matrices
ecmwf · Nov 27, 2024 · b4e3825 · b4e3825
1 parent 1b91b2f
commit b4e3825
Show file tree

Hide file tree

Showing 19 changed files with 1,289 additions and 4 deletions.
diff --git a/docs/conf.py b/docs/conf.py
@@ -33,10 +33,16 @@
 extensions = [
     "sphinx_rtd_theme",
     "nbsphinx",
+    "sphinx_issues",
     "earthkit.regrid.sphinxext.xref",
     "earthkit.regrid.sphinxext.module_output",
 ]
 
+
+# Path to GitHub repo {group}/{project}  (note that `group` is the GitHub user or organization)
+issues_github_path = "ecmwf/earthkit-regrid"
+
+
 # Add any paths that contain templates here, relative to this directory.
 templates_path = ["_templates"]
 

diff --git a/docs/examples/index.rst b/docs/examples/index.rst
@@ -11,3 +11,4 @@ Here is a list of example notebooks to illustrate how to use earthkit-regrid.
     numpy_arrays.ipynb
     healpix_fieldlist.ipynb
     octahedral_fieldlist.ipynb
+    memory_cache.ipynb
diff --git a/docs/examples/memory_cache.ipynb b/docs/examples/memory_cache.ipynb
@@ -0,0 +1,332 @@
+{
+ "cells": [
+  {
+   "cell_type": "markdown",
+   "id": "a57bee67-d094-4ec8-9175-33465f46e1d0",
+   "metadata": {
+    "editable": true,
+    "slideshow": {
+     "slide_type": ""
+    },
+    "tags": []
+   },
+   "source": [
+    "## Matrix memory cache"
+   ]
+  },
+  {
+   "cell_type": "raw",
+   "id": "ac1d768f-b386-4229-90c9-a5cc7c550643",
+   "metadata": {
+    "editable": true,
+    "raw_mimetype": "text/restructuredtext",
+    "slideshow": {
+     "slide_type": ""
+    },
+    "tags": []
+   },
+   "source": [
+    "This notebook demonstrates the usage of the in-memory interpolation matrix cache. By default this cache is enabled, for details see :ref:`memory_cache`."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 1,
+   "id": "3f3c0324-708c-4ea9-8eac-406cabc0a1e8",
+   "metadata": {
+    "editable": true,
+    "slideshow": {
+     "slide_type": ""
+    },
+    "tags": []
+   },
+   "outputs": [],
+   "source": [
+    "import earthkit.regrid as ekr\n",
+    "import numpy as np\n",
+    "\n",
+    "# create input data array for an O1280 grid\n",
+    "in_data = np.ones(6599680)\n",
+    "\n",
+    "# helper method for interpolation\n",
+    "def _run(n=10):\n",
+    "    for _ in range(n):\n",
+    "        res_array = ekr.interpolate(in_data, \n",
+    "                                {\"grid\": \"O1280\"}, \n",
+    "                                {\"grid\": [0.1,0.1]})"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "2aeaea49-cb22-42b0-8e34-65f2abf90ce3",
+   "metadata": {
+    "editable": true,
+    "slideshow": {
+     "slide_type": ""
+    },
+    "tags": []
+   },
+   "source": [
+    "#### The \"off\" cache policy"
+   ]
+  },
+  {
+   "cell_type": "raw",
+   "id": "4c2772b6-de75-4ec4-ae7d-5ff78a49a51c",
+   "metadata": {
+    "editable": true,
+    "raw_mimetype": "text/restructuredtext",
+    "slideshow": {
+     "slide_type": ""
+    },
+    "tags": []
+   },
+   "source": [
+    "We can control the matrix in-memory cache via :func:`set_memory_cache`. In this example we turn off the cache."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 2,
+   "id": "5b0734b7-272e-47a5-859f-2af38da13cc3",
+   "metadata": {
+    "editable": true,
+    "slideshow": {
+     "slide_type": ""
+    },
+    "tags": []
+   },
+   "outputs": [],
+   "source": [
+    "ekr.set_memory_cache(policy=\"off\")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 3,
+   "id": "22333b87-6005-43ec-9e86-d39833e7de35",
+   "metadata": {
+    "editable": true,
+    "slideshow": {
+     "slide_type": ""
+    },
+    "tags": []
+   },
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "CPU times: user 6.26 s, sys: 712 ms, total: 6.97 s\n",
+      "Wall time: 7.47 s\n"
+     ]
+    }
+   ],
+   "source": [
+    "%time _run()"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 4,
+   "id": "44c806ab-88f0-4364-a1c9-8d8a217a7295",
+   "metadata": {
+    "editable": true,
+    "slideshow": {
+     "slide_type": ""
+    },
+    "tags": []
+   },
+   "outputs": [
+    {
+     "data": {
+      "text/plain": [
+       "CacheInfo(hits=0, misses=0, maxsize=0, currsize=0, count=0, policy='off')"
+      ]
+     },
+     "execution_count": 4,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "ekr.memory_cache_info()"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "3485ddee-56fa-4346-8946-669286ef954d",
+   "metadata": {
+    "editable": true,
+    "slideshow": {
+     "slide_type": ""
+    },
+    "tags": []
+   },
+   "source": [
+    "#### The \"largest\" (default) cache policy"
+   ]
+  },
+  {
+   "cell_type": "raw",
+   "id": "e06d0739-ef6f-4a4e-b2de-7b141fcdb0c9",
+   "metadata": {
+    "editable": true,
+    "raw_mimetype": "text/restructuredtext",
+    "slideshow": {
+     "slide_type": ""
+    },
+    "tags": []
+   },
+   "source": [
+    "In this example we will use the default memory cache policy, which is \"largest\". For details see: :func:`set_memory_cache`. Notice the x9 speed-up we achieved with the caching."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 5,
+   "id": "7c8d2977-befc-47aa-a894-aab195aaf4db",
+   "metadata": {
+    "editable": true,
+    "slideshow": {
+     "slide_type": ""
+    },
+    "tags": []
+   },
+   "outputs": [],
+   "source": [
+    "ekr.set_memory_cache(policy=\"largest\")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 6,
+   "id": "74152589-4085-4c73-baef-d0c109f5e158",
+   "metadata": {
+    "editable": true,
+    "slideshow": {
+     "slide_type": ""
+    },
+    "tags": []
+   },
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "CPU times: user 750 ms, sys: 91.6 ms, total: 842 ms\n",
+      "Wall time: 847 ms\n"
+     ]
+    }
+   ],
+   "source": [
+    "%time _run()"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 7,
+   "id": "59137910-dfb2-4334-beb0-f5590b56d47e",
+   "metadata": {
+    "editable": true,
+    "slideshow": {
+     "slide_type": ""
+    },
+    "tags": []
+   },
+   "outputs": [
+    {
+     "data": {
+      "text/plain": [
+       "CacheInfo(hits=9, misses=1, maxsize=314572800, currsize=259170724, count=1, policy='largest')"
+      ]
+     },
+     "execution_count": 7,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "ekr.memory_cache_info()"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "a857e482-d056-4b21-b66d-0b08ff29720a",
+   "metadata": {
+    "editable": true,
+    "slideshow": {
+     "slide_type": ""
+    },
+    "tags": []
+   },
+   "source": [
+    "#### Clearing the cache"
+   ]
+  },
+  {
+   "cell_type": "raw",
+   "id": "4abb400b-f5a7-4dad-a802-a0a1ccc05e4b",
+   "metadata": {
+    "editable": true,
+    "raw_mimetype": "text/restructuredtext",
+    "slideshow": {
+     "slide_type": ""
+    },
+    "tags": []
+   },
+   "source": [
+    "We can clear the cache with :func:`clear_memory_cache`."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 8,
+   "id": "767f7420-0a80-4c3e-ac33-d8dfcd5c4cf7",
+   "metadata": {
+    "editable": true,
+    "slideshow": {
+     "slide_type": ""
+    },
+    "tags": []
+   },
+   "outputs": [
+    {
+     "data": {
+      "text/plain": [
+       "CacheInfo(hits=0, misses=0, maxsize=314572800, currsize=0, count=0, policy='largest')"
+      ]
+     },
+     "execution_count": 8,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "ekr.clear_memory_cache()\n",
+    "ekr.memory_cache_info()"
+   ]
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "dev_ecc",
+   "language": "python",
+   "name": "dev_ecc"
+  },
+  "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.10.13"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 5
+}
diff --git a/docs/index.rst b/docs/index.rst
@@ -23,6 +23,7 @@ The **earthkit-regrid** API features the :func:`interpolate` function to regrid
    :caption: Documentation
 
    interpolate
+   memory_cache
    gridspec
    inventory/index
 

diff --git a/docs/interpolate.rst b/docs/interpolate.rst
@@ -1,4 +1,4 @@
-interpolate
+Interpolation
 ==============
 
 .. py:function:: interpolate(values, in_grid=None, out_grid=None, matrix_source=None, method='linear')