#+PROPERTY: header-args :exports both :output-dir results :session xs :kernel python3
#+TITLE: Investigaton of Monte-Carlo Methods
#+AUTHOR: Valentin Boettcher

* Init
** Required Modules
#+NAME: e988e3f2-ad1f-49a3-ad60-bedba3863283
#+begin_src jupyter-python :exports both :tangle tangled/xs.py
  import numpy as np
  import matplotlib.pyplot as plt
  import monte_carlo
#+end_src

#+RESULTS: e988e3f2-ad1f-49a3-ad60-bedba3863283

** Utilities
#+NAME: 53548778-a4c1-461a-9b1f-0f401df12b08
#+BEGIN_SRC jupyter-python :exports both
%run ../utility.py
%load_ext autoreload
%aimport monte_carlo
%autoreload 1
#+END_SRC

#+RESULTS: 53548778-a4c1-461a-9b1f-0f401df12b08
: The autoreload extension is already loaded. To reload it, use:
:   %reload_ext autoreload

* Implementation
** Center of Mass Frame
#+NAME: 777a013b-6c20-44bd-b58b-6a7690c21c0e
#+BEGIN_SRC jupyter-python :exports both :results raw drawer :exports code :tangle tangled/xs.py
  """
  Implementation of the analytical cross section for q q_bar ->
  gamma gamma

  Author: Valentin Boettcher <hiro@protagon.space>
  """

  import numpy as np

  # NOTE: a more elegant solution would be a decorator
  def energy_factor(charge, esp):
      """
      Calculates the factor common to all other values in this module

      Arguments:
      esp -- center of momentum energy in GeV
      charge -- charge of the particle in units of the elementary charge
      """

      return charge ** 4 / (137.036 * esp) ** 2 / 6


  def diff_xs(θ, charge, esp):
      """
      Calculates the differential cross section as a function of the
      azimuth angle θ in units of 1/GeV².

      Here dΩ=sinθdθdφ

      Arguments:
      θ -- azimuth angle
      esp -- center of momentum energy in GeV
      charge -- charge of the particle in units of the elementary charge
      """

      f = energy_factor(charge, esp)
      return f * ((np.cos(θ) ** 2 + 1) / np.sin(θ) ** 2)


  def diff_xs_cosθ(cosθ, charge, esp):
      """
      Calculates the differential cross section as a function of the
      cosine of the azimuth angle θ in units of 1/GeV².

      Here dΩ=d(cosθ)dφ

      Arguments:
      cosθ -- cosine of the azimuth angle
      esp -- center of momentum energy in GeV
      charge -- charge of the particle in units of the elementary charge
      """

      f = energy_factor(charge, esp)
      return f * ((cosθ ** 2 + 1) / (1 - cosθ ** 2))


  def diff_xs_eta(η, charge, esp):
      """
      Calculates the differential cross section as a function of the
      pseudo rapidity of the photons in units of 1/GeV^2.

      This is actually the crossection dσ/(dφdη).

      Arguments:
      η -- pseudo rapidity
      esp -- center of momentum energy in GeV
      charge -- charge of the particle in units of the elementary charge
      """

      f = energy_factor(charge, esp)
      return f * (np.tanh(η) ** 2 + 1)


  def diff_xs_p_t(p_t, charge, esp):
      """
      Calculates the differential cross section as a function of the
      transverse momentum (p_t) of the photons in units of 1/GeV^2.

      This is actually the crossection dσ/(dφdp_t).

      Arguments:
      p_t -- transverse momentum in GeV
      esp -- center of momentum energy in GeV
      charge -- charge of the particle in units of the elementary charge
      """

      f = energy_factor(charge, esp)
      sqrt_fact = np.sqrt(1 - (2 * p_t / esp) ** 2)
      return f / p_t * (1 / sqrt_fact + sqrt_fact)


  def total_xs_eta(η, charge, esp):
      """
      Calculates the total cross section as a function of the pseudo
      rapidity of the photons in units of 1/GeV^2.  If the rapditiy is
      specified as a tuple, it is interpreted as an interval.  Otherwise
      the interval [-η, η] will be used.

      Arguments:
      η -- pseudo rapidity (tuple or number)
      esp -- center of momentum energy in GeV
      charge -- charge of the particle in units of the elementar charge
      """

      f = energy_factor(charge, esp)
      if not isinstance(η, tuple):
          η = (-η, η)

      if len(η) != 2:
          raise ValueError("Invalid η cut.")

      def F(x):
          return np.tanh(x) - 2 * x

      return 2 * np.pi * f * (F(η[0]) - F(η[1]))
#+END_SRC

#+RESULTS: 777a013b-6c20-44bd-b58b-6a7690c21c0e
* Calculations
First, set up the input parameters.
#+BEGIN_SRC jupyter-python :exports both :results raw drawer
η = 2.5
charge = 1/3
esp = 200  # GeV
#+END_SRC

#+RESULTS:

Set up the integration and plot intervals.
#+begin_src jupyter-python :exports both :results raw drawer
interval_η = [-η, η]
interval = η_to_θ([-η, η])
interval_cosθ = np.cos(interval)
interval_pt = np.sort(η_to_pt([0, η], esp/2))
plot_interval = [0.1, np.pi-.1]
#+end_src

#+RESULTS:

