Combining tiles

[1]:

import glob
import distributed
import numpy as np

dirname = "/glade/campaign/cgd/oce/people/bachman/ETP_1_20_tides/SHELF"

Request resources

[2]:

import dask_jobqueue

cluster = dask_jobqueue.PBSCluster(
    account="ncgd0011",
    #scheduler_options=dict(dashboard_address=":9797"),
    cores=1,  # The number of cores you want
    memory="12GB",  # Amount of memory
    processes=1,  # How many processes
    queue="casper",  # The queue to utilize
    local_directory="$TMPDIR",  # Use your local directory
    resource_spec="select=1:ncpus=1:mem=12GB",  # Specify resources
    walltime="04:00:00",
)
cluster

[3]:

client = distributed.Client(cluster)
client

[3]:

Client

Client-a3e4fad5-b2c2-11f1-84f3-a63dedc07991

Connection method: Cluster object Cluster type: dask_jobqueue.PBSCluster
Dashboard: http://172.17.0.2:8787/status

Cluster Info

 
[4]:

cluster.scale(2)

Note

The files are stored under paths that look like

'/glade/campaign/cgd/oce/people/bachman/ETP_1_20_tides/SHELF/ocean_shelf__1993_001.nc.0118',
'/glade/campaign/cgd/oce/people/bachman/ETP_1_20_tides/SHELF/ocean_shelf__1993_001.nc.0119',
'/glade/campaign/cgd/oce/people/bachman/ETP_1_20_tides/SHELF/ocean_shelf__1993_001.nc.0120',
`/glade/campaign/cgd/oce/people/bachman/ETP_1_20_tides/SHELF/ocean_shelf__1993_001.nc.0121',
...

Test out combining

Here’s what the files for a single day looks like.

[5]:

pattern = "1993_001"
files = sorted(glob.glob(f"{dirname}/*{pattern}*"))[:30]
files

[5]:

[]

Lets read the raw files using read_raw_files

[6]:

from mom6_tools.sections import read_raw_files
dsets = read_raw_files(files, parallel=True)

---------------------------------------------------------------------------
ModuleNotFoundError                       Traceback (most recent call last)
Cell In[6], line 1
----> 1 from mom6_tools.sections import read_raw_files
      2 dsets = read_raw_files(files, parallel=True)

File ~/checkouts/readthedocs.org/user_builds/mom6-tools/envs/latest/lib/python3.10/site-packages/mom6_tools/sections.py:6
      3 from functools import reduce
      5 import numpy as np
----> 6 import tqdm
      7 import xarray as xr
     10 def read_raw_files(
     11     paths, debug=False, parallel=False, use_cftime=True, engine="netcdf4", **kwargs
     12 ):

ModuleNotFoundError: No module named 'tqdm'

This returns a 1D list of Datasets

[7]:

dsets

---------------------------------------------------------------------------
NameError                                 Traceback (most recent call last)
Cell In[7], line 1
----> 1 dsets

NameError: name 'dsets' is not defined

Now visualize the bounding boxes for each tile using a single variable

[8]:

from mom6_tools.sections import visualize_tile

for ds in dsets:
    visualize_tile(ds.uo)

---------------------------------------------------------------------------
ModuleNotFoundError                       Traceback (most recent call last)
Cell In[8], line 1
----> 1 from mom6_tools.sections import visualize_tile
      3 for ds in dsets:
      4     visualize_tile(ds.uo)

File ~/checkouts/readthedocs.org/user_builds/mom6-tools/envs/latest/lib/python3.10/site-packages/mom6_tools/sections.py:6
      3 from functools import reduce
      5 import numpy as np
----> 6 import tqdm
      7 import xarray as xr
     10 def read_raw_files(
     11     paths, debug=False, parallel=False, use_cftime=True, engine="netcdf4", **kwargs
     12 ):

ModuleNotFoundError: No module named 'tqdm'

Now lets combine those tiles into a single Dataset. For that we will first reshape our 1D list of Datasets into 2D tiles matching the figure above (6 along x, 5 along y)

[9]:

from mom6_tools.sections import tile_raw_files

tiled = tile_raw_files(dsets, x=6, y=5)
print(f"ncols={len(tiled)}, nrows={len(tiled[0])}")

---------------------------------------------------------------------------
ModuleNotFoundError                       Traceback (most recent call last)
Cell In[9], line 1
----> 1 from mom6_tools.sections import tile_raw_files
      3 tiled = tile_raw_files(dsets, x=6, y=5)
      4 print(f"ncols={len(tiled)}, nrows={len(tiled[0])}")

File ~/checkouts/readthedocs.org/user_builds/mom6-tools/envs/latest/lib/python3.10/site-packages/mom6_tools/sections.py:6
      3 from functools import reduce
      5 import numpy as np
----> 6 import tqdm
      7 import xarray as xr
     10 def read_raw_files(
     11     paths, debug=False, parallel=False, use_cftime=True, engine="netcdf4", **kwargs
     12 ):

ModuleNotFoundError: No module named 'tqdm'

[10]:

from mom6_tools.sections import combine_manual

---------------------------------------------------------------------------
ModuleNotFoundError                       Traceback (most recent call last)
Cell In[10], line 1
----> 1 from mom6_tools.sections import combine_manual

File ~/checkouts/readthedocs.org/user_builds/mom6-tools/envs/latest/lib/python3.10/site-packages/mom6_tools/sections.py:6
      3 from functools import reduce
      5 import numpy as np
----> 6 import tqdm
      7 import xarray as xr
     10 def read_raw_files(
     11     paths, debug=False, parallel=False, use_cftime=True, engine="netcdf4", **kwargs
     12 ):

ModuleNotFoundError: No module named 'tqdm'

[11]:

combined = combine_manual(tiled)
combined

---------------------------------------------------------------------------
NameError                                 Traceback (most recent call last)
Cell In[11], line 1
----> 1 combined = combine_manual(tiled)
      2 combined

NameError: name 'combine_manual' is not defined

Make sure there are no artifacts

[12]:

import cf_xarray

combined.uo.cf.isel(Z=0, T=0).plot(robust=True)

---------------------------------------------------------------------------
ModuleNotFoundError                       Traceback (most recent call last)
Cell In[12], line 1
----> 1 import cf_xarray
      3 combined.uo.cf.isel(Z=0, T=0).plot(robust=True)

ModuleNotFoundError: No module named 'cf_xarray'

[13]:

combined.uo.cf.isel(Z=0, T=0).cf.diff("X").plot(robust=True)

---------------------------------------------------------------------------
NameError                                 Traceback (most recent call last)
Cell In[13], line 1
----> 1 combined.uo.cf.isel(Z=0, T=0).cf.diff("X").plot(robust=True)

NameError: name 'combined' is not defined

Combine files in parallel

  1. The idea is that we’ll use dask.delayed to execute one task per day i.e. we parallelize across output files

  2. This task will

    • synchronously read in all files for that day using read_raw_files,

    • assemble them to one xarray.Dataset using tile_raw_files and combine_manual

    • write that Dataset to a new file with compression enabled

The following function does all these steps

[14]:

def all_steps(pattern, dirname, scheduler="sync"):
    import glob
    import dask
    from mom_tools.sections import combine_manual, read_raw_files, tile_raw_files

    # compressor = zarr.Blosc(cname="zstd", clevel=3, shuffle=2)
    # encoding = {var: {"compressor": compressor} for var in concat}

    # complevel=4 makes no difference
    compr_dict = dict(zlib=True, complevel=1, _FillValue=None)

    globstr = f"{dirname}/*{pattern}*"
    files = sorted(glob.glob(f"{dirname}/*{pattern}*"))

    with dask.config.set(scheduler=scheduler):
        # only parallelize if not being executed in a delayed task
        raw_files_list = read_raw_files(files, parallel= scheduler != "sync")
        if not raw_files_list:
            raise ValueError(f"bad pattern: {pattern}; reading {globstr}")

        nfiles = len(raw_files_list)
        # make sure number of files is what I expect
        if nfiles != 30:
            raise ValueError(f"wrong number of files {nfiles} for pattern: {pattern}")

        # reshape into 2D list of lists
        # then concatenate along rows, then columns
        combined1 = combine_manual(tile_raw_files(raw_files_list[:30], 6, 5))

        # write to file
        name = f"{dirname}/compressed/ocean_shelf_{pattern}.nc"
        combined1.to_netcdf(
            name,
            unlimited_dims=["time"],
            encoding=dict.fromkeys(combined1.variables, compr_dict),
        )
    return pattern

Test

[15]:

ds = all_steps("1993_001", dirname)

---------------------------------------------------------------------------
ModuleNotFoundError                       Traceback (most recent call last)
Cell In[15], line 1
----> 1 ds = all_steps("1993_001", dirname)

Cell In[14], line 4, in all_steps(pattern, dirname, scheduler)
      2 import glob
      3 import dask
----> 4 from mom_tools.sections import combine_manual, read_raw_files, tile_raw_files
      6 # compressor = zarr.Blosc(cname="zstd", clevel=3, shuffle=2)
      7 # encoding = {var: {"compressor": compressor} for var in concat}
      8
      9 # complevel=4 makes no difference
     10 compr_dict = dict(zlib=True, complevel=1, _FillValue=None)

ModuleNotFoundError: No module named 'mom_tools'

Combine

Now determine all unique year_day files

[16]:

allfiles = sorted(glob.glob(f"{dirname}/*_*.nc.*"))
patterns = np.unique([file[-16:-8] for file in allfiles])
patterns

[16]:

array([], dtype=float64)

Now construct a list of delayed tasks

[17]:

tasks = [dask.delayed(all_steps)(pattern, dirname) for pattern in patterns]

And execute!

[18]:

results = dask.compute(*tasks, scheduler=client);

---------------------------------------------------------------------------
NameError                                 Traceback (most recent call last)
Cell In[18], line 1
----> 1 results = dask.compute(*tasks, scheduler=client);

NameError: name 'dask' is not defined