#!/usr/bin/env python
# coding: utf-8

# ## Analytic lanscape of a complex function

# The colored analytic landscape of a complex function, f, is the graph of the function modulus or log-modulus, colored according to the argument of  $f(z)$ at  any z in the function domain.
# 
# Here we illustrate how defining this surface as an instance of the Plotly class `plotly.graph_objects.Surface`, we can plot the contour  lines as well as the surface projection onto a z-plane, that amounts to plotting the domain coloring of the function f.
# We are using the cyclic HSV colorscale to color-encode the argument value. For details on HSV and domain coloring see this Jupyter notebook [https://nbviewer.jupyter.org/github/empet/Math/blob/master/DomainColoring.ipynb](https://nbviewer.jupyter.org/github/empet/Math/blob/master/DomainColoring.ipynb).

# In[1]:


import numpy as np
from numpy import pi
import plotly.graph_objs as go
import plotly.io as pio
pio.renderers.default ="notebook_connected"


# In[2]:


def eval_modulus(f, re=(-2.5, 2.5), im=(-2.5, 2.5), N=50, log = False):
    nrx = int(N * (re[1]-re[0]))
    nry = int(N * (im[1]-im[0]))
    x = np.linspace(re[0], re[1], nrx)
    y = np.linspace(im[0], im[1], nry)
    
    x, y = np.meshgrid(x, y)
    z = x + 1j*y
    w = f(z)
    w[np.isinf(w)] = np.nan
    if log:
        modf = np.log(np.absolute(w)) 
    else:     
        modf = np.absolute(w)  # |f|
    return x, y, np.angle(w), modf  #np.angle(w) is the argument of f


# HSV colorscale:

# In[3]:


pl_hsv = [[0.0, 'rgb(0, 242, 242)'],
          [0.083, 'rgb(0, 121, 242)'],
          [0.167, 'rgb(0, 0, 242)'],
          [0.25, 'rgb(121, 0, 242)'],
          [0.333, 'rgb(242, 0, 242)'],
          [0.417, 'rgb(242, 0, 121)'],
          [0.5, 'rgb(242, 0, 0)'],
          [0.583, 'rgb(242, 121, 0)'],
          [0.667, 'rgb(242, 242, 0)'],
          [0.75, 'rgb(121, 242, 0)'],
          [0.833, 'rgb(0, 242, 0)'],
          [0.917, 'rgb(0, 242, 121)'],
          [1.0, 'rgb(0, 242, 242)']]


# Define `tickvals` and `ticktext` for colorbar:

# In[4]:


tickvals=[-np.pi, -2*np.pi/3, -np.pi/3, 0, np.pi/3, 2*np.pi/3, np.pi]
#define the above values as strings with pi-unicode
ticktext=['-\u03c0', '-2\u03c0/3', '-\u03c0/3', '0', '\u03c0/3', '2\u03c0/3', '\u03c0']

coloraxis_settings = dict(colorscale= pl_hsv, 
                      colorbar_thickness=25, 
                      colorbar_len=0.7, 
                      colorbar_tickvals=tickvals, 
                      colorbar_ticktext=ticktext,
                      colorbar_title='arg(f)') 

def set_contours(min_mod, zrange_max, n=20, color = 'rgb(250, 250, 250)'):
    return dict(start=min_mod,
                end=zrange_max, highlight=True,
                size=(zrange_max-min_mod)/n, 
                width=1.5, #contour line width 
                color= color,
                project_z=True)


# Let us plot the  analytic landscape of the function f, defined below. 
# f has a zero multiple of order 3, and three simple poles.

# In[5]:


f = lambda z: z**3 / (z**3-1)

x, y, argf, modf = eval_modulus(f)

fig1  = go.Figure(go.Surface(x=x[0, :], y=y[:, 0], z=modf, 
                            surfacecolor=argf, coloraxis='coloraxis'))
z_range = (-4, 6)
fig1.update_layout(title_text = '$\\text{Analytic landscape of the function}\: f(z)=  \\displaystyle\\frac{z^3}{z^3-1}$',
                  title_x=0.5, font_family="Balto",
                  width=700, height=700,
                  coloraxis = coloraxis_settings,
                  scene_zaxis_range=z_range);


# Add the surface projection onto the z-plane of equation `z=z_range[0]`, and  the contour lines of the plotted surface, and their projection, to illustrate the domain coloring plot of the function f(z) onto this z-plane:

# In[6]:


min_mod = np.min(modf)
fig1.add_surface(x=x[0, :], y=y[:, 0], z= z_range[0]*np.ones(modf.shape), 
                 surfacecolor=argf, colorscale= pl_hsv, showscale=False)

fig1.data[0].update(contours_z=dict(show=True, **set_contours(min_mod, z_range[1], n=28, color = 'rgb(250, 250, 250)')));  
fig1.update_scenes(camera_eye_z=0.75);


# In[7]:


fig1.show()  


# Finally let us plot the log modulus of $h(z)= e^{1/z}$:

# In[8]:


h = lambda z: np.exp(1/z)
x, y, argh, modh = eval_modulus(h, log=True)

fig2 = go.Figure(go.Surface(x=x[0, :], y=y[:, 0], z=modh, 
                         surfacecolor=argh, 
                         coloraxis='coloraxis'))


z_range= (-6, 6)

fig2.update_layout(title_text = '$\\text{Analytic landscape of the function}\: f(z)= e^{1/z}$',
                  title_x=0.5, font_family="Balto",
                  width=700, height=700,
                  coloraxis = coloraxis_settings,
                  scene_zaxis_range=z_range);


# In[9]:


min_mod = np.min(modh)
fig2.add_surface(x=x[0, :], y=y[:, 0], z= z_range[0]*np.ones(modh.shape), 
                surfacecolor=argh, colorscale= pl_hsv, showscale=False)

fig2.data[0].update(contours_z=dict(show=True, 
                                  start=z_range[0],
                end=z_range[1], highlight=True,
                size=(z_range[1]-z_range[0])/26, 
                width=1.5, #contour line width 
                color= 'rgb(250,250,250)',
                project_z=True))
fig2.update_scenes(camera_eye_x=-1.55, camera_eye_y=1.55, camera_eye_z=0.6);


# In[10]:


fig2.show()


# Unlike the representation of modulus, $|f}$,  the log-modulus, $log(|f|$, has negative values where $0<|f|<1$

# In[ ]: