# Lecture 24 - Sequential Monte Carlo in `PyMC3`

```python
import numpy as np
import pymc3 as pm
import theano as T
from theano import shared, function, tensor as tt
from sample_smc import sample_smc

try:
    import sympy
except:
    _=!pip install sympy
    import sympy

import matplotlib.pyplot as plt
import seaborn as sns
sns.set()

import logging

%matplotlib inline
```

## Objectives

+ Compute the model evidence using `PyMC3`.
+ Do model selection with `PyMC3`.

## Sanity check - Does the calculation of the evidence with PySMC work?

Let
$$
p(\theta) = \mathcal{N}(\theta|0, 1),
$$
and
$$
p(y|\theta) = \mathcal{N}(y|\theta,0).
$$
The posterior of $\theta$ given $y$ is:
$$
p(\theta|y) = \frac{p(y|\theta)p(\theta)}{Z},
$$
where
$$
Z = \int_{-\infty}^{\infty} p(y|\theta)p(\theta)d\theta.
$$
Let's first calculate $Z$ analytically.

```python
import sympy.stats
sympy.init_printing()
y, t = sympy.symbols('y \\theta')
q = 1. / sympy.sqrt(2. * sympy.pi) * sympy.exp(-0.5 * (y - t) ** 2) * \
    1. / sympy.sqrt(2. * sympy.pi) * sympy.exp(-0.5 * t ** 2)
sympy.simplify(sympy.integrate(q, (t, -sympy.oo, sympy.oo)))
```

So, if the observed $y$ was zero, then the Z should be:
$$
Z = \frac{1}{2\sqrt{\pi}}.
$$

```python
Z = 1 /  2. / np.sqrt(np.pi)
print('log Z = {0:.3f}'.format(np.log(Z)))
```

    log Z = -1.266

All, right. Now, let's program this thing in pysmc and compare the results.

We start with the model:

```python
model = pm.Model()
yobs = 0.
with model:
    # prior over theta
    theta = pm.Normal('theta', mu=0., sigma=1.,testval=0.)

    # log likelihood
    llk = pm.Potential('llk', pm.Normal.dist(theta, 1.).logp(yobs))

    trace, smcres = sample_smc(1000)
```

    Sample initial stage: ...
    Stage:   0 Beta: 1.000 Steps:  25 Acce: 1.000

```python
# get the model evidence
log_evidence_smc = np.log(smcres.model.marginal_likelihood)
print('True log evidence: %.4f \nSMC log evidence: %.4f'%(np.log(Z), log_evidence_smc))
```

    True log evidence: -1.2655
    SMC log evidence: -1.2880

Which is close to the truth.

### Questions

+ Repeat the calculations above for a varying number of SMC particles. Start from 10 and go up to 10,000.

## Polynomial Regression

```python
def compute_design_matrix(X, phi):
    """
    Arguments:

    X   -  The observed inputs (1D array)
    phi -  The basis functions.
    """
    num_observations = X.shape[0]
    num_basis = phi.num_basis
    Phi = np.ndarray((num_observations, num_basis))
    for i in range(num_observations):
        Phi[i, :] = phi(X[i, :])
    return Phi

class PolynomialBasis(object):
    """
    A set of linear basis functions.

    Arguments:
    degree  -  The degree of the polynomial.
    """
    def __init__(self, degree):
        self.degree = degree
        self.num_basis = degree + 1
    def __call__(self, x):
        return np.array([x[0] ** i for i in range(self.degree + 1)])

class FourierBasis(object):
    """
    A set of linear basis functions.

    Arguments:
    num_terms  -  The number of Fourier terms.
    L          -  The period of the function.
    """
    def __init__(self, num_terms, L):
        self.num_terms = num_terms
        self.L = L
        self.num_basis = 2 * num_terms
    def __call__(self, x):
        res = np.ndarray((self.num_basis,))
        for i in range(num_terms):
            res[2 * i] = np.cos(2 * i * np.pi / self.L * x[0])
            res[2 * i + 1] = np.sin(2 * (i+1) * np.pi / self.L * x[0])
        return res

class RadialBasisFunctions(object):
    """
    A set of linear basis functions.

    Arguments:
    X   -  The centers of the radial basis functions.
    ell -  The assumed lengthscale.
    """
    def __init__(self, X, ell):
        self.X = X
        self.ell = ell
        self.num_basis = X.shape[0]
    def __call__(self, x):
        return np.exp(-.5 * (x - self.X) ** 2 / self.ell ** 2).flatten()
```

Let's generate some fake data.

```python
np.random.seed(12345)
def getdata(N, sigma2):
    X = 2 * np.random.rand(N) - 1.
    y = 0.5 * X  ** 3 - 0.3 * X ** 2 + np.sqrt(sigma2) * np.random.rand(N)
    return X, y

num_samples = 50
sigma2 = 1e-3
X, y = getdata(num_samples, sigma2)
plt.figure(figsize=(8, 6))
plt.plot(X, y, 'o', markeredgewidth=2, label='Data')
plt.xlabel('$x$', fontsize=20)
plt.ylabel('$y$', fontsize=20)
plt.legend(loc='best', fontsize=20)
plt.tight_layout()
```

We are going to implement a standard Bayesian linear regression and train it with `PyMC3`.
We will compute the evidence in order to select the best class of basis functions.
The model is as follows:

The output $y$ conditioned on the input $x$, the weights of the basis functions $w$ and
the noise variance $\sigma^2$ has likelihood:
$$
p(y|x,w,\sigma, \mathcal{M}) = \mathcal{N}(y|w^T\phi_{\mathcal{M}}(x), \sigma^2),
$$
where $\phi_{\mathcal{M},1}(\cdot), \dots, \phi_{\mathcal{M},m_{\mathcal{M}}}(\cdot)$ are the
$m_{\mathcal{M}}$ basis functions of the model $\mathcal{M}$.
We put a normal prior on the weights:
$$
p(w|\alpha) = \mathcal{N}(w|0, \alpha I_{m_{\mathcal{M}}}),
$$
and an inverse Gamma prior for $\sigma$ and $\alpha$:
$$
p(\sigma^2) = \mathrm{IG}(\sigma^2|1, 1),
$$
and
$$
p(\alpha) = \mathrm{IG}(\alpha|1,1).
$$

Assume that the data we have observed are:
$$
x_{1:n} = \{x_1,\dots,x_n\},\;\mathrm{and}\;y_{1:n} = \{y_1,\dots,y_n\}.
$$
Consider the design matrix $\Phi_{\mathcal{M}}\in\mathbb{R}^{n\times m}$:
$$
\Phi_{\mathcal{M},ij} = \phi_{\mathcal{M},j}(x_i).
$$
The likelihood of the data is:
$$
p(y_{1:n} | x_{1:n}, w, \sigma, \mathcal{M}) = \mathcal{N}(y_{1:n}|\Phi_{\mathcal{M}}w, \sigma^2I_n).
$$
Let's turn this into `PyMC3` code.

```python
def make_model(Phi, y):
    """
    INPUTS:
        Phi -> Design matrix.
        y   -> Target vector.

    RETURNS:
        model -> `pymc3.model` context.
    """
    num_data, num_features = Phi.shape

    # define the model
    with pm.Model() as model:
        # prior on the weights
        alpha = pm.InverseGamma('alpha', alpha=1., beta=1.)
        w = pm.Normal('w', mu=0., tau=alpha, shape=num_features)

        # prior on the likelihood noise variance
        sigma2 = pm.InverseGamma('sigma2', alpha=5., beta=0.1)

        # the data likelihood mean
        ymean = pm.Deterministic('ymean', tt.dot(Phi, w))

        # likelihood
        y = pm.Normal('y', ymean, sigma2, shape=num_data, observed=y)
        #llk = pm.Potential('llk', pm.Normal.dist(ymean, tt.sqrt(sigma2)).logp_sum(y))
    return model
```

Now, let's create a function that trains the model using pysmc for a polynomial basis with a given order.

