diff --git a/io/module_fv3_io_def.F90 b/io/module_fv3_io_def.F90 index fd9c129e0..2689ef1c2 100644 --- a/io/module_fv3_io_def.F90 +++ b/io/module_fv3_io_def.F90 @@ -29,7 +29,8 @@ module module_fv3_io_def real,dimension(:),allocatable :: cen_lon, cen_lat real,dimension(:),allocatable :: lon1, lat1, lon2, lat2, dlon, dlat real,dimension(:),allocatable :: stdlat1, stdlat2, dx, dy - integer,dimension(:),allocatable :: ideflate, nbits, zstandard_level + integer,dimension(:),allocatable :: ideflate, quantize_nsd, zstandard_level + character(len=esmf_maxstr),dimension(:),allocatable :: quantize_mode integer,dimension(:),allocatable :: ichunk2d, jchunk2d, ichunk3d, jchunk3d, kchunk3d end module module_fv3_io_def diff --git a/io/module_write_netcdf.F90 b/io/module_write_netcdf.F90 index 86650c6e7..e9670945b 100644 --- a/io/module_write_netcdf.F90 +++ b/io/module_write_netcdf.F90 @@ -10,7 +10,7 @@ module module_write_netcdf use mpi use esmf use netcdf - use module_fv3_io_def,only : ideflate, nbits, zstandard_level, & + use module_fv3_io_def,only : ideflate, quantize_mode, quantize_nsd, zstandard_level, & ichunk2d,jchunk2d,ichunk3d,jchunk3d,kchunk3d, & dx,dy,lon1,lat1,lon2,lat2, & time_unlimited @@ -21,11 +21,6 @@ module module_write_netcdf logical :: par - interface quantize_array - module procedure quantize_array_3d - module procedure quantize_array_4d - end interface - contains !---------------------------------------------------------------------------------------- @@ -86,6 +81,7 @@ subroutine write_netcdf(wrtfb, filename, & integer, dimension(:), allocatable :: dimids_2d, dimids_3d, dimids, chunksizes integer, dimension(:), allocatable :: varids integer :: xtype + integer :: quant_mode integer :: ishuffle logical shuffle @@ -358,10 +354,10 @@ subroutine write_netcdf(wrtfb, filename, & ncerr = nf90_def_var_chunking(ncid, varids(i), NF90_CHUNKED, chunksizes) ; NC_ERR_STOP(ncerr) end if - ishuffle = NF90_SHUFFLE - ! shuffle filter off for 3d fields using lossy compression - if (rank == 3 .and. nbits(grid_id) > 0) then - ishuffle = NF90_NOSHUFFLE + ishuffle = NF90_NOSHUFFLE + ! shuffle filter on when using lossy compression + if ( quantize_nsd(grid_id) > 0) then + ishuffle = NF90_SHUFFLE end if if (ideflate(grid_id) > 0) then ncerr = nf90_def_var_deflate(ncid, varids(i), ishuffle, 1, ideflate(grid_id)) ; NC_ERR_STOP(ncerr) @@ -370,6 +366,24 @@ subroutine write_netcdf(wrtfb, filename, & ncerr = nf90_def_var_zstandard(ncid, varids(i), zstandard_level(grid_id)) ; NC_ERR_STOP(ncerr) end if + ! turn on quantize only for 3d variables and if requested + if (rank == 3 .and. quantize_nsd(grid_id) > 0) then + ! nf90_quantize_bitgroom = 1 + ! nf90_quantize_granularbr = 2 + ! nf90_quantize_bitround = 3 (nsd is number of bits) + if (trim(quantize_mode(grid_id)) == 'quantize_bitgroom') then + quant_mode = 1 + else if (trim(quantize_mode(grid_id)) == 'quantize_granularbr') then + quant_mode = 2 + else if (trim(quantize_mode(grid_id)) == 'quantize_bitround') then + quant_mode = 3 + else + if (mype==0) write(0,*)'Unknown quantize_mode ', trim(quantize_mode(grid_id)) + call ESMF_Finalize(endflag=ESMF_END_ABORT) + endif + + ncerr = nf90_def_var_quantize(ncid, varids(i), quant_mode, quantize_nsd(grid_id)) ; NC_ERR_STOP(ncerr) + end if end if if (par) then @@ -626,14 +640,6 @@ subroutine write_netcdf(wrtfb, filename, & if (typekind == ESMF_TYPEKIND_R4) then if (par) then call ESMF_FieldGet(fcstField(i), localDe=0, farrayPtr=array_r4_3d, rc=rc); ESMF_ERR_RETURN(rc) - if ((ideflate(grid_id) > 0 .or. zstandard_level(grid_id) > 0) .and. nbits(grid_id) > 0) then - dataMax = maxval(array_r4_3d) - dataMin = minval(array_r4_3d) - call mpi_allreduce(mpi_in_place,dataMax,1,mpi_real4,mpi_max,mpi_comm,ierr) - call mpi_allreduce(mpi_in_place,dataMin,1,mpi_real4,mpi_min,mpi_comm,ierr) - call quantize_array(array_r4_3d, dataMin, dataMax, nbits(grid_id), compress_err(i)) - call mpi_allreduce(mpi_in_place,compress_err(i),1,mpi_real4,mpi_max,mpi_comm,ierr) - end if ncerr = nf90_put_var(ncid, varids(i), values=array_r4_3d, start=start_idx); NC_ERR_STOP(ncerr) else if (is_cubed_sphere) then @@ -642,17 +648,11 @@ subroutine write_netcdf(wrtfb, filename, & call ESMF_ArrayGather(array, array_r4_3d_cube(:,:,:,t), rootPet=0, tile=t, rc=rc); ESMF_ERR_RETURN(rc) end do if (mype==0) then - if ((ideflate(grid_id) > 0 .or. zstandard_level(grid_id) > 0) .and. nbits(grid_id) > 0) then - call quantize_array(array_r4_3d_cube, minval(array_r4_3d_cube), maxval(array_r4_3d_cube), nbits(grid_id), compress_err(i)) - end if ncerr = nf90_put_var(ncid, varids(i), values=array_r4_3d_cube, start=start_idx); NC_ERR_STOP(ncerr) end if else call ESMF_FieldGather(fcstField(i), array_r4_3d, rootPet=0, rc=rc); ESMF_ERR_RETURN(rc) if (mype==0) then - if ((ideflate(grid_id) > 0 .or. zstandard_level(grid_id) > 0) .and. nbits(grid_id) > 0) then - call quantize_array(array_r4_3d, minval(array_r4_3d), maxval(array_r4_3d), nbits(grid_id), compress_err(i)) - end if ncerr = nf90_put_var(ncid, varids(i), values=array_r4_3d, start=start_idx); NC_ERR_STOP(ncerr) end if end if @@ -688,17 +688,6 @@ subroutine write_netcdf(wrtfb, filename, & end do ! end fieldCount - if ((ideflate(grid_id) > 0 .or. zstandard_level(grid_id) > 0) .and. nbits(grid_id) > 0 .and. do_io) then - ncerr = nf90_redef(ncid=ncid); NC_ERR_STOP(ncerr) - do i=1, fieldCount - if (compress_err(i) > 0) then - ncerr = nf90_put_att(ncid, varids(i), 'max_abs_compression_error', compress_err(i)); NC_ERR_STOP(ncerr) - ncerr = nf90_put_att(ncid, varids(i), 'nbits', nbits(grid_id)); NC_ERR_STOP(ncerr) - end if - end do - ncerr = nf90_enddef(ncid=ncid); NC_ERR_STOP(ncerr) - end if - if (.not. par) then deallocate(array_r4) deallocate(array_r8) @@ -923,77 +912,5 @@ subroutine add_dim(ncid, dim_name, dimid, grid, mype, rc) end subroutine add_dim -!---------------------------------------------------------------------------------------- - subroutine quantize_array_3d(array, dataMin, dataMax, nbits, compress_err) - - real(4), dimension(:,:,:), intent(inout) :: array - real(4), intent(in) :: dataMin, dataMax - integer, intent(in) :: nbits - real(4), intent(out) :: compress_err - - real(4) :: scale_fact, offset - real(4), dimension(:,:,:), allocatable :: array_save - ! Lossy compression if nbits>0. - ! The floating point data is quantized to improve compression - ! See doi:10.5194/gmd-10-413-2017. The method employed - ! here is identical to the 'scaled linear packing' method in - ! that paper, except that the data are scaling into an arbitrary - ! range (2**nbits-1 not just 2**16-1) and are stored as - ! re-scaled floats instead of short integers. - ! The zlib algorithm does almost as - ! well packing the re-scaled floats as it does the scaled - ! integers, and this avoids the need for the client to apply the - ! rescaling (plus it allows the ability to adjust the packing - ! range). - scale_fact = (dataMax - dataMin) / (2**nbits-1) - offset = dataMin - if (scale_fact > 0.) then - allocate(array_save, source=array) - array = scale_fact*(nint((array_save - offset) / scale_fact)) + offset - ! compute max abs compression error - compress_err = maxval(abs(array_save-array)) - deallocate(array_save) - else - ! field is constant - compress_err = 0. - end if - end subroutine quantize_array_3d - - subroutine quantize_array_4d(array, dataMin, dataMax, nbits, compress_err) - - real(4), dimension(:,:,:,:), intent(inout) :: array - real(4), intent(in) :: dataMin, dataMax - integer, intent(in) :: nbits - real(4), intent(out) :: compress_err - - real(4) :: scale_fact, offset - real(4), dimension(:,:,:,:), allocatable :: array_save - - ! Lossy compression if nbits>0. - ! The floating point data is quantized to improve compression - ! See doi:10.5194/gmd-10-413-2017. The method employed - ! here is identical to the 'scaled linear packing' method in - ! that paper, except that the data are scaling into an arbitrary - ! range (2**nbits-1 not just 2**16-1) and are stored as - ! re-scaled floats instead of short integers. - ! The zlib algorithm does almost as - ! well packing the re-scaled floats as it does the scaled - ! integers, and this avoids the need for the client to apply the - ! rescaling (plus it allows the ability to adjust the packing - ! range). - scale_fact = (dataMax - dataMin) / (2**nbits-1) - offset = dataMin - if (scale_fact > 0.) then - allocate(array_save, source=array) - array = scale_fact*(nint((array_save - offset) / scale_fact)) + offset - ! compute max abs compression error - compress_err = maxval(abs(array_save-array)) - deallocate(array_save) - else - ! field is constant - compress_err = 0. - end if - end subroutine quantize_array_4d - !---------------------------------------------------------------------------------------- end module module_write_netcdf diff --git a/io/module_wrt_grid_comp.F90 b/io/module_wrt_grid_comp.F90 index ec8135217..bcca85ab4 100644 --- a/io/module_wrt_grid_comp.F90 +++ b/io/module_wrt_grid_comp.F90 @@ -36,7 +36,8 @@ module module_wrt_grid_comp n_group, num_files, & filename_base, output_grid, output_file, & imo,jmo,ichunk2d,jchunk2d, & - ichunk3d,jchunk3d,kchunk3d,nbits, & + ichunk3d,jchunk3d,kchunk3d, & + quantize_mode,quantize_nsd, & nsout => nsout_io, & cen_lon, cen_lat, & lon1, lat1, lon2, lat2, dlon, dlat, & @@ -361,7 +362,8 @@ subroutine wrt_initialize_p1(wrt_comp, imp_state_write, exp_state_write, clock, allocate(jchunk3d(ngrids)) allocate(kchunk3d(ngrids)) allocate(ideflate(ngrids)) - allocate(nbits(ngrids)) + allocate(quantize_mode(ngrids)) + allocate(quantize_nsd(ngrids)) allocate(zstandard_level(ngrids)) allocate(wrt_int_state%out_grid_info(ngrids)) @@ -472,28 +474,33 @@ subroutine wrt_initialize_p1(wrt_comp, imp_state_write, exp_state_write, clock, call ESMF_ConfigGetAttribute(config=CF,value=zstandard_level(n),default=0,label ='zstandard_level:',rc=rc) if (zstandard_level(n) < 0) zstandard_level(n)=0 - call ESMF_ConfigGetAttribute(config=CF,value=nbits(n),default=0,label ='nbits:',rc=rc) - ! zlib compression flag call ESMF_ConfigGetAttribute(config=CF,value=ideflate(n),default=0,label ='ideflate:',rc=rc) if (ideflate(n) < 0) ideflate(n)=0 - call ESMF_ConfigGetAttribute(config=CF,value=nbits(n),default=0,label ='nbits:',rc=rc) - if (ideflate(n) > 0 .and. zstandard_level(n) > 0) then write(0,*)"wrt_initialize_p1: zlib and zstd compression cannot be both enabled at the same time" call ESMF_LogWrite("wrt_initialize_p1: zlib and zstd compression cannot be both enabled at the same time",ESMF_LOGMSG_ERROR,rc=RC) call ESMF_Finalize(endflag=ESMF_END_ABORT) end if + ! quantize_mode and quantize_nsd + call ESMF_ConfigGetAttribute(config=CF,value=quantize_mode(n),default='quantize_bitgroom',label='quantize_mode:',rc=rc) + call ESMF_ConfigGetAttribute(config=CF,value=quantize_nsd(n),default=0,label='quantize_nsd:',rc=rc) + + if (.NOT. (trim(quantize_mode(n))=='quantize_bitgroom' & + .OR. trim(quantize_mode(n))=='quantize_granularbr' & + .OR. trim(quantize_mode(n))=='quantize_bitround') ) then + write(0,*)"wrt_initialize_p1: unknown quantize_mode ", trim(quantize_mode(n)) + call ESMF_LogWrite("wrt_initialize_p1: wrt_initialize_p1: unknown quantize_mode "//trim(quantize_mode(n)),ESMF_LOGMSG_ERROR,rc=RC) + call ESMF_Finalize(endflag=ESMF_END_ABORT) + end if + if (lprnt) then - print *,'ideflate=',ideflate(n),' nbits=',nbits(n) + print *,'ideflate=',ideflate(n) + print *,'quantize_mode=',trim(quantize_mode(n)),' quantize_nsd=',quantize_nsd(n) print *,'zstandard_level=',zstandard_level(n) end if - ! nbits quantization level for lossy compression (must be between 1 and 31) - ! 1 is most compression, 31 is least. If outside this range, set to zero - ! which means use lossless compression. - if (nbits(n) < 1 .or. nbits(n) > 31) nbits(n)=0 ! lossless compression (no quantization) if (cf_output_grid /= cf) then ! destroy the temporary config object created for nest domains @@ -2386,11 +2393,6 @@ subroutine wrt_run(wrt_comp, imp_state_write, exp_state_write,clock,rc) if (ESMF_LogFoundError(rcToCheck=rc, msg=ESMF_LOGERR_PASSTHRU, line=__LINE__, file=__FILE__)) return endif - if (nbits(grid_id) /= 0) then - call ESMF_LogWrite("wrt_run: lossy compression is not supported for regional grids",ESMF_LOGMSG_ERROR,rc=RC) - call ESMF_Finalize(endflag=ESMF_END_ABORT) - end if - call write_netcdf(wrt_int_state%wrtFB(nbdl),trim(filename), & use_parallel_netcdf, wrt_mpi_comm,wrt_int_state%mype, & grid_id,rc)