inferno.py

import numpy as np
import matplotlib.pyplot as plt
import numba

@numba.njit
def f(a_vals, b_vals, pat, itr, tran):
    N = a_vals.shape[0]
    M = b_vals.shape[0]
    out = np.empty((N, M))
    plen = len(pat)
    for i in range(N):
        for j in range(M):
            a = a_vals[i]
            b = b_vals[j]
            x = 0.5
            L = 0.0
            for k in range(tran):
                r = a if pat[k % plen] == 'A' else b
                x = r * x * (1 - x)
            for k in range(itr):
                r = a if pat[k % plen] == 'A' else b
                x = r * x * (1 - x)
                L += np.log(abs(r * (1 - 2 * x)))
            out[i, j] = L / itr
    return out

N = 800
a_vals = np.linspace(2.5, 4.0, N)
b_vals = np.linspace(2.5, 4.0, N)
pat = "ABBAAB"
Z = f(a_vals, b_vals, pat, 200, 100)
plt.imshow(Z.T, extent=[2.5, 4.0, 2.5, 4.0], cmap="inferno", origin="lower")
plt.axis("off")
plt.show()