{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "0",
   "metadata": {},
   "source": [
    "# Theory IV: Cycles in the graph"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1",
   "metadata": {},
   "source": [
    "Not all model collections are acyclic: models can, directly or indirectly through other models, depend on each other.\n",
    "\n",
    "We'll re-use the same Model1, Model2, Model3 as before, and add a cycle:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2",
   "metadata": {
    "tags": [
     "hide-input"
    ]
   },
   "outputs": [],
   "source": [
    "from math import sin, exp\n",
    "import inspect\n",
    "import networkx as nx\n",
    "\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 interface(f):\n",
    "    return {f.__name__ : {\"inputs\" : list(inspect.signature(f).parameters.keys()), \"outputs\" : [f.__name__.lower() + \"_output\"]}}\n",
    "\n",
    "\n",
    "def external_inputs(models : set):\n",
    "    total_inputs  = set()\n",
    "    model_outputs = set()\n",
    "\n",
    "    for model in models:\n",
    "        interface_i = interface(model)[model.__name__]\n",
    "        total_inputs.update(interface_i[\"inputs\"])\n",
    "        model_outputs.update(interface_i[\"outputs\"])\n",
    "\n",
    "    return total_inputs - model_outputs\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3",
   "metadata": {},
   "outputs": [],
   "source": [
    "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\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4",
   "metadata": {},
   "source": [
    "These form a chain: Model4 requires Model5, which requires Model6, which again requires Model4. \n",
    "\n",
    "There is now no order at all in which the models can be run one-by-one: the topological sort will simply error:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5",
   "metadata": {
    "tags": [
     "hide-input"
    ]
   },
   "outputs": [],
   "source": [
    "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": "code",
   "execution_count": null,
   "id": "6",
   "metadata": {},
   "outputs": [],
   "source": [
    "cycle_graph = models_to_graph([Model1, Model2, Model3, Model4, Model5, Model6])\n",
    "\n",
    "try:\n",
    "    list(nx.topological_sort(cycle_graph))\n",
    "except nx.NetworkXUnfeasible as e:\n",
    "    print(\"No valid ordering:\", e)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7",
   "metadata": {},
   "source": [
    "Such cycles can be resolved with a fixed point iteration. In order to do this, we do need to know what exactly the cycle is. In this case, it's straightforward, but it's not always the case.\n",
    "\n",
    "A standard graph algorithm that finds such cycles is the strongly connected component algorithm. \n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8",
   "metadata": {},
   "source": [
    "Instead of running model-by-model, this forces us to run 'block' by block. A block is a model or a set of models that are run together: in the case of a cycle, it's because they have to be solved together.\n",
    "\n",
    "If you have an acyclic graph, each model is itself just a block, but not here:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9",
   "metadata": {},
   "outputs": [],
   "source": [
    "def blocks_in_order(nx_digraph):\n",
    "    \"\"\"The models grouped into blocks, and the blocks put in a valid run order.\"\"\"\n",
    "    blocks   = list(nx.strongly_connected_components(nx_digraph))\n",
    "    block_of = {model: i for i, block in enumerate(blocks) for model in block}\n",
    "\n",
    "    # Order each block:\n",
    "    block_graph = nx.DiGraph()\n",
    "    block_graph.add_nodes_from(range(len(blocks)))\n",
    "    block_graph.add_edges_from((block_of[producer], block_of[consumer])\n",
    "                               for producer, consumer in nx_digraph.edges\n",
    "                               if block_of[producer] != block_of[consumer])\n",
    "\n",
    "    return [sorted(blocks[i]) for i in nx.topological_sort(block_graph)]\n",
    "\n",
    "\n",
    "blocks_in_order(cycle_graph)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "10",
   "metadata": {},
   "source": [
    "We see that we can run Model1 and Model2 first, then have to solve the block cycle, and then run Model3 (which could also be run *before* the cycle)."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "11",
   "metadata": {},
   "source": [
    "### Solving a block\n",
    "\n",
    "Running a single model is unchanged from the previous notebook:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "12",
   "metadata": {},
   "outputs": [],
   "source": [
    "def run_model(model, data_structure):\n",
    "\n",
    "    # Get I/O\n",
    "    interface_i = interface(model)[model.__name__]\n",
    "\n",
    "    # Prepare inputs:\n",
    "    inputs = {input_name: data_structure[input_name] for input_name in interface_i[\"inputs\"]}\n",
    "\n",
    "    # Return updated datastructure:\n",
    "    return {**data_structure, interface_i[\"outputs\"][0]: model(**inputs)}\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "13",
   "metadata": {},
   "source": [
    "A block of several models cannot be run like that: whichever model we start with is missing one of its inputs. The way out is to **assume** a value for one of the variables inside the cycle [^1] and make that variable the iteration variable.\n",
    "\n",
    "[^1]: In later notebooks we will formalize this notion in the sense of MDAO architectures; for now this demonstrates the mechanism as well.\n",
    "\n",
    "Given that assumption, in a block, that model output will no longer be needed and thus it can now be ordered:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "14",
   "metadata": {},
   "outputs": [],
   "source": [
    "def block_order(block, nx_digraph, assumed):\n",
    "    \"\"\"The order to run a block's models in, once `assumed` is taken as given.\"\"\"\n",
    "    block_graph = nx_digraph.subgraph(block).copy()\n",
    "\n",
    "    # Taking `assumed` as given means the edges carrying it are no longer dependencies.\n",
    "    block_graph.remove_edges_from([(producer, consumer)\n",
    "                                   for producer, consumer, edge in block_graph.edges(data=True)\n",
    "                                   if assumed in edge[\"inputs\"]])\n",
    "\n",
    "    return list(nx.topological_sort(block_graph))\n",
    "\n",
    "\n",
    "cycle = ['Model4', 'Model5', 'Model6']\n",
    "\n",
    "for assumed in ('model4_output', 'model5_output', 'model6_output'):\n",
    "    print(f\"assuming {assumed:15s} ->\", block_order(cycle, cycle_graph, assumed))\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "15",
   "metadata": {},
   "source": [
    "Each choice gives a different (valid) order. Now the block can be run: assume a value, run the models in that order, and see what comes back out. Convergence is achieved when \n",
    "the resulting value is the same as the assumed value:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "16",
   "metadata": {},
   "outputs": [],
   "source": [
    "def solve_block(block, data_structure, models, nx_digraph, assumed,\n",
    "                start = 1.0, tolerance = 1e-10, max_iterations = 100):\n",
    "    by_name = {model.__name__: model for model in models}\n",
    "    order   = block_order(block, nx_digraph, assumed)\n",
    "\n",
    "    # We have to start somewhere: nothing has produced `assumed` yet.\n",
    "    ds = {**data_structure, assumed: start}\n",
    "\n",
    "    for iteration in range(max_iterations):\n",
    "        old_value = ds[assumed]\n",
    "\n",
    "        for model_name in order:\n",
    "            ds = run_model(by_name[model_name], ds)\n",
    "\n",
    "        # Converged when running the block reproduces the value we assumed.\n",
    "        if abs(ds[assumed] - old_value) < tolerance:\n",
    "            print(f\"{block}: converged in {iteration + 1} iterations, assuming '{assumed}'\")\n",
    "            return ds\n",
    "\n",
    "    print(f\"{block}: did NOT converge in {max_iterations} iterations, assuming '{assumed}'\")\n",
    "    return ds\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "17",
   "metadata": {},
   "source": [
    "And that is enough to run any collection of models, cyclic or not:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "18",
   "metadata": {},
   "outputs": [],
   "source": [
    "def TotalModel(data_structure, models : set, assumed):\n",
    "    ds = data_structure.copy()\n",
    "\n",
    "    if not set(external_inputs(models)) == set(ds.keys()):\n",
    "        raise ValueError(f\"The provided data structure keys {set(ds.keys())} \" +\n",
    "                         f\"do not match the required external inputs {set(external_inputs(models))} for the given models.\")\n",
    "\n",
    "    nx_digraph = models_to_graph(models)\n",
    "    by_name    = {model.__name__: model for model in models}\n",
    "\n",
    "    for block in blocks_in_order(nx_digraph):\n",
    "\n",
    "        if len(block) == 1:\n",
    "            # An ordinary model: everything it reads is available already.\n",
    "            ds = run_model(by_name[block[0]], ds)\n",
    "        else:\n",
    "            # A cycle: it has to be solved. (note we start with a value of 1)\n",
    "            ds = solve_block(block, ds, models, nx_digraph, assumed)\n",
    "\n",
    "    return ds\n",
    "\n",
    "\n",
    "all_models = {Model1, Model2, Model3, Model4, Model5, Model6}\n",
    "ds_start   = {\"a\": 1.0, \"b\": 0.5, \"c\": 2.0, \"d\": 0.1, \"e\": 0.2}\n",
    "\n",
    "TotalModel(ds_start, all_models, assumed = \"model5_output\")\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "19",
   "metadata": {},
   "source": [
    "This works, and the nice thing is it works for *any model collection like this*! You just have to assume a model output\n",
    "and then the fixed point iteration will handle it.\n",
    "\n",
    "\n",
    "However, there were two big hidden assumptions:\n",
    "\n",
    "1. We chose which variable is the iteration variable. All three would have worked, so why this particular one?\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "20",
   "metadata": {},
   "outputs": [],
   "source": [
    "TotalModel(ds_start, all_models, assumed = \"model4_output\")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "21",
   "metadata": {},
   "outputs": [],
   "source": [
    "TotalModel(ds_start, all_models, assumed = \"model6_output\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "22",
   "metadata": {},
   "source": [
    "2. The algorithm: why should we use a fixed point iteration always? In some cases this might not converge, and we don't know whether the failure is algorithmic or whether the problem really has no answer: [^2]\n",
    "\n",
    "[^2]: Here it does have one. The three models in the cycle are affine, so the block has exactly one solution ($\\text{model5\\_output} = -37.70$). The sweep gain is $b$ whichever variable is assumed, so at $b = 1.1$ Gauss--Seidel diverges from all three choices; Newton, which is exact in one step on an affine residual, finds it immediately."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "23",
   "metadata": {},
   "outputs": [],
   "source": [
    "TotalModel({**ds_start, \"b\": 1.1}, all_models, assumed = \"model5_output\")\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "24",
   "metadata": {},
   "source": [
    "These choices cannot be made just on the basis of the graph (i.e. algorithmically); there's always some ambiguity."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "25",
   "metadata": {},
   "source": [
    "And there are more choices still: whether two cycles that happen to sit next to each other should be one block or two, or whether the cycle should not be solved here at all, but handed to an optimizer as a consistency constraint."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "26",
   "metadata": {},
   "source": [
    "## Summary\n",
    "\n",
    "A cyclic graph has no valid model ordering, but it does have a valid **block** ordering: the strongly connected components, ordered by sorting a graph of the blocks themselves. A block of one model is run; a block of several is a cycle and has to be solved. The graph-based representation enables us to look at the minimal cycles and automatically assign a solver (in this notebook a fixed point iteration) to it.\n",
    "\n",
    "Solving one requires assuming a variable inside it, which cuts the cycle and makes the block orderable again. Unlike the ordering itself, that choice (and the choice of solver, and of what is solved together) cannot be read off the graph.\n",
    "\n",
    "In other words, we need to, as users, take a look at the graph and make an informed decision. \n",
    "\n",
    "In the next notebook, we'll look at how to represent the graph in an easier-to-understand way, and look at other possible cycles."
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "beerpy",
   "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
}