# Gentle Introduction to Pytorch Autograd with Linear Regression
#### By Michael Przystupa

** By the end of this tutorial students will be able to:**:
- Analytically solve a linear regression
- Explain what pytorch automatically handles with it's autograd library
- Construct a linear model using the pytorch module

## Introduction

Linear regression is a fairly common practise and is a good starting place with introducing you building models using the pytorch library. We'll start by showing how analytically you can solve linear regression analytically which we can use to compare how our pytorch model compares.

First, let's import the libraries we'll need:

```python
import torch
import torch.nn as nn
import torch.optim as optim
import torch.nn.functional as F
import numpy as np
import matplotlib.pyplot as plt
#just your typical imports in python
```

## The Data Set

We'll be working with a toy dataset which is just a linear function with added guassian noise:
\begin{equation}
 f(x) = ax + b + N(0, 1)
\end{equation}
In real machine learning, your data won't be so nice, but we want something we can easily look at to get some intuition about how the model works.

** note ** : In general, we do not KNOW the function our data came from. That's why we want to build models to try and approximate something that looks like our data COULD have come from it.

```python
def f( x, a= 2.0, b = 1.0, add_noise=True):
    if add_noise:
        n = np.random.normal(0.0, 1.0)
    else:
            n = 0.0
    return a * x + b + n

X = np.linspace(-5.0, 5.0, num= 200)
y = np.array([f(x) for x in X])

plt.plot(X, y, 'g^')
plt.show()

```

## Analytically
Clearly, there is some relationship in this data, and we'd like to represent this some how with a function of the form
\begin{equation}
w_{1} x + w_{2} = \hat y
\end{equation}

In machine learning, to find these w's we generally want to have some objective to see how good the weights are. One example is minimizing the mean-squared error which for our equation looks like this:
\begin{equation}
Objective(w) = \sum_{i=1}^{n} (w_{1} x_{i} + w_{2} + - y_{i}) ^{2}
\end{equation}

You'll have to trust me on this, but we can solve this directly by doing the following manipulation:
\begin{equation}
w = (X^{t} X)^{-1} (X^{t} y)
\end{equation}

where X is our data matrix and y is our labels, and w is our weights we want to parameterize

```python
# Solving problem analytically:
X = np.stack((X, np.ones(X.shape[0])), axis=1) #this is so we have w_{2}
y = np.array([y]).transpose() # has to do with shape of data
X_t = X.transpose()

#Calculating each of the components
X_tX = np.matmul(X_t, X)
X_ty = np.matmul(X_t, y)
w = np.matmul(np.linalg.inv(X_tX) , (X_ty))

#this will show the best weights we can do with the data
print('w_0 = {}, w_1 = {}'.format(w[0], w[1]))

#Plotting our approximations of true values

#lambda is a key word in python for defining smaller functions
lin_model = lambda x : f(x, w[0], w[1], add_noise=False)
y_hat = [lin_model(x) for x in X[:,0]]

plt.plot(X[:,0], y, 'g^')
plt.plot(X[:,0], y_hat, 'r--')
plt.show()

```

This is all well and good...if you can calculate your inverse. In that case, you're going to have to get creative and do some optimization.

## Stochastic Gradient Descent  (S.G.D)

This can be fairly math heavy, but here's the gist of what we're going to do:
\begin{equation}
    w = w - \alpha * \nabla Objective(w)
\end{equation}

where $\alpha$ is referred to as the learning rate and is the amount we will update our current weights by at each step and $\nabla Objective(w)$ means we want to take the gradient with respect to our loss. One minor thing with this is that in S.G.D. we are making updates based off single examples in our data and this is done usually for performance reasons (although there is theory about why it works well too).

Maybe this all seems a little scary, particularly if you're not sure what  a gradient is. Thankfully, pytorch makes it so you don't need to pull out your old calculus textbook and handles calculating these for us as we'll see.

## Pytorch's autograd system

The most useful part of pytorch is the autograd system. The gist of it is that it will automatically calculate gradients for any operations done on a pytorch tensor, so we have to say good bye to our good friend numpy, which is easy to do:

```python
X = torch.from_numpy(X).float()
y = torch.from_numpy(y).float()

# Now, to do SGD we'll need a learning rate and an initial set of weights:
alpha = 1e-3
w = torch.randn(2, requires_grad=True) #requires grad is key to do this
```

To use pytorch's autograd feature, we're going to calculate the prediction with our current weights and call backward on our loss function. When you call .backward() on a tensor of interest, it calculates the gradient and we can access it with the .grad field of tensors to do our update

```python
order = torch.randperm(X.size()[0])

#redefine our linear model to use the tensor version of w
for _ in range(0, 10):
    total_loss = 0.0
    for i in order:
        x = X[i,:]
        y_hat = torch.matmul(x, w) #ax + b esentially
        loss = F.mse_loss(y_hat[0], y[i,0]) #our objective function
        loss.backward() #where our gradients get calculated
        total_loss += loss.item()
        with torch.no_grad():
            w -=  alpha * w.grad
            w.grad.zero_()
    print('Loss on epoch {}:  {}'.format(_, total_loss / X.size()[0]))

print(w)

```

    /home/mprzystupa/anaconda2/lib/python2.7/site-packages/ipykernel_launcher.py:9: UserWarning: invalid index of a 0-dim tensor. This will be an error in PyTorch 0.5. Use tensor.item() to convert a 0-dim tensor to a Python number
      if __name__ == '__main__':

    Loss on epoch 0:  15.95085129
    Loss on epoch 1:  1.26909185622
    Loss on epoch 2:  1.10887930466
    Loss on epoch 3:  1.05333386042
    Loss on epoch 4:  1.02935324583
    Loss on epoch 5:  1.01903245275
    Loss on epoch 6:  1.01468058052
    Loss on epoch 7:  1.01290969603
    Loss on epoch 8:  1.01223487893
    Loss on epoch 9:  1.01201125071
    tensor([1.9353, 1.0800], requires_grad=True)

```python

X = X.numpy()
y = y.numpy()
w = w.detach().numpy()
lin_model = lambda x : f(x, w[0], w[1], add_noise=False)
y_hat = [lin_model(x) for x in X[:,0]]

plt.plot(X[:,0], y, 'g^')
plt.plot(X[:,0], y_hat, 'r--')
plt.show()
```

As we can see, our model now performsabout the same as our analytical solution, but the advantage is that doing it this way is always possible and doesn't rely on our inverse having to exist.

## Going full Pytorch
Now the above loop was to illustrate what pytorch is doing. We cheated a little bit by also using the .backward() call instead of calculating it our selves, but that's the reason WHY these deep learning frameworks are great, because you don't HAVE to do this all by hand.

```python
X = torch.from_numpy(X).float()
y = torch.from_numpy(y).float()

class LinRegression(nn.Module):
    def __init__(self):
        super(LinRegression, self).__init__()
        self.linear = nn.Linear(2, 1, bias=False )
        #the bias term is generally set to True, but we added the 1's column to x

    def forward(self, x):
        return self.linear(x)

model = LinRegression()
sgd = optim.SGD(model.parameters(), lr=1e-3)
order = torch.randperm(X.size()[0])
for _ in range(0, 10):
    total_loss = 0.0
    for i in order:
        x = X[i,:]
        y_hat = model(x)
        loss = F.mse_loss(y_hat[0], y[i,0]) #our objective function
        loss.backward() #where our gradients get calculated
        total_loss += loss.item()
        sgd.step()
        sgd.zero_grad()
    print('Loss on epoch {}:  {}'.format(_, total_loss / X.size()[0]))

y_hat = model(X)

y_hat = y_hat.detach().numpy()
X = X.numpy()
y = y.numpy()

plt.plot(X[:,0], y, 'g^')
plt.plot(X[:,0], y_hat, 'r--')
plt.show()

```

## Conclusions

In this tutorial we've seen how to solve linear regression using pytorch's autograd system and how to build models with pytorch. We saw several pytorch packages including torch, torch.optim, torch.nn, and torch.nn.functional. We then saw how to use these modules together to train our model.

## Exercise: Logistic Regression
Logistic regression is a sort of extension of linear regression for the classification setting.

Using what you've seen in this tutorial and class, build a logistic regression model to train a model for the following toy dataset:

```python
def f(x):
    if x > 0.0:
        return 1.0
    else:
        return 0.0

X = np.linspace(-5.0, 5.0, num= 200)
y = np.array([f(x) for x in X])

plt.plot(X, y, 'g^')
plt.show()
```

```python
#starter code:
X = torch.from_numpy(X)
y = torch.from_numpy(y)
class LogisticRegression(nn.Module):
    def __init__(self):
        super(LogisticRegression, self).__init__()
        #insert code here
    def forward(self, x):
        #insert code
        return x

# some set-up
order = torch.randperm(X.size()[0])
#training loop
for _ in range(0, 10):
    total_loss = 0.0
    for i in order:
        x = X[i]
        y_hat = model(x)
        #insert call to loss function here

        loss.backward() #where our gradients get calculated
        total_loss += loss.item()
        sgd.step()
        sgd.zero_grad()
    print('Loss on epoch {}:  {}'.format(_, total_loss / X.size()[0]))

y_hat = model(X)
y_hat = y_hat.detach().numpy()
X = X.numpy()
y = y.numpy()

plt.plot(X[:,0], y, 'g^')
plt.plot(X[:,0], y_hat, 'r--')
plt.show()

```


# Conservative SDOF - Harmonic Balance (4)

- Good for strongly and weakly NL systems
- Requires significant knowledge about system studied
- Similar to Galerkin method, but in time

## Overview

Assume an harmonic expansion for $x(t)$

$$
\begin{gather*}
     x(t) = \sum^N_{m = 0} A_m \cos\left( m\omega t + m \beta_0\right)
\end{gather*}
$$

substitute this into the EOM and enforce zero projected residual, i.e. taking the inner product with orthogonal assumption.

$$
\begin{gather*}
    \left< EOM_{LHS},\phantom{-} \cos\left(m\omega t + m\beta_0\right) \right> = 0\\
    m = 0, 1, 2, \cdots, N
\end{gather*}
$$

which is

$$
\begin{gather*}
    \int_0^{\frac{2\pi}{\omega} } \left( EOM_{LHS}\bullet \cos\left(m\omega t + m\beta_0\right) \right)dt
\end{gather*}
$$

Then we solve for $A_0, A_2, A_3, \cdots, A_n$ and $\omega$ in terms of $A_1$ and other system parameters. In this example we solve the following equation

$$
\begin{gather*}
    \ddot{x} + \omega_0^2 x + \alpha_2 x^2 + \alpha_3 x^3 = 0.
\end{gather*}
$$

## Assumed Harmonic Expansion - up to and including $\cos\left(3\left(\omega t + \beta \right)\right)$ - fails

Assume a harmonic expansion for $x(t)$ - will use this to satisfy the cubic equation, so include harmonic terms at twice the linear natrual frequency

Note: phi = omega * t + beta0

```python
import sympy as sp
from sympy.simplify.fu import TR4, TR7, TR8, TR11
import numpy as np
from itertools import combinations_with_replacement
from math import factorial
```

```python
M = 3  # order of harmonic expansion
Ai = sp.symbols('A_(0:' + str(M+1) + ')', real=True)  # coefficient A
t, w, w0 = sp.symbols('t omega omega_0', real=True, positive=True)
b0 = sp.Symbol('beta_0')

# harmonic expansion of x(t)
phi = sp.Symbol('phi')
xh = sp.Function('x_h')(phi)
xh = sum([Ai[i] * sp.cos(i * phi) for i in range(M+1)])
xh
```

$\displaystyle A_{0} + A_{1} \cos{\left(\phi \right)} + A_{2} \cos{\left(2 \phi \right)} + A_{3} \cos{\left(3 \phi \right)}$

