docs/notebooks/python/sticky_notes.ipynb

{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "WX_VB4Hn8y78"
      },
      "source": [
        "# Sticky Notes: A mini-DSL for taking notes with symbolic objects.\n",
        "\n",
        "[![Open in Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/google/pyglove/blob/main/docs/notebooks/python/sticky_notes.ipynb)\n",
        "\n",
        "This notebook demostrates how to create domain-specific languages (DSL) in Python with symbolic placeholding and manipulation."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "BddZ9QqKppGC"
      },
      "outputs": [],
      "source": [
        "!pip install pyglove"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "aVmOS0BxpsN0"
      },
      "outputs": [],
      "source": [
        "import pyglove as pg"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "MZuKiFxOpue6"
      },
      "source": [
        "## Developing DSL primitives\n",
        "\n",
        "In this example, we introduce a `StickyNote` class, which can be passed to an arbitrary field within a symbolic object. By subclassing `pg.PureSymbolic`, `StickNote` objects are treated as pure symbolic values, thus delaying the construction of their owning objects. "
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "jrItbsu-4_ae"
      },
      "outputs": [],
      "source": [
        "@pg.members([\n",
        "  ('value', pg.typing.Any()),\n",
        "  ('notes', pg.typing.Str())\n",
        "])\n",
        "class StickyNote(pg.PureSymbolic, pg.Object):\n",
        "  \"\"\"StickyNote is a pure symbolic object that can be used as placeholders.\"\"\"\n",
        "\n",
        "# Transform to replace a node in the tree with a StickyNote object.\n",
        "def note(value, path, text):\n",
        "  return pg.patching.patch_on_path(\n",
        "      value, path, value_fn=lambda x: StickyNote(x, text))\n",
        "\n",
        "# Transform to replace StickyNote object with original value.\n",
        "def denote(value):\n",
        "  \"\"\"Remove notes from an symbolic object.\"\"\"\n",
        "  return pg.patching.patch_on_type(\n",
        "      value, StickyNote, value_fn=lambda x: x.value)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "83yXY6uyqL-2"
      },
      "source": [
        "## Use DSL primitives with any symbolic objects"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "executionInfo": {
          "elapsed": 56,
          "status": "ok",
          "timestamp": 1660956244202,
          "user": {
            "displayName": "",
            "userId": ""
          },
          "user_tz": 420
        },
        "id": "hqual-O6npKK",
        "outputId": "ba1fec28-eb1f-4740-b4d5-31f4fed31827"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "foo(\n",
            "  x = StickyNote(\n",
            "    value = foo(\n",
            "      x = 1,\n",
            "      y = 2\n",
            "    ),\n",
            "    notes = 'this value is copied somewhere.'\n",
            "  ),\n",
            "  y = foo(\n",
            "    x = StickyNote(\n",
            "      value = 3,\n",
            "      notes = 'this value is not copied.'\n",
            "    ),\n",
            "    y = 4\n",
            "  )\n",
            ")\n"
          ]
        }
      ],
      "source": [
        "@pg.symbolize\n",
        "def foo(x, y):\n",
        "  x = x() if callable(x) else x\n",
        "  y = y() if callable(y) else y\n",
        "  return x + y\n",
        "\n",
        "# Create a symbolic object.\n",
        "f = foo(foo(1, 2), foo(3, 4))\n",
        "\n",
        "# Take notes on sub-nodes 'x' and 'y'.\n",
        "note(f, 'x', 'this value is copied somewhere.')\n",
        "note(f, 'y.x', 'this value is not copied.')\n",
        "print(f)\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "executionInfo": {
          "elapsed": 53,
          "status": "ok",
          "timestamp": 1660956244380,
          "user": {
            "displayName": "",
            "userId": ""
          },
          "user_tz": 420
        },
        "id": "6FJDtk9oqbIf",
        "outputId": "bf173c1b-5897-4cc7-8344-bde9f049ddfa"
      },
      "outputs": [
        {
          "data": {
            "text/plain": [
              "{'x': StickyNote(value=foo(x=1, y=2), notes='this value is copied somewhere.'),\n",
              " 'y.x': StickyNote(value=3, notes='this value is not copied.')}"
            ]
          },
          "execution_count": 11,
          "metadata": {},
          "output_type": "execute_result"
        }
      ],
      "source": [
        "# Query noted nodes.\n",
        "pg.query(f, where=lambda v: isinstance(v, StickyNote))"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "executionInfo": {
          "elapsed": 54,
          "status": "ok",
          "timestamp": 1660956244555,
          "user": {
            "displayName": "",
            "userId": ""
          },
          "user_tz": 420
        },
        "id": "m4J2c6gzqXiq",
        "outputId": "b663272c-e963-4e0d-fa7d-650043d2dac6"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "10\n"
          ]
        }
      ],
      "source": [
        "# Remove notes so `f` can be evaluated.\n",
        "print(denote(f)())"
      ]
    }
  ],
  "metadata": {
    "colab": {
      "collapsed_sections": [],
      "last_runtime": {
        "build_target": "//learning/deepmind/dm_python:dm_notebook3",
        "kind": "private"
      },
      "name": "Sticky Notes: A mini-DSL for taking notes with symbolic objects",
      "provenance": []
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    },
    "language_info": {
      "name": "python"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}