def get_boundary_indices(function_space):
bc_map = dla.Function(function_space)
bc = dla.DirichletBC(function_space, dla.Constant(1.0), 'on_boundary')
bc.apply(bc_map.vector())
indices = np.arange(
bc_map.vector().size())[bc_map.vector().get_local() == 1.0]
return indices
def run_steady_state_model(function_space,
kappa,
forcing,
boundary_conditions=None,
velocity=None):
"""
Solve steady-state diffusion equation
-grad (k* grad u) = f
"""
mesh = function_space.mesh()
if boundary_conditions == None:
bndry_obj = dl.CompiledSubDomain("on_boundary")
boundary_conditions = [['dirichlet', bndry_obj, dla.Constant(0)]]
num_bndrys = len(boundary_conditions)
boundaries = mark_boundaries(mesh, boundary_conditions)
dirichlet_bcs = collect_dirichlet_boundaries(function_space,
boundary_conditions,
boundaries)
# To express integrals over the boundary parts using ds(i), we must first
# redefine the measure ds in terms of our boundary markers:
ds = dl.Measure('ds', domain=mesh, subdomain_data=boundaries)
dx = dl.Measure('dx', domain=mesh)
# Variational problem at each time
u = dl.TrialFunction(function_space)
v = dl.TestFunction(function_space)
a = kappa * dl.inner(dl.grad(u), dl.grad(v)) * dx
L = forcing * v * dx
if velocity is not None:
a += v * dl.dot(velocity, dl.grad(u)) * dx
beta_1_list = []
alpha_1_list = []
for ii in range(num_bndrys):
if (boundary_conditions[ii][0] == 'robin'):
alpha = boundary_conditions[ii][3]
a += alpha * u * v * ds(ii)
elif ((boundary_conditions[ii][0] == 'robin')
or (boundary_conditions[ii][0] == 'neumann')):
beta = boundary_conditions[ii][2]
L -= beta * v * ds(ii)
u = dla.Function(function_space)
# dl.assemble, apply and solve does not work with
# fenics adjoint
# A, b = dla.assemble_system(a, L, dirichlet_bcs)
# # apply boundary conditions
# for bc in dirichlet_bcs:
# bc.apply(A, b)
# dla.solve(A, u.vector(), b)
dla.solve(a == L, u, dirichlet_bcs)
return u
def fenics_cost(u, f):
x = ufl.SpatialCoordinate(mesh)
w = ufl.sin(ufl.pi * x[0]) * ufl.sin(ufl.pi * x[1])
d = 1 / (2 * ufl.pi ** 2) * w
alpha = fa.Constant(1e-6)
J_form = (0.5 * ufl.inner(u - d, u - d)) * ufl.dx + alpha / 2 * f ** 2 * ufl.dx
J = fa.assemble(J_form)
return J
def get_diffusivity(self, random_sample):
r"""
"""
assert random_sample.shape[0] == 1
Gamma = dla.Constant(self.Gamma*(1+random_sample[0]))
self.shallow_ice_diffusivity = ShallowIceDiffusivity(
Gamma, self.bed, self.beta, self.positivity_tol)
return self.shallow_ice_diffusivity
def get_forcing(self, random_sample):
r"""By Default the forcing is deterministic and set to
.. math:: (1.5+\cos(2\pi t))*cos(x_1)
where :math:`t` is time and :math:`x_1` is the first spatial dimension.
"""
forcing = dla.Constant(0.0)
return forcing
def solve_fenics(kappa0, kappa1):
f = fa.Expression(
"10*exp(-(pow(x[0] - 0.5, 2) + pow(x[1] - 0.5, 2)) / 0.02)", degree=2)
u = fa.Function(V)
bcs = [fa.DirichletBC(V, fa.Constant(0.0), "on_boundary")]
inner, grad, dx = ufl.inner, ufl.grad, ufl.dx
JJ = 0.5 * inner(kappa0 * grad(u), grad(u)) * dx - kappa1 * f * u * dx
v = fenics.TestFunction(V)
F = fenics.derivative(JJ, u, v)
fa.solve(F == 0, u, bcs=bcs)
return u
def get_num_subdomain_dofs(Vh, subdomain):
"""
Get the number of dofs on a subdomain
"""
temp = dla.Function(Vh)
bc = dla.DirichletBC(Vh, dla.Constant(1.0), subdomain)
# warning applying bc does not just apply subdomain.inside to all coordinates
# it does some boundary points more than once and other inside points not
# at all.
bc.apply(temp.vector())
vec = temp.vector().get_local()
dl.plot(temp)
import matplotlib.pyplot as plt
plt.show()
return np.where(vec > 0)[0].shape[0]
def test_fenics_forward():
numpy_output, _, _, _ = evaluate_primal(solve_fenics, templates, *inputs)
u = solve_fenics(fa.Constant(0.5), fa.Constant(0.6))
assert np.allclose(numpy_output, to_numpy(u))
f = fa.Expression(
"10*exp(-(pow(x[0] - 0.5, 2) + pow(x[1] - 0.5, 2)) / 0.02)", degree=2)
u = fa.Function(V)
bcs = [fa.DirichletBC(V, fa.Constant(0.0), "on_boundary")]
inner, grad, dx = ufl.inner, ufl.grad, ufl.dx
JJ = 0.5 * inner(kappa0 * grad(u), grad(u)) * dx - kappa1 * f * u * dx
v = fenics.TestFunction(V)
F = fenics.derivative(JJ, u, v)
fa.solve(F == 0, u, bcs=bcs)
return u
templates = (fa.Constant(0.0), fa.Constant(0.0))
inputs = (np.ones(1) * 0.5, np.ones(1) * 0.6)
def test_fenics_forward():
numpy_output, _, _, _ = evaluate_primal(solve_fenics, templates, *inputs)
u = solve_fenics(fa.Constant(0.5), fa.Constant(0.6))
assert np.allclose(numpy_output, to_numpy(u))
def test_fenics_vjp():
numpy_output, fenics_output, fenics_inputs, tape = evaluate_primal(
solve_fenics, templates, *inputs)
g = np.ones_like(numpy_output)
vjp_out = evaluate_pullback(fenics_output, fenics_inputs, tape, g)
check1 = np.isclose(vjp_out[0], np.asarray(-2.91792642))
def test_fenics_forward():
numpy_output, _, _, _, = evaluate_primal(assemble_fenics, templates, *inputs)
u1 = fa.interpolate(fa.Constant(1.0), V)
J = assemble_fenics(u1, fa.Constant(0.5), fa.Constant(0.6))
assert np.isclose(numpy_output, J)
V = fenics.FunctionSpace(mesh, "P", 1)
def assemble_fenics(u, kappa0, kappa1):
f = fa.Expression(
"10*exp(-(pow(x[0] - 0.5, 2) + pow(x[1] - 0.5, 2)) / 0.02)", degree=2
)
inner, grad, dx = ufl.inner, ufl.grad, ufl.dx
J_form = 0.5 * inner(kappa0 * grad(u), grad(u)) * dx - kappa1 * f * u * dx
J = fa.assemble(J_form)
return J
templates = (fa.Function(V), fa.Constant(0.0), fa.Constant(0.0))
inputs = (np.ones(V.dim()), np.ones(1) * 0.5, np.ones(1) * 0.6)
ff = lambda *args: evaluate_primal(assemble_fenics, templates, *args)[0] # noqa: E731
ff0 = lambda x: ff(x, inputs[1], inputs[2]) # noqa: E731
ff1 = lambda y: ff(inputs[0], y, inputs[2]) # noqa: E731
ff2 = lambda z: ff(inputs[0], inputs[1], z) # noqa: E731
def test_fenics_forward():
numpy_output, _, _, _, = evaluate_primal(assemble_fenics, templates, *inputs)
u1 = fa.interpolate(fa.Constant(1.0), V)
J = assemble_fenics(u1, fa.Constant(0.5), fa.Constant(0.6))
assert np.isclose(numpy_output, J)
def test_vjp_assemble_eval():
import ufl
import logging
from jaxfenics_adjoint import build_jax_fem_eval
from jaxfenics_adjoint import from_numpy
import matplotlib.pyplot as plt
config.update("jax_enable_x64", True)
fenics.set_log_level(fenics.LogLevel.ERROR)
logging.getLogger("FFC").setLevel(logging.WARNING)
logging.getLogger("UFL").setLevel(logging.WARNING)
mu = fenics_adjoint.Constant(1.0) # viscosity
alphaunderbar = 2.5 * mu / (100 ** 2) # parameter for \alpha
alphabar = 2.5 * mu / (0.01 ** 2) # parameter for \alpha
q = fenics_adjoint.Constant(
0.01
) # q value that controls difficulty/discrete-valuedness of solution
def alpha(rho):
"""Inverse permeability as a function of rho"""
return alphabar + (alphaunderbar - alphabar) * rho * (1 + q) / (rho + q)
N = 20
delta = 1.5 # The aspect ratio of the domain, 1 high and \delta wide
V = (
def run_model(function_space, time_step, final_time, forcing,
boundary_conditions, init_condition, nonlinear_diffusion,
second_order_timestepping=False, nlsparam=dict(),
positivity_tol=0):
if boundary_conditions is None:
bndry_obj = dl.CompiledSubDomain("on_boundary")
boundary_conditions = [['dirichlet', bndry_obj, dla.Constant(0)]]
dt = time_step
mesh = function_space.mesh()
num_bndrys = len(boundary_conditions)
assert num_bndrys > 0 # specify None for no boundaries
if (len(boundary_conditions) == 1 and
isinstance(boundary_conditions[0][2], dla.DirichletBC)):
ds = dl.Measure('ds', domain=mesh)
dirichlet_bcs = [boundary_conditions[0][2]]
else:
boundaries = mark_boundaries(mesh, boundary_conditions)
dirichlet_bcs = collect_dirichlet_boundaries(
function_space, boundary_conditions, boundaries)
ds = dl.Measure('ds', domain=mesh, subdomain_data=boundaries)
dx = dl.Measure('dx', domain=mesh)
u = dl.TrialFunction(function_space)
v = dl.TestFunction(function_space)
# Previous solution
u_1 = dla.interpolate(init_condition, function_space)
u_2 = dla.Function(function_space)
u_2.assign(u_1)
if not second_order_timestepping:
theta = 1
else:
theta = 0.5
if second_order_timestepping and hasattr(forcing, 't'):
forcing_1 = copy_expression(forcing)
else:
forcing_1 = forcing
kappa = nonlinear_diffusion(u)
a = u*v*dx + theta*dt*kappa*dl.inner(dl.grad(u), dl.grad(v))*dx
L = (u_1 + theta*dt*forcing)*v*dx
# subtract of positivity preserving part added to diffusion
if positivity_tol > 0:
a -= positivity_tol*dl.inner(dl.grad(u), dl.grad(v))*dx
if second_order_timestepping:
kappa_1 = nonlinear_diffusion(u_1)
L -= (1-theta)*dt*kappa_1*dl.inner(dl.grad(u_1), dl.grad(v))*dx
L += (1-theta)*dt*forcing_1*v*dx
beta_1_list = []
alpha_1_list = []
for ii in range(num_bndrys):
if (boundary_conditions[ii][0] == 'robin'):
alpha = boundary_conditions[ii][3]
a += theta*dt*alpha*u*v*ds(ii)
if second_order_timestepping:
if hasattr(alpha, 't'):
alpha_1 = copy_expression(alpha)
alpha_1_list.append(alpha_1)
else:
alpha_1 = alpha
L -= (1-theta)*dt*alpha_1*u_1*v*ds(ii)
if ((boundary_conditions[ii][0] == 'robin') or
(boundary_conditions[ii][0] == 'neumann')):
beta = boundary_conditions[ii][2]
print(type(theta), type(dt), type(beta), type(v))
L -= theta*dt*beta*v*ds(ii)
if second_order_timestepping:
if hasattr(beta, 't'):
beta_1 = copy_expression(beta)
beta_1_list.append(beta_1)
else:
# boundary condition is constant in time
beta_1 = beta
L -= (1-theta)*dt*beta_1*v*ds(ii)
if hasattr(init_condition, 't'):
t = init_condition.t
else:
t = 0.0
while t < final_time:
# print('TIME',t)
# Update current time
prev_t = t
forcing_1.t = prev_t
t += dt
t = min(t, final_time)
forcing.t = t
# set current time for time varying boundary conditions
for ii in range(num_bndrys):
if hasattr(boundary_conditions[ii][2], 't'):
boundary_conditions[ii][2].t = t
# set previous time for time varying boundary conditions when
# using second order timestepping. lists will be empty if using
# first order timestepping
for jj in range(len(beta_1_list)):
beta_1_list[jj].t = prev_t
for jj in range(len(alpha_1_list)):
alpha_1_list[jj].t = prev_t
# solver must be redefined at every timestep
F = a-L
F = dl.action(F, u_2)
J = dl.derivative(F, u_2, u)
dla.solve(F == 0, u_2, dirichlet_bcs, J=J, solver_parameters=nlsparam)
# import matplotlib.pyplot as plt
# pl = dl.plot(sol); plt.colorbar(pl); plt.show()
# import matplotlib.pyplot as plt
# pl = dl.plot(sol); plt.show()
# Update previous solution
u_1.assign(u_2)
return u_1
inner, dot, dx = ufl.inner, ufl.dot, ufl.dx
# Geometry and elasticity
t, h, L = 2.0, 1.0, 5.0 # Thickness, height and length
E, nu = 210e3, 0.3 # Young Modulus
G = E / (2.0 * (1.0 + nu)) # Shear Modulus
lmbda = E * nu / ((1.0 + nu) * (1.0 - 2.0 * nu)) # Lambda
def simp(x):
return eps + (1 - eps) * x**p
max_volume = 0.4 * L * h # Volume constraint
p = 4 # Exponent
eps = fa.Constant(1.0e-6) # Epsilon for SIMP
# Mesh, Control and Solution Spaces
nelx = 192
nely = 64
mesh = fa.RectangleMesh.create(
[fn.Point(0.0, 0.0), fn.Point(L, h)], [nelx, nely],
fn.CellType.Type.triangle)
V = fn.VectorFunctionSpace(mesh, "CG", 1) # Displacements
C = fn.FunctionSpace(mesh, "CG", 1) # Control
# Volumetric Load
q = -10.0 / t
b = fa.Constant((0.0, q))
def run_model(function_space,
kappa,
forcing,
init_condition,
dt,
final_time,
boundary_conditions=None,
second_order_timestepping=False,
exact_sol=None,
velocity=None,
point_sources=None,
intermediate_times=None):
"""
Use implicit euler to solve transient advection diffusion equation
du/dt = grad (k* grad u) - vel*grad u + f
WARNING: when point sources solution changes significantly when mesh is
varied
"""
mesh = function_space.mesh()
time_independent_boundaries = False
if boundary_conditions is None:
bndry_obj = dl.CompiledSubDomain("on_boundary")
boundary_conditions = [['dirichlet', bndry_obj, dla.Constant(0)]]
time_independent_boundaries = True
num_bndrys = len(boundary_conditions)
boundaries = mark_boundaries(mesh, boundary_conditions)
dirichlet_bcs = collect_dirichlet_boundaries(function_space,
boundary_conditions,
boundaries)
# To express integrals over the boundary parts using ds(i), we must first
# redefine the measure ds in terms of our boundary markers:
ds = dl.Measure('ds', domain=mesh, subdomain_data=boundaries)
dx = dl.Measure('dx', domain=mesh)
# Variational problem at each time
u = dl.TrialFunction(function_space)
v = dl.TestFunction(function_space)
# Previous solution
if hasattr(init_condition, 't'):
assert init_condition.t == 0
u_1 = dla.interpolate(init_condition, function_space)
if not second_order_timestepping:
theta = 1
else:
theta = 0.5
if hasattr(forcing, 't'):
forcing_1 = copy_expression(forcing)
else:
forcing_1 = forcing
def steady_state_form(u, v, f):
F = kappa * dl.inner(dl.grad(u), dl.grad(v)) * dx
F -= f * v * dx
if velocity is not None:
F += dl.dot(velocity, dl.grad(u)) * v * dx
return F
F = u*v*dx-u_1*v*dx + dt*theta*steady_state_form(u, v, forcing) + \
dt*(1.-theta)*steady_state_form(u_1, v, forcing_1)
a, L = dl.lhs(F), dl.rhs(F)
# a = u*v*dx + theta*dt*kappa*dl.inner(dl.grad(u), dl.grad(v))*dx
# L = (u_1 + dt*theta*forcing)*v*dx
# if velocity is not None:
# a += theta*dt*v*dl.dot(velocity,dl.grad(u))*dx
# if second_order_timestepping:
# L -= (1-theta)*dt*dl.inner(kappa*dl.grad(u_1), dl.grad(v))*dx
# L += (1-theta)*dt*forcing_1*v*dx
# if velocity is not None:
# L -= (1-theta)*dt*(v*dl.dot(velocity,dl.grad(u_1)))*dx
beta_1_list = []
alpha_1_list = []
for ii in range(num_bndrys):
if (boundary_conditions[ii][0] == 'robin'):
alpha = boundary_conditions[ii][3]
a += theta * dt * alpha * u * v * ds(ii)
if second_order_timestepping:
if hasattr(alpha, 't'):
alpha_1 = copy_expression(alpha)
alpha_1_list.append(alpha_1)
else:
alpha_1 = alpha
L -= (1 - theta) * dt * alpha_1 * u_1 * v * ds(ii)
if ((boundary_conditions[ii][0] == 'robin')
or (boundary_conditions[ii][0] == 'neumann')):
beta = boundary_conditions[ii][2]
L -= theta * dt * beta * v * ds(ii)
if second_order_timestepping:
if hasattr(beta, 't'):
beta_1 = copy_expression(beta)
beta_1_list.append(beta_1)
else:
# boundary condition is constant in time
beta_1 = beta
L -= (1 - theta) * dt * beta_1 * v * ds(ii)
if time_independent_boundaries:
# TODO this can be used if dirichlet and robin conditions are not
# time dependent.
A = dla.assemble(a)
for bc in dirichlet_bcs:
bc.apply(A)
solver = dla.LUSolver(A)
#solver.parameters["reuse_factorization"] = True
else:
solver = None
u_2 = dla.Function(function_space)
u_2.assign(u_1)
t = 0.0
dt_tol = 1e-12
n_time_steps = 0
if intermediate_times is not None:
intermediate_u = []
intermediate_cnt = 0
# assert in chronological order
assert np.allclose(intermediate_times, np.array(intermediate_times))
assert np.all(intermediate_times < final_time)
while t < final_time - dt_tol:
# Update current time
prev_t = t
forcing_1.t = prev_t
t += dt
t = min(t, final_time)
forcing.t = t
# set current time for time varying boundary conditions
for ii in range(num_bndrys):
if hasattr(boundary_conditions[ii][2], 't'):
boundary_conditions[ii][2].t = t
# set previous time for time varying boundary conditions when
# using second order timestepping. lists will be empty if using
# first order timestepping
for jj in range(len(beta_1_list)):
beta_1_list[jj].t = prev_t
for jj in range(len(alpha_1_list)):
alpha_1_list[jj].t = prev_t
#A, b = dl.assemble_system(a, L, dirichlet_bcs)
# for bc in dirichlet_bcs:
# bc.apply(A,b)
if boundary_conditions is not None:
A = dla.assemble(a)
for bc in dirichlet_bcs:
bc.apply(A)
b = dla.assemble(L)
for bc in dirichlet_bcs:
bc.apply(b)
if point_sources is not None:
ps_list = []
for ii in range(len(point_sources)):
point, expr = point_sources[ii]
ps_list.append((dl.Point(point[0], point[1]), expr(t)))
ps = dla.PointSource(function_space, ps_list)
ps.apply(b)
if solver is None:
dla.solve(A, u_2.vector(), b)
else:
solver.solve(u_2.vector(), b)
# tape = dla.get_working_tape()
# tape.visualise()
#print ("t =", t, "end t=", final_time)
# Update previous solution
u_1.assign(u_2)
# import matplotlib.pyplot as plt
# plt.subplot(131)
# pp=dl.plot(u_1)
# plt.subplot(132)
# dl.plot(forcing,mesh=mesh)
# plt.subplot(133)
# dl.plot(forcing_1,mesh=mesh)
# plt.colorbar(pp)
# plt.show()
# compute error
if exact_sol is not None:
exact_sol.t = t
error = dl.errornorm(exact_sol, u_2)
print('t = %.2f: error = %.3g' % (t, error))
# dl.plot(exact_sol,mesh=mesh)
# plt.show()
if (intermediate_times is not None
and intermediate_cnt < intermediate_times.shape[0]
and t >= intermediate_times[intermediate_cnt]):
# save solution closest to intermediate time
u_t = dla.Function(function_space)
u_t.assign(u_2)
intermediate_u.append(u_t)
intermediate_cnt += 1
n_time_steps += 1
# print ("t =", t, "end t=", final_time,"# time steps", n_time_steps)
if intermediate_times is None:
return u_2
else:
return intermediate_u + [u_2]
def eval_cost(u, x):
J_form = dot(b, u) * dx + fa.Constant(1.0e-8) * dot(grad(x), grad(x)) * dx
J = fa.assemble(J_form)
return J
f = fa.Expression(
"10*exp(-(pow(x[0] - 0.5, 2) + pow(x[1] - 0.5, 2)) / 0.02)", degree=2)
u = fa.Function(V)
bcs = [fa.DirichletBC(V, fa.Constant(0.0), "on_boundary")]
inner, grad, dx = ufl.inner, ufl.grad, ufl.dx
JJ = 0.5 * inner(kappa0 * grad(u), grad(u)) * dx - kappa1 * f * u * dx
v = fenics.TestFunction(V)
F = fenics.derivative(JJ, u, v)
fa.solve(F == 0, u, bcs=bcs)
return u
templates = (fa.Constant(0.0), fa.Constant(0.0))
# inputs = (np.ones(1) * 0.5, np.ones(1) * 0.6)
# templates = (fa.Function(DG), fa.Function(DG))
true_kappa0 = fa.Constant(1.25)
true_kappa1 = fa.Constant(0.55)
# true_kappa0 = fa.Function(DG)
# true_kappa0.interpolate(fa.Constant(1.25))
# true_kappa1 = fa.Function(DG)
# true_kappa1.interpolate(fa.Constant(0.55))
true_solution = solve_fenics(true_kappa0, true_kappa1)
true_solution_numpy = to_numpy(true_solution)
