from dolfin import *
import numpy as np
from matplotlib import cm
import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import Axes3D
plt.rc('text',usetex=True)
plt.rcParams['ytick.right']=True
plt.rcParams['xtick.top']=True
plt.rcParams['ytick.direction']='in'
plt.rcParams['ytick.labelsize']=26
plt.rcParams['xtick.direction']='in'
plt.rcParams['xtick.labelsize']=26
plt.rcParams['xtick.minor.visible']=True
plt.rcParams['ytick.minor.visible']=True
plt.rcParams['xtick.major.size']=8
plt.rcParams['xtick.minor.size']=4
plt.rcParams['ytick.major.size']=8
plt.rcParams['ytick.minor.size']=4
plt.rcParams['lines.markersize']=np.sqrt(36)
msh = UnitSquareMesh(10,10)
# elem = FiniteElement('CG',msh.ufl_cell(),2)
V =FunctionSpace(msh,'CG',1)
u = Function(V)
u.vector()[70] = 1.
x = y = np.linspace(0,1,101)
X,Y = np.meshgrid(x,y)
u_xy = np.zeros(X.shape)
for i in range(X.shape[0]):
for j in range(X.shape[1]):
u_xy[i,j] = u(x[i],y[j])
fig,ax = plt.subplots(1,1,figsize=(8,8))
funcplot = plot(u,cmap=cm.jet)
fig.colorbar(funcplot,shrink=0.75)
ax.tick_params(pad=10)
plot(msh)
fig.tight_layout()
fig.savefig(r'compac.png')
fig1 = plt.figure(figsize=(10,8))
ax1 = plt.gca(projection='3d')
pl2 = ax1.plot_surface(X,Y,u_xy,cmap=cm.jet)
fig1.colorbar(pl2)
ax1.tick_params(pad=10)
fig1.tight_layout()