# The Scientific Python Ecosystem

(Stolen from Jake Vanderplas.)

(Stolen from Jake Vanderplas.)

(Stolen from Jake Vanderplas.)

(Stolen from Jake Vanderplas.)

(Stolen from Jake Vanderplas.)

 =

```python
import uproot, numpy, matplotlib.pyplot
masses = uproot.open("data/Zmumu.root")["events"].array("M")

y, edges = numpy.histogram(masses, bins=100)
yerr = numpy.sqrt(y)

# the middle of each bin
x = (edges[1:] + edges[:-1])/2

matplotlib.pyplot.errorbar(x, y, yerr, fmt="o", capsize=3);
```

```python
import scipy.stats      # for functions like "cauchy"
import scipy.optimize   # for optimizations like "curve_fit"

f = lambda x, a, b, c, d: a*scipy.stats.cauchy.pdf(x, b, c) + d/x**2

# fit f to x, y, yerr, excluding points where y and yerr are zero
parameters, errors = scipy.optimize.curve_fit(f, x[y > 0], y[y > 0], sigma=yerr[y > 0])

matplotlib.pyplot.plot(x, f(x, *parameters))
matplotlib.pyplot.errorbar(x, y, yerr, fmt="o", capsize=3);
```

```python
xs = numpy.linspace(-10, 10, 100)

# Exercise: search through the SciPy docs and figure out how to convolve two functions.
lorentzian = scipy.stats.cauchy.pdf(xs)
gaussian   = scipy.stats.norm.pdf(xs)
voigtian   = ???

matplotlib.pyplot.plot(xs, lorentzian);
matplotlib.pyplot.plot(xs, gaussian);
matplotlib.pyplot.plot(xs, voigtian);
```

```python
import iminuit

def chi2(a, b, c, d):
    return ((y[y > 0] - f(x[y > 0], a, b, c, d))**2 / yerr[y > 0]**2).sum()

m = iminuit.Minuit(chi2, errordef=1,
                   a=3350, b=91, c=2, d=30,
                   error_a=100, error_b=1, error_c=0.1, error_d=10)
m.migrad()
m.hesse()
m.minos()
m.draw_mncontour("b", "c", nsigma=4);
```

```python
import sklearn.datasets
X1, y1 = sklearn.datasets.make_gaussian_quantiles(
    cov=2.0, n_samples=800, n_features=2, n_classes=2, random_state=1)
X2, y2 = sklearn.datasets.make_gaussian_quantiles(
    mean=(3, 3), cov=1.5, n_samples=1200, n_features=2, n_classes=2, random_state=1)
X = numpy.concatenate((X1, X2))
y = numpy.concatenate((y1, -y2 + 1))

# These points represent a hard classification problem.
matplotlib.pyplot.scatter(X[y == 0, 0], X[y == 0, 1], c="deepskyblue", edgecolor="black");
matplotlib.pyplot.scatter(X[y == 1, 0], X[y == 1, 1], c="orange", edgecolor="black");
```

```python
import sklearn.ensemble
import sklearn.tree

# Solution? Boosted decision tree!
bdt = sklearn.ensemble.AdaBoostClassifier(
    sklearn.tree.DecisionTreeClassifier(max_depth=1), algorithm="SAMME", n_estimators=200)

# Consistent interface: nearly every model has a fit method with this signature
bdt.fit(X, y)

print("\n\npredictions:\n", bdt.predict(X))
print("\n\ntruth:\n", y)
```

    predictions:
     [0 1 1 ... 1 1 1]

    truth:
     [0 1 1 ... 0 0 1]

```python
xx, yy = numpy.meshgrid(numpy.arange(-5, 8, 0.02), numpy.arange(-5, 8, 0.02))
Z = bdt.predict(numpy.c_[xx.ravel(), yy.ravel()])
Z = Z.reshape(xx.shape)

# The decision boundary is a close fit to the points.
matplotlib.pyplot.contourf(xx, yy, Z);
matplotlib.pyplot.scatter(X[y == 0, 0], X[y == 0, 1], c="deepskyblue", edgecolor="black", alpha=0.2);
matplotlib.pyplot.scatter(X[y == 1, 0], X[y == 1, 1], c="orange", edgecolor="black", alpha=0.2);
matplotlib.pyplot.xlim(-5, 8);
matplotlib.pyplot.ylim(-5, 8);
```

```python
twoclass_output = bdt.decision_function(X)

# Decision scores clearly separate the two populations.
matplotlib.pyplot.hist(twoclass_output[y == 0], bins=20, range=(-0.7, 0.7),
                       alpha=0.5, histtype="stepfilled", edgecolor="black");
matplotlib.pyplot.hist(twoclass_output[y == 1], bins=20, range=(-0.7, 0.7),
                       alpha=0.5, histtype="stepfilled", edgecolor="black");
```

```python
import torch

# Turn the Numpy training dataset into Torch-flavored arrays (relocatable to GPUs)
X2 = torch.from_numpy(X.astype(numpy.float32))
y2 = torch.from_numpy(y.astype(numpy.float32).reshape(-1, 1))

num_hidden = 8                                            # This is the neural network
model = torch.nn.Sequential(
    torch.nn.Linear(X2.shape[1], num_hidden),             # inputs → first hidden layer
    torch.nn.Sigmoid(),                                   # first sigmoid
    torch.nn.Linear(num_hidden, num_hidden),              # first hidden layer → second hidden layer
    torch.nn.Sigmoid(),                                   # second sigmoid
    torch.nn.Linear(num_hidden, y2.shape[1]))             # second hidden layer → outputs

loss_fn = torch.nn.MSELoss(reduction="sum")               # simple loss function
optimizer = torch.optim.Adam(model.parameters(), lr=0.01) # optimize with learning rate of 0.01

for i in range(1000):
    y2_pred = model(X2)                                   # compute prediction
    loss = loss_fn(y2_pred, y2)                           # compute prediction - training loss
    optimizer.zero_grad()                                 # reset gradient calculation
    loss.backward()                                       # backpropagation
    optimizer.step()                                      # optimize weights in the network
```

```python
Z2 = model(torch.from_numpy(numpy.c_[xx.ravel(), yy.ravel()].astype(numpy.float32)))
Z2 = Z2.detach().numpy().reshape(xx.shape)
Z2

# Same decision boundary plot as before (not a hard boundary)
matplotlib.pyplot.contourf(xx, yy, Z2);
matplotlib.pyplot.scatter(X[y == 0, 0], X[y == 0, 1], c="deepskyblue", edgecolor="black", alpha=0.1);
matplotlib.pyplot.scatter(X[y == 1, 0], X[y == 1, 1], c="orange", edgecolor="black", alpha=0.1);
matplotlib.pyplot.xlim(-5, 8);
matplotlib.pyplot.ylim(-5, 8);
```

**Exercises:**

   1. Change the Sigmoid to a RELU function (what is that?) and see what happens to the decision boundary.
   2. Add another hidden layer and change the number of nodes in each hidden layer.
   3. How would you monitor the convergence of the fit?

```python
import sympy

E, pt = sympy.symbols("E, pt", positive=True)
px, py, pz, eta, phi = sympy.symbols("px, py, pz, eta, phi", real=True)

# Who needs Mathematica?
sympy.solve([
    sympy.Eq(pt, sympy.sqrt(px**2 + py**2)),
    sympy.Eq(eta, sympy.atanh(pz / sympy.sqrt(px**2 + py**2 + pz**2))),
    sympy.Eq(phi, sympy.atan2(py, px))
], px, py, pz)
```

    [(-pt*Abs(cos(phi)), pt*sin(phi), -pt*sinh(eta)),
     (-pt*Abs(cos(phi)), pt*sin(phi), pt*sinh(eta)),
     (pt*Abs(cos(phi)), pt*sin(phi), -pt*sinh(eta)),
     (pt*Abs(cos(phi)), pt*sin(phi), pt*sinh(eta))]

```python
eta = sympy.atanh(pz / sympy.sqrt(px**2 + py**2 + pz**2))

# I most often use it for approximating a function with Taylor series
approximation = sympy.series(eta, pz, 0, 3)
approximation
```

$\displaystyle pz \left(- \frac{63 px^{10}}{256 px^{10} \sqrt{px^{2} + py^{2}} + 1280 px^{8} py^{2} \sqrt{px^{2} + py^{2}} + 2560 px^{6} py^{4} \sqrt{px^{2} + py^{2}} + 2560 px^{4} py^{6} \sqrt{px^{2} + py^{2}} + 1280 px^{2} py^{8} \sqrt{px^{2} + py^{2}} + 256 py^{10} \sqrt{px^{2} + py^{2}}} - \frac{315 px^{8} py^{2}}{256 px^{10} \sqrt{px^{2} + py^{2}} + 1280 px^{8} py^{2} \sqrt{px^{2} + py^{2}} + 2560 px^{6} py^{4} \sqrt{px^{2} + py^{2}} + 2560 px^{4} py^{6} \sqrt{px^{2} + py^{2}} + 1280 px^{2} py^{8} \sqrt{px^{2} + py^{2}} + 256 py^{10} \sqrt{px^{2} + py^{2}}} + \frac{385 px^{8}}{256 px^{8} \sqrt{px^{2} + py^{2}} + 1024 px^{6} py^{2} \sqrt{px^{2} + py^{2}} + 1536 px^{4} py^{4} \sqrt{px^{2} + py^{2}} + 1024 px^{2} py^{6} \sqrt{px^{2} + py^{2}} + 256 py^{8} \sqrt{px^{2} + py^{2}}} - \frac{315 px^{6} py^{4}}{128 px^{10} \sqrt{px^{2} + py^{2}} + 640 px^{8} py^{2} \sqrt{px^{2} + py^{2}} + 1280 px^{6} py^{4} \sqrt{px^{2} + py^{2}} + 1280 px^{4} py^{6} \sqrt{px^{2} + py^{2}} + 640 px^{2} py^{8} \sqrt{px^{2} + py^{2}} + 128 py^{10} \sqrt{px^{2} + py^{2}}} + \frac{385 px^{6} py^{2}}{64 px^{8} \sqrt{px^{2} + py^{2}} + 256 px^{6} py^{2} \sqrt{px^{2} + py^{2}} + 384 px^{4} py^{4} \sqrt{px^{2} + py^{2}} + 256 px^{2} py^{6} \sqrt{px^{2} + py^{2}} + 64 py^{8} \sqrt{px^{2} + py^{2}}} - \frac{495 px^{6}}{128 px^{6} \sqrt{px^{2} + py^{2}} + 384 px^{4} py^{2} \sqrt{px^{2} + py^{2}} + 384 px^{2} py^{4} \sqrt{px^{2} + py^{2}} + 128 py^{6} \sqrt{px^{2} + py^{2}}} - \frac{315 px^{4} py^{6}}{128 px^{10} \sqrt{px^{2} + py^{2}} + 640 px^{8} py^{2} \sqrt{px^{2} + py^{2}} + 1280 px^{6} py^{4} \sqrt{px^{2} + py^{2}} + 1280 px^{4} py^{6} \sqrt{px^{2} + py^{2}} + 640 px^{2} py^{8} \sqrt{px^{2} + py^{2}} + 128 py^{10} \sqrt{px^{2} + py^{2}}} + \frac{1155 px^{4} py^{4}}{128 px^{8} \sqrt{px^{2} + py^{2}} + 512 px^{6} py^{2} \sqrt{px^{2} + py^{2}} + 768 px^{4} py^{4} \sqrt{px^{2} + py^{2}} + 512 px^{2} py^{6} \sqrt{px^{2} + py^{2}} + 128 py^{8} \sqrt{px^{2} + py^{2}}} - \frac{1485 px^{4} py^{2}}{128 px^{6} \sqrt{px^{2} + py^{2}} + 384 px^{4} py^{2} \sqrt{px^{2} + py^{2}} + 384 px^{2} py^{4} \sqrt{px^{2} + py^{2}} + 128 py^{6} \sqrt{px^{2} + py^{2}}} + \frac{693 px^{4}}{128 px^{4} \sqrt{px^{2} + py^{2}} + 256 px^{2} py^{2} \sqrt{px^{2} + py^{2}} + 128 py^{4} \sqrt{px^{2} + py^{2}}} - \frac{315 px^{2} py^{8}}{256 px^{10} \sqrt{px^{2} + py^{2}} + 1280 px^{8} py^{2} \sqrt{px^{2} + py^{2}} + 2560 px^{6} py^{4} \sqrt{px^{2} + py^{2}} + 2560 px^{4} py^{6} \sqrt{px^{2} + py^{2}} + 1280 px^{2} py^{8} \sqrt{px^{2} + py^{2}} + 256 py^{10} \sqrt{px^{2} + py^{2}}} + \frac{385 px^{2} py^{6}}{64 px^{8} \sqrt{px^{2} + py^{2}} + 256 px^{6} py^{2} \sqrt{px^{2} + py^{2}} + 384 px^{4} py^{4} \sqrt{px^{2} + py^{2}} + 256 px^{2} py^{6} \sqrt{px^{2} + py^{2}} + 64 py^{8} \sqrt{px^{2} + py^{2}}} - \frac{1485 px^{2} py^{4}}{128 px^{6} \sqrt{px^{2} + py^{2}} + 384 px^{4} py^{2} \sqrt{px^{2} + py^{2}} + 384 px^{2} py^{4} \sqrt{px^{2} + py^{2}} + 128 py^{6} \sqrt{px^{2} + py^{2}}} + \frac{693 px^{2} py^{2}}{64 px^{4} \sqrt{px^{2} + py^{2}} + 128 px^{2} py^{2} \sqrt{px^{2} + py^{2}} + 64 py^{4} \sqrt{px^{2} + py^{2}}} - \frac{1155 px^{2}}{256 \left(px^{2} \sqrt{px^{2} + py^{2}} + py^{2} \sqrt{px^{2} + py^{2}}\right)} - \frac{63 py^{10}}{256 px^{10} \sqrt{px^{2} + py^{2}} + 1280 px^{8} py^{2} \sqrt{px^{2} + py^{2}} + 2560 px^{6} py^{4} \sqrt{px^{2} + py^{2}} + 2560 px^{4} py^{6} \sqrt{px^{2} + py^{2}} + 1280 px^{2} py^{8} \sqrt{px^{2} + py^{2}} + 256 py^{10} \sqrt{px^{2} + py^{2}}} + \frac{385 py^{8}}{256 px^{8} \sqrt{px^{2} + py^{2}} + 1024 px^{6} py^{2} \sqrt{px^{2} + py^{2}} + 1536 px^{4} py^{4} \sqrt{px^{2} + py^{2}} + 1024 px^{2} py^{6} \sqrt{px^{2} + py^{2}} + 256 py^{8} \sqrt{px^{2} + py^{2}}} - \frac{495 py^{6}}{128 px^{6} \sqrt{px^{2} + py^{2}} + 384 px^{4} py^{2} \sqrt{px^{2} + py^{2}} + 384 px^{2} py^{4} \sqrt{px^{2} + py^{2}} + 128 py^{6} \sqrt{px^{2} + py^{2}}} + \frac{693 py^{4}}{128 px^{4} \sqrt{px^{2} + py^{2}} + 256 px^{2} py^{2} \sqrt{px^{2} + py^{2}} + 128 py^{4} \sqrt{px^{2} + py^{2}}} - \frac{1155 py^{2}}{256 \left(px^{2} \sqrt{px^{2} + py^{2}} + py^{2} \sqrt{px^{2} + py^{2}}\right)} + \frac{693}{256 \sqrt{px^{2} + py^{2}}}\right) + O\left(pz^{3}\right)$

```python
import inspect

# ...and dropping code for the approximation into a numerical calculation
f = sympy.lambdify([px, py, pz], approximation.removeO())
print(inspect.getsource(f))
```

    def _lambdifygenerated(px, py, pz):
        return (pz*(-63*px**10/(256*px**10*sqrt(px**2 + py**2) + 1280*px**8*py**2*sqrt(px**2 + py**2) + 2560*px**6*py**4*sqrt(px**2 + py**2) + 2560*px**4*py**6*sqrt(px**2 + py**2) + 1280*px**2*py**8*sqrt(px**2 + py**2) + 256*py**10*sqrt(px**2 + py**2)) - 315*px**8*py**2/(256*px**10*sqrt(px**2 + py**2) + 1280*px**8*py**2*sqrt(px**2 + py**2) + 2560*px**6*py**4*sqrt(px**2 + py**2) + 2560*px**4*py**6*sqrt(px**2 + py**2) + 1280*px**2*py**8*sqrt(px**2 + py**2) + 256*py**10*sqrt(px**2 + py**2)) + 385*px**8/(256*px**8*sqrt(px**2 + py**2) + 1024*px**6*py**2*sqrt(px**2 + py**2) + 1536*px**4*py**4*sqrt(px**2 + py**2) + 1024*px**2*py**6*sqrt(px**2 + py**2) + 256*py**8*sqrt(px**2 + py**2)) - 315*px**6*py**4/(128*px**10*sqrt(px**2 + py**2) + 640*px**8*py**2*sqrt(px**2 + py**2) + 1280*px**6*py**4*sqrt(px**2 + py**2) + 1280*px**4*py**6*sqrt(px**2 + py**2) + 640*px**2*py**8*sqrt(px**2 + py**2) + 128*py**10*sqrt(px**2 + py**2)) + 385*px**6*py**2/(64*px**8*sqrt(px**2 + py**2) + 256*px**6*py**2*sqrt(px**2 + py**2) + 384*px**4*py**4*sqrt(px**2 + py**2) + 256*px**2*py**6*sqrt(px**2 + py**2) + 64*py**8*sqrt(px**2 + py**2)) - 495*px**6/(128*px**6*sqrt(px**2 + py**2) + 384*px**4*py**2*sqrt(px**2 + py**2) + 384*px**2*py**4*sqrt(px**2 + py**2) + 128*py**6*sqrt(px**2 + py**2)) - 315*px**4*py**6/(128*px**10*sqrt(px**2 + py**2) + 640*px**8*py**2*sqrt(px**2 + py**2) + 1280*px**6*py**4*sqrt(px**2 + py**2) + 1280*px**4*py**6*sqrt(px**2 + py**2) + 640*px**2*py**8*sqrt(px**2 + py**2) + 128*py**10*sqrt(px**2 + py**2)) + 1155*px**4*py**4/(128*px**8*sqrt(px**2 + py**2) + 512*px**6*py**2*sqrt(px**2 + py**2) + 768*px**4*py**4*sqrt(px**2 + py**2) + 512*px**2*py**6*sqrt(px**2 + py**2) + 128*py**8*sqrt(px**2 + py**2)) - 1485*px**4*py**2/(128*px**6*sqrt(px**2 + py**2) + 384*px**4*py**2*sqrt(px**2 + py**2) + 384*px**2*py**4*sqrt(px**2 + py**2) + 128*py**6*sqrt(px**2 + py**2)) + 693*px**4/(128*px**4*sqrt(px**2 + py**2) + 256*px**2*py**2*sqrt(px**2 + py**2) + 128*py**4*sqrt(px**2 + py**2)) - 315*px**2*py**8/(256*px**10*sqrt(px**2 + py**2) + 1280*px**8*py**2*sqrt(px**2 + py**2) + 2560*px**6*py**4*sqrt(px**2 + py**2) + 2560*px**4*py**6*sqrt(px**2 + py**2) + 1280*px**2*py**8*sqrt(px**2 + py**2) + 256*py**10*sqrt(px**2 + py**2)) + 385*px**2*py**6/(64*px**8*sqrt(px**2 + py**2) + 256*px**6*py**2*sqrt(px**2 + py**2) + 384*px**4*py**4*sqrt(px**2 + py**2) + 256*px**2*py**6*sqrt(px**2 + py**2) + 64*py**8*sqrt(px**2 + py**2)) - 1485*px**2*py**4/(128*px**6*sqrt(px**2 + py**2) + 384*px**4*py**2*sqrt(px**2 + py**2) + 384*px**2*py**4*sqrt(px**2 + py**2) + 128*py**6*sqrt(px**2 + py**2)) + 693*px**2*py**2/(64*px**4*sqrt(px**2 + py**2) + 128*px**2*py**2*sqrt(px**2 + py**2) + 64*py**4*sqrt(px**2 + py**2)) - 1155/256*px**2/(px**2*sqrt(px**2 + py**2) + py**2*sqrt(px**2 + py**2)) - 63*py**10/(256*px**10*sqrt(px**2 + py**2) + 1280*px**8*py**2*sqrt(px**2 + py**2) + 2560*px**6*py**4*sqrt(px**2 + py**2) + 2560*px**4*py**6*sqrt(px**2 + py**2) + 1280*px**2*py**8*sqrt(px**2 + py**2) + 256*py**10*sqrt(px**2 + py**2)) + 385*py**8/(256*px**8*sqrt(px**2 + py**2) + 1024*px**6*py**2*sqrt(px**2 + py**2) + 1536*px**4*py**4*sqrt(px**2 + py**2) + 1024*px**2*py**6*sqrt(px**2 + py**2) + 256*py**8*sqrt(px**2 + py**2)) - 495*py**6/(128*px**6*sqrt(px**2 + py**2) + 384*px**4*py**2*sqrt(px**2 + py**2) + 384*px**2*py**4*sqrt(px**2 + py**2) + 128*py**6*sqrt(px**2 + py**2)) + 693*py**4/(128*px**4*sqrt(px**2 + py**2) + 256*px**2*py**2*sqrt(px**2 + py**2) + 128*py**4*sqrt(px**2 + py**2)) - 1155/256*py**2/(px**2*sqrt(px**2 + py**2) + py**2*sqrt(px**2 + py**2)) + (693/256)/sqrt(px**2 + py**2)))

```python
# Exercise: factorize this expression by finding the relevant command in the documentation.
#
# (And if you're done with that, turn the result into a Python function and apply it to Numpy arrays.
# Is the factorized expression faster than the expanded expression?)

x, y, z = sympy.symbols("x y z")

x**2*z + 4*x*y*z + 4*y**2*z
```

$\displaystyle z \left(x + 2 y\right)^{2}$

# Libraries specifically for particle physics

```python
import numpythia
import numpythia.testcmnd
import pandas

pythia = numpythia.Pythia(numpythia.testcmnd.get_cmnd("w"), random_state=1)

events = list(pythia(events=1))    # pythia is a generator; evaluate it
pandas.DataFrame(events[0].all())  # event.all(SELECTION) returns a Numpy record array
```

      E
      px
      py
      pz
      pT
      mass
      rap
      eta
      theta
      phi
      prodx
      prody
      prodz
      prodt
      pdgid
      status

      0
      6500.000000
      0.000000e+00
      0.000000e+00
      6499.999932
      0.000000e+00
      9.382700e-01
      9.536422
      inf
      0.000000e+00
      0.000000
      0.000000
      0.000000
      0.000000
      0.000000
      2212
      4

      1
      156.337354
      -1.603964e+00
      1.789967e+00
      156.318878
      2.403474e+00
      1.388571e-05
      4.868189
      4.868189
      1.537424e-02
      2.301445
      0.000000
      0.000000
      0.000000
      0.000000
      21
      61

      2
      3157.193528
      1.632030e-01
      -1.049595e+00
      3157.193255
      1.062208e+00
      7.713300e-01
      8.478484
      8.690237
      3.364405e-04
      -1.416540
      0.000000
      0.000000
      0.000000
      0.000000
      2203
      63

      3
      3186.477199
      1.440761e+00
      -7.403724e-01
      3186.476770
      1.619859e+00
      3.300000e-01
      8.257147
      8.277479
      5.083543e-04
      -0.474687
      0.000000
      0.000000
      0.000000
      0.000000
      1
      63

      4
      156.328011
      0.000000e+00
      2.220446e-16
      156.328011
      2.220446e-16
      1.452593e-05
      16.879336
      inf
      1.420376e-18
      1.570796
      0.000000
      0.000000
      0.000000
      0.000000
      21
      53

      5
      3157.193528
      1.632030e-01
      -1.049595e+00
      3157.193255
      1.062208e+00
      7.713300e-01
      8.478484
      8.690237
      3.364405e-04
      -1.416540
      0.000000
      0.000000
      0.000000
      0.000000
      2203
      71

      6
      3186.477199
      1.440761e+00
      -7.403724e-01
      3186.476770
      1.619859e+00
      3.300000e-01
      8.257147
      8.277479
      5.083543e-04
      -0.474687
      0.000000
      0.000000
      0.000000
      0.000000
      1
      71

      7
      151.160792
      0.000000e+00
      1.110223e-16
      151.160792
      1.110223e-16
      1.027139e-05
      17.209104
      inf
      7.344649e-19
      1.570796
      0.000000
      0.000000
      0.000000
      0.000000
      21
      41

      8
      132.235959
      3.330669e-14
      -7.549517e-14
      132.235959
      8.251579e-14
      -5.394797e-06
      NaN
      inf
      6.240042e-16
      -1.155304
      0.000000
      0.000000
      0.000000
      0.000000
      21
      53

      9
      18.966263
      -1.438213e+00
      1.221482e+00
      18.872166
      1.886922e+00
      1.450244e-06
      2.998378
      2.998378
      9.965319e-02
      2.437502
      0.000000
      0.000000
      0.000000
      0.000000
      21
      43

      10
      129.079322
      1.665335e-14
      -3.774758e-14
      129.079322
      4.125790e-14
      -3.814697e-06
      NaN
      inf
      3.196321e-16
      -1.155304
      0.000000
      0.000000
      0.000000
      0.000000
      21
      42

      11
      129.079322
      8.881784e-16
      0.000000e+00
      129.079322
      8.881784e-16
      -1.078959e-05
      NaN
      inf
      6.880873e-18
      0.000000
      0.000000
      0.000000
      0.000000
      0.000000
      21
      41

      12
      80.761642
      3.552714e-14
      7.105427e-15
      80.761642
      3.623071e-14
      -4.046097e-06
      NaN
      inf
      4.486129e-16
      0.197396
      0.000000
      0.000000
      0.000000
      0.000000
      21
      41

      13
      50.773495
      8.705005e+00
      -2.136256e+01
      45.230650
      2.306808e+01
      -3.303625e-06
      1.425942
      1.425942
      4.716234e-01
      -1.183851
      0.000000
      0.000000
      0.000000
      0.000000
      21
      43

      14
      50.807348
      8.710641e+00
      -2.137533e+01
      45.261536
      2.308203e+01
      -3.568323e-06
      1.426011
      1.426011
      4.715918e-01
      -1.183833
      0.000000
      0.000000
      0.000000
      0.000000
      21
      44

      15
      13.604880
      3.594938e+00
      -7.168036e+00
      10.990379
      8.018998e+00
      3.371748e-07
      1.120740
      1.120740
      6.303416e-01
      -1.105931
      0.000000
      0.000000
      0.000000
      0.000000
      21
      51

      16
      40.359105
      5.115703e+00
      -1.420729e+01
      37.427795
      1.510025e+01
      4.768372e-07
      1.639262
      1.639262
      3.834771e-01
      -1.225174
      0.000000
      0.000000
      0.000000
      0.000000
      21
      51

      17
      40.456108
      5.538712e+00
      -1.456656e+01
      37.334097
      1.558403e+01
      2.977847e-06
      1.607769
      1.607769
      3.954335e-01
      -1.207444
      0.000000
      0.000000
      0.000000
      0.000000
      21
      44

      18
      3.251627
      7.572504e-01
      -8.581014e-01
      3.002218
      1.144450e+00
      5.000000e-01
      1.610928
      1.692069
      3.641966e-01
      -0.847750
      0.000000
      0.000000
      0.000000
      0.000000
      3
      51

      19
      43.935302
      6.591936e+00
      -1.655262e+01
      40.157415
      1.781692e+01
      5.000000e-01
      1.551378
      1.551737
      4.175833e-01
      -1.191807
      0.000000
      0.000000
      0.000000
      0.000000
      -3
      51

      20
      3.242875
      7.253165e-01
      -8.221281e-01
      3.010691
      1.096348e+00
      5.000000e-01
      1.646688
      1.734947
      3.492260e-01
      -0.847879
      0.000000
      0.000000
      0.000000
      0.000000
      3
      62

      21
      3.242875
      7.253165e-01
      -8.221281e-01
      3.010691
      1.096348e+00
      5.000000e-01
      1.646688
      1.734947
      3.492260e-01
      -0.847879
      0.000000
      0.000000
      0.000000
      0.000000
      3
      71

      22
      2.346952
      -8.393021e-02
      -5.233160e-01
      2.231516
      5.300037e-01
      4.976100e-01
      1.840202
      2.144513
      2.331877e-01
      -1.729824
      0.000000
      0.000000
      0.000000
      0.000000
      -311
      2

      23
      5.017068
      7.260278e-01
      -2.283182e-01
      4.957040
      7.610818e-01
      1.395700e-01
      2.556469
      2.572812
      1.523459e-01
      -0.304684
      0.000000
      0.000000
      0.000000
      0.000000
      -211
      1

      24
      11.570857
      -5.481886e-01
      5.102216e-02
      11.530791
      5.505579e-01
      7.889656e-01
      3.178571
      3.735561
      4.771053e-02
      3.048786
      0.000000
      0.000000
      0.000000
      0.000000
      223
      2

      25
      56.187180
      9.686375e-01
      -1.699006e-01
      56.178411
      9.834251e-01
      1.349800e-01
      4.729140
      4.738470
      1.750360e-02
      -0.173635
      0.000000
      0.000000
      0.000000
      0.000000
      111
      2

      26
      47.574596
      2.589860e-02
      4.131424e-02
      47.574366
      4.876068e-02
      1.395700e-01
      6.467052
      7.576273
      1.024936e-03
      1.010854
      0.000000
      0.000000
      0.000000
      0.000000
      211
      1

      27
      82.595915
      -4.117562e-01
      9.883136e-02
      82.593012
      4.234510e-01
      5.478500e-01
      5.474647
      5.966396
      5.126915e-03
      2.906025
      0.000000
      0.000000
      0.000000
      0.000000
      221
      2

      28
      56.017883
      -4.257959e-01
      -1.564910e-01
      56.013836
      4.536426e-01
      4.976100e-01
      5.114264
      5.509208
      8.098583e-03
      -2.789391
      0.000000
      0.000000
      0.000000
      0.000000
      311
      2

      29
      894.731258
      5.500858e-01
      -4.439534e-01
      894.730841
      7.068869e-01
      4.976100e-01
      7.635311
      7.836555
      7.900552e-04
      -0.679032
      0.000000
      0.000000
      0.000000
      0.000000
      -311
      2

      ...
      ...
      ...
      ...
      ...
      ...
      ...
      ...
      ...
      ...
      ...
      ...
      ...
      ...
      ...
      ...
      ...

      404
      14.070199
      6.235642e+00
      9.857060e+00
      -7.812991
      1.166383e+01
      9.395700e-01
      -0.625992
      -0.627790
      2.160998e+00
      1.006746
      914.334886
      1418.564505
      -1159.635058
      2052.420141
      -2112
      1

      405
      3.601545
      1.636962e+00
      2.357061e+00
      -2.171696
      2.869735e+00
      1.395700e-01
      -0.697833
      -0.698545
      2.218609e+00
      0.963776
      914.334886
      1418.564505
      -1159.635058
      2052.420141
      211
      1

      406
      1.011955
      3.654647e-01
      5.795201e-01
      -0.554103
      6.851336e-01
      4.976100e-01
      -0.614885
      -0.739487
      2.250851e+00
      1.008156
      0.000000
      0.000000
      0.000000
      0.000000
      311
      2

      407
      1.041324
      3.632195e-01
      9.422014e-01
      -0.212613
      1.009788e+00
      1.395700e-01
      -0.207086
      -0.209026
      1.778317e+00
      1.202851
      0.000000
      0.000000
      0.000000
      0.000000
      211
      1

      408
      4.816409
      9.318120e-01
      4.647537e+00
      -0.697281
      4.740029e+00
      4.936800e-01
      -0.145796
      -0.146579
      1.716854e+00
      1.372924
      0.000000
      0.000000
      0.000000
      0.000000
      -321
      1

      409
      1.483839
      2.430691e-01
      1.347858e+00
      -0.554757
      1.369600e+00
      1.349800e-01
      -0.392910
      -0.394720
      1.955649e+00
      1.392377
      0.000000
      0.000000
      0.000000
      0.000000
      111
      2

      410
      2.673447
      4.325854e-01
      2.613896e+00
      -0.329026
      2.649449e+00
      1.395700e-01
      -0.123699
      -0.123869
      1.694350e+00
      1.406788
      0.000000
      0.000000
      0.000000
      0.000000
      211
      1

      411
      2.544121
      -4.430991e-02
      2.525988e+00
      0.267866
      2.526377e+00
      1.349800e-01
      0.105680
      0.105830
      1.465163e+00
      1.588336
      0.000000
      0.000000
      0.000000
      0.000000
      111
      2

      412
      15.833678
      1.180977e+00
      1.553874e+01
      -2.498066
      1.558355e+01
      1.272000e+00
      -0.159098
      -0.159623
      1.729746e+00
      1.494940
      0.062971
      0.618065
      -0.091500
      0.629712
      10323
      2

      413
      9.232344
      1.325632e+00
      9.063685e+00
      -1.144123
      9.160114e+00
      1.395700e-01
      -0.124566
      -0.124580
      1.695055e+00
      1.425569
      0.062971
      0.618065
      -0.091500
      0.629712
      -211
      1

      414
      4.324876
      1.723379e+00
      2.662915e+00
      -2.939966
      3.171932e+00
      -1.332800e-07
      -0.828706
      -0.828706
      2.318259e+00
      0.996408
      0.000405
      0.000609
      -0.000678
      0.000998
      22
      1

      415
      0.436046
      2.093308e-01
      2.447012e-01
      -0.294003
      3.220218e-01
      -2.172199e-08
      -0.818492
      -0.818492
      2.310742e+00
      0.863144
      0.000405
      0.000609
      -0.000678
      0.000998
      22
      1

      416
      1.522633
      7.767609e-01
      8.311318e-01
      -1.012064
      1.137602e+00
      5.161914e-08
      -0.801152
      -0.801152
      2.297862e+00
      0.819200
      0.000046
      0.000050
      -0.000055
      0.000088
      22
      1

      417
      0.969276
      5.229614e-01
      5.920571e-01
      -0.561673
      7.899495e-01
      3.650024e-08
      -0.661674
      -0.661674
      2.188882e+00
      0.847287
      0.000046
      0.000050
      -0.000055
      0.000088
      22
      1

      418
      1.011955
      3.654647e-01
      5.795201e-01
      -0.554103
      6.851336e-01
      4.976100e-01
      -0.614885
      -0.739487
      2.250851e+00
      1.008156
      0.000000
      0.000000
      0.000000
      0.000000
      310
      2

      419
      1.287483
      1.745915e-01
      1.169569e+00
      -0.509155
      1.182529e+00
      -1.490116e-08
      -0.418262
      -0.418262
      1.977371e+00
      1.422612
      0.000013
      0.000072
      -0.000030
      0.000080
      22
      1

      420
      0.196355
      6.847765e-02
      1.782884e-01
      -0.045602
      1.909867e-01
      -6.452392e-09
      -0.236556
      -0.236556
      1.805176e+00
      1.204086
      0.000013
      0.000072
      -0.000030
      0.000080
      22
      1

      421
      1.432836
      2.026139e-02
      1.429064e+00
      0.101910
      1.429208e+00
      2.107342e-08
      0.071245
      0.071245
      1.499611e+00
      1.556619
      -0.000002
      0.000099
      0.000010
      0.000099
      22
      1

      422
      1.111285
      -6.457129e-02
      1.096924e+00
      0.165955
      1.098823e+00
      -2.107342e-08
      0.150462
      0.150462
      1.420899e+00
      1.629594
      -0.000002
      0.000099
      0.000010
      0.000099
      22
      1

      423
      8.461579
      5.658193e-01
      8.337937e+00
      -1.228553
      8.357114e+00
      4.976100e-01
      -0.146225
      -0.146482
      1.716758e+00
      1.503039
      0.062971
      0.618065
      -0.091500
      0.629712
      311
      2

      424
      7.372100
      6.151573e-01
      7.200799e+00
      -1.269513
      7.227028e+00
      7.115201e-01
      -0.173938
      -0.174771
      1.744684e+00
      1.485574
      0.062971
      0.618065
      -0.091500
      0.629712
      213
      2

      425
      0.710465
      1.442658e-01
      5.459319e-01
      -0.407955
      5.646718e-01
      1.395700e-01
      -0.653778
      -0.670972
      2.196440e+00
      1.312446
      24.321648
      38.567015
      -36.875481
      67.345538
      211
      1

      426
      0.301490
      2.211990e-01
      3.358821e-02
      -0.146148
      2.237346e-01
      1.395700e-01
      -0.529178
      -0.613923
      2.149432e+00
      0.150695
      24.321648
      38.567015
      -36.875481
      67.345538
      -211
      1

      427
      8.461579
      5.658193e-01
      8.337937e+00
      -1.228553
      8.357114e+00
      4.976100e-01
      -0.146225
      -0.146482
      1.716758e+00
      1.503039
      0.062971
      0.618065
      -0.091500
      0.629712
      310
      2

      428
      6.244888
      5.572298e-01
      6.083367e+00
      -1.288909
      6.108835e+00
      1.395700e-01
      -0.209402
      -0.209456
      1.778737e+00
      1.479452
      0.062971
      0.618065
      -0.091500
      0.629712
      211
      1

      429
      1.127212
      5.792751e-02
      1.117432e+00
      0.019396
      1.118933e+00
      1.349800e-01
      0.017209
      0.017333
      1.553464e+00
      1.519003
      0.062971
      0.618065
      -0.091500
      0.629712
      111
      2

      430
      5.794541
      2.045138e-01
      5.725449e+00
      -0.857107
      5.729100e+00
      1.395700e-01
      -0.149010
      -0.149053
      1.719301e+00
      1.535091
      0.062971
      0.618065
      -0.091500
      0.629712
      211
      1

      431
      2.667038
      3.613055e-01
      2.612489e+00
      -0.371446
      2.637354e+00
      1.395700e-01
      -0.140184
      -0.140379
      1.710716e+00
      1.433369
      0.062971
      0.618065
      -0.091500
      0.629712
      -211
      1

      432
      1.070853
      8.256109e-02
      1.067323e+00
      0.027047
      1.070512e+00
      1.490116e-08
      0.025263
      0.025263
      1.545536e+00
      1.493597
      0.062971
      0.618065
      -0.091500
      0.629712
      22
      1

      433
      0.056358
      -2.463358e-02
      5.010893e-02
      -0.007651
      5.583653e-02
      4.267856e-09
      -0.136606
      -0.136606
      1.706979e+00
      2.027702
      0.062971
      0.618065
      -0.091500
      0.629712
      22
      1

434 rows × 16 columns

```python
import pyjet

fourvectors = events[0].all()[["E", "px", "py", "pz"]]
for jet in pyjet.cluster(fourvectors, R=1.0, p=-1, ep=True).inclusive_jets():
    print(jet)
```

    PseudoJet(pt=2909.400, eta=-0.778, phi=0.947, mass=1292.519)
    PseudoJet(pt=2572.100, eta=-1.973, phi=-2.263, mass=1587.844)
    PseudoJet(pt=201.637, eta=1.507, phi=-1.185, mass=58.241)
    PseudoJet(pt=49.311, eta=-3.790, phi=2.058, mass=29.501)
    PseudoJet(pt=26.513, eta=-4.338, phi=-2.149, mass=89.196)
    PseudoJet(pt=16.728, eta=3.088, phi=2.557, mass=11.668)
    PseudoJet(pt=14.165, eta=-0.257, phi=-1.441, mass=13.213)
    PseudoJet(pt=9.817, eta=0.009, phi=1.695, mass=3.142)
    PseudoJet(pt=9.261, eta=-8.374, phi=-2.268, mass=7.053)
    PseudoJet(pt=8.531, eta=8.208, phi=-0.457, mass=6.606)
    PseudoJet(pt=6.348, eta=-2.649, phi=0.028, mass=2.832)
    PseudoJet(pt=3.513, eta=0.227, phi=-3.114, mass=2.351)
    PseudoJet(pt=2.876, eta=5.235, phi=2.406, mass=1.891)
    PseudoJet(pt=2.590, eta=2.299, phi=0.261, mass=1.236)
    PseudoJet(pt=2.579, eta=8.581, phi=-1.500, mass=2.023)
    PseudoJet(pt=2.478, eta=-6.952, phi=0.853, mass=-0.000)
    PseudoJet(pt=2.208, eta=5.565, phi=-0.178, mass=3.658)
    PseudoJet(pt=2.068, eta=-2.560, phi=2.793, mass=1.768)
    PseudoJet(pt=1.922, eta=1.835, phi=2.337, mass=0.044)
    PseudoJet(pt=1.713, eta=6.670, phi=0.492, mass=2.429)
    PseudoJet(pt=1.528, eta=5.579, phi=-2.839, mass=1.439)
    PseudoJet(pt=1.499, eta=7.787, phi=-1.366, mass=2.152)
    PseudoJet(pt=1.123, eta=4.394, phi=0.801, mass=0.725)
    PseudoJet(pt=0.978, eta=-3.338, phi=0.920, mass=0.750)
    PseudoJet(pt=0.870, eta=-1.678, phi=-0.786, mass=0.303)
    PseudoJet(pt=0.591, eta=0.326, phi=-0.322, mass=0.140)
    PseudoJet(pt=0.529, eta=-6.912, phi=-1.994, mass=1.216)
    PseudoJet(pt=0.447, eta=1.967, phi=-2.135, mass=0.140)
    PseudoJet(pt=0.366, eta=4.031, phi=2.918, mass=0.000)
    PseudoJet(pt=0.295, eta=6.775, phi=2.827, mass=0.140)
    PseudoJet(pt=0.250, eta=-2.234, phi=2.781, mass=0.888)
    PseudoJet(pt=0.221, eta=-7.481, phi=3.012, mass=0.140)
    PseudoJet(pt=0.217, eta=4.302, phi=-0.911, mass=0.753)
    PseudoJet(pt=0.189, eta=-0.868, phi=-2.367, mass=0.140)
    PseudoJet(pt=0.177, eta=-8.580, phi=2.723, mass=0.140)
    PseudoJet(pt=0.176, eta=4.122, phi=-2.363, mass=0.140)
    PseudoJet(pt=0.100, eta=3.703, phi=1.827, mass=-0.000)
    PseudoJet(pt=0.034, eta=3.283, phi=0.170, mass=-0.000)
    PseudoJet(pt=0.022, eta=-7.675, phi=-1.161, mass=0.000)
    PseudoJet(pt=0.000, eta=35.979, phi=-1.155, mass=-0.000)
    PseudoJet(pt=0.000, eta=-37.332, phi=-3.137, mass=-0.000)
    PseudoJet(pt=0.000, eta=-37.332, phi=-1.571, mass=-0.000)
    PseudoJet(pt=0.000, eta=-37.332, phi=1.984, mass=-0.000)
    PseudoJet(pt=0.000, eta=36.034, phi=0.197, mass=-0.000)
    PseudoJet(pt=0.000, eta=-37.332, phi=-0.876, mass=-0.000)
    PseudoJet(pt=0.000, eta=40.211, phi=0.000, mass=-0.000)
    PseudoJet(pt=0.000, eta=42.060, phi=1.571, mass=0.000)
    PseudoJet(pt=0.000, eta=-37.332, phi=-1.107, mass=-0.000)
    PseudoJet(pt=0.000, eta=100000.000, phi=0.000, mass=0.000)
    PseudoJet(pt=0.000, eta=100000.000, phi=0.000, mass=0.938)
    PseudoJet(pt=0.000, eta=100000.000, phi=0.000, mass=0.938)
    PseudoJet(pt=0.000, eta=100000.000, phi=0.000, mass=0.000)

