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

# In[1]:


import numpy as np
import dolfin


# In[8]:


mesh_sm = dolfin.UnitSquareMesh(3, 3)
ff_sm = dolfin.FacetFunctionSizet(mesh_sm)

mesh_lg = dolfin.UnitSquareMesh(50, 50)
ff_lg = dolfin.FacetFunctionSizet(mesh_lg)

for ff in (ff_sm, ff_lg):
    ff.array()[:] = np.random.randint(0,10,size=ff.array().shape)


# In[9]:


get_ipython().run_line_magic('matplotlib', 'inline')
import matplotlib.pyplot as plt


# In[11]:


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)


# In[ ]: