Refactor/test prognostic run steppers #1033

nbren12 · 2021-02-23T02:27:54Z

The steppers repeated code and were difficult to test. This PR refactors the nudging and machine learning steppers to enable more reuse, pushing the common/messy bits up to the TimeLoop class. This allows a better separation of concerns between the "Steppers" (see the docstrings of Stepper and TimeLoop for more information).

This more "pure" interface allowed regression testing more of the ML code. The nudging stepper is harder to test since it requires MPI and more data.

I tested each refactor using checksums, and fixed the non-determinism in the scikit-learn based tests.

Also added:

debugging conveniences
made DerivedFV3State an actual MutableMapping.

Refactor nudging logic to a more "pure" class.

...ognostic_c48_run/tests/_regtest_outputs/test_regression.test_fv3run_checksum_logs[keras].out

nbren12 · 2021-02-25T01:20:05Z

...ic_c48_run/tests/_regtest_outputs/test_regression.test_fv3run_checksum_restarts[sklearn].out

@@ -1 +1 @@
-f66f3591c848109b47439bdf6a616eeb


notice that the keras checksum is not updated. this suggests the refactors didn't break modify the state updates by the ML code.

nbren12 · 2021-02-25T01:21:36Z

...ic_c48_run/tests/_regtest_outputs/test_regression.test_fv3run_checksum_restarts[nudging].out

@@ -0,0 +1 @@
+566da3dda9be0fd64285cec8f3b93bed


notice that this checksum was introduce in 3a78a4b before any of the large refactors.

oliverwm1

Looks great! Definitely some more work to do in terms of unifying diagnostic calcs and perhaps generalizing how to know which tendencies should get applied to physics/dycore states. But this is a big step in the right direction I think.

Also we'll have to think about the right way to expand the number of "substeps" if we want to be able to apply updates between radiation/other physics computations. I could see the substeps getting a bit out of control.

oliverwm1 · 2021-02-25T20:16:48Z

workflows/prognostic_c48_run/runtime/derived_state.py

+        )
+        tracer_names = set(v for v in fv3gfs.wrapper.get_tracer_metadata())
+        # see __getitem__
+        local_names = {"latent_heat_flux", "total_water"}


Okay for now, but this does seem pretty bespoke. Should we have some kind of "diagnostics" table that these are a part of?

The wrapper can list all the available diagnostics withfv3gfs.wrapper._get_diagnostic_info. We should probably make this public api.

oliverwm1 · 2021-02-25T20:34:08Z

workflows/prognostic_c48_run/runtime/derived_state.py

+        for key in self.keys():
+            try:
+                data = self[key]
+            except:  # noqa: E722


Why not except KeyError?

I'm deleting this method since it is not hard to copy paste where appropriate.

workflows/prognostic_c48_run/runtime/loop.py

oliverwm1 · 2021-02-25T20:56:45Z

workflows/prognostic_c48_run/runtime/loop.py

+    Each time step of the model evolutions proceeds like this::
+
+        step_dynamics,
+        compute_physics,


presumably we will be splitting this up further if we want to apply radiation updates to the land surface model

True. This docstring is repetitive but it did help me think about the current implementation.

oliverwm1 · 2021-02-25T21:10:10Z

workflows/prognostic_c48_run/runtime/loop.py

+        try:
+            del diagnostics[TOTAL_PRECIP]
+        except KeyError:
+            pass


Suggested change

try:

del diagnostics[TOTAL_PRECIP]

except KeyError:

pass

diagnostics.pop(TOTAL_PRECIP, None)

?

The try is a little longer but more explicit in my opinion.

nbren12 · 2021-02-25T22:26:54Z

Definitely some more work to do in terms of unifying diagnostic calcs

For sure. I think "net_heating" and "total_water_path" are each computed in at least a couple places. We'll have to think more carefully about what "diagnostics" and ML scheme should save.

and perhaps generalizing how to know which tendencies should get applied to physics/dycore states.
Also we'll have to think about the right way to expand the number of "substeps" if we want to be able to apply updates between radiation/other physics computations. I could see the substeps getting a bit out of control.

I had some ideas about this. Generally speaking, these state updates are a DAG, where the edges are input/output variables and the nodes are individual schemes. If each scheme described it's inputs/outputs and gave a rough "priority" then the time looper could determine the precise order of computations.

frodre

Thanks, @nbren12. Everything looks good and overall, a nice simplification. I also learned a bit about testing/regression testing, which is always a bonus.

frodre · 2021-02-25T22:56:03Z

workflows/prognostic_c48_run/runtime/nudging.py

+    # TODO fix this bug in a follow-up non-refactor PR
+    # this logic replicates a bug in the previous nudged run
+    # the SSTs were never actually applied to the fv3gfs-wrapper state
+    # This bug could be scientifically significant.


Serious enough that we'd want to throw an error here until it's fixed?

This code is executed in every nudged run. It's a trivial fix now with this refactor. This scheme just needs to return the dict above rather than an empty dict, but I didn't want to make any bit-incompatible changes.

Related to #1037

The steppers repeated code and were difficult to test. This PR refactors the nudging and machine learning steppers to enable more reuse, pushing the common/messy bits up to the TimeLoop class. This allows a better separation of concerns between the "Steppers" (see the docstrings of `Stepper` and `TimeLoop` for more information). This more "pure" interface allowed regression testing more of the ML code. The nudging stepper is harder to test since it requires MPI and more data. I tested each refactor using checksums, and fixed the non-determinism in the scikit-learn based tests. Also added: - debugging conveniences - made `DerivedFV3State` an actual MutableMapping.

nbren12 added 15 commits February 22, 2021 13:47

Reset/require checksum tests

4d0940a

Check the logs checksum

0df2ebd

Check checksums for nudging

3a78a4b

Delete flaky log checker

a1943bf

Remove _fv3gfs requirement from machine_learning

c350308

Refactor ML code into a Pure ML Stepper

2391c76

Refactor model mocking to shared module

d8a42d9

Add DerivedFV3State.checkpoint for easy serialization

493137c

Add regression tests of PureMLStepper

a4ed599

Fix the scheme regression test

2248012

Simplify apply_python_to_dycore_state

2132b9f

Add module with debugging utilities

4e9844b

Lint machine_learning.py

3e4c945

Cleanup nudging stepper

c12b3ab

Refactor nudging logic to a more "pure" class.

Refactor MLStepper and NudgingStepper to more pure objects

9705380

nbren12 force-pushed the krasnopolsky-stepper branch from 9ed918f to 9705380 Compare February 24, 2021 01:49

nbren12 commented Feb 24, 2021

View reviewed changes

...ognostic_c48_run/tests/_regtest_outputs/test_regression.test_fv3run_checksum_logs[keras].out Outdated Show resolved Hide resolved

nbren12 added 6 commits February 24, 2021 11:34

(fv3fit) Make predict_columnwise deterministic

ebcc812

(progrun) test_machine_learning is now deterministic

d5cacf6

Reset sklearn regression checksum

c0b86cb

WIP Refactor stepper

850a672

Implement MutableMapping interface

d7d9d1e

Delete unused regression data

d073076

nbren12 commented Feb 25, 2021

View reviewed changes

nbren12 requested a review from oliverwm1 February 25, 2021 17:18

oliverwm1 approved these changes Feb 25, 2021

View reviewed changes

Respond to Oli's review

f8abe4d

Merge branch 'master' into krasnopolsky-stepper

3bb7f14

frodre approved these changes Feb 25, 2021

View reviewed changes

nbren12 merged commit 58da2ce into master Feb 25, 2021

nbren12 deleted the krasnopolsky-stepper branch February 25, 2021 23:53

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Refactor/test prognostic run steppers #1033

Refactor/test prognostic run steppers #1033

nbren12 commented Feb 23, 2021 •

edited

Loading

nbren12 Feb 25, 2021

nbren12 Feb 25, 2021

oliverwm1 left a comment

oliverwm1 Feb 25, 2021

nbren12 Feb 25, 2021

oliverwm1 Feb 25, 2021

nbren12 Feb 25, 2021

oliverwm1 Feb 25, 2021

nbren12 Feb 25, 2021

oliverwm1 Feb 25, 2021

nbren12 Feb 25, 2021

nbren12 commented Feb 25, 2021 •

edited

Loading

frodre left a comment

frodre Feb 25, 2021

nbren12 Feb 25, 2021

nbren12 Feb 25, 2021

		@@ -1 +1 @@
		f66f3591c848109b47439bdf6a616eeb

		@@ -0,0 +1 @@
		566da3dda9be0fd64285cec8f3b93bed

Refactor/test prognostic run steppers #1033

Refactor/test prognostic run steppers #1033

Conversation

nbren12 commented Feb 23, 2021 • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

oliverwm1 left a 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 commented Feb 25, 2021 • edited Loading

frodre left a 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 commented Feb 23, 2021 •

edited

Loading

nbren12 commented Feb 25, 2021 •

edited

Loading