```python
def fit_poly(phi, X, y, num_particles=100):
    """

    RETURNS:
        1. An instance of pymc3.Model for the SMC model.
        2. The SMC trace.
        3. An instance of pymc3.smc.SMC containing sampling information.
    """
    Phi = compute_design_matrix(X[:, None], phi)
    smcmodel = make_model(Phi, y)
    trace, res = sample_smc(draws=num_particles,
                         model=smcmodel,
                         progressbar=True,
                         threshold=0.8)
    return smcmodel, trace, res

phi = PolynomialBasis(3)
model, trace, res = fit_poly(phi, X, y)
```

    Sample initial stage: ...
    Stage:   0 Beta: 0.000 Steps:  25 Acce: 1.000

      0%|          | 0/100 [00:00<?, ?it/s]
    100%|██████████| 100/100 [00:00<00:00, 933.04it/s][A
    Stage:   1 Beta: 0.000 Steps:  25 Acce: 0.155

    100%|██████████| 100/100 [00:00<00:00, 1061.16it/s]
    Stage:   2 Beta: 0.000 Steps:  25 Acce: 0.159

    100%|██████████| 100/100 [00:00<00:00, 1048.84it/s]
    Stage:   3 Beta: 0.000 Steps:  25 Acce: 0.197

    100%|██████████| 100/100 [00:00<00:00, 1242.81it/s]
    Stage:   4 Beta: 0.001 Steps:  21 Acce: 0.257

    100%|██████████| 100/100 [00:00<00:00, 1725.92it/s]
    Stage:   5 Beta: 0.001 Steps:  15 Acce: 0.289

    100%|██████████| 100/100 [00:00<00:00, 1986.61it/s]
    Stage:   6 Beta: 0.002 Steps:  13 Acce: 0.276

    100%|██████████| 100/100 [00:00<00:00, 1770.82it/s]
    Stage:   7 Beta: 0.004 Steps:  14 Acce: 0.211

    100%|██████████| 100/100 [00:00<00:00, 1358.40it/s]
    Stage:   8 Beta: 0.006 Steps:  19 Acce: 0.228

    100%|██████████| 100/100 [00:00<00:00, 1510.53it/s]
    Stage:   9 Beta: 0.011 Steps:  17 Acce: 0.246

    100%|██████████| 100/100 [00:00<00:00, 1636.34it/s]
    Stage:  10 Beta: 0.017 Steps:  16 Acce: 0.196

    100%|██████████| 100/100 [00:00<00:00, 1238.94it/s]
    Stage:  11 Beta: 0.025 Steps:  21 Acce: 0.250

    100%|██████████| 100/100 [00:00<00:00, 1612.00it/s]
    Stage:  12 Beta: 0.035 Steps:  16 Acce: 0.244

    100%|██████████| 100/100 [00:00<00:00, 1618.49it/s]
    Stage:  13 Beta: 0.050 Steps:  16 Acce: 0.196

    100%|██████████| 100/100 [00:00<00:00, 1229.92it/s]
    Stage:  14 Beta: 0.068 Steps:  21 Acce: 0.285

    100%|██████████| 100/100 [00:00<00:00, 1966.60it/s]
    Stage:  15 Beta: 0.090 Steps:  13 Acce: 0.313

    100%|██████████| 100/100 [00:00<00:00, 2164.50it/s]
    Stage:  16 Beta: 0.116 Steps:  12 Acce: 0.338

    100%|██████████| 100/100 [00:00<00:00, 2285.18it/s]
    Stage:  17 Beta: 0.151 Steps:  11 Acce: 0.289

    100%|██████████| 100/100 [00:00<00:00, 1935.20it/s]
    Stage:  18 Beta: 0.192 Steps:  13 Acce: 0.278

    100%|██████████| 100/100 [00:00<00:00, 1840.10it/s]
    Stage:  19 Beta: 0.240 Steps:  14 Acce: 0.272

    100%|██████████| 100/100 [00:00<00:00, 1826.72it/s]
    Stage:  20 Beta: 0.298 Steps:  14 Acce: 0.264

    100%|██████████| 100/100 [00:00<00:00, 1681.72it/s]
    Stage:  21 Beta: 0.378 Steps:  15 Acce: 0.276

    100%|██████████| 100/100 [00:00<00:00, 1877.26it/s]
    Stage:  22 Beta: 0.479 Steps:  14 Acce: 0.289

    100%|██████████| 100/100 [00:00<00:00, 1988.47it/s]
    Stage:  23 Beta: 0.610 Steps:  13 Acce: 0.272

    100%|██████████| 100/100 [00:00<00:00, 1843.80it/s]
    Stage:  24 Beta: 0.816 Steps:  14 Acce: 0.252

    100%|██████████| 100/100 [00:00<00:00, 1749.14it/s]
    Stage:  25 Beta: 1.000 Steps:  15 Acce: 0.261

    100%|██████████| 100/100 [00:00<00:00, 1668.52it/s]

### Postprocessing

Once you have the `trace` object for the SMC simulation you can apply all the standard postprocessing tools from `PyMC3` as usual.

Here's the posterior distribution over the weights precision and the the likelihood noise, $\alpha$ and $\sigma^2$ respectively:

```python
_=pm.plot_posterior(trace, var_names=['alpha', 'sigma2'])
```

Here's the posterior predictive mean of the output $y$, i.e., $\mathbb{E}[y|x, w, \sigma]$:

```python
ppsamples = pm.sample_posterior_predictive(model=model,
                               trace=trace, var_names=['ymean'])['ymean']
```

    100%|██████████| 100/100 [00:00<00:00, 3874.47it/s]

```python
idx = np.argsort(X)
plt.figure(figsize=(10, 8))
plt.plot(X[idx], ppsamples.mean(0)[idx], linewidth=2.5, label='Posterior Predictive Mean' )
plt.plot(X, y, 'x', markeredgewidth=2.5, markersize=10, label='Observed data')

plt.legend(loc='best', fontsize=20)
plt.tight_layout()
```

SMC does a particle approximation of the posterior distribution. The particles themselves can be obtained from the `trace` object and the particle weights can be obtained from the  `res` object.

Recall that the approximate posterior distribution is of the form $p(\theta|\mathcal{D}) = \sum_{j=1}^{N} w_j \delta(\theta - \theta_j)$.

```python
particles_w = trace.w
particles_alpha = trace.alpha
particle_weights = res.weights   # <- these are the ws from the above equation
```

## Model comparison

Since SMC can approximate the model evidence it provides a principled way of comparing models. Let's compare 5 different polynomial regression models where we change the degree of the polynomial from 1 to 5.

```python
# Evaluate the evidence for the various degrees
log_Zs = []
D = [1, 2, 3, 4, 5]
for d in D:
    phi = PolynomialBasis(d)
    _, _, res = fit_poly(phi, X, y, num_particles=500)
    log_Z = np.log(res.model.marginal_likelihood)
    log_Zs.append(log_Z)
```

    Sample initial stage: ...
    Stage:   0 Beta: 0.000 Steps:  25 Acce: 1.000

      0%|          | 0/500 [00:00<?, ?it/s]
     21%|██▏       | 107/500 [00:00<00:00, 1069.53it/s]
     43%|████▎     | 214/500 [00:00<00:00, 1067.47it/s]
     64%|██████▍   | 319/500 [00:00<00:00, 1061.39it/s]
    100%|██████████| 500/500 [00:00<00:00, 1060.87it/s]
    Stage:   1 Beta: 0.000 Steps:  25 Acce: 0.288

      0%|          | 0/500 [00:00<?, ?it/s]
     41%|████      | 203/500 [00:00<00:00, 2027.15it/s]
    100%|██████████| 500/500 [00:00<00:00, 2003.37it/s]
    Stage:   2 Beta: 0.000 Steps:  13 Acce: 0.259

      0%|          | 0/500 [00:00<?, ?it/s]
     36%|███▌      | 180/500 [00:00<00:00, 1799.16it/s]
    100%|██████████| 500/500 [00:00<00:00, 1763.08it/s]
    Stage:   3 Beta: 0.000 Steps:  15 Acce: 0.271

      0%|          | 0/500 [00:00<?, ?it/s]
     37%|███▋      | 187/500 [00:00<00:00, 1869.01it/s]
    100%|██████████| 500/500 [00:00<00:00, 1860.91it/s]
    Stage:   4 Beta: 0.001 Steps:  14 Acce: 0.256

      0%|          | 0/500 [00:00<?, ?it/s]
     36%|███▌      | 179/500 [00:00<00:00, 1787.86it/s]
    100%|██████████| 500/500 [00:00<00:00, 1784.59it/s]
    Stage:   5 Beta: 0.001 Steps:  15 Acce: 0.259

      0%|          | 0/500 [00:00<?, ?it/s]
     35%|███▌      | 177/500 [00:00<00:00, 1761.36it/s]
    100%|██████████| 500/500 [00:00<00:00, 1753.73it/s]
    Stage:   6 Beta: 0.002 Steps:  15 Acce: 0.223

      0%|          | 0/500 [00:00<?, ?it/s]
     30%|███       | 151/500 [00:00<00:00, 1508.31it/s]
     60%|█████▉    | 298/500 [00:00<00:00, 1493.93it/s]
    100%|██████████| 500/500 [00:00<00:00, 1473.24it/s]
    Stage:   7 Beta: 0.003 Steps:  18 Acce: 0.242

      0%|          | 0/500 [00:00<?, ?it/s]
     34%|███▎      | 168/500 [00:00<00:00, 1676.20it/s]
     67%|██████▋   | 334/500 [00:00<00:00, 1670.55it/s]
    100%|██████████| 500/500 [00:00<00:00, 1644.80it/s]
    Stage:   8 Beta: 0.006 Steps:  16 Acce: 0.228

      0%|          | 0/500 [00:00<?, ?it/s]
     32%|███▏      | 159/500 [00:00<00:00, 1584.49it/s]
     63%|██████▎   | 315/500 [00:00<00:00, 1575.19it/s]
    100%|██████████| 500/500 [00:00<00:00, 1561.86it/s]
    Stage:   9 Beta: 0.009 Steps:  17 Acce: 0.235

      0%|          | 0/500 [00:00<?, ?it/s]
     32%|███▏      | 158/500 [00:00<00:00, 1576.07it/s]
     63%|██████▎   | 314/500 [00:00<00:00, 1570.26it/s]
    100%|██████████| 500/500 [00:00<00:00, 1550.86it/s]
    Stage:  10 Beta: 0.015 Steps:  17 Acce: 0.234

      0%|          | 0/500 [00:00<?, ?it/s]
     31%|███▏      | 157/500 [00:00<00:00, 1567.88it/s]
     62%|██████▏   | 311/500 [00:00<00:00, 1556.50it/s]
    100%|██████████| 500/500 [00:00<00:00, 1535.76it/s]
    Stage:  11 Beta: 0.024 Steps:  17 Acce: 0.246

      0%|          | 0/500 [00:00<?, ?it/s]
     33%|███▎      | 166/500 [00:00<00:00, 1654.25it/s]
     65%|██████▌   | 327/500 [00:00<00:00, 1638.17it/s]
    100%|██████████| 500/500 [00:00<00:00, 1618.74it/s]
    Stage:  12 Beta: 0.038 Steps:  16 Acce: 0.235

      0%|          | 0/500 [00:00<?, ?it/s]
     31%|███       | 154/500 [00:00<00:00, 1537.05it/s]
     62%|██████▏   | 309/500 [00:00<00:00, 1539.67it/s]
    100%|██████████| 500/500 [00:00<00:00, 1536.17it/s]
    Stage:  13 Beta: 0.061 Steps:  17 Acce: 0.260

      0%|          | 0/500 [00:00<?, ?it/s]
     35%|███▌      | 177/500 [00:00<00:00, 1762.24it/s]
    100%|██████████| 500/500 [00:00<00:00, 1745.45it/s]
    Stage:  14 Beta: 0.095 Steps:  15 Acce: 0.237

      0%|          | 0/500 [00:00<?, ?it/s]
     31%|███       | 154/500 [00:00<00:00, 1538.01it/s]
     60%|██████    | 302/500 [00:00<00:00, 1518.98it/s]
    100%|██████████| 500/500 [00:00<00:00, 1519.69it/s]
    Stage:  15 Beta: 0.149 Steps:  17 Acce: 0.251

      0%|          | 0/500 [00:00<?, ?it/s]
     35%|███▌      | 176/500 [00:00<00:00, 1751.92it/s]
    100%|██████████| 500/500 [00:00<00:00, 1734.64it/s]
    Stage:  16 Beta: 0.239 Steps:  15 Acce: 0.243

      0%|          | 0/500 [00:00<?, ?it/s]
     33%|███▎      | 167/500 [00:00<00:00, 1661.72it/s]
     66%|██████▌   | 330/500 [00:00<00:00, 1650.53it/s]
    100%|██████████| 500/500 [00:00<00:00, 1639.98it/s]
    Stage:  17 Beta: 0.381 Steps:  16 Acce: 0.258

      0%|          | 0/500 [00:00<?, ?it/s]
     35%|███▌      | 176/500 [00:00<00:00, 1756.82it/s]
    100%|██████████| 500/500 [00:00<00:00, 1755.52it/s]
    Stage:  18 Beta: 0.600 Steps:  15 Acce: 0.246

      0%|          | 0/500 [00:00<?, ?it/s]
     33%|███▎      | 167/500 [00:00<00:00, 1661.42it/s]
     66%|██████▌   | 328/500 [00:00<00:00, 1643.73it/s]
    100%|██████████| 500/500 [00:00<00:00, 1624.14it/s]
    Stage:  19 Beta: 0.947 Steps:  16 Acce: 0.291

      0%|          | 0/500 [00:00<?, ?it/s]
     40%|████      | 201/500 [00:00<00:00, 2009.38it/s]
    100%|██████████| 500/500 [00:00<00:00, 1991.53it/s]
    Stage:  20 Beta: 1.000 Steps:  13 Acce: 0.230

      0%|          | 0/500 [00:00<?, ?it/s]
     31%|███       | 156/500 [00:00<00:00, 1552.94it/s]
     61%|██████    | 306/500 [00:00<00:00, 1535.65it/s]
    100%|██████████| 500/500 [00:00<00:00, 1505.63it/s]
    Sample initial stage: ...
    Stage:   0 Beta: 0.000 Steps:  25 Acce: 1.000

      0%|          | 0/500 [00:00<?, ?it/s]
     22%|██▏       | 111/500 [00:00<00:00, 1102.46it/s]
     44%|████▍     | 219/500 [00:00<00:00, 1094.13it/s]
     65%|██████▌   | 327/500 [00:00<00:00, 1087.08it/s]
    100%|██████████| 500/500 [00:00<00:00, 1077.47it/s]
    Stage:   1 Beta: 0.000 Steps:  25 Acce: 0.222

      0%|          | 0/500 [00:00<?, ?it/s]
     30%|███       | 151/500 [00:00<00:00, 1501.80it/s]
     59%|█████▉    | 296/500 [00:00<00:00, 1482.92it/s]
    100%|██████████| 500/500 [00:00<00:00, 1464.64it/s]
    Stage:   2 Beta: 0.000 Steps:  18 Acce: 0.250

      0%|          | 0/500 [00:00<?, ?it/s]
     34%|███▍      | 170/500 [00:00<00:00, 1691.09it/s]
    100%|██████████| 500/500 [00:00<00:00, 1677.64it/s]
    Stage:   3 Beta: 0.000 Steps:  16 Acce: 0.237

      0%|          | 0/500 [00:00<?, ?it/s]
     32%|███▏      | 159/500 [00:00<00:00, 1586.73it/s]
     63%|██████▎   | 314/500 [00:00<00:00, 1573.64it/s]
    100%|██████████| 500/500 [00:00<00:00, 1559.97it/s]
    Stage:   4 Beta: 0.000 Steps:  17 Acce: 0.224

      0%|          | 0/500 [00:00<?, ?it/s]
     30%|███       | 150/500 [00:00<00:00, 1491.62it/s]
     59%|█████▉    | 297/500 [00:00<00:00, 1478.53it/s]
    100%|██████████| 500/500 [00:00<00:00, 1462.24it/s]
    Stage:   5 Beta: 0.001 Steps:  18 Acce: 0.225

      0%|          | 0/500 [00:00<?, ?it/s]
     30%|██▉       | 149/500 [00:00<00:00, 1482.93it/s]
     60%|█████▉    | 299/500 [00:00<00:00, 1485.79it/s]
    100%|██████████| 500/500 [00:00<00:00, 1482.70it/s]
    Stage:   6 Beta: 0.002 Steps:  18 Acce: 0.208

      0%|          | 0/500 [00:00<?, ?it/s]
     28%|██▊       | 142/500 [00:00<00:00, 1413.95it/s]
     56%|█████▋    | 282/500 [00:00<00:00, 1407.08it/s]
    100%|██████████| 500/500 [00:00<00:00, 1399.51it/s]
    Stage:   7 Beta: 0.003 Steps:  19 Acce: 0.201

      0%|          | 0/500 [00:00<?, ?it/s]
     26%|██▋       | 132/500 [00:00<00:00, 1315.07it/s]
     54%|█████▎    | 268/500 [00:00<00:00, 1325.73it/s]
    100%|██████████| 500/500 [00:00<00:00, 1335.19it/s]
    Stage:   8 Beta: 0.005 Steps:  20 Acce: 0.222

      0%|          | 0/500 [00:00<?, ?it/s]
     30%|██▉       | 149/500 [00:00<00:00, 1482.87it/s]
     60%|█████▉    | 298/500 [00:00<00:00, 1484.76it/s]
    100%|██████████| 500/500 [00:00<00:00, 1484.78it/s]
    Stage:   9 Beta: 0.008 Steps:  18 Acce: 0.233

      0%|          | 0/500 [00:00<?, ?it/s]
     32%|███▏      | 159/500 [00:00<00:00, 1588.49it/s]
     63%|██████▎   | 315/500 [00:00<00:00, 1576.88it/s]
    100%|██████████| 500/500 [00:00<00:00, 1554.79it/s]
    Stage:  10 Beta: 0.013 Steps:  17 Acce: 0.241

      0%|          | 0/500 [00:00<?, ?it/s]
     34%|███▎      | 168/500 [00:00<00:00, 1674.23it/s]
    100%|██████████| 500/500 [00:00<00:00, 1667.77it/s]
    Stage:  11 Beta: 0.021 Steps:  16 Acce: 0.243

      0%|          | 0/500 [00:00<?, ?it/s]
     34%|███▍      | 170/500 [00:00<00:00, 1694.05it/s]
    100%|██████████| 500/500 [00:00<00:00, 1670.03it/s]
    Stage:  12 Beta: 0.033 Steps:  16 Acce: 0.262

      0%|          | 0/500 [00:00<?, ?it/s]
     36%|███▌      | 179/500 [00:00<00:00, 1785.59it/s]
    100%|██████████| 500/500 [00:00<00:00, 1778.25it/s]
    Stage:  13 Beta: 0.052 Steps:  15 Acce: 0.257

      0%|          | 0/500 [00:00<?, ?it/s]
     36%|███▌      | 181/500 [00:00<00:00, 1802.70it/s]
    100%|██████████| 500/500 [00:00<00:00, 1776.71it/s]
    Stage:  14 Beta: 0.079 Steps:  15 Acce: 0.243

      0%|          | 0/500 [00:00<?, ?it/s]
     34%|███▍      | 169/500 [00:00<00:00, 1682.21it/s]
    100%|██████████| 500/500 [00:00<00:00, 1667.50it/s]
    Stage:  15 Beta: 0.120 Steps:  16 Acce: 0.275

      0%|          | 0/500 [00:00<?, ?it/s]
     38%|███▊      | 190/500 [00:00<00:00, 1893.38it/s]
    100%|██████████| 500/500 [00:00<00:00, 1881.68it/s]
    Stage:  16 Beta: 0.181 Steps:  14 Acce: 0.264

      0%|          | 0/500 [00:00<?, ?it/s]
     36%|███▌      | 178/500 [00:00<00:00, 1777.73it/s]
    100%|██████████| 500/500 [00:00<00:00, 1775.35it/s]
    Stage:  17 Beta: 0.264 Steps:  15 Acce: 0.276

      0%|          | 0/500 [00:00<?, ?it/s]
     39%|███▊      | 193/500 [00:00<00:00, 1927.86it/s]
    100%|██████████| 500/500 [00:00<00:00, 1895.63it/s]
    Stage:  18 Beta: 0.391 Steps:  14 Acce: 0.266

      0%|          | 0/500 [00:00<?, ?it/s]
     38%|███▊      | 190/500 [00:00<00:00, 1894.60it/s]
    100%|██████████| 500/500 [00:00<00:00, 1893.48it/s]
    Stage:  19 Beta: 0.597 Steps:  14 Acce: 0.259

      0%|          | 0/500 [00:00<?, ?it/s]
     36%|███▋      | 182/500 [00:00<00:00, 1814.00it/s]
    100%|██████████| 500/500 [00:00<00:00, 1784.45it/s]
    Stage:  20 Beta: 0.861 Steps:  15 Acce: 0.260

      0%|          | 0/500 [00:00<?, ?it/s]
     36%|███▌      | 178/500 [00:00<00:00, 1772.80it/s]
    100%|██████████| 500/500 [00:00<00:00, 1772.45it/s]
    Stage:  21 Beta: 1.000 Steps:  15 Acce: 0.269

      0%|          | 0/500 [00:00<?, ?it/s]
     37%|███▋      | 185/500 [00:00<00:00, 1848.93it/s]
    100%|██████████| 500/500 [00:00<00:00, 1843.60it/s]
    Sample initial stage: ...
    Stage:   0 Beta: 0.000 Steps:  25 Acce: 1.000

      0%|          | 0/500 [00:00<?, ?it/s]
     21%|██        | 106/500 [00:00<00:00, 1054.14it/s]
     42%|████▏     | 212/500 [00:00<00:00, 1054.90it/s]
     64%|██████▍   | 320/500 [00:00<00:00, 1059.56it/s]
    100%|██████████| 500/500 [00:00<00:00, 1063.03it/s]
    Stage:   1 Beta: 0.000 Steps:  25 Acce: 0.179

      0%|          | 0/500 [00:00<?, ?it/s]
     23%|██▎       | 116/500 [00:00<00:00, 1156.89it/s]
     46%|████▋     | 232/500 [00:00<00:00, 1157.22it/s]
     70%|██████▉   | 348/500 [00:00<00:00, 1155.89it/s]
    100%|██████████| 500/500 [00:00<00:00, 1150.37it/s]
    Stage:   2 Beta: 0.000 Steps:  23 Acce: 0.183

      0%|          | 0/500 [00:00<?, ?it/s]
     24%|██▍       | 121/500 [00:00<00:00, 1203.27it/s]
     48%|████▊     | 239/500 [00:00<00:00, 1193.19it/s]
     72%|███████▏  | 359/500 [00:00<00:00, 1194.53it/s]
    100%|██████████| 500/500 [00:00<00:00, 1187.77it/s]
    Stage:   3 Beta: 0.000 Steps:  22 Acce: 0.229

      0%|          | 0/500 [00:00<?, ?it/s]
     31%|███▏      | 157/500 [00:00<00:00, 1569.56it/s]
     62%|██████▏   | 310/500 [00:00<00:00, 1554.74it/s]
    100%|██████████| 500/500 [00:00<00:00, 1535.60it/s]
    Stage:   4 Beta: 0.000 Steps:  17 Acce: 0.205

      0%|          | 0/500 [00:00<?, ?it/s]
     27%|██▋       | 134/500 [00:00<00:00, 1331.19it/s]
     53%|█████▎    | 263/500 [00:00<00:00, 1317.43it/s]
    100%|██████████| 500/500 [00:00<00:00, 1313.46it/s]
    Stage:   5 Beta: 0.001 Steps:  20 Acce: 0.220

      0%|          | 0/500 [00:00<?, ?it/s]
     29%|██▉       | 145/500 [00:00<00:00, 1448.67it/s]
     57%|█████▋    | 287/500 [00:00<00:00, 1439.91it/s]
    100%|██████████| 500/500 [00:00<00:00, 1442.99it/s]
    Stage:   6 Beta: 0.002 Steps:  18 Acce: 0.209

      0%|          | 0/500 [00:00<?, ?it/s]
     28%|██▊       | 139/500 [00:00<00:00, 1384.92it/s]
     55%|█████▌    | 275/500 [00:00<00:00, 1375.67it/s]
    100%|██████████| 500/500 [00:00<00:00, 1364.39it/s]
    Stage:   7 Beta: 0.003 Steps:  19 Acce: 0.213

      0%|          | 0/500 [00:00<?, ?it/s]
     27%|██▋       | 136/500 [00:00<00:00, 1355.78it/s]
     55%|█████▍    | 274/500 [00:00<00:00, 1361.41it/s]
    100%|██████████| 500/500 [00:00<00:00, 1357.11it/s]
    Stage:   8 Beta: 0.004 Steps:  19 Acce: 0.244

      0%|          | 0/500 [00:00<?, ?it/s]
     32%|███▏      | 159/500 [00:00<00:00, 1587.49it/s]
     64%|██████▍   | 321/500 [00:00<00:00, 1594.92it/s]
    100%|██████████| 500/500 [00:00<00:00, 1603.08it/s]
    Stage:   9 Beta: 0.008 Steps:  16 Acce: 0.255

      0%|          | 0/500 [00:00<?, ?it/s]
     35%|███▍      | 174/500 [00:00<00:00, 1733.17it/s]
    100%|██████████| 500/500 [00:00<00:00, 1730.94it/s]
    Stage:  10 Beta: 0.012 Steps:  15 Acce: 0.250

      0%|          | 0/500 [00:00<?, ?it/s]
     35%|███▌      | 176/500 [00:00<00:00, 1754.75it/s]
    100%|██████████| 500/500 [00:00<00:00, 1729.70it/s]
    Stage:  11 Beta: 0.019 Steps:  15 Acce: 0.248

      0%|          | 0/500 [00:00<?, ?it/s]
     33%|███▎      | 164/500 [00:00<00:00, 1634.20it/s]
     66%|██████▌   | 329/500 [00:00<00:00, 1638.88it/s]
    100%|██████████| 500/500 [00:00<00:00, 1621.96it/s]
    Stage:  12 Beta: 0.027 Steps:  16 Acce: 0.228

      0%|          | 0/500 [00:00<?, ?it/s]
     31%|███       | 154/500 [00:00<00:00, 1532.29it/s]
     61%|██████    | 306/500 [00:00<00:00, 1525.70it/s]
    100%|██████████| 500/500 [00:00<00:00, 1520.49it/s]
    Stage:  13 Beta: 0.037 Steps:  17 Acce: 0.260

      0%|          | 0/500 [00:00<?, ?it/s]
     35%|███▌      | 176/500 [00:00<00:00, 1753.18it/s]
    100%|██████████| 500/500 [00:00<00:00, 1744.77it/s]
    Stage:  14 Beta: 0.049 Steps:  15 Acce: 0.269

      0%|          | 0/500 [00:00<?, ?it/s]
     38%|███▊      | 188/500 [00:00<00:00, 1873.02it/s]
    100%|██████████| 500/500 [00:00<00:00, 1856.05it/s]
    Stage:  15 Beta: 0.064 Steps:  14 Acce: 0.259

      0%|          | 0/500 [00:00<?, ?it/s]
     35%|███▌      | 176/500 [00:00<00:00, 1750.41it/s]
    100%|██████████| 500/500 [00:00<00:00, 1744.21it/s]
    Stage:  16 Beta: 0.082 Steps:  15 Acce: 0.301

      0%|          | 0/500 [00:00<?, ?it/s]
     44%|████▍     | 222/500 [00:00<00:00, 2216.88it/s]
    100%|██████████| 500/500 [00:00<00:00, 2160.81it/s]
    Stage:  17 Beta: 0.104 Steps:  12 Acce: 0.325

      0%|          | 0/500 [00:00<?, ?it/s]
     48%|████▊     | 238/500 [00:00<00:00, 2379.86it/s]
    100%|██████████| 500/500 [00:00<00:00, 2340.54it/s]
    Stage:  18 Beta: 0.130 Steps:  11 Acce: 0.306

      0%|          | 0/500 [00:00<?, ?it/s]
     44%|████▎     | 218/500 [00:00<00:00, 2169.89it/s]
    100%|██████████| 500/500 [00:00<00:00, 2133.08it/s]
    Stage:  19 Beta: 0.166 Steps:  12 Acce: 0.339

      0%|          | 0/500 [00:00<?, ?it/s]
     47%|████▋     | 237/500 [00:00<00:00, 2368.87it/s]
    100%|██████████| 500/500 [00:00<00:00, 2349.21it/s]
    Stage:  20 Beta: 0.209 Steps:  11 Acce: 0.346

      0%|          | 0/500 [00:00<?, ?it/s]
    100%|██████████| 500/500 [00:00<00:00, 2540.49it/s]
    Stage:  21 Beta: 0.264 Steps:  10 Acce: 0.299

      0%|          | 0/500 [00:00<?, ?it/s]
     44%|████▎     | 218/500 [00:00<00:00, 2172.21it/s]
    100%|██████████| 500/500 [00:00<00:00, 2135.84it/s]
    Stage:  22 Beta: 0.337 Steps:  12 Acce: 0.277

      0%|          | 0/500 [00:00<?, ?it/s]
     37%|███▋      | 186/500 [00:00<00:00, 1851.76it/s]
    100%|██████████| 500/500 [00:00<00:00, 1859.78it/s]
    Stage:  23 Beta: 0.433 Steps:  14 Acce: 0.283

      0%|          | 0/500 [00:00<?, ?it/s]
     40%|███▉      | 199/500 [00:00<00:00, 1983.40it/s]
    100%|██████████| 500/500 [00:00<00:00, 1983.67it/s]
    Stage:  24 Beta: 0.556 Steps:  13 Acce: 0.271

      0%|          | 0/500 [00:00<?, ?it/s]
     37%|███▋      | 186/500 [00:00<00:00, 1857.38it/s]
    100%|██████████| 500/500 [00:00<00:00, 1838.67it/s]
    Stage:  25 Beta: 0.716 Steps:  14 Acce: 0.246

      0%|          | 0/500 [00:00<?, ?it/s]
     33%|███▎      | 165/500 [00:00<00:00, 1646.21it/s]
     65%|██████▌   | 327/500 [00:00<00:00, 1638.23it/s]
    100%|██████████| 500/500 [00:00<00:00, 1622.28it/s]
    Stage:  26 Beta: 0.938 Steps:  16 Acce: 0.251

      0%|          | 0/500 [00:00<?, ?it/s]
     34%|███▍      | 171/500 [00:00<00:00, 1700.48it/s]
    100%|██████████| 500/500 [00:00<00:00, 1721.75it/s]
    Stage:  27 Beta: 1.000 Steps:  15 Acce: 0.245

      0%|          | 0/500 [00:00<?, ?it/s]
     32%|███▏      | 161/500 [00:00<00:00, 1609.55it/s]
     64%|██████▍   | 319/500 [00:00<00:00, 1598.80it/s]
    100%|██████████| 500/500 [00:00<00:00, 1577.64it/s]
    Sample initial stage: ...
    Stage:   0 Beta: 0.000 Steps:  25 Acce: 1.000

      0%|          | 0/500 [00:00<?, ?it/s]
     22%|██▏       | 108/500 [00:00<00:00, 1072.85it/s]
     42%|████▏     | 212/500 [00:00<00:00, 1060.37it/s]
     63%|██████▎   | 317/500 [00:00<00:00, 1056.47it/s]
    100%|██████████| 500/500 [00:00<00:00, 1054.29it/s]
    Stage:   1 Beta: 0.000 Steps:  25 Acce: 0.172

      0%|          | 0/500 [00:00<?, ?it/s]
     22%|██▏       | 112/500 [00:00<00:00, 1118.55it/s]
     44%|████▍     | 222/500 [00:00<00:00, 1111.77it/s]
     67%|██████▋   | 333/500 [00:00<00:00, 1111.16it/s]
    100%|██████████| 500/500 [00:00<00:00, 1102.35it/s]
    Stage:   2 Beta: 0.000 Steps:  24 Acce: 0.200

      0%|          | 0/500 [00:00<?, ?it/s]
     26%|██▋       | 132/500 [00:00<00:00, 1318.88it/s]
     52%|█████▏    | 262/500 [00:00<00:00, 1310.00it/s]
    100%|██████████| 500/500 [00:00<00:00, 1310.18it/s]
    Stage:   3 Beta: 0.000 Steps:  20 Acce: 0.225

      0%|          | 0/500 [00:00<?, ?it/s]
     29%|██▉       | 145/500 [00:00<00:00, 1449.41it/s]
     58%|█████▊    | 288/500 [00:00<00:00, 1440.46it/s]
    100%|██████████| 500/500 [00:00<00:00, 1432.79it/s]
    Stage:   4 Beta: 0.000 Steps:  18 Acce: 0.237

      0%|          | 0/500 [00:00<?, ?it/s]
     32%|███▏      | 162/500 [00:00<00:00, 1619.89it/s]
     64%|██████▍   | 322/500 [00:00<00:00, 1611.07it/s]
    100%|██████████| 500/500 [00:00<00:00, 1602.81it/s]
    Stage:   5 Beta: 0.001 Steps:  16 Acce: 0.239

      0%|          | 0/500 [00:00<?, ?it/s]
     33%|███▎      | 163/500 [00:00<00:00, 1629.53it/s]
     65%|██████▍   | 323/500 [00:00<00:00, 1619.02it/s]
    100%|██████████| 500/500 [00:00<00:00, 1607.04it/s]
    Stage:   6 Beta: 0.001 Steps:  16 Acce: 0.261

      0%|          | 0/500 [00:00<?, ?it/s]
     34%|███▍      | 172/500 [00:00<00:00, 1715.83it/s]
    100%|██████████| 500/500 [00:00<00:00, 1730.58it/s]
    Stage:   7 Beta: 0.002 Steps:  15 Acce: 0.272

      0%|          | 0/500 [00:00<?, ?it/s]
     37%|███▋      | 185/500 [00:00<00:00, 1847.38it/s]
    100%|██████████| 500/500 [00:00<00:00, 1831.29it/s]
    Stage:   8 Beta: 0.004 Steps:  14 Acce: 0.235

      0%|          | 0/500 [00:00<?, ?it/s]
     31%|███       | 155/500 [00:00<00:00, 1547.42it/s]
     62%|██████▏   | 308/500 [00:00<00:00, 1542.07it/s]
    100%|██████████| 500/500 [00:00<00:00, 1522.64it/s]
    Stage:   9 Beta: 0.007 Steps:  17 Acce: 0.234

      0%|          | 0/500 [00:00<?, ?it/s]
     31%|███       | 155/500 [00:00<00:00, 1548.53it/s]
     62%|██████▏   | 310/500 [00:00<00:00, 1546.09it/s]
    100%|██████████| 500/500 [00:00<00:00, 1531.79it/s]
    Stage:  10 Beta: 0.011 Steps:  17 Acce: 0.248

      0%|          | 0/500 [00:00<?, ?it/s]
     33%|███▎      | 163/500 [00:00<00:00, 1624.51it/s]
     65%|██████▌   | 325/500 [00:00<00:00, 1622.72it/s]
    100%|██████████| 500/500 [00:00<00:00, 1618.81it/s]
    Stage:  11 Beta: 0.017 Steps:  16 Acce: 0.226

      0%|          | 0/500 [00:00<?, ?it/s]
     31%|███       | 154/500 [00:00<00:00, 1536.28it/s]
     61%|██████    | 303/500 [00:00<00:00, 1521.00it/s]
    100%|██████████| 500/500 [00:00<00:00, 1511.12it/s]
    Stage:  12 Beta: 0.026 Steps:  17 Acce: 0.247

      0%|          | 0/500 [00:00<?, ?it/s]
     33%|███▎      | 164/500 [00:00<00:00, 1635.40it/s]
     65%|██████▌   | 325/500 [00:00<00:00, 1626.33it/s]
    100%|██████████| 500/500 [00:00<00:00, 1612.26it/s]
    Stage:  13 Beta: 0.037 Steps:  16 Acce: 0.277

      0%|          | 0/500 [00:00<?, ?it/s]
     38%|███▊      | 189/500 [00:00<00:00, 1885.08it/s]
    100%|██████████| 500/500 [00:00<00:00, 1849.35it/s]
    Stage:  14 Beta: 0.051 Steps:  14 Acce: 0.272

      0%|          | 0/500 [00:00<?, ?it/s]
     37%|███▋      | 185/500 [00:00<00:00, 1848.78it/s]
    100%|██████████| 500/500 [00:00<00:00, 1839.42it/s]
    Stage:  15 Beta: 0.066 Steps:  14 Acce: 0.262

      0%|          | 0/500 [00:00<?, ?it/s]
     31%|███       | 156/500 [00:00<00:00, 1552.66it/s]
     65%|██████▌   | 326/500 [00:00<00:00, 1593.78it/s]
    100%|██████████| 500/500 [00:00<00:00, 1647.52it/s]
    Stage:  16 Beta: 0.084 Steps:  15 Acce: 0.266

      0%|          | 0/500 [00:00<?, ?it/s]
     38%|███▊      | 191/500 [00:00<00:00, 1908.80it/s]
    100%|██████████| 500/500 [00:00<00:00, 1870.24it/s]
    Stage:  17 Beta: 0.105 Steps:  14 Acce: 0.279

      0%|          | 0/500 [00:00<?, ?it/s]
     37%|███▋      | 184/500 [00:00<00:00, 1839.96it/s]
    100%|██████████| 500/500 [00:00<00:00, 1839.83it/s]
    Stage:  18 Beta: 0.130 Steps:  14 Acce: 0.292

      0%|          | 0/500 [00:00<?, ?it/s]
     40%|████      | 200/500 [00:00<00:00, 1994.91it/s]
    100%|██████████| 500/500 [00:00<00:00, 1979.62it/s]
    Stage:  19 Beta: 0.159 Steps:  13 Acce: 0.289

      0%|          | 0/500 [00:00<?, ?it/s]
     41%|████      | 203/500 [00:00<00:00, 2024.84it/s]
    100%|██████████| 500/500 [00:00<00:00, 1999.22it/s]
    Stage:  20 Beta: 0.198 Steps:  13 Acce: 0.307

      0%|          | 0/500 [00:00<?, ?it/s]
     44%|████▍     | 219/500 [00:00<00:00, 2189.91it/s]
    100%|██████████| 500/500 [00:00<00:00, 2167.67it/s]
    Stage:  21 Beta: 0.244 Steps:  12 Acce: 0.317

      0%|          | 0/500 [00:00<?, ?it/s]
     43%|████▎     | 216/500 [00:00<00:00, 2152.30it/s]
    100%|██████████| 500/500 [00:00<00:00, 2141.15it/s]
    Stage:  22 Beta: 0.301 Steps:  12 Acce: 0.306

      0%|          | 0/500 [00:00<?, ?it/s]
     44%|████▎     | 218/500 [00:00<00:00, 2172.72it/s]
    100%|██████████| 500/500 [00:00<00:00, 2155.81it/s]
    Stage:  23 Beta: 0.379 Steps:  12 Acce: 0.285

      0%|          | 0/500 [00:00<?, ?it/s]
     39%|███▉      | 197/500 [00:00<00:00, 1962.11it/s]
    100%|██████████| 500/500 [00:00<00:00, 1962.52it/s]
    Stage:  24 Beta: 0.473 Steps:  13 Acce: 0.299

      0%|          | 0/500 [00:00<?, ?it/s]
     44%|████▍     | 221/500 [00:00<00:00, 2204.88it/s]
    100%|██████████| 500/500 [00:00<00:00, 2136.37it/s]
    Stage:  25 Beta: 0.609 Steps:  12 Acce: 0.262

      0%|          | 0/500 [00:00<?, ?it/s]
     35%|███▌      | 175/500 [00:00<00:00, 1742.79it/s]
    100%|██████████| 500/500 [00:00<00:00, 1725.43it/s]
    Stage:  26 Beta: 0.774 Steps:  15 Acce: 0.260

      0%|          | 0/500 [00:00<?, ?it/s]
     35%|███▍      | 174/500 [00:00<00:00, 1737.90it/s]
    100%|██████████| 500/500 [00:00<00:00, 1744.96it/s]
    Stage:  27 Beta: 1.000 Steps:  15 Acce: 0.255

      0%|          | 0/500 [00:00<?, ?it/s]
     34%|███▍      | 170/500 [00:00<00:00, 1696.58it/s]
    100%|██████████| 500/500 [00:00<00:00, 1673.65it/s]
    Sample initial stage: ...
    Stage:   0 Beta: 0.000 Steps:  25 Acce: 1.000

      0%|          | 0/500 [00:00<?, ?it/s]
     22%|██▏       | 109/500 [00:00<00:00, 1086.97it/s]
     43%|████▎     | 217/500 [00:00<00:00, 1083.80it/s]
     65%|██████▌   | 325/500 [00:00<00:00, 1081.61it/s]
    100%|██████████| 500/500 [00:00<00:00, 1071.84it/s]
    Stage:   1 Beta: 0.000 Steps:  25 Acce: 0.185

      0%|          | 0/500 [00:00<?, ?it/s]
     25%|██▌       | 125/500 [00:00<00:00, 1242.61it/s]
     49%|████▉     | 246/500 [00:00<00:00, 1232.08it/s]
     74%|███████▎  | 368/500 [00:00<00:00, 1226.11it/s]
    100%|██████████| 500/500 [00:00<00:00, 1214.16it/s]
    Stage:   2 Beta: 0.000 Steps:  22 Acce: 0.199

      0%|          | 0/500 [00:00<?, ?it/s]
     27%|██▋       | 133/500 [00:00<00:00, 1327.45it/s]
     53%|█████▎    | 266/500 [00:00<00:00, 1325.57it/s]
    100%|██████████| 500/500 [00:00<00:00, 1321.80it/s]
    Stage:   3 Beta: 0.000 Steps:  20 Acce: 0.202

      0%|          | 0/500 [00:00<?, ?it/s]
     27%|██▋       | 134/500 [00:00<00:00, 1338.78it/s]
     53%|█████▎    | 267/500 [00:00<00:00, 1333.01it/s]
    100%|██████████| 500/500 [00:00<00:00, 1330.21it/s]
    Stage:   4 Beta: 0.000 Steps:  20 Acce: 0.227

      0%|          | 0/500 [00:00<?, ?it/s]
     31%|███▏      | 157/500 [00:00<00:00, 1569.09it/s]
     63%|██████▎   | 316/500 [00:00<00:00, 1572.47it/s]
    100%|██████████| 500/500 [00:00<00:00, 1561.78it/s]
    Stage:   5 Beta: 0.001 Steps:  17 Acce: 0.248

      0%|          | 0/500 [00:00<?, ?it/s]
     34%|███▍      | 169/500 [00:00<00:00, 1681.45it/s]
     67%|██████▋   | 333/500 [00:00<00:00, 1666.36it/s]
    100%|██████████| 500/500 [00:00<00:00, 1641.69it/s]
    Stage:   6 Beta: 0.001 Steps:  16 Acce: 0.252

      0%|          | 0/500 [00:00<?, ?it/s]
     36%|███▌      | 178/500 [00:00<00:00, 1775.74it/s]
    100%|██████████| 500/500 [00:00<00:00, 1777.11it/s]
    Stage:   7 Beta: 0.002 Steps:  15 Acce: 0.272

      0%|          | 0/500 [00:00<?, ?it/s]
     38%|███▊      | 190/500 [00:00<00:00, 1898.56it/s]
    100%|██████████| 500/500 [00:00<00:00, 1889.00it/s]
    Stage:   8 Beta: 0.004 Steps:  14 Acce: 0.256

      0%|          | 0/500 [00:00<?, ?it/s]
     36%|███▌      | 181/500 [00:00<00:00, 1806.21it/s]
    100%|██████████| 500/500 [00:00<00:00, 1776.93it/s]
    Stage:   9 Beta: 0.006 Steps:  15 Acce: 0.254

      0%|          | 0/500 [00:00<?, ?it/s]
     36%|███▌      | 179/500 [00:00<00:00, 1789.59it/s]
    100%|██████████| 500/500 [00:00<00:00, 1762.80it/s]
    Stage:  10 Beta: 0.010 Steps:  15 Acce: 0.245

      0%|          | 0/500 [00:00<?, ?it/s]
     33%|███▎      | 166/500 [00:00<00:00, 1656.20it/s]
     67%|██████▋   | 333/500 [00:00<00:00, 1657.09it/s]
    100%|██████████| 500/500 [00:00<00:00, 1653.33it/s]
    Stage:  11 Beta: 0.016 Steps:  16 Acce: 0.232

      0%|          | 0/500 [00:00<?, ?it/s]
     31%|███       | 155/500 [00:00<00:00, 1547.77it/s]
     62%|██████▏   | 312/500 [00:00<00:00, 1552.81it/s]
    100%|██████████| 500/500 [00:00<00:00, 1543.57it/s]
    Stage:  12 Beta: 0.025 Steps:  17 Acce: 0.226

      0%|          | 0/500 [00:00<?, ?it/s]
     31%|███▏      | 157/500 [00:00<00:00, 1561.83it/s]
     63%|██████▎   | 313/500 [00:00<00:00, 1558.66it/s]
    100%|██████████| 500/500 [00:00<00:00, 1547.65it/s]
    Stage:  13 Beta: 0.037 Steps:  17 Acce: 0.257

      0%|          | 0/500 [00:00<?, ?it/s]
     35%|███▌      | 177/500 [00:00<00:00, 1765.20it/s]
    100%|██████████| 500/500 [00:00<00:00, 1759.77it/s]
    Stage:  14 Beta: 0.050 Steps:  15 Acce: 0.247

      0%|          | 0/500 [00:00<?, ?it/s]
     33%|███▎      | 166/500 [00:00<00:00, 1656.54it/s]
     66%|██████▌   | 330/500 [00:00<00:00, 1650.65it/s]
    100%|██████████| 500/500 [00:00<00:00, 1647.94it/s]
    Stage:  15 Beta: 0.065 Steps:  16 Acce: 0.268

      0%|          | 0/500 [00:00<?, ?it/s]
     38%|███▊      | 188/500 [00:00<00:00, 1876.47it/s]
    100%|██████████| 500/500 [00:00<00:00, 1866.43it/s]
    Stage:  16 Beta: 0.081 Steps:  14 Acce: 0.243

      0%|          | 0/500 [00:00<?, ?it/s]
     34%|███▎      | 168/500 [00:00<00:00, 1671.31it/s]
     66%|██████▌   | 331/500 [00:00<00:00, 1658.09it/s]
    100%|██████████| 500/500 [00:00<00:00, 1636.48it/s]
    Stage:  17 Beta: 0.100 Steps:  16 Acce: 0.276

      0%|          | 0/500 [00:00<?, ?it/s]
     38%|███▊      | 192/500 [00:00<00:00, 1917.25it/s]
    100%|██████████| 500/500 [00:00<00:00, 1884.04it/s]
    Stage:  18 Beta: 0.124 Steps:  14 Acce: 0.298

      0%|          | 0/500 [00:00<?, ?it/s]
     41%|████      | 203/500 [00:00<00:00, 2026.37it/s]
    100%|██████████| 500/500 [00:00<00:00, 2000.13it/s]
    Stage:  19 Beta: 0.153 Steps:  13 Acce: 0.300

      0%|          | 0/500 [00:00<?, ?it/s]
     44%|████▍     | 220/500 [00:00<00:00, 2191.91it/s]
    100%|██████████| 500/500 [00:00<00:00, 2165.52it/s]
    Stage:  20 Beta: 0.185 Steps:  12 Acce: 0.314

      0%|          | 0/500 [00:00<?, ?it/s]
     44%|████▍     | 221/500 [00:00<00:00, 2201.72it/s]
    100%|██████████| 500/500 [00:00<00:00, 2173.88it/s]
    Stage:  21 Beta: 0.221 Steps:  12 Acce: 0.298

      0%|          | 0/500 [00:00<?, ?it/s]
     44%|████▍     | 219/500 [00:00<00:00, 2184.89it/s]
    100%|██████████| 500/500 [00:00<00:00, 2168.14it/s]
    Stage:  22 Beta: 0.264 Steps:  12 Acce: 0.288

      0%|          | 0/500 [00:00<?, ?it/s]
     41%|████      | 203/500 [00:00<00:00, 2027.77it/s]
    100%|██████████| 500/500 [00:00<00:00, 2002.81it/s]
    Stage:  23 Beta: 0.317 Steps:  13 Acce: 0.297

      0%|          | 0/500 [00:00<?, ?it/s]
     40%|████      | 202/500 [00:00<00:00, 2013.44it/s]
    100%|██████████| 500/500 [00:00<00:00, 1997.53it/s]
    Stage:  24 Beta: 0.387 Steps:  13 Acce: 0.292

      0%|          | 0/500 [00:00<?, ?it/s]
     41%|████      | 205/500 [00:00<00:00, 2044.67it/s]
    100%|██████████| 500/500 [00:00<00:00, 2013.66it/s]
    Stage:  25 Beta: 0.472 Steps:  13 Acce: 0.268

      0%|          | 0/500 [00:00<?, ?it/s]
     38%|███▊      | 190/500 [00:00<00:00, 1895.87it/s]
    100%|██████████| 500/500 [00:00<00:00, 1874.79it/s]
    Stage:  26 Beta: 0.576 Steps:  14 Acce: 0.267

      0%|          | 0/500 [00:00<?, ?it/s]
     38%|███▊      | 189/500 [00:00<00:00, 1880.88it/s]
    100%|██████████| 500/500 [00:00<00:00, 1860.04it/s]
    Stage:  27 Beta: 0.723 Steps:  14 Acce: 0.274

      0%|          | 0/500 [00:00<?, ?it/s]
     38%|███▊      | 192/500 [00:00<00:00, 1911.78it/s]
    100%|██████████| 500/500 [00:00<00:00, 1880.73it/s]
    Stage:  28 Beta: 0.910 Steps:  14 Acce: 0.262

      0%|          | 0/500 [00:00<?, ?it/s]
     36%|███▌      | 179/500 [00:00<00:00, 1783.10it/s]
    100%|██████████| 500/500 [00:00<00:00, 1760.96it/s]
    Stage:  29 Beta: 1.000 Steps:  15 Acce: 0.251

      0%|          | 0/500 [00:00<?, ?it/s]
     35%|███▍      | 173/500 [00:00<00:00, 1727.79it/s]
    100%|██████████| 500/500 [00:00<00:00, 1705.16it/s]

```python
for d, log_Z in zip(D, log_Zs):
    print('degree %d gives %.4f'%(d, log_Z))
```

    degree 1 gives 25.0212
    degree 2 gives 41.0753
    degree 3 gives 142.4574
    degree 4 gives 139.3597
    degree 5 gives 137.5729

```python
plt.figure(figsize=(10, 5))
_=plt.bar(D, log_Zs, width=0.3)
_=plt.xticks(D)
plt.xticks(fontsize=20)
plt.yticks(fontsize=20)
plt.xlabel('Polynomial degree', fontsize=20)
plt.ylabel('Model Evidence', fontsize=20)
plt.tight_layout()
```

## Questions
+ The model with degree 3 polynomials has the gradest evidence. However, the degree 4 and 5 seem also very plausible. Is this a problem for the theory of Bayesian model selection? What complicates things here, is that model 3 is included in model 4 which is included in model 5. This requires us to design special priors for the models being right. They have to be consistent in some sense. For example, if model 3 is right then model 4 must be right, etc.

+ Revisit the motorcycle dataset problem. Evaluate the model evidence for a 1) Polynomial basis; 2) a Fourier basis; and 3) a Radial basis function basis.

## Revisiting Challenger Disaster Problem (Model Selection)

```python
challenger_data = np.genfromtxt("challenger_data.csv", skip_header=1,
                                usecols=[1, 2], missing_values="NA",
                                delimiter=",")
# drop the NA values
challenger_data = challenger_data[~np.isnan(challenger_data[:, 1])]

# plot it, as a function of temperature (the first column)
print("Temp (F), O-Ring failure?")
print(challenger_data)
```

    Temp (F), O-Ring failure?
    [[66.  0.]
     [70.  1.]
     [69.  0.]
     [68.  0.]
     [67.  0.]
     [72.  0.]
     [73.  0.]
     [70.  0.]
     [57.  1.]
     [63.  1.]
     [70.  1.]
     [78.  0.]
     [67.  0.]
     [53.  1.]
     [67.  0.]
     [75.  0.]
     [70.  0.]
     [81.  0.]
     [76.  0.]
     [79.  0.]
     [75.  1.]
     [76.  0.]
     [58.  1.]]

```python
# plot it, as a function of temperature (the first column)
plt.figure(figsize=(12, 5))
plt.plot(challenger_data[:, 0], challenger_data[:, 1], 'ro',
         markersize=15)
plt.ylabel("Damage Incident?",fontsize=20)
plt.xlabel("Outside temperature (Fahrenheit)",fontsize=20)
plt.title("Defects of the Space Shuttle O-Rings vs temperature",
          fontsize=20)
plt.yticks([0, 1], fontsize=15)
plt.xticks(fontsize=15)
plt.tight_layout()
```

