09_relationships.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Printed copies of *Elements of Data Science* are available now, with a **full color interior**, from [Lulu.com](https://www.lulu.com/shop/allen-downey/elements-of-data-science/paperback/product-9dyrwn.html)."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Relationships"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "tags": []
   },
   "source": [
    "[Click here to run this notebook on Colab](https://colab.research.google.com/github/AllenDowney/ElementsOfDataScience/blob/v1/09_relationships.ipynb)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "from os.path import basename, exists\n",
    "\n",
    "def download(url):\n",
    "    filename = basename(url)\n",
    "    if not exists(filename):\n",
    "        from urllib.request import urlretrieve\n",
    "\n",
    "        local, _ = urlretrieve(url, filename)\n",
    "        print(\"Downloaded \" + str(local))\n",
    "    return filename\n",
    "\n",
    "download('https://raw.githubusercontent.com/AllenDowney/ElementsOfDataScience/v1/utils.py')\n",
    "\n",
    "import utils"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This chapter explores relationships between variables.\n",
    "\n",
    "* We will visualize relationships using scatter plots, box plots, and violin plots,\n",
    "\n",
    "* And we will quantify relationships using correlation and simple regression.\n",
    "\n",
    "The most important lesson in this chapter is that you should always visualize the relationship between variables before you try to quantify it -- otherwise, you are likely to be misled."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Exploring relationships\n",
    "\n",
    "So far we have mostly considered one variable at a time.  Now it's time to explore relationships between variables.\n",
    "As a first example, we'll look at the relationship between height and weight.\n",
    "We'll use data from the Behavioral Risk Factor Surveillance System (BRFSS), which is run by the Centers for Disease Control.\n",
    "Based on the BRFSS data from 2021, I have created an extract with one row for each survey respondent and one column for each of the variables I selected."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "tags": []
   },
   "source": [
    "You can read about the BRFSS at  at <https://www.cdc.gov/brfss>."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "download('https://github.com/AllenDowney/ElementsOfDataScience/raw/v1/data/brfss_2021.hdf');"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "import pandas as pd\n",
    "\n",
    "brfss = pd.read_hdf('brfss_2021.hdf', 'brfss')\n",
    "brfss.shape"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Here are the first few rows."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 64,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "brfss.head()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 65,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# Solution goes here"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The BRFSS includes hundreds of variables.\n",
    "For the examples in this chapter, we'll work with just these nine.\n",
    "The ones we'll start with are `HTM4`, which records each respondent's height in centimeters, and `WTKG3`, which records weight in kilograms."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "height = brfss['HTM4']\n",
    "weight = brfss['WTKG3']"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "To visualize the relationship between these variables, we'll make a **scatter plot**, which shows one marker for each pair of values.\n",
    "Scatter plots are common and readily understood, but they are surprisingly hard to get right.\n",
    "As a first attempt, we'll use `plot` with the style string `o`, which plots a circle for each data point."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "import matplotlib.pyplot as plt\n",
    "\n",
    "plt.plot(height, weight, 'o')\n",
    "\n",
    "plt.xlabel('Height in cm')\n",
    "plt.ylabel('Weight in kg')\n",
    "plt.title('Scatter plot of weight versus height');"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Each marker represents the height and weight of one person.\n",
    "Based on the shape of the result, it looks like taller people are heavier, but there are a few things about this plot that make it hard to interpret.\n",
    "Most importantly, it is **overplotted**, which means that there are markers piled on top of each other so you can't tell where there are a lot of data points and where there is just one.\n",
    "When that happens, the picture can be misleading.\n",
    "\n",
    "One way to improve the plot is to use transparency, which we can do with the keyword argument `alpha`.\n",
    "The lower the value of alpha, the more transparent each data point is.  \n",
    "Here's what it looks like with `alpha=0.01`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.plot(height, weight, 'o', alpha=0.01)\n",
    "\n",
    "plt.xlabel('Height in cm')\n",
    "plt.ylabel('Weight in kg')\n",
    "plt.title('Scatter plot of weight versus height');"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This is better, but there are so many data points, the scatter plot is still overplotted.  The next step is to make the markers smaller.\n",
    "With `markersize=0.5` and a low value of alpha, the scatter plot is less saturated.\n",
    "Here's what it looks like."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.plot(height, weight, 'o', alpha=0.01, markersize=0.5)\n",
    "\n",
    "plt.xlabel('Height in cm')\n",
    "plt.ylabel('Weight in kg')\n",
    "plt.title('Scatter plot of weight versus height');"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Again, this is better, but now we can see that the points fall in discrete columns.\n",
    "That's because most heights were reported in inches and converted to centimeters.\n",
    "We can break up the columns by adding random noise to the values, which is called **jittering**.\n",
    "We'll can use NumPy to generate noise from a normal distribution with mean 0 and standard deviation 2."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "\n",
    "noise = np.random.normal(0, 2, size=len(brfss))\n",
    "height_jitter = height + noise"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Here's what the plot looks like with jittered heights."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.plot(height_jitter, weight, 'o', alpha=0.01, markersize=0.5)\n",
    "\n",
    "plt.xlabel('Height in cm')\n",
    "plt.ylabel('Weight in kg')\n",
    "plt.title('Scatter plot of weight versus height');"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The columns are gone, but now we can see that there are rows where people rounded off their weight.\n",
    "We can fix that by jittering weight, too."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [],
   "source": [
    "noise = np.random.normal(0, 2, size=len(brfss))\n",
    "weight_jitter = weight + noise"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Here's what it looks like."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.plot(height_jitter, weight_jitter, 'o', \n",
    "         alpha=0.01, markersize=0.5)\n",
    "\n",
    "plt.xlabel('Height in cm')\n",
    "plt.ylabel('Weight in kg')\n",
    "plt.title('Scatter plot of weight versus height');"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Finally, let's zoom in on the area where most of the data points are.\n",
    "The functions `xlim` and `ylim` set the lower and upper limits of the x-axis and y-axis; in this case, we plot heights from 140 to 200 centimeters and weights up to 160 kilograms.\n",
    "Here's what it looks like."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.plot(height_jitter, weight_jitter, 'o', alpha=0.01, markersize=0.5)\n",
    "\n",
    "plt.xlim([140, 200])\n",
    "plt.ylim([0, 160])\n",
    "plt.xlabel('Height in cm')\n",
    "plt.ylabel('Weight in kg')\n",
    "plt.title('Scatter plot of weight versus height');"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now we have a reliable picture of the relationship between height and weight.\n",
    "Below you can see the misleading plot we started with and the more reliable one we ended with.\n",
    "They are clearly different, and they suggest different relationships between these variables."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# Set the figure size\n",
    "plt.figure(figsize=(8.75, 3.2))\n",
    "\n",
    "# Create subplots with 2 rows, 1 column, and start plot 1\n",
    "plt.subplot(1, 2, 1)\n",
    "plt.plot(height, weight, 'o')\n",
    "\n",
    "plt.xlabel('Height in cm')\n",
    "plt.ylabel('Weight in kg')\n",
    "plt.title('Scatter plot of weight versus height')\n",
    "\n",
    "# Adjust the layout so the two plots don't overlap\n",
    "plt.tight_layout()\n",
    "\n",
    "# Start plot 2\n",
    "plt.subplot(1, 2, 2)\n",
    "\n",
    "plt.plot(height_jitter, weight_jitter, 'o', alpha=0.01, markersize=0.5)\n",
    "\n",
    "plt.xlim([140, 200])\n",
    "plt.ylim([0, 160])\n",
    "plt.xlabel('Height in cm')\n",
    "plt.ylabel('Weight in kg')\n",
    "plt.title('Scatter plot of weight versus height')\n",
    "plt.tight_layout()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The point of this example is that it takes some effort to make an effective scatter plot."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**Exercise:** Do people tend to gain weight as they get older?\n",
    "We can answer this question by visualizing the relationship between weight and age. \n",
    "But before we make a scatter plot, it is a good idea to visualize distributions one variable at a time.\n",
    "So let's look at the distribution of age.\n",
    "\n",
    "The BRFSS dataset includes a column, `AGE`, which represents each respondent's age in years. To protect respondents' privacy, ages are rounded off into 5-year bins.\n",
    "The values in `AGE` are the midpoints of the bins.\n",
    "\n",
    "- Extract the variable `'AGE'` from `brfss` and assign it to `age`.\n",
    "\n",
    "- Plot the PMF of `age` as a bar chart, using `Pmf` from `empiricaldist`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "try:\n",
    "    import empiricaldist\n",
    "except ImportError:\n",
    "    !pip install empiricaldist\n",
    "\n",
    "from empiricaldist import Pmf"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Solution goes here"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**Exercise:** Now let's look at the distribution of weight.\n",
    "The column that contains weight in kilograms is `WTKG3`.\n",
    "Because this column contains many unique values, displaying it as a PMF doesn't work very well.\n",
    "Instead, use `Cdf` from `empiricaldist` to plot the CDF of weight."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "Pmf.from_seq(weight).bar()\n",
    "\n",
    "plt.xlabel('Weight in kg')\n",
    "plt.ylabel('PMF')\n",
    "plt.title('Distribution of weight');"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Solution goes here"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**Exercise:** Now make a scatter plot of `weight` versus `age`.\n",
    "Adjust `alpha` and `markersize` to avoid overplotting.\n",
    "Use `ylim` to limit the y-axis from 0 to 200 kilograms."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Solution goes here"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**Exercise:** In the previous exercise, the ages fall in columns because they've been rounded into 5-year bins.\n",
    "If we jitter them, the scatter plot will show the relationship more clearly.\n",
    "\n",
    "- Add random noise to `age` with mean `0` and standard deviation `2.5`.\n",
    "- Make a scatter plot and adjust `alpha` and `markersize` again."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Solution goes here"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Visualizing relationships\n",
    "\n",
    "In the previous section we used scatter plots to visualize the relationship between weight and height.\n",
    "In this section, we'll explore the relationship between weight and *age*, and we'll see two new ways to visualize relationships: box plots and violin plots.\n",
    "\n",
    "Let's start with a scatter plot."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {},
   "outputs": [],
   "source": [
    "age = brfss['AGE']\n",
    "noise = np.random.normal(0, 1.0, size=len(brfss))\n",
    "age_jitter = age + noise\n",
    "\n",
    "plt.plot(age_jitter, weight_jitter, 'o', alpha=0.01, markersize=0.5)\n",
    "\n",
    "plt.xlabel('Age in years')\n",
    "plt.ylabel('Weight in kg')\n",
    "plt.title('Weight versus age');"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The ages are jittered, but the standard deviation of the noise is small enough that there's still space between the columns.\n",
    "That makes it possible to see the shape of the distribution in each age group, as well as the differences between groups.\n",
    "\n",
    "If we take this idea one step farther, we can use KDE to estimate the density function in each column and plot it.\n",
    "This way of visualizing the data is called a **violin plot**.\n",
    "Seaborn provides a function that makes violin plots, but before we can use it, we have to get rid of any rows with missing data.\n",
    "Here's how:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "metadata": {},
   "outputs": [],
   "source": [
    "data = brfss.dropna(subset=['AGE', 'WTKG3'])\n",
    "data.shape"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "`dropna` creates a new `DataFrame` that omits the rows in `brfss` where `AGE` or `WTKG3` are `NaN`.\n",
    "Now we can call `violinplot` like this:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "metadata": {},
   "outputs": [],
   "source": [
    "import seaborn as sns\n",
    "\n",
    "sns.violinplot(x='AGE', y='WTKG3', data=data, inner=None)\n",
    "\n",
    "plt.xlabel('Age in years')\n",
    "plt.ylabel('Weight in kg')\n",
    "plt.title('Weight versus age');"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "tags": []
   },
   "source": [
    "Documentation of `violinplot` is at <https://seaborn.pydata.org/generated/seaborn.violinplot.html>."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The `x` and `y` arguments are the names of columns from `data`, which is the `DataFrame` we just created.\n",
    "The argument `inner=None` simplifies the plot.\n",
    "In the result each shape shows the distribution of weight in one age group.\n",
    "The width of each shape is proportional to the estimated density, so it's like two vertical KDEs plotted back to back.\n",
    "\n",
    "Another, related way to look at data like this is called a **box plot**, which shows summary statistics of the values in each group.\n",
    "Seaborn provides a function that makes box plots -- we can call it like this:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "metadata": {},
   "outputs": [],
   "source": [
    "sns.boxplot(x='AGE', y='WTKG3', data=data, whis=10)\n",
    "\n",
    "plt.xlabel('Age in years')\n",
    "plt.ylabel('Weight in kg')\n",
    "plt.title('Weight versus age');"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "tags": []
   },
   "source": [
    "Documentation of `boxplot` is at <https://seaborn.pydata.org/generated/seaborn.boxplot.html>."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Each box represents the distribution of weight in an age group.\n",
    "The height of each box represents the interquartile range, which is the range from the 25th to the 75th percentile.\n",
    "The line in the middle of each box is the median.\n",
    "\n",
    "The keyword argument `whis` determines the behavior of the whiskers that extend above and below the boxes.\n",
    "With `whis=10`, they extend far enough to show the minimum and maximum values.\n",
    "\n",
    "In my opinion, this plot gives us the best view of the relationship between weight and age.\n",
    "\n",
    "* Looking at the medians, we can see that people in their 40s are the heaviest; younger and older people are lighter.\n",
    "\n",
    "* Looking at the sizes of the boxes, it seems like people in their 40s have the most variability in weight, too.\n",
    "\n",
    "* Looking at the whiskers, we can see that the distribution of weight is skewed -- that is, the heaviest people are farther from the median than the lightest people.\n",
    "\n",
    "When a distribution is skewed toward higher values, it is sometimes useful to plot it on a logarithmic scale.\n",
    "We can do that with the Pyplot function `yscale`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "metadata": {},
   "outputs": [],
   "source": [
    "sns.boxplot(x='AGE', y='WTKG3', data=data, whis=10)\n",
    "\n",
    "plt.yscale('log')\n",
    "plt.xlabel('Age in years')\n",
    "plt.ylabel('Weight in kg (log scale)')\n",
    "plt.title('Weight versus age');"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "On a log scale, the distributions are symmetric, so the whiskers extend about the same distance in both directions, the boxes are close to the middle of the figure, and we can see the relationship between age and weight clearly.\n",
    "\n",
    "In the following exercises, you can generate violin and box plots for other variables."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**Exercise:** Previously we looked at a scatter plot of height and weight, and saw that taller people tend to be heavier.\n",
    "Now let's take a closer look using a box plot.\n",
    "The `brfss` DataFrame contains a column named `_HTMG10` that represents height in centimeters, binned into 10 cm groups.\n",
    "\n",
    "- Make a box plot that shows the distribution of weight in each height group.\n",
    "\n",
    "- Plot the y-axis on a logarithmic scale.\n",
    "\n",
    "Suggestion: If the labels on the x-axis collide, you can rotate them with `plt.xticks(rotation=45)`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Solution goes here"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**Exercise:** As a second example, let's look at the relationship between income and height.  In the BRFSS, income is represented as a categorical variable -- that is, respondents are assigned to one of seven income categories.\n",
    "The column name is `_INCOMG1`.\n",
    "Before we connect income with anything else, let's look at the distribution by computing the PMF.\n",
    "Extract `_INCOMG1` from `brfss` and assign it to `income`.\n",
    "Then plot the PMF of `income` as a bar chart."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Solution goes here"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**Exercise:** Generate a violin plot that shows the distribution of height in each income group.\n",
    "Can you see a relationship between these variables?"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Solution goes here"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Quantifying Correlation\n",
    "\n",
    "In the previous section, we visualized relationships between pairs of variables.\n",
    "Now we'll quantify the strength of those relationships by computing their correlation.\n",
    "\n",
    "When people say \"correlation\" casually, they might mean any relationship between two variables.\n",
    "In statistics, it usually means a **correlation coefficient**, which is a number between -1 and 1 that quantifies the strength of a linear relationship between variables.\n",
    "To demonstrate, we'll select three columns from the BRFSS dataset:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "metadata": {},
   "outputs": [],
   "source": [
    "columns = ['HTM4', 'WTKG3', 'AGE']\n",
    "subset = brfss[columns]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The result is a `DataFrame` with just those columns.\n",
    "With this subset of the data, we can use the `corr` method, like this:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 30,
   "metadata": {},
   "outputs": [],
   "source": [
    "subset.corr()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The result is a **correlation matrix**.\n",
    "Reading across the first row, the correlation of `HTM4` with itself is 1.\n",
    "That's expected -- the correlation of anything with itself is 1.\n",
    "\n",
    "The next entry is more interesting: the correlation of height and weight is about 0.45.\n",
    "It's positive, which means taller people are heavier, and it's moderate in strength, which means it has some predictive value -- if you know someone's height, you can make a somewhat better guess about their weight.\n",
    "\n",
    "The correlation between height and age is about -0.09.\n",
    "It's negative, which means that older people tend to be shorter, but it's weak, which means that knowing someone's age would not help much if you were trying to guess their height."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Reading across the second row, we can see that the correlation of height and weight is the same as the correlation of weight and height -- because correlation is commutative.\n",
    "\n",
    "And we can see that the correlation of weight and age is only 0.001, which is very small.\n",
    "It is tempting to conclude that there is no relationship between age and weight, but we have already seen that there is.\n",
    "So why is the correlation so low?\n",
    "Remember that the relationship between weight and age looks like this."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 31,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "data = brfss.dropna(subset=['AGE', 'WTKG3'])\n",
    "sns.boxplot(x='AGE', y='WTKG3', data=data, whis=10)\n",
    "\n",
    "plt.xlabel('Age in years')\n",
    "plt.ylabel('Weight in kg')\n",
    "plt.title('Weight versus age');"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "As age increases, weight goes up and then down, so the relationship is nonlinear.\n",
    "But correlation only measures linear relationships.\n",
    "When the relationship is nonlinear, correlation generally underestimates how strong it is.\n",
    "\n",
    "To demonstrate this point more clearly, let's generate some fake data: `xs` contains equally-spaced points between -1 and 1; `ys` is `xs` squared plus some random noise."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 32,
   "metadata": {},
   "outputs": [],
   "source": [
    "xs = np.linspace(-1, 1)\n",
    "ys = xs**2 + np.random.normal(0, 0.05, len(xs))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Here's the scatter plot of these values."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 33,
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.plot(xs, ys, 'o', alpha=0.5)\n",
    "plt.xlabel('x')\n",
    "plt.ylabel('y')\n",
    "plt.title('Scatter plot of a fake dataset');"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This is a strong relationship in the sense that you can make a much better guess about `y` if you are given `x`.\n",
    "But here's the correlation matrix:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 34,
   "metadata": {},
   "outputs": [],
   "source": [
    "np.corrcoef(xs, ys)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Even though there is a strong non-linear relationship, the computed correlation is close to zero.\n",
    "In general, if correlation is high -- that is, close to 1 or -1 -- you can conclude that there is a strong linear relationship.\n",
    "But if correlation is close to zero, that doesn't mean there is no relationship; there might be a non-linear relationship.\n",
    "This is one reason correlation can be misleading.\n",
    "\n",
    "And there's another reason to be careful with correlation -- it doesn't mean what people take it to mean.\n",
    "Specifically, correlation says nothing about the slope of the line that fits the data.\n",
    "If we say that two variables are correlated, that means we can use one to predict the other.  But that might not be what we care about.\n",
    "\n",
    "For example, suppose we are concerned about the health effects of weight gain, so we plot weight versus age from 20 to 50 years old.\n",
    "I'll generate two fake datasets to demonstrate the point.\n",
    "In each dataset, `xs` represents age and `ys` represents weight."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "tags": []
   },
   "source": [
    "I use `np.random.seed` to initialize the random number generator so we get the same results every time we run."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 35,
   "metadata": {},
   "outputs": [],
   "source": [
    "np.random.seed(18)\n",
    "xs1 = np.linspace(20, 50)\n",
    "ys1 = 75 + 0.02 * xs1 + np.random.normal(0, 0.15, len(xs1))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 36,
   "metadata": {},
   "outputs": [],
   "source": [
    "np.random.seed(18)\n",
    "xs2 = np.linspace(20, 50)\n",
    "ys2 = 65 + 0.2 * xs2 + np.random.normal(0, 3, len(xs2))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Here's what the two scatter plots look like."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 37,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "plt.figure(figsize=(8.75, 3.2))\n",
    "\n",
    "plt.subplot(1, 2, 1)\n",
    "plt.plot(xs1, ys1, 'o', alpha=0.5)\n",
    "plt.xlabel('Age in years')\n",
    "plt.ylabel('Weight in kg')\n",
    "plt.title('Fake dataset #1')\n",
    "plt.tight_layout()\n",
    "\n",
    "plt.subplot(1, 2, 2)\n",
    "plt.plot(xs2, ys2, 'o', alpha=0.5)\n",
    "plt.xlabel('Age in years')\n",
    "plt.ylabel('Weight in kg')\n",
    "plt.title('Fake dataset #2')\n",
    "plt.tight_layout()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "I constructed these examples so they look similar, but they have substantially different correlations."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 38,
   "metadata": {},
   "outputs": [],
   "source": [
    "rho1 = np.corrcoef(xs1, ys1)[0, 1]\n",
    "rho1"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 39,
   "metadata": {},
   "outputs": [],
   "source": [
    "rho2 = np.corrcoef(xs2, ys2)[0, 1]\n",
    "rho2"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In the first dataset, the correlation is strong, close to 0.75.\n",
    "In the second dataset, the correlation is moderate, close to 0.5.\n",
    "So we might think the first relationship is more important.\n",
    "But look more closely at the y-axis in both figures.\n",
    "\n",
    "In the first example, the average weight gain over 30 years is less than 1 kilogram; in the second it is more than 5 kilograms!\n",
    "If we are concerned about the health effects of weight gain, the second relationship is probably more important, even though the correlation is lower.\n",
    "In this scenario, the statistic we really care about is the slope of the line that fits the data, not the coefficient of correlation.\n",
    "In the next section, we'll use linear regression to compute that slope, but first let's practice with correlation."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**Exercise:** The purpose of the BRFSS is to explore health risk factors, so it includes questions about diet.\n",
    "The column `_VEGESU1` represents the number of servings of vegetables respondents reported eating per day.\n",
    "Before we compute correlations, let's look at the distribution of this variable.\n",
    "Extract `_VEGESU1` from `brfss` and assign it to `vegesu` -- then plot the CDF of the values.\n",
    "\n",
    "The original dataset includes a small number of values greater than 10, some of them unreasonably large.\n",
    "For this extract, I have cut off the values at 10."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 40,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Solution goes here"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**Exercise:** Now let's visualize the relationship between age and vegetables.\n",
    "Make a box plot that summarizes the distribution of vegetable servings in each age group.\n",
    "How would you describe the relationship, if any?"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 41,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Solution goes here"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**Exercise:** Finally, let's look at correlations between age, income, and vegetable servings.\n",
    "\n",
    "- From `brfss`, select the columns `'AGE'`, `_INCOMG1`, and `_VEGESU1`.\n",
    "- Compute the correlation matrix for these variables.\n",
    "\n",
    "Is the correlation between age and vegetable servings what you expected based on the box plot?"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 42,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Solution goes here"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Simple Linear Regression\n",
    "\n",
    "In the previous section we saw that correlation does not always measure what we really want to know.\n",
    "In this section, we look at an alternative: simple linear regression.\n",
    "Here \"simple\" means there are only two variables, as opposed to \"multiple\" regression, which can work with any number of variables.\n",
    "\n",
    "In the previous section, I generated fake datasets showing two hypothetical relationships between weight and age.\n",
    "We computed correlations for both datasets, now let's compute lines of best fit.\n",
    "We can use `linregress` from the SciPy `stats` library, which takes two sequences as arguments.\n",
    "Here are the results for Fake Dataset #1."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 43,
   "metadata": {},
   "outputs": [],
   "source": [
    "from scipy.stats import linregress\n",
    "\n",
    "res1 = linregress(xs1, ys1)\n",
    "res1._asdict()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The result is a `LinregressResult` object that contains five values: `slope` is the slope of the line of best fit for the data, `intercept` is the intercept, and `rvalue` is correlation.\n",
    "We'll ignore the other values for now.\n",
    "\n",
    "For Fake Dataset #1, the estimated slope is about 0.019 kilograms per year or about 0.56 kilograms over the 30-year range."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 44,
   "metadata": {},
   "outputs": [],
   "source": [
    "res1.slope * 30"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Here are the results for Fake Dataset #2."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 45,
   "metadata": {},
   "outputs": [],
   "source": [
    "res2 = linregress(xs2, ys2)\n",
    "res2._asdict()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The estimated slope is almost 10 times higher, about 0.18 kilograms per year or about 5.3 kilograms per 30 years."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 46,
   "metadata": {},
   "outputs": [],
   "source": [
    "res2.slope * 30"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can use the results from `linregress` to plot the line of best fit and see how it relates to the data.\n",
    "First we get the minimum and maximum of the observed `xs`.\n",
    "Then we multiply by the slope and add the intercept."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 47,
   "metadata": {},
   "outputs": [],
   "source": [
    "low, high = xs1.min(), xs1.max()\n",
    "fx = np.array([low, high])\n",
    "fy = res1.intercept + res1.slope * fx"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Here's what the result looks like for the first example."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 48,
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.plot(xs1, ys1, 'o', alpha=0.5)\n",
    "plt.plot(fx, fy, '-')\n",
    "\n",
    "plt.xlabel('Age in years')\n",
    "plt.ylabel('Weight in kg')\n",
    "plt.title('Fake Dataset #1');"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can do the same thing for the second example."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 49,
   "metadata": {},
   "outputs": [],
   "source": [
    "low, high = xs2.min(), xs2.max()\n",
    "fx = np.array([low, high])\n",
    "fy = res2.intercept + res2.slope * fx"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "And here's what the result looks like."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 50,
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.plot(xs2, ys2, 'o', alpha=0.5)\n",
    "plt.plot(fx, fy, '-')\n",
    "\n",
    "plt.xlabel('Age in years')\n",
    "plt.ylabel('Weight in kg')\n",
    "plt.title('Fake Dataset #2');"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "At first glance it might seem like the slope is steeper in the first figure, but don't be fooled.\n",
    "If you look closely at the vertical scales, the slope in the second figure is almost 10 times higher."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Regression of Height and Weight\n",
    "\n",
    "Now let's look at an example of regression with real data.\n",
    "Here's the scatter plot of height and weight one more time."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 51,
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.plot(height_jitter, weight_jitter, 'o', \n",
    "         alpha=0.01, markersize=0.5)\n",
    "\n",
    "plt.xlim([140, 200])\n",
    "plt.ylim([0, 160])\n",
    "plt.xlabel('Height in cm')\n",
    "plt.ylabel('Weight in kg')\n",
    "plt.title('Scatter plot of weight versus height');"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "To compute the regression line, we'll use `linregress` again.\n",
    "But it can't handle `NaN` values, so we have to use `dropna` to remove rows that are missing the data we need."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 52,
   "metadata": {},
   "outputs": [],
   "source": [
    "data = brfss.dropna(subset=['WTKG3', 'HTM4'])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now we can compute the linear regression."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 53,
   "metadata": {},
   "outputs": [],
   "source": [
    "res_hw = linregress(data['HTM4'], data['WTKG3'])\n",
    "res_hw._asdict()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The slope is about 0.9 kilograms per centimeter, which means that we expect a person one centimeter taller to be almost a kilogram heavier.\n",
    "That's quite a lot.\n",
    "As before, we can compute the line of best fit:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 54,
   "metadata": {},
   "outputs": [],
   "source": [
    "low, high = data['HTM4'].min(), data['HTM4'].max()\n",
    "fx = np.array([low, high])\n",
    "fy = res_hw.intercept + res_hw.slope * fx"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "And here's what that looks like.  "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 55,
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.plot(height_jitter, weight_jitter, 'o', alpha=0.01, markersize=0.5)\n",
    "plt.plot(fx, fy, '-')\n",
    "\n",
    "plt.xlim([140, 200])\n",
    "plt.ylim([0, 160])\n",
    "plt.xlabel('Height in cm')\n",
    "plt.ylabel('Weight in kg')\n",
    "plt.title('Scatter plot of weight versus height');"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The slope of this line seems consistent with the scatter plot.\n",
    "As another example, here's the scatter plot of weight versus age, which we saw earlier."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 56,
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.plot(age_jitter, weight_jitter, 'o', \n",
    "         alpha=0.01, markersize=0.5)\n",
    "\n",
    "plt.ylim([0, 160])\n",
    "plt.xlabel('Age in years')\n",
    "plt.ylabel('Weight in kg')\n",
    "plt.title('Weight versus age');"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "As we have seen before, the relationship is nonlinear.\n",
    "Let's see what we get if we compute a linear regression."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 57,
   "metadata": {},
   "outputs": [],
   "source": [
    "subset = brfss.dropna(subset=['WTKG3', 'AGE']) \n",
    "res_aw = linregress(subset['AGE'], subset['WTKG3'])\n",
    "res_aw._asdict()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The estimated slope is close to zero.\n",
    "Here's what the line of best fit looks like."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 58,
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.plot(age_jitter, weight_jitter, 'o', \n",
    "         alpha=0.01, markersize=0.5)\n",
    "\n",
    "low, high = data['AGE'].min(), data['AGE'].max()\n",
    "fx = np.array([low, high])\n",
    "fy = res_aw.intercept + res_aw.slope * fx\n",
    "plt.plot(fx, fy, '-')\n",
    "\n",
    "plt.ylim([0, 160])\n",
    "plt.xlabel('Age in years')\n",
    "plt.ylabel('Weight in kg')\n",
    "plt.title('Weight versus age');"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "A straight line does not capture the relationship between these variables well.\n",
    "That's because linear regression has the same problem as correlation -- it only measures the strength of a linear relationship.\n",
    "\n",
    "In the next chapter, we'll see how to use multiple regression to quantify non-linear relationships.\n",
    "But first, let's practice using simple regression."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**Exercise:** Who do you think eats more vegetables, people with low income, or people with high income?\n",
    "Let's find out.\n",
    "As we've seen previously, the column `_INCOMG1` represents income level and `_VEGESU1` represents the number of vegetable servings respondents reported eating per day.\n",
    "Make a scatter plot with vegetable servings versus income, that is, with vegetable servings on the y-axis and income group on the x-axis.\n",
    "You might want to use `ylim` to zoom in on the bottom half of the y-axis."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 59,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Solution goes here"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**Exercise:** Now estimate the slope of the relationship between vegetable consumption and income.\n",
    "\n",
    "- Use `dropna` to select rows where `_INCOMG1` and `_VEGESU1` are not `NaN`.\n",
    "\n",
    "- Extract `_INCOMG1` and `_VEGESU1` and compute the simple linear regression of these variables.\n",
    "\n",
    "- Finally, plot the regression line on top of the scatter plot.\n",
    "\n",
    "What is the slope of the regression line?  What does this slope means in the context of the question we are exploring?"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 60,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Solution goes here"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 61,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Solution goes here"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 62,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Solution goes here"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 63,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Solution goes here"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Summary\n",
    "\n",
    "This chapter presents three ways to visualize the relationship between two variables: a scatter plot, violin plot, and box plot.\n",
    "A scatter plot is often a good choice when you are exploring a new data set, but it can take some attention to avoid overplotting.\n",
    "Violin plots and box plots are particularly useful when one of the variables has only a few unique values or the values have been rounded into bins.\n",
    "\n",
    "We considered two ways to quantify the strength of a relationship: the coefficient of correlation and the slope of a regression line.\n",
    "These statistics capture different aspect of what we might mean by \"strength\".\n",
    "The coefficient of correlation indicates how well we can predict one variable, given the other.\n",
    "The slope of the regression line indicates how much difference we expect in one variable as we vary the other.\n",
    "One or the other might be more relevant, depending on the context."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "tags": []
   },
   "source": [
    "*Elements of Data Science*\n",
    "\n",
    "Copyright 2021 [Allen B. Downey](https://allendowney.com)\n",
    "\n",
    "License: [Creative Commons Attribution-NonCommercial-ShareAlike 4.0 International](https://creativecommons.org/licenses/by-nc-sa/4.0/)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "celltoolbar": "Tags",
  "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.10.14"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}