"""
Solve
  - div(grad(p)) = -div(rho(x,y) g),
subject to the boundary conditions
  p = 0                for x in \Gamma_3,
  grad(p).n = rho g.n  for x in \Gamma_1 \cup \Gamma_2.
  
Assumes: variable density; constant gravity vector.
"""

from dolfin import *

# Define gravity vector
gravity = Expression(("0.0", "-3.3"), degree=2)

# Create mesh and define function space
mesh = Mesh("box-2domain.xml")
boundaries = MeshFunction('size_t', mesh, "box-2domain_facet_region.xml")
subdomains = MeshFunction('size_t', mesh, "box-2domain_physical_region.xml")

class density_from_regions(Expression):
  def __init__(self, regions, v0, v1, **kwargs):
    self.regions = regions
    self.v0 = v0
    self.v1 = v1

  def eval_cell(self, values, x, cell):
    if self.regions[cell.index] == 1:
      values[0] = self.v0
    elif self.regions[cell.index] == 2:
      values[0] = self.v1
    else:
      raise RuntimeError('Only two densities associated with region index 1 or 2 is expected')

# Define constant density from material regions
region = subdomains.array()[:]
rho = density_from_regions(region, 2.0, 1.1, degree=0)

# Define spaces
V = FunctionSpace(mesh, "Lagrange", 1)
q, p = TestFunction(V), TrialFunction(V)

# Define Dirichlet boundaries
pressure_reference_value = 0.0
bc = DirichletBC(V, pressure_reference_value, boundaries, 3)

# Variational problem
bform_A = inner(grad(q), grad(p)) * dx
lform_L = inner(grad(q), rho*gravity) * dx

# Compute solution
p = Function(V)
solve(bform_A == lform_L, p, bc)

# Save solution
file = File("box-2domain_p.pvd")
file << p

