{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 1 - ⚡ Zap for Network Utility Maximization Basics ⚡\n",
    "\n",
    "\n",
    "In this notebook, we introduce how to construct and solve NUM problems in Zap.\n",
    "We will manually construct a small network with links and routes, solve the NUM problem using both CVXPY and PMP (ADMM), and analyze the results.\n",
    "\n",
    "1. Creating a NUM Problem\n",
    "2. Solving NUM\n",
    "3. Analyzing Results"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "import zap\n",
    "import torch\n",
    "\n",
    "import numpy as np\n",
    "import cvxpy as cp\n",
    "from scipy.sparse import csc_matrix\n",
    "\n",
    "\n",
    "from experiments.resource_opt_solve.benchmarks.nu_opt_benchmark import NUOptBenchmarkSet\n",
    "from zap.resource_opt.nu_opt_bridge import NUOptBridge\n",
    "from zap.admm import ADMMSolver\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Creating NUM Problems\n",
    "\n",
    "NUM problems in Zap consist of a link-route matrix, which defines the underlying links and streams, link capacities, stream weights, and a list of which traffic streams have linear utilities (Zap assumes by default that all streams have logarithmic utilities).\n",
    "\n",
    "We provide a synthetic data generator which creates NUM test instances. We can, for example, generate a problem with the following parameters:\n",
    "\n",
    "- 2000 links\n",
    "- 1000 streams\n",
    "- 10 links in the average stream\n",
    "- Capacities uniformly in the range (0.1, 1)\n",
    "- 0% of links congested\n",
    "- Uniform stream weights of 1\n",
    "- Logarithmic utilities for all streams"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "base_seed = 42\n",
    "m = 2000 # Number of links\n",
    "n = 1000 # Number of streams\n",
    "avg_stream_length = 3 # Average stream uses 10 links\n",
    "capacity_range = (0.1, 1)\n",
    "link_congest_num_frac = 0\n",
    "lin_util_frac = 0\n",
    "\n",
    "benchmark = NUOptBenchmarkSet(num_problems=1, \n",
    "        m=m, \n",
    "        n=n, \n",
    "        avg_route_length=avg_stream_length, \n",
    "        capacity_range=capacity_range, \n",
    "        link_congest_num_frac=link_congest_num_frac,\n",
    "        base_seed=base_seed)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can extract the desired problem parameters from the synthetic data benchmark generator. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "R, capacities, w, linear_flow_idxs = benchmark.get_data(0)\n",
    "nu_opt_params = {\n",
    "    \"R\": R,\n",
    "    \"capacities\": capacities,\n",
    "    \"w\": w,\n",
    "    \"lin_device_idxs\": linear_flow_idxs,\n",
    "}"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "\n",
    "While we can use the synthetic data generator to create test instances, for demonstration purposes, we instead solve an instance of 'A Small Example' from \"Large-Scale Network Utility Maximization via GPU-Accelerated Proximal Message Passing\" (i.e. we use the same link-route matrix from Fig. 1).\n",
    "\n",
    "We generate a very small problem with the following parameters:\n",
    "\n",
    "- 3 links\n",
    "- 4 streams\n",
    "- Capacities uniformly in the range (0.1, 1)\n",
    "- 0% of links congested\n",
    "- Uniform stream weights of 1\n",
    "- Logarithmic utilities for all streams"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "m = 3\n",
    "n = 4\n",
    "rng = np.random.default_rng(42)\n",
    "capacities = rng.uniform(0.1, 1, size=m) # Capacity range (0.1, 1)\n",
    "w = np.ones(n)\n",
    "lin_device_idxs = None # No devices with linear utility components\n",
    "\n",
    "\n",
    "# Create link-route matrix R as CSC matrix\n",
    "data = np.array([1, 1, 1, 1, 1])\n",
    "row  = np.array([0, 1, 1, 2, 2])\n",
    "col  = np.array([0, 1, 2, 1, 3])\n",
    "R = csc_matrix((data, (row, col)), shape=(m, n))\n",
    "\n",
    "nu_opt_params = {\n",
    "    \"R\": R,\n",
    "    \"capacities\": capacities,\n",
    "    \"w\": w,\n",
    "    \"lin_device_idxs\": lin_device_idxs,\n",
    "}\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Before we construct the NUM problem to solve, we need to finally specify how we group streams. We create a separate group for each set of streams with the same number of terminals (i.e. streams that utilize the same number of links)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "grouping_params = {\"variable_grouping_strategy\": \"discrete_terminal_groups\"}"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We now create the NUM problem. `NUOptBridge` provides the interface for taking problem data and converting it into a NUM problem that we can solve. `NuOptBridge` converts the problem data into its corresponding bipartite graph representation which is used by the solver."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "nu_opt_bridge = NUOptBridge(nu_opt_params, grouping_params)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Solving a NUM Problem\n",
    "\n",
    "We can interface with the `nu_opt_bridge` object to solve our NUM problem in one of two ways:\n",
    "\n",
    "1. `cvxpy` - This method will build a model in `cvxpy` and send it to an off-the-shelf solver, such as Clarabel. You can also use commerical solvers like MOSEK. This approach is best for small to medium problems and finds highly accurate solutions.\n",
    "\n",
    "2. `admm` - This method will use a PyTorch implementation of the alternating direction method of multipliers (ADMM). This can be run on either a CPU or GPU and scales well to larger problems. In general, this method is only capable of finding medium accuracy solutions within a reasonable amount of time.\n",
    "\n",
    "\n",
    "### Solving with CVXPY\n",
    "\n",
    "To solve with `cvxpy`, we simply call `nu_opt_bridge.solve` and (optionally) specify a solver."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "outcome = nu_opt_bridge.solve(solver=cp.CLARABEL)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We will shortly address how to interpret the solver output.\n",
    "\n",
    "If you are using the benchmark set, you can also extract the benchmark problem via CVXPY and solve it directly as well."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "-1909.3491417000203"
      ]
     },
     "execution_count": 9,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "for i, prob in enumerate(benchmark):\n",
    "        problem = prob\n",
    "\n",
    "problem.solve(solver=cp.CLARABEL, verbose=False)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Solving with ADMM (Proximal Message Passing)\n",
    "\n",
    "Solving with ADMM is a little more complicated and requires two steps:\n",
    "\n",
    "1. Transfering device data to PyTorch\n",
    "2. Initializing an ADMM solver object."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "ADMM converged in 104 iterations.\n"
     ]
    }
   ],
   "source": [
    "machine = \"cpu\" # Pick \"cuda\" for a machine with an NVIDIA GPU\n",
    "dtype = torch.float32\n",
    "admm_devices = [d.torchify(machine=machine, dtype=dtype) for d in nu_opt_bridge.devices]\n",
    "admm = ADMMSolver(\n",
    "    machine=machine,\n",
    "    dtype=dtype,\n",
    "    atol=1e-6,\n",
    "    rtol=1e-6,\n",
    "    tau=2,\n",
    "    alpha=1.6,\n",
    "    num_iterations=1000\n",
    ")\n",
    "solution_admm, history_admm = admm.solve(nu_opt_bridge.net, admm_devices, nu_opt_bridge.time_horizon)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "DispatchOutcome(phase_duals=[None, None, None], local_equality_duals=None, local_inequality_duals=None, local_variables=[[tensor([[0.7966],\n",
       "        [0.2918],\n",
       "        [0.6695]])], [tensor([[0.2032]])], None], power=[[tensor([[0.7966],\n",
       "        [0.2918],\n",
       "        [0.6695]])], [tensor([[0.2032]]), tensor([[0.2032]])], [tensor([[-0.7966],\n",
       "        [-0.4950],\n",
       "        [-0.8727]])]], angle=[None, None, None], prices=tensor([[-1.2554],\n",
       "        [-3.4273],\n",
       "        [-1.4936]]), global_angle=None, problem=None, ground=None)"
      ]
     },
     "execution_count": 11,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# ADMM solutions need to be cast to a standard DispatchOutcome\n",
    "outcome_admm = solution_admm.as_outcome()\n",
    "\n",
    "outcome_admm"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Analyzing Results\n",
    "\n",
    "\n",
    "Results are packaged into a hierachically structured `DispatchOutcome` object. This allows us to directly access the terminal flows in the bipartite representation of the problem.\n",
    "\n",
    "1. At the top level, a `DispatchOutcome` has fields such as `power` and `prices`. You can access these fields like any other Python field, e.g., `outcome.power`.\n",
    "2. Each field contains either *stream-specific information* or *global information*. Stream-specific fields will contain a list of length `len(streams)`, where the `i`th entry in the list contains information specific to the `i`th stream. Global fields contain a 2d array of size `(num_nodes, 1)`. You can access the information for stream `i` using normal indexing, e.g., `outcome.power[i]`.\n",
    "3. For stream-specific information, each block of information is further broken down by the *terminal* of the stream. Slack streams have just a single terminal and will always be indexed as `outcome.power[i][0]`. Other utility streams may have two or more terminals. The data for terminal `j` is stored in `outcome.power[i][j]`.\n",
    "4. Finally, the data for terminal `j` of stream `i` is just a 2d array of size `(num_devices, 1)`, where `num_devices = devices[i].num_devices` is the number of devices in device group `i`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([[-1.25535188],\n",
       "       [-3.42724732],\n",
       "       [-1.4936407 ]])"
      ]
     },
     "execution_count": 12,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "outcome.prices  # Global, link duals"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[[array([[0.79656044],\n",
       "         [0.29177735],\n",
       "         [0.66952488]])],\n",
       " [array([[0.20321324]]), array([[0.20321324]])],\n",
       " [array([[-0.79656044],\n",
       "         [-0.4949906 ],\n",
       "         [-0.87273812]])]]"
      ]
     },
     "execution_count": 13,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "outcome.power # Stream-specific (utility and slack streams)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[array([[0.79656044],\n",
       "        [0.29177735],\n",
       "        [0.66952488]])]"
      ]
     },
     "execution_count": 14,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "outcome.power[0] # Terminal flows for the three single terminal utility streams (Traffic Streams 1, 3, 4) (x1, x3, x4)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[array([[0.20321324]]), array([[0.20321324]])]"
      ]
     },
     "execution_count": 15,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "outcome.power[1] # Terminal flows for the one two terminal utility stream (Traffic Stream 2) (x2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[array([[-0.79656044],\n",
       "        [-0.4949906 ],\n",
       "        [-0.87273812]])]"
      ]
     },
     "execution_count": 16,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "outcome.power[2] # Terminal flows for the three single terminal slack streams (Slack Streams 1, 2, 3) (s1-c1, s2-c2, s3-c3)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "zap-eFcx8M23-py3.12",
   "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.12.1"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}