```python
# Defining the EOM
N = 3  # order of powered terms of alpha
ai = sp.symbols('alpha_(2:' + str(N+1) + ')')
x = sp.Function('x')(phi)

# EOM
EOM = w**2 * sp.diff(x, phi, 2) + w0**2 * x + sum([ai[i-2] * x**i for i in range(2, N+1)])
EOM = sp.Eq(EOM, 0)
EOM
```

$\displaystyle \alpha_{2} x^{2}{\left(\phi \right)} + \alpha_{3} x^{3}{\left(\phi \right)} + \omega^{2} \frac{d^{2}}{d \phi^{2}} x{\left(\phi \right)} + \omega_{0}^{2} x{\left(\phi \right)} = 0$

```python
# Substitute x(t) for the harmonic expansion of x(t) into EOM
err = EOM.lhs.subs(x, xh).doit().expand()
err = TR7(err).expand()
err = TR8(err).expand()
err
```

$\displaystyle A_{0}^{3} \alpha_{3} + 3 A_{0}^{2} A_{1} \alpha_{3} \cos{\left(\phi \right)} + 3 A_{0}^{2} A_{2} \alpha_{3} \cos{\left(2 \phi \right)} + 3 A_{0}^{2} A_{3} \alpha_{3} \cos{\left(3 \phi \right)} + A_{0}^{2} \alpha_{2} + \frac{3 A_{0} A_{1}^{2} \alpha_{3} \cos{\left(2 \phi \right)}}{2} + \frac{3 A_{0} A_{1}^{2} \alpha_{3}}{2} + 3 A_{0} A_{1} A_{2} \alpha_{3} \cos{\left(\phi \right)} + 3 A_{0} A_{1} A_{2} \alpha_{3} \cos{\left(3 \phi \right)} + 3 A_{0} A_{1} A_{3} \alpha_{3} \cos{\left(2 \phi \right)} + 3 A_{0} A_{1} A_{3} \alpha_{3} \cos{\left(4 \phi \right)} + 2 A_{0} A_{1} \alpha_{2} \cos{\left(\phi \right)} + \frac{3 A_{0} A_{2}^{2} \alpha_{3} \cos{\left(4 \phi \right)}}{2} + \frac{3 A_{0} A_{2}^{2} \alpha_{3}}{2} + 3 A_{0} A_{2} A_{3} \alpha_{3} \cos{\left(\phi \right)} + 3 A_{0} A_{2} A_{3} \alpha_{3} \cos{\left(5 \phi \right)} + 2 A_{0} A_{2} \alpha_{2} \cos{\left(2 \phi \right)} + \frac{3 A_{0} A_{3}^{2} \alpha_{3} \cos{\left(6 \phi \right)}}{2} + \frac{3 A_{0} A_{3}^{2} \alpha_{3}}{2} + 2 A_{0} A_{3} \alpha_{2} \cos{\left(3 \phi \right)} + A_{0} \omega_{0}^{2} + \frac{3 A_{1}^{3} \alpha_{3} \cos{\left(\phi \right)}}{4} + \frac{A_{1}^{3} \alpha_{3} \cos{\left(3 \phi \right)}}{4} + \frac{3 A_{1}^{2} A_{2} \alpha_{3} \cos{\left(2 \phi \right)}}{2} + \frac{3 A_{1}^{2} A_{2} \alpha_{3} \cos{\left(4 \phi \right)}}{4} + \frac{3 A_{1}^{2} A_{2} \alpha_{3}}{4} + \frac{3 A_{1}^{2} A_{3} \alpha_{3} \cos{\left(\phi \right)}}{4} + \frac{3 A_{1}^{2} A_{3} \alpha_{3} \cos{\left(3 \phi \right)}}{2} + \frac{3 A_{1}^{2} A_{3} \alpha_{3} \cos{\left(5 \phi \right)}}{4} + \frac{A_{1}^{2} \alpha_{2} \cos{\left(2 \phi \right)}}{2} + \frac{A_{1}^{2} \alpha_{2}}{2} + \frac{3 A_{1} A_{2}^{2} \alpha_{3} \cos{\left(\phi \right)}}{2} + \frac{3 A_{1} A_{2}^{2} \alpha_{3} \cos{\left(3 \phi \right)}}{4} + \frac{3 A_{1} A_{2}^{2} \alpha_{3} \cos{\left(5 \phi \right)}}{4} + \frac{3 A_{1} A_{2} A_{3} \alpha_{3} \cos{\left(2 \phi \right)}}{2} + \frac{3 A_{1} A_{2} A_{3} \alpha_{3} \cos{\left(4 \phi \right)}}{2} + \frac{3 A_{1} A_{2} A_{3} \alpha_{3} \cos{\left(6 \phi \right)}}{2} + \frac{3 A_{1} A_{2} A_{3} \alpha_{3}}{2} + A_{1} A_{2} \alpha_{2} \cos{\left(\phi \right)} + A_{1} A_{2} \alpha_{2} \cos{\left(3 \phi \right)} + \frac{3 A_{1} A_{3}^{2} \alpha_{3} \cos{\left(\phi \right)}}{2} + \frac{3 A_{1} A_{3}^{2} \alpha_{3} \cos{\left(5 \phi \right)}}{4} + \frac{3 A_{1} A_{3}^{2} \alpha_{3} \cos{\left(7 \phi \right)}}{4} + A_{1} A_{3} \alpha_{2} \cos{\left(2 \phi \right)} + A_{1} A_{3} \alpha_{2} \cos{\left(4 \phi \right)} - A_{1} \omega^{2} \cos{\left(\phi \right)} + A_{1} \omega_{0}^{2} \cos{\left(\phi \right)} + \frac{3 A_{2}^{3} \alpha_{3} \cos{\left(2 \phi \right)}}{4} + \frac{A_{2}^{3} \alpha_{3} \cos{\left(6 \phi \right)}}{4} + \frac{3 A_{2}^{2} A_{3} \alpha_{3} \cos{\left(\phi \right)}}{4} + \frac{3 A_{2}^{2} A_{3} \alpha_{3} \cos{\left(3 \phi \right)}}{2} + \frac{3 A_{2}^{2} A_{3} \alpha_{3} \cos{\left(7 \phi \right)}}{4} + \frac{A_{2}^{2} \alpha_{2} \cos{\left(4 \phi \right)}}{2} + \frac{A_{2}^{2} \alpha_{2}}{2} + \frac{3 A_{2} A_{3}^{2} \alpha_{3} \cos{\left(2 \phi \right)}}{2} + \frac{3 A_{2} A_{3}^{2} \alpha_{3} \cos{\left(4 \phi \right)}}{4} + \frac{3 A_{2} A_{3}^{2} \alpha_{3} \cos{\left(8 \phi \right)}}{4} + A_{2} A_{3} \alpha_{2} \cos{\left(\phi \right)} + A_{2} A_{3} \alpha_{2} \cos{\left(5 \phi \right)} - 4 A_{2} \omega^{2} \cos{\left(2 \phi \right)} + A_{2} \omega_{0}^{2} \cos{\left(2 \phi \right)} + \frac{3 A_{3}^{3} \alpha_{3} \cos{\left(3 \phi \right)}}{4} + \frac{A_{3}^{3} \alpha_{3} \cos{\left(9 \phi \right)}}{4} + \frac{A_{3}^{2} \alpha_{2} \cos{\left(6 \phi \right)}}{2} + \frac{A_{3}^{2} \alpha_{2}}{2} - 9 A_{3} \omega^{2} \cos{\left(3 \phi \right)} + A_{3} \omega_{0}^{2} \cos{\left(3 \phi \right)}$