#+begin_note
Note that we could utilize the symetry of the integrand throughout,
but that doen't reduce variance and would complicate things now.
#+end_note

** Analytical Integration
 And now calculate the cross section in picobarn.
 #+BEGIN_SRC jupyter-python :exports both :results raw file :file xs.tex
   xs_gev = total_xs_eta(η, charge, esp)
   xs_pb = gev_to_pb(xs_gev)
   tex_value(xs_pb, unit=r'\pico\barn', prefix=r'\sigma = ',
             prec=6, save=('results', 'xs.tex'))
 #+END_SRC

 #+RESULTS:
 : \(\sigma = \SI{0.053793}{\pico\barn}\)

 Lets plot the total xs as a function of η.
 #+begin_src jupyter-python :exports both :results raw drawer
   fig, ax = set_up_plot()
   η_s = np.linspace(0, 3, 1000)
   ax.plot(η_s, gev_to_pb(total_xs_eta(η_s, charge, esp)))
   ax.set_xlabel(r'$\eta$')
   ax.set_ylabel(r'$\sigma$ [pb]')
   ax.set_xlim([0, max(η_s)])
   ax.set_ylim(0)
   save_fig(fig, 'total_xs', 'xs', size=[2.5, 2.5])
 #+end_src

 #+RESULTS:
 [[file:./.ob-jupyter/4522eb3fbeaa14978f9838371acb0650910b8dbf.png]]


 Compared to sherpa, it's pretty close.
 #+NAME: 81b5ed93-0312-45dc-beec-e2ba92e22626
 #+BEGIN_SRC jupyter-python :exports both :results raw drawer
   sherpa = 0.05380
   xs_pb - sherpa
 #+END_SRC

 #+RESULTS: 81b5ed93-0312-45dc-beec-e2ba92e22626
 : -6.7112594623469635e-06

 I had to set the runcard option ~EW_SCHEME: alpha0~ to use the pure
 QED coupling constant.

** Numerical Integration
Plot our nice distribution:
#+begin_src jupyter-python :exports both :results raw drawer
  plot_points = np.linspace(*plot_interval, 1000)

  fig, ax = set_up_plot()
  ax.plot(plot_points, gev_to_pb(diff_xs(plot_points, charge=charge, esp=esp)))
  ax.set_xlabel(r'$\theta$')
  ax.set_ylabel(r'$d\sigma/d\Omega$ [pb]')
  ax.set_xlim([plot_points.min(), plot_points.max()])
  ax.axvline(interval[0], color='gray', linestyle='--')
  ax.axvline(interval[1], color='gray', linestyle='--', label=rf'$|\eta|={η}$')
  ax.legend()
  save_fig(fig, 'diff_xs', 'xs', size=[2.5, 2.5])
#+end_src

#+RESULTS:
[[file:./.ob-jupyter/3dd905e7608b91a9d89503cb41660152f3b4b55c.png]]

Define the integrand.
#+begin_src jupyter-python :exports both :results raw drawer
  def xs_pb_int(θ):
      return 2*np.pi*gev_to_pb(np.sin(θ)*diff_xs(θ, charge=charge, esp=esp))

  def xs_pb_int_η(η):
      return 2*np.pi*gev_to_pb(diff_xs_eta(η, charge, esp))
#+end_src

#+RESULTS:

Plot the integrand. # TODO: remove duplication
#+begin_src jupyter-python :exports both :results raw drawer
  fig, ax = set_up_plot()
  ax.plot(plot_points, xs_pb_int(plot_points))
  ax.set_xlabel(r'$\theta$')
  ax.set_ylabel(r'$2\pi\cdot d\sigma/d\theta [pb]')
  ax.set_xlim([plot_points.min(), plot_points.max()])
  ax.axvline(interval[0], color='gray', linestyle='--')
  ax.axvline(interval[1], color='gray', linestyle='--', label=rf'$|\eta|={η}$')
  save_fig(fig, 'xs_integrand', 'xs', size=[3, 2.2])
#+end_src

#+RESULTS:
[[file:./.ob-jupyter/ccb6653162c81c3f3e843225cb8d759178f497e0.png]]
*** Integral over θ
Intergrate σ with the mc method.
#+begin_src jupyter-python :exports both :results raw drawer
  xs_pb_res = monte_carlo.integrate(xs_pb_int, interval, epsilon=1e-3)
  xs_pb_res
#+end_src

#+RESULTS:
: IntegrationResult(result=0.05341078157056901, sigma=0.0009403117477366878, N=2209)

We gonna export that as tex.
#+begin_src jupyter-python :exports both :results raw drawer
  tex_value(*xs_pb_res.combined_result, unit=r'\pico\barn',
            prefix=r'\sigma = ', save=('results', 'xs_mc.tex'))
  tex_value(xs_pb_res.N, prefix=r'N = ', save=('results', 'xs_mc_N.tex'))
#+end_src

#+RESULTS:
: \(N = 2209\)

*** Integration over η
Plot the intgrand of the pseudo rap.
#+begin_src jupyter-python :exports both :results raw drawer
  fig, ax = set_up_plot()
  points = np.linspace(-4, 4, 1000)
  ax.set_xlim([-4, 4])
  ax.plot(points, xs_pb_int_η(points))
  ax.set_xlabel(r'$\eta$')
  ax.set_ylabel(r'$2\pi\cdot d\sigma/d\eta$ [pb]')
  ax.axvline(interval_η[0], color='gray', linestyle='--')
  ax.axvline(interval_η[1], color='gray', linestyle='--', label=rf'$|\eta|={η}$')
  save_fig(fig, 'xs_integrand_eta', 'xs', size=[3, 2])
#+end_src

#+RESULTS:
[[file:./.ob-jupyter/87a932866f779a2a07abed4ca251fa98113beca7.png]]

#+begin_src jupyter-python :exports both :results raw drawer
  xs_pb_η = monte_carlo.integrate(xs_pb_int_η,
                                  interval_η, epsilon=1e-3)
  xs_pb_η
#+end_src

#+RESULTS:
: IntegrationResult(result=0.05233083353013518, sigma=0.0009844250201876254, N=143)

As we see, the result is a little better if we use pseudo rapidity,
because the differential cross section does not difverge anymore.  But
becase our η interval is covering the range where all the variance is
occuring, the improvement is rather marginal.

And yet again export that as tex.
#+begin_src jupyter-python :exports both :results raw drawer
  tex_value(*xs_pb_η.combined_result, unit=r'\pico\barn', prefix=r'\sigma = ',
            save=('results', 'xs_mc_eta.tex'))
  tex_value(xs_pb_η.N, prefix=r'N = ', save=('results', 'xs_mc_eta_N.tex'))
#+end_src

#+RESULTS:
: \(N = 143\)

*** Using =VEGAS=
Now we use =VEGAS= on the θ parametrisation and see what happens.
#+begin_src jupyter-python :exports both :results raw drawer
  num_increments = 11
  xs_pb_vegas = monte_carlo.integrate_vegas(
      xs_pb_int,
      interval,
      num_increments=num_increments,
      alpha=1,
      increment_epsilon=0.001,
      acumulate=False,
  )
  xs_pb_vegas
#+end_src

#+RESULTS:
: VegasIntegrationResult(result=0.05391881612132655, sigma=0.00013475254855627007, N=2805, increment_borders=array([0.16380276, 0.237527  , 0.34500054, 0.50924866, 0.7700416 ,
:        1.23164456, 1.91661518, 2.37353646, 2.6331313 , 2.79643305,
:        2.90461962, 2.9777899 ]), vegas_iterations=467)

This is pretty good, although the variance reduction may be achieved
partially by accumulating the results from all runns. Here this gives
us one order of magnitude more than we wanted.

And export that as tex.
#+begin_src jupyter-python :exports both :results raw drawer
  tex_value(*xs_pb_vegas.combined_result, unit=r'\pico\barn',
            prefix=r'\sigma = ', save=('results', 'xs_mc_θ_vegas.tex'))
  tex_value(xs_pb_vegas.N, prefix=r'N = ', save=('results', 'xs_mc_θ_vegas_N.tex'))
  tex_value(num_increments, prefix=r'K = ', save=('results', 'xs_mc_θ_vegas_K.tex'))
#+end_src

#+RESULTS:
: \(K = 11\)

Surprisingly, acumulation, the result ain't much different.
This depends, of course, on the iteration count.
#+begin_src jupyter-python :exports both :results raw drawer
  monte_carlo.integrate_vegas(
      xs_pb_int,
      interval,
      num_increments=num_increments,
      alpha=1,
      increment_epsilon=0.001,
      acumulate=True,
  )
#+end_src

#+RESULTS:
: VegasIntegrationResult(result=0.053783484288076215, sigma=2.2815161064775525e-05, N=2805, increment_borders=array([0.16380276, 0.23741945, 0.34524286, 0.50821323, 0.77035483,
:        1.22942806, 1.91564534, 2.37530622, 2.63425979, 2.79678665,
:        2.90425559, 2.9777899 ]), vegas_iterations=1204)

Let's define some little helpers.
#+begin_src jupyter-python :exports both :tangle tangled/plot_utils.py
  """
  Some shorthands for common plotting tasks related to the investigation
  of monte-carlo methods in one rimension.

  Author: Valentin Boettcher <hiro at protagon.space>
  """

  import matplotlib.pyplot as plt
  import numpy as np
  from utility import *


  def plot_increments(ax, increment_borders, label=None, *args, **kwargs):
      """Plot the increment borders from a list.  The first and last one

      :param ax: the axis on which to draw
      :param list increment_borders: the borders of the increments
      :param str label: the label to apply to one of the vertical lines
      """

      ax.axvline(x=increment_borders[1], label=label, *args, **kwargs)

      for increment in increment_borders[1:-1]:
          ax.axvline(x=increment, *args, **kwargs)


  def plot_vegas_weighted_distribution(
      ax, points, dist, increment_borders, *args, **kwargs
  ):
      """Plot the distribution with VEGAS weights applied.

      :param ax: axis
      :param points: points
      :param dist: distribution
      :param increment_borders: increment borders
      """

      num_increments = increment_borders.size
      weighted_dist = dist.copy()

      for left_border, right_border in zip(increment_borders[:-1], increment_borders[1:]):
          length = right_border - left_border
          mask = (left_border <= points) & (points <= right_border)
          weighted_dist[mask] = dist[mask] * num_increments * length

      ax.plot(points, weighted_dist, *args, **kwargs)


  def plot_stratified_rho(ax, points, increment_borders, *args, **kwargs):
      """Plot the weighting distribution resulting from the increment
      borders.

      :param ax: axis
      :param points: points
      :param increment_borders: increment borders

      """

      num_increments = increment_borders.size
      ρ = np.empty_like(points)
      for left_border, right_border in zip(increment_borders[:-1], increment_borders[1:]):
          length = right_border - left_border
          mask = (left_border <= points) & (points <= right_border)
          ρ[mask] = 1 / (num_increments * length)

      ax.plot(points, ρ, *args, **kwargs)
