Skip to content

Add option to use netcdf quantization #718

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
jkbk2004 merged 5 commits into from
Nov 27, 2023
Merged

Add option to use netcdf quantization #718

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
5 commits
Select commit Hold shift + click to select a range
9a1e367
Add options to specify netcdf quantize mode
DusanJovic-NOAA Nov 2, 2023
862e84e
Merge remote-tracking branch 'origin/develop' into nc_quantize
DusanJovic-NOAA Nov 8, 2023
dc1d43e
Remove old quantize routines
DusanJovic-NOAA Nov 9, 2023
6d38dea
Merge remote-tracking branch 'origin/develop' into nc_quantize
DusanJovic-NOAA Nov 9, 2023
652c2c6
Merge remote-tracking branch 'origin/develop' into nc_quantize
DusanJovic-NOAA Nov 22, 2023
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
3 changes: 2 additions & 1 deletion io/module_fv3_io_def.F90
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
131 changes: 24 additions & 107 deletions io/module_write_netcdf.F90
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand All @@ -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

!----------------------------------------------------------------------------------------
Expand Down Expand Up @@ -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

Expand Down Expand Up @@ -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
Comment thread
DusanJovic-NOAA marked this conversation as resolved.
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)
Expand All @@ -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
Expand Down Expand Up @@ -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
Expand All @@ -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
Expand Down Expand Up @@ -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)
Expand Down Expand Up @@ -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
34 changes: 18 additions & 16 deletions io/module_wrt_grid_comp.F90
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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, &
Expand Down Expand Up @@ -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))
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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)
Expand Down