simple fastplotlib example

demos/fastplotlib/README.md

@@ -0,0 +1,22 @@
+# fastplotlib demo
+
+This demo consists of a **generator** `actor` that generates random frames of size 512 * 512 that are sent via a queue to a **processor** `actor` that can be used to process the frames and send them via `zmq`. The `fastplotlib.ipynb` notebook then receives the most recent frame via `zmq` and displays it using [`fastplotlib`](https://github.com/kushalkolar/fastplotlib/).
+
+Usage:
+
+```bash
+# cd to this dir
+cd .../improv/demos/fastplotlib
+
+# start improv
+improv run ./fastplotlib.yaml
+
+# call `setup` in the improv TUI
+setup
+
+# Run the cells in the jupyter notebook until you receive
+# the dark blue square in the plot
+
+# once the plot is ready call `run` in the improv TUI
+run
+```

demos/fastplotlib/actors/sample_generator.py

@@ -0,0 +1,47 @@
+from improv.actor import Actor, RunManager
+from datetime import date #used for saving
+import numpy as np
+import time
+import logging; logger = logging.getLogger(__name__)
+logger.setLevel(logging.INFO)
+
+
+class Generator(Actor):
+    """
+    Generate data and puts it in the queue for the processor to take
+    """
+
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        self.data = None
+        self.name = "Generator"
+        self.frame_index = 0
+
+    def __str__(self):
+        return f"Name: {self.name}, Data: {self.data}"
+
+    def setup(self):
+        logger.info('Completed setup for Generator')
+
+    def stop(self):
+        print("Generator stopping")
+        return 0
+
+    def runStep(self):
+        """
+        Generates 512 x 512 frame and puts it in the queue for the processor
+        """
+
+        data = np.random.randint(0, 255, size=(512 * 512), dtype=np.uint16).reshape(512, 512)
+
+        frame_ix = np.array([self.frame_index], dtype=np.uint32)
+
+        # there must be a better way to do this
+        out = np.concatenate(
+            [data.ravel(), frame_ix],
+            dtype=np.uint32
+        )
+
+        self.q_out.put(out)
+
+        self.frame_index += 1

demos/fastplotlib/actors/sample_processor.py

@@ -0,0 +1,56 @@
+from improv.actor import Actor, RunManager
+import numpy as np
+from queue import Empty
+import logging; logger = logging.getLogger(__name__)
+import zmq
+logger.setLevel(logging.INFO)
+
+
+class Processor(Actor):
+    """
+    Process data and send it through zmq to be be visualized
+    """
+
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+
+    def setup(self):
+        """
+        Creates and binds the socket for zmq
+        """
+
+        self.name = "Processor"
+
+        context = zmq.Context()
+        self.socket = context.socket(zmq.PUB)
+        self.socket.bind("tcp://127.0.0.1:5555")
+
+        self.frame_index = 0
+
+        logger.info('Completed setup for Processor')
+
+    def stop(self):
+        logger.info("Processor stopping")
+        return 0
+
+    def runStep(self):
+        """
+        Gets the frame from the queue, take the mean, sends a memoryview
+        so the zmq subscriber can get the buffer to update the plot
+        """
+
+        frame = None
+
+        try:
+            frame = self.q_in.get(timeout=0.05)
+        except Empty:
+            pass
+        except:
+            logger.error("Could not get frame!")
+
+        if frame is not None:
+            self.frame_index += 1
+            # do some processing
+            frame.mean()
+            # send the buffer
+            self.socket.send(frame)

demos/fastplotlib/fastplotlib.ipynb

@@ -0,0 +1,208 @@
+{
+ "cells": [
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "275928f7-8ed1-496a-b841-c8625d755874",
+   "metadata": {
+    "tags": []
+   },
+   "outputs": [],
+   "source": [
+    "import zmq\n",
+    "import numpy as np\n",
+    "from fastplotlib import Plot"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "da757cf2-5de0-426f-a915-434244bbd970",
+   "metadata": {},
+   "source": [
+    "### Setup zmq subscriber client"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "6d4b4cad-175c-4de7-93fd-a78c00c74ab1",
+   "metadata": {
+    "tags": []
+   },
+   "outputs": [],
+   "source": [
+    "context = zmq.Context()\n",
+    "sub = context.socket(zmq.SUB)\n",
+    "sub.setsockopt(zmq.SUBSCRIBE, b\"\")\n",
+    "\n",
+    "# keep only the most recent message\n",
+    "sub.setsockopt(zmq.CONFLATE, 1)\n",
+    "\n",
+    "# address must match publisher in Processor actor\n",
+    "sub.connect(\"tcp://127.0.0.1:5555\")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "535c577e-07e3-4862-951d-3e35ee045df4",
+   "metadata": {},
+   "source": [
+    "for testing things, benchmark zmq"
+   ]
+  },
+  {
+   "cell_type": "raw",
+   "id": "df251a03-a0fa-4b84-b0f4-25617d90fcc9",
+   "metadata": {
+    "tags": []
+   },
+   "source": [
+    "%%timeit\n",
+    "try:\n",
+    "    a = sub.recv(zmq.NOBLOCK)\n",
+    "except zmq.Again:\n",
+    "    pass"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "e2d059cb-a7e5-46f0-8586-9fd39d6dd6fc",
+   "metadata": {
+    "tags": []
+   },
+   "outputs": [],
+   "source": [
+    "def get_buffer():\n",
+    "    \"\"\"\n",
+    "    Gets the buffer from the publisher\n",
+    "    \"\"\"\n",
+    "    try:\n",
+    "        b = sub.recv(zmq.NOBLOCK)\n",
+    "    except zmq.Again:\n",
+    "        pass\n",
+    "    else:\n",
+    "        return b\n",
+    "    \n",
+    "    return None"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "5b3f89a1-1a5d-45f2-8af7-56e55e2a0143",
+   "metadata": {},
+   "source": [
+    "### Live plot that updates using the most recent message :D "
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "6bf49577-06be-44ec-93b1-ad50b3d4d794",
+   "metadata": {
+    "tags": []
+   },
+   "outputs": [],
+   "source": [
+    "plot = Plot()\n",
+    "\n",
+    "# initialize image graphic with zeros\n",
+    "plot.add_image(\n",
+    "    np.zeros((512, 512), dtype=np.uint16),\n",
+    "    vmin=0,\n",
+    "    vmax=255,\n",
+    "    name=\"img\"\n",
+    ")\n",
+    "\n",
+    "def update_frame(p):\n",
+    "    # recieve memory with buffer\n",
+    "    buff = get_buffer()\n",
+    "    \n",
+    "    if buff is not None:\n",
+    "        # numpy array from buffer\n",
+    "        a = np.frombuffer(buff, dtype=np.uint32)\n",
+    "        ix = a[-1]\n",
+    "        # set graphic data\n",
+    "        p[\"img\"].data = a[:-1].reshape(512, 512)\n",
+    "        p.set_title(f\"frame: {ix}\")\n",
+    "\n",
+    "plot.add_animations(update_frame)\n",
+    "\n",
+    "plot.show()"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "bbb80338-49c3-4233-a98c-4c7f0e88825a",
+   "metadata": {},
+   "source": [
+    "## **fastplotlib is non blocking!**"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "e25aadce-bd4e-4597-85ca-b46d151009d3",
+   "metadata": {
+    "tags": []
+   },
+   "outputs": [],
+   "source": [
+    "plot.canvas.get_stats()"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "eae5cd5b-2657-40c9-993b-26052439235e",
+   "metadata": {
+    "tags": []
+   },
+   "outputs": [],
+   "source": [
+    "plot[\"img\"].cmap = \"viridis\""
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "b4fb841c-0634-4227-9a3b-aef32a1a3b45",
+   "metadata": {
+    "tags": []
+   },
+   "outputs": [],
+   "source": [
+    "plot[\"img\"].vmax=300"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "519f3a42-e684-4ca3-b301-8ef447ffd805",
+   "metadata": {},
+   "outputs": [],
+   "source": []
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "Python 3 (ipykernel)",
+   "language": "python",
+   "name": "python3"
+  },
+  "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.5"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 5
+}

demos/fastplotlib/fastplotlib.yaml

@@ -0,0 +1,12 @@
+actors:
+  Generator:
+    package: actors.sample_generator
+    class: Generator
+
+  Processor:
+    package: actors.sample_processor
+    class: Processor
+
+connections:
+  Generator.q_out: [Processor.q_in]
+