#+end_src

#+RESULTS:

And now we plot the integrand with the incremens.
#+begin_src jupyter-python :exports both :results raw drawer
  fig, ax = set_up_plot()
  ax.set_xlim(*interval)
  ax.set_xlabel(r"$\theta$")
  ax.set_ylabel(r"$2\pi\cdot d\sigma/d\theta$ [pb]")
  ax.set_ylim([0, 0.09])

  ax.plot(plot_points, xs_pb_int(plot_points), label="Distribution")

  plot_increments(
      ax,
      xs_pb_vegas.increment_borders,
      label="Increment Borders",
      color="gray",
      linestyle="--",
  )

  plot_vegas_weighted_distribution(
      ax,
      plot_points,
      xs_pb_int(plot_points),
      xs_pb_vegas.increment_borders,
      label="Weighted Distribution",
  )

  ax.legend(fontsize="small", loc="lower left")
  save_fig(fig, "xs_integrand_vegas", "xs", size=[5, 3])
#+end_src

#+RESULTS:
[[file:./.ob-jupyter/9cb9f40087d5d473cfa956e67f4055544037565d.png]]
*** Testing the Statistics
Let's battle test the statistics.
#+begin_src jupyter-python :exports both :results raw drawer
  num_runs = 1000
  num_within = 0

  for _ in range(num_runs):
      val, err = \
          monte_carlo.integrate(xs_pb_int, interval, epsilon=1e-3).combined_result
      if abs(xs_pb - val) <= err:
          num_within += 1

  num_within/num_runs
#+end_src

#+RESULTS:
: 0.687

So we see: the standard deviation is sound.

Doing the same thing with =VEGAS= works as well.
#+begin_src jupyter-python :exports both :results raw drawer
  num_runs = 1000
  num_within = 0
  for _ in range(num_runs):
      val, err = \
          monte_carlo.integrate_vegas(xs_pb_int, interval,
                                      num_increments=10, alpha=1,
                                      epsilon=1e-3, acumulate=False,
                                      vegas_point_density=100).combined_result

      if abs(xs_pb - val) <= err:
          num_within += 1
  num_within/num_runs
#+end_src

#+RESULTS:
: 0.677

** Sampling and Analysis
Define the sample number.
#+begin_src jupyter-python :exports both :results raw drawer
  sample_num = 1_000_000
  tex_value(
      sample_num, prefix="N = ", save=("results", "4imp-sample-size.tex"),
  )
#+end_src

#+RESULTS:
: \(N = 1000000\)

Let's define shortcuts for our distributions. The 2π are just there
for formal correctnes. Factors do not influecence the outcome.
#+begin_src jupyter-python :exports both :results raw drawer
  def dist_cosθ(x):
      return gev_to_pb(diff_xs_cosθ(x, charge, esp))

  def dist_η(x):
      return gev_to_pb(diff_xs_eta(x, charge, esp))
#+end_src

#+RESULTS:

*** Sampling the cosθ cross section

Now we monte-carlo sample our distribution. We observe that the efficiency his very bad!
#+begin_src jupyter-python :exports both :results raw drawer
  cosθ_sample, cosθ_efficiency = \
      monte_carlo.sample_unweighted_array(sample_num, dist_cosθ,
                                          interval_cosθ, report_efficiency=True,
                                          cache='cache/bare_cos_theta',
                                          proc='auto')
  cosθ_efficiency
#+end_src

#+RESULTS:
: 0.027352007278111257

Let's save that.
#+begin_src jupyter-python :exports both :results raw drawer
  tex_value(
      cosθ_efficiency * 100,
      prefix=r"\mathfrak{e} = ",
      suffix=r"\%",
      save=("results", "naive_th_samp.tex"),
  )
#+end_src

#+RESULTS:
: \(\mathfrak{e} = 3\%\)

Our distribution has a lot of variance, as can be seen by plotting it.
#+begin_src jupyter-python :exports both :results raw drawer
  pts = np.linspace(*interval_cosθ, 100)
  fig, ax = set_up_plot()
  ax.plot(pts, dist_cosθ(pts))
  ax.set_xlabel(r'$\cos\theta$')
  ax.set_ylabel(r'$\frac{d\sigma}{d\Omega}$')
#+end_src

#+RESULTS:
:RESULTS:
: Text(0, 0.5, '$\\frac{d\\sigma}{d\\Omega}$')
[[file:./.ob-jupyter/a9e1c809c0f72c09ab5e91022ecd407fcc833d95.png]]
:END:

We define a friendly and easy to integrate upper limit function.
#+begin_src jupyter-python :exports both :results raw drawer
  fig, ax = set_up_plot()
  upper_limit = dist_cosθ(interval_cosθ[0]) / interval_cosθ[0] ** 2
  upper_base = dist_cosθ(0)


  def upper(x):
      return upper_base + upper_limit * x ** 2


  def upper_int(x):
      return upper_base * x + upper_limit * x ** 3 / 3


  ax.plot(pts, upper(pts), label="upper bound")
  ax.plot(pts, dist_cosθ(pts), label=r"$f_{\cos\theta}$")

  ax.legend(fontsize='small')
  ax.set_xlabel(r"$\cos\theta$")
  ax.set_ylabel(r"$\frac{d\sigma}{d\cos\theta}$ [pb]")
  save_fig(fig, "upper_bound", "xs_sampling", size=(3, 2.5))
#+end_src

#+RESULTS:
[[file:./.ob-jupyter/647593b36e5170280820c31c63b884cae0ebbee6.png]]


To increase our efficiency, we have to specify an upper bound. That is
at least a little bit better. The numeric inversion is horribly inefficent.
#+begin_src jupyter-python :exports both :results raw drawer
  cosθ_sample_tuned, cosθ_efficiency_tuned = monte_carlo.sample_unweighted_array(
      sample_num,
      dist_cosθ,
      interval_cosθ,
      report_efficiency=True,
      proc="auto",
      cache="cache/bare_cos_theta_tuned",
      upper_bound=[upper, upper_int],
  )
  cosθ_efficiency_tuned
#+end_src

#+RESULTS:
: 0.07903687969629128
<<cosθ-bare-eff>>

#+begin_src jupyter-python :exports both :results raw drawer
  tex_value(
      cosθ_efficiency_tuned * 100,
      prefix=r"\mathfrak{e} = ",
      suffix=r"\%",
      save=("results", "tuned_th_samp.tex"),
  )
#+end_src

#+RESULTS:
: \(\mathfrak{e} = 8\%\)

# TODO: Looks fishy
Nice! And now draw some histograms.

We define an auxilliary method for convenience.
#+begin_src jupyter-python :exports both :results raw drawer :tangle tangled/plot_utils.py
  import matplotlib.gridspec as gridspec


  def draw_ratio_plot(histograms, normalize_to=1, **kwargs):
      fig, (ax_hist, ax_ratio) = set_up_plot(
          2, 1, sharex=True, gridspec_kw=dict(height_ratios=[3, 1], hspace=0), **kwargs
      )

      reference, edges = histograms[0]["hist"]
      reference_error = np.sqrt(reference)

      ref_int = hist_integral(histograms[0]["hist"])
      reference = reference / ref_int
      reference_error = reference_error / ref_int

      for histogram in histograms:
          heights, _ = histogram["hist"]
          integral = hist_integral([heights, edges])
          errors = np.sqrt(heights) / integral
          heights = heights / integral

          draw_histogram(
              ax_hist,
              [heights, edges],
              errorbars=errors,
              hist_kwargs=(
                  histogram["hist_kwargs"] if "hist_kwargs" in histogram else dict()
              ),
              errorbar_kwargs=(
                  histogram["errorbar_kwargs"]
                  if "errorbar_kwargs" in histogram
                  else dict()
              ),
              normalize_to=normalize_to,
          )

          set_up_axis(ax_ratio, pimp_top=False)
          ax_ratio.set_ylabel("ratio")
          draw_histogram(
              ax_ratio,
              [
                  np.divide(
                      heights, reference, out=np.ones_like(heights), where=reference != 0
                  ),
                  edges,
              ],
              errorbars=np.divide(
                  errors, reference, out=np.zeros_like(heights), where=reference != 0
              ),
              hist_kwargs=(
                  histogram["hist_kwargs"] if "hist_kwargs" in histogram else dict()
              ),
              errorbar_kwargs=(
                  histogram["errorbar_kwargs"]
                  if "errorbar_kwargs" in histogram
                  else dict()
              ),
              normalize_to=None,
          )

      return fig, (ax_hist, ax_ratio)


  def hist_integral(hist):
      heights, edges = hist
      return heights @ (edges[1:] - edges[:-1])


  def draw_histogram(
      ax,
      histogram,
      errorbars=True,
      hist_kwargs=dict(color="#1f77b4"),
      errorbar_kwargs=dict(),
      normalize_to=None,
  ):
      """Draws a histogram with optional errorbars using the step style.

      :param ax: axis to draw on
      :param histogram: an array of the form [heights, edges]
      :param hist_kwargs: keyword args to pass to `ax.step`
      :param errorbar_kwargs: keyword args to pass to `ax.errorbar`
      :param normalize_to: if set, the histogram will be normalized to the value
      :returns: the given axis
      """

      heights, edges = histogram
      centers = (edges[1:] + edges[:-1]) / 2
      deviations = (
          (errorbars if isinstance(errorbars, (np.ndarray, list)) else np.sqrt(heights))
          if errorbars is not False
          else None
      )

      if normalize_to is not None:
          integral = hist_integral(histogram)
          heights = heights / integral * normalize_to
          if errorbars is not False:
              deviations = deviations / integral * normalize_to

      hist_plot = ax.step(edges, [heights[0], *heights], **hist_kwargs)

      if errorbars is not False:
          if "color" not in errorbar_kwargs:
              errorbar_kwargs["color"] = hist_plot[0].get_color()

          ax.errorbar(centers, heights, deviations, linestyle="none", **errorbar_kwargs)

      ax.set_xlim(*[edges[0], edges[-1]])

      return ax


  def draw_histo_auto(points, xlabel, bins=50, range=None, rethist=False, **kwargs):
      """Creates a histogram figure from sample points, normalized to unity.

      :param points: samples
      :param xlabel: label of the x axis
      :param bins: number of bins
      :param range: the range of the values
      :param rethist: whether to return the histogram as third argument
      :returns: figure, axis
      """

      hist = np.histogram(points, bins, range=range, **kwargs)
      fig, ax = set_up_plot()
      draw_histogram(ax, hist, normalize_to=1)

      ax.set_xlabel(xlabel)
      ax.set_ylabel("Count")

      return (fig, ax, hist) if rethist else (fig, ax)
