Complex functions in Python

%matplotlib inline
import matplotlib.pyplot as plt
import numpy as np
from matplotlib.colors import hsv_to_rgb
rcdef = plt.rcParams.copy()

def g(x):
    return (1.0-1/(1+x**2))**0.2
    
def Hcomplex(z):# computes the hue corresponding to the complex number z
    H=np.angle(z)/(2*np.pi)+1
    return np.mod(H,1)

def func_vals(f, re, im,  N): #evaluates the complex function at the nodes of the grid
   #re and im are  tuples, re=(a,b) and im=(c,d), defining the rectangular region
   #N is the number of nodes per unit interval 
    
   l=re[1]-re[0]
   h=im[1]-im[0]
   resL=N*l #horizontal resolution
   resH=N*h#vertical resolution
   x=np.linspace(re[0], re[1],resL)
   y=np.linspace(im[0], im[1], resH)
   x,y=np.meshgrid(x,y)
   z=x+1j*y
   w=f(z)
   return w
   
def domaincol_c(w, s):#Classical domain coloring
    #w is the complex array of values f(z)
    #s is the constant saturation
    
    
    indi=np.where(np.isinf(w))#detects the values w=a+ib, with a or b or both =infinity
    indn=np.where(np.isnan(w))#detects nans
  
    H=Hcomplex(w)
    S = s*np.ones_like(H)
    modul=np.absolute(w)
    V= (1.0-1.0/(1+modul**2))**0.2
    # the points mapped to infinity are colored with white; hsv_to_rgb(0,0,1)=(1,1,1)=white
    H[indi]=0.0 
    S[indi]=0.0  
    V[indi]=1.0
    #hsv_to_rgb(0,0,0.5)=(0.5,0.5, 0.5)=gray  
    H[indn]=0
    S[indn]=0
    V[indn]=0.5
    HSV = np.dstack((H,S,V))
    RGB = hsv_to_rgb(HSV)
    return RGB
    
def plot_domain(color_func, f,   re=[-1,1], im= [-1,1], Title='',
                s=0.9, N=200, daxis=None):
    w=func_vals(f, re, im, N)
    domc=color_func(w, s)
    plt.xlabel("$\Re(z)$")
    plt.ylabel("$\Im(z)$")
    plt.title(Title)
    if(daxis):
         plt.imshow(domc, origin="lower", extent=[re[0], re[1], im[0], im[1]])
       
    else:
        plt.imshow(domc, origin="lower")
        plt.axis('off')
        
# EXAMPLE 1:
        
plt.rcParams['figure.figsize'] = 7, 7
ab=[-2,2]
cd=[-2,2]

f=lambda z: (z**2+1)
plot_domain(domaincol_c, f, re=ab, im= cd, Title='$z^2 + 1$', daxis=True)


# EXAMPLE 2:

f=lambda z: (z**2 - 1) * (z - 1 - 1j)**2 / (z**2 + 2 + 2j)
plot_domain(domaincol_c, f, re=ab, im= cd, Title='$f(z)=(z^2 - 1)(z - 1 - i)^2 /(z^2 + 2 + 2i)$', daxis=True)