def write_pfb_output(forcings_data, num_days, out_dir, dx, start_day_num=0):
varnames = ['APCP', 'DLWR', 'DSWR', 'Press', 'SPFH', 'Temp', 'UGRD', 'VGRD']
hours_in_file = 24
for var in tqdm(varnames):
file_time_start = start_day_num * hours_in_file
file_time_stop = file_time_start + hours_in_file
start = 0 # start hour (inclusive)
stop = hours_in_file # end hour (exclusive)
# start hour and stop hour for a day
# range is number of days contained in forcings file
for bulk_collect_times in range(0, num_days):
data_obj = PFData(np.nan_to_num(forcings_data[var].values[start:stop, :, :], nan=-9999.0))
data_obj.setDX(dx)
data_obj.setDY(dx)
data_obj.setDZ(dx)
data_obj.writeFile(os.path.join(out_dir, f'WRF.{var}.{file_time_start+1:06d}_to_{file_time_stop:06d}.pfb'))
del data_obj
start = stop
stop = stop + hours_in_file # size of day in hours
file_time_start = file_time_stop
file_time_stop = file_time_stop + hours_in_file
def write_pfb(data, outfile, x0=0, y0=0, z0=0, dx=1000, dz=1000):
"""Write a 3d numpy array to a PFB output file
Parameters
----------
data : ndarray
3d numpy data array to write to pfb !(x,y,z)!
outfile : str
filename and path to write output
x0 : int, optional
initial x location (Default value = 0)
y0 : int, optional
initial y location (Default value = 0)
z0 : int, optional
initial z location (Default value = 0)
dx : int, optional
horizontal resolution (Default value = 1000)
dz : int, optional
vertical resolution (Default value = 1000)
Returns
-------
None
"""
logging.info(f'wrote pfb file {outfile}, (z,y,x)={data.shape}')
pf_data = PFData()
pf_data.setDataArray(data)
pf_data.setDX(dx)
pf_data.setDY(dx)
pf_data.setDZ(dz)
pf_data.setX(x0)
pf_data.setY(y0)
pf_data.setZ(z0)
pf_data.writeFile(outfile)
def write_array_pfb(file_name, array, dx=1, dy=1, dz=1):
# Ensure this is 3 dimensions, since parflowio requires 3 dimensions.
while array.ndim < 3:
array = array[np.newaxis, :]
if array.ndim > 3:
raise Exception(f'Too many dimensions: {array.ndim}')
from parflowio.pyParflowio import PFData
data = PFData()
data.setDataArray(array)
data.setDX(dx)
data.setDY(dy)
data.setDZ(dz)
return data.writeFile(file_name)
def pfwrite(arr, pfbfile, dx=1, dy=1, dz=1, overwrite=False):
"""
Save an ndarray to a pfb file
:param arr: ndarray to save (must be 3-dimensional)
:param pfbfile: path to pfb file
:param overwrite: whether to overwrite the file if it exists
:return: None on success. Raises Exception on failure.
"""
if os.path.exists(pfbfile) and not overwrite:
raise RuntimeError(f'{pfbfile} already exists')
assert arr.ndim == 3, 'Only 3D arrays are supported'
pfb_data = PFData()
pfb_data.setDataArray(arr)
pfb_data.setDX(dx)
pfb_data.setDY(dy)
pfb_data.setDZ(dz)
pfb_data.writeFile(pfbfile)
pfb_data.close()
def test_assign_data(self):
test = PFData()
data = np.random.random_sample((50, 49, 31))
test.setDataArray(data)
self.assertEqual(31, test.getNX())
self.assertEqual(49, test.getNY())
self.assertEqual(50, test.getNZ())
test.setDX(1)
test.setDY(1)
test.setDZ(1)
test.setX(0)
test.setY(0)
test.setZ(0)
test.setP(1)
test.setQ(1)
test.setR(1)
test.writeFile(('test_write_raw.pfb'))
test_read = PFData(('test_write_raw.pfb'))
test_read.loadHeader()
test_read.loadData()
self.assertEqual(0, test_read.getX())
self.assertEqual(0, test_read.getY())
self.assertEqual(0, test_read.getZ())
self.assertEqual(31, test_read.getNX())
self.assertEqual(49, test_read.getNY())
self.assertEqual(50, test_read.getNZ())
self.assertEqual(1, test_read.getP())
self.assertEqual(1, test_read.getQ())
self.assertEqual(1, test_read.getR())
test_data = test_read.getDataAsArray()
self.assertIsNone(
np.testing.assert_array_equal(data, test_data),
'Data written to array should exist in '
'written PFB file.')
del data
test.close()
test_read.close()
os.remove(('test_write_raw.pfb'))