Update counts for prognostic and tracer variables

src/umfile_utils/um_fields_subset.py

@@ -15,7 +15,11 @@
 from textwrap import dedent
 from itertools import chain
 
-PROGNOSTIC_STASH_CODES = tuple(chain(range(1,999+1), range(33001,34999+1)))
+# Prognostic variables have section 0, 33 (tracers), or 34 (UKCA).
+# Tracer flux variables 3100 - 3129 are also treated as prognostic.
+PROGNOSTIC_STASH_CODES = tuple(chain(range(1,999+1), range(3100,3129+1), range(33001,34999+1), range(54000,54999+1)))
+# Tracer variables have section 33
+TRACER_STASH_CODES = tuple(range(33000, 33999+1))
 
 def convert_to_list(value: str):
     """
@@ -216,7 +220,111 @@ def filter_fieldsfile(input_file, prognostic, include_list, exclude_list):
 
     filtered_file.fields = include_fields(input_file.fields, include_list) if include_list is not None else exclude_fields(input_file.fields, exclude_list)
     return filtered_file
+
+
+def update_prognostic_count(fields_file):
+    """
+    Update the count of prognostic variables in a file's fixed length header.
+
+    Parameters
+    ----------
+    field: mule.ff.FieldsFile
+        mule fields file to be updated.
+
+    Returns
+    -------
+    None
+    """
+    previous_count = fields_file.fixed_length_header.total_prognostic_fields
+    updated_count = sum([is_prognostic(field) and is_instantaneous(field) for field in fields_file.fields])
+    if updated_count != previous_count:
+        print(f"Resetting no. of prognostic fields from {previous_count} to {updated_count}")
+        fields_file.fixed_length_header.total_prognostic_fields = updated_count
+
+
+def is_prognostic(field):
+    """
+    Check whether a field is a prognostic variable.
+
+    Parameters
+    ----------
+    field: mule.Field
+        The mule field to be checked
+
+    Returns
+    -------
+    bool
+        Whether the field is prognostic or not.
+    """
+    return field.lbuser4 in PROGNOSTIC_STASH_CODES
+
+
+def is_instantaneous(field):
+    """
+    Check that a field is instantaneous with no time aggregation or processing.
+
+    Parameters
+    ----------
+    field: mule.Field
+
+    Returns
+    -------
+    bool
+        whether a field is instantaneous.
+    """
+    # Check the field is instantaneous (lbtime < 10).
+    # Check the field has no time processing (lbproc == 0).
+    # Check the field is not a timeseries (lbcode < 30000).
+    return field.lbtim < 10 and field.lbproc == 0 and field.lbcode < 30000
+
+
+def update_tracer_count(fields_file):
+    """
+    Update the count of instantaneous tracer variables in a file's fixed length header.
+
+    Parameters
+    ----------
+    field: mule.ff.FieldsFile
+        mule fields file to be updated.
+
+    Returns
+    -------
+    None
+    """
+    previous_count = fields_file.integer_constants.num_passive_tracers
+    num_tracer_fields = sum([is_tracer(field)  and is_instantaneous(field) for field in fields_file.fields])
+    # Divide by number of levels for the number of tracer variables
+    if num_tracer_fields % fields_file.integer_constants.num_tracer_levels !=0:
+        raise ValueError(
+            f"Number of tracer levels {fields_file.integer_constants.num_tracer_levels} "
+            f"does not divide number of tracer fields {num_tracer_fields}. Number of "
+            "tracer variables cannot be determined."
+        )
+
+    updated_count = num_tracer_fields / fields_file.integer_constants.num_tracer_levels
+
+    if updated_count != previous_count:
+        print(f"Resetting no. of tracer fields from {previous_count} to {updated_count}")
+        fields_file.integer_constants.num_passive_tracers = updated_count
+
+
+
+def is_tracer(field):
+    """
+    Check that a field is a tracer field.
+    Parameters
+    ----------
+    field: mule.Field
+
+    Returns
+    -------
+    bool
+        whether a field is a tracer.
+    """
+    return field.lbuser4 in TRACER_STASH_CODES
 
+
+
 def main():
 
     # Parse the inputs and validate that they do not xlist or vlist are given.
@@ -234,6 +342,9 @@ def main():
     if not args.validate:
         filtered_file.validate = void_validation
 
+    update_prognostic_count(filtered_file)
+    update_tracer_count(filtered_file)
+
     filtered_file.to_file(output_filename)
 
 if __name__== "__main__":

tests/um_fields_subset_test.py

@@ -5,7 +5,7 @@
 from unittest.mock import MagicMock, patch
 
 import numpy as np
-from hypothesis import strategies as st
+from hypothesis import given, strategies as st
 from hypothesis.extra import numpy as stnp
 from umfile_utils.um_fields_subset import (
     PROGNOSTIC_STASH_CODES,
@@ -17,18 +17,23 @@
     void_validation,
     include_fields,
     exclude_fields,
+    is_prognostic,
+    update_prognostic_count,
+    is_tracer,
+    update_tracer_count,
+    is_instantaneous,
     main,
 )
 
-@pytest.fixture
+@pytest.fixture(scope="module")
 def create_mock_umfile():
     def _mock_umfile():
         """Factory function to create a mule UMfile mock object and initialize it with empty fields."""
         return MagicMock(fields=[])
 
     return _mock_umfile
 
-@pytest.fixture
+@pytest.fixture(scope="module")
 def create_mock_field():
     """Factory function to create a mule field mock object."""
 
@@ -288,12 +293,225 @@ def test_exclude_fields(create_mock_field):
             assert r.lbuser4 == er.lbuser4 # Ensure the stash codes are correct
 
 
+prognostic_codes = st.fixed_dictionaries(
+    {"lbuser4": st.integers(min_value=1, max_value=54999).filter(lambda n: n in PROGNOSTIC_STASH_CODES)}
+)
+non_prognostic_codes = st.fixed_dictionaries(
+    {"lbuser4": st.integers(min_value=1, max_value=34999).filter(lambda n: n not in PROGNOSTIC_STASH_CODES)}
+)
+@given(prognostic_codes, non_prognostic_codes)
+def test_is_prognostic(create_mock_field, prognostic_code, non_prognostic_code):
+    """Check that prognostic fields are correctly identified."""
+    mock_prognostic = create_mock_field(**prognostic_code)
+    assert is_prognostic(mock_prognostic)
+
+    mock_non_prognostic = create_mock_field(**non_prognostic_code)
+    assert not is_prognostic(mock_non_prognostic)
+
+
+instantaneous_constants = st.fixed_dictionaries(
+    {
+        "lbtim": st.integers(min_value=1, max_value=9),
+        "lbproc": st.just(0),
+        "lbcode": st.integers(min_value=0, max_value=29999)
+    }
+)
+@given(instantaneous_constants)
+def test_is_instantaneous(create_mock_field, integer_constants):
+    """Check that instantaneous fields are correctly identified."""
+    mock_field = create_mock_field(**integer_constants)
+    assert is_instantaneous(mock_field)
+
+
+non_instantaneous_constants = st.fixed_dictionaries(
+    {
+        "lbtim": st.integers(min_value=10),
+        "lbproc": st.integers().filter(lambda n: n != 0),
+        "lbcode": st.integers(min_value=30000)
+    }
+)
+@given(non_instantaneous_constants)
+def test_non_instantaneous(create_mock_field, integer_constants):
+    """Check that non-instantaneous fields are correctly identified."""
+    mock_field = create_mock_field(
+        lbtim=integer_constants["lbtim"],
+        lbproc=0,
+        lbcode=0
+    )
+    assert not is_instantaneous(mock_field)
+
+    mock_field = create_mock_field(
+        lbtim=1,
+        lbproc=integer_constants["lbproc"],
+        lbcode=0
+    )
+    assert not is_instantaneous(mock_field)
+
+    mock_field = create_mock_field(
+        lbtim=1,
+        lbproc=0,
+        lbcode=integer_constants["lbcode"]
+    )
+    assert not is_instantaneous(mock_field)
+
+
+def join_dict_strategies(*args):
+    """
+    Helper function for joining dictionary test strategies.
+    Given multiple dictionary strategies, returns a strategy
+    which samples from each dictionary strategy and joins the result.
+    """
+    def _join_func(*args):
+        combined = {}
+        for d in args:
+            combined.update(d)
+        return combined
+    return st.builds(_join_func, *args)
+
+
+prognostic_fields = join_dict_strategies(
+    prognostic_codes,
+    instantaneous_constants
+)
+
+# Non prognostic fields will either:
+#   1. Not have a prognostic STASH code
+#   2. Not be instantaneous
+non_prognostic_fields = (
+    join_dict_strategies(
+        prognostic_codes,
+        non_instantaneous_constants
+    ) |
+    join_dict_strategies(
+        non_prognostic_codes,
+        instantaneous_constants
+    ) |
+    join_dict_strategies(
+        non_prognostic_codes,
+        non_instantaneous_constants
+    )
+)
+
+prognostic_lists = st.lists(prognostic_fields, min_size=0, max_size=20)
+non_prognostic_lists = st.lists(non_prognostic_fields, min_size=0, max_size=20)
+@given(prognostic_lists, non_prognostic_lists, st.integers(min_value
+=0))
+def test_update_prognostic_count(create_mock_umfile,
+                                 create_mock_field,
+                                 prognostic_list,
+                                 non_prognostic_list,
+                                 initial_count):
+    """
+    Check that the update_prognostic_count function correctly calculates
+    the number of prognostic variables.
+    """
+    fields_file = create_mock_umfile()
+    fields_file.fixed_length_header.total_prognostic_fields = initial_count
+    for field in prognostic_list + non_prognostic_list:
+        fields_file.fields.append(
+            create_mock_field(**field)
+        )
+
+    update_prognostic_count(fields_file)
+
+    assert fields_file.fixed_length_header.total_prognostic_fields == len(prognostic_list)
+
+
+tracer_codes = st.fixed_dictionaries({"lbuser4": st.integers(min_value=33000, max_value=33999)})
+non_tracer_codes = st.fixed_dictionaries({"lbuser4": st.integers(min_value=1, max_value=32999)})
+@given(tracer_codes, non_tracer_codes)
+def test_is_tracer(create_mock_field, tracer_code, non_tracer_code):
+    """Check that tracer fields are correctly identified."""
+    mock_tracer = create_mock_field(**tracer_code)
+    assert is_tracer(mock_tracer)
+
+    mock_non_tracer = create_mock_field(**non_tracer_code)
+    assert not is_tracer(mock_non_tracer)
+
+
+tracer_fields = join_dict_strategies(
+    tracer_codes,
+    instantaneous_constants
+)
+
+# Non tracer fields will either:
+#   1. Not have a tracer STASH code
+#   2. Not be instantaneous
+non_tracer_fields = (
+    join_dict_strategies(
+        tracer_codes,
+        non_instantaneous_constants
+    ) |
+    join_dict_strategies(
+        non_tracer_codes,
+        instantaneous_constants
+    ) |
+    join_dict_strategies(
+        non_tracer_codes,
+        non_instantaneous_constants
+    )
+)
+
+@given(st.integers(min_value=1, max_value=40), st.integers(min_value=0, max_value=10), st.integers(min_value=0), st.data())
+def test_update_tracer_count(create_mock_umfile,
+                             create_mock_field,
+                             num_tracer_levels,
+                             num_tracer_vars,
+                             initial_count,
+                             data):
+    """
+    Check that the update_tracer_count function correctly calculates
+    the number of tracer variables.
+    """
+    fields_file = create_mock_umfile()
+    fields_file.integer_constants.num_passive_tracers = initial_count
+    fields_file.integer_constants.num_tracer_levels = num_tracer_levels
+
+    tracer_list = data.draw(st.lists(tracer_fields, min_size=num_tracer_levels * num_tracer_vars, max_size=num_tracer_levels * num_tracer_vars))
+    non_tracer_list = data.draw(st.lists(non_tracer_fields, min_size=1, max_size=20))
+
+    for field in tracer_list + non_tracer_list:
+        fields_file.fields.append(
+            create_mock_field(**field)
+        )
+
+    update_tracer_count(fields_file)
+
+    assert fields_file.integer_constants.num_passive_tracers == num_tracer_vars
+
+
+@given(st.data())
+def test_update_tracer_count_error(create_mock_umfile,
+                             create_mock_field,
+                             data):
+    """
+    Check that the update_tracer_count function produces an error
+    when the number of tracer fields and tracer levels are not consistent.
+    """
+    fields_file = create_mock_umfile()
+    fields_file.integer_constants.num_tracer_levels = 2
+    num_tracers = 7
+    tracer_list = data.draw(st.lists(tracer_fields, min_size=num_tracers, max_size=num_tracers))
+
+    for field in tracer_list:
+        fields_file.fields.append(
+            create_mock_field(**field)
+        )
+
+    with pytest.raises(ValueError, match="Number of tracer levels"):
+        update_tracer_count(fields_file)
+
+
 @patch("umfile_utils.um_fields_subset.parse_args")
 @patch("umfile_utils.um_fields_subset.create_default_outname")
 @patch("mule.DumpFile.from_file")
 @patch("umfile_utils.um_fields_subset.filter_fieldsfile")
 @patch("umfile_utils.um_fields_subset.void_validation")
+@patch("umfile_utils.um_fields_subset.update_prognostic_count")
+@patch("umfile_utils.um_fields_subset.update_tracer_count")
 def test_main(
+    mock_tracer_count,
+    mock_prognostic_count,
     mock_void_validation,
     mock_filter_fieldsfile,
     mock_mule_dumpfile_from_file,