class SawyerEliassenU(DiagnosticField): name = "SawyerEliassenU" def setup(self, state): space = state.spaces("HDiv") super(SawyerEliassenU, self).setup(state, space=space) u = state.fields("u") b = state.fields("b") v = inner(u, as_vector([0., 1., 0.])) # spaces V0 = FunctionSpace(state.mesh, "CG", 2) Vu = u.function_space() # project b to V0 self.b_v0 = Function(V0) btri = TrialFunction(V0) btes = TestFunction(V0) a = inner(btes, btri) * dx L = inner(btes, b) * dx projectbproblem = LinearVariationalProblem(a, L, self.b_v0) self.project_b_solver = LinearVariationalSolver( projectbproblem, solver_parameters={'ksp_type': 'cg'}) # project v to V0 self.v_v0 = Function(V0) vtri = TrialFunction(V0) vtes = TestFunction(V0) a = inner(vtes, vtri) * dx L = inner(vtes, v) * dx projectvproblem = LinearVariationalProblem(a, L, self.v_v0) self.project_v_solver = LinearVariationalSolver( projectvproblem, solver_parameters={'ksp_type': 'cg'}) # stm/psi is a stream function self.stm = Function(V0) psi = TrialFunction(V0) xsi = TestFunction(V0) f = state.parameters.f H = state.parameters.H L = state.parameters.L dbdy = state.parameters.dbdy x, y, z = SpatialCoordinate(state.mesh) bcs = [DirichletBC(V0, 0., "bottom"), DirichletBC(V0, 0., "top")] Mat = as_matrix([[b.dx(2), 0., -f * self.v_v0.dx(2)], [0., 0., 0.], [-self.b_v0.dx(0), 0., f**2 + f * self.v_v0.dx(0)]]) Equ = (inner(grad(xsi), dot(Mat, grad(psi))) - dbdy * inner(grad(xsi), as_vector([-v, 0., f * (z - H / 2)]))) * dx # fourth-order terms if state.parameters.fourthorder: eps = Constant(0.0001) brennersigma = Constant(10.0) n = FacetNormal(state.mesh) deltax = Constant(state.parameters.deltax) deltaz = Constant(state.parameters.deltaz) nn = as_matrix([[sqrt(brennersigma / Constant(deltax)), 0., 0.], [0., 0., 0.], [0., 0., sqrt(brennersigma / Constant(deltaz))]]) mu = as_matrix([[1., 0., 0.], [0., 0., 0.], [0., 0., H / L]]) # anisotropic form Equ += eps * ( div(dot(mu, grad(psi))) * div(dot(mu, grad(xsi))) * dx - (avg(dot(dot(grad(grad(psi)), n), n)) * jump(grad(xsi), n=n) + avg(dot(dot(grad(grad(xsi)), n), n)) * jump(grad(psi), n=n) - jump(nn * grad(psi), n=n) * jump(nn * grad(xsi), n=n)) * (dS_h + dS_v)) Au = lhs(Equ) Lu = rhs(Equ) stmproblem = LinearVariationalProblem(Au, Lu, self.stm, bcs=bcs) self.stream_function_solver = LinearVariationalSolver( stmproblem, solver_parameters={'ksp_type': 'cg'}) # solve for sawyer_eliassen u self.u = Function(Vu) utrial = TrialFunction(Vu) w = TestFunction(Vu) a = inner(w, utrial) * dx L = (w[0] * (-self.stm.dx(2)) + w[2] * (self.stm.dx(0))) * dx ugproblem = LinearVariationalProblem(a, L, self.u) self.sawyer_eliassen_u_solver = LinearVariationalSolver( ugproblem, solver_parameters={'ksp_type': 'cg'}) def compute(self, state): self.project_b_solver.solve() self.project_v_solver.solve() self.stream_function_solver.solve() self.sawyer_eliassen_u_solver.solve() sawyer_eliassen_u = self.u return self.field.project(sawyer_eliassen_u)
def __init__(self, diagnostic_variables, prognostic_variables, outputting, simulation_parameters): self.diagnostic_variables = diagnostic_variables self.prognostic_variables = prognostic_variables self.outputting = outputting self.simulation_parameters = simulation_parameters Dt = Constant(simulation_parameters['dt'][-1]) Ld = simulation_parameters['Ld'][-1] u = self.prognostic_variables.u Xi = self.prognostic_variables.dXi Vu = u.function_space() vector_u = True if Vu.ufl_element() == VectorElement else False ones = Function( VectorFunctionSpace(self.prognostic_variables.mesh, "CG", 1)).project(as_vector([Constant(1.0)])) self.to_update_constants = False self.interpolators = [] self.projectors = [] self.solvers = [] mesh = u.function_space().mesh() x, = SpatialCoordinate(mesh) alphasq = simulation_parameters['alphasq'][-1] periodic = simulation_parameters['periodic'][-1] # do peakon data checks here true_peakon_data = simulation_parameters['true_peakon_data'][-1] if true_peakon_data is not None: self.true_peakon_file = Dataset( 'results/' + true_peakon_data + '/data.nc', 'r') # check length of file is correct ndump = simulation_parameters['ndump'][-1] tmax = simulation_parameters['tmax'][-1] dt = simulation_parameters['dt'][-1] if len(self.true_peakon_file['time'][:]) != int(tmax / (ndump * dt)) + 1: raise ValueError( 'If reading in true peakon data, the dump frequency must be the same as that used for the true peakon data.' + ' Length of true peakon data as %i, but proposed length is %i' % (len(self.true_peakon_file['time'][:]), int(tmax / (ndump * dt)) + 1)) if self.true_peakon_file['p'][:].shape != (int(tmax / (ndump * dt)) + 1, ): raise ValueError( 'True peakon data shape %i must be the same shape as proposed data %i' % ((int(tmax / (ndump * dt)) + 1, ), self.true_peakon_file['p'][:].shape)) # do peakon data checks here true_mean_peakon_data = simulation_parameters['true_mean_peakon_data'][ -1] if true_mean_peakon_data is not None: self.true_mean_peakon_file = Dataset( 'results/' + true_mean_peakon_data + '/data.nc', 'r') # check length of file is correct ndump = simulation_parameters['ndump'][-1] tmax = simulation_parameters['tmax'][-1] dt = simulation_parameters['dt'][-1] if len(self.true_mean_peakon_file['time'][:]) != int(tmax / (ndump * dt)): raise ValueError( 'If reading in true peakon data, the dump frequency must be the same as that used for the true peakon data.' ) if self.true_mean_peakon_file['p'][:].shape != (int( tmax / (ndump * dt)), ): raise ValueError( 'True peakon data must have same shape as proposed data!') for key, value in self.diagnostic_variables.fields.items(): if key == 'uscalar': uscalar = self.diagnostic_variables.fields['uscalar'] u_interpolator = Interpolator(dot(ones, u), uscalar) self.interpolators.append(u_interpolator) elif key == 'Euscalar': Eu = self.prognostic_variables.Eu Euscalar = self.diagnostic_variables.fields['Euscalar'] Eu_interpolator = Interpolator(dot(ones, Eu), Euscalar) self.interpolators.append(Eu_interpolator) elif key == 'Xiscalar': Xi = self.prognostic_variables.dXi Xiscalar = self.diagnostic_variables.fields['Xiscalar'] Xi_interpolator = Interpolator(dot(ones, Xi), Xiscalar) self.interpolators.append(Xi_interpolator) elif key == 'du': if type(u.function_space().ufl_element()) == VectorElement: u_to_project = self.diagnostic_variables.fields['uscalar'] else: u_to_project = u du = self.diagnostic_variables.fields['du'] du_projector = Projector(u_to_project.dx(0), du) self.projectors.append(du_projector) elif key == 'jump_du': du = self.diagnostic_variables.fields['du'] jump_du = self.diagnostic_variables.fields['jump_du'] V = jump_du.function_space() jtrial = TrialFunction(V) psi = TestFunction(V) Lj = psi('+') * abs(jump(du)) * dS aj = psi('+') * jtrial('+') * dS jprob = LinearVariationalProblem(aj, Lj, jump_du) jsolver = LinearVariationalSolver(jprob) self.solvers.append(jsolver) elif key == 'du_smooth': du = self.diagnostic_variables.fields['du'] du_smooth = self.diagnostic_variables.fields['du_smooth'] projector = Projector(du, du_smooth) self.projectors.append(projector) elif key == 'u2_flux': gamma = simulation_parameters['gamma'][-1] u2_flux = self.diagnostic_variables.fields['u2_flux'] xis = self.prognostic_variables.pure_xi_list xis_x = [] xis_xxx = [] CG1 = FunctionSpace(mesh, "CG", 1) psi = TestFunction(CG1) for xi in xis: xis_x.append(Function(CG1).project(xi.dx(0))) for xi_x in xis_x: xi_xxx = Function(CG1) form = (psi * xi_xxx + psi.dx(0) * xi_x.dx(0)) * dx prob = NonlinearVariationalProblem(form, xi_xxx) solver = NonlinearVariationalSolver(prob) solver.solve() xis_xxx.append(xi_xxx) flux_expr = 0.0 * x for xi, xi_x, xi_xxx in zip(xis, xis_x, xis_xxx): flux_expr += (6 * u.dx(0) * xi + 12 * u * xi_x + gamma * xi_xxx) * (6 * u.dx(0) * xi + 24 * u * xi_x + gamma * xi_xxx) projector = Projector(flux_expr, u2_flux) self.projectors.append(projector) elif key == 'a': # find 6 * u_x * Xi + gamma * Xi_xxx mesh = u.function_space().mesh() gamma = simulation_parameters['gamma'][-1] a_flux = self.diagnostic_variables.fields['a'] xis = self.prognostic_variables.pure_xis xis_x = [] xis_xxx = [] CG1 = FunctionSpace(mesh, "CG", 1) psi = TestFunction(CG1) for xi in xis: xis_x.append(Function(CG1).project(xi.dx(0))) for xi_x in xis_x: xi_xxx = Function(CG1) form = (psi * xi_xxx + psi.dx(0) * xi_x.dx(0)) * dx prob = NonlinearVariationalProblem(form, xi_xxx) solver = NonlinearVariationalSolver(prob) solver.solve() xis_xxx.append(xi_xxx) x, = SpatialCoordinate(mesh) a_expr = 0.0 * x for xi, xi_x, xi_xxx in zip(xis, xis_x, xis_xxx): a_expr += 6 * u.dx(0) * xi + gamma * xi_xxx projector = Projector(a_expr, a_flux) self.projectors.append(projector) elif key == 'b': # find 12 * u * Xi_x mesh = u.function_space().mesh() gamma = simulation_parameters['gamma'][-1] b_flux = self.diagnostic_variables.fields['b'] xis = self.prognostic_variables.pure_xis x, = SpatialCoordinate(mesh) b_expr = 0.0 * x for xi, xi_x, xi_xxx in zip(xis, xis_x, xis_xxx): b_expr += 12 * u * xi.dx(0) projector = Projector(b_expr, b_flux) self.projectors.append(projector) elif key == 'kdv_1': # find the first part of the kdv form u0 = prognostic_variables.u0 uh = (u + u0) / 2 us = Dt * uh + sqrt(Dt) * Xi psi = TestFunction(Vu) du_1 = self.diagnostic_variables.fields['kdv_1'] eqn = psi * du_1 * dx - 6 * psi.dx(0) * uh * us * dx prob = NonlinearVariationalProblem(eqn, du_1) solver = NonlinearVariationalSolver(prob) self.solvers.append(solver) elif key == 'kdv_2': # find the second part of the kdv form u0 = prognostic_variables.u0 uh = (u + u0) / 2 us = Dt * uh + sqrt(Dt) * Xi psi = TestFunction(Vu) du_2 = self.diagnostic_variables.fields['kdv_2'] eqn = psi * du_2 * dx + 6 * psi * uh * us.dx(0) * dx prob = NonlinearVariationalProblem(eqn, du_2) solver = NonlinearVariationalSolver(prob) self.solvers.append(solver) elif key == 'kdv_3': # find the third part of the kdv form u0 = prognostic_variables.u0 uh = (u + u0) / 2 us = Dt * uh + sqrt(Dt) * Xi du_3 = self.diagnostic_variables.fields['kdv_3'] gamma = simulation_parameters['gamma'][-1] phi = TestFunction(Vu) F = Function(Vu) eqn = (phi * F * dx + phi.dx(0) * us.dx(0) * dx) prob = NonlinearVariationalProblem(eqn, F) solver = NonlinearVariationalSolver(prob) self.solvers.append(solver) self.projectors.append(Projector(-gamma * F.dx(0), du_3)) # nu = TestFunction(Vu) # back_eqn = nu * du_3 * dx - gamma * nu.dx(0) * F * dx # back_prob = NonlinearVariationalProblem(back_eqn, du_3) # back_solver = NonlinearVariationalSolver(back_prob) # self.solvers.append(solver) elif key == 'm': m = self.diagnostic_variables.fields['m'] phi = TestFunction(Vu) eqn = phi * m * dx - phi * u * dx - alphasq * phi.dx(0) * u.dx( 0) * dx prob = NonlinearVariationalProblem(eqn, m) solver = NonlinearVariationalSolver(prob) self.solvers.append(solver) elif key == 'u_xx': u_xx = self.diagnostic_variables.fields['u_xx'] phi = TestFunction(Vu) eqn = phi * u_xx * dx + phi.dx(0) * u_xx.dx(0) * dx prob = NonlinearVariationalProblem(eqn, u_xx) solver = NonlinearVariationalSolver(prob) self.solvers.append(solver) elif key == 'u_sde': self.to_update_constants = True self.Ld = Ld self.alphasq = alphasq self.p = Constant(1.0 * 0.5 * (1 + exp(-Ld / sqrt(alphasq))) / (1 - exp(-Ld / sqrt(alphasq)))) self.q = Constant(Ld / 2) u_sde = self.diagnostic_variables.fields['u_sde'] if periodic: expr = conditional( x < self.q - Ld / 2, self.p * ((exp(-(x - self.q + Ld) / sqrt(alphasq)) + exp(-Ld / sqrt(alphasq)) * exp( (x - self.q + Ld) / sqrt(alphasq))) / (1 - exp(-Ld / sqrt(alphasq)))), conditional( x < self.q + Ld / 2, self.p * ((exp(-sqrt((self.q - x)**2 / alphasq)) + exp(-Ld / sqrt(alphasq)) * exp(sqrt((self.q - x)**2 / alphasq))) / (1 - exp(-Ld / sqrt(alphasq)))), self.p * ((exp(-(self.q + Ld - x) / sqrt(alphasq)) + exp(-Ld / sqrt(alphasq) * exp( (self.q + Ld - x) / sqrt(alphasq)))) / (1 - exp(-Ld / sqrt(alphasq)))))) else: expr = conditional( x < self.q - Ld / 2, self.p * exp(-(x - self.q + Ld) / sqrt(alphasq)), conditional( x < self.q + Ld / 2, self.p * exp(-sqrt((self.q - x)**2 / alphasq)), self.p * exp(-(self.q + Ld - x) / sqrt(alphasq)))) self.interpolators.append(Interpolator(expr, u_sde)) elif key == 'u_sde_weak': u_sde = self.diagnostic_variables.fields['u_sde'] u_sde_weak = self.diagnostic_variables.fields['u_sde_weak'] psi = TestFunction(Vu) eqn = psi * u_sde_weak * dx - psi * (u - u_sde) * dx prob = NonlinearVariationalProblem(eqn, u_sde_weak) solver = NonlinearVariationalSolver(prob) self.solvers.append(solver) elif key == 'u_sde_mean': self.to_update_constants = True self.p = Constant(1.0) self.q = Constant(Ld / 2) if periodic: raise NotImplementedError( 'u_sde_mean not yet implemented for periodic peakon') u_sde = self.diagnostic_variables.fields['u_sde_mean'] expr = conditional( x < self.q - Ld / 2, self.p * exp(-(x - self.q + Ld) / sqrt(alphasq)), conditional( x < self.q + Ld / 2, self.p * exp(-sqrt((self.q - x)**2 / alphasq)), self.p * exp(-(self.q + Ld - x) / sqrt(alphasq)))) self.interpolators.append(Interpolator(expr, u_sde)) elif key == 'u_sde_weak_mean': u_sde = self.diagnostic_variables.fields['u_sde_mean'] u_sde_weak = self.diagnostic_variables.fields[ 'u_sde_weak_mean'] psi = TestFunction(Vu) eqn = psi * u_sde_weak * dx - psi * (u - u_sde) * dx prob = NonlinearVariationalProblem(eqn, u_sde_weak) solver = NonlinearVariationalSolver(prob) self.solvers.append(solver) elif key == 'pure_xi': pure_xi = 0.0 * x for xi in self.prognostic_variables.pure_xi_list: if vector_u: pure_xi += dot(ones, xi) else: pure_xi += xi Xiscalar = self.diagnostic_variables.fields['pure_xi'] Xi_interpolator = Interpolator(pure_xi, Xiscalar) self.interpolators.append(Xi_interpolator) elif key == 'pure_xi_x': pure_xi_x = 0.0 * x for xix in self.prognostic_variables.pure_xi_x_list: if vector_u: pure_xi_x += dot(ones, xix) else: pure_xi_x += xix Xiscalar = self.diagnostic_variables.fields['pure_xi_x'] Xi_interpolator = Interpolator(pure_xi_x, Xiscalar) self.interpolators.append(Xi_interpolator) elif key == 'pure_xi_xx': pure_xi_xx = 0.0 * x for xixx in self.prognostic_variables.pure_xi_xx_list: if vector_u: pure_xi_xx += dot(ones, xixx) else: pure_xi_xx += xixx Xiscalar = self.diagnostic_variables.fields['pure_xi_xx'] Xi_interpolator = Interpolator(pure_xi_xx, Xiscalar) self.interpolators.append(Xi_interpolator) elif key == 'pure_xi_xxx': pure_xi_xxx = 0.0 * x for xixxx in self.prognostic_variables.pure_xi_xxx_list: if vector_u: pure_xi_xxx += dot(ones, xixxx) else: pure_xi_xxx += xixxx Xiscalar = self.diagnostic_variables.fields['pure_xi_xxx'] Xi_interpolator = Interpolator(pure_xi_xxx, Xiscalar) self.interpolators.append(Xi_interpolator) elif key == 'pure_xi_xxxx': pure_xi_xxxx = 0.0 * x for xixxxx in self.prognostic_variables.pure_xi_xx_list: if vector_u: pure_xi_xxxx += dot(ones, xixxxx) else: pure_xi_xxxx += xixxxx Xiscalar = self.diagnostic_variables.fields['pure_xi_xxxx'] Xi_interpolator = Interpolator(pure_xi_xxxx, Xiscalar) self.interpolators.append(Xi_interpolator) else: raise NotImplementedError('Diagnostic %s not yet implemented' % key)