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

# # Lumping a mass matrix
# 
# Explicit dynamics simulations require the usage of lumped mass matrices i.e. diagonal mass matrices for which inversion can be done explicitly.
# 
# We show how to do this for P1 and P2 Lagrange elements. The former case can be done quite easily whereas the latter case requires a special integration scheme which is not present by default in FEniCS.
# 
# ## P1 elements
# 
# ### First method
# 
# This problem has been already tackled in the past in old forum posts, [see for instance here](https://fenicsproject.org/qa/4284/mass-matrix-lumping/).

# In[24]:


from dolfin import *
import numpy as np
mesh = UnitSquareMesh(1, 1)
V1 = FunctionSpace(mesh, "CG", 1)
v = TestFunction(V1)
u = TrialFunction(V1)

mass_form = v*u*dx
mass_action_form = action(mass_form, Constant(1))

M_consistent = assemble(mass_form)
print("Consistent mass matrix:\n", np.array_str(M_consistent.array(), precision=3))

M_lumped = assemble(mass_form)
M_lumped.zero()
M_lumped.set_diagonal(assemble(mass_action_form))
print("Lumped mass matrix:\n", np.array_str(M_lumped.array(), precision=3))


# For explicit dynamics simulation, the mass matrix can then be manipulated using the diagonal vector, to compute for instance $M^{-1}w$ where $w$ is some Function.

# In[25]:


M_vect = assemble(mass_action_form)
w = Function(V1)
iMw = Function(V1)
iMw.vector().set_local(w.vector().get_local()/M_vect.get_local())


# ### Second method
# 
# The above trick consists in summing all rows and affecting the corresponding mass to the corresponding diagonal degree of freedom. E.g. for element $T$ with number of vertices $n$, the diagonal lumped mass $i$ is obtained as:
# 
# $$\begin{equation}
# \overline{M}_i = \sum_j M_{ij} = \sum_j \int_{T} N_i(x)N_j(x) \text{ dx} = \int_T N_i(x) \text{ dx} = \dfrac{|T|}{n}
# \end{equation}$$
# 
# for linear shape functions.
# 
# We can obtain the same result by integrating the mass form with a "vertex" scheme (i.e. quadrature points are located at the vertices of the simplex with equal weights $\dfrac{|T|}{n}$:
# 
# $$\begin{align}
# M_{ij} &= \int_{T} N_i(x)N_j(x) \text{ dx('vertex')} \\
# &= \sum_{k=1}^n \dfrac{|T|}{n}  N_i(x_k)N_j(x_k) =  \sum_{k=1}^n \dfrac{|T|}{n} \delta_{ik}\delta_{jk} = \dfrac{|T|}{n} \delta_{ij} \label{diagonal-mass} \tag{1}
# \end{align}$$

# In[26]:


M_lumped2 = assemble(v*u*dx(scheme="vertex", metadata={"degree":1, "representation":"quadrature"}))
print("Lumped mass matrix ('vertex' scheme):\n", np.array_str(M_lumped2.array(), precision=3))


# ## P2 elements
# 
# The problem with the first method is that we obtain a singular matrix for P2 elements:

# In[27]:


V2 = FunctionSpace(mesh, "CG", 2)
v = TestFunction(V2)
u = TrialFunction(V2)

mass_form = v*u*dx
mass_action_form = action(mass_form, Constant(1))

M_consistent = assemble(mass_form)
print("Consistent mass matrix:\n", np.array_str(M_consistent.array(), precision=3))

M_lumped = assemble(mass_form)
M_lumped.zero()
M_lumped.set_diagonal(assemble(mass_action_form))
print("Lumped mass matrix:\n", np.array_str(M_lumped.array(), precision=3))


# We can however adapt the second method by considering a generalized "vertex" quadrature scheme with quadrature points located at the element degrees of freedom points i.e. at vertices and facets mid-points for P2 elements. In this case, we still have $\eqref{diagonal-mass}$. To do this we need to implement a custom quadrature scheme called `"lumped"` based on [the following this discussion](https://bitbucket.org/fenics-project/dolfin/issues/955/support-for-arbitrary-quadrature-rules).

# In[28]:


import lumping_scheme

M_lumped = assemble(v*u*dx(scheme="lumped", degree=2))
print("Lumped mass matrix ('lumped' scheme):\n", np.array_str(M_lumped.array(), precision=3))