Add test that prescribing SST alters model evolution (#256)

This PR adds a test confirming that with the use of the namelist flag added in ai2cm/fv3gfs-fortran#173, setting the "ocean_surface_temperature" modifies the evolution of the model's prognostic variables (e.g. air temperature).

In between ai2cm/fv3gfs-fortran#173 and ai2cm/fv3gfs-fortran#93, this functionality was broken (see the failing test associated with 5a6febb of this PR). In other words, since ai2cm/fv3gfs-fortran#93, Sfcprop%tsfco in the fortran model was modified within the physics driver prior to being used by any physics scheme; therefore when we set it via Python, it was immediately set to something else in the fortran before it could have any impact. The namelist parameter toggles the behavior back to what it was before ai2cm/fv3gfs-fortran#93, allowing us to prescribe the SST.

tests/test_all.py

@@ -40,6 +40,9 @@ def test_get_time_noleap(self):
     def test_flags(self):
         run_unittest_script("test_flags.py")
 
+    def test_set_ocean_surface_temperature(self):
+        run_unittest_script("test_set_ocean_surface_temperature.py")
+
 
 if __name__ == "__main__":
     unittest.main()

tests/test_overrides_for_surface_radiative_fluxes.py

@@ -1,12 +1,16 @@
 import unittest
 import os
-from copy import deepcopy
 from fv3gfs.wrapper._properties import OVERRIDES_FOR_SURFACE_RADIATIVE_FLUXES
 import numpy as np
 import fv3gfs.wrapper
 import fv3gfs.util
 from mpi4py import MPI
-from util import get_default_config, main
+from util import (
+    get_default_config,
+    get_state_single_variable,
+    main,
+    replace_state_with_random_values,
+)
 
 
 test_dir = os.path.dirname(os.path.abspath(__file__))
@@ -17,19 +21,6 @@
 ) = OVERRIDES_FOR_SURFACE_RADIATIVE_FLUXES
 
 
-def override_surface_radiative_fluxes_with_random_values():
-    old_state = fv3gfs.wrapper.get_state(names=OVERRIDES_FOR_SURFACE_RADIATIVE_FLUXES)
-    replace_state = deepcopy(old_state)
-    for name, quantity in replace_state.items():
-        quantity.view[:] = np.random.uniform(size=quantity.extent)
-    fv3gfs.wrapper.set_state(replace_state)
-    return replace_state
-
-
-def get_state_single_variable(name):
-    return fv3gfs.wrapper.get_state([name])[name].view[:]
-
-
 class OverridingSurfaceRadiativeFluxTests(unittest.TestCase):
     def __init__(self, *args, **kwargs):
         super(OverridingSurfaceRadiativeFluxTests, self).__init__(*args, **kwargs)
@@ -64,7 +55,7 @@ def test_overriding_fluxes_changes_model_state(self):
         # Restore state to original checkpoint; modify the radiative fluxes;
         # step the model again.
         fv3gfs.wrapper.set_state(checkpoint_state)
-        override_surface_radiative_fluxes_with_random_values()
+        replace_state_with_random_values(OVERRIDES_FOR_SURFACE_RADIATIVE_FLUXES)
         fv3gfs.wrapper.step()
         temperature_with_random_override = get_state_single_variable("air_temperature")
 
@@ -74,7 +65,9 @@ def test_overriding_fluxes_changes_model_state(self):
         )
 
     def test_overriding_fluxes_are_propagated_to_diagnostics(self):
-        replace_state = override_surface_radiative_fluxes_with_random_values()
+        replace_state = replace_state_with_random_values(
+            OVERRIDES_FOR_SURFACE_RADIATIVE_FLUXES
+        )
 
         # We need to step the model to fill the diagnostics buckets.
         fv3gfs.wrapper.step()

tests/test_set_ocean_surface_temperature.py

