diff --git a/examples/compressible/dcmip_3_1_meanflow_quads.py b/examples/compressible/dcmip_3_1_meanflow_quads.py index 0935eb82a..e37f790a9 100644 --- a/examples/compressible/dcmip_3_1_meanflow_quads.py +++ b/examples/compressible/dcmip_3_1_meanflow_quads.py @@ -78,9 +78,10 @@ # I/O dirname = 'dcmip_3_1_meanflow' -output = OutputParameters(dirname=dirname, - dumpfreq=dumpfreq, - log_level='INFO') +output = OutputParameters( + dirname=dirname, + dumpfreq=dumpfreq, +) diagnostic_fields = [Perturbation('theta'), Perturbation('rho'), CompressibleKineticEnergy()] io = IO(domain, output, diagnostic_fields=diagnostic_fields) diff --git a/examples/compressible/dry_bryan_fritsch.py b/examples/compressible/dry_bryan_fritsch.py index 265fef490..a6dba5a22 100644 --- a/examples/compressible/dry_bryan_fritsch.py +++ b/examples/compressible/dry_bryan_fritsch.py @@ -50,12 +50,13 @@ no_normal_flow_bc_ids=[1, 2]) # I/O -output = OutputParameters(dirname=dirname, - dumpfreq=int(tmax / (5*dt)), - dumplist=['rho'], - dump_vtus=False, - dump_nc=True, - log_level='INFO') +output = OutputParameters( + dirname=dirname, + dumpfreq=int(tmax / (5*dt)), + dumplist=['rho'], + dump_vtus=False, + dump_nc=True, +) diagnostic_fields = [Perturbation('theta')] io = IO(domain, output, diagnostic_fields=diagnostic_fields) diff --git a/examples/compressible/mountain_hydrostatic.py b/examples/compressible/mountain_hydrostatic.py index b0bd8a76f..b90630884 100644 --- a/examples/compressible/mountain_hydrostatic.py +++ b/examples/compressible/mountain_hydrostatic.py @@ -61,10 +61,11 @@ # I/O dirname = 'hydrostatic_mountain' -output = OutputParameters(dirname=dirname, - dumpfreq=dumpfreq, - dumplist=['u'], - log_level='INFO') +output = OutputParameters( + dirname=dirname, + dumpfreq=dumpfreq, + dumplist=['u'], +) diagnostic_fields = [CourantNumber(), VelocityZ(), HydrostaticImbalance(eqns), Perturbation('theta'), Perturbation('rho')] io = IO(domain, output, diagnostic_fields=diagnostic_fields) diff --git a/examples/compressible/skamarock_klemp_hydrostatic.py b/examples/compressible/skamarock_klemp_hydrostatic.py index 42d111f1e..90f2bd700 100644 --- a/examples/compressible/skamarock_klemp_hydrostatic.py +++ b/examples/compressible/skamarock_klemp_hydrostatic.py @@ -47,10 +47,11 @@ # I/O dirname = 'skamarock_klemp_hydrostatic' -output = OutputParameters(dirname=dirname, - dumpfreq=dumpfreq, - dumplist=['u'], - log_level='INFO') +output = OutputParameters( + dirname=dirname, + dumpfreq=dumpfreq, + dumplist=['u'], +) diagnostic_fields = [CourantNumber(), Perturbation('theta'), Perturbation('rho')] io = IO(domain, output, diagnostic_fields=diagnostic_fields) diff --git a/examples/compressible/skamarock_klemp_nonhydrostatic.py b/examples/compressible/skamarock_klemp_nonhydrostatic.py index 2b544e043..c6c32af7a 100644 --- a/examples/compressible/skamarock_klemp_nonhydrostatic.py +++ b/examples/compressible/skamarock_klemp_nonhydrostatic.py @@ -51,12 +51,13 @@ points_z = [H/2.] points = np.array([p for p in itertools.product(points_x, points_z)]) dirname = 'skamarock_klemp_nonlinear' -output = OutputParameters(dirname=dirname, - dumpfreq=dumpfreq, - pddumpfreq=dumpfreq, - dumplist=['u'], - point_data=[('theta_perturbation', points)], - log_level='INFO') +output = OutputParameters( + dirname=dirname, + dumpfreq=dumpfreq, + pddumpfreq=dumpfreq, + dumplist=['u'], + point_data=[('theta_perturbation', points)], +) diagnostic_fields = [CourantNumber(), Gradient("u"), Perturbation('theta'), Gradient("theta_perturbation"), Perturbation('rho'), RichardsonNumber("theta", parameters.g/Tsurf), Gradient("theta")] diff --git a/examples/compressible/straka_bubble.py b/examples/compressible/straka_bubble.py index 20442bf44..2d3388ccc 100644 --- a/examples/compressible/straka_bubble.py +++ b/examples/compressible/straka_bubble.py @@ -53,10 +53,11 @@ # I/O dirname = "straka_dx%s_dt%s" % (delta, dt) dumpfreq = int(tmax / (ndumps*dt)) - output = OutputParameters(dirname=dirname, - dumpfreq=dumpfreq, - dumplist=['u'], - log_level='INFO') + output = OutputParameters( + dirname=dirname, + dumpfreq=dumpfreq, + dumplist=['u'], + ) diagnostic_fields = [CourantNumber(), Perturbation('theta'), Perturbation('rho')] io = IO(domain, output, diagnostic_fields=diagnostic_fields) diff --git a/examples/compressible/unsaturated_bubble.py b/examples/compressible/unsaturated_bubble.py index d6d8afd4c..bf8224b91 100644 --- a/examples/compressible/unsaturated_bubble.py +++ b/examples/compressible/unsaturated_bubble.py @@ -51,9 +51,13 @@ # I/O dirname = 'unsaturated_bubble' -output = OutputParameters(dirname=dirname, dumpfreq=tdump, dump_nc=True, - dumplist=['cloud_water', 'rain'], log_level='INFO', - checkpoint=False) +output = OutputParameters( + dirname=dirname, + dumpfreq=tdump, + dump_nc=True, + dumplist=['cloud_water', 'rain'], + checkpoint=False +) diagnostic_fields = [RelativeHumidity(eqns), Perturbation('theta'), Perturbation('water_vapour'), Perturbation('rho')] io = IO(domain, output, diagnostic_fields=diagnostic_fields) diff --git a/examples/incompressible/skamarock_klemp_incompressible.py b/examples/incompressible/skamarock_klemp_incompressible.py index ce38a5230..0af31929e 100644 --- a/examples/incompressible/skamarock_klemp_incompressible.py +++ b/examples/incompressible/skamarock_klemp_incompressible.py @@ -44,10 +44,11 @@ eqns = IncompressibleBoussinesqEquations(domain, parameters) # I/O -output = OutputParameters(dirname='skamarock_klemp_incompressible', - dumpfreq=dumpfreq, - dumplist=['u'], - log_level='INFO') +output = OutputParameters( + dirname='skamarock_klemp_incompressible', + dumpfreq=dumpfreq, + dumplist=['u'], +) # list of diagnostic fields, each defined in a class in diagnostics.py diagnostic_fields = [CourantNumber(), Divergence(), Perturbation('b')] io = IO(domain, output, diagnostic_fields=diagnostic_fields) diff --git a/examples/shallow_water/linear_williamson_2.py b/examples/shallow_water/linear_williamson_2.py index d5c72828d..7c4285542 100644 --- a/examples/shallow_water/linear_williamson_2.py +++ b/examples/shallow_water/linear_williamson_2.py @@ -45,9 +45,10 @@ eqns = LinearShallowWaterEquations(domain, parameters, fexpr=fexpr) # I/O -output = OutputParameters(dirname='linear_williamson_2', - dumpfreq=dumpfreq, - log_level='INFO') +output = OutputParameters( + dirname='linear_williamson_2', + dumpfreq=dumpfreq, +) diagnostic_fields = [SteadyStateError('u'), SteadyStateError('D')] io = IO(domain, output, diagnostic_fields=diagnostic_fields) diff --git a/examples/shallow_water/moist_thermal_williamson5.py b/examples/shallow_water/moist_thermal_williamson5.py index 7980beac3..acb73e7a2 100644 --- a/examples/shallow_water/moist_thermal_williamson5.py +++ b/examples/shallow_water/moist_thermal_williamson5.py @@ -74,10 +74,11 @@ # I/O dirname = "moist_thermal_williamson5" -output = OutputParameters(dirname=dirname, - dumplist_latlon=['D'], - dumpfreq=dumpfreq, - log_level='INFO') +output = OutputParameters( + dirname=dirname, + dumplist_latlon=['D'], + dumpfreq=dumpfreq, +) diagnostic_fields = [Sum('D', 'topography'), CourantNumber()] io = IO(domain, output, diagnostic_fields=diagnostic_fields) diff --git a/examples/shallow_water/thermal_williamson2.py b/examples/shallow_water/thermal_williamson2.py index 62984f4c7..28f3a1d8f 100644 --- a/examples/shallow_water/thermal_williamson2.py +++ b/examples/shallow_water/thermal_williamson2.py @@ -43,10 +43,11 @@ # IO dirname = "thermal_williamson2" -output = OutputParameters(dirname=dirname, - dumpfreq=dumpfreq, - dumplist_latlon=['D', 'D_error'], - log_level='INFO') +output = OutputParameters( + dirname=dirname, + dumpfreq=dumpfreq, + dumplist_latlon=['D', 'D_error'], +) diagnostic_fields = [RelativeVorticity(), PotentialVorticity(), ShallowWaterKineticEnergy(), diff --git a/examples/shallow_water/williamson_2.py b/examples/shallow_water/williamson_2.py index 432dd7cc6..47563a632 100644 --- a/examples/shallow_water/williamson_2.py +++ b/examples/shallow_water/williamson_2.py @@ -52,11 +52,12 @@ # I/O dirname = "williamson_2_ref%s_dt%s" % (ref_level, dt) dumpfreq = int(tmax / (ndumps*dt)) - output = OutputParameters(dirname=dirname, - dumpfreq=dumpfreq, - dumplist_latlon=['D', 'D_error'], - log_level='INFO', - dump_nc=True) + output = OutputParameters( + dirname=dirname, + dumpfreq=dumpfreq, + dumplist_latlon=['D', 'D_error'], + dump_nc=True, + ) diagnostic_fields = [RelativeVorticity(), PotentialVorticity(), ShallowWaterKineticEnergy(), diff --git a/examples/shallow_water/williamson_5.py b/examples/shallow_water/williamson_5.py index 0bcb43a3d..b71fddab1 100644 --- a/examples/shallow_water/williamson_5.py +++ b/examples/shallow_water/williamson_5.py @@ -62,10 +62,11 @@ # I/O dirname = "williamson_5_ref%s_dt%s" % (ref_level, dt) dumpfreq = int(tmax / (ndumps*dt)) - output = OutputParameters(dirname=dirname, - dumplist_latlon=['D'], - dumpfreq=dumpfreq, - log_level='INFO') + output = OutputParameters( + dirname=dirname, + dumplist_latlon=['D'], + dumpfreq=dumpfreq, + ) diagnostic_fields = [Sum('D', 'topography')] io = IO(domain, output, diagnostic_fields=diagnostic_fields) diff --git a/gusto/__init__.py b/gusto/__init__.py index 9cf36fe52..772e94656 100644 --- a/gusto/__init__.py +++ b/gusto/__init__.py @@ -13,6 +13,10 @@ def perp(self, o, a): _monkey_patch_ufl() +# Start logging first, incase anything goes wrong +from gusto.logging import * # noqa +set_log_handler() + from gusto.active_tracers import * # noqa from gusto.common_forms import * # noqa from gusto.configuration import * # noqa diff --git a/gusto/active_tracers.py b/gusto/active_tracers.py index e261466cd..13eb08a8f 100644 --- a/gusto/active_tracers.py +++ b/gusto/active_tracers.py @@ -8,7 +8,8 @@ """ from enum import Enum -from gusto.configuration import TransportEquationType, logger +from gusto.configuration import TransportEquationType +from gusto.logging import logger __all__ = ["TracerVariableType", "Phases", "ActiveTracer", "WaterVapour", "CloudWater", "Rain"] diff --git a/gusto/configuration.py b/gusto/configuration.py index db9260ae1..4bf3b8b2b 100644 --- a/gusto/configuration.py +++ b/gusto/configuration.py @@ -1,35 +1,15 @@ """Some simple tools for configuring the model.""" from abc import ABCMeta, abstractproperty from enum import Enum -import logging -from logging import DEBUG, INFO, WARNING from firedrake import sqrt, Constant -__all__ = ["WARNING", "INFO", "DEBUG", "IntegrateByParts", - "TransportEquationType", "OutputParameters", - "CompressibleParameters", "ShallowWaterParameters", - "logger", "EmbeddedDGOptions", "RecoveryOptions", "SUPGOptions", - "SpongeLayerParameters", "DiffusionParameters"] - -logger = logging.getLogger("gusto") - - -def set_log_handler(comm): - """ - Sets the handler for logging. - - Args: - comm (:class:`MPI.Comm`): MPI communicator. - """ - handler = logging.StreamHandler() - handler.setFormatter(logging.Formatter(fmt="%(name)s:%(levelname)s %(message)s")) - if logger.hasHandlers(): - logger.handlers.clear() - if comm.rank == 0: - logger.addHandler(handler) - else: - logger.addHandler(logging.NullHandler()) +__all__ = [ + "IntegrateByParts", "TransportEquationType", "OutputParameters", + "CompressibleParameters", "ShallowWaterParameters", + "EmbeddedDGOptions", "RecoveryOptions", "SUPGOptions", + "SpongeLayerParameters", "DiffusionParameters" +] class IntegrateByParts(Enum): @@ -92,7 +72,7 @@ def __setattr__(self, name, value): # Almost all parameters should be Constants -- but there are some # specific exceptions which should be kept as integers - if type(value) in [float, int] and name not in ['dumpfreq', 'pddumpfreq', 'chkptfreq', 'log_level']: + if type(value) in [float, int] and name not in ['dumpfreq', 'pddumpfreq', 'chkptfreq']: object.__setattr__(self, name, Constant(value)) else: object.__setattr__(self, name, value) @@ -101,9 +81,6 @@ def __setattr__(self, name, value): class OutputParameters(Configuration): """Parameters for controlling outputting.""" - #: log_level for logger, can be DEBUG, INFO or WARNING. Takes - #: default value "warning" - log_level = WARNING dump_vtus = True dump_nc = False dumpfreq = 1 diff --git a/gusto/coordinates.py b/gusto/coordinates.py index af873d583..04a5cd8e7 100644 --- a/gusto/coordinates.py +++ b/gusto/coordinates.py @@ -3,7 +3,7 @@ Coordinate fields are stored in specified VectorFunctionSpaces. """ -from gusto.configuration import logger +from gusto.logging import logger from firedrake import (SpatialCoordinate, sqrt, atan_2, asin, Function) import numpy as np diff --git a/gusto/forcing.py b/gusto/forcing.py index 76667c5ed..662463515 100644 --- a/gusto/forcing.py +++ b/gusto/forcing.py @@ -1,10 +1,14 @@ """Discretisation of dynamic forcing terms, such as the pressure gradient.""" -from firedrake import (Function, TrialFunctions, DirichletBC, - LinearVariationalProblem, LinearVariationalSolver) -from gusto.configuration import logger, DEBUG -from gusto.labels import transport, diffusion, time_derivative, hydrostatic +from firedrake import ( + Function, TrialFunctions, DirichletBC, LinearVariationalProblem, + LinearVariationalSolver +) from gusto.fml import drop, replace_subject, name +from gusto.labels import ( + transport, diffusion, time_derivative, hydrostatic +) +from gusto.logging import logger, DEBUG, logging_ksp_monitor_true_residual __all__ = ["Forcing"] @@ -93,22 +97,20 @@ def __init__(self, equation, alpha): a.form, L_implicit.form, self.xF, bcs=bcs ) - solver_parameters = {} - if logger.isEnabledFor(DEBUG): - solver_parameters["ksp_monitor_true_residual"] = None - self.solvers = {} self.solvers["explicit"] = LinearVariationalSolver( explicit_forcing_problem, - solver_parameters=solver_parameters, options_prefix="ExplicitForcingSolver" ) self.solvers["implicit"] = LinearVariationalSolver( implicit_forcing_problem, - solver_parameters=solver_parameters, options_prefix="ImplicitForcingSolver" ) + if logger.isEnabledFor(DEBUG): + self.solvers["explicit"].snes.ksp.setMonitor(logging_ksp_monitor_true_residual) + self.solvers["implicit"].snes.ksp.setMonitor(logging_ksp_monitor_true_residual) + def apply(self, x_in, x_nl, x_out, label): """ Applies the discretisation for a forcing term F(x). diff --git a/gusto/initialisation_tools.py b/gusto/initialisation_tools.py index ad85578fb..e338c8b2e 100644 --- a/gusto/initialisation_tools.py +++ b/gusto/initialisation_tools.py @@ -9,7 +9,7 @@ sin, cos, sqrt, asin, atan_2, as_vector, min_value, max_value, FunctionSpace, \ errornorm, zero from gusto import thermodynamics -from gusto.configuration import logger +from gusto.logging import logger from gusto.recovery import Recoverer, BoundaryMethod diff --git a/gusto/io.py b/gusto/io.py index c26447042..f436607e4 100644 --- a/gusto/io.py +++ b/gusto/io.py @@ -10,7 +10,7 @@ from firedrake import (Function, functionspaceimpl, File, DumbCheckpoint, FILE_CREATE, FILE_READ, CheckpointFile) import numpy as np -from gusto.configuration import logger, set_log_handler +from gusto.logging import logger, update_logfile_location __all__ = ["pick_up_mesh", "IO"] @@ -36,6 +36,7 @@ def pick_up_mesh(output, mesh_name): chkfile = output.checkpoint_pickup_filename else: dumpdir = path.join("results", output.dirname) + update_logfile_location(dumpdir) chkfile = path.join(dumpdir, "chkpt.h5") with CheckpointFile(chkfile, 'r') as chk: mesh = chk.load_mesh(mesh_name) @@ -222,10 +223,6 @@ def __init__(self, domain, output, diagnostics=None, diagnostic_fields=None): self.dumpfile = None self.to_pick_up = None - # setup logger - logger.setLevel(output.log_level) - set_log_handler(self.mesh.comm) - def log_parameters(self, equation): """ Logs an equation's physical parameters that take non-default values. @@ -297,6 +294,7 @@ def setup_dump(self, state_fields, t, pick_up=False): self.output.point_data, self.output.checkpoint and not pick_up]): # setup output directory and check that it does not already exist self.dumpdir = path.join("results", self.output.dirname) + update_logfile_location(self.dumpdir) running_tests = '--running-tests' in sys.argv or "pytest" in self.output.dirname if self.mesh.comm.Get_rank() == 0: @@ -433,6 +431,7 @@ def pick_up_from_checkpoint(self, state_fields): # Set dumpdir if has not been done already if self.dumpdir is None: self.dumpdir = path.join("results", self.output.dirname) + update_logfile_location(self.dumpdir) # Need to pick up reference profiles, but don't know which are stored possible_ref_profiles = [] diff --git a/gusto/linear_solvers.py b/gusto/linear_solvers.py index d6c091373..68ddf1477 100644 --- a/gusto/linear_solvers.py +++ b/gusto/linear_solvers.py @@ -15,7 +15,7 @@ from pyop2.profiling import timed_function, timed_region from gusto.active_tracers import TracerVariableType -from gusto.configuration import logger, DEBUG +from gusto.logging import logger, DEBUG, logging_ksp_monitor_true_residual from gusto.labels import linearisation, time_derivative, hydrostatic from gusto import thermodynamics from gusto.fml.form_manipulation_language import Term, drop @@ -55,17 +55,26 @@ def __init__(self, equations, alpha=0.5, solver_parameters=None, solver_parameters = p self.solver_parameters = solver_parameters - if logger.isEnabledFor(DEBUG): - self.solver_parameters["ksp_monitor_true_residual"] = None + # ~ if logger.isEnabledFor(DEBUG): + # ~ self.solver_parameters["ksp_monitor_true_residual"] = None # setup the solver self._setup_solver() + @staticmethod + def log_ksp_residuals(ksp): + if logger.isEnabledFor(DEBUG): + ksp.setMonitor(logging_ksp_monitor_true_residual) + @abstractproperty def solver_parameters(self): """Solver parameters for this solver""" pass + @abstractmethod + def _setup_solver(self): + pass + @abstractmethod def solve(self): pass @@ -151,16 +160,6 @@ def __init__(self, equations, alpha=0.5, logger.warning("default quadrature degree most likely not sufficient for this degree element") self.quadrature_degree = (5, 5) - if logger.isEnabledFor(DEBUG): - # Set outer solver to FGMRES and turn on KSP monitor for the outer system - self.solver_parameters["ksp_type"] = "fgmres" - self.solver_parameters["mat_type"] = "aij" - self.solver_parameters["pmat_type"] = "matfree" - self.solver_parameters["ksp_monitor_true_residual"] = None - - # Turn monitor on for the trace system - self.solver_parameters["condensed_field"]["ksp_monitor_true_residual"] = None - super().__init__(equations, alpha, solver_parameters, overwrite_solver_parameters) @@ -350,6 +349,13 @@ def L_tr(f): # post-solve self.bcs = self.equations.bcs['u'] + # Log residuals on hybridized solver + self.log_ksp_residuals(self.hybridized_solver.snes.ksp) + # Log residuals on the trace system too + from gusto.logging import attach_custom_monitor + python_context = self.hybridized_solver.snes.ksp.pc.getPythonContext() + attach_custom_monitor(python_context, logging_ksp_monitor_true_residual) + @timed_function("Gusto:LinearSolve") def solve(self, xrhs, dy): """ @@ -506,6 +512,9 @@ def trace_nullsp(T): self.b) self.b_solver = LinearVariationalSolver(b_problem) + # Log residuals on hybridized solver + self.log_ksp_residuals(self.up_solver.snes.ksp) + @timed_function("Gusto:LinearSolve") def solve(self, xrhs, dy): """ diff --git a/gusto/logging.py b/gusto/logging.py new file mode 100644 index 000000000..6a5a5cbd2 --- /dev/null +++ b/gusto/logging.py @@ -0,0 +1,296 @@ +"""Gusto Logging Module + +All logging functionality for Gusto is controlled from +``gusto.logging``. A logger object ``logging.getLogger("gusto")`` is +created internally. + +The primary means of configuration is via environment variables, the +same way that the standard Python root logger is. See the +:mod:`logging` page for details. + +Set ``GUSTO_LOG_LEVEL`` to any of ``DEBUG``, ``INFO``, ``WARNING``, +``ERROR`` or ``CRITICAL`` (from most verbose to least verbose). + +Additionally the level of console (`stderr`) logging and logfile based +logging can be controlled separately. Console logging verbosity is set +using ``GUSTO_CONSOLE_LOG_LEVEL``. Logfile logging verbosity is set using +``GUSTO_LOGFILE_LOG_LEVEL``. + +By default a script that imports gusto will log only from rank 0 to the +console and to a file. This can be changed by setting the environment +variable ``GUSTO_PARALLEL_LOG`` to ``CONSOLE``, ``FILE`` or ``BOTH``. +Setting these will log from all ranks, not just rank 0. Console output +will be interleaved, but log files contain the rank number as part of +the logfile name. + +""" + +import logging +import sys +import os +import shutil + +from datetime import datetime +from logging import NOTSET, DEBUG, INFO, WARNING, ERROR, CRITICAL # noqa: F401 +from pathlib import Path + +from firedrake.slate.static_condensation import scpc, hybridization +from petsc4py import PETSc +from pyop2.mpi import COMM_WORLD + +__all__ = [ + "logger", "set_log_handler", "LoggingError", "NOTSET", "DEBUG", + "INFO", "WARNING", "ERROR", "CRITICAL" +] + + +class LoggingError(Exception): + pass + + +logging.captureWarnings(True) +logger = logging.getLogger("gusto") + + +def capture_exceptions(exception_type, exception_value, traceback, logger=logger): + """ This function allows all unhandled exceptions to be logged to + Gusto's logs + """ + logger.error( + "Gusto is logging this unhandled exception:", + exc_info=(exception_type, exception_value, traceback) + ) + + +sys.excepthook = capture_exceptions + +# Set the log level based on environment variables +log_level = os.environ.get("GUSTO_LOG_LEVEL", WARNING) +logfile_level = os.environ.get("GUSTO_FILE_LOG_LEVEL", DEBUG) +logconsole_level = os.environ.get("GUSTO_CONSOLE_LOG_LEVEL", INFO) +log_level_list = [log_level, logfile_level, logconsole_level] +log_levels = [ + logging.getLevelNamesMapping().get(x) if isinstance(x, str) + else x + for x in log_level_list +] +logger.setLevel(min(log_levels)) + +# Setup parallel logging based on environment variables +parallel_log = os.environ.get("GUSTO_PARALLEL_LOG", None) +options = ["CONSOLE", "FILE", "BOTH"] +if parallel_log is not None: + parallel_log = parallel_log.upper() + if parallel_log.upper() not in options: + parallel_log = None + + +def create_logfile_handler(path, mode="w"): + """ Handler for logfiles. + + Args: + path: path to log file + + """ + logfile = logging.FileHandler(filename=path, mode=mode) + logfile.setLevel(logfile_level) + logfile_formatter = logging.Formatter( + '%(asctime)s %(levelname)-8s %(message)s' + ) + logfile.setFormatter(logfile_formatter) + return logfile + + +def create_console_handler(fmt): + """ Handler for console logging. + + Args: + fmt: format string for log output + """ + console = logging.StreamHandler() + console.setLevel(logconsole_level) + console_formatter = logging.Formatter(fmt) + console.setFormatter(console_formatter) + return console + + +def set_log_handler(comm=COMM_WORLD): + """ Set all handlers for logging. + + Args: + comm (:class:`MPI.Comm`): MPI communicator. + """ + # Set up logging + timestamp = datetime.now() + logfile_name = f"temp-gusto-{timestamp.strftime('%Y-%m-%dT%H%M%S')}" + if parallel_log in ["FILE", "BOTH"]: + logfile_name += f"_{comm.rank}" + # PID is required here for running pytest with xdist + logfile_name += f"_{os.getpid()}.log" + if comm.rank == 0: + os.makedirs("results", exist_ok=True) + logfile_path = os.path.join("results", logfile_name) + + console_format_str = "" + if parallel_log in ["CONSOLE", "BOTH"]: + console_format_str += f"[{comm.rank}] " + console_format_str += '%(levelname)-8s %(message)s' + + if comm.rank == 0: + # Always log on rank 0 + lfh = create_logfile_handler(logfile_path) + lfh.name = "gusto-temp-file-log" + logger.addHandler(lfh) + ch = create_console_handler(console_format_str) + ch.name = "gusto-console-log" + logger.addHandler(ch) + else: + # Only log on other ranks if enabled + if parallel_log in ["FILE", "BOTH"]: + lfh = create_logfile_handler(logfile_path) + lfh.name = "gusto-temp-file-log" + logger.addHandler(lfh) + if parallel_log in ["CONSOLE", "BOTH"]: + ch = create_console_handler(console_format_str) + ch.name = "gusto-console-log" + logger.addHandler(ch) + if not logger.handlers: + logger.addHandler(logging.NullHandler()) + + logger.info("Running %s" % " ".join(sys.argv)) + + +def update_logfile_location(new_path): + """ Update the location of the logfile. + + This is used to move the temporary log file created in the results + directory to the appropriate model directory. + + """ + new_path = Path(new_path) + fh = [*filter(lambda x: x.name == "gusto-temp-file-log", logger.handlers)] + + if len(fh) == 1: + fh = fh[0] + logger.debug("Closing temporary logger and moving logfile") + old_path = Path(fh.baseFilename) + filename = Path(old_path.name.removeprefix("temp-")) + fh.flush() + fh.close() + logger.removeHandler(fh) + + os.makedirs(new_path, exist_ok=True) + # Use shutil.move and not os.rename as new path may be on a + # different file system. Notably, this is the case for CI. + shutil.move(old_path, new_path/filename) + + new_fh = create_logfile_handler(new_path/filename, mode="a") + new_fh.name = "gusto-file-log" + logger.addHandler(new_fh) + logger.debug("Re-opening logger") + elif len(fh) > 1: + raise LoggingError( + "More than one log handler with name `gusto-temp-file-log`\n" + "Logging has been set up incorrectly" + ) + + +# We want a map from ENUM to Norm names +_norm_to_enum = {k: v for k, v in PETSc.KSP.NormType.__dict__.items() if isinstance(v, int)} +_enum_to_norm = {v: k.lower() for k, v in _norm_to_enum.items() if 'NORM_' not in k} + + +# The logging monitors will only log at level debug, but you should avoid +# adding an expensive Python callback the log level is not DEBUG by +# checking the logger like so: +# ``` +# if logger.isEnabledFor(DEBUG): +# ksp.setMonitor(logging_ksp_monitor) +# ``` +def logging_ksp_monitor(ksp, iteration, residual_norm): + ''' + Clone of C code at: + https://petsc.org/main/src/ksp/ksp/interface/iterativ.c.html#KSPMonitorResidual + Example output: + Residual norms for firedrake_0_ solve + 0 KSP Residual norm 3.175267221735e+00 + + ''' + tab_level = ksp.getTabLevel() + tab = ' ' + if iteration == 0: + logger.debug(tab*tab_level + f'Residual norms for {ksp.prefix} solve') + logger.debug( + tab*(tab_level - 1) + + f'{iteration: 5d} KSP Residual norm {residual_norm:14.12e}' + ) + + +def logging_ksp_monitor_true_residual(ksp, iteration, residual_norm): + ''' + Clone of C code: + https://petsc.org/main/src/ksp/ksp/interface/iterativ.c.html#KSPMonitorTrueResidual + Example output: + Residual norms for firedrake_0_ solve + 0 KSP preconditioned resid norm 3.175267221735e+00 true resid norm 3.175267221735e+00 ||r(i)||/||b|| 1.000000000000e+00 + + ''' + tab_level = ksp.getTabLevel() + tab = ' ' + residual = ksp.buildResidual() + true_norm = residual.norm(PETSc.NormType.NORM_2) + bnorm = ksp.vec_rhs.norm(PETSc.NormType.NORM_2) + if bnorm == 0: + residual_over_b = float('inf') + else: + residual_over_b = true_norm / bnorm + if iteration == 0: + logger.debug(tab*tab_level + f'Residual norms for {ksp.prefix} solve') + logger.debug( + tab*(tab_level - 1) + + f'{iteration: 5d} KSP {_enum_to_norm[ksp.norm_type]} resid norm {residual_norm:14.12e}' + + f' true resid norm {true_norm:14.12e}' + + f' ||r(i)||/||b|| {residual_over_b:14.12e}' + ) + + +def _wrap_method(obj, method_str, ksp_str, monitor): + ''' + Used to patch the method with name `method_str` of the object `obj`, + by setting the monitor of the solver with name `ksp_str` to `monitor`. + + Intended use: + ``` + foo.initialize = _wraps_initialize( + context + "initialize", + "my_ksp", + my_custom_monitor + ) + ``` + + If this is confusing, do not try and call this function! + ''' + old_init = getattr(obj, method_str) + + def new_init(pc): + old_init(pc) + getattr(obj, ksp_str).setMonitor(monitor) + return new_init + + +def attach_custom_monitor(context, monitor): + if isinstance(context, scpc.SCPC): + context.initialize = _wrap_method( + context, + "initialize", + "condensed_ksp", + monitor + ) + elif isinstance(context, hybridization.HybridizationPC): + context.initialize = _wrap_method( + context, + "initialize", + "trace_ksp", + monitor + ) diff --git a/gusto/physics.py b/gusto/physics.py index 203dc9a55..9ababeb4e 100644 --- a/gusto/physics.py +++ b/gusto/physics.py @@ -13,7 +13,7 @@ from gusto.equations import CompressibleEulerEquations from gusto.fml import identity, Term, subject from gusto.labels import physics, transporting_velocity -from gusto.configuration import logger +from gusto.logging import logger from firedrake import (Interpolator, conditional, Function, dx, min_value, max_value, Constant, pi, Projector) from gusto import thermodynamics diff --git a/gusto/time_discretisation.py b/gusto/time_discretisation.py index 54396e527..98e8d53f5 100644 --- a/gusto/time_discretisation.py +++ b/gusto/time_discretisation.py @@ -6,14 +6,20 @@ """ from abc import ABCMeta, abstractmethod, abstractproperty -from firedrake import (Function, TestFunction, NonlinearVariationalProblem, - NonlinearVariationalSolver, DirichletBC) + +from firedrake import ( + Function, TestFunction, NonlinearVariationalProblem, + NonlinearVariationalSolver, DirichletBC +) from firedrake.formmanipulation import split_form from firedrake.utils import cached_property -from gusto.configuration import (logger, DEBUG, EmbeddedDGOptions, RecoveryOptions) + +from gusto.configuration import EmbeddedDGOptions, RecoveryOptions +from gusto.fml import ( + replace_subject, replace_test_function, Term, all_terms, drop +) from gusto.labels import time_derivative, prognostic, physics -from gusto.fml import (replace_subject, replace_test_function, Term, - all_terms, drop) +from gusto.logging import logger, DEBUG, logging_ksp_monitor_true_residual from gusto.wrappers import * @@ -92,8 +98,6 @@ def __init__(self, domain, field_name=None, solver_parameters=None, 'sub_pc_type': 'ilu'} else: self.solver_parameters = solver_parameters - if logger.isEnabledFor(DEBUG): - self.solver_parameters["ksp_monitor_true_residual"] = None def setup(self, equation, apply_bcs=True, *active_labels): """ @@ -209,7 +213,14 @@ def solver(self): # setup solver using lhs and rhs defined in derived class problem = NonlinearVariationalProblem(self.lhs-self.rhs, self.x_out, bcs=self.bcs) solver_name = self.field_name+self.__class__.__name__ - return NonlinearVariationalSolver(problem, solver_parameters=self.solver_parameters, options_prefix=solver_name) + solver = NonlinearVariationalSolver( + problem, + solver_parameters=self.solver_parameters, + options_prefix=solver_name + ) + if logger.isEnabledFor(DEBUG): + solver.snes.ksp.setMonitor(logging_ksp_monitor_true_residual) + return solver @abstractmethod def apply(self, x_out, x_in): diff --git a/gusto/timeloop.py b/gusto/timeloop.py index d1f1f2664..3f73913b1 100644 --- a/gusto/timeloop.py +++ b/gusto/timeloop.py @@ -3,14 +3,16 @@ from abc import ABCMeta, abstractmethod, abstractproperty from firedrake import Function, Projector, Constant, split from pyop2.profiling import timed_stage -from gusto.configuration import logger from gusto.equations import PrognosticEquationSet +from gusto.fml import drop, Label, Term +from gusto.fields import TimeLevelFields, StateFields from gusto.forcing import Forcing -from gusto.fml.form_manipulation_language import drop, Label, Term -from gusto.labels import (transport, diffusion, time_derivative, linearisation, - prognostic, physics, transporting_velocity) +from gusto.labels import ( + transport, diffusion, time_derivative, linearisation, prognostic, + physics, transporting_velocity +) from gusto.linear_solvers import LinearTimesteppingSolver -from gusto.fields import TimeLevelFields, StateFields +from gusto.logging import logger from gusto.time_discretisation import ExplicitTimeDiscretisation from gusto.transport_methods import TransportMethod import ufl diff --git a/gusto/transport_methods.py b/gusto/transport_methods.py index 5e21c4964..30e6fa974 100644 --- a/gusto/transport_methods.py +++ b/gusto/transport_methods.py @@ -5,9 +5,10 @@ from firedrake import (dx, dS, dS_v, dS_h, ds_t, ds_b, ds_v, dot, inner, outer, jump, grad, div, FacetNormal, Function, sign, avg, cross, curl) -from gusto.configuration import IntegrateByParts, TransportEquationType, logger +from gusto.configuration import IntegrateByParts, TransportEquationType from gusto.fml import Term, keep, drop from gusto.labels import prognostic, transport, transporting_velocity, ibp_label +from gusto.logging import logger from gusto.spatial_methods import SpatialMethod __all__ = ["DefaultTransport", "DGUpwind"] diff --git a/integration-tests/equations/test_linear_sw_wave.py b/integration-tests/equations/test_linear_sw_wave.py index a19897aa9..47b03e347 100644 --- a/integration-tests/equations/test_linear_sw_wave.py +++ b/integration-tests/equations/test_linear_sw_wave.py @@ -33,9 +33,10 @@ def run_linear_sw_wave(tmpdir): eqns = LinearShallowWaterEquations(domain, parameters, fexpr=fexpr) # I/O - output = OutputParameters(dirname=str(tmpdir)+"/linear_sw_wave", - dumpfreq=1, - log_level='INFO') + output = OutputParameters( + dirname=str(tmpdir)+"/linear_sw_wave", + dumpfreq=1, + ) io = IO(domain, output) transport_methods = [DefaultTransport(eqns, "D")] diff --git a/integration-tests/equations/test_sw_fplane.py b/integration-tests/equations/test_sw_fplane.py index 3dba7da41..ccb78ac75 100644 --- a/integration-tests/equations/test_sw_fplane.py +++ b/integration-tests/equations/test_sw_fplane.py @@ -29,9 +29,10 @@ def run_sw_fplane(tmpdir): eqns = ShallowWaterEquations(domain, parameters, fexpr=fexpr) # I/O - output = OutputParameters(dirname=str(tmpdir)+"/sw_fplane", - dumpfreq=1, - log_level='INFO') + output = OutputParameters( + dirname=str(tmpdir)+"/sw_fplane", + dumpfreq=1, + ) io = IO(domain, output, diagnostic_fields=[CourantNumber()]) diff --git a/integration-tests/model/test_checkpointing.py b/integration-tests/model/test_checkpointing.py index be7c1eda8..f8a39b69c 100644 --- a/integration-tests/model/test_checkpointing.py +++ b/integration-tests/model/test_checkpointing.py @@ -113,12 +113,18 @@ def test_checkpointing(tmpdir, stepper_type, checkpoint_method): dirname_2 = str(tmpdir)+'/checkpointing_2' dirname_3 = str(tmpdir)+'/checkpointing_3' - output_1 = OutputParameters(dirname=dirname_1, dumpfreq=1, - checkpoint_method=checkpoint_method, - chkptfreq=4, log_level='INFO') - output_2 = OutputParameters(dirname=dirname_2, dumpfreq=1, - checkpoint_method=checkpoint_method, - chkptfreq=2, log_level='INFO') + output_1 = OutputParameters( + dirname=dirname_1, + dumpfreq=1, + checkpoint_method=checkpoint_method, + chkptfreq=4, + ) + output_2 = OutputParameters( + dirname=dirname_2, + dumpfreq=1, + checkpoint_method=checkpoint_method, + chkptfreq=2, + ) stepper_1, eqns_1 = set_up_model_objects(mesh, dt, output_1, stepper_type) stepper_2, eqns_2 = set_up_model_objects(mesh, dt, output_2, stepper_type) @@ -144,10 +150,13 @@ def test_checkpointing(tmpdir, stepper_type, checkpoint_method): chkpt_filename = 'chkpt' if checkpoint_method == 'dumbcheckpoint' else 'chkpt.h5' chkpt_2_path = path.join(stepper_2.io.dumpdir, chkpt_filename) - output_3 = OutputParameters(dirname=dirname_3, dumpfreq=1, - chkptfreq=2, log_level='INFO', - checkpoint_method=checkpoint_method, - checkpoint_pickup_filename=chkpt_2_path) + output_3 = OutputParameters( + dirname=dirname_3, + dumpfreq=1, + chkptfreq=2, + checkpoint_method=checkpoint_method, + checkpoint_pickup_filename=chkpt_2_path, + ) if checkpoint_method == 'checkpointfile': mesh = pick_up_mesh(output_3, mesh_name) @@ -180,10 +189,13 @@ def test_checkpointing(tmpdir, stepper_type, checkpoint_method): # Run *new* timestepper for 2 time steps # ------------------------------------------------------------------------ # - output_3 = OutputParameters(dirname=dirname_3, dumpfreq=1, - chkptfreq=2, log_level='INFO', - checkpoint_method=checkpoint_method, - checkpoint_pickup_filename=chkpt_2_path) + output_3 = OutputParameters( + dirname=dirname_3, + dumpfreq=1, + chkptfreq=2, + checkpoint_method=checkpoint_method, + checkpoint_pickup_filename=chkpt_2_path + ) if checkpoint_method == 'checkpointfile': mesh = pick_up_mesh(output_3, mesh_name) stepper_3, _ = set_up_model_objects(mesh, dt, output_3, stepper_type) diff --git a/integration-tests/model/test_nc_outputting.py b/integration-tests/model/test_nc_outputting.py index ca9519a6c..ef704cacf 100644 --- a/integration-tests/model/test_nc_outputting.py +++ b/integration-tests/model/test_nc_outputting.py @@ -73,8 +73,13 @@ def test_nc_outputting(tmpdir, geometry, domain_and_mesh_details): eqn = AdvectionEquation(domain, V, 'f') transport_scheme = ForwardEuler(domain) transport_method = DGUpwind(eqn, 'f') - output = OutputParameters(dirname=dirname, dumpfreq=1, dump_nc=True, - dumplist=['f'], log_level='INFO', checkpoint=False) + output = OutputParameters( + dirname=dirname, + dumpfreq=1, + dump_nc=True, + dumplist=['f'], + checkpoint=False + ) # Make velocity components for this geometry if geometry == "interval":