Allow reading fields with a constituent dimension

src/data/registry.xml

@@ -530,8 +530,10 @@
       <dimensions>horizontal_dimension vertical_layer_dimension</dimensions>
       <ic_file_input_names>cp_or_cv_dycore</ic_file_input_names>
     </variable>
+
     <!-- Constituent Variables -->
     <!-- These are only used to set possible IC file input names, as the constituents object handles allocation. -->
+    <!-- Note: the first IC file input name should correspond to the short name of the constituent used in CAM, in order to facilitate reading constituent-dimensioned input fields. -->
     <variable local_name="q"
               standard_name="water_vapor_mixing_ratio_wrt_moist_air_and_condensed_water"
               units="kg kg-1" type="real" constituent="true">

src/data/write_init_files.py

@@ -739,14 +739,21 @@ 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 = ""
+
     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.
+    #    dimensional variables).
+    # The second dimension must either be vertical or number of constituents.
     # XXgoldyXX: If we ever need to read scalars, it would have to be
     #            done using global attributes, not 'infld'.
     ldims = len(dims)
@@ -758,7 +765,17 @@ def get_dimension_info(hvar):
         fail_reason += f"{suff}{lname} has no horizontal dimension"
         suff = "; "
     # end if
-    if (ldims > 2) or ((ldims > 1) and (not levnm)):
+
+    if has_constituent_dim:
+        # A special case where any dimensions include number_of_ccpp_constituents,
+        # in this case the variable needs to be suffixed by an underscore plus the constituent name
+        # and read and reassembled separately into host model constituent indices
+        # based on constituent name.
+        # This case will be handled separately.
+        legal_dims = True
+    elif (ldims > 2) or ((ldims > 1) and (not levnm)):
+        # The regular case where the second dimension must be vertical,
+        # and higher dimensions are unsupported.
         legal_dims = False
         unsupp = []
         for dim in dims:
@@ -785,6 +802,7 @@ def get_dimension_info(hvar):
         # end if
         suff = "; "
     # end if
+
     if legal_dims and levnm:
         # <hvar> should be legal, find the correct local name for the
         #    vertical dimension
@@ -808,7 +826,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_dim
 
 def write_phys_read_subroutine(outfile, host_dict, host_vars, host_imports,
                                phys_check_fname_str, constituent_set,
@@ -850,7 +869,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,14 +879,23 @@ 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 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), "
+            else:
+                # Replace vertical dimension with local name
+                call_str = "call read_field(file, " +                             \
+                           f"'{var_stdname}', input_var_names(:,name_idx), "
+            # end if
+
             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.
@@ -903,7 +931,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 +996,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 +1137,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)
@@ -1138,7 +1171,7 @@ def write_phys_read_subroutine(outfile, host_dict, host_vars, host_imports,
     outfile.write("call const_props(constituent_idx)%minimum(constituent_min_value, constituent_errflg, constituent_errmsg)", 5)
     outfile.write("field_data_ptr(:,:,constituent_idx) = constituent_min_value", 5)
     outfile.write("if (masterproc) then", 5)
-    outfile.write("write(iulog,*) 'Constituent ', trim(std_name), ' default value not configured.  Setting to 0.'", 6)
+    outfile.write("write(iulog,*) 'Constituent ', trim(std_name), ' default value not configured. Setting to min value of ', constituent_min_value", 6)
     outfile.write("end if", 5)
     outfile.write("end if", 4)
     outfile.write("end if", 3)
@@ -1191,7 +1224,12 @@ 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)
+
+        # If this is a constituent-indexed field, do not check it for now.
+        if has_constituent_read:
+            continue
+        # end if
 
         # Set "check_field" call string:
         if call_check_field:

src/physics/utils/physics_data.F90

@@ -8,6 +8,7 @@ module physics_data
 
    public :: find_input_name_idx
    public :: read_field
+   public :: read_constituent_dimensioned_field
    public :: check_field
 
    !Non-standard variable indices:
@@ -26,6 +27,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 +330,198 @@ 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, error_on_not_found)
+      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, check_allocate
+      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   !Parallel I/O (PIO) file type.
+      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)
+      logical, optional,                     intent(in)    :: error_on_not_found  ! Flag to error and exit 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
+      logical                          :: error_on_not_found_local
+      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=256)               :: errmsg
+
+      character(len=*), parameter      :: subname = 'read_constituent_dimensioned_field: '
+
+      if (present(error_on_not_found)) then
+         error_on_not_found_local = error_on_not_found
+      else
+         error_on_not_found_local = .true.
+      end if
+
+      ! Initialize tracking variables
+      any_missing = .false.
+      missing_vars = ''
+
+      ! Allocate temporary buffer
+      allocate(buffer(size(field_array, 1)), stat=ierr, errmsg=errmsg)
+      call check_allocate(ierr, subname, 'buffer', errmsg=errmsg)
+
+      !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, errmsg=errmsg)
+      call check_allocate(ierr, subname, 'constituent_short_names', errmsg=errmsg)
+
+      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 specified in the registry for this constituent
+         ! 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 phys_inputvar_loop
+
+         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
+
+      ! 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
+
+      if(.not. var_found .and. error_on_not_found_local) then
+         call endrun(subname//'Required constituent-dimensioned variables not found: No match for ' // trim(std_name))
+      end if
+
+      ! 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 (.not. error_on_not_found_local) then
+               ! Use default value (already set at initialization)
+
+               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. error_on_not_found_local) then
+         call endrun(subname//'Required constituent-dimensioned variables not found: ' // trim(missing_vars) // &
+                     'Make sure the constituent short name is the first in the <ic_file_input_names> list in the registry.')
+      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