Top-Level API Discussion for Initial Release

[jthielen]

UPDATE 2022-09-28: This discussion is out-of-date, with a new top level API centered around a single generate_mosaic function for simple use cases, and a Pangeo Forge inspired recipe framework for more complicated ones.

---

Over the next two months, OpenMosaic will be seeing more intensive development under grant-sponsored funding. This work will be focused on making the package ready-to-use for research applications. A large part of that will be designing and implementing the API of the first released version (the other being testing). And so, what follows are my current high-level API plans for OpenMosaic after some initial brainstorming. This discussion is mostly placed here for reference's sake, however, if you happen to come across it before implementation occurs, feel free to contribute your thoughts and suggestions!

Existing API

Very provisional prototype that leaves way too much to the user to handle (see https://github.com/jthielen/OpenMosaic/blob/main/examples/process_full_conus.py). While (for now) low-level routines seem to be decent, the user-level API needs to be completely redone.

Mockups of Draft API for Primary Use-Cases

Example 1

"I want a mosaic over a given region at a single time using default data sources as simply as possible"

# Define grid
cf_attrs = {
    'grid_mapping_name': 'lambert_conformal_conic',
    'latitude_of_projection_origin': 38.5,
    'longitude_of_central_meridian': 262.5,
    'standard_parallel': 38.5,
    'earth_radius': 6371229.0
}
grid_params = {
    'x_min': -2.3e6,
    'x_max': 2.4e6,
    'dx': 2.0e3,
    'y_min': -1.5e6,
    'y_max': 1.5e6,
    'dy': 2.0e3,
    'z_min': 1.0e3,
    'nz': 24,
    'dz': 1.0e3
}
grid = openmosaic.GridDefinition(cf_attrs, grid_params)

# Define time
time = '2020-08-10 16:00Z'

# Set up dask
cluster = LocalCluster()
cluster.scale(6)
client = Client(cluster)

# Generate 3D grid of reflectivity
ds = openmosiac.generate_mosaic_with_autoload(
    grid=grid,
    time=time,
    fields=['reflectivity'],
    include_auxiliary_coords=False,
    dask_client=client,
    verbose=False
).compute()

# Output composite reflectivity
ds.max('altitude').to_netcdf('comp_refl_20200810T1600Z.nc')

Example 2

"What if I want that single time mosaic, but for lots of fields and not just composite reflectivity?"

grid, time, client = ...  # as in example 1
era5_loader = openmosaic.loaders.RDA.ERA5Loader(cache=True)

# Generate 3D fields
ds = openmosaic.generate_mosaic_with_autoload(
    grid=grid,
    time=time,
    fields=[
        'reflectivity',
        'differential_reflectivity',
        'cross_correlation_ratio',
        'spectrum_width'
    ],
    compute_fields=[
        openmosaic.calculations.VelocityDerivatives(
            azimuthal_shear=True,
            radial_divergence=False,
            dealias_velocities=True,
            dealias_profile_data=era5_loader
        ),
        openmosaic.calculations.SpecificDifferentialPhase()
    ]
)

# Derived 2D fields
ds = ds.pipe(
    openmosaic.calculations.MESH(
        analysis_data=era5_loader
    )
).pipe(
    openmosaic.calculations.LayerMaximum(
        fields=['azimuthal_shear']
        layers=[(0, 2.0e3), (3.0e3, 6.0e3)]
        layer_labels=['low_level', 'mid_level'],
        drop=True
    )
).pipe(
    openmosaic.calculations.LayerMaximum(
        fields=['reflectivity'],
        layer_labels='composite',
        drop=False
    )
)

# Output all fields
ds.to_netcdf('mosaic_combined_20200810T1600Z.nc')

What if I already have my files on disk? And no dask? And lat/lon grid?

xref #10

nexrad_files = [
    'KILX_20210119_080351.bz2',
    'KLOT_20210119_075832.bz2'
]
time = pd.Timestamp('20210119 08:00')

cf_attrs = {
    'grid_mapping_name': 'latitude_longitude',
}
grid_params = {
    'dx': 0.02,
    'dy': 0.02,
    'z_min': 1.0e3,
    'nz': 24,
    'dz': 1.0e3
}
grid = openmosaic.GridDefinition(
    cf_attrs,
    grid_params,
    set_missing_arguments='bounding'
)

# Generate 3D grid of reflectivity
ds = openmosiac.generate_mosaic(
    grid=grid,
    radars=nexrad_files,
    analysis_time=time,
    fields=['reflectivity'],
    include_auxiliary_coords=False,
    dask_client=False,
    verbose=True
).compute()

# Output composite reflectivity
ds.max('altitude').to_netcdf('comp_refl_202101190800.nc')

What if I want a workflow that distributes mosaicing across many times intelligently with dask? And how about a visual progress indicator?

(fields taken from the intermediate data target for SVRIMG for Detailed Morphology)

grid, client = ...  # as in example 1
times = pd.date_range('2020-08-10 06:00Z', '2020-08-11 00:00Z', freq='5min')
era5_loader = openmosaic.loaders.RDA.ERA5Loader(
    cache='path/to/worker-shared-storage/analysis/'
)

openmosaic.write_mosiacs_with_autoload(
    grid=grid,
    times=times,
    output='my_zarr_store.zarr',
    fields=['reflectivity'],
    compute_fields=[
        openmosaic.calculations.VelocityDerivatives(
            azimuthal_shear=True,
            radial_divergence=False,
            dealias_velocities=True,
            dealias_profile_data=era5_loader
        )
    ],
    post_gridding_operations=[
        openmosaic.calculations.EchoTopHeight(threshold=18),
        openmosaic.calculations.MESH(
            analysis_data=era5_loader
        ),
        openmosaic.calculations.LayerMaximum(
            fields=['azimuthal_shear']
            layers=[(0, 2.0e3), (3.0e3, 6.0e3)]
            layer_labels=['low_level', 'mid_level'],
            drop=True
        ),
        openmosaic.calculations.LayerMaximum(
            fields=['reflectivity'],
            layer_labels='composite',
            drop=True
        )
    ],
    radar_source_cache='path/to/worker-shared-storage/radar/',
    gridding_cache='path/to/worker-shared-storage/gridding/',
    include_auxiliary_coords=True,
    dask_client=client,
    verbose='html'
)

Core Concepts of Draft API

Typical Script Outline

• Define a grid and a time (or times)
• Set local files to use or choose autoloading of Level II volumes
• Call the mosaicing function that best fits your use case
• Apply post gridding operations and save/use output (if not already done in packaged function)

Main Components/Objects

"The Grid"

Define the target points on which to map Level II data. Can use a memory-performant (but regular grid only) openmosaic.GridDefinition or supply a custom grid xarray.Dataset.

Canned/default grids (such as those matching other common datasets) may be included in future releases.

"Data Loaders"

We need some way to fetch (and possibly munge and/or cache) remote data such as Level II volumes from AWS S3 or ERA5 reanalyses from NCAR RDA, given input requirements from the gridding operations. Simple use cases can rely on the defaults, but many other use cases will need some degree of custom configuration in how these data are fetched and operated on. These also need to be Dask serializable to enable full parallelization.

Examples of loaders to be available in initial release

• RDA.ERA5Loader (default analysis dataset...in a module since I anticipate more RDA stuff down the road that can share useful things)
• S3.LevelIILoader (default Level II dataset...again in a module since I anticipate more things loaded through S3 later on)

"Calculations"

We need to be able to calculate fields on both the source data's radial grid/volumes and the output cartesian grid. Due to the complexities involved in presenting these with the simplest user-level options, while also handling all the ways they can be operated with internally, they will be classes rather than functions. They will be able to applied with or without Dask, and can themselves have data loaders they use.

A key concept for users: specification order of calculations in all (or most) cases will be important! For example, later calculations can take previously computed fields as inputs, and dropped fields are no longer available for later calculations.

Examples of calculations to be available in initial release:

• Radial Grid
  • VelocityDerivatives (given relation, handle both azimuthal shear and radial divergence together...also includes dealiasing)
  • SpecificDifferentialPhase (KDP)
  • PyARTFilter (wrap a pyart filter)
• Output/Cartesian Grid
  • MESH (and so also SHI on which it depends)
  • LayerMaximum (e.g., for both composite reflectivity and low-/mid-level azishear)
  • EchoTopHeight

Future calculations that may be supported include:

• SL3D

Gridding Functions

The core operations of OpenMosaic will be the gridding/mosaicing functions. The Gridder object-based approach of the prior API proved unsatisfactory, and instead, different functions for different common workflows will be provided, all of which are controlled by the previously specified objects. The internals are still likely to involve some kind of "task manager" class to assign dask futures with an understanding of the workflows, but I no longer have plans for that to be part of the user-facing API.

Examples of gridding functions to be available in initial release

• generate_mosaic (mosaic(s) at given time(s) using existing NEXRAD files or PyART radar objects, output in-memory as xarray.Dataset. Times, if multiple, are looped over.)
• generate_mosaic_with_autoload (mosaic(s) at given time(s) using Level II data automatic loaded from a remote source, output in-memory as xarray.Dataset. Times, if multiple, are looped over.)
• write_mosiacs_with_autoload (mosaics at given times using automatically loaded Level II data, output directly to disk. Includes option to specify post gridding calculations to be applied before writing to disk. Dispatches all tasks with Dask, rather than looping over time.)