{
 "cells": [
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Customzie Formulation"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Encapsuled the optimization problem calss, AMS provides direct access to the optimization formulation, where users have the option to customize the formulation without playing with the source code.\n",
    "\n",
    "In this example, we will walk through the optimization problem structure and show how to customize the formulation."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "\n",
    "import ams"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "ams.config_logger(stream_level=20)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Inspect the Optimization Problem Structure"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Parsing input file \"/Users/jinningwang/work/miniconda3/envs/amsre/lib/python3.12/site-packages/ams/cases/5bus/pjm5bus_demo.xlsx\"...\n",
      "Input file parsed in 0.0853 seconds.\n",
      "Zero line rates detacted in rate_b, rate_c, adjusted to 999.\n",
      "System set up in 0.0015 seconds.\n"
     ]
    }
   ],
   "source": [
    "sp = ams.load(ams.get_case('5bus/pjm5bus_demo.xlsx'),\n",
    "              setup=True,\n",
    "              no_output=True,)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In AMS, a routine collects the descriptive dispatch formulations.\n",
    "`DCOPF`, `RTED`, etc, are the subclasses of `RoutineBase`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Building system matrices\n",
      "Parsing OModel for <DCOPF>\n",
      "Evaluating OModel for <DCOPF>\n",
      "Finalizing OModel for <DCOPF>\n",
      "<DCOPF> initialized in 0.0215 seconds.\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "True"
      ]
     },
     "execution_count": 4,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "sp.DCOPF.init()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "After successful initialization, the attribute `om` is populated with CVXPY-based optimization problem.\n",
    "\n",
    "The user can even hack to the source `prob` attribute to customize it if necessary."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "cvxpy.problems.problem.Problem"
      ]
     },
     "execution_count": 5,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "type(sp.DCOPF.om.prob)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Customize Built-in Formulation"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Here we extend DCOPF with consideration of CO2 emission, where the original formulation can be found in the documentation\n",
    "[Routine Reference - DCOPF](https://ltb.readthedocs.io/projects/ams/en/latest/typedoc/DCED.html#dcopf).\n",
    "To simplify the demonstration, following assumptions are made:\n",
    "1. Variable $\\boxed{e_g}$ is the CO2 emission of each generator. It is proportional to the generation, described by a parameter $\\boxed{k_e}$ in the unit t/p.u..\n",
    "1. Total CO2 emmission is limited by a constant cap $\\boxed{t_e}$, in the unit $t$.\n",
    "1. A tax $\\boxed{c_e}$ is imposed on each unit of CO2 emission in the unit of $/p.u., and the tax is included in the objective function.\n",
    "\n",
    "Thus, the revised formulation is as follows, where box indicates the revision:\n",
    "\n",
    "min. $\\sum ( c_{2} p_{g}^2 + c_{1} p_{g} + u_{g} c_{0} + \\boxed{c_{e} e_{g}} )$\n",
    "\n",
    "s.t.\n",
    "\n",
    "$\\boxed{ e_{g} - k_{e} p_{g} = 0}$\n",
    "\n",
    "$\\boxed{ \\sum e_{g} - t_{e} \\leq 0}$\n",
    "\n",
    "$-p_g + c_{trl,ne}p_{g,0} + c_{trl,e}p_{g,\\min} \\leq 0$\n",
    "\n",
    "$p_g - c_{trl,ne}p_{g,0} - c_{trl,e}p_{g,\\max} \\leq 0$\n",
    "\n",
    "$B_{bus}\\theta_{bus} + p^{inj}_{bus} + C_{lpd} + C_{sh}g_{sh} - C_{p}g_{p} = 0$\n",
    "\n",
    "$-B_f\\theta_{bus} - p^{inj}_f - R_{ATEA} \\leq 0$\n",
    "\n",
    "$B_f\\theta_{bus} + p^{inj}_f - R_{ATEA} \\leq 0$\n",
    "\n",
    "$-C^T_f\\theta_{bus} - \\theta_{\\max} \\leq 0$\n",
    "\n",
    "$C^T_f\\theta_{bus} - \\theta_{\\max} \\leq 0$\n",
    "\n",
    "Decision variables: $p_g$, $\\theta_{bus}$, $\\boxed{e_g}$\n",
    "\n",
    "Note that line flow $p_{lf}$ is calculated as $B_f\\theta_{bus} + p^{inj}_f$ after solving the problem."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Add services\n",
    "\n",
    "Services are used to store values or build matrix for easier formulation."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "ValueService: DCOPF.te"
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "sp.DCOPF.addService(name='te', tex_name='t_e',\n",
    "                    unit='t', info='emission cap',\n",
    "                    value=12)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Add parameters\n",
    "\n",
    "We need the following parameters to be defined as `RParam`: `ke` and `ce`. They should be 1D array in the same length as the number of generators and `te` is a scalar.\n",
    "\n",
    "For a general `RParam`, it has attributes `model`, `indexer`, and `imodel` to describe its source model and index model. The definition of `c2` in DCOPF source code is a good example.\n",
    "However, for ones defined through API, since there is no model containing it, all above attributes are not applicable, and the user should be aware of the sequence of the parameters.\n",
    "\n",
    "Considering the sequence can be indexed by the generator index, it is used to reference the variables order.\n",
    "Assuming `ke` is reciprocal to the generator capacity, and `ce` is the same for each generator, we can define the parameters as follows:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "# get the generator indices\n",
    "stg_idx = sp.DCOPF.pg.get_all_idxes()\n",
    "\n",
    "# get the value of pmax\n",
    "pmax = sp.DCOPF.get(src='pmax', attr='v', idx=stg_idx)\n",
    "\n",
    "# assume the emission factor is 1 for all generators\n",
    "ke = np.ones_like(pmax)\n",
    "\n",
    "# assume tax is reciprocal of pmax\n",
    "ce = np.reciprocal(pmax)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "RParam: DCOPF.ke"
      ]
     },
     "execution_count": 8,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "sp.DCOPF.addRParam(name='ke', tex_name='k_e',\n",
    "                   info='gen emission factor',\n",
    "                   model=None, src=None, unit=None,\n",
    "                   v=ke)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "RParam: DCOPF.ce"
      ]
     },
     "execution_count": 9,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "sp.DCOPF.addRParam(name='ce', tex_name='c_e',\n",
    "                   info='gen emission tax',\n",
    "                   model=None, src=None, unit=None,\n",
    "                   v=ce)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Add variables\n",
    "\n",
    "The gerator emission `eg` is added as a new variable."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Var: StaticGen.eg"
      ]
     },
     "execution_count": 10,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "sp.DCOPF.addVars(name='eg', tex_name='e_g',\n",
    "                 info='Gen emission', unit='t',\n",
    "                 model='StaticGen', src=None)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Add constraints\n",
    "\n",
    "The CO2 emission is an equality constraint, and the CO2 emission cap is a simple linear inequality constraint.\n",
    "\n",
    "If wish to revise an existing built-in constraint, you can redefine the constraint `e_str` attribute."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Constraint: egb [ON]"
      ]
     },
     "execution_count": 11,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "sp.DCOPF.addConstrs(name='egb', info='Gen emission balance',\n",
    "                    e_str='eg - mul(ke, pg)', is_eq=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Constraint: eub [ON]"
      ]
     },
     "execution_count": 12,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "sp.DCOPF.addConstrs(name='eub', info='emission upper bound',\n",
    "                    e_str='sum(eg) - te', is_eq=False,)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Revise the objective function\n",
    "\n",
    "The `e_str` can be revised to include the CO2 emission tax.\n",
    "Here we only need to append the tax term to the original objective function."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [],
   "source": [
    "sp.DCOPF.obj.e_str += '+ sum(mul(ce, pg))'"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Finalize the Customization\n",
    "\n",
    "After revising the problem, remember to initialize it before solving."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Parsing OModel for <DCOPF>\n",
      "Evaluating OModel for <DCOPF>\n",
      "Finalizing OModel for <DCOPF>\n",
      "<DCOPF> initialized in 0.0061 seconds.\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "True"
      ]
     },
     "execution_count": 14,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "sp.DCOPF.init()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Solve it and Check the Results"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "<DCOPF> solved as optimal in 0.0243 seconds, converged in 8 iterations with CLARABEL.\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "True"
      ]
     },
     "execution_count": 15,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "sp.DCOPF.run(solver='CLARABEL')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Inspect the results."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([0.2       , 2.35272305, 0.6       , 6.64727695, 0.2       ])"
      ]
     },
     "execution_count": 16,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "sp.DCOPF.eg.v"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([0.2       , 2.35272305, 0.6       , 6.64727695, 0.2       ])"
      ]
     },
     "execution_count": 17,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "sp.DCOPF.pg.v"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "np.float64(1.8761187311381662)"
      ]
     },
     "execution_count": 18,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "sp.DCOPF.obj.v"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Load the original problem as a baseline for comparison."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Parsing input file \"/Users/jinningwang/work/miniconda3/envs/amsre/lib/python3.12/site-packages/ams/cases/5bus/pjm5bus_demo.xlsx\"...\n",
      "Input file parsed in 0.0288 seconds.\n",
      "Zero line rates detacted in rate_b, rate_c, adjusted to 999.\n",
      "System set up in 0.0018 seconds.\n"
     ]
    }
   ],
   "source": [
    "sp0 = ams.load(ams.get_case('5bus/pjm5bus_demo.xlsx'),\n",
    "               setup=True,\n",
    "               no_output=True,)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Building system matrices\n",
      "Parsing OModel for <DCOPF>\n",
      "Evaluating OModel for <DCOPF>\n",
      "Finalizing OModel for <DCOPF>\n",
      "<DCOPF> initialized in 0.0075 seconds.\n",
      "<DCOPF> solved as optimal in 0.0060 seconds, converged in 8 iterations with CLARABEL.\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "True"
      ]
     },
     "execution_count": 20,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "sp0.DCOPF.run(solver='CLARABEL')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "From the comparasion, we can see that the generation schedule changes."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([0.2       , 1.43998388, 0.6       , 5.76001612, 2.        ])"
      ]
     },
     "execution_count": 21,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "sp0.DCOPF.pg.v"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "np.float64(0.9585953247653323)"
      ]
     },
     "execution_count": 22,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "sp0.DCOPF.obj.v"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "amsre",
   "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.0"
  },
  "orig_nbformat": 4
 },
 "nbformat": 4,
 "nbformat_minor": 2
}