{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "eeb9ab4e",
   "metadata": {},
   "source": [
    "# The shape of the equations\n",
    "\n",
    "Matching and BLT sorting are described in the literature as graph algorithms, and\n",
    "implemented here as ones — but what they *do* is permute a matrix until it is lower\n",
    "triangular. That is a picture, so this notebook draws it.\n",
    "\n",
    "Three stages get looked at: the incidence matrix a flat model comes out with, what alias\n",
    "elimination deletes from it, and the staircase that causalization leaves behind. The last\n",
    "cell is a dropdown over every example model in the repository, and it needs a live kernel. The\n",
    "file itself is attached to this page — [structure.ipynb](structure.ipynb) — but it reads those\n",
    "models from `examples/`, so it wants a checkout rather than a bare download:\n",
    "\n",
    "```bash\n",
    "git clone https://gitlab.com/jorgeecardona/pymodelica && cd pymodelica\n",
    "pip install \"modelica[sim]\" matplotlib ipywidgets\n",
    "jupyter lab docs/notebooks/structure.ipynb\n",
    "```"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "11873e27",
   "metadata": {},
   "outputs": [],
   "source": [
    "from pathlib import Path\n",
    "\n",
    "import matplotlib as mpl\n",
    "import matplotlib.pyplot as plt\n",
    "import numpy as np\n",
    "from matplotlib.colors import ListedColormap\n",
    "\n",
    "# The models come from `examples/` in the repository rather than from strings in here: a copy\n",
    "# pasted into a notebook is a copy that goes stale, and the reason this page is executed when the\n",
    "# site is built is that it then cannot show output from a model the project does not have.\n",
    "root = next(\n",
    "    (p for p in [Path.cwd(), *Path.cwd().parents] if (p / \"examples\" / \"models\").is_dir()), None\n",
    ")\n",
    "if root is None:\n",
    "    message = (\n",
    "        \"This notebook reads its models from `examples/` in the pymodelica repository. \"\n",
    "        \"Run it from a checkout:\\n\"\n",
    "        \"    git clone https://gitlab.com/jorgeecardona/pymodelica\\n\"\n",
    "        \"    cd pymodelica && jupyter lab docs/notebooks/\"\n",
    "    )\n",
    "    raise RuntimeError(message)\n",
    "\n",
    "SURFACE, INK, MUTED, LINE = \"#fcfcfb\", \"#0b0b0b\", \"#52514e\", \"#d4d3cd\"\n",
    "BLUE, ORANGE, AQUA = \"#2a78d6\", \"#eb6834\", \"#1baf7a\"\n",
    "\n",
    "mpl.rcParams.update(\n",
    "    {\n",
    "        \"figure.figsize\": (8.4, 3.6),\n",
    "        \"figure.dpi\": 110,\n",
    "        \"figure.facecolor\": SURFACE,\n",
    "        \"axes.facecolor\": SURFACE,\n",
    "        \"savefig.facecolor\": SURFACE,\n",
    "        \"axes.edgecolor\": LINE,\n",
    "        \"axes.labelcolor\": MUTED,\n",
    "        \"axes.spines.top\": False,\n",
    "        \"axes.spines.right\": False,\n",
    "        \"axes.grid\": False,\n",
    "        \"axes.titlelocation\": \"left\",\n",
    "        \"axes.titlepad\": 10,\n",
    "        \"axes.titlesize\": 11,\n",
    "        \"lines.linewidth\": 2.0,\n",
    "        \"legend.frameon\": False,\n",
    "        \"text.color\": INK,\n",
    "        \"xtick.color\": MUTED,\n",
    "        \"ytick.color\": MUTED,\n",
    "        \"font.size\": 10,\n",
    "    }\n",
    ")\n",
    "%config InlineBackend.figure_formats = [\"svg\"]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "83981105",
   "metadata": {},
   "source": [
    "## A model big enough to have a shape\n",
    "\n",
    "The example models are small on purpose, and a 5×5 matrix has no texture to look at. So\n",
    "this one is generated: an RC ladder of *n* stages, written as Modelica text and parsed like\n",
    "any other source. The component library it draws on is the same\n",
    "`examples/library/Electrical.mo` the hand-written examples use."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c9b42815",
   "metadata": {},
   "outputs": [],
   "source": [
    "from modelica.lang import parse, unparse_equation\n",
    "\n",
    "electrical = parse((root / \"examples\" / \"library\" / \"Electrical.mo\").read_text())\n",
    "\n",
    "\n",
    "def ladder(stages: int) -> str:\n",
    "    \"\"\"Modelica source for a resistor–capacitor ladder of `stages` stages.\"\"\"\n",
    "    head = [f'model Ladder \"An RC ladder of {stages} stages\"']\n",
    "    head += [\"  Electrical.ConstantVoltage source(V = 1);\", \"  Electrical.Ground ground;\"]\n",
    "    for i in range(stages):\n",
    "        head += [\n",
    "            f\"  Electrical.Resistor r{i}(R = 1);\",\n",
    "            f\"  Electrical.Capacitor c{i}(C = 1, v(start = 0, fixed = true));\",\n",
    "        ]\n",
    "    body = [\"equation\", \"  connect(source.p, r0.p);\"]\n",
    "    for i in range(stages):\n",
    "        body += [f\"  connect(r{i}.n, c{i}.p);\", f\"  connect(c{i}.n, source.n);\"]\n",
    "        if i + 1 < stages:\n",
    "            body += [f\"  connect(c{i}.p, r{i + 1}.p);\"]\n",
    "    body += [\"  connect(source.n, ground.p);\", \"end Ladder;\"]\n",
    "    return \"\\n\".join([*head, *body]) + \"\\n\"\n",
    "\n",
    "\n",
    "print(ladder(2))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "17ccc949",
   "metadata": {},
   "source": [
    "## 1. What flattening hands over\n",
    "\n",
    "The incidence matrix has one row per equation and one column per unknown, and a mark where\n",
    "the equation mentions the unknown. Nothing in it says which equation *solves for* which\n",
    "unknown — Modelica's `=` is acausal, and this matrix is exactly that ignorance written\n",
    "down."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7f0fda7e",
   "metadata": {},
   "outputs": [],
   "source": [
    "from modelica.ir.rewrite import names\n",
    "from modelica.sim import System, causalize, flatten\n",
    "\n",
    "\n",
    "def flat_of(stages: int):\n",
    "    return flatten([electrical, parse(ladder(stages))], \"Ladder\")\n",
    "\n",
    "\n",
    "flat = flat_of(6)\n",
    "unknowns = [v.name for v in flat.unknowns]\n",
    "column = {name: i for i, name in enumerate(unknowns)}\n",
    "\n",
    "incidence = np.zeros((len(flat.equations), len(unknowns)), dtype=int)\n",
    "for row, equation in enumerate(flat.equations):\n",
    "    for name in {*names(equation.lhs), *names(equation.rhs)}:\n",
    "        if name in column:\n",
    "            incidence[row, column[name]] = 1\n",
    "\n",
    "density = incidence.sum() / incidence.size\n",
    "print(f\"{incidence.shape[0]} equations × {incidence.shape[1]} unknowns, {density:.1%} filled\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fd5af658",
   "metadata": {},
   "source": [
    "Drawn, with the marks in blue and the empty entries left as surface:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c3187ce4",
   "metadata": {},
   "outputs": [],
   "source": [
    "CELLS = ListedColormap([SURFACE, BLUE])\n",
    "\n",
    "\n",
    "def spy(ax, matrix, title: str, cmap=CELLS) -> None:\n",
    "    \"\"\"One binary matrix, equations down, unknowns across.\"\"\"\n",
    "    ax.imshow(matrix, cmap=cmap, interpolation=\"nearest\", aspect=\"equal\", vmin=0, vmax=cmap.N - 1)\n",
    "    ax.set_title(title)\n",
    "    ax.set_xlabel(\"unknowns →\")\n",
    "    ax.set_ylabel(\"← equations\")\n",
    "    ax.set_xticks([])\n",
    "    ax.set_yticks([])\n",
    "    for side in (\"top\", \"right\"):\n",
    "        ax.spines[side].set_visible(True)\n",
    "    for spine in ax.spines.values():\n",
    "        spine.set_color(LINE)\n",
    "\n",
    "\n",
    "fig, ax = plt.subplots(figsize=(4.6, 4.6))\n",
    "spy(ax, incidence, f\"Ladder(6), as flattened — {incidence.sum()} marks\")\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d30ec1a3",
   "metadata": {},
   "source": [
    "## 2. Most of it is `a = b`\n",
    "\n",
    "The band down the diagonal and the scatter around it are two different things. A\n",
    "`connect(a, b)` between potential variables becomes `a.v = b.v`, and a large fraction of\n",
    "what a flattener emits is exactly that: equations whose only content is that two names are\n",
    "the same variable. Alias elimination spends them, and the system keeps a map from every name\n",
    "it removed to what replaced it — so a user can still ask for `r3.p.v` afterwards.\n",
    "\n",
    "In orange, the columns whose variable did not survive:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7543e2bb",
   "metadata": {
    "lines_to_next_cell": 1
   },
   "outputs": [],
   "source": [
    "from modelica.ir import nodes as ir\n",
    "\n",
    "system = causalize(flat)\n",
    "survives = np.array([name not in system.aliases for name in unknowns])\n",
    "\n",
    "kept = np.zeros_like(incidence)\n",
    "kept[(incidence == 1) & survives] = 1\n",
    "kept[(incidence == 1) & ~survives] = 2\n",
    "\n",
    "fig, ax = plt.subplots(figsize=(4.6, 4.6))\n",
    "spy(\n",
    "    ax,\n",
    "    kept,\n",
    "    \"Blue survives causalization, orange is aliased away\",\n",
    "    ListedColormap([SURFACE, BLUE, ORANGE]),\n",
    ")\n",
    "plt.show()\n",
    "\n",
    "trivial = sum(1 for e in flat.equations if isinstance(e.lhs, ir.Ref) and isinstance(e.rhs, ir.Ref))\n",
    "computed = len(unknowns) - len(system.aliases) - len(system.states)\n",
    "print(f\"{trivial} of {len(flat.equations)} equations are literally `a = b`\\n\")\n",
    "print(f\"{len(unknowns)} unknowns in the flat model:\")\n",
    "print(f\"  {len(system.aliases):3} eliminated as aliases\")\n",
    "print(f\"  {len(system.states):3} became states — integrated, not computed\")\n",
    "print(f\"  {computed:3} are computed by the {len(system.blocks)} blocks\")\n",
    "print(\n",
    "    f\"\\n(the blocks also compute {len(system.derivatives)} derivative variables the flat model never had)\"\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "43f47487",
   "metadata": {},
   "source": [
    "## 3. The staircase\n",
    "\n",
    "What causalization returns is a sequence of **blocks**: each one is a set of equations and\n",
    "the same number of unknowns, sorted so that a block only depends on blocks before it. Drawn\n",
    "with the blocks in evaluation order, that ordering *is* the lower-triangular shape — every\n",
    "mark above the diagonal would be a dependency on something not computed yet.\n",
    "\n",
    "Marks on the diagonal are the unknown a block solves for; marks below it are the values it\n",
    "reads. The diagonal is orange to separate the two."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7006d715",
   "metadata": {},
   "outputs": [],
   "source": [
    "def staircase(system: System) -> np.ndarray:\n",
    "    \"\"\"Block-ordered incidence: 2 on the diagonal a block solves, 1 where it reads.\"\"\"\n",
    "    order = [u for block in system.blocks for u in block.unknowns]\n",
    "    at = {name: i for i, name in enumerate(order)}\n",
    "    matrix = np.zeros((len(order), len(order)), dtype=int)\n",
    "    row = 0\n",
    "    for block in system.blocks:\n",
    "        for equation in block.equations:\n",
    "            for name in {*names(equation.lhs), *names(equation.rhs)}:\n",
    "                if name in at:\n",
    "                    matrix[row, at[name]] = 1\n",
    "            for name in block.unknowns:\n",
    "                matrix[row, at[name]] = 2\n",
    "            row += 1\n",
    "    return matrix\n",
    "\n",
    "\n",
    "sorted_blocks = staircase(system)\n",
    "\n",
    "fig, ax = plt.subplots(figsize=(4.6, 4.6))\n",
    "spy(\n",
    "    ax,\n",
    "    sorted_blocks,\n",
    "    f\"Ladder(6) causalized — {len(system.blocks)} blocks\",\n",
    "    ListedColormap([SURFACE, BLUE, ORANGE]),\n",
    ")\n",
    "plt.show()\n",
    "\n",
    "print(\n",
    "    f\"largest block: {system.largest_block}   explicit: {system.is_explicit}   states: {system.states}\"\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b8ad57e6",
   "metadata": {},
   "source": [
    "`is_explicit` is the claim that matters: every block is 1×1 and carries an assignment, so\n",
    "evaluating the right-hand side is straight-line arithmetic with no inner solve anywhere.\n",
    "The generated code for this model has no `while` in it.\n",
    "\n",
    "## How the three counts scale\n",
    "\n",
    "The interesting number is not how many equations flattening produces — it is how few\n",
    "survive. For the ladder, the equations grow at twelve per stage and the blocks at four,\n",
    "because the rest were connect statements:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fad7619b",
   "metadata": {},
   "outputs": [],
   "source": [
    "rows = []\n",
    "for n in range(1, 11):\n",
    "    model = flat_of(n)\n",
    "    causalized = causalize(model)\n",
    "    rows.append((n, len(model.equations), len(causalized.blocks), len(causalized.states)))\n",
    "\n",
    "stages, equations, blocks, states = (np.array(values) for values in zip(*rows, strict=True))\n",
    "\n",
    "fig, ax = plt.subplots(figsize=(6.2, 3.4))\n",
    "for series, colour, label in (\n",
    "    (equations, BLUE, \"flat equations\"),\n",
    "    (blocks, ORANGE, \"blocks after sorting\"),\n",
    "    (states, AQUA, \"states\"),\n",
    "):\n",
    "    ax.plot(stages, series, color=colour, marker=\"o\", markersize=5)\n",
    "    ax.annotate(\n",
    "        label,\n",
    "        xy=(stages[-1], series[-1]),\n",
    "        xytext=(6, 0),\n",
    "        textcoords=\"offset points\",\n",
    "        color=colour,\n",
    "        fontsize=9,\n",
    "        va=\"center\",\n",
    "    )\n",
    "\n",
    "ax.set_xlim(0.6, 12.8)\n",
    "ax.set_xticks(range(1, 11))\n",
    "ax.set_title(\"An RC ladder, stage by stage\")\n",
    "ax.set_xlabel(\"stages\")\n",
    "ax.set_ylabel(\"count\")\n",
    "ax.grid(visible=True, color=\"#e8e7e2\", linewidth=0.8)\n",
    "ax.set_axisbelow(True)\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "62564b2f",
   "metadata": {},
   "source": [
    "## When the staircase cannot be built: the pendulum\n",
    "\n",
    "Everything above sorts completely because every block came out 1×1. The pendulum does not.\n",
    "In Cartesian coordinates its constraint `x² + y² = L²` involves no derivative at all, so no\n",
    "amount of permuting produces a triangular form — the matrix is structurally singular until\n",
    "index reduction differentiates the constraint and adds the result to the system."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4c4a7c1e",
   "metadata": {},
   "outputs": [],
   "source": [
    "pendulum = causalize(\n",
    "    flatten([parse((root / \"examples\" / \"models\" / \"Pendulum.mo\").read_text())], \"Pendulum\")\n",
    ")\n",
    "\n",
    "print(\n",
    "    f\"index: {pendulum.index}   blocks: {len(pendulum.blocks)}   largest: {pendulum.largest_block}   explicit: {pendulum.is_explicit}\"\n",
    ")\n",
    "for block in pendulum.blocks:\n",
    "    kind = \"assignment\" if block.is_linear_scalar else f\"{len(block)}×{len(block)} solve\"\n",
    "    print(f\"\\n  {kind}: {', '.join(block.unknowns)}\")\n",
    "    for equation in block.equations:\n",
    "        print(f\"    {unparse_equation(equation)}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "abfaf4e3",
   "metadata": {},
   "source": [
    "Read that from the top. The first block solves `x` out of `x² + y² = L²` — 1×1, and still\n",
    "no assignment, because the equation is quadratic in `x` and there is no expression to write\n",
    "down; that block is a scalar Newton iteration at run time. The second gets `vx` from the\n",
    "once-differentiated constraint, which *is* linear in `vx`, so it is an assignment.\n",
    "\n",
    "The last one is the interesting one: `F` and both accelerations, three unknowns determined\n",
    "together by three equations, one of them the twice-differentiated constraint. That is a\n",
    "linear solve inside every evaluation of the right-hand side, and no permutation removes it —\n",
    "it is a property of the model, which is what `largest_block` reports.\n",
    "\n",
    "`pendulum.index` is 3, because the constraint had to be differentiated twice before the\n",
    "accelerations appeared at all. That is also why the model has two admissible state sets\n",
    "rather than one, which is [the next notebook](../events/)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b072a24d",
   "metadata": {},
   "outputs": [],
   "source": [
    "fig, ax = plt.subplots(figsize=(3.4, 3.4))\n",
    "spy(\n",
    "    ax,\n",
    "    staircase(pendulum),\n",
    "    \"Pendulum — the 3×3 block, bottom right\",\n",
    "    ListedColormap([SURFACE, BLUE, ORANGE]),\n",
    ")\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0e3623c2",
   "metadata": {},
   "source": [
    "## Every example, on a dropdown\n",
    "\n",
    "The same two pictures for anything in `examples/models`. `DrainingTank` and `BouncingBall`\n",
    "are tiny; `RLCCircuit` and `MassSpringDamper` show the alias story at hand-written scale;\n",
    "`Pendulum` is the one that does not sort."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b4cb9db6",
   "metadata": {},
   "outputs": [],
   "source": [
    "LIBRARIES = [\n",
    "    parse(path.read_text()) for path in sorted((root / \"examples\" / \"library\").glob(\"*.mo\"))\n",
    "]\n",
    "MODELS = sorted(path.stem for path in (root / \"examples\" / \"models\").glob(\"*.mo\"))\n",
    "\n",
    "\n",
    "def show(model: str) -> None:\n",
    "    \"\"\"Raw incidence beside the causalized staircase, for one example model.\"\"\"\n",
    "    flat = flatten(\n",
    "        [*LIBRARIES, parse((root / \"examples\" / \"models\" / f\"{model}.mo\").read_text())], model\n",
    "    )\n",
    "    system = causalize(flat)\n",
    "    unknowns = [v.name for v in flat.unknowns]\n",
    "    column = {name: i for i, name in enumerate(unknowns)}\n",
    "    raw = np.zeros((len(flat.equations), len(unknowns)), dtype=int)\n",
    "    for row, equation in enumerate(flat.equations):\n",
    "        for name in {*names(equation.lhs), *names(equation.rhs)}:\n",
    "            if name in column:\n",
    "                raw[row, column[name]] = 1\n",
    "\n",
    "    fig, (left, right) = plt.subplots(1, 2, figsize=(8.4, 4.2))\n",
    "    spy(left, raw, f\"{model}: {raw.shape[0]}×{raw.shape[1]} as flattened\")\n",
    "    spy(\n",
    "        right,\n",
    "        staircase(system),\n",
    "        f\"{len(system.blocks)} blocks, largest {system.largest_block}\",\n",
    "        ListedColormap([SURFACE, BLUE, ORANGE]),\n",
    "    )\n",
    "    plt.tight_layout()\n",
    "    plt.show()\n",
    "\n",
    "    print(f\"index {system.index}   explicit {system.is_explicit}   states {system.states}\")\n",
    "    print(f\"state sets {system.candidates}\")\n",
    "\n",
    "\n",
    "show(\"RLCCircuit\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "316c88ab",
   "metadata": {},
   "outputs": [],
   "source": [
    "import ipywidgets as widgets\n",
    "\n",
    "_ = widgets.interact(\n",
    "    show, model=widgets.Dropdown(options=MODELS, value=\"RLCCircuit\", description=\"model\")\n",
    ")"
   ]
  }
 ],
 "metadata": {
  "jupytext": {
   "cell_metadata_filter": "-all",
   "main_language": "python",
   "notebook_metadata_filter": "-all"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}