A Two-Asset Savings Model with an Income-Contribution Scheme

Mateo Velásquez-Giraldo

 mvelasq2@jhu.edu

 Johns Hopkins University

This notebook demonstrates the use of the `RiskyContrib` agent type
of the [HARK Toolkit](https://econ-ark.org/toolkit). The model represents an agent who can
save using two different assets---one risky and the other risk-free---to insure
against fluctuations in his income, but faces frictions to transferring funds between
assets. The flexibility of its implementation and its inclusion in the HARK
toolkit will allow users to adapt the model to realistic life-cycle calibrations, and
also to embedded it in heterogeneous-agents macroeconomic models.

```python
# Preamble

from HARK.ConsumptionSaving.ConsRiskyContribModel import (
    RiskyContribConsumerType,
    init_risky_contrib_lifecycle,
)
from time import time
from copy import copy

import numpy as np
import seaborn as sns
import pandas as pd
from Simulations.tools import pol_funcs_dframe, age_profiles
import os

# This is a jupytext paired notebook that autogenerates a .py file
# which can be executed from a terminal command line
# But a terminal does not permit inline figures, so we need to test jupyter vs terminal
# Google "how can I check if code is executed in the ipython notebook"
from IPython import get_ipython # In case it was run from python instead of ipython

# If the ipython process contains 'terminal' assume not in a notebook
def in_ipynb():
    try:
        if 'terminal' in str(type(get_ipython())):
            return False
        else:
            return True
    except NameError:
        return False

# Determine whether to make the figures inline (for spyder or jupyter)
# vs whatever is the automatic setting that will apply if run from the terminal
if in_ipynb():
    # %matplotlib inline generates a syntax error when run from the shell
    # so do this instead
    get_ipython().run_line_magic('matplotlib', 'inline')
else:
    get_ipython().run_line_magic('matplotlib', 'auto')
```

## Model Description

I now discuss the main components of the model informally, and leave its full
recursive mathematical representation for Section \ref{sec:recursive}.

### Time, mortality, and utility

Time advances in discrete steps that I will index with $t$. The model can
be used in both infinite and finite-horizon versions.

Agents face an exogenous risk of death $\delta_t$ each period, which becomes certain at the
maximum age of the finite-horizon version. There are no intentional bequests; agents
will consume all of their resources if they reach the last period, but they can leave
accidental bequests upon premature death.

In each period, agents derive utility from consumption only. Their utility function
follows a constant relative risk aversion specification. Formally, for a level of
consumption $C$, the agent derives instant utility
\begin{equation}
	u(C) = \frac{C^{1-\rho}}{1- \rho}.
\end{equation}

#### Income process

Agents supply labor inelastically. Their labor earnings $Y_{i,t}$ are the product of a
permanent component $P_{i,t}$ and a transitory stochastic component $\theta_{i,t}$ as
in \cite{Carroll1997qje},  where $i$ indexes different agents. Formally,
\begin{equation*}
\begin{split}
\ln Y_{i,t} &= \ln P_{i,t} + \ln \theta_{i,t} \\
\ln P_{i,t} &= \ln P_{i,t-1} + \ln \Gamma_{i,t} + \ln \psi_{i,t}
\end{split}
\end{equation*}
where $\Gamma_{i,t}$ is a deterministic growth factor that can capture
life-cycle patterns in earnings, and
$\ln \psi_{i,t}\sim \mathcal{N}(-\sigma^2_{\psi,t}/2, \sigma_{\psi,t})$
is a multiplicative shock to permanent income\footnote{The mean of the shock is set so that $E[\psi_{i,t}] = 1$.}.

The transitory component $\theta_{i,t}$ is a mixture that models unemployment and
other temporal fluctuations in income as
\begin{equation*}
\ln\theta_{i,t} = \begin{cases}
\ln \mathcal{U}, & \text{With probability } \mho\\
\ln \tilde{\theta}_{i,t}\sim\mathcal{N}(-\sigma^2_{\theta,t}/2, \sigma_{\theta,t}), & \text{With probability } 1-\mho,
\end{cases}
\end{equation*}
with $\mho$ representing the probability of unemployment and $\mathcal{U}$ the replacement
factor of unemployment benefits.

This specification of the income process is parsimonious and flexible enough to accommodate
life-cycle patterns in income growth and volatility, transitory unemployment and exogenous
retirement. Introduced by \cite{Carroll1997qje}, this income specification is common in studies
of life-cycle wealth accumulation and portfolio choice; see e.g.,
\cite{Cagetti2003jbes,Cocco2005rfs,Fagereng2017jof}. The specification has
also been used in studies of income volatility such as \cite{Carroll1992bpea,Carroll1997jme,Sabelhaus2010jme}, which have yielded calibrations of its stochastic shocks' distributions.

#### Financial assets and frictions

Agents smooth their consumption by saving and have two assets
available for this purpose. The first is a risk-free liquid account with
constant per-period return factor $R$. The second has a stochastic return
factor $\tilde{R}$ that is log-normally distributed and independent across
time. Various interpretations such as stocks, a retirement fund, or entrepreneurial
capital could be given to the risky asset. Importantly, consumption must be paid for
using funds from the risk-free account. The levels of risk-free and risky assets
owned by the agent will both be state variables, denoted with $M_{i,t}$ and $N_{i,t}$
respectively.

Portfolio rebalancing takes place by moving funds between the risk-free
and risky accounts. These flows are one of the agents' control variables
and are denoted as $D_{i,t}$, with $D_{i,t}>0$ representing a movement of
funds from the risk-free to the risky account. Withdrawals from the risky
account are subject to a constant-rate tax $\tau$ which can represent, for
instance, capital-gains realization taxes or early retirement-fund withdrawal
penalties. In sum, denoting post-rebalancing asset levels with $\tilde{\cdot}$,
\begin{equation*}
\begin{split}
\tilde{M}_{i,t} &= M_{i,t} - D_{i,t}(1 - 1_{[D_{i,t}\leq0]}\tau)\\
\tilde{N}_{i,t} &= N_{i,t} + D_{i,t}.
\end{split}
\end{equation*}

At any given period, an agent might not be able to rebalance his portfolio.
This ability is governed by an exogenous stochastic shock that is realized
at the start of the period

\begin{equation*}
\Adj_t \sim \text{Bernoulli}(p_t),
\end{equation*}

with $\Adj_t=1$ meaning that the agent can rebalance and $\NAdj_t=1$ ($\Adj_t = 0$)
forcing him to set $D_{i,t} = 0$. This friction is a parsimonious way to capture
the fact that portfolio rebalancing is costly and households do it sporadically.
Recent studies have advocated for \cite{Giglio2021aer} and used
\cite{Luetticke2021aej_macro} this kind of rebalancing friction.

To partially evade the possibility of being unable to rebalance their accounts, agents
can use an income deduction scheme. By default, labor income ($Y_{i,t}$) is deposited to
the risk-free liquid account at the start of every period. However, agents can pre-commit
to have a fraction  $\Contr_t\in[0,1]$ of their income diverted to their risky account instead.
This fraction can be tweaked by the agent whenever $\Adj_t = 1$; otherwise it stays at its
previous value, $\Contr_{t+1} = \Contr_t$.

#### Timing

The previous figure summarizes the timing of stochastic shocks and
optimizing decisions that occur within a period of the life cycle model.

### Recursive representation of the model

Individual subscripts $i$ are dropped for simplicity. The value function for
an agent who is not allowed to rebalance his portfolio at time $t$ is

\begin{equation*}
\begin{split}
V^{\NAdj}_{t}(M_t, N_t, P_t, \Contr_t) = \max_{C_t} u(C_t)
+ p_{t+1} &\beta\delta_{t+1} E_t \left[  V^{\Adj}_{t+1}\left( M_{t+1}, N_{t+1},
P_{t+1} \right)\right] +\\
\left(1-p_{t+1}\right) &\beta\delta_{t+1} E_t\left[V^{\NAdj}_{t+1}\left(M_{t+1},
N_{t+1}, P_{t+1}, \Contr_{t+1}\right) \right]\\
\text{Subject to:} \quad &\\
0\leq& C_t \leq M_t \\
A_t =& M_t - C_t \\
M_{t+1} =& R A_t + (1-\Contr_{t+1}) Y_{t+1}\\
N_{t+1} =& \tilde{R}_{t+1}N_t + \Contr_{t+1}Y_{t+1}\\
P_{t+1} =& \Gamma_{t+1} \psi_{t+1} P_{t}\\
Y_{t+1} =& \theta_{t+1} P_{t+1}\\
\Contr_{t+1} =& \Contr_t
\end{split}
\end{equation*}

and that of agent who is allowed to rebalance is

\begin{equation*}
\begin{split}
V^{\Adj}_{t}(M_t, N_t, P_t) = \max_{C_t,D_t,\Contr_{t+1}}
u(C_t) + p_{t+1} &\beta\delta_{t+1} E_t \left[  V^{\Adj}_{t+1}\left( M_{t+1},
N_{t+1}, P_{t+1} \right)\right] +\\
\left(1-p_{t+1}\right) &\beta\delta_{t+1} E_t\left[V^{\NAdj}_{t+1}\left(M_{t+1},
N_{t+1}, P_{t+1}, \Contr_{t+1}\right) \right]\\
\text{Subject to:} \quad &\\
\quad -N_t \leq D_t \leq M_t, \quad \Contr_{t+1} \in& [0,1], \quad 0 \leq C_t \leq \tilde{M}_t\\
\hfill\\
\tilde{M}_t =& M_t - D_t\left(1-1_{[D_t\leq0]}\tau\right)\\
\tilde{N}_t =& N_t + D_t\\
A_t =& \tilde{M}_t - C_t \\
M_{t+1} =& R A_t + (1-\Contr_{t+1}) Y_{t+1}\\
N_{t+1} =& \tilde{R}_{t+1} \tilde{N}_t + \Contr_{t+1}Y_{t+1}\\
P_{t+1} =& \Gamma_{t+1}\psi_{t+1} P_{t}\\
Y_{t+1} =& \theta_{t+1} P_{t+1}
\end{split}
\end{equation*}

## Parametrizations

This notebook will only demonstrate a life-cycle calibration of the model. However, the current implementation of the model is able to find and use the solution to infinite-horizon formulations (see the document in this repository for details).

For the present exercise, I calibrate the model's mortality and income paths to represent  individuals who enter the model at age 25, retire exogenously at 65, and live to a maximum age of 90. Survival probabilities ($\delta$) come from the 2004 SSA life-tables for males. Income growth factors and volatilities ($\Gamma$, $\sigma_\psi$, $\sigma_\theta$) come from the calibration for high-school graduates in \cite{Cagetti2003jbes}.

To highlight the effect of different financial frictions on wealth accumulation and portfolio choice, I consider different configurations for the risky-withdrawal tax $\tau$ and the probability of being able to rebalance $p$:

\begin{itemize}
\item \textbf{Base}: the probability of being able to rebalance is $p = 1$
and the risky withdrawal tax rate is $\tau = 0$, both constant throughout the agents' lives.

\item \textbf{Tax}: the risky withdrawal tax is constant at $10\%$ and the agents
can always rebalance their portfolios.

\item \textbf{Calvo}: there is no risky withdrawal tax, but there is only a $25\%$ chance
that agents can rebalance their portfolios every year.

\item \textbf{Retirement}: there is no risky withdrawal tax, but the agents' ability
to rebalance their portfolio is time-varying; they can rebalance their assets and pick
their income-deduction scheme for certain when they enter the model at age 25, but
then they have no chance of doing so again ($p=0$) until they retire. After retirement
at age 65, they can always rebalance their portfolio ($p=1$).
\end{itemize}

The rest of the parameters take the following values

| Name in HARK | Mathematical Symbol   | Value   |
|--------------|-----------------------|---------|
| `CRRA`       | $\rho$                | $5.0$   |
| `Rfree`      | $R$                   | $1.03$  |
| `DiscFac`    | $\beta$               | $0.9$   |
| `UnempPrb`   | $\mho$                | $0.05$  |
| `IncUnemp`   | $\mathcal{U}$         | $0.3$   |
| `RiskyAvg`   | $E[\tilde{R}]$        | $1.08$  |
| `RiskyStd`   | $\sqrt{V[\tilde{R}]}$ | $0.18$  |

```python
# Base calibration setup

# Make the problem life-cycle
par_LC_base = init_risky_contrib_lifecycle.copy()

# Turn off aggregate growth
par_LC_base['PermGroFacAgg'] = 1.0

# and frictionless to start
par_LC_base["AdjustPrb"] = 1.0
par_LC_base["tau"] = 0.0

# Make contribution shares a continuous choice
par_LC_base["DiscreteShareBool"] = False
par_LC_base["vFuncBool"] = False

# Make grids go up to higher levels of resources
# (one of the calibration produces high levels of nNrm)
par_LC_base.update({"mNrmMax": 500, "nNrmMax":1000})
```

```python
# Alternative calibrations

# Tax
par_LC_tax = copy(par_LC_base)
par_LC_tax["tau"] = 0.1

# Calvo
par_LC_calvo = copy(par_LC_base)
par_LC_calvo["AdjustPrb"] = 0.25

# Retirement
par_LC_retirement = copy(par_LC_base)
par_LC_retirement["AdjustPrb"] = [1.0] + [0.0]*39 + [1.0]*25
par_LC_retirement["tau"] = [0.0]*41 + [0.0]*24
par_LC_retirement["UnempPrb"] = 0.0
```

## Solution

With the parametrizations created, I now create and solve the agents.

```python
agents = {
    "Base": RiskyContribConsumerType(**par_LC_base),
    "Tax": RiskyContribConsumerType(**par_LC_tax),
    "Calvo": RiskyContribConsumerType(**par_LC_calvo),
    "Retirement": RiskyContribConsumerType(**par_LC_retirement),
}

for agent in agents:

    print("Now solving " + agent)
    t0 = time()
    agents[agent].solve(verbose=True)
    t1 = time()
    print("Solving " + agent + " took " + str(t1 - t0) + " seconds.")
```

    Now solving Base
    Finished cycle #1 of 1 in 37.66172480583191 seconds.
    Solving Base took 37.66172480583191 seconds.
    Now solving Tax
    Finished cycle #1 of 1 in 36.05218529701233 seconds.
    Solving Tax took 36.05218529701233 seconds.
    Now solving Calvo
    Finished cycle #1 of 1 in 84.95266103744507 seconds.
    Solving Calvo took 84.95266103744507 seconds.
    Now solving Retirement
    Finished cycle #1 of 1 in 53.0900604724884 seconds.
    Solving Retirement took 53.0900604724884 seconds.

## Policy function inspection

Once agents have been solved, we can use their policy functions $\dFrac_t(m,n)$, $\Contr_t(\tilde{m},\tilde{n})$, and $c_t(\tilde{m},\tilde{n}, \Contr)$ for any period $t$.

```python
# Example
calib = "Base"
t = 0
print(agents[calib].solution[t].stage_sols["Reb"].dfracFunc_Adj(1,1))
print(agents[calib].solution[t].stage_sols["Sha"].ShareFunc_Adj(1.2,0.8))
print(agents[calib].solution[t].stage_sols["Cns"].cFunc(1.2,0.8,0.5))
```

    0.1280581440929615
    0.0
    0.7103217030091933

I now illustrate the policy functions of different calibrations for an arbitrary period graphically.

Note that the solution algorithm represents the three simultaneous decisions that an agent can take as happening sequentially in ''stages'', in the order `Rebalancing` -> `Income deduction share` -> `Consumption`. See the document in this repository for more details.

```python
# Setup
from HARK.utilities import (
    determine_platform,
    test_latex_installation,
    setup_latex_env_notebook,
)

pf = determine_platform()
try:
    latexExists = test_latex_installation(pf)
except ImportError:  # windows and MacOS requires manual install
    latexExists = False

setup_latex_env_notebook(pf, latexExists)

# General aesthetics
sns.set_style("whitegrid")
sns.set_context("paper", font_scale=1.5, rc={"lines.linewidth": 2.5})

# Parameters to feed to policy functions
t = 10
age = t + 24
mNrmGrid = np.linspace(0, 40, 100)
nNrm_vals = np.array([0.0, 20.0, 40])
Share_vals = np.array([0.0, 0.5])

# Evaluate policy functions
polfuncs = pol_funcs_dframe(agents, t, mNrmGrid, nNrm_vals, Share_vals)
```

### Rebalancing

The solution to this stage problem will be the policy function $d_t(\cdot, \cdot)$
that gives the optimal flow from risk-free to risky assets, which can be negative.
However, it is convenient to define a normalized policy function $\dFrac_t$ as
\begin{equation*}
\dFrac_t(m, n) = \begin{cases}
d_t(m,n)/m, & \text{ if } d_t(m,n) \geq 0 \\
d_t(m,n)/n, & \text{ if } d_t(m,n) < 0
\end{cases}
\end{equation*}
so that $-1 \leq \dFrac(m,n) \leq 1$ for all $(m,n)$.

```python
# Rebalancing fraction
polfuncs["$n$"] = polfuncs["n"]
g = sns.FacetGrid(
    polfuncs[polfuncs.control == "dfrac"],
    col="$n$",
    hue="model",
    height=3,
    aspect=(7 / 3) * 1 / 3,
)
g.map(sns.lineplot, "m", "value", alpha=0.7)
g.add_legend(bbox_to_anchor=[0.5, 0.0], ncol=4, title="")
g.set_axis_labels("$m$", "Rebalancing fraction: $\dFrac$")
```

### Income deduction share

```python
# After rebalancing, m and n turn to their "tilde" versions. Create ntilde
# just for seaborn's grid labels.
polfuncs["$\\tilde{n}$"] = polfuncs["n"]
polfuncs["$\\Contr$"] = polfuncs["Share"]

# Share fraction
g = sns.FacetGrid(
    polfuncs[polfuncs.control == "Share"],
    col="$\\tilde{n}$",
    hue="model",
    height=3,
    aspect=(7 / 3) * 1 / 3,
)
g.map(sns.lineplot, "m", "value", alpha=0.7)
g.add_legend(bbox_to_anchor=[0.5, 0.0], ncol=4, title="")
g.set_axis_labels("$\\tilde{m}$", r"Deduction Share: $\Contr$")
```

### Consumption

```python
# Consumption
g = sns.FacetGrid(
    polfuncs[polfuncs.control == "c"],
    col="$\\tilde{n}$",
    row="$\\Contr$",
    hue="model",
    height=3,
    aspect=(7 / 3) * 1 / 3,
)
g.map(sns.lineplot, "m", "value", alpha=0.7)
g.add_legend(bbox_to_anchor=[0.5, 0.0], ncol=4, title="")
g.set_axis_labels("$\\tilde{m}$", "Consumption: $c$")
```

## Simulation and average life-cycle profiles

With the policy functions, it is easy to simulate populations of agents. I now simulate many agents for every calibration to obtain the average lifetime profiles of relevant variables like consumption, income, and wealth in its different components.

```python
# Simulation

n_agents = 100
t_sim    = 200
profiles = []
for agent in agents:
    agents[agent].AgentCount = n_agents
    agents[agent].T_sim = t_sim
    agents[agent].track_vars = ['pLvl','t_age','Adjust',
                                'mNrm','nNrm','mNrmTilde','nNrmTilde','aNrm',
                                'cNrm', 'Share', 'dfrac']
    agents[agent].initialize_sim()
    agents[agent].simulate()
    profile = age_profiles(agents[agent])
    profile['Model'] = agent
    profiles.append(profile)
```

```python
# Plot

simdata = pd.concat(profiles)

# Latex names
simdata = simdata.rename(columns = {'pLvl': 'Perm. Income $P$',
                 'Mtilde': 'Risk-free Assets $\\tilde{M}$',
                 'Ntilde': 'Risky Assets $\\tilde{N}$',
                 'C': 'Consumption  $C$',
                 'StockShare': 'Risky Share of Savings',
                 'Share': 'Deduct. Share $\\Contr$'})

lc_means = pd.melt(simdata,
                   id_vars = ['t_age', 'Model'],
                   value_vars = ['Perm. Income $P$',
                                 'Risk-free Assets $\\tilde{M}$',
                                 'Risky Assets $\\tilde{N}$',
                                 'Consumption  $C$',
                                 'Risky Share of Savings','Deduct. Share $\\Contr$'])

lc_means['Age'] = lc_means['t_age'] + 24

# Drop the last year, as people's behavior is substantially different.
lc_means = lc_means[lc_means['Age']<max(lc_means['Age'])]

# General aesthetics
sns.set_style("whitegrid")
sns.set_context("paper", font_scale=1.5, rc={"lines.linewidth": 2.5})

g = sns.FacetGrid(
    lc_means,
    col="variable",
    col_wrap = 3,
    hue="Model",
    height=3,
    aspect=(7 / 3) * 1 / 3,
    sharey=False
)
g.map(sns.lineplot, "Age", "value", alpha=0.7)
g.add_legend(bbox_to_anchor=[0.5, 0.0], ncol=4, title="")
g.set_axis_labels("Age", "")
g.set_titles(col_template = '{col_name}')
```

# References

(Carroll, 1997) Carroll Christopher D., ``_Buffer-Stock Saving and the Life Cycle/Permanent Income Hypothesis*_'', The Quarterly Journal of Economics, vol. 112, number 1, pp. 1-55, 02 1997.  [online](https://doi.org/10.1162/003355397555109)

(Cagetti, 2003) Cagetti Marco, ``_Wealth Accumulation Over the Life Cycle and Precautionary Savings_'', Journal of Business \& Economic Statistics, vol. 21, number 3, pp. 339-353,  2003.  [online](https://doi.org/10.1198/073500103288619007

)

(Cocco, Gomes et al., 2005) Cocco Jo\~ao F., Gomes Francisco J. and Maenhout Pascal J., ``_Consumption and Portfolio Choice over the Life Cycle_'', The Review of Financial Studies, vol. 18, number 2, pp. 491-533, 02 2005.  [online](https://doi.org/10.1093/rfs/hhi017)

(Fagereng, Gottlieb et al., 2017) Fagereng Andreas, Gottlieb Charles and Guiso Luigi, ``_Asset Market Participation and Portfolio Choice over the
	Life-Cycle_'', The Journal of Finance, vol. 72, number 2, pp. 705-750,  2017.  [online](https://onlinelibrary.wiley.com/doi/abs/10.1111/jofi.12484)

(Carroll, 1992) Carroll Christopher D., ``_The Buffer-Stock Theory of Saving: Some Macroeconomic Evidence_'', Brookings Papers on Economic Activity, vol. 1992, number 2, pp. 61--156,  1992.  [online](http://www.jstor.org/stable/2534582)

(Carroll and Samwick, 1997) Carroll Christopher D. and Samwick Andrew A., ``_The nature of precautionary wealth_'', Journal of Monetary Economics, vol. 40, number 1, pp. 41-71,  1997.  [online](https://www.sciencedirect.com/science/article/pii/S0304393297000366)

(Sabelhaus and Song, 2010) Sabelhaus John and Song Jae, ``_The great moderation in micro labor earnings_'', Journal of Monetary Economics, vol. 57, number 4, pp. 391-403,  2010.  [online](https://www.sciencedirect.com/science/article/pii/S0304393210000358)

(Giglio, Maggiori et al., 2021) Giglio Stefano, Maggiori Matteo, Stroebel Johannes et al., ``_Five Facts about Beliefs and Portfolios_'', American Economic Review, vol. 111, number 5, pp. 1481-1522, May 2021.  [online](https://www.aeaweb.org/articles?id=10.1257/aer.20200243)

(Luetticke, 2021) Luetticke Ralph, ``_Transmission of Monetary Policy with Heterogeneity in Household Portfolios_'', American Economic Journal: Macroeconomics, vol. 13, number 2, pp. 1-25, April 2021.  [online](https://www.aeaweb.org/articles?id=10.1257/mac.20190064)


# Clothing Recognition Task
If this is your first time running a notebook - welcome!! Notebooks are awesome because they let us play around and experiment
with code with near-instant feedback. Some pointers:

1. To execute a cell, click on it and hit SHIFT-Enter
2. Once something is executed, the variables are in memory - inspect them!

## Getting Started
This first cell imports the necessary libraries so we can get started:

```python
import torch
import numpy as np
from PIL import Image
import torch.nn as nn
import torch.onnx as onnx
import torch.nn.functional as F
import matplotlib.pyplot as plt
from torch.utils.data import DataLoader
from torchvision import datasets, transforms

```

## Data!
 Without data we really can't do anything with machine learning. At this point we have our sharp question **can we predict a digit given a 28x28 vector of numbers?**

 Once that is all squared away we need to take a look at our data. The next cell is a helper function that visualizes the the digits (a sanity check).

The next cell downloads the standard digit dataset (called FashionMNIST). The `transform` and `target_transform` parts of this call add some conversion steps to make the data more suitable for the models we will try.

```python
clothing = datasets.FashionMNIST('data', train=True, download=True,
                        transform=transforms.Compose([
                            transforms.ToTensor(),
                            transforms.Lambda(lambda x: x.reshape(28*28))
                        ]),
                        target_transform=transforms.Compose([
                            transforms.Lambda(lambda y: torch.zeros(10, dtype=torch.float).scatter_(0, torch.tensor(y), value=1))
                        ])
                     )
```

Now to get for the actual clothing classes

```python
classes = ["T-shirt/top", "Trouser", "Pullover", "Dress", "Coat", "Sandal", "Shirt", "Sneaker", "Bag", "Ankle boot"]
```

A little helper to draw the clothing

```python
def draw_clothes(clothing):
    fig, axes = plt.subplots(7, 10, figsize=(18, 7),
                            subplot_kw={'xticks':[], 'yticks':[]},
                            gridspec_kw=dict(hspace=0.1, wspace=2.5))
    for i, ax in enumerate(axes.flat):
        X, y = clothing[i]
        ax.imshow(255 - X.reshape(28,28) * 255, cmap='gray')
        ax.set_title('{0}'.format(classes[torch.argmax(y).item()]))
```

Here is our sanity check!

```python
draw_clothes(clothing)
```

Feel free to take a look at the data by printing out `x` and/or `y`. They really are just numbers! A couple of things that might seem strange when you print them out:

1. "I thought you said the images were 784 sized vectors of 0-255???" - They were! We just normalized the vectors by dividing by 255 so the new range is 0-1 (it makes the numbers more tidy)
2. "I thought the `y` was a numerical answer?? Instead it's a 10 sized vector!" - Yes - this is called a one-hot encoding of the answer. Now there's a `1` in the index of the right answer. Again, this makes the math work a lot better for the models we will be creating.

```python
x, y = clothing[0]
print(x)
print(y)
```

# Choosing Models
Now that we have some data it's time to start picking models we think might work. This is where the science part of data-science comes in: we guess and then check if our assumptions were right. Imagine models like water pipes that have to distribute water to 10 different hoses depending on 784 knobs. These 784 knobs represent the individual pixels in the digit and the 10 hoses at the end represent the actual number (or at least the index of the one with the most water coming out of it). Our job now is to pick the plumbing in between.

The next three cells represent three different constructions in an increasingly more complex order:

1. The first is a simple linear model,
2. The second is a 3 layer Neural Network,
3. and the last is a full convolutional neural network

While it is out of the scope of this tutorial to fully explain how they work, just imagine they are basically plumbing with internal knobs that have to be tuned to produce the right water pressure at the end to push the most water out of the right
index. As you go down each cell the plumbing and corresponding internal knobs just get more complicated.

```python
class SimpleLinear(nn.Module):
    def __init__(self):
        super(SimpleLinear, self).__init__()
        self.layer1 = nn.Linear(28*28, 10)

    def forward(self, x):
        x = self.layer1(x)
        return F.softmax(x, dim=1)
```

```python
class NeuralNework(nn.Module):
    def __init__(self):
        super(NeuralNework, self).__init__()
        self.layer1 = nn.Linear(28*28, 512)
        self.layer2 = nn.Linear(512, 512)
        self.output = nn.Linear(512, 10)

    def forward(self, x):
        x = F.relu(self.layer1(x))
        x = F.relu(self.layer2(x))
        x = self.output(x)
        return F.softmax(x, dim=1)
```

```python
class CNN(nn.Module):
    def __init__(self):
        super(CNN, self).__init__()
        self.conv1 = nn.Conv2d(1, 10, kernel_size=5)
        self.conv2 = nn.Conv2d(10, 20, kernel_size=5)
        self.conv2_drop = nn.Dropout2d()
        self.fc1 = nn.Linear(320, 50)
        self.fc2 = nn.Linear(50, 10)

    def forward(self, x):
        x = x.view(-1, 1, 28, 28)
        x = F.relu(F.max_pool2d(self.conv1(x), 2))
        x = F.relu(F.max_pool2d(self.conv2_drop(self.conv2(x)), 2))
        x = x.view(-1, 320)
        x = F.relu(self.fc1(x))
        x = F.dropout(x, training=self.training)
        x = self.fc2(x)
        return F.softmax(x, dim=1)
```

# Optimizing Model Parameters
Now that we have some models it's time to optimize the internal parameters to see if it can do a good job at recognizing digits! It turns out there are some parameters that we can give the optimization algorithm to tune how it trains - these are called hyper-parameters. That's what the two variables represent below:

```python
learning_rate = 1e-3
batch_size = 64
epochs = 5
```

The `learning_rate` basically specifies how fast the algorithm will learn the model parameters. Right now you're probably thinking "let's set it to fifty million #amirite?" The best analogy for why this is a bad idea is golf. I'm a terrible golfist (is that right?) so I don't really know anything - but pretend you are trying to sink a shot (again sorry) but can only hit the ball the same distance every time. Easy right? Hit it the exact length from where you are to the hole! Done! Now pretend you don't know where the hole is but just know the general direction. Now the distance you choose actually matters. If it is too long a distance you'll miss the hole, and then when you hit it back you'll overshoot again. If the distance is too small then it will take forever to get there but for sure you'll eventually get it in. Basically you have to guess what the right distance per shot should be and then try it out. That is basically what the learning rate does for finding the "hole in one" for the right parameters (ok, I'm done with the golf stuff).

Below there are three things that make this all work:
1. **The Model** - this is the function we're making that takes in the digit vector and should return the right number
2. **The Cost Function** (sometimes called the loss function). I know I promised I was done with golf but I lied. Remember how I said in our screwy golf game you knew the general direction of the hole? The cost function tells us the distance to the hole - when it's zero we're there! In actual scientific terms, the cost function tells us how bad the model is at getting the right answer. As we take shots you should see the cost function decreasing. If this does not happen then something is wrong. At this point I would change the shot distance (or `learning_rate`) to something smaller and try again. If that doesn't work maybe change the model!
3. **The Optimizer** - this part is the bit that actually changes the model parameters. It has a sense for the direction we should be shooting and updates all of the internal numbers inside the model to find the best internal knobs to predict the right digits. In this case I am using the Binary Cross Entropy cost function because, well, I know it works. There are a ton of different cost functions you can choose from that fit a variety of different scenarios.

```python
# where to run
device = 'cuda' if torch.cuda.is_available() else 'cpu'
print('Using {} device'.format(device))

# selected model - you can uncomment
# the other models for future runs
model = SimpleLinear().to(device)
#model = NeuralNework().to(device)
#model = CNN().to(device)
print(model)

# cost function used to determine best parameters
cost = torch.nn.BCELoss()

# used to create optimal parameters
optimizer = torch.optim.Adam(model.parameters(), lr=learning_rate)
```

I forgot to mention the whole `cuda` thing. GPUs speed this whole process up because this is basically all a big Matrix multiplcation problem in a `for` loop. PyTorch is great because you basically need to tell the model where to run (either on the CPU or using CUDA - which is a platform for moving computations to the GPU).

Now for the learning part! The `dataloader`'s job is to iterate through the entire dataset (in this case 60,000 examples of digits and their corresponding label) but to take chunks of size `batch_size` of the data to process. This is another hyperparameter that needs to be chosen. `epochs` is the number of times we want to loop through the dataset in its entirety (again something we choose based upon how our experiment goes).

A word on how long this takes - it takes a while.

```python
# loader
dataloader = DataLoader(clothing, batch_size=batch_size, num_workers=0, pin_memory=True)

# golfing time!
for t in range(epochs):
    print('Epoch {}'.format(t))
    print('-------------------------------')
    # go through the entire dataset once
    for batch, (X, Y) in enumerate(dataloader):
        X, Y = X.to(device), Y.to(device)
        # zero out gradient
        optimizer.zero_grad()
        # make a prediction on this batch!
        pred = model(X)
        # how bad is it?
        loss = cost(pred, Y)
        # compute gradients
        loss.backward()
        # update parameters
        optimizer.step()

        if batch % 100 == 0:
            print('loss: {:>10f}  [{:>5d}/{:>5d}]'.format(loss.item(), batch * len(X), len(dataloader.dataset)))

    print('loss: {:>10f}  [{:>5d}/{:>5d}]\n'.format(loss.item(), len(dataloader.dataset), len(dataloader.dataset)))

print('Done!')
```

So what the heck did this actually do? Great question. I will explain this `SimpleLinear` model since its the on we ran first. We are basically learning one matrix and one vector. That's it (it's really anti-climacttic if you ask me). Let me show you why that actually does anything:

\begin{align}
prediction = W \cdot x + b
\end{align}

\begin{align}
\begin{bmatrix}
\hat{y}_1 \\ \vdots \\ \hat{y}_{10}
\end{bmatrix}_{10 \times 1} =
\begin{bmatrix}
w_{1,1} &  \ldots &  w_{1,784} \\
\vdots    & \ddots &  \\
w_{10,1}  & \ldots & w_{10, 784} \\
\end{bmatrix}_{10 \times 784}
\cdot
\begin{bmatrix}x_1 \\ x_2 \\ \vdots \\ x_{783} \\ x_{784}\end{bmatrix}_{784 \times 1}
 +
\begin{bmatrix}
b_1 \\ \vdots \\ b_{10}
\end{bmatrix}_{10 \times 1}
\end{align}

The matrix `W`and the vector `b` is what the `SimpleLinear` model learns. The output is a 10 by 1 matrix whose largest value is the index of the thing number we want to predict. Take a look at what the algorithm learned for the two variables (as well as the sizes):

```python
for p in model.parameters():
    print(p.shape)
    print(p)
```

# Is it working???
The best way to figure out if it is working is to test the model on data the learning process hasn't used. Luckily we have such a dataset (it is basically a held out section of the data we already have used). I'm loading it all up the same way as before and printing them out to show you that they're different.

```python
test_clothing = datasets.FashionMNIST('data', train=False, download=True,
                        transform=transforms.Compose([
                            transforms.ToTensor(),
                            transforms.Lambda(lambda x: x.reshape(28*28))
                        ]),
                        target_transform=transforms.Compose([
                            transforms.Lambda(lambda y: torch.zeros(10, dtype=torch.float).scatter_(0, torch.tensor(y), value=1))
                        ])
                     )
draw_clothes(test_clothing)
```

Let's test it on the first trousers above (index 2).

```python
x, y = test_clothing[2]
x = x.to(device).view(1, 28*28)
```

We'll tell the model we are using it for evaluation (sometimes called inference) and pass in our previously unseen digit.

```python
model.eval()
with torch.no_grad():
    pred = model(x)
    pred = pred.to('cpu').detach()[0]

print(pred)
```

Now let's see if the predicted clothing item matches the actual clothing item:

```python
classes[pred.argmax(0)], classes[y.argmax(0)]
```

Let's see how well we do over *all* of the test data!

```python
# data loader for test digits
test_dataloader = DataLoader(test_clothing, batch_size=64, num_workers=0, pin_memory=True)

# set model to evaluation mode
model.eval()
test_loss = 0
correct = 0

# loop!
with torch.no_grad():
    for batch, (X, Y) in enumerate(dataloader):
        X, Y = X.to(device), Y.to(device)
        pred = model(X)

        test_loss += cost(pred, Y).item()
        correct += (pred.argmax(1) == Y.argmax(1)).type(torch.float).sum().item()

test_loss /= len(dataloader.dataset)
correct /= len(dataloader.dataset)
print('\nTest Error:')
print('acc: {:>0.1f}%, avg loss: {:>8f}'.format(100*correct, test_loss))
```

# Saving the Model
Every framework is different - in this case PyTorch let's us save the model (which you remember is just a big matrix `W` and a vector `b`) to an internal format as well as to the ONNX format. These can then be loaded up as an asset to a program that is executed every time you need to recognize a digit!

```python
# create dummy variable to traverse graph
x = torch.randint(255, (1, 28*28), dtype=torch.float).to(device) / 255
onnx.export(model, x, 'model.onnx')
print('Saved onnx model to model.onnx')

# saving PyTorch Model Dictionary
torch.save(model.state_dict(), 'model.pth')
print('Saved PyTorch Model to model.pth')
```

# Play!
Now that you've gone through the whole process, please go back up to play around! Try changing:
1. The actual model! The other models are almost identical with the exception that they learn additional Matrices (W's and b's) that the images pass through to get the final answer.
2. The hyperparameters like `learning_rate`, `batch_size`, and `epoch`. Does it make things better or worse? 92% is ok, does any other combination of model's and hyperparamters fare better?

## Final Thoughts
Would love your feedback! Was this helpful? Any parts confusing? Drop me a line!

```python

```


# PHY321: Harmonic Oscillations, Damping, Resonances and time-dependent Forces

**[Morten Hjorth-Jensen](http://mhjgit.github.io/info/doc/web/)**, Department of Physics and Astronomy and Facility for Rare Ion Beams (FRIB), Michigan State University, USA and Department of Physics, University of Oslo, Norway

Date: **Mar 13, 2021**

Copyright 1999-2021, [Morten Hjorth-Jensen](http://mhjgit.github.io/info/doc/web/). Released under CC Attribution-NonCommercial 4.0 license

## Aims and Overarching Motivation

### Monday

Damped oscillations. Analytical and numerical solutions
**Reading suggestion**: Taylor sections 5.4-5.5.

### Wednesday

No lecture, study day

### Friday

Driven oscillations and resonances with examples.
**Reading suggestion**: Taylor sections 5.5-5.6.

## Damped Oscillators

We consider only the case where the damping force is proportional to
the velocity. This is counter to dragging friction, where the force is
proportional in strength to the normal force and independent of
velocity, and is also inconsistent with wind resistance, where the
magnitude of the drag force is proportional the square of the
velocity. Rolling resistance does seem to be mainly proportional to
the velocity. However, the main motivation for considering damping
forces proportional to the velocity is that the math is more
friendly. This is because the differential equation is linear,
i.e. each term is of order $x$, $\dot{x}$, $\ddot{x}\cdots$, or even
terms with no mention of $x$, and there are no terms such as $x^2$ or
$x\ddot{x}$. The equations of motion for a spring with damping force
$-b\dot{x}$ are

$$
\begin{equation}
m\ddot{x}+b\dot{x}+kx=0.
\label{_auto1} \tag{1}
\end{equation}
$$

## Harmonic Oscillator,  Damping

Just to make the solution a bit less messy, we rewrite this equation as

$$
\begin{equation}
\label{eq:dampeddiffyq} \tag{2}
\ddot{x}+2\beta\dot{x}+\omega_0^2x=0,~~~~\beta\equiv b/2m,~\omega_0\equiv\sqrt{k/m}.
\end{equation}
$$

Both $\beta$ and $\omega$ have dimensions of inverse time. To find solutions (see appendix C in the text) you must make an educated guess at the form of the solution. To do this, first realize that the solution will need an arbitrary normalization $A$ because the equation is linear. Secondly, realize that if the form is

$$
\begin{equation}
x=Ae^{rt}
\label{_auto2} \tag{3}
\end{equation}
$$

that each derivative simply brings out an extra power of $r$. This
means that the $Ae^{rt}$ factors out and one can simply solve for an
equation for $r$. Plugging this form into Eq. ([2](#eq:dampeddiffyq)),

$$
\begin{equation}
r^2+2\beta r+\omega_0^2=0.
\label{_auto3} \tag{4}
\end{equation}
$$

## Harmonic Oscillator,  Solutions of Damped Motion

Because this is a quadratic equation there will be two solutions,

$$
\begin{equation}
r=-\beta\pm\sqrt{\beta^2-\omega_0^2}.
\label{_auto4} \tag{5}
\end{equation}
$$

We refer to the two solutions as $r_1$ and $r_2$ corresponding to the
$+$ and $-$ roots. As expected, there should be two arbitrary
constants involved in the solution,

$$
\begin{equation}
x=A_1e^{r_1t}+A_2e^{r_2t},
\label{_auto5} \tag{6}
\end{equation}
$$

where the coefficients $A_1$ and $A_2$ are determined by initial
conditions.

The roots listed above, $\sqrt{\omega_0^2-\beta_0^2}$, will be
imaginary if the damping is small and $\beta<\omega_0$. In that case,
$r$ is complex and the factor $\exp{(rt)}$ will have some oscillatory
behavior. If the roots are real, there will only be exponentially
decaying solutions. There are three cases:

## Underdamped: $\beta<\omega_0$

$$
\begin{eqnarray}
x&=&A_1e^{-\beta t}e^{i\omega't}+A_2e^{-\beta t}e^{-i\omega't},~~\omega'\equiv\sqrt{\omega_0^2-\beta^2}\\
\nonumber
&=&(A_1+A_2)e^{-\beta t}\cos\omega't+i(A_1-A_2)e^{-\beta t}\sin\omega't.
\end{eqnarray}
$$

Here we have made use of the identity
$e^{i\omega't}=\cos\omega't+i\sin\omega't$. Because the constants are
arbitrary, and because the real and imaginary parts are both solutions
individually, we can simply consider the real part of the solution
alone:

$$
\begin{eqnarray}
\label{eq:homogsolution} \tag{7}
x&=&B_1e^{-\beta t}\cos\omega't+B_2e^{-\beta t}\sin\omega't,\\
\nonumber
\omega'&\equiv&\sqrt{\omega_0^2-\beta^2}.
\end{eqnarray}
$$

## Critical dampling: $\beta=\omega_0$

In this case the two terms involving $r_1$ and $r_2$ are identical
because $\omega'=0$. Because we need to arbitrary constants, there
needs to be another solution. This is found by simply guessing, or by
taking the limit of $\omega'\rightarrow 0$ from the underdamped
solution. The solution is then

$$
\begin{equation}
\label{eq:criticallydamped} \tag{8}
x=Ae^{-\beta t}+Bte^{-\beta t}.
\end{equation}
$$

The critically damped solution is interesting because the solution
approaches zero quickly, but does not oscillate. For a problem with
zero initial velocity, the solution never crosses zero. This is a good
choice for designing shock absorbers or swinging doors.

## Overdamped: $\beta>\omega_0$

$$
\begin{eqnarray}
x&=&A_1\exp{-(\beta+\sqrt{\beta^2-\omega_0^2})t}+A_2\exp{-(\beta-\sqrt{\beta^2-\omega_0^2})t}
\end{eqnarray}
$$

This solution will also never pass the origin more than once, and then
only if the initial velocity is strong and initially toward zero.

Given $b$, $m$ and $\omega_0$, find $x(t)$ for a particle whose
initial position is $x=0$ and has initial velocity $v_0$ (assuming an
underdamped solution).

The solution is of the form,

$$
\begin{eqnarray*}
x&=&e^{-\beta t}\left[A_1\cos(\omega' t)+A_2\sin\omega't\right],\\
\dot{x}&=&-\beta x+\omega'e^{-\beta t}\left[-A_1\sin\omega't+A_2\cos\omega't\right].\\
\omega'&\equiv&\sqrt{\omega_0^2-\beta^2},~~~\beta\equiv b/2m.
\end{eqnarray*}
$$

From the initial conditions, $A_1=0$ because $x(0)=0$ and $\omega'A_2=v_0$. So

$$
x=\frac{v_0}{\omega'}e^{-\beta t}\sin\omega't.
$$

## Harmonic Oscillator,  Solutions

Consider a single solution with no arbitrary constants, which we will
call a **particular solution**, $x_p(t)$. It should be emphasized
that this is **A** particular solution, because there exists an
infinite number of such solutions because the general solution should
have two arbitrary constants. Now consider solutions to the same
equation without the driving term, which include two arbitrary
constants. These are called either **homogenous solutions** or
**complementary solutions**, and were given above,
e.g. Eq. ([7](#eq:homogsolution)) for the underdamped case. The
homogenous solution already incorporates the two arbitrary constants,
so any sum of a homogenous solution and a particular solution will
represent the **general solution** of the equation. The general
solution incorporates the two arbitrary constants $A$ and $B$ to
accommodate the two initial conditions. One could have picked a
different particular solution, i.e. the original particular solution
plus any homogenous solution with the arbitrary constants $A_p$ and
$B_p$ chosen at will. When one adds in the homogenous solution, which
has adjustable constants with arbitrary constants $A'$ and $B'$, to
the new particular solution, one can get the same general solution by
simply adjusting the new constants such that $A'+A_p=A$ and
$B'+B_p=B$. Thus, the choice of $A_p$ and $B_p$ are irrelevant, and
when choosing the particular solution it is best to make the simplest
choice possible.

## Harmonic Oscillator,  Particular Solution

To find a particular solution, one first guesses at the form,

$$
\begin{equation}
\label{eq:partform} \tag{9}
x_p(t)=D\cos(\omega t-\delta),
\end{equation}
$$

and rewrite the differential equation as

$$
\begin{equation}
D\left\{-\omega^2\cos(\omega t-\delta)-2\beta\omega\sin(\omega t-\delta)+\omega_0^2\cos(\omega t-\delta)\right\}=\frac{F_0}{m}\cos(\omega t).
\label{_auto6} \tag{10}
\end{equation}
$$

One can now use angle addition formulas to get

$$
\begin{eqnarray}
D\left\{(-\omega^2\cos\delta+2\beta\omega\sin\delta+\omega_0^2\cos\delta)\cos(\omega t)\right.&&\\
\nonumber
\left.+(-\omega^2\sin\delta-2\beta\omega\cos\delta+\omega_0^2\sin\delta)\sin(\omega t)\right\}
&=&\frac{F_0}{m}\cos(\omega t).
\end{eqnarray}
$$

Both the $\cos$ and $\sin$ terms need to equate if the expression is to hold at all times. Thus, this becomes two equations

$$
\begin{eqnarray}
D\left\{-\omega^2\cos\delta+2\beta\omega\sin\delta+\omega_0^2\cos\delta\right\}&=&\frac{F_0}{m}\\
\nonumber
-\omega^2\sin\delta-2\beta\omega\cos\delta+\omega_0^2\sin\delta&=&0.
\end{eqnarray}
$$

After dividing by $\cos\delta$, the lower expression leads to

$$
\begin{equation}
\tan\delta=\frac{2\beta\omega}{\omega_0^2-\omega^2}.
\label{_auto7} \tag{11}
\end{equation}
$$

## Solving with Driven Oscillations

Using the identities $\tan^2+1=\csc^2$ and $\sin^2+\cos^2=1$, one can also express $\sin\delta$ and $\cos\delta$,

$$
\begin{eqnarray}
\sin\delta&=&\frac{2\beta\omega}{\sqrt{(\omega_0^2-\omega^2)^2+4\omega^2\beta^2}},\\
\nonumber
\cos\delta&=&\frac{(\omega_0^2-\omega^2)}{\sqrt{(\omega_0^2-\omega^2)^2+4\omega^2\beta^2}}
\end{eqnarray}
$$

Inserting the expressions for $\cos\delta$ and $\sin\delta$ into the expression for $D$,

$$
\begin{equation}
\label{eq:Ddrive} \tag{12}
D=\frac{F_0/m}{\sqrt{(\omega_0^2-\omega^2)^2+4\omega^2\beta^2}}.
\end{equation}
$$

For a given initial condition, e.g. initial displacement and velocity,
one must add the homogenous solution then solve for the two arbitrary
constants. However, because the homogenous solutions decay with time
as $e^{-\beta t}$, the particular solution is all that remains at
large times, and is therefore the steady state solution. Because the
arbitrary constants are all in the homogenous solution, all memory of
the initial conditions are lost at large times, $t>>1/\beta$.

The amplitude of the motion, $D$, is linearly proportional to the
driving force ($F_0/m$), but also depends on the driving frequency
$\omega$. For small $\beta$ the maximum will occur at
$\omega=\omega_0$. This is referred to as a resonance. In the limit
$\beta\rightarrow 0$ the amplitude at resonance approaches infinity.

## Alternative Derivation for Driven Oscillators

Here, we derive the same expressions as in Equations ([9](#eq:partform)) and ([12](#eq:Ddrive)) but express the driving forces as

$$
\begin{eqnarray}
F(t)&=&F_0e^{i\omega t},
\end{eqnarray}
$$

rather than as $F_0\cos\omega t$. The real part of $F$ is the same as before. For the differential equation,

$$
\begin{eqnarray}
\label{eq:compdrive} \tag{13}
\ddot{x}+2\beta\dot{x}+\omega_0^2x&=&\frac{F_0}{m}e^{i\omega t},
\end{eqnarray}
$$

one can treat $x(t)$ as an imaginary function. Because the operations
$d^2/dt^2$ and $d/dt$ are real and thus do not mix the real and
imaginary parts of $x(t)$, Eq. ([13](#eq:compdrive)) is effectively 2
equations. Because $e^{\omega t}=\cos\omega t+i\sin\omega t$, the real
part of the solution for $x(t)$ gives the solution for a driving force
$F_0\cos\omega t$, and the imaginary part of $x$ corresponds to the
case where the driving force is $F_0\sin\omega t$. It is rather easy
to solve for the complex $x$ in this case, and by taking the real part
of the solution, one finds the answer for the $\cos\omega t$ driving
force.

We assume a simple form for the particular solution

$$
\begin{equation}
x_p=De^{i\omega t},
\label{_auto8} \tag{14}
\end{equation}
$$

where $D$ is a complex constant.

From Eq. ([13](#eq:compdrive)) one inserts the form for $x_p$ above to get

$$
\begin{eqnarray}
D\left\{-\omega^2+2i\beta\omega+\omega_0^2\right\}e^{i\omega t}=(F_0/m)e^{i\omega t},\\
\nonumber
D=\frac{F_0/m}{(\omega_0^2-\omega^2)+2i\beta\omega}.
\end{eqnarray}
$$

The norm and phase for $D=|D|e^{-i\delta}$ can be read by inspection,

$$
\begin{equation}
|D|=\frac{F_0/m}{\sqrt{(\omega_0^2-\omega^2)^2+4\beta^2\omega^2}},~~~~\tan\delta=\frac{2\beta\omega}{\omega_0^2-\omega^2}.
\label{_auto9} \tag{15}
\end{equation}
$$

This is the same expression for $\delta$ as before. One then finds $x_p(t)$,

$$
\begin{eqnarray}
\label{eq:fastdriven1} \tag{16}
x_p(t)&=&\Re\frac{(F_0/m)e^{i\omega t-i\delta}}{\sqrt{(\omega_0^2-\omega^2)^2+4\beta^2\omega^2}}\\
\nonumber
&=&\frac{(F_0/m)\cos(\omega t-\delta)}{\sqrt{(\omega_0^2-\omega^2)^2+4\beta^2\omega^2}}.
\end{eqnarray}
$$

This is the same answer as before.
If one wished to solve for the case where $F(t)= F_0\sin\omega t$, the imaginary part of the solution would work

$$
\begin{eqnarray}
\label{eq:fastdriven2} \tag{17}
x_p(t)&=&\Im\frac{(F_0/m)e^{i\omega t-i\delta}}{\sqrt{(\omega_0^2-\omega^2)^2+4\beta^2\omega^2}}\\
\nonumber
&=&\frac{(F_0/m)\sin(\omega t-\delta)}{\sqrt{(\omega_0^2-\omega^2)^2+4\beta^2\omega^2}}.
\end{eqnarray}
$$

## Damped and Driven Oscillator

Consider the damped and driven harmonic oscillator worked out above. Given $F_0, m,\beta$ and $\omega_0$, solve for the complete solution $x(t)$ for the case where $F=F_0\sin\omega t$ with initial conditions $x(t=0)=0$ and $v(t=0)=0$. Assume the underdamped case.

The general solution including the arbitrary constants includes both the homogenous and particular solutions,

$$
\begin{eqnarray*}
x(t)&=&\frac{F_0}{m}\frac{\sin(\omega t-\delta)}{\sqrt{(\omega_0^2-\omega^2)^2+4\beta^2\omega^2}}
+A\cos\omega't e^{-\beta t}+B\sin\omega't e^{-\beta t}.
\end{eqnarray*}
$$

The quantities $\delta$ and $\omega'$ are given earlier, $\omega'=\sqrt{\omega_0^2-\beta^2},
\delta=\tan^{-1}(2\beta\omega/(\omega_0^2-\omega^2)$. Here, solving
the problem means finding the arbitrary constants $A$ and
$B$. Satisfying the initial conditions for the initial position and
velocity:

$$
\begin{eqnarray*}
x(t=0)=0&=&-\eta\sin\delta+A,\\
v(t=0)=0&=&\omega\eta\cos\delta-\beta A+\omega'B,\\
\eta&\equiv&\frac{F_0}{m}\frac{1}{\sqrt{(\omega_0^2-\omega^2)^2+4\beta^2\omega^2}}.
\end{eqnarray*}
$$

The problem is now reduced to 2 equations and 2 unknowns, $A$ and $B$. The solution is

$$
\begin{eqnarray}
A&=& \eta\sin\delta ,~~~B=\frac{-\omega\eta\cos\delta+\beta\eta\sin\delta}{\omega'}.
\end{eqnarray}
$$

## Resonance Widths; the $Q$ factor

From above, the particular solution for a driving force, $F=F_0\cos\omega t$, is

$$
\begin{eqnarray}
x_p(t)&=&\frac{F_0/m}{\sqrt{(\omega_0^2-\omega^2)^2+4\omega^2\beta^2}}\cos(\omega_t-\delta),\\
\nonumber
\delta&=&\tan^{-1}\left(\frac{2\beta\omega}{\omega_0^2-\omega^2}\right).
\end{eqnarray}
$$

If one fixes the driving frequency $\omega$ and adjusts the
fundamental frequency $\omega_0=\sqrt{k/m}$, the maximum amplitude
occurs when $\omega_0=\omega$ because that is when the term from the
denominator $(\omega_0^2-\omega^2)^2+4\omega^2\beta^2$ is at a
minimum. This is akin to dialing into a radio station. However, if one
fixes $\omega_0$ and adjusts the driving frequency one minimize with
respect to $\omega$, e.g. set

$$
\begin{equation}
\frac{d}{d\omega}\left[(\omega_0^2-\omega^2)^2+4\omega^2\beta^2\right]=0,
\label{_auto10} \tag{18}
\end{equation}
$$

and one finds that the maximum amplitude occurs when
$\omega=\sqrt{\omega_0^2-2\beta^2}$. If $\beta$ is small relative to
$\omega_0$, one can simply state that the maximum amplitude is

$$
\begin{equation}
x_{\rm max}\approx\frac{F_0}{2m\beta \omega_0}.
\label{_auto11} \tag{19}
\end{equation}
$$

$$
\begin{eqnarray}
\frac{4\omega^2\beta^2}{(\omega_0^2-\omega^2)^2+4\omega^2\beta^2}=\frac{1}{2}.
\end{eqnarray}
$$

For small damping this occurs when $\omega=\omega_0\pm \beta$, so the $FWHM\approx 2\beta$. For the purposes of tuning to a specific frequency, one wants the width to be as small as possible. The ratio of $\omega_0$ to $FWHM$ is known as the [quality factor](https://en.wikipedia.org/wiki/Q_factor), or $Q$ factor,

$$
\begin{equation}
Q\equiv \frac{\omega_0}{2\beta}.
\label{_auto12} \tag{20}
\end{equation}
$$

## Numerical Studies of Driven Oscillations

Solving the problem of driven oscillations numerically gives us much
more flexibility to study different types of driving forces. We can
reuse our earlier code by simply adding a driving force. If we stay in
the $x$-direction only this can be easily done by adding a term
$F_{\mathrm{ext}}(x,t)$. Note that we have kept it rather general
here, allowing for both a spatial and a temporal dependence.

Before we dive into the code, we need to briefly remind ourselves
about the equations we started with for the case with damping, namely

$$
m\frac{d^2x}{dt^2} + b\frac{dx}{dt}+kx(t) =0,
$$

with no external force applied to the system.

Let us now for simplicty assume that our external force is given by

$$
F_{\mathrm{ext}}(t) = F_0\cos{(\omega t)},
$$

where $F_0$ is a constant (what is its dimension?) and $\omega$ is the frequency of the applied external driving force.
**Small question:** would you expect energy to be conserved now?

Introducing the external force into our lovely differential equation
and dividing by $m$ and introducing $\omega_0^2=\sqrt{k/m}$ we have

$$
\frac{d^2x}{dt^2} + \frac{b}{m}\frac{dx}{dt}+\omega_0^2x(t) =\frac{F_0}{m}\cos{(\omega t)},
$$

Thereafter we introduce a dimensionless time $\tau = t\omega_0$
and a dimensionless frequency $\tilde{\omega}=\omega/\omega_0$. We have then

$$
\frac{d^2x}{d\tau^2} + \frac{b}{m\omega_0}\frac{dx}{d\tau}+x(\tau) =\frac{F_0}{m\omega_0^2}\cos{(\tilde{\omega}\tau)},
$$

Introducing a new amplitude $\tilde{F} =F_0/(m\omega_0^2)$ (check dimensionality again) we have

$$
\frac{d^2x}{d\tau^2} + \frac{b}{m\omega_0}\frac{dx}{d\tau}+x(\tau) =\tilde{F}\cos{(\tilde{\omega}\tau)}.
$$

Our final step, as we did in the case of various types of damping, is
to define $\gamma = b/(2m\omega_0)$ and rewrite our equations as

$$
\frac{d^2x}{d\tau^2} + 2\gamma\frac{dx}{d\tau}+x(\tau) =\tilde{F}\cos{(\tilde{\omega}\tau)}.
$$

This is the equation we will code below using the Euler-Cromer method.

```
%matplotlib inline

# Common imports
import numpy as np
import pandas as pd
from math import *
import matplotlib.pyplot as plt
import os

# Where to save the figures and data files
PROJECT_ROOT_DIR = "Results"
FIGURE_ID = "Results/FigureFiles"
DATA_ID = "DataFiles/"

if not os.path.exists(PROJECT_ROOT_DIR):
    os.mkdir(PROJECT_ROOT_DIR)

if not os.path.exists(FIGURE_ID):
    os.makedirs(FIGURE_ID)

if not os.path.exists(DATA_ID):
    os.makedirs(DATA_ID)

def image_path(fig_id):
    return os.path.join(FIGURE_ID, fig_id)

def data_path(dat_id):
    return os.path.join(DATA_ID, dat_id)

def save_fig(fig_id):
    plt.savefig(image_path(fig_id) + ".png", format='png')

from pylab import plt, mpl
plt.style.use('seaborn')
mpl.rcParams['font.family'] = 'serif'

DeltaT = 0.001
#set up arrays
tfinal = 20 # in dimensionless time
n = ceil(tfinal/DeltaT)
# set up arrays for t, v, and x
t = np.zeros(n)
v = np.zeros(n)
x = np.zeros(n)
# Initial conditions as one-dimensional arrays of time
x0 =  1.0
v0 = 0.0
x[0] = x0
v[0] = v0
gamma = 0.2
Omegatilde = 0.5
Ftilde = 1.0
# Start integrating using Euler-Cromer's method
for i in range(n-1):
    # Set up the acceleration
    # Here you could have defined your own function for this
    a =  -2*gamma*v[i]-x[i]+Ftilde*cos(t[i]*Omegatilde)
    # update velocity, time and position
    v[i+1] = v[i] + DeltaT*a
    x[i+1] = x[i] + DeltaT*v[i+1]
    t[i+1] = t[i] + DeltaT
# Plot position as function of time
fig, ax = plt.subplots()
ax.set_ylabel('x[m]')
ax.set_xlabel('t[s]')
ax.plot(t, x)
fig.tight_layout()
save_fig("ForcedBlockEulerCromer")
plt.show()
```

In the above example we have focused on the Euler-Cromer method. This
method has a local truncation error which is proportional to $\Delta t^2$
and thereby a global error which is proportional to $\Delta t$.
We can improve this by using the Runge-Kutta family of
methods. The widely popular Runge-Kutta to fourth order or just **RK4**
has indeed a much better truncation error. The RK4 method has a global
error which is proportional to $\Delta t$.

Let us revisit this method and see how we can implement it for the above example.

## Differential Equations, Runge-Kutta methods

Runge-Kutta (RK) methods are based on Taylor expansion formulae, but yield
in general better algorithms for solutions of an ordinary differential equation.
The basic philosophy is that it provides an intermediate step in the computation of $y_{i+1}$.

To see this, consider first the following definitions

$$
\begin{equation}
\frac{dy}{dt}=f(t,y),
\label{_auto13} \tag{21}
\end{equation}
$$

and

$$
\begin{equation}
y(t)=\int f(t,y) dt,
\label{_auto14} \tag{22}
\end{equation}
$$

and

$$
\begin{equation}
y_{i+1}=y_i+ \int_{t_i}^{t_{i+1}} f(t,y) dt.
\label{_auto15} \tag{23}
\end{equation}
$$

To demonstrate the philosophy behind RK methods, let us consider
the second-order RK method, RK2.
The first approximation consists in Taylor expanding $f(t,y)$
around the center of the integration interval $t_i$ to $t_{i+1}$,
that is, at $t_i+h/2$, $h$ being the step.
Using the midpoint formula for an integral,
defining $y(t_i+h/2) = y_{i+1/2}$ and
$t_i+h/2 = t_{i+1/2}$, we obtain

$$
\begin{equation}
\int_{t_i}^{t_{i+1}} f(t,y) dt \approx hf(t_{i+1/2},y_{i+1/2}) +O(h^3).
\label{_auto16} \tag{24}
\end{equation}
$$

This means in turn that we have

$$
\begin{equation}
y_{i+1}=y_i + hf(t_{i+1/2},y_{i+1/2}) +O(h^3).
\label{_auto17} \tag{25}
\end{equation}
$$

However, we do not know the value of   $y_{i+1/2}$. Here comes thus the next approximation, namely, we use Euler's
method to approximate $y_{i+1/2}$. We have then

$$
\begin{equation}
y_{(i+1/2)}=y_i + \frac{h}{2}\frac{dy}{dt}=y(t_i) + \frac{h}{2}f(t_i,y_i).
\label{_auto18} \tag{26}
\end{equation}
$$

This means that we can define the following algorithm for
the second-order Runge-Kutta method, RK2.

4
6

<
<
<
!
!
M
A
T
H
_
B
L
O
C
K

$$
\begin{equation}
k_2=hf(t_{i+1/2},y_i+k_1/2),
\label{_auto20} \tag{28}
\end{equation}
$$

with the final value

$$
\begin{equation}
y_{i+i}\approx y_i + k_2 +O(h^3).
\label{_auto21} \tag{29}
\end{equation}
$$

The difference between the previous one-step methods
is that we now need an intermediate step in our evaluation,
namely $t_i+h/2 = t_{(i+1/2)}$ where we evaluate the derivative $f$.
This involves more operations, but the gain is a better stability
in the solution.

The fourth-order Runge-Kutta, RK4, has the following algorithm

4
9

<
<
<
!
!
M
A
T
H
_
B
L
O
C
K

$$
k_3=hf(t_i+h/2,y_i+k_2/2)\hspace{0.5cm}   k_4=hf(t_i+h,y_i+k_3)
$$

with the final result

$$
y_{i+1}=y_i +\frac{1}{6}\left( k_1 +2k_2+2k_3+k_4\right).
$$

Thus, the algorithm consists in first calculating $k_1$
with $t_i$, $y_1$ and $f$ as inputs. Thereafter, we increase the step
size by $h/2$ and calculate $k_2$, then $k_3$ and finally $k_4$. The global error goes as $O(h^4)$.

However, at this stage, if we keep adding different methods in our
main program, the code will quickly become messy and ugly. Before we
proceed thus, we will now introduce functions that enbody the various
methods for solving differential equations. This means that we can
separate out these methods in own functions and files (and later as classes and more
generic functions) and simply call them when needed. Similarly, we
could easily encapsulate various forces or other quantities of
interest in terms of functions. To see this, let us bring up the code
we developed above for the simple sliding block, but now only with the simple forward Euler method. We introduce
two functions, one for the simple Euler method and one for the
force.

Note that here the forward Euler method does not know the specific force function to be called.
It receives just an input the name. We can easily change the force by adding another function.

```
def ForwardEuler(v,x,t,n,Force):
    for i in range(n-1):
        v[i+1] = v[i] + DeltaT*Force(v[i],x[i],t[i])
        x[i+1] = x[i] + DeltaT*v[i]
        t[i+1] = t[i] + DeltaT
```

```
def SpringForce(v,x,t):
#   note here that we have divided by mass and we return the acceleration
    return  -2*gamma*v-x+Ftilde*cos(t*Omegatilde)
```

It is easy to add a new method like the Euler-Cromer

```
def ForwardEulerCromer(v,x,t,n,Force):
    for i in range(n-1):
        a = Force(v[i],x[i],t[i])
        v[i+1] = v[i] + DeltaT*a
        x[i+1] = x[i] + DeltaT*v[i+1]
        t[i+1] = t[i] + DeltaT
```

and the Velocity Verlet method (be careful with time-dependence here, it is not an ideal method for non-conservative forces))

```
def VelocityVerlet(v,x,t,n,Force):
    for i in range(n-1):
        a = Force(v[i],x[i],t[i])
        x[i+1] = x[i] + DeltaT*v[i]+0.5*a*DeltaT*DeltaT
        anew = Force(v[i],x[i+1],t[i+1])
        v[i+1] = v[i] + 0.5*DeltaT*(a+anew)
        t[i+1] = t[i] + DeltaT
```

Finally, we can now add the Runge-Kutta2 method via a new function

```
def RK2(v,x,t,n,Force):
    for i in range(n-1):
# Setting up k1
        k1x = DeltaT*v[i]
        k1v = DeltaT*Force(v[i],x[i],t[i])
# Setting up k2
        vv = v[i]+k1v*0.5
        xx = x[i]+k1x*0.5
        k2x = DeltaT*vv
        k2v = DeltaT*Force(vv,xx,t[i]+DeltaT*0.5)
# Final result
        x[i+1] = x[i]+k2x
        v[i+1] = v[i]+k2v
	t[i+1] = t[i]+DeltaT
```

Finally, we can now add the Runge-Kutta2 method via a new function

```
def RK4(v,x,t,n,Force):
    for i in range(n-1):
# Setting up k1
        k1x = DeltaT*v[i]
        k1v = DeltaT*Force(v[i],x[i],t[i])
# Setting up k2
        vv = v[i]+k1v*0.5
        xx = x[i]+k1x*0.5
        k2x = DeltaT*vv
        k2v = DeltaT*Force(vv,xx,t[i]+DeltaT*0.5)
# Setting up k3
        vv = v[i]+k2v*0.5
        xx = x[i]+k2x*0.5
        k3x = DeltaT*vv
        k3v = DeltaT*Force(vv,xx,t[i]+DeltaT*0.5)
# Setting up k4
        vv = v[i]+k3v
        xx = x[i]+k3x
        k4x = DeltaT*vv
        k4v = DeltaT*Force(vv,xx,t[i]+DeltaT)
# Final result
        x[i+1] = x[i]+(k1x+2*k2x+2*k3x+k4x)/6.
        v[i+1] = v[i]+(k1v+2*k2v+2*k3v+k4v)/6.
        t[i+1] = t[i] + DeltaT
```

The Runge-Kutta family of methods are particularly useful when we have a time-dependent acceleration.
If we have forces which depend only the spatial degrees of freedom (no velocity and/or time-dependence), then energy conserving methods like the Velocity Verlet or the Euler-Cromer method are preferred. As soon as we introduce an explicit time-dependence and/or add dissipitave forces like friction or air resistance, then methods like the family of Runge-Kutta methods are well suited for this.
The code below uses the Runge-Kutta4 methods.

```
DeltaT = 0.001
#set up arrays
tfinal = 20 # in dimensionless time
n = ceil(tfinal/DeltaT)
# set up arrays for t, v, and x
t = np.zeros(n)
v = np.zeros(n)
x = np.zeros(n)
# Initial conditions (can change to more than one dim)
x0 =  1.0
v0 = 0.0
x[0] = x0
v[0] = v0
gamma = 0.2
Omegatilde = 0.5
Ftilde = 1.0
# Start integrating using Euler's method
# Note that we define the force function as a SpringForce
RK4(v,x,t,n,SpringForce)

# Plot position as function of time
fig, ax = plt.subplots()
ax.set_ylabel('x[m]')
ax.set_xlabel('t[s]')
ax.plot(t, x)
fig.tight_layout()
save_fig("ForcedBlockRK4")
plt.show()
```

## Principle of Superposition and Periodic Forces (Fourier Transforms)

If one has several driving forces, $F(t)=\sum_n F_n(t)$, one can find
the particular solution to each $F_n$, $x_{pn}(t)$, and the particular
solution for the entire driving force is

$$
\begin{equation}
x_p(t)=\sum_nx_{pn}(t).
\label{_auto22} \tag{30}
\end{equation}
$$

This is known as the principal of superposition. It only applies when
the homogenous equation is linear. If there were an anharmonic term
such as $x^3$ in the homogenous equation, then when one summed various
solutions, $x=(\sum_n x_n)^2$, one would get cross
terms. Superposition is especially useful when $F(t)$ can be written
as a sum of sinusoidal terms, because the solutions for each
sinusoidal (sine or cosine)  term is analytic, as we saw above.

Driving forces are often periodic, even when they are not
sinusoidal. Periodicity implies that for some time $\tau$

$$
\begin{eqnarray}
F(t+\tau)=F(t).
\end{eqnarray}
$$

One example of a non-sinusoidal periodic force is a square wave. Many
components in electric circuits are non-linear, e.g. diodes, which
makes many wave forms non-sinusoidal even when the circuits are being
driven by purely sinusoidal sources.

The code here shows a typical example of such a square wave generated using the functionality included in the **scipy** Python package. We have used a period of $\tau=0.2$.

```
import numpy as np
import math
from scipy import signal
import matplotlib.pyplot as plt

# number of points
n = 500
# start and final times
t0 = 0.0
tn = 1.0
# Period
t = np.linspace(t0, tn, n, endpoint=False)
SqrSignal = np.zeros(n)
SqrSignal = 1.0+signal.square(2*np.pi*5*t)
plt.plot(t, SqrSignal)
plt.ylim(-0.5, 2.5)
plt.show()
```

For the sinusoidal example studied above the
period is $\tau=2\pi/\omega$. However, higher harmonics can also
satisfy the periodicity requirement. In general, any force that
satisfies the periodicity requirement can be expressed as a sum over
harmonics,

$$
\begin{equation}
F(t)=\frac{f_0}{2}+\sum_{n>0} f_n\cos(2n\pi t/\tau)+g_n\sin(2n\pi t/\tau).
\label{_auto23} \tag{31}
\end{equation}
$$

One can write down the answer for
$x_{pn}(t)$, by substituting $f_n/m$ or $g_n/m$ for $F_0/m$ into Eq.s
([16](#eq:fastdriven1)) or ([17](#eq:fastdriven2)) respectively. By
writing each factor $2n\pi t/\tau$ as $n\omega t$, with $\omega\equiv
2\pi/\tau$,

$$
\begin{equation}
\label{eq:fourierdef1} \tag{32}
F(t)=\frac{f_0}{2}+\sum_{n>0}f_n\cos(n\omega t)+g_n\sin(n\omega t).
\end{equation}
$$

The solutions for $x(t)$ then come from replacing $\omega$ with
$n\omega$ for each term in the particular solution in Equations
([9](#eq:partform)) and ([12](#eq:Ddrive)),

$$
\begin{eqnarray}
x_p(t)&=&\frac{f_0}{2k}+\sum_{n>0} \alpha_n\cos(n\omega t-\delta_n)+\beta_n\sin(n\omega t-\delta_n),\\
\nonumber
\alpha_n&=&\frac{f_n/m}{\sqrt{((n\omega)^2-\omega_0^2)+4\beta^2n^2\omega^2}},\\
\nonumber
\beta_n&=&\frac{g_n/m}{\sqrt{((n\omega)^2-\omega_0^2)+4\beta^2n^2\omega^2}},\\
\nonumber
\delta_n&=&\tan^{-1}\left(\frac{2\beta n\omega}{\omega_0^2-n^2\omega^2}\right).
\end{eqnarray}
$$

Because the forces have been applied for a long time, any non-zero
damping eliminates the homogenous parts of the solution, so one need
only consider the particular solution for each $n$.

The problem will considered solved if one can find expressions for the
coefficients $f_n$ and $g_n$, even though the solutions are expressed
as an infinite sum. The coefficients can be extracted from the
function $F(t)$ by

$$
\begin{eqnarray}
\label{eq:fourierdef2} \tag{33}
f_n&=&\frac{2}{\tau}\int_{-\tau/2}^{\tau/2} dt~F(t)\cos(2n\pi t/\tau),\\
\nonumber
g_n&=&\frac{2}{\tau}\int_{-\tau/2}^{\tau/2} dt~F(t)\sin(2n\pi t/\tau).
\end{eqnarray}
$$

To check the consistency of these expressions and to verify
Eq. ([33](#eq:fourierdef2)), one can insert the expansion of $F(t)$ in
Eq. ([32](#eq:fourierdef1)) into the expression for the coefficients in
Eq. ([33](#eq:fourierdef2)) and see whether

$$
\begin{eqnarray}
f_n&=?&\frac{2}{\tau}\int_{-\tau/2}^{\tau/2} dt~\left\{
\frac{f_0}{2}+\sum_{m>0}f_m\cos(m\omega t)+g_m\sin(m\omega t)
\right\}\cos(n\omega t).
\end{eqnarray}
$$

Immediately, one can throw away all the terms with $g_m$ because they
convolute an even and an odd function. The term with $f_0/2$
disappears because $\cos(n\omega t)$ is equally positive and negative
over the interval and will integrate to zero. For all the terms
$f_m\cos(m\omega t)$ appearing in the sum, one can use angle addition
formulas to see that $\cos(m\omega t)\cos(n\omega
t)=(1/2)(\cos[(m+n)\omega t]+\cos[(m-n)\omega t]$. This will integrate
to zero unless $m=n$. In that case the $m=n$ term gives

$$
\begin{equation}
\int_{-\tau/2}^{\tau/2}dt~\cos^2(m\omega t)=\frac{\tau}{2},
\label{_auto24} \tag{34}
\end{equation}
$$

and

$$
\begin{eqnarray}
f_n&=?&\frac{2}{\tau}\int_{-\tau/2}^{\tau/2} dt~f_n/2\\
\nonumber
&=&f_n~\checkmark.
\end{eqnarray}
$$

The same method can be used to check for the consistency of $g_n$.

Consider the driving force:

$$
\begin{equation}
F(t)=At/\tau,~~-\tau/2<t<\tau/2,~~~F(t+\tau)=F(t).
\label{_auto25} \tag{35}
\end{equation}
$$

Find the Fourier coefficients $f_n$ and $g_n$ for all $n$ using Eq. ([33](#eq:fourierdef2)).

Only the odd coefficients enter by symmetry, i.e. $f_n=0$. One can find $g_n$ integrating by parts,

$$
\begin{eqnarray}
\label{eq:fouriersolution} \tag{36}
g_n&=&\frac{2}{\tau}\int_{-\tau/2}^{\tau/2}dt~\sin(n\omega t) \frac{At}{\tau}\\
\nonumber
u&=&t,~dv=\sin(n\omega t)dt,~v=-\cos(n\omega t)/(n\omega),\\
\nonumber
g_n&=&\frac{-2A}{n\omega \tau^2}\int_{-\tau/2}^{\tau/2}dt~\cos(n\omega t)
+\left.2A\frac{-t\cos(n\omega t)}{n\omega\tau^2}\right|_{-\tau/2}^{\tau/2}.
\end{eqnarray}
$$

The first term is zero because $\cos(n\omega t)$ will be equally
positive and negative over the interval. Using the fact that
$\omega\tau=2\pi$,

$$
\begin{eqnarray}
g_n&=&-\frac{2A}{2n\pi}\cos(n\omega\tau/2)\\
\nonumber
&=&-\frac{A}{n\pi}\cos(n\pi)\\
\nonumber
&=&\frac{A}{n\pi}(-1)^{n+1}.
\end{eqnarray}
$$

## Fourier Series

More text will come here, chpater 5.7-5.8 of Taylor are discussed
during the lectures. The code here uses the Fourier series discussed
in chapter 5.7 for a square wave signal. The equations for the
coefficients are are discussed in Taylor section 5.7, see Example
5.4. The code here visualizes the various approximations given by
Fourier series compared with a square wave with period $T=0.2$, witth
$0.1$ and max value $F=2$. We see that when we increase the number of
components in the Fourier series, the Fourier series approximation gets closes and closes to the square wave signal.

```
import numpy as np
import math
from scipy import signal
import matplotlib.pyplot as plt

# number of points
n = 500
# start and final times
t0 = 0.0
tn = 1.0
# Period
T =0.2
# Max value of square signal
Fmax= 2.0
# Width of signal
Width = 0.1
t = np.linspace(t0, tn, n, endpoint=False)
SqrSignal = np.zeros(n)
FourierSeriesSignal = np.zeros(n)
SqrSignal = 1.0+signal.square(2*np.pi*5*t+np.pi*Width/T)
a0 = Fmax*Width/T
FourierSeriesSignal = a0
Factor = 2.0*Fmax/np.pi
for i in range(1,500):
    FourierSeriesSignal += Factor/(i)*np.sin(np.pi*i*Width/T)*np.cos(i*t*2*np.pi/T)
plt.plot(t, SqrSignal)
plt.plot(t, FourierSeriesSignal)
plt.ylim(-0.5, 2.5)
plt.show()
```


# CHEM 1000 - Spring 2022
Prof. Geoffrey Hutchison, University of Pittsburgh

## 5 Scalar and Vector Operators

Chapter 5 in [*Mathematical Methods for Chemists*](http://sites.bu.edu/straub/mathematical-methods-for-molecular-science/)

By the end of this session, you should be able to:
- Understand the concept of vector-valued functions and vector fields
- Identify sinks, sources, and saddle points of vector fields
- Understand the concept of vector operators
- Understand the gradient and applications to chemistry (e.g., forces)

### Scalars vs. Vectors

Reminder...

**Scalars** are just numbers - they have a magnitude, a size. The mass of a molecule would be an example, e.g., 120 amu.

**Vectors** have both a magnitude and a direction:
- velocity $\mathbf{v}$
- acceleration $\mathbf{a}$
- force $\mathbf{F}$
- electric field $\mathbf{E}$

### Scalar Functions and Vector Functions

A **function** takes in a number, a vector, etc. and returns a number:

$$
\sin 0 = 0
$$

Notice that $\sin x$ is a scalar function. You give it something, and it returns a **scalar**.

By extension, there must be **vector functions** too - one that returns a vector for every point.

Right now, you're experiencing force due to gravity. If you stand up, the forces acting on your body change over time. (Consider if you go on a roller coaster or fly in an airplane.)

$$
\overrightarrow{\boldsymbol{F}}(t)
$$

Notice that time is a scalar - it's just a number. So a vector function returns a vector regardless of what the input is. It might be one-dimensional (e.g., the force we feel at a given time $t$) or 2D or 3D, etc. (e.g., the forces on a satellite in space .. we probably care about the position of the satellite but maybe also time.)

### Vector Fields

When we have a vector function in 2D or 3D, we usually call these **vector fields**.

A **vector field** is a function that return a vector for every (x,y) or (x, y, z) point.

It sounds abstract, but we're actually already familiar with the concept. Consider a weather map showing wind:

Depending on our location, the wind speed and vector will differ. Let's look at a [tropical cyclone (Hurricane
Sally)](https://en.wikipedia.org/wiki/Tropical_Storm_Sally_(2020)) in the Gulf of Mexico.

Obviously, the further away from the hurricane, the wind speed (magnitude) decreases, and there's an *eye* in the center, where there's no wind at all. Also, the vector direction differs depending on where you are.

In chemistry, we encounter a range of vector functions and vector fields. For example:

- The force acting on a charge in an electric field. At each point in space, the force acting on the charge will have a specific magnitude and a direction.
- Fluid flow for which each element of the fluid at some point in space has a given speed and direction of flow.

There are a few key terms with vector fields:
- a **sink** is a point in which all vectors flow inward
- a **source** is a point at which all vectors flow outward
- a **saddle-point** has no net inward or outward flow (i.e., they balance exactly)

(Image from [*Mathematical Methods for Chemists*](http://sites.bu.edu/straub/mathematical-methods-for-molecular-science/))

Here, the red point is a **source** and all the arrows point away from it (imagine the forces from a negative charge towards an electron).

The black point is a **sink** and all the arrows point towards it (e.g., a positive charge that will attract an electron).

The two "X" points are **saddle points**. Notice that they come inward along one direction and outward along another direction. (Think about a horse saddle that rises up to the neck and head and also to the tail, and slopes down along the sides.)

### Example Vector Field

Let's try plotting

$$
\vec{F}=\frac{x}{5} \hat{\mathbf{x}}+\frac{y}{5} \hat{\mathbf{y}}
$$

We'll use `numpy` and `matplotlib`. The code is a bit different because we're creating a "quiverplot".

```python
# Let's plot some vector fields with numpy and matplotlib
# this is just our normal 'import numpy and matplotlib' code
import numpy as np
import matplotlib.pyplot as plt
%matplotlib inline
%config InlineBackend.figure_format = 'retina'
plt.style.use('./chem1000.mplstyle')
```

```python
# Example adapted from:
# https://pythonforundergradengineers.com/quiver-plot-with-matplotlib-and-jupyter-notebooks.html

# we're going to create a set of points from -5 to +5
# along x and y axes
x = np.arange(-5,6,1) # remember that np.arange() doesn't include the end value
y = np.arange(-5,6,1)
# remember, this takes the numpy arrays above and makes a mesh
X, Y = np.meshgrid(x, y)

# here's how we express the function
F_x = X/5 # x component
F_y = Y/5 # y component

# here's how we create a "quiver plot" with matplotlib
# if you want to know more, please ask
fig, ax = plt.subplots() # create a figure

ax.quiver(X,Y,F_x,F_y) # axis.quiver( X, Y mesh, Func X, Func Y )

# We'll go from -6 to +6 on each axis
# to see the arrows
ax.axis([-6, 6, -6, 6])
ax.set_aspect('equal') # make sure it's exactly square

plt.show()
```

**In the plot above, what kind of a point is at the origin?**

Let's try another example. This time, we'll plot:

$$
\vec{F}=\frac{x}{5} \hat{\mathbf{x}} - \frac{y}{5} \hat{\mathbf{y}}
$$

```python
# here's how we express the function
F_x = +X/5 # x component
F_y = -Y/5 # y component

fig, ax = plt.subplots() # create a figure
ax.quiver(X,Y,F_x,F_y) # axis.quiver( X, Y mesh, Func X, Func Y )

# We'll go from -6 to +6 on each axis
# to see the arrows
ax.axis([-6, 6, -6, 6])
ax.set_aspect('equal') # make sure it's exactly square

plt.show()
```

**Now what kind of point is at the origin?**

**Gradients**

Perhaps one of the most useful vector functions / vector fields in chemistry comes from the **gradient** of a scalar function.

The **gradient** operator in 2D Cartesian coordinates (x, y) is
$$
\boldsymbol{\nabla} \equiv \hat{\mathbf{x}} \frac{\partial}{\partial x}+\hat{\mathbf{y}} \frac{\partial}{\partial y}
$$

We use $\boldsymbol{\nabla}$ as a short-hand for the gradient operator, no matter what coordinate system is. Particularly in polar or spherical coordinates the expression can get complicated.

So what does it do?

$$
\boldsymbol{\nabla} V(x, y) = \left(\hat{\mathbf{x}} \frac{\partial}{\partial x}+\hat{\mathbf{y}} \frac{\partial}{\partial y}\right) V(x, y)
$$

At every point (x,y) we take the partial derivative of $V(x,y)$ (a scalar function) with respect to x and y, and use those as the x-component and y-component of a vector.

In other words, the gradient operator returns a **vector** from a **scalar** function.

**Laplace Operator**

The **Laplace** operator (sometimes called the Laplacian) in 2D Cartesian coordinates (x, y) is

$$
\nabla^{2} \equiv \frac{\partial^{2}}{\partial x^{2}}+\frac{\partial^{2}}{\partial y^{2}}
$$

We use $\nabla^2$ as a short-hand for the Laplace operator, no matter what coordinate system we use. At every point, we take the second partial derivatives of a function, and add the components.

The Laplace operator takes a scalar function and returns a new scalar function.

**Potential Energy and Forces**

Consider the interaction of two atoms according to the [Lennard-Jones potential](https://en.wikipedia.org/wiki/Lennard-Jones_potential):

$$
V(r)=4 \varepsilon\left[\left(\frac{\sigma}{r}\right)^{12}-\left(\frac{\sigma}{r}\right)^{6}\right]
$$

where $r$ is the distance between the two atoms, $\sigma$ represents the atomic diameter (or the sum of the two atomic radii), and $\epsilon$ represents the binding energy. In this case, we know that there will be a minimum energy at:

$$
r_{\min }=\sqrt[6]{2} \sigma
$$

Image from [*Mathematical Methods for Chemists*](http://sites.bu.edu/straub/mathematical-methods-for-molecular-science/):

Now the force will depend on the potential energy:

$$
\mathbf{F}=-\left[\hat{\mathbf{x}} \frac{\partial}{\partial x}+\hat{\mathbf{y}} \frac{\partial}{\partial y}+\hat{\mathbf{z}} \frac{\partial}{\partial z}\right] V(x, y, z)=-\nabla V(x, y, z)
$$

Put simply, the forces are the gradient of the potential energy!

$$
\mathbf{F}=-\boldsymbol{\nabla} V
$$

We can even do the work in this case - it's one-dimensional in $\hat{\mathbf{r}}$:
$$
\mathbf{F}(r)=-\frac{d V(r)}{d r} \hat{\mathbf{r}}=4 \varepsilon\left[12\left(\frac{\sigma^{12}}{r^{13}}\right)-6\left(\frac{\sigma^6}{r^7}\right)\right] \hat{\mathbf{r}}
$$

So when $r < r_{min}$ the repulsive force proportional to $1/r^{12}$ dominates the total force and pushes in a positive direction, while at distances $r > r_{min}$ the attractive force proportional to $1/r^6$ is dominant and pushes in a negative direction.

Notice that the forces on the atom always aim to minimize the energy:

- If the two atoms are close together, the repulsive force pushes in a direction towards $r_{min}$.
- If the two atoms are far apart, the attractive force pulls in a direction towards $r_{min}$.

While the Lennard-Jones potential has a very simple one-dimensional form (e.g., one atom is at the origin and the other is some distance $r$ away), the concept of the gradient and forces applies regardless of the potential energy.

If we have some method to calculate the potential energy (e.g., quantum chemistry, etc.) we can take the gradient, get the forces on the atoms, and move them accordingly to get a minimum energy.

In this case, we can see that the central carbon-carbon bond is too short and the carbon atoms are pulling apart, resulting in the hydrogen atoms moving in different directions.

Eventually as we repeat the process (find the gradient and forces, move the atoms a bit, re-calculate), we can minimize the potential energy and find an optimized geometry.

**Example problem:**

For a scalar potential energy, $V(x, y, z)=x^{2}+y^{2}+z^{2}$, derive the force defined as the negative gradient of the potential.

```python
from sympy import init_session
init_session()
```

```python
V = x**2 + y**2 + z**2
# get the x-component
diff(V, x)
```

```python
# the y-component
diff(V, y)
```

```python
# the z-component
diff(V, z)
```

Okay, we probably could have done that by inspection.

$$
\boldsymbol{\nabla} V(x,y,z) = 2x \hat{\mathbf{x}} + 2y \hat{\mathbf{y}} + 2z \hat{\mathbf{z}}
$$

**If that's the gradient, what are the forces?**

### Gradient in Spherical Coordinates

One last thing.. so far, we've expressed the gradient in 2D or 3D Cartesian coordinates, and it's not too bad.

In spherical coordinates, it's a little messier:

$$
\boldsymbol{\nabla} V=\hat{\mathbf{r}} \frac{\partial V}{\partial r}+\hat{\boldsymbol{\theta}} \frac{1}{r} \frac{\partial V}{\partial \theta}+\hat{\boldsymbol{\varphi}} \frac{1}{r \sin \theta} \frac{\partial V}{\partial \varphi}
$$

The Laplace operator is similarly messy:

$$
\nabla^{2} V=\frac{1}{r^{2}} \frac{\partial}{\partial r}\left(r^{2} \frac{\partial V}{\partial r}\right)+\frac{1}{r^{2} \sin \theta} \frac{\partial}{\partial \theta}\left(\sin \theta \frac{\partial V}{\partial \theta}\right)+\frac{1}{r^{2} \sin ^{2} \theta} \frac{\partial^{2} V}{\partial \varphi^{2}}
$$

We'll come back to these as we need them, but it's a nice illustration of why we use operators. We can write a short symbol and it represents a longer, messier operator.

-------
This notebook is from Prof. Geoffrey Hutchison, University of Pittsburgh
https://github.com/ghutchis/chem1000


# Column Fraction Function

- Using **scipy.interpolate.RectBivariateSpline** to interpolate function:
- Linear interpolation (kx=1, ky=1)

```python
import numpy as np
import pandas as pd
import sympy as s
import math
import glob

import matplotlib
import matplotlib.pyplot as plt
from cycler import cycler

from scipy import interpolate

from astropy import units as u
from astropy import constants as const
from astropy.units import imperial
imperial.enable()

%matplotlib inline
import os
#home_dir = os.environ['/Users/eriksolhaug'] + '/'
import pyCloudy as pc
```

```python
pd.set_option("display.max_rows", None, "display.max_columns", None)
pd.set_option('precision', 16)
pd.set_option('display.float_format', lambda x: '%.3e' % x)
```

```python
# The directory in which we will have the model
# You may want to change this to a different place so that the current directory
# will not receive all the Cloudy files.
dir_ = '/Users/eriksolhaug/cloudy/c17.02/cloudoutputs/'

user_dir = '/Users/eriksolhaug/cloudy/c17.02/'

# Define verbosity to high level (will print errors, warnings and messages)
pc.log_.level = 3
```

```python
def calccolfrac(element, ion_number, model_name, model_dir):
    """
    Returns FUNCTION object f(X, Y) where X is a hydrogen density (LOG) and Y is a temperature (LOG). The output fraction is in LOG.

    The temperature range given to the Cloudy input file is 10^4 - 10^7 K
    and the hydrogen density range given to the Cloudy input file is 10^-5.0 - 10^-3.0 cm^-3

    Inputs:
    element - a string, element pertaining to the ion to be analyzed
    ion_number - an integer, denoting the ionization state of the ion (e.g. a request for CIV would need ion_number=4)
    model_name - a string, name of model run in CLOUDY (f.ex. 'model_41')
    model_dir - a string, directory containing model output files

    Output:
    f - an object, function yielding fractional column for the requested ionic species as a fraction in LOG of its column relative
    to all other available ions for the requested element as a function of temperature and hydrogen density
    To get the column fraction for a certain hden X and temperature Y, simply call f(X, Y) and take the exponent 10**f(X, Y) of this to find
    the column fraction of the ion species relative to all ions of that element.

    An example:
                In[]:
                f = calccolfrac('O', 6, 'model_47', '/Users/eriksolhaug/cloudy/c17.02/cloudoutputs/')
                10**f(-5.0, 5.5)
                Out[]:
                array([[0.05054196]]) -- where 0.05054196 is the column fraction of OVI for hydrogen density of 10**-5.0 and a temperature of 10**5.5 Kelvin
    """

    # Interpolating function for the data computed in Cloudy
    grid_df = pd.read_csv(f'{model_dir}/{model_name}.grid', sep='\t')
    hdengrid = grid_df.loc[:, 'HDEN=%f L']
    tempgrid = grid_df.loc[:, 'CONSTANT ']
    x = hdengrid
    y = tempgrid

    # Loading .ele_ files
    col_df = pd.read_csv(f'{model_dir}{model_name}.ele_{element}', sep = '\t')

    # Getting fractional columns
    z_array = []
    for index in range(0, x.size):
        frac_col = col_df.iloc[index*2].iloc[ion_number]
        z_array.append(np.log10(frac_col+1e-20))  #Adding a small value to avoide case of log(0)
    z = pd.DataFrame(z_array, columns=['z'])

    if model_name == 'model_43':
        step = 5
    elif model_name == 'model_45':
        step = 41
    elif model_name == 'model_46':
        step = 401
    elif model_name == 'model_47':
        step = 101
    else:
        step = 11

    # Putting vectors in dataframe representation
    xyz = pd.DataFrame({'x' : x, 'y' : y, 'z' : z['z']})

    # Simplifying x and y inputs
    xi = xyz['x'][:step]
    yi = xyz['y'][::step]

    # Preparing spline arrays
    twoDarray = []
    for i in range(len(xi)):
        array = []
        for j in range(len(yi)):
            idx = i + j*step
            array.append(xyz['z'][idx])
        twoDarray.append(array)

    # Simplifying z inputs
    zi = twoDarray

    print(xi, yi, zi)
    print(len(xi), len(yi), len(zi))

    # INTERPOLATION
    f = interpolate.RectBivariateSpline(xi, yi, zi, kx=1, ky=1) # Linear determined by kx, ky

    # Displaying match between old fractions and interpolated function
    interpolated_z = []
    for temp in yi:
        for hden in xi:
            interpolated_z.append(f(hden, temp))
    interpolated_z = np.concatenate(interpolated_z)

    print(interpolated_z)

    return f
```

```python
def colplot(keyword, second_val, element, ion_numbers, model_name, model_dir, plot_dir):
    '''
    Function used to make and save plots of the fractional columns for different elements

    Input:
    keyword - a string, either 'temp' or 'hden' for what needs to be plotted against
    second_val - a float in LOG, value for either temp or hden (whatever not requested by keyword) - the plot needs a set temp or hden and this is set by this parameter
    element - a string, chemical symbol for the requested element (f.ex. 'C')
    ion_numbers - an list, denoting the ionization states of the ions in the list (e.g. a request for CIV would need ion_number=4)
    model_name - a string, name of model run in CLOUDY (f.ex. 'model_42')
    model_dir - a string, directory containing model output files
    plot_dir - a string, directory for where to save plot
    '''

    # Defining Roman Numerals
    def int_to_Roman(num):
        val = [
            1000, 900, 500, 400,
            100, 90, 50, 40,
            10, 9, 5, 4,
            1
            ]
        syb = [
            "M", "CM", "D", "CD",
            "C", "XC", "L", "XL",
            "X", "IX", "V", "IV",
            "I"
            ]
        roman_num = ''
        i = 0
        while  num > 0:
            for _ in range(num // val[i]):
                roman_num += syb[i]
                num -= val[i]
            i += 1
        return roman_num

    # Plotting
    tick_fontsize = 16
    axis_fontsize = 22
    lwidth = 3
    lstyle = '-'

    fig, ax = plt.subplots(1, 1)
    fig.set_size_inches(8, 5)
    fig.tight_layout(w_pad = 10.0)

    if keyword == 'temp':
        constant = second_val
        vary = np.arange(4, 7, 0.01)
        xaxis = 10**vary
    elif keyword == 'hden':
        vary = np.arange(-5.0, -3.0, 0.01)
        constant = second_val
        xaxis = 10**vary
    else:
        print('Not a valid keyword. Needs either "temp" or "hden".')

    tot = np.zeros(len(vary))
    for ion_number in ion_numbers:
        f = calccolfrac(element, ion_number, model_name, model_dir)
        ion_name = element+int_to_Roman(ion_number)
        if keyword == 'temp':
            other_keyword = 'hden'
            ax.plot(xaxis, 10**f(constant, vary)[0], linewidth=lwidth, label=f'{ion_name}', linestyle=lstyle)
            for i in range(0, len(tot)):
                tot[i] += 10**f(constant, vary[i])[0]
        elif keyword == 'hden':
            other_keyword = 'temp'
            ax.plot(xaxis, 10**f(vary, constant), linewidth=lwidth, label=f'{ion_name}', linestyle=lstyle)
            for i in range(0, len(tot)):
                tot[i] += 10**f(vary[i], constant)
    print(f'Tot: {tot}')
    ax.plot(xaxis, tot, c='k', linewidth=lwidth, linestyle=':', label = 'tot')
    ax.set_xscale('log')
    ax.set_yscale('log')
    if keyword == 'temp':
        ax.set_xlabel(r'Temperature ($K$)', fontsize = 12)
    else:
        ax.set_xlabel(r'Hydrogen density ($cm^{-3}$)', fontsize = 12)
    ax.set_ylabel('Column Fraction', fontsize = 12)
    ax.set_ylim(1e-6, 5e0)
    ax.set_title(f'Fractional Columns for {element} | Constant {other_keyword}: 10^{second_val}', fontsize=18, fontweight='bold', pad=15)
    for tick in ax.xaxis.get_major_ticks():
        tick.label.set_fontsize(fontsize=tick_fontsize)
    for tick in ax.yaxis.get_major_ticks():
        tick.label.set_fontsize(fontsize=tick_fontsize)
    ax.xaxis.label.set_fontsize(axis_fontsize)
    ax.yaxis.label.set_fontsize(axis_fontsize)
    ax.tick_params(which='major', width=2, length=8)
    ax.tick_params(which='minor', width=1, length=5)
    ax.grid(linestyle='--')
    ax.legend(bbox_to_anchor=(1.05, 1), loc='upper left', prop={'size': 18})
    ax.set_prop_cycle(cycler('color', ['c', 'm', 'y', 'k']))

    fig.savefig(f'{plot_dir}/{element}_{keyword}vary_{second_val}_columnfractions.pdf', bbox_inches="tight")
```

### Example Executions:

```python
element = 'O'
ion_number = 6
model_name = 'model_47'
model_dir = '/Users/eriksolhaug/cloudy/c17.02/cloudoutputs/'
f = calccolfrac(element, ion_number, model_name, model_dir)
```

```python
10**f(-5.0,5.5)
```

$\: \uparrow$ This is the fractional column of the input ion relative to all other ions for that element.

One can use the same procedure for other ions:

```python
element = 'C'
ion_number = 4
model_name = 'model_47'
model_dir = '/Users/eriksolhaug/cloudy/c17.02/cloudoutputs/'
f = calccolfrac(element, ion_number, model_name, model_dir)
```

### Plotting fractional columns for all ions (C, N, O, Si)

```python
model_name = 'model_47'
plot_dir = '/Users/eriksolhaug/cloudy/c17.02/es/es_data/column_fractions/' + model_name + '/'
element_array = ['C', 'N', 'O', 'Si']
keyword_array = ['temp', 'hden']
temp_val_array = np.arange(-5.0, -2.5, 0.5)
hden_val_array = np.arange(4.0, 6.5, 0.5)
for element in element_array:
    if element == 'C':
        ion_numbers = np.arange(1, 6+1+1)
    elif element == 'N':
        ion_numbers = np.arange(1, 7+1+1)
    elif element == 'O':
        ion_numbers = np.arange(1, 8+1+1)
    elif element == 'Si':
        ion_numbers = np.arange(1, 14+1+1)
    for keyword in keyword_array:
        if keyword == 'temp':
            second_val_array = temp_val_array
        elif keyword == 'hden':
            second_val_array = hden_val_array
        for second_val in second_val_array:
            plot = colplot(keyword, second_val, element, ion_numbers, model_name, model_dir, plot_dir)
```

```python

```