writing_good_code.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "c5013d4d",
   "metadata": {},
   "source": [
    "\n",
    "<a id='writing-good-code'></a>\n",
    "<div id=\"qe-notebook-header\" align=\"right\" style=\"text-align:right;\">\n",
    "        <a href=\"https://quantecon.org/\" title=\"quantecon.org\">\n",
    "                <img style=\"width:250px;display:inline;\" width=\"250px\" src=\"https://assets.quantecon.org/img/qe-menubar-logo.svg\" alt=\"QuantEcon\">\n",
    "        </a>\n",
    "</div>"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0ef85ba6",
   "metadata": {},
   "source": [
    "# Writing Good Code\n",
    "\n",
    "\n",
    "<a id='index-0'></a>\n",
    "> “Any fool can write code that a computer can understand. Good programmers write code that humans can understand.” – Martin Fowler"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "77df4d5b",
   "metadata": {},
   "source": [
    "## Overview\n",
    "\n",
    "When computer programs are small, poorly written code is not overly costly.\n",
    "\n",
    "But more data, more sophisticated models, and more computer power are enabling us to take on more challenging problems that involve writing longer programs.\n",
    "\n",
    "For such programs, investment in good coding practices will pay high returns.\n",
    "\n",
    "The main payoffs are higher productivity and faster code.\n",
    "\n",
    "In this lecture, we review some elements of good coding practice.\n",
    "\n",
    "We also touch on modern developments in scientific computing — such as just in time compilation — and how they affect good program design."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f83010aa",
   "metadata": {},
   "source": [
    "## An Example of Poor Code\n",
    "\n",
    "Let’s have a look at some poorly written code.\n",
    "\n",
    "The job of the code is to generate and plot time series of the simplified Solow model\n",
    "\n",
    "\n",
    "<a id='equation-gc-solmod'></a>\n",
    "$$\n",
    "k_{t+1} = s k_t^{\\alpha} + (1 - \\delta) k_t,\n",
    "\\quad t = 0, 1, 2, \\ldots \\tag{20.1}\n",
    "$$\n",
    "\n",
    "Here\n",
    "\n",
    "- $ k_t $ is capital at time $ t $ and  \n",
    "- $ s, \\alpha, \\delta $ are parameters (savings, a productivity parameter and depreciation)  \n",
    "\n",
    "\n",
    "For each parameterization, the code\n",
    "\n",
    "1. sets $ k_0 = 1 $  \n",
    "1. iterates using [(20.1)](#equation-gc-solmod) to produce a sequence $ k_0, k_1, k_2 \\ldots , k_T $  \n",
    "1. plots the sequence  \n",
    "\n",
    "\n",
    "The plots will be grouped into three subfigures.\n",
    "\n",
    "In each subfigure, two parameters are held fixed while another varies"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "44806d2e",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "\n",
    "# Allocate memory for time series\n",
    "k = np.empty(50)\n",
    "\n",
    "fig, axes = plt.subplots(3, 1, figsize=(8, 16))\n",
    "\n",
    "# Trajectories with different α\n",
    "δ = 0.1\n",
    "s = 0.4\n",
    "α = (0.25, 0.33, 0.45)\n",
    "\n",
    "for j in range(3):\n",
    "    k[0] = 1\n",
    "    for t in range(49):\n",
    "        k[t+1] = s * k[t]**α[j] + (1 - δ) * k[t]\n",
    "    axes[0].plot(k, 'o-', label=rf\"$\\alpha = {α[j]},\\; s = {s},\\; \\delta={δ}$\")\n",
    "\n",
    "axes[0].grid(lw=0.2)\n",
    "axes[0].set_ylim(0, 18)\n",
    "axes[0].set_xlabel('time')\n",
    "axes[0].set_ylabel('capital')\n",
    "axes[0].legend(loc='upper left', frameon=True)\n",
    "\n",
    "# Trajectories with different s\n",
    "δ = 0.1\n",
    "α = 0.33\n",
    "s = (0.3, 0.4, 0.5)\n",
    "\n",
    "for j in range(3):\n",
    "    k[0] = 1\n",
    "    for t in range(49):\n",
    "        k[t+1] = s[j] * k[t]**α + (1 - δ) * k[t]\n",
    "    axes[1].plot(k, 'o-', label=rf\"$\\alpha = {α},\\; s = {s[j]},\\; \\delta={δ}$\")\n",
    "\n",
    "axes[1].grid(lw=0.2)\n",
    "axes[1].set_xlabel('time')\n",
    "axes[1].set_ylabel('capital')\n",
    "axes[1].set_ylim(0, 18)\n",
    "axes[1].legend(loc='upper left', frameon=True)\n",
    "\n",
    "# Trajectories with different δ\n",
    "δ = (0.05, 0.1, 0.15)\n",
    "α = 0.33\n",
    "s = 0.4\n",
    "\n",
    "for j in range(3):\n",
    "    k[0] = 1\n",
    "    for t in range(49):\n",
    "        k[t+1] = s * k[t]**α + (1 - δ[j]) * k[t]\n",
    "    axes[2].plot(k, 'o-', label=rf\"$\\alpha = {α},\\; s = {s},\\; \\delta={δ[j]}$\")\n",
    "\n",
    "axes[2].set_ylim(0, 18)\n",
    "axes[2].set_xlabel('time')\n",
    "axes[2].set_ylabel('capital')\n",
    "axes[2].grid(lw=0.2)\n",
    "axes[2].legend(loc='upper left', frameon=True)\n",
    "\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2a7554c7",
   "metadata": {},
   "source": [
    "True, the code more or less follows [PEP8](https://www.python.org/dev/peps/pep-0008/).\n",
    "\n",
    "At the same time, it’s very poorly structured.\n",
    "\n",
    "Let’s talk about why that’s the case, and what we can do about it."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0f6fba02",
   "metadata": {},
   "source": [
    "## Good Coding Practice\n",
    "\n",
    "There are usually many different ways to write a program that accomplishes a given task.\n",
    "\n",
    "For small programs, like the one above, the way you write code doesn’t matter too much.\n",
    "\n",
    "But if you are ambitious and want to produce useful things, you’ll  write medium to large programs too.\n",
    "\n",
    "In those settings, coding style matters **a great deal**.\n",
    "\n",
    "Fortunately, lots of smart people have thought about the best way to write code.\n",
    "\n",
    "Here are some basic precepts."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "79bcff6d",
   "metadata": {},
   "source": [
    "### Don’t Use Magic Numbers\n",
    "\n",
    "If you look at the code above, you’ll see numbers like `50` and `49` and `3` scattered through the code.\n",
    "\n",
    "These kinds of numeric literals in the body of your code are sometimes called “magic numbers”.\n",
    "\n",
    "This is not a compliment.\n",
    "\n",
    "While numeric literals are not all evil, the numbers shown in the program above\n",
    "should certainly be replaced by named constants.\n",
    "\n",
    "For example, the code above could declare the variable `time_series_length = 50`.\n",
    "\n",
    "Then in the loops, `49` should be replaced by `time_series_length - 1`.\n",
    "\n",
    "The advantages are:\n",
    "\n",
    "- the meaning is much clearer throughout  \n",
    "- to alter the time series length, you only need to change one value  "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9c513880",
   "metadata": {},
   "source": [
    "### Don’t Repeat Yourself\n",
    "\n",
    "The other mortal sin in the code snippet above is repetition.\n",
    "\n",
    "Blocks of logic (such as the loop to generate time series) are repeated with only minor changes.\n",
    "\n",
    "This violates a fundamental tenet of programming: Don’t repeat yourself (DRY).\n",
    "\n",
    "- Also called DIE (duplication is evil).  \n",
    "\n",
    "\n",
    "Yes, we realize that you can just cut and paste and change a few symbols.\n",
    "\n",
    "But as a programmer, your aim should be to **automate** repetition, **not** do it yourself.\n",
    "\n",
    "More importantly, repeating the same logic in different places means that eventually one of them will likely be wrong.\n",
    "\n",
    "If you want to know more, read the excellent summary found on [this page](https://code.tutsplus.com/tutorials/3-key-software-principles-you-must-understand--net-25161).\n",
    "\n",
    "We’ll talk about how to avoid repetition below."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f23e6537",
   "metadata": {},
   "source": [
    "### Minimize Global Variables\n",
    "\n",
    "Sure, global variables (i.e., names assigned to values outside of any function or class) are convenient.\n",
    "\n",
    "Rookie programmers typically use global variables with abandon — as we once did ourselves.\n",
    "\n",
    "But global variables are dangerous, especially in medium to large size programs, since\n",
    "\n",
    "- they can affect what happens in any part of your program  \n",
    "- they can be changed by any function  \n",
    "\n",
    "\n",
    "This makes it much harder to be certain about what some  small part of a given piece of code actually commands.\n",
    "\n",
    "Here’s a [useful discussion on the topic](http://wiki.c2.com/?GlobalVariablesAreBad).\n",
    "\n",
    "While the odd global in small scripts is no big deal, we recommend that you teach yourself to avoid them.\n",
    "\n",
    "(We’ll discuss how just below)."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c5adaeed",
   "metadata": {},
   "source": [
    "#### JIT Compilation\n",
    "\n",
    "For scientific computing, there is another good reason to avoid global variables.\n",
    "\n",
    "As [we’ve seen in previous lectures](https://python-programming.quantecon.org/numba.html), JIT compilation can generate excellent performance for scripting languages like Python.\n",
    "\n",
    "But the task of the compiler used for JIT compilation becomes harder when global variables are present.\n",
    "\n",
    "Put differently, the type inference required for JIT compilation is safer and\n",
    "more effective when variables are sandboxed inside a function."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a5ed573c",
   "metadata": {},
   "source": [
    "### Use Functions or Classes\n",
    "\n",
    "Fortunately, we can easily avoid the evils of global variables and WET code.\n",
    "\n",
    "- WET stands for “we enjoy typing” and is the opposite of DRY.  \n",
    "\n",
    "\n",
    "We can do this by making frequent use of functions or classes.\n",
    "\n",
    "In fact, functions and classes are designed specifically to help us avoid shaming ourselves by repeating code or excessive use of global variables."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "46aec21a",
   "metadata": {},
   "source": [
    "#### Which One, Functions or Classes?\n",
    "\n",
    "Both can be useful, and in fact they work well with each other.\n",
    "\n",
    "We’ll learn more about these topics over time.\n",
    "\n",
    "(Personal preference is part of the story too)\n",
    "\n",
    "What’s really important is that you use one or the other or both."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0b3a2b44",
   "metadata": {},
   "source": [
    "## Revisiting the Example\n",
    "\n",
    "Here’s some code that reproduces the plot above with better coding style."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "418ffac2",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "from itertools import product\n",
    "\n",
    "def plot_path(ax, αs, s_vals, δs, time_series_length=50):\n",
    "    \"\"\"\n",
    "    Add a time series plot to the axes ax for all given parameters.\n",
    "    \"\"\"\n",
    "    k = np.empty(time_series_length)\n",
    "\n",
    "    for (α, s, δ) in product(αs, s_vals, δs):\n",
    "        k[0] = 1\n",
    "        for t in range(time_series_length-1):\n",
    "            k[t+1] = s * k[t]**α + (1 - δ) * k[t]\n",
    "        ax.plot(k, 'o-', label=rf\"$\\alpha = {α},\\; s = {s},\\; \\delta = {δ}$\")\n",
    "\n",
    "    ax.set_xlabel('time')\n",
    "    ax.set_ylabel('capital')\n",
    "    ax.set_ylim(0, 18)\n",
    "    ax.legend(loc='upper left', frameon=True)\n",
    "\n",
    "fig, axes = plt.subplots(3, 1, figsize=(8, 16))\n",
    "\n",
    "# Parameters (αs, s_vals, δs)\n",
    "set_one = ([0.25, 0.33, 0.45], [0.4], [0.1])\n",
    "set_two = ([0.33], [0.3, 0.4, 0.5], [0.1])\n",
    "set_three = ([0.33], [0.4], [0.05, 0.1, 0.15])\n",
    "\n",
    "for (ax, params) in zip(axes, (set_one, set_two, set_three)):\n",
    "    αs, s_vals, δs = params\n",
    "    plot_path(ax, αs, s_vals, δs)\n",
    "\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8d4ce22d",
   "metadata": {},
   "source": [
    "If you inspect this code, you will see that\n",
    "\n",
    "- it uses a function to avoid repetition.  \n",
    "- Global variables are quarantined by collecting them together at the end, not the start of the program.  \n",
    "- Magic numbers are avoided.  \n",
    "- The loop at the end where the actual work is done is short and relatively simple.  "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "800d9bc9",
   "metadata": {},
   "source": [
    "## Exercises"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a0e26afc",
   "metadata": {},
   "source": [
    "## Exercise 20.1\n",
    "\n",
    "Here is some code that needs improving.\n",
    "\n",
    "It involves a basic supply and demand problem.\n",
    "\n",
    "Supply is given by\n",
    "\n",
    "$$\n",
    "q_s(p) = \\exp(\\alpha p) - \\beta.\n",
    "$$\n",
    "\n",
    "The demand curve is\n",
    "\n",
    "$$\n",
    "q_d(p) = \\gamma p^{-\\delta}.\n",
    "$$\n",
    "\n",
    "The values $ \\alpha $, $ \\beta $, $ \\gamma $ and\n",
    "$ \\delta $ are **parameters**\n",
    "\n",
    "The equilibrium $ p^* $ is the price such that\n",
    "$ q_d(p) = q_s(p) $.\n",
    "\n",
    "We can solve for this equilibrium using a root finding algorithm.\n",
    "Specifically, we will find the $ p $ such that $ h(p) = 0 $,\n",
    "where\n",
    "\n",
    "$$\n",
    "h(p) := q_d(p) - q_s(p)\n",
    "$$\n",
    "\n",
    "This yields the equilibrium price $ p^* $. From this we get the\n",
    "equilibrium quantity by $ q^* = q_s(p^*) $\n",
    "\n",
    "The parameter values will be\n",
    "\n",
    "- $ \\alpha = 0.1 $  \n",
    "- $ \\beta = 1 $  \n",
    "- $ \\gamma = 1 $  \n",
    "- $ \\delta = 1 $  "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f8af5db1",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "from scipy.optimize import brentq\n",
    "\n",
    "# Compute equilibrium\n",
    "def h(p):\n",
    "    return p**(-1) - (np.exp(0.1 * p) - 1)  # demand - supply\n",
    "\n",
    "p_star = brentq(h, 2, 4)\n",
    "q_star = np.exp(0.1 * p_star) - 1\n",
    "\n",
    "print(f'Equilibrium price is {p_star: .2f}')\n",
    "print(f'Equilibrium quantity is {q_star: .2f}')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b606f29c",
   "metadata": {},
   "source": [
    "Let’s also plot our results."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "98e6deb0",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "# Now plot\n",
    "grid = np.linspace(2, 4, 100)\n",
    "fig, ax = plt.subplots()\n",
    "\n",
    "qs = np.exp(0.1 * grid) - 1\n",
    "qd = grid**(-1)\n",
    "\n",
    "\n",
    "ax.plot(grid, qd, 'b-', lw=2, label='demand')\n",
    "ax.plot(grid, qs, 'g-', lw=2, label='supply')\n",
    "\n",
    "ax.set_xlabel('price')\n",
    "ax.set_ylabel('quantity')\n",
    "ax.legend(loc='upper center')\n",
    "\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "01bfcc8e",
   "metadata": {},
   "source": [
    "We also want to consider supply and demand shifts.\n",
    "\n",
    "For example, let’s see what happens when demand shifts up, with $ \\gamma $ increasing to $ 1.25 $:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b55ba614",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "# Compute equilibrium\n",
    "def h(p):\n",
    "    return 1.25 * p**(-1) - (np.exp(0.1 * p) - 1)\n",
    "\n",
    "p_star = brentq(h, 2, 4)\n",
    "q_star = np.exp(0.1 * p_star) - 1\n",
    "\n",
    "print(f'Equilibrium price is {p_star: .2f}')\n",
    "print(f'Equilibrium quantity is {q_star: .2f}')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8f9dcd72",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "# Now plot\n",
    "p_grid = np.linspace(2, 4, 100)\n",
    "fig, ax = plt.subplots()\n",
    "\n",
    "qs = np.exp(0.1 * p_grid) - 1\n",
    "qd = 1.25 * p_grid**(-1)\n",
    "\n",
    "\n",
    "ax.plot(grid, qd, 'b-', lw=2, label='demand')\n",
    "ax.plot(grid, qs, 'g-', lw=2, label='supply')\n",
    "\n",
    "ax.set_xlabel('price')\n",
    "ax.set_ylabel('quantity')\n",
    "ax.legend(loc='upper center')\n",
    "\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "571f8af1",
   "metadata": {},
   "source": [
    "Now we might consider supply shifts, but you already get the idea that there’s\n",
    "a lot of repeated code here.\n",
    "\n",
    "Refactor and improve clarity in the code above using the principles discussed\n",
    "in this lecture."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "31482396",
   "metadata": {},
   "source": [
    "## Solution to[ Exercise 20.1](https://python-programming.quantecon.org/#wgc-exercise-1)\n",
    "\n",
    "Here’s one solution, that uses a class:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a2e07339",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "class Equilibrium:\n",
    "\n",
    "    def __init__(self, α=0.1, β=1, γ=1, δ=1):\n",
    "        self.α, self.β, self.γ, self.δ = α, β, γ, δ\n",
    "\n",
    "    def qs(self, p):\n",
    "        return np.exp(self.α * p) - self.β\n",
    "\n",
    "    def qd(self, p):\n",
    "        return self.γ * p**(-self.δ)\n",
    "\n",
    "    def compute_equilibrium(self):\n",
    "        def h(p):\n",
    "            return self.qd(p) - self.qs(p)\n",
    "        p_star = brentq(h, 2, 4)\n",
    "        q_star = np.exp(self.α * p_star) - self.β\n",
    "\n",
    "        print(f'Equilibrium price is {p_star: .2f}')\n",
    "        print(f'Equilibrium quantity is {q_star: .2f}')\n",
    "\n",
    "    def plot_equilibrium(self):\n",
    "        # Now plot\n",
    "        grid = np.linspace(2, 4, 100)\n",
    "        fig, ax = plt.subplots()\n",
    "\n",
    "        ax.plot(grid, self.qd(grid), 'b-', lw=2, label='demand')\n",
    "        ax.plot(grid, self.qs(grid), 'g-', lw=2, label='supply')\n",
    "\n",
    "        ax.set_xlabel('price')\n",
    "        ax.set_ylabel('quantity')\n",
    "        ax.legend(loc='upper center')\n",
    "\n",
    "        plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "52064557",
   "metadata": {},
   "source": [
    "Let’s create an instance at the default parameter values."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d4639827",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "eq = Equilibrium()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f26be470",
   "metadata": {},
   "source": [
    "Now we’ll compute the equilibrium and plot it."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7e05a822",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "eq.compute_equilibrium()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "42dead63",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "eq.plot_equilibrium()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1013d988",
   "metadata": {},
   "source": [
    "One of the nice things about our refactored code is that, when we change\n",
    "parameters, we don’t need to repeat ourselves:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7076a076",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "eq.γ = 1.25"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f9cc20c3",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "eq.compute_equilibrium()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ed48810e",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "eq.plot_equilibrium()"
   ]
  }
 ],
 "metadata": {
  "date": 1722309857.0648944,
  "filename": "writing_good_code.md",
  "kernelspec": {
   "display_name": "Python",
   "language": "python3",
   "name": "python3"
  },
  "title": "Writing Good Code"
 },
 "nbformat": 4,
 "nbformat_minor": 5
}