Skip to content
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

parallel computing for storing output #716

Merged
merged 5 commits into from
Dec 22, 2023
Merged

parallel computing for storing output #716

Show file tree
Hide file tree
Changes from 3 commits
Commits
Show all changes
5 commits
Select commit Hold shift + click to select a range
fd87f90
parallel computing
AminTorabi-NOAA Dec 21, 2023
77dceb2
Add LOG and option to do or avoid multiproccessing
AminTorabi-NOAA Dec 22, 2023
2b130b3
deleting pdb
AminTorabi-NOAA Dec 22, 2023
366d2a8
Addressing comments
AminTorabi-NOAA Dec 22, 2023
85defcb
Addressing comment
AminTorabi-NOAA Dec 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
194 changes: 114 additions & 80 deletions src/troute-network/troute/nhd_io.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1948,14 +1948,104 @@ def write_waterbody_netcdf(
}
)

def helper_write_flowveldepth(flowveldepth, subset_df, current_time_step,
stream_output_directory, stream_output_type, gage,
nudge_timesteps, time_dim, stream_output_internal_frequency, time_steps,
counter, t0, stream_output_timediff):
'''
Just a helper function for 'write_flowveldepth_netcdf' function for parallel computing.
'''

if stream_output_directory:
if stream_output_type =='.nc':
file_name = f"{current_time_step}.flowveldepth.nc"

elif stream_output_type=='.csv':
file_name = f"{current_time_step}.flowveldepth.csv"
# Save the data to CSV file
subset_df.to_csv(f"{stream_output_directory}/{file_name}", index=True)
LOG.debug(f"Flowveldepth data saved as CSV files in {stream_output_directory}")

elif stream_output_type=='.pkl':
file_name = f"{current_time_step}.flowveldepth.pkl"
# Save the data to Pickle file
subset_df.to_pickle(f"{stream_output_directory}/{file_name}")
LOG.debug(f"Flowveldepth data saved as PICKLE files in {stream_output_directory}")

else:
print('WRONG FORMAT')

if stream_output_directory:
if (stream_output_type =='.nc'):
# Open netCDF4 Dataset in write mode
with netCDF4.Dataset(
filename=f"{stream_output_directory}/{file_name}",
mode='w',
format='NETCDF4'
) as ncfile:

# ============ DIMENSIONS ===================
_ = ncfile.createDimension('feature_id', None)
_ = ncfile.createDimension('time_step (sec)', subset_df.iloc[:, 0::4].shape[1])
_ = ncfile.createDimension('gage', gage)
_ = ncfile.createDimension('nudge_timestep', nudge_timesteps) # Add dimension for nudge time steps

# =========== q,v,d,ndg VARIABLES ===============
for counters, var in enumerate(['flowrate', 'velocity', 'depth', 'nudge']):
QVD = ncfile.createVariable(
varname=var,
datatype=np.float32,
dimensions=('feature_id', 'time_step (sec)',),
)

QVD.units = 'm3/s m/s m m3/s'
QVD.description = f'Data for {var}'

# Prepare data for writing
data_array = subset_df.iloc[:, counters::4].to_numpy(dtype=np.float32)

# Set data for each feature_id and time_step
ncfile.variables[var][:] = data_array
feature_id = ncfile.createVariable(
varname='feature_id',
datatype=np.int32,
dimensions=('feature_id',),
)
feature_id[:] = flowveldepth.index.to_numpy(dtype=np.int32)
feature_id.units = 'None'
feature_id.description = 'Feature IDs'
###
time_step = ncfile.createVariable(
varname='time_step (sec)',
datatype=np.int32,
dimensions=('time_step (sec)',),
)
time_step[:] = np.array(time_dim[:subset_df.iloc[:, 0::4].shape[1]], dtype=np.int32)
time_step.units = 'sec'
time_step.description = 'time stamp'
# =========== GLOBAL ATTRIBUTES ===============
ncfile.setncatts(
{
'TITLE': 'OUTPUT FROM T-ROUTE',
'Time step (sec)': f'{stream_output_internal_frequency}',
'model_initialization_time': t0.strftime('%Y-%m-%d_%H:%M:%S'),
'model_reference_time': time_steps[counter].strftime('%Y-%m-%d_%H:%M:%S'),
'comment': f'The file includes {stream_output_timediff} hour data which includes {len(time_dim)} timesteps',
'code_version': '',
}
)
LOG.debug(f"Flowveldepth data saved as NetCDF files in {stream_output_directory}")

def write_flowveldepth_netcdf(stream_output_directory,
flowveldepth,
nudge,
usgs_positions_id,
t0,
stream_output_timediff,
stream_output_type,
stream_output_internal_frequency = 5):
stream_output_internal_frequency = 5,
cpu_pool = 6,
use_parallel = True):
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

It looks like this use_parallel has not been added to output.py when the function gets called? If you want to do it this way, you also need to determine when/how use_parallel is determined based on user input. I think you might be able to just drop this though and use only cpu_pool in the if/else statements. What do you think?

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Yes sure. So I can just say if cpu_pul > 1 do parallel otherwise don't. The only caveat is user have to set cpu_pool = 1 in output.py not here (the value 6 here doesn't matter)

Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

cpu_pool is already an input into nwm_output_generator in output.py, so it just needs to be added as in input into nhd_io.write_flowveldepth_netcdf()

