# Create a decay dataset suitable for radioactivedecay from PyNE
This notebook creates a set of decay dataset files suitable for radioactivedecay `v0.4.5+` from the decay data in PyNE `v0.7.6`. The PyNE data is based on the [191004 ENDSF](https://github.com/pyne/pyne/pull/1216) release.

First import the necessary modules.

```python
import math, pickle

from pyne import nucname, data
from pyne.material import Material
import pyne

import numpy as np
import pandas as pd
from scipy import sparse
from sympy import Integer, S, log, Matrix
from sympy.matrices import SparseMatrix

print("Using PyNE version:", pyne.__version__)
```

    Using PyNE version: 0.7.6

### Create a DataFrame containing the PyNE decay data
Create a list of all the ground state (non-metastable) radionuclides in PyNE. Note we exclude the metastable states as the PyNE treatment for decay chains passing through metastable states is [incorrect](https://github.com/pyne/pyne/issues/739) as of `v0.7.6`. We also exclude radionuclides with undefined half-lives.

```python
pyne_nonmetastable_ids = []
for z in range(1,120):
    for a in range(1,300):
        try:
            id = z*10000000+a*10000
            hl = data.half_life(id)
        except:
            continue
        if hl == float("inf"): continue  # ignore stable nuclides
        elif math.isnan(hl): continue  # ignore nuclides where the half-life is undefined half-lives
        pyne_nonmetastable_ids.append(id)
print("Total number of radionuclides:", len(pyne_nonmetastable_ids))
```

    Total number of radionuclides: 2920

Define functions to fill a Pandas DataFrame with the decay data from PyNE.

```python
def add_hyphen(name):
    """Add hypen to radionuclide name string e.g. H3 to H-3."""

    for i in range(1, len(name)):
        if not name[i].isdigit():
            continue
        name = name[:i] + "-" + name[i:]
        break
    return name

def create_rows(ids):
    """Create a list of dictionaries which will become rows of the DataFrame of decay data."""

    rows = []
    for id in ids:
        name = add_hyphen(nucname.name(id))
        Z, A = nucname.znum(id), nucname.anum(id)
        hl = data.half_life(id)
        children = list(data.decay_children(id))
        bf = []
        modes = []
        atomic_mass = data.atomic_mass(id)
        for c in children:
            bf.append(data.branch_ratio(id, c))
            cZ, cA = nucname.znum(c), nucname.anum(c)
            if Z == cZ and A == cA: modes.append("IT")
            elif Z-2 == cZ and A-4 == cA: modes.append("α")
            elif Z+1 == cZ and A == cA: modes.append("β-")
            elif Z-1 == cZ and A == cA: modes.append("β+ or EC")
            else: modes.append("SF or other")
        rows.append({"Radionuclide": name, "id": id, "Z": Z, "A": A, "Half-life_s": hl,
                    "Num_decay_modes": len(children), "Progeny": children, "Branching_fractions": bf,
                    "Modes": modes, "Atomic_mass": atomic_mass})
    return rows
```

Add all the PyNE decay data to a DataFrame.

```python
col_names = ["Radionuclide", "id", "Z", "A", "Half-life_s", "Num_decay_modes",
             "Progeny", "Branching_fractions", "Modes", "Atomic_mass"]
pyne_full = pd.DataFrame(create_rows(pyne_nonmetastable_ids), columns=col_names)
pyne_full.set_index("Radionuclide", inplace=True)
pyne_full.to_csv("pyne_full.csv", index=True)
pyne_full.head(n=10)
```

      id
      Z
      A
      Half-life_s
      Num_decay_modes
      Progeny
      Branching_fractions
      Modes
      Atomic_mass

      Radionuclide

      H-3
      10030000
      1
      3
      3.887896e+08
      1
      [20030000]
      [1.0]
      [β-]
      3.016049

      H-6
      10060000
      1
      6
      2.016584e-23
      1
      [10050000]
      [1.0]
      [SF or other]
      6.044955

      He-5
      20050000
      2
      5
      8.167164e-24
      0
      []
      []
      []
      5.012057

      He-6
      20060000
      2
      6
      8.067000e-01
      1
      [30060000]
      [1.0]
      [β-]
      6.018886

      He-7
      20070000
      2
      7
      1.890547e-18
      0
      []
      []
      []
      7.027991

      He-8
      20080000
      2
      8
      1.191000e-01
      2
      [30070000, 30080000]
      [0.16, 0.84]
      [SF or other, β-]
      8.033934

      He-10
      20100000
      2
      10
      3.781095e-18
      1
      [20090000]
      [1.0]
      [SF or other]
      10.052815

      Li-5
      30050000
      3
      5
      1.550249e-23
      0
      []
      []
      []
      5.012538

      Li-8
      30080000
      3
      8
      8.399000e-01
      1
      [40080000]
      [1.0]
      [β-]
      8.022486

      Li-9
      30090000
      3
      9
      1.783000e-01
      2
      [40080000, 40090000]
      [0.5080000000000001, 0.4919999999999999]
      [SF or other, β-]
      9.026790

### Order the DataFrame so all progeny are located below their parent
The radionuclides in the DataFrame need to be ordered so that progeny (decay children) are always located lower than their parent. This is so the subsequent matrices that we create are lower triangular.

To achieve this we first count how many times each radioactive decay mode occurs in the dataset.

```python
modes = pd.Series(np.concatenate(pyne_full.Modes))
print("β+ or electron capture:", modes.value_counts()["β+ or EC"])
print("β-:", modes.value_counts()["β-"])
print("α:", modes.value_counts()["α"])
print("Spontaneous Fission or other:", modes.value_counts()["SF or other"])
print("Total number of decay modes:", pyne_full.Num_decay_modes.sum())
```

    β+ or electron capture: 1143
    β-: 1133
    α: 580
    Spontaneous Fission or other: 1257
    Total number of decay modes: 4113

We order by decreasing mass number (A), followed by decreasing atomic number (Z), as there are more β+ and EC decays than β- decays.

```python
pyne_full.sort_values(by=["A", "Z"], inplace=True, ascending=[False, False])
pyne_full.head(n=10)
```

      id
      Z
      A
      Half-life_s
      Num_decay_modes
      Progeny
      Branching_fractions
      Modes
      Atomic_mass

      Radionuclide

      Fl-285
      1142850000
      114
      285
      0.1500
      1
      [1122810000]
      [1.0]
      [α]
      285.183579

      Cn-285
      1122850000
      112
      285
      34.0000
      1
      [1102810000]
      [1.0]
      [α]
      285.177321

      Fl-284
      1142840000
      114
      284
      0.0025
      0
      []
      []
      []
      284.181344

      Cn-284
      1122840000
      112
      284
      0.1010
      0
      []
      []
      []
      284.174499

      Cn-283
      1122830000
      112
      283
      4.0000
      0
      []
      []
      []
      283.173362

      Cn-282
      1122820000
      112
      282
      0.0005
      0
      []
      []
      []
      282.170668

      Rg-282
      1112820000
      111
      282
      0.5000
      1
      [1092780000]
      [1.0]
      [α]
      282.169405

      Cn-281
      1122810000
      112
      281
      0.1300
      1
      [1102770000]
      [1.0]
      [α]
      281.169641

      Rg-281
      1112810000
      111
      281
      26.0000
      0
      []
      []
      []
      281.166718

      Ds-281
      1102810000
      110
      281
      9.6000
      0
      []
      []
      []
      281.164715

Now it is necessary to correct the positions of the remaining radionuclides that are not ordered correctly. We do this by looping over all the radionuclides in the DataFrame, and checking if their progeny are located below. If not, the positions of the parent and progeny rows in the DataFrame are switched. This process takes a few passes until all the parents and progeny are correctly ordered.

```python
nuclide_list = list(pyne_full.index)
id_list = list(pyne_full.id)
swapping = 1
while swapping >= 1:
    swaps = 0
    for parent in nuclide_list:
        for c, mode, bf in zip(pyne_full.at[parent, "Progeny"],
                               pyne_full.at[parent, "Modes"],
                               pyne_full.at[parent, "Branching_fractions"]):
            if data.decay_const(c) == 0.0 or c not in id_list:
                continue
            j = nuclide_list.index(parent)
            k = id_list.index(c)
            if  j > k:
                nuclide_list[j], nuclide_list[k] = nuclide_list[k], nuclide_list[j]
                id_list[j], id_list[k] = id_list[k], id_list[j]
                pyne_full = pyne_full.reindex(index=nuclide_list)
                swaps +=1
    print("Iteration", swapping, "number of swaps:", swaps)
    swapping += 1
    if swaps == 0: swapping = 0
pyne_full.head(n=10)
```

    Iteration 1 number of swaps: 901
    Iteration 2 number of swaps: 632
    Iteration 3 number of swaps: 425
    Iteration 4 number of swaps: 262
    Iteration 5 number of swaps: 135
    Iteration 6 number of swaps: 53
    Iteration 7 number of swaps: 16
    Iteration 8 number of swaps: 1
    Iteration 9 number of swaps: 0

      id
      Z
      A
      Half-life_s
      Num_decay_modes
      Progeny
      Branching_fractions
      Modes
      Atomic_mass

      Radionuclide

      Fl-285
      1142850000
      114
      285
      0.1500
      1
      [1122810000]
      [1.0]
      [α]
      285.183579

      Cn-285
      1122850000
      112
      285
      34.0000
      1
      [1102810000]
      [1.0]
      [α]
      285.177321

      Fl-284
      1142840000
      114
      284
      0.0025
      0
      []
      []
      []
      284.181344

      Cn-284
      1122840000
      112
      284
      0.1010
      0
      []
      []
      []
      284.174499

      Cn-283
      1122830000
      112
      283
      4.0000
      0
      []
      []
      []
      283.173362

      Cn-282
      1122820000
      112
      282
      0.0005
      0
      []
      []
      []
      282.170668

      Rg-282
      1112820000
      111
      282
      0.5000
      1
      [1092780000]
      [1.0]
      [α]
      282.169405

      Cn-281
      1122810000
      112
      281
      0.1300
      1
      [1102770000]
      [1.0]
      [α]
      281.169641

      Rg-281
      1112810000
      111
      281
      26.0000
      0
      []
      []
      []
      281.166718

      Ds-281
      1102810000
      110
      281
      9.6000
      0
      []
      []
      []
      281.164715

### Now make the dataset files for radioactivedecay
The process of making datasets for radioactivedecay is as follows. We first make the sparse lower triangular matrix *Λ*, which captures the decay relationships and branching fractions between parents and their immediate (first) progeny. We then make the sparse matrix _C_, which is used in decay calculations, and from this make its inverse *C-1*.

First we define some functions used for making *Λ*, _C_ and *C-1*.

```python
def make_lambda_mat(df):
    """Make the lambda matrix and a list of the decay constants."""

    rows = np.array([], dtype=np.int64)
    cols = np.array([], dtype=np.int64)
    values = np.array([], dtype=np.float64)
    lambdas = []
    ln2 = np.log(2)

    nuclide_list = list(df.index)
    id_list = list(df.id)

    for parent in nuclide_list:
        j = nuclide_list.index(parent)
        rows = np.append(rows, [j])
        cols = np.append(cols, [j])
        lambd = ln2/df.at[parent, "Half-life_s"]
        values = np.append(values, -lambd)
        lambdas = np.append(lambdas, lambd)
        for progeny, bf in zip(df.at[parent, "Progeny"], df.at[parent, "Branching_fractions"]):
            if (progeny not in id_list): continue
            i = id_list.index(progeny)
            rows = np.append(rows, [i])
            cols = np.append(cols, [j])
            values = np.append(values, [lambd*bf])

    return sparse.csc_matrix((values, (rows, cols))), lambdas

def prepare_C_inv_C(df):
    """Prepare data structures needed to make C and inv_C."""

    nuclide_list = list(df.index)
    num_nuclides = len(nuclide_list)

    rows_dict = {}
    for i in range(num_nuclides-1, -1, -1):
        a,_ = lambda_mat[:,i].nonzero()
        b = a
        for j in a:
            if j > i:
                b = np.unique(np.concatenate((b,rows_dict[j])))
        rows_dict[i] = b

    rows_C = np.array([], dtype=np.int64)
    cols_C = np.array([], dtype=np.int64)
    for i in range(0, num_nuclides):
        rows_C = np.concatenate((rows_C,rows_dict[i]))
        cols_C = np.concatenate((cols_C,np.array([i]*len(rows_dict[i]))))

    C = sparse.csc_matrix((np.array([0.0]*rows_C.size, dtype=np.float64), (rows_C, cols_C)))
    inv_C = sparse.csc_matrix((np.array([0.0]*rows_C.size, dtype=np.float64), (rows_C, cols_C)))

    return rows_dict, rows_C, cols_C, C, inv_C

def make_C(rows_dict, rows_C, cols_C, C, lambda_mat, df):
    """Calculate C. Report cases of radionuclides with identical or similar half-lives in the same decay chain."""

    nuclide_list = list(df.index)
    for index in range(0, rows_C.size):
        i = rows_C[index]
        j = cols_C[index]
        if i == j: C[i,i] = 1.0
        else:
            sigma = 0.0
            for k in rows_dict[j]:
                if k == i: break
                sigma += lambda_mat[i,k]*C[k,j]
            if lambda_mat[j,j]==lambda_mat[i,i]:
                print("equal decay constants:", nuclide_list[i], nuclide_list[j])
            C[i,j] = sigma/(lambda_mat[j,j]-lambda_mat[i,i])
            if abs((lambda_mat[j,j]-lambda_mat[i,i])/lambda_mat[j,j]) < 1E-4:
                print("rel_diff of decay constants < 1E-4:", nuclide_list[i], nuclide_list[j])
    return C

def make_inv_C(rows_dict, rows_C, cols_C, C, inv_C):
    """Calculate inv_C."""

    for index in range(0, rows_C.size):
        i = rows_C[index]
        j = cols_C[index]
        if i == j: inv_C[i,i] = 1.0
        else:
            sigma = 0.0
            for k in rows_dict[j]:
                if k == i: break
                sigma -= C[i,k]*inv_C[k,j]
            inv_C[i,j] = sigma
    return inv_C
```

The process of making _Λ_, _C_ and *C-1* is complicated as PyNE includes some decay chains where two radionuclides have identical half-lives. PyNE has [special routines](https://pyne.io/theorymanual/decay.html) to cope with this, but radioactivedecay currently does not. Fortunately these cases are limited to some fairly obscure radionuclides which are unlikely to be relevant to most users.

The following is a first pass through at making *Λ* and *C*. It highlights the cases where radionuclides in the same chain have identical half-lives, and also cases where radionuclides in the same chain have similar half-lives (relative difference < 1E-4).

```python
lambda_mat, lambdas = make_lambda_mat(pyne_full)
rows_dict, rows_C, cols_C, C, inv_C = prepare_C_inv_C(pyne_full)
C = make_C(rows_dict, rows_C, cols_C, C, lambda_mat, pyne_full)
```

    equal decay constants: Os-179 Pt-183
    rel_diff of decay constants < 1E-4: Os-179 Pt-183

    <ipython-input-8-266c7a1f1871>:70: RuntimeWarning: divide by zero encountered in double_scalars
      C[i,j] = sigma/(lambda_mat[j,j]-lambda_mat[i,i])
    <ipython-input-8-266c7a1f1871>:67: RuntimeWarning: invalid value encountered in double_scalars
      sigma += lambda_mat[i,k]*C[k,j]

    equal decay constants: Re-168 Ir-172
    rel_diff of decay constants < 1E-4: Re-168 Ir-172
    equal decay constants: Tm-149 Lu-153
    rel_diff of decay constants < 1E-4: Tm-149 Lu-153

So there are radionuclides with identical half-lives in the chains containing 183Pt, 172Ir and 153Lu. These cases withstanding, there are no other chains containing radionuclides with decay constants with a relative difference of less than 1E-4.

It turns out that there is a [bug](https://github.com/pyne/pyne/issues/1342) in PyNE `v0.7.6` causing it to incorrectly calculate decayed activities for the chains containing 183Pt, 172Ir and 153Lu. The bug affects all radionuclides in the chains upwards from these three radionuclides. Because of this bug and the fact that radioactivedecay does not support chains with radionuclides with equal half-lives, we remove the affected radionuclides from the decay dataset. Note even by doing this, decay calculation results are unaffected for chains starting with radionuclides below 183Pt, 172Ir and 153Lu.

This function finds the radionuclides to remove.

```python
def find_affected_radionuclides(nuclide_list, lambda_mat, nuclide):
    """Find radionuclides higher in decay chain than nuclide."""

    s1 = {nuclide_list.index(nuclide)}
    index = 0
    while index < len(nuclide_list):
        s2 = set(lambda_mat.getcol(index).indices)
        if len(s1.intersection(s2)) > 0:
            s2 = set([s for s in list(s2) if s <= index])
            if s2.issubset(s1):
                index += 1
                continue
            s1 = s2.union(s1)
            index = 0
            continue
        index +=1
    return [nuclide_list[nuclide] for nuclide in s1]

nuclide_list = list(pyne_full.index)
affected = find_affected_radionuclides(nuclide_list, lambda_mat, "Pt-183")
print("Radionuclides affected for Pt-183:", affected)
remove = affected

affected = find_affected_radionuclides(nuclide_list, lambda_mat, "Ir-172")
print("Radionuclides affected for Ir-172:", affected)
remove.extend(affected)

affected = find_affected_radionuclides(nuclide_list, lambda_mat, "Lu-153")
print("Radionuclides affected for Lu-153:", affected)
remove.extend(affected)
```

    Radionuclides affected for Pt-183: ['Po-191', 'Pt-183', 'Bi-191', 'Pb-187', 'Tl-187', 'Rn-195', 'At-195', 'Hg-183', 'Au-183']
    Radionuclides affected for Ir-172: ['Pb-180', 'Hg-176', 'Tl-177', 'Pt-172', 'Ir-172']
    Radionuclides affected for Lu-153: ['Lu-153', 'Ta-157']

In total there are 16 radionuclides to be removed from the decay dataset.

```python
pyne_truncated = pyne_full.copy()
pyne_truncated = pyne_truncated.drop(labels=remove)
pyne_truncated.to_csv("pyne_truncated.csv", index=True)
```

Now this is done, we can make the matrices _C_ and *C-1* used by radioactivedecay.

```python
lambda_mat, lambdas = make_lambda_mat(pyne_truncated)
rows_dict, rows_C, cols_C, C, inv_C = prepare_C_inv_C(pyne_truncated)
C = make_C(rows_dict, rows_C, cols_C, C, lambda_mat, pyne_truncated)
inv_C = make_inv_C(rows_dict, rows_C, cols_C, C, inv_C)
```

### Calculate SymPy versions of the matrices for arbitrary-precision calculations

We now calculate SymPy versions of _C_ and *C-1* for arbitrary-precision calculations. First define some functions for processing the data into SymPy objects:

```python
year_sympy = S(36525)/10000

def to_rational(number):
    """
    Converts half-life string to SymPy object.
    """

    if number == '0.0':
        return S(0)

    if 'e' in number or 'E' in number:
        if 'e' in number:
            end = number.split('e')[1]
            number = number.split('e')[0]
        else:
            end = number.split('E')[1]
            number = number.split('E')[0]
        parts = number.split('.')
        if len(parts) == 1: parts.append('')
        if end[0] == '+':
            multiply = 1
            factor = S(10**int(end.lstrip('+')))
        else:
            multiply = 0
            factor = S(10**int(end.lstrip('-')))
        denom = S(10**len(parts[1]))
        parts[0] = parts[0].lstrip('0')
        if len(parts[0]) == 0: parts[1] = parts[1].lstrip('0')
        if multiply == 1:
            return S(parts[0]+parts[1])*factor/denom
        else: return S(parts[0]+parts[1])/(denom*factor)
    parts = number.split('.')
    if len(parts) == 1: parts.append('')
    denom = S(10**len(parts[1]))
    parts[0] = parts[0].lstrip('0')
    if len(parts[0]) == 0: parts[1] = parts[1].lstrip('0')
    return S(parts[0]+parts[1])/denom
```

Now make a SymPy version of the *Λ* matrix:

```python
num_nuclides = len(pyne_truncated)
lambda_mat_sympy = SparseMatrix.zeros(num_nuclides, num_nuclides)
lambdas_sympy = Matrix.zeros(num_nuclides, 1)
nuclide_list = list(pyne_truncated.index)
id_list = list(pyne_truncated.id)
masses_sympy = Matrix.zeros(num_nuclides, 1)

for parent in nuclide_list:
    j = nuclide_list.index(parent)
    hl_sympy = to_rational(str(pyne_truncated.at[parent, "Half-life_s"]))
    lambd = log(2)/hl_sympy
    lambda_mat_sympy[j, j] = -lambd
    lambdas_sympy[j] = lambd
    for progeny, bf in zip(pyne_truncated.at[parent, "Progeny"], pyne_truncated.at[parent, "Branching_fractions"]):
        if (progeny not in id_list): continue
        i = id_list.index(progeny)
        lambda_mat_sympy[i, j] = lambd*to_rational(str(bf))
    masses_sympy[j] = to_rational(str(pyne_truncated.at[parent, "Atomic_mass"]))
```

Now make a SymPy version of the _C_ and *C-1* matrix:

```python
c_sympy = SparseMatrix.zeros(num_nuclides, num_nuclides)
c_inv_sympy = SparseMatrix.zeros(num_nuclides, num_nuclides)

for index in range(0, rows_C.size):
    i = rows_C[index]
    j = cols_C[index]
    if i == j: c_sympy[i, i] = Integer(1)
    else:
        sigma = Integer(0)
        for k in rows_dict[j]:
            if k == i: break
            sigma += lambda_mat_sympy[i, k]*c_sympy[k, j]
        c_sympy[i, j] = sigma/(lambda_mat_sympy[j, j]-lambda_mat_sympy[i, i])

for index in range(0, rows_C.size):
    i = rows_C[index]
    j = cols_C[index]
    if i == j: c_inv_sympy[i, i] = Integer(1)
    else:
        sigma = Integer(0)
        for k in rows_dict[j]:
            if k == i: break
            sigma -= c_sympy[i, k]*c_inv_sympy[k, j]
        c_inv_sympy[i, j] = sigma
```

### Save the outputs

Write output files containing _C_ and *C-1* in SciPy and SymPy sparse format, and other files needed to create a dataset suitable for radioactive decay `v0.4.0+`.

```python
hldata = np.array([(np.float64(hl), 's', str(hl) + ' s') for hl in pyne_truncated["Half-life_s"]], dtype=object)

prog_bfs_modes = np.array([{}]*len(pyne_truncated.index))
i = 0
for parent in list(pyne_truncated.index):
    progeny = [add_hyphen(nucname.name(id)) for id in pyne_truncated.at[parent, "Progeny"]]
    bfs = dict(zip(progeny, pyne_truncated.at[parent, "Branching_fractions"]))
    modes = dict(zip(progeny, pyne_truncated.at[parent, "Modes"]))
    bfs = {key: value for key, value in sorted(bfs.items(), key=lambda x: x[1], reverse=True)}
    prog_bfs_modes[i] = {progeny: [bf, modes[progeny]] for progeny, bf in bfs.items()}
    i += 1

np.savez_compressed("./decay_data.npz", radionuclides=np.array(nuclide_list),
                    masses=np.array(list(pyne_truncated["Atomic_mass"])),
                    hldata=hldata, prog_bfs_modes=prog_bfs_modes,
                    year_conv=365.25)

# Write out SciPy sparse matrices (convert to CSR format)
sparse.save_npz("./c_scipy.npz", C.tocsr())
sparse.save_npz("./c_inv_scipy.npz", inv_C.tocsr())

import pkg_resources, sympy
if pkg_resources.parse_version(sympy.__version__) >= pkg_resources.parse_version('1.9'):
    pickle_type = '1.9'
else:
    pickle_type = '1.8'

# Write out SymPy objects to pickle files
with open(f"c_sympy_{pickle_type}.pickle", "wb") as outfile:
    outfile.write(pickle.dumps(c_sympy))
with open(f"c_inv_sympy_{pickle_type}.pickle", "wb") as outfile:
    outfile.write(pickle.dumps(c_inv_sympy))
with open(f"atomic_masses_sympy_{pickle_type}.pickle", "wb") as outfile:
    outfile.write(pickle.dumps(masses_sympy))
with open(f"decay_consts_sympy_{pickle_type}.pickle", "wb") as outfile:
    outfile.write(pickle.dumps(lambdas_sympy))
with open(f"year_conversion_sympy_{pickle_type}.pickle", "wb") as outfile:
    outfile.write(pickle.dumps(year_sympy))
```


```python
#@markdown Configuración Inicial
from IPython.utils import io

from google.colab.data_table import DataTable
from IPython.display import Javascript
from IPython.display import display, Markdown, Latex
from IPython.display import Audio, Image
from IPython.display import IFrame, HTML
with io.capture_output() as capt:
  !pip install gradio

import gradio as gr
import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
import random

def dLatex(self):
  return display(Latex(self))

def dMarkdown(self):
  return display(Markdown(self))

# Avoids scroll-in-the-scroll in the entire Notebook
# https://stackoverflow.com/a/66891328
def resize_colab_cell():
  display(Javascript(
      'google.colab.output.setIframeHeight(0, true, {maxHeight: 5000})'
      ))
get_ipython().events.register('pre_run_cell', resize_colab_cell)

def dLatex(self):
  return display(Latex(self))

def dMarkdown(self):
  return display(Markdown(self))
```

# Deriva Instrumental

Deriva *(drift)*: cambios con el tiempo en la lectura de los gravímetros en una estación

**Causas:**
* Constante elástica del resorte (resortes y fibras de torsión no son perfectamente elásticos)
* Temperatura
* Mareas terrestres

Al repetir una medición en el mismo lugar y bajo las mismas circunstancias, no  se obtiene el mismo valor (variaciones de aprox. $0.001\,\mathrm{mGal}$)

**¿De qué depende la deriva?**
* Tipo de instrumento (varía de un instrumento a otro)
* Características del resorte
* Edad y uso del gravímetro
* Fluctuaciones de temperatura externa
* Tipo de transporte
* Cambios de presión
* Cambios de voltaje en la batería
* Valores de deriva suelen ser menores a $1\,\mathrm{mGal}$

**Deriva según el gravímetro:**

* Gravímetros LC&R:
  * Instrumentos nuevos en promedio menor a $1\,\mathrm{mGal/mes}$
  * en un equipo con algunos años de uso, deriva menor a $0.5\,\mathrm{mGal/mes}$ (Manual LC&R)
* Gravímetros Scintrex: Modelo CG5, deriva residual menor a $20\,\mu\mathrm{Gal/dia}$
  * El instrumento realiza correcciones automáticas
* Gravímetros Burries deriva menor a $0.3\,\mathrm{mGal/mes}$ (Jentzsch et al., 2017)

**¿Cómo se corrige la deriva?**
* Realizar mediciones en la estación base con diferencia
de una o dos horas
* Graficar las lecturas en función del tiempo
* Por ello, primera y última medición del día se realizan
en la estación base

\begin{align}
d
&=
\frac{L_{f}-L_{i}}{t_{f}-t_{i}}
\\
{L^{c}}_{j}
&=
L_{j}-d\left(t_{j}-t_{A i}\right)
\end{align}

```python
#@markdown Tabla
pd.options.display.float_format = '{:.2f}'.format
df0 = pd.DataFrame(columns=["ESTACIÓN", "Hora", "Lectura"],
                data=[[ 1, "08:00", 520],
                      [ 2, "08:30", 490],
                      [ 3, "09:05", 450],
                      [ 4, "09:40", 420],
                      [ 1, "10:10", 522],
                      [ 5, "10:45", 470],
                      [ 6, "11:20", 510],
                      [ 7, "11:55", 580],
                      [ 4, "12:40", 424],
                      [ 8, "13:40", 570],
                      [ 9, "14:15", 650],
                      [10, "14:50", 510],
                      [ 7, "15:15", 575]]
                )

@np.vectorize
def hour2min(hour):
  h, m = hour.split(':')
  return 60*int(h)+int(m)

df = df0.copy()
df['Minutos'] = hour2min(df0['Hora'])
df['Minutos'] = df['Minutos']-min(df['Minutos'])
df = df[['ESTACIÓN', 'Minutos', 'Lectura']]

display(df0.set_index("ESTACIÓN").T)
```

    <IPython.core.display.Javascript object>

      ESTACIÓN
      1
      2
      3
      4
      1
      5
      6
      7
      4
      8
      9
      10
      7

      Hora
      08:00
      08:30
      09:05
      09:40
      10:10
      10:45
      11:20
      11:55
      12:40
      13:40
      14:15
      14:50
      15:15

      Lectura
      520
      490
      450
      420
      522
      470
      510
      580
      424
      570
      650
      510
      575

```python
k = 0.728 #@param {type:"slider", min:0.5, max:1, step:0.001}
g_0 = 951253 #@param {type:"slider", min:9e5, max:10e5, step:1}
```

    <IPython.core.display.Javascript object>

```python
k   = 0.882
g_0 = 979166.3
```

    <IPython.core.display.Javascript object>

```python
def find_loop(df, base, drop_index=None):
  dfd = df.loc[df['ESTACIÓN'] == base]
  min, max = tuple(dfd.index)[:2]
  if drop_index==None:
    return df.loc[min:max]
  else:
    return df.loc[min:max].drop(drop_index)

def rev_loop(dfv, g):
  ind = dfv.index
  delta_l, delta_m = tuple(dfv[['Lectura','Minutos']].\
                           iloc[[0,-1]].diff().iloc[-1])
  d = delta_l/delta_m
  dfv['Lectura_rev'] = [dfv['Lectura'][i] -
                        d * (dfv['Minutos'][i] -
                             dfv['Minutos'][min(ind)])
                        for i in ind]
  dfv['delta_g'] = k * dfv['Lectura_rev'].diff().fillna(0)
  dfv['g'] = g + dfv['delta_g'].cumsum()
  return d, dfv

df_dup = df.loc[df['ESTACIÓN'].duplicated()]

bases = df_dup['ESTACIÓN'].to_list()
dup_index = [None] + df_dup.index.to_list()[:-1]

print('    bases:', bases)
print('dup_index:', dup_index)

dfs = [find_loop(df, base, dup) for base, dup in zip(bases, dup_index)]
for dfi in dfs:
  display(dfi)

fig, ax = plt.subplots()
for l, dfi in enumerate(dfs):
  dfi.plot(
      x='Minutos',
      y='ESTACIÓN',
      style='o-',
      alpha=.75,
      label=f'Loop {l:d}',
      zorder=len(dfs)-l,
      ax=ax)
ax.set_yticks(range(1+max(df['ESTACIÓN'])))
ax.grid(axis='y')
plt.show()
```

```python
df_0 = dfs[0].copy()
d_1, df_1 = rev_loop(df_0, g_0)
dMarkdown(f'$d_0={d_1:g}$')
display(df_0)
```

    <IPython.core.display.Javascript object>

$d_0=0.0153846$

      ESTACIÓN
      Minutos
      Lectura
      Lectura_rev
      delta_g
      g

      0
      1
      0
      520
      520.00
      0.00
      979166.30

      1
      2
      30
      490
      489.54
      -26.87
      979139.43

      2
      3
      65
      450
      449.00
      -35.75
      979103.68

      3
      4
      100
      420
      418.46
      -26.93
      979076.74

      4
      1
      130
      522
      520.00
      89.56
      979166.30

```python
lrev_1, g_1 = df_0[['Lectura_rev', 'g']].iloc[-2]
print(lrev_1, g_1, sep='\n')

df_1 = dfs[1].copy()
df_1.loc[min(df_1.index), 'Lectura'] = lrev_1

d_1, df_1 = rev_loop(df_1, g_1)
dMarkdown(f'$d_1={d_1:g}$')
display(df_1)
```

    <IPython.core.display.Javascript object>

    418.46153846153845
    979076.7430769231

$d_1=0.0307692$

      ESTACIÓN
      Minutos
      Lectura
      Lectura_rev
      delta_g
      g

      3
      4
      100
      418.46
      418.46
      0.00
      979076.74

      5
      5
      165
      470.00
      468.00
      43.69
      979120.44

      6
      6
      200
      510.00
      506.92
      34.33
      979154.77

      7
      7
      235
      580.00
      575.85
      60.79
      979215.56

      8
      4
      280
      424.00
      418.46
      -138.81
      979076.74

```python
lrev_2, g_2 = df_1[['Lectura_rev', 'g']].iloc[-2]
print(lrev_2, g_2, sep='\n')

df_2 = dfs[2].copy()
df_2.loc[min(df_2.index), 'Lectura'] = lrev_2

d_2, df_2 = rev_loop(df_2, g_2)
dMarkdown(f'$d_2={d_2:g}$')
display(df_2)
```

    <IPython.core.display.Javascript object>

    575.8461538461538
    979215.5563076923

$d_2=-0.00423077$

      ESTACIÓN
      Minutos
      Lectura
      Lectura_rev
      delta_g
      g

      7
      7
      235
      575.85
      575.85
      0.00
      979215.56

      9
      8
      340
      570.00
      570.44
      -4.76
      979210.79

      10
      9
      375
      650.00
      650.59
      70.69
      979281.48

      11
      10
      410
      510.00
      510.74
      -123.35
      979158.13

      12
      7
      435
      575.00
      575.85
      57.42
      979215.56

```python
df_rev = pd.concat([
                    df_0.iloc[:-1],
                    df_1.iloc[1:-1],
                    df_2.iloc[1:-1]
                    ], ignore_index=True
                   )

df_rev['bad'] = (k * df_rev['Lectura_rev'].diff().fillna(0) !=
                     df_rev['delta_g'])
display(df_rev)
print('Total bad:', df_rev['bad'].sum())
```

    <IPython.core.display.Javascript object>

      ESTACIÓN
      Minutos
      Lectura
      Lectura_rev
      delta_g
      g
      bad

      0
      1
      0
      520.00
      520.00
      0.00
      979166.30
      False

      1
      2
      30
      490.00
      489.54
      -26.87
      979139.43
      False

      2
      3
      65
      450.00
      449.00
      -35.75
      979103.68
      False

      3
      4
      100
      420.00
      418.46
      -26.93
      979076.74
      False

      4
      5
      165
      470.00
      468.00
      43.69
      979120.44
      False

      5
      6
      200
      510.00
      506.92
      34.33
      979154.77
      False

      6
      7
      235
      580.00
      575.85
      60.79
      979215.56
      False

      7
      8
      340
      570.00
      570.44
      -4.76
      979210.79
      False

      8
      9
      375
      650.00
      650.59
      70.69
      979281.48
      False

      9
      10
      410
      510.00
      510.74
      -123.35
      979158.13
      False

    Total bad: 0

```python
df_drop = df.drop_duplicates(subset='ESTACIÓN',
                             ignore_index=True)
df_rev['rev_diff'] = (df_drop['Lectura']-df_rev['Lectura_rev'])

display(pd.concat([
                   df_drop[['ESTACIÓN','Lectura']],
                   df_rev[['ESTACIÓN','Lectura_rev']]
                   ], axis=1
                  )
)

fig, ax = plt.subplots()
df_rev.plot(
    x='Minutos',
    y='rev_diff',
    ax=ax)
plt.show()

fig, ax = plt.subplots()
df_rev.plot(
    x='Minutos',
    y='g',
    ax=ax)
plt.show()
```


# `rlplay`-ing around with the Double-Dee QN

```python
import torch
import numpy

import matplotlib.pyplot as plt
%matplotlib inline

import gym
```

See example.ipynb for the overview of `rlplay`

## Tabular CartPole with Q-learning

### The environment

A version of the Taxi Environment that disassembles the observation into pieces

```python
from gym.spaces import Dict, Discrete

class StructuredTaxiEnv(gym.ObservationWrapper):
    def __init__(self, env):
        super().__init__(env)
        self.observation_space = Dict(dict(
            row=Discrete(5),
            col=Discrete(5),
            location=Discrete(5),
            destination=Discrete(4)
        ))

    @staticmethod
    def observation(n):
        # observation = r x c x p x d \in 5 x 5 x 5 x 4
        # n, d = divmod(observation, 4)
        # n, p = divmod(n, 5)
        # r, c = divmod(n, 5)
        d = n % 4
        n = n // 4
        p = n % 5
        n = n // 5
        c = n % 5
        r = n // 5
        return dict(row=r, col=c, location=p, destination=d)
```

A wrapper that makes Taxi's `.render` compatible with `core_evaluate`

```python
class RenderPatch(gym.Wrapper):
    def render(self, mode='human'):
        result = self.env.render(mode)
        if result is None:
            return True

        return result
```

The environment factory

```python
from gym.envs.toy_text import TaxiEnv

def base_factory(struct=False, seed=None):
    env = gym.make("Taxi-v3")
    if struct:
        env = StructuredTaxiEnv(env)
    return RenderPatch(env)
```

### the Actor

A procedure and a layer, which converts the input integer data into its
little-endian binary representation as float $\{0, 1\}^m$ vectors.

```python
def onehotbits(input, n_bits=63, dtype=torch.float):
    """Encode integers to fixed-width binary floating point vectors"""
    assert not input.dtype.is_floating_point
    assert 0 < n_bits < 64  # torch.int64 is signed, so 64-1 bits max

    # n_bits = {torch.int64: 63, torch.int32: 31, torch.int16: 15, torch.int8 : 7}

    # get mask of set bits
    pow2 = torch.tensor([1 << j for j in range(n_bits)]).to(input.device)
    x = input.unsqueeze(-1).bitwise_and(pow2).to(bool)

    # upcast bool to float to get one-hot
    return x.to(dtype)

class OneHotBits(torch.nn.Module):
    def __init__(self, n_bits=63, dtype=torch.float):
        assert 1 <= n_bits < 64
        super().__init__()
        self.n_bits, self.dtype = n_bits, dtype

    def forward(self, input):
        return onehotbits(input, n_bits=self.n_bits, dtype=self.dtype)
```

A special module dictionary, which applies itself to the input dict of tensors

```python
from typing import Optional, Mapping
from torch.nn import Module, ModuleDict as BaseModuleDict

class ModuleDict(BaseModuleDict):
    """The ModuleDict, that applies itself to the input dicts."""
    def __init__(
        self,
        modules: Optional[Mapping[str, Module]] = None,
        dim: Optional[int]=-1
    ) -> None:
        super().__init__(modules)
        self.dim = dim

    def forward(self, input):
        # enforce concatenation in the order of the declaration in  __init__
        return torch.cat([
            m(input[k]) for k, m in self.items()
        ], dim=self.dim)
```

A simple tabular q-learner actor

```python
from rlplay.engine import BaseActorModule

class TaxiActor(BaseActorModule):
    def __init__(self, duelling=False, struct=False, epsilon=0.1):
        super().__init__()
        self.duelling = duelling

        # the table is actually an embedding of the state index
        n_dim = 6 + (1 if duelling else 0)
        if not struct:
            self.table = torch.nn.Embedding(500, n_dim)

        else:
            self.table = torch.nn.Sequential(
                ModuleDict(dict(
                    row=torch.nn.Embedding(5, 4),
                    col=torch.nn.Embedding(5, 4),
                    location=torch.nn.Embedding(5, 4),
                    destination=torch.nn.Embedding(4, 4),
                )),
                torch.nn.Linear(16, n_dim)
            )

        # for updating the exploration epsilon in the clones
        self.register_buffer('epsilon', torch.tensor(epsilon))

    def forward(self, obs, act, rew, fin, *, hx=None, stepno=None, virtual=False):
        qv, hx = self.table(obs), ()
        if self.duelling:
            # dueling net [Wang et al. (2016)](https://arxiv.org/abs/1511.06581)
            val, adv= torch.split_with_sizes(qv, [1, 6], dim=-1)
            qv = adv + (val - adv.mean(dim=-1, keepdim=True))

            val, actions = val.squeeze(-1), qv.max(dim=-1).indices

        else:
            val, actions = qv.max(dim=-1)

        if self.training:
            *head, n_actions = qv.shape
            actions = actions.where(
                torch.rand(head).gt(self.epsilon),
                torch.randint(n_actions, size=head))

        return actions, hx, dict(q=qv, value=val)
```

### D-DQN loss

Service functions for the algorithms

```python
from plyr import apply, suply, xgetitem

def timeshift(state, *, shift=1):
    """Get current and shifted slices of nested objects."""
    # use `xgetitem` to let None through
    # XXX `curr[t]` = (x_t, a_{t-1}, r_t, d_t), t=0..T-H
    curr = suply(xgetitem, state, index=slice(None, -shift))

    # XXX `next[t]` = (x_{t+H}, a_{t+H-1}, r_{t+H}, d_{t+H}), t=0..T-H
    next = suply(xgetitem, state, index=slice(shift, None))

    return curr, next
```

Double DQN loss for contiguous trajectory fragments.

* `.state[t+1].rew` -- $r_{t+1}$
* `.state[t+1].fin` -- $d_{t+1}$
* `.state[t+1].act` -- $a_t$ which caused $
    (s_t, z_t, a_t) \longrightarrow (s_{t+1}, z_{t+1}, r_{t+1}, d_{t+1})
$
* `_, _, fragment.actor[t] = actor(.state[t])` -- $
    q(z_t, h_t, \cdot; \theta_{\text{old}})
$ -- used for rollout collection
* `_, _, info_module[t] = module(.state[t])` -- $
    q(z_t, h_t, \cdot; \theta)
$ -- the current q-function, producing $(h_t)_{t=0}^{T+1}$
* `_, _, info_target[t] = target(.state[t])` -- $
    q(z_t, h_t, \cdot; \theta_-)
$ -- the q-target with $(h^-_t)_{t=0}^{T+1}$

The current q-network minimizes the $\mathrm{TD}(0)$-error is

$$
\delta_t(\theta)
    = \bigl(
        r_{t+1}
        + \gamma 1_{\{\neg d_{t+1}\}} v^*(z_{t+1})
    \bigr) - q(z_t, h_t, a_t; \theta)
    \,, $$

where the approximate state value estimate $
    v^*(z_{t+1})
$ is one of
* `Q-learning`: $
    \max_a q(z_{t+1}, h_{t+1}, a; \theta)
$
* `DQN`: $
    \max_a q(z_{t+1}, h_{t+1}, a; \theta_-)
$
* `double DQN`: $
    q(z_{t+1}, h_{t+1}, \hat{a}_{t+1}; \theta_-)
$ for $
    \hat{a}_{t+1} = \arg \max_a Qq(z_{t+1}, h_{t+1}, a; \theta)
$

One of the works related to DQN learning of recurrent agents is [Kapturowski et al. (2018)](https://openreview.net/forum?id=r1lyTjAqYX), who propose to use a burn-in period in
the contiguous trajectory fragment in order to compensate for the representation drift,
due to different RNN parameters $\theta_-$ and $\theta$.

There is no clear-cut evidence suggesting that the hidden recurrent sequences $h_t$
and $h^-_t$ yield significantly different results.

```python
import torch.nn.functional as F

# @torch.enable_grad()
def ddq_learn(fragment, module, *, gamma=0.95, target=None, double=False):
    r"""Compute the Double-DQN loss over a _contiguous_ fragment of a trajectory.

    Details
    -------
    In Q-learning the action value function minimizes the TD-error

    $$
        r_{t+1}
            + \gamma 1_{\neg d_{t+1}} v^*(z_{t+1})
            - q(z_t, a_t; \theta)
        \,, $$

    w.r.t. Q-network parameters $\theta$ where $z_t$ is the actionable state,
    $r_{t+1}$ is the reward for $s_t \to s_{t+1}$ transition. The value of
    $z_t$ include the current observation $x_t$ and the recurrent state $h_t$,
    the last action $a_{t-1}$, the last reward $r_t$, and termination flag
    $d_t$.

    In the classic Q-learning there is no target network and the next state
    optimal state value function is bootstrapped using the current Q-network
    (`module`):

    $$
        v^*(z_{t+1})
            \approx \max_a q(z_{t+1}, a; \theta)
        \,. $$

    The DQN method, proposed by

        [Minh et al. (2013)](https://arxiv.org/abs/1312.5602),

    uses a secondary Q-network (`target`) to estimate the value of the next
    state:

    $$
        v^*(z_{t+1})
            \approx \max_a q(z_{t+1}, a; \theta^-)
        \,, $$

    where $\theta^-$ are the frozen parameters of the target Q-network. The
    Double DQN algorithm of

        [van Hasselt et al. (2015)](https://arxiv.org/abs/1509.06461)

    unravels the $\max$ operator as
    $
        \max_k u_k \equiv u_{\arg \max_k u_k}
    $
    and replaces the outer $u$ with the Q-values of the target Q-network, while
    computing the inner $u$ (inside the $\arg\max$) with the current Q-network.
    Specifically, the Double DQN value estimate is

    $$
        v^*(z_{t+1})
            \approx q(z_{t+1}, \hat{a}_{t+1}; \theta^-)
            \,,
            \hat{a}_{t+1}
                = \arg \max_a q(z_{t+1}, a; \theta)
        \,, $$

    for $
        \hat{a}_{t+1}
            = \arg \max_a q(s_{t+1}, a; \theta)
    $ being the action taken by the current Q-network $\theta$ at $z_{t+1}$.

    Recurrent DQN
    -------------
    The key problem with the recurrent state $h_t$ in $z_t$ is its representation
    drift: the endogenous states used for collecting trajectory data during the
    rollout are produced by an actor with stale parameters $\theta_{\text{old}}$,
    and thus might have high discrepancy with the recurrent state produced by
    the current Q-network $\theta$ or the target $\theta-_$. To mitigate this

        [Kapturowski et al. (2018)](https://openreview.net/forum?id=r1lyTjAqYX)

    proposed to spend a slice `burn-in` of the recorded trajectory on
    aligning the recurrent representation. Specifically, starting with $h_0$
    (contained in `fragment.hx`) they propose to launch two sequences $h_t$
    and $h^-_t$ from the same $h^-_0 = h_0$ using $q(\cdot; \theta)$ and
    $q(\cdot; \theta^-)$, respectively.
    """

    trajectory, hx = fragment.state, fragment.hx
    obs, act, rew, fin = trajectory.obs, trajectory.act, trajectory.rew, trajectory.fin

    # get $Q(z_t, h_t, \cdot; \theta)$ for all t=0..T
    _, _, info_module = module(
        obs, act, rew, fin, hx=hx, stepno=trajectory.stepno)

    # get the next state `state[t+1]` $z_{t+1}$ to access $a_t$
    state_next = suply(xgetitem, trajectory, index=slice(1, None))

    # $\hat{A}_t$, the module's response to current and next state,
    #  contains the q-values. `curr` is $q(z_t, h_{t+1}, \cdot; \theta)$
    #  and `next` is $q(z_{t+1}, h_{t+1}, \cdot; \theta)$ is `next`.
    info_module_curr, info_module_next = timeshift(info_module)

    # get $q(z_t, h_t, a_t; \theta)$ for all t=0..T-1
    q_replay = info_module_curr['q'].gather(-1, state_next.act.unsqueeze(-1))

    # get $\hat{v}_{t+1}(z_{t+1}) = ...$
    with torch.no_grad():
        if target is None:
            # get $... = \max_a Q(z_{t+1}, h_{t+1}, a; \theta)$
            q_value = info_module_next['q'].max(dim=-1, keepdim=True).values

        else:
            _, _, info_target = target(
                obs, act, rew, fin, hx=hx, stepno=trajectory.stepno)

            info_target_next = suply(xgetitem, info_target, index=slice(1, None))
            if not double:
                # get $... = \max_a Q(z_{t+1}, h^-_{t+1}, a; \theta^-)$
                q_value = info_target_next['q'].max(dim=-1, keepdim=True).values

            else:
                # get $\hat{a}_{t+1} = \arg \max_a Q(z_{t+1}, h_{t+1}, a; \theta)$
                hat_act = info_module_next['q'].max(dim=-1).indices.unsqueeze(-1)

                # get $... = Q(z_{t+1}, h^-_{t+1}, \hat{a}_{t+1}; \theta^-)$
                q_value = info_target_next['q'].gather(-1, hat_act)

        # get $r_{t+1} + \gamma 1_{d_{t+1}} \hat{v}_{t+1}(z_{t+1})$ using inplace ops
        q_value.masked_fill_(state_next.fin.unsqueeze(-1), 0.)
        q_value.mul_(gamma).add_(state_next.rew.unsqueeze(-1))

    # td-error ell-2 loss
    return F.mse_loss(q_replay, q_value, reduction='sum'), {}
```

The following is my incomplete take on [Bellemare et al. (2017)](http://proceedings.mlr.press/v70/bellemare17a.html)

Then the distributional Bellman operator on a policy $\pi_\theta$ at $(s, a)$ is
defined as a random variable on $\mathbb{R}$ with the following law

$$
(T Q_\theta)(s, a)
    \overset{D}{=}
        r + \gamma Q_\theta(s', a')
    \,, \text{ for }
    (r, s') \sim p(r, s'\mid s, a)
    \,, \text{ and }
    a' \sim \pi_\theta(a \mid s')
    \,.
$$

Formally, this means the following (omitting the dependency of on $(s, a)$):

\begin{align}
  \mathbb{P}\bigl(
    T Q_\theta \in U
  \bigr)
    &= \mathbb{P}\bigl(
      R + \gamma Q_\theta(S', A') \in U
    \bigr)
    \\
    &= \mathbb{E}_{r, s'\sim p(r, s'\mid s, a)}
        \mathbb{E}_{a'\sim \pi_\theta(a \mid s')}
          \mathbb{P}\bigl(
            R + \gamma Q_\theta(S', A') \in U
            \,\big \vert\, S'=s, R=r, A'=a'
          \bigr)
    \\
    &= \mathbb{E}_{r, s'\sim p(r, s'\mid s, a)}
        \mathbb{E}_{a'\sim \pi_\theta(a \mid s')}
          \mathbb{P}\bigl(
            Q_\theta(s', a') \in \frac{U - r} \gamma
          \bigr)
    \,,
\end{align}

or, in other words, $T$ acts on the conditional distribution $
    Q_\theta(\cdot \mid s, a)
$ thus: for any bounded Borel measurable $f$

\begin{align}
  \int (T Q_\theta)(dv \mid s, a) f(v)
    &= \int
        p(ds', dr\mid s, a)  % clump p and R into one cond-distrib
        \pi_\theta(da' \mid s')
        Q_\theta(dv\mid s', a')
        f(r + \gamma v)
    \\
    &=  % Fubini
        \mathbb{E}_{r, s'\sim p(r, s'\mid s, a)}
        \mathbb{E}_{a'\sim \pi_\theta(a \mid s')}
        \mathbb{E}_{v\sim Q_\theta(v\mid s', a')}
            f(r + \gamma v)
    \,.
