# Chasing the power of optimal length increase

(c) 2021 Tom Röschinger. This work is licensed under a [Creative Commons Attribution License CC-BY 4.0](https://creativecommons.org/licenses/by/4.0/). All code contained herein is licensed under an [MIT license](https://opensource.org/licenses/MIT).

***

```julia
using LinearAlgebra, Distributions, Plots, Polynomials, LinearAlgebra, Turing, Printf
# Custom package
using Jevo, Jedi

# Set plotting style
Jedi.default_gr!()
```

    ┌ Info: Precompiling Jedi [b681c197-c997-42fd-b5bb-d7d7839f617e]
    └ @ Base loading.jl:1278
    WARNING: using StatsBase.ecdf in module Jedi conflicts with an existing identifier.
    ┌ Warning: Package Jedi does not have Measures in its dependencies:
    │ - If you have Jedi checked out for development and have
    │   added Measures as a dependency but haven't updated your primary
    │   environment's manifest file, try `Pkg.resolve()`.
    │ - Otherwise you may need to report an issue with Jedi
    └ Loading Measures into Jedi from project dependency, future warnings for Jedi are suppressed.

    Plots.GRBackend()

```julia
@model function fit_exponent(x, y)
    sigma ~ truncated(Normal(0, 10^-6), 0, Inf)
    a ~ truncated(Normal(0, 1), 0, Inf)
    b ~ Normal(0, 1)
    _y = (x .* a) .+ b
     y ~ MvNormal(_y, sigma)
end
```

    fit_exponent (generic function with 1 method)

## 1. Full Gamma Dynamics

First, let's have a look at the full $\Gamma$ dynamics. Therefore, we use the Kimura substitution probability to compute a distribution. Let's write down the rates for a match turning into a mismatch, $u_+(\Gamma, l)$, and the rate for a mismatch turning into a match, $u_-(\Gamma, l)$,

\begin{align}
u_+(\Gamma) &= \left(1 - \frac{\Gamma}{\epsilon l}\right) \left[p(F(\Gamma + \epsilon, l) - F(\Gamma, l) +\frac{\kappa}{N}\right],\\
u_-(\Gamma) &= \frac{\Gamma}{\epsilon l}\frac{1}{1-n} \left[p(F(\Gamma - \epsilon, l) - F(\Gamma, l) +\frac{\kappa}{N}\right],
\end{align}

where $p(s)$ is the Kimura substitution probability. The prefactors determine the probability of a match/mismatch mutating, which depends on the current state of the population. The extra prefactor $\frac{1}{1-n}$ comes from the consideration that a mismatch can mutate to another mismatch, which is not detectable on the level of quantitative traits. A similar factor is absent from the rate turning a match into am mismatch, since every mutation on a match leads to a mismatch.

Then we use these rates to determine the probability distribution over the binding energies,

$$
Q(\Gamma, l) \propto \prod_{\Gamma^\prime=\epsilon}^\Gamma \frac{u_+(\Gamma-\epsilon, l)}{u_-(\Gamma, l)}.
$$

It might be more useful to compute the free fitness $\Psi$ instead, which is simply the log of the distribution,

$$
\Psi(\Gamma, l) \propto \sum_{\Gamma^\prime=\epsilon}^\Gamma \log\left[\frac{u_+(\Gamma-\epsilon, l)}{u_-(\Gamma, l)}\right].
$$

```julia
# Neutral Expectation
γ_0(n) = (n-1)/n

# Binding Threshold
γ_1(l, n, l_0) = γ_0(n) - l_0/l

# Binding probability
pb(γ, l, n, l_0, gap) = 1 / (1 + exp(gap * (l / l_0) * (γ - γ_1(l, n, l_0))))

# Fitness component for functional binding
F_b(γ, l, n, l_0, gap, f0) = f0 * pb(γ, l, n, l_0, gap)

# Fitness component of genomic constraint
F_c(l, l_0, fl) = - fl * l / l_0

# Total Fitness
F(γ, l, n, l_0, gap, f0, fl) = F_b(γ, l, n, l_0, gap, f0) + F_c(l, l_0, fl)

function kimura(s)
    if s^2 < 10^(-20)
        return 1
    elseif s < -10^7
        return 0
    else
        return s/(1 - exp(-s))
    end
end

# Substitution rate for trait changes
# (k, l) -> (k + 1, l)
up(γ, l, n, l_0, gap, f0, fl, κ) =  (1 - γ) * (kimura(F(γ + 1/l, l, n, l_0, gap, f0, fl)
        - F(γ, l, n, l_0, gap, f0, fl)) + κ )

# (k, l) -> (k - 1, l)
um(γ, l, n, l_0, gap, f0, fl, κ) = γ * (1 - (n-2)/(n-1)) * (kimura(F(γ - 1/l, l, n, l_0, gap, f0, fl)
        - F(γ, l, n, l_0, gap, f0, fl)) + κ )

# Free Fitness
function Ψ(l, n, l_0, gap, f0, fl, κ)
    # Start at neutral γ
    γ_i = floor(γ_0(n) * l) / l
    return collect(0:1/l:γ_i), push!(Float64[0], cumsum([log(up(γ - 1/l, l,  n, l_0, gap, f0, fl, κ) / um(γ, l,  n, l_0, gap, f0, fl, κ)) for γ in 1/l:1/l:γ_i])...)

end
```

    Ψ (generic function with 1 method)

Using the free fitness, we can also compute the resulting distribution.

```julia
# Probability distribution (from free Fitness)
function Q(l, n, l_0, gap, f0, fl, κ)
    x, y = Ψ(l, n, l_0, gap, f0, fl, κ)
    return x, exp.(y)
end

```

    Q (generic function with 1 method)

Let's make some plots. Therefore we first determine some parameters.

```julia
gap = 10
l_0 = 20
f0 = 40l_0
fl = 1l_0
n = 4;
```

First we look at the fitness landscape and the free fitness as a sanity check, to determine that we actually implemented them correctly.

```julia
x, y = Ψ(100, n, l_0, gap, f0, fl, 0)
f = [F(γ, 100, n, l_0, gap, f0, fl) for γ in x]
println(x[argmax(y)])
plot(x, (y .- minimum(y))/(maximum(y)-minimum(y)), label="Free Fitness")

plot!(x, (f .- minimum(f))/(maximum(f)-minimum(f)), color="orange", label="Fitness")
```

    0.44

Looking good so far. Let's have a look at a couple at a distribution and how it changes with some non-equilibrium.

```julia
x1, y1 = Q(100, n, l_0, gap, f0, 0, 0)
x2, y2 = Q(100, n, l_0, gap, f0, 0, 20)
plot(x1, y1/sum(y1), title="Distributions at l=100", titlefontsize=12, label="κ = 0")
plot!(x2, y2/sum(y2), label="κ = 20")
```

We can nicely see how the distribution moves to less specificity and already gets a little mass at the neutral peak, which we have to consider later. Let's have a look at how the mean fitness looks for increasing length.

```julia
function mean_fitness(l, n, l_0, gap, f0, fl, κ)
    x, y = Q(l, n, l_0, gap, f0, fl, κ)
    y = y/sum(y)
    sum([F(γ, l, n, l_0, gap, f0, fl) for γ in x] .* y)
end

p_mf_0 = [mean_fitness(l, n, l_0, gap, f0, fl, 0) for l in 40:200]
p_mf_20 = [mean_fitness(l, n, l_0, gap, f0, fl, 20) for l in 40:200]
plot(40:200, p_mf_0, label="κ=0", title="Mean Fitness")
plot!(40:200, p_mf_20, label="κ=20")
```

As we expected, for increased non-equilibrium, it is not possible for some lengths to adapt. Now let's find the optimal length. Therefore we first compute some stuff. We start of with writing functions to compute the new intensive energy, where we have to consider the discreteness of the sequence, therefore the intensive energy changes for all types of length mutations (match/mismatch addition/removal). After computing the new intensive binding energy, we compute the selection coefficient of a length mutation at a given $\gamma$.

```julia
# (k,l) -> (k+1,l+1)
γ_pp(γ, l) = min((γ + 1/l)/(1 + 1/l), 1)

# (k,l) -> (k,l+1)
γ_pm(γ, l) = max(γ /(1 + 1/l), 0)

# (k,l) -> (k,l-1)
γ_mp(γ, l) = min(γ/(1 - 1/l), 1)

# (k,l) -> (k-1,l-1)
γ_mm(γ, l) = max((γ - 1/l)/(1 - 1/l), 0)

# selection coeffcients for length increase and decrease
# (k,l) -> (k+1,l+1)
s_pp(γ, l, n, l_0, gap, f0, fl) =
    F(γ_pp(γ, l), l + 1, n, l_0, gap, f0, fl) -
    F(γ, l, n, l_0, gap, f0, fl)

# (k,l) -> (k,l+1)
s_pm(γ, l, n, l_0, gap, f0, fl) =
    F(γ_pm(γ, l), l + 1, n, l_0, gap, f0, fl) -
    F(γ, l, n, l_0, gap, f0, fl)

# (k,l) -> (k-1,l-1)
s_mm(γ, l, n, l_0, gap, f0, fl) =
 F(γ_mm(γ, l), l - 1, n, l_0, gap, f0, fl) -
  F(γ, l, n, l_0, gap, f0, fl)

# (k,l) -> (k,l-1)
s_mp(γ, l, n, l_0, gap, f0, fl) =
 F(γ_mp(γ, l), l - 1, n, l_0, gap, f0, fl) -
  F(γ, l, n, l_0, gap, f0, fl)
```

    s_mp (generic function with 1 method)

Now that we can compute the selection coefficient, we compute the length substitution rates. For that we have to consider the difference of mutations that increase the length, where we add a random position, hence a match with probability $1/n$, and mutations where we decrease length, where the probability of removing a match depends on the intensive energy of the site, $(1-\gamma)$. The total length increase and decrease mutation rates are both averaged over the steady state distribution at that length, which comes from the assumption of separation of time scales.

```julia
function p_l_th(κ, n, l_0, gap, f0, fl, l_range=collect(40:200))
    vp = []
    vm = []
    for l in l_range
        x, y = Q(l, n, l_0, gap, f0, fl, κ)
        y = y/sum(y)
        s_mm_p = [s_pp(γ, l, n, l_0, gap, f0, fl) for γ in x]
        s_m_p = [s_pm(γ, l, n, l_0, gap, f0, fl) for γ in x]
        s_mm_m = [s_mm(γ, l, n, l_0, gap, f0, fl) for γ in x]
        s_m_m = [s_mp(γ, l, n, l_0, gap, f0, fl) for γ in x]
        v_mm_p = [(n-1)/n * kimura(s) for s in s_mm_p]
        v_m_p = [1/n * kimura(s) for s in s_m_p]
        v_mm_m = [ γ * kimura(s) for (s, γ) in zip(s_mm_m, x)]
        v_m_m = [(1 - γ)  * kimura(s) for (s, γ) in zip(s_m_m, x)]
        push!(vp, sum((v_m_p .+ v_mm_p) .* y))
        push!(vm, sum((v_mm_m .+ v_m_m) .* y))
    end
    return vp, vm, push!([0.01], (cumsum([log(vp[i]/vm[i+1]) for i in 1:length(vp)-1]).+0.01)...)
end
```

    p_l_th (generic function with 2 methods)

Now let's have a look at the length increase and decrease substitution rates and the resulting effective fitness.

```julia
l_arr = collect(30:150)
vp, vm, p_l = p_l_th(0, n, l_0, gap, f0, fl, l_arr)
p1 = plot(l_arr, vp, label="vp", legend=:topright, title="κ=0", xlabel="l")
plot!(p1, l_arr, vm, label="vm")
p2 = plot(l_arr, p_l ./ maximum(p_l), label="Effective Fitness", legend=:bottomleft, title="Fitnesses (both scaled)", xlabel="l")
mf = [mean_fitness(l, n, l_0, gap, f0, fl, 0) for l in l_arr]
plot!(p2, l_arr, (mf .- mf[1]) ./ (maximum(mf)-mf[1]) , label="Mean Fitness")
plot([p1, p2]..., size=(700, 300))
```

We can nicely see how the effective fitness is shifted to higher lengths due to the asymmetry of length mutations. When we consider optimal lengths in non-equilibrium, we have to consider that for smaller lengths sites will not able to adapt at all, and therefore be non-functional. We want to exclude these non-functional sites for now. Hence we use a function to find the minimal length of functional sites.

```julia
function min_func_length(n, l_0, gap, f0, fl, κ, l_arr=collect(20:200))
    gamma_max_list = zeros(length(l_arr))
    for (i, l) in enumerate(l_arr)
        x, y = Q(l, n, l_0, gap, f0, fl, κ)
        gamma_max_list[i] = x[argmax(y)]
    end
    iterator = [[x, y] for (x, y) in zip(gamma_max_list, l_arr)]
    return findfirst( x -> x[1] < γ_1(x[2], n, l_0), iterator) + l_arr[1]-1
end

lmin_arr = [min_func_length(n, l_0, gap, f0, fl, κ) for κ in 0:20]
κ_arr = collect(0:20)

plot(
    κ_arr .+ 1,
    lmin_arr,
    xlabel="κ",
    ylabel="l_m",
    title="Minimal Functional Length",
    linewidth=2
)
```

This looks nicely linear. Now let's look at the length mutation rates and the effective fitness for high non-equilibrium.

```julia
lm = min_func_length(n, l_0, gap, f0, fl, 20)
l_arr = collect(lm:200)
vp, vm, p_l = p_l_th(20, n, l_0, gap, f0, fl, l_arr)

p1 = plot(l_arr, vp, label="vp", legend=:topright, title="κ=20", xlabel="l")
plot!(p1, l_arr, vm, label="vm")
p2 = plot(l_arr, p_l ./ maximum(p_l), label="Effective Fitness", legend=:bottomleft, title="Both scaled", xlabel="l")
mf = [mean_fitness(l, n, l_0, gap, f0, fl, 20) for l in l_arr]
plot!(p2, l_arr, (mf .- mf[1]) ./ (maximum(mf)-mf[1]) , label="Mean Fitness")
plot([p1, p2]..., size=(700, 300))
```

Now we have all the tools to compute the optimal length as a function of non-equilibrium.

```julia
function optimal_length_full(n, l_0, gap, f0, fl, κ_arr=collect(0:20))

    l_opt_th = []
    for κ in κ_arr
        lm = min_func_length(n, l_0, gap, f0, fl, κ)
        l_arr = collect(lm:250)
        p_l = p_l_th(κ, n, l_0, gap, f0, fl, l_arr)[3]
        l_opt = argmax(p_l) + lm - 1
        push!(l_opt_th, l_opt)
    end
    return l_opt_th
end

function plot_opt_length_full(n, l_0, gap, f0, fl, κ_arr=collect(0:20))
    l_opt_lin = optimal_length_full(n, l_0, gap, f0, fl, κ_arr)

    p_lin = plot(
        κ_arr,
        l_opt_lin,
        xlabel="κ",
        ylabel="l_opt",
        title="f0=$(f0/l_0)l_0, fl=$(fl/l_0)l_0"
    )
    #=
    plot!(p_lin,
        κ_arr[1:100],
        (1 .+κ_arr[1:100]) .^(1/2) * l_opt_lin[1],
        color="gray",
        linestyle=:dash
    )
    =#
    κ_arr_log = exp.(range(0, length=length(κ_arr), stop=log(κ_arr[end] + 1))) .- 1
    l_opt_log = optimal_length_full(n, l_0, gap, f0, fl, κ_arr_log)

    logx = log.(1 .+ κ_arr_log)
    logy = log.(l_opt_log)

    chn = sample(fit_exponent(logx[end-15:end], logy[end-15:end]), NUTS(0.75), 5_000)

    p_log = scatter(
        log.(κ_arr_log .+ 1),
        log.(l_opt_log),
        legend=:topleft,
        xlabel="log(1+κ)",
        ylabel="log(l_opt)"
    )

    plot!(
        p_log,
        log.(κ_arr_log .+ 1),
        log.((κ_arr_log .+ 1).^(mean(chn[:a]))) .+ mean(chn[:b]),
        linestyle=:dash,
        title=@sprintf "LogLog with slope %.5f" mean(chn[:a])
    )

    return plot([p_lin, p_log]..., size=(700, 300))
end
```

    plot_opt_length_full (generic function with 2 methods)

```julia
plot_opt_length_full(n, l_0, gap, 40l_0, l_0, collect(0:30))
```

    ┌ Warning: The current proposal will be rejected due to numerical error(s).
    │   isfinite.((θ, r, ℓπ, ℓκ)) = (true, false, false, false)
    └ @ AdvancedHMC /Users/tomroschinger/.julia/packages/AdvancedHMC/MIxdK/src/hamiltonian.jl:47
    ┌ Warning: The current proposal will be rejected due to numerical error(s).
    │   isfinite.((θ, r, ℓπ, ℓκ)) = (true, false, false, false)
    └ @ AdvancedHMC /Users/tomroschinger/.julia/packages/AdvancedHMC/MIxdK/src/hamiltonian.jl:47
    ┌ Warning: The current proposal will be rejected due to numerical error(s).
    │   isfinite.((θ, r, ℓπ, ℓκ)) = (true, false, false, false)
    └ @ AdvancedHMC /Users/tomroschinger/.julia/packages/AdvancedHMC/MIxdK/src/hamiltonian.jl:47
    ┌ Warning: The current proposal will be rejected due to numerical error(s).
    │   isfinite.((θ, r, ℓπ, ℓκ)) = (true, false, false, false)
    └ @ AdvancedHMC /Users/tomroschinger/.julia/packages/AdvancedHMC/MIxdK/src/hamiltonian.jl:47
    ┌ Warning: The current proposal will be rejected due to numerical error(s).
    │   isfinite.((θ, r, ℓπ, ℓκ)) = (true, false, false, false)
    └ @ AdvancedHMC /Users/tomroschinger/.julia/packages/AdvancedHMC/MIxdK/src/hamiltonian.jl:47
    ┌ Warning: The current proposal will be rejected due to numerical error(s).
    │   isfinite.((θ, r, ℓπ, ℓκ)) = (true, false, false, false)
    └ @ AdvancedHMC /Users/tomroschinger/.julia/packages/AdvancedHMC/MIxdK/src/hamiltonian.jl:47
    ┌ Warning: The current proposal will be rejected due to numerical error(s).
    │   isfinite.((θ, r, ℓπ, ℓκ)) = (true, false, false, false)
    └ @ AdvancedHMC /Users/tomroschinger/.julia/packages/AdvancedHMC/MIxdK/src/hamiltonian.jl:47
    ┌ Warning: The current proposal will be rejected due to numerical error(s).
    │   isfinite.((θ, r, ℓπ, ℓκ)) = (true, false, false, false)
    └ @ AdvancedHMC /Users/tomroschinger/.julia/packages/AdvancedHMC/MIxdK/src/hamiltonian.jl:47
    ┌ Warning: The current proposal will be rejected due to numerical error(s).
    │   isfinite.((θ, r, ℓπ, ℓκ)) = (true, false, false, false)
    └ @ AdvancedHMC /Users/tomroschinger/.julia/packages/AdvancedHMC/MIxdK/src/hamiltonian.jl:47
    ┌ Warning: The current proposal will be rejected due to numerical error(s).
    │   isfinite.((θ, r, ℓπ, ℓκ)) = (true, false, false, false)
    └ @ AdvancedHMC /Users/tomroschinger/.julia/packages/AdvancedHMC/MIxdK/src/hamiltonian.jl:47
    ┌ Warning: The current proposal will be rejected due to numerical error(s).
    │   isfinite.((θ, r, ℓπ, ℓκ)) = (true, false, false, false)
    └ @ AdvancedHMC /Users/tomroschinger/.julia/packages/AdvancedHMC/MIxdK/src/hamiltonian.jl:47
    ┌ Warning: The current proposal will be rejected due to numerical error(s).
    │   isfinite.((θ, r, ℓπ, ℓκ)) = (true, false, false, false)
    └ @ AdvancedHMC /Users/tomroschinger/.julia/packages/AdvancedHMC/MIxdK/src/hamiltonian.jl:47
    ┌ Info: Found initial step size
    │   ϵ = 6.103515625e-6
    └ @ Turing.Inference /Users/tomroschinger/.julia/packages/Turing/O1Pn0/src/inference/hmc.jl:195
    Sampling: 100%|█████████████████████████████████████████| Time: 0:00:02

```julia
plot_opt_length_full(n, l_0, gap, 200l_0, 1.6l_0, collect(0:40))
```

    ┌ Warning: The current proposal will be rejected due to numerical error(s).
    │   isfinite.((θ, r, ℓπ, ℓκ)) = (true, false, false, false)
    └ @ AdvancedHMC /Users/tomroschinger/.julia/packages/AdvancedHMC/MIxdK/src/hamiltonian.jl:47
    ┌ Warning: The current proposal will be rejected due to numerical error(s).
    │   isfinite.((θ, r, ℓπ, ℓκ)) = (true, false, false, false)
    └ @ AdvancedHMC /Users/tomroschinger/.julia/packages/AdvancedHMC/MIxdK/src/hamiltonian.jl:47
    ┌ Warning: The current proposal will be rejected due to numerical error(s).
    │   isfinite.((θ, r, ℓπ, ℓκ)) = (true, false, false, false)
    └ @ AdvancedHMC /Users/tomroschinger/.julia/packages/AdvancedHMC/MIxdK/src/hamiltonian.jl:47
    ┌ Warning: The current proposal will be rejected due to numerical error(s).
    │   isfinite.((θ, r, ℓπ, ℓκ)) = (true, false, false, false)
    └ @ AdvancedHMC /Users/tomroschinger/.julia/packages/AdvancedHMC/MIxdK/src/hamiltonian.jl:47
    ┌ Warning: The current proposal will be rejected due to numerical error(s).
    │   isfinite.((θ, r, ℓπ, ℓκ)) = (true, false, false, false)
    └ @ AdvancedHMC /Users/tomroschinger/.julia/packages/AdvancedHMC/MIxdK/src/hamiltonian.jl:47
    ┌ Warning: The current proposal will be rejected due to numerical error(s).
    │   isfinite.((θ, r, ℓπ, ℓκ)) = (true, false, false, false)
    └ @ AdvancedHMC /Users/tomroschinger/.julia/packages/AdvancedHMC/MIxdK/src/hamiltonian.jl:47
    ┌ Warning: The current proposal will be rejected due to numerical error(s).
    │   isfinite.((θ, r, ℓπ, ℓκ)) = (true, false, false, false)
    └ @ AdvancedHMC /Users/tomroschinger/.julia/packages/AdvancedHMC/MIxdK/src/hamiltonian.jl:47
    ┌ Warning: The current proposal will be rejected due to numerical error(s).
    │   isfinite.((θ, r, ℓπ, ℓκ)) = (true, false, false, false)
    └ @ AdvancedHMC /Users/tomroschinger/.julia/packages/AdvancedHMC/MIxdK/src/hamiltonian.jl:47
    ┌ Warning: The current proposal will be rejected due to numerical error(s).
    │   isfinite.((θ, r, ℓπ, ℓκ)) = (true, false, false, false)
    └ @ AdvancedHMC /Users/tomroschinger/.julia/packages/AdvancedHMC/MIxdK/src/hamiltonian.jl:47
    ┌ Warning: The current proposal will be rejected due to numerical error(s).
    │   isfinite.((θ, r, ℓπ, ℓκ)) = (true, false, false, false)
    └ @ AdvancedHMC /Users/tomroschinger/.julia/packages/AdvancedHMC/MIxdK/src/hamiltonian.jl:47
    ┌ Warning: The current proposal will be rejected due to numerical error(s).
    │   isfinite.((θ, r, ℓπ, ℓκ)) = (true, false, false, false)
    └ @ AdvancedHMC /Users/tomroschinger/.julia/packages/AdvancedHMC/MIxdK/src/hamiltonian.jl:47
    ┌ Info: Found initial step size
    │   ϵ = 6.103515625e-6
    └ @ Turing.Inference /Users/tomroschinger/.julia/packages/Turing/O1Pn0/src/inference/hmc.jl:195
    Sampling: 100%|█████████████████████████████████████████| Time: 0:00:04

```julia
plot_opt_length_full(n, l_0, gap, 300l_0, 2l_0, collect(0:40))
```

    ┌ Warning: The current proposal will be rejected due to numerical error(s).
    │   isfinite.((θ, r, ℓπ, ℓκ)) = (true, false, false, false)
    └ @ AdvancedHMC /Users/tomroschinger/.julia/packages/AdvancedHMC/MIxdK/src/hamiltonian.jl:47
    ┌ Warning: The current proposal will be rejected due to numerical error(s).
    │   isfinite.((θ, r, ℓπ, ℓκ)) = (true, false, false, false)
    └ @ AdvancedHMC /Users/tomroschinger/.julia/packages/AdvancedHMC/MIxdK/src/hamiltonian.jl:47
    ┌ Warning: The current proposal will be rejected due to numerical error(s).
    │   isfinite.((θ, r, ℓπ, ℓκ)) = (true, false, false, false)
    └ @ AdvancedHMC /Users/tomroschinger/.julia/packages/AdvancedHMC/MIxdK/src/hamiltonian.jl:47
    ┌ Warning: The current proposal will be rejected due to numerical error(s).
    │   isfinite.((θ, r, ℓπ, ℓκ)) = (true, false, false, false)
    └ @ AdvancedHMC /Users/tomroschinger/.julia/packages/AdvancedHMC/MIxdK/src/hamiltonian.jl:47
    ┌ Warning: The current proposal will be rejected due to numerical error(s).
    │   isfinite.((θ, r, ℓπ, ℓκ)) = (true, false, false, false)
    └ @ AdvancedHMC /Users/tomroschinger/.julia/packages/AdvancedHMC/MIxdK/src/hamiltonian.jl:47
    ┌ Warning: The current proposal will be rejected due to numerical error(s).
    │   isfinite.((θ, r, ℓπ, ℓκ)) = (true, false, false, false)
    └ @ AdvancedHMC /Users/tomroschinger/.julia/packages/AdvancedHMC/MIxdK/src/hamiltonian.jl:47
    ┌ Warning: The current proposal will be rejected due to numerical error(s).
    │   isfinite.((θ, r, ℓπ, ℓκ)) = (true, false, false, false)
    └ @ AdvancedHMC /Users/tomroschinger/.julia/packages/AdvancedHMC/MIxdK/src/hamiltonian.jl:47
    ┌ Warning: The current proposal will be rejected due to numerical error(s).
    │   isfinite.((θ, r, ℓπ, ℓκ)) = (true, false, false, false)
    └ @ AdvancedHMC /Users/tomroschinger/.julia/packages/AdvancedHMC/MIxdK/src/hamiltonian.jl:47
    ┌ Warning: The current proposal will be rejected due to numerical error(s).
    │   isfinite.((θ, r, ℓπ, ℓκ)) = (true, false, false, false)
    └ @ AdvancedHMC /Users/tomroschinger/.julia/packages/AdvancedHMC/MIxdK/src/hamiltonian.jl:47
    ┌ Warning: The current proposal will be rejected due to numerical error(s).
    │   isfinite.((θ, r, ℓπ, ℓκ)) = (true, false, false, false)
    └ @ AdvancedHMC /Users/tomroschinger/.julia/packages/AdvancedHMC/MIxdK/src/hamiltonian.jl:47
    ┌ Warning: The current proposal will be rejected due to numerical error(s).
    │   isfinite.((θ, r, ℓπ, ℓκ)) = (true, false, false, false)
    └ @ AdvancedHMC /Users/tomroschinger/.julia/packages/AdvancedHMC/MIxdK/src/hamiltonian.jl:47
    ┌ Warning: The current proposal will be rejected due to numerical error(s).
    │   isfinite.((θ, r, ℓπ, ℓκ)) = (true, false, false, false)
    └ @ AdvancedHMC /Users/tomroschinger/.julia/packages/AdvancedHMC/MIxdK/src/hamiltonian.jl:47
    ┌ Warning: The current proposal will be rejected due to numerical error(s).
    │   isfinite.((θ, r, ℓπ, ℓκ)) = (true, false, false, false)
    └ @ AdvancedHMC /Users/tomroschinger/.julia/packages/AdvancedHMC/MIxdK/src/hamiltonian.jl:47
    ┌ Warning: The current proposal will be rejected due to numerical error(s).
    │   isfinite.((θ, r, ℓπ, ℓκ)) = (true, false, false, false)
    └ @ AdvancedHMC /Users/tomroschinger/.julia/packages/AdvancedHMC/MIxdK/src/hamiltonian.jl:47
    ┌ Warning: The current proposal will be rejected due to numerical error(s).
    │   isfinite.((θ, r, ℓπ, ℓκ)) = (true, false, false, false)
    └ @ AdvancedHMC /Users/tomroschinger/.julia/packages/AdvancedHMC/MIxdK/src/hamiltonian.jl:47
    ┌ Info: Found initial step size
    │   ϵ = 3.814697265625e-7
    └ @ Turing.Inference /Users/tomroschinger/.julia/packages/Turing/O1Pn0/src/inference/hmc.jl:195
    Sampling: 100%|█████████████████████████████████████████| Time: 0:00:03

# Only use the maximum of the energy distribution

Now we want to compare using the full distribution as we did above to simply using the maximum point of the distribution to compare. Therefore we first compute the maximum of the distribution.

```julia
function γ_star(l, n, l_0, gap, f0, fl, κ)
    x, y = Ψ(l, n, l_0, gap, f0, fl, κ)
    max_ind = argmax(y)
    if max_ind == length(x)
        return x[end]
    elseif max_ind == 1
        return x[1]
    else
        f = Polynomials.fit(x[max_ind-1:max_ind+1], y[max_ind-1:max_ind+1], 2)
        return -f[1]/2f[2]
    end
end
```

    γ_star (generic function with 1 method)

```julia
# (k,l) -> (k+1,l+1)
γ_pp(γ, l) = min((γ + 1/l)/(1 + 1/l), 1)

# (k,l) -> (k,l+1)
γ_pm(γ, l) = max(γ /(1 + 1/l), 0)

# (k,l) -> (k,l-1)
γ_mp(γ, l) = min(γ/(1 - 1/l), 1)

# (k,l) -> (k-1,l-1)
γ_mm(γ, l) = max((γ - 1/l)/(1 - 1/l), 0)

# selection coeffcients for length increase and decrease
# (k,l) -> (k+1,l+1)
s_pp(γ, l, n, l_0, gap, f0, fl) =
    F(γ_pp(γ, l), l + 1, n, l_0, gap, f0, fl) -
    F(γ, l, n, l_0, gap, f0, fl)

# (k,l) -> (k,l+1)
s_pm(γ, l, n, l_0, gap, f0, fl) =
    F(γ_pm(γ, l), l + 1, n, l_0, gap, f0, fl) -
    F(γ, l, n, l_0, gap, f0, fl)

# (k,l) -> (k-1,l-1)
s_mm(γ, l, n, l_0, gap, f0, fl) =
 F(γ_mm(γ, l), l - 1, n, l_0, gap, f0, fl) -
  F(γ, l, n, l_0, gap, f0, fl)

# (k,l) -> (k,l-1)
s_mp(γ, l, n, l_0, gap, f0, fl) =
 F(γ_mp(γ, l), l - 1, n, l_0, gap, f0, fl) -
  F(γ, l, n, l_0, gap, f0, fl)

# substitution rates of length increase and decrease
v_pp(γ, l, n, l_0, gap, f0, fl) = γ_0(n) * kimura(s_pp(γ, l, n, l_0, gap, f0, fl))

v_pm(γ, l, n, l_0, gap, f0, fl) = (1 - γ_0(n)) * kimura(s_pm(γ, l, n, l_0, gap, f0, fl))

v_plus(γ, l, n, l_0, gap, f0, fl) = v_pp(γ, l, n, l_0, gap, f0, fl) + v_pm(γ, l, n, l_0, gap, f0, fl)

v_plus_star(l, n, l_0, gap, f0, fl, κ) = v_plus(γ_star(l, n, l_0, gap, f0, fl, κ), l, n, l_0, gap, f0, fl)

v_mm(γ, l, n, l_0, gap, f0, fl) = γ * kimura(s_mm(γ, l, n, l_0, gap, f0, fl))

v_mp(γ, l, n, l_0, gap, f0, fl) = (1 - γ) * kimura(s_mp(γ, l, n, l_0, gap, f0, fl))

v_minus(γ, l, n, l_0, gap, f0, fl) = v_mm(γ, l, n, l_0, gap, f0, fl) + v_mp(γ, l, n, l_0, gap, f0, fl)

v_minus_star(l, n, l_0, gap, f0, fl, κ) = v_minus(γ_star(l, n, l_0, gap, f0, fl, κ), l, n, l_0, gap, f0, fl)

# Effective Fitness
F_eff(l, n, l_0, gap, f0, fl, κ, lmin) = sum([log(v_plus_star(l_, n, l_0, gap, f0, fl, κ)/
    v_minus_star(l_+1, n, l_0, gap, f0, fl, κ)) for l_  in lmin:l])

γ_star_list(n, l_0, gap, f0, fl, κ, l_max=200) = [l > 10 ? γ_star(l, n, l_0, gap, f0, fl, κ) : missing  for l in 1:l_max]

function l_star(n, l_0, gap, f0, fl, κ, l_max=200)
    gstarlist = γ_star_list(n, l_0, gap, f0, fl, κ, l_max)
    funct_pos_list = []
    for l in 1:l_max
        if (~ismissing(gstarlist[l])) && (gstarlist[l] < γ_1(l, n, l_0))
            push!(funct_pos_list, l)
        end
    end
    γ_fun_list = gstarlist[funct_pos_list]
    lm = funct_pos_list[1]
    f_list = zeros(l_max)

    f_list[lm+1:l_max] = cumsum([log(v_plus(gstarlist[i-1], i-1, n, l_0, gap, f0, fl)/
                v_minus(gstarlist[i], i, n, l_0, gap, f0, fl)) for i in lm+1:l_max])
    max_ind = argmax(f_list)

    fmax = f_list[max_ind]
    lst = max_ind
    return gstarlist[max_ind], γ_fun_list, f_list, fmax, lst, lm
end
```

    l_star (generic function with 2 methods)

```julia
function p_l_star(κ, n, l_0, gap, f0, fl, l_range=collect(20:250))
    vp = []
    vm = []
    for l in l_range
        γ = γ_star(l, n, l_0, gap, f0, fl, κ)
        s_mm_p = s_pp(γ, l, n, l_0, gap, f0, fl)
        s_m_p = s_pm(γ, l, n, l_0, gap, f0, fl)
        s_mm_m = s_mm(γ, l, n, l_0, gap, f0, fl)
        s_m_m = s_mp(γ, l, n, l_0, gap, f0, fl)
        v_mm_p = γ_0(n) * kimura(s_mm_p)
        v_m_p = (1 - γ_0(n)) * kimura(s_m_p)
        v_mm_m = γ * kimura(s_mm_m)
        v_m_m = (1 - γ)  * kimura(s_m_m)
        push!(vp, v_m_p + v_mm_p)
        push!(vm, v_mm_m .+ v_m_m)
    end
    return vp, vm, push!([0.01], (cumsum([log(vp[i-1]/vm[i]) for i in 2:length(vp)]).+0.01)...)
end
```

    p_l_star (generic function with 2 methods)

```julia
κ = 40
lm = min_func_length(n, l_0, gap, f0, fl, κ)
plot(lm:200, p_l_star(κ, n, l_0, gap, f0, fl, collect(lm:200))[3])
plot!(twinx(), lm:200, [F_eff(l, n, l_0, gap, f0, fl, κ, lm) for l in lm:200], color = "orange")
```

```julia
function optimal_length_star(n, l_0, gap, f0, fl, κ_arr=collect(0:20))
    l_opt_star = []
    for κ in κ_arr
        push!(l_opt_star, l_star(n, l_0, gap, f0, fl, κ)[end-1])
    end
    return l_opt_star
end

function plot_opt_length_star(n, l_0, gap, f0, fl, κ_arr=collect(0:20))
    l_opt_lin = optimal_length_star(n, l_0, gap, f0, fl, κ_arr)

    p_lin = plot(
        κ_arr,
        l_opt_lin,
        xlabel="κ",
        ylabel="l_opt",
        title="f0=$(f0/l_0)l_0, fl=$(fl/l_0)l_0"
    )

    κ_arr_log = exp.(range(0, length=length(κ_arr), stop=log(κ_arr[end] + 1))) .- 1
    l_opt_log = optimal_length_star(n, l_0, gap, f0, fl, κ_arr_log)

    logx = log.(1 .+ κ_arr_log)
    logy = log.(l_opt_log)

    chn = sample(fit_exponent(logx[end-15:end], logy[end-15:end]), NUTS(0.75), 5_000)

    p_log = scatter(
        log.(κ_arr_log .+ 1),
        log.(l_opt_log),
        legend=:topleft,
        xlabel="log(1+κ)",
        ylabel="log(l_opt)"
    )

    plot!(
        p_log,
        log.(κ_arr_log .+ 1),
        log.((κ_arr_log .+ 1).^(mean(chn[:a]))) .+ mean(chn[:b]),
        linestyle=:dash,
        title=@sprintf "LogLog with slope %.3f" mean(chn[:a])
    )

    return plot([p_lin, p_log]..., size=(900, 450))
end
```

    plot_opt_length_star (generic function with 2 methods)

```julia
plot_opt_length_star(n, l_0, gap, 40l_0, l_0, collect(0:40))
```

    ┌ Warning: The current proposal will be rejected due to numerical error(s).
    │   isfinite.((θ, r, ℓπ, ℓκ)) = (true, false, false, false)
    └ @ AdvancedHMC /Users/tomroschinger/.julia/packages/AdvancedHMC/MIxdK/src/hamiltonian.jl:47
    ┌ Warning: The current proposal will be rejected due to numerical error(s).
    │   isfinite.((θ, r, ℓπ, ℓκ)) = (true, false, false, false)
    └ @ AdvancedHMC /Users/tomroschinger/.julia/packages/AdvancedHMC/MIxdK/src/hamiltonian.jl:47
    ┌ Warning: The current proposal will be rejected due to numerical error(s).
    │   isfinite.((θ, r, ℓπ, ℓκ)) = (true, false, false, false)
    └ @ AdvancedHMC /Users/tomroschinger/.julia/packages/AdvancedHMC/MIxdK/src/hamiltonian.jl:47
    ┌ Warning: The current proposal will be rejected due to numerical error(s).
    │   isfinite.((θ, r, ℓπ, ℓκ)) = (true, false, false, false)
    └ @ AdvancedHMC /Users/tomroschinger/.julia/packages/AdvancedHMC/MIxdK/src/hamiltonian.jl:47
    ┌ Warning: The current proposal will be rejected due to numerical error(s).
    │   isfinite.((θ, r, ℓπ, ℓκ)) = (true, false, false, false)
    └ @ AdvancedHMC /Users/tomroschinger/.julia/packages/AdvancedHMC/MIxdK/src/hamiltonian.jl:47
    ┌ Warning: The current proposal will be rejected due to numerical error(s).
    │   isfinite.((θ, r, ℓπ, ℓκ)) = (true, false, false, false)
    └ @ AdvancedHMC /Users/tomroschinger/.julia/packages/AdvancedHMC/MIxdK/src/hamiltonian.jl:47
    ┌ Warning: The current proposal will be rejected due to numerical error(s).
    │   isfinite.((θ, r, ℓπ, ℓκ)) = (true, false, false, false)
    └ @ AdvancedHMC /Users/tomroschinger/.julia/packages/AdvancedHMC/MIxdK/src/hamiltonian.jl:47
    ┌ Warning: The current proposal will be rejected due to numerical error(s).
    │   isfinite.((θ, r, ℓπ, ℓκ)) = (true, false, false, false)
    └ @ AdvancedHMC /Users/tomroschinger/.julia/packages/AdvancedHMC/MIxdK/src/hamiltonian.jl:47
    ┌ Warning: The current proposal will be rejected due to numerical error(s).
    │   isfinite.((θ, r, ℓπ, ℓκ)) = (true, false, false, false)
    └ @ AdvancedHMC /Users/tomroschinger/.julia/packages/AdvancedHMC/MIxdK/src/hamiltonian.jl:47
    ┌ Warning: The current proposal will be rejected due to numerical error(s).
    │   isfinite.((θ, r, ℓπ, ℓκ)) = (true, false, false, false)
    └ @ AdvancedHMC /Users/tomroschinger/.julia/packages/AdvancedHMC/MIxdK/src/hamiltonian.jl:47
    ┌ Warning: The current proposal will be rejected due to numerical error(s).
    │   isfinite.((θ, r, ℓπ, ℓκ)) = (true, false, false, false)
    └ @ AdvancedHMC /Users/tomroschinger/.julia/packages/AdvancedHMC/MIxdK/src/hamiltonian.jl:47
    ┌ Warning: The current proposal will be rejected due to numerical error(s).
    │   isfinite.((θ, r, ℓπ, ℓκ)) = (true, false, false, false)
    └ @ AdvancedHMC /Users/tomroschinger/.julia/packages/AdvancedHMC/MIxdK/src/hamiltonian.jl:47
    ┌ Warning: The current proposal will be rejected due to numerical error(s).
    │   isfinite.((θ, r, ℓπ, ℓκ)) = (true, false, false, false)
    └ @ AdvancedHMC /Users/tomroschinger/.julia/packages/AdvancedHMC/MIxdK/src/hamiltonian.jl:47
    ┌ Info: Found initial step size
    │   ϵ = 3.0517578125e-6
    └ @ Turing.Inference /Users/tomroschinger/.julia/packages/Turing/O1Pn0/src/inference/hmc.jl:195
    Sampling: 100%|█████████████████████████████████████████| Time: 0:00:02

```julia
plot_opt_length_star(n, l_0, gap, 200l_0, 1.6fl, collect(0:40))
```

    ┌ Warning: The current proposal will be rejected due to numerical error(s).
    │   isfinite.((θ, r, ℓπ, ℓκ)) = (true, false, false, false)
    └ @ AdvancedHMC /Users/tomroschinger/.julia/packages/AdvancedHMC/MIxdK/src/hamiltonian.jl:47
    ┌ Warning: The current proposal will be rejected due to numerical error(s).
    │   isfinite.((θ, r, ℓπ, ℓκ)) = (true, false, false, false)
    └ @ AdvancedHMC /Users/tomroschinger/.julia/packages/AdvancedHMC/MIxdK/src/hamiltonian.jl:47
    ┌ Warning: The current proposal will be rejected due to numerical error(s).
    │   isfinite.((θ, r, ℓπ, ℓκ)) = (true, false, false, false)
    └ @ AdvancedHMC /Users/tomroschinger/.julia/packages/AdvancedHMC/MIxdK/src/hamiltonian.jl:47
    ┌ Warning: The current proposal will be rejected due to numerical error(s).
    │   isfinite.((θ, r, ℓπ, ℓκ)) = (true, false, false, false)
    └ @ AdvancedHMC /Users/tomroschinger/.julia/packages/AdvancedHMC/MIxdK/src/hamiltonian.jl:47
    ┌ Warning: The current proposal will be rejected due to numerical error(s).
    │   isfinite.((θ, r, ℓπ, ℓκ)) = (true, false, false, false)
    └ @ AdvancedHMC /Users/tomroschinger/.julia/packages/AdvancedHMC/MIxdK/src/hamiltonian.jl:47
    ┌ Warning: The current proposal will be rejected due to numerical error(s).
    │   isfinite.((θ, r, ℓπ, ℓκ)) = (true, false, false, false)
    └ @ AdvancedHMC /Users/tomroschinger/.julia/packages/AdvancedHMC/MIxdK/src/hamiltonian.jl:47
    ┌ Warning: The current proposal will be rejected due to numerical error(s).
    │   isfinite.((θ, r, ℓπ, ℓκ)) = (true, false, false, false)
    └ @ AdvancedHMC /Users/tomroschinger/.julia/packages/AdvancedHMC/MIxdK/src/hamiltonian.jl:47
    ┌ Warning: The current proposal will be rejected due to numerical error(s).
    │   isfinite.((θ, r, ℓπ, ℓκ)) = (true, false, false, false)
    └ @ AdvancedHMC /Users/tomroschinger/.julia/packages/AdvancedHMC/MIxdK/src/hamiltonian.jl:47
    ┌ Warning: The current proposal will be rejected due to numerical error(s).
    │   isfinite.((θ, r, ℓπ, ℓκ)) = (true, false, false, false)
    └ @ AdvancedHMC /Users/tomroschinger/.julia/packages/AdvancedHMC/MIxdK/src/hamiltonian.jl:47
    ┌ Warning: The current proposal will be rejected due to numerical error(s).
    │   isfinite.((θ, r, ℓπ, ℓκ)) = (true, false, false, false)
    └ @ AdvancedHMC /Users/tomroschinger/.julia/packages/AdvancedHMC/MIxdK/src/hamiltonian.jl:47
    ┌ Warning: The current proposal will be rejected due to numerical error(s).
    │   isfinite.((θ, r, ℓπ, ℓκ)) = (true, false, false, false)
    └ @ AdvancedHMC /Users/tomroschinger/.julia/packages/AdvancedHMC/MIxdK/src/hamiltonian.jl:47
    ┌ Info: Found initial step size
    │   ϵ = 1.220703125e-5
    └ @ Turing.Inference /Users/tomroschinger/.julia/packages/Turing/O1Pn0/src/inference/hmc.jl:195
    Sampling: 100%|█████████████████████████████████████████| Time: 0:00:02

```julia
plot_opt_length_star(n, l_0, gap, 300l_0, 2fl, collect(0:40))
```

    ┌ Warning: The current proposal will be rejected due to numerical error(s).
    │   isfinite.((θ, r, ℓπ, ℓκ)) = (true, false, false, false)
    └ @ AdvancedHMC /Users/tomroschinger/.julia/packages/AdvancedHMC/MIxdK/src/hamiltonian.jl:47
    ┌ Warning: The current proposal will be rejected due to numerical error(s).
    │   isfinite.((θ, r, ℓπ, ℓκ)) = (true, false, false, false)
    └ @ AdvancedHMC /Users/tomroschinger/.julia/packages/AdvancedHMC/MIxdK/src/hamiltonian.jl:47
    ┌ Warning: The current proposal will be rejected due to numerical error(s).
    │   isfinite.((θ, r, ℓπ, ℓκ)) = (true, false, false, false)
    └ @ AdvancedHMC /Users/tomroschinger/.julia/packages/AdvancedHMC/MIxdK/src/hamiltonian.jl:47
    ┌ Warning: The current proposal will be rejected due to numerical error(s).
    │   isfinite.((θ, r, ℓπ, ℓκ)) = (true, false, false, false)
    └ @ AdvancedHMC /Users/tomroschinger/.julia/packages/AdvancedHMC/MIxdK/src/hamiltonian.jl:47
    ┌ Warning: The current proposal will be rejected due to numerical error(s).
    │   isfinite.((θ, r, ℓπ, ℓκ)) = (true, false, false, false)
    └ @ AdvancedHMC /Users/tomroschinger/.julia/packages/AdvancedHMC/MIxdK/src/hamiltonian.jl:47
    ┌ Warning: The current proposal will be rejected due to numerical error(s).
    │   isfinite.((θ, r, ℓπ, ℓκ)) = (true, false, false, false)
    └ @ AdvancedHMC /Users/tomroschinger/.julia/packages/AdvancedHMC/MIxdK/src/hamiltonian.jl:47
    ┌ Warning: The current proposal will be rejected due to numerical error(s).
    │   isfinite.((θ, r, ℓπ, ℓκ)) = (true, false, false, false)
    └ @ AdvancedHMC /Users/tomroschinger/.julia/packages/AdvancedHMC/MIxdK/src/hamiltonian.jl:47
    ┌ Warning: The current proposal will be rejected due to numerical error(s).
    │   isfinite.((θ, r, ℓπ, ℓκ)) = (true, false, false, false)
    └ @ AdvancedHMC /Users/tomroschinger/.julia/packages/AdvancedHMC/MIxdK/src/hamiltonian.jl:47
    ┌ Warning: The current proposal will be rejected due to numerical error(s).
    │   isfinite.((θ, r, ℓπ, ℓκ)) = (true, false, false, false)
    └ @ AdvancedHMC /Users/tomroschinger/.julia/packages/AdvancedHMC/MIxdK/src/hamiltonian.jl:47
    ┌ Warning: The current proposal will be rejected due to numerical error(s).
    │   isfinite.((θ, r, ℓπ, ℓκ)) = (true, false, false, false)
    └ @ AdvancedHMC /Users/tomroschinger/.julia/packages/AdvancedHMC/MIxdK/src/hamiltonian.jl:47
    ┌ Warning: The current proposal will be rejected due to numerical error(s).
    │   isfinite.((θ, r, ℓπ, ℓκ)) = (true, false, false, false)
    └ @ AdvancedHMC /Users/tomroschinger/.julia/packages/AdvancedHMC/MIxdK/src/hamiltonian.jl:47
    ┌ Warning: The current proposal will be rejected due to numerical error(s).
    │   isfinite.((θ, r, ℓπ, ℓκ)) = (true, false, false, false)
    └ @ AdvancedHMC /Users/tomroschinger/.julia/packages/AdvancedHMC/MIxdK/src/hamiltonian.jl:47
    ┌ Warning: The current proposal will be rejected due to numerical error(s).
    │   isfinite.((θ, r, ℓπ, ℓκ)) = (true, false, false, false)
    └ @ AdvancedHMC /Users/tomroschinger/.julia/packages/AdvancedHMC/MIxdK/src/hamiltonian.jl:47
    ┌ Warning: The current proposal will be rejected due to numerical error(s).
    │   isfinite.((θ, r, ℓπ, ℓκ)) = (true, false, false, false)
    └ @ AdvancedHMC /Users/tomroschinger/.julia/packages/AdvancedHMC/MIxdK/src/hamiltonian.jl:47
    ┌ Info: Found initial step size
    │   ϵ = 1.52587890625e-6
    └ @ Turing.Inference /Users/tomroschinger/.julia/packages/Turing/O1Pn0/src/inference/hmc.jl:195
    Sampling: 100%|█████████████████████████████████████████| Time: 0:00:02

## Gamma Star

```julia
f0=100l_0
fl=1.6l_0

κ_arr = collect(0:30)
κ_log = exp.(range(0, length=length(κ_arr), stop=log(κ_arr[end] + 1))) .- 1

p_lin = plot(
    κ_arr,
    [l_star(n, l_0, gap, f0, fl, κ)[1] for κ in κ_arr],
    xlabel="κ",
    ylabel="γ_star"
)
p_log = plot(
    log.(1 .+ κ_log),
    [log(l_star(n, l_0, gap, f0, fl, κ)[1]) for κ in κ_log],
    xlabel="log(1+κ)",
    ylabel="log(sstar)"
)

y = 0:0.01:1
x = 10:150
ff(x, y) = begin
        F(y, x, n, l_0, gap, f0, fl)/l_0
    end
X = repeat(reshape(x, 1, :), length(y), 1)
Y = repeat(y, 1, length(x))
Z = map(ff, X, Y)
p1 = contour(x, y, ff, fill = true, color=:viridis, levels=20)

gs = [l_star(n, l_0, gap, f0, fl, κ)[1] for κ in κ_arr]
ls = [l_star(n, l_0, gap, f0, fl, κ)[end-1] for κ in κ_arr]

scatter!(p1, ls, gs)

plot([p_lin, p_log, p1]..., size=(1100, 300), layout=(1, 3))
```

## S_Star

```julia
sstar(l, n, l_0, gap, f0, fl, κ) = F(γ_star(l, n, l_0, gap, f0, fl, κ) - 1/l, l, n, l_0, gap, f0, fl) - F(γ_star(l, n, l_0, gap, f0, fl, κ), l, n, l_0, gap, f0, fl)

```

    sstar (generic function with 1 method)

```julia
function plot_sstar(n, l_0, gap, f0, fl, κ_arr=collect(0:30))

    l_opt = optimal_length_star(n, l_0, gap, f0, fl, κ_arr)
    star = [sstar(l, n, l_0, gap, f0, fl, κ) for (l, κ) in zip(l_opt, κ_arr)]

    κ_log = exp.(range(0, length=length(κ_arr), stop=log(κ_arr[end] + 1))) .- 1
    star_log = [sstar(l, n, l_0, gap, f0, fl, κ) for (l, κ) in zip(l_opt, κ_log)]

    logx = log.(1 .+ κ_log)
    logy = log.(star_log)

    chn = sample(fit_exponent(logx[2:end], logy[2:end]),NUTS(0.65), 3_000)

    p_lin = plot(
        κ_arr,
        star,
        xlabel="κ",
        ylabel="s_star",
        title="f0=$(f0/l_0)l_0, fl=$(fl/l_0)l_0"
    )
    p_log = plot(log.(1 .+ κ_log), log.(star_log))
    plot!(
        p_log,
        log.(1 .+ κ_log),
        log.((1 .+ κ_log) .^(mean(chn[:a]))) .+ mean(chn[:b]),
        linestyle=:dash,
        xlabel="log(1+κ)",
        ylabel="log(s_star)",
        title=@sprintf "LogLog with slope %.3f" mean(chn[:a])

    )

    return plot([p_lin, p_log]..., size=(900, 450))
end
```

    plot_sstar (generic function with 2 methods)

```julia
plot_sstar(n, l_0, gap, 40l_0, l_0)
```

    ┌ Warning: The current proposal will be rejected due to numerical error(s).
    │   isfinite.((θ, r, ℓπ, ℓκ)) = (true, false, false, false)
    └ @ AdvancedHMC /Users/tomroschinger/.julia/packages/AdvancedHMC/MIxdK/src/hamiltonian.jl:47
    ┌ Warning: The current proposal will be rejected due to numerical error(s).
    │   isfinite.((θ, r, ℓπ, ℓκ)) = (true, false, false, false)
    └ @ AdvancedHMC /Users/tomroschinger/.julia/packages/AdvancedHMC/MIxdK/src/hamiltonian.jl:47
    ┌ Warning: The current proposal will be rejected due to numerical error(s).
    │   isfinite.((θ, r, ℓπ, ℓκ)) = (true, false, false, false)
    └ @ AdvancedHMC /Users/tomroschinger/.julia/packages/AdvancedHMC/MIxdK/src/hamiltonian.jl:47
    ┌ Warning: The current proposal will be rejected due to numerical error(s).
    │   isfinite.((θ, r, ℓπ, ℓκ)) = (true, false, false, false)
    └ @ AdvancedHMC /Users/tomroschinger/.julia/packages/AdvancedHMC/MIxdK/src/hamiltonian.jl:47
    ┌ Warning: The current proposal will be rejected due to numerical error(s).
    │   isfinite.((θ, r, ℓπ, ℓκ)) = (true, false, false, false)
    └ @ AdvancedHMC /Users/tomroschinger/.julia/packages/AdvancedHMC/MIxdK/src/hamiltonian.jl:47
    ┌ Warning: The current proposal will be rejected due to numerical error(s).
    │   isfinite.((θ, r, ℓπ, ℓκ)) = (true, false, false, false)
    └ @ AdvancedHMC /Users/tomroschinger/.julia/packages/AdvancedHMC/MIxdK/src/hamiltonian.jl:47
    ┌ Warning: The current proposal will be rejected due to numerical error(s).
    │   isfinite.((θ, r, ℓπ, ℓκ)) = (true, false, false, false)
    └ @ AdvancedHMC /Users/tomroschinger/.julia/packages/AdvancedHMC/MIxdK/src/hamiltonian.jl:47
    ┌ Warning: The current proposal will be rejected due to numerical error(s).
    │   isfinite.((θ, r, ℓπ, ℓκ)) = (true, false, false, false)
    └ @ AdvancedHMC /Users/tomroschinger/.julia/packages/AdvancedHMC/MIxdK/src/hamiltonian.jl:47
    ┌ Warning: The current proposal will be rejected due to numerical error(s).
    │   isfinite.((θ, r, ℓπ, ℓκ)) = (true, false, false, false)
    └ @ AdvancedHMC /Users/tomroschinger/.julia/packages/AdvancedHMC/MIxdK/src/hamiltonian.jl:47
    ┌ Warning: The current proposal will be rejected due to numerical error(s).
    │   isfinite.((θ, r, ℓπ, ℓκ)) = (true, false, false, false)
    └ @ AdvancedHMC /Users/tomroschinger/.julia/packages/AdvancedHMC/MIxdK/src/hamiltonian.jl:47
    ┌ Warning: The current proposal will be rejected due to numerical error(s).
    │   isfinite.((θ, r, ℓπ, ℓκ)) = (true, false, false, false)
    └ @ AdvancedHMC /Users/tomroschinger/.julia/packages/AdvancedHMC/MIxdK/src/hamiltonian.jl:47
    ┌ Warning: The current proposal will be rejected due to numerical error(s).
    │   isfinite.((θ, r, ℓπ, ℓκ)) = (true, false, false, false)
    └ @ AdvancedHMC /Users/tomroschinger/.julia/packages/AdvancedHMC/MIxdK/src/hamiltonian.jl:47
    ┌ Warning: The current proposal will be rejected due to numerical error(s).
    │   isfinite.((θ, r, ℓπ, ℓκ)) = (true, false, false, false)
    └ @ AdvancedHMC /Users/tomroschinger/.julia/packages/AdvancedHMC/MIxdK/src/hamiltonian.jl:47
    ┌ Warning: The current proposal will be rejected due to numerical error(s).
    │   isfinite.((θ, r, ℓπ, ℓκ)) = (true, false, false, false)
    └ @ AdvancedHMC /Users/tomroschinger/.julia/packages/AdvancedHMC/MIxdK/src/hamiltonian.jl:47
    ┌ Warning: The current proposal will be rejected due to numerical error(s).
    │   isfinite.((θ, r, ℓπ, ℓκ)) = (true, false, false, false)
    └ @ AdvancedHMC /Users/tomroschinger/.julia/packages/AdvancedHMC/MIxdK/src/hamiltonian.jl:47
    ┌ Info: Found initial step size
    │   ϵ = 7.62939453125e-7
    └ @ Turing.Inference /Users/tomroschinger/.julia/packages/Turing/O1Pn0/src/inference/hmc.jl:195
    ┌ Warning: The current proposal will be rejected due to numerical error(s).
    │   isfinite.((θ, r, ℓπ, ℓκ)) = (true, true, true, false)
    └ @ AdvancedHMC /Users/tomroschinger/.julia/packages/AdvancedHMC/MIxdK/src/hamiltonian.jl:47
    Sampling: 100%|█████████████████████████████████████████| Time: 0:00:00

```julia
plot_sstar(n, l_0, gap, 200l_0, 1.6l_0)
```

    ┌ Warning: The current proposal will be rejected due to numerical error(s).
    │   isfinite.((θ, r, ℓπ, ℓκ)) = (true, false, false, false)
    └ @ AdvancedHMC /Users/tomroschinger/.julia/packages/AdvancedHMC/MIxdK/src/hamiltonian.jl:47
    ┌ Warning: The current proposal will be rejected due to numerical error(s).
    │   isfinite.((θ, r, ℓπ, ℓκ)) = (true, false, false, false)
    └ @ AdvancedHMC /Users/tomroschinger/.julia/packages/AdvancedHMC/MIxdK/src/hamiltonian.jl:47
    ┌ Warning: The current proposal will be rejected due to numerical error(s).
    │   isfinite.((θ, r, ℓπ, ℓκ)) = (true, false, false, false)
    └ @ AdvancedHMC /Users/tomroschinger/.julia/packages/AdvancedHMC/MIxdK/src/hamiltonian.jl:47
    ┌ Warning: The current proposal will be rejected due to numerical error(s).
    │   isfinite.((θ, r, ℓπ, ℓκ)) = (true, false, false, false)
    └ @ AdvancedHMC /Users/tomroschinger/.julia/packages/AdvancedHMC/MIxdK/src/hamiltonian.jl:47
    ┌ Warning: The current proposal will be rejected due to numerical error(s).
    │   isfinite.((θ, r, ℓπ, ℓκ)) = (true, false, false, false)
    └ @ AdvancedHMC /Users/tomroschinger/.julia/packages/AdvancedHMC/MIxdK/src/hamiltonian.jl:47
    ┌ Warning: The current proposal will be rejected due to numerical error(s).
    │   isfinite.((θ, r, ℓπ, ℓκ)) = (true, false, false, false)
    └ @ AdvancedHMC /Users/tomroschinger/.julia/packages/AdvancedHMC/MIxdK/src/hamiltonian.jl:47
    ┌ Warning: The current proposal will be rejected due to numerical error(s).
    │   isfinite.((θ, r, ℓπ, ℓκ)) = (true, false, false, false)
    └ @ AdvancedHMC /Users/tomroschinger/.julia/packages/AdvancedHMC/MIxdK/src/hamiltonian.jl:47
    ┌ Warning: The current proposal will be rejected due to numerical error(s).
    │   isfinite.((θ, r, ℓπ, ℓκ)) = (true, false, false, false)
    └ @ AdvancedHMC /Users/tomroschinger/.julia/packages/AdvancedHMC/MIxdK/src/hamiltonian.jl:47
    ┌ Warning: The current proposal will be rejected due to numerical error(s).
    │   isfinite.((θ, r, ℓπ, ℓκ)) = (true, false, false, false)
    └ @ AdvancedHMC /Users/tomroschinger/.julia/packages/AdvancedHMC/MIxdK/src/hamiltonian.jl:47
    ┌ Warning: The current proposal will be rejected due to numerical error(s).
    │   isfinite.((θ, r, ℓπ, ℓκ)) = (true, false, false, false)
    └ @ AdvancedHMC /Users/tomroschinger/.julia/packages/AdvancedHMC/MIxdK/src/hamiltonian.jl:47
    ┌ Warning: The current proposal will be rejected due to numerical error(s).
    │   isfinite.((θ, r, ℓπ, ℓκ)) = (true, false, false, false)
    └ @ AdvancedHMC /Users/tomroschinger/.julia/packages/AdvancedHMC/MIxdK/src/hamiltonian.jl:47
    ┌ Info: Found initial step size
    │   ϵ = 6.103515625e-6
    └ @ Turing.Inference /Users/tomroschinger/.julia/packages/Turing/O1Pn0/src/inference/hmc.jl:195
    Sampling: 100%|█████████████████████████████████████████| Time: 0:00:00

```julia
plot_sstar(n, l_0, gap, 300l_0, 2l_0)
```

    ┌ Warning: The current proposal will be rejected due to numerical error(s).
    │   isfinite.((θ, r, ℓπ, ℓκ)) = (true, false, false, false)
    └ @ AdvancedHMC /Users/tomroschinger/.julia/packages/AdvancedHMC/MIxdK/src/hamiltonian.jl:47
    ┌ Warning: The current proposal will be rejected due to numerical error(s).
    │   isfinite.((θ, r, ℓπ, ℓκ)) = (true, false, false, false)
    └ @ AdvancedHMC /Users/tomroschinger/.julia/packages/AdvancedHMC/MIxdK/src/hamiltonian.jl:47
    ┌ Warning: The current proposal will be rejected due to numerical error(s).
    │   isfinite.((θ, r, ℓπ, ℓκ)) = (true, false, false, false)
    └ @ AdvancedHMC /Users/tomroschinger/.julia/packages/AdvancedHMC/MIxdK/src/hamiltonian.jl:47
    ┌ Warning: The current proposal will be rejected due to numerical error(s).
    │   isfinite.((θ, r, ℓπ, ℓκ)) = (true, false, false, false)
    └ @ AdvancedHMC /Users/tomroschinger/.julia/packages/AdvancedHMC/MIxdK/src/hamiltonian.jl:47
    ┌ Warning: The current proposal will be rejected due to numerical error(s).
    │   isfinite.((θ, r, ℓπ, ℓκ)) = (true, false, false, false)
    └ @ AdvancedHMC /Users/tomroschinger/.julia/packages/AdvancedHMC/MIxdK/src/hamiltonian.jl:47
    ┌ Warning: The current proposal will be rejected due to numerical error(s).
    │   isfinite.((θ, r, ℓπ, ℓκ)) = (true, false, false, false)
    └ @ AdvancedHMC /Users/tomroschinger/.julia/packages/AdvancedHMC/MIxdK/src/hamiltonian.jl:47
    ┌ Warning: The current proposal will be rejected due to numerical error(s).
    │   isfinite.((θ, r, ℓπ, ℓκ)) = (true, false, false, false)
    └ @ AdvancedHMC /Users/tomroschinger/.julia/packages/AdvancedHMC/MIxdK/src/hamiltonian.jl:47
    ┌ Warning: The current proposal will be rejected due to numerical error(s).
    │   isfinite.((θ, r, ℓπ, ℓκ)) = (true, false, false, false)
    └ @ AdvancedHMC /Users/tomroschinger/.julia/packages/AdvancedHMC/MIxdK/src/hamiltonian.jl:47
    ┌ Warning: The current proposal will be rejected due to numerical error(s).
    │   isfinite.((θ, r, ℓπ, ℓκ)) = (true, false, false, false)
    └ @ AdvancedHMC /Users/tomroschinger/.julia/packages/AdvancedHMC/MIxdK/src/hamiltonian.jl:47
    ┌ Warning: The current proposal will be rejected due to numerical error(s).
    │   isfinite.((θ, r, ℓπ, ℓκ)) = (true, false, false, false)
    └ @ AdvancedHMC /Users/tomroschinger/.julia/packages/AdvancedHMC/MIxdK/src/hamiltonian.jl:47
    ┌ Warning: The current proposal will be rejected due to numerical error(s).
    │   isfinite.((θ, r, ℓπ, ℓκ)) = (true, false, false, false)
    └ @ AdvancedHMC /Users/tomroschinger/.julia/packages/AdvancedHMC/MIxdK/src/hamiltonian.jl:47
    ┌ Warning: The current proposal will be rejected due to numerical error(s).
    │   isfinite.((θ, r, ℓπ, ℓκ)) = (true, false, false, false)
    └ @ AdvancedHMC /Users/tomroschinger/.julia/packages/AdvancedHMC/MIxdK/src/hamiltonian.jl:47
    ┌ Warning: The current proposal will be rejected due to numerical error(s).
    │   isfinite.((θ, r, ℓπ, ℓκ)) = (true, false, false, false)
    └ @ AdvancedHMC /Users/tomroschinger/.julia/packages/AdvancedHMC/MIxdK/src/hamiltonian.jl:47
    ┌ Info: Found initial step size
    │   ϵ = 3.0517578125e-6
    └ @ Turing.Inference /Users/tomroschinger/.julia/packages/Turing/O1Pn0/src/inference/hmc.jl:195
    Sampling: 100%|█████████████████████████████████████████| Time: 0:00:00

```julia

```


# Solve cartpole with REINFORCE

> By Jonas Busk ([jbusk@dtu.dk](mailto:jbusk@dtu.dk))

**2019 update:** Changes have been made to the display of environments due to the previous `viewer` being incompatible with newer versions of Gym.

In this part, we will create an agent that can learn to solve the [cartpole problem](https://gym.openai.com/envs/CartPole-v0/) from OpenAI Gym by applying a simple policy gradient method called REINFORCE.
In the cartpole problem, we need to balance a pole on a cart that moves along a track by applying left and right forces to the cart.

We will implement a probabilistic policy, that given a state of the environment, $s$, outputs a probability distribution over available actions, $a$:

$$
p_\theta(a|s)
$$

The policy is a neural network with parameters $\theta$ that can be trained with gradient descent.
When the set of available actions is discrete, we can use a network with softmax output do describe the distribution.
The core idea of training the policy network is quite simple: *we want to maximize the expected total reward by increasing the probability of good actions and decreasing the probability of bad actions*.

To achieve this, we apply the gradient of the expected discounted total reward (return):

$$
\begin{align}
\nabla_\theta \mathbb{E}[R|\theta] &= \nabla_\theta \int p_\theta(a|s) R({a}) \, da \\
&= \int \nabla_\theta p_\theta(a|s) R(a)  \, da \\
&= \int p_\theta(a|s) \nabla_\theta \log p_\theta(a|s) R(a) \, da \\
&= \mathbb{E}[\nabla_\theta \log p_\theta(a|s) R(a)]
\end{align}
$$

by definition of expectation and using the identity

$$
\nabla_\theta p_\theta(a|s) = p_\theta(a|s) \nabla_\theta \log p_\theta(a|s) \ .
$$

The expectation cannot be evaluated analytically, but we have an environment simulator that when supplied with our current policy $p_\theta(a|s)$ can return a sequence of *actions*, *states* and *rewards*. This allows us to replace the integral with a Monte Carlo average:

$$
\nabla_\theta \mathbb{E}[R|\theta] \approx \frac{1}{T} \sum_{t=0}^T \nabla_\theta \log p_\theta(a_t|s_t) R_t \ ,
$$

which is our final gradient estimator, also known as REINFORCE. In the Monte Carlo estimator we run the environment simulator for a predefined number of steps with actions chosen stochastically according to the current stochastic action network $p_\theta(a|s)$.

*Note: For simple reinforcement learning problems (like the one we will address in this exercise) there are simpler methods that work just fine. However, the Policy Gradient method (with some extensions) has been shown to also work well for complex problems with high dimensional inputs and many parameters, where simple methods become inadequate.*

## Policy gradient

```python
%matplotlib inline
import matplotlib.pyplot as plt
import numpy as np
import torch
import torch.nn as nn
import torch.optim as optim
import torch.nn.functional as F
import gym
from gym import wrappers
```

First we create the environment:

```python
env = gym.make('CartPole-v0') # Create environment
```

A state in this environment is four numbers describing the position of the cart along with the angle and speed of the pole.
There are two available actions: push the cart *left* or *right* encoded as 0 and 1.

```python
s = env.reset()
a = env.action_space.sample()
print('sample state:', s)
print('sample action:', a )
```

    sample state: [-0.01494115  0.04810255 -0.03759179  0.04219313]
    sample action: 0

**Note:** you will likely not be able to render environments in a Google Colab instance. Therefore, it may be beneficial for you to run this week's notebooks locally and/or team up with another student if you do not have a local environment.

The function below is used for display:

```python
def show_replay():
    """
    Not-so-elegant way to display the MP4 file generated by the Monitor wrapper inside a notebook.
    The Monitor wrapper dumps the replay to a local file that we then display as a HTML video object.
    """
    import io
    import base64
    from IPython.display import HTML
    video = io.open('./gym-results/openaigym.video.%s.video000000.mp4' % env.file_infix, 'r+b').read()
    encoded = base64.b64encode(video)
    return HTML(data='''
        '''
    .format(encoded.decode('ascii')))
```

Let us see how the environment looks when we just take random actions. Note that the episode ends when the pole either 1) is more than 15 degrees from vertical, 2) more outside of the frame or 3) the pole is successfully balanced for some fixed duration.

```python
env = gym.make('CartPole-v0') # Create environment
env = wrappers.Monitor(env, "./gym-results", force=True) # Create wrapper to display environment
env.reset() # Reset environment

# Run environment
while True:
    env.render() # Render environment
    action = env.action_space.sample() # Get a random action
    _, _, done, _ = env.step(action) # Take a step
    if done: break # Break if environment is done

env.close() # Close environment
show_replay()
```

Taking random actions does not do a very good job at balancing the pole. Let us now apply the Policy Gradient method described above to solve this task!

Let's first define our network and helper functions.

```python
class PolicyNet(nn.Module):
    """Policy network"""

    def __init__(self, n_inputs, n_hidden, n_outputs, learning_rate):
        super(PolicyNet, self).__init__()
        # network
        self.hidden1 = nn.Linear(n_inputs, n_hidden)
        self.hidden2 = nn.Linear(n_hidden, 48)

        self.relu = nn.ReLU()
        self.out = nn.Linear(48, n_outputs)
        # training
        self.optimizer = optim.Adam(self.parameters(), lr=learning_rate)

    def forward(self, x):
        x = self.relu(self.hidden1(x))
        x = self.relu(self.hidden2(x))
        x = self.out(x)
        return F.softmax(x, dim=1)

    def loss(self, action_probabilities, returns):
        return -torch.mean(torch.mul(torch.log(action_probabilities), returns))
```

```python
def compute_returns(rewards, discount_factor):
    """Compute discounted returns."""
    returns = np.zeros(len(rewards))
    returns[-1] = rewards[-1]
    for t in reversed(range(len(rewards)-1)):
        returns[t] = rewards[t] + discount_factor * returns[t+1]
    return returns
```

To start with, our policy will be a rather simple neural network with one hidden layer. We can retrieve the shape of the state space (input) and action space (output) from the environment.

```python
n_inputs = env.observation_space.shape[0]
n_hidden = 24
n_outputs = env.action_space.n

print('state shape:', n_inputs)
print('action shape:', n_outputs)
```

    state shape: 4
    action shape: 2

```python
# training settings

num_episodes = 1000
rollout_limit = 500 # max rollout length
discount_factor = 1.0 # reward discount factor (gamma), 1.0 = no discount
learning_rate = 0.001 # you know this by now
val_freq = 100 # validation frequency

# setup policy network

policy = PolicyNet(n_inputs, n_hidden, n_outputs, learning_rate)

# train policy network

try:
    training_rewards, losses = [], []
    print('start training')
    for i in range(num_episodes):
        rollout = []
        s = env.reset()
        for j in range(rollout_limit):
            # generate rollout by iteratively evaluating the current policy on the environment
            with torch.no_grad():
                a_prob = policy(torch.from_numpy(np.atleast_2d(s)).float())
                a = torch.multinomial(a_prob, num_samples=1).squeeze().numpy()
            s1, r, done, _ = env.step(a)
            rollout.append((s, a, r))
            s = s1
            if done: break
        # prepare batch
        rollout = np.array(rollout, dtype=object)
        states = np.vstack(rollout[:,0])
        actions = np.vstack(rollout[:,1])
        rewards = np.array(rollout[:,2], dtype=float)
        returns = compute_returns(rewards, discount_factor)

        # policy gradient update
        policy.optimizer.zero_grad()
        a_probs = policy(torch.from_numpy(states).float()).gather(1, torch.from_numpy(actions)).view(-1)
        loss = policy.loss(a_probs, torch.from_numpy(returns).float())
        loss.backward()
        policy.optimizer.step()
        # bookkeeping
        training_rewards.append(sum(rewards))
        losses.append(loss.item())
        # print
        if (i+1) % val_freq == 0:
            # validation
            validation_rewards = []
            for _ in range(10):
                s = env.reset()
                reward = 0
                for _ in range(rollout_limit):
                    with torch.no_grad():
                        a_prob = policy(torch.from_numpy(np.atleast_2d(s)).float())
                        a = a_prob.argmax().item()
                    s, r, done, _ = env.step(a)
                    reward += r
                    if done: break
                validation_rewards.append(reward)
            print('{:4d}. mean training reward: {:6.2f}, mean validation reward: {:6.2f}, mean loss: {:7.4f}'.format(i+1, np.mean(training_rewards[-val_freq:]), np.mean(validation_rewards), np.mean(losses[-val_freq:])))
    print('done')
except KeyboardInterrupt:
    print('interrupt')
```

    start training
     100. mean training reward:  23.93, mean validation reward:  75.20, mean loss:  8.6052
     200. mean training reward:  30.77, mean validation reward: 152.20, mean loss: 10.7694
     300. mean training reward:  45.00, mean validation reward:  63.20, mean loss: 14.9648
     400. mean training reward: 104.37, mean validation reward: 200.00, mean loss: 32.2687
     500. mean training reward: 135.34, mean validation reward: 105.80, mean loss: 40.1283
     600. mean training reward: 148.13, mean validation reward: 197.20, mean loss: 42.8639
     700. mean training reward: 191.08, mean validation reward: 200.00, mean loss: 55.1002
     800. mean training reward: 192.69, mean validation reward: 200.00, mean loss: 55.4097
     900. mean training reward: 194.71, mean validation reward: 200.00, mean loss: 55.1190
    1000. mean training reward: 188.10, mean validation reward: 197.90, mean loss: 53.3239
    done

```python
# plot results
def moving_average(a, n=10) :
    ret = np.cumsum(a, dtype=float)
    ret[n:] = ret[n:] - ret[:-n]
    return ret / n

plt.figure(figsize=(16,6))
plt.subplot(211)
plt.plot(range(1, len(training_rewards)+1), training_rewards, label='training reward')
plt.plot(moving_average(training_rewards))
plt.xlabel('episode'); plt.ylabel('reward')
plt.xlim((0, len(training_rewards)))
plt.legend(loc=4); plt.grid()
plt.subplot(212)
plt.plot(range(1, len(losses)+1), losses, label='loss')
plt.plot(moving_average(losses))
plt.xlabel('episode'); plt.ylabel('loss')
plt.xlim((0, len(losses)))
plt.legend(loc=4); plt.grid()
plt.tight_layout(); plt.show()
```

Now let's review the solution!

```python
env = wrappers.Monitor(env, "./gym-results", force=True) # Create wrapper to display environment
s = env.reset()

for _ in range(500):
    env.render()
    a = policy(torch.from_numpy(np.atleast_2d(s)).float()).argmax().item()
    s, r, done, _ = env.step(a)
    if done: break

env.close()
show_replay()
```

## Reducing variance

By default, this gradient estimator has high variance and therefore variance reduction becomes important to learn more complex tasks.
We can reduce variance by subtracting a baseline from the returns, which is unbiased in expectation:

$$
\nabla_\theta \mathbb{E}[R|\theta] \approx \frac{1}{T} \sum_{t=0}^T \nabla_\theta \log p_\theta(a_t|s_t) (R_t-b_t) \ ,
$$

where the baseline, $b_t$, is estimated by the return a timestep $t$ averaged over $V$ rollouts.

$$
b_t = \frac{1}{V} \sum_{v=1}^V R_t^{(v)} \ .
$$

## Exercises

Now it is your turn! Make sure you read and understand the code, then play around with it and try to make it learn better and faster.

Experiment with the:

* number of episodes
* discount factor
* learning rate
* network layers

### Exercise 1

*Describe any changes you made to the code and why you think they improve the agent. Are you able to get solutions consistently?*

**Answer:**

The first thing I looked at, was the number of layers in the neural network. Here I ended with two linear layers with 24, and 48 hidden units.

Next I looked at the number of episodes. Here the model began reaching a solve (195 points) at around 1000 episodes, dropping down, and then almost consistently hitting a reward of 200. Adding to the number of episodes after 1600, did not seem to change much, and 1600 might even be a bit excessive.

Next I looked at the discount factor. Changing the value from 1.0 saw very poor results, and droped both training and validation reward significantly. This makes sense, as the further from 1 the discount factor deviates, the less previous events matter, and in the case of this problem, the actions are very pinpoint, and a small mistake, can lead to irreversible results in further states. This problem is pretty apparent, if the reader has ever played "Daddy Longlegs", which has the same mechanic.

Lastly I looked at the learning rate. Here the already established value of 0.001. My thought here was, that it seemed the model pretty quickly found good solutions, and making the learning rate aggresive, might lead to even quicker training results, but this backfired, and the model did not learn anything at all. Making the learning rate even lower, made the model - as expected - use significantly more time, but arrived at the same performance, as the original value of .001, so I kept it there.

### Exercise 2

*Consider the following sequence of rewards produced by an agent interacting with an environment for 10 timesteps:*

[0, 1, 1, 1, 0, 1, 1, 0, 0, 0]

* *What is the total reward?*
* *What is the total future reward in each timestep?*
* *What is the discounted future reward in each timestep if $\gamma = 0.9$?*

*Hint: See introdution notebook.*

**Answer:**

* `Total reward: 5`
* `total future reward in each timestep: [5, 5, 4, 3, 2, 2, 1, 0, 0, 0]`
* `Discounted future reward in each timestep: [3.560931, 3.95659, 3.2851, 2.539, 1.71, 1.9, 1.0, 0.0, 0.0, 0.0]`

### Exercise 3

*In the training output, you will sometimes observe the validation reward starts out lower than the training reward but as training progresses they cross over and the validation reward becomes higher than the training reward. How can you explain this behavior?*

*Hint: Do we use the policy network in the same way during training and validation?*

**Answer:**

This is as far as i can see due to the use of torch.multinomial, and argmax. When using torch.multinomial, instead of choosing the best estimated action, which argmax would return, it instead makes a multinomial distribution of the 2-indices-vector, and draws a random sample from this. Early on, as the output is binary (right or left), this random action might in some cases, as mentioned in the question, lead to a better result, as the model is not well polished yet. Later on, the randomness however, reduces the success of the training model. Although one would expect the model to be more sure in its actions (maybe the a_prob goes from [.4, .6] to [.999, .001]), the randomness, can still hinder the perfomance of the model.

### Exercise 4

*How does the policy gradient method we have used address the exploration-exploitation dilemma?*

*Hint: See the introduction notebook about exploration-exploitation.*

**Answer:**

*Answer here...*

### Exercise 5 [optional]

Extend the code above to reduce variance of the gradient estimator by computing and subtracting the baseline estimate.

*Hint: You need to sample a batch of rollouts (now we sample just one) for each update in order to compute the baseline, $b_t$.*

```python
# Calculator
rewards = [0, 1, 1, 1, 0, 1, 1, 0, 0, 0]
gamma = .9

total_future = sum(rewards)
future_rewards = [sum(rewards[i:]) for i in range(len(rewards))]
discounted_rewards = []
for i in range(len(rewards)):
    discounted_rewards.append(sum([x*(gamma**t) for t, x in enumerate(rewards[i:])]))

print(f"{total_future=}")
print(f"{future_rewards=}")
print(f"{discounted_rewards=}")
```

    total_future=5
    future_rewards=[5, 5, 4, 3, 2, 2, 1, 0, 0, 0]
    discounted_rewards=[3.560931, 3.95659, 3.2851, 2.539, 1.71, 1.9, 1.0, 0.0, 0.0, 0.0]


.. _nb_custom:

## Custom Variable Type

In the following, we describe a custom variable problem. The variable is a string with a fixed length in our case.
However, we formulate the problem to be easily extended to have a variable length.
The objective function looks as follows:

\begin{align}
\begin{split}
\max f_1(x) & = & \, \# a \\[2mm]
\max f_2(x) & = & \, \# b
\end{split}
\end{align}

The first objective is the number of a's in the string and the second the number of b's.
For instance, for the variable "abdfgdgabb" the $f_1(x)=2$ and $f_2(x)=3$.

```python
import numpy as np
from pymoo.model.problem import Problem

class MyProblem(Problem):

    def __init__(self, n_characters=10):
        super().__init__(n_var=1, n_obj=2, n_constr=0, elementwise_evaluation=True)
        self.n_characters = n_characters
        self.ALPHABET = [c for c in string.ascii_lowercase]

    def _evaluate(self, x, out, *args, **kwargs):
        n_a, n_b = 0, 0
        for c in x[0]:
            if c == 'a':
                n_a += 1
            elif c == 'b':
                n_b += 1

        out["F"] = np.array([- n_a, - n_b], dtype=np.float)

```

The problem definition above define a problem with just one variable. This variable can be considered as a complex object which is in our case a string. The same principle can be used to use other data structure such as trees or lists with variable lengths. Because both objectives have to be maximized, we are minimizing their negative values.

To solve the problem the evolutionary operators sampling, crossover, mutation and duplication check needs to be implemented.
Each of the modules will be shown in the following.

### Sampling

Our sampling method just generates a random string which is equivalent of choosing a random letter from the alphabet (only lower case).
Because of the implementation of having only one variable, we return a matrix with the shape (n,1).

```python
from pymoo.model.sampling import Sampling

class MySampling(Sampling):

    def _do(self, problem, n_samples, **kwargs):
        X = np.full((n_samples, 1), None, dtype=np.object)

        for i in range(n_samples):
            X[i, 0] = "".join([np.random.choice(problem.ALPHABET) for _ in range(problem.n_characters)])

        return X

```

### Crossover

The crossover operator combines parents to create offsprings. In our framework the crossover operator retrieves the input already with predefined matings.
Our crossover randomly picks a character from the first or from the second parent.

```python
from pymoo.model.crossover import Crossover

class MyCrossover(Crossover):
    def __init__(self):

        # define the crossover: number of parents and number of offsprings
        super().__init__(2, 2)

    def _do(self, problem, X, **kwargs):

        # The input of has the following shape (n_parents, n_matings, n_var)
        _, n_matings, n_var = X.shape

        # The output owith the shape (n_offsprings, n_matings, n_var)
        # Because there the number of parents and offsprings are equal it keeps the shape of X
        Y = np.full_like(X, None, dtype=np.object)

        # for each mating provided
        for k in range(n_matings):

            # get the first and the second parent
            a, b = X[0, k, 0], X[1, k, 0]

            # prepare the offsprings
            off_a = ["_"] * problem.n_characters
            off_b = ["_"] * problem.n_characters

            for i in range(problem.n_characters):
                if np.random.random() < 0.5:
                    off_a[i] = a[i]
                    off_b[i] = b[i]
                else:
                    off_a[i] = b[i]
                    off_b[i] = a[i]

            # join the character list and set the output
            Y[0, k, 0], Y[1, k, 0] = "".join(off_a), "".join(off_b)

        return Y
```

### Mutation

The mutation needs to be implemented for our string object as well. We either change the order of the string (40%), randomly pick a new character with a given probability (40%) or leave the string unmodified (20%).

```python
from pymoo.model.mutation import Mutation

class MyMutation(Mutation):
    def __init__(self):
        super().__init__()

    def _do(self, problem, X, **kwargs):

        # for each individual
        for i in range(len(X)):

            r = np.random.random()

            # with a probabilty of 40% - change the order of characters
            if r < 0.4:
                perm = np.random.permutation(problem.n_characters)
                X[i, 0] = "".join(np.array([e for e in X[i, 0]])[perm])

            # also with a probabilty of 40% - change a character randomly
            elif r < 0.8:
                prob = 1 / problem.n_characters
                mut = [c if np.random.random() > prob
                       else np.random.choice(problem.ALPHABET) for c in X[i, 0]]
                X[i, 0] = "".join(mut)

        return X
```

### Duplicates

Moreover, do not underestimate the importance of filtering out duplicates during the evolution. This helps to make sure a mating produces an offspring that does not already exist in the current population. If it does another mating process is triggered until unique individual do exist.

```python
from pymoo.model.duplicate import ElementwiseDuplicateElimination

class MyDuplicateElimination(ElementwiseDuplicateElimination):

    def is_equal(self, a, b):
        return a.X[0] == b.X[0]

```

### Optimize

Finally, we create an algorithm object with all the modules defined above.

```python
import string
import numpy as np

from pymoo.algorithms.nsga2 import NSGA2
from pymoo.optimize import minimize

algorithm = NSGA2(pop_size=11,
                  sampling=MySampling(),
                  crossover=MyCrossover(),
                  mutation=MyMutation(),
                  eliminate_duplicates=MyDuplicateElimination())

res = minimize(MyProblem(),
               algorithm,
               ('n_gen', 100),
               seed=1,
               verbose=False)

```

```python
from pymoo.visualization.scatter import Scatter
Scatter().add(res.F).show()
```

```python
results = res.X[np.argsort(res.F[:, 0])]
count = [np.sum([e == "a" for e in r]) for r in results[:, 0]]
print(np.column_stack([results, count]))
```

    [['aaaaaaaaaa' 10]
     ['baaaaaaaaa' 9]
     ['baaaabaaaa' 8]
     ['aaabbabaaa' 7]
     ['abaabbaaba' 6]
     ['babbaababa' 5]
     ['abbbabbbaa' 4]
     ['abbbaabbbb' 3]
     ['abbbbabbbb' 2]
     ['bbbbbbabbb' 1]
     ['bbbbbbbbbb' 0]]

```python

```


# Hello Entangled World

### Importing all the Necessary qiskit libraries

```python
%matplotlib inline

from qiskit import *
from qiskit.visualization import *
from qiskit.tools.monitor import *
```

```python
circ_bell = QuantumCircuit(2,2)

#we have a two Qubit and two classical bit circuit

circ_bell.h(0)

#then we apply the haadmard gate on the first qubit

circ_bell.cx(0,1)

#now we apply the controlled not gate which have control on the
#first qubit and target on the second qubit

circ_bell.measure(range(2),range(2))

#here we do the measurement on both the states

circ_bell.draw(output='mpl')

#we draw the circuit
```

## Executing on the qasm simulator

```python
backend = Aer.get_backend('qasm_simulator')
job = execute(circ_bell, backend, shots = 1024)

counts = job.result().get_counts()
```

```python
print(counts)
```

    {'00': 534, '11': 490}

It tells us out of the total number of shots we sent, how many were the |00> states and how many were the |11> state

## Finding the State Vector Amplitude

```python
circ_bell2 = QuantumCircuit(2)
circ_bell2.h(0)
circ_bell2.cx(0,1)

backend = Aer.get_backend('statevector_simulator')
job = execute(circ_bell2, backend)

state = job.result().get_statevector()

for i in range(4):
    s = format(i,"b")
    #converting to binary
    s = (2-len(s))*"0"+s
    #prepend zeroes

    print("Amplitude of",s,"=",state[i])

print()
```

    Amplitude of 00 = (0.7071067811865476+0j)
    Amplitude of 01 = 0j
    Amplitude of 10 = 0j
    Amplitude of 11 = (0.7071067811865476+0j)

## Finding the Probabilities

```python
for i in range(4):
    s = format(i,"b")
    #converting to binary
    s = (2-len(s))*"0"+s
    #prepend zeroes

    print("Probability of",s,"=",abs(state[i])**2)
```

    Probability of 00 = 0.5000000000000001
    Probability of 01 = 0.0
    Probability of 10 = 0.0
    Probability of 11 = 0.5000000000000001

### Details about all the Devices available to us

```python
provider = IBMQ.load_account()
backend_overview()
```

    C:\Users\Monit Sharma\Documents\Anaconda\lib\site-packages\qiskit\providers\ibmq\ibmqfactory.py:192: UserWarning: Timestamps in IBMQ backend properties, jobs, and job results are all now in local time instead of UTC.
      warnings.warn('Timestamps in IBMQ backend properties, jobs, and job results '

    ibmq_santiago                ibmq_athens                  ibmq_armonk
    -------------                -----------                  -----------
    Num. Qubits:  5              Num. Qubits:  5              Num. Qubits:  1
    Pending Jobs: 25             Pending Jobs: 32             Pending Jobs: 22
    Least busy:   False          Least busy:   False          Least busy:   False
    Operational:  True           Operational:  True           Operational:  True
    Avg. T1:      107.7          Avg. T1:      79.0           Avg. T1:      184.7
    Avg. T2:      89.9           Avg. T2:      95.6           Avg. T2:      195.4

    ibmq_valencia               ibmq_ourense                 ibmq_vigo
    -------------               ------------                 ---------
    Num. Qubits:  5             Num. Qubits:  5              Num. Qubits:  5
    Pending Jobs: 0             Pending Jobs: 1290           Pending Jobs: 533
    Least busy:   True          Least busy:   False          Least busy:   False
    Operational:  True          Operational:  True           Operational:  True
    Avg. T1:      96.2          Avg. T1:      94.5           Avg. T1:      113.0
    Avg. T2:      54.4          Avg. T2:      68.1           Avg. T2:      83.7

    ibmq_16_melbourne            ibmqx2
    -----------------            ------
    Num. Qubits:  15             Num. Qubits:  5
    Pending Jobs: 3222           Pending Jobs: 30
    Least busy:   False          Least busy:   False
    Operational:  True           Operational:  True
    Avg. T1:      52.0           Avg. T1:      60.2
    Avg. T2:      57.1           Avg. T2:      31.6

```python
backend_monitor(provider.get_backend("ibmq_armonk"))
```

    ibmq_armonk
    ===========
    Configuration
    -------------
        n_qubits: 1
        operational: True
        status_msg: active
        pending_jobs: 23
        backend_version: 1.1.3
        basis_gates: ['id', 'u1', 'u2', 'u3']
        local: False
        simulator: False
        dynamic_reprate_enabled: False
        meas_kernels: ['hw_boxcar']
        sample_name: SPARROW
        discriminators: ['linear_discriminator', 'quadratic_discriminator']
        backend_name: ibmq_armonk
        max_shots: 8192
        url: None
        n_uchannels: 0
        supported_instructions: ['setf', 'measure', 'x', 'cx', 'id', 'u2', 'shiftf', 'acquire', 'u1', 'delay', 'u3', 'play']
        meas_lo_range: [[6.493427855, 7.493427855]]
        n_registers: 1
        conditional_latency: []
        quantum_volume: None
        meas_map: [[0]]
        coupling_map: None
        rep_times: [1000.0]
        conditional: False
        qubit_channel_mapping: [['m0', 'd0']]
        u_channel_lo: []
        parametric_pulses: ['gaussian', 'gaussian_square', 'drag', 'constant']
        qubit_lo_range: [[4.474446855948306, 5.474446855948306]]
        dtm: 0.2222222222222222
        description: 1 qubit device
        dt: 0.2222222222222222
        open_pulse: True
        allow_q_object: True
        allow_object_storage: True
        acquisition_latency: []
        credits_required: True
        uchannels_enabled: True
        hamiltonian: {'description': 'Qubits are modeled as Duffing oscillators. In this case, the system includes higher energy states, i.e. not just |0> and |1>. The Pauli operators are generalized via the following set of transformations:\n\n$(\\mathbb{I}-\\sigma_{i}^z)/2 \\rightarrow O_i \\equiv b^\\dagger_{i} b_{i}$,\n\n$\\sigma_{+} \\rightarrow b^\\dagger$,\n\n$\\sigma_{-} \\rightarrow b$,\n\n$\\sigma_{i}^X \\rightarrow b^\\dagger_{i} + b_{i}$.\n\nQubits are coupled through resonator buses. The provided Hamiltonian has been projected into the zero excitation subspace of the resonator buses leading to an effective qubit-qubit flip-flop interaction. The qubit resonance frequencies in the Hamiltonian are the cavity dressed frequencies and not exactly what is returned by the backend defaults, which also includes the dressing due to the qubit-qubit interactions.\n\nQuantities are returned in angular frequencies, with units 2*pi*GHz.\n\nWARNING: Currently not all system Hamiltonian information is available to the public, missing values have been replaced with 0.\n', 'h_latex': '\\begin{align} \\mathcal{H}/\\hbar = & \\sum_{i=0}^{0}\\left(\\frac{\\omega_{q,i}}{2}(\\mathbb{I}-\\sigma_i^{z})+\\frac{\\Delta_{i}}{2}(O_i^2-O_i)+\\Omega_{d,i}D_i(t)\\sigma_i^{X}\\right) \\\\ \\end{align}', 'h_str': ['_SUM[i,0,0,wq{i}/2*(I{i}-Z{i})]', '_SUM[i,0,0,delta{i}/2*O{i}*O{i}]', '_SUM[i,0,0,-delta{i}/2*O{i}]', '_SUM[i,0,0,omegad{i}*X{i}||D{i}]'], 'osc': {}, 'qub': {'0': 3}, 'vars': {'delta0': -2.189518576306432, 'omegad0': 0.0721186503632053, 'wq0': 31.25537139664009}}
        memory: True
        max_experiments: 75
        multi_meas_enabled: False
        online_date: 2019-10-16 04:00:00+00:00
        meas_levels: [1, 2]

    Qubits [Name / Freq / T1 / T2 / U1 err / U2 err / U3 err / Readout err]
    -----------------------------------------------------------------------
        Q0 / 4.97445 GHz / 184.73939 us / 195.43515 us / 0 / 0.00046 / 0.00092 / -0.34847

    Multi-Qubit Gates [Name / Type / Gate Error]
    --------------------------------------------

### Getting the least Busy Quantum Computer for our Work

```python
from qiskit.providers.ibmq import least_busy

```

```python
backend= least_busy(provider.backends(operational=True, simulator=False, status_msg='active',filters=lambda x: x.configuration().n_qubits >1))

print("We are executing on....", backend)

print("It has ",backend.status().pending_jobs,"pending jobs")

job_exp = execute(circ_bell, backend= backend)

job_monitor(job_exp)
```

    We are executing on.... ibmq_valencia
    It has  1 pending jobs
    Job Status: job has successfully run

## Plotting the Results

```python
result_exp = job_exp.result()

counts_exp = result_exp.get_counts(circ_bell)

plot_histogram([counts_exp,counts], legend=['Device', 'Simulator'])
```

```python

```


```python
from tectosaur2.nb_config import setup
setup()
```

```python
import warnings
import numpy as np
import sympy as sp
import matplotlib.pyplot as plt
from tectosaur2 import (
    integrate_term,
    pts_grid,
    gauss_rule,
    upsample,
    panelize_symbolic_surface,
    refine_surfaces,
)
from tectosaur2.laplace2d import DoubleLayer, Hypersingular

```

```python
sp_t = sp.var('t')
x, y = sp_t, 0.3 * sp.sin((sp_t + 1) * 2 * sp.pi)
surface = (sp_t, x, y)
# kernel = DoubleLayer(d_cutoff=4.0, d_up=2.0, d_qbx=0.4, default_tol=1e-13, max_p=100)
kernel = Hypersingular(d_cutoff=5.0, d_up=2.5, d_qbx=0.5, default_tol=1e-12, max_p=300)
```

```python
def slip_fnc(xhat):
    # This must be zero at the endpoints for potential gradient to be finite.
    # return np.ones_like(xhat)
    return np.cos(xhat * np.pi) + 1.0

surface_low = refine_surfaces([surface], gauss_rule(6), max_curvature=1.5)
# surface_low = panelize_symbolic_surface(*surface, gauss_rule(10), n_panels=30)
slip_low = slip_fnc(surface_low.quad_pts)
surface_low.n_panels
```

    32

```python

nobs = 200
zoomx = [-1.5, 1.5]
zoomy = [-1.5, 1.5]
xs = np.linspace(*zoomx, nobs)
ys = np.linspace(*zoomy, nobs)
obs_pts = pts_grid(xs, ys)

singularities = np.array(
    [
        [surface[1].subs(surface[0], -1), surface[2].subs(surface[0], -1)],
        [surface[1].subs(surface[0], 1), surface[2].subs(surface[0], 1)],
    ]
)
qbx_mat, report = integrate_term(
        kernel, obs_pts, surface_low, tol=1e-8, singularities=singularities, return_report=True
    )
qbx_vals = qbx_mat[:, 1, :, 0].dot(slip_low).reshape((nobs, nobs))

```

    /Users/tbent/Dropbox/active/eq/tectosaur2/tectosaur2/integrate.py:210: UserWarning: Some expanded integrals reached maximum expansion order. These integrals may be inaccurate.
      warnings.warn(

```python
report['p'][report['p']>=50]
```

    array([ 54,  64,  54,  64,  64,  64,  64,  54,  54, 300, 300, 194,  54,
           144,  64,  84,  94,  54,  74,  54, 224, 184,  94,  94, 104,  54,
            74,  54,  74, 244,  54,  54, 174,  94,  64,  64,  54,  54, 104,
            74,  74, 124,  64,  54,  54,  84, 284, 284,  84,  54,  54,  64,
           124,  74,  74, 104,  54,  54,  64,  64,  94, 174,  54,  54, 244,
            74,  54,  74,  54, 104,  94,  94, 184, 224,  54,  74,  54,  94,
            84,  64, 144,  54, 194, 300, 300,  54,  54,  64,  64,  64,  64,
            54,  64,  54], dtype=int32)

```python
report['exp_rs'][report['p']>=50]
```

    array([0.03383182, 0.03383182, 0.03383182, 0.03383182, 0.03383182,
           0.03383182, 0.03383182, 0.03383182, 0.0452643 , 0.03383182,
           0.03383182, 0.0452643 , 0.03383182, 0.03383182, 0.0452643 ,
           0.0452643 , 0.0452643 , 0.0452643 , 0.0452643 , 0.0452643 ,
           0.0452643 , 0.0452643 , 0.0452643 , 0.0452643 , 0.05786859,
           0.05786859, 0.05786859, 0.05786859, 0.05786859, 0.05786859,
           0.05786859, 0.05786859, 0.05786859, 0.05786859, 0.05786859,
           0.06536074, 0.05786859, 0.06536074, 0.06536074, 0.03921644,
           0.06536074, 0.06536074, 0.06536074, 0.06536074, 0.06536074,
           0.06536074, 0.06536074, 0.06536074, 0.06536074, 0.06536074,
           0.06536074, 0.06536074, 0.06536074, 0.06536074, 0.03921644,
           0.06536074, 0.06536074, 0.05786859, 0.06536074, 0.05786859,
           0.05786859, 0.05786859, 0.05786859, 0.05786859, 0.05786859,
           0.05786859, 0.05786859, 0.05786859, 0.05786859, 0.05786859,
           0.0452643 , 0.0452643 , 0.0452643 , 0.0452643 , 0.0452643 ,
           0.0452643 , 0.0452643 , 0.0452643 , 0.0452643 , 0.0452643 ,
           0.03383182, 0.03383182, 0.0452643 , 0.03383182, 0.03383182,
           0.0452643 , 0.03383182, 0.03383182, 0.03383182, 0.03383182,
           0.03383182, 0.03383182, 0.03383182, 0.03383182])

```python
report.keys()
```

    dict_keys(['combined_src', 'srcs', 'obs_pts', 'n_qbx', 'p', 'integration_failed', 'n_subsets', 'max_order_reached', 'use_qbx', 'exp_centers', 'exp_rs', 'closest_src_pts', 'direction', 'on_surface', 'n_nearfield'])

```python
from tectosaur2.debug import plot_max_order_reached
plt.figure(figsize=(8,8))
plot_max_order_reached(report)
# plt.xlim([-1, -0.3])
# plt.ylim([0.5, 1.3])
plt.show()
```

```python

vmin = -1.0
vmax = 1.0
levels = np.linspace(vmin, vmax, 11)

obsx = obs_pts[:, 0].reshape((nobs, nobs))
obsy = obs_pts[:, 1].reshape((nobs, nobs))

plt.figure(figsize=(16, 12))
plt.subplot(2, 3, 4)
plt.title("QBX solution")
cntf = plt.contourf(obsx, obsy, qbx_vals, levels=levels, extend="both")
plt.contour(
    obsx,
    obsy,
    qbx_vals,
    colors="k",
    linestyles="-",
    linewidths=0.5,
    levels=levels,
    extend="both",
)
plt.plot(surface_low.pts[:, 0], surface_low.pts[:, 1], "k-", linewidth=1.5)
plt.xlim(zoomx)
plt.ylim(zoomy)
plt.tight_layout()
plt.show()

```

```python

```