My Tetration (complex-number) function has to be better vectorized (Python)

Question

Inspired by 3blue1brown I'm trying to graph the escape (divergence) of the Tetration function with Python –– something akin to this beautiful graphic on Wikipedia.

def tetration_com(base, tol=10**-15, max_step=10**6):
  # returns t, the infinite tetration of base.
  # if t does not converge, the function returns an escape value, aka how fast it diverges..

  t = 1.0
  step = 0
  escape = None
  ln_base = cmath.log(base)

  t_last = 0
  try:
    while(abs(t - t_last) > tol):
      if(step > max_step):
        raise OverflowError
      t_last = t
      t = cmath.exp(ln_base*t)   # [ base^t == e^(ln(base)*t) ]
      step += 1

  except(OverflowError):
    t = None
    escape = 1000/step
    # the escape value is is inversely related to the number of steps it took
    # us to diverge to infinity

  return t, escape

I'm trying to make it work with a meshgrid in order to graph the escape on the x-y plane. Python doesn't like that the output is 2 unpacked variables (the limit, or the escape) – this I can definitely solve by splitting into two functions.

But another issue is that the complex-math operations (cmath.log, cmath.exp) are only working well with scalars...

I tried to vectorize the function:

nx, ny = 700, 500
x, y = np.linspace(-3.5, 3.5, nx), np.linspace(-2.5, 2.5, ny)
xv, yv = np.meshgrid(x, y)

tetration_vec = np.vectorize(tetration_com)
t, escape = tetration_vec(xv + yv*1j, max_step=500)

But it's running forever.

Any advice on how to deal with the complex-math operations and the vectorization?

Answer 1

Here is how I graphed my escape in the end:

def tetration_com(base, tol=10**-15, max_step=10**6, max_val=10**2):
  # returns t, the infinite tetration of base.
  # if t does not converge, the function returns an escape value
  # aka how fast it diverges..

  t = 1.0
  step = 0
  escape = None

  t_last = 0
  try:
    while(abs(t - t_last) > tol):
      if(step > max_step or abs(t) > max_val):
        raise OverflowError
      t_last = t
      t = pow(base, t)
      step += 1
  except(OverflowError):
    t = None
    escape = 1000/step
    # the escape value is is inversely related to the number of steps it took
    # us to diverge to infinity

  return t, escape

Vectorized helper function:

def tetra_graph_escape(real, imag, tol=10**-15, max_step=10**3, max_val=10**2):
  return np.array([np.array([tetration_com(r + im*1j, tol=tol, max_step=max_step, max_val=max_val)[1]
                             for im in imag]) for r in real])

Graphing:

# graph our escape:
nx, ny = 700, 500
x, y = np.linspace(-3.5, 3.5, nx), np.linspace(-2.5, 2.5, ny)

val, escape = tetra_graph_conv(x, y), tetra_graph_escape(x, y)

import matplotlib.pyplot as plt

for r in range(len(escape)):
  for c in range(len(escape[0])):
    if escape[r][c] is None:
      escape[r][c] = -100

escape[460][250]

plt.contour(escape)

Contour graph of Tetration escape: