{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "0",
   "metadata": {},
   "source": [
    "# Sellar"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1",
   "metadata": {},
   "source": [
    "The Sellar problem is a classic problem with one coupled discipline:\n",
    "\n",
    "$$\n",
    "\\text{minimize }_{x_1, z_1, z_2} f\n",
    "\n",
    "$$\n",
    "subject to:\n",
    "\n",
    "$$ \n",
    "\\begin{aligned}\n",
    "y_1 &\\geq 3.6\n",
    "y_2 &\\leq 24 \\\\\n",
    "-10 \\leq z_1 &\\leq 10 \\\\\n",
    "0 \\leq z_2 &\\leq 10 \\\\\n",
    "0 \\leq x_1 &\\leq 10\n",
    "\\end{aligned}\n",
    "$$\n",
    "and with \n",
    "\n",
    "$$ \n",
    "\\begin{aligned}\n",
    "y_1 &= z_1^2 + z_2 + x_1 - 0.2 y_2\\\\\n",
    "y_2 &= \\sqrt{y_1} + z_1 + z_2\\\\\n",
    "f   &= x_1^2 + z_2 + y_1 + \\exp(-y_2)\n",
    "\\end{aligned}\n",
    "$$\n",
    "\n",
    "We'll look a bit on how Cottax would handle this internally, but first let's just look at what you'd do in practice:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2",
   "metadata": {},
   "outputs": [],
   "source": [
    "import jax\n",
    "jax.config.update('jax_enable_x64', True)\n",
    "import jax.numpy as jnp\n",
    "import cottax.assemblies.jax as cx\n",
    "from cottax.assemblies.jax import Start\n",
    "flat = cx.flat"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Declaring\n",
    "@flat.function(\"y_1\")\n",
    "def D1(x_1, z_1, z_2, y_2):\n",
    "    return z_1**2 + z_2 + x_1 - 0.2 * y_2\n",
    "\n",
    "@flat.function(\"y_2\")\n",
    "def D2(y_1, z_1, z_2):\n",
    "    return jnp.sqrt(y_1) + z_1 + z_2\n",
    "\n",
    "@flat.function(\"f\")\n",
    "def obj(x_1, z_2, y_1, y_2):\n",
    "    return x_1**2 + z_2 + y_1 + jnp.exp(-y_2)\n",
    "\n",
    "@flat.inequality_constraint(\"g_1\")\n",
    "def C1(y_1):\n",
    "    return 3.6 - y_1\n",
    "\n",
    "@flat.inequality_constraint(\"g_2\")\n",
    "def C2(y_2):\n",
    "    return y_2 -  24\n",
    "\n",
    "opt      = flat.optimise(design = {\"z_1\" : (-10 ,10), \"z_2\": (0, 10),  \"x_1\": (0, 10)}, minimise = obj , subject_to = (C1, C2))\n",
    "declared = flat.declare(D1, D2, obj, C1, C2, opt)\n",
    "\n",
    "# Running\n",
    "\n",
    "drivers = (cx.SLSQP(), cx.Newton(), cx.Picard())\n",
    "mdf      = declared + flat.mdf(opt) + flat.drive(*drivers)\n",
    "\n",
    "start_val = {Start: dict(z_1 = 5., z_2 = 2., x_1 = 1., y_2 = 0.4)}\n",
    "\n",
    "out = flat.run(mdf, read_from = start_val)\n",
    "\n",
    "for name in (\"z_1\", \"z_2\", \"x_1\", \"f\", \"g_1\", \"g_2\"):\n",
    "    print(f'{name:<3}', out[name])\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4",
   "metadata": {},
   "outputs": [],
   "source": [
    "cx.render_html(mdf, blocks=True, formatter = cx.LatexFormatter())\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5",
   "metadata": {},
   "source": [
    "## Breaking it down"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6",
   "metadata": {},
   "source": [
    "The cottax graph itself doesn't use any implicit conversions and is totally unambiguous. This enables the mathematics to work without having to continuously convert meanings or even have two meanings silently be the same. \n",
    "\n",
    "Our goal is then to provide the surface interface to the core, which does have these implicit conversions. These surfaces can be swapped if another conversion is desired without affecting any of the core functionality. \n",
    "\n",
    "This similarly holds for the output side: a visualization in most cases requires a formatter to make the cottax native types legible or transform their representation to nicely-looking latex.\n",
    "\n",
    "The full representation can be shown below:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7",
   "metadata": {
    "tags": [
     "hide-input"
    ]
   },
   "outputs": [],
   "source": [
    "import dataclasses\n",
    "\n",
    "def place(p):    \n",
    "    keys = ', '.join(repr(k)\n",
    "                     for k in p.segments)\n",
    "    return f'{type(p).__name__}({keys})'\n",
    "\n",
    "def full(x, depth=1):\n",
    "    pad, inner = '    ' * (depth - 1), '    ' * depth\n",
    "    if hasattr(x, 'segments'):                        # a VarPath or NodePath, on one line\n",
    "        return place(x)\n",
    "    if isinstance(x, tuple):\n",
    "        if not x: return '()'\n",
    "        return '(\\n' + ''.join(f'{inner}{full(v, depth + 1)},\\n' for v in x) + pad + ')'\n",
    "    if dataclasses.is_dataclass(x) and not isinstance(x, type) and dataclasses.fields(x):\n",
    "        body = ''.join(f'{inner}{f.name} = {full(getattr(x, f.name), depth + 1)},\\n'\n",
    "                       for f in dataclasses.fields(x))\n",
    "        return f'{type(x).__name__}(\\n{body}{pad})'\n",
    "    return getattr(x, '__qualname__', None) or repr(x)  # Eq, Min, a function\n",
    "\n",
    "def show_full(g):\n",
    "    for n in g.nodes:\n",
    "        print(f'{place(n)}: {full(g[n])}')\n",
    "    for inner, outer in g.within.items():\n",
    "        print(f'{place(inner)} within {place(outer)}')\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8",
   "metadata": {
    "tags": [
     "hide-output"
    ]
   },
   "outputs": [],
   "source": [
    "show_full(mdf.graph)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "markdown",
   "id": "10",
   "metadata": {},
   "source": [
    "This is clearly not easy to write by hand! The surfaces and assemblies that can be imported make a much more convenient way to write this. In fact, the spelling itself is usually shortened a lot to keep it more readable. \n",
    "\n",
    "A quick explanation of the above:\n",
    "\n",
    "- `VarPath` is the variable PyTree, `Node-Path`. These are different: a node label is not a variable place unless you explicitly convert!\n",
    "- `DictKey` is a dictionary key in the datastructure. `^` indicates a name that is created (NOT in the global datastructure by definition!)\n",
    "- relations are comparisons between left-hand-side and right-hand-sides: None means zero. \n",
    "- `inputs` are reads of a function, `outputs` the outputs of a function. `unknowns` are the variables a condition variable owns and `condition` is the condition it needs to satisfy. `parameters` are variables that affect the answer but are not conditions: data is driver data that is needed for the driver but does not participate in what makes an answer correct (an initial condition!)\n",
    "- a `contract` is the driver contract which adds the edges for the initial value (the driver data)\n",
    "- a `implementation` refers to the function or some implementation of the driver.\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "11",
   "metadata": {},
   "source": [
    "## Given values"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "12",
   "metadata": {},
   "source": [
    "Also, our philosophy is that we don't really want any actual values in the graph itself: it must be the structure only. Then, the graph gives a concrete view on what has to be provided to run it: data hidden in the graph obfuscates this. Furthermore, the graph defines the program (through e.g. a JIT cache key), and if values are changing, the compilation has to be redone. [^1] \n",
    "\n",
    "This does lead to some user-friendlyness problems: if you declare an optimizer, you usually give it bounds, and you don't want to index into some place to give them. The declaration and surface consumers help here. They contain both the graph and a parallel datastructure\n",
    "with those given values and their names in the graph. Additional tags can be used as well (which also have a parallel datastructure).\n",
    "\n",
    "[^1]: In fact, in cottax, this is not necessarily enforced to give you the flexibility to i.e. give some solver setting or hard-coded tolerances, which are non-array-like and therefore hash with it (the above drivers have default tolerances)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "13",
   "metadata": {},
   "source": [
    "The MDF Declared object has the following given:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "14",
   "metadata": {},
   "outputs": [],
   "source": [
    "print(mdf.given)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "15",
   "metadata": {},
   "source": [
    "And we can automatically do that with function reads as well:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "16",
   "metadata": {},
   "outputs": [],
   "source": [
    "@flat.function(\"k\", reads = {\"b\" :3})\n",
    "def f3(a, b ,c):\n",
    "    ...\n",
    "\n",
    "f3.given"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "17",
   "metadata": {},
   "source": [
    "These are given a special place: this wouldn't be necessary in the flat space, but in hierarchical spaces there may not be a value in the datastructure to express it easily and it is convenient to give. \n",
    "\n",
    "\n",
    "Then, the above problem requires a set of start values: they are not the same as the variables output (a graph cannot overwrite boundary inputs by definition) and have therefore to be given as well. Again, this is not so convenient in general, because\n",
    "they might have these minted ^ names."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "18",
   "metadata": {},
   "outputs": [],
   "source": [
    "mdf.graph.graph.boundary_inputs"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "19",
   "metadata": {},
   "source": [
    "The assembly run can extract these ^start values and their copies (^start^copy['y_2']) from a datastructure:\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "20",
   "metadata": {},
   "outputs": [],
   "source": [
    "flat.valued(mdf, read_from={Start :  dict(z_1 = 5., z_2 = 2., x_1 = 1., y_2 = 0.4)}).values"
   ]
  }
 ],
 "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
}