```python
# This cell is tagged for removal

import os
import math
import pickle
from glob import glob

import numpy as np
import pandas as pd
import sympy
sympy.init_printing(use_unicode=False, wrap_line=False)
from IPython import display as idisplay

from sklearn.metrics import r2_score

import aliases # important this goes first to configure PATH

from everest.window import Canvas, plot, raster, DataChannel as Channel, get_cmap
from everest.window.colourmaps import *
from everest.window import image, imop

from referencing import search
from myst_nb import glue
from analysis import isovisc, arrhenius, utilities, common, analysis, visualisation
%matplotlib inline

osjoin = os.path.join

import warnings

warnings.filterwarnings("ignore",category=UserWarning)
```

# Isoviscous rheology in cylindrical geometry

## Defining a coordinate system

In any convection model, gravity defines the natural down direction and gives us our first most important scale: the depth $z$ from the surface, or its complement, the height from the model base $h=1-z$.

If the domain is allowed to curve around a certain locus, a cylindrical or annular geometry is obtained which is more appropriate for planetary mantles. While we retain $h$ and $z$ as terms relevant to any action within the domain, we must also introduce a concept of radial height $r$, understood here to represent the distance from the planetary centre of gravity. The cylindrical domain, for us representing the mantle, is thus bounded by the inner radius $r_{i}$ and the outer radius $r_{o}$, defining an area of $\pi(r_o^2 - r_i^2)$.

Our choice of radii implies a degree of curvature $f$:

$$ f \equiv \frac{r_o}{r_i} $$

Where $f=1$ is equivalent to an infinitely wide Cartesian box, $f\to0$ represents a full disk, and the values $\sim 0.5$ and $\sim 0.9$ would be appropriate for the whole mantle and upper mantle respectively. The ratio of radii $f$ is identical to the ratio of circumferences, so that $f=0.5$ represents a system where the arc length of the base is half that of the surface. (Note that this would imply infinite planetary radii at $f=1$ - hence the planar-like endmember $f=1$ is not strictly reachable under an assumption of curvature, though arbitrarily high values can be set to reproduce that behaviour {cite}`Jarvis1993-cb`.)

If we further stipulate:

$$
\Delta r = r_{o} - r_{i} = 1 \\
r_{o} \to 1 \quad as \quad f \to 0
$$

Then:

$$
r_i = \frac{f}{1 - f}, \quad r_o = \frac{1}{1 - f} \\
r(h) = r_i + h
$$

Which suggests non-dimensionalising the planetary radial height $r$ as $r^{*}$ (*r-star*) such that ${r^{*}}_o = 1$:

$$
{r^{*}}_i = f, \quad {r^{*}}_o = 1 \\
\Delta r^{*} = 1 - f \\
r^{*}(h) = \frac{h + r_i}{r_0} = \frac{\Delta r^{*}}{r} = h(1-f) + f
$$

This leaves us with four different terms to describe radial position: $h$, the dimensionless height from the mantle base; $z$, its complement; $r$, the radial scale such that $r_o - r_i = 1$; and $r^{*}$, the radial scale such that $r_0 = 1$. Each of these scales will prove natural in some contexts and less so in others, and all find use in our analysis.

To complete our coordinate system, we require an angular coordinate as well: the angle $\theta$ in radians anticlockwise from an arbitrary origin. Often we will only want to reproduce a small wedge of the annulus, as time-dependence and numerical workload scale exponentially with aspect ratio; we can define our wedge selection $\Theta$ in radians:

$$ \theta: 0 \to \Theta, \quad \Theta \leq 2\pi $$

If the simulation is to represent periodic flow around the annulus, values of $\Theta$ must satsify $\pi / l$, where $l$ is a positive integer representing the number of convection cell pairs it would take to populate the full annulus at equivalent curvature (i.e. the number of upwellings or downwellings).

The selection of $\Theta$ also selects an aspect ratio for the domain, but only once we choose a depth at which to measure it. It is most convenient to use the arc length through the mid-depth, which then relates to $\Theta$ and $f$ via a new term: the radial distance from the planetary centre of gravity to the mantle mid-depth, $r_m$:

$$
r_m \equiv \frac{r_{i} + r_{o}}{2} = \frac{1 + f}{2 \left( 1 - f \right)} \\
A = r_m \Theta
$$

If radius has been non-dimensionalised as above, this is equivalent to stating that the aspect ratio is equal to the area of the domain - regardless of $f$ - just as it is in the Cartesian case.

Such a scheme leaves us with two competing claims for a 'natural' denominator of the angular coordinate - $\Theta$ and $r_m$. While authors have sometimes preferred to keep $\Theta$ and $r_m$ constant and allow $A$ to vary {cite}`Jarvis1994-np`, we have for the most part chosen to fix $A$ and $r_m$ with $\Theta$ as the free parameter, as in {cite}`Jarvis1993-cb`: this simplifies comparisons with plane-layer simulations at the cost of producing planforms which would be unstable if scaled to the full annulus.

Some system forcings, like internal heat, scale with area. The proportion of the annulus lying below a particular height $h$ - which we shall call $D$ for 'disk' - is a function of the inner and mid-radii:

$$
D(h) = \frac{r^2 - {r_i}^2}{2 r_m}, \quad r = h + r_i
$$

Having declared that the total area shall always equal the aspect ratio $A$, the true area under any depth $h$ is simply $D \cdot A$.

From $r_m$ we are also able to derive $s$, a statement of the angular length in 'real' terms (i.e. in the same metric as $r$). It is convenient to non-dimensionalise $s$ as $s^{*} = s / A$, such that the dimensionless length through the mid-depth ${s^{*}}_m = 1$. We can then write $s^{*}$ very simply as a function of $r^{*}$, and the inner and outer lengths accordingly:

$$
s^{*} = 2 \frac{r^{*}}{1+f} \\
{s^{*}}_i = 2 \frac{f}{1+f}, \quad {s^{*}}_o = 2 \frac{1}{1+f}
$$

The length $s$ is, among other things, the factor by which an average measurement of some variable taken across a layer can be converted into a total value for that layer. It is vital to account for varying $s$ whenever comparing between different layers in a given system, or between equivalent layers in systems of differing $f$.

```python
_h, _A, _H = sympy.symbols(('h', 'A', 'H'), positive = True, real = True)
_f = sympy.Symbol('f', positive = True)
_ro = 1 / (1 - _f)
_ri = _ro * _f
_r = _h + _ri
_rstar = (_h + _ri) / _ro
_rm = (_ri + _ro) / 2
_Theta = _A / _rm
_D = (_r**2 - _ri**2) / (2 * _rm)
_sstar = 2 * _rstar / (1 + _f)
_sstari = 2 * _f / (1 + _f)
_sstaro = 2 / (1 + _f)
```

```python
def r_outer(f):
    return 1 / (1 - f)
def r_inner(f):
    return r_outer(f) * f
def radius(h, f):
    return h + r_inner(f)
def r_mid(f):
    return (r_inner(f) + r_outer(f)) / 2

def r_star(h, f):
    rinner, router = r_inner(f), r_outer(f)
    rstar = (h + rinner) / router # equiv to (h * (1 - f) + f)
    return rstar

def s_star(h, f):
    return 2 * r_star(h, f) / (1 + f)

def sub_area(h, f):
    if f == 1:
        return h
    return (radius(h, f)**2 - r_inner(f)**2) / (2 * r_mid(f))

#     for h in [0, 0.3, 0.5, 0.8, 0.99]:
#         print(f"    h = {h}")
#         print(f"        {sub_area(h, f)}")
```

## Basal heating

### Conductive solution

```python
# fs = [0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7, 0.8, 0.9, 1.0]
# geotherms = []
# condavts = []
# for f in fs:
#     system = Isovisc(alpha = 1., f = f, res = 32)
#     system.initialise()
#     temp = system.locals.temperatureField
#     diff = system.locals.diffusivityFn
#     cond = pfn.conduction.default(temp, 0, diff)
#     condavt = pfn.integral.volume(cond).data
#     left = pfn.surface.left(cond)
#     fig = QuickFig(cond, background = 'white', edgecolour = 'white')
#     geotherms.append(left.data)
#     condavts.append(condavt)
#     fig.save('cond_f' + str(f).replace('.', '-'))
# out = {
#     'fs': fs,
#     'geotherms': geotherms,
#     'avts': condavts,
#     }
# with open('cond.pkl', mode = 'wb') as file:
#     file.write(pickle.dumps(out))

impaths = sorted(os.path.relpath(path) for path in glob(osjoin(aliases.storagedir, 'cond_f*.png')))
ims = tuple(image.fromfile(path) for path in impaths)
thumbs = imop.vstack(
    imop.hstack(*ims[:5]),
    imop.hstack(*ims[5:]),
    )

with open(osjoin(aliases.storagedir, 'condf.pkl'), mode = 'rb') as file:
    conddata = pickle.loads(file.read())
condgeotherms, condavts, condfs = (conddata[key] for key in ('geotherms', 'avts', 'fs'))

# Canvas 0

canvas0 = Canvas(size = (8, 8/3), shape = (1, 3))
ax1 = canvas0.make_ax((0, 0))
ax2 = canvas0.make_ax((0, 1))
ax3 = canvas0.make_ax((0, 2))
for f, T in zip(condfs, condgeotherms):
    f = min(0.999, f)
    h = np.linspace(0, 1, len(T))
    dT, hdT = analysis.derivative(T, h, n = 1)
    sstar = s_star(hdT, f)
    phi = dT * sstar
    ax1.line(
        Channel(T, label = 'T'),
        Channel(h, label = 'h', lims = (0, 1)),
        c = cmap(f, condfs, style = 'turbo'),
        )
    ax2.line(
        Channel(dT, label = '\delta T / \delta h', lims = (-4, 0), capped = (True, True)),
        Channel(hdT, label = 'h', lims = (0, 1)),
        c = cmap(f, condfs, style = 'turbo'),
        )
    ax3.line(
        Channel(phi, label = '\phi_q', lims = (-1, -0.6)),
        Channel(hdT, label = 'h', lims = (0, 1)),
        c = cmap(f, condfs, style = 'turbo'),
        )
for ax in (ax2, ax3):
    ax.props.edges.y.ticks.major.labels = []
    ax.props.edges.y.label.visible = False

# Canvas 1

canvas1 = Canvas(shape = (1, 2), size = (5.5, 3))

ax1 = canvas1.make_ax((0, 0))
ax2 = canvas1.make_ax((0, 1))

fslopes = []

for f, T in zip(condfs, condgeotherms):

    f = min(0.999, f)

    h = np.linspace(0, 1, len(T))
    rstar = r_star(h, f)

    ax1.line(
        Tchan := Channel(T, label = 'T'),
        Channel(rstar, label = 'r^{*}'),
        c = cmap(f, condfs, style = 'turbo'),
        )

    ax2.line(
        Tchan,
        Channel(rstar, label = r"r^{*}", capped = (True, True), log = True),
        c = cmap(f, condfs, style = 'turbo'),
        )

    fslopes.append(np.mean(np.gradient(T, np.log(rstar), edge_order = 2)))

fslopes = np.array(fslopes)

# ax2.props.edges.y.swap()

ax2.props.legend.set_handles_labels(
    (row[0] for row in ax1.collections),
    (str(f) for f in condfs),
    )
ax2.props.legend.title.text = 'f'
ax2.props.legend.title.visible = True
ax2.props.legend.mplprops['bbox_to_anchor'] = (1.75, 1.05)
# ax1.props.legend.mplprops['ncol'] = 2
ax2.props.legend.frame.colour = 'black'
ax2.props.legend.frame.visible = True

# Canvas 2

canvas2 = Canvas(size = (2.5, 5), shape = (2, 1))

ax1 = canvas2.make_ax(place = (0, 0))
ax1.line(
    Channel(condfs, label = 'f', lims = (0, 1.), capped = (True, True)),
    Tchan := Channel(condavts, label = 'T_{av}', lims = (0.2, 0.5), capped = (True, True)),
    )

def func(f):
    return 0.5 * f ** (1. / math.e)
predf = np.array(list(map(func, condfs)))
ax2 = canvas2.make_ax(place = (1, 0))
ax2.line(
    predfchan := Channel(predf, label = r'\frac{1}{2}f^{1/e}', lims = (0.2, 0.5), capped = (True, True)),
    Tchan,
    )
linscore = r2_score(predf, condavts)
ax2.line(
    predfchan,
    Channel(predfchan.data, lims = Tchan.lims, capped = Tchan.capped),
    linestyle = '--'
    )
trendlabel = r'y=x, \\ R^2 =' + str(round(linscore, 3))
ax2.annotate(
    predf[3],
    predf[3],
    label = trendlabel,
    points = (30, -30),
    arrowProps = dict(arrowstyle = "->", color = '#ff7f0e'),
    )

# Assembly

# fig = imop.hstack(imop.vstack(canvas1, thumbs), canvas2)
fig = imop.paste(
    imop.vstack(canvas0, imop.hstack(canvas1, canvas2, pad = (255, 255, 255))),
    imop.resize(thumbs, size = 0.178),
    coord = (0.01, 0.96),
    corner = 'bl',
    )

# canvas = Canvas(size = (3, 5))
# ax = canvas.make_ax()
# for condgeotherm in condgeotherms:
#     ax.line(
#         Channel(np.diff(condgeotherms[0]) / np.diff(h), label = r'\frac{dT}{dh}'),
#         Channel(h[:-1], label = 'h', lims = (0, 1), capped = (True, True)),
#         )
# canvas

# Display

glue("isocondf", fig, display = False)
glue("isocondfr2", linscore, display = False)

fig
```

