{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "0",
   "metadata": {},
   "source": [
    "# Theory III:  Computational Graphs"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1",
   "metadata": {},
   "source": [
    "The previous notebooks have illustrated that if you have a set of pure functions with explicit interfaces, you can reason about the dependencies between them. \n",
    "\n",
    "Such a set of functions thus defines a *computational graph*: nodes are the functions and the directed edges connecting them are the outputs and inputs of the functions.\n",
    "\n",
    "We'll re-use the functions from before and draw the computational graph:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2",
   "metadata": {
    "tags": [
     "hide-input"
    ]
   },
   "outputs": [],
   "source": [
    "from math import sin, exp\n",
    "import inspect\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",
    "def Model4(model1_output, f):\n",
    "    return model1_output + f\n",
    "\n",
    "def interface(f):\n",
    "    return {f.__name__ : {\"inputs\" : list(inspect.signature(f).parameters.keys()), \"outputs\" : [f.__name__.lower() + \"_output\"]}}\n",
    "\n",
    "\n",
    "\n",
    "import cottax.assemblies.jax as cx\n",
    "from cottax.assemblies.jax import flat\n",
    "from cottax.consumers.visualization.xdsm import PlainFormat\n",
    "from cottax.consumers.visualization import draw\n",
    "\n",
    "def to_graph(*models):\n",
    "    return cx.absolute.declare(*(flat.function(m, m.__name__.lower() + \"_output\") for m in models)).graph\n",
    "\n",
    "_ = draw(to_graph(Model1, Model2, Model3, Model4), formatter = PlainFormat(), boundary = True)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3",
   "metadata": {},
   "source": [
    "If an order exists in which the functions can be executed one-by-one, the computational graph is a *directed acyclic graph* (DAG). This notebook only considers these graphs; for cyclic graphs, see the next notebook.\n",
    "\n",
    "Standard algorithms exist to determine a valid ordering of a DAG. To determine this order, we'll use the `networkx` library. Below is some code to transfer our models to the networkx graph format."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4",
   "metadata": {
    "tags": [
     "hide-input"
    ]
   },
   "outputs": [],
   "source": [
    "import networkx as nx\n",
    "\n",
    "\n",
    "def models_to_graph(models):\n",
    "    graph = nx.DiGraph()\n",
    "\n",
    "    # Map each output to the model that produces it\n",
    "\n",
    "    models = [interface(i) for i in models]\n",
    "    \n",
    "    output_to_model = {}\n",
    "\n",
    "    for model_dict in models:\n",
    "        for model_name, config in model_dict.items():\n",
    "            graph.add_node(model_name)\n",
    "\n",
    "            for output in config.get(\"outputs\", []):\n",
    "                output_to_model[output] = model_name\n",
    "\n",
    "    # Create model -> model edges based on inputs\n",
    "    for model_dict in models:\n",
    "        for model_name, config in model_dict.items():\n",
    "\n",
    "            for input_name in config.get(\"inputs\", []):\n",
    "                producer = output_to_model.get(input_name)\n",
    "\n",
    "                if producer is not None:\n",
    "                    # Several variables can join the same pair of models, so the edge\n",
    "                    # keeps all of them rather than only the last one written.\n",
    "                    if graph.has_edge(producer, model_name):\n",
    "                        graph[producer][model_name][\"inputs\"].append(input_name)\n",
    "                    else:\n",
    "                        graph.add_edge(producer, model_name, inputs=[input_name])\n",
    "\n",
    "    return graph\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5",
   "metadata": {},
   "source": [
    "Then, we can compute a valid ordering [^1]\n",
    "\n",
    "[^1]: In fact there are multiple valid ones: Model2 and Model4 can be executed in any order after Model1, and Model3 can be executed after both of them. This ordering is not unique, but it is valid.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6",
   "metadata": {},
   "outputs": [],
   "source": [
    "nx_digraph = models_to_graph([Model4, Model3, Model2, Model1])\n",
    "\n",
    "list(nx.topological_sort(nx_digraph))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7",
   "metadata": {},
   "source": [
    "The basic evaluation that is valid for *any combination of models that form a DAG* is then as follows (we also add a check that the input data structure matches what we'd expect):"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8",
   "metadata": {
    "tags": [
     "hide-input"
    ]
   },
   "outputs": [],
   "source": [
    "def external_inputs(models : set):\n",
    "    # Get what base inputs are required and what outputs are produced by the models\n",
    "    total_inputs = set()\n",
    "    model_outputs= set()\n",
    "\n",
    "    for model in models:\n",
    "        interface_i = interface(model)[model.__name__]\n",
    "\n",
    "        total_inputs.update(interface_i[\"inputs\"])\n",
    "        model_outputs.update(interface_i[\"outputs\"])\n",
    "\n",
    "    # These are the inputs that have to be defined externally: no model produces them!\n",
    "    external_inputs = total_inputs - model_outputs\n",
    "\n",
    "    return external_inputs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9",
   "metadata": {},
   "outputs": [],
   "source": [
    "def TotalModel(data_structure, models : set):\n",
    "    ds_start = data_structure.copy() # Keeping it pure: the original data structure is not modified, but a copy is used instead.\n",
    "\n",
    "    # Let's also check that the provided data structure has all the required external inputs for the given models:\n",
    "    if not set(external_inputs(models)) == set(ds_start.keys()):\n",
    "        raise ValueError(f\"The provided data structure keys {set(ds_start.keys())}\" + \n",
    "                         f\"do not match the required external inputs {set(external_inputs(models))} for the given models.\")\n",
    "\n",
    "    # Create the graph and order it:\n",
    "    nx_digraph = models_to_graph(list(models))\n",
    "    str_to_model = {model.__name__: model for model in models}\n",
    "    models_in_order = list(nx.topological_sort(nx_digraph))\n",
    "\n",
    "    # Rest the same as before:\n",
    "    for model_name in models_in_order:\n",
    "        model = str_to_model[model_name]\n",
    "        # Get the interface for the current model\n",
    "        interface_i = interface(model)[model.__name__]\n",
    "        \n",
    "        # Prepare the inputs for the current model based on its interface\n",
    "        inputs = {input_name: ds_start[input_name] for input_name in interface_i[\"inputs\"]}\n",
    "        \n",
    "        # Call the model with the prepared inputs\n",
    "        output = model(**inputs)\n",
    "        \n",
    "        # Update the data structure with the output of the current model\n",
    "        ds_start.update({interface_i[\"outputs\"][0]: output})\n",
    "    return ds_start\n",
    "\n",
    "ds_start = {\"a\": 1.0, \"b\": 0.5, \"c\": 2.0, \"d\": 0.1, \"e\": 0.2, \"f\" : 3.0}\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "10",
   "metadata": {},
   "outputs": [],
   "source": [
    "TotalModel(ds_start, {Model1, Model2, Model3, Model4})"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "11",
   "metadata": {},
   "source": [
    "and it will error if we put something there that is not an external input for the given models:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "12",
   "metadata": {},
   "outputs": [],
   "source": [
    "ds_wrong = {\"a\": 1.0, \"b\": 0.5, \"c\": 2.0, \"d\": 0.1, \"e\": 0.2, \"f\" : 3.0, \"model1_output\" : 4.0}\n",
    "\n",
    "try:\n",
    "    TotalModel(ds_wrong, {Model1, Model2, Model3, Model4})\n",
    "except ValueError as ve:\n",
    "    print(f\"ValueError: {ve}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "13",
   "metadata": {},
   "source": [
    "Great! Let's use an alternative for Model4 that uses Model3 output instead of Model1 output[^3]\n",
    "\n",
    "[^3]: The choice that the model's output is the name is not very nice here; but we'll keep that for brevity. In a more realistic setting you could make the input and output anything you want. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "14",
   "metadata": {},
   "outputs": [],
   "source": [
    "def Model4(a, model3_output):\n",
    "    return a + model3_output\n",
    "\n",
    "\n",
    "\n",
    "nx_digraph = models_to_graph([Model4, Model3, Model2, Model1])\n",
    "print(list(nx.topological_sort(nx_digraph)))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "15",
   "metadata": {},
   "source": [
    "The topological order has immediately changed to reflect this new dependency!\n",
    "\n",
    "We can also immediately run it ( after removing the now-useless `f` input):"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "16",
   "metadata": {},
   "outputs": [],
   "source": [
    "# remove `f` now!\n",
    "ds_start = {\"a\": 1.0, \"b\": 0.5, \"c\": 2.0, \"d\": 0.1, \"e\": 0.2} \n",
    "TotalModel(ds_start, {Model1, Model2, Model3, Model4})"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "17",
   "metadata": {},
   "source": [
    "### Alternative models"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "18",
   "metadata": {},
   "source": [
    "We can now solve the alternative problem. Let's define a way to switch between two models:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "19",
   "metadata": {},
   "outputs": [],
   "source": [
    "def make_model4(model4_i : bool):\n",
    "    if model4_i:\n",
    "        def Model4(a, model3_output):\n",
    "            return a + model3_output\n",
    "    else:\n",
    "        def Model4(model1_output, f):\n",
    "            return model1_output + f\n",
    "    return Model4\n",
    "\n",
    "def make_model5(model5_i : bool):\n",
    "    if model5_i:\n",
    "        def Model5(model4_output, f):\n",
    "            return model4_output + f\n",
    "    else:\n",
    "        def Model5(model2_output, model3_output, f):\n",
    "            return model2_output + model3_output + f\n",
    "    return Model5"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "20",
   "metadata": {},
   "source": [
    "We can immediately create the graph and see the various possibilities:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "21",
   "metadata": {},
   "outputs": [],
   "source": [
    "def create_graph(model4_i : bool, model5_i : bool):\n",
    "    Model4 = make_model4(model4_i)\n",
    "    Model5 = make_model5(model5_i)\n",
    "    nx_digraph = models_to_graph([Model4, Model3, Model2, Model1, Model5])\n",
    "    print(\"Topological ordering: \", list(nx.topological_sort(nx_digraph)))\n",
    "    return draw(to_graph(Model1, Model2, Model3, Model4, Model5), formatter = PlainFormat(), boundary = True)\n",
    "\n",
    "_ = create_graph(True, True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "22",
   "metadata": {},
   "outputs": [],
   "source": [
    "_ = create_graph(False, True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "23",
   "metadata": {},
   "outputs": [],
   "source": [
    "_ = create_graph(False, False)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "24",
   "metadata": {},
   "source": [
    "These graphs are rather cluttered and difficult to analyse by hand; this is precisely the point! An algorithm can analyse these very fast and create a valid ordering and do the analysing for you, as long as you have the pure functions and explicit data interfaces![^4]\n",
    "\n",
    "[^4]: A reasonable question is if this isn't much slower in execution than hand-ordering: the answer is no. A jax compilation traces through this graph once and creates a compiled function which is literally this graph. After compilation there's *no overhead at all!* All overhead is in the graph construction and tracing, but since these are not large graphs, this is a very small cost. (a 1600 model graph takes under a few ms)."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "25",
   "metadata": {},
   "source": [
    "## Summary\n",
    "\n",
    "Explicit data interfaces and pure functions define a computational graph, which standard algorithms can use to define a valid ordering between the models.\n",
    "\n",
    "Thus, given a collection of models, you don't need to hand-order them, this can be done automatically!\n",
    "\n",
    "In the next notebook, we'll look at collections of models that *cannot* be ordered because of a dependency cycle."
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "process_port",
   "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
}