Note there are terms with frequency content as high as $6\phi$

To find $A_0$, $\omega$, $A_2$, we project the error back on the assumed expansion and set the projections to zero

Note: same as setting the coefficients in front of each $\cos(n\phi)$ to zero for $n = 0, 1, 2$

```python
proj_eqn = {}
for i in range(M+1):
    proj_eqn['eq'+str(i)] = sp.integrate(sp.expand_trig(err * sp.cos(i * phi)), (phi, 0, 2*sp.pi))
```

```python
proj_eqn['eq0']
```

$\displaystyle 2 \pi A_{0}^{3} \alpha_{3} + 2 \pi A_{0}^{2} \alpha_{2} + 3 \pi A_{0} A_{1}^{2} \alpha_{3} + 3 \pi A_{0} A_{2}^{2} \alpha_{3} + 3 \pi A_{0} A_{3}^{2} \alpha_{3} + 2 \pi A_{0} \omega_{0}^{2} + \frac{3 \pi A_{1}^{2} A_{2} \alpha_{3}}{2} + \pi A_{1}^{2} \alpha_{2} + 3 \pi A_{1} A_{2} A_{3} \alpha_{3} + \pi A_{2}^{2} \alpha_{2} + \pi A_{3}^{2} \alpha_{2}$

```python
proj_eqn['eq1']
```