```python
canvas = Canvas(size = (3, 3))
ax = canvas.make_ax()
allT, allr = [], []
for f, T in zip(condfs, condgeotherms):
    f = min(0.999, f)
    h = np.linspace(0, 1, len(T))
    rstar = r_star(h, f)
    ax.line(
        rchan := Channel(np.log(rstar) / np.log(f), lims = (0, 1), label = r'\frac{\ln{r^{*}}}{\ln{f}}'),
        Channel(T, label = 'T'),
        c = cmap(f, condfs, style = 'turbo'),
        )
    allT.extend(T)
    allr.extend(rchan.data)
linscore = r2_score(allT, allr)
ax.line(
    np.linspace(0, 1, 10),
    np.linspace(0, 1, 10),
    color = '#ff7f0e',
    linestyle = '--',
    )
trendlabel = r'y=x, \\ R^2 =' + str(round(linscore, 8))
ax.annotate(
    rchan.data[15],
    rchan.data[15],
    label = trendlabel,
    points = (30, -30),
    arrowProps = dict(arrowstyle = "->", color = '#ff7f0e'),
    )
glue("isocondffit", canvas, display = False)
canvas
```

```{glue:figure} isocondf
:figwidth: 800px
:name: "isocondf"

Summary of the scaling behaviours of isoviscous conduction for varying curvature parameter $f$. We obtain a natural scaling for $f$ versus $T_{av}$ with an $R^2$ of {glue:text}`isocondfr2:.3f`.
```

```{glue:figure} isocondffit
:figwidth: 300px
:name: "isocondffit"

The analytical scaling of conductive temperature with $\ln{r^{*}}/\ln{f}$ holds empirically with extreme precision.
```

It is a requirement of a conductive steady state ($Nu=1$) that the thermal flux must be the same through every layer. In the planar case this results in a linear geotherm which, in a model with fixed and unitless boundary temperatures, results in a simple function of $T = z$ where $z$ is dimensionless depth from the top of the model. The average temperature is then trivially $T_{av}=0.5$. (For any system in pure conduction the *Nusselt* number is by definition $1$.)

In a cylindrical domain, however, the length of each layer $s$ is a function of depth and curvature as we have shown; consequently, shallower layers may transmit the same flux with a smaller temperature drop:

$$
\phi_q \propto s \cdot \frac{dT}{dh}
$$

To define the flux, we need the geothermal gradient. The conductive geotherm can be elegantly stated in terms of $r^{*}$ {numref}`isocondf` {numref}`isocondffit`:

$$
T(h) = \frac{\ln{r^{*}}}{\ln{f}}
$$

And so the geothermal gradient:

$$
\frac{dT}{dh} = \frac{f-1}{r^{*}\ln{f}}
$$

And finally the flux itself can be written as:

$$ \begin{align*}
\phi_q &\propto \frac{s^{*}(f-1)}{r^{*}\ln{f}} \\
&= \frac{2(1-f)}{(f+1)\ln{f}}
\end{align*} $$

$$ \begin{align*}
&\to -1 &as \quad f \to 1 \\
&\to 0 &as \quad f \to 0
\end{align*} $$

Or very succinctly in terms of the 'true' radius of the mid-depth:

$$
\phi_q \propto \frac{1}{r_m \ln{f}}
$$

To facilitate comparison between systems of different curvature, we can then use the above to define a dimensionless planetary flux ${\phi_q}^{*}$ - which is really just another name for the *Nusselt* number $Nu$:

$$ \begin{align*}
{\phi_q}^{*} &= \frac{ {\phi_q} }{ {\phi_q}_c } \\
&\equiv Nu
\end{align*} $$

Where the subscript $c$, here and elsewhere, denotes a purely conductive endmember. Because $Nu$ now inherits a dependency on $f$, it is no longer equivalent to the dimensionless surface temperature gradient, and so it is important always to present and discuss it in its proper terms as a ratio of fluxes.

Just as the flux now scales with $f$, so must the average mantle temperature. In the planar case, the average temperature of the system is always half the temperature drop. In the cylindrical case, however:

$$ \begin{align*}
T_{av} &= \dfrac{1}{2} \large{\root \huge{e} \of f} \\
&\equiv T_{c}
\end{align*} $$

The relationship is apparent in the numerical results {numref}`isocondf`.

### Instability and convection

An implication of $Nu$'s dependency on curvature is that the upper and lower boundaries must no longer be symmetrical. This invalidates many of the assumptions that made the planar case amenable to analysis. The additional space at the top of the model now allows more room for downwellings relative to basal upwellings, tending to promote instability {cite}`Jarvis1991-ir`; on the other hand, the curved geotherm and the increased surface for radiating heat would tend to permit a comparatively thicker upper boundary layer. The effect of these countervailing forcings on the fundamental scalings of $Nu$, $Ra$, $Ra_{cr}$, and the all-important relation $Nu \propto R^{\beta}$ is not obvious.

To begin to unpack the complexities of convection in the annulus, we can start with the assumption that - as in the planar case - the convective steady state will eventually result in a broad intracellular region of uniform temperature $T_{cell}$. Assuming a unit temperature drop $\Delta T = 1$, we can write:

$$ \begin{align*}
{\Delta T}_o &= T_{cell} \\
{\Delta T}_i &= 1 - T_{cell}
\end{align*} $$

Knowing that the inner and outer fluxes ${\phi_q}_i$ and ${\phi_q}_o$ must be equal at steady state, and that the outer boundary - due to its greater length - can sustain that flux with a gradient shallower by a factor of $f$, we can deduce a relation between the outer and inner thermal gradients, and thence between $T_{cell}$ and the inner and outer boundary layer thicknesses ${\Delta r}_i$ and ${\Delta r}_o$:

$$ \begin{align*}
f \frac{{\Delta T}_i}{{\Delta r}_i} &= \frac{{\Delta T}_o}{{\Delta r}_o} \\
\frac{{\Delta r}_i}{{\Delta r}_o} &= f \frac{1 - T_{cell}}{T_{cell}}
\end{align*} $$

For each of the two layers, we can prescribe a layer-specific *Rayleigh* number accordingly:

$$ \begin{align*}
Ra_o &\propto T_{cell} {{\Delta r}_o}^3 \\
Ra_i &\propto (1 - T_{cell}) {{\Delta r}_i}^3
\end{align*} $$

Having maintained non-dimensionality throughout, it is simple relate these two boundary *Rayleigh* numbers to the bulk $Ra$ value:

$$
Ra_{layer} = Ra \cdot {\Delta T}_{layer} \cdot {{\Delta r}_{layer}}^3
$$

At this point, however, we have exhausted the insight we can obtain without making further assumptions. If we provide that the inner and outer boundary thicknesses must be the same, as they are in the planar case, we can see that:

$$
T_{cell} = \frac{f}{f + 1} \quad \leftarrow {\Delta r}_i = {\Delta r}_o
$$

This, however, would imply that the inner and outer *Rayleigh* numbers are divergent. If we instead choose to conserve $Ra$, then: ({cite}`Jarvis1993-cb`)

$$
T_{cell} = \frac{1}{1 + f^{-3/4}} \quad \leftarrow Ra_i = Ra_o
$$

Both possibilities converge on $0.5$ when $f\to1$ and $0$ when $f\to0$, as we would expect.

However it is estimated, it is clear that, as $Ra$ increases and boundaries thin, more of the mantle will fall in the intracellular region and global temperatures as a whole will approach $T_{cell}$. Conversely, if $Ra$ slips below its critical value, the boundary layers will disapper and the entire domain will enter the conductive regime: $T^{av} = T_{c}$. These two temperatures therefore make up respectively the lower and upper endmembers of global temperature:

$$ \begin{align*}
T_{av} &\approx T_{c}, \quad Ra < Ra_{cr} \\
&\to T_{cell}, \quad Ra \to \infty
\end{align*} $$

