Runfile timeloop prephysics computations and configuration #1081

brianhenn · 2021-03-11T05:47:55Z

In order to override the coarse model radiative fluxes with predicted values prior to the land surface model call during a prognostic run, this PR implement a prephysics steps in the sklearn_runfile's loop, after dynamics. It allows for configuring this optional prephysics computation by specifying a machine learning model, in the style of a MachineLearningConfig used by the existing python computations after the physics routine (renamed as the postphysics step by this PR.

Added public API:

a PrephysicsConfig specification for prognostic runs, an optional MachineLearningConfig specification

Significant internal changes:

added an MLStateStepper object to runtime.steppers.machine_learning which returns ML-predicted states (as opposed to tendencies)
runtime.loop adds a _prephysics steps to the loop
runtime.loop supports having two Stepper objects, one for the prephysics and one for the previous python computations after physics, now called postphysics
Tests added

brianhenn · 2021-03-12T19:28:11Z

workflows/prognostic_c48_run/runtime/steppers/prephysics.py

+class Prescriber:
+    """A pre-physics stepper which obtains prescribed values from an external source
+
+        TODO: Implement methods


This is a stub which does nothing, for the follow-on PR

Can you remove it then? I think this should be unit tested before merging to master.

brianhenn · 2021-03-12T19:29:07Z

..._regtest_outputs/test_machine_learning.test_MLStepper_regression_checksum[PureMLStepper].out

@@ -14,4 +14,6 @@
      - fcc46bebe36ea131688f8e15700e18d4


These checksums were preserved but the test name was not

brianhenn · 2021-03-12T19:32:00Z

Would also be nice to get the fv3gfs-wrapper PR merged first so that we can point at master here.

spencerkclark

In my reading this looks like a pretty clean approach to solving the multiple stepper issue. I just have a couple comments / suggestions, but I'll let @nbren12 give the final approval.

I agree it would be good if the fv3gfs-wrapper PR was merged; the changes there are pretty trivial but there seems to be an issue with the CI (I recall I was able to get the tests to pass locally).

workflows/prognostic_c48_run/runtime/loop.py

nbren12

Thanks for putting this together. I have a fair number of comments, and am happy to chat over zoom to explain any that don't have clear motivations. Some overall thoughts:

we should wait until fv3gfs-wrapper is merged until master
Revert changes to machine_learning.py and refactor unique "MLStateStepper" logic to either MLStepper or Loop.
Remove Prescriber class and related configs until that feature is functional.
Simplify coupled fixture hierarchy in tests.

nbren12 · 2021-03-12T22:56:39Z

workflows/prognostic_c48_run/runtime/steppers/prephysics.py

+class Prescriber:
+    """A pre-physics stepper which obtains prescribed values from an external source
+
+        TODO: Implement methods


Can you remove it then? I think this should be unit tested before merging to master.

nbren12 · 2021-03-12T22:57:33Z

workflows/prognostic_c48_run/runtime/steppers/prephysics.py

+
+
+@dataclasses.dataclass
+class PrephysicsConfig:


Why make a new object for this? You can use Union directly in UserConfig. It's odd to have a config object with one attribute named "config".

nbren12 · 2021-03-12T23:02:14Z

workflows/prognostic_c48_run/runtime/steppers/prephysics.py

+
+    Attributes:
+        variables: list variable names to prescribe
+        data_source: path to the source of the data to prescribe


Can you mention what kind of data must be at this path? Zarr?

Moot since reverting

nbren12 · 2021-03-12T23:07:06Z

workflows/prognostic_c48_run/prepare_config.py

+    prephysics: Optional[PrephysicsConfig]
+    if "prephysics" in config_dict:
+        prephysics = dacite.from_dict(
+            PrephysicsConfig, {"config": config_dict["prephysics"]}


This renaming breaks down the 1-1 correspondence between the UserConfig attribute names and the yaml representation. This makes the documentation less useful. Again, PrephysicsConfig doesn't seem like it should be a class.

nbren12 · 2021-03-12T23:07:34Z

workflows/prognostic_c48_run/docs/config-usage.rst

@@ -109,6 +109,22 @@ It can be used multiple times to specify multiple models. For example::
        --model_url path/to_another/model
        > fv3config.yaml

+Prephysics


You can document the prephysics config in the docstring of the UserConfig class.

workflows/prognostic_c48_run/runtime/steppers/machine_learning.py

nbren12 · 2021-03-12T23:32:10Z

workflows/prognostic_c48_run/tests/machine_learning_mocks.py

@@ -60,6 +80,33 @@ def get_mock_sklearn_model() -> fv3fit.Predictor:
    return model


+def get_mock_rad_flux_model() -> fv3fit.Predictor:


Looks like a lot of copy paste here. Can we add some parameters to get_mock_sklearn_model?

nbren12 · 2021-03-12T23:39:42Z

workflows/prognostic_c48_run/runtime/loop.py

@@ -191,30 +218,64 @@ def _get_states_to_output(self, config: UserConfig) -> Sequence[str]:
                    states_to_output = diagnostic.variables  # type: ignore
        return states_to_output

-    def _get_stepper(self, config: UserConfig) -> Optional[Stepper]:
+    def _get_steppers(self, config: UserConfig) -> Mapping[str, Optional[Stepper]]:


We heavily rely on mypy to detect errors in this un-testable Loop class, and the old design was more amenable to static analysis . The Mapping[str, Optional[Stepper]] return type will prevent mypy from detecting errors statically (e.g. if "_compute_physics" is not inside this dict). It will also add more indirection for the reader.

Can you revert these changes, and put the prephysics specific code into _get_prephysics_stepper method and store the stepper in self._prephysics_stepper?

nbren12 · 2021-03-12T23:40:56Z

workflows/prognostic_c48_run/runtime/loop.py

@@ -177,7 +204,7 @@ def __init__(

        self._states_to_output: Sequence[str] = self._get_states_to_output(config)
        self._log_debug(f"States to output: {self._states_to_output}")
-        self.stepper = self._get_stepper(config)
+        self.steppers = self._get_steppers(config)


Suggested change

self.steppers = self._get_steppers(config)

self.post_physics_stepper = self._get_stepper(config)

self.pre_physics_stepper = self._get_pre_physics_stepper(config)

nbren12 · 2021-03-12T23:44:54Z

workflows/prognostic_c48_run/tests/machine_learning_mocks.py

@@ -27,6 +27,26 @@ def _model_dataset() -> xr.Dataset:
    return data


+def _rad_model_dataset() -> xr.Dataset:


Copy paste here. Can you add these new variables to _model_dataset and delete this function?

nbren12

Looking good. Thanks. There was one place below I briefly got confused.

nbren12 · 2021-03-15T22:37:14Z

workflows/prognostic_c48_run/runtime/loop.py

+        self._log_info("Downloading ML Model")
+        if self.rank == 0:
+            local_model_paths = download_model(
+                ml_config, os.path.join(step, "ml_model")


I got a little worried seeing this path logic. At first glance it looks like it could be a backwards incompatible change to how the user configuration is interpreted, but really it is just a temporary local cache. Can you note this in a comment or better yet...modify download_model to use tempfile.NamedTemporaryDirectory?

nbren12 · 2021-03-15T22:39:46Z

workflows/prognostic_c48_run/runtime/loop.py

+            self._compute_prephysics,
+            self._apply_prephysics,


Since these functions occur in sequence would it be simpler to group them in one function?

spencerkclark

Thanks for the updates @brianhenn! It seems like we should be able to merge the wrapper PR today too.

…ephysics-steps

spencerkclark

ai2cm/fv3gfs-wrapper#244 has now been merged, so we should be able to merge this PR soon, with a few more minor updates.

spencerkclark · 2021-03-22T13:43:59Z

workflows/prognostic_c48_run/runtime/diagnostics/machine_learning.py

+        "total_sky_downward_shortwave_flux_at_surface_override",
+        "total_sky_net_shortwave_flux_at_surface_override",
+        "total_sky_downward_longwave_flux_at_surface_override",


These have been renamed in the merged version of ai2cm/fv3gfs-wrapper#244:

Suggested change

"total_sky_downward_shortwave_flux_at_surface_override",

"total_sky_net_shortwave_flux_at_surface_override",

"total_sky_downward_longwave_flux_at_surface_override",

"override_for_time_adjusted_total_sky_downward_shortwave_flux_at_surface",

"override_for_time_adjusted_total_sky_net_shortwave_flux_at_surface",

"override_for_time_adjusted_total_sky_downward_longwave_flux_at_surface",

spencerkclark · 2021-03-22T13:44:49Z

workflows/prognostic_c48_run/runtime/loop.py

+            "total_sky_downward_shortwave_flux_at_surface_override",
+            "total_sky_net_shortwave_flux_at_surface_override",
+            "total_sky_downward_longwave_flux_at_surface_override",


Suggested change

"total_sky_downward_shortwave_flux_at_surface_override",

"total_sky_net_shortwave_flux_at_surface_override",

"total_sky_downward_longwave_flux_at_surface_override",

"override_for_time_adjusted_total_sky_downward_shortwave_flux_at_surface",

"override_for_time_adjusted_total_sky_net_shortwave_flux_at_surface",

"override_for_time_adjusted_total_sky_downward_longwave_flux_at_surface",

brianhenn added 11 commits March 10, 2021 05:24

prephysics call, mock prescriber and PrephysicsConfig

ebf22c7

tests passing

e6543eb

updated wrapper and fortran model for radiative flux setting

1914f47

add prephysics ML subclass

381b76f

update to wrapper with albedo

4a0fd33

MLStateStepper tests

c687809

cleanup

427f4c5

Merge branch 'master' into feature/timeloop-prephysics-steps

8fc2498

Merge branch 'master' into feature/timeloop-prephysics-steps

0cabbaa

updated docs

cd9ec8f

updated fortran model external

bf163ad

brianhenn commented Mar 12, 2021

View reviewed changes

brianhenn changed the title ~~Feature/timeloop prephysics steps~~ Runfile timeloop prephysics computations and configuration Mar 12, 2021

spencerkclark reviewed Mar 12, 2021

View reviewed changes

workflows/prognostic_c48_run/runtime/loop.py Outdated Show resolved Hide resolved

workflows/prognostic_c48_run/runtime/loop.py Outdated Show resolved Hide resolved

workflows/prognostic_c48_run/runtime/loop.py Show resolved Hide resolved

nbren12 suggested changes Mar 12, 2021

View reviewed changes

brianhenn added 5 commits March 14, 2021 06:05

revert Prescriber

b6f6b45

prephysics and postphysics steppers

bb19f9e

prephysics diagnostic ml

436e036

modifying tests per PR review

e7be263

cleanup

dcd1e05

nbren12 approved these changes Mar 15, 2021

View reviewed changes

brianhenn added 2 commits March 15, 2021 23:50

Merge branch 'master' into feature/timeloop-prephysics-steps

a9e498e

address addl PR comments

e0e13a7

spencerkclark approved these changes Mar 16, 2021

View reviewed changes

Merge remote-tracking branch 'origin/master' into feature/timeloop-pr…

ac21fb3

…ephysics-steps

nbren12 enabled auto-merge (squash) March 19, 2021 19:49

nbren12 disabled auto-merge March 19, 2021 19:51

spencerkclark reviewed Mar 22, 2021

View reviewed changes

brianhenn added 3 commits March 23, 2021 20:25

update to use master wrapper

73a80d0

Merge branch 'master' into feature/timeloop-prephysics-steps

1c05d1d

fix regression test model caching problem

f8fa385

brianhenn merged commit 149b545 into master Mar 23, 2021

brianhenn deleted the feature/timeloop-prephysics-steps branch March 23, 2021 23:58

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Runfile timeloop prephysics computations and configuration #1081

Runfile timeloop prephysics computations and configuration #1081

brianhenn commented Mar 11, 2021 •

edited

Loading

brianhenn Mar 12, 2021

nbren12 Mar 12, 2021

brianhenn Mar 12, 2021

brianhenn commented Mar 12, 2021

spencerkclark left a comment

nbren12 left a comment •

edited

Loading

nbren12 Mar 12, 2021

nbren12 Mar 12, 2021

nbren12 Mar 12, 2021

brianhenn Mar 15, 2021

nbren12 Mar 12, 2021

nbren12 Mar 12, 2021

nbren12 Mar 12, 2021

nbren12 Mar 12, 2021

nbren12 Mar 12, 2021

nbren12 Mar 12, 2021 •

edited

Loading

nbren12 left a comment

nbren12 Mar 15, 2021

nbren12 Mar 15, 2021

spencerkclark left a comment

spencerkclark left a comment

spencerkclark Mar 22, 2021

spencerkclark Mar 22, 2021

		@@ -60,6 +80,33 @@ def get_mock_sklearn_model() -> fv3fit.Predictor:
		return model


		def get_mock_rad_flux_model() -> fv3fit.Predictor:

	self.steppers = self._get_steppers(config)
	self.post_physics_stepper = self._get_stepper(config)
	self.pre_physics_stepper = self._get_pre_physics_stepper(config)

		@@ -27,6 +27,26 @@ def _model_dataset() -> xr.Dataset:
		return data


		def _rad_model_dataset() -> xr.Dataset:



		@dataclasses.dataclass
		class PrephysicsConfig:

Runfile timeloop prephysics computations and configuration #1081

Runfile timeloop prephysics computations and configuration #1081

Conversation

brianhenn commented Mar 11, 2021 • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

brianhenn commented Mar 12, 2021

spencerkclark left a comment

Choose a reason for hiding this comment

nbren12 left a comment • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

nbren12 Mar 12, 2021 • edited Loading

Choose a reason for hiding this comment

nbren12 left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

spencerkclark left a comment

Choose a reason for hiding this comment

spencerkclark left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

brianhenn commented Mar 11, 2021 •

edited

Loading

nbren12 left a comment •

edited

Loading

nbren12 Mar 12, 2021 •

edited

Loading