CauchyRiemann_Sympy.py

import numpy as np 
import matplotlib.pyplot as plt 
import sympy as sb

# カラーバー用
from mpl_toolkits.axes_grid1 import make_axes_locatable

# フォントサイズ
plt.rcParams["font.size"] = 10

# sympyの変数
x = sb.Symbol('x',real=True)
y = sb.Symbol('y',real=True)
z = sb.Symbol('z')
I = sb.I

# 関数
# f = sb.exp(z)
f = z**2
# f = z**3
# f = 1/(z-1)
# f = sb.conjugate(z)

# z=x+iyを代入
F = f.subs(z,(x+I*y))

# f(z) = u + iv
u = sb.re(sb.expand(F))
v = sb.im(sb.expand(F))

# ∂u/∂x
dudx = u.diff(x)
# ∂u/∂y
dudy = u.diff(y)

# ∂v/∂x
dvdx = v.diff(x)
# ∂v/∂y
dvdy = v.diff(y)

lm = [-1,1]
xx = np.linspace(lm[0],lm[1],100)
yy = np.linspace(lm[0],lm[1],100)

X,Y = np.meshgrid(xx,yy) 

U    = sb.lambdify((x,y),u)(X,Y)
V    = sb.lambdify((x,y),v)(X,Y)
dUdx = sb.lambdify((x,y),dudx)(X,Y)
dUdy = sb.lambdify((x,y),dudy)(X,Y)

dVdx = sb.lambdify((x,y),dvdx)(X,Y)
dVdy = sb.lambdify((x,y),dvdy)(X,Y)

plt.close('all')

plt.figure(figsize=(6,8))

ZZ = [U,V,dUdx,dVdy,dUdy,dVdx,]
ttls = ['$u(x,y)$','$v(x,y)$','$u_x(x,y)$','$v_y(x,y)$','$u_y(x,y)$','$v_x(x,y)$']
F = [u,v,dudx,dvdy,dudy,dvdx,]

N = len(Y)

axes = []
for i,Z in enumerate(ZZ): 
    ax = plt.subplot(3,2,i+1)

    # 定数の時
    if type(Z) is not np.ndarray:
        Z = Z * np.ones_like(X)
        
    im = ax.pcolormesh(X,Y,Z,cmap='bwr')
    
    # タイトル
    plt.title(ttls[i]+'=$'+sb.latex(F[i])+'$')
    
    # カラーバー
    divider = make_axes_locatable(ax)
    cax = divider.append_axes("right", size="5%", pad=0.05)

    plt.colorbar(im, cax=cax,format='%.1f')
    axes.append(ax)

    im.set_clim(lm)

# 軸の設定
for i,axis in enumerate(axes):
    plt.sca(axis)
    
    plt.xlim(lm)
    plt.ylim(lm)    
    plt.xlabel('x')
    plt.ylabel('y')
    plt.gca().set_aspect(1/plt.gca().get_data_ratio())

plt.tight_layout()
plt.text(0.5, 0.96, '$f(z) = ' + sb.latex(f)+'$', ha='center',fontsize=14, transform=plt.gcf().transFigure)
plt.subplots_adjust(top=0.93)

# 保存
# plt.savefig(str(f)+'.png')