def __init__(self, mesh, conditions, timestepping, params, output, solver_params): super().__init__(mesh, conditions, timestepping, params, output, solver_params) self.u0 = Function(self.V) self.u1 = Function(self.V) self.sigma0 = Function(self.S) self.sigma1 = Function(self.S) theta = conditions.theta uh = (1-theta) * self.u0 + theta * self.u1 a = Function(self.U) h = Function(self.U) p = TestFunction(self.V) q = TestFunction(self.S) self.initial_condition((self.u0, conditions.ic['u']), (a, conditions.ic['a']), (h, conditions.ic['h'])) ep_dot = self.strain(grad(uh)) zeta = self.zeta(h, a, self.delta(uh)) eta = zeta * params.e ** (-2) rheology = 2 * eta * ep_dot + (zeta - eta) * tr(ep_dot) * Identity(2) - 0.5 * self.Ice_Strength(h, a) * Identity(2) self.initial_condition((self.sigma0, rheology),(self.sigma1, self.sigma0)) def sigma_next(timestep, zeta, ep_dot, sigma, P): A = 1 + 0.25 * (timestep * params.e ** 2) / params.T B = timestep * 0.125 * (1 - params.e ** 2) / params.T rhs = (1 - (timestep * params.e ** 2) / (4 * params.T)) * sigma - timestep / params.T * ( 0.125 * (1 - params.e ** 2) * tr(sigma) * Identity(2) - 0.25 * P * Identity(2) + zeta * ep_dot) C = (rhs[0, 0] - rhs[1, 1]) / A D = (rhs[0, 0] + rhs[1, 1]) / (A + 2 * B) sigma00 = 0.5 * (C + D) sigma11 = 0.5 * (D - C) sigma01 = rhs[0, 1] sigma = as_matrix([[sigma00, sigma01], [sigma01, sigma11]]) return sigma s = sigma_next(self.timestep, zeta, ep_dot, self.sigma0, self.Ice_Strength(h, a)) sh = (1-theta) * s + theta * self.sigma0 eqn = self.momentum_equation(h, self.u1, self.u0, p, sh, params.rho, uh, conditions.ocean_curr, params.rho_a, params.C_a, params.rho_w, params.C_w, conditions.geo_wind, params.cor, self.timestep, ind=self.ind) tensor_eqn = inner(self.sigma1-s, q) * dx if conditions.stabilised['state']: alpha = conditions.stabilised['alpha'] eqn += stabilisation_term(alpha=alpha, zeta=avg(zeta), mesh=mesh, v=uh, test=p) bcs = DirichletBC(self.V, conditions.bc['u'], "on_boundary") uprob = NonlinearVariationalProblem(eqn, self.u1, bcs) self.usolver = NonlinearVariationalSolver(uprob, solver_parameters=solver_params.bt_params) sprob = NonlinearVariationalProblem(tensor_eqn, self.sigma1) self.ssolver = NonlinearVariationalSolver(sprob, solver_parameters=solver_params.bt_params)
def __init__(self, F, butcher_tableau, t, dt, u0, bcs=None, solver_parameters=None, update_solver_parameters=None, splitting=AI, nullspace=None, appctx=None): self.u0 = u0 self.t = t self.dt = dt self.num_fields = len(u0.function_space()) self.num_stages = len(butcher_tableau.b) self.butcher_tableau = butcher_tableau Fbig, update_stuff, UU, bigBCs, gblah, nsp = getFormStage( F, butcher_tableau, u0, t, dt, bcs, splitting=splitting) self.UU = UU self.bigBCs = bigBCs self.bcdat = gblah self.update_stuff = update_stuff self.prob = NonlinearVariationalProblem(Fbig, UU, bigBCs) appctx_irksome = { "F": F, "butcher_tableau": butcher_tableau, "t": t, "dt": dt, "u0": u0, "bcs": bcs, "nullspace": nullspace } if appctx is None: appctx = appctx_irksome else: appctx = {**appctx, **appctx_irksome} self.solver = NonlinearVariationalSolver( self.prob, appctx=appctx, nullspace=nsp, solver_parameters=solver_parameters) unew, Fupdate, update_bcs, update_bcs_gblah = self.update_stuff self.update_problem = NonlinearVariationalProblem( Fupdate, unew, update_bcs) self.update_solver = NonlinearVariationalSolver( self.update_problem, solver_parameters=update_solver_parameters) self._update = self._update_general
def setup_solver(self, up_init=None): """ Setup the solvers """ self.up0 = Function(self.W) if up_init is not None: chk_in = checkpointing.HDF5File(up_init, file_mode='r') chk_in.read(self.up0, "/up") chk_in.close() self.u0, self.p0 = split(self.up0) self.up = Function(self.W) if up_init is not None: chk_in = checkpointing.HDF5File(up_init, file_mode='r') chk_in.read(self.up, "/up") chk_in.close() self.u1, self.p1 = split(self.up) self.up.sub(0).rename("velocity") self.up.sub(1).rename("pressure") v, q = TestFunctions(self.W) h = CellVolume(self.mesh) u_norm = sqrt(dot(self.u0, self.u0)) if self.has_nullspace: nullspace = MixedVectorSpaceBasis( self.W, [self.W.sub(0), VectorSpaceBasis(constant=True)]) else: nullspace = None tau = ((2.0 / self.dt)**2 + (2.0 * u_norm / h)**2 + (4.0 * self.nu / h**2)**2)**(-0.5) # temporal discretization F = (1.0 / self.dt) * inner(self.u1 - self.u0, v) * dx # weak form F += (+inner(dot(self.u0, nabla_grad(self.u1)), v) * dx + self.nu * inner(grad(self.u1), grad(v)) * dx - (1.0 / self.rho) * self.p1 * div(v) * dx + div(self.u1) * q * dx - inner(self.forcing, v) * dx) # residual form R = (+(1.0 / self.dt) * (self.u1 - self.u0) + dot(self.u0, nabla_grad(self.u1)) - self.nu * div(grad(self.u1)) + (1.0 / self.rho) * grad(self.p1) - self.forcing) # GLS F += tau * inner( +dot(self.u0, nabla_grad(v)) - self.nu * div(grad(v)) + (1.0 / self.rho) * grad(q), R) * dx self.problem = NonlinearVariationalProblem(F, self.up, self.bcs) self.solver = NonlinearVariationalSolver( self.problem, nullspace=nullspace, solver_parameters=self.solver_parameters)
def __init__(self, mesh, conditions, timestepping, params, output, solver_params): super().__init__(mesh, conditions, timestepping, params, output, solver_params) self.w0 = Function(self.W3) self.w1 = Function(self.W3) u0, s0, h0, a0 = self.w0.split() p, q, r, m = TestFunctions(self.W3) self.initial_condition((u0, conditions.ic['u']), (s0, conditions.ic['s']), (a0, conditions.ic['a']), (h0, conditions.ic['h'])) self.w1.assign(self.w0) u1, s1, h1, a1 = split(self.w1) u0, s0, h0, a0 = split(self.w0) theta = conditions.theta uh = (1-theta) * u0 + theta * u1 sh = (1-theta) * s0 + theta * s1 hh = (1-theta) * h0 + theta * h1 ah = (1-theta) * a0 + theta * a1 ep_dot = self.strain(grad(uh)) zeta = self.zeta(hh, ah, self.delta(uh)) rheology = params.e ** 2 * sh + Identity(2) * 0.5 * ((1 - params.e ** 2) * tr(sh) + self.Ice_Strength(hh, ah)) eqn = self.momentum_equation(hh, u1, u0, p, sh, params.rho, uh, conditions.ocean_curr, params.rho_a, params.C_a, params.rho_w, params.C_w, conditions.geo_wind, params.cor, self.timestep, ind=self.ind) eqn += self.transport_equation(uh, hh, ah, h1, h0, a1, a0, r, m, self.n, self.timestep) eqn += inner(self.ind * (s1 - s0) + 0.5 * self.timestep * rheology / params.T, q) * dx eqn -= inner(q * zeta * self.timestep / params.T, ep_dot) * dx if conditions.stabilised['state']: alpha = conditions.stabilised['alpha'] eqn += self.stabilisation_term(alpha=alpha, zeta=avg(zeta), mesh=mesh, v=uh, test=p) bcs = DirichletBC(self.W3.sub(0), conditions.bc['u'], "on_boundary") uprob = NonlinearVariationalProblem(eqn, self.w1, bcs) self.usolver = NonlinearVariationalSolver(uprob, solver_parameters=solver_params.bt_params) self.u1, self.s0, self.h1, self.a1 = self.w1.split()
def solve(self, file_path='ttip_result/solution.pvd'): """ Setup and solve the nonlinear problem. Save value to file given. Any additional keyword arguments are passed to the iteration method. Args: file_path (string, optional): The path to save the pvd file to. vtk files will be generated in the same directory as the pvd. It is recommended that this is a separate drectory per run. Defaults to 'TTiP_result/solution.pvd'. """ F = self.problem.a - self.problem.L steady_state = self.is_steady_state() if isinstance(self.problem, BoundaryMixin): var_prob = NonlinearVariationalProblem(F, self.u, bcs=self.problem.bcs) else: var_prob = NonlinearVariationalProblem(F, self.u) solver = NonlinearVariationalSolver(problem=var_prob, solver_parameters=self.params) outfile = File(file_path) outfile.write(self.u, target_degree=1, target_continuity=H1) if steady_state: solver.solve() outfile.write(self.u, target_degree=1, target_continuity=H1) else: self.problem.T_.assign(self.u) last_perc = 0 for i in range(self.problem.steps): solver.solve() perc = int(100 * (i + 1) / self.problem.steps) if perc > last_perc: print(f'{perc}%') last_perc = perc self.problem.T_.assign(self.u) outfile.write(self.u, target_degree=1, target_continuity=H1)
def __init__(self, mesh, conditions, timestepping, params, output, solver_params): super().__init__(mesh, conditions, timestepping, params, output, solver_params) self.w0 = Function(self.W2) self.w1 = Function(self.W2) u0, h0, a0 = self.w0.split() p, q, r = TestFunctions(self.W2) self.initial_condition((u0, conditions.ic['u']), (h0, conditions.ic['h']), (a0, conditions.ic['a'])) self.w1.assign(self.w0) u1, h1, a1 = split(self.w1) u0, h0, a0 = split(self.w0) theta = conditions.theta uh = (1-theta) * u0 + theta * u1 ah = (1-theta) * a0 + theta * a1 hh = (1-theta) * h0 + theta * h1 ep_dot = self.strain(grad(uh)) zeta = self.zeta(hh, ah, self.delta(uh)) eta = zeta * params.e ** (-2) sigma = 2 * eta * ep_dot + (zeta - eta) * tr(ep_dot) * Identity(2) - self.Ice_Strength(hh, ah) * 0.5 * Identity( 2) eqn = self.momentum_equation(hh, u1, u0, p, sigma, params.rho, uh, conditions.ocean_curr, params.rho_a, params.C_a, params.rho_w, params.C_w, conditions.geo_wind, params.cor, self.timestep) eqn += self.transport_equation(uh, hh, ah, h1, h0, a1, a0, q, r, self.n, self.timestep) if conditions.stabilised['state']: alpha = conditions.stabilised['alpha'] eqn += self.stabilisation_term(alpha=alpha, zeta=avg(zeta), mesh=mesh, v=uh, test=p) bcs = DirichletBC(self.W2.sub(0), conditions.bc['u'], "on_boundary") uprob = NonlinearVariationalProblem(eqn, self.w1, bcs) self.usolver = NonlinearVariationalSolver(uprob, solver_parameters=solver_params.bt_params) self.u1, self.h1, self.a1 = self.w1.split()
def wrapper(self, **kwargs): """To disable the annotation, just pass :py:data:`annotate=False` to this routine, and it acts exactly like the Firedrake solve call. This is useful in cases where the solve is known to be irrelevant or diagnostic for the purposes of the adjoint computation (such as projecting fields to other function spaces for the purposes of visualisation).""" annotate = annotate_tape(kwargs) if annotate: tape = get_working_tape() problem = self._ad_problem sb_kwargs = NonlinearVariationalSolveBlock.pop_kwargs(kwargs) sb_kwargs.update(kwargs) block = NonlinearVariationalSolveBlock( problem._ad_F == 0, problem._ad_u, problem._ad_bcs, problem_J=problem._ad_J, solver_params=self.parameters, solver_kwargs=self._ad_kwargs, **sb_kwargs) if not self._ad_nlvs: from firedrake import NonlinearVariationalSolver self._ad_nlvs = NonlinearVariationalSolver( self._ad_problem_clone(self._ad_problem, block.get_dependencies()), **self._ad_kwargs) block._ad_nlvs = self._ad_nlvs tape.add_block(block) with stop_annotating(): out = solve(self, **kwargs) if annotate: block.add_output( self._ad_problem._ad_u.create_block_variable()) return out
def __init__(self, prognostic_variables, simulation_parameters): mesh = simulation_parameters['mesh'][-1] x, = SpatialCoordinate(mesh) Ld = simulation_parameters['Ld'][-1] self.scheme = simulation_parameters['scheme'][-1] self.dt = simulation_parameters['dt'][-1] self.num_Xis = simulation_parameters['num_Xis'][-1] self.Xi_family = simulation_parameters['Xi_family'][-1] self.dXi = prognostic_variables.dXi self.dWs = [Constant(0.0) for dw in range(self.num_Xis)] self.dW_nums = prognostic_variables.dW_nums self.Xi_functions = [] self.nXi_updates = simulation_parameters['nXi_updates'][-1] self.smooth_t = simulation_parameters['smooth_t'][-1] self.fixed_dW = simulation_parameters['fixed_dW'][-1] if self.smooth_t is not None and self.nXi_updates > 1: raise ValueError('Prescribing forcing and including multiple Xi updates are not compatible.') if self.smooth_t is not None or self.fixed_dW is not None: print('WARNING: Remember to change sigma to sigma * sqrt(dt) with the prescribed forcing option or the fixed_dW option.') seed = simulation_parameters['seed'][-1] np.random.seed(seed) # make sure sigma is a Constant if self.num_Xis != 0: if isinstance(simulation_parameters['sigma'][-1], Constant): self.sigma = simulation_parameters['sigma'][-1] else: self.sigma = Constant(simulation_parameters['sigma'][-1]) else: self.sigma = Constant(0.0) self.pure_xi_list = prognostic_variables.pure_xi_list self.pure_xi_x_list = prognostic_variables.pure_xi_x_list self.pure_xi_xx_list = prognostic_variables.pure_xi_xx_list self.pure_xi_xxx_list = prognostic_variables.pure_xi_xxx_list self.pure_xi_xxxx_list = prognostic_variables.pure_xi_xxxx_list for xi in range(self.num_Xis): self.pure_xi_list.append(Function(self.dXi.function_space())) self.pure_xi_x_list.append(Function(self.dXi.function_space())) self.pure_xi_xx_list.append(Function(self.dXi.function_space())) self.pure_xi_xxx_list.append(Function(self.dXi.function_space())) self.pure_xi_xxxx_list.append(Function(self.dXi.function_space())) if self.Xi_family == 'sines': for n in range(self.num_Xis): if (n+1) % 2 == 1: self.Xi_functions.append(self.sigma * sin(2*(n+1)*pi*x/Ld)) else: self.Xi_functions.append(self.sigma * cos(2*(n+1)*pi*x/Ld)) elif self.Xi_family == 'double_sines': for n in range(self.num_Xis): if (n+1) % 2 == 1: self.Xi_functions.append(self.sigma * sin(4*(n+1)*pi*x/Ld)) else: self.Xi_functions.append(self.sigma * cos(4*(n+1)*pi*x/Ld)) elif self.Xi_family == 'high_freq_sines': for n in range(self.num_Xis): if (n+1) % 2 == 1: self.Xi_functions.append(self.sigma * sin((2*(n+1)+10)*pi*x/Ld)) else: self.Xi_functions.append(self.sigma * cos((2*(n+1)+10)*pi*x/Ld)) elif self.Xi_family == 'gaussians': for n in range(self.num_Xis): self.Xi_functions.append(self.sigma * 0.5*self.num_Xis*exp(-((x-Ld*(n+1)/(self.num_Xis +1.0))/2.)**2)) elif self.Xi_family == 'quadratic': if self.num_Xis > 1: raise NotImplementedError('Quadratic Xi not yet implemented for more than one Xi') else: self.Xi_functions.append(32/(Ld*Ld)*conditional(x > Ld/4, conditional(x > 3*Ld/8, conditional(x > 5*Ld/8, conditional(x < 3*Ld/4, self.sigma * (x - 3*Ld/4)**2, 0.0), (x-Ld/2)**2+Ld**2/32), (x-Ld/4)**2), 0.0)) elif self.Xi_family == 'proper_peak': if self.num_Xis > 1: raise NotImplementedError('Quadratic Xi not yet implemented for more than one Xi') else: self.Xi_functions.append(self.sigma * 0.5*2/(exp(x-Ld/2)+exp(-x+Ld/2))) elif self.Xi_family == 'constant': if self.num_Xis > 1: raise NotImplementedError('Constant Xi not yet implemented for more than one Xi') else: self.Xi_functions.append(self.sigma * (sin(0*pi*x/Ld)+1)) else: raise NotImplementedError('Xi_family %s not implemented' % self.Xi_family) # make lists of functions for xi_x, xi_xx and xi_xxx if self.scheme in ['hydrodynamic', 'LASCH_hydrodynamic']: self.dXi_x = prognostic_variables.dXi_x self.dXi_xx = prognostic_variables.dXi_xx self.Xi_x_functions = [] self.Xi_xx_functions = [] for Xi_expr in self.Xi_functions: Xi_x_function = Function(self.dXi_x.function_space()) Xi_xx_function = Function(self.dXi_xx.function_space()) phi_x = TestFunction(self.dXi_x.function_space()) phi_xx = TestFunction(self.dXi_xx.function_space()) Xi_x_eqn = phi_x * Xi_x_function * dx + phi_x.dx(0) * Xi_expr * dx Xi_xx_eqn = phi_xx * Xi_xx_function * dx + phi_xx.dx(0) * Xi_x_function * dx Xi_x_problem = NonlinearVariationalProblem(Xi_x_eqn, Xi_x_function) Xi_xx_problem = NonlinearVariationalProblem(Xi_xx_eqn, Xi_xx_function) Xi_x_solver = NonlinearVariationalSolver(Xi_x_problem) Xi_xx_solver = NonlinearVariationalSolver(Xi_xx_problem) # for some reason these solvers don't work for constant Xi functions # so just manually make the derivatives be zero if self.Xi_family == 'constant': Xi_x_function.interpolate(0.0*x) Xi_xx_function.interpolate(0.0*x) else: Xi_x_solver.solve() Xi_xx_solver.solve() self.Xi_x_functions.append(Xi_x_function) self.Xi_xx_functions.append(Xi_xx_function) # now make a master xi Xi_expr = 0.0*x for dW, Xi_function, pure_xi, pure_xi_x, pure_xi_xx, pure_xi_xxx, pure_xi_xxxx in zip(self.dWs, self.Xi_functions, self.pure_xi_list, self.pure_xi_x_list, self.pure_xi_xx_list, self.pure_xi_xxx_list, self.pure_xi_xxxx_list): Xi_expr += dW * Xi_function if self.scheme in ['upwind', 'LASCH']: pure_xi.interpolate(as_vector([Xi_function])) pure_xi_x.project(as_vector([Xi_function.dx(0)])) CG1 = FunctionSpace(mesh, "CG", 1) psi = TestFunction(CG1) xixx_scalar = Function(CG1) xixx_eqn = psi * xixx_scalar * dx + psi.dx(0) * Xi_function.dx(0) * dx prob = NonlinearVariationalProblem(xixx_eqn, xixx_scalar) solver = NonlinearVariationalSolver(prob) solver.solve() pure_xi_xx.interpolate(as_vector([xixx_scalar])) else: pure_xi.interpolate(Xi_function) # I guess we can't take the gradient of constants if self.Xi_family != 'constant': pure_xi_x.project(Xi_function.dx(0)) pure_xi_xx.project(pure_xi_x.dx(0)) pure_xi_xxx.project(pure_xi_xx.dx(0)) pure_xi_xxxx.project(pure_xi_xxx.dx(0)) if self.scheme in ['upwind', 'LASCH']: self.dXi_interpolator = Interpolator(as_vector([Xi_expr]), self.dXi) else: self.dXi_interpolator = Interpolator(Xi_expr, self.dXi) if self.scheme in ['hydrodynamic', 'LASCH_hydrodynamic']: # initialise blank expressions Xi_x_expr = 0.0*x Xi_xx_expr = 0.0*x # make full expressions by adding all dW * Xi_xs for dW, Xi_x_function, Xi_xx_function in zip(self.dWs, self.Xi_x_functions, self.Xi_xx_functions): Xi_x_expr += dW * Xi_x_function Xi_xx_expr += dW * Xi_xx_function self.dXi_x_interpolator = Interpolator(Xi_x_expr, self.dXi_x) self.dXi_xx_interpolator = Interpolator(Xi_xx_expr, self.dXi_xx)
def solver(self): # setup solver using lhs and rhs defined in derived class problem = NonlinearVariationalProblem(self.lhs-self.rhs, self.dq, bcs=self.bcs) solver_name = self.field_name+self.__class__.__name__ return NonlinearVariationalSolver(problem, solver_parameters=self.solver_parameters, options_prefix=solver_name)
def assemble(self, eqn, func, bcs, params): uprob = NonlinearVariationalProblem(eqn, func, bcs) self.usolver = NonlinearVariationalSolver(uprob, solver_parameters=params)
rho_averaged = Function(Vt) rho_recoverer = Recoverer(rho0, rho_averaged, VDG=FunctionSpace(mesh, BrokenElement(Vt.ufl_element())), boundary_method=physics_boundary_method) rho_recoverer.project() exner = thermodynamics.exner_pressure(state.parameters, rho_averaged, theta0) p = thermodynamics.p(state.parameters, exner) T = thermodynamics.T(state.parameters, theta0, exner, r_v=w_v) w_sat = thermodynamics.r_sat(state.parameters, T, p) w_functional = (phi * w_v * dxp - phi * w_sat * dxp) w_problem = NonlinearVariationalProblem(w_functional, w_v) w_solver = NonlinearVariationalSolver(w_problem) w_solver.solve() water_v0.assign(w_v) water_c0.assign(water_t - water_v0) state.set_reference_profiles([('rho', rho_b), ('theta', theta_b), ('vapour_mixing_ratio', water_vb)]) rho_opts = None theta_opts = EmbeddedDGOptions() u_transport = ImplicitMidpoint(state, "u") transported_fields = [ SSPRK3(state, "rho", options=rho_opts), SSPRK3(state, "theta", options=theta_opts),
def compressible_hydrostatic_balance(state, theta0, rho0, pi0=None, top=False, pi_boundary=Constant(1.0), water_t=None, solve_for_rho=False, params=None): """ Compute a hydrostatically balanced density given a potential temperature profile. :arg state: The :class:`State` object. :arg theta0: :class:`.Function`containing the potential temperature. :arg rho0: :class:`.Function` to write the initial density into. :arg top: If True, set a boundary condition at the top. Otherwise, set it at the bottom. :arg pi_boundary: a field or expression to use as boundary data for pi on the top or bottom as specified. :arg water_t: the initial total water mixing ratio field. """ # Calculate hydrostatic Pi VDG = state.spaces("DG") Vv = state.spaces("Vv") W = MixedFunctionSpace((Vv, VDG)) v, pi = TrialFunctions(W) dv, dpi = TestFunctions(W) n = FacetNormal(state.mesh) cp = state.parameters.cp # add effect of density of water upon theta theta = theta0 if water_t is not None: theta = theta0 / (1 + water_t) alhs = ( (cp*inner(v, dv) - cp*div(dv*theta)*pi)*dx + dpi*div(theta*v)*dx ) if top: bmeasure = ds_t bstring = "bottom" else: bmeasure = ds_b bstring = "top" arhs = -cp*inner(dv, n)*theta*pi_boundary*bmeasure g = state.parameters.g arhs -= g*inner(dv, state.k)*dx bcs = [DirichletBC(W.sub(0), 0.0, bstring)] w = Function(W) PiProblem = LinearVariationalProblem(alhs, arhs, w, bcs=bcs) if params is None: params = {'pc_type': 'fieldsplit', 'pc_fieldsplit_type': 'schur', 'ksp_type': 'gmres', 'ksp_monitor_true_residual': True, 'ksp_max_it': 100, 'ksp_gmres_restart': 50, 'pc_fieldsplit_schur_fact_type': 'FULL', 'pc_fieldsplit_schur_precondition': 'selfp', 'fieldsplit_0_ksp_type': 'richardson', 'fieldsplit_0_ksp_max_it': 5, 'fieldsplit_0_pc_type': 'gamg', 'fieldsplit_1_pc_gamg_sym_graph': True, 'fieldsplit_1_mg_levels_ksp_type': 'chebyshev', 'fieldsplit_1_mg_levels_ksp_chebyshev_esteig': True, 'fieldsplit_1_mg_levels_ksp_max_it': 5, 'fieldsplit_1_mg_levels_pc_type': 'bjacobi', 'fieldsplit_1_mg_levels_sub_pc_type': 'ilu'} PiSolver = LinearVariationalSolver(PiProblem, solver_parameters=params) PiSolver.solve() v, Pi = w.split() if pi0 is not None: pi0.assign(Pi) if solve_for_rho: w1 = Function(W) v, rho = w1.split() rho.interpolate(thermodynamics.rho(state.parameters, theta0, Pi)) v, rho = split(w1) dv, dpi = TestFunctions(W) pi = thermodynamics.pi(state.parameters, rho, theta0) F = ( (cp*inner(v, dv) - cp*div(dv*theta)*pi)*dx + dpi*div(theta0*v)*dx + cp*inner(dv, n)*theta*pi_boundary*bmeasure ) F += g*inner(dv, state.k)*dx rhoproblem = NonlinearVariationalProblem(F, w1, bcs=bcs) rhosolver = NonlinearVariationalSolver(rhoproblem, solver_parameters=params) rhosolver.solve() v, rho_ = w1.split() rho0.assign(rho_) else: rho0.interpolate(thermodynamics.rho(state.parameters, theta0, Pi))
def moist_hydrostatic_balance(state, theta_e, water_t, pi_boundary=Constant(1.0)): """ Given a wet equivalent potential temperature, theta_e, and the total moisture content, water_t, compute a hydrostatically balance virtual potential temperature, dry density and water vapour profile. :arg state: The :class:`State` object. :arg theta_e: The initial wet equivalent potential temperature profile. :arg water_t: The total water pseudo-mixing ratio profile. :arg pi_boundary: the value of pi on the lower boundary of the domain. """ theta0 = state.fields('theta') rho0 = state.fields('rho') water_v0 = state.fields('water_v') # Calculate hydrostatic Pi Vt = theta0.function_space() Vr = rho0.function_space() Vv = state.fields('u').function_space() n = FacetNormal(state.mesh) g = state.parameters.g cp = state.parameters.cp R_d = state.parameters.R_d p_0 = state.parameters.p_0 VDG = state.spaces("DG") if any(deg > 2 for deg in VDG.ufl_element().degree()): state.logger.warning("default quadrature degree most likely not sufficient for this degree element") quadrature_degree = (5, 5) params = {'ksp_type': 'preonly', 'ksp_monitor_true_residual': True, 'ksp_converged_reason': True, 'snes_converged_reason': True, 'ksp_max_it': 100, 'mat_type': 'aij', 'pc_type': 'lu', 'pc_factor_mat_solver_type': 'mumps'} theta0.interpolate(theta_e) water_v0.interpolate(water_t) Pi = Function(Vr) epsilon = 0.9 # relaxation constant # set up mixed space Z = MixedFunctionSpace((Vt, Vt)) z = Function(Z) gamma, phi = TestFunctions(Z) theta_v, w_v = z.split() # give first guesses for trial functions theta_v.assign(theta0) w_v.assign(water_v0) theta_v, w_v = split(z) # define variables T = thermodynamics.T(state.parameters, theta_v, Pi, r_v=w_v) p = thermodynamics.p(state.parameters, Pi) w_sat = thermodynamics.r_sat(state.parameters, T, p) dxp = dx(degree=(quadrature_degree)) # set up weak form of theta_e and w_sat equations F = (-gamma * theta_e * dxp + gamma * thermodynamics.theta_e(state.parameters, T, p, w_v, water_t) * dxp - phi * w_v * dxp + phi * w_sat * dxp) problem = NonlinearVariationalProblem(F, z) solver = NonlinearVariationalSolver(problem, solver_parameters=params) theta_v, w_v = z.split() Pi_h = Function(Vr).interpolate((p / p_0) ** (R_d / cp)) # solve for Pi with theta_v and w_v constant # then solve for theta_v and w_v with Pi constant for i in range(5): compressible_hydrostatic_balance(state, theta0, rho0, pi0=Pi_h, water_t=water_t) Pi.assign(Pi * (1 - epsilon) + epsilon * Pi_h) solver.solve() theta0.assign(theta0 * (1 - epsilon) + epsilon * theta_v) water_v0.assign(water_v0 * (1 - epsilon) + epsilon * w_v) # now begin on Newton solver, setup up new mixed space Z = MixedFunctionSpace((Vt, Vt, Vr, Vv)) z = Function(Z) gamma, phi, psi, w = TestFunctions(Z) theta_v, w_v, pi, v = z.split() # use previous values as first guesses for newton solver theta_v.assign(theta0) w_v.assign(water_v0) pi.assign(Pi) theta_v, w_v, pi, v = split(z) # define variables T = thermodynamics.T(state.parameters, theta_v, pi, r_v=w_v) p = thermodynamics.p(state.parameters, pi) w_sat = thermodynamics.r_sat(state.parameters, T, p) F = (-gamma * theta_e * dxp + gamma * thermodynamics.theta_e(state.parameters, T, p, w_v, water_t) * dxp - phi * w_v * dxp + phi * w_sat * dxp + cp * inner(v, w) * dxp - cp * div(w * theta_v / (1.0 + water_t)) * pi * dxp + psi * div(theta_v * v / (1.0 + water_t)) * dxp + cp * inner(w, n) * pi_boundary * theta_v / (1.0 + water_t) * ds_b + g * inner(w, state.k) * dxp) bcs = [DirichletBC(Z.sub(3), 0.0, "top")] problem = NonlinearVariationalProblem(F, z, bcs=bcs) solver = NonlinearVariationalSolver(problem, solver_parameters=params) solver.solve() theta_v, w_v, pi, v = z.split() # assign final values theta0.assign(theta_v) water_v0.assign(w_v) # find rho compressible_hydrostatic_balance(state, theta0, rho0, water_t=water_t, solve_for_rho=True)
def __init__(self, prognostic_variables, simulation_parameters): mesh = simulation_parameters['mesh'][-1] self.scheme = simulation_parameters['scheme'][-1] self.timestepping = simulation_parameters['timestepping'][-1] alphasq = simulation_parameters['alphasq'][-1] c0 = simulation_parameters['c0'][-1] gamma = simulation_parameters['gamma'][-1] Dt = Constant(simulation_parameters['dt'][-1]) self.solvers = [] if alphasq.values()[0] > 0.0 and gamma.values()[0] == 0.0: self.setup = 'ch' if self.scheme == 'upwind' and self.timestepping == 'ssprk3': Vm = prognostic_variables.Vm Vu = prognostic_variables.Vu self.m = prognostic_variables.m self.u = prognostic_variables.u self.Xi = prognostic_variables.dXi self.m0 = Function(Vm).assign(self.m) # now make problem for the actual problem psi = TestFunction(Vm) self.m_trial = Function(Vm) self.dm = Function( Vm ) # introduce this as the advection operator for a single step us = Dt * self.u + self.Xi nhat = FacetNormal(mesh) un = 0.5 * (dot(us, nhat) + abs(dot(us, nhat))) ones = Function(Vu).project(as_vector([Constant(1.)])) Lm = (psi * self.dm * dx - psi.dx(0) * self.m_trial * dot(ones, us) * dx + psi * self.m_trial * dot(ones, us.dx(0)) * dx + jump(psi) * (un('+') * self.m_trial('+') - un('-') * self.m_trial('-')) * dS) mprob = NonlinearVariationalProblem(Lm, self.dm) self.msolver = NonlinearVariationalSolver(mprob, solver_parameters={ 'ksp_type': 'preonly', 'pc_type': 'bjacobi', 'sub_pc_type': 'ilu' }) phi = TestFunction(Vu) Lu = (dot(phi, ones) * self.m * dx - dot(phi, self.u) * dx - alphasq * dot(self.u.dx(0), phi.dx(0)) * dx) uprob = NonlinearVariationalProblem(Lu, self.u) self.usolver = NonlinearVariationalSolver(uprob, solver_parameters={ 'ksp_type': 'preonly', 'pc_type': 'lu' }) elif self.scheme == 'hydrodynamic' and self.timestepping == 'midpoint': Vu = prognostic_variables.Vu self.u = prognostic_variables.u W = MixedFunctionSpace((Vu, ) * 3) psi, phi, zeta = TestFunctions(W) w1 = Function(W) self.u1, dFh, dGh = split(w1) uh = (self.u1 + self.u) / 2 dXi = prognostic_variables.dXi dXi_x = prognostic_variables.dXi_x dXi_xx = prognostic_variables.dXi_xx dvh = Dt * uh + dXi Lu = (psi * (self.u1 - self.u) * dx + psi * uh.dx(0) * dvh * dx - psi.dx(0) * dFh * dx + psi * dGh * dx + phi * dFh * dx + alphasq * phi.dx(0) * dFh.dx(0) * dx - phi * uh * uh * Dt * dx - 0.5 * alphasq * phi * uh.dx(0) * uh.dx(0) * Dt * dx + zeta * dGh * dx + alphasq * zeta.dx(0) * dGh.dx(0) * dx - 2 * zeta * uh * dXi_x * dx - alphasq * zeta * uh.dx(0) * dXi_xx * dx) self.u1, dFh, dGh = w1.split() uprob = NonlinearVariationalProblem(Lu, w1) self.usolver = NonlinearVariationalSolver(uprob, solver_parameters={ 'mat_type': 'aij', 'ksp_type': 'preonly', 'pc_type': 'lu' }) elif self.scheme == 'no_gradient' and self.timestepping == 'midpoint': # a version of the hydrodynamic form but without exploiting the gradient Vu = prognostic_variables.Vu self.u = prognostic_variables.u W = MixedFunctionSpace((Vu, ) * 3) psi, phi, zeta = TestFunctions(W) w1 = Function(W) self.u1, dFh, dGh = split(w1) uh = (self.u1 + self.u) / 2 dXi = prognostic_variables.dXi dXi_x = prognostic_variables.dXi_x dXi_xx = prognostic_variables.dXi_xx dvh = Dt * uh + dXi Lu = (psi * (self.u1 - self.u) * dx + psi * uh.dx(0) * dvh * dx + psi * dFh.dx(0) * dx + psi * dGh * dx + phi * dFh * dx + alphasq * phi.dx(0) * dFh.dx(0) * dx - phi * uh * uh * Dt * dx - 0.5 * alphasq * phi * uh.dx(0) * uh.dx(0) * Dt * dx + zeta * dGh * dx + alphasq * zeta.dx(0) * dGh.dx(0) * dx - 2 * zeta * uh * dXi_x * dx - alphasq * zeta * uh.dx(0) * dXi_xx * dx) self.u1, dFh, dGh = w1.split() uprob = NonlinearVariationalProblem(Lu, w1) self.usolver = NonlinearVariationalSolver(uprob, solver_parameters={ 'mat_type': 'aij', 'ksp_type': 'preonly', 'pc_type': 'lu' }) elif self.scheme == 'test' and self.timestepping == 'midpoint': self.u = prognostic_variables.u Vu = prognostic_variables.Vu psi = TestFunction(Vu) self.u1 = Function(Vu) uh = (self.u1 + self.u) / 2 dvh = Dt * uh + prognostic_variables.dXi eqn = (psi * (self.u1 - self.u) * dx - psi * uh * dvh.dx(0) * dx) prob = NonlinearVariationalProblem(eqn, self.u1) self.usolver = NonlinearVariationalSolver(prob, solver_parameters={ 'mat_type': 'aij', 'ksp_type': 'preonly', 'pc_type': 'lu' }) else: raise ValueError( 'Scheme %s and timestepping %s either not compatible or not recognised.' % (self.scheme, self.timestepping)) elif alphasq.values()[0] == 0.0 and gamma.values()[0] > 0.0: self.setup = 'kdv' if self.scheme == 'upwind' and self.timestepping == 'ssprk3': raise NotImplementedError( 'Scheme %s and timestepping %s not yet implemented.' % (self.scheme, self.timestepping)) elif self.scheme == 'upwind' and self.timestepping == 'midpoint': raise NotImplementedError( 'Scheme %s and timestepping %s not yet implemented.' % (self.scheme, self.timestepping)) elif self.scheme == 'hydrodynamic' and self.timestepping == 'midpoint': raise NotImplementedError( 'Scheme %s and timestepping %s not yet implemented.' % (self.scheme, self.timestepping)) else: raise ValueError( 'Scheme %s and timestepping %s either not compatible or not recognised.' % (self.scheme, self.timestepping)) else: raise NotImplementedError( 'Schemes for your values of alpha squared %.3f and gamma %.3f are not yet implemented.' % (alphasq, gamma))
def compressible_hydrostatic_balance(state, theta0, rho0, exner0=None, top=False, exner_boundary=Constant(1.0), mr_t=None, solve_for_rho=False, params=None): """ Compute a hydrostatically balanced density given a potential temperature profile. By default, this uses a vertically-oriented hybridization procedure for solving the resulting discrete systems. :arg state: The :class:`State` object. :arg theta0: :class:`.Function`containing the potential temperature. :arg rho0: :class:`.Function` to write the initial density into. :arg top: If True, set a boundary condition at the top. Otherwise, set it at the bottom. :arg exner_boundary: a field or expression to use as boundary data for exner on the top or bottom as specified. :arg mr_t: the initial total water mixing ratio field. """ # Calculate hydrostatic Pi VDG = state.spaces("DG") Vu = state.spaces("HDiv") Vv = FunctionSpace(state.mesh, Vu.ufl_element()._elements[-1]) W = MixedFunctionSpace((Vv, VDG)) v, exner = TrialFunctions(W) dv, dexner = TestFunctions(W) n = FacetNormal(state.mesh) cp = state.parameters.cp # add effect of density of water upon theta theta = theta0 if mr_t is not None: theta = theta0 / (1 + mr_t) alhs = ((cp * inner(v, dv) - cp * div(dv * theta) * exner) * dx + dexner * div(theta * v) * dx) if top: bmeasure = ds_t bstring = "bottom" else: bmeasure = ds_b bstring = "top" arhs = -cp * inner(dv, n) * theta * exner_boundary * bmeasure # Possibly make g vary with spatial coordinates? g = state.parameters.g arhs -= g * inner(dv, state.k) * dx bcs = [DirichletBC(W.sub(0), zero(), bstring)] w = Function(W) exner_problem = LinearVariationalProblem(alhs, arhs, w, bcs=bcs) if params is None: params = { 'ksp_type': 'preonly', 'pc_type': 'python', 'mat_type': 'matfree', 'pc_python_type': 'gusto.VerticalHybridizationPC', # Vertical trace system is only coupled vertically in columns # block ILU is a direct solver! 'vert_hybridization': { 'ksp_type': 'preonly', 'pc_type': 'bjacobi', 'sub_pc_type': 'ilu' } } exner_solver = LinearVariationalSolver(exner_problem, solver_parameters=params, options_prefix="exner_solver") exner_solver.solve() v, exner = w.split() if exner0 is not None: exner0.assign(exner) if solve_for_rho: w1 = Function(W) v, rho = w1.split() rho.interpolate(thermodynamics.rho(state.parameters, theta0, exner)) v, rho = split(w1) dv, dexner = TestFunctions(W) exner = thermodynamics.exner_pressure(state.parameters, rho, theta0) F = ((cp * inner(v, dv) - cp * div(dv * theta) * exner) * dx + dexner * div(theta0 * v) * dx + cp * inner(dv, n) * theta * exner_boundary * bmeasure) F += g * inner(dv, state.k) * dx rhoproblem = NonlinearVariationalProblem(F, w1, bcs=bcs) rhosolver = NonlinearVariationalSolver(rhoproblem, solver_parameters=params, options_prefix="rhosolver") rhosolver.solve() v, rho_ = w1.split() rho0.assign(rho_) else: rho0.interpolate(thermodynamics.rho(state.parameters, theta0, exner))
def __init__(self, mesh: object, kappa: float, comm_space: MPI.Comm, mu: float = 5., *args, **kwargs): """ Constructor :param mesh: spatial domain :param kappa: diffusion coefficient :param mu: penalty weighting function """ super(Diffusion2D, self).__init__(*args, **kwargs) # Spatial domain and function space self.mesh = mesh V = FunctionSpace(self.mesh, "DG", 1) self.function_space = V self.comm_space = comm_space # Placeholder for time step - will be updated in the update method self.dt = Constant(0.) # Things we need for the form gamma = TestFunction(V) phi = TrialFunction(V) self.f = Function(V) n = FacetNormal(mesh) # Set up the rhs and bilinear form of the equation a = (inner(gamma, phi) * dx + self.dt * (inner(grad(gamma), grad(phi) * kappa) * dx - inner(2 * avg(outer(phi, n)), avg(grad(gamma) * kappa)) * dS - inner(avg(grad(phi) * kappa), 2 * avg(outer(gamma, n))) * dS + mu * inner(2 * avg(outer(phi, n)), 2 * avg(outer(gamma, n) * kappa)) * dS)) rhs = inner(gamma, self.f) * dx # Function to hold the solution self.soln = Function(V) # Setup problem and solver prob = LinearVariationalProblem(a, rhs, self.soln) self.solver = NonlinearVariationalSolver(prob) # Set the data structure for any user-defined time point self.vector_template = VectorDiffusion2D(size=len(self.function_space), comm_space=self.comm_space) # Set initial condition: # Setting up a Gaussian blob in the centre of the domain. self.vector_t_start = VectorDiffusion2D(size=len(self.function_space), comm_space=self.comm_space) x = SpatialCoordinate(self.mesh) initial_tracer = exp(-((x[0] - 5)**2 + (x[1] - 5)**2)) tmp = Function(self.function_space) tmp.interpolate(initial_tracer) self.vector_t_start.set_values(np.copy(tmp.dat.data))
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)
def build_initial_conditions(prognostic_variables, simulation_parameters): """ Initialises the prognostic variables based on the initial condition string. :arg prognostic_variables: a PrognosticVariables object. :arg simulation_parameters: a dictionary containing the simulation parameters. """ mesh = simulation_parameters['mesh'][-1] ic = simulation_parameters['ic'][-1] alphasq = simulation_parameters['alphasq'][-1] c0 = simulation_parameters['c0'][-1] gamma = simulation_parameters['gamma'][-1] x, = SpatialCoordinate(mesh) Ld = simulation_parameters['Ld'][-1] deltax = Ld / simulation_parameters['resolution'][-1] w = simulation_parameters['peak_width'][-1] epsilon = 1 ic_dict = { 'two_peaks': (0.2 * 2 / (exp(x - 403. / 15. * 40. / Ld) + exp(-x + 403. / 15. * 40. / Ld)) + 0.5 * 2 / (exp(x - 203. / 15. * 40. / Ld) + exp(-x + 203. / 15. * 40. / Ld))), 'gaussian': 0.5 * exp(-((x - 10.) / 2.)**2), 'gaussian_narrow': 0.5 * exp(-((x - 10.) / 1.)**2), 'gaussian_wide': 0.5 * exp(-((x - 10.) / 3.)**2), 'peakon': conditional(x < Ld / 2., exp((x - Ld / 2) / sqrt(alphasq)), exp(-(x - Ld / 2) / sqrt(alphasq))), 'one_peak': 0.5 * 2 / (exp(x - 203. / 15. * 40. / Ld) + exp(-x + 203. / 15. * 40. / Ld)), 'proper_peak': 0.5 * 2 / (exp(x - Ld / 4) + exp(-x + Ld / 4)), 'new_peak': 0.5 * 2 / (exp((x - Ld / 4) / w) + exp((-x + Ld / 4) / w)), 'flat': Constant(2 * pi**2 / (9 * 40**2)), 'fast_flat': Constant(0.1), 'coshes': Constant(2000) * cosh((2000**0.5 / 2) * (x - 0.75))**(-2) + Constant(1000) * cosh(1000**0.5 / 2 * (x - 0.25))**(-2), 'd_peakon': exp(-sqrt((x - Ld / 2)**2 + epsilon * deltax**2) / sqrt(alphasq)), 'zero': Constant(0.0), 'two_peakons': conditional( x < Ld / 4, exp((x - Ld / 4) / sqrt(alphasq)) - exp(-(x + Ld / 4) / sqrt(alphasq)), conditional( x < 3 * Ld / 4, exp(-(x - Ld / 4) / sqrt(alphasq)) - exp( (x - 3 * Ld / 4) / sqrt(alphasq)), exp((x - 5 * Ld / 4) / sqrt(alphasq)) - exp(-(x - 3 * Ld / 4) / sqrt(alphasq)))), 'twin_peakons': conditional( x < Ld / 4, exp((x - Ld / 4) / sqrt(alphasq)) + 0.5 * exp( (x - Ld / 2) / sqrt(alphasq)), conditional( x < Ld / 2, exp(-(x - Ld / 4) / sqrt(alphasq)) + 0.5 * exp( (x - Ld / 2) / sqrt(alphasq)), conditional( x < 3 * Ld / 4, exp(-(x - Ld / 4) / sqrt(alphasq)) + 0.5 * exp(-(x - Ld / 2) / sqrt(alphasq)), exp((x - 5 * Ld / 4) / sqrt(alphasq)) + 0.5 * exp(-(x - Ld / 2) / sqrt(alphasq))))), 'periodic_peakon': (conditional( x < Ld / 2, 0.5 / (1 - exp(-Ld / sqrt(alphasq))) * (exp((x - Ld / 2) / sqrt(alphasq)) + exp(-Ld / sqrt(alphasq)) * exp(-(x - Ld / 2) / sqrt(alphasq))), 0.5 / (1 - exp(-Ld / sqrt(alphasq))) * (exp(-(x - Ld / 2) / sqrt(alphasq)) + exp(-Ld / sqrt(alphasq)) * exp((x - Ld / 2) / sqrt(alphasq))))), 'cos_bell': conditional(x < Ld / 4, (cos(pi * (x - Ld / 8) / (2 * Ld / 8)))**2, 0.0), 'antisymmetric': 1 / (exp((x - Ld / 4) / Ld) + exp((-x + Ld / 4) / Ld)) - 1 / (exp( (Ld - x - Ld / 4) / Ld) + exp((Ld + x + Ld / 4) / Ld)) } ic_expr = ic_dict[ic] if prognostic_variables.scheme in ['upwind', 'LASCH']: VCG5 = FunctionSpace(mesh, "CG", 5) smooth_condition = Function(VCG5).interpolate(ic_expr) prognostic_variables.u.project(as_vector([smooth_condition])) # need to find initial m by solving helmholtz problem CG1 = FunctionSpace(mesh, "CG", 1) u0 = prognostic_variables.u p = TestFunction(CG1) m_CG = Function(CG1) ones = Function(prognostic_variables.Vu).project( as_vector([Constant(1.)])) Lm = (p * m_CG - p * dot(ones, u0) - alphasq * p.dx(0) * dot(ones, u0.dx(0))) * dx mprob0 = NonlinearVariationalProblem(Lm, m_CG) msolver0 = NonlinearVariationalSolver(mprob0, solver_parameters={ 'ksp_type': 'preonly', 'pc_type': 'lu' }) msolver0.solve() prognostic_variables.m.interpolate(m_CG) if prognostic_variables.scheme == 'LASCH': prognostic_variables.Eu.assign(prognostic_variables.u) prognostic_variables.Em.assign(prognostic_variables.m) elif prognostic_variables.scheme in ('conforming', 'hydrodynamic', 'test', 'LASCH_hydrodynamic', 'LASCH_hydrodynamic_m', 'no_gradient'): if ic == 'peakon': Vu = prognostic_variables.Vu # delta = Function(Vu) # middle_index = int(len(delta.dat.data[:]) / 2) # delta.dat.data[middle_index] = 1 # u0 = prognostic_variables.u # phi = TestFunction(Vu) # # eqn = phi * u0 * dx + alphasq * phi.dx(0) * u0.dx(0) * dx - phi * delta * dx # prob = NonlinearVariationalProblem(eqn, u0) # solver = NonlinearVariationalSolver(prob) # solver.solve() # W = MixedFunctionSpace((Vu, Vu)) # psi, phi = TestFunctions(W) # w = Function(W) # u, F = w.split() # u.interpolate(ic_expr) # u, F = split(w) # # eqn = (psi * u * dx - psi * (0.5 * u * u + F) * dx # + phi * F * dx + alphasq * phi.dx(0) * F.dx(0) * dx # - phi * u * u * dx - 0.5 * alphasq * phi * u.dx(0) * u.dx(0) * dx) # # u, F = w.split() # # prob = NonlinearVariationalProblem(eqn, w) # solver = NonlinearVariationalSolver(prob) # solver.solve() # prognostic_variables.u.assign(u) prognostic_variables.u.project(ic_expr) # prognostic_variables.u.interpolate(ic_expr) else: VCG5 = FunctionSpace(mesh, "CG", 5) smooth_condition = Function(VCG5).interpolate(ic_expr) prognostic_variables.u.project(smooth_condition) if prognostic_variables.scheme in [ 'LASCH_hydrodynamic', 'LASCH_hydrodynamic_m' ]: prognostic_variables.Eu.assign(prognostic_variables.u) else: raise NotImplementedError('Other schemes not yet implemented.')
def run(steady=False): """ solve CdT/dt = S + div(k*grad(T)) => C*v*(dT/dt)/k*dx - S*v/k*dx + grad(v)*grad(T)*dx = v*dot(grad(T), n)*ds """ steps = 250 dt = 1e-10 timescale = (0, steps * dt) if steady: print('Running steady state.') else: print(f'Running with time step {dt:.2g}s on time interval: ' f'{timescale[0]:.2g}s - {timescale[1]:.2g}s') dt_invc = Constant(1 / dt) extent = [40e-6, 40e-6, 40e-6] mesh = BoxMesh(20, 20, 20, *extent) V = FunctionSpace(mesh, 'CG', 1) print(V.dim()) T = Function(V) # temperature at time i+1 (electron for now) T_ = Function(V) # temperature at time i v = TestFunction(V) # test function S = create_S(mesh, V, extent) C = create_heat_capacity(mesh, V, extent) k = create_conductivity(mesh, V, T) set_initial_value(mesh, T_, extent) # Mass matrix section M = C * T * dt_invc * v * dx M_ = C * T_ * dt_invc * v * dx # Stiffness matrix section A = k * dot(grad(T), grad(v)) * dx # function section f = S * v * dx # boundaries bcs, R, b = create_dirichlet_bounds(mesh, V, T, v, k, g=100, boundary=[1, 2, 3, 4, 5, 6]) # bcs += create_dirichlet_bounds(mesh, V, T, v, k, 500, [6])[0] # bcs, R, b = create_robin_bounds(mesh, T, v, k, 1e8/(100), 1e8) if steady: steps = 1 a = A + R L = f + b else: a = M + A + R L = M_ + f + b prob = NonlinearVariationalProblem(a - L, T, bcs=bcs) solver = NonlinearVariationalSolver(prob, solver_parameters=SOLVE_PARAMS) T.assign(T_) timestamp = datetime.now().strftime("%d-%b-%Y-%H-%M-%S") outfile = File(f'{timestamp}/first_output.pvd') outfile.write(T_, target_degree=1, target_continuity=H1) last_perc = 0 for i in range(steps): solver.solve() perc = int(100 * (i + 1) / steps) if perc > last_perc: print(f'{perc}%') last_perc = perc T_.assign(T) outfile.write(T_, target_degree=1, target_continuity=H1)