Unstructured grid files axis (NOAA-GFDL#1114)

* adds the compress attribute and axis data to the netcdf file

* Fix some typos, change axis_length to a required argument, change the structured_axis to allocatables

diag_manager/diag_axis.F90

@@ -138,9 +138,11 @@ INTEGER FUNCTION diag_axis_init(name, DATA, units, cart_name, long_name, directi
     ENDIF
 
     if (use_modern_diag) then
-      diag_axis_init = fms_diag_object%fms_diag_axis_init(name, DATA, units, cart_name, long_name=long_name,&
-       & direction=direction, set_name=set_name, edges=edges, Domain=Domain, Domain2=Domain2, DomainU=DomainU, &
-       & aux=aux, req=req, tile_count=tile_count, domain_position=domain_position )
+      !TODO Passing in the axis_length because of a gnu issue where inside fms_diag_axis_init, the size of DATA
+      !was 2 which was causing the axis_data to not be written correctly...
+      diag_axis_init = fms_diag_object%fms_diag_axis_init(name, DATA, units, cart_name, size(DATA(:)),  &
+       & long_name=long_name, direction=direction, set_name=set_name, edges=edges, Domain=Domain, Domain2=Domain2, &
+       & DomainU=DomainU, aux=aux, req=req, tile_count=tile_count, domain_position=domain_position)
       return
     endif
     IF ( PRESENT(tile_count)) THEN

diag_manager/fms_diag_axis_object.F90

@@ -49,7 +49,7 @@ module fms_diag_axis_object_mod
   public :: fmsDiagAxis_type, fms_diag_axis_object_init, fms_diag_axis_object_end, &
           & get_domain_and_domain_type, diagDomain_t, &
           & DIAGDOMAIN2D_T, fmsDiagSubAxis_type, fmsDiagAxisContainer_type, fmsDiagFullAxis_type, DIAGDOMAINUG_T
-  public :: define_new_axis, define_subaxis
+  public :: define_new_axis, define_subaxis, parse_compress_att, get_axis_id_from_name
 
   !> @}
 
@@ -96,6 +96,9 @@ module fms_diag_axis_object_mod
        procedure :: get_axis_name
        procedure :: write_axis_metadata
        procedure :: write_axis_data
+       procedure :: add_structured_axis_ids
+       procedure :: get_structured_axis
+       procedure :: is_unstructured_grid
   END TYPE fmsDiagAxis_type
 
   !> @brief Type to hold the subaxis
@@ -138,6 +141,8 @@ module fms_diag_axis_object_mod
      TYPE(fmsDiagAttribute_type),allocatable , private :: attributes(:) !< Array to hold user definable attributes
      INTEGER                        , private :: num_attributes  !< Number of defined attibutes
      INTEGER                        , private :: domain_position !< The position in the doman (NORTH, EAST or CENTER)
+     integer, allocatable           , private :: structured_ids(:) !< If the axis is in the unstructured grid,
+                                                                   !! this is the axis ids of the structured axis
 
      contains
 
@@ -160,7 +165,7 @@ module fms_diag_axis_object_mod
   !!!!!!!!!!!!!!!!! DIAG AXIS PROCEDURES !!!!!!!!!!!!!!!!!
   !> @brief Initialize the axis
   subroutine register_diag_axis_obj(this, axis_name, axis_data, units, cart_name, long_name, direction,&
-  & set_name, Domain, Domain2, DomainU, aux, req, tile_count, domain_position )
+  & set_name, Domain, Domain2, DomainU, aux, req, tile_count, domain_position, axis_length )
     class(fmsDiagFullAxis_type),INTENT(out)  :: this            !< Diag_axis obj
     CHARACTER(len=*),   INTENT(in)           :: axis_name       !< Name of the axis
     class(*),           INTENT(in)           :: axis_data(:)    !< Array of coordinate values