$\displaystyle 3 \pi A_{0}^{2} A_{1} \alpha_{3} + 3 \pi A_{0} A_{1} A_{2} \alpha_{3} + 2 \pi A_{0} A_{1} \alpha_{2} + 3 \pi A_{0} A_{2} A_{3} \alpha_{3} + \frac{3 \pi A_{1}^{3} \alpha_{3}}{4} + \frac{3 \pi A_{1}^{2} A_{3} \alpha_{3}}{4} + \frac{3 \pi A_{1} A_{2}^{2} \alpha_{3}}{2} + \pi A_{1} A_{2} \alpha_{2} + \frac{3 \pi A_{1} A_{3}^{2} \alpha_{3}}{2} - \pi A_{1} \omega^{2} + \pi A_{1} \omega_{0}^{2} + \frac{3 \pi A_{2}^{2} A_{3} \alpha_{3}}{4} + \pi A_{2} A_{3} \alpha_{2}$

```python
proj_eqn['eq2']
```

$\displaystyle 3 \pi A_{0}^{2} A_{2} \alpha_{3} + \frac{3 \pi A_{0} A_{1}^{2} \alpha_{3}}{2} + 3 \pi A_{0} A_{1} A_{3} \alpha_{3} + 2 \pi A_{0} A_{2} \alpha_{2} + \frac{3 \pi A_{1}^{2} A_{2} \alpha_{3}}{2} + \frac{\pi A_{1}^{2} \alpha_{2}}{2} + \frac{3 \pi A_{1} A_{2} A_{3} \alpha_{3}}{2} + \pi A_{1} A_{3} \alpha_{2} + \frac{3 \pi A_{2}^{3} \alpha_{3}}{4} + \frac{3 \pi A_{2} A_{3}^{2} \alpha_{3}}{2} - 4 \pi A_{2} \omega^{2} + \pi A_{2} \omega_{0}^{2}$

```python
proj_eqn['eq3']
```

$\displaystyle 3 \pi A_{0}^{2} A_{3} \alpha_{3} + 3 \pi A_{0} A_{1} A_{2} \alpha_{3} + 2 \pi A_{0} A_{3} \alpha_{2} + \frac{\pi A_{1}^{3} \alpha_{3}}{4} + \frac{3 \pi A_{1}^{2} A_{3} \alpha_{3}}{2} + \frac{3 \pi A_{1} A_{2}^{2} \alpha_{3}}{4} + \pi A_{1} A_{2} \alpha_{2} + \frac{3 \pi A_{2}^{2} A_{3} \alpha_{3}}{2} + \frac{3 \pi A_{3}^{3} \alpha_{3}}{4} - 9 \pi A_{3} \omega^{2} + \pi A_{3} \omega_{0}^{2}$

if $A_1$ is small neglect terms with $A_0^2, A_2^2, A_0A_2, A_1^2A_0, A_0^3, A_1^2A_2, \cdots$

This is because there is the assumption of
$$
\begin{gather*}
    A_1 > A_i \phantom{---} i = 0, 2, 3, \cdots
\end{gather*}
$$

