def func_Bs(*args):
v_names = args[-3]
time_limit_in_min = args[-2]
return_shape = args[-1]
d = {v_names[i]: args[i].reshape((-1, )) for i in range(len(args[:-3]))}
# print([v.shape for v in d.values()])
# print('lat:', d['lat'], 'lon:', d['lon'], 'prob:', d['prob'], flush=True)
# print('time:', d['time'], flush=True)
# print([v.shape for v in d.values()], flush=True)
Bs = -np.inf * np.ones(return_shape)
start_time = time.time()
timeout_msg = ""
try:
Bs = custom_timeout(
time_limit_in_min * 60,
CARDAMOMlib.load_Bs_greg_dict,
d,
# integration_method="trapezoidal",
# nr_nodes=11
)
except TimeoutError:
duration = (time.time() - start_time) / 60
timeout_msg = 'Timeout after %2.2f minutes' % duration
print(timeout_msg)
write_to_logfile(logfilename, "finished single,", "lat:", d["lat"], "lon:",
d["lon"], "prob:", d["prob"], timeout_msg)
return Bs.reshape(return_shape)
def func_pwc_mr_fd(ds_single):
# print(ds_single)
ds_res = CARDAMOMlib.compute_ds_pwc_mr_fd_greg(ds_single, comp_dict)
write_to_logfile("finished single,", "lat:", ds_single.lat.data, "lon:",
ds_single.lon.data, "prob:", ds_single.prob.data)
return ds_res
def func_Bs(*args):
v_names = args[-3]
time_limit_in_min = args[-2]
return_shape = args[-1]
d = {v_names[i]: args[i].reshape((-1, )) for i in range(len(args[:-3]))}
Bs = -np.inf * np.ones(return_shape)
start_time = time.time()
timeout_msg = ''
try:
Bs = custom_timeout(
time_limit_in_min * 60,
CARDAMOMlib.load_Bs_greg_dict,
d,
# integration_method="trapezoidal",
# nr_nodes=1001
)
except TimeoutError:
duration = (time.time() - start_time) / 60
timeout_msg = 'Timeout after %2.2f min' % duration
print(timeout_msg)
write_to_logfile(logfilename, "finished single,", "lat:", d["lat"], "lon:",
d["lon"], "prob:", d["prob"], timeout_msg)
#return Bs.reshape(1, len(d['time']), 6, 6)
return Bs.reshape(return_shape)
l = slice(s[0][0], s[0][-1] + 1, 1)
p = slice(s[1][0], s[1][-1] + 1, 1)
z[:, l, p, ...] = arr[:, l, p, ...].compute()
# -
# write header to logfile
c = Bs.chunks
nr_chunks = np.prod([len(val) for val in c])
print('nr_chunks:', nr_chunks)
nr_singles = np.prod(Bs.shape[:3])
print('nr_singles:', nr_singles)
write_to_logfile(
logfilename,
'starting:',
# nr_chunks, "chunks, ",
nr_singles,
"singles")
# +
# %%time
batchwise_to_zarr_CARDAMOM(Bs,
data_folder + filestem + output_folder + "Bs",
rm=True)
write_to_logfile(logfilename, 'done')
print('done')
# -
Bs_zarr = zarr.open(data_folder + filestem + output_folder + "Bs")
da_Bs_restricted = da.from_zarr(Bs_zarr)[:, :, :, 0, 0, 0]
def remove_timeouts(timeout):
z = zarr.open(zarr_dir_name)
da_Bs_restricted = da.from_zarr(z)[slices['lat'], slices['lon'],
slices['prob'], 0, 0, 0]
timeout_coords = np.where(da_Bs_restricted.compute() == -np.inf)
timeout_nrs = len(timeout_coords[0])
if timeout_nrs > 0:
print('number of timeouts:', timeout_nrs)
else:
print('timeouts successfully removed')
return
timeout_variables = []
for v, name in zip(variables, variable_names):
if name not in non_data_vars:
v_stack_list = []
for lat, lon, prob in zip(*[timeout_coords[k] for k in range(3)]):
v_stack_list.append(v[lat, lon, prob])
timeout_variables.append(da.stack(v_stack_list))
timeout_variables.append(
da.from_array(timeout_coords[0].reshape(-1, 1), chunks=(1, 1))) # lat
timeout_variables.append(
da.from_array(timeout_coords[1].reshape(-1, 1), chunks=(1, 1))) # lon
timeout_variables.append(
da.from_array(timeout_coords[2].reshape(-1, 1), chunks=(1, 1))) # prob
timeout_variables.append(variables[variable_names.index('time')].reshape(
1, -1)) # time
return_shape = (1, nr_times, nr_pools, nr_pools)
timeout_Bs = timeout_variables[0].map_blocks(
func_Bs,
*timeout_variables[1:],
variable_names,
timeout, # time_limit_in_min
return_shape,
new_axis=[2, 3],
chunks=return_shape,
dtype=np.float64,
meta=np.ndarray(return_shape, dtype=np.float64))
# write header to logfile
c = timeout_Bs.chunks
nr_chunks = np.prod([len(val) for val in c])
print('nr_chunks:', nr_chunks)
nr_singles = timeout_Bs.shape[0]
print('nr_singles:', nr_singles)
write_to_logfile(
logfilename,
'starting',
# nr_chunks, "chunks, ",
nr_singles,
"timeout singles,",
"timeout =",
timeout,
"min")
print('timeout =', timeout, 'min')
# do the computation
timeout_Bs_computed = timeout_Bs.compute()
# compute the z coords from the sliced Bs coords
f_lat = lambda x: slices['lat'].start + x * slices['lat'].step
f_lon = lambda x: slices['lon'].start + x * slices['lon'].step
f_prob = lambda x: slices['prob'].start + x * slices['prob'].step
timeout_coords_z_lat = np.array([f_lat(x) for x in timeout_coords[0]])
timeout_coords_z_lon = np.array([f_lon(x) for x in timeout_coords[1]])
timeout_coords_z_prob = np.array([f_prob(x) for x in timeout_coords[2]])
timeout_coords_z = (timeout_coords_z_lat, timeout_coords_z_lon,
timeout_coords_z_prob)
# update zarr file
for nr, lat, lon, prob in zip(range(timeout_nrs),
*[timeout_coords_z[k] for k in range(3)]):
z[lat, lon, prob, ...] = timeout_Bs_computed[nr, ...]
write_to_logfile(logfilename, 'done, timeout =', timeout, 'min')
print('done, timeout = ', timeout, 'min')
for s_arr, s_z in zip(arr_slices_tuples, z_slices_tuples):
z[s_z[0], s_z[1], s_z[2], ...] = arr[s_arr[0], s_arr[1], s_arr[2],
...].compute()
# write header to logfile
c = Bs.chunks
nr_chunks = np.prod([len(val) for val in c])
print('nr_chunks:', nr_chunks)
nr_singles = np.prod(Bs.shape[:3])
print('nr_singles:', nr_singles)
write_to_logfile(
logfilename,
'starting',
# nr_chunks, "chunks, ",
nr_singles,
"singles,",
"timeout =",
timeout,
"min")
print('timeout:', timeout, ' min')
# +
# %%time
batchwise_to_zarr_CARDAMOM(Bs, z, slices)
write_to_logfile(logfilename, 'done')
print('done')
# -