def get_trace_nullspace_vecs(forward, nullspace, V, V_d, TraceSpace): """Gets the nullspace vectors corresponding to the Schur complement system for the multipliers. :arg forward: A Slate expression denoting the forward elimination operator. :arg nullspace: A nullspace for the original mixed problem :arg V: The original "unbroken" space. :arg V_d: The broken space. :arg TraceSpace: The space of approximate traces. Returns: A list of vectors describing the nullspace of the multiplier system """ from firedrake import project, assemble, Function vecs = nullspace.getVecs() tmp = Function(V) tmp_b = Function(V_d) tnsp_tmp = Function(TraceSpace) forward_action = forward * tmp_b new_vecs = [] for v in vecs: with tmp.dat.vec as t: v.copy(t) project(tmp, tmp_b) assemble(forward_action, tensor=tnsp_tmp) with tnsp_tmp.dat.vec_ro as v: new_vecs.append(v.copy()) return new_vecs
def test_diagnostic_solver(): Nx, Ny = 32, 32 mesh2d = firedrake.RectangleMesh(Nx, Ny, Lx, Ly) mesh = firedrake.ExtrudedMesh(mesh2d, layers=1) Q = firedrake.FunctionSpace(mesh, family='CG', degree=2, vfamily='DG', vdegree=0) x, y, ζ = firedrake.SpatialCoordinate(mesh) h = firedrake.interpolate(h0 - dh * x / Lx, Q) s = firedrake.interpolate(d + h0 - dh + ds * (1 - x / Lx), Q) u_expr = firedrake.as_vector(((0.95 + 0.05 * ζ) * exact_u(x), 0)) A = firedrake.Constant(icepack.rate_factor(254.15)) C = firedrake.Constant(0.001) model = icepack.models.HybridModel() opts = {'dirichlet_ids': [1, 3, 4], 'tol': 1e-12} max_degree = 5 Nz = 32 xs = np.array([(Lx / 2, Ly / 2, k / Nz) for k in range(Nz + 1)]) us = np.zeros((max_degree + 1, Nz + 1)) for vdegree in range(max_degree, 0, -1): solver = icepack.solvers.FlowSolver(model, **opts) V = firedrake.VectorFunctionSpace(mesh, dim=2, family='CG', degree=2, vfamily='GL', vdegree=vdegree) u0 = firedrake.interpolate(u_expr, V) u = solver.diagnostic_solve(velocity=u0, thickness=h, surface=s, fluidity=A, friction=C) V0 = firedrake.VectorFunctionSpace(mesh, dim=2, family='CG', degree=2, vfamily='DG', vdegree=0) depth_avg_u = firedrake.project(u, V0) shear_u = firedrake.project(u - depth_avg_u, V) assert icepack.norm(shear_u, norm_type='Linfty') > 1e-2 us_center = np.array(u.at(xs, tolerance=1e-6)) us[vdegree, :] = np.sqrt(np.sum(us_center**2, 1)) norm = np.linalg.norm(us[max_degree, :]) error = np.linalg.norm(us[vdegree, :] - us[max_degree, :]) / norm print(error, flush=True) assert error < 1e-2
def initial_conditions(self): ic1 = fd.project(fd.Expression([0., 0., 0.]), self.V) ic2 = fd.project( fd.Expression(["0.1*(1-cos(pi*x[2]/Lz/2.))", 0., 0.], Lz=self.Lz), self.V) # ic3 = fd.project( fd.Expression(["0.1*x[2]/Lz_B", 0., 0.], Lz_B=self.Lz), self.V ) self.X.assign(ic2) # self.static_solver() self.U.assign(ic1)
def test_damage_transport(): nx, ny = 32, 32 Lx, Ly = 20e3, 20e3 mesh = firedrake.RectangleMesh(nx, ny, Lx, Ly) x, y = firedrake.SpatialCoordinate(mesh) V = firedrake.VectorFunctionSpace(mesh, "CG", 2) Q = firedrake.FunctionSpace(mesh, "CG", 2) u0 = 100.0 h0, dh = 500.0, 100.0 T = 268.0 ρ = ρ_I * (1 - ρ_I / ρ_W) Z = icepack.rate_factor(T) * (ρ * g * h0 / 4)**n q = 1 - (1 - (dh / h0) * (x / Lx))**(n + 1) du = Z * q * Lx * (h0 / dh) / (n + 1) u = interpolate(as_vector((u0 + du, 0)), V) h = interpolate(h0 - dh * x / Lx, Q) A = firedrake.Constant(icepack.rate_factor(T)) S = firedrake.TensorFunctionSpace(mesh, "DG", 1) ε = firedrake.project(sym(grad(u)), S) M = firedrake.project(membrane_stress(strain_rate=ε, fluidity=A), S) degree = 1 Δ = firedrake.FunctionSpace(mesh, "DG", degree) D_inflow = firedrake.Constant(0.0) D = firedrake.Function(Δ) damage_model = icepack.models.DamageTransport() damage_solver = icepack.solvers.DamageSolver(damage_model) final_time = Lx / u0 max_speed = u.at((Lx - 1.0, Ly / 2), tolerance=1e-10)[0] δx = Lx / nx timestep = δx / max_speed / (2 * degree + 1) num_steps = int(final_time / timestep) dt = final_time / num_steps for step in range(num_steps): D = damage_solver.solve( dt, damage=D, velocity=u, strain_rate=ε, membrane_stress=M, damage_inflow=D_inflow, ) Dmax = D.dat.data_ro[:].max() assert 0 < Dmax < 1
def initial_conditions(self): # old: ic1 = fd.project( fd.Expression([0.,0.,0.]), self.V ) # old: ic2 = fd.project( fd.Expression(["0.1*(1-cos(pi*x[2]/Lz/2.))",0.,0.], Lz=self.Lz), self.V ) ic1 = fd.project(fd.Constant([0., 0., 0.]), self.V) x = fd.SpatialCoordinate(self.mesh) ic2 = fd.project( fd.as_vector( [0.1 * (1 - fd.cos(fd.pi * x[2] / self.Lz / 2.)), 0., 0.]), self.V) # ic3 = fd.project( fd.Expression(["0.1*x[2]/Lz_B", 0., 0.], Lz_B=self.Lz), self.V ) self.X.assign(ic2) # self.static_solver() self.U.assign(ic1)
def test_eigenvalues(): nx, ny = 32, 32 mesh = firedrake.UnitSquareMesh(nx, ny) x, y = firedrake.SpatialCoordinate(mesh) V = firedrake.VectorFunctionSpace(mesh, family='CG', degree=2) u = interpolate(as_vector((x, 0)), V) Q = firedrake.FunctionSpace(mesh, family='DG', degree=2) ε = sym(grad(u)) Λ1, Λ2 = eigenvalues(ε) λ1 = firedrake.project(Λ1, Q) λ2 = firedrake.project(Λ2, Q) assert norm(λ1 - Constant(1)) < norm(u) / (nx * ny) assert norm(λ2) < norm(u) / (nx * ny)
def test_converting_fields(): δT = 5.0 T_surface = Tm - δT T_expr = firedrake.min_value(Tm, T_surface + 2 * δT * (1 - ζ)) f_expr = firedrake.max_value(0, 0.0033 * (1 - 2 * ζ)) model = icepack.models.HeatTransport3D() E = firedrake.project(model.energy_density(T_expr, f_expr), Q) f = firedrake.project(model.meltwater_fraction(E), Q) T = firedrake.project(model.temperature(E), Q) avg_meltwater = firedrake.assemble(f * ds_b) / (Lx * Ly) assert avg_meltwater > 0 avg_temp = firedrake.assemble(T * h * dx) / firedrake.assemble(h * dx) assert (avg_temp > T_surface) and (avg_temp < Tm)
def __init__(self, V, gcur, u0, u0_mult, i, t, dt): if V.component is not None: # bottommost space is bit of VFS if V.parent.index is None: # but not part of a MFS sub = V.component try: gdat = interpolate(gcur - u0_mult[i] * u0, V) gmethod = lambda g, u: gdat.interpolate(g - u0_mult[i] * u. sub(sub)) except: # noqa: E722 gdat = project(gcur - u0_mult[i] * u0, V) gmethod = lambda g, u: gdat.project(g - u0_mult[i] * u.sub( sub)) else: # V is a bit of a VFS inside an MFS sub0 = V.parent.index sub1 = V.component try: gdat = interpolate( gcur - u0_mult[i] * u0.sub(sub0).sub(sub1), V) gmethod = lambda g, u: gdat.interpolate(g - u0_mult[ i] * u.sub(sub0).sub(sub1)) except: # noqa: E722 gdat = project(gcur - u0_mult[i] * u0.sub(sub0).sub(sub1), V) gmethod = lambda g, u: gdat.project(g - u0_mult[i] * u.sub( sub0).sub(sub1)) else: # V is not a bit of a VFS if V.index is None: # not part of MFS, either try: gdat = interpolate(gcur - u0_mult[i] * u0, V) gmethod = lambda g, u: gdat.interpolate(g - u0_mult[i] * u) except: # noqa: E722 gdat = project(gcur - u0_mult[i] * u0, V) gmethod = lambda g, u: gdat.project(g - u0_mult[i] * u) else: # part of MFS sub = V.index try: gdat = interpolate(gcur - u0_mult[i] * u0.sub(sub), V) gmethod = lambda g, u: gdat.interpolate(g - u0_mult[i] * u. sub(sub)) except: # noqa: E722 gdat = project(gcur - u0_mult[i] * u0.sub(sub), V) gmethod = lambda g, u: gdat.project(g - u0_mult[i] * u.sub( sub)) self.gstuff = (gdat, gcur, gmethod)
def apply(self, pc, x, y): """We solve the forward eliminated problem for the approximate traces of the scalar solution (the multipliers) and reconstruct the "broken flux and scalar variable." Lastly, we project the broken solutions into the mimetic non-broken finite element space. """ from firedrake import project # Transfer non-broken x into a firedrake function with self.unbroken_rhs.dat.vec as v: x.copy(v) # Transfer unbroken_rhs into broken_rhs field0, field1 = self.unbroken_rhs.split() bfield0, bfield1 = self.broken_rhs.split() # This updates broken_rhs project(field0, bfield0) field1.dat.copy(bfield1.dat) # Compute the rhs for the multiplier system self._assemble_Srhs() # Solve the system for the Lagrange multipliers with self.schur_rhs.dat.vec_ro as b: with self.trace_solution.dat.vec as x: self.ksp.solve(b, x) # Assemble the pressure and velocity (in that order) # using the Lagrange multipliers self._assemble_pressure() self._assemble_velocity() # Project the broken solution into non-broken spaces sigma_h, u_h = self.broken_solution.split() sigma_u, u_u = self.unbroken_solution.split() u_h.dat.copy(u_u.dat) self.projector.project() with self.unbroken_solution.dat.vec_ro as v: v.copy(y)
def mult(self, mat, x, y): # Copy function data into the fivredrake function self.x_fntn.dat.data[:] = x[:] # Transfer the function to meshmode self.meshmode_src_connection.from_firedrake(project( self.x_fntn, dgfspace), out=self.x_mm_fntn) # Restrict to boundary x_mm_fntn_on_bdy = src_bdy_connection(self.x_mm_fntn) # Apply the operation potential_int_mm = self.pyt_op(self.actx, u=x_mm_fntn_on_bdy, k=self.k) grad_potential_int_mm = self.pyt_grad_op(self.actx, u=x_mm_fntn_on_bdy, k=self.k) # Store in firedrake self.potential_int.dat.data[target_indices] = potential_int_mm.get( ) for dim in range(grad_potential_int_mm.shape[0]): self.grad_potential_int.dat.data[ target_indices, dim] = grad_potential_int_mm[dim].get() # Integrate the potential r""" Compute the inner products using firedrake. Note this will be subtracted later, hence appears off by a sign. .. math:: \langle n(x) \cdot \nabla( \int_\Gamma( u(y) \partial_n H_0^{(1)}(\kappa |x - y|) )d\gamma(y) ), v \rangle_\Sigma - \langle i \kappa \cdot \int_\Gamma( u(y) \partial_n H_0^{(1)}(\kappa |x - y|) )d\gamma(y), v \rangle_\Sigma """ self.pyt_result = assemble( inner(inner(self.grad_potential_int, self.n), self.v) * ds(outer_bdy_id) - inner(self.potential_int, self.v) * ds(outer_bdy_id)) # y <- Ax - evaluated potential self.A.mult(x, y) with self.pyt_result.dat.vec_ro as ep: y.axpy(-1, ep)
def siahorizontalvelocity(mesh): hbase = surfaceelevation(mesh) if mesh._base_mesh.cell_dimension() == 2: if mesh._base_mesh.ufl_cell() == fd.quadrilateral: Vvectorbase = fd.VectorFunctionSpace(mesh._base_mesh, 'DQ', 0) VvectorR = fd.VectorFunctionSpace(mesh, 'DQ', 0, vfamily='R', vdegree=0, dim=2) else: Vvectorbase = fd.VectorFunctionSpace(mesh._base_mesh, 'DP', 0) VvectorR = fd.VectorFunctionSpace(mesh, 'DP', 0, vfamily='R', vdegree=0, dim=2) gradhbase = fd.project(fd.grad(hbase), Vvectorbase) Vvector = fd.VectorFunctionSpace(mesh, 'DQ', 0, dim=2) elif mesh._base_mesh.cell_dimension() == 1: Vvectorbase = fd.FunctionSpace(mesh._base_mesh, 'DP', 0) gradhbase = fd.project(hbase.dx(0), Vvectorbase) VvectorR = fd.FunctionSpace(mesh, 'DP', 0, vfamily='R', vdegree=0) Vvector = fd.FunctionSpace(mesh, 'DQ', 0) else: raise ValueError('base mesh of extruded input mesh must be 1D or 2D') gradh = fd.Function(VvectorR) gradh.dat.data[:] = gradhbase.dat.data_ro[:] h = extend(mesh, hbase) DQ0 = fd.FunctionSpace(mesh, 'DQ', 0) h0 = fd.project(h, DQ0) x = fd.SpatialCoordinate(mesh) z0 = fd.project(x[mesh._base_mesh.cell_dimension()], DQ0) # FIXME following only valid in flat bed case uvsia = -Gamma * (h0**(n + 1) - (h0 - z0)**(n + 1)) * abs(gradh)**(n - 1) * gradh uv = fd.Function(Vvector).interpolate(uvsia) uv.rename('velocitySIA') return uv
def create_schur_nullspace(P, forward, V, V_d, TraceSpace, comm): """Gets the nullspace vectors corresponding to the Schur complement system for the multipliers. :arg P: The mixed operator from the ImplicitMatrixContext. :arg forward: A Slate expression denoting the forward elimination operator. :arg V: The original "unbroken" space. :arg V_d: The broken space. :arg TraceSpace: The space of approximate traces. Returns: A nullspace (if there is one) for the Schur-complement system. """ from firedrake import assemble, Function, project nullspace = P.getNullSpace() if nullspace.handle == 0: # No nullspace return None vecs = nullspace.getVecs() tmp = Function(V) tmp_b = Function(V_d) tnsp_tmp = Function(TraceSpace) forward_action = forward * tmp_b new_vecs = [] for v in vecs: with tmp.dat.vec_wo as t: v.copy(t) project(tmp, tmp_b) assemble(forward_action, tensor=tnsp_tmp) with tnsp_tmp.dat.vec_ro as v: new_vecs.append(v.copy()) # Normalize for v in new_vecs: v.normalize() schur_nullspace = PETSc.NullSpace().create(vectors=new_vecs, comm=comm) return schur_nullspace
def tripcolor(function, *args, **kwargs): r"""Plot a discontinuous finite element field""" axes = kwargs.pop('axes', plt.gca()) mesh = function.ufl_domain() coordinates = _get_coordinates(mesh) coords = coordinates.dat.data_ro x, y = coords[:, 0], coords[:, 1] triangles = coordinates.cell_node_map().values triangulation = matplotlib.tri.Triangulation(x, y, triangles) Q = firedrake.FunctionSpace(mesh, family='DG', degree=0) fn = firedrake.project(function, Q) return axes.tripcolor(triangulation, fn.dat.data_ro, *args, **kwargs)
def wave(n, deg, butcher_tableau, splitting=AI): N = 2**n msh = UnitIntervalMesh(N) params = { "snes_type": "ksponly", "ksp_type": "preonly", "mat_type": "aij", "pc_type": "lu" } V = FunctionSpace(msh, "CG", deg) W = FunctionSpace(msh, "DG", deg - 1) Z = V * W x, = SpatialCoordinate(msh) t = Constant(0.0) dt = Constant(2.0 / N) up = project(as_vector([0, sin(pi * x)]), Z) u, p = split(up) v, w = TestFunctions(Z) F = (inner(Dt(u), v) * dx + inner(u.dx(0), w) * dx + inner(Dt(p), w) * dx - inner(p, v.dx(0)) * dx) E = 0.5 * (inner(u, u) * dx + inner(p, p) * dx) stepper = TimeStepper(F, butcher_tableau, t, dt, up, solver_parameters=params, splitting=splitting) energies = [] while (float(t) < 1.0): if (float(t) + float(dt) > 1.0): dt.assign(1.0 - float(t)) stepper.advance() t.assign(float(t) + float(dt)) energies.append(assemble(E)) return np.array(energies)
def heat(n, deg, time_stages, stage_type="deriv", splitting=IA): N = 2**n msh = UnitIntervalMesh(N) params = { "snes_type": "ksponly", "ksp_type": "preonly", "mat_type": "aij", "pc_type": "lu" } V = FunctionSpace(msh, "CG", deg) x, = SpatialCoordinate(msh) t = Constant(0.0) dt = Constant(2.0 / N) uexact = exp(-t) * sin(pi * x) rhs = expand_derivatives(diff(uexact, t)) - div(grad(uexact)) butcher_tableau = GaussLegendre(time_stages) u = project(uexact, V) v = TestFunction(V) F = (inner(Dt(u), v) * dx + inner(grad(u), grad(v)) * dx - inner(rhs, v) * dx) bc = DirichletBC(V, Constant(0), "on_boundary") stepper = TimeStepper(F, butcher_tableau, t, dt, u, bcs=bc, solver_parameters=params, stage_type=stage_type, splitting=splitting) while (float(t) < 1.0): if (float(t) + float(dt) > 1.0): dt.assign(1.0 - float(t)) stepper.advance() t.assign(float(t) + float(dt)) return errornorm(uexact, u) / norm(uexact)
def to_2nd_order(mesh, circle_bdy_id=None, rad=1.0): # make new coordinates function V = VectorFunctionSpace(mesh, 'CG', 2) new_coordinates = project(mesh.coordinates, V) # If we have a circle, move any nodes on the circle bdy # onto the circle. Note circle MUST be centered at origin if circle_bdy_id is not None: nodes_on_circle = V.boundary_nodes(circle_bdy_id, 'geometric') #Force all cell nodes to have given radius :arg:`rad` for node in nodes_on_circle: scale = rad / la.norm(new_coordinates.dat.data[node]) new_coordinates.dat.data[node] *= scale # Make a new mesh with given coordinates return Mesh(new_coordinates)
def test_strain_heating(): E_initial = firedrake.interpolate(E_surface + q_bed / α * h * (1 - ζ), Q) u0 = 100.0 du = 100.0 u_expr = as_vector((u0 + du * x / Lx, 0)) u = firedrake.interpolate(u_expr, V) w = firedrake.interpolate((-du / Lx + dh / Lx / h * u[0]) * ζ, W) E_q = E_initial.copy(deepcopy=True) E_0 = E_initial.copy(deepcopy=True) from icepack.models.hybrid import ( horizontal_strain, vertical_strain, stresses ) model = icepack.models.HeatTransport3D() T = model.temperature(E_q) A = icepack.rate_factor(T) ε_x, ε_z = horizontal_strain(u, s, h), vertical_strain(u, h) τ_x, τ_z = stresses(ε_x, ε_z, A) q = firedrake.project(inner(τ_x, ε_x) + inner(τ_z, ε_z), Q) fields = { 'velocity': u, 'vertical_velocity': w, 'thickness': h, 'surface': s, 'heat_bed': Constant(q_bed), 'energy_inflow': E_initial, 'energy_surface': Constant(E_surface) } dt = 10.0 final_time = Lx / u0 num_steps = int(final_time / dt) + 1 solver_strain = icepack.solvers.HeatTransportSolver(model) solver_no_strain = icepack.solvers.HeatTransportSolver(model) for step in range(num_steps): E_q = solver_strain.solve(dt, energy=E_q, heat=q, **fields) E_0 = solver_no_strain.solve(dt, energy=E_0, heat=Constant(0), **fields) assert assemble(E_q * h * dx) > assemble(E_0 * h * dx)
def test_strain_heating(): E_initial = firedrake.interpolate(E_surface + q_bed / α * h * (1 - ζ), Q) u0 = 100.0 du = 100.0 u_expr = as_vector((u0 + du * x / Lx, 0)) u = firedrake.interpolate(u_expr, V) w = firedrake.interpolate((-du / Lx + dh / Lx / h * u[0]) * ζ, W) E_q = E_initial.copy(deepcopy=True) E_0 = E_initial.copy(deepcopy=True) model = icepack.models.HeatTransport3D() T = model.temperature(E_q) A = icepack.rate_factor(T) ε_x = horizontal_strain_rate(velocity=u, surface=s, thickness=h) ε_z = vertical_strain_rate(velocity=u, thickness=h) τ_x, τ_z = stresses(strain_rate_x=ε_x, strain_rate_z=ε_z, fluidity=A) q = firedrake.project(inner(τ_x, ε_x) + inner(τ_z, ε_z), Q) fields = { "velocity": u, "vertical_velocity": w, "thickness": h, "surface": s, "heat_bed": Constant(q_bed), "energy_inflow": E_initial, "energy_surface": Constant(E_surface), } dt = 10.0 final_time = Lx / u0 num_steps = int(final_time / dt) + 1 solver_strain = icepack.solvers.HeatTransportSolver(model) solver_no_strain = icepack.solvers.HeatTransportSolver(model) for step in range(num_steps): E_q = solver_strain.solve(dt, energy=E_q, heat=q, **fields) E_0 = solver_no_strain.solve(dt, energy=E_0, heat=Constant(0), **fields) assert assemble(E_q * h * dx) > assemble(E_0 * h * dx)
def test_inflow_boundary(scheme): start = 16 finish = 3 * start incr = 4 num_points = np.array(list(range(start, finish + incr, incr))) errors = np.zeros_like(num_points, dtype=np.float64) for k, nx in enumerate(num_points): mesh = firedrake.UnitSquareMesh(nx, nx, diagonal='crossed') degree = 1 Q = firedrake.FunctionSpace(mesh, 'DG', 1) x = firedrake.SpatialCoordinate(mesh) U = Constant(1.) u = as_vector((U, 0.)) q_in = Constant(1.) s = Constant(0.) final_time = 0.5 min_diameter = mesh.cell_sizes.dat.data_ro[:].min() max_speed = 1. multiplier = multipliers[scheme] / 2 timestep = multiplier * min_diameter / max_speed / (2 * degree + 1) num_steps = int(final_time / timestep) dt = final_time / num_steps q_0 = firedrake.project(q_in - x[0], Q) equation = plumes.models.advection.make_equation(u, s, q_in) integrator = scheme(equation, q_0) for step in range(num_steps): integrator.step(dt) q = integrator.state T = Constant(final_time) expr = q_in - firedrake.conditional(x[0] > U * T, x[0] - U * T, 0) errors[k] = assemble(abs(q - expr) * dx) / assemble(abs(expr) * dx) slope, intercept = np.polyfit(np.log2(1 / num_points), np.log2(errors), 1) print(f'log(error) ~= {slope:5.3f} * log(dx) {intercept:+5.3f}') assert slope > degree - 0.1
def betaInit(s, h, speed, V, Q, Q1, grounded, inversionParams): """Compute intitial beta using 0.5 taud. Parameters ---------- s : firedrake function model surface elevation h : firedrake function model thickness speed : firedrake function modelled speed V : firedrake vector function space vector function space Q : firedrake function space scalar function space grounded : firedrake function Mask with 1s for grounded 0 for floating. """ # Use a result from prior inversion checkFile = inversionParams['initFile'] Quse = Q if inversionParams['initWithDeg1']: checkFile = f'{inversionParams["inversionResult"]}.deg1' Quse = Q1 if checkFile is not None: betaTemp = mf.getCheckPointVars(checkFile, 'betaInv', Quse)['betaInv'] beta1 = icepack.interpolate(betaTemp, Q) return beta1 # No prior result, so use fraction of taud tauD = firedrake.project(-rhoI * g * h * grad(s), V) # stress = firedrake.sqrt(firedrake.inner(tauD, tauD)) Print('stress', firedrake.assemble(stress * firedrake.dx)) fraction = firedrake.Constant(0.95) U = max_value(speed, 1) C = fraction * stress / U**(1/m) if inversionParams['friction'] == 'schoof': mExp = 1/m + 1 U0 = firedrake.Constant(inversionParams['uThresh']) C = C * (m/(m+1)) * (U0**mExp + U**mExp)**(1/(m+1)) beta = firedrake.interpolate(firedrake.sqrt(C) * grounded, Q) return beta
def regularization_form(r): mesh = UnitSquareMesh(2 ** r, 2 ** r) x = SpatialCoordinate(mesh) S = VectorFunctionSpace(mesh, "CG", 1) beta = 4.0 reg_solver = RegularizationSolver(S, mesh, beta=beta, gamma=0.0, dx=dx) # Exact solution with free Neumann boundary conditions for this domain u_exact = Function(S) u_exact_component = cos(x[0] * pi * 2) * cos(x[1] * pi * 2) u_exact.interpolate(as_vector((u_exact_component, u_exact_component))) f = Function(S) theta = TestFunction(S) f_component = (1 + beta * 8 * pi * pi) * u_exact_component f.interpolate(as_vector((f_component, f_component))) rhs_form = inner(f, theta) * dx velocity = Function(S) rhs = assemble(rhs_form) reg_solver.solve(velocity, rhs) File("solution_vel_unitsquare.pvd").write(velocity) return norm(project(u_exact - velocity, S))
def test_2D_dirichlet_regions(): # Can't do mesh refinement because MeshHierarchy ruins the domain tags mesh = Mesh("./2D_mesh.msh") dim = mesh.geometric_dimension() x = SpatialCoordinate(mesh) S = VectorFunctionSpace(mesh, "CG", 1) beta = 1.0 reg_solver = RegularizationSolver( S, mesh, beta=beta, gamma=0.0, dx=dx, design_domain=0 ) # Exact solution with free Neumann boundary conditions for this domain u_exact = Function(S) u_exact_component = (-cos(x[0] * pi * 2) + 1) * (-cos(x[1] * pi * 2) + 1) u_exact.interpolate( as_vector(tuple(u_exact_component for _ in range(dim))) ) f = Function(S) f_component = ( -beta * ( 4.0 * pi * pi * cos(2 * pi * x[0]) * (-cos(2 * pi * x[1]) + 1) + 4.0 * pi * pi * cos(2 * pi * x[1]) * (-cos(2 * pi * x[0]) + 1) ) + u_exact_component ) f.interpolate(as_vector(tuple(f_component for _ in range(dim)))) theta = TestFunction(S) rhs_form = inner(f, theta) * dx velocity = Function(S) rhs = assemble(rhs_form) reg_solver.solve(velocity, rhs) assert norm(project(domainify(u_exact, x) - velocity, S)) < 1e-3
def depth_average(q3d, weight=firedrake.Constant(1)): r"""Return the weighted depth average of a function on an extruded mesh""" element3d = q3d.ufl_element() # Create the element `E x DG0` where `E` is the horizontal element for the # input field element_z = firedrake.FiniteElement(family='DG', cell='interval', degree=0) shape = q3d.ufl_shape if len(shape) == 0: element_xy = element3d.sub_elements()[0] element_avg = firedrake.TensorProductElement(element_xy, element_z) element2d = element_xy elif len(shape) == 1: element_xy = element3d.sub_elements()[0].sub_elements()[0] element_u = firedrake.TensorProductElement(element_xy, element_z) element_avg = firedrake.VectorElement(element_u, dim=shape[0]) element2d = firedrake.VectorElement(element_xy, dim=shape[0]) else: raise NotImplementedError('Depth average of tensor fields not yet ' 'implemented!') # Project the weighted 3D field onto vertical DG0 mesh3d = q3d.ufl_domain() Q_avg = firedrake.FunctionSpace(mesh3d, element_avg) q_avg = firedrake.project(weight * q3d, Q_avg) # Create a function space on the 2D mesh and a 2D function defined on this # space, then copy the raw vector of expansion coefficients from the 3D DG0 # field into the coefficients of the 2D field. TODO: Get some assurance # from the firedrake folks that this will always work. mesh2d = mesh3d._base_mesh Q2D = firedrake.FunctionSpace(mesh2d, element2d) q2d = firedrake.Function(Q2D) q2d.dat.data[:] = q_avg.dat.data_ro[:] return q2d
def nonlocal_integral_eq( mesh, scatterer_bdy_id, outer_bdy_id, wave_number, options_prefix=None, solver_parameters=None, fspace=None, vfspace=None, true_sol_grad_expr=None, actx=None, dgfspace=None, dgvfspace=None, meshmode_src_connection=None, qbx_kwargs=None, ): r""" see run_method for descriptions of unlisted args args: gamma and beta are used to precondition with the following equation: \Delta u - \kappa^2 \gamma u = 0 (\partial_n - i\kappa\beta) u |_\Sigma = 0 """ # make sure we get outer bdy id as tuple in case it consists of multiple ids if isinstance(outer_bdy_id, int): outer_bdy_id = [outer_bdy_id] outer_bdy_id = tuple(outer_bdy_id) # away from the excluded region, but firedrake and meshmode point # into pyt_inner_normal_sign = -1 ambient_dim = mesh.geometric_dimension() # {{{ Build src and tgt # build connection meshmode near src boundary -> src boundary inside meshmode from meshmode.discretization.poly_element import \ InterpolatoryQuadratureSimplexGroupFactory from meshmode.discretization.connection import make_face_restriction factory = InterpolatoryQuadratureSimplexGroupFactory( dgfspace.finat_element.degree) src_bdy_connection = make_face_restriction(actx, meshmode_src_connection.discr, factory, scatterer_bdy_id) # source is a qbx layer potential from pytential.qbx import QBXLayerPotentialSource disable_refinement = (fspace.mesh().geometric_dimension() == 3) qbx = QBXLayerPotentialSource(src_bdy_connection.to_discr, **qbx_kwargs, _disable_refinement=disable_refinement) # get target indices and point-set target_indices, target = get_target_points_and_indices( fspace, outer_bdy_id) # }}} # build the operations from pytential import bind, sym r""" ..math: x \in \Sigma grad_op(x) = \nabla( \int_\Gamma( u(y) \partial_n H_0^{(1)}(\kappa |x - y|) )d\gamma(y) ) """ grad_op = pyt_inner_normal_sign * sym.grad( ambient_dim, sym.D(HelmholtzKernel(ambient_dim), sym.var("u"), k=sym.var("k"), qbx_forced_limit=None)) r""" ..math: x \in \Sigma op(x) = i \kappa \cdot \int_\Gamma( u(y) \partial_n H_0^{(1)}(\kappa |x - y|) )d\gamma(y) """ op = pyt_inner_normal_sign * 1j * sym.var("k") * (sym.D( HelmholtzKernel(ambient_dim), sym.var("u"), k=sym.var("k"), qbx_forced_limit=None)) # bind the operations pyt_grad_op = bind((qbx, target), grad_op) pyt_op = bind((qbx, target), op) # }}} class MatrixFreeB(object): def __init__(self, A, pyt_grad_op, pyt_op, actx, kappa): """ :arg kappa: The wave number """ self.actx = actx self.k = kappa self.pyt_op = pyt_op self.pyt_grad_op = pyt_grad_op self.A = A self.meshmode_src_connection = meshmode_src_connection # {{{ Create some functions needed for multing self.x_fntn = Function(fspace) # CG self.potential_int = Function(fspace) self.potential_int.dat.data[:] = 0.0 self.grad_potential_int = Function(vfspace) self.grad_potential_int.dat.data[:] = 0.0 self.pyt_result = Function(fspace) self.n = FacetNormal(mesh) self.v = TestFunction(fspace) # some meshmode ones self.x_mm_fntn = self.meshmode_src_connection.discr.empty( self.actx, dtype='c') # }}} def mult(self, mat, x, y): # Copy function data into the fivredrake function self.x_fntn.dat.data[:] = x[:] # Transfer the function to meshmode self.meshmode_src_connection.from_firedrake(project( self.x_fntn, dgfspace), out=self.x_mm_fntn) # Restrict to boundary x_mm_fntn_on_bdy = src_bdy_connection(self.x_mm_fntn) # Apply the operation potential_int_mm = self.pyt_op(self.actx, u=x_mm_fntn_on_bdy, k=self.k) grad_potential_int_mm = self.pyt_grad_op(self.actx, u=x_mm_fntn_on_bdy, k=self.k) # Store in firedrake self.potential_int.dat.data[target_indices] = potential_int_mm.get( ) for dim in range(grad_potential_int_mm.shape[0]): self.grad_potential_int.dat.data[ target_indices, dim] = grad_potential_int_mm[dim].get() # Integrate the potential r""" Compute the inner products using firedrake. Note this will be subtracted later, hence appears off by a sign. .. math:: \langle n(x) \cdot \nabla( \int_\Gamma( u(y) \partial_n H_0^{(1)}(\kappa |x - y|) )d\gamma(y) ), v \rangle_\Sigma - \langle i \kappa \cdot \int_\Gamma( u(y) \partial_n H_0^{(1)}(\kappa |x - y|) )d\gamma(y), v \rangle_\Sigma """ self.pyt_result = assemble( inner(inner(self.grad_potential_int, self.n), self.v) * ds(outer_bdy_id) - inner(self.potential_int, self.v) * ds(outer_bdy_id)) # y <- Ax - evaluated potential self.A.mult(x, y) with self.pyt_result.dat.vec_ro as ep: y.axpy(-1, ep) # {{{ Compute normal helmholtz operator u = TrialFunction(fspace) v = TestFunction(fspace) r""" .. math:: \langle \nabla u, \nabla v \rangle - \kappa^2 \cdot \langle u, v \rangle - i \kappa \langle u, v \rangle_\Sigma """ a = inner(grad(u), grad(v)) * dx \ - Constant(wave_number**2) * inner(u, v) * dx \ - Constant(1j * wave_number) * inner(u, v) * ds(outer_bdy_id) # get the concrete matrix from a general bilinear form A = assemble(a).M.handle # }}} # {{{ Setup Python matrix B = PETSc.Mat().create() # build matrix context Bctx = MatrixFreeB(A, pyt_grad_op, pyt_op, actx, wave_number) # set up B as same size as A B.setSizes(*A.getSizes()) B.setType(B.Type.PYTHON) B.setPythonContext(Bctx) B.setUp() # }}} # {{{ Create rhs # Remember f is \partial_n(true_sol)|_\Gamma # so we just need to compute \int_\Gamma\partial_n(true_sol) H(x-y) sigma = sym.make_sym_vector("sigma", ambient_dim) r""" ..math: x \in \Sigma grad_op(x) = \nabla( \int_\Gamma( f(y) H_0^{(1)}(\kappa |x - y|) )d\gamma(y) ) """ grad_op = pyt_inner_normal_sign * \ sym.grad(ambient_dim, sym.S(HelmholtzKernel(ambient_dim), sym.n_dot(sigma), k=sym.var("k"), qbx_forced_limit=None)) r""" ..math: x \in \Sigma op(x) = i \kappa \cdot \int_\Gamma( f(y) H_0^{(1)}(\kappa |x - y|) )d\gamma(y) ) """ op = 1j * sym.var("k") * pyt_inner_normal_sign * \ sym.S(HelmholtzKernel(ambient_dim), sym.n_dot(sigma), k=sym.var("k"), qbx_forced_limit=None) rhs_grad_op = bind((qbx, target), grad_op) rhs_op = bind((qbx, target), op) # Transfer to meshmode metadata = {'quadrature_degree': 2 * fspace.ufl_element().degree()} dg_true_sol_grad = project(true_sol_grad_expr, dgvfspace, form_compiler_parameters=metadata) true_sol_grad_mm = meshmode_src_connection.from_firedrake(dg_true_sol_grad, actx=actx) true_sol_grad_mm = src_bdy_connection(true_sol_grad_mm) # Apply the operations f_grad_convoluted_mm = rhs_grad_op(actx, sigma=true_sol_grad_mm, k=wave_number) f_convoluted_mm = rhs_op(actx, sigma=true_sol_grad_mm, k=wave_number) # Transfer function back to firedrake f_grad_convoluted = Function(vfspace) f_convoluted = Function(fspace) f_grad_convoluted.dat.data[:] = 0.0 f_convoluted.dat.data[:] = 0.0 for dim in range(f_grad_convoluted_mm.shape[0]): f_grad_convoluted.dat.data[target_indices, dim] = f_grad_convoluted_mm[dim].get() f_convoluted.dat.data[target_indices] = f_convoluted_mm.get() r""" \langle f, v \rangle_\Gamma + \langle i \kappa \cdot \int_\Gamma( f(y) H_0^{(1)}(\kappa |x - y|) )d\gamma(y), v \rangle_\Sigma - \langle n(x) \cdot \nabla( \int_\Gamma( f(y) H_0^{(1)}(\kappa |x - y|) )d\gamma(y) ), v \rangle_\Sigma """ rhs_form = inner(inner(true_sol_grad_expr, FacetNormal(mesh)), v) * ds(scatterer_bdy_id, metadata=metadata) \ + inner(f_convoluted, v) * ds(outer_bdy_id) \ - inner(inner(f_grad_convoluted, FacetNormal(mesh)), v) * ds(outer_bdy_id) rhs = assemble(rhs_form) # {{{ set up a solver: solution = Function(fspace, name="Computed Solution") # {{{ Used for preconditioning if 'gamma' in solver_parameters or 'beta' in solver_parameters: gamma = complex(solver_parameters.pop('gamma', 1.0)) import cmath beta = complex(solver_parameters.pop('beta', cmath.sqrt(gamma))) p = inner(grad(u), grad(v)) * dx \ - Constant(wave_number**2 * gamma) * inner(u, v) * dx \ - Constant(1j * wave_number * beta) * inner(u, v) * ds(outer_bdy_id) P = assemble(p).M.handle else: P = A # }}} # Set up options to contain solver parameters: ksp = PETSc.KSP().create() if solver_parameters['pc_type'] == 'pyamg': del solver_parameters['pc_type'] # We are using the AMG preconditioner pyamg_tol = solver_parameters.get('pyamg_tol', None) if pyamg_tol is not None: pyamg_tol = float(pyamg_tol) pyamg_maxiter = solver_parameters.get('pyamg_maxiter', None) if pyamg_maxiter is not None: pyamg_maxiter = int(pyamg_maxiter) ksp.setOperators(B) ksp.setUp() pc = ksp.pc pc.setType(pc.Type.PYTHON) pc.setPythonContext( AMGTransmissionPreconditioner(wave_number, fspace, A, tol=pyamg_tol, maxiter=pyamg_maxiter, use_plane_waves=True)) # Otherwise use regular preconditioner else: ksp.setOperators(B, P) options_manager = OptionsManager(solver_parameters, options_prefix) options_manager.set_from_options(ksp) import petsc4py.PETSc petsc4py.PETSc.Sys.popErrorHandler() with rhs.dat.vec_ro as b: with solution.dat.vec as x: ksp.solve(b, x) # }}} return ksp, solution
# # Now we're ready to solve the variational problem. We define `w` to be a # function to hold the solution on the mixed space. w = fd.Function(W) # solve fd.solve(a == L, w, bcs=[bc0, bc1]) # %% # Post-process and plot interesting fields print('* post-process') # collect fields vel, press = w.split() press.rename('pressure') # Projecting velocity field to a continuous space (visualization purporse) P1 = fd.VectorFunctionSpace(mesh, "CG", 1) vel_proj = fd.project(vel, P1) vel_proj.rename('velocity_proj') # project permeability Kxx kxx = fd.Function(DG) kxx.rename('Kxx') kxx.dat.data[...] = 1 / Kinv.dat.data[:, 0, 0] # print results fd.File("plots/darcy_rt.pvd").write(vel_proj, press, kxx) print('* normal termination')
def test_linear_bed(scheme): start = 16 finish = 2 * start incr = 4 num_points = np.array(list(range(start, finish + incr, incr))) errors = np.zeros_like(num_points, dtype=np.float64) # TODO: Test with different degrees + function spaces degree = 1 for k, nx in enumerate(num_points): Lx, Ly = 20.0, 20.0 mesh = firedrake.RectangleMesh(nx, nx, Lx, Ly, diagonal='crossed') Q = firedrake.FunctionSpace(mesh, 'DG', degree) V = firedrake.VectorFunctionSpace(mesh, 'DG', degree) Z = Q * V x = firedrake.SpatialCoordinate(mesh) b_0 = Constant(0.2) δb = Constant(0.0) b = b_0 - δb * x[0] / Lx g = Constant(gravity) h_in = Constant(0.1) u_in = Constant(2.) q_in = h_in * u_in # Make the exact steady-state thickness. h = firedrake.project(h_in, Q) φ = firedrake.TestFunction(Q) F = (h + 0.5 * q_in**2 / (g * h**2) - (h_in + 0.5 * q_in**2 / (g * h_in**2)) - δb * x[0] / Lx) * φ * dx firedrake.solve(F == 0, h) # The exact steady-state momentum is a constant. q = firedrake.project(as_vector((q_in, 0.0)), V) # Get the maximum wave speed and choose a timestep that will satisfy # the CFL condition. c = firedrake.project(q[0] / h + sqrt(g * h), Q) max_speed = c.dat.data_ro[:].max() min_diameter = mesh.cell_sizes.dat.data_ro[:].min() multiplier = multipliers[scheme] timestep = multiplier * min_diameter / max_speed / (2 * degree + 1) # Run the simulation for a full residence time of the system. final_time = 2 * Lx / float(u_in) num_steps = int(final_time / timestep) dt = final_time / num_steps z_0 = firedrake.Function(Z) h_0, q_0 = z_0.split() q_0.assign(q) # Add a small perturbation to the initial state. δh = Constant(h_in / 10) r = Constant(Lx / 8) y = Constant((Lx / 2, Ly / 2)) h_0.project(h + δh * max_value(0, 1 - inner(x - y, x - y) / r**2)) bcs = { 'thickness_in': h_in, 'momentum_in': as_vector((q_in, 0)), 'inflow_ids': (1, ), 'outflow_ids': (2, ) } equation = plumes.models.shallow_water.make_equation(g, b, **bcs) integrator = scheme(equation, z_0) for step in range(num_steps): integrator.step(dt) h_n, q_n = integrator.state.split() errors[k] = assemble(abs(h_n - h) * dx) / assemble(h * dx) slope, intercept = np.polyfit(np.log2(1 / num_points), np.log2(errors), 1) print(f'log(error) ~= {slope:5.3f} * log(dx) {intercept:+5.3f}') assert slope > degree - 0.05
Vvec = fd.VectorFunctionSpace(mesh, "CG", 1) # We'll also need the test and trial functions corresponding to this # function space:: u = fd.TrialFunction(V) v = fd.TestFunction(V) # We declare a function over our function space and give it the # value of our right hand side function:: # Permeability tensor print("START: Read in reservoir fields") Delta = np.array([Delta_x, Delta_y]) coords = fd.project(mesh.coordinates, Vvec).dat.data kx_array = np.load('./spe10/spe10_kx.npy') ky_array = np.load('./spe10/spe10_ky.npy') kz_array = np.load('./spe10/spe10_kz.npy') phi_array = np.load('./spe10/spe10_po.npy') Kx = fd.Function(V) Ky = fd.Function(V) Kz = fd.Function(V) phi = fd.Constat(0.2) coords2ijk = np.vectorize(coords2ijk, excluded=['data_array', 'Delta']) Kx.dat.data[...] = coords2ijk(coords[:, 0], coords[:, 1], coords[:, 2], Delta=Delta, data_array=kx_array) Ky.dat.data[...] = coords2ijk(coords[:, 0], coords[:, 1],
def test_3D_dirichlet_regions(): # Can't do mesh refinement because MeshHierarchy ruins the domain tags mesh = Mesh("./3D_mesh.msh") dim = mesh.geometric_dimension() x = SpatialCoordinate(mesh) solver_parameters_amg = { "ksp_type": "cg", "ksp_converged_reason": None, "ksp_rtol": 1e-7, "pc_type": "hypre", "pc_hypre_type": "boomeramg", "pc_hypre_boomeramg_max_iter": 5, "pc_hypre_boomeramg_coarsen_type": "PMIS", "pc_hypre_boomeramg_agg_nl": 2, "pc_hypre_boomeramg_strong_threshold": 0.95, "pc_hypre_boomeramg_interp_type": "ext+i", "pc_hypre_boomeramg_P_max": 2, "pc_hypre_boomeramg_relax_type_all": "sequential-Gauss-Seidel", "pc_hypre_boomeramg_grid_sweeps_all": 1, "pc_hypre_boomeramg_truncfactor": 0.3, "pc_hypre_boomeramg_max_levels": 6, } S = VectorFunctionSpace(mesh, "CG", 1) beta = 1.0 reg_solver = RegularizationSolver( S, mesh, beta=beta, gamma=0.0, dx=dx, design_domain=0, solver_parameters=solver_parameters_amg, ) # Exact solution with free Neumann boundary conditions for this domain u_exact = Function(S) u_exact_component = ( (-cos(x[0] * pi * 2) + 1) * (-cos(x[1] * pi * 2) + 1) * (-cos(x[2] * pi * 2) + 1) ) u_exact.interpolate( as_vector(tuple(u_exact_component for _ in range(dim))) ) f = Function(S) f_component = ( -beta * ( 4.0 * pi * pi * cos(2 * pi * x[0]) * (-cos(2 * pi * x[1]) + 1) * (-cos(2 * pi * x[2]) + 1) + 4.0 * pi * pi * cos(2 * pi * x[1]) * (-cos(2 * pi * x[0]) + 1) * (-cos(2 * pi * x[2]) + 1) + 4.0 * pi * pi * cos(2 * pi * x[2]) * (-cos(2 * pi * x[1]) + 1) * (-cos(2 * pi * x[0]) + 1) ) + u_exact_component ) f.interpolate(as_vector(tuple(f_component for _ in range(dim)))) theta = TestFunction(S) rhs_form = inner(f, theta) * dx velocity = Function(S) rhs = assemble(rhs_form) reg_solver.solve(velocity, rhs) error = norm( project( domainify(u_exact, x) - velocity, S, solver_parameters=solver_parameters_amg, ) ) assert error < 5e-2
- fd.conditional(fd.dot(u0, n) < 0, r * fd.dot(u0, n)*Fw(s_in), 0.0)*fd.ds - (r('+') - r('-'))*(un('+')*Fw(s_mid)('+') - un('-')*Fw(s_mid)('-'))*fd.dS) # Res. R = F_p + F_s prob = fd.NonlinearVariationalProblem(R, U, bcs=[bc0, bc1, bc3]) solver2ph = fd.NonlinearVariationalSolver(prob) # %% # 4) Solve problem # initialize variables xv, xp, xs = U.split() vel = fd.project(u0, P1) xp.rename('pressure') xs.rename('saturation') vel.rename('velocity') outfile = fd.File(oFile) it = 0 t = 0.0 while t < sim_time: t += dt print("*iteration= {:3d}, dtime= {:8.6f}, time={:8.6f}".format(it, dt, t)) solver2ph.solve() # update solution U0.assign(U)
# function space:: u = fd.TrialFunction(V) v = fd.TestFunction(V) # We declare a function over our function space and give it the # value of our right hand side function:: # Permeability tensor print("START: Read in reservoir fields") Delta = np.array([Delta_x, Delta_y, Delta_z]) ct = 79.08e-11 # (1.0+0.2*3+0.8*4.947)*14.2 * 10**-6 kgf/cm2 mu = 0.003 # Pa-s # coords = fd.project(mesh.coordinates, C).dat.data coords = fd.project(mesh.coordinates, Vvec).dat.data Kx = fd.interpolate(fd.Constant(98.7e-15), V) Ky = fd.interpolate(fd.Constant(98.7e-15), V) Kz = fd.interpolate(fd.Constant(98.7e-16), V) phi = fd.interpolate(fd.Constant(0.25), V) # well location xw = np.array([118.872, 181.356, Lz / 2]) Tx = fd.Function(V) Ty = fd.Function(V) Tz = fd.Function(V) w = fd.Function(V) I = fd.Function(V)
def getFormStage(F, butch, u0, t, dt, bcs=None, splitting=None, nullspace=None): """Given a time-dependent variational form and a :class:`ButcherTableau`, produce UFL for the s-stage RK method. :arg F: UFL form for the semidiscrete ODE/DAE :arg butch: the :class:`ButcherTableau` for the RK method being used to advance in time. :arg u0: a :class:`Function` referring to the state of the PDE system at time `t` :arg t: a :class:`Constant` referring to the current time level. Any explicit time-dependence in F is included :arg dt: a :class:`Constant` referring to the size of the current time step. :arg splitting: a callable that maps the (floating point) Butcher matrix a to a pair of matrices `A1, A2` such that `butch.A = A1 A2`. This is used to vary between the classical RK formulation and Butcher's reformulation that leads to a denser mass matrix with block-diagonal stiffness. Only `AI` and `IA` are currently supported. :arg bcs: optionally, a :class:`DirichletBC` object (or iterable thereof) containing (possible time-dependent) boundary conditions imposed on the system. :arg nullspace: A list of tuples of the form (index, VSB) where index is an index into the function space associated with `u` and VSB is a :class: `firedrake.VectorSpaceBasis` instance to be passed to a `firedrake.MixedVectorSpaceBasis` over the larger space associated with the Runge-Kutta method On output, we return a tuple consisting of several parts: - Fnew, the :class:`Form` - possibly a 4-tuple containing information needed to solve a mass matrix to update the solution (this is empty for RadauIIA methods for which there is a trivial update function. - UU, the :class:`firedrake.Function` holding all the stage time values. It lives in a :class:`firedrake.FunctionSpace` corresponding to the s-way tensor product of the space on which the semidiscrete form lives. - `bcnew`, a list of :class:`firedrake.DirichletBC` objects to be posed on the stages, - 'nspnew', the :class:`firedrake.MixedVectorSpaceBasis` object that represents the nullspace of the coupled system - `gblah`, a list of tuples of the form (f, expr, method), where f is a :class:`firedrake.Function` and expr is a :class:`ufl.Expr`. At each time step, each expr needs to be re-interpolated/projected onto the corresponding f in order for Firedrake to pick up that time-dependent boundary conditions need to be re-applied. The interpolation/projection is encoded in method, which is either `f.interpolate(expr-c*u0)` or `f.project(expr-c*u0)`, depending on whether the function space for f supports interpolation or not. """ v = F.arguments()[0] V = v.function_space() assert V == u0.function_space() num_stages = butch.num_stages num_fields = len(V) # s-way product space for the stage variables Vbig = reduce(mul, (V for _ in range(num_stages))) VV = TestFunction(Vbig) UU = Function(Vbig) vecconst = np.vectorize(Constant) C = vecconst(butch.c) A = vecconst(butch.A) # set up the pieces we need to work with to do our substitutions nsxnf = (num_stages, num_fields) if num_fields == 1: u0bits = np.array([u0], dtype="O") vbits = np.array([v], dtype="O") if num_stages == 1: # single-stage method VVbits = np.array([[VV]], dtype="O") UUbits = np.array([[UU]], dtype="O") else: # multi-stage methods VVbits = np.reshape(split(VV), nsxnf) UUbits = np.reshape(split(UU), nsxnf) else: u0bits = np.array(list(split(u0)), dtype="O") vbits = np.array(list(split(v)), dtype="O") VVbits = np.reshape(split(VV), nsxnf) UUbits = np.reshape(split(UU), nsxnf) split_form = extract_terms(F) Fnew = Zero() # first, process terms with a time derivative. I'm # assuming we have something of the form inner(Dt(g(u0)), v)*dx # For each stage i, this gets replaced with # inner((g(stages[i]) - g(u0))/dt, v)*dx # but we have to carefully handle the cases where g indexes into # pieces of u dtless = strip_dt_form(split_form.time) if splitting is None or splitting == AI: for i in range(num_stages): repl = {t: t + C[i] * dt} for j in range(num_fields): repl[u0bits[j]] = UUbits[i][j] - u0bits[j] repl[vbits[j]] = VVbits[i][j] # Also get replacements right for indexing. for j in range(num_fields): for ii in np.ndindex(u0bits[j].ufl_shape): repl[u0bits[j][ii]] = UUbits[i][j][ii] - u0bits[j][ii] repl[vbits[j][ii]] = VVbits[i][j][ii] Fnew += replace(dtless, repl) # Now for the non-time derivative parts for i in range(num_stages): # replace test function repl = {} for k in range(num_fields): repl[vbits[k]] = VVbits[i][k] for ii in np.ndindex(vbits[k].ufl_shape): repl[vbits[k][ii]] = VVbits[i][k][ii] Ftmp = replace(split_form.remainder, repl) # replace the solution with stage values for j in range(num_stages): repl = {t: t + C[j] * dt} for k in range(num_fields): repl[u0bits[k]] = UUbits[j][k] for ii in np.ndindex(u0bits[k].ufl_shape): repl[u0bits[k][ii]] = UUbits[j][k][ii] # and sum the contribution Fnew += A[i, j] * dt * replace(Ftmp, repl) elif splitting == IA: Ainv = np.vectorize(Constant)(np.linalg.inv(butch.A)) # time derivative part gets inverse of Butcher matrix. for i in range(num_stages): repl = {} for k in range(num_fields): repl[vbits[k]] = VVbits[i][k] for ii in np.ndindex(vbits[k].ufl_shape): repl[vbits[k][ii]] = VVbits[i][k][ii] Ftmp = replace(dtless, repl) for j in range(num_stages): repl = {t: t + C[j] * dt} for k in range(num_fields): repl[u0bits[k]] = (UUbits[j][k] - u0bits[k]) for ii in np.ndindex(u0bits[k].ufl_shape): repl[u0bits[k][ii]] = UUbits[j][k][ii] - u0bits[k][ii] Fnew += Ainv[i, j] * replace(Ftmp, repl) # rest of the operator: just diagonal! for i in range(num_stages): repl = {t: t + C[i] * dt} for j in range(num_fields): repl[u0bits[j]] = UUbits[i][j] repl[vbits[j]] = VVbits[i][j] # Also get replacements right for indexing. for j in range(num_fields): for ii in np.ndindex(u0bits[j].ufl_shape): repl[u0bits[j][ii]] = UUbits[i][j][ii] repl[vbits[j][ii]] = VVbits[i][j][ii] Fnew += dt * replace(split_form.remainder, repl) else: raise NotImplementedError("Can't handle that splitting type") if bcs is None: bcs = [] bcsnew = [] gblah = [] # For each BC, we need a new BC for each stage # so we need to figure out how the function is indexed (mixed + vec) # and then set it to have the value of the original argument at # time t+C[i]*dt. for bc in bcs: if num_fields == 1: # not mixed space comp = bc.function_space().component if comp is not None: # check for sub-piece of vector-valued Vsp = V.sub(comp) Vbigi = lambda i: Vbig[i].sub(comp) else: Vsp = V Vbigi = lambda i: Vbig[i] else: # mixed space sub = bc.function_space_index() comp = bc.function_space().component if comp is not None: # check for sub-piece of vector-valued Vsp = V.sub(sub).sub(comp) Vbigi = lambda i: Vbig[sub + num_fields * i].sub(comp) else: Vsp = V.sub(sub) Vbigi = lambda i: Vbig[sub + num_fields * i] bcarg = bc._original_arg for i in range(num_stages): try: gdat = interpolate(bcarg, Vsp) gmethod = lambda gd, gc: gd.interpolate(gc) except: # noqa: E722 gdat = project(bcarg, Vsp) gmethod = lambda gd, gc: gd.project(gc) gcur = replace(bcarg, {t: t + C[i] * dt}) bcsnew.append(DirichletBC(Vbigi(i), gdat, bc.sub_domain)) gblah.append((gdat, gcur, gmethod)) nspacenew = getNullspace(V, Vbig, butch, nullspace) # For RIIA, we have an optimized update rule and don't need to # build the variational form for doing updates. # But something's broken with null spaces, so that's a TO-DO. unew = Function(V) Fupdate = inner(unew - u0, v) * dx B = vectorize(Constant)(butch.b) C = vectorize(Constant)(butch.c) for i in range(num_stages): repl = {t: t + C[i] * dt} for k in range(num_fields): repl[u0bits[k]] = UUbits[i][k] for ii in np.ndindex(u0bits[k].ufl_shape): repl[u0bits[k][ii]] = UUbits[i][k][ii] eFFi = replace(split_form.remainder, repl) Fupdate += dt * B[i] * eFFi # And the BC's for the update -- just the original BC at t+dt update_bcs = [] update_bcs_gblah = [] for bc in bcs: if num_fields == 1: # not mixed space comp = bc.function_space().component if comp is not None: # check for sub-piece of vector-valued Vsp = V.sub(comp) else: Vsp = V else: # mixed space sub = bc.function_space_index() comp = bc.function_space().component if comp is not None: # check for sub-piece of vector-valued Vsp = V.sub(sub).sub(comp) else: Vsp = V.sub(sub) bcarg = bc._original_arg try: gdat = interpolate(bcarg, Vsp) gmethod = lambda gd, gc: gd.interpolate(gc) except: # noqa: E722 gdat = project(bcarg, Vsp) gmethod = lambda gd, gc: gd.project(gc) gcur = replace(bcarg, {t: t + dt}) update_bcs.append(DirichletBC(Vsp, gdat, bc.sub_domain)) update_bcs_gblah.append((gdat, gcur, gmethod)) return (Fnew, (unew, Fupdate, update_bcs, update_bcs_gblah), UU, bcsnew, gblah, nspacenew)