def cmd(args): input_ds = Dataset(args.input_file, 'r') depth_var = input_ds.variables[args.depth_data_var] create_dir_for_file(args.output_file) output_ds = Dataset(args.output_file, 'w') add_missing_dim_vars(input_ds, output_ds, depth_var.dimensions) depth_data = depth_var[:] depth_data *= args.depth_factor mask = np.ma.masked_where(depth_data <= 0.0, np.ones(depth_data.shape), copy=False) output_var = output_ds.createVariable(names.VAR_MASK, 'u1', dimensions=depth_var.dimensions, fill_value=0) output_var[:] = mask input_ds.close() output_ds.close()
def cmd(args): in_ds = Dataset(args.input_file, 'r') in_grid, in_grid_dim_names = \ init_grid_from_vars(in_ds.variables[args.lon_name], in_ds.variables[args.lat_name]) scalar_vars, vector_vars = split_scalar_and_vector_vars(args.var_names) create_dir_for_file(args.output_file) out_ds = Dataset(args.output_file, 'w') add_missing_dim_vars(in_ds, out_ds, in_grid_dim_names) grid_ds = Dataset(args.grid_file, 'r') grid_proj_var = grid_ds.variables[names.VAR_PROJECTION] converter = init_converter_from_proj_var(grid_proj_var) print 'Calculating coordinates of grid points:' xx, yy = [], [] for i in xrange(in_grid.shape[0]): _progress(i, in_grid.shape[0]) row_xx, row_yy = converter.convert_points(in_grid[i, :, 1], in_grid[i, :, 0]) xx.append(row_xx) yy.append(row_yy) xx = np.concatenate(xx) yy = np.concatenate(yy) _progress(in_grid.shape[0], in_grid.shape[0]) out_proj_var = out_ds.createVariable(names.VAR_PROJECTION, grid_proj_var.dtype) copy_nc_attributes(grid_proj_var, out_proj_var) out_x_var = out_ds.createVariable(args.x_name, xx.dtype, dimensions=in_grid_dim_names) copy_nc_attributes(grid_ds.variables[names.DIMVAR_X], out_x_var) out_x_var[:, :] = xx out_y_var = out_ds.createVariable(args.y_name, yy.dtype, dimensions=in_grid_dim_names) copy_nc_attributes(grid_ds.variables[names.DIMVAR_Y], out_y_var) out_y_var[:, :] = yy grid_ds.close() if len(scalar_vars) > 0: print 'Processing scalar fields:' for var_name in scalar_vars: print var_name in_var = in_ds.variables[var_name] add_missing_dim_vars(in_ds, out_ds, in_var.dimensions) out_var = out_ds.createVariable(var_name, in_var.dtype, dimensions=in_var.dimensions) copy_nc_attributes(in_var, out_var) iter_mask = np.ones((len(in_var.shape, )), dtype=bool) iter_mask[-MAX_COPY_DIM_COUNT:] = False dim_iterator = DimIterator(in_var.shape, None, iter_mask) write_op_count = len(dim_iterator) for write_op_num, slc in enumerate(dim_iterator.slice_tuples()): _progress(write_op_num, write_op_count) out_var[slc] = in_var[slc] _progress(write_op_count, write_op_count) if len(vector_vars) > 0: print 'Processing vector fields:' for var_name_pair in vector_vars: print var_name_pair in_u_var = in_ds.variables[var_name_pair[0]] in_v_var = in_ds.variables[var_name_pair[1]] if in_u_var.dimensions != in_v_var.dimensions: raise Exception() grid_dim_indices = find_dim_indices(in_grid_dim_names, in_u_var.dimensions) if any(idx is None for idx in grid_dim_indices): raise Exception() add_missing_dim_vars(in_ds, out_ds, in_u_var.dimensions) out_x_var = out_ds.createVariable('_'.join(var_name_pair) + '_' + args.x_name, in_u_var.dtype, dimensions=in_u_var.dimensions) out_x_var.projection = out_proj_var.grid_mapping_name out_y_var = out_ds.createVariable('_'.join(var_name_pair) + '_' + args.y_name, in_v_var.dtype, dimensions=in_u_var.dimensions) out_y_var.projection = out_proj_var.grid_mapping_name swap_axes = grid_dim_indices[0] > grid_dim_indices[1] iter_mask = np.ones((len(in_u_var.shape, )), dtype=bool) iter_mask[grid_dim_indices] = False dim_iterator = DimIterator(in_u_var.shape, None, iter_mask) write_op_count = len(dim_iterator) for write_op_num, slc in enumerate(dim_iterator.slice_tuples()): _progress(write_op_num, write_op_count) in_u_field = in_u_var[slc] in_v_field = in_v_var[slc] if swap_axes: in_u_field = np.swapaxes(in_u_field, grid_dim_indices[0], grid_dim_indices[1]) in_v_field = np.swapaxes(in_v_field, grid_dim_indices[0], grid_dim_indices[1]) out_x_field, out_y_field, = \ converter.convert_vectors(in_u_field, in_v_field, in_grid[..., 1], in_grid[..., 0]) if swap_axes: out_x_field = np.swapaxes(out_x_field, grid_dim_indices[0], grid_dim_indices[1]) out_y_field = np.swapaxes(out_y_field, grid_dim_indices[0], grid_dim_indices[1]) out_x_var[slc] = out_x_field out_y_var[slc] = out_y_field _progress(write_op_count, write_op_count) add_history(out_ds, get_history(in_ds)) in_ds.close() out_ds.close()
def cmd(args): scalar_vars, vector_vars = split_scalar_and_vector_vars(args.var_names) if len(vector_vars) < 1: raise Exception('No input vector fields specified.') if args.angle_file is None: args.angle_file = args.input_file angle_ds = Dataset(args.angle_file, 'r') if args.angle_name not in angle_ds.variables: raise Exception('Angle file \'%s\' does not contain angle variable ' '\'%s\'' % (args.angle_file, args.angle_name)) angle_var = angle_ds.variables[args.angle_name] angle_shape = angle_var.shape in_ds = Dataset(args.input_file, 'r') if args.dim_names is None: args.dim_names = angle_var.dimensions else: if len(args.dim_names) != len(angle_var.dimensions): raise Exception('Number of dimensions specified by the parameter ' '--dim-names is not equal to the number of ' 'dimensions of the --angle-name variable.') for dim_idx, dim_name in enumerate(args.dim_names): if dim_name not in in_ds.dimensions: raise Exception('Input file does not contain dimension ' '\'%s\'.' % dim_name) if in_ds.dimensions[dim_name].size != angle_shape[dim_idx]: raise Exception('Size of the dimension \'%s\' of the input ' 'file is not equal to the size of the ' 'dimension \'%s\' of the angle file.' % (dim_name, angle_var.dimensions[dim_idx])) rot_matrices = gen_rot_matrices_rad(angle_var[:]) \ if args.angle_units == 'rad' \ else gen_rot_matrices_deg(angle_var[:]) angle_ds.close() create_dir_for_file(args.output_file) out_ds = Dataset(args.output_file, 'w') if len(scalar_vars) > 0: print 'Processing scalar fields:' for var_name in scalar_vars: print var_name in_var = in_ds.variables[var_name] add_missing_dim_vars(in_ds, out_ds, in_var.dimensions) out_var = out_ds.createVariable(var_name, in_var.dtype, dimensions=in_var.dimensions) copy_nc_attributes(in_var, out_var) iter_mask = np.ones((len(in_var.shape, )), dtype=bool) iter_mask[-MAX_COPY_DIM_COUNT:] = False dim_iterator = DimIterator(in_var.shape, None, iter_mask) write_op_count = len(dim_iterator) for write_op_num, slc in enumerate(dim_iterator.slice_tuples()): _progress(write_op_num, write_op_count) out_var[slc] = in_var[slc] _progress(write_op_count, write_op_count) if len(vector_vars) > 0: print 'Processing vector fields:' angle_iter_mask = np.ones((len(args.dim_names),), dtype=bool) dims_per_operation = len(args.dim_names) \ if len(args.dim_names) < MAX_COPY_DIM_COUNT \ else MAX_COPY_DIM_COUNT angle_iter_mask[-dims_per_operation:] = False angle_dim_iterator = DimIterator(angle_shape, None, angle_iter_mask) angle_iter_count = len(angle_dim_iterator) for var_name_pair in vector_vars: print var_name_pair in_u_var = in_ds.variables[var_name_pair[0]] in_v_var = in_ds.variables[var_name_pair[1]] if in_u_var.dimensions != in_v_var.dimensions: raise Exception() angle_dim_indices = find_dim_indices(args.dim_names, in_u_var.dimensions) for idx, dim_name in enumerate(args.dim_names): if angle_dim_indices[idx] is None: raise Exception('Variable \'%s\' is not specified along ' 'dimension \'%s\'.' % (var_name_pair[0], dim_name)) add_missing_dim_vars(in_ds, out_ds, in_u_var.dimensions) out_u_var = out_ds.createVariable( var_name_pair[0], in_u_var.dtype, dimensions=in_u_var.dimensions) out_v_var = out_ds.createVariable( var_name_pair[1], in_v_var.dtype, dimensions=in_u_var.dimensions) fixed_indices = angle_dim_indices[-dims_per_operation:] dim_var_order = np.arange(len(fixed_indices)) dim_angle_order = np.argsort(fixed_indices) var_iter_mask = np.ones((len(in_u_var.shape, )), dtype=bool) var_iter_mask[fixed_indices] = False for angle_iter_num, angle_slc in \ enumerate(angle_dim_iterator.slice_tuples()): var_fixed_slices = np.repeat(None, len(in_u_var.shape)) var_fixed_slices[angle_dim_indices] = angle_slc var_iter = DimIterator(in_u_var.shape, var_fixed_slices, var_iter_mask) var_iter_count = len(var_iter) total_iter_count = angle_iter_count * var_iter_count for var_iter_num, var_slc in enumerate(var_iter.slice_tuples()): _progress(var_iter_num + angle_iter_num * var_iter_count, total_iter_count) in_u_field = in_u_var[var_slc] in_v_field = in_v_var[var_slc] in_u_field = np.moveaxis(in_u_field, dim_var_order, dim_angle_order) in_v_field = np.moveaxis(in_v_field, dim_var_order, dim_angle_order) out_u_field, out_v_field = \ apply_rot_matrices(in_u_field, in_v_field, rot_matrices[(slice(None), slice(None)) + angle_slc]) out_u_field = np.moveaxis(out_u_field, dim_angle_order, dim_var_order) out_v_field = np.moveaxis(out_v_field, dim_angle_order, dim_var_order) out_u_var[var_slc] = out_u_field out_v_var[var_slc] = out_v_field _progress(total_iter_count, total_iter_count) add_history(out_ds, get_history(in_ds)) in_ds.close() out_ds.close()
def cmd(args): in_ds = Dataset(args.input_file, 'r') # Latitude variable is mandatory. in_lat_var = in_ds.variables[args.lat_name] if 1 != len(in_lat_var.dimensions): raise Exception('\'%s\' is not 1D variable.' % args.lat_name) in_lat_dim_name = in_lat_var.dimensions[0] # Longitude variable is optional (but only for scalar fields). if args.lon_name is not None: in_lon_var = in_ds.variables[args.lon_name] if 1 != len(in_lon_var.dimensions): raise Exception('\'%s\' is not 1D variable.' % args.lon_name) in_lon_dim_name = in_lon_var.dimensions[0] if in_lat_dim_name == in_lon_dim_name: raise Exception('Latitude and longitude dimension variables ' 'can not be specified along the same dimension.') scalar_vars, vector_vars = split_scalar_and_vector_vars(args.var_names) if len(vector_vars) > 0 and args.lon_name is None: raise Exception('Vector fields cannot be processed without longitude ' 'variable.') add_north_pole = (args.add == 'north' or args.add == 'both') add_south_pole = (args.add == 'south' or args.add == 'both') lat_list = in_lat_var[:] pole_tol = lat_list.dtype.type(POLE_TOLERANCE) np_lat = lat_list.dtype.type(90) min_lat_idx, max_lat_idx = 0, -1 lat_list_ascending = True if lat_list.shape[0] > 1: if np.all(lat_list[1:] > lat_list[:-1]): pass elif np.all(lat_list[1:] < lat_list[:-1]): min_lat_idx, max_lat_idx = max_lat_idx, min_lat_idx lat_list_ascending = False else: raise Exception('Latitudes must be sorted and not be duplicated.') elif lat_list.shape[0] != 1: raise Exception('Array of latitudes must not be empty.') append_row = prepend_row = False if add_north_pole: if np.abs(lat_list[max_lat_idx] - np_lat) <= pole_tol: raise Exception('Input grid already contains grid points for the ' 'North Pole.') if lat_list_ascending: append_row = True else: prepend_row = True lat_list = _extend_axis(lat_list, np_lat, not lat_list_ascending) if add_south_pole: if np.abs(lat_list[min_lat_idx] + np_lat) <= pole_tol: raise Exception('Input grid already contains grid points for the ' 'South Pole.') if lat_list_ascending: prepend_row = True else: append_row = True lat_list = _extend_axis(lat_list, -np_lat, lat_list_ascending) create_dir_for_file(args.output_file) out_ds = Dataset(args.output_file, 'w') out_ds.createDimension( in_lat_dim_name, None if in_ds.dimensions[in_lat_dim_name].isunlimited() else lat_list.shape[0]) out_lat_var = out_ds.createVariable(args.lat_name, in_lat_var.dtype, dimensions=(in_lat_dim_name, )) copy_nc_attributes(in_lat_var, out_lat_var) out_lat_var[:] = lat_list if args.lon_name is not None: lon_list = in_lon_var[:] out_ds.createDimension( in_lon_dim_name, None if in_ds.dimensions[in_lon_dim_name].isunlimited() else lon_list.shape[0]) out_lon_var = out_ds.createVariable(args.lon_name, in_lon_var.dtype, dimensions=(in_lon_dim_name, )) copy_nc_attributes(in_lon_var, out_lon_var) out_lon_var[:] = lon_list if len(scalar_vars) > 0: print 'Processing scalar fields:' for var_name in scalar_vars: print var_name in_var = in_ds.variables[var_name] add_missing_dim_vars(in_ds, out_ds, in_var.dimensions) lat_idx, lon_idx = find_dim_indices([args.lat_name, args.lon_name], in_var.dimensions) out_var = out_ds.createVariable(var_name, in_var.dtype, dimensions=in_var.dimensions) iter_mask = np.ones((len(in_var.shape, )), dtype=bool) if lat_idx is not None: iter_mask[lat_idx] = False swap_axes = False if lon_idx is not None: iter_mask[lon_idx] = False swap_axes = lon_idx < lat_idx read_iter = DimIterator(in_var.shape, None, iter_mask) write_iter = DimIterator(out_var.shape, None, iter_mask) write_op_count = len(read_iter) for write_op_num, (read_slc, write_slc) in enumerate( izip(read_iter.slice_tuples(), write_iter.slice_tuples())): _progress(write_op_num, write_op_count) in_field = in_var[read_slc] if swap_axes: in_field = np.swapaxes(in_field, lat_idx, lon_idx) out_field = in_field if prepend_row: out_field = _extend_scalar_field(out_field, True) if append_row: out_field = _extend_scalar_field(out_field, False) if swap_axes: out_field = np.swapaxes(out_field, lat_idx, lon_idx) out_var[write_slc] = out_field _progress(write_op_count, write_op_count) if len(vector_vars) > 0: print 'Processing vector fields:' to_zero, from_zero = gen_rot_matrices_deg(lon_list, True) for var_name_pair in vector_vars: print var_name_pair in_u_var = in_ds.variables[var_name_pair[0]] in_v_var = in_ds.variables[var_name_pair[1]] if in_u_var.dimensions != in_v_var.dimensions: raise Exception() lat_idx, lon_idx = find_dim_indices([args.lat_name, args.lon_name], in_u_var.dimensions) if lat_idx is None or lon_idx is None: raise Exception() add_missing_dim_vars(in_ds, out_ds, in_u_var.dimensions) out_u_var = out_ds.createVariable(var_name_pair[0], in_u_var.dtype, dimensions=in_u_var.dimensions) out_v_var = out_ds.createVariable(var_name_pair[1], in_v_var.dtype, dimensions=in_v_var.dimensions) swap_axes = lon_idx < lat_idx iter_mask = np.ones((len(in_u_var.shape, )), dtype=bool) iter_mask[lat_idx] = iter_mask[lon_idx] = False read_iter = DimIterator(in_u_var.shape, None, iter_mask) write_iter = DimIterator(out_u_var.shape, None, iter_mask) write_op_count = len(read_iter) for write_op_num, (read_slc, write_slc) in enumerate( izip(read_iter.slice_tuples(), write_iter.slice_tuples())): _progress(write_op_num, write_op_count) in_u_field = in_u_var[read_slc] in_v_field = in_v_var[read_slc] if swap_axes: in_u_field = np.swapaxes(in_u_field, lat_idx, lon_idx) in_v_field = np.swapaxes(in_v_field, lat_idx, lon_idx) out_u_field, out_v_field = in_u_field, in_v_field if prepend_row: out_u_field, out_v_field = \ _extend_vector_field(out_u_field, out_v_field, to_zero, from_zero, True) if append_row: out_u_field, out_v_field = \ _extend_vector_field(out_u_field, out_v_field, to_zero, from_zero, False) if swap_axes: out_u_field = np.swapaxes(out_u_field, lat_idx, lon_idx) out_v_field = np.swapaxes(out_v_field, lat_idx, lon_idx) out_u_var[write_slc] = out_u_field out_v_var[write_slc] = out_v_field _progress(write_op_count, write_op_count) add_history(out_ds, get_history(in_ds)) in_ds.close() out_ds.close()
def cmd(args): weight_ds = Dataset(args.weight_file, 'r') expected_in_dims = tuple(weight_ds.variables[names.VAR_INPUT_DIMS][:]) expected_in_shape = tuple(weight_ds.variables[names.VAR_INPUT_SHAPE][:]) in_ds = Dataset(args.input_file, 'r') for dim_idx, dim_name in enumerate(expected_in_dims): if dim_name not in in_ds.dimensions or \ in_ds.dimensions[dim_name].size != \ expected_in_shape[dim_idx]: raise Exception('Weight file does not match the input file.') weight_var = weight_ds.variables[names.VAR_WEIGHTS] out_dims = weight_var.dimensions[:-1] weights = weight_var[:] indices = _split_and_squeeze( np.ma.filled(weight_ds.variables[names.VAR_INDICES][:], 0), -2) weight_ds.close() create_dir_for_file(args.output_file) out_ds = Dataset(args.output_file, 'w') dim_rename_dict = {} for dim_idx, dim_name in enumerate(out_dims): out_ds.createDimension(dim_name, weights.shape[dim_idx]) dim_rename_dict[expected_in_dims[dim_idx]] = dim_name for var_name in args.var_names: print var_name in_var = in_ds.variables[var_name] in_var_dim_indices = find_dim_indices(expected_in_dims, in_var.dimensions) in_field_dim_indices = np.argsort(in_var_dim_indices) iter_mask = np.ones((len(in_var.shape, )), dtype=bool) for dim_idx in in_var_dim_indices: if dim_idx is None: raise Exception() iter_mask[dim_idx] = False out_dim_tuple = rename_dimensions(in_var.dimensions, dim_rename_dict) add_missing_dim_vars(in_ds, out_ds, out_dim_tuple) out_var = out_ds.createVariable(var_name, in_var.dtype, dimensions=out_dim_tuple) read_iter = DimIterator(in_var.shape, None, iter_mask) write_iter = DimIterator(out_var.shape, None, iter_mask) write_op_count = len(read_iter) for write_op_num, (read_slc, write_slc) in enumerate( izip(read_iter.slice_tuples(), write_iter.slice_tuples())): _progress(write_op_num, write_op_count) in_field = in_var[read_slc] in_field, undo_order = reorder_axes(in_field, in_field_dim_indices) in_field = in_field[indices] out_field = np.ma.masked_where( np.ma.count_masked(in_field, axis=2) > 0, np.ma.sum(in_field * weights, axis=-1), copy=False) out_field, _ = reorder_axes(out_field, undo_order) out_var[write_slc] = out_field _progress(write_op_count, write_op_count) add_history(out_ds, get_history(in_ds)) in_ds.close() out_ds.close()