From d38cf3aa154b94b9dd3acffc005a853dbfd78cef Mon Sep 17 00:00:00 2001 From: Spencer Wong Date: Wed, 17 Jun 2026 17:07:25 +1000 Subject: [PATCH 1/4] Update prognostic and tracer counts in the headers --- src/umfile_utils/um_fields_subset.py | 108 ++++++++++++++++++++++++++- 1 file changed, 107 insertions(+), 1 deletion(-) diff --git a/src/umfile_utils/um_fields_subset.py b/src/umfile_utils/um_fields_subset.py index 55cd628c..8af01af3 100644 --- a/src/umfile_utils/um_fields_subset.py +++ b/src/umfile_utils/um_fields_subset.py @@ -16,6 +16,7 @@ from itertools import chain PROGNOSTIC_STASH_CODES = tuple(chain(range(1,999+1), range(33001,34999+1))) +TRACER_STASH_CODES = tuple(range(33001, 33999+1)) def convert_to_list(value: str): """ @@ -216,7 +217,109 @@ 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) 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. + """ + section = field.lbuser4//1000 + # Prognostic variables have section 0, 33 (tracers), or 34 (UKCA). + # Tracer fluxes 3100 - 3129 are also treated as prognostic. + code_condition = section in [0, 33, 34] or 3100 <= field.lbuser4 <= 3129 + + return code_condition and is_instantaneous(field) + + +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 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.integer_constants.num_passive_tracers + num_tracer_fields = sum([is_tracer(field) for field in fields_file.fields]) + # Divide by number of levels for the number of tracer variables + 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 + + #TODO: Do we need to reset the tracer levels to equal p levels? + + +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 and is_instantaneous(field) + + def main(): # Parse the inputs and validate that they do not xlist or vlist are given. @@ -234,6 +337,9 @@ def main(): if not args.validate: filtered_file.validate = void_validation + update_prognostic_count(filtered_file) + # update_tracer_count(fields_file) + filtered_file.to_file(output_filename) if __name__== "__main__": From 6bad0f237c9d81645804804a62c107291cf2ddd1 Mon Sep 17 00:00:00 2001 From: Spencer Wong Date: Fri, 19 Jun 2026 15:40:08 +1000 Subject: [PATCH 2/4] Reset prognostic and tracer counts --- src/umfile_utils/um_fields_subset.py | 31 ++-- tests/um_fields_subset_test.py | 223 ++++++++++++++++++++++++++- 2 files changed, 238 insertions(+), 16 deletions(-) diff --git a/src/umfile_utils/um_fields_subset.py b/src/umfile_utils/um_fields_subset.py index 8af01af3..ccee9a37 100644 --- a/src/umfile_utils/um_fields_subset.py +++ b/src/umfile_utils/um_fields_subset.py @@ -15,8 +15,10 @@ from textwrap import dedent from itertools import chain -PROGNOSTIC_STASH_CODES = tuple(chain(range(1,999+1), range(33001,34999+1))) -TRACER_STASH_CODES = tuple(range(33001, 33999+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))) +TRACER_STASH_CODES = tuple(range(33000, 33999+1)) def convert_to_list(value: str): """ @@ -233,7 +235,7 @@ def update_prognostic_count(fields_file): None """ previous_count = fields_file.fixed_length_header.total_prognostic_fields - updated_count = sum([is_prognostic(field) for field in fields_file.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 @@ -253,12 +255,7 @@ def is_prognostic(field): bool Whether the field is prognostic or not. """ - section = field.lbuser4//1000 - # Prognostic variables have section 0, 33 (tracers), or 34 (UKCA). - # Tracer fluxes 3100 - 3129 are also treated as prognostic. - code_condition = section in [0, 33, 34] or 3100 <= field.lbuser4 <= 3129 - - return code_condition and is_instantaneous(field) + return field.lbuser4 in PROGNOSTIC_STASH_CODES def is_instantaneous(field): @@ -282,7 +279,7 @@ def is_instantaneous(field): def update_tracer_count(fields_file): """ - Update the count of prognostic variables in a file's fixed length header. + Update the count of instantaneous tracer variables in a file's fixed length header. Parameters ---------- @@ -294,8 +291,15 @@ def update_tracer_count(fields_file): None """ previous_count = fields_file.integer_constants.num_passive_tracers - num_tracer_fields = sum([is_tracer(field) for field in fields_file.fields]) + 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: @@ -317,7 +321,8 @@ def is_tracer(field): bool whether a field is a tracer. """ - return field.lbuser4 in TRACER_STASH_CODES and is_instantaneous(field) + return field.lbuser4 in TRACER_STASH_CODES + def main(): @@ -338,7 +343,7 @@ def main(): filtered_file.validate = void_validation update_prognostic_count(filtered_file) - # update_tracer_count(fields_file) + update_tracer_count(filtered_file) filtered_file.to_file(output_filename) diff --git a/tests/um_fields_subset_test.py b/tests/um_fields_subset_test.py index b0e31194..89355c3a 100644 --- a/tests/um_fields_subset_test.py +++ b/tests/um_fields_subset_test.py @@ -5,10 +5,11 @@ 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, + TRACER_STASH_CODES, convert_to_list, create_default_outname, warn_if_stash_not_present, @@ -17,10 +18,15 @@ 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.""" @@ -28,7 +34,7 @@ def _mock_umfile(): return _mock_umfile -@pytest.fixture +@pytest.fixture(scope="module") def create_mock_field(): """Factory function to create a mule field mock object.""" @@ -288,11 +294,222 @@ 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=34999).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 variables 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 variables 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 strattegies. + Given multiple dictionary strategies, returns a strategy + which joins examples from each strategy. + """ + 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 prognostic variables 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_ function correctly calculates + the number of prognostic variables. + """ + 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_void_validation, mock_filter_fieldsfile, From bd19a209cc9c5854bd71d77cdfd0b20ab3b69bdc Mon Sep 17 00:00:00 2001 From: Spencer Wong Date: Mon, 22 Jun 2026 11:21:52 +1000 Subject: [PATCH 3/4] Reset the number of tracer and prognostic fields in the header (#31) --- src/umfile_utils/um_fields_subset.py | 2 +- tests/um_fields_subset_test.py | 19 ++++++++++--------- 2 files changed, 11 insertions(+), 10 deletions(-) diff --git a/src/umfile_utils/um_fields_subset.py b/src/umfile_utils/um_fields_subset.py index ccee9a37..7c451800 100644 --- a/src/umfile_utils/um_fields_subset.py +++ b/src/umfile_utils/um_fields_subset.py @@ -18,6 +18,7 @@ # 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))) +# Tracer variables have section 33 TRACER_STASH_CODES = tuple(range(33000, 33999+1)) def convert_to_list(value: str): @@ -306,7 +307,6 @@ def update_tracer_count(fields_file): print(f"Resetting no. of tracer fields from {previous_count} to {updated_count}") fields_file.integer_constants.num_passive_tracers = updated_count - #TODO: Do we need to reset the tracer levels to equal p levels? def is_tracer(field): diff --git a/tests/um_fields_subset_test.py b/tests/um_fields_subset_test.py index 89355c3a..05e215e0 100644 --- a/tests/um_fields_subset_test.py +++ b/tests/um_fields_subset_test.py @@ -9,7 +9,6 @@ from hypothesis.extra import numpy as stnp from umfile_utils.um_fields_subset import ( PROGNOSTIC_STASH_CODES, - TRACER_STASH_CODES, convert_to_list, create_default_outname, warn_if_stash_not_present, @@ -302,7 +301,7 @@ def test_exclude_fields(create_mock_field): ) @given(prognostic_codes, non_prognostic_codes) def test_is_prognostic(create_mock_field, prognostic_code, non_prognostic_code): - """Check that prognostic variables are correctly identified""" + """Check that prognostic fields are correctly identified.""" mock_prognostic = create_mock_field(**prognostic_code) assert is_prognostic(mock_prognostic) @@ -319,7 +318,7 @@ def test_is_prognostic(create_mock_field, prognostic_code, non_prognostic_code): ) @given(instantaneous_constants) def test_is_instantaneous(create_mock_field, integer_constants): - """Check that instantaneous variables are correctly identified""" + """Check that instantaneous fields are correctly identified.""" mock_field = create_mock_field(**integer_constants) assert is_instantaneous(mock_field) @@ -333,7 +332,7 @@ def test_is_instantaneous(create_mock_field, integer_constants): ) @given(non_instantaneous_constants) def test_non_instantaneous(create_mock_field, integer_constants): - """Check that non-instantaneous fields are correctly identified""" + """Check that non-instantaneous fields are correctly identified.""" mock_field = create_mock_field( lbtim=integer_constants["lbtim"], lbproc=0, @@ -358,9 +357,9 @@ def test_non_instantaneous(create_mock_field, integer_constants): def join_dict_strategies(*args): """ - Helper function for joining dictionary test strattegies. + Helper function for joining dictionary test strategies. Given multiple dictionary strategies, returns a strategy - which joins examples from each strategy. + which samples from each dictionary strategy and joins the result. """ def _join_func(*args): combined = {} @@ -422,7 +421,7 @@ def test_update_prognostic_count(create_mock_umfile, 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 prognostic variables are correctly identified""" + """Check that tracer fields are correctly identified.""" mock_tracer = create_mock_field(**tracer_code) assert is_tracer(mock_tracer) @@ -486,8 +485,8 @@ def test_update_tracer_count_error(create_mock_umfile, create_mock_field, data): """ - Check that the update_ function correctly calculates - the number of prognostic variables. + 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 @@ -511,6 +510,8 @@ def test_update_tracer_count_error(create_mock_umfile, @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, From 09077ef57e552f5028767f38284b56d0de98c6fd Mon Sep 17 00:00:00 2001 From: Spencer Wong Date: Thu, 25 Jun 2026 15:09:48 +1000 Subject: [PATCH 4/4] Include section 54 prognostics --- src/umfile_utils/um_fields_subset.py | 2 +- tests/um_fields_subset_test.py | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/src/umfile_utils/um_fields_subset.py b/src/umfile_utils/um_fields_subset.py index 7c451800..a00518aa 100644 --- a/src/umfile_utils/um_fields_subset.py +++ b/src/umfile_utils/um_fields_subset.py @@ -17,7 +17,7 @@ # 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))) +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)) diff --git a/tests/um_fields_subset_test.py b/tests/um_fields_subset_test.py index 05e215e0..f0d92101 100644 --- a/tests/um_fields_subset_test.py +++ b/tests/um_fields_subset_test.py @@ -294,7 +294,7 @@ def test_exclude_fields(create_mock_field): prognostic_codes = st.fixed_dictionaries( - {"lbuser4": st.integers(min_value=1, max_value=34999).filter(lambda n: n in PROGNOSTIC_STASH_CODES)} + {"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)}