```python
# gather the data and apply preprocessing if any
temp = challenger_data[:, 0]
temp_scaled = (temp - np.mean(temp))/np.std(temp)
data = challenger_data[:, 1]

# instantiate the pymc3 model
challenger_model = pm.Model()

# define the graph
with challenger_model:
    # define the prior
    alpha = pm.Normal('alpha', mu=0., sigma=10.)
    beta = pm.Normal('beta', mu=0., sigma=10.)

    # get the probabilities of failure at each observed temp
    p = pm.Deterministic('p', 1./(1. + tt.exp(alpha + beta*temp_scaled)))

    # define the likelihood
    x = pm.Bernoulli('x', p=p, observed=data)
print("Challenger space shuttle disaster model:")
challenger_model
```

    Challenger space shuttle disaster model:

$$
            \begin{array}{rcl}
            \text{alpha} &\sim & \text{Normal}(\mathit{mu}=0.0,~\mathit{sigma}=10.0)\\\text{beta} &\sim & \text{Normal}(\mathit{mu}=0.0,~\mathit{sigma}=10.0)\\\text{p} &\sim & \text{Deterministic}(\text{Constant},~\text{Constant},~\text{alpha},~\text{beta},~\text{Constant})\\\text{x} &\sim & \text{Bernoulli}(\mathit{p}=\text{p})
            \end{array}
            $$

```python
num_particles = 500
trace, smc = sample_smc(model=challenger_model,
                           draws=num_particles,
                          threshold=0.8,
                          progressbar=True)
```

    Sample initial stage: ...
    Stage:   0 Beta: 0.000 Steps:  25 Acce: 1.000

      0%|          | 0/500 [00:00<?, ?it/s]
     24%|██▍       | 121/500 [00:00<00:00, 1204.65it/s]
     48%|████▊     | 239/500 [00:00<00:00, 1195.81it/s]
     72%|███████▏  | 360/500 [00:00<00:00, 1199.98it/s]
    100%|██████████| 500/500 [00:00<00:00, 1192.56it/s]
    Stage:   1 Beta: 0.011 Steps:  25 Acce: 0.549

    100%|██████████| 500/500 [00:00<00:00, 5654.91it/s]
    Stage:   2 Beta: 0.028 Steps:   5 Acce: 0.406

      0%|          | 0/500 [00:00<?, ?it/s]
    100%|██████████| 500/500 [00:00<00:00, 3532.90it/s]
    Stage:   3 Beta: 0.056 Steps:   8 Acce: 0.346

      0%|          | 0/500 [00:00<?, ?it/s]
    100%|██████████| 500/500 [00:00<00:00, 2808.64it/s]
    Stage:   4 Beta: 0.106 Steps:  10 Acce: 0.301

      0%|          | 0/500 [00:00<?, ?it/s]
     48%|████▊     | 240/500 [00:00<00:00, 2394.73it/s]
    100%|██████████| 500/500 [00:00<00:00, 2361.99it/s]
    Stage:   5 Beta: 0.182 Steps:  12 Acce: 0.275

      0%|          | 0/500 [00:00<?, ?it/s]
     41%|████      | 205/500 [00:00<00:00, 2048.77it/s]
    100%|██████████| 500/500 [00:00<00:00, 2041.58it/s]
    Stage:   6 Beta: 0.313 Steps:  14 Acce: 0.266

      0%|          | 0/500 [00:00<?, ?it/s]
     42%|████▏     | 212/500 [00:00<00:00, 2117.50it/s]
    100%|██████████| 500/500 [00:00<00:00, 2076.29it/s]
    Stage:   7 Beta: 0.529 Steps:  14 Acce: 0.250

      0%|          | 0/500 [00:00<?, ?it/s]
     39%|███▉      | 196/500 [00:00<00:00, 1959.82it/s]
    100%|██████████| 500/500 [00:00<00:00, 1943.86it/s]
    Stage:   8 Beta: 0.897 Steps:  15 Acce: 0.270

      0%|          | 0/500 [00:00<?, ?it/s]
     42%|████▏     | 209/500 [00:00<00:00, 2085.22it/s]
    100%|██████████| 500/500 [00:00<00:00, 2060.88it/s]
    Stage:   9 Beta: 1.000 Steps:  14 Acce: 0.261

      0%|          | 0/500 [00:00<?, ?it/s]
     38%|███▊      | 192/500 [00:00<00:00, 1912.39it/s]
    100%|██████████| 500/500 [00:00<00:00, 1860.62it/s]

```python
ppsamples = pm.sample_posterior_predictive(model=challenger_model,
                                           trace=trace,
                                           var_names=['p'])['p']
```

      0%|          | 0/500 [00:00<?, ?it/s]
    100%|██████████| 500/500 [00:00<00:00, 3612.98it/s]

```python
ppmean = ppsamples.mean(0)
pp_lower, pp_upper = np.percentile(ppsamples, axis=0, q=[2.5, 97.5])

plt.figure(figsize=(15, 8))
plt.plot(temp, data, 'ro', markersize=12, label='Observed data')
idx=np.argsort(temp)
plt.plot(temp[idx], ppmean[idx], linestyle='--', linewidth=2.5,
         label='Post. pred. mean prob.')
plt.fill_between(temp[idx], pp_lower[idx], pp_upper[idx],
                color='purple', alpha=0.25, label='95% Confidence')
plt.ylabel("Probability estimate",fontsize=20)
plt.xlabel("Outside temperature (Fahrenheit)",fontsize=20)
plt.title("Defects of the Space Shuttle O-Rings vs temperature",
          fontsize=20)
plt.yticks(np.arange(0., 1.01, 0.2), fontsize=20)
plt.xticks(fontsize=20)
plt.legend(loc='best', fontsize=20)
plt.tight_layout()
```

```python
logZ_temp = np.log(smc.model.marginal_likelihood)
```

```python

```


## Linear Population Differential Equations

In this notebook we will solve a simple population equation [1]

```python
## Library
import numpy as np
import matplotlib.pyplot as plt
```

## Linear Population Equation
The linear population equation is a first order linear differential equation given by,
\begin{equation}
\frac{d p}{dt}=\alpha_{Birth} p-\alpha_{Death}p=\alpha p,
\end{equation}
where $p$ is the population at time $t$,
for $a\leq t \leq b$ years, where $\alpha$ is the birth rate $\alpha_{Birth}$ minus death rate $\alpha_{Death}$, with the initial condition
$$ p(a)=A. $$
The linear population equation has the exact solution
\begin{equation}
p(t)=Ae^{(\alpha (t-a))}.
\end{equation}
In this simulation we set the parameters of the different equations:
*  as $\alpha=0.3$, with the inital condition $p(0)=3$, over the time period $0\leq t \leq 10$.

```python
alpha=0.3
```

## Discrete Time Domain
To numerically approximate the solution we discretise the continuous domain using a step size, $h=\frac{10-0}{10}=1$,
which gives
\begin{equation}
t_i=0+i h,
\end{equation}
for $i=0,...10$.

The Figure below illustrate the discrete domain.

```python
h=1
a=0
b=10
t=np.arange(a,b+h/2,h)
fig = plt.figure(1,figsize=(5,4))
ax = fig.add_subplot(1,1,1)
ax.plot(t,0*t,'o:',color='k')
ax.set_xlabel('Time (yrs)')
plt.title('Discrete Time')
plt.show()
```

```python
N=int((b-a)/h) # Number of Steps
P=np.zeros(N+1) # Numerical Solution P
P[0]=3
p=3*np.exp(alpha*(t-a)) # Exact Solution
```

We numerically approximate the differential equation using Eulers method to give the difference equation
\begin{equation}
P_{i}=P_{i-1}+h \big(\alpha P_{i-1} \big),
\end{equation}
where $h$ is the stepsize and where $P_i$ is the numerical approximation of $p(t_i)$ at distance $t_i$, for $i=0,...,10.$

```python
for i in range (1,N+1):
    P[i]=P[i-1]+h*(alpha*P[i-1])

```

## Results
The plot bellow shows the Numerical solution $P_i$ (left), the analytic solution $p(x)$ (middle) and the error $P(x_i)-P_i$ (right) as a function of time $t$.

```python
## Plotting Figure
fig = plt.figure(1,figsize=(12,4))

# --- left hand plot
ax = fig.add_subplot(1,3,1)
ax.plot(t,P,'o:',color='red',label='Population')
ax.legend()

ax.set_xlabel('Time (yrs)')
ax.set_ylabel('Population (billions)')
#ax.legend(loc='best')
ax.set_title('Numerical Solution, P')

# --- middle plot
ax = fig.add_subplot(1,3,2)
ax.plot(t,p,'-',color='red',label='Population')

ax.legend()

ax.set_xlabel('Time (yrs)')
ax.set_ylabel('Population (billions)')

ax.set_title('Analytic (exact) Solution p(t)')

# --- right hand plot
ax = fig.add_subplot(1,3,3)
ax.plot(t,P-p,'o:',color='red',label='ERROR')
ax.legend()

ax.set_title(r'Error $p(x_i)-P_i$')
ax.set_ylabel('Error(Population)')
ax.set_xlabel('Time (yrs)')

# --- title, explanatory text and save
fig.suptitle('Linear Population Equation', fontsize=20)
plt.tight_layout()
plt.subplots_adjust(top=0.85)

```

# References

[1] Stover, Christopher and Weisstein, Eric W. "Population Growth." From MathWorld--A Wolfram Web Resource. http://mathworld.wolfram.com/PopulationGrowth.html


```python
import sympy as sy
import numpy as np
import control.matlab as cm
from mpmath import polyroots
sy.init_printing()
```

```python
kb, k, f, m, M = sy.symbols('k_b, k, f, m, M', real=True)
W = 1/m + 1/M
```

```python
s = sy.symbols('s', real=False)
```

```python
A = sy.Matrix([[0, 1, 0, 0], [-kb/M, 0, k/M, f/M], [0,0,0,1], [kb/M, 0, -k*W, -f*W]])
B = sy.Matrix([[0],[1/M], [0], [-W]])
Bv = sy.Matrix([[0],[1/M], [0], [-1/M]])
C = sy.Matrix([[1., 0, 0, 0]])

A

```

```python
G = C*(s*sy.eye(4)- A).inv()*B
Gv = C*(s*sy.eye(4)- A).inv()*Bv

```

```python
G
```

```python
# Constants
Mn = 100
mn = 6
kbn = 100
kn = 3
fn = 3

Gn = G.subs([(M,Mn), (m,mn), (kb, kbn), (k, kn), (f, fn)])
Gn = sy.simplify(Gn[0])
(num, den) = sy.fraction(Gn)
numP = sy.Poly(num, s)
numcoeffs = np.asarray(numP.all_coeffs(), dtype=float)
denP = sy.Poly(den, s)
dencoeffs = np.asarray(denP.all_coeffs(), dtype=float)

Gvn = Gv.subs([(M,Mn), (m,mn), (kb, kbn), (k, kn), (f, fn)])
Gvn = sy.simplify(Gvn[0])
(numv, denv) = sy.fraction(Gvn)
numPv = sy.Poly(numv, s)
numcoeffsv = np.asarray(numPv.all_coeffs(), dtype=float)
denPv = sy.Poly(denv, s)
dencoeffsv = np.asarray(denPv.all_coeffs(), dtype=float)

```

```python
#sy.factor(denP, extension=[sy.I])
poles = polyroots(dencoeffs)
poles
```

    [mpc(real='-0.23408260199788464', imag='-0.66645226904530763'),
     mpc(real='-0.23408260199788464', imag='0.66645226904530763'),
     mpc(real='-0.030917398002115359', imag='1.000570798665283'),
     mpc(real='-0.030917398002115359', imag='-1.000570798665283')]

```python
Gs = cm.tf(numcoeffs, dencoeffs)
Gsv = cm.tf(numcoeffsv, dencoeffsv)
cm.bode(Gs, Gsv)

```

```python
cm.pole(Gs)
```

    array([-0.0309174+1.0005708j , -0.0309174-1.0005708j ,
           -0.2340826+0.66645227j, -0.2340826-0.66645227j])

```python
Gsv
```

               s^2 + 0.5 s + 0.5
    --------------------------------------
    100 s^4 + 53 s^3 + 153 s^2 + 50 s + 50

```python
numcoeffs
```

    array([ 1.,  0.,  0.])

```python
Gs2 = cm.tf(1./100*np.array([0,0,1.]), [1, 0.5, 1.5, 0.5, 0.5])
Gs2
```

                     0.01
    -------------------------------------
    s^4 + 0.5 s^3 + 1.5 s^2 + 0.5 s + 0.5

```python
cm.pole(Gs2)
```

    array([  5.27355937e-16+1.j        ,   5.27355937e-16-1.j        ,
            -2.50000000e-01+0.66143783j,  -2.50000000e-01-0.66143783j])

```python
sy.simplify(sy.expand((s**2 + 2*0.1*1*s + 1)*(s**2 + 5*s + 7)))
```

```python
(num3, den3) = cm.zpk2tf(k=0.5, p=[-0.1+1j, -0.1-1j, -0.25 + 0.6*1j, -0.25-0.6*1j], z=[0, 0])
Gs3 = cm.tf(num3, den3)
Gs3
```

                       0.5 s^2
    ---------------------------------------------
    s^4 + 0.7 s^3 + 1.533 s^2 + 0.5895 s + 0.4267

```python
cm.rlocus(Gs3)
```

```python
1+ 1J
```

    (1+1j)

```python

```

notebooks/Active mass damper system.ipynb


```python
import numpy as np
from sympy import *
```

```python
def newton_integration(function, a, b, m):
    h = (b - a)/(3*m)
    x_values = np.arange(a, b, h)
    sum = function(a) + function(x_values[-1]*3)
    subsum1 = 0
    subsum2 = 0
    for i in range(m):
        subsum1 += function(x_values[3*i-2]) + function(x_values[3*i-1])
        if i != m - 1:
            subsum2 += function(x_values[3*i])
    sum += 3*subsum1
    sum += 2*subsum2
    return 3*h/8*sum
```

```python
x = Symbol('x')
function = lambdify(x, sin(x), 'numpy')
a = 0
b = np.pi
m = 1000
```

```python
from scipy import integrate
exact_integral = -cos(np.pi) + cos(0)
newton_integral = newton_integration(function, a, b, m)
print('Вычисленный интеграл: ', newton_integral)
print('Точное значение:', exact_integral)
print('Ошибка:', exact_integral-newton_integral)
```

    Вычисленный интеграл:  1.999995065203693
    Точное значение: 2.00000000000000
    Ошибка: 4.93479630692129e-6


# New Algorithms for Simulating Dynamical Friction

## Dan T. Abell, David L. Bruhwiler, Stephen D. Webb — *RadiaSoft, LLC*

This notebook describes—and documents in code—algorithms for simulating
the dynamical friction experienced by ions in the presence of magnetized electrons.