'''
Write the results of flowveldepth and nudge to netcdf- break.
Arguments
Expand Down Expand Up @@ -2006,8 +2096,11 @@ def write_flowveldepth_netcdf(stream_output_directory,
# Create time step values based on t0
time_steps = [t0 + timedelta(hours= (i * stream_output_timediff)) for i in range(nsteps//nstep_nc)]
time_dim = [t * stream_output_internal_frequency*60 for t in range(1, int(stream_output_timediff * 60 / stream_output_internal_frequency) + 1)]
jobs = []

for counter, i in enumerate(range(0, nsteps, nstep_nc)):


# Define the range of columns for this file
start_col = i * 4
end_col = min((i + nstep_nc) * 4 , nsteps * 4)
Expand All @@ -2021,83 +2114,24 @@ def write_flowveldepth_netcdf(stream_output_directory,

# Create the file name based on the current time step
current_time_step = time_steps[counter].strftime('%Y%m%d%H%M')
if stream_output_directory:
if stream_output_type =='.nc':
file_name = f"{current_time_step}.flowveldepth.nc"

elif stream_output_type=='.csv':
file_name = f"{current_time_step}.flowveldepth.csv"
# Save the data to CSV file
subset_df.to_csv(f"{stream_output_directory}/{file_name}", index=True)
LOG.debug(f"Flowveldepth data saved as CSV files in {stream_output_directory}")

elif stream_output_type=='.pkl':
file_name = f"{current_time_step}.flowveldepth.pkl"
# Save the data to Pickle file
subset_df.to_pickle(f"{stream_output_directory}/{file_name}")
LOG.debug(f"Flowveldepth data saved as PICKLE files in {stream_output_directory}")

else:
print('WRONG FORMAT')

if stream_output_directory:
if (stream_output_type =='.nc'):
# Open netCDF4 Dataset in write mode
with netCDF4.Dataset(
filename=f"{stream_output_directory}/{file_name}",
mode='w',
format='NETCDF4'
) as ncfile:

# ============ DIMENSIONS ===================
_ = ncfile.createDimension('feature_id', None)
_ = ncfile.createDimension('time_step (sec)', subset_df.iloc[:, 0::4].shape[1])
_ = ncfile.createDimension('gage', gage)
_ = ncfile.createDimension('nudge_timestep', nudge_timesteps) # Add dimension for nudge time steps

# =========== q,v,d,ndg VARIABLES ===============
for counters, var in enumerate(['flowrate', 'velocity', 'depth', 'nudge']):
QVD = ncfile.createVariable(
varname=var,
datatype=np.float32,
dimensions=('feature_id', 'time_step (sec)',),
)

QVD.units = 'm3/s m/s m m3/s'
QVD.description = f'Data for {var}'

# Prepare data for writing
data_array = subset_df.iloc[:, counters::4].to_numpy(dtype=np.float32)

# Set data for each feature_id and time_step
ncfile.variables[var][:] = data_array
feature_id = ncfile.createVariable(
varname='feature_id',
datatype=np.int32,
dimensions=('feature_id',),
)
feature_id[:] = flowveldepth.index.to_numpy(dtype=np.int32)
feature_id.units = 'None'
feature_id.description = 'Feature IDs'
###
time_step = ncfile.createVariable(
varname='time_step (sec)',
datatype=np.int32,
dimensions=('time_step (sec)',),
)
time_step[:] = np.array(time_dim[:subset_df.iloc[:, 0::4].shape[1]], dtype=np.int32)
time_step.units = 'sec'
time_step.description = 'time stamp'
# =========== GLOBAL ATTRIBUTES ===============
ncfile.setncatts(
{
'TITLE': 'OUTPUT FROM T-ROUTE',
'Time step (sec)': f'{stream_output_internal_frequency}',
'model_initialization_time': t0.strftime('%Y-%m-%d_%H:%M:%S'),
'model_reference_time': time_steps[counter].strftime('%Y-%m-%d_%H:%M:%S'),
'comment': f'The file includes {stream_output_timediff} hour data which includes {len(time_dim)} timesteps',
'code_version': '',
}
)
LOG.debug(f"Flowveldepth data saved as NetCDF files in {stream_output_directory}")

args = (flowveldepth, subset_df, current_time_step,
stream_output_directory, stream_output_type, gage,
nudge_timesteps, time_dim, stream_output_internal_frequency, time_steps,
counter, t0, stream_output_timediff)

if use_parallel and cpu_pool > 1:
jobs.append(delayed(helper_write_flowveldepth)(*args))
LOG.debug(f"Job for step {counter} added for parallel processing.")
else:
helper_write_flowveldepth(*args)

if use_parallel and cpu_pool > 1:
try:
# Execute all jobs in parallel
with Parallel(n_jobs=cpu_pool) as parallel:
parallel(jobs)
except Exception as e:
LOG.error("Error during parallel processing: %s", e)

LOG.debug("Completed the write_flowveldepth_netcdf function")
3 changes: 2 additions & 1 deletion src/troute-nwm/src/nwm_routing/output.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -209,7 +209,8 @@ def nwm_output_generator(
t0,
int(stream_output_timediff),
stream_output_type,
stream_output_internal_frequency)
stream_output_internal_frequency,
cpu_pool)

if test:
flowveldepth.to_pickle(Path(test))
Expand Down