Allow reading fields with a constituent dimension

Author: jimmielin

src/data/write_init_files.py

@@ -739,12 +739,19 @@ def get_dimension_info(hvar):
        - The local variable name of the vertical dimension (or None)
        - True if <hvar> has one dimension which is a horizontal dimension or
             if <hvar> has two dimensions (horizontal and vertical)
+       - Flag if any dimensions are number_of_ccpp_constituents and needs
+            reading separate variables by constituent and reassembled into
+            host model indices.
     """
     vdim_name = None
     legal_dims = False
     fail_reason = ""
+    has_constituent_read = False
+
     dims = hvar.get_dimensions()
     levnm = hvar.has_vertical_dimension()
+    has_constituent_dim = any('number_of_ccpp_constituents' in dim for dim in dims)
+
     # <hvar> is only 'legal' for 2 or 3 dimensional fields (i.e., 1 or 2
     #    dimensional variables). The second dimension must be vertical.
     # XXgoldyXX: If we ever need to read scalars, it would have to be
@@ -785,6 +792,19 @@ def get_dimension_info(hvar):
         # end if
         suff = "; "
     # end if
+
+    # A special case where any dimensions include number_of_ccpp_constituents,
+    # in this case the variable needs to be suffixed by _ + constituent name
+    # and read and reassembled separately into host model constituent indices
+    # based on constituent name.
+    #
+    # In this case, override legal_dims as this case will be handled separately
+    if has_constituent_dim:
+        has_constituent_read = True
+        legal_dims = True
+        fail_reason = ""
+    # end if
+
     if legal_dims and levnm:
         # <hvar> should be legal, find the correct local name for the
         #    vertical dimension
@@ -808,7 +828,8 @@ def get_dimension_info(hvar):
             raise ValueError(f"Vertical dimension, '{levnm}', not found")
         # end if
     # end if
-    return vdim_name, legal_dims, fail_reason
+
+    return vdim_name, legal_dims, fail_reason, has_constituent_read
 
 def write_phys_read_subroutine(outfile, host_dict, host_vars, host_imports,
                                phys_check_fname_str, constituent_set,
@@ -850,7 +871,7 @@ def write_phys_read_subroutine(outfile, host_dict, host_vars, host_imports,
         call_string_key = f"case ('{var_stdname}')"
 
         # Extract vertical level variable:
-        levnm, call_read_field, reason = get_dimension_info(hvar)
+        levnm, call_read_field, reason, has_constituent_read = get_dimension_info(hvar)
         if hvar.get_prop_value('protected'):
             call_read_field = False
             if reason:
@@ -860,20 +881,39 @@ def write_phys_read_subroutine(outfile, host_dict, host_vars, host_imports,
             # end if
             lvar = hvar.get_prop_value('local_name')
             reason += f"{suff}{lvar} is a protected variable"
+        # end if
+
         # Set "read_field" call string:
         if call_read_field:
-            # Replace vertical dimension with local name
-            call_str = "call read_field(file, " +                             \
-                       f"'{var_stdname}', input_var_names(:,name_idx), "
-            if levnm is not None:
-                call_str += f"'{levnm}', "
-            # end if
-            err_on_not_found_string = ""
-            if var_stdname in vars_init_value:
-                # if initial value is available, do not throw error when not found in initial condition file.
-                err_on_not_found_string = ", error_on_not_found=.false."
+            if has_constituent_read:
+                # Special case for constituent-dimension variables.
+                call_str = f"call read_constituent_dimensioned_field(const_props, file, '{var_stdname}', input_var_names(:,name_idx), "
+                if levnm is not None:
+                    call_str += f"'{levnm}', "
+                # end if
+
+                initial_value_string = ""
+                if var_stdname in vars_init_value:
+                    # if initial value is available, pass it to the read routine
+                    # so it can be used for this variable when data for some
+                    # constituent cannot be found.
+                    initial_value_string = f", initial_value={vars_init_value[var_stdname]}_kind_phys"
+                # end if
+                call_str += f"timestep, {var_locname}{initial_value_string})"
+            else:
+                # Replace vertical dimension with local name
+                call_str = "call read_field(file, " +                             \
+                           f"'{var_stdname}', input_var_names(:,name_idx), "
+                if levnm is not None:
+                    call_str += f"'{levnm}', "
+                # end if
+                err_on_not_found_string = ""
+                if var_stdname in vars_init_value:
+                    # if initial value is available, do not throw error when not found in initial condition file.
+                    err_on_not_found_string = ", error_on_not_found=.false."
+                # end if
+                call_str += f"timestep, {var_locname}{err_on_not_found_string})"
             # end if
-            call_str += f"timestep, {var_locname}{err_on_not_found_string})"
         else:
             # if initial value is assigned, then it can be ignored
             if var_stdname in vars_init_value:
@@ -903,7 +943,8 @@ def write_phys_read_subroutine(outfile, host_dict, host_vars, host_imports,
                  ["shr_kind_mod", ["SHR_KIND_CS, SHR_KIND_CL, SHR_KIND_CX"]],
                  ["physics_data", ["read_field", "find_input_name_idx",
                                    "no_exist_idx", "init_mark_idx",
-                                   "prot_no_init_idx", "const_idx"]],
+                                   "prot_no_init_idx", "const_idx",
+                                   "read_constituent_dimensioned_field"]],
                  ["cam_ccpp_cap", ["ccpp_physics_suite_variables", 
                                    "cam_constituents_array", 
                                    "cam_model_const_properties"]],
@@ -967,8 +1008,13 @@ def write_phys_read_subroutine(outfile, host_dict, host_vars, host_imports,
     outfile.write("logical                    :: use_init_variables", 2)
     outfile.blank_line()
 
+    # Prepare constituent properties pointer for later usage:
+    outfile.comment("Get constituent properties pointer:", 2)
+    outfile.write("const_props => cam_model_const_properties()", 2)
+    outfile.blank_line()
+
     # Initialize variables:
-    outfile.comment("Initalize missing and non-initialized variables strings:",
+    outfile.comment("Initialize missing and non-initialized variables strings:",
                     2)
     outfile.write("missing_required_vars = ' '", 2)
     outfile.write("protected_non_init_vars = ' '", 2)
@@ -1103,7 +1149,6 @@ def write_phys_read_subroutine(outfile, host_dict, host_vars, host_imports,
     # Read in constituent data
     outfile.comment("Read in constituent variables if not using init variables", 2)
     outfile.write("field_data_ptr => cam_constituents_array()", 2)
-    outfile.write("const_props => cam_model_const_properties()", 2)
     outfile.blank_line()
     outfile.comment("Iterate over all registered constituents", 2)
     outfile.write("do constituent_idx = 1, size(const_props)", 2)
@@ -1191,7 +1236,7 @@ def write_phys_check_subroutine(outfile, host_dict, host_vars, host_imports,
         call_string_key = f"case ('{var_stdname}')"
 
         # Extract vertical level variable:
-        levnm, call_check_field, reason = get_dimension_info(hvar)
+        levnm, call_check_field, reason, has_constituent_read = get_dimension_info(hvar)
 
         # Set "check_field" call string:
         if call_check_field:

src/physics/utils/physics_data.F90

@@ -26,6 +26,10 @@ module physics_data
       module procedure check_field_3d
    end interface check_field
 
+   interface read_constituent_dimensioned_field
+      module procedure read_constituent_dimensioned_field_2d
+   end interface read_constituent_dimensioned_field
+
 !==============================================================================
 CONTAINS
 !==============================================================================
@@ -325,6 +329,197 @@ subroutine read_field_3d(file, std_name, var_names, vcoord_name,           &
       end if
    end subroutine read_field_3d
 
+   subroutine read_constituent_dimensioned_field_2d(const_props, file, std_name, base_var_names, timestep, field_array, initial_value)
+      use shr_assert_mod,       only: shr_assert_in_domain
+      use shr_sys_mod,          only: shr_sys_flush
+      use pio,                  only: file_desc_t, var_desc_t
+      use spmd_utils,           only: masterproc
+      use cam_pio_utils,        only: cam_pio_find_var
+      use cam_abortutils,       only: endrun
+      use cam_logfile,          only: iulog
+      use cam_field_read,       only: cam_read_field
+      use ccpp_constituent_prop_mod, only: ccpp_constituent_prop_ptr_t
+      use phys_vars_init_check, only: mark_as_read_from_file
+      use phys_vars_init_check, only: phys_var_stdnames, input_var_names, phys_var_num
+
+      ! Dummy arguments
+      type(ccpp_constituent_prop_ptr_t),     intent(in)    :: const_props(:)      ! Constituent properties
+      type(file_desc_t),                     intent(inout) :: file
+      character(len=*),  intent(in)                        :: std_name            ! Standard name of base variable.
+      character(len=*),                      intent(in)    :: base_var_names(:)   ! "Base" name(s) used to construct variable name (base_constname)
+      integer,                               intent(in)    :: timestep            ! Timestep to read [count]
+      real(kind_phys),                       intent(inout) :: field_array(:,:)    ! Output field array (ncol, pcnst)
+      real(kind_phys), optional,             intent(in)    :: initial_value       ! Default value if not found
+
+      ! Local variables
+      logical                          :: var_found
+      character(len=128)               :: constituent_name
+      character(len=256)               :: file_var_name
+      character(len=256)               :: found_name
+      character(len=512)               :: missing_vars
+      type(var_desc_t)                 :: vardesc
+      real(kind_phys), allocatable     :: buffer(:)
+      integer                          :: const_idx, base_idx
+      integer                          :: ierr
+      real(kind_phys)                  :: default_value
+      logical                          :: has_initial_value
+      logical                          :: any_missing
+
+      ! For construction of constituent short name mapping
+      character(len=128), allocatable  :: constituent_short_names(:)
+      character(len=128)               :: constituent_std_name
+      integer                          :: n
+      integer                          :: const_input_idx
+
+      character(len=*), parameter      :: subname = 'read_constituent_dimensioned_field: '
+
+      ! Check if initial value was provided
+      has_initial_value = present(initial_value)
+
+      ! Initialize tracking variables
+      any_missing = .false.
+      missing_vars = ''
+
+      ! Allocate temporary buffer
+      allocate(buffer(size(field_array, 1)), stat=ierr)
+      if (ierr /= 0) then
+         call endrun(subname//'Failed to allocate buffer')
+      end if
+
+      !REMOVECAM:
+      ! Because the constituent properties pointer contains standard names, and not input constituent names
+      ! (e.g., Q, CLDLIQ, ...) which are used in the input file names,
+      ! we have to construct a mapping of the standard names to the short input IC file names
+      ! When CAM is retired and only standard names are used for constituents, this mapping can be removed.
+      allocate(constituent_short_names(size(const_props)), stat=ierr)
+      if (ierr /= 0) then
+         call endrun(subname//'Failed to allocate constituent_short_names')
+      end if
+
+      const_shortmap_loop: do const_idx = 1, size(const_props)
+         ! Get constituent standard name.
+         call const_props(const_idx)%standard_name(constituent_std_name)
+
+         ! Check if constituent standard name is in the registry to look up its IC name
+         ! n.b. this assumes that the first IC name is the short name
+         const_input_idx = -1
+         phys_inputvar_loop: do n = 1, phys_var_num
+            if (trim(phys_var_stdnames(n)) == trim(constituent_std_name)) then
+               const_input_idx = n
+               exit phys_inputvar_loop
+            end if
+         end do
+
+         if (const_input_idx > 0) then
+            ! Use the first entry from the input_var_names -- assumed to be short name.
+            constituent_short_names(const_idx) = trim(input_var_names(1, const_input_idx))
+         else
+            ! Use the standard name itself if not found in registry.
+            constituent_short_names(const_idx) = trim(constituent_std_name)
+         end if
+      end do const_shortmap_loop
+      !END REMOVECAM
+
+      ! Initialize field array to default value (only if initial_value provided)
+      if (has_initial_value) then
+         field_array(:,:) = initial_value
+      end if
+
+      ! Loop through all possible base names to find correct base name.
+      ! Note this assumes that the same base name is used for all constituents.
+      ! i.e., there cannot be something like cam_in_cflx_Q & cflx_CLDLIQ in one file.
+      base_idx_loop: do base_idx = 1, size(base_var_names)
+         ! Loop through all constituents
+         const_idx_loop: do const_idx = 1, size(const_props)
+            ! Get constituent short name
+            constituent_name = constituent_short_names(const_idx)
+
+            ! Create file variable name: <base_var_name>_<constituent_name>
+            file_var_name = trim(base_var_names(base_idx)) // '_' // trim(constituent_name)
+
+            ! Try to find variable in file
+            var_found = .false.
+            call cam_pio_find_var(file, [file_var_name], found_name, vardesc, var_found)
+
+            if(var_found) then
+               exit base_idx_loop
+            endif
+         end do const_idx_loop
+      end do base_idx_loop
+
+      ! Once base_idx is identified, use it in the actual constituent loop:
+      const_read_loop: do const_idx = 1, size(const_props)
+         ! Get constituent short name
+         constituent_name = constituent_short_names(const_idx)
+
+         ! Create file variable name: <base_var_name>_<constituent_name>
+         file_var_name = trim(base_var_names(base_idx)) // '_' // trim(constituent_name)
+
+         ! Try to find variable in file
+         var_found = .false.
+         call cam_pio_find_var(file, [file_var_name], found_name, vardesc, var_found)
+
+         if (var_found) then
+            ! Read the variable
+            if (masterproc) then
+               write(iulog, *) 'Reading constituent-dimensioned input field, ', trim(found_name)
+               call shr_sys_flush(iulog)
+            end if
+
+            call cam_read_field(found_name, file, buffer, var_found, timelevel=timestep)
+
+            if (var_found) then
+               ! Copy to correct constituent index in field array
+               field_array(:, const_idx) = buffer(:)
+
+               ! Check for NaN values
+               call shr_assert_in_domain(field_array(:, const_idx), is_nan=.false., &
+                    varname=trim(found_name), &
+                    msg=subname//'NaN found in '//trim(found_name))
+            else
+               ! Failed to read even though variable was found
+               any_missing = .true.
+               if (len_trim(missing_vars) > 0) then
+                  missing_vars = trim(missing_vars) // ', ' // trim(file_var_name)
+               else
+                  missing_vars = trim(file_var_name)
+               end if
+            end if
+         else
+            ! Variable not found in file
+            any_missing = .true.
+            if (len_trim(missing_vars) > 0) then
+               missing_vars = trim(missing_vars) // ', ' // trim(file_var_name)
+            else
+               missing_vars = trim(file_var_name)
+            end if
+
+            if (has_initial_value) then
+               ! Use default value (already set above)
+
+               if (masterproc) then
+                  write(iulog, *) 'Constituent-dimensioned field ', trim(file_var_name), &
+                                 ' not found, using default value for constituent ', trim(constituent_name)
+                  call shr_sys_flush(iulog)
+               end if
+            end if
+         end if
+      end do const_read_loop
+
+      ! Check if we should fail due to missing variables
+      if (any_missing .and. .not. has_initial_value) then
+         call endrun(subname//'Required constituent-dimensioned variables not found: ' // trim(missing_vars))
+      end if
+
+      ! Mark the base variable as read from file (only if no errors)
+      call mark_as_read_from_file(std_name)
+
+      ! Clean up
+      deallocate(constituent_short_names)
+      deallocate(buffer)
+
+   end subroutine read_constituent_dimensioned_field_2d
+
    subroutine check_field_2d(file, var_names, timestep, current_value,        &
       stdname, min_difference, min_relative_value, is_first, diff_found)
       use pio,            only: file_desc_t, var_desc_t