Commit 7e32a62c authored by Dion Häfner's avatar Dion Häfner

first draft of plugin system

parent c81a69dd
......@@ -6,10 +6,10 @@ from .io_tools import hdf5 as h5tools
@veros_method
def create_diagnostics(vs):
def create_default_diagnostics(vs):
return {Diag.name: Diag(vs) for Diag in (averages.Averages, cfl_monitor.CFLMonitor,
energy.Energy, overturning.Overturning,
snapshot.Snapshot, tracer_monitor.TracerMonitor)}
energy.Energy, overturning.Overturning,
snapshot.Snapshot, tracer_monitor.TracerMonitor)}
@veros_method
......
from collections import namedtuple
from .variables import Variable
from .settings import Setting
from .diagnostics.diagnostic import VerosDiagnostic
VerosPlugin = namedtuple('VerosPlugin', [
'name',
'module',
'setup_entrypoint',
'run_entrypoint',
'settings',
'variables',
'conditional_variables',
'diagnostics',
])
def load_plugin(module):
if not hasattr(module, '__VEROS_INTERFACE__'):
raise RuntimeError('module {} is not a valid Veros plugin'.format(module.__name__))
interface = module.__VEROS_INTERFACE__
setup_entrypoint = interface.get('setup_entrypoint')
if not callable(setup_entrypoint):
raise RuntimeError('module {} is missing a valid setup entrypoint'.format(module.__name__))
run_entrypoint = interface.get('run_entrypoint')
if not callable(run_entrypoint):
raise RuntimeError('module {} is missing a valid run entrypoint'.format(module.__name__))
name = interface.get('name', module.__name__)
settings = interface.get('settings', [])
for setting, val in settings.items():
if not isinstance(val, Setting):
raise TypeError('got unexpected type {} for setting {}'.format(type(val), setting))
variables = interface.get('variables', [])
for variable, val in variables.items():
if not isinstance(val, Variable):
raise TypeError('got unexpected type {} for variable {}'.format(type(val), variable))
conditional_variables = interface.get('conditional_variables', [])
for _, sub_variables in conditional_variables.items():
for variable, val in sub_variables.items():
if not isinstance(val, Variable):
raise TypeError('got unexpected type {} for variable {}'.format(type(val), variable))
diagnostics = interface.get('diagnostics', [])
for diagnostic in diagnostics:
if not issubclass(diagnostic, VerosDiagnostic):
raise TypeError('got unexpected type {} for diagnostic {}'.format(type(diagnostic), diagnostic))
return VerosPlugin(
name=name,
module=module,
setup_entrypoint=setup_entrypoint,
run_entrypoint=run_entrypoint,
settings=settings,
variables=variables,
conditional_variables=conditional_variables,
diagnostics=diagnostics
)
......@@ -133,7 +133,11 @@ SETTINGS = OrderedDict([
def set_default_settings(vs):
for key, setting in SETTINGS.items():
update_settings(vs, SETTINGS)
def update_settings(vs, settings):
for key, setting in settings.items():
setattr(vs, key, setting.type(setting.default))
......
......@@ -9,6 +9,8 @@ from veros import VerosSetup, veros_method, distributed
from veros.variables import Variable
import veros.tools
import veros_bgc
BASE_PATH = os.path.dirname(os.path.realpath(__file__))
DATA_FILES = veros.tools.get_assets(
'global_4deg',
......@@ -19,6 +21,8 @@ DATA_FILES = veros.tools.get_assets(
class GlobalFourDegreeBGC(VerosSetup):
"""Global 4 degree model with 15 vertical levels and biogeochemistry.
"""
__veros_plugins__ = (veros_bgc,)
@veros_method
def set_parameter(self, vs):
vs.identifier = '4deg'
......@@ -83,23 +87,23 @@ class GlobalFourDegreeBGC(VerosSetup):
vs.enable_kappaH_profile = True
# eke
self.K_gm_0 = 1000.0
self.enable_eke = False
self.eke_k_max = 1e4
self.eke_c_k = 0.4
self.eke_c_eps = 0.5
self.eke_cross = 2.
self.eke_crhin = 1.0
self.eke_lmin = 100.0
self.enable_eke_superbee_advection = False
self.enable_eke_isopycnal_diffusion = False
vs.K_gm_0 = 1000.0
vs.enable_eke = False
vs.eke_k_max = 1e4
vs.eke_c_k = 0.4
vs.eke_c_eps = 0.5
vs.eke_cross = 2.
vs.eke_crhin = 1.0
vs.eke_lmin = 100.0
vs.enable_eke_superbee_advection = False
vs.enable_eke_isopycnal_diffusion = False
# idemix
self.enable_idemix = False
self.enable_idemix_hor_diffusion = False
self.enable_eke_diss_surfbot = False
self.eke_diss_surfbot_frac = 0.2 # fraction which goes into bottom
self.enable_idemix_superbee_advection = False
vs.enable_idemix = False
vs.enable_idemix_hor_diffusion = False
vs.enable_eke_diss_surfbot = False
vs.eke_diss_surfbot_frac = 0.2 # fraction which goes into bottom
vs.enable_idemix_superbee_advection = False
vs.eq_of_state_type = 5
......
import math
from . import variables, settings
from . import variables, settings, timer, plugins, diagnostics
class VerosState:
......@@ -14,21 +14,50 @@ class VerosState:
rho_0 = 1024. # Boussinesq reference density in :math:`kg/m^3`
grav = 9.81 # Gravitational constant in :math:`m/s^2`
def __init__(self):
def __init__(self, use_plugins=None):
self.variables = {}
self.diagnostics = {}
self.poisson_solver = None
self.nisle = 0 # to be overriden during streamfunction_init
self.taum1, self.tau, self.taup1 = 0, 1, 2 # pointers to last, current, and next time step
self.time, self.itt = 0., 0 # current time and iteration
if use_plugins is not None:
self._plugin_interfaces = tuple(plugins.load_plugin(p) for p in use_plugins)
else:
self._plugin_interfaces = tuple()
settings.set_default_settings(self)
for plugin in self._plugin_interfaces:
settings.update_settings(self, plugin.settings)
self.timers = {k: timer.Timer() for k in (
'setup', 'main', 'momentum', 'temperature', 'eke', 'idemix',
'tke', 'diagnostics', 'pressure', 'friction', 'isoneutral',
'vmix', 'eq_of_state', 'plugins'
)}
for plugin in self._plugin_interfaces:
self.timers[plugin.name] = timer.Timer()
def allocate_variables(self):
self.variables.update(variables.get_standard_variables(self))
for plugin in self._plugin_interfaces:
plugin_vars = variables.get_active_variables(self, plugin.variables, plugin.conditional_variables)
self.variables.update(plugin_vars)
for key, var in self.variables.items():
setattr(self, key, variables.allocate(self, var.dims, dtype=var.dtype))
def create_diagnostics(self):
self.diagnostics.update(diagnostics.create_default_diagnostics(self))
for plugin in self._plugin_interfaces:
for diagnostic in plugin.diagnostics:
self.diagnostics[diagnostic.name] = diagnostic(self)
def to_xarray(self):
import xarray as xr
......
......@@ -3,7 +3,7 @@ from loguru import logger
from .state import VerosState
class DistributedVerosState(VerosState):
class DistributedVerosState:
"""A proxy wrapper to temporarily synchronize a distributed state.
Use `gather_arrays` to retrieve distributed variables from parent VerosState object,
......@@ -65,3 +65,6 @@ class DistributedVerosState(VerosState):
return self._vs.__setattr__(attr, val)
raise AttributeError('Cannot access distributed variable %s since it was not retrieved' % attr)
def __repr__(self):
return '{}(parent_state={})'.format(self.__class__.__name__, repr(self._vs))
......@@ -2,8 +2,7 @@ import timeit
class Timer:
def __init__(self, name):
self.name = name
def __init__(self):
self.total_time = 0
self.last_time = 0
......@@ -24,9 +23,6 @@ class Timer:
self.last_time = timeit.default_timer() - self.start_time
self.total_time += self.last_time
def print_time(self):
print('[{}]: {}s'.format(self.name, self.get_time()))
def get_time(self):
return self.total_time
......
......@@ -728,24 +728,35 @@ CONDITIONAL_VARIABLES = OrderedDict([
@veros_method
def get_standard_variables(vs):
def get_active_variables(vs, main_variables=None, conditional_variables=None):
variables = {}
for var_name, var in MAIN_VARIABLES.items():
variables[var_name] = var
if main_variables is not None:
for var_name, var in main_variables.items():
variables[var_name] = var
for condition, var_dict in CONDITIONAL_VARIABLES.items():
if condition.startswith('not '):
eval_condition = not bool(getattr(vs, condition[4:]))
else:
eval_condition = bool(getattr(vs, condition))
if eval_condition:
for var_name, var in var_dict.items():
variables[var_name] = var
if conditional_variables is not None:
for condition, var_dict in conditional_variables.items():
if condition.startswith('not '):
eval_condition = not bool(getattr(vs, condition[4:]))
else:
eval_condition = bool(getattr(vs, condition))
if eval_condition:
for var_name, var in var_dict.items():
variables[var_name] = var
return variables
@veros_method
def get_standard_variables(vs):
return get_active_variables(
vs,
main_variables=MAIN_VARIABLES,
conditional_variables=CONDITIONAL_VARIABLES
)
@veros_method(inline=True)
def allocate(vs, dimensions, dtype=None, include_ghosts=True, local=True, fill=0):
if dtype is None:
......
......@@ -8,7 +8,7 @@ from veros import (
runtime_settings as rs, runtime_state as rst
)
from veros.state import VerosState
from veros.timer import Timer
from veros.plugins import load_plugin
from veros.core import (
momentum, numerics, thermodynamics, eke, tke, idemix,
isoneutral, streamfunction, advection, utilities
......@@ -46,19 +46,28 @@ class VerosSetup(metaclass=abc.ABCMeta):
>>> plt.show()
"""
__veros_plugins__ = tuple()
def __init__(self, state=None, override=None):
def __init__(self, state=None, override=None, plugins=None):
self.override_settings = override or {}
logs.setup_logging(loglevel=rs.loglevel)
if plugins is not None:
self.__veros_plugins__ = tuple(plugins)
self._plugin_interfaces = tuple(load_plugin(p) for p in self.__veros_plugins__)
if state is None:
self.state = VerosState()
self.state = VerosState(use_plugins=self.__veros_plugins__)
self.state.timers = {k: Timer(k) for k in (
'setup', 'main', 'momentum', 'temperature', 'eke', 'idemix',
'tke', 'diagnostics', 'pressure', 'friction', 'isoneutral',
'vmix', 'eq_of_state'
)}
this_plugins = set(p.module for p in self.state._plugin_interfaces)
state_plugins = set(p.module for p in self._plugin_interfaces)
if this_plugins != state_plugins:
raise ValueError(
'VerosState was created with plugin modules {}, but this setup uses {}'
.format(state_plugins, this_plugins)
)
@abc.abstractmethod
def set_parameter(self, vs):
......@@ -201,7 +210,10 @@ class VerosSetup(metaclass=abc.ABCMeta):
streamfunction.streamfunction_init(vs)
eke.init_eke(vs)
vs.diagnostics = diagnostics.create_diagnostics(vs)
for plugin in self._plugin_interfaces:
plugin.setup_entrypoint(vs)
vs.create_diagnostics()
self.set_diagnostics(vs)
diagnostics.initialize(vs)
diagnostics.read_restart(vs)
......@@ -284,6 +296,11 @@ class VerosSetup(metaclass=abc.ABCMeta):
momentum.vertical_velocity(vs)
with vs.timers['plugins']:
for plugin in self._plugin_interfaces:
with vs.timers[plugin.name]:
plugin.run_entrypoint(vs)
vs.itt += 1
vs.time += vs.dt_tracer
pbar.advance_time(vs.dt_tracer)
......@@ -316,7 +333,7 @@ class VerosSetup(metaclass=abc.ABCMeta):
finally:
diagnostics.write_restart(vs, force=True)
logger.debug('\n'.join([
timing_summary = [
'',
'Timing summary:',
' setup time = {:.2f}s'.format(vs.timers['setup'].get_time()),
......@@ -332,7 +349,15 @@ class VerosSetup(metaclass=abc.ABCMeta):
' IDEMIX = {:.2f}s'.format(vs.timers['idemix'].get_time()),
' TKE = {:.2f}s'.format(vs.timers['tke'].get_time()),
' diagnostics and I/O = {:.2f}s'.format(vs.timers['diagnostics'].get_time()),
]))
' plugins = {:.2f}s'.format(vs.timers['plugins'].get_time()),
]
timing_summary.extend([
' {:<22} = {:.2f}s'.format(plugin.name, vs.timers[plugin.name].get_time())
for plugin in vs._plugin_interfaces
])
logger.debug('\n'.join(timing_summary))
if profiler is not None:
diagnostics.stop_profiler(profiler)
Markdown is supported
0%
or
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment