From 2cae8a5da8ac1ddacd61cd86f832abd72c381a6d Mon Sep 17 00:00:00 2001 From: Austin Raney Date: Fri, 9 Aug 2024 15:32:14 -0400 Subject: [PATCH 01/11] feat: add ngen_cal_model_configure hook --- python/ngen_cal/src/ngen/cal/_hookspec.py | 14 ++++++++++++++ 1 file changed, 14 insertions(+) diff --git a/python/ngen_cal/src/ngen/cal/_hookspec.py b/python/ngen_cal/src/ngen/cal/_hookspec.py index ef1e5bf6..dbda5d88 100644 --- a/python/ngen_cal/src/ngen/cal/_hookspec.py +++ b/python/ngen_cal/src/ngen/cal/_hookspec.py @@ -13,6 +13,7 @@ from hypy.nexus import Nexus from ngen.cal.configuration import General + from ngen.cal.model import ModelExec from ngen.cal.meta import JobMeta hookspec = pluggy.HookspecMarker(PROJECT_SLUG) @@ -51,6 +52,19 @@ def ngen_cal_finish(exception: Exception | None) -> None: class ModelHooks: + @hookspec + def ngen_cal_model_configure(self, config: ModelExec) -> None: + """ + Called before calibration begins. + This allow plugins to perform initial configuration. + + Plugins' configuration data should be provided using the + `plugins_settings` field in the `model` section of an `ngen.cal` + configuration file. + By convention, the name of the plugin should be used as top level key in + the `plugin_settings` dictionary. + """ + @hookspec(firstresult=True) def ngen_cal_model_observations( self, From 17badd54b9675418bc03354711efc37dbe491fec Mon Sep 17 00:00:00 2001 From: Austin Raney Date: Fri, 9 Aug 2024 15:32:45 -0400 Subject: [PATCH 02/11] test: ngen_cal_model_configure --- python/ngen_cal/tests/test_plugin_system.py | 6 ++++++ 1 file changed, 6 insertions(+) diff --git a/python/ngen_cal/tests/test_plugin_system.py b/python/ngen_cal/tests/test_plugin_system.py index b76e9baf..36401b2e 100644 --- a/python/ngen_cal/tests/test_plugin_system.py +++ b/python/ngen_cal/tests/test_plugin_system.py @@ -15,6 +15,8 @@ from hypy.nexus import Nexus from ngen.cal.configuration import General + from ngen.cal.model import ModelExec + from pathlib import Path def test_setup_plugin_manager(): @@ -74,6 +76,10 @@ def ngen_cal_start(self) -> None: def ngen_cal_finish(self) -> None: """Called after exiting the calibration loop.""" + @hookimpl + def ngen_cal_model_configure(self, config: ModelExec) -> None: + """Test model model configure plugin""" + @hookimpl def ngen_cal_model_output(self) -> None: """Test model output plugin""" From 5e3ad147a7d72b76b7caf760d7aa594b729534e7 Mon Sep 17 00:00:00 2001 From: Austin Raney Date: Fri, 9 Aug 2024 16:37:22 -0400 Subject: [PATCH 03/11] refactor: Agent dep injects a Model instead of taking a raw config --- python/ngen_cal/src/ngen/cal/agent.py | 19 ++++++++++--------- 1 file changed, 10 insertions(+), 9 deletions(-) diff --git a/python/ngen_cal/src/ngen/cal/agent.py b/python/ngen_cal/src/ngen/cal/agent.py index 4f124c8d..fa529c27 100644 --- a/python/ngen_cal/src/ngen/cal/agent.py +++ b/python/ngen_cal/src/ngen/cal/agent.py @@ -2,7 +2,7 @@ from abc import ABC, abstractmethod from ngen.cal.meta import JobMeta -from ngen.cal.configuration import Model +from ngen.cal.configuration import Model, NoModel from ngen.cal.utils import pushd from pathlib import Path from typing import TYPE_CHECKING @@ -59,9 +59,13 @@ def best_params(self) -> str: class Agent(BaseAgent): - def __init__(self, model_conf, workdir: Path, log: bool=False, restart: bool=False, parameters: Mapping[str, Any] | None = {}): + def __init__(self, model: Model, workdir: 'Path', log: bool=False, restart: bool=False, parameters: 'Optional[Mapping[str, Any]]' = {}): self._workdir = workdir self._job = None + assert not isinstance(model.model, NoModel), "invariant" + # NOTE: if support for new models is added, this will need to be modified + model_inner = model.model.unwrap() + self._model = model if restart: # find prior ngen workdirs # FIXME if a user starts with an independent calibration strategy @@ -72,17 +76,14 @@ def __init__(self, model_conf, workdir: Path, log: bool=False, restart: bool=Fal # 0 correctly since not all basin params can be loaded. # There are probably some similar issues with explicit and independent, since they have # similar data semantics - workdirs = list(Path.glob(workdir, model_conf['type']+"_*_worker")) + workdirs = list(Path.glob(workdir, model_inner.type+"_*_worker")) if len(workdirs) > 1: print("More than one existing workdir, cannot restart") elif len(workdirs) == 1: - self._job = JobMeta(model_conf['type'], workdir, workdirs[0], log=log) + self._job = JobMeta(model_inner.type, workdir, workdirs[0], log=log) if self._job is None: - self._job = JobMeta(model_conf['type'], workdir, log=log) - resolved_binary = Path(model_conf['binary']).resolve() - model_conf['workdir'] = self.job.workdir - self._model = Model(model=model_conf, binary=resolved_binary) + self._job = JobMeta(model_inner.type, workdir, log=log) self._model.model.resolve_paths(self.job.workdir) self._params = parameters @@ -117,4 +118,4 @@ def duplicate(self) -> Agent: data = self.model.__root__.copy(deep=True) #return a new agent, which has a unique Model instance #and its own Job/workspace - return Agent(data.dict(by_alias=True), self._workdir) + return Agent(data, self._workdir) From c101bf075873d92da5657929197b0f36871d72e0 Mon Sep 17 00:00:00 2001 From: Austin Raney Date: Fri, 9 Aug 2024 16:38:03 -0400 Subject: [PATCH 04/11] feat: Ngen.unwrap convenience method. returns __root__ --- python/ngen_cal/src/ngen/cal/ngen.py | 5 +++++ 1 file changed, 5 insertions(+) diff --git a/python/ngen_cal/src/ngen/cal/ngen.py b/python/ngen_cal/src/ngen/cal/ngen.py index 27c7db13..bf2e1ea3 100644 --- a/python/ngen_cal/src/ngen/cal/ngen.py +++ b/python/ngen_cal/src/ngen/cal/ngen.py @@ -595,6 +595,11 @@ def get_binary(self) -> str: return self.__root__.get_binary() def update_config(self, *args, **kwargs): return self.__root__.update_config(*args, **kwargs) + + def unwrap(self) -> NgenBase: + """convenience method that returns the underlying __root__ instance""" + return self.__root__ + #proxy methods for model @property def adjustables(self): From b1a2cea350ed0916be6f3ce9c4f94b55a024a9a3 Mon Sep 17 00:00:00 2001 From: Austin Raney Date: Fri, 9 Aug 2024 16:38:55 -0400 Subject: [PATCH 05/11] refactor: construct Model and inject into Agent --- python/ngen_cal/src/ngen/cal/__main__.py | 15 ++++++++++++--- 1 file changed, 12 insertions(+), 3 deletions(-) diff --git a/python/ngen_cal/src/ngen/cal/__main__.py b/python/ngen_cal/src/ngen/cal/__main__.py index 1d369f3c..67634251 100644 --- a/python/ngen_cal/src/ngen/cal/__main__.py +++ b/python/ngen_cal/src/ngen/cal/__main__.py @@ -4,7 +4,8 @@ import yaml from os import chdir from pathlib import Path -from ngen.cal.configuration import General +from ngen.cal.configuration import General, Model +from ngen.cal.ngen import Ngen from ngen.cal.search import dds, dds_set, pso_search from ngen.cal.strategy import Algorithm from ngen.cal.agent import Agent @@ -42,6 +43,13 @@ def main(general: General, model_conf: Mapping[str, Any]): import numpy as np np.random.seed(general.random_seed) + # model scope plugins setup in constructor + model = Model(model=model_conf) + + # NOTE: if support for new models is added, this will need to be modified + assert isinstance(model.model, Ngen), f"ngen.cal.ngen.Ngen expected, got {type(model.model)}" + model_inner = model.model.unwrap() + plugins = cast(List[Union[Callable, ModuleType]], general.plugins) plugin_manager = setup_plugin_manager(plugins) @@ -57,8 +65,9 @@ def main(general: General, model_conf: Mapping[str, Any]): into a single variable vector and calibrating a set of heterogenous formultions... """ start_iteration = 0 - #Initialize the starting agent - agent = Agent(model_conf, general.workdir, general.log, general.restart, general.strategy.parameters) + + # Initialize the starting agent + agent = Agent(model, general.workdir, general.log, general.restart, general.strategy.parameters) if general.strategy.algorithm == Algorithm.dds: func = dds_set #FIXME what about explicit/dds start_iteration = general.start_iteration From 56e96dbb2c9b13600a08073e68ba6ce01f6e438f Mon Sep 17 00:00:00 2001 From: Austin Raney Date: Fri, 9 Aug 2024 16:39:19 -0400 Subject: [PATCH 06/11] feat: call ngen_cal_model_configure plugin --- python/ngen_cal/src/ngen/cal/__main__.py | 1 + 1 file changed, 1 insertion(+) diff --git a/python/ngen_cal/src/ngen/cal/__main__.py b/python/ngen_cal/src/ngen/cal/__main__.py index 67634251..7355cafa 100644 --- a/python/ngen_cal/src/ngen/cal/__main__.py +++ b/python/ngen_cal/src/ngen/cal/__main__.py @@ -57,6 +57,7 @@ def main(general: General, model_conf: Mapping[str, Any]): # setup plugins plugin_manager.hook.ngen_cal_configure(config=general) + model_inner._plugin_manager.hook.ngen_cal_model_configure(config=model_inner) print("Starting calib") From 88f5ae3127f4ba3596a32acbb8d051910a783efa Mon Sep 17 00:00:00 2001 From: Austin Raney Date: Fri, 9 Aug 2024 17:00:27 -0400 Subject: [PATCH 07/11] refactor(test): inject Model instance --- python/ngen_cal/tests/conftest.py | 9 +++++---- 1 file changed, 5 insertions(+), 4 deletions(-) diff --git a/python/ngen_cal/tests/conftest.py b/python/ngen_cal/tests/conftest.py index 64881ce9..678b1cf8 100644 --- a/python/ngen_cal/tests/conftest.py +++ b/python/ngen_cal/tests/conftest.py @@ -5,7 +5,7 @@ import json import pandas as pd # type: ignore import geopandas as gpd # type: ignore -from ngen.cal.configuration import General +from ngen.cal.configuration import General, Model from ngen.cal.ngen import Ngen from ngen.cal.meta import JobMeta from ngen.cal.calibration_cathment import CalibrationCatchment @@ -104,9 +104,10 @@ def meta(ngen_config, general_config, mocker) -> Generator[JobMeta, None, None]: yield m @pytest.fixture -def agent(ngen_config, general_config) -> Generator['Agent', None, None]: - a = Agent(ngen_config.__root__.dict(), general_config.workdir, general_config.log) - yield a +def agent(ngen_config, general_config) -> Agent: + model = Model(model=ngen_config) + a = Agent(model, general_config.workdir, general_config.log) + return a @pytest.fixture def eval(ngen_config) -> Generator[EvaluationOptions, None, None]: From 29c3528602d221e2f2b0f33f2073f7c671665ba2 Mon Sep 17 00:00:00 2001 From: Austin Raney Date: Thu, 15 Aug 2024 13:30:33 -0400 Subject: [PATCH 08/11] chore: update comments and improve naming --- python/ngen_cal/src/ngen/cal/agent.py | 11 ++++++----- 1 file changed, 6 insertions(+), 5 deletions(-) diff --git a/python/ngen_cal/src/ngen/cal/agent.py b/python/ngen_cal/src/ngen/cal/agent.py index fa529c27..4f09cf3a 100644 --- a/python/ngen_cal/src/ngen/cal/agent.py +++ b/python/ngen_cal/src/ngen/cal/agent.py @@ -63,8 +63,9 @@ def __init__(self, model: Model, workdir: 'Path', log: bool=False, restart: bool self._workdir = workdir self._job = None assert not isinstance(model.model, NoModel), "invariant" - # NOTE: if support for new models is added, this will need to be modified - model_inner = model.model.unwrap() + # NOTE: if support for new models is added, support for other model + # type variants will be required + ngen_model = model.model.unwrap() self._model = model if restart: # find prior ngen workdirs @@ -76,14 +77,14 @@ def __init__(self, model: Model, workdir: 'Path', log: bool=False, restart: bool # 0 correctly since not all basin params can be loaded. # There are probably some similar issues with explicit and independent, since they have # similar data semantics - workdirs = list(Path.glob(workdir, model_inner.type+"_*_worker")) + workdirs = list(Path.glob(workdir, ngen_model.type+"_*_worker")) if len(workdirs) > 1: print("More than one existing workdir, cannot restart") elif len(workdirs) == 1: - self._job = JobMeta(model_inner.type, workdir, workdirs[0], log=log) + self._job = JobMeta(ngen_model.type, workdir, workdirs[0], log=log) if self._job is None: - self._job = JobMeta(model_inner.type, workdir, log=log) + self._job = JobMeta(ngen_model.type, workdir, log=log) self._model.model.resolve_paths(self.job.workdir) self._params = parameters From dd18e0ea3e2ba2756708a6037b5c5dbe77e3b126 Mon Sep 17 00:00:00 2001 From: Austin Raney Date: Thu, 15 Aug 2024 13:49:45 -0400 Subject: [PATCH 09/11] style: type hint --- python/ngen_cal/src/ngen/cal/agent.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/python/ngen_cal/src/ngen/cal/agent.py b/python/ngen_cal/src/ngen/cal/agent.py index 4f09cf3a..34370107 100644 --- a/python/ngen_cal/src/ngen/cal/agent.py +++ b/python/ngen_cal/src/ngen/cal/agent.py @@ -59,7 +59,7 @@ def best_params(self) -> str: class Agent(BaseAgent): - def __init__(self, model: Model, workdir: 'Path', log: bool=False, restart: bool=False, parameters: 'Optional[Mapping[str, Any]]' = {}): + def __init__(self, model: Model, workdir: Path, log: bool=False, restart: bool=False, parameters: Mapping[str, Any] | None = {}): self._workdir = workdir self._job = None assert not isinstance(model.model, NoModel), "invariant" From 813def977f18bfbc6405b4c4cde556e116984e74 Mon Sep 17 00:00:00 2001 From: Austin Raney Date: Thu, 15 Aug 2024 14:38:51 -0400 Subject: [PATCH 10/11] fix(pr-review): call ngen_cal_model_configure after Agent init Agent mutates the model config; If ngen_cal_model_configure is called before the Agent is initialized, it won't get the final view of the config before calibration. Co-authored-by: Nels --- python/ngen_cal/src/ngen/cal/__main__.py | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/python/ngen_cal/src/ngen/cal/__main__.py b/python/ngen_cal/src/ngen/cal/__main__.py index 7355cafa..230951c7 100644 --- a/python/ngen_cal/src/ngen/cal/__main__.py +++ b/python/ngen_cal/src/ngen/cal/__main__.py @@ -57,7 +57,6 @@ def main(general: General, model_conf: Mapping[str, Any]): # setup plugins plugin_manager.hook.ngen_cal_configure(config=general) - model_inner._plugin_manager.hook.ngen_cal_model_configure(config=model_inner) print("Starting calib") @@ -69,6 +68,10 @@ def main(general: General, model_conf: Mapping[str, Any]): # Initialize the starting agent agent = Agent(model, general.workdir, general.log, general.restart, general.strategy.parameters) + + # Agent mutates the model config, so `ngen_cal_model_configure` is called afterwards + model_inner._plugin_manager.hook.ngen_cal_model_configure(config=model_inner) + if general.strategy.algorithm == Algorithm.dds: func = dds_set #FIXME what about explicit/dds start_iteration = general.start_iteration From 94cd2aa7ea0a3e75fd53283a06109c83e3ad451c Mon Sep 17 00:00:00 2001 From: Austin Raney Date: Thu, 15 Aug 2024 14:40:40 -0400 Subject: [PATCH 11/11] fix(pr-review): sync model and job meta workdirs Co-authored-by: Nels --- python/ngen_cal/src/ngen/cal/agent.py | 1 + 1 file changed, 1 insertion(+) diff --git a/python/ngen_cal/src/ngen/cal/agent.py b/python/ngen_cal/src/ngen/cal/agent.py index 34370107..6979094a 100644 --- a/python/ngen_cal/src/ngen/cal/agent.py +++ b/python/ngen_cal/src/ngen/cal/agent.py @@ -85,6 +85,7 @@ def __init__(self, model: Model, workdir: Path, log: bool=False, restart: bool=F if self._job is None: self._job = JobMeta(ngen_model.type, workdir, log=log) + ngen_model.workdir = self.job.workdir self._model.model.resolve_paths(self.job.workdir) self._params = parameters