"""
LINEAR ISOTROPIC ELASTICITY

Firedrake program for 2D one elastic layer with a reverse fault dipping 45° to the west.

A numerical solver of elastostatic equation div(Ce(u)) = f with u = 0 on the boundary
and discontinuity of displacement on fault (:= Gamma_I) (rupture)
The lowest order Arnold-Falk-Winther element is used.

MIXED METHOD: both stress and displacement fields are approximated as primary variables.

Written in mixed form:
A sigma = e(u)							# costitutive equation
-div(sigma) = f 						# equilibrium equation
with weak symmetry elements of AFW (denoted by V) where
gamma = - skw(grad(u))

where A = compliance matric ---> inverse of elasticity stiffness matrix C (fourth-order tensor)


(stress, deformation, pressure, skew):
(sigma, u, p, gamma) in V : trial functions with sigma n = g on Gamma_N (if exists)
(tau, v, q, eta) in V : test functions with tau n = 0 on Gamma_N

A variational form is derived from the Hellinger-Reissner principle:
(A*sigma, tau) + (div(tau), u) +/-(?) (gamma, tau) = <u0 dot(tau, n)>
- (div(sigma), v) + alpha * (grad(p), v) = 0
(Se*p, q)_t - alpha*(u,grad(q))_t + ((k/mu)*grad(p), grad(q)) = 0
(sigma, eta) = 0

The slip boundary condition across the fault is defined as:
<u0 dot(tau, n)> = int_{Gamma_D} (u0, (tau*n))*ds + (f, w) + int_{Gamma_I} (disp, tau*n))*dS

where disp is the jump of displacement at fault.

The traction BC is essential (it is imposed a priori) on the space where the stress tensor is sought.
The displacement BC arises naturally from the variational form.


Simone Puel
06/24/2019
"""


from __future__ import print_function
from firedrake import *
import numpy as np
import time
import matplotlib.pyplot as plt
import math


"""

"""


# monitor time to solution
start = time.time()


#################################################################
###################### USER INPUT ###############################
#################################################################
space_dim = 2 			# 2 --> 2D, 3 --> 3D

#################################################################
#################################################################


# Physical parameters
rho = 2670.0			# [kg/m3]
vs = 3550.0				# [m/s]
vp = 6040.0				# [m/s]

# find elastic parameters: mu (shear modulus), lmbda (Lame parameter), nu (Poisson's ratio), K (bulk modulus) and E (Young's modulus)
mu = rho*pow(vs,2)
lmbda = rho*pow(vp,2)-2.0*mu
nu = lmbda/(2.0*(lmbda+mu))
K = lmbda+2.0/3.0*mu
E = mu*(3.0*lmbda+2.0*mu)/(lmbda+mu)



##################### DEFINE FUNCTIONS ##########################
# strain --> symmetric part of the displacement gradient tensor
def strain(u): # strain = 1/2 (grad u + grad u^T)
	return 0.5*(nabla_grad(u) + nabla_grad(u).T)

def stress(u): # stress = 2 mu strain + lambda tr(strain) I
	return 2.0*mu*strain(u) + lmbda*nabla_div(u)*Identity(space_dim)

# invert constitutive equation for linear isotropic elasticity: sigma = c * epsilon
def Asigma(s): # A sigma = 1/2mu * (sigma - lmbda/(3*lambda + 2*mu) * tr(sigma) * I)
	return (1.0+nu)/E*(s - nu/(1.0+nu)*tr(s)*Identity(space_dim))

def asym(z):
	return z[0,1] - z[1,0]

def tangent(n):
    tangent = as_vector([n[1], -n[0]])						# as_vector gives a row vector (usually the result is a column vector)
    return tangent



######################## MESH ###################################
mesh = Mesh("Mesh_NormalFault.msh")

zerodispl_id = 1
topright_id = 2
topleft_id = 3
fault_id = 4


####################### FUNCTION SPACES ###########################
BDM = VectorFunctionSpace(mesh, "BDM", 1)
DGv = VectorFunctionSpace(mesh, "DG", 0)
DGs = FunctionSpace(mesh, "DG", 0)

# Mixed function spaces
Vh_element = MixedElement([BDM.ufl_element(), DGv.ufl_element(), DGs.ufl_element()])
Vh = FunctionSpace(mesh, Vh_element)						# new function space as a mixed one


x = SpatialCoordinate(mesh)
#u0 = project(as_vector([0.1-x[0]*0.1, 0]), DGv)
u0 = Constant((0.0, 0.0))


################## VARIATIONAL FORM (WEAK FORM) ###################
(sigma, u, r) = TrialFunctions(Vh)							# define the trial functions (the unknowns: displacement (u), stress (sigma) and rotation (p))
(tau, v, q) = TestFunctions(Vh)								# define the test functions on the mixed space
n = FacetNormal(mesh)										# define the facet normal for given mesh


# Write the weak form: lhs and rhs
lhs_varf = inner(Asigma(sigma), tau)*dx + inner(u, div(tau))*dx + inner(r, asym(tau))*dx \
        + inner(div(sigma), v)*dx + inner(asym(sigma), q)*dx

rhs_varf = inner(u0, tau*n)*ds(zerodispl_id)							# natural BCs: displacement

# This is the slip term and it is added in the rhs
slip = -1
rhs_varf2 = inner(slip, dot(tau('+')*n('+'), tangent(n)('+')))*dS(fault_id)
rhs_varf += rhs_varf2

# Solution
sol = Function(Vh)



############################ IMPOSE BCs ##############################
# create zero stress boundary condition and therefore the weak form is written in term of stress (compliance tensor) rather than displacement
zero_stress = as_vector([Constant((0., 0.)), Constant((0., 0.))])
bc1 = DirichletBC(Vh.sub(0), zero_stress, topleft_id)		# zero stress essential BCs y = 0 and y = 1
bc2 = DirichletBC(Vh.sub(0), zero_stress, topright_id)      # zero stress essential boundary condition y = 0 and y = 1




####################### SOLVE THE SYSTEM ############################
solve(lhs_varf == rhs_varf, sol, bcs=[bc1,bc2])

(sigma, u, p) = sol.split(deepcopy=True)
#sid = File("stress.pvd") << sigma
#uid = File("displacement.pvd") << u
#pid = File("rotation.pvd") << p
#mid = File("mesh.pvd") << facet_data


#d = u.geometric_dimension() # number of space dimensions


# print the time elapsed
print("Time elapsed", time.time()-start, "seconds")



# PLOT
plot(u)
plt.show()


# end of the file
print("Finished!")