@@ -177,6 +182,7 @@ subroutine register_diag_axis_obj(this, axis_name, axis_data, units, cart_name,
     CHARACTER(len=*),   INTENT(in), OPTIONAL :: req             !< Required field names.
     INTEGER,            INTENT(in), OPTIONAL :: tile_count      !< Number of tiles
     INTEGER,            INTENT(in), OPTIONAL :: domain_position !< Domain position, "NORTH" or "EAST"
+    integer,            intent(in), optional :: axis_length     !< The length of the axis size(axis_data(:))
 
     this%axis_name = trim(axis_name)
     this%units = trim(units)
@@ -187,12 +193,14 @@ subroutine register_diag_axis_obj(this, axis_name, axis_data, units, cart_name,
 
     select type (axis_data)
     type is (real(kind=r8_kind))
-      allocate(real(kind=r8_kind) :: this%axis_data(size(axis_data)))
+      allocate(real(kind=r8_kind) :: this%axis_data(axis_length))
       this%axis_data = axis_data
+      this%length = axis_length
       this%type_of_data = "double" !< This is what fms2_io expects in the register_field call
     type is (real(kind=r4_kind))
-      allocate(real(kind=r4_kind) :: this%axis_data(size(axis_data)))
+      allocate(real(kind=r4_kind) :: this%axis_data(axis_length))
       this%axis_data = axis_data
+      this%length = axis_length
       this%type_of_data = "float" !< This is what fms2_io expects in the register_field call
     class default
       call mpp_error(FATAL, "The axis_data in your diag_axis_init call is not a supported type. &
@@ -226,8 +234,6 @@ subroutine register_diag_axis_obj(this, axis_name, axis_data, units, cart_name,
     if (present(domain_position)) this%domain_position = domain_position
     call check_if_valid_domain_position(this%domain_position)
 
-    this%length = size(axis_data)
-
     this%direction = 0
     if (present(direction)) this%direction = direction
     call check_if_valid_direction(this%direction)
@@ -309,15 +315,15 @@ subroutine write_axis_metadata(this, fileobj, parent_axis)
         end select
       type is (FmsNetcdfUnstructuredDomainFile_t)
         select case (diag_axis%type_of_domain)
-        case (NO_DOMAIN)
-          !< Here the fileobj is in the unstructured domain, but the axis is not
-          !< Unstructured domain fileobjs can have axis that are not domain decomposed (i.e "Z" axis)
-          call register_axis(fileobj, axis_name, axis_length)
-          call register_field(fileobj, axis_name, diag_axis%type_of_data, (/axis_name/))
         case (UG_DOMAIN)
           !< Here the axis is in a unstructured domain
           call register_axis(fileobj, axis_name)
           call register_field(fileobj, axis_name, diag_axis%type_of_data, (/axis_name/))
+        case default
+          !< Here the fileobj is in the unstructured domain, but the axis is not
+          !< Unstructured domain fileobjs can have axis that are not domain decomposed (i.e "Z" axis)
+          call register_axis(fileobj, axis_name, axis_length)
+          call register_field(fileobj, axis_name, diag_axis%type_of_data, (/axis_name/))
         end select
     end select
 
@@ -383,6 +389,44 @@ subroutine write_axis_data(this, fileobj, parent_axis)
     end select
   end subroutine write_axis_data
 
+  !< @brief Determine if the axis is in the unstructured grid
+  !! @return .True. if the axis is in unstructured grid
+  pure logical function is_unstructured_grid(this)
+    class(fmsDiagAxis_type),           target, INTENT(in)    :: this        !< diag_axis obj
+
+    is_unstructured_grid = .false.
+    select type (this)
+    type is (fmsDiagFullAxis_type)
+      is_unstructured_grid = trim(this%cart_name) .eq. "U"
+    end select
+  end function is_unstructured_grid
+
+  !< @brief Adds the structured axis ids to the axis object
+  subroutine add_structured_axis_ids(this, axis_ids)
+    class(fmsDiagAxis_type),           target, INTENT(inout) :: this        !< diag_axis obj
+    integer,                                   intent(in)    :: axis_ids(2) !< axis ids to add to the axis object
+
+    select type (this)
+    type is (fmsDiagFullAxis_type)
+      allocate(this%structured_ids(2))
+      this%structured_ids = axis_ids
+    end select
+  end subroutine add_structured_axis_ids
+
+  !< @brief Get the structured axis ids from the axis object
+  !! @return the structured axis ids
+  pure function get_structured_axis(this) &
+  result(rslt)
+    class(fmsDiagAxis_type),           target, INTENT(in) :: this !< diag_axis obj
+    integer :: rslt(2)
+
+    rslt = diag_null
+    select type (this)
+    type is (fmsDiagFullAxis_type)
+      rslt = this%structured_ids
+    end select
+  end function get_structured_axis
+
   !> @brief Get the starting and ending indices of the global io domain of the axis
   subroutine get_global_io_domain(this, global_io_index)
     class(fmsDiagFullAxis_type), intent(in)  :: this               !< diag_axis obj
@@ -972,6 +1016,48 @@ pure function get_subaxes_id(this) &
 
   end function
 
+  !< @brief Parses the "compress" attribute to get the names of the two axis
+  !! @return the names of the structured axis
+  pure function parse_compress_att(compress_att) &
+  result(axis_names)
+    class(*), intent(in) :: compress_att(:) !< The compress attribute to parse
+    character(len=120)   :: axis_names(2)
+
+    integer            :: ios           !< Errorcode after parsing the compress attribute
+
+    select type (compress_att)
+      type is (character(len=*))
+        read(compress_att(1),*, iostat=ios) axis_names
+        if (ios .ne. 0) axis_names = ""
+      class default
+        axis_names = ""
+    end select
+  end function parse_compress_att
+
+  !< @brief Determine the axis id of a axis
+  !! @return Axis id
+  pure function get_axis_id_from_name(axis_name, diag_axis, naxis) &
+  result(axis_id)
+    class(fmsDiagAxisContainer_type), intent(in) :: diag_axis(:) !< Array of axis object
+    character(len=*),                 intent(in) :: axis_name    !< Name of the axis
+    integer,                          intent(in) :: naxis        !< Number of axis that have been registered
+    integer                                      :: axis_id
+
+    integer :: i !< For do loops
+
+    axis_id = diag_null
+    do i = 1, naxis
+      select type(axis => diag_axis(i)%axis)
+      type is (fmsDiagFullAxis_type)
+        if (trim(axis%axis_name) .eq. trim(axis_name)) then
+          axis_id = i
+          return
+        endif
+      end select
+    enddo
+
+  end function get_axis_id_from_name
+
 #endif
 end module fms_diag_axis_object_mod
 !> @}

diag_manager/fms_diag_file_object.F90

@@ -1118,24 +1118,33 @@ end subroutine increase_unlimited_dimension
 !< @brief Writes the axis metadata for the file
 subroutine write_axis_metadata(this, diag_axis)
   class(fmsDiagFileContainer_type), intent(inout), target :: this            !< The file object
-  class(fmsDiagAxisContainer_type), intent(in)            :: diag_axis(:)    !< Diag_axis object
+  class(fmsDiagAxisContainer_type), intent(in),    target :: diag_axis(:)    !< Diag_axis object
 
   class(fmsDiagFile_type), pointer     :: diag_file      !< Diag_file object to open
   class(FmsNetcdfFile_t),  pointer     :: fileobj        !< The fileobj to write to
-  integer                              :: i              !< For do loops
-  integer                              :: j              !< diag_file%axis_ids(i) (for less typing)
+  integer                              :: i,k            !< For do loops
   integer                              :: parent_axis_id !< Id of the parent_axis
+  integer                              :: structured_ids(2) !< Ids of the uncompress axis
+
+  class(fmsDiagAxisContainer_type), pointer :: axis_ptr !< pointer to the axis object currently writing
 
   diag_file => this%FMS_diag_file
   fileobj => diag_file%fileobj
 
   do i = 1, diag_file%number_of_axis
-    j = diag_file%axis_ids(i)
-    parent_axis_id = diag_axis(j)%axis%get_parent_axis_id()
+    axis_ptr => diag_axis(diag_file%axis_ids(i))
+    parent_axis_id = axis_ptr%axis%get_parent_axis_id()
     if (parent_axis_id .eq. DIAG_NULL) then
-      call diag_axis(j)%axis%write_axis_metadata(fileobj)
+      call axis_ptr%axis%write_axis_metadata(fileobj)
     else
-      call diag_axis(j)%axis%write_axis_metadata(fileobj, diag_axis(parent_axis_id)%axis)
+      call axis_ptr%axis%write_axis_metadata(fileobj, diag_axis(parent_axis_id)%axis)
+    endif
+
+    if (axis_ptr%axis%is_unstructured_grid()) then
+      structured_ids = axis_ptr%axis%get_structured_axis()
+      do k = 1, size(structured_ids)
+        call diag_axis(structured_ids(k))%axis%write_axis_metadata(fileobj)
+      enddo
     endif
   enddo
 
@@ -1188,9 +1197,10 @@ subroutine write_axis_data(this, diag_axis)
 
   class(fmsDiagFile_type), pointer     :: diag_file      !< Diag_file object to open
   class(FmsNetcdfFile_t),  pointer     :: fileobj        !< The fileobj to write to
-  integer                              :: i              !< For do loops
+  integer                              :: i, k           !< For do loops
   integer                              :: j              !< diag_file%axis_ids(i) (for less typing)
   integer                              :: parent_axis_id !< Id of the parent_axis
+  integer                              :: structured_ids(2) !< Ids of the uncompress axis
 
   diag_file => this%FMS_diag_file
   fileobj => diag_file%fileobj
@@ -1203,6 +1213,13 @@ subroutine write_axis_data(this, diag_axis)
     else
       call diag_axis(j)%axis%write_axis_data(fileobj, diag_axis(parent_axis_id)%axis)
     endif
+
+    if (diag_axis(j)%axis%is_unstructured_grid()) then
+      structured_ids = diag_axis(j)%axis%get_structured_axis()
+      do k = 1, size(structured_ids)
+        call diag_axis(structured_ids(k))%axis%write_axis_data(fileobj)
+      enddo
+    endif
   enddo
 
 end subroutine write_axis_data

diag_manager/fms_diag_object.F90

@@ -30,7 +30,8 @@ module fms_diag_object_mod
                             & get_diag_files_id, diag_yaml
 use fms_diag_axis_object_mod, only: fms_diag_axis_object_init, fmsDiagAxis_type, fmsDiagSubAxis_type, &
                                    &diagDomain_t, get_domain_and_domain_type, diagDomain2d_t, &
-                                   &fmsDiagAxisContainer_type, fms_diag_axis_object_end, fmsDiagFullAxis_type
+                                   &fmsDiagAxisContainer_type, fms_diag_axis_object_end, fmsDiagFullAxis_type, &
+                                   &parse_compress_att, get_axis_id_from_name
 use fms_diag_buffer_mod
 #endif
 #if defined(_OPENMP)
@@ -373,7 +374,7 @@ end function fms_register_static_field
 !> @brief Wrapper for the register_diag_axis subroutine. This is needed to keep the diag_axis_init
 !! interface the same
 !> @return Axis id
-FUNCTION fms_diag_axis_init(this, axis_name, axis_data, units, cart_name, long_name, direction,&
+FUNCTION fms_diag_axis_init(this, axis_name, axis_data, units, cart_name, axis_length, long_name, direction,&
   & set_name, edges, Domain, Domain2, DomainU, aux, req, tile_count, domain_position ) &
   & result(id)
 
@@ -382,6 +383,7 @@ FUNCTION fms_diag_axis_init(this, axis_name, axis_data, units, cart_name, long_n
   CLASS(*),           INTENT(in)           :: axis_data(:)    !< Array of coordinate values
   CHARACTER(len=*),   INTENT(in)           :: units           !< Units for the axis
   CHARACTER(len=1),   INTENT(in)           :: cart_name       !< Cartesian axis ("X", "Y", "Z", "T", "U", "N")
+  integer,            intent(in)           :: axis_length     !< The length of the axis size(axis_data(:))
   CHARACTER(len=*),   INTENT(in), OPTIONAL :: long_name       !< Long name for the axis.
   CHARACTER(len=*),   INTENT(in), OPTIONAL :: set_name        !< Name of the parent axis, if it is a subaxis
   INTEGER,            INTENT(in), OPTIONAL :: direction       !< Indicates the direction of the axis
@@ -423,7 +425,7 @@ FUNCTION fms_diag_axis_init(this, axis_name, axis_data, units, cart_name, long_n
     endif
     call axis%register(axis_name, axis_data, units, cart_name, long_name=long_name, &
       & direction=direction, set_name=set_name, Domain=Domain, Domain2=Domain2, DomainU=DomainU, aux=aux, &
-      & req=req, tile_count=tile_count, domain_position=domain_position)
+      & req=req, tile_count=tile_count, domain_position=domain_position, axis_length=axis_length)
 
     id = this%registered_axis
     call axis%set_axis_id(id)
@@ -629,6 +631,9 @@ subroutine fms_diag_axis_add_attribute(this, axis_id, att_name, att_value)
   character(len=*), intent(in) :: att_name     !< Name of the attribute
   class(*),         intent(in) :: att_value(:) !< The attribute value to add
 
+  character(len=20) :: axis_names(2) !< Names of the uncompress axis
+  integer           :: uncmx_ids(2)  !< Ids of the uncompress axis
+  integer           :: j             !< For do loops
 #ifndef use_yaml
 CALL MPP_ERROR(FATAL,"You can not use the modern diag manager without compiling with -Duse_yaml")
 #else
@@ -638,6 +643,25 @@ subroutine fms_diag_axis_add_attribute(this, axis_id, att_name, att_value)
   select type (axis => this%diag_axis(axis_id)%axis)
   type is (fmsDiagFullAxis_type)
     call axis%add_axis_attribute(att_name, att_value)
+
+    !! Axis that are in the "unstructured" domain require a "compress" attribute for the
+    !! combiner and PP. This attribute is passed in via a diag_axis_add_attribute call in the model code
+    !! The compress attribute indicates the names of the axis that were compressed
+    !! For example grid_index:compress = "grid_yt grid_xt"
+    !! The metadata and the data for these axis also needs to be written to the file
+    if (trim(att_name) .eq. "compress") then
+      !< If the attribute is the "compress" attribute, get the axis names,
+      !! and the ids of the axis and add it to the axis object so it can be written to netcdf files
+      !! that use this axis
+      axis_names = parse_compress_att(att_value)
+      do j = 1, size(axis_names)
+        uncmx_ids(j) = get_axis_id_from_name(axis_names(j), this%diag_axis, this%registered_axis)
+        if (uncmx_ids(j) .eq. diag_null) call mpp_error(FATAL, &
+          &"Error parsing the compress attribute for axis: "//trim(axis%get_axis_name())//&
+          &". Be sure that the axes in the compress attribute are registered")
+      enddo
+      call axis%add_structured_axis_ids(uncmx_ids)
+    endif
   end select
 #endif
 end subroutine fms_diag_axis_add_attribute

test_fms/diag_manager/test_modern_diag.F90

@@ -109,11 +109,12 @@ program test_modern_diag
                          set_name="land", DomainU=land_domain, aux="geolon_t geolat_t")
 
 id_z  = diag_axis_init('z',  z,  'point_Z', 'z', long_name='point_Z')
-!TODO call diag_axis_add_attribute (id_z, 'formula', 'p(n,k,j,i) = ap(k) + b(k)*ps(n,j,i)')
+call diag_axis_add_attribute (id_z, 'formula', 'p(n,k,j,i) = ap(k) + b(k)*ps(n,j,i)')
 call diag_axis_add_attribute (id_z, 'integer', 10)
 call diag_axis_add_attribute (id_z, '1d integer', (/10, 10/))
 call diag_axis_add_attribute (id_z, 'real', 10.)
 call diag_axis_add_attribute (id_x, '1d real', (/10./))
+call diag_axis_add_attribute (id_ug, 'compress', 'x y')
 
 if (id_x  .ne. 1) call mpp_error(FATAL, "The x axis does not have the expected id")
 if (id_y  .ne. 2) call mpp_error(FATAL, "The y axis does not have the expected id")