def run(self, t, tmax, pickup=False):
logger.info("Profiling one time-step.")
self.semi_implicit_step()
logger.info("Profile complete for one time-step.")
def run(self, t, tmax, pickup=False):
"""
This is the timeloop. After completing the semi implicit step
any passively advected fields are updated, implicit diffusion and
physics updates are applied (if required).
"""
state = self.state
t = self.setup_timeloop(state, t, tmax, pickup)
dt = state.timestepping.dt
while t < tmax - 0.5*dt:
logger.info("at start of timestep, t=%s, dt=%s" % (t, dt))
t += dt
state.t.assign(t)
state.xnp1.assign(state.xn)
for name, evaluation in self.prescribed_fields:
state.fields(name).project(evaluation(t))
self.semi_implicit_step()
for name, advection in self.passive_advection:
field = getattr(state.fields, name)
# first computes ubar from state.xn and state.xnp1
advection.update_ubar(state.xn, state.xnp1, state.timestepping.alpha)
# advects a field from xn and puts result in xnp1
advection.apply(field, field)
state.xb.assign(state.xn)
state.xn.assign(state.xnp1)
with timed_stage("Diffusion"):
for name, diffusion in self.diffused_fields:
field = getattr(state.fields, name)
diffusion.apply(field, field)
with timed_stage("Physics"):
for physics in self.physics_list:
physics.apply()
with timed_stage("Dump output"):
state.dump(t)
if state.output.checkpoint:
state.chkpt.close()
logger.info("TIMELOOP complete. t=%s, tmax=%s" % (t, tmax))
def run(self, t, tmax, pickup=False):
"""
This is the timeloop. After completing the semi implicit step
any passively transported fields are updated, implicit diffusion and
physics updates are applied (if required).
"""
state = self.state
if pickup:
t = state.pickup_from_checkpoint()
state.setup_diagnostics()
with timed_stage("Dump output"):
state.setup_dump(t, tmax, pickup)
state.t.assign(t)
self.x.initialise(state)
while float(state.t) < tmax - 0.5*float(state.dt):
logger.info(f'at start of timestep, t={float(state.t)}, dt={float(state.dt)}')
self.x.update()
self.timestep()
for field in self.x.np1:
state.fields(field.name()).assign(field)
with timed_stage("Physics"):
for physics in self.physics_list:
physics.apply()
# TODO: Hack to ensure that xnp1 fields are updated
for field in self.x.np1:
field.assign(state.fields(field.name()))
state.t.assign(state.t + state.dt)
with timed_stage("Dump output"):
state.dump(float(state.t))
if state.output.checkpoint:
state.chkpt.close()
logger.info(f'TIMELOOP complete. t={float(state.t)}, tmax={tmax}')
def semi_implicit_step(self):
self.state.xrhs -= self.state.xn
if self.hybridization:
solver = self.linear_solver.hybridized_solver
else:
solver = self.linear_solver.urho_solver
logger.info("Warming up linear solver.")
with PETSc.Log.Stage("Warm-up stage"):
solver.solve()
solver._problem.u.assign(0.0)
solver.snes.setConvergenceHistory()
solver.snes.ksp.setConvergenceHistory()
logger.info("Warm up finished. Timing linear solver.")
with PETSc.Log.Stage("linear_solve"):
solver.solve()
if not self.suppress_data_output:
self.extract_ksp_info(solver)
def __init__(self,
mesh,
dt,
output=None,
parameters=None,
diagnostics=None,
diagnostic_fields=None):
if output is None:
raise RuntimeError(
"You must provide a directory name for dumping results")
else:
self.output = output
self.parameters = parameters
if diagnostics is not None:
self.diagnostics = diagnostics
else:
self.diagnostics = Diagnostics()
if diagnostic_fields is not None:
self.diagnostic_fields = diagnostic_fields
else:
self.diagnostic_fields = []
# The mesh
self.mesh = mesh
self.spaces = SpaceCreator(mesh)
if self.output.dumplist is None:
self.output.dumplist = []
self.fields = StateFields(*self.output.dumplist)
self.dumpdir = None
self.dumpfile = None
self.to_pickup = None
# figure out if we're on a sphere
try:
self.on_sphere = (mesh._base_mesh.geometric_dimension() == 3
and mesh._base_mesh.topological_dimension() == 2)
except AttributeError:
self.on_sphere = (mesh.geometric_dimension() == 3
and mesh.topological_dimension() == 2)
# build the vertical normal and define perp for 2d geometries
dim = mesh.topological_dimension()
if self.on_sphere:
x = SpatialCoordinate(mesh)
R = sqrt(inner(x, x))
self.k = interpolate(x / R, mesh.coordinates.function_space())
if dim == 2:
outward_normals = CellNormal(mesh)
self.perp = lambda u: cross(outward_normals, u)
else:
kvec = [0.0] * dim
kvec[dim - 1] = 1.0
self.k = Constant(kvec)
if dim == 2:
self.perp = lambda u: as_vector([-u[1], u[0]])
# setup logger
logger.setLevel(output.log_level)
set_log_handler(mesh.comm)
if parameters is not None:
logger.info("Physical parameters that take non-default values:")
logger.info(", ".join("%s: %s" % (k, float(v))
for (k, v) in vars(parameters).items()))
# Constant to hold current time
self.t = Constant(0.0)
if type(dt) is Constant:
self.dt = dt
elif type(dt) in (float, int):
self.dt = Constant(dt)
else:
raise TypeError(
f'dt must be a Constant, float or int, not {type(dt)}')
def __init__(self,
mesh,
vertical_degree=None,
horizontal_degree=1,
family="RT",
Coriolis=None,
sponge_function=None,
hydrostatic=None,
timestepping=None,
output=None,
parameters=None,
diagnostics=None,
fieldlist=None,
diagnostic_fields=None,
u_bc_ids=None):
self.family = family
self.vertical_degree = vertical_degree
self.horizontal_degree = horizontal_degree
self.Omega = Coriolis
self.mu = sponge_function
self.hydrostatic = hydrostatic
self.timestepping = timestepping
if output is None:
raise RuntimeError(
"You must provide a directory name for dumping results")
else:
self.output = output
self.parameters = parameters
if fieldlist is None:
raise RuntimeError(
"You must provide a fieldlist containing the names of the prognostic fields"
)
else:
self.fieldlist = fieldlist
if diagnostics is not None:
self.diagnostics = diagnostics
else:
self.diagnostics = Diagnostics(*fieldlist)
if diagnostic_fields is not None:
self.diagnostic_fields = diagnostic_fields
else:
self.diagnostic_fields = []
if u_bc_ids is not None:
self.u_bc_ids = u_bc_ids
else:
self.u_bc_ids = []
# The mesh
self.mesh = mesh
# Build the spaces
self._build_spaces(mesh, vertical_degree, horizontal_degree, family)
# Allocate state
self._allocate_state()
if self.output.dumplist is None:
self.output.dumplist = fieldlist
self.fields = FieldCreator(fieldlist, self.xn, self.output.dumplist)
# set up bcs
V = self.fields('u').function_space()
self.bcs = []
if V.extruded:
self.bcs.append(DirichletBC(V, 0.0, "bottom"))
self.bcs.append(DirichletBC(V, 0.0, "top"))
for id in self.u_bc_ids:
self.bcs.append(DirichletBC(V, 0.0, id))
self.dumpfile = None
# figure out if we're on a sphere
try:
self.on_sphere = (mesh._base_mesh.geometric_dimension() == 3
and mesh._base_mesh.topological_dimension() == 2)
except AttributeError:
self.on_sphere = (mesh.geometric_dimension() == 3
and mesh.topological_dimension() == 2)
# build the vertical normal and define perp for 2d geometries
dim = mesh.topological_dimension()
if self.on_sphere:
x = SpatialCoordinate(mesh)
R = sqrt(inner(x, x))
self.k = interpolate(x / R, mesh.coordinates.function_space())
if dim == 2:
outward_normals = CellNormal(mesh)
self.perp = lambda u: cross(outward_normals, u)
else:
kvec = [0.0] * dim
kvec[dim - 1] = 1.0
self.k = Constant(kvec)
if dim == 2:
self.perp = lambda u: as_vector([-u[1], u[0]])
# project test function for hydrostatic case
if self.hydrostatic:
self.h_project = lambda u: u - self.k * inner(u, self.k)
else:
self.h_project = lambda u: u
# Constant to hold current time
self.t = Constant(0.0)
# setup logger
logger.setLevel(output.log_level)
set_log_handler(mesh.comm)
logger.info("Timestepping parameters that take non-default values:")
logger.info(", ".join("%s: %s" % item
for item in vars(timestepping).items()))
if parameters is not None:
logger.info("Physical parameters that take non-default values:")
logger.info(", ".join("%s: %s" % item
for item in vars(parameters).items()))