Comformal color mapping of images of two variable complex functions-CodePudding

I have plotted a the image of a two-variable complex function in the Argand plane, but I would like to use color to see where in the input space the points of that image came from.

I colored the Argand plane (input space):

import numpy as np
import sympy as sym
import cmath
import matplotlib.pyplot as plt
%matplotlib inline 

a = np.arange(-np.pi, np.pi, 0.1)
b = np.arange(-np.pi, np.pi, 0.1)

A, B = np.meshgrid(a, b)

fx = np.sqrt(A**2   B**2)
plt.figure(figsize=(10,10))
plt.pcolormesh(A,B, fx)
plt.colorbar()
plt.show()

Yielding:

And the code:

import numpy as np
# import sympy as sym
# import cmath
import matplotlib.pyplot as plt
# %matplotlib inline 

a = np.arange(-np.pi, np.pi, 0.1)
b = np.arange(-np.pi, np.pi, 0.1)

A, B = np.meshgrid(a, b)

fx = np.sqrt(A**2   B**2)

fig, (ax1, ax2) = plt.subplots(1, 2,figsize=(10,5))
ax1.pcolormesh(A,B, fx)
# ax1.colorbar()

y = np.exp(1j*A) - np.exp(-1j*B)

real = y.real
# extract imaginary part using numpy
imag = y.imag

# plot the complex numbers
ax2.scatter(real, imag, c=fx, s=7, edgecolor='k', linewidths=0.1)
#plt.pcolormesh(cmap=fx)

ax2.set_xlim(-2.05, 2.05)
ax2.set_ylim(-2.05, 2.05)
ax2.set_ylabel('Imaginary')
ax2.set_xlabel('Real')

plt.show()