The $\LaTeX$ preamble is *here*.
$$
%% math text
\newcommand{\hmhsp}{\mspace{1mu}}% math hair space
\newcommand{\mhsp}{\mspace{2mu}}% math hair space
\newcommand{\ud}{\mathop{}\!\mathrm{d}}% upright d for differential
\newcommand{\ui}{\mathrm{i}}% upright i for imaginary unit
\newcommand{\ue}{\mathrm{e}}% upright e for Euler number
%%
\newcommand{\Mion}{m_\text{ion}}
\newcommand{\Me}{m_\text{e}}
%%
\newcommand{\vQion}{\vec{q}_\text{ion}}
\newcommand{\vPion}{\vec{p}_\text{ion}}
\newcommand{\Qion}[1]{#1_\text{ion}}
\newcommand{\Pion}[1]{p_{\text{ion},\hmhsp#1}}
%%
\newcommand{\vQe}{\vec{q}_\text{e}}
\newcommand{\vPe}{\vec{p}_\text{e}}
\newcommand{\Qe}[1]{#1_\text{e}}
\newcommand{\Pe}[1]{p_{\text{e},\hmhsp#1}}
%%
\newcommand{\Map}[2][]{\mathcal{#2}^{#1}}
%%
\newcommand{\pgc}{p_\text{gc}}
\newcommand{\xgc}{x_\text{gc}}
\newcommand{\ygc}{y_\text{gc}}
$$

```python
""" Python preamble """
%matplotlib inline
```

```python
print mp
```

```python
""" Python preamble (cont.) """
from __future__ import division

import numpy as np
import math
import matplotlib.pyplot as plt
import matplotlib.cm as cm
from matplotlib.colors import LogNorm
import matplotlib as mpl

from scipy.constants import pi
from scipy.constants import speed_of_light as clight
from scipy.constants import epsilon_0 as eps0
from scipy.constants import mu_0 as mu0
from scipy.constants import elementary_charge as qe
from scipy.constants import electron_mass as me
from scipy.constants import proton_mass as mp
from scipy.constants import Boltzmann as kB
fourPiEps0 = 4 * pi * eps0
invFourPiEps0 = 1 / fourPiEps0

""" reset some default options """
np.set_printoptions(linewidth=96)

""" indexing """
(Ix, Ipx, Iy, Ipy, Iz, Ipz) = range(6)

""" prefixes """
(femto, pico, nano, micro, milli, one, kilo, mega, giga, tera, peta) = \
    10. ** np.asarray(range(-15, 15+1, 3))
```

We define the ion charge and mass here as global parameters.
We do the same for the magnetic field strength $B$ and the
thermal velocity $v_\text{th}$.

Then we compute various related derived quantities.

```python
"""
angular frequency of Larmor rotations
NB: This is a *signed* quantity, which means
that for electrons, say, you must set Z = -1.
"""
def omega_Larmor(mass, B, Z = 1):
    return Z * qe * B / mass

Z_ion = 1
M_ion = mp
B_mag = 1. # Tesla
e_temp = 300. # Kelvin

N_gyro = 100 # a somewhat arbitrary choice, range [100, 160]

""" derived quantities """
V_th = math.sqrt(2 * kB * e_temp / me)
rho_gc = me * V_th / (qe * B_mag)
Omega_e = omega_Larmor(me, B_mag, Z = -1)
T_e = (2 * pi) / abs(Omega_e)
T_intxn = N_gyro * T_e

print "V_th = ", V_th
print "rho_gc / µm = ", rho_gc / micro
print "Omega_e / s^(-1) = ", Omega_e
print "frequency / GHz = ", Omega_e / (2 * pi) / giga
print "T_e / ns = ", T_e / nano
print "T_intxn / ns = ", T_intxn / nano
```

    V_th =  95361.4171888
    rho_gc / µm =  0.542189740332
    Omega_e / s^(-1) =  -1.7588200236e+11
    frequency / GHz =  -27.9924900765
    T_e / ns =  0.0357238672682
    T_intxn / ns =  3.57238672682

## Two-body Magnetized Collisions

The Hamiltonian for a two-body interaction between an ion and a magnetized electron is
$$
\vphantom{\Big]}
  H(\vQion, \vPion, \vQe, \vPe)
    = H_0(\vPion, \Qe{y}, \vPe)
    + H_\text{C}(\vQion, \vQe)
$$
where
$$
\begin{align}
  H_0(\vPion, \Qe{y}, \vPe)
    &= \frac{1}{2\Mion}\bigl(\Pion{x}^2 + \Pion{y}^2 + \Pion{z}^2\bigr)
     + \frac{1}{2\Me}\bigl((\Pe{x} + e B \Qe{y})^2 + \Pe{y}^2 + \Pe{z}^2\bigr),\\[1ex]
  H_\text{C}(\vQion, \vQe)
    &= -\frac{Ze^2}{4\pi\varepsilon_0} \big/
                {\sqrt{(\Qion{x}-\Qe{x})^2 + (\Qion{y}-\Qe{y})^2 + (\Qion{z}-\Qe{z})^2}},
\end{align}\\[1ex]
$$
and $e$ denotes the elementary quantum of charge.

The simplest second-order scheme for integrating this system uses
a *split-operator* approach: We approximate the total map $\Map{M}$ for a
time step of size $h$ by the symmetric form
$$
\vphantom{\Big]}
  \Map{M}(h) \approx \Map{M}_0(h/2) \Map{M}_C(h) \Map{M}_0(h/2)
$$
where $\Map{M}_0$ and $\Map{M}_C$ are the exact maps for the Hamiltonians $H_0$
and $H_C$ respectively. The map $\Map{M}_0$ is a simple linear map. The map
$\Map{M}_C$ generates a nonlinear kick of both ion and electron momenta.

### Hamiltonians for Two-body Magnetized Collisions

```python
"""
Hamiltonian for free ion and electron in a magnetic field, under
the assuption that the ion is unaffected by that magnetic field.

Arguments:
  z_i (ndArray): 6 x N array of canonical coördinates
                 and conjugate momenta for the ions
  z_e (ndArray): 6 x N array of canonical coördinates
                 and conjugate momenta for the electrons
  In both of the above arrays, the six phase-space variables
  are given in the order(x, px, y, py, z, pz)

Return:
  the total 'free' energy of each ion-electron pair

"""
def H_twobody_0(z_i, z_e):
    ham_i = ((z_i[Ipx,:] ** 2 + z_i[Ipy,:] ** 2 + z_i[Ipz,:] ** 2)
             / (2 * M_ion))
    ham_e = ((z_e[Ipx,:] + (-qe) * B_mag * z_e[Iy,:]) ** 2
             + z_e[Ipy,:] ** 2 + z_e[Ipz,:] ** 2) / (2 * me)
    return ham_i + ham_e

"""
Hamiltonian for the interaction of each ion-electron pair.
"""
def H_twobody_C(z_i, z_e):
    g_ie = -(Z_ion * qe ** 2) / (4 * pi * eps0)
    intxn = g_ie / np.sqrt(
        + (z_i[Ix,:] - z_e[Ix,:]) ** 2
        + (z_i[Iy,:] - z_e[Iy,:]) ** 2
        + (z_i[Iz,:] - z_e[Iz,:]) ** 2)
    return intxn

"""
Total Hamiltonian for each ion-electron pair.
"""
def H_twobody(z_i, z_e):
    ham_0 = H_twobody_0(z_i, z_e)
    ham_C = H_twobody_C(z_i, z_e)
    return ham_0 + ham_C
```

### Maps for Two-body Magnetized Collisions

```python
"""
define transfer maps for ions and electrons
There are three maps to define here: one each
for ions and electrons under H_0, and another
"""

""" matrix for a linear drift """
def MatD(mass, h):
    Mdrift = np.identity(6)
    for i in (Ix, Iy, Iz):
        Mdrift[i, i + 1] = h / mass
    return Mdrift

""" matrix for linear electron dynamics in a solenoidal field """
def MatK0_e(h):
    mw = me * Omega_e
    wh = Omega_e * h
    cwh = math.cos(wh)
    swh = math.sin(wh)
    cwh1m = 2 * math.sin(wh / 2) ** 2  # 1 - cos(a) = 2 sin^2(a / 2)
    MK0 = np.identity(6)
    MK0[Iy,  Iy ] = cwh
    MK0[Ipy, Ipy] = cwh
    MK0[Iy,  Ipy] = swh / mw
    MK0[Ipy, Iy ] = -mw * swh
    MK0[Iz,  Ipz] = h / me
    MK0[Ix,  Ipx] = swh / mw
    MK0[Ix,  Iy ] = swh
    MK0[Ix,  Ipy] = cwh1m / mw
    MK0[Iy,  Ipx] = -cwh1m / mw
    MK0[Ipy, Ipx] = -swh
    return MK0

"""
map phase-space coördinates forward in time by amount h
based on the Hamiltonian H_0, which describes the free
motion of ions and the motion of electrons in a solenoidal
magnetic field
"""
def MapZ_0(h, z_i, z_e):
    mat = MatD(M_ion, h)
    zf_i = mat.dot(z_i)
    mat = MatK0_e(h)
    zf_e = mat.dot(z_e)
    return zf_i, zf_e

"""
map phase-space coördinates forward in time by amount h
based on the Hamiltonian H_C, which describes the collision
between a single ion-electron pair
"""
def MapZ_C(h, z_i, z_e):
    g = h * Z_ion * qe ** 2 / (4 * pi * eps0)
    dz = z_i - z_e
    denom = (dz[Ix,:] ** 2 + dz[Iy,:] ** 2 + dz[Iz,:] ** 2) ** (3/2)
    zf_i = z_i.copy()
    zf_e = z_e.copy()
    for ip in (Ipx, Ipy, Ipz):
        zf_i[ip,:] = z_i[ip,:] - g * dz[ip - 1] / denom
        zf_e[ip,:] = z_e[ip,:] + g * dz[ip - 1] / denom
    return zf_i, zf_e

def apply_MapZ_0(h, n, z_i, z_e):
    mat_i = MatD(M_ion, h)
    mat_e = MatK0_e(h)
    zf_i = [z_i]
    zf_e = [z_e]
    for i in range(n):
        z_i = mat_i.dot(z_i)
        z_e = mat_e.dot(z_e)
        zf_i.append(z_i)
        zf_e.append(z_e)
    return np.asarray(zf_i), np.asarray(zf_e)

""" second-order split-operator integration for the total Hamiltonian """
def apply_MapZ(h, n, z_i, z_e):
    hh = 0.5 * h
    mat_i = MatD(M_ion, hh)
    mat_e = MatK0_e(hh)
    zf_i = [z_i]
    zf_e = [z_e]
    for i in range(n):
        z_i = mat_i.dot(z_i)
        z_e = mat_e.dot(z_e)
        z_i, z_e = MapZ_C(h, z_i, z_e)
        z_e = mat_e.dot(z_e)
        z_i = mat_i.dot(z_i)
        zf_i.append(z_i)
        zf_e.append(z_e)
    return np.asarray(zf_i), np.asarray(zf_e)
```

## Guiding-center Coördinates and $\Theta$-J Coördinates
### Transformations To and From Guiding-center Coördinates and $\Theta$-J Coördinates

We transform the electron's transverse phase-space coördinates
using the type-1 generating function
$$
  F_1(x,y;\, \phi,\ygc)
    = m\Omega\Bigl[\frac{1}{2}(y - \ygc)^2\cot\phi - y \ygc\Bigr].
$$
This yields the following transformation rules:
*to* guiding-center coördinates
$$
\begin{align}
  m\Omega &= qB_0,
               \quad\text{(this is a signed quantity)}\\[1ex]
  \phi    &= \arctan\Bigl(\frac{p_x + e B y}{p_y}\Bigr),\\[1ex]
  p_\phi  &= \frac{1}{2m\Omega}\bigl[(p_x + m\Omega y)^2 + p_y^2\bigr],\\[1ex]
  \ygc    &= -\frac{p_x}{m\Omega},\\[1ex]
  \pgc    &= p_y + m\Omega x.
\end{align}
$$
*from* guiding-center coördinates
$$
\begin{align}
  r_L &= \frac{1}{m\Omega}\sqrt{2m\Omega\,p_\phi},
           \quad\text{(this is a signed quantity)}\\[1ex]
  x   &= \frac{\pgc}{m\Omega} - r_L\cos\phi,\\[1ex]
  p_x &= -m\Omega\,\ygc,\\[1ex]
  y   &= \ygc + r_L\sin\phi,\\[1ex]
  p_y &= m\Omega\,r_L\cos\phi.
\end{align}
$$

We also require the transformation to and from the coördinates $\Theta$-J:
$$
\begin{align}
  \Theta &= \dotsb, \\
  J &= p_\phi + \frac{Ze^2}{4\pi\varepsilon_0} \frac{r_L}{\Omega}
      \frac{(\Qion{x}-\xgc)\cos\phi - (\Qion{y}-\ygc)\sin\phi}{%
        \bigl[(\Qion{x}-\Qe{x})^2 + (\Qion{y}-\Qe{y})^2 + (\Qion{z}-\Qe{z})^2
              + r_L^2\bigr]^{3/2}}.
\end{align}
$$

$$
\begin{align}
  \phi &= \dotsb, \\
  p_\phi &=  \dotsb.
\end{align}
$$

```python
""" convert to guiding-center coordinates """
def toGuidingCenter(z_e):
    mOmega = me * Omega_e
    zgc = z_e.copy()
    zgc[Ix,:]  = np.arctan2(z_e[Ipx,:] + mOmega * z_e[Iy,:], z_e[Ipy,:])
    zgc[Ipx,:] = (((z_e[Ipx,:] + mOmega * z_e[Iy,:]) ** 2
                   + z_e[Ipy,:] ** 2) / (2 * mOmega))
    zgc[Iy,:]  = - z_e[Ipx,:] / mOmega
    zgc[Ipy,:] = z_e[Ipy,:] + mOmega * z_e[Ix,:]
    return zgc

""" convert from guiding-center coordinates """
def fromGuidingCenter(zgc):
    mOmega = me * Omega_e
    rhoL = np.sqrt(2 * mOmega * zgc[Ipx,:]) / mOmega
    z_e = zgc.copy()
    z_e[Ix,:]  = zgc[Ipy,:] / mOmega - rhoL * np.cos(zgc[Ix,:])
    z_e[Ipx,:] = - mOmega * zgc[Iy,:]
    z_e[Iy,:]  = zgc[Iy,:] + rhoL * np.sin(zgc[Ix,:])
    z_e[Ipy,:] = mOmega * rhoL * np.cos(zgc[Ix,:])
    return z_e

""" return J(z_gc) coordinates using the (probably correct) minus sign """
def actionJ(z_i, zgc):
    g = Z_ion * qe ** 2 / (4 * pi * eps0)
    mOmega = me * Omega_e
    rhoL = np.sqrt(2 * zgc[Ipx,:] / mOmega)
    num = ((z_i[Ix,:] - zgc[Ipy,:] / mOmega) * np.cos(zgc[Ix,:])
        -  (z_i[Iy,:] - zgc[Iy,:]) * np.sin(zgc[Ix,:]))
    den = ((z_i[Ix,:] - zgc[Ipy,:] / mOmega) ** 2
        +  (z_i[Iy,:] - zgc[Iy,:]) ** 2
        +  (z_i[Iz,:] - zgc[Iz,:]) ** 2 + rhoL ** 2) ** (3/2)
    return zgc[Ipx,:] + g * (rhoL / Omega_e) * num / den
```

```python
""" return the Larmor radius """
def rLarmor(z_e):
    mOmega = me * Omega_e
    return np.sqrt((z_e[Ipx,:] + mOmega * z_e[Iy,:]) ** 2 + z_e[Ipy,:] ** 2) / mOmega

""" return the Larmor radius """
def rLarmor_gc(zgc):
    mOmega = me * Omega_e
    return np.sqrt(2 * mOmega * zgc[Ipx,:]) / mOmega

""" return the perturbation ratio (uses the minus sign) """
def pertubationRatio(z_i, zgc):
    mOmega = me * Omega_e
    rhoL = np.sqrt(2 * mOmega * zgc[Ipx,:]) / mOmega
    num = ((z_i[Ix,:] - zgc[Ipy,:] / mOmega) * np.cos(zgc[Ix,:])
           - (z_i[Iy,:] - zgc[Iy,:]) * np.sin(zgc[Ix,:]))
    den = ((z_i[Ix,:] - zgc[Ipy,:] / mOmega) ** 2
        +  (z_i[Iy,:] - zgc[Iy,:]) ** 2
        +  (z_i[Iz,:] - zgc[Iz,:]) ** 2 + rhoL ** 2)
    return 2 * rhoL * num / den

""" return the ratio (impact parameter) / (Larmor radius) """
def impactParamOverRL(z_i, zgc):
    mOmega = me * Omega_e
    rhoL = np.sqrt(2 * mOmega * zgc[Ipx,:]) / mOmega
    b = np.sqrt((z_i[Ix,:] - zgc[Ipy,:] / mOmega) ** 2
               + (z_i[Iy,:] - zgc[Iy,:]) ** 2)
    return b / rhoL
```

### Hamiltonians using Guiding-center Coördinates

```python
"""
Hamiltonian for free ion and electron in a magnetic field, under
the assuption that the ion is unaffected by that magnetic field.
"""
def H_gc_0(z_i, zgc):
    ham_i = ((z_i[Ipx,:] ** 2 + z_i[Ipy,:] ** 2 + z_i[Ipz,:] ** 2)
             / (2 * M_ion))
    ham_e = Omega_e * actionJ(z_i, zgc) + zgc[Ipz,:] ** 2 / (2 * me)
    return ham_i + ham_e

"""
Hamiltonian for the ion-electron interaction
in guiding-center (Θ,J) coördinates.
"""
def H_gc_C(z_i, zgc):
    g_ie = (Z_ion * qe ** 2) / (4 * pi * eps0)
    mOmega = me * Omega_e
    intxn = -g_ie / np.sqrt(
        + (z_i[Ix,:] - zgc[Ipy,:] / mOmega) ** 2
        + (z_i[Iy,:] - zgc[Iy,:]) ** 2
        + (z_i[Iz,:] - zgc[Iz,:]) ** 2
        + 2 * actionJ(z_i, zgc) / mOmega)
    return intxn

"""
total Hamiltonian for the ion-electron system in GC coördinates
"""
def H_gc(z_i, zgc):
    ham_0 = H_gc_0(z_i, zgc)
    ham_C = H_gc_C(z_i, zgc)
    return ham_0 + ham_C

"""
Hamiltonian for the ion-electron interaction
in guiding-center (Θ,J) coördinates.
"""
def H_gc_Cp(z_i, zgc):
    g_ie = (Z_ion * qe ** 2) / (4 * pi * eps0)
    mOmega = me * Omega_e
    intxn = -g_ie / np.sqrt(
        + (z_i[Ix,:] - zgc[Ipy,:] / mOmega) ** 2
        + (z_i[Iy,:] - zgc[Iy,:]) ** 2
        + (z_i[Iz,:] - zgc[Iz,:]) ** 2
        + 2 * actionJp(z_i, zgc) / mOmega)
    return intxn

"""
total Hamiltonian for the ion-electron system in GC coördinates
"""
def H_gc_m(z_i, zgc):
    ham_0 = H_gc_0(z_i, zgc)
    ham_C = H_gc_Cm(z_i, zgc)
    return ham_0 + ham_C
```

### Maps using Guiding-center Coördinates

```python
"""
define transfer maps for ions and electrons
There are three maps to define here: one each
for ions and electrons under H_0, and another
for the interaction under H_c
"""

"""
Map phase-space coördinates forward in time by amount h.
This map is based on the Hamiltonian H_gc_0, which describes
the free motion of ions and the averaged motion of electrons
in a solenoidal magnetic field.
NB: We do not update the \Theta coördinate, as it does not
contribute to the dynamics of any other variables.
"""
def MapZgc_0(h, z_i, zgc):
    matD = MatD(M_ion, h)
    zf_i = matD.dot(z_i)
    zf_e = zgc.copy()
    zf_e[Iz,:] += (h / me) * zgc[Ipz,:]
    return zf_i, zf_e

"""
Map phase-space coördinates forward in time by amount h.
This map is based on the Hamiltonian H_gc_C, which describes
the collision between a single ion-electron pair in guiding-center
coördinates.
NB: We do not update the \Theta coördinate, as it does not
contribute to the dynamics of any other variables.
"""
def MapZgc_C(h, z_i, zgc):
    g_ie = Z_ion * qe ** 2 / (4 * pi * eps0)
    mOmega = me * Omega_e
    dr3 = ((z_i[Ix,:] - zgc[Ipy,:] / mOmega) ** 2
           + (z_i[Iy,:] - zgc[Iy,:]) ** 2
           + (z_i[Iz,:] - zgc[Iz,:]) ** 2
           + (2 / mOmega) * actionJ(z_i, zgc)) ** (3/2)
    Omega_gc = (g_ie / mOmega) / dr3
    S = np.sin(Omega_gc * h)
    C1 = 2 * np.sin(Omega_gc * (h / 2)) ** 2
    zf_i = z_i.copy()
    zf_e = zgc.copy()
    Dxgc = ((z_i[Ix,:] - zgc[Ipy,:] / mOmega) * C1
            + (z_i[Iy,:] - zgc[Iy,:]) * S)
    Dygc = ((z_i[Iy,:] - zgc[Iy,:]) * C1
            - (z_i[Ix,:]- zgc[Ipy,:] / mOmega) * S)
    Dpz = (Omega_gc * h) * mOmega * (z_i[Iz,:] - zgc[Iz,:])
    zf_i[Ipx,:] += mOmega * Dygc
    zf_i[Ipy,:] -= mOmega * Dxgc
    zf_i[Ipz,:] -= Dpz
    zf_e[Iy,:]  += Dygc
    zf_e[Ipy,:] += mOmega * Dxgc
    zf_e[Ipz,:] += Dpz
    return zf_i, zf_e

def apply_MapZgc_0(h, n, z_i, zgc):
    mat_i = MatD(M_ion, h)
    mat_e = np.identity(6)
    mat_e[Iz, Ipz] = h / me
    zf_i = [z_i]
    zf_e = [zgc]
    for i in range(n):
        z_i = mat_i.dot(z_i)
        zgc = mat_e.dot(zgc)
        zf_i.append(z_i)
        zf_e.append(zgc)
    return np.asarray(zf_i), np.asarray(zf_e)

def apply_MapZgc(h, n, z_i, zgc):
    hh = 0.5 * h
    mat_i = MatD(M_ion, hh)
    mat_e = np.identity(6)
    mat_e[Iz, Ipz] = hh / me
    zf_i = [z_i]
    zf_e = [zgc]
    for i in range(n):
        z_i = mat_i.dot(z_i)
        zgc = mat_e.dot(zgc)
        z_i, zgc = MapZgc_C(h, z_i, zgc)
        zgc = mat_e.dot(zgc)
        z_i = mat_i.dot(z_i)
        zf_i.append(z_i)
        zf_e.append(zgc)
    return np.asarray(zf_i), np.asarray(zf_e)
```

### Maps using the Magnus Expansion

```python
"""
compute $\Delta P_\text{ion}$ using the Magnus expansion
"""
def magnus_deltaPIon(h, z_i, zgc):
    g_ie = (Z_ion * qe ** 2) / (4 * pi * eps0)
    mOmega = me * Omega_e
    xgc = zgc[Ipy,:] / mOmega
    C1 = ((z_i[Ix,:] - xgc) ** 2
        + (z_i[Iy,:] - zgc[Iy,:]) ** 2
        + (z_i[Iz,:] - zgc[Iz,:]) ** 2
        + (2 / mOmega) * actionJ(z_i, zgc))
    C2 = (2 * ((z_i[Ix,:] - xgc) * z_i[Ipx,:] / M_ion
             + (z_i[Iy,:] - zgc[Iy,:]) * z_i[Ipy,:] / M_ion
             + (z_i[Iz,:] - zgc[Iz,:]) * (z_i[Ipz,:] / M_ion - zgc[Ipz,:] / me)))
    C3 = ((z_i[Ipx,:] / M_ion) ** 2
        + (z_i[Ipy,:] / M_ion) ** 2
        + (z_i[Ipz,:] / M_ion - zgc[Ipz,:] / me) ** 2)
    B = np.sqrt(C1 + (C2 + C3 * h) * h)
    Delta = 4 * C1 * C3 - C2 ** 2
    D1 = (2 * C3 * h + C2) / B - C2 / np.sqrt(C1)
    D2 = (C2 * h + 2 * C1) / B - 2 * np.sqrt(C1)
    dPx = - ((2 * g_ie / Delta) *
             ((z_i[Ix,:] - xgc) * D1 - (z_i[Ipx,:] / M_ion) * D2))
    dPy = - ((2 * g_ie / Delta) *
             ((z_i[Iy,:] - zgc[Iy,:]) * D1 - (z_i[Ipy,:] / M_ion) * D2))
    dPz = - ((2 * g_ie / Delta) *
             ((z_i[Iz,:] - zgc[Iz,:]) * D1 - (z_i[Ipz,:] / M_ion - zgc[Ipz,:] / me) * D2))
    return np.asarray(( dPx, dPy, dPz)).T
```

```python
"""
For the Magnus computation to work well, the interaction must be perturbative.
This function return a value for the minimum impact parameter, above which the
interaction becomes perturbative.
"""
def minImpactParam(magB, EkOverEV, bOverRL):
    numer = 2 * EkOverEV * (invFourPiEps0 * Z_ion * me)
    denom = ((1 / bOverRL) ** 2 + (tan_alpha / (N_gyro * pi)) ** 2) * magB ** 2
    return (numer / denom) ** (1/3)
```

### Conversion from (*Q*,*V*) to (*Q*,*P*)

```python
"""
define matrix that transforms ion coördinate-velocity data
to canonically conjugate phase-space variables
"""
QVtoZion = np.identity(6)
for i in (Ipx, Ipy, Ipz):
    QVtoZion[i, i] = M_ion

"""
define matrix that transforms electron coördinate-velocity data
to canonically conjugate phase-space variables
NB: This depends on the local magnetic vector potential,
hence on the local magnetic field.
"""
QVtoZe = np.identity(6)
for i in (Ipx, Ipy, Ipz):
    QVtoZe[i, i] = me
QVtoZe[Ipx, Iy] = (-qe) * (-B_mag)
```

# Explore a range of values for the perturbation parameters

```python
Z_ion = 1
M_ion = mp
e_temp = 300. # Kelvin
B_mag = 1. # Tesla

N_gyro = 100    # enforce adequate averaging
tan_alpha = 5.0 # and an adequate opening angle

""" derived quantities """
Omega_e = omega_Larmor(me, B_mag, Z = -1)
T_e = (2 * pi) / abs(Omega_e)
```

```python
bOverLarmorR = 20.0 # 20 ==> max perturbation ratio of about 0.1
EkinOverVmax = 10.0 # 10 ==> eV_closest_approach / E_kinetic = 0.1
print minImpactParam(B_mag, EkinOverVmax, bOverLarmorR) / micro
```

    3.90333173895

```python
decades = 2
pts_per_decade = 3
logBs = np.linspace(0., 1. * float(decades), num = 1 + pts_per_decade * decades, endpoint = True)
bvals = np.power(10, logBs)
impactParameterB = micro * bvals

print "b / μm =    ", impactParameterB / micro
print "b_min / μm =", minImpactParam(B_mag, EkinOverVmax, bOverLarmorR) / micro
```

    b / μm =     [   1.            2.15443469    4.64158883   10.           21.5443469    46.41588834  100.        ]
    b_min / μm = 3.90333173895

```python
""" (the following depend on the impact parameter) """
LarmorR = impactParameterB / bOverLarmorR  # (this version is defined positive)
V_perp = - LarmorR * Omega_e
L_intxn = tan_alpha * impactParameterB #[-1] * np.ones_like(impactParameterB)
V_z = 2 * L_intxn / (N_gyro * T_e)
T_intxn = 2 * L_intxn / V_z
```

```python
num_steps_per_gyro = 40
delta_Te = T_e / num_steps_per_gyro
print T_intxn / delta_Te
num_steps = int(np.floor(T_intxn[0] / delta_Te))
print "num_steps =", num_steps
```

    [ 4000.  4000.  4000.  4000.  4000.  4000.  4000.]
    num_steps = 4000

```python
""" initial condition for the ion --- just one, for now """
QVion = np.array([
          ( 0.0 * rho_gc,  0.000 * V_th, 0.0 * rho_gc,  0.000 * V_th, 0.0,  0.000 * V_th)
        ]).transpose()
Zion = QVtoZion.dot(QVion)

""" initial conditions for the electrons """
QVelec = np.asarray([
           np.zeros(impactParameterB.shape),
           V_perp,
           impactParameterB - LarmorR,
           np.zeros(impactParameterB.shape),
           - L_intxn,
           L_intxn * abs(Omega_e) / (pi * N_gyro)
        ])
Zelec = QVtoZe.dot(QVelec)

num_elec0 = Zelec.shape[1]
num_ion0 = Zion.shape[1]

"""
=== create arrays of electrons and ions ===
Here we arrange them so that we can pair each ion with each
electron, and compute the \Delta{p} for each interaction.
"""
ZI_elec = np.hstack([Zelec for e in range(num_ion0)])
ZI_ion = Zion[:, np.arange(num_ion0 * num_elec0) // num_elec0]

num_elec = ZI_elec.shape[1]
num_ion = ZI_ion.shape[1]
```

```python
ZF_i, ZF_e = apply_MapZ(delta_Te, num_steps, ZI_ion, ZI_elec)
```

```python
navg = 1 # number of gyrotron oscillations over which to average; set to 1, 2, 4, 5, or 10
ZI_elec_gc = toGuidingCenter(ZI_elec)
ZFgc_i, ZFgc_e = apply_MapZgc(navg * 40*delta_Te, num_steps//40 // navg, ZI_ion, ZI_elec_gc)
```

```python
deltaP_exp = np.array(ZF_i[-1, [Ipx, Ipy, Ipz], :] - ZF_i[0, [Ipx, Ipy, Ipz], :]).T
deltaP_avg = np.array(ZFgc_i[-1, [Ipx, Ipy, Ipz], :] - ZFgc_i[0, [Ipx, Ipy, Ipz], :]).T
deltaP_mgn = magnus_deltaPIon(T_intxn, ZI_ion, ZI_elec_gc)

print deltaP_exp
print
print deltaP_avg
print
print deltaP_mgn
```

    [[  1.08030578e-26   5.72808106e-26   1.45007842e-27]
     [  4.03458322e-28   1.66767979e-26   1.38486416e-27]
     [  2.76032939e-29   6.40599505e-27   9.05956974e-29]
     [  7.85620796e-31   1.58638251e-27   2.60010730e-30]
     [  1.75039675e-32   3.47567664e-28   5.80271684e-32]
     [  3.78408946e-34   7.50099341e-29   1.25468090e-33]
     [  8.15538018e-36   1.61631831e-29   2.70410697e-35]]

    [[  1.05836708e-26   5.67504416e-26   8.51768805e-28]
     [  4.00940540e-28   1.66262914e-26   1.38431451e-27]
     [  2.74641067e-29   6.39066436e-27   9.05285901e-29]
     [  7.81439841e-31   1.58238970e-27   2.59745819e-30]
     [  1.74099270e-32   3.46684860e-28   5.79654896e-32]
     [  3.76373889e-34   7.48192282e-29   1.25334141e-33]
     [  8.11151685e-36   1.61220857e-29   2.70121883e-35]]

    [[ -0.00000000e+00   1.61222627e-25  -1.01894869e-40]
     [ -0.00000000e+00   3.47345583e-26  -3.78365367e-41]
     [ -0.00000000e+00   7.48333795e-27  -0.00000000e+00]
     [ -0.00000000e+00   1.61223638e-27  -0.00000000e+00]
     [ -0.00000000e+00   3.47345800e-28  -0.00000000e+00]
     [ -0.00000000e+00   7.48333842e-29  -4.49591691e-44]
     [ -0.00000000e+00   1.61223639e-29  -0.00000000e+00]]

```python
dDeltaP_ax = deltaP_avg - deltaP_exp
dDeltaP_mx = deltaP_mgn - deltaP_exp

relErr_avg = (np.linalg.norm(dDeltaP_ax, axis = 1)
             / np.linalg.norm(deltaP_exp, axis = 1))
relErr_mgn = (np.linalg.norm(dDeltaP_mx, axis = 1)
             / np.linalg.norm(deltaP_exp, axis = 1))

eV_closest_approach = (invFourPiEps0 * Z_ion * qe ** 2 / impactParameterB)
E_kinetic_e = (me / 2) * (V_perp ** 2 + V_z ** 2)
eVcaOverEkin = eV_closest_approach / E_kinetic_e

llres = plt.figure()
plt.loglog(eVcaOverEkin, relErr_avg, 'bo')
plt.loglog(eVcaOverEkin, relErr_mgn, 'rs')
plt.annotate(s="Averging",xy=(1.e-1,1.e-3))
plt.annotate(s="Magnus",xy=(1.e-2,5.e-1))

plt.xlabel("$eV_{\mathrm{ca}} / E_{\mathrm{kin}}$")
plt.ylabel("relative error in $\Delta{P}_{\mathrm{ion}}$")

plt.show()

llres.savefig("/Users/dabell/Desktop/foo.pdf")
```

```python
mae_compare.savefig("/Users/dabell/RadiaSoft/MCool/MgnAvgExpCompare.pdf")
```


# Illustration of importance sampling for GHZ states

The script [PurityImportanceSampling](PurityImportanceSampling.py) shows how to access the purity $p_2=\mathrm{Tr}(\rho^2)$ based on importance sampling of randomized Measurements [link to the paper](https://arxiv.org/pdf/2102.13524.pdf).

### Description of the experimental situation
In this example we consider an experiment where we would like to implement a large $N$-qubit GHZ state
\begin{equation}
|\mathrm{GHZ}\rangle = \frac{1}{\sqrt{2}}\left(|0\dots 0\rangle + |1\dots 1\rangle \right)
\end{equation}
However, due to decoherence, we realize instead a mixed-state of the form
\begin{equation}
\rho = (1-p)|\mathrm{GHZ}\rangle \langle \mathrm{GHZ}| + p \mathbb{1}/2^N
\end{equation}
We would like here to measure the purity $p_2$ of the state $\rho$.

### Measurement protocol with importance sampling of randomized measurements
The script is based on randomized measurements. The purity can be expressed as an average over randomized measurements
\begin{equation}
p_2=\int X(u) du \quad \mathrm{with}  \quad X(u) = 2^N\sum_{s,s'} (-2)^{-D[s,s']} P_u(s) P_u(s')
\end{equation}
with $P_u(s)=\langle s |u \rho u^\dagger |s \rangle$ measurements in the computational basis.  Each CUE local random unitary  $u = \otimes_i u_i = \otimes_i R_y(\theta_i)R_z(\phi_i)$ is made of single qubit gates.

In order to reduce the required number of measuremens, we can use importance sampling and sample unitaries non-uniformly. As we know that the state $\rho$ is closed to the pure state state $|\mathrm{GHZ}\rangle$, which was targeted, we can define a probability distribution
\begin{equation}
p_\mathrm{IS}(u) = \frac{|X_\mathrm{GHZ}(u)|}{\int |X_\mathrm{GHZ}(u)|} du \quad \mathrm{with}
\quad
X_\mathrm{GHZ}(u)=2^N\sum_{s,s'} (-2)^{-D[s,s']} P^{\mathrm{GHZ}}_u(s) P^\mathrm{GHZ}_u(s')
\end{equation}
where the probabilities $P^\mathrm{GHZ}_u(s)=|\langle s u|\mathrm{GHZ}\rangle|^2$ can be calculated easily. Then the protocol consists first in sampling unitaries $u$ according to $p_\mathrm{IS}(u)$ from a Metropolis algorithm (before the experiment takes place).
where the probabilities $P_\mathrm{GHZ}(s)=|\langle s |u\mathrm{GHZ}\rangle|^2$ can be calculated easily. The protocol consists first then in sampling unitaries $u$ according to $p_\mathrm{IS}(u)$ from a Metropolis algorithm (before the experiment takes place).

These single qubit unitaries are then applied on the quantum computer on the state $\rho$. One measures for each unitary $u^{(r)}$ $N_M$ bistrings $\{s_m^{(r)}\}$ and access an estimation of $X(u)$ via
\begin{equation}
X_e(u^{(r)}) = \frac{2^N}{N_M(N_M-1)}
    \sum_{m \neq m'} (-2) ^{-D[s_m^{(r)},s^{(r)}_{m'}]}
\end{equation}
We obtain the purity of the state $\rho$ based on the formula
\begin{equation}
[p_2]_\mathrm{IS} = \frac{ 1}{N_s}
    \sum_{r = 1}^{N_u} \frac{n^{(r)} X_e(u^{(r)}) }{p_\mathrm{IS}(u^{(r)})}
\end{equation}
 where $N_s$ is the total number of unitaries $u$ collected by the metropolis algorithm which contains $N_u$ distinct ones and $n^{(r)}$ takes into account the occurence of each of unitary $u^{(r)}$ of the $N_u$ distinct samples and satisfies $N_s=\sum_r n^{(r)}$.

 ### Example
 with $N=16$, $N_M=10\times 2^{16}$, $N_u=100$, $p=0.25$, the script returns for three different simulated experiments

Run 1----

p2 (uniform sampling) =  0.6197747395756498

p2 (Importance sampling) =  0.5633602285126839

Error uniform:  10.18 %

Error IS:  0.15 %

Run 2----

p2 (uniform sampling) =  0.5684924515878843

p2 (Importance sampling) =  0.5665709762430637

Error uniform:  1.06 %

Error IS:  0.72 %

Run 3----

p2 (uniform sampling) =  0.3738835630369831

p2 (Importance sampling) =  0.5560702642126601

Error uniform:  33.53 %

Error IS:  1.14 %

```python

```