```python
import particle
from hepunits.units import cm

import IPython.display

# Find all strange baryons with c*tau > 1 cm
for x in particle.Particle.findall(lambda p: p.pdgid.is_baryon and p.pdgid.has_strange and
                                             p.width > 0 and p.ctau > 1 * cm):

    IPython.display.display(IPython.display.Latex("$" + x.latex_name + "$"))
    print(repr(x), end="\n\n")
```

$\Sigma^{-}$

    <Particle: name="Sigma-", pdgid=3112, mass=1197.45 ± 0.03 MeV>

$\bar{\Sigma}^{+}$

    <Particle: name="Sigma~+", pdgid=-3112, mass=1197.45 ± 0.03 MeV>

$\Lambda$

    <Particle: name="Lambda", pdgid=3122, mass=1115.683 ± 0.006 MeV>

$\bar{\Lambda}$

    <Particle: name="Lambda~", pdgid=-3122, mass=1115.683 ± 0.006 MeV>

$\Sigma^{+}$

    <Particle: name="Sigma+", pdgid=3222, mass=1189.37 ± 0.07 MeV>

$\bar{\Sigma}^{-}$

    <Particle: name="Sigma~-", pdgid=-3222, mass=1189.37 ± 0.07 MeV>

$\Xi^{-}$

    <Particle: name="Xi-", pdgid=3312, mass=1321.71 ± 0.07 MeV>

$\bar{\Xi}^{+}$

    <Particle: name="Xi~+", pdgid=-3312, mass=1321.71 ± 0.07 MeV>

$\Xi^{0}$

    <Particle: name="Xi0", pdgid=3322, mass=1314.86 ± 0.20 MeV>

$\bar{\Xi}^{0}$

    <Particle: name="Xi~0", pdgid=-3322, mass=1314.86 ± 0.20 MeV>

$\Omega^{-}$

    <Particle: name="Omega-", pdgid=3334, mass=1672.5 ± 0.3 MeV>

$\bar{\Omega}^{+}$

    <Particle: name="Omega~+", pdgid=-3334, mass=1672.5 ± 0.3 MeV>

$\Omega_{b}^{-}$

    <Particle: name="Omega(b)-", pdgid=5332, mass=6046.1 ± 1.7 MeV>

$\bar{\Omega}_{b}^{+}$

    <Particle: name="Omega(b)~+", pdgid=-5332, mass=6046.1 ± 1.7 MeV>

```python
import uncertainties
from uncertainties.umath import *

x = uncertainties.ufloat(1, 0.1)   # 1 ± 0.1

print(f"\n2*x              = {2*x}")

print(f"\nsin(2*x)         = {sin(2*x)}")

print(f"\nsqrt(x) - x**0.5 = {sqrt(x) - x**0.5}")   # correlations!
```

    2*x              = 2.00+/-0.20

    sin(2*x)         = 0.91+/-0.08

    sqrt(x) - x**0.5 = 0.0+/-0

```python
import pyhf       # like CmsCombine, but built using Python tools
import awkward

s, b, obs = [[0.0, 0.0, 0.1, 0.2, 0.5, 1.0, 1.8, 2.7, 2.0, 0.3],   # signal peaks around bin 8
             [6.0, 5.5, 4.5, 3.8, 3.3, 2.6, 2.0, 1.8, 1.5, 1.0],   # background steadily falls
             [  6,   7,   4,   4,   4,   3,   4,   3,   4,   0]]   # observations perfer 3/4 signal
model = pyhf.simplemodels.hepdata_like(signal_data=s, bkg_data=b, bkg_uncerts=numpy.sqrt(b).tolist())
def hypotest(mu):
    return pyhf.utils.hypotest(mu, obs + model.config.auxdata, model, return_expected_set=True)

mus = numpy.linspace(0, 2, 30)
CLs = awkward.fromiter([hypotest(mu) for mu in mus])
CLs_observed = CLs[:, 0, 0]
CLs_minus2, CLs_minus1, CLs_expected, CLs_plus1, CLs_plus2 = [CLs[:, 1, i, 0] for i in range(5)]

matplotlib.pyplot.fill_between(mus, CLs_minus2, CLs_plus2, facecolor="yellow");
matplotlib.pyplot.fill_between(mus, CLs_minus1, CLs_plus1, facecolor="limegreen");
matplotlib.pyplot.plot(mus, CLs_expected, c="black", linestyle="dotted");
matplotlib.pyplot.plot(mus, CLs_observed, c="black", marker="o");
```


# Neural Nets using [JAX](https://github.com/google/jax#readme)
**JAX is NumPy on the CPU, GPU, and TPU, with great [automatic differentiation](https://en.wikipedia.org/wiki/Automatic_differentiation) for high-performance machine learning research.**

Version 0.1, in `nn-jax`

Should do [Working efficiently with jupyter lab](https://florianwilhelm.info/2018/11/working_efficiently_with_jupyter_lab/)

When this was a notebook with integrated tests, we did: \
`
%load_ext autoreload
%autoreload 2
%matplotlib widget
#%matplotlib inline`

```python
import numpy as np
import jax.numpy as jnp
from jax import grad, jit, vmap
from jax import random
```

We'll be generating random data in the following examples. One big difference between NumPy and JAX is how you generate random numbers. For more details, see [Common Gotchas in JAX].

[Common Gotchas in JAX]: https://jax.readthedocs.io/en/latest/notebooks/Common_Gotchas_in_JAX.html#%F0%9F%94%AA-Random-Numbers

```python
key = random.PRNGKey(0)
x = random.normal(key, (10,))
```

A network built of components which:
1. accept an ordered set of reals (we'll use `numpy.array`, and  call them vectors) at the input port and produce another at the output port - this is forward propagation. ${\displaystyle f\colon \mathbf {R} ^{n}\to \mathbf {R} ^{m}}$
1. accept an ordered set of reals at the output port, representing the gradient of the loss function at the output, and produce the gradient of the loss function at the input port - this is back propagation, aka backprop. ${\displaystyle b\colon \mathbf {R} ^{m}\to \mathbf {R} ^{n}}$
1. from the gradient of the loss function at the output, calculate the partial of the loss function w.r.t the internal parameters ${\displaystyle \frac{\partial E}{\partial w} }$
1. accept a scalar $\eta$ to control the adjustment of internal parameters. _Or is this effected by scaling the loss gradient before passing? YES_
1. update internal parameters ${\displaystyle w \leftarrow w - \eta \frac{\partial E}{\partial w} }$

```python
class Layer:
    def __init__(self):
        pass

    def __call__(self, x):
        """Compute response to input"""
        raise NotImplementedError

    def backprop(self, output_delE):
        """Use output error gradient to adjust internal parameters, return gradient of error at input"""
        raise NotImplementedError

    def state_vector(self):
        """Provide the layer's learnable state as a vector"""
        raise NotImplementedError

    def set_state_from_vector(self, sv):
        """Set the layer's learnable state from a vector"""
        raise NotImplementedError
```

A network built of a cascade of layers:

```python
class Network:
    def __init__(self):
        self.layers = []
        self.eta = 0.1 #FIXME

    def extend(self, net):
        self.layers.append(net)
        return self

    def __call__(self, input):
        v = input
        for net in self.layers:
            v = net(v)
        return v

    def learn(self, facts, eta=None):
        self.eta = eta or self.eta
        for x, ideal in facts:
            y = self(x)
            e = y - ideal
            egrad = e * self.eta / e.shape[0]
            for net in reversed(self.layers):
                egrad = net.backprop(egrad)
        #loss = float(e.dot(e.T))/2.0
        loss = jnp.einsum('...ij,...ij', e, e) / (2.0 * e.shape[0])
        return loss

    def losses(self, facts):
        return [jnp.einsum('...ij,...ij', e, e) / (2.0 * e.shape[0]) \
                for e in (self(x) - ideal for x, ideal in facts)]

    def state_vector(self):
        """Provide the network's learnable state as a vector"""
        return jnp.concatenate([layer.state_vector() for layer in self.layers])

    def set_state_from_vector(self, sv):
        """Set the layer's learnable state from a vector"""
        i = 0
        for layer in self.layers:
            lsvlen = len(layer.state_vector())
            layer.set_state_from_vector(sv[i:i+lsvlen])
            i += lsvlen
```

___

## Useful Layers

### Identify

```python
class IdentityLayer(Layer):
    def __call__(self, x):
        return x

    def backprop(self, output_delE):
        return output_delE

    def state_vector(self):
        return jnp.array([])

    def set_state_from_vector(self, sv):
        pass
```

### Affine
A layer that does an [affine transformation](https://mathworld.wolfram.com/AffineTransformation.html) aka affinity, which is the classic fully-connected layer with output offsets.

$$ \mathbf{M} \mathbf{x} + \mathbf{b} = \mathbf{y} $$
where
$$
\mathbf{x} = \sum_{j=1}^{n} x_j \mathbf{\hat{x}}_j \\
\mathbf{b} = \sum_{i=1}^{m} b_i \mathbf{\hat{y}}_i \\
\mathbf{y} = \sum_{i=1}^{m} y_i \mathbf{\hat{y}}_i
$$
and $\mathbf{M}$ can be written
$$
\begin{bmatrix}
    m_{1,1} & \dots & m_{1,n} \\
    \vdots & \ddots & \vdots \\
    m_{m,1} & \dots & m_{m,n}
\end{bmatrix} \\
$$

#### Error gradient back-propagation
$$
\begin{align}
 \frac{\partial loss}{\partial\mathbf{x}}
  &= \frac{\partial loss}{\partial\mathbf{y}} \frac{\partial\mathbf{y}}{\partial\mathbf{x}} \\
  &= \mathbf{M}^\mathsf{T}\frac{\partial loss}{\partial\mathbf{y}}
\end{align}
$$

#### Parameter adjustment
$$
 \frac{\partial loss}{\partial\mathbf{M}}
 = \frac{\partial loss}{\partial\mathbf{y}} \frac{\partial\mathbf{y}}{\partial\mathbf{M}}
 = \frac{\partial loss}{\partial\mathbf{y}} \mathbf{x} \\
 \frac{\partial loss}{\partial\mathbf{b}}
 = \frac{\partial loss}{\partial\mathbf{y}} \frac{\partial\mathbf{y}}{\partial\mathbf{b}}
 = \frac{\partial loss}{\partial\mathbf{y}}
$$

#### Adapting to `numpy`

In `numpy` it is more convenient to use row vectors, particularly for calculating the transform on multiple inputs in one operation. We use the identity $ \mathbf{M} \mathbf{x} = (\mathbf{x} \mathbf{M}^\mathsf{T})^\mathsf{T}.$ To avoid cluttering names, we will use `M` in the code below to hold $\mathbf{M}^\mathsf{T}$.

```python
class AffineLayer(Layer):
    """An affine transformation, which is the classic fully-connected layer with offsets.

    The layer has n inputs and m outputs, which numbers must be supplied
    upon creation. The inputs and outputs are marshalled in numpy arrays, 1-D
    in the case of a single calculation, and 2-D when calculating the outputs
    of multiple inputs in one call.
    If called with 1-D array having shape == (n,), e.g numpy.arange(n), it will
    return a 1-D numpy array of shape (m,).
    If called with a 2-D numpy array, input shall have shape (k,n) and will return
    a 2-D numpy array of shape (k,m), suitable as input to a subsequent layer
    that has input width m.
    """
    def __init__(self, n, m):
        self.M = jnp.empty((n, m))
        self.b = jnp.empty(m)
        self.randomize()

    def randomize(self):
        self.M = random.normal(key, self.M.shape, dtype=jnp.float32)
        self.b = random.normal(key, self.b.shape, dtype=jnp.float32)

    def __call__(self, x):
        self.input = x
        self.output = x @ self.M + self.b
        return self.output

    def backprop(self, output_delE):
        input_delE = output_delE @ self.M.T
        o_delE = jnp.atleast_2d(output_delE)
        self.M -= jnp.einsum('...ki,...kj->...ji', o_delE, jnp.atleast_2d(self.input))
        self.b -= jnp.sum(o_delE, 0)
        return input_delE

    def state_vector(self):
        return jnp.concatenate((self.M.ravel(), self.b.ravel()))

    def set_state_from_vector(self, sv):
        """Set the layer's learnable state from a vector"""
        l_M = len(self.M.ravel())
        l_b = len(self.b.ravel())
        self.M = sv[:l_M].reshape(self.M.shape)
        self.b = sv[l_M : l_M + l_b].reshape(self.b.shape)
```

### Map
Maps a scalar function on the inputs, for e.g. activation layers.

```python
class MapLayer(Layer):
    """Map a scalar function on the input taken element-wise"""
    def __init__(self, fun, dfundx):
        self.vfun = jnp.vectorize(fun)
        self.vdfundx = jnp.vectorize(dfundx)

    def __call__(self, x):
        self.input = x
        return self.vfun(x)

    def backprop(self, output_delE):
        input_delE = self.vdfundx(self.input) * output_delE
        return input_delE

    def state_vector(self):
        return jnp.array([])

    def set_state_from_vector(self, sv):
        pass
```

---

# Tests
*Incomplete* \
Also `unittest` the `.py` version with a separate test script, see `test-nn_v3.py`.

Make a few test arrays:

```python
if __name__ == '__main__':
    one_wide = jnp.atleast_2d(jnp.arange(1*4)).reshape(-1,1)
    print(f"one_wide is:\n{one_wide}")
    two_wide = jnp.arange(2*4).reshape(-1,2)
    print(f"two_wide is:\n{two_wide}")
    three_wide = jnp.arange(3*4).reshape(-1,3)
    print(f"three_wide is:\n{three_wide}\n")
```

    one_wide is:
    [[0]
     [1]
     [2]
     [3]]
    two_wide is:
    [[0 1]
     [2 3]
     [4 5]
     [6 7]]
    three_wide is:
    [[ 0  1  2]
     [ 3  4  5]
     [ 6  7  8]
     [ 9 10 11]]

## Tooling for Testing

```python
if __name__ == '__main__':
    import sympy

    class VC():
        def grad(f, x, eps=1e-3):
            #epsihat = jnp.eye(x.size) * eps
            epsihat = jnp.eye(x.shape[-1]) * eps
            yp = jnp.apply_along_axis(f, 1, x + epsihat)
            ym = jnp.apply_along_axis(f, 1, x - epsihat)
            return (yp - ym)/(2 * eps)

        def tensor_grad(f, x, eps=1e-3):
            return jnp.apply_along_axis(lambda v: VC.grad(f, v, eps), 1, x)

    def closenuf(a, b, tol=0.001):
        return jnp.allclose(a, b, rtol=tol)

    def arangep(n, starting_index=0):
        sympy.sieve.extend_to_no(starting_index + n)
        return jnp.array(sympy.sieve._list[starting_index:starting_index + n])
```

```python
#VC.grad(lambda x:x**2, three_wide[1])
```

```python
#VC.tensor_grad(lambda x:x**2, three_wide)
```

---

Input to a layer can be a single (row) vector, or a vertical stack of row vectors,
a 2-d array that resembles a matrix. We need to test each layer class with both single and stacked input.

## Identity layer

```python
if __name__ == '__main__':
    iL = IdentityLayer()

    # It's transparent from input to output
    assert jnp.equal(iL(jnp.arange(5)), jnp.arange(5)).all()
    assert (iL(three_wide) == three_wide).all()

    # It back-propagates the loss gradient without alteration
    assert jnp.equal(iL.backprop(jnp.arange(7)), jnp.arange(7)).all()
    assert (iL.backprop(three_wide) == three_wide).all()

    # It works for stacked input
    # (see above)
```

## Map layer

#### Test single vector input behavior

```python
if __name__ == '__main__':
    mL = MapLayer(lambda x:x**2, lambda d:2*d)

    # It applies the forward transformation
    assert jnp.equal(mL(jnp.array([-2,1,3])), jnp.array([4,1,9])).all()

    # It back-propagages the loss gradient
    x = jnp.array([1,2,2])
    y = mL(x)

    # for loss function, use L2-distance from some ideal
    # (divided by 2, for convenient gradient = error)
    ideal = jnp.array([2,3,5])
    loss = lambda v: (v - ideal).dot(v - ideal) / 2.0
    loss_at_y = loss(y)
    print(f"x = {x}, y = {y}, loss at y = {loss_at_y}")

    # find numerical gradient of loss function at y, the layer output
    grad_y = VC.grad(loss, y)
    print(f"∇𝑙𝑜𝑠𝑠(𝑦) = {grad_y}")

    # find the numerical gradient of the loss w.r.t. the input of the layer
    grad_x = VC.grad(lambda x:loss(mL(x)), x)
    print(f"∇𝑙𝑜𝑠𝑠(𝑥) = {grad_x}")

    # The backprop method does the same
    _ = mL(x) # Make sure the last x is in the right place
    in_delE = mL.backprop(grad_y)
    print(f"backprop({grad_y}) = {in_delE}")
    assert closenuf(in_delE, grad_x)

    # The backprop operation did not change the behavior
    assert jnp.equal(mL(x), y).all()
```

    x = [1 2 2], y = [1 4 4], loss at y = 1.5
    ∇𝑙𝑜𝑠𝑠(𝑦) = [-0.99998707  0.99992746 -0.99992746]
    ∇𝑙𝑜𝑠𝑠(𝑥) = [-1.9999741  3.9999483 -3.9999483]
    backprop([-0.99998707  0.99992746 -0.99992746]) = [-1.9999741  3.9997098 -3.9997098]

#### Test stacked-vectors input:

```python
if __name__ == '__main__':
    mL = MapLayer(lambda x:x**2, lambda d:2*d)

    two_wide_sq = jnp.array([[ 0,  1],
                            [ 4,  9],
                            [16, 25],
                            [36, 49]])
    # It applies the forward transformation
    assert jnp.equal(mL(two_wide), two_wide_sq).all()

    # It back-propagages the loss gradient
    x = two_wide
    y = mL(x)

    # for loss function, use L2-distance from some ideal
    # (divided by 2, for convenient gradient = error)
    ideal = two_wide * 2 + 11
    #print(y - ideal)
    #loss = lambda v: (v - ideal).dot(v - ideal) / 2.0
    loss = lambda v: jnp.einsum('ij,ij', v-ideal, v-ideal) / (2 * v.shape[0])
    loss_at_y = loss(y)
    print(f"x =\n{x}\ny =\n{y}, loss = {loss_at_y}\n")

    # find numerical gradient of loss function at y, the layer output
    grad_y = VC.tensor_grad(loss, y)
    print(f"∇𝑙𝑜𝑠𝑠(𝑦) =\n{grad_y}\n")

    # find the numerical gradient of the loss w.r.t. the input of the layer
    grad_x = VC.tensor_grad(lambda x:loss(mL(x)), x)
    print(f"∇𝑙𝑜𝑠𝑠(𝑥) =\n{grad_x}\n")

    # The backprop method does the same
    _ = mL(x) # Make sure the last x is in the right place
    in_delE = mL.backprop(grad_y)
    print(f"backprop({grad_y}) =\n{in_delE}")
    assert closenuf(in_delE, grad_x)

    # The backprop operation did not change the behavior
    assert jnp.equal(mL(x), y).all()
```

    x =
    [[0 1]
     [2 3]
     [4 5]
     [6 7]]
    y =
    [[ 0  1]
     [ 4  9]
     [16 25]
     [36 49]], loss = 152.5

    ∇𝑙𝑜𝑠𝑠(𝑦) =
    [[-33.99658   -36.01074  ]
     [-26.000975  -19.989012 ]
     [ -1.9989012  11.993407 ]
     [ 37.963863   59.93652  ]]

    ∇𝑙𝑜𝑠𝑠(𝑥) =
    [[   0.        -71.99096 ]
     [-104.0039   -119.99511 ]
     [ -15.995025  119.991295]
     [ 455.99362   839.9047  ]]

    backprop([[-33.99658   -36.01074  ]
     [-26.000975  -19.989012 ]
     [ -1.9989012  11.993407 ]
     [ 37.963863   59.93652  ]]) =
    [[  -0.        -72.02148 ]
     [-104.0039   -119.93407 ]
     [ -15.99121   119.934074]
     [ 455.56635   839.11127 ]]

## Affine layer

#### Test single vector input behavior

Test, for single input-vector operations:
* input and output widths
* state vector setting and getting
* forward calculation

```python
if __name__ == '__main__':
    # Affine
    a = AffineLayer(2,3)
    key = random.PRNGKey(0)
    #x = random.normal(key, (10,))

    # The input and output widths are correct
    assert a(jnp.arange(2)).shape == (3,)

    # Its internal state can be set
    a.set_state_from_vector(jnp.arange(9))
    # and read back
    assert (a.state_vector() == jnp.arange(9)).all()
    # NOTE: The two assertions below are commented out because they depend
    # on white-box knowledge, and are duplicative of other tests
    #assert jnp.equal(a.M, jnp.array([[0, 1, 2],
    #                               [3, 4, 5]])).all()
    #assert jnp.equal(a.b, jnp.array([6, 7, 8])).all()

    # Its internal state observed using numerical gradient is correct
    x = random.uniform(key, (2,))
    y = a(x)
    dydx = VC.grad(a, x)
    b = y - x.dot(dydx)
    #print(dydx, b)
    #print(dydx, jnp.arange(6).reshape(2,-1))
    assert closenuf(dydx, jnp.arange(6).reshape(2, -1))
    #print(b, jnp.arange(6, 9))
    assert closenuf(b, jnp.arange(6, 9))

    # It performs a single-input forward calculation correctly
    x = jnp.array([2, 1])
    y = a(x)
    #print(f"a.M is:\n{a.M}\na.b is {a.b}\nx is: {x}\ny is: {y}\n")
    assert (y == jnp.array([9, 13, 17])).all()

    # It performs a different single-input forward calculation correctly
    a.set_state_from_vector(jnp.array([ 2,  3,  5,  7, 11, 13, 17, 19, 23]))
    x = jnp.array([[29, 31]])
    y = a(x)
    assert (y == jnp.array([[292, 447, 571]])).all()
```

Test, for single input-vector operations:
* back-propagation of the loss gradient
* learning (change in forward function) from the back-prop operation

```python
if __name__ == '__main__':
    # Affine
    a = AffineLayer(2,3)
    a.set_state_from_vector(jnp.arange(9))

    # Doing a single-input-vector calculation
    x = jnp.array([2, 1])
    y = a(x)
    assert jnp.equal(y, jnp.array([9, 13, 17])).all()

    # It back-propagages the loss gradient
    ideal = jnp.array([11,12,10])
    loss = lambda v: (v - ideal).dot(v - ideal) / 2.0
    loss_at_y = loss(y)
    print(f"x = {x}, y = {y}, loss = {loss_at_y}")
    grad_y = VC.grad(loss, y)
    print(f"∇𝑙𝑜𝑠𝑠(𝑦) = {grad_y}")
    grad_x = VC.grad(lambda x:loss(a(x)), x)
    print(f"∇𝑙𝑜𝑠𝑠(𝑥) = {grad_x}")

    # Back-propagate the loss gradient from layer output to input
    _ = a(x) # Make sure the last x is in the right place
    out_delE = grad_y * 0.1 # Backprop one-tenth of the loss gradient
    in_delE = a.backprop(out_delE)
    print(f"backprop({out_delE}) = {in_delE}")

    # The loss gradient back-propagated to the layer input is correct
    assert closenuf(in_delE / 0.1, grad_x)

    # And how did the learning affect the layer?
    print(f"Now a({x}) = {a(x)}, loss = {loss(a(x))}")
    print(f"state_vector is {a.state_vector()}")
    # FIXME: Check the change is correct
```

    x = [2 1], y = [ 9 13 17], loss = 27.0
    ∇𝑙𝑜𝑠𝑠(𝑦) = [-2.0008085  1.0004042  6.9961543]
    ∇𝑙𝑜𝑠𝑠(𝑥) = [14.991759 33.006664]
    backprop([-0.20008086  0.10004043  0.6996154 ]) = [1.4992713 3.2979963]
    Now a([2 1]) = [10.200485 12.399757 12.802307], loss = 4.325977325439453
    state_vector is [0.4001617  0.7999191  0.60076916 3.2000809  3.8999596  4.3003845
     6.200081   6.8999596  7.3003845 ]

#### Test batch operations

Test, for batch operations:
* input and output widths
* forward calculation

```python
if __name__ == '__main__':
    # Affine
    a = AffineLayer(2,3)
    a.set_state_from_vector(jnp.arange(9))

    # The input and output widths for the forward calculation are correct
    x = two_wide
    y = a(two_wide)
    assert y.shape[0] == x.shape[0]
    assert y.shape[1] == 3

    # The input and output widths for the backprop calculation are correct
    bp = a.backprop(three_wide * 0.001)
    assert bp.shape[0] == three_wide.shape[0]
    assert bp.shape[1] == x.shape[1]

    # The forward calculation is correct (in at least two instances)
    a.set_state_from_vector(jnp.arange(9))
    x = jnp.array([[0, 1],
                  [2, 3],
                  [4, 5],
                  [6, 7]])
    assert (a(x) == jnp.array([[ 9, 11, 13],
                              [15, 21, 27],
                              [21, 31, 41],
                              [27, 41, 55]])).all()
    #print(f"a.M is:\n{a.M}\na.b is {a.b}\nx is: {x}\ny is: {y}")
    a.set_state_from_vector(jnp.array([ 2,  3,  5,  7, 11, 13, 17, 19, 23]))
    y = a(x)
    #print(f"x is: {x}\ny is: {y}")
    assert (y == jnp.array([[ 24,  30,  36],
                           [ 42,  58,  72],
                           [ 60,  86, 108],
                           [ 78, 114, 144]])).all()
```

Test, for batch operations:
* back-propagation of the loss gradient
* learning (change in forward function) from the back-prop operation

```python
if __name__ == '__main__':
    # Affine
    a = AffineLayer(2,3)
    a.set_state_from_vector(jnp.arange(9))
    x = jnp.array([[0, 1],
                  [2, 3],
                  [4, 5],
                  [6, 7]])
    y = a(x)

    # It back-propagages the loss gradient

    # for loss function, use L2-distance from some ideal
    # (divided by 2, for convenient gradient = error)
    ideal = x @ arangep(2*3).reshape(2,3) + arangep(3,6) # A known, different parameter setting
    print(f"y - ideal =\n{y - ideal}")
    #loss = lambda v: (v - ideal).dot(v - ideal) / 2.0
    loss = lambda v: jnp.einsum('ij,ij', v-ideal, v-ideal) / (2 * v.shape[0])
    loss_at_y = loss(y)
    print(f"x =\n{x}\nideal =\n{ideal}\ny =\n{y}, loss = {loss_at_y}\n")

    # find numerical gradient of loss function at y, the layer output
    grad_y = VC.tensor_grad(loss, y)
    print(f"∇𝑙𝑜𝑠𝑠(𝑦) =\n{grad_y}")

    # find the numerical gradient of the loss w.r.t. the input of the layer
    grad_x = VC.tensor_grad(lambda x:loss(a(x)), x)
    print(f"∇𝑙𝑜𝑠𝑠(𝑥) =\n{grad_x}")

    # Back-propagate the loss gradient from layer output to input
    _ = a(x) # Make sure the last x is in the right place
    out_delE = grad_y * 0.01 # Backprop one percent of the loss gradient
    in_delE = a.backprop(out_delE)
    print(f"backprop({out_delE}) = {in_delE}")

    # The loss gradient back-propagated to the layer input is correct
    #assert closenuf(in_delE / 0.1, grad_x)

    # And how did the learning affect the layer?
    print(f"Now a({x}) = {a(x)}, loss = {loss(a(x))}")
    print(f"state_vector is {a.state_vector()}")
    # FIXME: Check the change is correct
```

    y - ideal =
    [[-15 -19 -23]
     [-27 -37 -45]
     [-39 -55 -67]
     [-51 -73 -89]]
    x =
    [[0 1]
     [2 3]
     [4 5]
     [6 7]]
    ideal =
    [[ 24  30  36]
     [ 42  58  72]
     [ 60  86 108]
     [ 78 114 144]]
    y =
    [[ 9 11 13]
     [15 21 27]
     [21 31 41]
     [27 41 55]], loss = 3765.5

    ∇𝑙𝑜𝑠𝑠(𝑦) =
    [[ -55.66406   -81.05468  -102.539055]
     [ -48.33984   -67.871086  -83.98437 ]
     [ -40.03906   -54.687496  -65.18554 ]
     [ -31.98242   -41.25976   -46.630856]]
    ∇𝑙𝑜𝑠𝑠(𝑥) =
    [[ -286.13278 -1006.8359 ]
     [ -235.83983  -835.44916]
     [ -185.30272  -665.039  ]
     [ -134.52147  -494.38474]]
    backprop([[-0.55664057 -0.81054676 -1.0253905 ]
     [-0.48339838 -0.6787108  -0.83984363]
     [-0.40039057 -0.54687494 -0.65185535]
     [-0.3198242  -0.4125976  -0.46630853]]) = [[ -2.8613276 -10.039061 ]
     [ -2.358398   -8.364256 ]
     [ -1.8505857  -6.6479483]
     [ -1.3452146  -4.9414053]]
    Now a([[0 1]
     [2 3]
     [4 5]
     [6 7]]) = [[ 17.007812  21.917969  26.051758]
     [ 44.47754   60.897453  74.35839 ]
     [ 71.947266  99.876945 122.66502 ]
     [ 99.41699  138.85643  170.97166 ]], loss = 323.4553527832031
    state_vector is [ 4.487304  7.020507  9.084959  9.247559 12.469237 15.068358  7.760254
      9.44873  10.983398]

#### Test batch operations when the affine layer has only one input

Test, for batch operations:
* input and output widths
* forward calculation

```python
if __name__ == '__main__':
    # Affine
    a = AffineLayer(1,3)
    a.set_state_from_vector(jnp.arange(6))

    # The input and output widths for the forward calculation are correct
    x = one_wide
    y = a(one_wide)
    assert y.shape[0] == x.shape[0]
    assert y.shape[1] == 3

    # The input and output widths for the backprop calculation are correct
    bp = a.backprop(three_wide * 0.001)
    assert bp.shape[0] == three_wide.shape[0]
    assert bp.shape[1] == x.shape[1]

    # The forward calculation is correct (in at least two instances)
    a.set_state_from_vector(jnp.arange(6))
    x = jnp.array([[0],
                  [1],
                  [2],
                  [3]])
    assert (y == jnp.array([[ 3.,  4.,  5.],
                           [ 3.,  5.,  7.],
                           [ 3.,  6.,  9.],
                           [ 3.,  7., 11.]])).all()
    #print(f"a.M is:\n{a.M}\na.b is {a.b}\nx is: {x}\ny is: {y}")
    a.set_state_from_vector(jnp.array([ 2,  3,  5,  7, 11, 13]))
    y = a(x)
    #print(f"x is: {x}\ny is: {y}")
    assert (a(x) == jnp.array([[ 7, 11, 13],
                              [ 9, 14, 18],
                              [11, 17, 23],
                              [13, 20, 28]])).all()
```

Test, for batch operations:
* back-propagation of the loss gradient
* learning (change in forward function) from the back-prop operation

```python
if __name__ == '__main__':
    # Affine
    a = AffineLayer(1,3)
    a.set_state_from_vector(jnp.arange(6))
    x = jnp.array([[0],
                  [1],
                  [2],
                  [3]])
    y = a(x)
    #print(f"x =\n{x}\ny =\n{y}")

    # It back-propagages the loss gradient

    # for loss function, use L2-distance from some ideal
    # (divided by 2, for convenient gradient = error)
    ideal = x @ arangep(1*3).reshape(1,3) + arangep(3,6) # A known, different parameter setting
    print(f"y - ideal =\n{y - ideal}")
    #loss = lambda v: (v - ideal).dot(v - ideal) / 2.0
    loss = lambda v: jnp.einsum('ij,ij', v-ideal, v-ideal) / (2 * v.shape[0])
    loss_at_y = loss(y)
    print(f"x =\n{x}\nideal =\n{ideal}\ny =\n{y}, loss = {loss_at_y}\n")

    # find numerical gradient of loss function at y, the layer output
    grad_y = VC.tensor_grad(loss, y)
    print(f"∇𝑙𝑜𝑠𝑠(𝑦) =\n{grad_y}")

    # find the numerical gradient of the loss w.r.t. the input of the layer
    grad_x = VC.tensor_grad(lambda x:loss(a(x)), x)
    print(f"∇𝑙𝑜𝑠𝑠(𝑥) =\n{grad_x}")

    # Back-propagate the loss gradient from layer output to input
    _ = a(x) # Make sure the last x is in the right place
    out_delE = grad_y * 0.01 # Backprop one percent of the loss gradient
    in_delE = a.backprop(out_delE)
    print(f"backprop({out_delE}) = {in_delE}")

    # The loss gradient back-propagated to the layer input is correct
    #assert closenuf(in_delE / 0.1, grad_x)

    # And how did the learning affect the layer?
    print(f"Now a({x}) = {a(x)}, loss = {loss(a(x))}")
    print(f"state_vector is {a.state_vector()}")
    # FIXME: Check the change is correct
```

    y - ideal =
    [[-14 -15 -18]
     [-16 -17 -21]
     [-18 -19 -24]
     [-20 -21 -27]]
    x =
    [[0]
     [1]
     [2]
     [3]]
    ideal =
    [[17 19 23]
     [19 22 28]
     [21 25 33]
     [23 28 38]]
    y =
    [[ 3  4  5]
     [ 3  5  7]
     [ 3  6  9]
     [ 3  7 11]], loss = 570.25

    ∇𝑙𝑜𝑠𝑠(𝑦) =
    [[-22.613523 -25.970457 -33.935543]
     [-22.644041 -24.627684 -31.311033]
     [-22.644041 -23.254393 -28.625486]
     [-22.674559 -21.972654 -26.000975]]
    ∇𝑙𝑜𝑠𝑠(𝑥) =
    [[-93.90258]
     [-87.3413 ]
     [-80.62743]
     [-73.9746 ]]
    backprop([[-0.22613522 -0.25970456 -0.3393554 ]
     [-0.2264404  -0.24627683 -0.31311032]
     [-0.2264404  -0.23254392 -0.28625485]
     [-0.22674558 -0.21972653 -0.26000974]]) = [[-0.9384154 ]
     [-0.87249744]
     [-0.8050536 ]
     [-0.739746  ]]
    Now a([[0]
     [1]
     [2]
     [3]]) = [[ 3.9057617  4.958252   6.1987305]
     [ 5.26532    7.3287964  9.86438  ]
     [ 6.6248775  9.699341  13.530029 ]
     [ 7.984435  12.069884  17.195679 ]], loss = 389.4485778808594
    state_vector is [1.3595579 2.3705442 3.6656492 3.9057617 4.958252  6.1987305]

## Network

### Network assembly

The simplest, the empty network, does identity:

```python
if __name__ == '__main__':
    net = Network()
    assert all(x == net(x) for x in [0, 42, 'cows in trouble'])
    assert all((x == net(x)).all() for x in [np.arange(7), jnp.arange(3*4*5).reshape(3,4,5)])
```

A stack of maps composes the operations:

```python
if __name__ == '__main__':
    net = Network()
    net.extend(MapLayer(lambda x: x**3, lambda d: 3*d**2))
    x = jnp.array([0, 2, 3, 42, -3.14])
    assert (net(x) == x**3).all()
    net.extend(MapLayer(lambda x: 7-x, lambda d: -1))
    assert (net(x) == 7-x**3).all()

    # It operates on each element of an input vector separately
    assert (net(jnp.arange(4)) == 7 - jnp.arange(4) ** 3).all()
```

A composition of affine transformations

_[to do someday]_

### Network Learning

Test simple batch learning of a single affine layer

```python
if __name__ == '__main__':
    net = Network()
    a = AffineLayer(2,3)
    a.set_state_from_vector(jnp.arange(9)) # A well-known initial state
    net.extend(a)
    print(f"\nNet has state {net.state_vector()}")

    x = jnp.array([[0, 1],
                  [2, 3],
                  [4, 5],
                  [6, 7]])

    # The net wraps the layer
    y = a(x)
    assert (net(x) == y).all()

    # Make the training batch.
    # We use a separate affine layer, initialized differently, to determine the ideal
    t_a = AffineLayer(2,3)
    t_a.set_state_from_vector(arangep(9)) # A known different initial state (of primes)
    ideal = t_a(x)

    facts = [(x, ideal)]
    print(f"facts are:\n{facts}\n")
    print(f"net(x) =\n{net(x)}")

    net.eta = 0.01
    for i in range(10):
        print(f"net.learn(facts) = {net.learn(facts)}")
    print(f"net(x) =\n{net(x)}")

    # A simple fact yielder:
    def fact_ory(facts, n):
        for i in range(n):
            yield facts

    #print(f"list(fact_ory(facts[0], 3)) =\n{list(fact_ory(facts[0], 3))}\n")
    print(f"net.learn(fact_ory(facts[0],10)) = {net.learn(fact_ory(facts[0],10))}")
    print(f"net(x) =\n{net(x)}")
    for i in range(1000):
        loss = net.learn(fact_ory(facts[0],10))
        if loss < 1e-25:
            break
    print(f"did {(i+1)*10} more learnings of fact. Now loss is {loss}")
    print(f"net(x) =\n{net(x)}")

    print(f"net.state_vector() = {net.state_vector()}")

    # The network has learned the target transform
    assert closenuf(net(x), facts[0][1])

    # Save prior results and learn again, with different batch clustering
    prev_run_loss = loss
    prev_y = net(x)
    net.set_state_from_vector(jnp.arange(9)) # A well-known initial state
    print(f"\nReset net to state {net.state_vector()}")

    # Try multiple batches in each call to Network.learn
    def multibatch_fact_ory(facts, n):
        for i in range(n//2):
            yield facts * 2

    for i in range(1000):
        loss = net.learn(fact_ory(facts[0],10))
        if loss < 1e-25:
            break
    print(f"did {(i+1)*10} learnings of fact. Now loss is {loss}")
    print(f"net(x) =\n{net(x)}")

    # The results should match exactly
    assert loss == prev_run_loss
    assert (net(x) == prev_y).all()
```

    Net has state [0 1 2 3 4 5 6 7 8]
    facts are:
    [(DeviceArray([[0, 1],
                 [2, 3],
                 [4, 5],
                 [6, 7]], dtype=int32), DeviceArray([[ 24,  30,  36],
                 [ 42,  58,  72],
                 [ 60,  86, 108],
                 [ 78, 114, 144]], dtype=int32))]

    net(x) =
    [[ 9 11 13]
     [15 21 27]
     [21 31 41]
     [27 41 55]]
    net.learn(facts) = 3765.5
    net.learn(facts) = 1605.49853515625
    net.learn(facts) = 708.707275390625
    net.learn(facts) = 336.169189453125
    net.learn(facts) = 181.206787109375
    net.learn(facts) = 116.54439544677734
    net.learn(facts) = 89.3607177734375
    net.learn(facts) = 77.7341537475586
    net.learn(facts) = 72.56697082519531
    net.learn(facts) = 70.0833511352539
    net(x) =
    [[ 14.699151  18.972603  22.707714]
     [ 36.72516   51.692673  64.392944]
     [ 58.75116   84.41274  106.078186]
     [ 80.77717  117.13282  147.76343 ]]
    net.learn(fact_ory(facts[0],10)) = 63.1659049987793
    net(x) =
    [[ 15.146448  19.51728   23.365414]
     [ 37.197094  52.312637  65.145096]
     [ 59.24774   85.107994 106.924774]
     [ 81.29839  117.90334  148.70447 ]]
    did 10000 more learnings of fact. Now loss is 1.9352683011675254e-08
    net(x) =
    [[ 23.999866  29.99986   35.999733]
     [ 41.999924  57.99992   71.99985 ]
     [ 59.99998   85.999985 107.99996 ]
     [ 78.00004  114.00004  144.00008 ]]
    net.state_vector() = [-0.999908   0.6666431  1.6668558  9.999937  13.333387  16.333202
     13.999928  16.666473  19.666529 ]

    Reset net to state [0 1 2 3 4 5 6 7 8]
    did 10000 learnings of fact. Now loss is 1.9352683011675254e-08
    net(x) =
    [[ 23.999866  29.99986   35.999733]
     [ 41.999924  57.99992   71.99985 ]
     [ 59.99998   85.999985 107.99996 ]
     [ 78.00004  114.00004  144.00008 ]]

### Test Network.losses

```python
if __name__ == '__main__':
    # Make a network. Leave it with the default identity behavior.
    net = Network()

    x = jnp.array([[0, 1],
                  [2, 3],
                  [4, 5],
                  [6, 7]])
    ideal = net(x)
    facts = [(x, ideal), (x, ideal-np.array([1,-1])), (x, 2*x)]
    assert (net.losses(facts) == [0, 1, 17.5])

    # Add some layers
    net.extend(AffineLayer(2,3)).extend(MapLayer(jnp.sin, jnp.cos)).extend(AffineLayer(3,2))
    # Place it in a known state for test repeatability
    net.set_state_from_vector(jnp.arange(len(net.state_vector())))
    ideal = net(x)
    facts = [(x, ideal), (x, ideal-np.array([1,-1]))]
    #print(net.losses(facts))
    assert (net.losses(facts) == [0, 1])
```

---

To produce an importable `nn.py`:
1. Save this notebook
1. Uncomment the `jupyter nbconvert` line below
1. Execute it.
1. Comment out the convert again
1. Save the notebook again in that form

```python
###!jupyter nbconvert --to script nn-jax.ipynb
```

```python

```