equation 0 tells us that $A_0$ is proportional to $A_1$ squared - We have a constant term multiplying $A_0$ (in addition to others) and no $A_1$ terms but only $A_1^2$ terms.

```python
deglist = list(map(int, sp.degree_list(proj_eqn['eq0'], gens=Ai)))
A1_deg = deglist[1]
removing_combos = []
Ai_ = tuple(filter(lambda A: A != Ai[1], Ai))

# A1 ** n -> n can be up to O(A1**4) which is square of A1**2
keep_terms = []
for i in range(A1_deg+1, A1_deg**2+1):
    temp = []
    for j in range(i):
        temp.append(Ai[1])
    keep_terms.append(tuple(temp))

temp = list(combinations_with_replacement(Ai_, A1_deg))
for t in temp:
    if t not in keep_terms:
        removing_combos.append(t)

for i in range(A1_deg+1, max(deglist)+1):
    temp = list(combinations_with_replacement(Ai, i))
    for t in temp:
        if t not in keep_terms:
            removing_combos.append(t)

eq_copy = proj_eqn['eq0']
for term in removing_combos:
    eq_copy = eq_copy.subs(np.prod(term), 0)

eq_copy = sp.Eq(eq_copy, 0)

eq0_new = eq_copy
eq0_ordered = eq0_new
epsilon = sp.Symbol('epsilon')
for A in Ai_:
    eq0_ordered = eq0_ordered.subs(A, epsilon**A1_deg * A)
eq0_ordered = sp.collect(eq0_ordered.subs(Ai[1], epsilon * Ai[1]).lhs, epsilon)
eq0_ordered = sp.Eq(eq0_ordered, 0)

eq_copy = sp.Eq(Ai[0], sp.solve(eq_copy, Ai[0])[0])
eq_copy
```

$\displaystyle A_{0} = - \frac{A_{1}^{2} \alpha_{2}}{2 \omega_{0}^{2}}$

equation 2 tells $A_2$ is also proportional to $A_1$ squared

```python
deglist = list(map(int, sp.degree_list(proj_eqn['eq2'], gens=Ai)))
A1_deg = deglist[1]
Ai_ = tuple(filter(lambda A: A != Ai[1], Ai))

for i in range(A1_deg+1, max(deglist)+1):
    temp = list(combinations_with_replacement(Ai, i))
    for t in temp:
        if t not in removing_combos:
            if t not in keep_terms:
                removing_combos.append(t)

eq_copy = proj_eqn['eq2']
for term in removing_combos:
    eq_copy = eq_copy.subs(np.prod(term), 0)

eq_copy = sp.Eq(eq_copy, 0)

eq2_new = eq_copy
eq2_ordered = eq2_new
for A in Ai_:
    eq2_ordered = eq2_ordered.subs(A, epsilon**A1_deg * A)
eq2_ordered = sp.collect(eq2_ordered.subs(Ai[1], epsilon * Ai[1]).lhs, epsilon)
eq2_ordered = sp.Eq(eq2_ordered, 0)

eq_copy = sp.Eq(Ai[2], sp.solve(eq_copy, Ai[2])[0])
eq_copy
```

$\displaystyle A_{2} = \frac{A_{1} \alpha_{2} \left(A_{1} + 2 A_{3}\right)}{2 \left(4 \omega^{2} - \omega_{0}^{2}\right)}$

equation 3 tells us $A_3$ is also proportional to $A_1^2$

```python
deglist = list(map(int, sp.degree_list(proj_eqn['eq3'], gens=Ai)))
A1_deg = deglist[1]
Ai_ = tuple(filter(lambda A: A != Ai[1], Ai))

for i in range(A1_deg+1, max(deglist)+1):
    temp = list(combinations_with_replacement(Ai, i))
    for t in temp:
        if t not in removing_combos:
            if t not in keep_terms:
                removing_combos.append(t)

eq_copy = proj_eqn['eq3']
for term in removing_combos:
    eq_copy = eq_copy.subs(np.prod(term), 0)

eq_copy = sp.Eq(eq_copy, 0)

eq3_new = eq_copy
eq3_ordered = eq2_new
for A in Ai_:
    eq3_ordered = eq3_ordered.subs(A, epsilon**A1_deg * A)
eq3_ordered = sp.collect(eq3_ordered.subs(Ai[1], epsilon * Ai[1]).lhs, epsilon)
eq3_ordered = sp.Eq(eq3_ordered, 0)

eq_copy = sp.Eq(Ai[3], sp.solve(eq_copy, Ai[3])[0])
eq_copy
```

