{ "cells": [ { "cell_type": "markdown", "id": "0", "metadata": {}, "source": [ "# Theory V: xDSM and problem types\n", "\n", "In this notebook we'll take a look at a standard way to visualize graphs and look at different ways a cycle can emerge." ] }, { "cell_type": "markdown", "id": "1", "metadata": {}, "source": [ "Let's first re-use the 6 models from before. " ] }, { "cell_type": "code", "execution_count": null, "id": "2", "metadata": { "tags": [ "hide-input" ] }, "outputs": [], "source": [ "import inspect\n", "from math import sin, exp\n", "import cottax.assemblies.jax as cx\n", "from cottax.assemblies.jax import flat\n", "from cottax.consumers.visualization.xdsm import PlainFormat\n", "\n", "\n", "def to_graph(*models):\n", " # a plain function writes `_output`; anything already declared is taken as it is\n", " return cx.absolute.declare(*(flat.function(m, m.__name__.lower() + \"_output\")\n", " if inspect.isfunction(m) else m for m in models)).graph\n", "\n", "\n", "def Model1(a,b):\n", " \n", " # Complicated *pure* internal calculation\n", " x = a**2\n", " y = sin(b)\n", " z = exp(x * y)\n", " output = z / (1 + x)\n", "\n", " #--------- Outputs\n", " return output\n", "\n", " \n", "def Model2(model1_output, c):\n", " \n", " # something else complicated and *pure*\n", " iteration_var = 0.531\n", " for i in range(1, 4):\n", " iteration_var = c * iteration_var + model1_output\n", "\n", " #--------- Outputs\n", " return iteration_var\n", "\n", "\n", "def Model3(model2_output, e, c, d):\n", " \n", " # something else also complicated and *pure*\n", " tolerance = 1e-6\n", " error = float('inf')\n", " x = model2_output\n", "\n", " while error > tolerance:\n", " old_x = x\n", "\n", " # Some nonlinear update\n", " x = (old_x - e) / (1 + abs(old_x - e))\n", " x = x + c * sin(x) + d\n", " error = abs(x - old_x)\n", "\n", " #--------- Outputs\n", " return x\n", "\n", "\n", "def Model4(model5_output, a, model1_output):\n", " return model5_output + a + model1_output\n", "\n", "def Model5(model1_output, b, model6_output):\n", " return model1_output + b * model6_output\n", "\n", "def Model6(d, model4_output):\n", " return d + model4_output" ] }, { "cell_type": "markdown", "id": "3", "metadata": {}, "source": [ "Let's draw that again using the graph representation (we'll use a graphviz version since the matplotlib gets too cluttered:)" ] }, { "cell_type": "code", "execution_count": null, "id": "4", "metadata": { "tags": [ "hide-input" ] }, "outputs": [], "source": [ "from pathlib import Path\n", "from IPython.display import SVG\n", "from cottax.consumers.visualization import render\n", "\n", "# graphviz is an optional dependency, and the docs CI has neither the package nor\n", "# `dot`: render the picture where they are installed (refreshing the copy stored\n", "# beside this notebook), and fall back to that stored copy where they are not.\n", "stored = next((p for p in (Path(\"theory_v_graphviz.svg\"),\n", " Path(\"docs/source/theory/theory_v_graphviz.svg\"))\n", " if p.exists()), Path(\"theory_v_graphviz.svg\"))\n", "\n", "try:\n", " svg = render(to_graph(Model1, Model2, Model3, Model4, Model5, Model6),\n", " formatter = PlainFormat()).pipe(format = \"svg\", encoding = \"utf-8\")\n", " stored.write_text(svg)\n", "except Exception: # graphviz not importable, or `dot` not on PATH\n", " svg = stored.read_text()\n", "\n", "SVG(svg)" ] }, { "cell_type": "markdown", "id": "5", "metadata": {}, "source": [ "It's already quite cluttered, and this is only with a few models. Moreover, it's not so clear which ones are cycles and which ones are not. \n", "\n", "A different method to visualize such graphs is using an adjacency matrix.\n", "\n", "The convention is that the models are on the diagonal, and off-diagonal elements indicate edges.\n", "\n", "An off-diagonal edge *above* the diagonal means an edge pointing from the higher model (more top-left) to the lower model (more down-right) and an edge *below* the diagonal points from the lower model to the higher model." ] }, { "cell_type": "code", "execution_count": null, "id": "6", "metadata": {}, "outputs": [], "source": [ "from cottax.consumers.visualization import render_html\n", "from functools import partial \n", "xdsm = partial(render_html, formatter = PlainFormat(), xdsm = {\"input_label\": \"Input\", \"output_label\": \"Output\"})\n", "xdsm(to_graph(Model1, Model2, Model3, Model4, Model5, Model6))" ] }, { "cell_type": "markdown", "id": "7", "metadata": {}, "source": [ "In other words model1_outputs sits above the diagonal, so it points from Model1 (is an output of) and towards Model2 (is an input of).\n", "\n", "The nice thing about this diagram is that if we see that a matrix is fully upper triangular, it means that the ordering is valid! Each lower model only depends on higher models; running the models from high to low is a valid ordering for the model collection. \n", "\n", "With cycles and the strongly connected components algorithm, we see the cycles immediately; they are the parts of the graph where there is a lower triangular edge (like above). Ordering the models by those blocks -- each cycle one block, and everything else a block of its own -- visualizes a model collection clearly:" ] }, { "cell_type": "code", "execution_count": null, "id": "8", "metadata": {}, "outputs": [], "source": [ "def Model7(model1_output, model4_output, model8_output):\n", " return model1_output**model4_output + model8_output\n", "\n", "\n", "def Model8(model3_output, a, model1_output, model9_output):\n", " return a + model3_output / model1_output + model9_output\n", "\n", "def Model9(model7_output, b):\n", " return b- model7_output\n", "\n", "model1_to_9 = [Model1, Model2, Model3, Model6, Model9, Model4, Model5, Model7, Model8]\n", "xdsm(to_graph(*model1_to_9), blocks=True)" ] }, { "cell_type": "markdown", "id": "9", "metadata": {}, "source": [ "We've now got two cycles, which the SCC algorithm blocks and is clearly visible. These are minimal and can't be further reduced [^1]\n", "\n", "[^1]: We could re-order though; Model3 can be run before the Model456 cycle! But for correctness this is irrelevant.\n", "\n", "The nice thing is it allows us to make a *decision*: what solver seems appropriate here for each cycle? And should we take the two cycles together? Inspecting the graph gives a scalable way to look at such cycles immediately [^2].\n", "\n", "\n", "[^2]: You could also just *print* the SCC blocks and make a decision based on that; but this figure gives more information at a quick glance. " ] }, { "cell_type": "markdown", "id": "10", "metadata": {}, "source": [ "## Problem Types\n", "\n", "A cycle *between* models is not the only cycle that can be present. \n", "\n", "An example is an implicit model:\n", "\n", "Solve $x^3 + x^2 +\\sqrt{x - y} = 0.0$\n", "\n", "Why is this a cycle? Well, if you define a particular solver, *it's not*: it's just an explicit model obtaining x from y or vice-versa, but with a solve step inbetween.\n", "However, we've been arguing for flexibility the past couple notebooks: *a priori* we don't know what the best solver is, and we want to remain flexible. Furthermore, we might actually want to solve this together with another model (if there's strong coupling between an input/ output). \n", "\n", "\n", "In some sense, a cycle is precisely something that *needs solving* or a *problem type*.\n", "\n", "Let's look at an implicit function, which is defined by its residual being 0." ] }, { "cell_type": "code", "execution_count": null, "id": "11", "metadata": {}, "outputs": [], "source": [ "def Model10(model3_output, model10_output):\n", " return model10_output**3 + model10_output**2 + (model10_output - model3_output)**0.5\n", "\n", "\n", "\n", "xdsm(to_graph(*model1_to_9, flat.residual(Model10, \"model10_output\")), blocks=True)" ] }, { "cell_type": "markdown", "id": "12", "metadata": {}, "source": [ "A root-find problem type, signalling that this needs to be resolved, is now added to the Model10. Both names it minted are drawn as *what they stand for* rather than as their spelling: $\\textrm{cond}(\\texttt{model10\\_output})$ is the residual on `model10_output`, and $\\textrm{solve}(\\texttt{Model10})$ is the problem over `Model10`.\n", "\n", "Such problem types also exist for fixed point iterations (like a cycle) or for an internal optimizer.\n", "\n", "These have an algebra: you can combine problem types as long as they make sense together: two root-finds are a bigger root-find, a root-find and an optimizer are a set of equality constraints added to the optimizer." ] }, { "cell_type": "markdown", "id": "13", "metadata": {}, "source": [ "All three problem types are declared the same way. The places after the function name what the problem *owns*, and\n", "which wrapper takes the function is what says what makes its answer right:\n", "\n", "| | wraps | right when | drawn as |\n", "|---|---|---|---|\n", "| `residual` | the residual | the residual vanishes | $\\textrm{cond}(u)$ |\n", "| `fixed_point` | the next value | the next value equals the current one | $\\textrm{cond}(u)$ |\n", "| `optimise` | nothing: an objective and constraints | the objective is smallest, and the constraints hold | — |\n", "\n", "The first two declare a *pair* (a body that runs, and a problem that closes it), which is why they\n", "are always a cycle. An optimiser declares one node and no body: an objective and a constraint are\n", "ordinary variables some model already computes, so it mints no name of its own and there is\n", "nothing extra to notate." ] }, { "cell_type": "code", "execution_count": null, "id": "14", "metadata": {}, "outputs": [], "source": [ "def Model11(model11_output, model10_output):\n", " \"\"\"A damped iteration: the condition *is* the next iterate, so nothing is subtracted.\"\"\"\n", " return 0.5 * (model11_output + model10_output)\n", "\n", "\n", "def Cost(model11_output, a, b):\n", " return (model11_output - 1.0)**2 + 0.1 * (a**2 + b**2)\n", "\n", "\n", "@flat.inequality_constraint(\"margin_output\") # a constraint: its output held <= 0\n", "def Margin(model3_output):\n", " return model3_output - 5.0\n", "\n", "\n", "# An optimiser owns the design variables, and reads what the models compute.\n", "# `a` and `b` are no longer external inputs: the optimiser is what produces them.\n", "Design = cx.absolute.optimise(minimise=flat.place(\"cost_output\"), # the objective\n", " design=(flat.place(\"a\"), flat.place(\"b\")),\n", " subject_to=(Margin,)) # g(x) <= 0" ] }, { "cell_type": "code", "execution_count": null, "id": "15", "metadata": {}, "outputs": [], "source": [ "design_graph = to_graph(Model1, Model2, Model3, flat.residual(Model10, \"model10_output\"),\n", " flat.fixed_point(Model11, \"model11_output\"), Cost, Margin, {\"Design\": Design})\n", "\n", "xdsm(design_graph)" ] }, { "cell_type": "markdown", "id": "16", "metadata": {}, "source": [ "For a large graph the XDSM stops being readable; the N2 matrix keeps one row per model. A mark at (row, column) says the row produces what the column reads, so marks below the diagonal are feedback, and every block solved together is boxed in the colour of its problem type." ] }, { "cell_type": "code", "execution_count": null, "id": "17", "metadata": {}, "outputs": [], "source": [ "from cottax.consumers.visualization import render_html\n", "\n", "render_html(design_graph, views=(\"n2\",))" ] }, { "cell_type": "markdown", "id": "18", "metadata": {}, "source": [ "In this situation, the three problem types are all present, and they are all coupled. Now, the graph alone really doesn't determine the nesting: which solver sits inside the other? Can they be combined? In this case, it's important to look at the graph and make an informed decision. This cannot happen with only explicit functions: there's only one, total cycle, not three nested ones. You can make it one cycle by putting everything in the optimizer naturally, but this is not always desired. \n", "\n", "This'll be in the form of operations in the graph: as we've already seen in the previous notebook, you can e.g. assume a particular variable during an iteration to make it acyclic, but there are multiple such operations and they can be combined. " ] }, { "cell_type": "markdown", "id": "19", "metadata": {}, "source": [ "## Summary\n", "\n", "The matrix representation provides a scalable way to inspect graphs and in combination with the strongly connected components algorithm easily locates and visualizes cycles, which demand a solver.\n", "\n", "Beyond model-dependency cycles, a problem can also be declared outright. These always have to be assigned a solver, and can lead to a nested problem statement. \n", "\n", "In the next notebook, we'll look at ways to untangle this and assign solvers using graph operations." ] } ], "metadata": { "kernelspec": { "display_name": "Python 3", "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.14.4" } }, "nbformat": 4, "nbformat_minor": 5 }