{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "f2ce45bd-0bd5-4618-ac41-5a4e82fc99b1",
   "metadata": {},
   "source": [
    "# Building Custom Curves (Nelson-Siegel)\n",
    "\n",
    "The *rateslib* curves objects are structured in such a way that it is\n",
    "strightforward to create new, custom curve objects useable throughout the library, provided they implement\n",
    "the necessary abstract base objects.\n",
    "\n",
    "## Nelson-Siegel Parametrization\n",
    "\n",
    "This example will construct a parametric **Nelson-Siegel Curve**, whose continuously\n",
    "compounded zero rate, $r$, at time, $T$, is given by the following\n",
    "equation of **four** parameters:\n",
    "\n",
    "$$\n",
    "r(T) = \\begin{bmatrix} \\beta_0 & \\beta_1 & \\beta_2 \\end{bmatrix} \\begin{bmatrix} 1 \\\\ \\lambda (1- e^{-T/ \\lambda}) / T \\\\ \\lambda (1- e^{-T/ \\lambda})/ T - e^{-T/ \\lambda} \\end{bmatrix}\n",
    "$$\n",
    "\n",
    "This leads to the **discount factors** on that curve equaling:\n",
    "\n",
    "$$\n",
    "v(T) = e^{-T r(T)}\n",
    "$$\n",
    "\n",
    "## _BaseCurve ABCs\n",
    "\n",
    "First we setup the skeletal structure of our custom curve. We will inherit from\n",
    "[_BaseCurve][basecurve] and setup the necessary abstract base class (ABC) properties.\n",
    "\n",
    "Some items we know in advance regarding our custom curve:\n",
    "\n",
    "- it will return discount factors, defining its [_CurveType][curvetype],\n",
    "- it has `4` parameters,\n",
    "- we must define a ``start`` date (which makes its initial node) and ``end`` in its [_CurveNodes][curvenodes]\n",
    "- we will just inherit ``meta`` property intialisation from the [_BaseCurve][basecurve]\n",
    "\n",
    "[basecurve]: ./api/rateslib.curves._BaseCurve.rst\n",
    "[curvetype]: ./api/rateslib.curves._CurveType.rst\n",
    "[curvenodes]: ./api/rateslib.curves._CurveNodes.rst\n",
    "[withops]: ./api/rateslib.curves._WithOperations.rst"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f8200b93-0ea8-408e-81d5-dad9a25f8503",
   "metadata": {},
   "outputs": [],
   "source": [
    "from rateslib import *\n",
    "import numpy as np\n",
    "\n",
    "from rateslib.curves import _BaseCurve, _CurveType, _CurveNodes\n",
    "\n",
    "class NelsonSiegelCurve(_BaseCurve):\n",
    "    \n",
    "    # ABC properties \n",
    "    \n",
    "    _base_type = _CurveType.dfs\n",
    "    _id = None\n",
    "    _meta = None\n",
    "    _nodes = None\n",
    "    _ad = None\n",
    "    _interpolator = None\n",
    "    _n = 4\n",
    "   \n",
    "    def __init__(self, start, end, beta0, beta1, beta2, lambd):\n",
    "        self._id = super()._id\n",
    "        self._meta = super()._meta\n",
    "        self._ad = 0\n",
    "        self._nodes = _CurveNodes({start: 0.0, end: 0.0})\n",
    "        self._params = (beta0, beta1, beta2, lambd)\n",
    "\n",
    "    # ABC required methods\n",
    "    \n",
    "    def _set_ad_order(self, order):\n",
    "        raise NotImplementedError()\n",
    "   \n",
    "    def __getitem__(self, date):\n",
    "        raise NotImplementedError()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "60b2dae0-037b-46bc-a42a-72282a65e8d7",
   "metadata": {},
   "source": [
    "This curve can now be initialised without raising any errors relating to *Abstract Base Classes*.\n",
    "However, it doesn't do much without implementing the ``__getitem__`` method."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "70973c5c-5b30-47a6-a216-483d6c062886",
   "metadata": {},
   "outputs": [],
   "source": [
    "curve = NelsonSiegelCurve(dt(2000, 1, 1), dt(2010, 1, 1), 0.025, 0.0, 0.0, 0.5)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "56d3c045-6247-4720-a345-819b53c9ae15",
   "metadata": {},
   "source": [
    "## The `__getitem__` method\n",
    "\n",
    "This method will return discount factors. If the requested date is prior to the curve: return zero as usual. If the date is the\n",
    "same as the initial node date: return one, else use continuously compounded rates to\n",
    "derive a discount factor."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "aaa9c1e8-fb67-4164-b963-95ebc838a6d4",
   "metadata": {},
   "outputs": [],
   "source": [
    "class NelsonSiegelCurve(_BaseCurve):\n",
    "    _base_type = _CurveType.dfs\n",
    "    _id = None\n",
    "    _meta = None\n",
    "    _nodes = None\n",
    "    _ad = None\n",
    "    _interpolator = None\n",
    "    _n = 4\n",
    "    \n",
    "    def __init__(self, start, end, beta0, beta1, beta2, lambd):\n",
    "        self._id = super()._id\n",
    "        self._meta = super()._meta\n",
    "        self._ad = 0\n",
    "        self._nodes = _CurveNodes({start: 0.0, end: 0.0})\n",
    "        self._params = (beta0, beta1, beta2, lambd)\n",
    "    \n",
    "    def _set_ad_order(self, order):\n",
    "        raise NotImplementedError()\n",
    "    \n",
    "    def __getitem__(self, date):\n",
    "        if date < self.nodes.initial:\n",
    "            return 0.0\n",
    "        elif date == self.nodes.initial:\n",
    "            return 1.0\n",
    "        b0, b1, b2, l0 = self._params\n",
    "        T = dcf(self.nodes.initial, date, convention=self.meta.convention, calendar=self.meta.calendar)\n",
    "        a1 = l0 * (1 - dual_exp(-T / l0)) / T\n",
    "        a2 =  a1 - dual_exp(-T / l0)\n",
    "        r = b0 + a1 * b1 + a2 * b2\n",
    "        return dual_exp(-T * r)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8b4b87aa-3188-4674-a0d6-7b07e0fc0581",
   "metadata": {},
   "source": [
    "Once this method is added to the class and the discount factors are available,\n",
    "all of the provided methods are also available. The following snippet of code demonstrates each of the [rate()][basecurve], [plot()][basecurve], [roll()][basecurve], [shift()][basecurve] and [translate()][basecurve] methods without the user having implemented any of them:\n",
    "\n",
    "[basecurve]: ./api/rateslib.curves._BaseCurve.rst"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "80905b10-9b8f-4e5b-8df6-47785431748c",
   "metadata": {},
   "outputs": [],
   "source": [
    "ns_curve = NelsonSiegelCurve(\n",
    "    start=dt(2000, 1, 1), \n",
    "    end=dt(2010, 1, 1), \n",
    "    beta0=0.03, \n",
    "    beta1=-0.01, \n",
    "    beta2=0.01, \n",
    "    lambd=0.75\n",
    ")\n",
    "ns_curve.plot(\"1b\", comparators=[ns_curve.shift(100), ns_curve.roll(\"6m\")])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "53960f99-3aa7-4c94-9ff9-5affbe0530b4",
   "metadata": {},
   "source": [
    "## Mutatbility and the ``Solver``\n",
    "\n",
    "In order to allow this curve to be calibrated by a [Solver][solver],\n",
    "we need to add some elements that allows the [Solver][solver] to interact\n",
    "with it. We will also set the `NelsonSeigelCurve` to inherit\n",
    "[_WithMutability][withmut].\n",
    "\n",
    "Firstly, we can add the ``getter`` methods:\n",
    "\n",
    "[solver]: ./api/rateslib.solver.Solver.rst\n",
    "[withmut]: ./api/rateslib.curves._WithMutability.rst"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9ba30bac-d804-436a-96ce-fed21d171b39",
   "metadata": {},
   "outputs": [],
   "source": [
    "_ini_solve = 0  #  <-  this tells the `Solver` not to 'skip' any parameters (like the initial DF of 1.0 on a `Curve`)\n",
    "\n",
    "def _get_node_vector(self):\n",
    "   # this maps the curve's parameters to a consistent NumPy array consumed by the `Solver`.\n",
    "   return np.array(self._params)\n",
    "\n",
    "def _get_node_vars(self):\n",
    "   # this get the ordered variable names of the curve's parameters for AD management.\n",
    "   return tuple(f\"{self._id}{i}\" for i in range(self._ini_solve, self._n))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "73ce8408-7f97-44e3-be8b-736cab3e0fde",
   "metadata": {},
   "source": [
    "The ``setter`` methods that the [Solver][solver] needs are slightly\n",
    "more complicated. These allow the [Solver][solver] to mutate the *Curve* directly and provide mappings.\n",
    "\n",
    "[solver]: ./api/rateslib.solver.Solver.rst"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ac7ae1e0-e4a6-41ba-ae4c-70e084f9c64b",
   "metadata": {},
   "outputs": [],
   "source": [
    "from rateslib.curves import _WithMutability\n",
    "from rateslib.mutability import _new_state_post\n",
    "from rateslib.dual import set_order_convert\n",
    "from rateslib.dual.utils import _dual_float\n",
    "\n",
    "@_new_state_post  # <- this sets a new state every time the curve is mutated\n",
    "def _set_node_vector(self, vector, ad):\n",
    "    # this maps the `Solver's` output new node values to the curve's parameters attribute with the appropriate AD number type\n",
    "    if ad == 0:\n",
    "        self._params = tuple(_dual_float(_) for _ in vector)\n",
    "    elif ad == 1:\n",
    "        self._params = tuple(\n",
    "            Dual(_dual_float(_), [f\"{self._id}{i}\"], []) for i, _ in enumerate(vector)\n",
    "        )\n",
    "    else: # ad == 2\n",
    "        self._params = tuple(\n",
    "            Dual2(_dual_float(_), [f\"{self._id}{i}\"], [], []) for i, _ in enumerate(vector)\n",
    "        )\n",
    "\n",
    "def _set_ad_order(self, order):\n",
    "    # this allows the `Solver` to change the AD number type of the curve for `delta` and `gamma` calculations.\n",
    "    if self.ad == order:\n",
    "        return None\n",
    "    else:\n",
    "        self._ad = order\n",
    "        self._params = tuple(\n",
    "            set_order_convert(_, order, [f\"{self._id}{i}\"]) for i, _ in enumerate(self._params)\n",
    "        )"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6563e9e5-3ce4-4dd6-9bc8-32f93328795b",
   "metadata": {},
   "source": [
    "Adding these elements yields the **final code** to implement this user custom *Curve* class:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "eb6585d8-2bbe-477a-bcce-ec44bd0d15a4",
   "metadata": {},
   "outputs": [],
   "source": [
    "class NelsonSiegelCurve(_WithMutability, _BaseCurve):\n",
    "\n",
    "    # ABC properties\n",
    "    \n",
    "    _ini_solve = 0\n",
    "    _base_type = _CurveType.dfs\n",
    "    _id = None\n",
    "    _meta = None\n",
    "    _nodes = None\n",
    "    _ad = None\n",
    "    _interpolator = None\n",
    "    _n = 4\n",
    "    \n",
    "    def __init__(self, start, end, beta0, beta1, beta2, lamb):\n",
    "        self._id = super()._id\n",
    "        self._meta = super()._meta\n",
    "        self._ad = 0\n",
    "        self._nodes = _CurveNodes({start: 0.0, end: 0.0})\n",
    "        self._params = (beta0, beta1, beta2, lamb)\n",
    "\n",
    "    def __getitem__(self, date):\n",
    "        if date < self.nodes.initial:\n",
    "            return 0.0\n",
    "        elif date == self.nodes.initial:\n",
    "            return 1.0\n",
    "        b0, b1, b2, l0 = self._params\n",
    "        T = dcf(self.nodes.initial, date, convention=self.meta.convention, calendar=self.meta.calendar)\n",
    "        a1 = l0 * (1 - dual_exp(-T / l0)) / T\n",
    "        a2 =  a1 - dual_exp(-T / l0)\n",
    "        r = b0 + a1 * b1 + a2 * b2\n",
    "        return dual_exp(-T * r)\n",
    "\n",
    "    # Solver mutability methods\n",
    "    \n",
    "    def _get_node_vector(self):\n",
    "        return np.array(self._params)\n",
    "    \n",
    "    def _get_node_vars(self):\n",
    "        return tuple(f\"{self._id}{i}\" for i in range(self._ini_solve, self._n))\n",
    "    \n",
    "    def _set_node_vector(self, vector, ad):\n",
    "        if ad == 0:\n",
    "            self._params = tuple(_dual_float(_) for _ in vector)\n",
    "        elif ad == 1:\n",
    "            self._params = tuple(\n",
    "                 Dual(_dual_float(_), [f\"{self._id}{i}\"], []) for i, _ in enumerate(vector)\n",
    "            )\n",
    "        else: # ad == 2\n",
    "            self._params = tuple(\n",
    "                Dual2(_dual_float(_), [f\"{self._id}{i}\"], [], []) for i, _ in enumerate(vector)\n",
    "            )\n",
    "    \n",
    "    def _set_ad_order(self, order):\n",
    "        if self.ad == order:\n",
    "            return None\n",
    "        else:\n",
    "            self._ad = order\n",
    "            self._params = tuple(\n",
    "                set_order_convert(_, order, [f\"{self._id}{i}\"]) for i, _ in enumerate(self._params)\n",
    "            )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1bf5b775-0521-4f87-b4f0-aa7d9bb1037e",
   "metadata": {},
   "outputs": [],
   "source": [
    "ns_curve = NelsonSiegelCurve(dt(2000, 1, 1), dt(2010, 1, 1), 0.03, -0.01, 0.01, 0.75)\n",
    "solver = Solver(\n",
    "    curves=[ns_curve],\n",
    "    instruments=[\n",
    "        IRS(dt(2000, 1, 1), \"2y\", \"A\", curves=ns_curve),\n",
    "        IRS(dt(2000, 1, 1), \"4y\", \"A\", curves=ns_curve),\n",
    "        IRS(dt(2000, 1, 1), \"7y\", \"A\", curves=ns_curve),\n",
    "    ],\n",
    "    s=[2.45, 2.90, 2.66],\n",
    "    id=\"NS\",\n",
    "    instrument_labels=[\"2y\", \"4y\", \"7y\"],\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "de8906d5-5d82-4f1d-b610-a6bd60845fb3",
   "metadata": {},
   "outputs": [],
   "source": [
    "ns_curve.plot(\"1b\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1d5484c5-8ca4-4317-941f-5b3b779b4a3e",
   "metadata": {},
   "source": [
    "## Risk\n",
    "\n",
    "Since the [Solver][solver] has been invoked all typically delta and gamma methods can now also be used against this curve risk model\n",
    "\n",
    "[solver]: ./api/rateslib.solver.Solver.rst"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2854108d-3247-45a9-80e8-0f566f675f88",
   "metadata": {},
   "outputs": [],
   "source": [
    "portfolio = IRS(dt(2003, 7, 1), \"4y\", \"A\", notional=50e6, fixed_rate=4.5, curves=ns_curve)\n",
    "portfolio.delta(solver=solver)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "754534bd-646e-472d-b316-193390d37e18",
   "metadata": {},
   "outputs": [],
   "source": [
    "portfolio.gamma(solver=solver)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "28b956f7-b1b1-426e-87ff-38d2c620805b",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "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.13.0"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}