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 fenics_solve(f):
u = fa.Function(V, name="State")
v = fn.TestFunction(V)
F = (ufl.inner(ufl.grad(u), ufl.grad(v)) - f * v) * ufl.dx
bcs = [fa.DirichletBC(V, 0.0, "on_boundary")]
fa.solve(F == 0, u, bcs)
return u
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 forward(x):
u = fn.TrialFunction(V)
w = fn.TestFunction(V)
sigma = lmbda * tr(sym(grad(u))) * Identity(2) + 2 * G * sym(
grad(u)) # Stress
R = simp(x) * inner(sigma, grad(w)) * dx - dot(b, w) * dx
a, L = ufl.lhs(R), ufl.rhs(R)
u = fa.Function(V)
fa.solve(a == L, u, bcs)
return u
def compute_convergence_rates(run_model,
u_e,
max_degree=1,
num_levels=5,
min_n=8,
min_degree=1):
"""Compute convergences rates for various error norms
Adapted from https://fenicsproject.org/pub/tutorial/html/._ftut1020.html
"""
h = {} # discretization parameter: h[degree][level]
E = {} # error measure(s): E[degree][level][error_type]
# Iterate over degrees and mesh refinement levels
degrees = range(min_degree, max_degree + 1)
for degree in degrees:
n = min_n # coarsest mesh division
h[degree] = []
E[degree] = []
for ii in range(num_levels):
h[degree].append(1.0 / n)
u = run_model(n, degree)
if (hasattr(u_e, 'function_space')
and u.function_space() != u_e.function_space()):
V = dl.FunctionSpace(u.function_space().mesh(),
u_e.ufl_element().family(),
u_e.ufl_element().degree())
u_e_interp = dla.Function(V)
u_e_interp.interpolate(u_e)
errors = compute_errors(u_e_interp, u)
else:
# if not hasattr(u_e,'function_space') the u_e is an expression
errors = compute_errors(u_e, u)
E[degree].append(errors)
print('2 x (%d x %d) P%d mesh, %d unknowns, E1 = %g' %
(n, n, degree, u.function_space().dim(), errors['u - u_e']))
n *= 2
# Compute convergence rates
from math import log as ln # log is a fenics name too
etypes = list(E[min_degree][0].keys())
rates = {}
for degree in degrees:
rates[degree] = {}
for error_type in sorted(etypes):
rates[degree][error_type] = []
for i in range(1, num_levels):
Ei = E[degree][i][error_type]
Eim1 = E[degree][i - 1][error_type]
r = ln(Ei / Eim1) / ln(h[degree][i] / h[degree][i - 1])
rates[degree][error_type].append(round(r, 2))
return etypes, degrees, rates, E
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 forward(rho):
"""Solve the forward problem for a given fluid distribution rho(x)."""
w = fenics_adjoint.Function(W)
(u, p) = fenics.split(w)
(v, q) = fenics.TestFunctions(W)
inner, grad, dx, div = ufl.inner, ufl.grad, ufl.dx, ufl.div
F = (
alpha(rho) * inner(u, v) * dx
+ inner(grad(u), grad(v)) * dx
+ inner(grad(p), v) * dx
+ inner(div(u), q) * dx
)
bcs = [
fenics_adjoint.DirichletBC(W.sub(0).sub(1), 0, "on_boundary"),
fenics_adjoint.DirichletBC(W.sub(0).sub(0), inflow_outflow_bc, "on_boundary"),
]
fenics_adjoint.solve(F == 0, w, bcs=bcs)
return w
def get_surface_of_3d_function(Vh_2d, z, function):
for V in Vh_2d.split():
assert V.ufl_element().degree() == 1
mesh_coords = Vh_2d.mesh().coordinates().reshape((-1, 2))
v_2d = dl.vertex_to_dof_map(Vh_2d)
#v_2d = Vh_2d.dofmap().dofs()
values = np.zeros(Vh_2d.dim(), dtype=float)
for ii in range(mesh_coords.shape[0]):
y = function(mesh_coords[ii, 0], mesh_coords[ii, 1], z)
if np.isscalar(y):
stride = 1
else:
stride = len(y)
dofs = [v_2d[stride * ii + jj] for jj in range(stride)]
#print(dofs,y)
values[dofs] = y
function_2d = dla.Function(Vh_2d)
function_2d.vector()[:] = values
return function_2d
C = fn.FunctionSpace(mesh, "CG", 1) # Control
# Volumetric Load
q = -10.0 / t
b = fa.Constant((0.0, q))
def Left_boundary(x, on_boundary):
return on_boundary and abs(x[0]) < fn.DOLFIN_EPS
u_L = fa.Constant((0.0, 0.0))
bcs = [fa.DirichletBC(V, u_L, Left_boundary)]
@build_jax_fem_eval((fa.Function(C), ))
def forward(x):
u = fn.TrialFunction(V)
w = fn.TestFunction(V)
sigma = lmbda * tr(sym(grad(u))) * Identity(2) + 2 * G * sym(
grad(u)) # Stress
R = simp(x) * inner(sigma, grad(w)) * dx - dot(b, w) * dx
a, L = ufl.lhs(R), ufl.rhs(R)
u = fa.Function(V)
fa.solve(a == L, u, bcs)
return u
@build_jax_fem_eval((fa.Function(V), fa.Function(C)))
def eval_cost(u, x):
J_form = dot(b, u) * dx + fa.Constant(1.0e-8) * dot(grad(x), grad(x)) * dx
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]
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():
U_h = fenics.VectorElement("CG", mesh.ufl_cell(), 2)
P_h = fenics.FiniteElement("CG", mesh.ufl_cell(), 1)
W = fenics.FunctionSpace(mesh, U_h * P_h) # mixed Taylor-Hood function space
# Define the boundary condition on velocity
(x, y) = ufl.SpatialCoordinate(mesh)
l = 1.0 / 6.0 # noqa: E741
gbar = 1.0
cond1 = ufl.And(ufl.gt(y, (1.0 / 4 - l / 2)), ufl.lt(y, (1.0 / 4 + l / 2)))
val1 = gbar * (1 - (2 * (y - 0.25) / l) ** 2)
cond2 = ufl.And(ufl.gt(y, (3.0 / 4 - l / 2)), ufl.lt(y, (3.0 / 4 + l / 2)))
val2 = gbar * (1 - (2 * (y - 0.75) / l) ** 2)
inflow_outflow = ufl.conditional(cond1, val1, ufl.conditional(cond2, val2, 0))
inflow_outflow_bc = fenics_adjoint.project(inflow_outflow, W.sub(0).sub(0).collapse())
solve_templates = (fenics_adjoint.Function(A),)
assemble_templates = (fenics_adjoint.Function(W), fenics_adjoint.Function(A))
@build_jax_fem_eval(solve_templates)
def forward(rho):
"""Solve the forward problem for a given fluid distribution rho(x)."""
w = fenics_adjoint.Function(W)
(u, p) = fenics.split(w)
(v, q) = fenics.TestFunctions(W)
inner, grad, dx, div = ufl.inner, ufl.grad, ufl.dx, ufl.div
F = (
alpha(rho) * inner(u, v) * dx
+ inner(grad(u), grad(v)) * dx
+ inner(grad(p), v) * dx
def main():
parser = argparse.ArgumentParser()
parser.add_argument('algNO', nargs='?', type=int, default=3)
parser.add_argument('seedNO', nargs='?', type=int, default=2017)
parser.add_argument('num_samp', nargs='?', type=int, default=1000)
parser.add_argument('num_burnin', nargs='?', type=int, default=0)
parser.add_argument('step_sizes',
nargs='?',
type=float,
default=[.02, .3, .15, 4.5, 1.5])
parser.add_argument('step_nums',
nargs='?',
type=int,
default=[1, 1, 4, 1, 4])
parser.add_argument('algs',
nargs='?',
type=str,
default=('pCN', 'infMALA', 'infHMC', 'DRinfmMALA',
'DRinfmHMC'))
args = parser.parse_args()
# set the (global) random seed
np.random.seed(args.seedNO)
print('Random seed is set to %d.' % args.seedNO)
# define the model
nozz_w = 1.
nx = 40
ny = 80
rans = RANS(nozz_w=nozz_w, nx=nx, ny=ny)
rans.setup(args.seedNO, src4init='solution')
# (randomly) initialize parameter
# noise = dl.Vector()
# rans.prior.init_vector(noise,"noise")
# Random.normal(noise, 1., True)
PARAMETER = 1
# parameter = rans.model_stat.generate_vector(PARAMETER)
# rans.prior.sample(noise,parameter)
# read from MAP
parameter = dl.Function(rans.Vh[PARAMETER])
MAP_file = os.path.join(
os.getcwd(),
'analysis-L25W8-nofreshinit-yesparacont-2000/_samp_DRinfmMALA_dim3321_2018-06-06-16-24-32.h5'
)
if os.path.isfile(MAP_file):
f = dl.HDF5File(rans.mpi_comm, MAP_file, "r")
f.read(parameter, 'sample_{0}'.format(999))
f.close()
else:
parameter = rans.get_MAP(SAVE=True)
parameter = parameter.vector()
# forward solver: whether do continuation
fresh_init = False
para_cont = True
if rans.rank == 0:
print("Forward solving with" + {
True: "",
False: "out"
}[fresh_init] + " repeated fresh initialization and with" + {
True: "",
False: "out"
}[para_cont] + " parameter continuation...")
# run MCMC to generate samples
if rans.rank == 0:
print("Preparing %s sampler with " + {
True: "initial",
False: "fixed"
}[adpt_stepsz] + " step size %g for %d step(s)..." %
(args.algs[args.algNO], args.step_sizes[args.algNO],
args.step_nums[args.algNO]))
inf_MC=geoinfMC(parameter,rans,args.step_sizes[args.algNO],args.step_nums[args.algNO],args.algs[args.algNO],True, \
target_acpt=0.7,fresh_init=fresh_init,para_cont=para_cont,src4init='solution')
inf_MC.setup(args.num_samp, args.num_burnin, 1, mpi_comm=rans.mpi_comm)
if rans.rank != 0:
inf_MC.parameters['print_level'] = -1
inf_MC.sample(num_retry_bad=0)
# append PDE information including the count of solving
filename = os.path.join(inf_MC.savepath, inf_MC.filename + '.pckl')
f = open(filename, 'ab')
# soln_count=rans.soln_count.copy()
soln_count = [
rans.pde.solveFwd.count, rans.pde.solveAdj.count,
rans.pde.solveIncremental.count
]
pickle.dump([nozz_w, nx, ny, fresh_init, para_cont, soln_count, args], f)
f.close()
# rename
f_newname = os.path.join(inf_MC.savepath,
'RANS_seed' + str(args.seedNO) + '_' +
inf_MC.filename + '.pckl') # change filename
if rans.rank == 0:
os.rename(filename, f_newname)
C = fn.FunctionSpace(mesh, "CG", 1) # Control
# Volumetric Load
q = -10.0 / t
b = fa.Constant((0.0, q))
def Left_boundary(x, on_boundary):
return on_boundary and abs(x[0]) < fn.DOLFIN_EPS
u_L = fa.Constant((0.0, 0.0))
bcs = [fa.DirichletBC(V, u_L, Left_boundary)]
@build_jax_fem_eval((fa.Function(C), ))
def forward(x):
u = fn.TrialFunction(V)
w = fn.TestFunction(V)
sigma = lmbda * tr(sym(grad(u))) * Identity(2) + 2 * G * sym(
grad(u)) # Stress
R = simp(x) * inner(sigma, grad(w)) * dx - dot(b, w) * dx
a, L = ufl.lhs(R), ufl.rhs(R)
u = fa.Function(V)
fa.solve(a == L, u, bcs)
return u
@build_jax_fem_eval((fa.Function(V), fa.Function(C)))
def eval_cost(u, x):
J_form = dot(b, u) * dx + fa.Constant(1.0e-8) * dot(grad(x), grad(x)) * dx
def constrained_newton_energy_solve(F,
uh,
dirichlet_bcs=None,
bc0=None,
linear_solver=None,
opts=dict(),
C=None,
constraint_vec=None):
"""
See https://uvilla.github.io/inverse15/UnconstrainedMinimization.html
F: dl.Expression
The energy functional.
uh : dl.Function
Final solution. The initial state on entry to the function
will be used as initial guess and then overwritten
dirichlet_bcs : list
The Dirichlet boundary conditions on the unknown u.
bc0 : list
The Dirichlet boundary conditions for the step (du) in the Newton
iterations.
"""
max_iter = opts.get("max_iter", 20)
# exit when sqrt(g,g)/sqrt(g_0,g_0) <= rel_tolerance"
rtol = opts.get("rel_tolerance", 1e-8)
# exit when sqrt(g,g) <= abs_tolerance
atol = opts.get("abs_tolerance", 1e-9)
# exit when (g,du) <= gdu_tolerance
gdu_tol = opts.get("gdu_tolerance", 1e-14)
# define armijo sufficient decrease
c_armijo = opts.get("c_armijo", 1e-4)
# exit if max backtracking steps reached
max_backtrack = opts.get("max_backtracking_iter", 20)
# define verbosity
prt_level = opts.get("print_level", 0)
termination_reasons = [
"Maximum number of Iteration reached", #0
"Norm of the gradient less than tolerance", #1
"Maximum number of backtracking reached", #2
"Norm of (g, du) less than tolerance" #3
]
it = 0
total_cg_iter = 0
converged = False
reason = 0
L = F
if C is not None:
L += C
if prt_level > 0:
print("Solving Constrained Nonlinear Problem")
else:
if prt_level > 0:
print("Solving Unconstrained Nonlinear Problem")
# Compute gradient and hessian
grad = dla.derivative(L, uh)
H = dla.derivative(grad, uh)
# Applying boundary conditions
if dirichlet_bcs is not None:
if type(dirichlet_bcs) is dla.DirichletBC:
bcsl = [dirichlet_bcs]
else:
bcsl = dirichlet_bcs
[bc.apply(uh.vector()) for bc in bcsl]
if constraint_vec is not None:
assert C is not None
dcd_state = dla.assemble(dla.derivative(C, u))
dcd_lagrangeMult = dcd_state * constraint_vec
if not dcd_lagrangeMult.norm("l2") < 1.e-14:
msg = "The initial guess does not satisfy the constraint."
raise ValueError(msg)
# Setting variables
Fn = dla.assemble(F)
gn = dla.assemble(grad)
g0_norm = gn.norm("l2")
gn_norm = g0_norm
tol = max(g0_norm * rtol, atol)
du = dla.Function(uh.function_space()).vector()
#if linear_solver =='PETScLU':
# linear_solver = dl.PETScLUSolver(uh.function_space().mesh().mpi_comm())
#else:
# assert linear_solver is None
if prt_level > 0:
print("{0:>3} {1:>6} {2:>15} {3:>15} {4:>15} {5:>15}".format(
"Nit", "CGit", "Energy", "||g||", "(g,du)", "alpha"))
print("{0:3d} {1:6d} {2:15e} {3:15e} {4:15} {5:15}".format(
0, 0, Fn, g0_norm, " NA ", " NA"))
converged = False
reason = 0
for it in range(max_iter):
if bc0 is not None:
[Hn, gn] = dla.assemble_system(H, grad, bc0)
else:
Hn = dla.assemble(H)
gn = dla.assemble(grad)
Hn.init_vector(du, 1)
if linear_solver is None:
lin_it = dla.solve(Hn, du, -gn, "cg", "petsc_amg")
else:
print('a')
lin_it = dla.solve(Hn, du, -gn, "lu")
#linear_solver.set_operator(Hn)
#lin_it = linear_solver.solve(du, -gn)
total_cg_iter += lin_it
du_gn = du.inner(gn)
alpha = 1.0
if (np.abs(du_gn) < gdu_tol):
converged = True
reason = 3
uh.vector().axpy(alpha, du)
Fn = dla.assemble(F)
gn_norm = gn.norm("l2")
break
uh_backtrack = uh.copy(deepcopy=True)
bk_converged = False
#Backtrack
for j in range(max_backtrack):
uh.assign(uh_backtrack)
uh.vector().axpy(alpha, du)
Fnext = dla.assemble(F)
#print(Fnext,Fn + alpha*c_armijo*du_gn)
if Fnext < Fn + alpha * c_armijo * du_gn:
Fn = Fnext
bk_converged = True
break
alpha /= 2.
if not bk_converged:
reason = 2
break
gn_norm = gn.norm("l2")
if prt_level > 0:
print("{0:3d} {1:6d} {2:15e} {3:15e} {4:15e} {5:15e}".format(
it + 1, lin_it, Fn, gn_norm, du_gn, alpha))
if gn_norm < tol:
converged = True
reason = 1
break
if prt_level > 0:
if reason is 3:
print("{0:3d} {1:6d} {2:15e} {3:15e} {4:15e} {5:15e}".format(
it + 1, lin_it, Fn, gn_norm, du_gn, alpha))
print(termination_reasons[reason])
if converged:
print("Newton converged in ", it, \
"nonlinear iterations and ", total_cg_iter,
"linear iterations." )
else:
print("Newton did NOT converge in ", it, "iterations.")
print("Final norm of the gradient: ", gn_norm)
print("Value of the cost functional: ", Fn)
if reason in [0, 2]:
raise Exception(termination_reasons[reason])
return uh
def __call__(self, coef):
assert coef.shape[1] == 1
np_field = super().__call__(coef)[:, 0]
field = dla.Function(self.function_space)
field.vector()[:] = np_field
return field
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
def load_fenics_function(function_space, filename):
function = dla.Function(function_space)
fFile = dl.HDF5File(function_space.mesh().mpi_comm(), filename, "r")
fFile.read(function, "/f")
fFile.close()
return function
# Create mesh, refined in the center
n = 64
mesh = fn.UnitSquareMesh(n, n)
cf = fn.MeshFunction("bool", mesh, mesh.geometry().dim())
subdomain = fn.CompiledSubDomain(
"std::abs(x[0]-0.5) < 0.25 && std::abs(x[1]-0.5) < 0.25"
)
subdomain.mark(cf, True)
mesh = fa.Mesh(fn.refine(mesh, cf))
# Define discrete function spaces and functions
V = fn.FunctionSpace(mesh, "CG", 1)
W = fn.FunctionSpace(mesh, "DG", 0)
solve_templates = (fa.Function(W),)
assemble_templates = (fa.Function(V), fa.Function(W))
# Define and solve the Poisson equation
@build_jax_fem_eval(solve_templates)
def fenics_solve(f):
u = fa.Function(V, name="State")
v = fn.TestFunction(V)
F = (ufl.inner(ufl.grad(u), ufl.grad(v)) - f * v) * ufl.dx
bcs = [fa.DirichletBC(V, 0.0, "on_boundary")]
fa.solve(F == 0, u, bcs)
return u
# Define functional of interest and the reduced functional
@build_jax_fem_eval(assemble_templates)
def unconstrained_newton_solve(F,
J,
uh,
dirichlet_bcs=None,
bc0=None,
linear_solver=None,
opts=dict()):
"""
F: dl.Expression
The variational form.
uh : dl.Function
Final solution. The initial state on entry to the function
will be used as initial guess and then overwritten
dirichlet_bcs : list
The Dirichlet boundary conditions on the unknown u.
bc0 : list
The Dirichlet boundary conditions for the step (du) in the Newton
iterations.
"""
max_iter = opts.get("max_iter", 50)
# exit when sqrt(g,g)/sqrt(g_0,g_0) <= rel_tolerance"
rtol = opts.get("rel_tolerance", 1e-8)
# exit when sqrt(g,g) <= abs_tolerance
atol = opts.get("abs_tolerance", 1e-9)
# exit when (g,du) <= gdu_tolerance
gdu_tol = opts.get("gdu_tolerance", 1e-14)
# define armijo sufficient decrease
c_armijo = opts.get("c_armijo", 1e-4)
# exit if max backtracking steps reached
max_backtrack = opts.get("max_backtracking_iter", 20)
# define verbosity
prt_level = opts.get("print_level", 0)
termination_reasons = [
"Maximum number of Iteration reached", #0
"Norm of the gradient less than tolerance", #1
"Maximum number of backtracking reached", #2
"Norm of (g, du) less than tolerance", #3
"Norm of residual less than tolerance"
] #4
it = 0
total_cg_iter = 0
converged = False
reason = 0
if prt_level > 0:
print("Solving Nonlinear Problem")
# Applying boundary conditions
if dirichlet_bcs is not None:
if type(dirichlet_bcs) is dla.DirichletBC:
bcsl = [dirichlet_bcs]
else:
bcsl = dirichlet_bcs
[bc.apply(uh.vector()) for bc in bcsl]
if type(bc0) is dla.DirichletBC:
bc0 = [bc0]
# Setting variables
gn = dla.assemble(F)
res_func = dla.Function(uh.function_space())
res_func.assign(dla.Function(uh.function_space(), gn))
res = res_func.vector()
if bc0 is not None:
for bc in bc0:
bc.apply(res)
Fn = res.norm("l2")
g0_norm = gn.norm("l2")
gn_norm = g0_norm
tol = max(g0_norm * rtol, atol)
res_tol = max(Fn * rtol, atol)
du = dla.Function(uh.function_space()).vector()
if linear_solver == 'PETScLU':
linear_solver = dla.PETScLUSolver(
uh.function_space().mesh().mpi_comm())
else:
assert linear_solver is None
if prt_level > 0:
print("{0:>3} {1:>6} {2:>15} {3:>15} {4:>15} {5:>15} {6:>6}".format(
"Nit", "CGit", "||r||", "||g||", "(g,du)", "alpha", "Nbt"))
print("{0:3d} {1:6d} {2:15e} {3:15e} {4:15} {5:10} {6:s}".
format(0, 0, Fn, g0_norm, " NA ", " NA", "NA"))
converged = False
reason = 0
nbt = 0
for it in range(max_iter):
if bc0 is not None:
[Hn, gn] = dla.assemble_system(J, F, bc0)
else:
Hn = dla.assemble(J)
gn = dla.assemble(F)
Hn.init_vector(du, 1)
if linear_solver is None:
lin_it = dla.solve(Hn, du, -gn, "cg", "petsc_amg")
else:
linear_solver.set_operator(Hn)
lin_it = linear_solver.solve(du, -gn)
total_cg_iter += lin_it
du_gn = du.inner(gn)
alpha = 1.0
if (np.abs(du_gn) < gdu_tol):
converged = True
reason = 3
uh.vector().axpy(alpha, du)
gn_norm = gn.norm("l2")
Fn = gn_norm
break
uh_backtrack = uh.copy(deepcopy=True)
bk_converged = False
#Backtrack
for nbt in range(max_backtrack):
uh.assign(uh_backtrack)
uh.vector().axpy(alpha, du)
res = dla.assemble(F)
if bc0 is not None:
for bc in bc0:
bc.apply(res)
Fnext = res.norm("l2")
#print(Fn,Fnext,Fn + alpha*c_armijo*du_gn)
if Fnext < Fn + alpha * c_armijo * du_gn:
#if True:
Fn = Fnext
bk_converged = True
break
alpha /= 2.
if not bk_converged:
reason = 2
break
gn_norm = gn.norm("l2")
if prt_level > 0:
print("{0:3d} {1:6d} {2:15e} {3:15e} {4:15e} {5:15e} {6:3d}".
format(it + 1, lin_it, Fn, gn_norm, du_gn, alpha, nbt + 1))
if gn_norm < tol:
converged = True
reason = 1
break
if Fn < res_tol:
converged = True
reason = 4
break
if prt_level > 0:
if reason is 3:
print("{0:3d} {1:6d} {2:15e} {3:15e} {4:15e} {5:15e} {6:3d}".
format(it + 1, lin_it, Fn, gn_norm, du_gn, alpha, nbt + 1))
print(termination_reasons[reason])
if converged:
print("Newton converged in ", it, \
"nonlinear iterations and ", total_cg_iter,
"linear iterations." )
else:
print("Newton did NOT converge in ", it, "iterations.")
print("Final norm of the gradient: ", gn_norm)
print("Value of the cost functional: ", Fn)
if reason in [0, 2]:
raise Exception(termination_reasons[reason])
return uh