Newer
Older
import netCDF4 as nc
import numpy as np
def split_dataset(input_file, output_pattern, dim_name, chunk_size):
# Load the input dataset
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
ds = nc.Dataset(input_file, 'r', format='NETCDF4')
dim_size = len(ds.dimensions[dim_name])
# Calculate the number of chunks
num_chunks = int(np.ceil(dim_size / chunk_size))
# Loop through each chunk
for i in range(num_chunks-1):
# Determine the start and end indices of this chunk
start = i * chunk_size
end = min((i + 1) * chunk_size, dim_size)
# Slicing along our dimension of interest
slice_indices = slice(start, end)
# Create a new output file for this chunk
output_file = output_pattern.format(i)
rootgrp = nc.Dataset(output_file, 'w', format='NETCDF4')
# Copy dimensions
for name, dim in ds.dimensions.items():
# Adjust the dimension size for the split dimension
if name == dim_name:
dim_size = len(range(start, end))
else:
dim_size = len(dim) if not dim.isunlimited() else None
rootgrp.createDimension(name, dim_size)
# Copy variables
for name, var in ds.variables.items():
var.set_auto_maskandscale(False)
outVar = rootgrp.createVariable(name, var.datatype, var.dimensions)
outVar.set_auto_maskandscale(False)
# Copy variable attributes
if name != 'gs_1c_spect': # The original file has bad metadata for this, and possibly other fields
outVar.setncatts({k: var.getncattr(k) for k in var.ncattrs()})
# Divide variable data for the split dimension, keep others as is
if dim_name in var.dimensions:
outVar[:,] = var[slice_indices,]
else:
outVar[:,] = var[:,]
# Copy global attributes
rootgrp.setncatts({k: ds.getncattr(k) for k in ds.ncattrs()})
rootgrp.close()
ds.close()
# Call the function
# split_dataset('input.nc', 'output_{}.nc', 'time', 10)