It makes intuitive sense that the effect of increasing $Ra$ should be to decrease global temperatures, since that is exactly why convection is preferred wherever possible - though this intuition may not hold for all rheologies.

Of course, what we desire most of all is a cylindrical scaling for the mantle convection power law $Nu \propto R^{\beta}$. Following {cite}`Jarvis1993-cb` and mandating equality of inner and outer $Ra_{layer}$, it is possible to construct a 'geometric correction' $g(f)$ that functions as a coefficient of the *beta* scaling:

$$
g(f) = \frac{Nu_{c}}{{T_{cell}}^{4/3}} \quad \leftarrow Ra_i = Ra_o \\
Nu = g(f) \cdot R^{\frac{1}{3}}
$$

Using this scaling, Jarvis was able to obtain a *beta* exponent of $0.321 \pm 0.001$ across four values of $f$ from $(1.0 - 0.1)$ {cite}`Jarvis1993-cb`.

## Internal heating in the annulus

### Conductive solution

```python
with open(osjoin(aliases.storagedir, 'condhfinsulating.pkl'), mode = 'rb') as file:
    conddata = pickle.loads(file.read())
condhs, condfs = zip(*conddata['hfs'])
condhs = tuple(round(val, 1) for val in condhs)
frm = pd.DataFrame(dict(
    H = condhs, f = condfs, T = conddata['avts'], geotherm = conddata['geotherms']
    ))
frm = frm.loc[frm['H'] > 0]
frm = frm.set_index(['H', 'f'])
Hs, fs = (np.array(sorted(set(frm.index.get_level_values(level)))) for level in ('H', 'f'))
frm['h'] = frm['geotherm'].apply(lambda x: np.linspace(0, 1, len(x)))
# frm['rstar'] = frm.apply(lambda fr: r_star(fr['h'], fr.name[1]), axis = 1)
# frm['sstar'] = frm.apply(lambda fr: s_star(fr['h'], fr.name[1]), axis = 1)
```

```python
canvas1 = Canvas(size = (8, 8/3), shape = (1, 3))
ax1 = canvas1.make_ax((0, 0))
ax2 = canvas1.make_ax((0, 1))
ax3 = canvas1.make_ax((0, 2))
# ax4 = canvas.make_ax((1, 1))
# extract = []

allx, ally = [], []

for (H, f), values in frm.iterrows():

    if f == 1:
        f = 0.99999
    h, T = values['h'], values['geotherm']

    ax1.line(
        Channel(T / H, lims = (0, 0.5), label = 'T/H'),
        Channel(h, lims = (0, 1), label = 'h'),
        color = cmap(f, fs, style = 'turbo'),
        linewidth = H / 5
        )

    dT, hdT = analysis.derivative(T, h, n = 1)
    ax2.line(
        phichan := Channel(dT * s_star(hdT, f) / H, lims = (-1, 0), label = r"\phi_q/H"),
        Channel(hdT, lims = (0, 1), label = 'h'),
        color = cmap(f, fs, style = 'turbo'),
        linewidth = H / 5
        )
    ax3.line(
        phichan,
        dchan := Channel(-sub_area(hdT, f), label = r"-D"),
        color = cmap(f, fs, style = 'turbo'),
        linewidth = H / 5
        )
    ally.extend(dchan.data)
    allx.extend(phichan.data)

# ax2.props.edges.y.label.visible = False
# ax2.props.edges.y.ticks.major.labels = []

linscore = r2_score(ally, allx)
ax3.line(
    np.linspace(-1, 0, 10),
    np.linspace(-1, 0, 10),
    color = '#ff7f0e',
    linestyle = '--',
    )
trendlabel = r'y=x, \\ R^2 =' + str(round(linscore, 8))
ax3.annotate(
    -0.5,
    -0.5,
    label = trendlabel,
    points = (15, -45),
    arrowProps = dict(arrowstyle = "->", color = '#ff7f0e'),
    )

canvas2 = Canvas(size = (8, 4), shape = (1, 2))
ax1 = canvas2.make_ax((0, 0))
ax2 = canvas2.make_ax((0, 1))
allx, ally = [], []
for (H, f), values in frm.iterrows():
    if f == 1:
        f = 0.99999
    h, T = values['h'], values['geotherm']
    dT, hdT = analysis.derivative(T, h, n = 1)
    ax1.line(
        xchan := Channel(dT / H, lims = (-1, 0), label = r'\frac{\delta T}{\delta h H}'),
        Channel(hdT, lims = (0, 1), label = 'h'),
        color = cmap(f, fs, style = 'turbo'),
        linewidth = H / 5
        )
    ax2.line(
        xchan,
        ychan := Channel(
            -sub_area(hdT, f) / s_star(hdT, f),
            lims = (-1, 0), label = r"-\frac{D}{s^{*}}"
            ),
        color = cmap(f, fs, style = 'turbo'),
        linewidth = H / 5
        )
    ally.extend(ychan.data)
    allx.extend(xchan.data)
linscore = r2_score(ally, allx)
ax2.line(
    np.linspace(-1, 0, 10),
    np.linspace(-1, 0, 10),
    color = '#ff7f0e',
    linestyle = '--',
    )
trendlabel = r'y=x, \\ R^2 =' + str(round(linscore, 10))
ax2.annotate(
    -0.5,
    -0.5,
    label = trendlabel,
    points = (15, -45),
    arrowProps = dict(arrowstyle = "->", color = '#ff7f0e'),
    )
# ax2.props.edges.y.label.visible = False
# ax2.props.edges.y.ticks.major.labels = []

ax1.props.legend.set_handles_labels(
    (row[0] for row in ax1.collections[10::11]),
    (str(f) for f in fs),
    )
ax1.props.legend.title.text = 'f'
ax1.props.legend.title.visible = True
# ax1.props.legend.mplprops['bbox_to_anchor'] = (1.75, 1.05)
# ax1.props.legend.mplprops['ncol'] = 2
ax1.props.legend.frame.colour = 'black'
ax1.props.legend.frame.visible = True

fig = imop.vstack(canvas1, canvas2)

glue("isocondinternal", fig, display = False)
fig

# ax2.props.legend.set_handles_labels(
#     (row[0] for row in ax1.collections[10::len(Hs)]),
#     (str(f) for f in fs),
#     )
# ax2.props.legend.title.text = 'f'
# ax2.props.legend.title.visible = True
# #     ax2.props.legend.mplprops['bbox_to_anchor'] = (1.75, 1.05)
# # ax1.props.legend.mplprops['ncol'] = 2
# ax2.props.legend.frame.colour = 'black'
# ax2.props.legend.frame.visible = True
```

```{glue:figure} isocondinternal
:figwidth: 800px
:name: "isocondinternal"

Summary of the scaling behaviours of isoviscous conduction under internal heating $H$ for varying curvature parameter $f$ (colours as in previous charts). While samples of varying heat have been plotted, they do not appear in these charts due to the intentional factoring out of $H$, demonstrating that this parameter is a simple coefficient.
```

It was established previously that, for an internally heated system, the geotherm and geothermal gradient are represented by:

$$ \begin{align*}
{T(h)}_{c(internal)} &= \frac{H}{2} \left( 1 - h^2 \right) \\
{{T(h)}_{c(internal)}}^{'} &= H\cdot h
\end{align*} $$

This is intuitive because the source flux visible to each layer is proportional to the area below that layer, which goes linearly with height $h$ in a planar domain.

In the annulus, though, the proportion of the domain beneath a given height $h$ is instead represented by $D$, as we have shown. If we further assume that $H$ is non-dimensionalised so as to represent the total flux of the model (i.e. it equals $1$ for all geometries), then the flux through each layer height $h$ of the annulus must simply be:

$$
{\phi_q}(h) = -H \cdot D(h)
$$

We show that this holds exactly {numref}`isocondinternal`.

As before, the geothermal gradient required to transmit this flux must account for the varying layer length $s^{*}$ - a function of $h$ and the $f$ parameter. Thus:

$$
\frac{dT}{dh} \propto \frac{\phi_q}{s^{*}} = -\frac{HD}{s^{*}}
$$

The integral with respect to $h$ yields the geotherm:

```python
canvas = Canvas(size = (8, 4), shape = (1, 2))
ax1 = canvas.make_ax((0, 0))
ax2 = canvas.make_ax((0, 1))
for (H, f), values in frm.iterrows():
    if f == 1:
        continue
    h, T = values['h'], values['geotherm']
    dT, hdT = analysis.derivative(T, h, n = 1)
    sstar = s_star(hdT, f)
#     phi = dT * sstar / H
    D = -sub_area(hdT, f)
    ax1.line(
        hchan := Channel(h, lims = (0, 1), label = 'h'),
        T / H,
        color = cmap(f, fs, style = 'turbo'),
        linewidth = H / 5
        )
    ax2.line(
        hchan,
        T / H ,
        color = cmap(f, fs, style = 'turbo'),
        linewidth = H / 5
        )
# ax2.props.edges.y.label.visible = False
# ax2.props.edges.y.ticks.major.labels = []
canvas
```

## Mixed heating in the annulus

### Conductive solution

```python
with open(osjoin(aliases.storagedir, 'condhfmixed.pkl'), mode = 'rb') as file:
    conddata = pickle.loads(file.read())
condhs, condfs = zip(*conddata['hfs'])
condhs = tuple(round(val, 1) for val in condhs)
frm = pd.DataFrame(dict(H = condhs, f = condfs, T = conddata['avts'], geotherm = conddata['geotherms']))
frm = frm.set_index(['H', 'f'])
Hs, fs = (np.array(sorted(set(frm.index.get_level_values(level)))) for level in ('H', 'f'))
```


# "Bayesian Model Selection"
> "We illustrate the Model Selection problem of figuring out the number of peaks in a given signal"

- toc:true
- branch: master
- badges: true
- comments: true
- author: John J. Molina
- categories: [Data Analysis, Model Selection]

# Motivation

This is the third post, dealing with the problem of Model Selection and corresponding to Chapter 4 of Sivia and Skilling's book. The problem we wish to tackle is the following: Given a signal, how many peaks is there most evidence for, and what are the positions and amplitudes of those peaks. If our prior information included knowledge of the number of peaks, figuring out their positions and amplitudes would be a simple excercise in parameter estimation. However, here we are faced with the more challenging task of choosing between different models (i.e., the one with $M=2$, $M=3$,... number of peaks).

# References
This topic is discussed in some detail in Sivia's book (Ch. 4, Least-Squares Extensions), as well as Jaynes' (Ch. 20, Model Comparison).
- Data Analysis : A Bayesian Tutorial, second edition. D.S. Sivia with J. Skilling, Oxford, Oxford University Press (2006)
- Probability Theory: The Logic of Science. E. T. Jaynes, Cambridge, Cambridge Unviversity Press (2003)

In addition, Sivia also has several papers covering this material, with an emphasis on analyzing spectral data
- [Molecular spectroscopy and Bayesian spectral analysis—how many lines are there? ](https://aip.scitation.org/doi/10.1063/1.462505)
D. S. Sivia and C. J. Carlile, J. Chem. Phys. 96, 170 (1992)
- [Bayesian analysis of quasieleastic neutron scattering data](https://www.sciencedirect.com/science/article/abs/pii/092145269290036R)
 D. S. Sivia, C. J. Carlile, W. S. Howells, and S. Konig, Physica B 182, 341 (1992)
- [A Bayesian approach to extracting structure-factor amplitudes from powder diffraction data](http://scripts.iucr.org/cgi-bin/paper?S0108767394003235)
 D. S. Sivia and W. I. F. David, Acta Cryst. A 50, 703 (1994)
- [The Bayesian approach to reflectivity data](https://www.sciencedirect.com/science/article/abs/pii/S0921452698002592)
D. S. Sivia and J. R. P. Webster, Physica B 248, 327 (1998)

# Introduction

Let's say we have to decide between two models, $M=1$ and $M=2$, given some experimental data $D$. How to decide between them? We should look at the posterior ratio

\begin{align}
O_{21} &\equiv \frac{P(M_2 \lvert D, I)}{P(M_1 \lvert D, I)} \\
&= \frac{\frac{P(D\lvert M_1, I) P(M_1\lvert I)}{P(D\lvert I)}}{ \frac{P(D\lvert M_2, I) P(M_2\lvert I)}{P(D\lvert I)}}\\
&= \frac{P(D | M_2)}{P(D|M_1)}\cdot \frac{P(M_2)}{P(M_1)} \\
&= \frac{Z_2}{Z_1} \cdot \frac{\Pi_2}{\Pi_1}
\end{align}
where $Z_i = P(D|M_i,I)$ is the evidence for model $M_i$, and $\Pi_i = P(M_i\lvert I)$ the prior. This ratio of likelihoods $Z_2/Z_1$ is also known as the 'Bayes Factor'.

We can compute the likelihood $P(D|M_i, I)$ by marginalizing over the parameters $\Theta=\Theta_i$ of model $M_i$. To simplify the notation we drop the index $i$, but it should be clear that all quantities are defined on a per/model basis

\begin{align}
Z\equiv P(D | M, I) &= \int\mathrm{d}\Theta\, P(D, \Theta\lvert M, I) \\
&= \int\mathrm{d}\Theta\, P(D \lvert \Theta, M, I) \cdot P(\Theta\lvert M, I) \\
&= \int\mathrm{d}\Theta\, L(\Theta_i) \cdot \Pi(\Theta)
\end{align}
with $L(\Theta) = P(D| \Theta, M, I)$ and $\Pi(\Theta) = P(\Theta\lvert M, I)$ the likelihood of the data and prior for the parameters of model $M$, respectively.

Recall that the integrand that appears here is proportional to posterior for $\Theta$, which is the main quantity of interest in a parameter estimation problem
\begin{align}
P(\Theta\lvert D, M, I) &= \frac{\overbrace{P(D\lvert \Theta, M, I)}^{\text{Likelihood}\, L} \overbrace{P(\Theta\lvert M, I)}^{\text{Prior} \,\Pi}}{\underbrace{P(D\lvert M, I)}_{\textrm{Evidence} \,Z}} \\
&= \frac{L(\Theta) \Pi(\Theta)}{Z}
\end{align}
Where the evidence integral was usually ignored since it is just a normalizing factor. However, we now see that for a Model Selection problem these $Z$ are the main objects of interest.

Evaluating this integral by brute force will only be possible for very simple models (low dimensions). However, following Jaynes (ch. 20), we can rewrite these factors into a more telling form.

Let $L^\star = L(\Theta^\star) = \max_{\Theta}{L}$, the evidence is now
\begin{align}
Z &= L^\star\underbrace{\int\text{d}\Theta\, \frac{L(\Theta)}{L^\star} \cdot \Pi(\Theta)}_{\textrm{Ockham factor} \,W} = L^\star W\\
\end{align}
where $W$ is the so-called 'Ockham factor'. It is this term which will protect us from over-fitting. By definition, the prior $\Pi(\Theta) = P(\Theta | M)$ integrates to one, so we see that the factor of $L/L^\star$ is picking out the fraction of prior probability mass contained in the high-likelihood regions. Quoting from Jaynes:

 "...the Ockham factor is the ratio by which the parameter space is contracted by the information in the data, which expresses how much the vagueness of our prior infomration deteriorates the performance of [the model], by placing prior probability outside its high-likelihood region"

Finally, we can write the Bayes factor appearing in the posterior ratio as
\begin{align}
O_{21} &= \underbrace{\frac{L_2^\star}{L_1^\star}}_{\textrm{Goodness of fit}}\cdot\underbrace{\frac{W_2^{\phantom{\star}}}{W_1^{\phantom{\star}}}}_{\textrm{Model robustness}}\cdot\underbrace{\frac{\Pi_2^{\phantom{\star}}}{\Pi_1^{\phantom{\star}}}}_{\textrm{Prior ratio}}
\end{align}

To see how this could favor one model over another, let's use the Laplace approximation to try and evaluate the integral in the Ockham factor.
\begin{align}
W &= \int\textrm{d}\Theta e^{\mathcal{L}} \Pi(\Theta) \\
\mathcal{L}&= \ln{\frac{L(\Theta)}{L^\star}}
\end{align}
expanding $\mathcal{L}$ around the maximum, $\Theta^\star$, we have
\begin{align}
\mathcal{L}&\simeq\mathcal{L}^\star + \left.\nabla\mathcal{L}\right\lvert_{\Theta^\star}\cdot (\Theta - \Theta^\star) + \frac{1}{2}(\Theta-\Theta^\star)^{t}\cdot\left.\nabla\nabla\mathcal{L}\right\lvert_{\Theta^*}\cdot(\Theta-\Theta^\star) \\
&= \frac{1}{2}(\Theta - \Theta^\star)^t\cdot \left.\nabla\nabla\mathcal{L}\right\lvert_{\Theta^\star}(\Theta-\Theta^\star)
\end{align}
since $\mathcal{L}^\star = \ln{L^\star/L^\star} = 0 = \left.\nabla\mathcal{L}\right\lvert_{\Theta^\star}$ by definition.

Thus, we can write the Ockham factor as
\begin{align}
W &= \int\text{d}\Theta\, \exp{\left[-\frac{1}{2}(\Theta-\Theta^\star)^t\cdot\Sigma^{-1}\cdot(\Theta-\Theta^\star)\right]} \Pi(\Theta) \\
\Sigma &= \left(\left.-\nabla\nabla\mathcal{L}\right\lvert_{\Theta^\star}\right)^{-1}
\end{align}

Now assume that $\Pi(\Theta)$ is constant, $\Pi(\Theta) = 1/\Omega$, with $\Omega$ the allowed volume in parameter space.
\begin{align}
W &\simeq \frac{\sqrt{(2\pi)^p \det{\Sigma}}}{\Omega}
\end{align}
with $p$ the number of parameters for this model, $\Theta = (\Theta_1, \Theta_2, \ldots, \Theta_p)$.

Recall that $\Sigma$ is the covariance matrix, so that $\det\Sigma$ is related to the volume of the parameter region that is consistent with the data.
In this case, larger error bars are actually a good thing! They are telling us that the model allows for more values that can explain the data, and thus is more robust, and should be preferred.

# How many lines are there?

Let our parameters be denoted as

\begin{align}
M   &: \textrm{Number of Line Peaks}\\
A_i &: \textrm{Amplitude of peak} \,i\\
X_i &: \textrm{Position of peak}\, i\\
b_1 &: \textrm{Intercept for linear background signal}\\
b_2 &: \textrm{Slope for linear background signal}\\
D=\{x_k,y_k\}   &: \textrm{Experimental data} \\
\end{align}

With the following prior information

\begin{align}
I &: \textrm{Lines peaks are measured as gaussians of fixed width } W\textrm{, subject to Gaussian noise of variance }\sigma, \textrm{ in the presence of a linear background signal}.
\end{align}

The inference problem to solve is then, given a spectra, how many peaks are there (at what positions and with what magnitudes)?

The posterior probability we are interested in is $P(M\lvert D I)$, i.e., the probability that the signal corresponds to an $M$-peak model. From Bayes theorem, we have the following

\begin{align}
P(M, D \lvert I) &= P(M\lvert D,I)P(D\lvert I) = P(D\lvert M,I)P(M\lvert I) \\
P(M \lvert D,I) &= \frac{P(D\lvert M,I) P(M\lvert I)}{P(D\lvert I)} \\
\end{align}

Now, if we assign uniform priors for $P(M\lvert I) = 1/M_{\text{max}}$, with $M=1,\ldots, M_{\text{max}}$, the posterior becomes proportional to the evidence
\begin{align}
P(M\lvert D, I) &\propto P(D\lvert M, I) = Z_M
\end{align}
where the normalization constant can be evaluated later, since
\begin{align}
\sum_{M=0}^{M_{\text{Max}}} P(M\lvert D, I)&= \alpha \sum_M P(D\lvert M, I)= 1 \\
\alpha^{-1} \equiv M_{\text{max}} P(D\lvert I) &= \sum_M P(D\lvert M, I)
\end{align}

To evaluate the evidence , we "simply" marginalize over the model parameters,

\begin{align}
Z_M &= \int\textrm{d}b_1\int \textrm{d} b_2 \int \textrm{d}A^{M}\int\textrm{d}X^M P(D, b_1, b_2, \{A_i, X_i\}\lvert M, I) \\
&= \int\textrm{d}b_1\int\textrm{d}b_2\int\textrm{d}A^{M}\int\textrm{d}X^M P(D\lvert b_1, b_2, \{A_i,X_i\}, M, I) P(b_1, b_2, \{A_i,X_i\}\lvert M, I)
\end{align}

The prior can be decomposed as
\begin{align}
P(b_1, b_2, \{A_i, X_i\}\lvert M, I) &= P(b_1, b_2\lvert \{A_i, X_i\}, M, I) \, P(\{A_i, X_i\}\lvert M, I) \\
&= P(b_1, b_2\lvert I) \, P(\{A_i, X_i\}\lvert M, I)
\end{align}
where we use the fact that the background signal parameters do not depend on the number of lines (or peak parameters). Thus, this prior $P(b_1,b_2\lvert I)$ can be absorbed into the normalization constant $\alpha$.

For the peak position and amplitude parameters we assume independent uniform priors, such that $X_{\textrm{min}}\le X \le X_{\textrm{max}}$ and $0\le A \le A_{\textrm{max}}$
\begin{align}
P(\{A_i,X_i\}\lvert M, I) &= \prod_{i} P(A_i,X_i\lvert M,I) \\
&= \prod_i P(A_i\lvert M,I) P(X_i\lvert M, I) \\
&= \left(A_{\text{max}} \left(X_{\text{max}} - X_{\text{min}}\right)\right)^{-M}
\end{align}

Therefore, the posterior becomes

\begin{align}
P(M\lvert D, I) &\propto \frac{1}{\left[A_{\text{max}} \left(X_\text{max} - X_\text{min}\right)\right]^M}\int\textrm{d}b_1\int\textrm{d}b_2 \int\text{d}A^M\int\text{d}X^M P(D\lvert b_1, b_2, \{A_i,X_i\}, M, I)
\end{align}

Our prior information says that the measurements are independent, subject to Gaussian noise of amplitude $\sigma_k$. Then, the Likelihood for the data, given the M-peak model and its parameters becomes

\begin{align}
P(D\lvert b_1, b_2, \{A_i, x_i\}, M, I) &\propto P(\{D_k\}\lvert b_1, b_2, \{A_i, x_i\}, M, I)\\
&= \prod_k P(D_k\lvert b_1, b_2, \{A_i, x_i\}, M, I) \\
&= \Pi_k \frac{1}{\sqrt{2\pi} \sigma_k} \exp{\left[-\frac{1}{\sqrt{2\pi}} \left(\frac{y_k - f(x_k; b_1, b_2, \{A_i, X_i\})}{\sigma_k}\right)^2\right]} \\
&\propto \exp{\left(-\frac{\chi^2}{2}\right)}\\
\chi^2 &= \sum_k\left(\frac{y_k - f_k}{\sigma_k}\right)^2
\end{align}
where the factors of $\sqrt{2\pi\sigma_k^2}$ can be absorbed into the normalization constant since they don't depend on $M$, but only the number of data points $N$ (the $\sigma_k$ are not model parameters).

We therefore have the following for the posterior (evidence)

\begin{align}
P(M\lvert D, I) \propto \frac{1}{\left[A_{\text{max}} \left(X_{\text{max}} - X_{\text{min}}\right)\right]^M}\int\textrm{d}b_1\int\textrm{d}b_2\int\textrm{d}A^M\int\textrm{d}X^M \exp{\left(-\frac{\chi^2}{2}\right)}
\end{align}

A brute force approach to this $2M + 2$ dimensional integral is out of the question. However, we can proceed analytically by using the Laplace approximation, expanding around the minimum of $\chi^2$.
Let $\Theta$ denote the model parameters, $\Theta = (b_1, b_2, \{A_i\}, \{X_i\})$, with $\Theta_0$ the minimum, such that

\begin{align}
\chi^2&\simeq \chi^2_0 + \frac{1}{2}\left(\Theta - \Theta_0\right)^t\cdot \left.\nabla\nabla \chi^2\right\lvert_{\Theta_0} \cdot\left(\Theta - \Theta_0\right) \\
&= \chi^2_0 + \left(\Theta-\Theta_0\right)^t\cdot\Sigma^{-1}\cdot\left(\Theta-\Theta_0\right)\\
\Sigma^{-1} &= \frac{1}{2}\left.\nabla\nabla\chi^2\right\lvert_{\Theta_0}
\end{align}

Putting this expression back into the evidence integral gives us

\begin{align}
P(M\lvert D, I) &\propto \frac{1}{\left[A_{\text{max}}\left(X_{\text{max}} - X_{\text{min}}\right)\right]^M} e^{-\chi^2_0/2}
\int\textrm{d}\Theta \exp{\left[-\frac{1}{2}\left(\Theta-\Theta_0\right)^t\cdot\Sigma^{-1}\cdot\left(\Theta-\Theta_0\right)\right]} \\
&\propto \frac{M!}{\left[A_{\text{max}}\left(X_{\text{max}}- X_{\text{min}}\right)\right]^M} e^{-\chi^2_0/2} \sqrt{(2\pi)^{2M+2}\det{\Sigma}}
\end{align}
where the factor of $M!$ is needed to account for all possible permutations of the peaks (i.e., the labels we use are arbitrary and should not matter).

The determinant of the Hessian matrix can be simplified as follows,
\begin{align}
\det{\Sigma} &= \frac{1}{\det{\Sigma^{-1}}} = \frac{1}{\det{\left(\frac{1}{2} \left.\nabla\nabla\chi^2\right\lvert_{\Theta}\right)}} = \frac{2^{2M+2}}{\det{\left(\left.\nabla\nabla\chi^2\right\lvert_{\Theta_0}\right)}}
\end{align}
from which we finally obtain
\begin{align}
P(M\lvert D, I) &\propto \frac{M!}{\left[A_{\text{max}}\left(X_{\text{max}}- X_{\text{min}}\right)\right]^M} e^{-\chi^2_0/2} \sqrt{\frac{(2\pi)^{2M+2} 2^{2M + 2}}{\det{\left(\left.\nabla\nabla\chi^2\right\lvert_{\Theta_0}\right)}}} \\
&=  \frac{M! (4\pi)^{M+1}}{\left[A_{\text{max}}\left(X_{\text{max}}- X_{\text{min}}\right)\right]^M \sqrt{\det{\left(\left.\nabla\nabla\chi^2\right\lvert_{\Theta_0}\right)}}} e^{-\chi^2_0/2}
\end{align}

We dropped all the proportionality constants that did not depend on $M$, but we can now recover them by enforcing the normalization condition, $\sum_k P(M\lvert D, I) = 1$.

As always, it is more convenient to work with the logarithm of the posterior,

\begin{align}
\ln{P(M\lvert D, I)}&= -\frac{\chi_0^2}{2} + \ln{M!} + M\Big[\ln{4\pi} - \ln\big(A_{\text{max}}\left(X_{\text{max}}- X_{\text{min}}\right)\big)\Big] - \frac{1}{2}\ln{\left[\det{\left(\left.\nabla\nabla\chi^2\right\lvert_{\Theta_0}\right)}\right]} + \ln{4\pi}
\end{align}

Thus, within the approximations we have made, the model selection problem can be effectively reduced to a series of minimization problems (one for each of the models). Due to the non-linearity, Sivia (Ch. 4.2.1) advocates for the following numerical algorithm

* ($M=0$)
    1. Find the optimum values for the 2d problem (b_1, b_2)
* ($M=1$)
    1. Perform a 1d scan over the new peak position $x_1$, by minimizing with respect to $(b_1, b_2, A_1)$, using the optimum values for ($M=0$) as initial guess. For the initial $A_1$, we can use the minimum of signal as a guide. Take the $x_1$ value resulting in the lowest $\chi^2$ as the best guess.
    1. Perform a full minimization $(b_1, b_2, A_1, x_1)$, using the optimum values obtained at step (1.) as an initial guess.
* ($M=2$)
    1. Perform a 1d scan over the new peak position $x_2$, by minimizing with respect to the background and amplitude parameters, $(b_1, b_2, A_1, A_2)$. At each point in this scan, the positions of the peaks are fixed to ($x_1^{(M=1)}, x_2$), with the initial guess taken from the optimum value for $M=1$. Take the $x_2$ value resulting in the lowest $\chi^2$ as the best guess.
    1. Perform a full minimization $(b_1, b_2, A_1, A_2, x_1, x_2)$, using the optimum values obtained at step (1.) as an initial guess.

$\vdots$

* ($M=m$)
    1. Perform a 1d scan over the new peak position $x_m$, by minimizing with respect to the background and amplitude parameters, $(b_1, b_2, A_1, A_2, \ldots A_m)$. At each point in this scan, the positions of the peaks are fixed to ($x_1^{(M=m-1)}, \ldots, x_{m-1}^{(M=m-1)}, x_m$), with the initial guess taken from the optimum value for $M=m-1$. Take the $x_m$ value resulting int he lowest $\chi^2$ as the best guess.
    1. Perform a full minimization $(b_1, b_2, A_1, \ldots , A_m, x_1, \ldots, x_m)$, using the optimum values obtained at step (1.) as an initial guess.

Additional hints and best practices:
* Scale all parameters to be unity and dimensionless
* Add small multiples of identity to the Hessian matrix: $\nabla\nabla\chi^2 \longrightarrow \nabla\nabla\chi^2 + \epsilon I$

## Numerical Example

```python
import jax
import jax.numpy as np
import numpy as onp
import matplotlib as mpl
import matplotlib.pyplot as plt
import pandas as pd
import pickle

from numpy import random
from scipy import integrate
from scipy import optimize
from scipy.optimize import minimize
from jax import grad, jit, vmap, jacfwd, jacrev
from jax.config import config; config.update("jax_enable_x64", True)
from functools import partial,reduce
import matplotlib.patheffects as PathEffects

mpl.style.use(['seaborn-poster', 'seaborn-muted'])

#betanalphas' colormap
colors = ["#DCBCBC","#C79999","#B97C7C","#A25050","#8F2727", "#7C0000","#DCBCBC20", "#8F272720","#00000060"]
color  = {i[0]:i[1] for i in zip(['light','light_highlight','mid','mid_highlight','dark','dark_highlight','light_trans','dark_trans','superfine'],colors)}

def addtxt(ax, x, y, txt, fs=8, lw=3, clr='k', bclr='w', rot=0):
    """Add text to figure axis"""
    return ax.text(x, y, txt, color=clr, ha='left', transform=ax.transAxes, rotation=rot, weight='bold',
                   path_effects=[PathEffects.withStroke(linewidth=lw, foreground=bclr)], fontsize=fs)

def hessian(f):
    """Returns a function which computes the Hessian of a function f
           if f(x) gives the values of the function at x, and J = hessian(f)
           J(x) gives the Hessian at x"""
    return jit(jacfwd(jacrev(f)))
```

Let's define our ideal peak and background signals

\begin{align}
f_{\mathrm{background}}(x; b_1, b_2) &= b_1 + b_2 x \\
f_{\text{peak}}(x; A, X) &= A \exp{\left(-\frac{(x-X)^2}{2 w^2}\right)}
\end{align}

with $w=0.01$ the width of the peaks. Then, for an $M$-peak model, the ideal signal would be given as
\begin{align}
f(x; b_1, b_2, \{A_i, X_i\}) &= f_{\mathrm{background}}(x;b_1, b_2) + \sum_{i=1}^M f_{\text{peak}}(x; A_i, X_i)
\end{align}

```python
def ideal_bckg(b1,b2,x):
    """Linear Backbround Signal"""
    return b1 + b2*x

def ideal_peak(xi,ai,x):
    """Ideal Gaussian Peak"""
    return ai*np.exp(-(x-xi)**2/(2*0.01**2))

def squash(b1,b2,pos,amp):
    """Return signal parameters as single array"""
    n = len(pos)
    p = np.empty(2 + 2*n, dtype=pos.dtype)
    p = jax.ops.index_update(p, jax.ops.index[0],b1)
    p = jax.ops.index_update(p, jax.ops.index[1],b2)
    p = jax.ops.index_update(p, jax.ops.index[2::2],pos)
    p = jax.ops.index_update(p, jax.ops.index[3::2],amp)
    return p
def unsquash(p):
    """Return signal parameters independently"""
    return p[0], p[1], p[2::2], p[3::2]

def ideal_signal0(p, x):
    """Ideal (noiseless) signal"""
    b1,b2,xi,ai = unsquash(p)
    return ideal_bckg(b1,b2,x) + np.sum(ideal_peak(xi,ai,x), axis=-1)

ideal_signal = jit(vmap(ideal_signal0, in_axes=(None, 0), out_axes=0))
```

Let's generate some synthetic data for a five peak signal (roughly corresponding to Fig. 4.4 in Sivia's book)

```python
peaks_xi = (np.array([11.5, 11.55, 12.15, 12.8, 13.5])-11.0)/3.0
peaks_ai = np.array([1.8, 7.0, 3.5, 3.8, 8.9])/10
peaks    = np.vstack([peaks_xi, peaks_ai]).transpose()
line_b1  = 0.3
line_b2  = 0.02
params   = squash(line_b1, line_b2, peaks_xi, peaks_ai)
```

    /opt/anaconda3/envs/ML/lib/python3.7/site-packages/jax/lib/xla_bridge.py:123: UserWarning: No GPU/TPU found, falling back to CPU.
      warnings.warn('No GPU/TPU found, falling back to CPU.')

```python
fig, [ax,bx] = plt.subplots(figsize=(18,9),ncols=2)
theta        = np.linspace(0, 1., num=300)
signal       = ideal_signal(params, theta)

ax.plot(theta, signal, color=color['mid'])
ax.plot(theta, ideal_bckg(line_b1, line_b2, theta), color=color['mid_highlight'], ls=':', lw=4)
bx.plot(theta-0.5, ideal_peak(theta, 1, 0.5), color=color['mid'])
for x0,a0 in zip(peaks_xi, peaks_ai):
    baseline = ideal_bckg(line_b1, line_b2, x0)
    ax.vlines(x0, baseline, baseline + a0, alpha=0.8, color=color['superfine'])
ax.set_title(r'Pure Signal')
bx.set_title(r'Ideal Peak')
bx.set_xlim(-0.06, 0.06)
ax.set_ylim(0.2, 1.3)
fig.tight_layout()
plt.show()
```

Now let's add some noise to this ideal signal, to see how our peak finding algorithm behaves.
We generate 5 noisy datasets, with increasing amplitude of the signal noise.

```python
def noisy_signal(μ, σ, n):
    """Given ideal signal μ, return simulated measurement, averaged over n samples, with variance σ"""
    counts = onp.random.normal(μ, σ, size=(n,)+μ.shape)
    avg    = onp.average(counts, axis=0)
    err    = onp.std(counts, axis=0)
    return pd.DataFrame({'avg':avg, 'err':err})

random.seed(12345)
σs    = np.linspace(0.02, 0.1, num=5)
datas = [noisy_signal(signal, σ, 100) for σ in σs]
```

```python
fig, axes = plt.subplots(figsize=(18,18), ncols=2, nrows=3, sharex=True, sharey=True)
for ax,data,σ in zip(axes.flatten(), datas, σs):
    ax.errorbar(theta, data['avg'], data['err'], ls='None', lw=1.5, alpha=0.8, color=color['mid_highlight'])
    ax.plot(theta, data['avg'] , marker='o', ms=6, mfc='None', mew=1.2, color=color['dark'], ls='None')
for ax,σ in zip(axes.flatten()[:-1], σs):
    for x0,a0 in zip(peaks_xi, peaks_ai):
        baseline = ideal_bckg(line_b1, line_b2, x0)
        ax.vlines(x0, baseline, baseline + a0, alpha=0.8, color=color['superfine'])
        addtxt(ax, 0.1, 0.9, f'$\sigma = {σ:3.2f}$', fs=18)
fig.tight_layout()
```

To find the optimum parameters, we minimize the $\chi^2$,
\begin{align}
\chi^2 &= \sum_k\left(\frac{y_k - f_k}{\sigma_k}\right)^2
\end{align}

However, keep in mind that the algorithm outlined above requires three seperate minimizations:
* Background only
    \begin{align}
    f_k = f_{\text{background}}(x_k; b_1, b_2)
    \end{align}
* Background and peak amplitudes
    \begin{align}
    f_k = f(x_k; b_1, b_2, \{A_i, \underline{X_i}\})
    \end{align}
    where the underline refers to parameters that are held fixed (i.e., the peak positions)
* Full minimization
    \begin{align}
    f_k = f(x_k; b_1, b_2, \{A_i,X_i\})
    \end{align}

In practice we minimize the average deviations, i.e., $\chi^2/N$. So in the end, we should not forget to rescale both $\chi^2$ and the Hessian $\nabla\nabla\chi^2$.

Let's define the corresponding loss functions. Thanks to the fact that we are using jax, we get the gradients and the hessians for free (basically). This greatly improves the convergence of the minimization routines.

```python
@jit
def loss_bckg(p, x, y, dy):
    b1,b2 = p
    ymodel = ideal_bckg(b1, b2, x)
    return np.mean(((y-ymodel)/dy)**2)
dloss_bckg = jit(grad(loss_bckg))

@jit
def loss(p, x, y, dy):
    ymodel = ideal_signal(p, x)
    return np.mean(((y - ymodel)/dy)**2)
dloss    = jit(grad(loss))
hess_loss= hessian(loss)

@jit
def loss_fixpos(p, pos, x, y, dy):
    param  = squash(p[0], p[1], pos, p[2:])
    return loss(param, x, y, dy)
dloss_fixpos = jit(grad(loss_fixpos))
```

Thanks to the fact that we are using jax, we get the gradients and the hessians for free (basically). This greatly improves the convergence of the minimization routines.

The peak finding algorithm can then be written as follows

```python
def findNewPeak(x,y,dy,old_guess):
    """Determine number of positions of peaks in a given signal
    x           : measured x data
    y           : measured y data
    dy          : measured y uncertainty
    """
    ### 1D scan over x to locate peak candidate
    def fitBckAmp(init, peak_pos, x, y, dy):
        """Optimize background and amplitude parameters given position of the peaks
        x    : measured x data
        y    : measured y data
        dy   : measured y uncertainty
        peak_pos : fixed position of peaks
        init     : inital guess for background (b, m) and peak amplitudes (ai) """
        opt0 = optimize.minimize(loss_fixpos, init, args=(peak_pos, x, y, dy), method='BFGS', jac=dloss_fixpos, options={'disp':0})
        return opt0
    old_b1, old_b2, old_pos, old_amp = unsquash(old_guess)
    old_n = len(old_pos)

    ####
    #### initial guess for background and amplitude parameters
    ####
    init0 = np.zeros(2 + old_n + 1)
    init0 = jax.ops.index_update(init0, jax.ops.index[:2], old_guess[:2])     # b_1,b_2
    init0 = jax.ops.index_update(init0, jax.ops.index[2:2+old_n], old_amp)    # a_i (i=1,n-1)
    init0 = jax.ops.index_update(init0, jax.ops.index[old_n], np.minimum(np.min(old_amp)*0.25, np.min(y))) #a_n
    # fixed position parameters
    pos0  = np.zeros(old_n + 1)
    pos0  = jax.ops.index_update(pos0, jax.ops.index[:old_n], old_pos)
    opt0  = [fitBckAmp(init0,jax.ops.index_update(pos0, jax.ops.index[old_n], x0),x,y,dy) for x0 in x]  # scan over x, optimizing (b1,b2,ai)
    idx0  = np.argmin(np.array([dmy['fun'] for dmy in opt0]))                                           # best guess for new peak @ min(chi2)
    p0    = opt0[idx0]['x']                                                                             # optimized (b1,b2,ai) parameters

    ###
    ### Full optimization in 2n + 2 parameters space
    ###
    pos0 = jax.ops.index_update(pos0, jax.ops.index[old_n], x[idx0])
    init = squash(p0[0], p0[1], pos0, p0[2:])

    opt1 = optimize.minimize(loss, init, args=(x,y,dy), method='Nelder-Mead')
    opt  = optimize.minimize(loss, opt1['x'], args=(x,y,dy), method='BFGS', jac=dloss)

    opt['chi2'] = opt['fun']*len(x)
    opt['H']    = hess_loss(opt['x'], x, y, dy)*len(x)
    opt['H_inv']= np.linalg.inv(opt['H'])

    #rescale approximate hessian given by bfgs
    opt['hess_inv'] = opt['hess_inv']/len(x)
    opt['hess']     = np.linalg.inv(opt['hess_inv'])

    return opt

def fitBackground(x,y,dy):
    """ Fit background signal """
    opt = optimize.minimize(loss_bckg, np.array([np.min(y), 0.0]), args=(x,y,dy), method='Nelder-Mead')
    return opt['x']

def findPeaks(x,y,dy,Mmax):
    opts     = []
    oldguess = fitBackground(x,y,dy)
    print('*** Searching for Peaks ***')
    for M in range(1,Mmax+1):
        opt  = findNewPeak(x,y,dy,oldguess)
        opts.append(opt)

        oldguess = opt['x']
        print(f"\t M = {M}:\t χ2 = {opt['chi2']:8.3e}, niter = {opt['nit']:4d}, Converged = {opt['success']:3b} : {opt['message']}")
    return opts
```

```python
opts  = [findPeaks(theta, data['avg'].values, data['err'].values, 10) for data in datas]
```

    *** Searching for Peaks ***
    M = 1:	 χ2 = 9.849e+03, niter =    6, Converged =   1 : Optimization terminated successfully.
    M = 2:	 χ2 = 3.237e+03, niter =    9, Converged =   1 : Optimization terminated successfully.
    M = 3:	 χ2 = 1.626e+03, niter =   12, Converged =   1 : Optimization terminated successfully.
    M = 4:	 χ2 = 1.466e+02, niter =   14, Converged =   1 : Optimization terminated successfully.
    M = 5:	 χ2 = 3.382e+00, niter =   16, Converged =   1 : Optimization terminated successfully.
    M = 6:	 χ2 = 3.285e+00, niter =   19, Converged =   1 : Optimization terminated successfully.
    M = 7:	 χ2 = 3.214e+00, niter =   20, Converged =   1 : Optimization terminated successfully.
    M = 8:	 χ2 = 3.159e+00, niter =   24, Converged =   1 : Optimization terminated successfully.
    M = 9:	 χ2 = 3.097e+00, niter =   25, Converged =   1 : Optimization terminated successfully.
    M = 10:	 χ2 = 3.034e+00, niter =   28, Converged =   1 : Optimization terminated successfully.
    *** Searching for Peaks ***
    M = 1:	 χ2 = 2.663e+03, niter =    7, Converged =   1 : Optimization terminated successfully.
    M = 2:	 χ2 = 8.824e+02, niter =   10, Converged =   1 : Optimization terminated successfully.
    M = 3:	 χ2 = 4.246e+02, niter =   12, Converged =   1 : Optimization terminated successfully.
    M = 4:	 χ2 = 3.677e+01, niter =   15, Converged =   1 : Optimization terminated successfully.
    M = 5:	 χ2 = 3.118e+00, niter =   16, Converged =   1 : Optimization terminated successfully.
    M = 6:	 χ2 = 3.051e+00, niter =   18, Converged =   1 : Optimization terminated successfully.
    M = 7:	 χ2 = 3.006e+00, niter =   19, Converged =   1 : Optimization terminated successfully.
    M = 8:	 χ2 = 2.964e+00, niter =   20, Converged =   1 : Optimization terminated successfully.
    M = 9:	 χ2 = 2.928e+00, niter =   20, Converged =   1 : Optimization terminated successfully.
    M = 10:	 χ2 = 2.898e+00, niter =   22, Converged =   1 : Optimization terminated successfully.
    *** Searching for Peaks ***
    M = 1:	 χ2 = 1.181e+03, niter =    6, Converged =   1 : Optimization terminated successfully.
    M = 2:	 χ2 = 4.212e+02, niter =    9, Converged =   1 : Optimization terminated successfully.
    M = 3:	 χ2 = 2.119e+02, niter =   11, Converged =   1 : Optimization terminated successfully.
    M = 4:	 χ2 = 2.007e+01, niter =   14, Converged =   1 : Optimization terminated successfully.
    M = 5:	 χ2 = 3.538e+00, niter =   17, Converged =   1 : Optimization terminated successfully.
    M = 6:	 χ2 = 3.467e+00, niter =   18, Converged =   1 : Optimization terminated successfully.
    M = 7:	 χ2 = 3.424e+00, niter =   18, Converged =   1 : Optimization terminated successfully.
    M = 8:	 χ2 = 3.383e+00, niter =   18, Converged =   1 : Optimization terminated successfully.
    M = 9:	 χ2 = 3.343e+00, niter =   17, Converged =   1 : Optimization terminated successfully.
    M = 10:	 χ2 = 3.304e+00, niter =   18, Converged =   1 : Optimization terminated successfully.
    *** Searching for Peaks ***
    M = 1:	 χ2 = 6.756e+02, niter =    6, Converged =   1 : Optimization terminated successfully.
    M = 2:	 χ2 = 2.144e+02, niter =   10, Converged =   1 : Optimization terminated successfully.
    M = 3:	 χ2 = 1.084e+02, niter =   13, Converged =   1 : Optimization terminated successfully.
    M = 4:	 χ2 = 1.188e+01, niter =   12, Converged =   1 : Optimization terminated successfully.
    M = 5:	 χ2 = 2.853e+00, niter =   13, Converged =   1 : Optimization terminated successfully.
    M = 6:	 χ2 = 2.758e+00, niter =   17, Converged =   1 : Optimization terminated successfully.
    M = 7:	 χ2 = 2.687e+00, niter =   15, Converged =   1 : Optimization terminated successfully.
    M = 8:	 χ2 = 2.630e+00, niter =   14, Converged =   1 : Optimization terminated successfully.
    M = 9:	 χ2 = 2.577e+00, niter =   16, Converged =   1 : Optimization terminated successfully.
    M = 10:	 χ2 = 2.539e+00, niter =   15, Converged =   1 : Optimization terminated successfully.
    *** Searching for Peaks ***
    M = 1:	 χ2 = 4.331e+02, niter =    6, Converged =   1 : Optimization terminated successfully.
    M = 2:	 χ2 = 1.453e+02, niter =   10, Converged =   1 : Optimization terminated successfully.
    M = 3:	 χ2 = 7.171e+01, niter =   10, Converged =   1 : Optimization terminated successfully.
    M = 4:	 χ2 = 8.578e+00, niter =   15, Converged =   1 : Optimization terminated successfully.
    M = 5:	 χ2 = 2.993e+00, niter =   14, Converged =   1 : Optimization terminated successfully.
    M = 6:	 χ2 = 2.903e+00, niter =   15, Converged =   1 : Optimization terminated successfully.
    M = 7:	 χ2 = 2.825e+00, niter =   16, Converged =   1 : Optimization terminated successfully.
    M = 8:	 χ2 = 2.757e+00, niter =   15, Converged =   1 : Optimization terminated successfully.
    M = 9:	 χ2 = 2.713e+00, niter =   17, Converged =   1 : Optimization terminated successfully.
    M = 10:	 χ2 = 2.679e+00, niter =   18, Converged =   1 : Optimization terminated successfully.

\begin{align}
\ln{P(M\lvert D, I)}&= -\frac{\chi_0^2}{2} + \ln{M!} + M\Big[\ln{4\pi} - \ln\big(A_{\text{max}}\left(X_{\text{max}}- X_{\text{min}}\right)\big)\Big] - \frac{1}{2}\ln{\left[\det{\left(\left.\nabla\nabla\chi^2\right\lvert_{\Theta_0}\right)}\right]} + \ln{4\pi}
\end{align}

```python
def dump(fname, obj):
    with open(fname, 'wb') as fp:
        pickle.dump(obj, fp, protocol=pickle.HIGHEST_PROTOCOL)
def load(fname):
    with open(fname, 'rb') as fp:
        return pickle.load(fp)
```

```python
def logposterior(vol, opt):
    """Calculate log posterior probability
    vol: prior volume term (Amax (Xmax - Xmin))
    opt: result of minimization procedure
    Returns:
    log(P(M|D,I)) = -0.5 χ^2 + log(M!) + M[log(4π) - log(vol)] - 0.5log(det(\nabla\nabla χ^2))
    """
    def logfact(n):
        """Computes log(n!)"""
        return np.sum([np.log(i) for i in range(1,n+1)])

    hess = opt['H']
    M    = (len(opt['x']) - 2)//2
    logp = -opt['chi2']/2 + logfact(M) + M*np.log(4*np.pi/vol) - 0.5*np.log(np.linalg.det(hess))
    return logp

def normalize(logp):
    """Computed normalized posteriors"""
    α = reduce(np.logaddexp, logp)
    return logp - α

def estimates(opt):
    """Compute optimum position/amplitude and background parameters"""
    Θ = opt['x']
    Σ = opt['H_inv']
    sig = np.sqrt(2.0*np.diagonal(Σ))

    b1,b2,amp,pos     = unsquash(opt['x'])
    db1,db2,damp,dpos = unsquash(sig)
    return [b1,db1], [b2, db2], [amp, damp], [pos, dpos]
```

```python
logps, bests = [], []
for opt, data in zip(opts, datas):
    y      = data['avg'].values
    arange = integrate.simps(y,x=theta) / integrate.simps(ideal_signal(squash(np.min(y), 0.0, np.array([0.5]), np.array([1.0])), theta), x=theta)
    xrange = np.max(theta) - np.min(theta)
    logp   = normalize(np.array([logposterior(arange*xrange, opt_m) for opt_m in opt]))

    bests.append(np.argmax(logp))                             # for each signal (corresponding to different σ) find the model M with the hight posterior probability
    logps.append(np.array([lp / np.log(10) for lp in logp]))  # report log posteriors in log10 basis
```

```python
fig, ax = plt.subplots(figsize=(12,6))
for lp,sig,c in zip(logps, σs, mpl.cm.viridis(np.linspace(0,1,len(σs)))):
    ax.plot(np.arange(1,11), lp, marker='o', color=c, label=f'$\sigma = {sig:.2f}$', alpha=0.9)
ax.set_xlim(3,9)
ax.set_ylim(-12,1)
ax.set_ylabel(r'$\log_{10}{P\left(M| D I\right)}$', fontsize=22)
ax.set_xlabel(r'Number of peaks $M$', fontsize=22)
ax.legend(fontsize=22)
plt.show()
```

We see that for the three signals with the larger signal/noise ratio, we are able to correctly infer the number of peaks, i.e., the posterior $P(M=5\lvert D, I)$ maximum.
For the two signals with the large noise amplitude, we have a flatter posterior, with $P(M=4\lvert D, I)\gtrsim P(M=5\lvert D, I)$.

Now, let's plot the peak positions and amplitudes for the model $M$ with the highest posterior, for each of the five signals.

```python
fig, ax = plt.subplots(figsize=(12,9))
print(f'Background parameters:')
colors = mpl.cm.viridis(np.linspace(0, 1, len(σs)))
for i,opt in enumerate(opts):
    best,sig,c = bests[i], σs[i], colors[i]
    best_inter, best_slope, best_pos, best_amp = estimates(opt[best])
    print(f'σ = {sig:.2f}  m = {best_slope[0]:.3f} +/- {best_slope[1]:.3f}, b = {best_inter[0]:.3f} +/- {best_inter[1]:.3f}')
    ax.errorbar(best_pos[0], best_amp[0], xerr=best_pos[1], yerr=best_amp[1], \
                ls='none', marker='s', color=c, mew=3, mfc='None', alpha=0.6, label=f'$\sigma = {sig:.2f}$')

ax.plot(peaks_xi, peaks_ai, marker='x', ms=22, mew=3, ls='None', color='C2')
ax.legend(fontsize=22)
ax.set_xlabel(r'Peak Position $x$', fontsize=22)
ax.set_ylabel(r'Amplitude', fontsize=22)
plt.show()
```

As expected, for the three signals with the smallest noise we are able to correctly infer not only the number of peaks, but their positions and amplitudes.
For the signals with large noise, the two peaks at $x\sim 0.2$ become indistinguishable, merging into one peak with a larger amplitude.


# Nonlinear elasticity

```python
%matplotlib inline
%config InlineBackend.figure_format = 'svg'
import matplotlib as mpl
mpl.rcParams['font.size'] = 8
figsize =(8,4)
mpl.rcParams['figure.figsize'] = figsize

import numpy as np
from scipy.optimize import fsolve
import matplotlib.pyplot as plt
from utils import riemann_tools
from ipywidgets import interact
from ipywidgets import widgets
from clawpack import riemann
from exact_solvers import nonlinear_elasticity
```

In this chapter we investigate a nonlinear model of elastic strain in heterogeneous materials.  This system is equivalent to the $p$-system of gas dynamics, although the stress-strain relation we will use here is very different from the pressure-density relation typically used in gas dynamics.  The equations we consider are:
\begin{align}
\epsilon_t(x,t) - u_x(x,t) & = 0 \\
(\rho(x)u(x,t))_t - \sigma(\epsilon(x,t),x)_x & = 0.
\end{align}
Here $\epsilon$ is the strain, $u$ is the velocity, $\rho$ is the material density, $\sigma$ is the stress,
and ${\mathcal M}=\rho u$ is the momentum.
The first equation is a kinematic relation, while the second represents Newton's second law.  This is a nonlinear
conservation law with spatially varying flux function, in which
\begin{align}
q & = \begin{bmatrix} \epsilon \\ \rho u \end{bmatrix}, & f(q,x) & = \begin{bmatrix} -{\mathcal M}/\rho(x) \\ -\sigma(\epsilon,x) \end{bmatrix}.
\end{align}
If we take a linear stress strain relationship $\sigma(\epsilon,x)=K(x)\epsilon$, then this system is equivalent to the acoustics equations that we have
studied previously.  Here we consider instead a quadratic stress response:

\begin{align}
\sigma(\epsilon,x) = K_1(x) \epsilon + K_2(x) \epsilon^2.
\end{align}

We assume that the spatially-varying functions $\rho, K_1, K_2$ are piecewise constant, taking values
$(\rho_l, K_{1l}, K_{2l})$ for $x<0$ and values $(\rho_r, K_{1r}, K_{2r})$ for $x>0$.  This system has been investigated numerically in (LeVeque, 2002),(LeVeque & Yong, 2003), and (Ketcheson & LeVeque, 2012).

## Hyperbolic structure
The flux jacobian is
\begin{align}
f'(q) = \begin{bmatrix} 0 & -1/\rho(x) \\ -\sigma_\epsilon(\epsilon,x) & 0 \end{bmatrix},
\end{align}
with eigenvalues (characteristic speeds)
\begin{align}
\lambda^\pm(x) = \pm \sqrt{\frac{\sigma_\epsilon(\epsilon,x)}{\rho(x)}} = \pm c(\epsilon, x).
\end{align}
Here for the stress-strain relation we have chosen, $\sigma_\epsilon = K_1(x) + 2 K_2(x)\epsilon$.

It's also convenient to define the nonlinear impedance $Z(\epsilon, x) = \rho(x) c(\epsilon,x)$.  Then the eigenvectors of the flux Jacobian are
\begin{align}
R & = \begin{bmatrix} 1 & 1 \\ Z(\epsilon,x) & -Z(\epsilon,x) \end{bmatrix}.
\end{align}
Both characteristic fields are genuinely nonlinear.  Furthermore, since the characteristic speeds each have a definite sign, this system does not admit transonic rarefactions.

### Structure of the Riemann solution
Based on the eigenstructure of the flux jacobian above, the Riemann solution will always include a left-going and a right-going wave, each of which may be a shock or rarefaction (since both fields are genuinely nonlinear).  We will see -- similarly to our analysis in [the chapter on variable-speed-limit traffic](Traffic_variable_speed.ipynb) that the jump in $\rho$ and $K$ at $x=0$ induces a stationary wave there.  See the figure below for the overall structure of the Riemann solution.

### Hugoniot loci
From the Rankine-Hugoniot jump conditions for the system we obtain the 1-Hugoniot locus for a state $(\epsilon^*_l, u^*_l)$ connected by a 1-shock to a state $(\epsilon_l, u_l)$:
\begin{align}
u^*_l & = u_l - \left( \frac{\left(\sigma_l(\epsilon^*_l)-\sigma_l(\epsilon_l)\right)(\epsilon^*_l-\epsilon_l)}{\rho_l}  \right)^{1/2}
\end{align}
Here $\sigma_l(\epsilon)$ is shorthand for the stress relation in the left medium.
Similarly, a state $(\epsilon^*_r,u^*_r)$ that is connected by a 2-shock to a state $(\epsilon_r, u_r)$ must satisfy
\begin{align}
u^*_r & = u_r - \left( \frac{\left(\sigma_r(\epsilon^*_r)-\sigma_r(\epsilon_r)\right)(\epsilon^*_r-\epsilon_r)}{\rho_r}  \right)^{1/2}.
\end{align}

### Integral curves
The integral curves can be found by writing $\tilde{q}'(\xi) = r^{1,2}(\tilde{q}(\xi))$ and integrating.  This leads to
\begin{align}
u^*_l  & = u_l + \frac{1}{3 K_{2l} \sqrt{\rho_l}} \left( \sigma_{l,\epsilon}(\epsilon^*_l)^{3/2} - \sigma_{l,\epsilon}(\epsilon)^{3/2} \right) \label{NE:integral-curve-1} \\
u^*_r  & = u_r - \frac{1}{3 K_{2r} \sqrt{\rho_r}} \left( \sigma_{r,\epsilon}(\epsilon^*_r)^{3/2} - \sigma_{r,\epsilon}(\epsilon)^{3/2} \right)\label{NE:integral-curve-2}
\end{align}
Here $\sigma_{l,\epsilon}$ is the derivative of the stress function w.r.t $\epsilon$ in the left medium.

### The entropy condition
For a 1-shock, we need that $\lambda^-(\epsilon_l,x<0) > \lambda^-(\epsilon^*_l,x<0)$, which is equivalent to the condition $\epsilon^*_l>\epsilon_l$.  Similarly, a 2-shock is entropy-satisfying if $\epsilon^*_r > \epsilon_r$.

### Interface conditions
Because the flux depends explicitly on $x$, we do not necessarily have continuity of $q$ at $x=0$; i.e. in general $q^*_l \ne q^*_r$.  Instead, the flux must be continuous: $f(q^*_l)=f(q^*_r)$.  For the present system, this means that the stress and velocity must be continuous, which makes sense from a physical point of view.  If the velocity were not continuous, the material would fracture (or overlap itself).  Continuity of the stress is required by Newton's laws.

### Structure of rarefaction waves
For this system, the structure of a centered rarefaction wave can be determined very simply.  Since the characteristic velocity must be equal to $\xi = x/t$ at each point along the wave, we have $\xi = \pm\sqrt{\sigma_\epsilon/\rho}$, or
\begin{align}
\xi^2 = \frac{K_1 + 2K_2\epsilon}{\rho}
\end{align}
which leads to $\epsilon = (\rho\xi^2 - K_1)/(2K_2)$.  Once the value of $\epsilon$ is known, $u$ can be determined using the integral curves \eqref{NE:integral-curve-1} or \eqref{NE:integral-curve-2}.

### Solution of the Riemann problem
Below we show the solution of the Riemann problem.  *To view the code that computes this exact solution, uncomment and execute the next cell.*

```python
# %load exact_solvers/nonlinear_elasticity.py
```

```python
def dsigma(eps, K1, K2):
    "Derivative of stress w.r.t. strain."
    return K1 + 2*K2*eps

def lambda1(q, xi, aux):
    eps = q[0]
    rho, K1, K2 = aux
    return -np.sqrt(dsigma(eps, K1, K2)/rho)

def lambda2(q, xi, aux):
    return -lambda1(q,xi,aux)
```

```python
def make_plot_function(q_l, q_r, aux_l, aux_r):
    states, speeds, reval, wave_types = \
        nonlinear_elasticity.exact_riemann_solution(q_l,q_r,aux_l,aux_r)

    def plot_function(t,which_char):
        ax = riemann_tools.plot_riemann(states,speeds,reval,wave_types,
                                        t=t,t_pointer=0,
                                        extra_axes=True,
                                        variable_names=['Strain','Momentum'])
        if which_char == 1:
            riemann_tools.plot_characteristics(reval,lambda1,(aux_l,aux_r),ax[0])
        elif which_char == 2:
            riemann_tools.plot_characteristics(reval,lambda2,(aux_l,aux_r),ax[0])
        nonlinear_elasticity.phase_plane_plot(q_l, q_r, aux_l, aux_r, ax[3])
        plt.show()
    return plot_function

def plot_riemann_nonlinear_elasticity(rho_l,rho_r,v_l,v_r):
    plot_function = make_plot_function(rho_l,rho_r,v_l,v_r)
    interact(plot_function, t=widgets.FloatSlider(value=0.,min=0,max=1.,step=0.1),
             which_char=widgets.Dropdown(options=[None,1,2],
                                         description='Show characteristics'));
```

```python
aux_l = np.array((1., 5., 1.))
aux_r = np.array((1., 2., 1.))
q_l = np.array([2.1, 0.])
q_r = np.array([0.0, 0.])

plot_riemann_nonlinear_elasticity(q_l, q_r, aux_l, aux_r)
```

## Approximate solution of the Riemann problem using $f$-waves
The exact solver above requires a nonlinear iterative rootfinder and is relatively expensive.  A very cheap approximate Riemann solver for this system was developed in (LeVeque, 2002) using the $f$-wave approach.  One simply approximates both nonlinear waves as shocks, with speeds equal to the characteristic speeds of the left and right states:
\begin{align}
s^1 & = - \sqrt{\frac{\sigma_{\epsilon,l}(\epsilon_l)}{\rho_l}} \\
s^2 & = + \sqrt{\frac{\sigma_{\epsilon,r}(\epsilon_r)}{\rho_r}}.
\end{align}
Meanwhile, the waves are obtained by decomposing the jump in the flux $f(q_r,x>0) - f(q_l,x<0)$ in terms of the eigenvectors of the flux jacobian.

```python
solver = riemann.nonlinear_elasticity_1D_py.nonlinear_elasticity_1D

problem_data = {'stress_relation' : 'quadratic'}
fw_states, fw_speeds, fw_reval = \
            riemann_tools.riemann_solution(solver,q_l,q_r,aux_l,aux_r,
                                           problem_data=problem_data,
                                           verbose=False,
                                           stationary_wave=True,
                                           fwave=True)
```

```python
plot_function = \
    riemann_tools.make_plot_function(fw_states,fw_speeds, fw_reval,
                                     layout='vertical',
                                     variable_names=('Strain','Momentum'))
interact(plot_function, t=widgets.FloatSlider(value=0.4,min=0,max=.9,step=.1));
```

## Comparison of exact and approximate solutions

```python
ex_states, ex_speeds, ex_reval, wave_types = \
            nonlinear_elasticity.exact_riemann_solution(q_l,q_r,aux_l,aux_r)

varnames = nonlinear_elasticity.conserved_variables
plot_function = riemann_tools.make_plot_function([ex_states,fw_states],
                                                 [ex_speeds,fw_speeds],
                                                 [ex_reval,fw_reval],
                                                 [wave_types,['contact']*3],
                                                 ['Exact','$f$-wave'],
                                                 layout='vertical',
                                                 variable_names=varnames)
interact(plot_function, t=widgets.FloatSlider(value=0.4,min=0, max=0.9, step=0.1));
```

As we can see, there are significant discrepancies between the approximate solution and the exact one.  But in a full numerical discretization, the left- and right-going waves are averaged over neighboring cells at the next step, and the approximate solver yields an effective result quite close to that of the exact solver.