Add test that prescribing SST alters model evolution

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,73 @@
+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 mask_non_ocean_values(field):
+    is_ocean = np.isclose(get_state_single_variable("land_sea_mask"), 0.0)
+    return np.where(is_ocean, field, np.nan)
+
+
+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())
+
+        fv3gfs.wrapper.step()
+        air_temperature_from_default_ocean_temperature = get_state_single_variable(
+            "air_temperature"
+        )
+
+        # Restore state to original checkpoint; modify the SST;
+        # step the model again.
+        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"
+        )
+
+        # We expect these states to differ.
+        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[:]
+
+        # Over the ocean we expect these results to be equal.
+        result = mask_non_ocean_values(surface_temperature_diagnostic)
+        expected = mask_non_ocean_values(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[:]