@@ -0,0 +1,74 @@
+import unittest
+import os
+import numpy as np
+import fv3gfs.wrapper
+import fv3gfs.util
+from mpi4py import MPI
+from util import (
+    get_default_config,
+    get_state_single_variable,
+    main,
+    replace_state_with_random_values,
+)
+
+
+test_dir = os.path.dirname(os.path.abspath(__file__))
+
+
+def select_ocean_values(*fields):
+    is_ocean = np.isclose(get_state_single_variable("land_sea_mask"), 0.0)
+    return (field[is_ocean] for field in fields)
+
+
+class PrescribeSSTTests(unittest.TestCase):
+    def __init__(self, *args, **kwargs):
+        super(PrescribeSSTTests, self).__init__(*args, **kwargs)
+
+    def setUp(self):
+        pass
+
+    def tearDown(self):
+        MPI.COMM_WORLD.barrier()
+
+    def test_prescribing_sst_changes_model_state(self):
+        checkpoint_state = fv3gfs.wrapper.get_state(fv3gfs.wrapper.get_restart_names())
+
+        # If we do not set the sea surface temperature and
+        # use_climatological_sst is set to .false., the sea surface temperature
+        # will remain at what it was set to in the initial conditions for the
+        # duration of the run.
+        fv3gfs.wrapper.step()
+        air_temperature_from_default_ocean_temperature = get_state_single_variable(
+            "air_temperature"
+        )
+
+        fv3gfs.wrapper.set_state(checkpoint_state)
+        replace_state_with_random_values(["ocean_surface_temperature"])
+        fv3gfs.wrapper.step()
+        air_temperature_from_prescribed_ocean_temperature = get_state_single_variable(
+            "air_temperature"
+        )
+
+        assert not np.allclose(
+            air_temperature_from_default_ocean_temperature,
+            air_temperature_from_prescribed_ocean_temperature,
+        )
+
+    def test_prescribing_sst_changes_surface_temperature_diagnostic(self):
+        replaced_state = replace_state_with_random_values(["ocean_surface_temperature"])
+        prescribed_sst = replaced_state["ocean_surface_temperature"].view[:]
+        fv3gfs.wrapper.step()
+        surface_temperature_diagnostic = fv3gfs.wrapper.get_diagnostic_by_name(
+            "tsfc", module_name="gfs_sfc"
+        ).view[:]
+
+        result, expected = select_ocean_values(
+            surface_temperature_diagnostic, prescribed_sst
+        )
+        np.testing.assert_allclose(result, expected)
+
+
+if __name__ == "__main__":
+    config = get_default_config()
+    config["namelist"]["gfs_physics_nml"]["use_climatological_sst"] = False
+    main(test_dir, config)

tests/test_setters.py

@@ -11,7 +11,13 @@
 )
 import fv3gfs.util
 from mpi4py import MPI
-from util import get_current_config, get_default_config, generate_data_dict, main
+from util import (
+    get_current_config,
+    get_default_config,
+    generate_data_dict,
+    main,
+    replace_state_with_random_values,
+)
 
 
 test_dir = os.path.dirname(os.path.abspath(__file__))
@@ -128,12 +134,7 @@ def _set_names_helper(self, name_list):
         self._set_names_one_at_a_time_helper(name_list)
 
     def _set_all_names_at_once_helper(self, name_list):
-        old_state = fv3gfs.wrapper.get_state(names=name_list)
-        self._check_gotten_state(old_state, name_list)
-        replace_state = deepcopy(old_state)
-        for name, quantity in replace_state.items():
-            quantity.view[:] = np.random.uniform(size=quantity.extent)
-        fv3gfs.wrapper.set_state(replace_state)
+        replace_state = replace_state_with_random_values(name_list)
         new_state = fv3gfs.wrapper.get_state(names=name_list)
         self._check_gotten_state(new_state, name_list)
         for name, new_quantity in new_state.items():

tests/util.py

@@ -7,8 +7,10 @@
 import yaml
 import unittest
 import shutil
+import numpy as np
 import fv3config
 import fv3gfs.wrapper
+from copy import deepcopy
 from mpi4py import MPI
 
 libc = ctypes.CDLL(None)
@@ -113,3 +115,16 @@ def get_current_config():
 
 def generate_data_dict(properties):
     return {entry["name"]: entry for entry in properties}
+
+
+def replace_state_with_random_values(names):
+    old_state = fv3gfs.wrapper.get_state(names=names)
+    replace_state = deepcopy(old_state)
+    for name, quantity in replace_state.items():
+        quantity.view[:] = np.random.uniform(size=quantity.extent)
+    fv3gfs.wrapper.set_state(replace_state)
+    return replace_state
+
+
+def get_state_single_variable(name):
+    return fv3gfs.wrapper.get_state([name])[name].view[:]