import dolfin
import matplotlib.pyplot as plt
from scipy.interpolate import CloughTocher2DInterpolator
def facet_plot2d(facet_func, mesh=True, kind='auto'):
# get scalar-valued facet function
z = facet_func.array()
if kind == 'auto':
if len(z) < 32:
kind = 'scatter'
else:
kind = 'heatmap'
if kind not in {'scatter', 'heatmap'}:
raise ValueError('kind must be "scatter" or "heatmap", not %s' % kind)
# construct X,Y coordinates of facets
coords = np.zeros((z.shape[0], 2), dtype=float)
for i, facet in enumerate(dolfin.facets(facet_func.mesh())):
mp = facet.midpoint()
coords[i,:] = facet.midpoint().array()[:2]
if kind == 'heatmap':
# interpolate points for heatmap
x, y = coords.T
X, Y = np.meshgrid(
np.linspace(x.min(), x.max(), 1000),
np.linspace(y.min(), y.max(), 1000),
)
f = CloughTocher2DInterpolator(coords, z)
Z = f(X, Y)
plt.pcolormesh(X, Y, Z)
if mesh:
dolfin.plot(facet_func.mesh())
plt.grid(False)
if kind == 'scatter':
plt.scatter(*coords.T, c=z,
s=400,
)
plt.colorbar()
facet_plot2d(ff_sm, kind='scatter')
plt.figure()
facet_plot2d(ff_lg, kind='heatmap', mesh=False)