#+end_src

#+RESULTS:

The histogram for cosθ.
#+begin_src jupyter-python :exports both :results raw drawer
  fig, _ = draw_histo_auto(cosθ_sample, r'$\cos\theta$')
  save_fig(fig, 'histo_cos_theta', 'xs', size=(4,3))
  hist_cosθ = np.histogram(cosθ_sample, bins=50, range=interval_cosθ)
#+end_src

#+RESULTS:
[[file:./.ob-jupyter/dde553030cdb96c1f0a0b223abf9bdd4602119af.png]]

*** Observables
Now we define some utilities to draw real 4-momentum samples.
#+begin_src jupyter-python :exports both :tangle tangled/xs.py
  @numpy_cache("momentum_cache")
  def sample_momenta(sample_num, interval, charge, esp, seed=None, **kwargs):
      """Samples `sample_num` unweighted photon 4-momenta from the
      cross-section. Superflous kwargs are passed on to
      `sample_unweighted_array`.

      :param sample_num: number of samples to take
      :param interval: cosθ interval to sample from
      :param charge: the charge of the quark
      :param esp: center of mass energy
      :param seed: the seed for the rng, optional, default is system
          time

      :returns: an array of 4 photon momenta

      :rtype: np.ndarray

      """

      cosθ_sample = monte_carlo.sample_unweighted_array(
          sample_num, lambda x: diff_xs_cosθ(x, charge, esp), interval_cosθ, **kwargs
      )

      φ_sample = np.random.uniform(0, 1, sample_num)

      def make_momentum(esp, cosθ, φ):
          sinθ = np.sqrt(1 - cosθ ** 2)
          return np.array([1, sinθ * np.cos(φ), sinθ * np.sin(φ), cosθ],) * esp / 2

      momenta = np.array(
          [make_momentum(esp, cosθ, φ) for cosθ, φ in np.array([cosθ_sample, φ_sample]).T]
      )
      return momenta
#+end_src

#+RESULTS:

To generate histograms of other obeservables, we have to define them
as functions on 4-impuleses. Using those to transform samples is
analogous to transforming the distribution itself.
#+begin_src jupyter-python :session obs :exports both :results raw drawer :tangle tangled/observables.py
  """This module defines some observables on arrays of 4-pulses."""
  import numpy as np
  from utility import minkowski_product

  def p_t(p):
      """Transverse momentum

      :param p: array of 4-momenta
      """

      return np.linalg.norm(p[:,1:3], axis=1)

  def η(p):
      """Pseudo rapidity.

      :param p: array of 4-momenta
      """

      return np.arccosh(np.linalg.norm(p[:,1:], axis=1)/p_t(p))*np.sign(p[:, 3])

  def inv_m(p_1, p_2):
      """Invariant mass off the final state system.

      :param p_1: array of 4-momenta, first fs particle
      :param p_2: array of 4-momenta, second fs particle
      """

      total_p = p_1 + p_2
      return np.sqrt(minkowski_product(total_p, total_p))

  def cosθ(p):
      return p[:, 3] / p[:, 0]
#+end_src

#+RESULTS:

And import them.
#+begin_src jupyter-python :exports both :results raw drawer
  %aimport tangled.observables
  obs = tangled.observables
#+end_src

#+RESULTS:

Lets try it out.
#+begin_src jupyter-python :exports both :results raw drawer
  momentum_sample = sample_momenta(
      sample_num,
      interval_cosθ,
      charge,
      esp,
      proc='auto',
      momentum_cache="cache/momenta_bare_cos_theta",
  )
  momentum_sample
#+end_src

#+RESULTS:
: array([[100.        ,  16.04646249,  12.03145593,  97.96813313],
:        [100.        ,  51.41640893,  11.55602144,  84.98712409],
:        [100.        ,  40.75310071,  39.90715071,  82.13771426],
:        ...,
:        [100.        ,  20.83112183,   7.06626328,  97.55066523],
:        [100.        ,  33.23340199,   1.86567636, -94.29772131],
:        [100.        ,  32.9373831 ,  47.67830405, -81.49790254]])


Now let's make a histogram of the η distribution.
#+begin_src jupyter-python :exports both :results raw drawer
  η_sample = obs.η(momentum_sample)
  fig, ax, hist_obs_η = draw_histo_auto(
      η_sample, r"$eta$", range=interval_η, rethist=True
  )
  save_fig(fig, "histo_eta", "xs_sampling", size=[3, 3])
#+end_src

