# Hankel Transforms on power spectra / correlation functions

The purpose of this notebook is to develop a robust way to do the hankel transforms of power-correlation functions. The general method has been established for a while (using Ogata's method a la the `hankel` package, but with a few optimizations), but some loose ends still need tying up -- eg. how best to do the necessary inter/extrapolation under different circumstances, and what limits to use on $k$ and $r$.

We have the formulas

\begin{align}
    \xi(r) &= \frac{1}{2\pi^2} \int_0^\infty P(k) k^2 \frac{\sin(k r)}{kr} dk \\
    P(k) &= 4\pi \int_0^\infty \xi(r) r^2 \frac{\sin(k r)}{kr} dr.
\end{align}

Note that for physical power spectra / correlation functions, these always converge. For linear power spectra, both ends go to zero, thus ensuring convergence. And that means by definition that $\xi(r)$ converges.

## The Transform Function

I'm going to copy-paste the Hankel Transform function from `halomod`, as it appears in this particular version (see below), to enable working on it more easily.

```python
import numpy as np
from scipy.interpolate import InterpolatedUnivariateSpline as spline
import warnings
```

```python
def pfunc_linear(logk, p, lnk, power_pos):
    spl = spline(lnk, p, k=3)
    result = np.zeros_like(logk)
    inner_mask = (lnk.min() <= logk) & (logk <= lnk.max())
    result[inner_mask] = spl(logk[inner_mask])

    lower_mask = logk < lnk.min()
    if power_pos[0]:
        result[lower_mask] = np.exp((np.log(p[1]) - np.log(p[0])) * (logk[lower_mask] - lnk[0]) / (
                lnk[1] - lnk[0]) + np.log(p[0]))
    else:
        result[lower_mask] = (p[1] - p[0]) * (logk[lower_mask] - lnk[0]) / (lnk[1] - lnk[0]) + p[0]

    upper_mask = logk > lnk.max()
    if power_pos[1]:
        if p[-1] <=0 or p[-2] <= 0:
            raise ValueError("Something went horrible wrong")
        result[upper_mask] = np.exp(
            (np.log(p[-1]) - np.log(p[-2])) * (logk[upper_mask] - lnk[-1]) / (
                lnk[-1] - lnk[-2]) + np.log(p[-1])
        )
    else:
        result[upper_mask] = (p[-1] - p[-2]) * (logk[upper_mask] - lnk[-1]) / (
            lnk[-1] - lnk[-2]) + p[-1]

    return result

def get_sumspace(h, x_max=None):
    # Optimal value of nmax, given h.
    nmax = int(3.2 / h)

    x_nmax = np.pi * nmax * np.tanh(np.pi*np.sinh(h*nmax))

    if not x_max or x_nmax <= x_max:
        roots = np.arange(1, nmax + 1)
        t = h * roots
        s = np.pi * np.sinh(t)
        x = np.pi * roots * np.tanh(s / 2)
    else:
        t = []
        s = []
        x = []
        chunksize = 100
        i = 1
        while (not x) or (np.max(x) < x_max):
            roots = np.arange(i, i + chunksize)
            t.append(h * roots)
            s.append(np.pi * np.sinh(t[-1]))
            x.append(np.pi * roots * np.tanh(s[-1] / 2))
            i += chunksize
        t = np.array(t).flatten()
        s = np.array(s).flatten()
        x = np.array(x).flatten()

    dpsi = 1 + np.cosh(s)
    dpsi[dpsi != 0] = (np.pi * t * np.cosh(t) + np.sinh(s)) / dpsi[dpsi != 0]

    return x, np.pi * np.sin(x) * dpsi * x

def power_to_corr_ogata(
    power: np.ndarray, k: np.ndarray, r: np.ndarray, h=0.005,
    power_pos=(True, True), rtol=1e-3, atol=1e-15, _reverse=False,
    integration_max=None
):
    """
    Convert a 3D power spectrum to a correlation function.

    Uses Ogata's method (Ogata 2005) for Hankel Transforms in 3D.

    Parameters
    ----------
    power : np.ndarray
        The power spectrum to convert -- either 1D or 2D. If 1D, it is the power as a
        function of k. If 2D, the first dimension should be ``len(r)``, and the power
        to integrate is considered to be different for each ``r``.
    k : np.ndarray
        Array of same length as the last axis of ``power``, giving the fourier-space
        co-ordinates.
    r : np.ndarray
        The real-space co-ordinates to which to transform.
    h : int, optional
        Controls the spacing of the intervals (note the intervals are not equispaced).
        Smaller numbers give smaller intervals.
    power_pos : tuple of bool, optional
        Whether 'power' is definitely positive, at either end. If so, a slightly better
        extrapolation can be achieved.
    Notes
    -----
    See the `hankel <https://hankel.readthedocs.io>`_ documentation for details on the
    implementation here. This particular function is restricted to a spherical transform.
    """
    v = 'kr'
    if _reverse:
        v = v[::-1]

    func = 'corr' if _reverse else "power"

    lnk = np.log(k)

    if integration_max:
        x_max = integration_max * r.max()
    else:
        x_max = None

    x, sumparts = get_sumspace(h, x_max)

    if power_pos[0] and not np.all(power.T[:2] > 0):
        power_pos = (False, power_pos[1])

    if power_pos[1] and not np.all(power.T[-2:] > 0):
        power_pos = (power_pos[0], False)

    out = np.zeros(len(r))

    warn_upper = True
    warn_lower = True
    warn_conv = True
    for ir, rr in enumerate(r):
        kk = x / rr

        summand = sumparts * pfunc(np.log(kk), power[ir] if power.ndim == 2 else power, lnk, power_pos)
        cumsum = np.cumsum(summand)

        if kk.min() < k.min() and warn_lower:
            warnings.warn(
                f"In hankel transform, {func} at {v[1]}={rr:.2e} was extrapolated to "
                f"{v[0]}={kk.min():.2e}. Minimum provided was {k.min():.2e}. "
                f"Lowest value required is {x.min() / r.max():.2e}",
                stacklevel=2
            )
            warn_lower = False

        # If all k values accessed weren't interpolated, just return it.
        if kk.max() <= k.max():
            out[ir] = cumsum[-1]
        else:
            # Check whether we have convergence at k.max
            indx = np.where(kk > k.max())[0][0]

            if np.isclose(cumsum[indx], cumsum[indx - 1], atol=atol, rtol=rtol):
                # If it converged in the non-extrapolated part, return that.
                out[ir] = cumsum[indx]
            else:
                # Otherwise, warn the user, and just return the full sum.
                if warn_upper:
                    warnings.warn(
                        f"In hankel transform, {func} at {v[1]}={rr:.2e} was extrapolated to "
                        f"{v[0]}={kk.max():.2e}. Maximum provided was {k.max():.2e}. ",
                        stacklevel=2
                    )
                    warn_upper = False

                if not np.isclose(cumsum[-1], cumsum[-2], atol=atol, rtol=rtol) and warn_conv:
                    warnings.warn(
                        f"Hankel transform of {func} did not converge for {v[1]}={rr:.2e}. "
                        f"It is likely that higher {v[1]} will also not converge. "
                        f"Absolute error estimate = {cumsum[-1] - cumsum[-2]:.2e}. "
                        f"Relative error estimate = {cumsum[-1]/cumsum[-2] - 1:.2e}",
                        stacklevel=2
                    )
                    warn_conv = False

                out[ir] = cumsum[-1]

    return out / (2 * np.pi ** 2 * r ** 3)

def corr_to_power_ogata(corr, r, k, h=0.005, power_pos=(True, False), atol=1e-15, rtol=1e-3, integration_max=None):
    """
    Convert an isotropic 3D correlation function to a power spectrum.

    Uses Ogata's method (Ogata 2005) for Hankel Transforms in 3D.

    Parameters
    ----------
    corr : np.ndarray
        The correlation function to convert as a function of ``r``.
    r : np.ndarray
        Array of same length as ``corr``, giving the real-space co-ordinates.
    k : np.ndarray
        The fourier-space co-ordinates to which to transform.
    n : int, optional
        The number of subdivisions in the integral.
    h : int, optional
        Controls the spacing of the intervals (note the intervals are not equispaced).
        Smaller numbers give smaller intervals.

    Notes
    -----
    See the `hankel <https://hankel.readthedocs.io>`_ documentation for details on the
    implementation here. This particular function is restricted to a spherical transform.
    """
    return 8 * np.pi ** 3 * power_to_corr_ogata(corr, r, k, h, power_pos=power_pos, atol=atol, rtol=rtol, _reverse=True, integration_max=integration_max)
```

## The Main Idea

The function as implemented here does the standard Ogata transform, but it is passed a simple vector of power/correlation, rather than a callable function. This means it has to interpolate that vector to create a callable function, and will often have to extrapolate, if the extent of the transformed variable is too large.

In this case, the function will try to extrapolate linearly in log-space, unless the function it's trying to extrapolate is negative (or zero) on that particular end, in which case it will do simple linear extrapolation.
The slope of this extrapolation is gotten simply from the two last points on the curve, which is probably prone to error, and should be investigated.

The idea then would be to enforce that both the underlying $k$ and $r$ vectors extend well into the regime in which they essentially become power-laws, so that this extrapolation is a reasonable thing to do. This should have the benefit that the transform is well-defined over any extent (i.e. the round-trip from $k$ to $r$ to $k$ again would work perfectly fine, even if the latter $k$ were extended).

The main things to investigate, then, are:

* Relevant limits on the $k$ vector that need to be observed to get accurate results (the user is able to choose $k$ in `hmf`, so we can't bank on it entirely, but it should already be broad enough that `sigma` is accurate, so we can probably count on it being reasonable).
* Relevant limits on the $r$ vector -- $r_{\rm min}$, $r_{\rm max}$ and $dr$.
* How to do the extrapolation most robustly.
* What the relevant value of $h$ is for power spectra and correlation functions

## Matter Correlation

```python
from halomod import DMHaloModel

import matplotlib.pyplot as plt
%matplotlib inline
```

```python
hm = DMHaloModel()
power = hm.power
```

```python
r = np.logspace(-3, 3, 100)
```

```python
pfunc = pfunc_linear
corr = power_to_corr_ogata(power, hm.k, r, h=0.005)
```

    /home/steven/miniconda3/envs/halomod/lib/python3.7/site-packages/ipykernel_launcher.py:1: UserWarning: In hankel transform, power at r=1.00e-03 was extrapolated to k=2.01e+06. Maximum provided was 1.95e+04.
      """Entry point for launching an IPython kernel.

```python
plt.plot(r, corr)
plt.xscale('log')
plt.yscale('symlog', linthreshy=1e-7)
```

Okay, so we should check the lower and upper bounds of $k$ and how they're performing under the extrapolation:

```python
hm.lnk_min
```

    -18.420680743952367

```python
corr_default = corr
```

```python
lnk_min_corrs = []
for lnk_min in [-30, -25, -20, -18.42, -15]:
    hm.lnk_min = lnk_min

    corr = power_to_corr_ogata(hm.power, hm.k, r, h=0.005)

    plt.plot(r, corr / corr_default - 1, label=str(lnk_min))

    lnk_min_corrs.append(corr)

plt.xscale('log')
plt.legend()
#plt.yscale('symlog', linthreshy=1e-7)
plt.ylim(-0.5, 0.5)
```

Good, so the default value is good to a percent or two over the whole range. Let's try lnk_max:

```python
hm.lnk_min = -18.42

lnk_max_corrs = []
for lnk_max in [5, 7, 9.9, 12, 18]:
    hm.lnk_max = lnk_max

    corr = power_to_corr_ogata(hm.power, hm.k, r, h=0.005)

    plt.plot(r, corr / corr_default - 1, label=str(lnk_max))

    lnk_max_corrs.append(corr)

plt.xscale('log')
plt.legend()
#plt.yscale('symlog', linthreshy=1e-7)
plt.ylim(-0.5, 0.5)
```

Again, we're good to a couple of percent for the default value.

Let's try dlnk.

```python
hm.dlnk
```

    0.05

```python
hm.lnk_max = 9.9

dlnk_corrs = []
for dlnk in [0.01, 0.1]:
    hm.dlnk = dlnk

    corr = power_to_corr_ogata(hm.power, hm.k, r, h=0.005)

    plt.plot(r, corr / corr_default - 1, label=str(dlnk))

    dlnk_corrs.append(corr)

plt.xscale('log')
plt.legend()
#plt.yscale('symlog', linthreshy=1e-7)
plt.ylim(-0.5, 0.5)
```

Going to dlnk=0.1 certainly seems to affect things here quite significantly (though only on scales we don't care about _directly_... we'll have to see how they affect the inverted power spectrum). Even at 0.05 (the default), it is affected at the 10% level, but again, only on fairly large scales.

Note that this could a resolution thing, or it could be a numerical instability in the extrapolation due to closer points.

Finally let's look at h.

```python
hm.dlnk = 0.05

h_corrs = []
for h in [0.0005, 0.001, 0.005, 0.01]:
    corr = power_to_corr_ogata(hm.power, hm.k, r, h=h)

    plt.plot(r, corr / corr_default - 1, label=str(h))

    dlnk_corrs.append(corr)

plt.xscale('log')
plt.legend()
#plt.yscale('symlog', linthreshy=1e-7)
plt.ylim(-0.5, 0.5)
```

Here we do see some systematic differences, but only at very small magnitude. It seems like 0.005 is good enough.

## Correlation to power.

Now that we're fairly confident we got the correlation OK, we need to get the power spectrum back. The weird things here are that the correlation function goes negative and has a negative power-law-ish thing going on. It's at a fairly low level, but not getting it right will lead to bad errors in the power spectrum at small $k$. Also, the correlation function is much more peaked at small $r$ and thus may need a smaller $h$.

```python
r_hires = np.logspace(-3, 3, 1000)
corr_hires = power_to_corr_ogata(hm.power, hm.k, r_hires, h=0.005)
```

    /home/steven/miniconda3/envs/halomod/lib/python3.7/site-packages/ipykernel_launcher.py:2: UserWarning: In hankel transform, power at r=1.00e-03 was extrapolated to k=2.01e+06. Maximum provided was 1.95e+04.

Our test will be to get back the power spectrum at the full range we input it. This is only possible if our extrapolation is reasonable. Unfortunately, it looks like there's a bit of a spike in the correlation that might make our extrapolation a bit crazy. We'll see.

```python
power_roundtrip = corr_to_power_ogata(corr_hires, r_hires, hm.k)
```

    /home/steven/miniconda3/envs/halomod/lib/python3.7/site-packages/ipykernel_launcher.py:166: UserWarning: In hankel transform, corr at k=1.00e-08 was extrapolated to r=2.01e+11. Maximum provided was 1.00e+03.
    /home/steven/miniconda3/envs/halomod/lib/python3.7/site-packages/ipykernel_launcher.py:166: UserWarning: In hankel transform, corr at k=2.53e+01 was extrapolated to r=9.76e-04. Minimum provided was 1.00e-03. Lowest value required is 1.26e-06

```python
plt.plot(hm.k, hm.power)
plt.plot(hm.k, power_roundtrip)
plt.xscale('log')
plt.yscale('log')
```

OK, it looks like the large-scale extrapolation did _not_ do well here. That's kinda understandable since we did a linear extrapolation instead of log. Let's try a log extrapolation. All we have to do is redefine the `pfunc`:

```python
def pfunc_log(logk, p, lnk, power_pos):
    spl = spline(lnk, p, k=3)
    result = np.zeros_like(logk)
    inner_mask = (lnk.min() <= logk) & (logk <= lnk.max())
    result[inner_mask] = spl(logk[inner_mask])

    lower_mask = logk < lnk.min()
    if power_pos[0]:
        result[lower_mask] = np.exp((np.log(p[1]) - np.log(p[0])) * (logk[lower_mask] - lnk[0]) / (
                lnk[1] - lnk[0]) + np.log(p[0]))
    else:
        # Do a -log extrapolation
        result[lower_mask] = -np.exp(
            (np.log(-p[1]) - np.log(-p[0])) * (logk[lower_mask] - lnk[0]) / (
                lnk[1] - lnk[0]) + np.log(-p[0])
        )

    upper_mask = logk > lnk.max()
    if power_pos[1]:
        result[upper_mask] = np.exp(
            (np.log(p[-1]) - np.log(p[-2])) * (logk[upper_mask] - lnk[-1]) / (
                lnk[-1] - lnk[-2]) + np.log(p[-1])
        )
    else:
        # Do a -log extrapolation
        result[upper_mask] = -np.exp(
            (np.log(-p[-1]) - np.log(-p[-2])) * (logk[upper_mask] - lnk[-1]) / (
                lnk[-1] - lnk[-2]) + np.log(-p[-1])
        )

    return result
```

First let's look at this extrapolated function:

```python
r_super_large = np.logspace(-5, 10, 1000)
corr_extended = pfunc_log(np.log(r_super_large), corr_hires, np.log(r_hires), (True, False))
```

```python
plt.plot(r_super_large, corr_extended)
plt.xscale('log')
plt.yscale('symlog', linthreshy=1e-8)
```

The problem here is that because of the blip at the end (or numerical noise?), the extrapolation super fails. One way around this may actually be to assume that stuff over $r=1000$ just doesn't contribute like ever, and call it zero. Let's try that:

```python
def pfunc_zero(logk, p, lnk, power_pos):
    spl = spline(lnk, p, k=3)
    result = np.zeros_like(logk)
    inner_mask = (lnk.min() <= logk) & (logk <= lnk.max())
    result[inner_mask] = spl(logk[inner_mask])

    lower_mask = logk < lnk.min()
    if power_pos[0]:
        result[lower_mask] = np.exp((np.log(p[1]) - np.log(p[0])) * (logk[lower_mask] - lnk[0]) / (
                lnk[1] - lnk[0]) + np.log(p[0]))
    else:
        # Do a -log extrapolation
        result[lower_mask] = 0

    upper_mask = logk > lnk.max()
    if power_pos[1]:
        result[upper_mask] = np.exp(
            (np.log(p[-1]) - np.log(p[-2])) * (logk[upper_mask] - lnk[-1]) / (
                lnk[-1] - lnk[-2]) + np.log(p[-1])
        )
    else:
        # Do a -log extrapolation
        result[upper_mask] = 0

    return result
```

```python
pfunc = pfunc_zero
power_roundtrip0 = corr_to_power_ogata(corr_hires, r_hires, hm.k)
```

    /home/steven/miniconda3/envs/halomod/lib/python3.7/site-packages/ipykernel_launcher.py:167: UserWarning: In hankel transform, corr at k=2.53e+01 was extrapolated to r=9.76e-04. Minimum provided was 1.00e-03. Lowest value required is 1.26e-06

```python
plt.plot(hm.k, hm.power)
plt.plot(hm.k, power_roundtrip0)
plt.xscale('log')
plt.yscale('log')
```

OK, this also didn't do well. What if we try to increase the range of $r$?

```python
pfunc = pfunc_log
r_hires_larger = np.logspace(-3, 4, 1000)
corr_hires_larger = power_to_corr_ogata(hm.power, hm.k, r_hires_larger)
```

    /home/steven/miniconda3/envs/halomod/lib/python3.7/site-packages/ipykernel_launcher.py:3: UserWarning: In hankel transform, power at r=1.00e-03 was extrapolated to k=2.01e+06. Maximum provided was 1.95e+04.
      This is separate from the ipykernel package so we can avoid doing imports until

```python
plt.plot(r_hires_larger, corr_hires_larger)
plt.xscale('log')
plt.yscale('symlog', linthreshy=1e-10)
```

```python
power_roundtrip_larger = corr_to_power_ogata(corr_hires_larger, r_hires_larger, hm.k)
```

    /home/steven/miniconda3/envs/halomod/lib/python3.7/site-packages/ipykernel_launcher.py:27: RuntimeWarning: invalid value encountered in log
    /home/steven/miniconda3/envs/halomod/lib/python3.7/site-packages/ipykernel_launcher.py:28: RuntimeWarning: invalid value encountered in log
    /home/steven/miniconda3/envs/halomod/lib/python3.7/site-packages/ipykernel_launcher.py:167: UserWarning: In hankel transform, corr at k=1.00e-08 was extrapolated to r=2.01e+11. Maximum provided was 1.00e+04.
    /home/steven/miniconda3/envs/halomod/lib/python3.7/site-packages/ipykernel_launcher.py:167: UserWarning: Hankel transform of corr did not converge for k=1.00e-08. It is likely that higher k will also not converge. Absolute error estimate = nan. Relative error estimate = nan
    /home/steven/miniconda3/envs/halomod/lib/python3.7/site-packages/ipykernel_launcher.py:167: UserWarning: In hankel transform, corr at k=2.53e+01 was extrapolated to r=9.76e-04. Minimum provided was 1.00e-03. Lowest value required is 1.26e-06

```python
plt.plot(hm.k, hm.power)
plt.plot(hm.k, power_roundtrip_larger)
plt.xscale('log')
plt.yscale('log')
```

OK, so something went really wrong here too. But, it really does look like the tail of the correlation function should be just a negative power-law. If we can somehow model that, we could check.

```python
from scipy.interpolate import UnivariateSpline
```

```python
UnivariateSpline?
```

```python
def pfunc_logspline(logk, p, lnk, power_pos):
    spl = spline(lnk, p, k=3)
    result = np.zeros_like(logk)
    inner_mask = (lnk.min() <= logk) & (logk <= lnk.max())
    result[inner_mask] = spl(logk[inner_mask])

    lower_mask = logk < lnk.min()
    if power_pos[0]:
        result[lower_mask] = np.exp((np.log(p[1]) - np.log(p[0])) * (logk[lower_mask] - lnk[0]) / (
                lnk[1] - lnk[0]) + np.log(p[0]))
    else:
        # Do a -log extrapolation
        result[lower_mask] = -np.exp(
            (np.log(-p[1]) - np.log(-p[0])) * (logk[lower_mask] - lnk[0]) / (
                lnk[1] - lnk[0]) + np.log(-p[0])
        )

    upper_mask = logk > lnk.max()
    if power_pos[1]:
        result[upper_mask] = np.exp(
            (np.log(p[-1]) - np.log(p[-2])) * (logk[upper_mask] - lnk[-1]) / (
                lnk[-1] - lnk[-2]) + np.log(p[-1])
        )
    else:
        indx = np.argmin(p)
        this_lnk = lnk[indx:]
        this_p = p[indx:]

        this_lnk = this_lnk[this_p < 0]
        this_p = this_p[this_p <0]
        logSpline = UnivariateSpline(this_lnk, np.log(-this_p), k=1)

        # Do a -log extrapolation
        result[upper_mask] = -np.exp(logSpline(logk[upper_mask]))

    return result
```

```python
pfunc = pfunc_logspline
power_roundtrip_larger = corr_to_power_ogata(corr_hires_larger, r_hires_larger, hm.k)
```

    /home/steven/miniconda3/envs/halomod/lib/python3.7/site-packages/ipykernel_launcher.py:167: UserWarning: In hankel transform, corr at k=3.65e-06 was extrapolated to r=5.50e+08. Maximum provided was 1.00e+04.
    /home/steven/miniconda3/envs/halomod/lib/python3.7/site-packages/ipykernel_launcher.py:167: UserWarning: In hankel transform, corr at k=2.53e+01 was extrapolated to r=9.76e-04. Minimum provided was 1.00e-03. Lowest value required is 1.26e-06

```python
plt.plot(hm.k, hm.power)
plt.plot(hm.k, power_roundtrip_larger)
plt.xscale('log')
plt.yscale('log')
```

```python
r_super_large = np.logspace(-5, 10, 1000)
corr_extended = pfunc_logspline(np.log(r_super_large), corr_hires, np.log(r_hires), (True, False))
```

```python
plt.plot(r_super_large, corr_extended)
plt.xscale('log')
plt.yscale('symlog', linthreshy=1e-10)
```

OK, this still didn't work somehow. Let's try something excessive.

```python
r_enormous = np.logspace(-3, 10, 2000)
corr_enormous = power_to_corr_ogata(hm.power, hm.k, r_enormous)
```

    /home/steven/miniconda3/envs/halomod/lib/python3.7/site-packages/ipykernel_launcher.py:2: UserWarning: In hankel transform, power at r=1.00e-03 was extrapolated to k=2.01e+06. Maximum provided was 1.95e+04.

    /home/steven/miniconda3/envs/halomod/lib/python3.7/site-packages/ipykernel_launcher.py:2: UserWarning: In hankel transform, power at r=2.50e+06 was extrapolated to k=9.89e-09. Minimum provided was 1.00e-08. Lowest value required is 2.47e-12

```python
plt.plot(r_super_large, corr_extended)
plt.plot(r_enormous, corr_enormous)
plt.xscale('log')
plt.yscale('symlog', linthreshy=1e-10)
```

```python
pfunc = pfunc_logspline
power_roundtrip_enormous = corr_to_power_ogata(corr_enormous, r_enormous, hm.k)
```

    /home/steven/miniconda3/envs/halomod/lib/python3.7/site-packages/ipykernel_launcher.py:167: UserWarning: In hankel transform, corr at k=2.53e+01 was extrapolated to r=9.76e-04. Minimum provided was 1.00e-03. Lowest value required is 1.26e-06

```python
plt.plot(hm.k, hm.power)
plt.plot(hm.k, power_roundtrip_enormous)
plt.xscale('log')
plt.yscale('log')
```

OK, still didn't work. Higher res?

```python
r_ginormous = np.logspace(-3, 10, 10000)
corr_ginormous = power_to_corr_ogata(hm.power, hm.k, r_ginormous)
```

    /home/steven/miniconda3/envs/halomod/lib/python3.7/site-packages/ipykernel_launcher.py:2: UserWarning: In hankel transform, power at r=1.00e-03 was extrapolated to k=2.01e+06. Maximum provided was 1.95e+04.

    /home/steven/miniconda3/envs/halomod/lib/python3.7/site-packages/ipykernel_launcher.py:2: UserWarning: In hankel transform, power at r=2.47e+06 was extrapolated to k=1.00e-08. Minimum provided was 1.00e-08. Lowest value required is 2.47e-12

```python
pfunc = pfunc_logspline
power_roundtrip_ginormous = corr_to_power_ogata(corr_ginormous, r_ginormous, hm.k)
```

    /home/steven/miniconda3/envs/halomod/lib/python3.7/site-packages/ipykernel_launcher.py:167: UserWarning: In hankel transform, corr at k=2.53e+01 was extrapolated to r=9.76e-04. Minimum provided was 1.00e-03. Lowest value required is 1.26e-06

```python
plt.plot(hm.k, hm.power)
plt.plot(hm.k, power_roundtrip_ginormous)
plt.xscale('log')
plt.yscale('log')
```

Nope. Looks like you just can't do it? Might just have to restrict the range of output $k$. Maybe we need to use the smaller $h$ to get these scales back...

```python
corr_ginormous_large_h = power_to_corr_ogata(hm.power, hm.k, r_ginormous, h=1e-4)
```

    /home/steven/miniconda3/envs/halomod/lib/python3.7/site-packages/ipykernel_launcher.py:1: UserWarning: In hankel transform, power at r=1.00e-03 was extrapolated to k=1.01e+08. Maximum provided was 1.95e+04.
      """Entry point for launching an IPython kernel.
    /home/steven/miniconda3/envs/halomod/lib/python3.7/site-packages/ipykernel_launcher.py:1: UserWarning: In hankel transform, power at r=4.95e+04 was extrapolated to k=9.98e-09. Minimum provided was 1.00e-08. Lowest value required is 4.93e-14
      """Entry point for launching an IPython kernel.

```python
power_roundtrip_ginormous_large_h = corr_to_power_ogata(corr_ginormous_large_h, r_ginormous, hm.k, h=1e-6)
power_roundtrip_ginormous_mid_h = corr_to_power_ogata(corr_ginormous_large_h, r_ginormous, hm.k, h=1e-5)

```

    /home/steven/miniconda3/envs/halomod/lib/python3.7/site-packages/ipykernel_launcher.py:167: UserWarning: In hankel transform, corr at k=5.13e-02 was extrapolated to r=9.62e-04. Minimum provided was 1.00e-03. Lowest value required is 2.53e-09

```python
plt.plot(hm.k, hm.power)
plt.plot(hm.k, power_roundtrip_ginormous_mid_h)
plt.plot(hm.k, power_roundtrip_ginormous_large_h)
plt.xscale('log')
plt.yscale('log')
```

OK, so this seems to be an issue with too low of a $h$ in the corr-to-power transform. Why?

Perhaps we can plot out the summand to see what's going on, at say $k=10^{-3}$.

```python

```

```python
fig, ax = plt.subplots(4, 1, figsize=(12, 9))

use_roots = True

for i, h in enumerate([0.005, 0.001, 0.0001, 0.00001]):
    x, summand_parts = get_sumspace(h)
    R = x / 1e-3

    fnc = pfunc(np.log(R), corr_ginormous, np.log(r_ginormous), (True, False))

    if not use_roots:
        R = np.arange(len(R))

    cumsum = np.cumsum(fnc)

    ax[i].scatter(R, summand_parts/summand_parts.max(), label='Extra Factors')
    ax[i].scatter(R, fnc/fnc.max(), label="Correlation")
    ax[i].scatter(R, (summand_parts * fnc) /( summand_parts*fnc).max(), label="Product")
    ax[i].scatter(R, cumsum / cumsum.max(), label='CumSum')

    if use_roots:
        ax[i].set_xscale('log')
ax[0].legend()
```

OK, so what this is saying is that the cumulative sum is most influenced by the peak in the product around $r = x/k = 20-30$. When $h$ is largeish (0.005), the smallest $x \sim 10 h$ is already too large to properly capture this peak. Unfortunately, it does not seem to be a peak in either the extra factors themselves, or the correlation function, just their product. Let's try zoomin in on this region.

```python
fig, ax = plt.subplots(4, 1, figsize=(12, 9))

use_roots = True

for i, h in enumerate([0.005, 0.001, 0.0001, 0.00001]):
    x, summand_parts = get_sumspace(h)
    R = x / 1e-3

    summand_parts = summand_parts [R < 150]
    R = R[R<150]

    fnc = pfunc(np.log(R), corr_ginormous, np.log(r_ginormous), (True, False))

    if not use_roots:
        R = np.arange(len(R))

    cumsum = np.cumsum(fnc)

    ax[i].scatter(R[R<150], summand_parts/summand_parts.max(), label='Extra Factors')
    ax[i].scatter(R, fnc/fnc.max(), label="Correlation")
    ax[i].scatter(R, (summand_parts * fnc) /( summand_parts*fnc).max(), label="Product")
    ax[i].scatter(R, cumsum / cumsum.max(), label='CumSum')

    if use_roots:
        ax[i].set_xscale('log')
ax[0].legend()
```

I think the point here is that we will never need scales beyond probably the BAO peak. We can test this theory out with an even smaller $k$:

```python
fig, ax = plt.subplots(5, 1, figsize=(12, 9))

use_roots = True
pfunc = pfunc_logspline

Rmax = 1000
k = 1e-5

for i, h in enumerate([1e-4, 1e-5, 1e-6, 1e-7, 1e-8]):
    x, summand_parts = get_sumspace(h, x_max=Rmax * k)
    R = x / k

    summand_parts = summand_parts
    R = R

    fnc = pfunc(np.log(R), corr_ginormous, np.log(r_ginormous), (True, False))

    if not use_roots:
        R = np.arange(len(R))

    cumsum = np.cumsum(fnc)

    print(h, cumsum[-1])

    ax[i].scatter(R, summand_parts/summand_parts.max(), label='Extra Factors')
    ax[i].scatter(R, fnc/fnc.max(), label="Correlation")
    ax[i].scatter(R, (summand_parts * fnc) /( summand_parts*fnc).max(), label="Product")
    ax[i].scatter(R, cumsum / cumsum.max(), label='CumSum')

    if use_roots:
        ax[i].set_xscale('log')
ax[0].legend()
```

In fact, the peak here seems to always be at about $r=30$, regardless of the $k$.

One weird thing is that the cumulative sum keeps getting bigger with decreasing h... I'm not sure I understand that. Let's try doing the actual computation with a very small $h$, and using the fact that we don't need the whole range of $x$.

This would take an awful long time under the current implementation -- we'll need to change it so that the range of $x$ is different for each $k$:

```python
def power_to_corr_ogata_each(
    power: np.ndarray, k: np.ndarray, r: np.ndarray, h=0.005,
    power_pos=(True, True), rtol=1e-3, atol=1e-15, _reverse=False,
    integration_max=None
):
    v = 'kr'
    if _reverse:
        v = v[::-1]

    func = 'corr' if _reverse else "power"

    lnk = np.log(k)

    if power_pos[0] and not np.all(power.T[:2] > 0):
        power_pos = (False, power_pos[1])

    if power_pos[1] and not np.all(power.T[-2:] > 0):
        power_pos = (power_pos[0], False)

    out = np.zeros(len(r))

    warn_upper = True
    warn_lower = True
    warn_conv = True
    for ir, rr in enumerate(r):
        if integration_max:
            x_max = integration_max * rr
        else:
            x_max = None

        x, sumparts = get_sumspace(h, x_max)

        kk = x / rr

        summand = sumparts * pfunc(np.log(kk), power[ir] if power.ndim == 2 else power, lnk, power_pos)
        cumsum = np.cumsum(summand)

        if kk.min() < k.min() and warn_lower:
            warnings.warn(
                f"In hankel transform, {func} at {v[1]}={rr:.2e} was extrapolated to "
                f"{v[0]}={kk.min():.2e}. Minimum provided was {k.min():.2e}. "
                f"Lowest value required is {x.min() / r.max():.2e}",
                stacklevel=2
            )
            warn_lower = False

        # If all k values accessed weren't interpolated, just return it.
        if kk.max() <= k.max():
            out[ir] = cumsum[-1]
        else:
            # Check whether we have convergence at k.max
            indx = np.where(kk > k.max())[0][0]

            if np.isclose(cumsum[indx], cumsum[indx - 1], atol=atol, rtol=rtol):
                # If it converged in the non-extrapolated part, return that.
                out[ir] = cumsum[indx]
            else:
                # Otherwise, warn the user, and just return the full sum.
                if warn_upper:
                    warnings.warn(
                        f"In hankel transform, {func} at {v[1]}={rr:.2e} was extrapolated to "
                        f"{v[0]}={kk.max():.2e}. Maximum provided was {k.max():.2e}. ",
                        stacklevel=2
                    )
                    warn_upper = False

                if not np.isclose(cumsum[-1], cumsum[-2], atol=atol, rtol=rtol) and warn_conv:
                    warnings.warn(
                        f"Hankel transform of {func} did not converge for {v[1]}={rr:.2e}. "
                        f"It is likely that higher {v[1]} will also not converge. "
                        f"Absolute error estimate = {cumsum[-1] - cumsum[-2]:.2e}. "
                        f"Relative error estimate = {cumsum[-1]/cumsum[-2] - 1:.2e}",
                        stacklevel=2
                    )
                    warn_conv = False

                out[ir] = cumsum[-1]

    return out / (2 * np.pi ** 2 * r ** 3)

def corr_to_power_ogata_each(corr, r, k, h=0.005, power_pos=(True, False), atol=1e-15, rtol=1e-3, integration_max=None):
    return 8 * np.pi ** 3 * power_to_corr_ogata_each(corr, r, k, h, power_pos=power_pos, atol=atol, rtol=rtol, _reverse=True, integration_max=integration_max)
```

```python
power_roundtrip_ginormous_h = corr_to_power_ogata_each(corr_ginormous_large_h, r_ginormous, hm.k, h=1e-5, integration_max=600)
```

    /home/steven/miniconda3/envs/halomod/lib/python3.7/site-packages/ipykernel_launcher.py:84: UserWarning: In hankel transform, corr at k=5.13e-02 was extrapolated to r=9.62e-04. Minimum provided was 1.00e-03. Lowest value required is 2.53e-09

```python
plt.plot(hm.k, hm.power)
plt.plot(hm.k, power_roundtrip_ginormous_mid_h)
plt.plot(hm.k, power_roundtrip_ginormous_large_h)
plt.plot(hm.k, power_roundtrip_ginormous_h)

plt.xscale('log')
plt.yscale('log')
```

It doesn't look like I will easily solve this. Perhaps the best course of action is to restrict the halo model range of $k$ to something more modest, and use $h=0.005$. What's the lowest $k$ we can trust for $h=0.005$?

```python
pfunc = pfunc_logspline
power_roundtrip0 = corr_to_power_ogata(corr_ginormous_large_h, r_ginormous, hm.k, h=0.001, integration_max=600)
```

    /home/steven/miniconda3/envs/halomod/lib/python3.7/site-packages/ipykernel_launcher.py:195: UserWarning: In hankel transform, corr at k=5.10e+00 was extrapolated to r=9.67e-04. Minimum provided was 1.00e-03. Lowest value required is 2.53e-07

```python
plt.plot(hm.k, hm.power)
plt.plot(hm.k, power_roundtrip0)

plt.xscale('log')
plt.yscale('log')
plt.ylim(1e3, 1e5)
```

```python
%timeit power_roundtrip0 = corr_to_power_ogata(corr_ginormous_large_h, r_ginormous, hm.k, h=0.001, integration_max=600)
```

    /home/steven/miniconda3/envs/halomod/lib/python3.7/site-packages/ipykernel_launcher.py:195: UserWarning: In hankel transform, corr at k=5.10e+00 was extrapolated to r=9.67e-04. Minimum provided was 1.00e-03. Lowest value required is 2.53e-07

    1.83 s ± 85.3 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

```python
%timeit power_roundtrip0 = corr_to_power_ogata(corr_ginormous_large_h, r_ginormous, hm.k, h=0.001)
```

    /home/steven/miniconda3/envs/halomod/lib/python3.7/site-packages/ipykernel_launcher.py:195: UserWarning: In hankel transform, corr at k=5.10e+00 was extrapolated to r=9.67e-04. Minimum provided was 1.00e-03. Lowest value required is 2.53e-07

    1.89 s ± 211 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

```python
plt.plot(hm.k,  np.exp(hm._unnormalised_lnT))
plt.plot(hm.k, hm.k**-1.5 * np.log(hm.k) * np.exp(hm._unnormalised_lnT)[-1] / (hm.k**-1.5 * np.log(hm.k))[-1])
plt.xscale('log')
plt.yscale('log')
```

## Use ExtendedSpline and updated function

```python
from halomod.tools import hankel_transform, ExtendedSpline, _zero
```

```python
hm = DMHaloModel()
```

```python
r_table = np.logspace(-3, 4, 300)
```

```python
corr = hankel_transform(hm.linear_power_fnc, r_table, 'r')
```

```python
corr_fnc = ExtendedSpline(r_table, corr, upper_func=_zero, lower_func = 'power_law', match_upper=False)
```

```python

plt.plot(r_table, corr_fnc(r_table))
plt.xscale('log')
plt.yscale('log')
```

```python
power = hankel_transform(corr_fnc, hm.k, 'k', h=0.0001)
```

```python
plt.plot(hm.k, hm.power)
plt.plot(hm.k, power)
plt.xscale('log')
plt.yscale('log')
```

Let's try making the low-k power simply the bias times the linear power...

```python
power[hm.k < 1e-2] =hm.power[hm.k < 1e-2]
```

```python
plt.plot(hm.k, power/hm.power -1)
#plt.plot(hm.k, power)
plt.xscale('log')
#plt.yscale('log')
```

```python
%%time

d_rmin = -3
d_rmax = 3.5
d_dr = 0.01
d_h_forward = 0.005
d_h_back = 0.0001
mask = hm.k > 1e-2

rmins = [-5, -4, -3]
rmaxs = [2.5, 3.5, 4.5]
drs =[0.005, 0.01, 0.05]
hfs = [0.0001, 0.001, 0.005]
hbs=[0.00001, 0.0001, 0.001]

def get_fb_power(rmin, rmax, dr, h_forward, h_back):
    r_table = 10**np.arange(rmin, rmax, dr)
    corr = hankel_transform(hm.linear_power_fnc, r_table, 'r', h= h_forward)
    corr_fnc = ExtendedSpline(r_table, corr, upper_func=_zero, lower_func = 'power_law', match_upper=False)
    return hankel_transform(corr_fnc, hm.k[mask], 'k', h=h_back)

rmin_power = []
for rmin in rmins:
    rmin_power.append(get_fb_power(rmin,d_rmax,d_dr, d_h_forward, d_h_back))

rmax_power = []
for rmax in rmaxs:
    rmax_power.append(get_fb_power(d_rmin,rmax,d_dr, d_h_forward, d_h_back))

dr_power = []
for dr in drs:
    dr_power.append(get_fb_power(d_rmin,d_rmax,dr, d_h_forward, d_h_back))

hf_power = []
for h_forward in hfs:
    hf_power.append(get_fb_power(d_rmin,d_rmax,d_dr, h_forward, d_h_back))

hb_power = []
for h_back in hbs:
    hb_power.append(get_fb_power(d_rmin,d_rmax,d_dr, d_h_forward, h_back))
```

    CPU times: user 2min 5s, sys: 87.4 ms, total: 2min 6s
    Wall time: 2min 6s

```python
for i, power in enumerate(rmin_power):
    plt.plot(hm.k[mask], power/hm.power[mask] -1, label=f'rmin={rmins[i]}')
    plt.xscale('log')
plt.legend()
```

```python
for i, power in enumerate(rmax_power):
    plt.plot(hm.k[mask], power/hm.power[mask] -1, label=f'rmax={rmaxs[i]}')
    plt.xscale('log')
plt.legend()
```

```python
for i, power in enumerate(dr_power):
    plt.plot(hm.k[mask], power/hm.power[mask] -1, label=f'dr={drs[i]}')
    plt.xscale('log')
plt.legend()
```

```python
for i, power in enumerate(hf_power):
    plt.plot(hm.k[mask], power/hm.power[mask] -1, label=f'hf={hfs[i]}')
    plt.xscale('log')
plt.legend()
```

```python
for i, power in enumerate(hb_power):
    plt.plot(hm.k[mask], power/hm.power[mask] -1, label=f'hb={hbs[i]}')
    plt.xscale('log')
plt.legend()
```

After all this, it looks like $h$ has to be decreased for the forward transform (probably for the large scales of the correlationo function). Let's run the whole thing again, but with hf=1e-4 as default:

```python
%%time

d_rmin = -3
d_rmax = 3.5
d_dr = 0.01
d_h_forward = 0.0001
d_h_back = 0.0001
mask = hm.k > 1e-2

rmins = [-5, -4, -3]
rmaxs = [2.5, 3.5, 4.5]
drs =[0.005, 0.01, 0.05]
hfs = [0.0001, 0.001, 0.005]
hbs=[0.00001, 0.0001, 0.001]

def get_fb_power(rmin, rmax, dr, h_forward, h_back):
    r_table = 10**np.arange(rmin, rmax, dr)
    corr = hankel_transform(hm.linear_power_fnc, r_table, 'r', h= h_forward)
    corr_fnc = ExtendedSpline(r_table, corr, upper_func=_zero, lower_func = 'power_law', match_upper=False)
    return hankel_transform(corr_fnc, hm.k[mask], 'k', h=h_back)

rmin_power = []
for rmin in rmins:
    rmin_power.append(get_fb_power(rmin,d_rmax,d_dr, d_h_forward, d_h_back))

rmax_power = []
for rmax in rmaxs:
    rmax_power.append(get_fb_power(d_rmin,rmax,d_dr, d_h_forward, d_h_back))

dr_power = []
for dr in drs:
    dr_power.append(get_fb_power(d_rmin,d_rmax,dr, d_h_forward, d_h_back))

hf_power = []
for h_forward in hfs:
    hf_power.append(get_fb_power(d_rmin,d_rmax,d_dr, h_forward, d_h_back))

hb_power = []
for h_back in hbs:
    hb_power.append(get_fb_power(d_rmin,d_rmax,d_dr, d_h_forward, h_back))
```

    CPU times: user 9min 37s, sys: 562 ms, total: 9min 38s
    Wall time: 9min 40s

```python
for i, power in enumerate(rmin_power):
    plt.plot(hm.k[mask], power/hm.power[mask] -1, label=f'rmin={rmins[i]}')
    plt.xscale('log')
plt.legend()
```

```python
for i, power in enumerate(rmax_power):
    plt.plot(hm.k[mask], power/hm.power[mask] -1, label=f'rmax={rmaxs[i]}')
    plt.xscale('log')
plt.legend()
```

```python
for i, power in enumerate(dr_power):
    plt.plot(hm.k[mask], power/hm.power[mask] -1, label=f'dr={drs[i]}')
    plt.xscale('log')
plt.legend()
```

```python
for i, power in enumerate(hb_power):
    plt.plot(hm.k[mask], power/hm.power[mask] -1, label=f'hb={hbs[i]}')
    plt.xscale('log')
plt.legend()
```

From this, it looks like to get <1% errors on the round-trip, we need $h=10^{-4}$ for the forward transform, but we can use $h \sim 10^{-3}$ for the backwards one.

We should also time this method versus doing a standard quadrature rule:

```python
from scipy.integrate import quad
```

```python
def hankel_scipy(f, k, trns_var):
    intg = lambda x: f(x) * x

    out = np.zeros_like(k)
    for i, kk in enumerate(k):
        out[i] = quad(intg, 1e-6 / kk, 1e4 / kk, weight='sin', wvar=kk, points=[0], limit=500)[0]

    if trns_var == 'r':
        return out / k/ 2*np.pi**2
    else:
        return 4 * np.pi * out / k
```

```python
def get_fb_power_sp(rmin, rmax, dr):
    r_table = 10**np.arange(rmin, rmax, dr)
    corr = hankel_scipy(hm.linear_power_fnc, r_table, 'r')
    corr_fnc = ExtendedSpline(r_table, corr, upper_func=_zero, lower_func = 'power_law', match_upper=False)
    return hankel_scipy(corr_fnc, hm.k[mask], 'k'), corr
```

```python
sp,sp_corr = get_fb_power_sp(-3, 4, 0.05)
```

```python
sp_corr
```

    array([ 8.69764552e+04,  8.01515522e+04,  7.38613954e+04,  6.80641693e+04,
            6.27211542e+04,  5.77968265e+04,  5.32583531e+04,  4.90754833e+04,
            4.52203877e+04,  4.16673633e+04,  3.83927459e+04,  3.53747141e+04,
            3.25931706e+04,  3.00295730e+04,  2.76668630e+04,  2.54892775e+04,
            2.34823237e+04,  2.16326326e+04,  1.99278774e+04,  1.83566967e+04,
            1.69086339e+04,  1.55740370e+04,  1.43440112e+04,  1.32103707e+04,
            1.21655582e+04,  1.12026141e+04,  1.03151221e+04,  9.49717333e+03,
            8.74331539e+03,  8.04852638e+03,  7.40818171e+03,  6.81801272e+03,
            6.27408568e+03,  5.77278017e+03,  5.31075642e+03,  4.88493372e+03,
            4.49247880e+03,  4.13077589e+03,  3.79741383e+03,  3.49017432e+03,
            3.20700975e+03,  2.94603291e+03,  2.70550552e+03,  2.48382691e+03,
            2.27951872e+03,  2.09122026e+03,  1.91767767e+03,  1.75773512e+03,
            1.61032693e+03,  1.47447197e+03,  1.34926540e+03,  1.23387366e+03,
            1.12752846e+03,  1.02952269e+03,  9.39204058e+02,  8.55972518e+02,
            7.79275173e+02,  7.08603360e+02,  6.43488815e+02,  5.83501106e+02,
            5.28244756e+02,  4.77356804e+02,  4.30504527e+02,  3.87383715e+02,
            3.47716601e+02,  3.11250521e+02,  2.77756469e+02,  2.47027864e+02,
            2.18879158e+02,  1.93144684e+02,  1.69677045e+02,  1.48345093e+02,
            1.29031367e+02,  1.11628520e+02,  9.60347899e+01,  8.21486915e+01,
            6.98636552e+01,  5.90640016e+01,  4.96244391e+01,  4.14147974e+01,
            3.43095318e+01,  2.81971697e+01,  2.29815642e+01,  1.85712197e+01,
            1.48687381e+01,  1.17804013e+01,  9.23423249e+00,  7.16281133e+00,
            5.48263136e+00,  4.13753913e+00,  3.08227243e+00,  2.25437351e+00,
            1.62322969e+00,  1.14129488e+00,  7.80628654e-01,  5.19411546e-01,
            3.34818777e-01,  2.00595793e-01,  1.01720030e-01,  7.85375979e-02,
            1.67722784e-01,  7.33921978e-02, -1.74998879e-02, -3.25665440e-02,
           -2.74652045e-02, -2.00340167e-02, -1.45000922e-02, -1.08451390e-02,
           -7.08699485e-03, -5.72383200e-03, -3.54202380e-03, -1.37415096e-03,
           -1.39096996e-03, -1.06505596e-03, -5.02001332e-04, -3.10903194e-04,
           -2.91818465e-04, -1.42518288e-04, -7.74697777e-05, -7.86306994e-05,
           -7.10652926e-05,  4.53088493e-06, -1.13673820e-07, -5.13238918e-06,
           -6.43444846e-06,  1.56931671e-07,  5.97133685e-06,  8.82160729e-06,
            6.77282732e-06,  7.29473181e-06,  6.96767208e-06,  6.46206704e-06,
            5.87001816e-06,  5.46606402e-06,  5.16155051e-06,  4.73720919e-06,
            4.39794298e-06,  4.03457230e-06,  3.74692221e-06,  3.44370952e-06])

```python
plt.plot(hm.k[mask], sp/hm.power[mask] - 1)
plt.xscale('log')
```

Nope, not working.


```python
%pylab inline
```

    Populating the interactive namespace from numpy and matplotlib

```python
%%html
<link href="https://fonts.googleapis.com/css?family=Open+Sans" rel="stylesheet">
<style>#notebook-container{font-size: 13pt;font-family:'Open Sans', sans-serif;} div.text_cell{max-width: 104ex;}</style>
```

<link href="https://fonts.googleapis.com/css?family=Open+Sans" rel="stylesheet">
<style>#notebook-container{font-size: 13pt;font-family:'Open Sans', sans-serif;} div.text_cell{max-width: 104ex;}</style>

# Complex Numbers

## Definition

A _complex number_ $z$ is represented with an expression of the form $a+bi$:
* There is the _real_ part ($\mathrm{Re}\ z$) which is $a$.
* There is the _imaginary_ part ($\mathrm{Im}\ z$) which is $b$.

Where $i$ is defined as $i^2=-1$.

### Sum of complex numbers

Rules for addition and subtraction with complex numbers:

$$\tag{1}(a+bi)+(c+di)=(a+c)+(b+d)i$$
$$\tag{2}(a+bi)-(c+di)=(a-c)+(b-d)i$$

### Product of complex numbers

Rules for multiplication with complex numbers:

$$\begin{align}(a+bi)(c+di)&=ac+bdi+bci+bdi^2\\\tag{3}&=(ac-bd)+(ad+bc)i\end{align}$$

### Complex conjugate

The _conjugate_ for a complex number $z$ is defined as:

$$\tag{4}\bar{z}=(a-bi)$$

#### Properties of conjugates

Conjugates adhere to the following properties:

1. $\bar{z+w}=\bar{z}+\bar{w}$
2. $\bar{zw}=\bar{z}\bar{w}$
3. $\bar{z^n}=\bar{z}^n$

### Modulus

The _modulus_ or _absolute value_ for a complex number $z$ can be interpreted as the distance from the origin. It is defined as:

$$\tag{5}|\ z\ |\ = \sqrt{a^2+b^2}$$

An important thing to notice is that:

$$\begin{align}z\bar{z} &= (a+bi)(a-bi) \\ &= a^2+abi-abi-b^2i^2 \\ &= a^2 + b^2 \end{align}$$

Therefore, $z\bar{z} = |\ z\ |^2$.

### Division of complex numbers

Because $|\ z\ |^2 \in \mathbb{R}$, we can explain why division works. It is similar to rationalizing the denominator with roots:

$$\tag{6}\dfrac{z}{w} = \dfrac{z\bar{w}}{w\bar{w}} = \dfrac{z\bar{w}}{|\ w\ |^2}$$

### Roots of complex numbers

Since $i^2=-1$, we have $i=\sqrt{-1}$. But also $(-i)^2=i^2=-1$, here $-i$ is also a square root of $-1$. We call $i=\sqrt{-1}$ if $c \in \mathbb{N}^+$ the _principal square root of $-1$_, and we write this as:

$$\tag{7}\sqrt{-c}=\sqrt{c}i$$

### Python implementation

Now that we have defined the behaviour of a complex number, we can implement it in Python:

```python
class C():
    def __init__(self, real, imag):
        self.real = real
        self.imag = imag

    def __repr__(self):
        return '{} {} {}i'.format(self.real, '+' if self.imag >= 0 else '-', abs(self.imag))

    def __add__(self, z):
        if not isinstance(z, C): return TypeError('z must be C.')
        return C(self.real + z.real, self.imag + z.imag)

    def __sub__(self, z):
        if not isinstance(z, C): return TypeError('z must be C.')
        return C(self.real - z.real, self.imag - z.imag)

    def __mul__(self, z):
        if not isinstance(z, C): return TypeError('z must be C.')
        return C(self.real * z.real - self.imag * z.imag,
                 self.real * z.imag + self.imag * z.real)

    def __truediv__(self, z):
        if not isinstance(z, C): return TypeError('z must be C.')
        return C((self.real * z.real + self.imag * z.imag) / (z.real**2 + z.imag**2),
                 (self.imag * z.real - self.real * z.imag) / (z.real**2 + z.imag**2))

    def __eq__(self, z):
        if isinstance(z, C):
            return self.real is z.real and self.imag is z.imag
        return False

    def __ne__(self, z):
        if not isinstance(z, C): return TypeError('z must be C.')
        return not self == z

    def __lt__(self, z):
        if not isinstance(z, C): return TypeError('z must be C.')
        return self.mod() < z.mod()

    def __gt__(self, z):
        if not isinstance(z, C): return TypeError('z must be C.')
        return self.mod() > z.mod()

    def mod(self):
        return sqrt(self.real**2 + self.imag**2)

    def conjugate(self):
        return C(self.real, -self.imag)
```

```python
z = C(2, 3)
w = C(5, -6)
```

```python
print('z = {}'.format(z))
print('w = {}'.format(w))
```

    z = 2 + 3i
    w = 5 - 6i

```python
z+w
```

    7 - 3i

```python
z-w
```

    -3 + 9i

```python
z==w
```

    False

```python
z!=w
```

    True

```python
z*w
```

    28 + 3i

```python
C(5,2) * C(5,-2)
```

    29 + 0i

```python
(z-w)
```

    -3 + 9i

```python
z/w
```

    -0.13114754098360656 + 0.4426229508196721i

```python
C(2,1)/C(4,2)
```

    0.5 + 0.0i

```python
C(1,1)/C(1,1)
```

    1.0 + 0.0i

```python
C(3,4).mod()
```

    5.0

```python
C(1,1) > C(1,1.001)
```

    False

```python
C(1,1) < C(1,1.001)
```

    True

```python
C(1,1).conjugate()
```

    1 - 1i

```python
C(1,1) + 5
```

    TypeError('z must be C.')

### Quadratic formula with complex numbers

Now that we can take the square root of negative numbers, we can implement the quadratic formula which also handles a negative discriminant.

$$x = \dfrac{-b\pm\sqrt{b^2-4ac}}{2a}$$

Where the discriminant $d$ defined as $b^2-4ac$; gives the following three cases:

1. $d>0$: There are two solutions $x_1, x_2 \in \mathbb{R}$.
2. $d=0$: There is one solution $x \in \mathbb{R}$.
3. $d<0$: There are two solutions $x_1, x_2 \in \mathbb{C}$.

```python
def quadratic_formula(a, b, c):
    if a == 0: return ValueError('a=0 does not yield a second-degree polynomial.')
    d = b**2 - 4*a*c
    if d > 0:
        x1 = (-b + math.sqrt(d)) / (2 * a)
        x2 = (-b - math.sqrt(d)) / (2 * a)
        return (x1, x2)
    elif d == 0:
        x = (-b) / (2 * a)
        return x
    else:
        x1 = C((-b) / (2 * a),   math.sqrt(-d)  / (2 * a))
        x2 = C((-b) / (2 * a), (-math.sqrt(-d)) / (2 * a))
        return (x1, x2)
```

```python
quadratic_formula(0, 1, 2)
```

    ValueError('a=0 does not yield a second-degree polynomial.')

```python
quadratic_formula(1, 1, 1)
```

    (-0.5 + 0.8660254037844386i, -0.5 - 0.8660254037844386i)

```python
quadratic_formula(2, 2, 0)
```

    (0.0, -1.0)

## Polar Coordinates

Any complex number $z = a+bi$ can be represented by polar coordinates ($r$, $\theta$) with $r\geq 0$.

$$\begin{cases}a = r \cos \theta \\ b = r \sin \theta\end{cases}$$

Thus we can write any complex number $z$ in the form

$$\tag{8}z = r(\cos \theta + i \sin \theta)$$

where $r=|\ z\ |\ =\sqrt{a^2+b^2}$ and $\tan\theta=b/a$.

The angle $\theta$ is called the _argument_ of $z$ and we write $\theta=\mathrm{arg}(z)
$. Note that $\mathrm{arg}(z)$ is not unique; it differs by an integer multiple of $2\pi$.

```python

```


# Understanding the FFT Algorithm

*This notebook first appeared as a post by Jake Vanderplas on [Pythonic Perambulations](http://jakevdp.github.io/blog/2013/08/28/understanding-the-fft/).  The notebook content is BSD-licensed.*

<!-- PELICAN_BEGIN_SUMMARY -->
The Fast Fourier Transform (FFT) is one of the most important algorithms in signal processing and data analysis. I've used it for years, but having no formal computer science background, It occurred to me this week that I've never thought to ask *how* the FFT computes the discrete Fourier transform so quickly.  I dusted off an old algorithms book and looked into it, and enjoyed reading about the deceptively simple computational trick that JW Cooley and John Tukey outlined in their classic [1965 paper](http://www.ams.org/journals/mcom/1965-19-090/S0025-5718-1965-0178586-1/) introducing the subject.

The goal of this post is to dive into the Cooley-Tukey FFT algorithm, explaining the symmetries that lead to it, and to show some straightforward Python implementations putting the theory into practice.  My hope is that this exploration will give data scientists like myself a more complete picture of what's going on in the background of the algorithms we use.
<!-- PELICAN_END_SUMMARY -->

## The Discrete Fourier Transform

The FFT is a fast, $\mathcal{O}[N\log N]$ algorithm to compute the Discrete Fourier Transform (DFT), which
naively is an $\mathcal{O}[N^2]$ computation.  The DFT, like the more familiar continuous version of the Fourier transform, has a forward and inverse form which are defined as follows:

**Forward Discrete Fourier Transform (DFT):**
$$X_k = \sum_{n=0}^{N-1} x_n \cdot e^{-i~2\pi~k~n~/~N}$$

**Inverse Discrete Fourier Transform (IDFT):**
$$x_n = \frac{1}{N}\sum_{k=0}^{N-1} X_k e^{i~2\pi~k~n~/~N}$$

The transformation from $x_n \to X_k$ is a translation from configuration space to frequency space, and can be very useful in both exploring the power spectrum of a signal, and also for transforming certain problems for more efficient computation.  For some examples of this in action, you can check out Chapter 10 of our upcoming Astronomy/Statistics book, with figures and Python source code available [here](http://www.astroml.org/book_figures/chapter10/).  For an example of the FFT being used to simplify an otherwise difficult differential equation integration, see my post on [Solving the Schrodinger Equation in Python](http://jakevdp.github.io/blog/2012/09/05/quantum-python/).

Because of the importance of the FFT in so many fields, Python contains many standard tools and wrappers to compute this.  Both NumPy and SciPy have wrappers of the extremely well-tested FFTPACK library, found in the submodules ``numpy.fft`` and ``scipy.fftpack`` respectively.  The fastest FFT I am aware of is in the [FFTW](http://www.fftw.org/) package, which is also available in Python via the [PyFFTW](https://pypi.python.org/pypi/pyFFTW) package.

For the moment, though, let's leave these implementations aside and ask how we might compute the FFT in Python from scratch.

## Computing the Discrete Fourier Transform

For simplicity, we'll concern ourself only with the forward transform, as the inverse transform can be implemented in a very similar manner.  Taking a look at the DFT expression above, we see that it is nothing more than a straightforward linear operation: a matrix-vector multiplication of $\vec{x}$,

$$\vec{X} = M \cdot \vec{x}$$

with the matrix $M$ given by

$$M_{kn} = e^{-i~2\pi~k~n~/~N}.$$

With this in mind, we can compute the DFT using simple matrix multiplication as follows:

```python
import numpy as np
def DFT_slow(x):
    """Compute the discrete Fourier Transform of the 1D array x"""
    x = np.asarray(x, dtype=float)
    N = x.shape[0]
    n = np.arange(N)
    k = n.reshape((N, 1))
    M = np.exp(-2j * np.pi * k * n / N)
    return np.dot(M, x)
```

We can double-check the result by comparing to numpy's built-in FFT function:

```python
x = np.random.random(16)
np.allclose(DFT_slow(x), np.fft.fft(x))
```

    True

Just to confirm the sluggishness of our algorithm, we can compare the execution times
of these two approaches:

```python
%timeit DFT_slow(x)
%timeit np.fft.fft(x)
```

    62.5 µs ± 3.27 µs per loop (mean ± std. dev. of 7 runs, 10000 loops each)
    3.72 µs ± 296 ns per loop (mean ± std. dev. of 7 runs, 100000 loops each)

We are over 1000 times slower, which is to be expected for such a simplistic implementation.  But that's not the worst of it.  For an input vector of length $N$, the FFT algorithm scales as $\mathcal{O}[N\log N]$, while our slow algorithm scales as $\mathcal{O}[N^2]$.  That means that for $N=10^6$ elements, we'd expect the FFT to complete in somewhere around 50 ms, while our slow algorithm would take nearly 20 hours!

So how does the FFT accomplish this speedup?  The answer lies in exploiting symmetry.

## Symmetries in the Discrete Fourier Transform

One of the most important tools in the belt of an algorithm-builder is to exploit symmetries of a problem.  If you can show analytically that one piece of a problem is simply related to another, you can compute the subresult
only once and save that computational cost.  Cooley and Tukey used exactly this approach in deriving the FFT.

We'll start by asking what the value of $X_{N+k}$ is. From our above expression:

$$
\begin{align*}
X_{N + k} &=  \sum_{n=0}^{N-1} x_n \cdot e^{-i~2\pi~(N + k)~n~/~N}\\
          &= \sum_{n=0}^{N-1} x_n \cdot e^{- i~2\pi~n} \cdot e^{-i~2\pi~k~n~/~N}\\
          &= \sum_{n=0}^{N-1} x_n \cdot e^{-i~2\pi~k~n~/~N}
\end{align*}
$$

where we've used the identity $\exp[2\pi~i~n] = 1$ which holds for any integer $n$.

The last line shows a nice symmetry property of the DFT:

$$X_{N+k} = X_k.$$

By a simple extension,

$$X_{(k + i) \cdot N} = X_k$$

for any integer $i$.  As we'll see below, this symmetry can be exploited to compute the DFT much more quickly.

## DFT to FFT: Exploiting Symmetry

Cooley and Tukey showed that it's possible to divide the DFT computation into two smaller parts.  From
the definition of the DFT we have:

$$
\begin{align}
X_k &= \sum_{n=0}^{N-1} x_n \cdot e^{-i~2\pi~k~n~/~N} \\
    &= \sum_{m=0}^{N/2 - 1} x_{2m} \cdot e^{-i~2\pi~k~(2m)~/~N} + \sum_{m=0}^{N/2 - 1} x_{2m + 1} \cdot e^{-i~2\pi~k~(2m + 1)~/~N} \\
    &= \sum_{m=0}^{N/2 - 1} x_{2m} \cdot e^{-i~2\pi~k~m~/~(N/2)} + e^{-i~2\pi~k~/~N} \sum_{m=0}^{N/2 - 1} x_{2m + 1} \cdot e^{-i~2\pi~k~m~/~(N/2)}
\end{align}
$$

We've split the single Discrete Fourier transform into two terms which themselves look very similar to smaller Discrete Fourier Transforms, one on the odd-numbered values, and one on the even-numbered values. So far, however, we haven't saved any computational cycles.  Each term consists of $(N/2)*N$ computations, for a total of $N^2$.

The trick comes in making use of symmetries in each of these terms.  Because the range of $k$ is $0 \le k < N$, while the range of $n$ is $0 \le n < M \equiv N/2$, we see from the symmetry properties above that we need only perform half the computations for each sub-problem.  Our $\mathcal{O}[N^2]$ computation has become $\mathcal{O}[M^2]$, with $M$ half the size of $N$.

But there's no reason to stop there: as long as our smaller Fourier transforms have an even-valued $M$, we can reapply this divide-and-conquer approach, halving the computational cost each time, until our arrays are small enough that the strategy is no longer beneficial.  In the asymptotic limit, this recursive approach scales as $\mathcal{O}[N\log N]$.

This recursive algorithm can be implemented very quickly in Python, falling-back on our slow DFT code when the size of the sub-problem becomes suitably small:

```python
def FFT(x):
    """A recursive implementation of the 1D Cooley-Tukey FFT"""
    x = np.asarray(x, dtype=float)
    N = x.shape[0]

    if N % 2 > 0:
        raise ValueError("size of x must be a power of 2")
    elif N <= 32:  # this cutoff should be optimized
        return DFT_slow(x)
    else:
        X_even = FFT(x[::2])
        X_odd = FFT(x[1::2])
        factor = np.exp(-2j * np.pi * np.arange(N) / N)
        return np.concatenate([X_even + factor[:int(N / 2)] * X_odd,
                               X_even + factor[int(N / 2):] * X_odd])
```

Here we'll do a quick check that our algorithm produces the correct result:

```python
x = np.random.random(1024)
np.allclose(FFT(x), np.fft.fft(x))
```

    True

And we'll time this algorithm against our slow version:

```python
%timeit DFT_slow(x)
%timeit FFT(x)
%timeit np.fft.fft(x)
```

    73.4 ms ± 108 µs per loop (mean ± std. dev. of 7 runs, 10 loops each)
    5.13 ms ± 196 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)
    3.88 µs ± 144 ns per loop (mean ± std. dev. of 7 runs, 100000 loops each)

Our calculation is faster than the naive version by over an order of magnitude!  What's more, our recursive algorithm is asymptotically $\mathcal{O}[N\log N]$: we've implemented the Fast Fourier Transform.

Note that we still haven't come close to the speed of the built-in FFT algorithm in numpy, and this is to be expected.  The FFTPACK algorithm behind numpy's ``fft`` is a Fortran implementation which has received years of tweaks and optimizations.  Furthermore, our NumPy solution involves both Python-stack recursions and the allocation of many temporary arrays, which adds significant computation time.

A good strategy to speed up code when working with Python/NumPy is to vectorize repeated computations where possible.  We can do this, and in the process remove our recursive function calls, and make our Python FFT even more efficient.

## Vectorized Numpy Version

Notice that in the above recursive FFT implementation, at the lowest recursion level we perform $N~/~32$ identical matrix-vector products.  The efficiency of our algorithm would benefit by computing these matrix-vector products all at once as a single matrix-matrix product.  At each subsequent level of recursion, we also perform duplicate operations which can be vectorized.  NumPy excels at this sort of operation, and we can make use of that fact to create this vectorized version of the Fast Fourier Transform:

```python
def FFT_vectorized(x):
    """A vectorized, non-recursive version of the Cooley-Tukey FFT"""
    x = np.asarray(x, dtype=float)
    N = x.shape[0]

    if np.log2(N) % 1 > 0:
        raise ValueError("size of x must be a power of 2")

    # N_min here is equivalent to the stopping condition above,
    # and should be a power of 2
    N_min = min(N, 32)

    # Perform an O[N^2] DFT on all length-N_min sub-problems at once
    n = np.arange(N_min)
    k = n[:, None]
    M = np.exp(-2j * np.pi * n * k / N_min)
    X = np.dot(M, x.reshape((N_min, -1)))

    # build-up each level of the recursive calculation all at once
    while X.shape[0] < N:
        X_even = X[:, :int(X.shape[1] / 2)]
        X_odd = X[:, int(X.shape[1] / 2):]
        factor = np.exp(-1j * np.pi * np.arange(X.shape[0])/ X.shape[0])[:, None]
        X = np.vstack([X_even + factor * X_odd,
                       X_even - factor * X_odd])

    return X.ravel()
```

Though the algorithm is a bit more opaque, it is simply a rearrangement of the operations used in the recursive version with one exception: we exploit a symmetry in the ``factor`` computation and construct only half of the array.  Again, we'll confirm that our function yields the correct result:

```python
x = np.random.random(2048)
np.allclose(FFT_vectorized(x), np.fft.fft(x))
```

    True

Because our algorithms are becoming much more efficient, we can use a larger array to compare the timings,
leaving out ``DFT_slow``:

```python
x = np.random.random(1024 * 16)
%timeit FFT(x)
%timeit FFT_vectorized(x)
%timeit np.fft.fft(x)
```

    101 ms ± 799 µs per loop (mean ± std. dev. of 7 runs, 10 loops each)
    2.29 ms ± 29.9 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)
    256 µs ± 3.05 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)

We've improved our implementation by another order of magnitude!  We're now within about a factor of 10 of the FFTPACK benchmark, using only a couple dozen lines of pure Python + NumPy.  Though it's still no match computationally speaking, readibility-wise the Python version is far superior to the FFTPACK source, which you can browse [here](http://www.netlib.org/fftpack/fft.c).

So how does FFTPACK attain this last bit of speedup?  Well, mainly it's just a matter of detailed bookkeeping.  FFTPACK spends a lot of time making sure to reuse any sub-computation that can be reused.  Our numpy version still involves an excess of memory allocation and copying; in a low-level language like Fortran it's easier to control and minimize memory use.  In addition, the Cooley-Tukey algorithm can be extended to use splits of size other than 2 (what we've implemented here is known as the *radix-2* Cooley-Tukey FFT).  Also, other more sophisticated FFT algorithms may be used, including fundamentally distinct approaches based on convolutions (see, e.g. Bluestein's algorithm and Rader's algorithm).  The combination of the above extensions and techniques can lead to very fast FFTs even on arrays whose size is not a power of two.

Though the pure-Python functions are probably not useful in practice, I hope they've provided a bit of an intuition into what's going on in the background of FFT-based data analysis.  As data scientists, we can make-do with black-box implementations of fundamental tools constructed by our more algorithmically-minded colleagues, but I am a firm believer that the more understanding we have about the low-level algorithms we're applying to our data, the better practitioners we'll be.

*This blog post was written entirely in the IPython Notebook.  The full notebook can be downloaded
[here](http://jakevdp.github.io/downloads/notebooks/UnderstandingTheFFT.ipynb),
or viewed statically
[here](http://nbviewer.ipython.org/url/jakevdp.github.io/downloads/notebooks/UnderstandingTheFFT.ipynb).*

```python

```


```python
import sympy as sm
import galgebra.ga as ga
import IPython as ipy
#ga.printer.latex?
#ga.Ga?
#ipy.display.Math?
sm.init_printing(latex_printer=ga.printer.latex, use_latex='mathjax')
ipy.display.Math(f"y = {ga.printer.latex(sm.S('n')**2)}")
```

$\displaystyle y = n^{2}$

```python
x,y,z = sm.symbols('x y z',real=True)
xyz = (x,y,z)
# e_x, e_y, e_z =  unit vectors
# x,y,z = cordinaters
# g = metric tensor
o3d = ga.Ga('e_x e_y e_z', g=[1,1,1],coords=xyz)
ga.Ga('e_x e_y e_z', g=[1,1,1],coords=xyz).grad
ga.Ga('e_x e_y e_z', g=[1,1,1],coords=xyz).g
```

$\displaystyle \left[\begin{array}{ccc}1 & 0 & 0\\0 & 1 & 0\\0 & 0 & 1\end{array}\right]$

```python
# to define a
o3d.mv('S', 'scalar')
```

\begin{equation*} S = S \end{equation*}

```python
o3d.mv('V', 'vector')
```

\begin{equation*} V = V^{x} \boldsymbol{e}_{x} + V^{y} \boldsymbol{e}_{y} + V^{z} \boldsymbol{e}_{z} \end{equation*}

```python
o3d.mv('B', 'bivector')
```

\begin{equation*} B = B^{xy} \boldsymbol{e}_{x}\wedge \boldsymbol{e}_{y} + B^{xz} \boldsymbol{e}_{x}\wedge \boldsymbol{e}_{z} + B^{yz} \boldsymbol{e}_{y}\wedge \boldsymbol{e}_{z} \end{equation*}

```python
o3d.mv('I', 'pseudo')
```

\begin{equation*} I = I^{xyz} \boldsymbol{e}_{x}\wedge \boldsymbol{e}_{y}\wedge \boldsymbol{e}_{z} \end{equation*}

```python
o3d.mv()
```

$\displaystyle \left(  \boldsymbol{e}_{x}, \   \boldsymbol{e}_{y}, \   \boldsymbol{e}_{z}\right)$

```python
a = o3d.mv('a', 'vector')
b = o3d.mv('b','vector')
# geometric product
b*a
# wedge product = outer product
a^b
# inner product
a|b
# left contraction
a<b
# right contraction
a>b
```

\begin{equation*} a^{x} b^{x} + a^{y} b^{y} + a^{z} b^{z} \end{equation*}

```python
A = o3d.mv('A', 'bivector')
B = o3d.mv('B', 'bivector')
A>B
```

\begin{equation*} - A^{xy} B^{xy} - A^{xz} B^{xz} - A^{yz} B^{yz} \end{equation*}

```python
A|B
```

\begin{equation*} - A^{xy} B^{xy} - A^{xz} B^{xz} - A^{yz} B^{yz} \end{equation*}

```python
ex,ey,ez=o3d.mv()
ey
```

\begin{equation*}  \boldsymbol{e}_{y} \end{equation*}

```python
v = x**2*o3d.mv()[0] + y**2*o3d.mv()[1] + z**2*o3d.mv()[2]
x*(o3d.mv()[1]^o3d.mv()[2]) + y*(o3d.mv()[0]^o3d.mv()[2]) + z*(o3d.mv()[0]^o3d.mv()[1])
```

\begin{equation*} z \boldsymbol{e}_{x}\wedge \boldsymbol{e}_{y} + y \boldsymbol{e}_{x}\wedge \boldsymbol{e}_{z} + x \boldsymbol{e}_{y}\wedge \boldsymbol{e}_{z} \end{equation*}

```python
o3d.mv('v','mv')
```

\begin{equation*} v = v  + v^{x} \boldsymbol{e}_{x} + v^{y} \boldsymbol{e}_{y} + v^{z} \boldsymbol{e}_{z} + v^{xy} \boldsymbol{e}_{x}\wedge \boldsymbol{e}_{y} + v^{xz} \boldsymbol{e}_{x}\wedge \boldsymbol{e}_{z} + v^{yz} \boldsymbol{e}_{y}\wedge \boldsymbol{e}_{z} + v^{xyz} \boldsymbol{e}_{x}\wedge \boldsymbol{e}_{y}\wedge \boldsymbol{e}_{z} \end{equation*}

```python
o3d.mv('w','mv').Fmt(3)
```

\begin{equation*}  \begin{aligned}[t]  & w  \\  &  + w^{x} \boldsymbol{e}_{x} \\  &  + w^{y} \boldsymbol{e}_{y} \\  &  + w^{z} \boldsymbol{e}_{z} \\  &  + w^{xy} \boldsymbol{e}_{x}\wedge \boldsymbol{e}_{y} \\  &  + w^{xz} \boldsymbol{e}_{x}\wedge \boldsymbol{e}_{z} \\  &  + w^{yz} \boldsymbol{e}_{y}\wedge \boldsymbol{e}_{z} \\  &  + w^{xyz} \boldsymbol{e}_{x}\wedge \boldsymbol{e}_{y}\wedge \boldsymbol{e}_{z}  \end{aligned}  \end{equation*}

```python

```


# 분포의 대표값

분포(distribution)은 집합 전체의 특성을 나타내는 것이므로 히스토그램 등의 시각화 또는 분포 함수 등의 수식을 사용하는 것이 일반적이다. 하지만 전체적인 특성을 나타낼 수 있는 몇 개의 대표값을 구할 수 있다면 분포의 특성을 빠르게 파악한다거나 두 개의 분포를 서로 비교하는데 도움이 될 수 있다.

가장 먼저 생각할 수 있는 것은 분포를 하나의 숫자로 대표할 수 있는 수치, 즉 대푯값이다.

대푯값으로 널리 쓰이는 것들은 다음과 같은 것을 있다.

* 평균
* 기댓값
* 중간값

# 샘플 평균

우리가 보통 평균(mean, average)라고 할 때 의미하는 것은 **샘플 평균(sample mean, sample average)**을 말한다.

샘플(sample)은 정확히는 이미 가지고 있는 자료 값의 집합이 어떤 확률 모형 또는 모집합으로 부터 생성된 것이라는 가정하에 부르는 용어이지만 유한한 갯수의 자료 값의 집합으로 부터 다음과 같은 식을 사용하여 계산한 값을 샘플 평균이라고 한다.

$$ m = \dfrac{1}{N}\sum_{i=1}^{N} x_i $$

이 식에서 $N$는 자료의 갯수, $i = 1, \cdots, N$는 자료의 순서 혹은 인덱스(index)를 뜻하고, $x_i$는 $i$번째 자료의 값을 뜻한다.

이번에는 자료가 어떤 확률 분포를 따르는 확률 변수 $X$의 샘플이라고 가정하자. 이러한 경우에는 확률 변수 위에 bar를 추가하여 $\bar{X}$와 같이 샘플 평균을 나타내기도 한다. $X_i$는 $i$번째로 실현된 샘플을 의미한다.

$$ \bar{X} = \dfrac{1}{N}\sum_{i=1}^{N} X_i $$

샘플링을 한 번 해서 나온 자료 값의 집합에서 구한 샘플 평균 $\bar{X}$이 있을 때, 샘플링을 다시 해서 또 다른 자료 값의 집합을 구하면 이 자료는 먼저 자료와 다르기 때문에 여기에서 구한 샘플 평균은 앞서 구한 샘플 평균과 갑이 달라진다.

이렇게 샘플링 할 때마다 샘플 평균이 예측할 수 없는 값이 나오기 때문에 **샘플 평균 $\bar{X}$도 하나의 확률 변수**라고 할 수 있다. 물론 이 확률 변수는 원래의 확률 변수와는 다른 분포를 가진다.

## 확률 분포의 기댓값

만일 확률 변수가 따르고 있는 확률 모형, 정확히는 확률 밀도 함수를 알고 있을 경우에는 다음과 같은 수식을 사용하여 이론적인 평균을 구할 수 있다. 이러한 이론적 평균을 확률 변수의 **기댓값(expectation)**이라고 한다. 확률 모형이 존재한다는 것이 문맥상으로 확실한 경우에는 단순히 평균(mean)이라고 말하기도 한다.

확률 변수의 기대값을 구하는 연산자(operator)는 영어 Expection의 첫글자를 사용하여 $\operatorname{E}[\cdot]$로 표기한다. 기대값은 그리스 문자 $\mu$로 표기한다.

$$ \mu = \operatorname{E}[X] = \int_{-\infty}^{\infty} x f(x) dx $$

이산 확률 변수의 경우에는 다음과 같이 확률 질량 함수 $P(x)$를 사용하여 기댓값을 구한다.

$$ \mu = \operatorname{E}[X] = \sum xP(x) $$

## 기댓값의 성질

기댓값은 확률 모형이라는 수식을 사용한 것이므로 다음과 같은 성질을 가진다는 것을 수학적으로 증명할 수 있다.
복잡한 수식을 가지는 확률 모형의 경우에는 이러한 성질을 이용하면 기댓값을 계산하는데 도움이 된다.

* 랜덤 변수가 아닌 고정된 값 $c$에 대해
$$ \operatorname{E}[c] = c$$

* 선형성
$$
\begin{align}
  \operatorname{E}[cX]    &= c \operatorname{E}[X] \\
  \operatorname{E}[X + Y] &=   \operatorname{E}[X] + \operatorname{E}[Y]
\end{align}
$$

## 기댓값과 샘플 평균의 관계

앞서 샘플 평균도 하나의 확률 변수라고 했다. 따라서 샘플 평균에도 기댓값이 존재한다. 또한 샘플 평균의 기댓값은 원래의 확률 변수의 기댓값과 일치함을 수학적으로 증명할 수 있다.

$$ \operatorname{E}[ \bar{X} ] = \operatorname{E}[X] $$

(증명)

$$
\begin{eqnarray}
\operatorname{E}[\bar{X}]
&=& \operatorname{E}\left[\dfrac{1}{N}\sum_{1}^{N}X_i \right]  \\
&=& \dfrac{1}{N}\sum_{1}^{N}\operatorname{E}[X_i]  \\
&=& \dfrac{1}{N}\sum_{1}^{N}\operatorname{E}[X]  \\
&=& \dfrac{1}{N} N \operatorname{E}[X] \\
&=& \operatorname{E}[X]  \\
\end{eqnarray}
$$

## 중앙값

**중앙값(median)**은 전체 자료를 크기별로 정렬하였을 때 가장 중앙에 위치하게 되는 값을 말한다. 중앙값에도 샘플 중앙값과 확률 분포로부터 계산한 이론적 중앙값이 존재한다.

샘플 중앙값은 몇 개의 정의가 존재하지만 많이 쓰이는 정의는 다음과 같다.

전체 샘플의 갯수가 $N$인 경우,

* $N$이 홀수이면: 중간값은 $(n + 1)/2$번째 샘플의 값
* $N$이 짝수이면: 중간값은 $n/2$번째 샘플의 값과 $n/2 + 1$번째 샘플의 값의 평균

확률 분포로부터 이론적 중앙값은 그 값보다 큰 값이 나올 확률과 작은 값이 나올 확률이 동일하게 0.5이여야 하므로 다음과 같이 누적 확률 분포 $F(x)$에서 계산할 수 있다.

$$ \text{median} = F^{-1}(0.5) $$

$$ 0.5 = F(\text{median}) $$

## 최빈값

최빈값(mode)은 직역하자면 가장 빈번하게 나올 수 있는 값을 말한다. 하지만 연속 확률 분포인 경우에는 어느 값에 대해서나 특정한 값이 나올 확률은 0(zero)이므로 다음과 같이 확률 밀도 함수의 값이 가장 큰 확률 변수의 값으로 정의한다. 즉, 확률 밀도 함수의 최댓값의 위치이다.

$$ \text{mode} = \arg \max_x f(x) $$

이산 샘플 자료에 대해서는 최빈값을 구할 수 있으나 연속 샘플 자료는 최빈값이 존재하지 않거나 의미가 없기 때문에 히스토그램과 같이 구간을 분할하여 각 구간의 대표값으로 카테고리화하여 최빈값을 구해야 한다.

## 기댓값, 중앙값, 최빈값의 비교

확률 분포 즉, 확률 밀도 함수가 대칭인 경우에는 기댓값, 중앙값, 최빈값이 모두 같다. 그러나 분포가 어느 한쪽으로 찌그러진(skewed) 경우에는 다음 그림과 같이 달라질 수 있다.

계산량으로 비교하면 기댓값이 가장 계산하기 쉽고 중앙값은 계산량이 기댓값보다 증가하며 최빈값은 최적화 과정을 통해서만 구할 수 있으므로 계산량이 가장 많으며 오차가 크다.

그러나 기댓값은 이상값(outliner)이나 한쪽으로 찌그러진 상태에 큰 영향을 받지만 중앙값이나 최빈값은 이에 대한 영향이 적다는 장점이 있다.

## Python을 사용한 계산

Python에서는 다음과 같은 numpy 함수를 사용하여 계산할 수 있다.

* [`mean()`](http://docs.scipy.org/doc/numpy/reference/generated/numpy.mean.html)
* [`nanmean()`](http://docs.scipy.org/doc/numpy/reference/generated/numpy.nanmean.html)
* [`median()`](http://docs.scipy.org/doc/numpy/reference/generated/numpy.median.html)
* [`argmax()`](http://docs.scipy.org/doc/numpy/reference/generated/numpy.argmax.html)
* [`histogram()`](http://docs.scipy.org/doc/numpy/reference/generated/numpy.histogram.html)

```python
np.random.seed(0)
x = np.random.normal(size=1000)
```

```python
np.mean(x), np.median(x)
```

    (-0.045256707490195384, -0.058028034799627745)

```python
ns, bins = np.histogram(x, bins=np.linspace(-10,10,20))
```

```python
M = np.argmax(ns)
bins[M], bins[M+1]
```

    (-0.52631578947368496, 0.52631578947368318)

기술 통계 방법으로 Scipy나 Pandas를 사용하여 구하는 방법은 다음 노트북을 참조한다.

[[school_notebook:69b85ec9afb7415d8da4c8cd5e46d22e]]

09. 기초 확률론2 - 확률 변수/06. 분포의 대표값.ipynb

09. 기초 확률론2 - 확률 변수/06. 분포의 대표값.ipynb

09. 기초 확률론2 - 확률 변수/06. 분포의 대표값.ipynb


# Euler Problem 243

A positive fraction whose numerator is less than its denominator is called a proper fraction.
For any denominator, d, there will be d−1 proper fractions; for example, with d = 12:
1/12 , 2/12 , 3/12 , 4/12 , 5/12 , 6/12 , 7/12 , 8/12 , 9/12 , 10/12 , 11/12 .

We shall call a fraction that cannot be cancelled down a resilient fraction.
Furthermore we shall define the resilience of a denominator, R(d), to be the ratio of its proper fractions that are resilient; for example, R(12) = 4/11 .
In fact, d = 12 is the smallest denominator having a resilience R(d) < 4/10 .

Find the smallest denominator d, having a resilience R(d) < 15499/94744 .

```python
import heapq
from sympy import primerange
smallprimes = list(primerange(1, 100))
queue = []
heapq.heappush(queue, (2, 1, 1))
r = 1
while queue and r >= 15499/94744:
    n, i, phi = heapq.heappop(queue)
    r = phi / (n - 1)
    p, q = smallprimes[i-1:i+1]
    heapq.heappush(queue, (p*n, i, p*phi))
    heapq.heappush(queue, (q*n, i+1, (q-1)*phi))
print(n)
```

    892371480

```python

```


# V2: SCF optimization with VAMPyR

## V2.1: Hydrogen atom

In order to solve the one-electron Schr\"{o}dinger equation in MWs we reformulate them in an integral form [1].

\begin{equation}
 \phi = -2\hat{G}_{\mu}\hat{V}\phi
\end{equation}

Where $\hat{V}$ is the potential acting on the system, $\phi$ is the wavefunction, $\hat{G}$ is the Helmholtz integral operator, where its kernel is defined as $G_\mu(r - r') = \frac{\exp(-\mu |r - r'|)}{4\pi |r - r'|}$
and $\mu$ is a parameter defined above through the energy.

The Helmholtz operator is already implemented in vampyr, therefore the only things you need are the integral KS equation and the definition of $\mu$

\begin{equation}
     \mu = \sqrt{-2E}
\end{equation}

The way you initialize the Helmholtz operator is as follows
```
H = vp.HelmholtzOperator( mra, exp=mu, prec=eps )
```
where `mu` is the $\mu$ is the parameter defined above, mra you have seen before, and `eps` is the desired threshold precision. This operator is applied the same way you applied the vp.ScalingProjector earlier.

In this exercise you will be solving the KS equation iteratively for a simple system, the Hydrogen atom. This means that you only have the nuclear potential to take into account for the potential term in the KS equation.

\begin{align}
    V_{nuc}(\mathbf{r}) = -\frac{1}{|\mathbf{r}|}
\end{align}

We will also be working with a single orbital, of which the initial guess is

\begin{align}
    \phi_0(\mathbf{r}) = e^{-|\mathbf{r}|^2}
\end{align}
where
\begin{align}
    |\mathbf{r}| = \sqrt{x^2 + y^2 + z^2}
\end{align}

The orbital update is defined as follows
\begin{align}
    \Delta\tilde{\phi}^n &= -2\hat{G}[V_{nuc}\phi^n] - \phi^n \\
    \Delta\tilde{\phi}^n &= \tilde{\phi}^{n+1} - \phi^n
\end{align}
where we use \~ to denote a function that is **not** normalized, and $n$ is the iteration index.

What you need to do in this exercise is
1. Make a nuclear potential as a python function
2. Make an initial guess for the orbital as a python function (hint use `np.exp` )
3. Create a Helmholtz operator with $\mu$ as shown above, use the exact value of $\epsilon$ for a hydrogen atom ($-0.5$ A.U.)
4. Project both nuclear potential and orbital through vampyr
5. Compute new orbital through application of the Helmholtz operator
6. Compute the size of the orbital update $||\tilde{\phi}^{n+1} - \phi^n||$
7. Normalize the orbital $\phi^{n+1} = \tilde{\phi}^{n+1}/||\tilde{\phi}^{n+1}||$
8. Update orbital $\phi^{n+1} \rightarrow \phi^{n}$ for next iteration
9. Repeat steps 5-8 until your orbital energies have converged

The convergence criterion is the norm of $\Delta \phi^n$, but you should start by looping a set amount of times before trying the threshold.

```python
from vampyr import vampyr3d as vp
import numpy as np
import matplotlib.pyplot as plt

r_x = np.linspace(-0.99, 0.99, 1000) # create an evenly spaced set of points between -0.99 and 0.99
r_y = np.zeros(1000)
r_z = np.zeros(1000)
r = [r_x, r_y, r_z]

# Analytic nuclear potential
def f_nuc(r):
    R = np.sqrt(r[0]*r[0] + r[1]*r[1] + r[2]*r[2])
    return -1.0 / R

# Analytic guess for solution
def f_phi(r):
    R = np.sqrt(r[0]*r[0] + r[1]*r[1] + r[2]*r[2])
    return np.exp(-R*R)

# Global parameters
k = 7                        # Polynomial order
L = [-20,20]                 # Simulation box size
epsilon = 1.0e-4             # Relative precision

# Define MRA and multiwavelet projector
MRA = vp.MultiResolutionAnalysis(order=k, box=L)
P_eps = vp.ScalingProjector(mra=MRA, prec=epsilon)

# Project analytic nuclear potential
V_nuc = P_eps(f_nuc)

# Initial guess for the wavefunction
phi_n = P_eps(f_phi)
phi_n.normalize()

# Prepare Helmholtz operator
E = -0.5
mu = np.sqrt(-2*E)
H = vp.HelmholtzOperator(mra=MRA, exp=mu, prec=epsilon)

# Minimization loop
thrs = 1.0e-3
update = 1.0
i = 0
while (update > thrs):
    # Apply Helmholtz operator
    Vphi = V_nuc * phi_n
    phi_np1 = -2*H(Vphi)
    norm = phi_np1.norm()

    # Compute orbital and energy updates
    dPhi_n = phi_np1 - phi_n
    update = dPhi_n.norm()

    # Prepare for next iteration
    phi_n = phi_np1
    phi_n.normalize()
    phi_n.crop(epsilon)  # Truncate MW exansion based on precision

    phi_n_plt = [phi_n([x, 0.0, 0.0]) for x in r_x]
    plt.plot(r_x, phi_n_plt)

    print("iteration: {}    Norm: {}   Update: {}".format(i, norm, update))
    i += 1

plt.show()
```

I don't know if I want to add more than this right now

## 4 Sources

[1] Stig Rune Jensen, Santanu Saha, José A. Flores-Livas, William Huhn, Volker Blum, Stefan Goedecker, and Luca Frediani The Elephant in the Room of Density Functional Theory Calculations. The Journal of Physical Chemistry Letters 2017 8 (7), 1449-1457
DOI: 10.1021/acs.jpclett.7b00255

```python

```


```python
# https://colab.research.google.com/github/kassbohm/wb-snippets/blob/master/ipynb/HTM_03/Selbst/1.5.ipynb

from sympy.physics.units import *
from sympy import *

phi, beta, l1, b, w1 = var("phi, beta, l1, b, omega1")

sub_list = [
    (phi,   45 *pi/180),
    (beta, 150 *pi/180),
    (l1, 1*m / 10),
    (b,  3*m / 10),
    (w1,  1 / s),
    ]

# Unknowns:
vx, vy, ax, ay = var("vx, vy, ax, ay")

# Vectors:
vA = Matrix([vx, vy, 0])
aA = Matrix([ax, ay, 0])
W1 = Matrix([0, 0, w1])
cp, sp = cos(phi), sin(phi)
d1 = l1 * Matrix([cp, sp, 0])

# Equations:
eq1 = Eq(vA, W1.cross(d1))
eq2 = Eq(aA, W1.cross(W1.cross(d1)))

pprint("\nEquation 1:")
pprint(eq1)
pprint("\nEquation 2:")
pprint(eq2)

pprint("\nSolution:")
sol = solve([eq1, eq2], [vx, vy, ax, ay ])
pprint(sol)

vxs, vys, axs, ays = sol[vx], sol[vy], sol[ax], sol[ay]
vA = Matrix([vxs, vys, 0])
aA = Matrix([axs, ays, 0])

pprint("\nv / (m/s):")
tmp = vA.subs(sub_list)
tmp /= (m/s)
pprint(tmp)

pprint("\na / (m/s²):")
tmp = aA.subs(sub_list)
tmp /= (m/s/s)
pprint(tmp)

# Unknowns:
bp, bpp = var("b', b''")

# Angular velocity
w3, a3 = var("omega3, alpha3")

# Vectors:
O3 = Matrix([0, 0, w3])
A3 = Matrix([0, 0, a3])
cb, sb = cos(beta), sin(beta)
e = Matrix([cb, sb, 0])
f = Matrix([-sb, cb, 0])
d3 = b * e
vrel = bp * e
arel = bpp * e

# Equations:
eq3 = Eq( vA, O3.cross(d3) + vrel )
# eq4 = Eq( aA, A3.cross(d3) + O3.cross(O3.cross(d3)) + arel )
eq4 = Eq( aA, arel + bp*w3*f + A3.cross(d3) + O3.cross(vrel + b*w3*f) )

pprint("\nEquation 3:")
pprint(eq3)
pprint("\nEquation 4:")
pprint(eq4)

sol = solve([eq3, eq4], [bp, bpp, w3, a3], dict=True)
sol = sol[0]

pprint("\nSolution:")
bps, bpps, w3s, a3s = sol[bp], sol[bpp], sol[w3], sol[a3]
pprint("\nb':")
tmp =  bps
tmp = tmp.simplify()
pprint(tmp)
pprint("\nb'':")
tmp =  bpps
tmp = tmp.simplify()
pprint(tmp)
pprint("\nω₃:")
tmp =  w3s
tmp = tmp.simplify()
pprint(tmp)
pprint("\nw₃':")
tmp =  a3s
tmp = tmp.simplify()
pprint(tmp)

pprint("\nb' / (m/s):")
tmp = bps.subs(sub_list)
tmp /= (m/s)
tmp = tmp.simplify()
pprint(tmp)

pprint("\nb'' / (m/s²):")
tmp = bpps.subs(sub_list)
tmp /= (m/s/s)
tmp = tmp.simplify()
pprint(tmp)

pprint("\nω₃ / (1/s):")
tmp = w3s.subs(sub_list)
tmp /= (1/s)
tmp = tmp.simplify()
pprint(tmp)

pprint("\nw₃' / (1/s²):")
tmp = a3s.subs(sub_list)
tmp /= (1/s/s)
tmp = tmp.simplify()
pprint(tmp)

# Equation 1:
# ⎡vx⎤   ⎡-l₁⋅ω₁⋅sin(φ)⎤
# ⎢  ⎥   ⎢             ⎥
# ⎢vy⎥ = ⎢l₁⋅ω₁⋅cos(φ) ⎥
# ⎢  ⎥   ⎢             ⎥
# ⎣0 ⎦   ⎣      0      ⎦
#
# Equation 2:
#        ⎡      2       ⎤
# ⎡ax⎤   ⎢-l₁⋅ω₁ ⋅cos(φ)⎥
# ⎢  ⎥   ⎢              ⎥
# ⎢ay⎥ = ⎢      2       ⎥
# ⎢  ⎥   ⎢-l₁⋅ω₁ ⋅sin(φ)⎥
# ⎣0 ⎦   ⎢              ⎥
#        ⎣      0       ⎦
#
# Solution:
# ⎧          2                   2                                            ⎫
# ⎨ax: -l₁⋅ω₁ ⋅cos(φ), ay: -l₁⋅ω₁ ⋅sin(φ), vx: -l₁⋅ω₁⋅sin(φ), vy: l₁⋅ω₁⋅cos(φ)⎬
# ⎩                                                                           ⎭
#
# v / (m/s):
# ⎡-√2 ⎤
# ⎢────⎥
# ⎢ 20 ⎥
# ⎢    ⎥
# ⎢ √2 ⎥
# ⎢ ── ⎥
# ⎢ 20 ⎥
# ⎢    ⎥
# ⎣ 0  ⎦
#
# a / (m/s²):
# ⎡-√2 ⎤
# ⎢────⎥
# ⎢ 20 ⎥
# ⎢    ⎥
# ⎢-√2 ⎥
# ⎢────⎥
# ⎢ 20 ⎥
# ⎢    ⎥
# ⎣ 0  ⎦
#
# Equation 3:
# ⎡-l₁⋅ω₁⋅sin(φ)⎤   ⎡-b⋅ω₃⋅sin(β) + b'⋅cos(β)⎤
# ⎢             ⎥   ⎢                        ⎥
# ⎢l₁⋅ω₁⋅cos(φ) ⎥ = ⎢b⋅ω₃⋅cos(β) + b'⋅sin(β) ⎥
# ⎢             ⎥   ⎢                        ⎥
# ⎣      0      ⎦   ⎣           0            ⎦
#
# Equation 4:
# ⎡      2       ⎤
# ⎢-l₁⋅ω₁ ⋅cos(φ)⎥   ⎡-α₃⋅b⋅sin(β) - b'⋅ω₃⋅sin(β) + b''⋅cos(β) - ω₃⋅(b⋅ω₃⋅cos(β) + b'⋅sin(β))⎤
# ⎢              ⎥   ⎢                                                                       ⎥
# ⎢      2       ⎥ = ⎢α₃⋅b⋅cos(β) + b'⋅ω₃⋅cos(β) + b''⋅sin(β) + ω₃⋅(-b⋅ω₃⋅sin(β) + b'⋅cos(β))⎥
# ⎢-l₁⋅ω₁ ⋅sin(φ)⎥   ⎢                                                                       ⎥
# ⎢              ⎥   ⎣                                   0                                   ⎦
# ⎣      0       ⎦
#
# Solution:
#
# b':
# l₁⋅ω₁⋅sin(β - φ)
#
# b'':
#       2 ⎛                     2                                            2        ⎞
# -l₁⋅ω₁ ⋅⎝b⋅cos(β - φ) + l₁⋅sin (β) - 2⋅l₁⋅sin(β)⋅sin(φ)⋅cos(β - φ) + l₁⋅sin (φ) - l₁⎠
# ──────────────────────────────────────────────────────────────────────────────────────
#                                           b
#
# ω₃:
# l₁⋅ω₁⋅cos(β - φ)
# ────────────────
#        b
#
# w₃':
#      2
# l₁⋅ω₁ ⋅(b - 2⋅l₁⋅cos(β - φ))⋅sin(β - φ)
# ───────────────────────────────────────
#                     2
#                    b
#
# b' / (m/s):
# √2   √6
# ── + ──
# 40   40
#
# b'' / (m/s²):
#   √2    √3   1    √6
# - ── - ─── + ── + ──
#   40   120   60   40
#
# ω₃ / (1/s):
#   √6   √2
# - ── + ──
#   12   12
#
# w₃' / (1/s²):
# 1    √2   √6
# ── + ── + ──
# 18   12   12

```


```python
from datascience import *
import sympy
solve = lambda x,y: sympy.solve(x-y)[0] if len(sympy.solve(x-y))==1 else "Not Single Solution"
import matplotlib.pyplot as plt
import numpy as np
import pandas as pd
import datetime as dt
import warnings
warnings.simplefilter("ignore")
%matplotlib inline
```

# Bertrand Competition

Another model we consider is **Bertrand competition**, named for Joseph Louis Francois Bertrand, that is similar to Cournot competition but that firms compete using _prices_ rather than quantity. Under the assumptions of this model, consumers want to buy everything at the lowest price, and if the price is the same then demand is evenly split between those producers. One fundamental assumption is that all firms have the same unit cost of production, which means that as long as the price the firm sets is above the unit cost, it is willing to supply any amount that is demanded.

An example of a Bertrand oligopoly comes form the soft drink industry: Coke and Pepsi (which form a **duopoly**, a market with only two participants). Both firms compete by changing their prices based on a function that takes into account the price charged by their competitor. This model predicts that even this small competition will result in prices being reduced to the marginal cost level, the same outcome as perfect competition.

## Bertrand Equilibrium

To find the Bertrand equilibrium, let $c$ be the (constant) marginal cost, $p_1$ be firm 1's price level, $p_2$ be firm 2's price level, and $p_m$ be the monopoly price level. Firm 1's price depends on what it believes firm 2 will set its prices to be. Because consumers always buy at the lowest price and the firm will fulfill any level of demand, pricing just below firm 2 will obtain full market demand for firm 1. Why might this not be a good idea? If firm 2 is pricing below the level of marginal cost, then firm 1 will incur losses because they would need to sell at a price lower than the cost of production.

Let $p'_1(p_2)$ be firm 1's optimal price based on price $p_2$ set by firm 2. The graph below shows $p'_1(p_2)$. Note that when $p_2 < c$, $p'_1$ is equal to $c$, that $p'_1$ rises linearly along but _just below_ the line $p_1 = p_2$ with $p_2$ until $p_2$ reaches $p_m$ (the monopoly price level), and that it then levels off at $p_m$. In this way, firm 1's price stays below firm 2's price when it is not operating at a loss and does not exceed $p_m$ (because $p_m$ is the profit-maximizing amount for a monopoly and producing more actually results in less profit). This piecewise function has the formula

$$
p'_1(p_2) = \begin{cases}
c & \text{if } p_2 < c + h \\
p_2 - h & \text{if } c + h \le p_2 < p_m + h \\
p_m & \text{otherwise}
\end{cases}
$$

where $h$ is a small positive value and indicates the vertical distance between $p'_1$ and the line $p_1 = p_2$. We can think of $h$ as the amount by which firm 1 will undercut firm 2: as long as firm 1 will be operating at a profit and not exceeding $p_m$, they will sell at $h$ dollars below $p_2$.

```python
p_m = 18
c = 6
dist = 0.4 # the distance from p_1 = p_2 to p'_1(p_2)

xs1 = np.linspace(0, c, 100)
ys1 = c * np.ones_like(xs1)

xs2 = np.linspace(c + dist, p_m + dist, 100)
ys2 = xs2.copy() - dist

xs3 = np.linspace(p_m + dist, 25, 100)
ys3 = p_m * np.ones_like(xs3)

xs = np.append(np.append(xs1, xs2), xs3)
ys = np.append(np.append(ys1, ys2), ys3)

plt.figure(figsize=[8,8])

plt.plot(xs, ys, label=r"$p'_1(p_2)$")

y_equals_x_xs = np.linspace(0, 25, 100)
y_equals_x_ys = y_equals_x_xs.copy()

plt.plot(y_equals_x_xs, y_equals_x_ys, label=r"$p_1 = p_2$")

plt.vlines(c, 0, c, linestyle="dashed")
plt.vlines(p_m, 0, p_m, linestyle="dashed")
plt.hlines(p_m, 0, p_m, linestyle="dashed")

plt.text(2, 6.5, r"$p'_1(p_2)$", size=16)
plt.text(17, 20, r"$p_1 = p_2$", size=16)

plt.xlim([0,22])
plt.ylim([0,22])
plt.xlabel(r"$p_2$", size=16)
plt.ylabel(r"$p_1$", size=16)
plt.xticks([c, p_m], [r"$c$", r"$p_m$"], size=14)
plt.yticks([c, p_m], [r"$c$", r"$p_m$"], size=14)
plt.legend();
```

Because firm 2 has the same marginal cost $c$ as firm 1, its reaction function $p'_2(p_1)$ is symmetrical to firm 1's about the line $p_1 = p_2$:

```python
p_m = 18
c = 6
dist = 0.4 # the distance from p_1 = p_2 to p'_1(p_2)

xs1 = np.linspace(0, c, 100)
ys1 = c * np.ones_like(xs1)

xs2 = np.linspace(c + dist, p_m + dist, 100)
ys2 = xs2.copy() - dist

xs3 = np.linspace(p_m + dist, 25, 100)
ys3 = p_m * np.ones_like(xs3)

xs_1 = np.append(np.append(xs1, xs2), xs3)
ys_1 = np.append(np.append(ys1, ys2), ys3)

ys1 = np.linspace(0, c, 100)
xs1 = c * np.ones_like(xs1)

ys2 = np.linspace(c + dist, p_m + dist, 100)
xs2 = xs2.copy() - dist

ys3 = np.linspace(p_m + dist, 25, 100)
xs3 = p_m * np.ones_like(xs3)

xs_2 = np.append(np.append(xs1, xs2), xs3)
ys_2 = np.append(np.append(ys1, ys2), ys3)

plt.figure(figsize=[8,8])

plt.plot(xs_1, ys_1, label=r"$p'_1(p_2)$")
plt.plot(xs_2, ys_2, label=r"$p'_2(p_1)$")

y_equals_x_xs = np.linspace(0, 25, 100)
y_equals_x_ys = y_equals_x_xs.copy()

plt.plot(y_equals_x_xs, y_equals_x_ys, label=r"$p_1 = p_2$")

plt.vlines(p_m, 0, p_m, linestyle="dashed")
plt.hlines(p_m, 0, p_m, linestyle="dashed")

plt.text(2, 6.5, r"$p'_1(p_2)$", size=16)
plt.text(6.5, 2, r"$p'_2(p_1)$", size=16)
plt.text(18.2, 21, r"$p_1 = p_2$", size=16)

plt.xlim([0,22])
plt.ylim([0,22])
plt.xlabel(r"$p_2$", size=16)
plt.ylabel(r"$p_1$", size=16)
plt.xticks([c, p_m], [r"$c$", r"$p_m$"], size=14)
plt.yticks([c, p_m], [r"$c$", r"$p_m$"], size=14)
plt.legend();
```

These two strategies form a Nash equilibrium because neither firm can increase profits by changing their own strategy unilaterally. The equilibrium occurs where $p_1 = p'_1(p_2)$ and $p_2 = p'_2(p_1)$, at the intersection of the two reaction curves. Notably, this means that the Bertrand equilibrium occurs when both firms are producing _at marginal cost_.

This makes intuitive sense: say that the two firms both set equal prices at a price above $c$ where they split demand equally. Then both firms have incentive to reduce their price slightly and take the other half of the market share from their competitor. Thus, both firms are tempted to lower prices as much as possible, but lowering below the level of marginal cost makes no sense because then they're operating at a loss. Thus, both firms sell at the price level $c$.

## Implications

The Bertrand model implies that even a duopoly in a market is enough to push prices down to the level of perfect competition. It does, however, rely on some serious assumptions. For example, there are many reasons why consumers might not buy the lowest-priced item (e.g. non-price competition, search costs). When these factors are included in the Bertrand model, the same result is no longer reached. It also ignores the fact that firms may not be able to supply the entire market demand; including these capacity constraints in the model can result in the system having no Nash equilibrium. Lastly, the Bertrand model demonstrates big incentives to cooperate and raise prices to the monopoly level; however, this state is not a Nash equilibrium, and in fact, the only Nash equilibrium of this model is the non-cooperative one with prices at marginal cost.

## Applying Bertrand

Now that we have derived the Bertrand equilibrium, let's apply it to a problem. Consider the Coke-Pepsi duopoly we mentioned above. Suppose that the only product in the soft-drink market is the 12-oz. can, that the market demand for cans is given by $P = -0.05 Q + 5.05$, and that the marginal cost for Coke and Pepsi is constant at $c = 0.25$. To find the equilibrium price for Coke based in it's belief that Pepsi will sell at $p_2 = 1$, we need to start by finding the monopoly price level $p_m$; recall from Cournot that this occurs when the marginal revenue curve of the market demand intersects the marginal cost. The marginal revenue is $r(q) = -0.1 q + 5.05$:

```python
P_fn = lambda x: -0.05 * x + 5.05
r_fn = lambda x: 2 * P_fn(x) - P_fn(0)
Qs = np.linspace(-10, 1000, 1000)
Ps = P_fn(Qs)
rs = r_fn(Qs)

plt.figure(figsize=[7,7])
plt.plot(Qs, Ps, label=r"$P(Q)$")
plt.plot(Qs, rs, label=r"$r(Q)$")
plt.hlines(.25, -10, 150, color="r", label=r"$c$")
plt.xlim([0, 110])
plt.ylim([0, 5.5])
plt.xlabel(r"Quantity $Q$", size=16)
plt.ylabel(r"Price $P$", size=16)
plt.legend();
```

Using SymPy, we can find the quantity $q$ at which $r(q) = c$. This value, denoted $q_m$, is the monopoly quantity.

```python
c = 0.25
q = sympy.Symbol("q")
r = -.1 * q + 5.05

q_m = solve(r, c)
q_m
```

    48.0000000000000

The monopoly price $p_m$ is the price from the market demand curve that this level of output:

```python
Q = sympy.Symbol("Q")
P = -.05 * Q + 5.05

p_m = P.subs(Q, q_m)
p_m
```

    2.65000000000000

Now that we have found $p_m$, we can use this to construct Coke's reaction function $p'_1(p_2)$ to Pepsi's choice of price. Assuming Coke selects $h=0.1$ (that is, Coke will sell at \$0.10 below Pepsi as long as they operate at a profit), the formula for $p'_1$ is

$$
p'_1(p_2) = \begin{cases}
0.25 & \text{if } p_2 < 0.25 + 0.1 \\
p_2 - 0.1 & \text{if } 0.25 + 0.1 \le p_2 < 2.65 + 0.1 \\
2.65 & \text{otherwise}
\end{cases}
$$

```python
c = 0.25
def p1_fn(p_2, h):
    if p_2 < c + h:
        return c
    elif c + h <= p_2 < p_m + h:
        return p_2 - h
    else:
        return p_m

p1_fn = np.vectorize(p1_fn)

xs = np.linspace(-1, 4, 1000)
ys = xs.copy()
p1s =  p1_fn(xs, .1)

plt.figure(figsize=[7,7])
plt.plot(xs, p1s, label=r"$p'_1(p_2)$")
plt.plot(xs, ys, label=r"$p_1 = p_2$")
plt.hlines(c, 0, 5, color="r", label=r"$c$")
plt.vlines(c, 0, 5, color="r")
plt.hlines(p_m, 0, 5, color="g", label=r"$p_m$")
plt.vlines(p_m, 0, 5, color="g")
plt.xlim([0,3])
plt.ylim([0,3])
plt.xlabel(r"$p_2$", size=16)
plt.ylabel(r"$p_1$", size=16)
plt.legend();
```

Finally, to find Coke's selling price, we find $p'_1(1)$, since Coke believes Pepsi will sell at \$1.

```python
c = 0.25
h = 0.1
p_2 = sympy.Symbol("p_2")
p_1_prime = sympy.Piecewise(
    (c, p_2 < c + h),
    (p_2 - h, p_2 < p_m + h),
    (p_m, True)
)

p_1_prime.subs(p_2, 1)
```

    0.900000000000000

Thus, if Coke believes that Pepsi will sell cans at \$1, it should sell at \$0.90. This should make intuitive sense: we showed that $p'_1(p_2)$ was below the line $p_1=p_2$ by a vertical distance of $h=0.1$, so it makes sense that Coke would sell at \$0.10 below Pepsi. If Coke had believed that Pepsi was going to sell at, say, \$2.70, then it would have been better for them to sell at the monopoly price level $p_m = 2.65$. If Pepsi was selling below margin cost, at \$0.20 maybe, then Coke's best bet would have been to sell at $c = 0.25$, although they would have sold 0 units of output because consumers buy from the lowest-priced vendor.