\end{align}

Thus suggests several approximations, depending on the tractability of the inner expectations.

Ideally we would like to find such distribution, that
$$
(T Q_\theta)(s, a)
    \overset{D}{=}
    Q_\theta(s, a)
    \,, $$
for all or most $s$ and $a$.

If $f$ is $L$-Lipschitz, then
\begin{align}
\biggl\lvert
    \int \bigl(
        (T H)(dv \mid s, a) - (T G)(dv \mid s, a)
        \bigr) f(v)
    \biggr\rvert
    &\leq \int
        p(ds', dr\mid s, a)
        \pi_\theta(da' \mid s')
        H(dh\mid s', a')
        G(dg \mid s', a')
        \bigl\lvert
            f(r + \gamma h) - f(r + \gamma g)
        \bigr\rvert
    \\
    &\leq L \gamma \int
        p(ds', dr\mid s, a)
        \pi_\theta(da' \mid s')
        H(dh \mid s', a')
        G(dg \mid s', a')
        \bigl\lvert h - g \bigr\rvert
    \\
    &\leq L \gamma \int
        p(ds', dr\mid s, a)
        \pi_\theta(da' \mid s')
        \int
            H(dh \mid s', a')
            G(dg \mid s', a')
            \bigl\lvert h - g \bigr\rvert
    \\
    &\leq L \gamma \sup_{s, a}
        \int
            H(dh \mid s, a)
            G(dg \mid s, a)
            \bigl\lvert h - g \bigr\rvert
    \,,
\end{align}

so there is hope for contraction or even convergence, however, it is better to
consult [Bellemare et al. (2017)](http://proceedings.mlr.press/v70/bellemare17a.html)
first.

Below is the loss for a version of the distributional DQN, specific to Gaussian approximations
and maxent regularization.

We minimize (in expectation over transition $s, a \to s'$) the Kullback-Leibler
divergence between distributions $
    (T Q_\theta)(dv \mid s, a)
$ and $
    Q_\theta(dv \mid s, a)
$, while at the same time encouraging $
    Q_\theta(dv \mid s, a)
$ to exhibit high entropy. The actions $a' \sim \pi_\theta(a\mid s')$
are sampled greedily as $
    \arg \max_a \int Q_\theta(dv \mid s', a) v
$ (assuming finite action space).

```python
def gdqn(fragment, module, *, gamma, C=1e-3):
    trajectory, hx = fragment.state, fragment.hx
    obs, act, rew, fin = trajectory.obs, trajectory.act, trajectory.rew, trajectory.fin

    # get $Q(z_t, h_t, \cdot; \theta)$ for all t=0..T
    act, _, info_module = module(
        obs, act, rew, fin, hx=hx, stepno=trajectory.stepno)

    # get the next state `state[t+1]` $z_{t+1}$ to access $a_t$
    state_next, act = suply(xgetitem, (trajectory, act.unsqueeze(-1)), index=slice(1, None))

    # $\hat{A}_t$, the module's response to current and next state,
    #  contains the q-values. `curr` is $q(z_t, h_{t+1}, \cdot; \theta)$
    #  and `next` is $q(z_{t+1}, h_{t+1}, \cdot; \theta)$ is `next`.
    info_module_curr, info_module_next = timeshift(info_module)

    # get T Z_\theta(\cdot \mid z_t, a_t)
    #  \overset{=}{D} N(\cdot \mid r_{t+1} + \gamma 1_{d_{t+1}} \hat{v}_{t+1}(z_{t+1}), ...)
    # use action a' \sim \arg \max_a z_\theta(a \mid z_{t+1}, a_{t+1})
    factor = torch.where(state_next.fin, 1e-4, gamma)  # XXX some leakage from the next traj.
    with torch.no_grad():
        # get $r_{t+1} + \gamma 1_{d_{t+1}} \hat{v}_{t+1}(z_{t+1})$
        zed = info_module_next['loc'].gather(-1, act).squeeze(-1)
        loc = state_next.rew + factor * zed
        # XXX should the termination does not affect the scale?
        scl = factor * info_module_next['scl'].gather(-1, act).squeeze(-1)
        q = Normal(loc, scl)

    # get Z_\theta(\cdot \mid s_t, a_t)
    act = state_next.act.unsqueeze(-1)
    p = Normal(
        info_module_curr['loc'].gather(-1, act).squeeze(-1),
        info_module_curr['scl'].gather(-1, act).squeeze(-1),
    )

    # distance from T z_\theta to z_\theta on s-a-r-s' samples
    dst = kl_divergence(q, p).mean()  # q is fixed target.
    # XXX fwd kl means covering, not mode-seeking.

    # use kl from N(0, 1) prior, but can we also use entropy?
    # pi = Normal(*torch.tensor([0., 1.])).expand(loc.shape)
    # kl_reg = kl_divergence(p, pi).mean()
    reg = -p.entropy().mean()

    return dst + C * reg, {
        'dst': float(dst),
        'reg': float(reg),
    }
```

Let's define an actor for the distributional Q-learning with the Gaussian approximation:

$$
Q_\theta(
    v \mid z, a
    ) = \mathcal{N}\bigl(
        v \,\big\vert \,
        \mu(z, a),
        \sigma^2(z, a)
    \bigr)
    \,. $$

The greedy action sampling must take into account the distribution of $
    j^* = \arg\max_i z_i
$ for a collection of independent $
    z_i \sim \mathcal{N}(Z_i \mid \mu_i, \sigma^2_i)
$, not just the means. Specifically, we must make sure that $j^*=j$ with probability

$$
\mathbb{P}(Z_j \geq  \max_{i\neq j} Z_i)
    = \mathbb{E}
        \mathbb{P}(Z_j \geq \max_{i\neq j} Z_i \big \vert Z_j)
    = \mathbb{E}_{\xi \sim \mathcal{N}(0, 1)}
        \prod_{i\neq j} \Phi\biggl(
            \frac{\mu_j - \mu_i}{\sigma_i}
            + \frac{\sigma_j}{\sigma_i} \xi
        \biggr)
    \,, $$

wherein we used independence and the fact that the max of independent Gaussian rvs satisfies

$$
\mathbb{P}(\max_i Z_i \leq c)
    = \prod_i \mathbb{P}(Z_i \leq c)
    = \prod_i \Phi\biggl(
        \frac{c - \mu_i}{\sigma_i}
    \biggr)
    \,, $$

for $\Phi$ -- the univariate standard Gaussian CDF.

But this is hard to compute... So we use the largest mean heuristic.
Maybe we can use extreme value asymptotics here?

```python
from rlplay.engine import BaseActorModule
from torch.distributions import Independent, Normal, kl_divergence

class GaussTaxiActor(BaseActorModule):
    def __init__(self, *, struct=False, epsilon=0.1):
        super().__init__()
        self.struct = struct

        if not struct:
            self.table = torch.nn.Embedding(500, 2 * 6)

        else:
            self.table = torch.nn.Sequential(
                ModuleDict(dict(
                    row=torch.nn.Embedding(5, 3),
                    col=torch.nn.Embedding(5, 3),
                    location=torch.nn.Embedding(5, 3),
                    destination=torch.nn.Embedding(4, 3),
                )),
                torch.nn.Linear(12, 32),
                torch.nn.ReLU(),
                torch.nn.Linear(32, 2 * 6),
            )

        # for updating the exploration epsilon in the clones
        self.register_buffer('epsilon', torch.tensor(epsilon))

    def forward(self, obs, act, rew, fin, *, hx=None, stepno=None, virtual=False):
        out, hx = self.table(obs), ()

        loc, out = torch.chunk(out, 2, dim=-1)
        scl = torch.clamp(F.softplus(out), min=1e-3)

        # val, act = Normal(loc, scl).sample().max(dim=-1)
        val, act = loc.max(dim=-1)  # unfounded heuristic :(

        return act, hx, dict(loc=loc, scl=scl, q=loc, value=val)
```

### Run!

prepare the optimizer for the learner

```python
gamma = 0.99
use_dist = True
use_target = False
use_double = False
use_duelling = False

# `target` does not work for some reason at all with taxi, maybe the freeze schedule is off?
#  or contiguous fragments work to the detriment of learning

# `duelling` also fails for both target and double, and sorta for ordinary q
```

Initialize the learner and the environment factories

```python
from functools import partial

structured = False

factory_eval = partial(base_factory, struct=structured)
factory = partial(base_factory, struct=structured)

if use_dist:
    learner = GaussTaxiActor(struct=structured)
else:
    learner = TaxiActor(duelling=use_duelling, epsilon=1., struct=structured)

learner.train()
device_ = torch.device('cpu')  # torch.device('cuda:0')
learner.to(device=device_)

optim = torch.optim.Adam(learner.parameters(), lr=1e-2)
```

Initialize the sampler

```python
T, B = 25, 8
```
from rlplay.engine.rollout import same

batchit = same.rollout(
    [factory() for _ in range(B)],
    learner,
    n_steps=T,
    sticky=False,
    device=device_,
)from rlplay.engine.rollout import single

batchit = single.rollout(
    factory,
    learner,
    n_steps=T,
    n_envs=B,
    sticky=False,
    clone=True,
    close=False,
    device=device_,
    start_method='fork',  # fork in notebook for macos, spawn in linux
)

```python
from rlplay.engine.rollout import multi

batchit = multi.rollout(
    factory,
    learner,
    n_steps=T,
    n_actors=8,
    n_per_actor=B,
    n_buffers=16,
    n_per_batch=2,
    sticky=False,
    pinned=False,
    clone=True,
    close=False,
    device=device_,
    start_method='fork',  # fork in notebook for macos, spawn in linux
)
```

A generator of evaluation rewards

```python
from rlplay.engine.rollout.evaluate import evaluate

test_it = evaluate(factory_eval, learner, n_envs=4, n_steps=500,
                   clone=False, device=device_, start_method='fork')
```

Implement your favourite training method

```python
import tqdm
import copy
from math import log, exp
from torch.nn.utils import clip_grad_norm_

torch.set_num_threads(1)

# the training loop
losses, rewards = [], []
decay = -log(2) / 50  # exploration epsilon halflife
for epoch in tqdm.tqdm(range(400)):
    # freeze the target for q
    target = copy.deepcopy(learner) if use_target else None

    for j, batch in zip(range(100), batchit):
        if use_dist:
            loss, info = gdqn(batch, learner, gamma=gamma)

        else:
            loss, info = ddq_learn(batch, learner, target=target,
                                   gamma=gamma, double=use_double)

        optim.zero_grad()
        loss.backward()
        grad = clip_grad_norm_(learner.parameters(), max_norm=1e3)
        optim.step()

        losses.append(dict(
            loss=float(loss), grad=float(grad), **info
        ))

    learner.epsilon.mul_(exp(decay)).clip_(0.1, 1.0)

    # fetch the evaluation results lagged by one inner loop!
    rewards.append(next(test_it))
```

```python
# close the generators
batchit.close()
test_it.close()
```
import pdb; pdb.pm()

```python
def collate(records):
    """collate identically keyed dicts"""
    out, n_records = {}, 0
    for record in records:
        for k, v in record.items():
            out.setdefault(k, []).append(v)

    return out

data = {k: numpy.array(v) for k, v in collate(losses).items()}
```

```python
plt.semilogy(data['loss'])
```

```python
plt.semilogy(data['grad'])
```

```python
if 'dst' in data:
    plt.semilogy(data['dst'])
```

```python
if 'reg' in data:
    plt.semilogy(data['reg'])
```

```python
rewards = numpy.stack(rewards, axis=0)
```

```python
rewards
```

```python
m, s = numpy.median(rewards, axis=-1), rewards.std(axis=-1)
```

```python
fi, ax = plt.subplots(1, 1, figsize=(4, 2), dpi=300)

ax.plot(numpy.mean(rewards, axis=-1))
ax.plot(numpy.median(rewards, axis=-1))
ax.plot(numpy.min(rewards, axis=-1))
ax.plot(numpy.std(rewards, axis=-1))
# ax.plot(m+s * 1.96)
# ax.plot(m-s * 1.96)

plt.show()
```

The ultimate evaluation run

```python
from rlplay.engine import core

with factory_eval() as env:
    learner.eval()
    eval_rewards, info = core.evaluate([
        env
    ], learner, render=True, n_steps=1e2, device=device_)

print(sum(eval_rewards))
```
import pdb; pdb.pm()

Let's analyze the performance

```python
(info['q'] - numpy.expand_dims(info['value'], -1)).shape
```

```python
plt.plot((info['q'] - numpy.expand_dims(info['value'], -1))[:, 0])
```

```python
from rlplay.algo.returns import npy_returns, npy_deltas

td_target = eval_rewards + gamma * info['value'][1:]
td_error = td_target - info['value'][:-1]
# td_error = npy_deltas(
#     eval_rewards, numpy.zeros_like(eval_rewards, dtype=bool), info['value'][:-1],
#     gamma=gamma, bootstrap=info['value'][-1])

fig, ax = plt.subplots(1, 1, figsize=(4, 2), dpi=300)
ax.semilogy(abs(td_error) / abs(td_target))
ax.set_title('relative td(1)-error');
```

```python
from rlplay.algo.returns import npy_returns, npy_deltas

# plt.plot(
#     npy_returns(eval_rewards, numpy.zeros_like(eval_rewards, dtype=bool),
#                 gamma=gamma, bootstrap=info['value'][-1]))
fig, ax = plt.subplots(1, 1, figsize=(4, 2), dpi=300)
ax.plot(info['value']);
```

```python
with torch.no_grad():
    obs = torch.arange(500).unsqueeze(1)
    if structured:
        obs = StructuredTaxiEnv.observation(obs)

    _, _, info = learner(obs, act=None, rew=None, fin=None)
    if 'loc' in info:
        # shuffle dims to `a d p r c`
        loc = info['loc'].reshape(5, 5, 5, 4, -1).permute(4, 3, 2, 0, 1)
        scl = info['scl'].reshape(5, 5, 5, 4, -1).permute(4, 3, 2, 0, 1)
        p = Normal(loc, scl)
        ent = p.entropy()

    else:
        loc = info['q'].reshape(5, 5, 5, 4, -1).permute(4, 3, 2, 0, 1)
```

```python
d = 'R'
dst = {'R': 0, 'G': 1, 'Y': 2, 'B': 3}

# +---------+
# |R: | : :G|
# | : | : : |
# | : : : : |
# | | : | : |
# |Y| : |B: |
# +---------+

x = loc[:, dst[d]].permute(1, 2, 0, 3)  # .reshape(5 * 5, 6 * 5)
plt.imshow(torch.cat([
    x, x[:, :, :4].max(2, keepdim=True).values,
], dim=2).reshape(5 * 5, 7 * 5))

# cols: action snewpd, rows: location RGYBT
plt.xlabel('action + pool')
plt.xticks(list(range(2, 35, 5)), [*'snewpd', 'max'])
plt.ylabel('location')
plt.yticks(list(range(2, 25, 5)), list('RGYBT'))
plt.title(f'destination: {d}');
```

```python
assert False
```