wrf-python

WRF-Python uses a layered decorator pattern to bridge the gap between high-level Python/xarray objects and low-level Fortran routines. When a diagnostic is called, the decorators are executed from top to bottom (outermost to innermost) to prepare data, and then from bottom to top to clean up the result.

A typical decorator chain for a compiled routine includes:

1. check_args: Verifies input array shapes to provide clear error messages.
2. left_iteration: Handles multidimensional arrays by iterating over 'leftmost' dimensions and passing slices to the Fortran routine.
3. cast_type: Converts input arrays to the precision required by Fortran (e.g., REAL(KIND=8)) and casts the result back to the original type.
4. extract_and_transpose: Converts xarray.DataArray or C-ordered numpy.ndarray into Fortran-ordered arrays and handles the final transposition back to C-order for the user.

When xarray is enabled, wrf.getvar uses an internal thread-local cache to store XLAT and XLONG coordinate variables and moving nest metadata. This prevents repeated extraction when processing sequences of WRF files.

By default, the cache holds up to 20 items, keyed by the object ID of the file or sequence. If you are creating new file objects in a loop (e.g., using netCDF4.Dataset), the cache will fill up, leading to higher memory usage until it reaches the limit. This is not a memory leak; memory usage will eventually stabilize.

You can manage this behavior using:

• wrf.disable_xarray(): Disables xarray functionality entirely.
• wrf.set_cache_size(size): Adjusts the number of items held in the cache. Setting the size to 0 disables the cache completely.

from netCDF4 import Dataset
import wrf

# To disable xarray functionality and reduce memory overhead:
wrf.disable_xarray()

for i in range(150):
    f = Dataset('wrfout_d01_2005-07-17_12_00_00.nc')
    p = wrf.getvar(f, 'pressure')
    f.close()

WRF-Python utilizes compiled Fortran 90 routines for high-performance computations. These are exposed to Python via f2py.

The Extension Layer

Raw Fortran routines are compiled into the wrf._wrffortran extension module. However, users should not interact with wrf._wrffortran directly. Instead, these routines are wrapped in the extension.py module to provide a user-friendly interface.

How the wrapper handles data

When you call a routine exported via extension.py, the following logic is applied to ensure compatibility and performance:

1. Argument Validation: The wrapper verifies argument shapes to provide clearer error messages than standard f2py errors.
2. Output Allocation: It allocates an output array based on the algorithm's requirements and the input data dimensions.
3. Dimension Iteration: It iterates over the "leftmost" dimensions. For example, if a 2D Fortran algorithm receives a 5D array, the wrapper iterates over the 3 leftmost dimensions to process slices.
4. Type Casting: It casts input arrays to the required dtype (typically converting 4-byte floats to 8-byte doubles for WRF data).
5. Memory Efficiency: It extracts arrays from xarray into numpy arrays and transposes them into Fortran ordering. This reorders the shape tuple and sets the Fortran ordering flag without copying data, allowing the Fortran routine to work directly on the data pointers.

The vertcross function computes a vertical cross-section interpolation between two points.

Key Parameters:

• variable: The 3D field to interpolate.
• vertical_coord: The vertical coordinate field (e.g., geopotential height z).
• wrfin: The original NetCDF file object.
• start_point: A CoordPair object defining the starting latitude and longitude.
• end_point: A CoordPair object defining the ending latitude and longitude.
• latlon=True: If set, includes latitude/longitude points in the metadata.
• meta=True: If set, includes metadata in the returned object.

Usage Pattern:

1. Define points using CoordPair(lat=..., lon=...).
2. Call vertcross for each variable of interest.
3. Use the resulting object's coordinates (e.g., coords['xy_loc'] or coords['vertical']) to set axis ticks for plotting.

from wrf import vertcross, CoordPair

# Set the start point and end point for the cross section
start_point = CoordPair(lat=26.76, lon=-80.0)
end_point = CoordPair(lat=26.76, lon=-77.8)

# Compute the vertical cross-section interpolation
z_cross = vertcross(Z, z, wrfin=ncfile, start_point=start_point, 
                    end_point=end_point, latlon=True, meta=True)