# Deep dive into Tensorflow example “Partial Differential Equations”

## Problem statement

Tensorflow documentation provides very [nice tutorial examples](https://www.tensorflow.org/tutorials/). One of the non machine learning examples called ["Partial Differential Equations"](https://www.tensorflow.org/tutorials/non-ml/pdes), which models a surface of a pond as few raindrops land on it.

Tensorflow code for this example is straightforward, but mathematics and physics behind are not. I can only imagine present time physicist students, who can say from top of their head on what is going on in this example. So far i'm pretty sure they still would spend some time refreshing some necessary knowledge.

I was a physicist student some time ago, so let me try to shed some light. As a discamiler i want to say that in some statements i potentially can be wrong, so far in general i have a sense of overall picture now (after 3 nights :) ).

## Quick source code explanation

Original source code: https://www.tensorflow.org/tutorials/non-ml/pdes

Extra changes:
- added speed of the waves $c$
- Laplacian kernel divided by 2 to be aligned with kernel from Wikipedia and other papers

If you already know the code, then feel free to skip to [Mathematics](#Mathematics) section.

In this example author tries to model a physical process that happens when a rain drop hit a water surface. Tensorflow framework mostly used here to do a parallel linear algebra computations. As you know Tensorflow has two things in it's core: tensor and operation. Tensor is usually some vector (matrix) or scalar. Operation can be some linear algebra operation (e.g. matrix multiplication). Both of them are part (nodes) of internal TensorFlow graph. This graph can be executed within a TensorFlow session.

So let's see what tensors and operations we need to model our physical process in TensorFlow.

First import necessary libraries and define function that will draw a pond.

```python
#Import libraries for simulation
import tensorflow as tf
import numpy as np

#Imports for visualization
import PIL.Image
from io import BytesIO
from IPython.display import clear_output, Image, display

def DisplayArray(a, fmt='jpeg', rng=[0,1]):
  """Display an array as a picture."""
  a = (a - rng[0])/float(rng[1] - rng[0])*255
  a = np.uint8(np.clip(a, 0, 255))
  f = BytesIO()
  PIL.Image.fromarray(a).save(f, fmt)
  clear_output(wait = True)
  display(Image(data=f.getvalue()))
```

Then initialize Tensorflow session.

```python
sess = tf.InteractiveSession()
```

Following two functions are necessary to perform operation called convolution. Convolutional kernel `k` slides the image `x` and produces new matrix `y`. Every `y` value is a sum of multiplication of area (covered by `k`) on `x` and kernel `k`.

[image source](https://towardsdatascience.com/intuitively-understanding-convolutions-for-deep-learning-1f6f42faee1)

NOTE: In our case `y` will be the same shape as `x`, because padding set to `SAME`.

```python
def make_kernel(a):
  """Transform a 2D array into a convolution kernel"""
  a = np.asarray(a)
  a = a.reshape(list(a.shape) + [1,1])
  return tf.constant(a, dtype=1)

def simple_conv(x, k):
  """A simplified 2D convolution operation"""
  x = tf.expand_dims(tf.expand_dims(x, 0), -1)
  y = tf.nn.depthwise_conv2d(x, k, [1, 1, 1, 1], padding='SAME')
  return y[0, :, :, 0]
```

Discrete Laplace operator described below.

```python
def laplace(x):
  """Compute the 2D laplacian of an array"""
  laplace_k = make_kernel([[0.25, 0.5, 0.25],
                           [0.5, -3., 0.5],
                           [0.25, 0.5, 0.25]])
  return simple_conv(x, laplace_k)
```

Initialize pond `u_init` (500x500x1) with zeros. Also initialize matrix that will store first derivative `ut_unit`.

```python
N = 500

# Initial Conditions -- some rain drops hit a pond

# Set everything to zero
u_init = np.zeros([N, N], dtype=np.float32)
ut_init = np.zeros([N, N], dtype=np.float32)

# Some rain drops hit a pond at random points
for n in range(40):
  a,b = np.random.randint(0, N, 2)
  u_init[a,b] = np.random.uniform()

DisplayArray(u_init, rng=[-0.1, 0.1])
```

Define TensorFlow input, variables and finally two operations. Then group them and run session. All details below.

```python
# Parameters:
# eps -- time resolution
# damping -- wave damping
# c -- wave speed # Not a part of original code version
eps = tf.placeholder(tf.float32, shape=())
damping = tf.placeholder(tf.float32, shape=())
c = tf.placeholder(tf.float32, shape=())

# Create variables for simulation state
U  = tf.Variable(u_init)
Ut = tf.Variable(ut_init)

# Discretized PDE update rules
U_ = U + eps * Ut
Ut_ = Ut + eps * ((c ** 2) * laplace(U) - damping * Ut)

# Operation to update the state
step = tf.group(
  U.assign(U_),
  Ut.assign(Ut_))

# Initialize state to initial conditions
tf.global_variables_initializer().run()
```

Finally run 1000 steps and display pond after every step.

```python
# Run 1000 steps of PDE
for i in range(1000):
  # Step simulation
  step.run({eps: 0.03, damping: 0.04, c: 3.0})
  DisplayArray(U.eval(), rng=[-0.1, 0.1])

sess.close()
```

### Now feel free to run the code. Mathematical explanation is coming.

## Mathematics

Let's try formulate this code from mathematical point of view.

This example is a numerical solution ([Euler's method](https://en.wikipedia.org/wiki/Numerical_methods_for_ordinary_differential_equations#Euler_method)) of [wave equation](https://en.wikipedia.org/wiki/Wave_equation) (second order Partial Differencian Equation - PDE). It's quite hard to model PDEs (and computationary expensive?), that's why PDE was transformed into ODE (oridinary differential equation), which only has one differentiable variable - time. And can be easily computed with Euler's method.

Transformation of PDE to ODE is done with Laplacian transformation. In particular, [discrete Laplacian operator](https://en.wikipedia.org/wiki/Discrete_Laplace_operator) (very rough approximation of second derivative) was used. More precisely isotropic discrete Laplacian operator. Isotropic discrete Laplacian operator is just a 3x3 matrix (convolutional kernel) in our case.

### Wave equation with Euler's method

In general, this problem has analytical solutions, e.g. if we restrict our waves to be monochomatic (single frequency), cylindrical waves in cylindrical coordinate system ([example](http://farside.ph.utexas.edu/teaching/315/Waves/node47.html)). So instead of brute force approximation of second derivatives, we can compute bunch of sin/cos functions and then project cylindrical coordinates to euclidean. But this at least requires careful analytical equations preparation, which can take a time. So let's look on numerical solution instead.

We want to modify a wave function $u(x,y)$, so it's values can be computed using Euler's method. For Euler we need first order differential equations, but wave function is a second order. This means we need to apply Euler's method twice.

Initial [wave equation](https://en.wikipedia.org/wiki/Wave_equation#Introduction) looks like

\begin{equation}
 \frac{{\partial ^2 u}}{{\partial t^2 }} = {{c^2 \Delta u}}
\end{equation}

Which is a second direvative $u''$. Let's write discrete equations for $u$, where $\Delta t$ is a time difference

\begin{equation}
 u_0(x,y) = C\\
 u_{n+1}(x,y) = u_n(x,y) + u'_n(x,y) \Delta t\\
 u'_{n+1}(x,y) = u'_n(x,y) + u''_n(x,y) \Delta t\\
\end{equation}

We know initial conditions $C$ and we also can replace $u''$ with right part of wave equation.

\begin{equation}
 u'_{n+1}(x,y) = u'_n(x,y) + c^2 \Delta u \Delta t\\
\end{equation}

Let's also add reasoning behind wave dampling. Assuming this is a friction force, which for small speeds proportional to speed.

\begin{equation}
 u'_{n+1}(x,y) = u'_n(x,y) + (c^2 \Delta u - k_{dampling} u'_n(x,y)) \Delta t\\
\end{equation}

Final form that is used in code is

\begin{equation}
 u_0(x,y) = C\\
 u_{n+1}(x,y) = u_n(x,y) + u'_n(x,y) \Delta t\\
 u'_{n+1}(x,y) = u'_n(x,y) + (c^2 \Delta u - k_{dampling} u'_n(x,y)) \Delta t\\
\end{equation}

### Discrete Laplacian operator

Laplacian can be a tough thing to understand actually. But let me try to explain in simple words. For detailed explanation look [here](https://ocw.mit.edu/courses/mathematics/18-03-differential-equations-spring-2010/video-lectures/lecture-19-introduction-to-the-laplace-transform/) (MIT lecture), for intuitions look [here](https://dsp.stackexchange.com/questions/11008/intuitive-interpretation-of-laplace-transform).

Small context about [Fourier transform](https://en.wikipedia.org/wiki/Fourier_transform) (easier intuition).

In our case we have some signal (wave), it was proven that **every signal can be broken into parts**. In case, of Fourier transform signal can be described as a **sum of harmonics (sinusoids)**. In other words, Fourier transform breaks **ANY** complex signal into bunch of simple sinusoids. Fourier transform is a special case of Laplace transform.

In case of Laplace transform signal can be described by it’s "moments" (i.e. mean and variance, first and second derivative respectively). Put simpler (for our case), we can extract enough information from neighbours spatial points to estimate our next value (i.e. move from spatial x,y variables to time variable only). Just to remember, Laplacian transformation allows partial differential equation (PDE) to become ordinary differential equation (ODE), which is much simpler to solve.

For numerical solutions discrete Laplace operator usually used with convolution. Most frequent use case is image edge detection. Basically we slide the image with special filter (kernel). Kernel is just a 3x3 matrix that is an rough approximation of analytical Laplace transform. Multiplying neighbour pixels values with kernel and then taking a sum of all values, gives us second derivative (i.e. speed of intensity/color change) that is then used in wave equation.

Plots of non-isotropic (left) and isotropic Laplacian kernels.

[source wikipedia](https://de.wikipedia.org/wiki/Laplace-Filter)

Small note, author of Tensorflow example used kernel matrix that for some reason was multiplied by 2 (may be i'm missing smth?). According to [wikipedia](https://en.wikipedia.org/wiki/Discrete_Laplace_operator#Implementation_via_operator_discretization), isotropic Laplace convolution kernel is

\begin{bmatrix}
    0.25 & 0.5 & 0.25 \\
    0.5 & -3 & 0.5 \\
    0.25 & 0.5 & 0.25
\end{bmatrix}

One more time, **Laplacian kernel can be used to get rough approximation of the second partial derivative** of the function.

## Physical parameters of waves

Just a few notes about physical parameters of the waves (wave length, frequency and amplitude). Because of numerical approximation limiations it seems impossible to get a single harmonic wave.

For example, single harmonic cylindrical wave should have constant frequency (constant distance between wave crests), but as experiments show distance is changing. I assume this is because that initial wave contains multiple harmonics, so eventually they are splitting up.

## Conclusion

In this blog post, i've described results of my small research regarding math background of Tensorflow example [Partial Differential Equations](https://www.tensorflow.org/tutorials/non-ml/pdes). As a result, i've added speed parameter and corrected Laplace kernel for isotropic case.


# Combining measurements

When we do a fit, we can have additional knowledge about a parameter from other measurements. This can be taken into account either through a simultaneous fit or by adding a constraint (subsidiary measurement).

## Adding a constraint

If we know a parameters value from a different measurement and want to constraint this using its uncertainty, a Gaussian constraint can be added to the likelihood as

\begin{align}
\mathrm{constr_i} = \mathrm{Gauss}(\mu_{measured}; \theta_i, \sigma_{measured})
\end{align}

In general, additional terms can be added to the likelihood arbitrarily in zfit, be it to incorporate other shaped measurements or to add penalty terms to confine a fit within boundaries.

```python
import zfit
import numpy as np
import tensorflow as tf
import matplotlib.pyplot as plt
import mplhep
import hepunits as u
import particle.literals as lp

plt.rcParams['figure.figsize'] = (8,6)
```

```python
mu_true =  lp.B_plus.mass * u.MeV
sigma_true = 50 * u.MeV

# number of signal and background
n_sig_rare = 120
n_bkg_rare = 700
```

```python
# create some data
signal_np = np.random.normal(loc=mu_true, scale=sigma_true, size=n_sig_rare)
bkg_np_raw = np.random.exponential(size=20000, scale=700)
bkg_np = bkg_np_raw[bkg_np_raw<1000][:n_bkg_rare] + 5000  # just cutting right, but zfit could also cut
```

```python
# Firstly, the observable and its range is defined
obs = zfit.Space('Bmass', (5000, 6000))  # for whole range
```

```python
# load data into zfit
data = zfit.Data.from_numpy(obs=obs, array=np.concatenate([signal_np, bkg_np], axis=0))
```

```python

# Parameters are specified:  (name (unique), initial, lower, upper) whereas lower, upper are optional
mu = zfit.Parameter('mu', 2790, 5100, 5400)
sigma = zfit.Parameter('sigma', 20, 1, 200)
signal = zfit.pdf.Gauss(mu=mu, sigma=sigma, obs=obs)

lam = zfit.Parameter('lambda', -0.002, -0.1, -0.00001, step_size=0.001)  # floating, also without limits
comb_bkg = zfit.pdf.Exponential(lam, obs=obs)

sig_yield = zfit.Parameter('sig_yield', n_sig_rare + 30,
                                step_size=3)  # step size: default is small, use appropriate
bkg_yield = zfit.Parameter('bkg_yield', n_bkg_rare - 40, step_size=1)
# Create extended PDFs
extended_sig = signal.create_extended(sig_yield)
extended_bkg = comb_bkg.create_extended(bkg_yield)

# The final model is the combination of the signal and backgrond PDF
model = zfit.pdf.SumPDF([extended_bkg, extended_sig])
```

```python
constraint = zfit.constraint.GaussianConstraint(mu, observation=5275 * u.MeV, uncertainty=15 * u.MeV)
```

```python
nll = zfit.loss.ExtendedUnbinnedNLL(model, data, constraints=constraint)
minimizer = zfit.minimize.Minuit(use_minuit_grad=True)
result = minimizer.minimize(nll)
result.hesse();
```

    WARNING:tensorflow:From /home/jonas/anaconda3/envs/zfit38/lib/python3.8/site-packages/tensorflow_probability/python/distributions/distribution.py:284: MultivariateNormalFullCovariance.__init__ (from tensorflow_probability.python.distributions.mvn_full_covariance) is deprecated and will be removed after 2019-12-01.
    Instructions for updating:
    `MultivariateNormalFullCovariance` is deprecated, use `MultivariateNormalTriL(loc=loc, scale_tril=tf.linalg.cholesky(covariance_matrix))` instead.
    WARNING:tensorflow:From /home/jonas/anaconda3/envs/zfit38/lib/python3.8/site-packages/tensorflow/python/ops/linalg/linear_operator_lower_triangular.py:151: calling LinearOperator.__init__ (from tensorflow.python.ops.linalg.linear_operator) with graph_parents is deprecated and will be removed in a future version.
    Instructions for updating:
    Do not pass `graph_parents`.  They will  no longer be used.
    ------------------------------------------------------------------
    | FCN = 865.7                   |     Ncalls=183 (183 total)     |
    | EDM = 5.28e-05 (Goal: 0.001)  |            up = 0.5            |
    ------------------------------------------------------------------
    |  Valid Min.   | Valid Param.  | Above EDM | Reached call limit |
    ------------------------------------------------------------------
    |     True      |     True      |   False   |       False        |
    ------------------------------------------------------------------
    | Hesse failed  |   Has cov.    | Accurate  | Pos. def. | Forced |
    ------------------------------------------------------------------
    |     False     |     True      |   True    |   True    | False  |
    ------------------------------------------------------------------

```python
print(result.params)
```

    name           value    minuit_hesse    at limit
    ---------  ---------  --------------  ----------
    bkg_yield      724.9     +/-      33       False
    sig_yield      95.13     +/-      22       False
    lambda     -0.001627     +/- 0.00014       False
    mu              5286     +/-     8.9       False
    sigma          38.01     +/-      11       False

## Simultaneous fits

Sometimes, we don't want to fit a single channel but multiple with the same likelihood and having shared parameters between them. In this example, we will fit the decay simultaneously to its resonant control channel.

A simultaneous likelihood is the product of different likelihoods defined by
\begin{align}
\mathcal{L}_{f(x)}(\theta | {data_0, data_1, ..., data_n}) &= \prod_{i=1}^{n} \mathcal{L}(\theta_i, data_i)
\end{align}

and becomes in the NLL a sum

\begin{align}
\mathrm{NLL}_{f(x)}(\theta | {data_0, data_1, ..., data_n}) &= \sum_{i=1}^{n} \mathrm{NLL}(\theta_i, data_i)
\end{align}

```python
n_sig_reso = 40000
n_bkg_reso = 3000
```

```python
# create some data
signal_np_reso = np.random.normal(loc=mu_true, scale=sigma_true * 0.7, size=n_sig_reso)
bkg_np_raw_reso = np.random.exponential(size=20000, scale=900)
bkg_np_reso = bkg_np_raw_reso[bkg_np_raw_reso<1000][:n_bkg_reso] + 5000

# load data into zfit
obs_reso = zfit.Space('Bmass_reso', (5000, 6000))
data_reso = zfit.Data.from_numpy(obs=obs_reso, array=np.concatenate([signal_np_reso, bkg_np_reso], axis=0))
```

### Sharing and composing parameters

As an example, the same signal shape will be used with the identical `mu` yet a scaled `sigma`. This means that the `sigma` of our control mode corresponds to the `sigma` of our signal times a scaling parameter. Therefore, the `scaled_sigma` is a function of two other parameters, `sigma` and `sigma_scaling` and _cannot_ change it's value independently. There are two fundamentally distinct types to represent this behavior in zfit, independent (`Parameter`) and dependent parameters (`ComplexParameter`, `ComposedParameter`,...)

#### Independent parameter

An independent parameter has, as a distinctive method, a `set_value` that _changes_ the value of the parameter. In a fit, or in general, these are _the only object_ that can directly change their value and therefore do not depend on other objects while most other objects depend on Parameters.
As a consequence, this parameters can have limits (which effectively restrict the possible values a `Parameter` can be assigned to) and have a `step_size`, a hint to any minimization algorithm about the order of magnitude that a change in the parameter will have on the loss.

Another attribute is a `floating` flag: if set to `False`, the parameter won't be floating in the fit, whether explicitly given or implicitly inferred from the dependencies.

#### Dependent parameter

These are single-valued functions effectively that depend on other objects, usually other parameters. Therefore, a dependent parameter does not have a `set_value` function and also does not posses limits. The latter is preferred to be set with the `Parameter` it depends on, however, if a hard limit is required, this can always be enforced in the definition of a `ComposedParameter`.

The most notable parameter is the `ComposedParameter`, which returns an arbitrary function of its input arguments, the latter which can be specified with the `params` argument.

While this parameters *cannot* change there value explicitly and therefore won't be used by a minimizer, the zfit minimizers automatically extract the independent parameters that a dependent parameter depends on (if this is given as an argument.)

As a consequence, these parameters also miss a `step_size` attribute. Furthermore, `floating` can't be used, neither set nor retrieved; it is rather advised to check directly with its dependencies.

#### Sharing parameters

Since in zfit, every parameter object is unique, also defined by its name, it is straightforward to know when a parameter is shared in the loss and when it is not: if the same object is used in two places, it is shared. This can be arbitrarily mixed.

```python
# Firstly, we create a free scaling parameter
sigma_scaling = zfit.Parameter('sigma_scaling', 0.9, 0.1, 10, step_size=0.1)

def sigma_scaled_fn(sigma, sigma_scaling):
    return sigma * sigma_scaling  # this can be an arbitrary function

sigma_scaled = zfit.ComposedParameter('sigma scaled',  # name
                                      sigma_scaled_fn,  # function
                                      params=[sigma, sigma_scaling]  # the objects used inside the function
                                      )
```

```python
signal_reso = zfit.pdf.Gauss(mu=mu,  # the same as for the rare mode
                             sigma=sigma_scaled,
                             obs=obs_reso
                            )

lambda_reso = zfit.Parameter('lambda_reso', -0.002, -0.01, 0.0001)  # floating
comb_bkg_reso_pdf = zfit.pdf.Exponential(lambda_reso, obs=obs_reso)

reso_sig_yield = zfit.Parameter('reso_sig_yield', n_sig_reso - 100, 0, n_sig_reso * 3,
                                step_size=1)  # step size: default is small, use appropriate
reso_bkg_yield = zfit.Parameter('reso_bkg_yield', n_bkg_reso + 70, 0, 2e5, step_size=1)

# Create the extended models
extended_sig_reso = signal_reso.create_extended(reso_sig_yield)
extended_bkg_reso = comb_bkg_reso_pdf.create_extended(reso_bkg_yield)
model_reso = zfit.pdf.SumPDF([extended_bkg_reso, extended_sig_reso])
```

To implement the simultaneous fit, there are two ways to achieve this in zfit. As an important distinction to other frameworks, zfit translates the above equation
\begin{align}
\mathrm{NLL}_{f(x)}(\theta | {data_0, data_1, ..., data_n}) &= \sum_{i=1}^{n} \mathrm{NLL}(\theta_i, data_i)
\end{align}

directly into code.

We can build two losses and add them directly or give a list of models and data, which build a loss each one and add up.

```python
nll_rare = zfit.loss.ExtendedUnbinnedNLL(model, data)
nll_reso = zfit.loss.ExtendedUnbinnedNLL(model_reso, data_reso)
nll_simultaneous = nll_rare + nll_reso
```

    /home/jonas/Documents/physics/software/zfit_project/zfit_repo/zfit/core/loss.py:99: UserWarning: The fit_range argument is depreceated and will maybe removed in future releases. It is preferred to define the range in the space when creating the data and the model.
      warnings.warn("The fit_range argument is depreceated and will maybe removed in future releases. "

```python
result_simultaneous = minimizer.minimize(nll_simultaneous)
```

    ------------------------------------------------------------------
    | FCN = -1.879e+05              |     Ncalls=204 (204 total)     |
    | EDM = 0.0103 (Goal: 0.001)    |            up = 0.5            |
    ------------------------------------------------------------------
    |  Valid Min.   | Valid Param.  | Above EDM | Reached call limit |
    ------------------------------------------------------------------
    |     False     |     True      |   True    |       False        |
    ------------------------------------------------------------------
    | Hesse failed  |   Has cov.    | Accurate  | Pos. def. | Forced |
    ------------------------------------------------------------------
    |     False     |     True      |   True    |   True    | False  |
    ------------------------------------------------------------------

```python
result_simultaneous.hesse()
```

    OrderedDict([(<zfit.Parameter 'bkg_yield' floating=True value=724.3>,
                  {'error': 30.057057240887225}),
                 (<zfit.Parameter 'sig_yield' floating=True value=96.1>,
                  {'error': 16.733265730998557}),
                 (<zfit.Parameter 'lambda' floating=True value=-0.001617>,
                  {'error': 0.00013918478595347154}),
                 (<zfit.Parameter 'mu' floating=True value=5279>,
                  {'error': 0.18154673209572916}),
                 (<zfit.Parameter 'sigma' floating=True value=38.49>,
                  {'error': 2.134483909511519}),
                 (<zfit.Parameter 'reso_bkg_yield' floating=True value=3016>,
                  {'error': 65.59422345128854}),
                 (<zfit.Parameter 'reso_sig_yield' floating=True value=3.999e+04>,
                  {'error': 203.16808935583364}),
                 (<zfit.Parameter 'lambda_reso' floating=True value=-0.001063>,
                  {'error': 6.862439953704262e-05}),
                 (<zfit.Parameter 'sigma_scaling' floating=True value=0.9111>,
                  {'error': 0.05054266401515093})])

```python
print(result_simultaneous.params)
```

    name                value    minuit_hesse    at limit
    --------------  ---------  --------------  ----------
    bkg_yield           724.3     +/-      30       False
    sig_yield            96.1     +/-      17       False
    lambda          -0.001617     +/- 0.00014       False
    mu                   5279     +/-    0.18       False
    sigma               38.49     +/-     2.1       False
    reso_bkg_yield       3016     +/-      66       False
    reso_sig_yield      39990     +/-   2e+02       False
    lambda_reso     -0.001063     +/- 6.9e-05       False
    sigma_scaling      0.9111     +/-   0.051       False

### Plotting a simultaneous loss

Since the definition of a simultaneous fit is as above, it is simple to plot each component separately: either my using the attributes of the loss to access the models and plot in a general fashion or directly reuse the model and data from before; we created them manually before.

```python
# Sets the values of the parameters to the result of the simultaneous fit
# in case they were modified.
zfit.param.set_values(nll_simultaneous.get_params(), result_simultaneous)

def plot_fit_projection(model, data, nbins=30, ax=None):
    # The function will be reused.
    if ax is None:
        ax = plt.gca()

    lower, upper = data.data_range.limit1d

    # Creates and histogram of the data and plots it with mplhep.
    counts, bin_edges = np.histogram(data.unstack_x(), bins=nbins)
    mplhep.histplot(counts, bins=bin_edges, histtype="errorbar", yerr=True,
                    label="Data", ax=ax, color="black")

    binwidth = np.diff(bin_edges)[0]
    x = tf.linspace(lower, upper, num=1000)  # or np.linspace

    # Line plots of the total pdf and the sub-pdfs.
    y = model.ext_pdf(x) * binwidth
    ax.plot(x, y, label="total", color="royalblue")
    for m, l, c in zip(model.get_models(), ["background", "signal"], ["forestgreen", "crimson"]):
        ym = m.ext_pdf(x) * binwidth
        ax.plot(x, ym, label=l, color=c)

    ax.set_title(data.data_range.obs[0])
    ax.set_xlim(lower, upper)
    ax.legend(fontsize=15)

    return ax

fig, axs = plt.subplots(1, 2, figsize=(16, 6))

for mod, dat, ax, nb in zip(nll_simultaneous.model, nll_simultaneous.data, axs, [30, 60]):
    plot_fit_projection(mod, dat, nbins=nb, ax=ax)
```

## Discovery test

We observed an excess of our signal:

```python
print(result_simultaneous.params[sig_yield])
```

    {'value': 96.10030901659354, 'minuit_hesse': {'error': 16.733265730998557}}

Now we would like to compute the significance of this observation or in other words the probabilty that this observation is the result of the statistical fluctuation. To do so we have to perform an hypothesis test where the null and alternative hypotheses are defined as:

* $H_{0}$, the null or background only hypothesis, i.e. $N_{sig} = 0$;
* $H_{1}$, the alternative hypothesis, i.e $N_{sig} = \hat{N}_{sig}$, where $\hat{N}_{sig}$ is the fitted value of $N_{sig}$ printed above.

In `hepstats` to formulate a hypothesis you have to use the `POI` (Parameter Of Interest) class.

```python
from hepstats.hypotests.parameters import POI

# the null hypothesis
sig_yield_poi = POI(sig_yield, 0)
```

What the `POI` class does is to take as input a `zfit.Parameter` instance and a value corresponding to a given hypothesis. You can notice that we didn't define here the alternative hypothesis as in the discovery test the value of POI for alternate is set to the best fit value.

The test statistic used is the profile likelihood ratio and defined as:

\begin{equation}
q_{0} = \left\{
    \begin{array}{ll}
        -2 \ln \frac{\mathcal{L}(N_{sig}=0, \; \hat{\hat{\theta}})}{\mathcal{L}(N_{sig}=\hat{N}_{sig}, \; \hat{\theta})} & \mbox{if } \; \hat{N}_{sig} \geq 0 \\
        0 & \mbox{if } \; \hat{N}_{sig} < 0
    \end{array}
\right.
\end{equation}

where $\hat{\theta}$ are the best fitted values of the nuisances parameters (i.e. background yield, exponential slope...), while $\hat{\hat{\theta}}$ are the fitted values of the nuisances when ${N}_{sig} = 0$.

From the test statistic distribution a p-value can computed as

\begin{equation}
p_{0} = \int_{q_{0}^{obs}}^{\infty} f(q_{0} |H_{0}) dq_{0}
\end{equation}

where $q_{0}^{obs}$ is the value of the test statistic evaluated on observed data.

The construction of the test statistic and the computation of the p-value is done in a `Calculator` object in `hepstats`. In this example we will use in this example the `AsymptoticCalculator` calculator which assumes that $q_{0}$ follows a $\chi^2(ndof=1)$ which simplifies the p-value computation to

\begin{equation}
p_{0} = 1 - \Phi\bigg({\sqrt{q_{0}^{obs}}}\bigg).
\end{equation}

The calculator objects takes as input the likelihood function and a minimizer to profile the likelihood.

```python
from hepstats.hypotests.calculators import AsymptoticCalculator

# construction of the calculator instance
calculator = AsymptoticCalculator(input=nll_simultaneous, minimizer=minimizer)
calculator.bestfit = result_simultaneous

# equivalent to above
calculator = AsymptoticCalculator(input=result_simultaneous, minimizer=minimizer)
```

There is another calculator in `hepstats` called `FrequentistCalculator` which constructs the test statistic distribution $f(q_{0} |H_{0})$ with pseudo-experiments (toys), but it takes more time.

The `Discovery` class is a high-level class that takes as input a calculator and a `POI` instance representing the null hypothesis, it basically asks the calculator to compute the p-value and also computes the signifance as

\begin{equation}
Z = \Phi^{-1}(1 - p_0).
\end{equation}

```python
from hepstats.hypotests import Discovery

discovery = Discovery(calculator=calculator, poinull=sig_yield_poi)
discovery.result()
```

    p_value for the Null hypothesis = 1.619449019329977e-11
    Significance (in units of sigma) = 6.635278774759537

    (1.619449019329977e-11, 6.635278774759537)

So we get a significance of about $7\sigma$ which is well above the $5 \sigma$ threshold for discoveries 😃.

## Upper limit calculation

Let's try to compute the discovery significance with a lower number of generated signal events.

```python
# Sets the values of the parameters to the result of the simultaneous fit
zfit.param.set_values(nll_simultaneous.get_params(), result_simultaneous)
sigma_scaling.floating=False

# Creates a sampler that will draw events from the model
sampler = model.create_sampler()

# Creates new simultaneous loss
nll_simultaneous_low_sig = zfit.loss.ExtendedUnbinnedNLL(model, sampler) + nll_reso
```

```python
# Samples with sig_yield = 10. Since the model is extended the number of
# signal generated is drawn from a poisson distribution with lambda = 10.
sampler.resample({sig_yield: 10})
```

```python
calculator_low_sig = AsymptoticCalculator(input=nll_simultaneous_low_sig, minimizer=minimizer)

discovery_low_sig = Discovery(calculator=calculator_low_sig, poinull=sig_yield_poi)
discovery_low_sig.result()
print(f"\n {calculator_low_sig.bestfit.params} \n")
```

We might consider computing an upper limit on the signal yield instead. The test statistic for an upper limit calculation is

\begin{equation}
q_{N_{sig}} = \left\{
    \begin{array}{ll}
        -2 \ln \frac{\mathcal{L}(N_{sig}, \; \hat{\hat{\theta}})}{\mathcal{L}(N_{sig}=\hat{N}_{sig}, \; \hat{\theta})} & \mbox{if } \; \hat{N}_{sig} \leq N_{sig} \\
        0 & \mbox{if } \; \hat{N}_{sig} > N_{sig}.
    \end{array}
\right.
\end{equation}

and the p-value is

\begin{equation}
p_{N_{sig}} = \int_{q_{N_{sig}}^{obs}}^{\infty} f(q_{N_{sig}} |N_{sig}) dq_{N_{sig}}.
\end{equation}

The upper limit on $N_{sig}$, $N_{sig, \uparrow}$ is found for $p_{N_{sig, \uparrow}} = 1 - \alpha$, $\alpha$ being the confidence level (typically $95 \%$). The upper limit is found by interpolation of the p-values as a function of $N_{sig}$, which is done the `UpperLimit` class. We have to give the range of values of $N_{sig}$ to scan using the `POIarray` class which as the `POI` class takes as input the parameter but takes several values to evaluate the parameter instead of one.

```python
from hepstats.hypotests import UpperLimit
from hepstats.hypotests.parameters import POIarray

# Background only hypothesis.
bkg_only = POI(sig_yield, 0)
# Range of Nsig values to scan.
sig_yield_scan = POIarray(sig_yield, np.linspace(0, 70, 10))

ul = UpperLimit(calculator=calculator_low_sig, poinull=sig_yield_scan, poialt=bkg_only)
ul.upperlimit(alpha=0.05);
```

```python
from utils import plotlimit

plotlimit(ul, CLs=False)
```

## Splot

This is now an demonstration of the **sPlot** algorithm, described in [Pivk:2004ty](https://arxiv.org/pdf/physics/0402083.pdf).

If a data sample is populated by different sources of events, like signal and background, **sPlot** is able to unfold the contributions of the different sources for a given variable.

Let's construct a dataset with two variables, the invariant mass and lifetime, for the resonant signal defined above and the combinatorial background.

```python
# Signal distributions.
nsig_sw = 20000
np_sig_m_sw = signal_reso.sample(nsig_sw).numpy().reshape(-1,)
np_sig_t_sw = np.random.exponential(size=nsig_sw, scale=1)

# Background distributions.
nbkg_sw = 150000
np_bkg_m_sw = comb_bkg_reso_pdf.sample(nbkg_sw).numpy().reshape(-1,)
np_bkg_t_sw = np.random.normal(size=nbkg_sw, loc=2.0, scale=2.5)

# Lifetime cut.
t_cut = np_bkg_t_sw > 0
np_bkg_t_sw = np_bkg_t_sw[t_cut]
np_bkg_m_sw = np_bkg_m_sw[t_cut]

# Mass distribution
np_m_sw = np.concatenate([np_sig_m_sw, np_bkg_m_sw])

# Lifetime distribution
np_t_sw = np.concatenate([np_sig_t_sw, np_bkg_t_sw])

# Plots the mass and lifetime distribution.
fig, axs = plt.subplots(1, 2, figsize=(16, 6))
axs[0].hist([np_bkg_m_sw, np_sig_m_sw], bins=50, stacked=True, label=("background", "signal"), alpha=.7)
axs[0].set_xlabel("m")
axs[0].legend(fontsize=15)
axs[1].hist([np_bkg_t_sw, np_sig_t_sw], bins=50, stacked=True, label=("background", "signal"), alpha=.7)
axs[1].set_xlabel("t")
axs[1].legend(fontsize=15);
```

In this particular example we want to unfold the signal lifetime distribution. To do so **sPlot** needs a discriminant variable to determine the yields of the various sources using an <ins>extended</ins> maximum likelihood fit.

```python
# Builds the loss.
data_sw = zfit.Data.from_numpy(obs=obs_reso, array=np_m_sw)
nll_sw = zfit.loss.ExtendedUnbinnedNLL(model_reso, data_sw)

# This parameter was useful in the simultaneous fit but not anymore so we fix it.
sigma_scaling.floating=False

# Minimizes the loss.
result_sw = minimizer.minimize(nll_sw)
print(result_sw.params)
```

```python
# Visualization of the result.
zfit.param.set_values(nll_sw.get_params(), result_sw)
plot_fit_projection(model_reso, data_sw, nbins=100)
```

**sPlot** will use the fitted yield for each sources to derive the so-called **sWeights** for each data point:

\begin{equation}
W_{n}(x) = \frac{\sum_{j=1}^{N_S} V_{nj} f_j(x)}{\sum_{k=1}^{N_S} N_{k}f_k(x)}
\end{equation}

with

\begin{equation}
V_{nj}^{-1} = \sum_{e=1}^{N} \frac{f_n(x_e) f_j(x_e)}{(\sum_{k=1}^{N_S} N_{k}f_k(x))^2}
\end{equation}

where ${N_S}$ is the number of sources in the data sample, here 2. The index $n$ represents the source, for instance $0$ is the signal and $1$ is the background, then $f_0$ and $N_0$ are the pdf and yield for the signal.

In `hepstats` the **sWeights** are computed with the `compute_sweights` function which takes as arguments the <ins>fitted</ins> extended model and the discrimant data (on which the fit was performed).

```python
from hepstats.splot import compute_sweights

weights = compute_sweights(model_reso, data_sw)

print(weights)
```

```python
print("Sum of signal sWeights: ", np.sum(weights[reso_sig_yield]))
print("Sum of background sWeights: ", np.sum(weights[reso_bkg_yield]))
```

Now we can apply the signal **sWeights** on the lifetime distribution and retrieve its signal components.

```python
fig, axs = plt.subplots(1, 2, figsize=(16, 6))
nbins = 40

sorter = np_m_sw.argsort()

axs[0].plot(np_m_sw[sorter], weights[reso_sig_yield][sorter], label="$w_\\mathrm{sig}$")
axs[0].plot(np_m_sw[sorter], weights[reso_bkg_yield][sorter], label="$w_\\mathrm{bkg}$")
axs[0].plot(np_m_sw[sorter], weights[reso_sig_yield][sorter] + weights[reso_bkg_yield][sorter],
            "-k", label="$w_\\mathrm{sig} + w_\\mathrm{bkg}$")
axs[0].axhline(0, color="0.5")
axs[0].legend(fontsize=15)
axs[0].set_xlim(5000, 5600)

axs[1].hist(np_t_sw, bins=nbins, range=(0, 6), weights=weights[reso_sig_yield], label="weighted histogram", alpha=.5)
axs[1].hist(np_sig_t_sw, bins=nbins, range=(0, 6), histtype="step", label="true histogram", lw=1.5)
axs[1].set_xlabel("t")
axs[1].legend(fontsize=15);
```

Be careful the **sPlot** technique works only on variables that are uncorrelated with the discriminant variable.

```python
print(f"Correlation between m and t: {np.corrcoef(np_m_sw, np_t_sw)[0, 1]}")
```

Let's apply to signal **sWeights** on the mass distribution to see how bad the results of **sPlot** is when applied on a variable that is correlated with the discrimant variable.

```python
plt.hist(np_m_sw, bins=100, range=(5000, 6000), weights=weights[reso_sig_yield]);
```

```python

```


```python
#/usr/bin/python
# -*- coding: utf-8 -*-
import sys
import h5py
import numpy as np
import functools
from mpl_toolkits.mplot3d import Axes3D
from matplotlib import pyplot as plt
from sympy.physics.hydrogen import R_nl
from sympy import var
import gc

f = h5py.File("static_results.h5",'r')
dset1 = f.get("/numres")
dset2 = f.get("/params1d")
dset3=f.get("/levels")
p = np.array(dset2)
a = np.array(dset1)
ind=np.array(dset3)
print(len(ind))
f.close()
dset1=0;
dset2=0;
dset3=0
nx=int(p[4])
ne = int(p[6])

x=np.linspace(0,p[0],nx-1)
dx=p[0]/nx
fig=plt.figure()
ax=fig.add_subplot(111)
# a[int(ind[i])*nx:int(ind[i])*nx+nx] = normalize(a[int(ind[i])*nx:int(ind[i])*nx+nx],p[0],nx)
for i in range(0,len(ind)):
    ax.plot(x,a[int(ind[i])*nx+1:int(ind[i])*nx+nx],label="$E_{"+str(i)+"}$")
#    ax.plot(x,a[i*nx+1:(i+1)*nx])
ax.set_xlabel("$x$ $(a_0)$",size=20)
ax.set_ylabel("$|\Psi_n|^2$ (not normed)",size=20)
#plt.legend()
i=0
plt.legend()
plt.show()
x=0
a=0
p=0
```

```python
8*1e7*1e3/1024**3
```

    74.50580596923828

```python

```


# 14 Linear Algebra: String Problem – Students (1)
## Motivating problem: Two masses on three strings
Two masses $M_1$ and $M_2$ are hung from a horizontal rod with length $L$ in such a way that a rope of length $L_1$ connects the left end of the rod to $M_1$, a rope of length $L_2$ connects $M_1$ and $M_2$, and a rope of length $L_3$ connects $M_2$ to the right end of the rod. The system is at rest (in equilibrium under gravity).

Find the angles that the ropes make with the rod and the tension forces in the ropes.

## Theoretical background
Treat $\sin\theta_i$ and $\cos\theta_i$ together with $T_i$, $1\leq i \leq 3$, as unknowns that have to simultaneously fulfill the nine equations
\begin{align}
-T_1 \cos\theta_1 + T_2\cos\theta_2 &= 0\\
 T_1 \sin\theta_1 - T_2\sin\theta_2 - W_1 &= 0\\
                   -T_2\cos\theta_2 + T_3\cos\theta_3 &= 0\\
                    T_2\sin\theta_2 + T_3\sin\theta_3 - W_2 &= 0\\
 L_1\cos\theta_1 + L_2\cos\theta_2 + L_3\cos\theta_3 - L &= 0\\
-L_1\sin\theta_1 - L_2\sin\theta_2 + L_3\sin\theta_3  &= 0\\
\sin^2\theta_1 + \cos^2\theta_1 - 1 &= 0\\
\sin^2\theta_2 + \cos^2\theta_2 - 1 &= 0\\
\sin^2\theta_3 + \cos^2\theta_3 - 1 &= 0
\end{align}

Consider the nine equations a vector function $\mathbf{f}$ that takes a 9-vector $\mathbf{x}$ of the unknowns as argument:
\begin{align}
\mathbf{f}(\mathbf{x}) &= 0\\
\mathbf{x} &= \begin{pmatrix}
x_0 \\ x_1 \\ x_2 \\
x_3 \\ x_4 \\ x_5 \\
x_6 \\ x_7 \\ x_8
\end{pmatrix} =
\begin{pmatrix}
\sin\theta_1 \\ \sin\theta_2 \\ \sin\theta_3 \\
\cos\theta_1 \\ \cos\theta_2 \\ \cos\theta_3 \\
T_1 \\ T_2 \\ T_3
\end{pmatrix} \\
\mathbf{L} &= \left(\begin{array}{c}
L \\ L_1 \\ L_2 \\ L_3
\end{array}\right), \quad
\mathbf{W} = \left(\begin{array}{c}
W_1 \\ W_2
\end{array}\right)
\end{align}

Using the unknowns from above, our system of 9 coupled equations is:

\begin{align}
-x_6 x_3 + x_7 x_4 &= 0\\
 x_6 x_0 - x_7 x_1 - W_1 &= 0\\
-x_7x_4 + x_8 x_5 &= 0\\
 x_7x_1 + x_8 x_2 - W_2 &= 0\\
 L_1x_3 + L_2 x_4 + L_3 x_5 - L &= 0\\
-L_1x_0 - L_2 x_1 + L_3 x_2  &= 0\\
x_{0}^{2} + x_{3}^{2} - 1 &= 0\\
x_{1}^{2} + x_{4}^{2} - 1 &= 0\\
x_{2}^{2} + x_{5}^{2} - 1 &= 0
\end{align}

Solve the root-finding problem $\mathbf{f}(\mathbf{x}) = 0$ with the **generalized Newton-Raphson** algorithm:

$$
\mathsf{J}(\mathbf{x}) \Delta\mathbf{x} = -\mathbf{f}(\mathbf{x})
$$

and

$$
\mathbf{x} \leftarrow \mathbf{x} + \Delta\mathbf{x}.
$$

## Problem setup
Set the problem parameters and the objective function $\mathbf{f}(\mathbf{x})$

```python
import numpy as np

# problem parameters
W = np.array([10, 20])
L = np.array([8, 3, 4, 4])

def f_2masses(x, L, W):
    return np.array([
            -x[6]*x[3] + x[7]*x[4],
             x[6]*x[0] - x[7]*x[1] - W[0],
            -x[7]*x[4] + x[8]*x[5],
            # Eq 4
            # Eq 5
            -L[1]*x[0] - L[2]*x[1] + L[3]*x[2],
            x[0]**2 + x[3]**2 - 1,
            # Eq 8
            x[2]**2 + x[5]**2 - 1,
        ])

def fLW(x, L=L, W=W):
    return f_2masses(x, L, W)
```

### Initial values
Guess some initial values (they don't have to fullfil the equations!):

```python
# initial parameters x0
# ...
x0 =
```

Check that we can calculate $\mathbf{f}(\mathbf{x}_0)$:

```python
f_2masses(x0, L, W)
```

### Visualization
Plot the positions of the 2 masses and the 3 strings for any solution vector $\mathbf{x}$:

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

```python
def plot_2masses(x, L, W, **kwargs):
    """Plot 2 mass/3 string problem for parameter vector x and parameters L and W"""

    kwargs.setdefault('linestyle', '-')
    kwargs.setdefault('marker', 'o')
    kwargs.setdefault('linewidth', 1)

    ax = kwargs.pop('ax', None)
    if ax is None:
        ax = plt.subplot(111)

    r0 = np.array([0, 0])
    r1 = r0 + np.array([L[0], 0])
    rod = np.transpose([r0, r1])

    L1 = r0 + np.array([L[1]*x[3], -L[1]*x[0]])
    L2 = L1 + np.array([L[2]*x[4], -L[2]*x[1]])
    L3 = L2 + np.array([L[3]*x[5], L[3]*x[2]])
    strings = np.transpose([r0, L1, L2, L3])

    ax.plot(rod[0], rod[1], color="black", marker="d", linewidth=4)
    ax.plot(strings[0], strings[1], **kwargs)
    ax.set_aspect(1)
    return ax
```

What does the initial guess look like?

```python
plot_2masses(x0, L, W)
```

## Jacobian
Write a function `Jacobian(f, x, h=1e-5)` that computes the Jacobian matrix numerically (use the central difference algorithm).

\begin{align}
\mathbf{J} &= \frac{\partial \mathbf{f}}{\partial\mathbf{x}} \\
J_{ij} &= \frac{\partial f_i(x_1, \dots, x_j, \dots)}{\partial x_j} \\
       &\approx \frac{f_i(x_1, \dots, x_j + \frac{h}{2}, \dots) - f_i(x_1, \dots, x_j - \frac{h}{2}, \dots)}{h}
\end{align}

```python
def Jacobian(f, x, h=1e-5):
    """df_i/dx_j with central difference (fi(xj+h/2)-fi(xj-h/2))/h"""
    J = np.zeros((len(f(x)), len(x)), dtype=np.float64)

    raise NotImplementedError

    return J
```

Test `Jacobian()` on
$$
\mathbf{g}(\mathbf{x}) = \begin{pmatrix}
  x_0^2 - x_1 \\ x_0
  \end{pmatrix}
$$
with analytical result

$$
\mathsf{J} = \left[\frac{\partial g_i}{\partial x_j}\right] =
\begin{pmatrix}
  \frac{\partial g_0}{\partial x_0} & \frac{\partial g_0}{\partial x_1}\\
  \frac{\partial g_1}{\partial x_0} & \frac{\partial g_1}{\partial x_1}
\end{pmatrix}
=
\begin{pmatrix}
  2 x_0 & -1\\
  1 & 0
\end{pmatrix}
$$

Given a test vector $\mathbf{x}_\text{test} = (1, 0)$, what is the numerical answer for $\mathsf{J}(\mathbf{x}_\text{test})$?

$$
\mathsf{J}(\mathbf{x}_\text{test}) = \text{?}
$$

Test your `Jacobian()` function with $\mathbf{x}_\text{test}$ and check that you get the same answer:

```python
def g(x):
    return np.array([
           # Eq Test 1
           # Eq Test 2
        ])
x_test = # choose a simple test vector
# run function and print result
# compare to analytical result
```

Just testing that it also works for our starting vector:

```python
Jacobian(fLW, x0)
```

## n-D Newton-Raphson Root Finding
Write a function `newton_raphson(f, x, Nmax=100, tol=1e-8, h=1e-5)` to find a root for a vector function `f(x)=0`. (See also [13 Root-finding by trial-and-error](https://asu-compmethodsphysics-phy494.github.io/ASU-PHY494/2020/03/26/13_Root_finding/) and the _1D Newton-Raphson algorithm_ in [13-Root-finding.ipynb](https://github.com/ASU-CompMethodsPhysics-PHY494/PHY494-resources/blob/master/13_root_finding/13-Root-finding.ipynb).) As a convergence criterion we demand that the length of the vector `f(x)` (the norm --- see `np.linalg.norm`) be less than the tolerance.

```python
def newton_raphson(f, x, Nmax=100, tol=1e-8, h=1e-5):
    """n-D Newton-Raphson: solves f(x) = 0.

    Iterate until |f(x)| < tol or nmax steps.
    """
    x = x.copy()

    raise NotImplementedError

    else:
        print("Newton-Raphson: no root found after {0} iterations (eps={1}); "
            "best guess is {2} with error {3}".format(Nmax, tol, x, fx))
    return x
```

### Solve 2 masses/3 strings problem

#### Solution

```python
# solve the string problem
x =
```

Plot the starting configuration and the solution:

```python
plot_2masses(x0, L, W)
plot_2masses(x, L, W)
```

Pretty-print the solution (angles in degrees):

```python
def pretty_print(x):
    theta = np.rad2deg(np.arcsin(x[0:3]))
    tensions = x[6:]
    print("theta1 = {0[0]:.1f} \t theta2 = {0[1]:.1f} \t theta3 = {0[2]:.1f}".format(theta))
    print("T1     = {0[0]:.1f} \t T2     = {0[1]:.1f} \t T3     = {0[2]:.1f}".format(tensions))
```

```python
print("Starting values")
pretty_print(x0)
print()
print("Solution")
pretty_print(x)
```

#### Show intermediate steps
Create a new function `newton_raphson_intermediates()` based on `newtopn_raphson()` that returns *all* trial `x` values including the last one.

```python
def newton_raphson_intermediates(f, x, Nmax=100, tol=1e-8, h=1e-5):
    """n-D Newton-Raphson: solves f(x) = 0.

    Iterate until |f(x)| < tol or nmax steps.

    Returns all intermediates.
    """
    intermediates = []
    x = x.copy()

    raise NotImplementedError

    else:
        print("Newton-Raphson: no root found after {0} iterations (eps={1}); "
            "best guess is {2} with error {3}".format(Nmax, tol, x, fx))
    return np.array(intermediates)
```

Visualize the intermediate configurations:

```python
x_series = newton_raphson_intermediates(fLW, x0)
```

```python
ax = plt.subplot(111)
ax.set_prop_cycle("color", [plt.cm.viridis_r(i) for i in np.linspace(0, 1, len(x_series))])
for iteration, x in enumerate(x_series):
    plot_2masses(x, L, W, label=str(iteration), ax=ax)
ax.legend(loc="best");
```

It's convenient to turn the above plotting code into a function that we can reuse:

```python
def plot_series(x_series, L, W):
    """Plot all N masses/strings solution vectors in x_series (N, 9) array"""
    ax = plt.subplot(111)
    ax.set_prop_cycle("color", [plt.cm.viridis_r(i) for i in np.linspace(0, 1, len(x_series))])
    for iteration, x in enumerate(x_series):
        plot_2masses(x, L, W, label=str(iteration), ax=ax)
    ax.legend(loc="best")
    return ax
```

## Additional work
Try different masses, e.g. M1 = M2 = 10, or M1= 0 , M2 = 10.

### M1 = M2 = 10

### M1 = 0, M2 = 10

```python

```