$\displaystyle A_{3} = \frac{A_{1} \left(A_{1}^{2} \alpha_{3} + 4 A_{2} \alpha_{2}\right)}{4 \left(9 \omega^{2} - \omega_{0}^{2}\right)}$

```python
eq_copy = proj_eqn['eq1']
for term in removing_combos:
    eq_copy = eq_copy.subs(np.prod(term), 0)

eq_copy = sp.Eq(eq_copy, 0)
eq1_new = eq_copy
eq1_ordered = eq1_new
for A in Ai_:
    eq1_ordered = eq1_ordered.subs(A, epsilon**A1_deg * A)
eq1_ordered = sp.collect(eq1_ordered.subs(Ai[1], epsilon * Ai[1]).lhs, epsilon)
eq1_ordered = sp.Eq(eq1_ordered, 0)
```

To facilitate the solution, we order the coefficients using epsilon and then ignore all terms quadratic and higher in $A_0$ and $A_2$ - this means we keep terms up to tand including $O(\epsilon^3)$ -> small amplitude assumption

$A_1$ is ordered at $\epsilon$ and $A_0$ and $A_2$ are order at $\epsilon^2$

```python
# Eq0
eq0_ordered
```

$\displaystyle \epsilon^{2} \left(2 \pi A_{0} \omega_{0}^{2} + \pi A_{1}^{2} \alpha_{2}\right) = 0$

```python
# Eq1
eq1_ordered
```

$\displaystyle \frac{3 \pi A_{1}^{3} \alpha_{3} \epsilon^{3}}{4} + \epsilon^{4} \left(2 \pi A_{0} A_{1} \alpha_{2} + \pi A_{1} A_{2} \alpha_{2}\right) + \epsilon \left(- \pi A_{1} \omega^{2} + \pi A_{1} \omega_{0}^{2}\right) = 0$

```python
# Eq2
eq2_ordered
```

$\displaystyle \pi A_{1} A_{3} \alpha_{2} \epsilon^{3} + \epsilon^{2} \left(\frac{\pi A_{1}^{2} \alpha_{2}}{2} - 4 \pi A_{2} \omega^{2} + \pi A_{2} \omega_{0}^{2}\right) = 0$

```python
# Eq3
eq3_ordered
```

$\displaystyle \frac{\pi A_{1}^{2} \alpha_{2} \epsilon^{2}}{2} + \pi A_{1} A_{3} \alpha_{2} \epsilon^{4} + \epsilon^{3} \left(- 4 \pi A_{2} \omega^{2} + \pi A_{2} \omega_{0}^{2}\right) = 0$

```python
A0_sol = sp.solve(eq0_ordered, Ai[0])[0]
A0_sol
```

$\displaystyle - \frac{A_{1}^{2} \alpha_{2}}{2 \omega_{0}^{2}}$

```python
A2_sol_temp = sp.solve(eq2_ordered, Ai[2])[0]
A2_sol_temp
```

$\displaystyle \frac{A_{1} \alpha_{2} \left(A_{1} + 2 A_{3} \epsilon\right)}{2 \left(4 \omega^{2} - \omega_{0}^{2}\right)}$

```python
eq3_ordered_update = eq3_ordered.subs(Ai[2], A2_sol_temp).expand()
A3_sol = sp.solve(eq3_ordered_update, Ai[3])[0]
A3_sol
```

$\displaystyle \tilde{\infty} A_{1}^{2} \alpha_{2} \epsilon^{2} \left(\epsilon + 1\right)$

#### failed