#+RESULTS:
[[file:./.ob-jupyter/e2b510b9e200304cea662510e2bb1448cddf5055.png]]


And the same for the p_t (transverse momentum) distribution.
#+begin_src jupyter-python :exports both :results raw drawer
  p_t_sample = obs.p_t(momentum_sample)
  fig, ax, hist_obs_pt = draw_histo_auto(
      p_t_sample, r"$p_T$ [GeV]", range=interval_pt, rethist=True
  )
  save_fig(fig, "histo_pt", "xs_sampling", size=[3, 3])
#+end_src

#+RESULTS:
[[file:./.ob-jupyter/fba21aa6168c255a5523d865bace1ed6cfd2cab6.png]]

That looks somewhat fishy, but it isn't.
#+begin_src jupyter-python :exports both :results raw drawer
  fig, ax = set_up_plot()
  points = np.linspace(interval_pt[0], interval_pt[1] - .01, 1000)
  ax.plot(points, gev_to_pb(diff_xs_p_t(points, charge, esp)))
  ax.set_xlabel(r'$p_T$')
  ax.set_xlim(interval_pt[0], interval_pt[1] + 1)
  ax.set_ylim([0, gev_to_pb(diff_xs_p_t(interval_pt[1] -.01, charge, esp))])
  ax.set_ylabel(r'$\frac{d\sigma}{dp_t}$ [pb]')
  save_fig(fig, 'diff_xs_p_t', 'xs_sampling', size=[4, 2])
#+end_src

#+RESULTS:
[[file:./.ob-jupyter/29724b8c1f2b0005a05f64f999cf95d248ee0082.png]]
this is strongly peaked at p_t=100GeV. (The jacobian goes like 1/x there!)

*** Sampling the η cross section
An again we see that the efficiency is way, way! better...
#+begin_src jupyter-python :exports both :results raw drawer
  η_sample, η_efficiency = monte_carlo.sample_unweighted_array(
      sample_num,
      dist_η,
      interval_η,
      report_efficiency=True,
      proc="auto",
      cache="cache/sample_bare_eta",
  )
  tex_value(
      η_efficiency * 100,
      prefix=r"\mathfrak{e} = ",
      suffix=r"\%",
      save=("results", "eta_eff.tex"),
  )
#+end_src

#+RESULTS:
: \(\mathfrak{e} = 41\%\)
<<η-eff>>

Let's draw a histogram to compare with the previous results.
#+begin_src jupyter-python :exports both :results raw drawer
  η_hist = np.histogram(η_sample, bins=50)
  fig, (ax_hist, ax_ratio) = draw_ratio_plot(
      [
          dict(hist=η_hist, hist_kwargs=dict(label=r"sampled from $d\sigma / d\eta$"),),
          dict(
              hist=hist_obs_η,
              hist_kwargs=dict(
                  label=r"sampled from $d\sigma / d\cos\theta$", color="black"
              ),
          ),
      ],
  )

  ax_hist.legend(loc="upper center", fontsize="small")
  ax_ratio.set_xlabel(r"$\eta$")
  save_fig(fig, "comparison_eta", "xs_sampling", size=(4, 4))
#+end_src

#+RESULTS:
[[file:./.ob-jupyter/931d0c4522606eb420d0fb674a918a79e6244ce2.png]]

Looks good to me :).

*** Sampling with =VEGAS=
To get the increments, we have to let =VEGAS= loose on our
distribution. We throw away the integral, but keep the increments.

#+begin_src jupyter-python :exports both :results raw drawer
  K = 10
  increments = monte_carlo.integrate_vegas(
      dist_cosθ, interval_cosθ, num_increments=K, alpha=1, increment_epsilon=0.001
  ).increment_borders
  tex_value(
      K, prefix=r"K = ", save=("results", "vegas_samp_num_increments.tex"),
  )
  increments
#+end_src

#+RESULTS:
: array([-0.9866143 , -0.96961793, -0.93102927, -0.83928851, -0.60522124,
:         0.00165114,  0.60381197,  0.83945077,  0.93141193,  0.9698008 ,
:         0.9866143 ])

Visualizing the increment borders gives us the information we want.
#+begin_src jupyter-python :exports both :results raw drawer
  pts = np.linspace(*interval_cosθ, 100)
  fig, ax = set_up_plot()
  ax.plot(pts, dist_cosθ(pts))
  ax.set_xlabel(r'$\cos\theta$')
  ax.set_ylabel(r'$\frac{d\sigma}{d\Omega}$')
  ax.set_xlim(*interval_cosθ)
  plot_increments(ax, increments,
                  label='Increment Borderds', color='gray', linestyle='--')
  ax.legend()
#+end_src

#+RESULTS:
:RESULTS:
: <matplotlib.legend.Legend at 0x7f7205eb4340>
[[file:./.ob-jupyter/b4b1e7c332c55259afcda37c371f1edc4a56c0ea.png]]
:END:

We can now plot the reweighted distribution to observe the variance
reduction visually.

#+begin_src jupyter-python :exports both :results raw drawer
  pts = np.linspace(*interval_cosθ, 1000)
  fig, ax = set_up_plot()
  ax.plot(pts, dist_cosθ(pts), label="Distribution")
  plot_vegas_weighted_distribution(
      ax, pts, dist_cosθ(pts), increments, label="Weighted Distribution"
  )
  ax.set_xlabel(r"$\cos\theta$")
  ax.set_ylabel(r"$\frac{d\sigma}{d\cos\theta}$")
  ax.set_xlim(*interval_cosθ)
  plot_increments(
      ax, increments, label="Increment Borderds", color="gray", linestyle="--"
  )
  ax.legend(fontsize="small")
  save_fig(fig, "vegas_strat_dist", "xs_sampling", size=(3, 2.3))
#+end_src

#+RESULTS:
[[file:./.ob-jupyter/b2f71eaa54629fa9e6f00976eccb567a71ff0f54.png]]


I am batman! Let's plot the weighting distribution.
#+begin_src jupyter-python :exports both :results raw drawer
  pts = np.linspace(*interval_cosθ, 1000)
  fig, ax = set_up_plot()
  plot_stratified_rho(ax, pts, increments)
  ax.set_xlabel(r"$\cos\theta$")
  ax.set_ylabel(r"$\rho")
  ax.set_xlim(*interval_cosθ)
  save_fig(fig, "vegas_rho", "xs_sampling", size=(3, 2.3))
#+end_src

#+RESULTS:
[[file:./.ob-jupyter/4a8e255a72be03dc50b269e14143043b5d9472f5.png]]

Now, draw a sample and look at the efficiency.

#+begin_src jupyter-python :exports both :results raw drawer
  cosθ_sample_strat, cosθ_efficiency_strat = monte_carlo.sample_unweighted_array(
      sample_num,
      dist_cosθ,
      increment_borders=increments,
      report_efficiency=True,
      proc="auto",
      cache="cache/sample_bare_cos_theta_vegas",
  )
  cosθ_efficiency_strat
#+end_src

#+RESULTS:
: 0.5875845864661654

#+begin_src jupyter-python :exports both :results raw drawer
  tex_value(
      cosθ_efficiency_strat * 100,
      prefix=r"\mathfrak{e} = ",
      suffix=r"\%",
      save=("results", "strat_th_samp.tex"),
  )
#+end_src

#+RESULTS:
: \(\mathfrak{e} = 59\%\)

If we compare that to [[cosθ-bare-eff]], we can see the improvement :P.
It is even better the [[η-eff]].  The histogram looks just the same.

#+begin_src jupyter-python :exports both :results raw drawer
fig, _ = draw_histo_auto(cosθ_sample_strat, r'$\cos\theta$')
save_fig(fig, 'histo_cos_theta_strat', 'xs', size=(4,3))
#+end_src

#+RESULTS:
[[file:./.ob-jupyter/1d90af7d456726ad3780a203acc8938f1894f6b1.png]]

*** Some Histograms with Rivet
**** Init
#+begin_src jupyter-python :exports both :results raw drawer
  import yoda
#+end_src

#+RESULTS:
: Welcome to JupyROOT 6.20/04

**** Plot the Histos
#+RESULTS:

#+begin_src jupyter-python :exports both :results raw drawer :tangle tangled/plot_utils.py
  def yoda_to_numpy(histo):
      histo.normalize(
          histo.numEntries() * ((histo.xMax() - histo.xMin()) / histo.numBins())
      )
      edges = np.append(histo.xMins(), histo.xMax())
      heights = histo.yVals().astype(int)

      return heights, edges


  def draw_yoda_histo_auto(h, xlabel, **kwargs):
      hist = yoda_to_numpy(h)
      fig, ax = set_up_plot()
      draw_histogram(ax, hist, errorbars=True, normalize_to=1, **kwargs)

      ax.set_xlabel(xlabel)
      return fig, ax
#+end_src

#+RESULTS:

#+begin_src jupyter-python :exports both :results raw drawer
  yoda_file = yoda.read("../../runcards/qqgg/analysis/Analysis.yoda")
  sherpa_histos = {
      "pT": dict(reference=hist_obs_pt, label="$p_T$ [GeV]"),
      "eta": dict(reference=hist_obs_η, label=r"$\eta$"),
      "cos_theta": dict(reference=hist_cosθ, label=r"$\cos\theta$"),
  }

  for key, sherpa_hist in sherpa_histos.items():
      yoda_hist = yoda_to_numpy(yoda_file["/MC_DIPHOTON_SIMPLE/" + key])
      label = sherpa_hist["label"]
      fig, (ax_hist, ax_ratio) = draw_ratio_plot(
          [
              dict(
                  hist=yoda_hist,
                  hist_kwargs=dict(
                      label="Sherpa Reference"
                  ),
                  errorbars=True,
              ),
              dict(
                  hist=sherpa_hist["reference"],
                  hist_kwargs=dict(label="Own Implementation"),
              ),
          ],
      )
      ax_ratio.set_xlabel(label)
      ax_hist.legend(fontsize='small')
      save_fig(fig, "histo_sherpa_" + key, "xs_sampling", size=(4, 3.5))
#+end_src

#+RESULTS:
:RESULTS:
[[file:./.ob-jupyter/29fb6b4200f4f70eb8cbb96b498114caa819d146.png]]
[[file:./.ob-jupyter/b86e1bcd4461345fb9bfe7114c2626bb966c9d8b.png]]
[[file:./.ob-jupyter/f8c1f07465d43e58bace241d748c45e21acd2852.png]]
:END: