In [1]:

# Useful for debugging
%load_ext autoreload
%autoreload 2
%config InlineBackend.figure_format = 'retina'

# Useful for debugging %load_ext autoreload %autoreload 2 %config InlineBackend.figure_format = 'retina'

In [2]:

from distgen.dist import SuperGaussian
from distgen.dist import Norm
from distgen.physical_constants import unit_registry

from matplotlib import pyplot as plt

import numpy as np

from distgen.dist import SuperGaussian from distgen.dist import Norm from distgen.physical_constants import unit_registry from matplotlib import pyplot as plt import numpy as np

Metrics for comparing different distributions¶

---

In some cases, it is desired to target a particular distribution shape in an optimization. To facilitate this the following two metrics are implemented:

Kullback-Leibler (Relative Entropy)¶

Defined as: $D_{KL}(P|Q) = \int_{-\infty}^{\infty}p(x)\ln{\left(\frac{p(x)}{q(x)}\right)} dx$

Note that this is not well defined in regions where the PDF $q(x)=0$. This causes trouble for distributions like a uniform distribution. For this, its suggested to use a SuperGaussian to approximate the uniform dist (see below).

https://en.wikipedia.org/wiki/Kullback–Leibler_divergence

In [3]:

L = 2 * unit_registry("ps")
avg_t_sg = 0 * unit_registry("ps")
sigma_t_sg = L / np.sqrt(12)

L = 2 * unit_registry("ps") avg_t_sg = 0 * unit_registry("ps") sigma_t_sg = L / np.sqrt(12)

In [4]:

sg = SuperGaussian("t", avg_t=avg_t_sg, sigma_t=sigma_t_sg, p=12)

sg = SuperGaussian("t", avg_t=avg_t_sg, sigma_t=sigma_t_sg, p=12)

In [5]:

tsg, Psg = sg.get_x_pts(), sg.pdf()
plt.plot(tsg, Psg);

tsg, Psg = sg.get_x_pts(), sg.pdf() plt.plot(tsg, Psg);

In [6]:

norm = Norm("t", avg_t=avg_t_sg, sigma_t=sigma_t_sg)

norm = Norm("t", avg_t=avg_t_sg, sigma_t=sigma_t_sg)

In [7]:

tn, Pn = norm.get_x_pts(), norm.pdf()
plt.plot(tn, Pn);

tn, Pn = norm.get_x_pts(), norm.pdf() plt.plot(tn, Pn);

In [8]:

from distgen.metrics import kullback_liebler_div

from distgen.metrics import kullback_liebler_div

In [9]:

kullback_liebler_div(tn, Pn, tsg, Psg, as_float=False)

kullback_liebler_div(tn, Pn, tsg, Psg, as_float=False)

Out[9]:

4.301370109274345 dimensionless

The functions work with raw NumPy arrays, and support both float output or Pint Quantity outputs:

In [10]:

kullback_liebler_div(
    tn.magnitude, Pn.magnitude, tsg.magnitude, Psg.magnitude, as_float=True
)

kullback_liebler_div( tn.magnitude, Pn.magnitude, tsg.magnitude, Psg.magnitude, as_float=True )

Out[10]:

np.float64(4.301370109274345)

In addition to the Kullback Liebler Divergence, the residual squared between two distributions is implemented:

In [11]:

from distgen.metrics import res2

from distgen.metrics import res2

In [12]:

res2(tn, Pn, tsg, Psg, as_float=False, normalize=True)

res2(tn, Pn, tsg, Psg, as_float=False, normalize=True)

Out[12]:

0.11983271504891024 dimensionless

In [13]:

res2(tn.magnitude, Pn.magnitude, tsg.magnitude, Psg.magnitude, as_float=True)

res2(tn.magnitude, Pn.magnitude, tsg.magnitude, Psg.magnitude, as_float=True)

Out[13]:

np.float64(0.058441901838167956)

Helper Functions¶

---

In [14]:

from distgen.metrics import resample_pq

from distgen.metrics import resample_pq

In [15]:

resample_pq(tn, Pn, tsg, Psg, plot=True);

resample_pq(tn, Pn, tsg, Psg, plot=True);

In [16]:

dist_yaml = """
n_particle: 30000
species: electron
r_dist:
  truncation_fraction:
    units: dimensionless
    value: 0.5
  truncation_radius:
    units: mm
    value: 2.3319043122
  type: rg
random:
  type: hammersley
start:
  MTE:
    units: meV
    value: 130
  type: cathode
t_dist:
  p:
    units: ''
    value: 1
  sigma_t:
    units: ps
    value: 10
  type: sg
total_charge:
  units: pC
  value: 100
"""

dist_yaml = """ n_particle: 30000 species: electron r_dist: truncation_fraction: units: dimensionless value: 0.5 truncation_radius: units: mm value: 2.3319043122 type: rg random: type: hammersley start: MTE: units: meV value: 130 type: cathode t_dist: p: units: '' value: 1 sigma_t: units: ps value: 10 type: sg total_charge: units: pC value: 100 """

In [17]:

from distgen import Generator

from distgen import Generator

In [18]:

D = Generator(dist_yaml)

D = Generator(dist_yaml)

In [19]:

P = D.run()

P = D.run()

In [20]:

P.plot("t")

P.plot("t")

In [21]:

from distgen.metrics import get_current_profile
from distgen.metrics import rms_equivalent_current_nonuniformity

from distgen.metrics import get_current_profile from distgen.metrics import rms_equivalent_current_nonuniformity

In [22]:

t, current_profile = get_current_profile(P)

t, current_profile = get_current_profile(P)

In [23]:

plt.plot(t, current_profile);

plt.plot(t, current_profile);

In [24]:

ps = np.linspace(1, 12, 50)

ps = np.linspace(1, 12, 50)

In [25]:

kldivs = np.zeros(ps.shape)
res2s = np.zeros(ps.shape)

for ii, p in enumerate(ps):
    D["t_dist:p"] = p

    P = D.run()

    t, current_profile = get_current_profile(P)

    plt.plot(t, current_profile)

    kldivs[ii] = rms_equivalent_current_nonuniformity(P, method="kl_div", p=12)

    res2s[ii] = rms_equivalent_current_nonuniformity(P, method="res2")

plt.xlabel("t (s)")
plt.ylabel("$\\rho$ ($s^{-1}$)");

kldivs = np.zeros(ps.shape) res2s = np.zeros(ps.shape) for ii, p in enumerate(ps): D["t_dist:p"] = p P = D.run() t, current_profile = get_current_profile(P) plt.plot(t, current_profile) kldivs[ii] = rms_equivalent_current_nonuniformity(P, method="kl_div", p=12) res2s[ii] = rms_equivalent_current_nonuniformity(P, method="res2") plt.xlabel("t (s)") plt.ylabel("$\\rho$ ($s^{-1}$)");

1.0 dimensionless
1.0 dimensionless
1.2244897959183674 dimensionless
1.4489795918367347 dimensionless

1.6734693877551021 dimensionless
1.8979591836734695 dimensionless
2.1224489795918364 dimensionless
2.3469387755102042 dimensionless

2.571428571428571 dimensionless
2.795918367346939 dimensionless
3.020408163265306 dimensionless
3.2448979591836733 dimensionless
3.4693877551020407 dimensionless

3.693877551020408 dimensionless
3.9183673469387754 dimensionless
4.142857142857142 dimensionless
4.36734693877551 dimensionless

4.591836734693878 dimensionless
4.816326530612245 dimensionless
5.040816326530612 dimensionless
5.26530612244898 dimensionless

5.489795918367347 dimensionless
5.714285714285714 dimensionless
5.938775510204081 dimensionless
6.163265306122449 dimensionless

6.387755102040816 dimensionless
6.612244897959184 dimensionless
6.836734693877551 dimensionless
7.061224489795919 dimensionless

7.285714285714286 dimensionless
7.510204081632653 dimensionless
7.73469387755102 dimensionless
7.959183673469388 dimensionless

8.183673469387756 dimensionless
8.408163265306122 dimensionless
8.63265306122449 dimensionless
8.857142857142858 dimensionless

9.081632653061224 dimensionless
9.306122448979592 dimensionless
9.53061224489796 dimensionless
9.755102040816327 dimensionless

9.979591836734693 dimensionless
10.204081632653061 dimensionless
10.428571428571429 dimensionless
10.653061224489797 dimensionless

10.877551020408163 dimensionless
11.10204081632653 dimensionless
11.326530612244898 dimensionless
11.551020408163266 dimensionless

11.775510204081632 dimensionless

$D_{KL} = \int_{-\infty}^{\infty} P\ln(P/Q)dt$

Note, not defined for uniform beam where $Q_u = \frac{1}{t_2-t_1}\left[\theta(t-t_1)-\theta(t-t_2)\right]$. If comparing to uniform beam, replace target distribution with rms equivalent super-Gaussian with power $p$: $Q=Q_{SG}(t; p)$. So KL-div nonuniformity:

$\lim_{p\rightarrow\infty}\int_{-\infty}^{\infty} P\ln[P/Q_{SG}(t;p)]dt$

In [26]:

plt.plot(ps, kldivs)
plt.xlabel("super-Gaussian power")
plt.ylabel("KL-Divergence");

plt.plot(ps, kldivs) plt.xlabel("super-Gaussian power") plt.ylabel("KL-Divergence");

$\frac{\int_{-\infty}^{\infty} (P-Q)^2dt}{\int_{-\infty}^{\infty} Q^2(t)dt}$

In [27]:

plt.plot(ps, res2s)
plt.xlabel("super-Gaussian power")
plt.ylabel("integrated squared residuals");

plt.plot(ps, res2s) plt.xlabel("super-Gaussian power") plt.ylabel("integrated squared residuals");

In [28]:

D["t_dist:p"] = 4
P = D.run()
P.plot("t")

D["t_dist:p"] = 4 P = D.run() P.plot("t")

12.0 dimensionless

In [29]:

D["t_dist:p"] = 6
P = D.run()
P.plot("t")

D["t_dist:p"] = 6 P = D.run() P.plot("t")

4 dimensionless

In [ ]:

In [ ]: