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)