Other utility modules

NEDAS.utils.conversion module

NEDAS.utils.conversion.units_convert(units_from, units_to, var)

Convert units for a given variable.

Parameters:

• units_from (str or numbers.Number) – Source units for the input variable

• units_to (str or numbers.Number) – Target units to convert to

• var (np.ndarray) – The input variable

Returns:

Variable with converted units

Return type:

np.ndarray

NEDAS.utils.conversion.proj2dict(proj: Proj) → dict

Convert map projection name in pyproj.Proj to a dictionary of human-readable parameters

Parameters:

proj (pyproj.Proj) – Map projection object.

Returns:

A dictionary of projection parameters, such as name, lat_0, lon_0, etc.

Return type:

dict

NEDAS.utils.conversion.t2h(t: datetime) → float

Convert datetime object to hours since 1900-1-1 00:00

NEDAS.utils.conversion.h2t(h: float) → datetime

Convert hours since 1900-1-1 00:00 to datetime object

NEDAS.utils.conversion.t2s(t: datetime) → str

Convert datetime object to a time string 'ccyymmddHHMM'

NEDAS.utils.conversion.s2t(s: str) → datetime

Convert a time string 'ccyymmddHHMM' to a datetime object

NEDAS.utils.conversion.seconds_to_timestr(seconds: int) → str

Convert from seconds to time duration string ‘HH:MM:SS’

NEDAS.utils.conversion.ensure_list(v) → list

If the input v is a list, return itself; if not, return [v].

NEDAS.utils.shell_utils module

NEDAS.utils.netcdf_lib module

NEDAS.utils.netcdf_lib.nc_open(filename: str, mode: Literal['r', 'w', 'a', 'r+'], comm: Comm | None = None) → Dataset

Open a netCDF file.

Parameters:

• filename (str) – Path to the netCDF file.

• mode (str) – Open mode, (e.g. 'r', 'w').

• comm (Comm, optional) – MPI communicator object. If None, open the netCDF4.Dataset normally. If communicator is available, try parallel=True open when opening the file. If that’s not supported, acquire a file lock in communicator for blocking serial access of the file.

Returns:

netCDF file handle.

Return type:

netCDF4.Dataset

NEDAS.utils.netcdf_lib.nc_close(filename: str, f: Dataset, comm: Comm | None = None) → None

Close the netCDF file handle.

Parameters:

• filename (str) – Path to the opened netCDF file.

• f (netCDF4.Dataset) – netCDF4 Dataset file handle.

• comm (Comm, optional) – MPI communicator object. If None, just close the file directly. If communicator is specified, release the file lock after closing it.

NEDAS.utils.netcdf_lib.nc_write_var(filename: str, dim: Mapping[str, int | None], varname: str, dat: ndarray, dtype: str | None = None, recno: dict[str, int] | None = None, attr: dict[str, str] | None = None, comm: Comm | None = None) → None

Write a variable to a netCDF file.

Parameters:

• filename (str) – Path to the output netCDF file.

• dim (dict) – Dictionary {dimension name (str): dimension size (int)} of each dimension. The dimension size can be None if it is unlimited dimension (can append more records afterwards)

• varname (str) – Name of the output variable. Variable groups are supported, use 'group/varname' as varname.

• dat (np.ndarray) – Data for output, number of its dimensions must match dim (excluding unlimited dimensions).

• dtype (str, optional) – Data type to convert the input data to.

• recno (dict, optional) – Dictionary {dimension name (str): record number (int)}, indicating the index in unlimited dimensions for the data to be written to. Each unlimited dimension defined in dim should have a corresponding recno entry.

• attr (dict, optional) – Dictionary {name of attribute (str): value (str)}, additional attributes to be added.

• comm (Comm, optional) – MPI communicator object, handling parallel I/O and make sure thread-safe writing of data.

NEDAS.utils.netcdf_lib.nc_read_var(filename: str, varname: str, comm: Comm | None = None) → ndarray

Read a variable from a netCDF file.

This function by default reads the entire variable, if you only want a slice, it is more efficient to use netCDF4.Dataset handle directly.

Parameters:

• filename (str) – Path to the netCDF file for reading.

• varname (str) – Name of the variable to read.

• comm (Comm, optional) – MPI communicator object.

Returns:

Variable read from the file.

Return type:

np.ndarray

NEDAS.utils.fft_lib module

NEDAS.utils.fft_lib.fft2(f: ndarray) → ndarray

2D FFT implemented by pyFFTW. If pyFFTW is not available, will switch to np.fft.fft2

Parameters:

f (np.ndarray) – Input field, of shape (…, ny, nx), of float32 type, the last two dimensions will be transformed by the 2D FFT.

Returns:

Output field in spectral space, of complex64 type.

Return type:

np.ndarray

NEDAS.utils.fft_lib.ifft2(fh: ndarray) → ndarray

Inverse 2D FFT implemented by pyFFTW. If pyFFTW is not available, will switch to np.fft.ifft2

Parameters:

fh (np.ndarray) – Input field, of shape (…, nky, nkx), of complex64 type, the last two dimensions will be inverse transformed by 2D FFT.

Returns:

Output field in physical space, of float32 type.

Return type:

np.ndarray

NEDAS.utils.fft_lib.fftwn(n)

Wavenumber sequence corresponding to the FFT output in 1 dimension.

Parameters:

n (int) – The size of the dimension

Returns:

The sequence of wavenumber (0,1,2,…-2,-1) for this dimension

Return type:

np.ndarray

NEDAS.utils.fft_lib.get_wn(fld)

Generates a meshgrid of wavenumbers corresponding to the input field dimensions.

Parameters:

fld (np.ndarray) – Input field, the last two dimensions are the horizontal directions (ny, nx)

Returns:

Wavenumbers corresponding to the input fields dimensions, relative to the domain size (nx or ny, whichever is larger).

Return type:

np.ndarray

NEDAS.utils.njit module

NEDAS.utils.parallel module

class NEDAS.utils.parallel.Comm

Bases: object

Communicator class supporting both serial and MPI programs.

When the python program is started with MPI environment, for example:

$ mpirun -n 10 python -m mpi4py program.py

A communicator can be obtained from the mpi4py package:

>>> from mpi4py import MPI
>>> comm = MPI.COMM_WORLD

However, when the program is run in

parallel_io

If netCDF4.Dataset is built with parallel I/O support.

Type:

bool

mpi_ready: bool = False

parallel_io: bool

init_file_lock(filename)

Initialize file locks for thread-safe I/O.

Parameters:

filename (str) – Path to the file.

check_parallel_io() → bool

Check if netCDF4 is built with parallel I/O support.

Returns:

True if netCDF4 module support parallel I/O mode.

Return type:

bool

cleanup_file_locks()

acquire_file_lock(filename)

release_file_lock(filename)

finalize()

Clean up MPI resources cleanly to avoid hangs on exit.

class NEDAS.utils.parallel.DummyComm

Bases: object

Dummy communicator for python without mpi

Get_size()

Get_rank()

Barrier()

Abort(code: int)

Split(color=0, key=0)

bcast(obj, root=0)

send(obj, dest, tag)

recv(source, tag)

allgather(obj)

gather(obj, root=0)

allreduce(obj)

reduce(obj, root=0)

NEDAS.utils.parallel.by_rank(comm: Comm, rank: int) → Callable[[Callable[[P], T]], Callable[[P], T | None]]

Decorator for func() to be run only by rank 0 in comm

NEDAS.utils.parallel.bcast_by_root(comm: Comm) → Callable[[Callable[[P], T]], Callable[[P], T]]

Decorator for func() to be run only by rank 0 in comm, and result of func() is then broadcasted to all other ranks.

NEDAS.utils.parallel.distribute_tasks(comm: Comm, tasks: ndarray | Sequence, load: ndarray | Sequence | None = None) → dict[int, list]

Divide a list of task indices and assign a subset to each rank in comm

Parameters:

• comm (Comm) – MPI communicator

• tasks (ArrayLike) – List of task indices (to be distributed over the processors)

• load (np.ndarray, optional) – Amount of workload for each task element The default is None, we will let tasks have equal workload

Returns:

Dictionary {rank:list}, list is the subset of tasks for the processor rank

calling this function to work on

Return type:

dict

class NEDAS.utils.parallel.OfflineScheduler(c, nworker: int, walltime: int | None = None, check_dt: float = 0.1, debug: bool = False)

Bases: object

An offline scheduler class for queuing and running multiple jobs on available workers (group of processors). The jobs are submitted by one processor with the scheduler, while the job.run code is calling subprocess to be run on the worker

submit_job(name: str, job: Callable, *args, **kwargs) → None

Submit a job to the scheduler, hold info in jobs dict Input: - name (str): is a unique name to identify this job - job (Callable), is_running and kill methods - *args, **kwargs are to be passed into job()

monitor_job_queue() → None

Monitor the available_workers and pending_jobs, assign a job to a worker if possible Monitor the running_jobs for jobs that are finished, kill jobs that exceed walltime, and move the finished jobs to completed_jobs

start_queue()

Start the job queue, and wait for jobs to complete

shutdown()

NEDAS.utils.progress module

NEDAS.utils.progress.print_with_cache(msg: str, prev_msg: str) → str

NEDAS.utils.progress.watch_files(files, timeout=1000, check_dt=1)

NEDAS.utils.progress.watch_log(logfile: str, keyword: str, timeout: int = 1000, check_dt: int = 1) → None

NEDAS.utils.progress.find_keyword_in_file(file: str, keyword: str) → bool

NEDAS.utils.progress.count_lines_in_file(file: str) → int

class NEDAS.utils.progress.Formatter(interactive: bool = True, is_notebook: bool = False, cols=80, anchor=50, tabspace=4, progress_bar_width=10)

Bases: object

Formatter of the progress display.

Parameters:

• interactive (bool, optional) – Whether the output is interactive (supports ansi escape code). Defaults to True.

• anchor (int, optional) – Characters to anchor the left part of status line. Defaults to 50.

• tabspace (int, optional) – Number of spaces for one call stack level indentation. Defaults to 4.

• progress_bar_width (int, optional) – Width of the progress bar in characters. Defaults to 10.

strip_escape_code(text: str) → str

truncate(text: str) → str

indent(level: int, branch: bool = True) → str

Generate the indent string to form call stack tree structure in log.

Parameters:

• level (int) – The current call stack level.

• branch (bool) – Whether a branch is needed at the end.

Returns:

The indent string

Return type:

str

padding(level: int, name: str) → str

Generate the padding string to align the status line.

Parameters:

• level (int) – The current call stack level.

• name (str) – The name of the current function or task.

Returns:

The padding string

Return type:

str

dimmer(msg)

progress_bar(task_id: int, ntask: int) → str

Generate a progress bar based on task_id and ntask.

Parameters:

• task_id (int) – Current task index, from 0 to ntask-1

• ntask (int) – Total number of tasks

Returns:

The progress bar msg to be shown.

Return type:

str

Note: Will require the print command with end=”” option so that new line updated is overwritting the old line.

class NEDAS.utils.progress.Progress(interactive: bool = True, is_notebook: bool = False, cols: int = 80, debug: bool = False, call_stack: list[dict] | None = None, call_stack_max_level: int | None = None, anchor: int = 50, tabspace: int = 4, progress_bar_width: int = 10, io_interval: float = 0.1)

Bases: object

Progress tracker and displayer. Used by Context.logger to show runtime progress.

formatter: Formatter

interactive: bool

debug: bool

call_stack: list[dict[str, Any]]

call_stack_max_level: int | None

new_node(func_name: str | None = None) → dict

property node: dict

within_max_level(level: int) → bool

is_leaf(node) → bool

property level: int

set_flag(flag: str)

get_timer_msg(node)

push(func_name: str)

pop()

update() → str

log(msg: str, flag: str) → str

Safely injects a global message without breaking the tree.

NEDAS.utils.random_perturb module

NEDAS.utils.random_perturb.random_field_gaussian(nx, ny, amp, hcorr)

Random field with a Gaussian spectrum.

Parameters:

• nx (int) – Number of grid points in x direction.

• ny (int) – Number of grid points in y direction.

• amp (float) – Amplitude (standard deviation) of the random field.

• hcorr (float) – Horizontal decorrelation length (number of grid points).

Returns:

The output random field with shape (ny, nx).

Return type:

np.ndarray

NEDAS.utils.random_perturb.gaussian(k, sig)

gaussian spectrum

NEDAS.utils.random_perturb.get_sig_in_gaussian(nx, ny, hcorr)

Derive the sig parameter in gaussian spectrum. Called by random_field_gaussian, cached the result since in external for loop there will be fixed input

NEDAS.utils.random_perturb.random_field_powerlaw(nx, ny, amp, pwrlaw)

NEDAS.utils.random_perturb.random_displacement(grid, mask, amp, hcorr)

NEDAS.utils.random_perturb.get_velocity_from_press(grid, pres, scale_wind=False, press_amp=None, press_hcorr=None, wind_amp=None)

derive wind velocity from pressure field

NEDAS.utils.spatial_operation module

NEDAS.utils.spatial_operation.gradx(fld, dx, cyclic_dim=None)

Gradient of input field in x direction

Parameters:

• fld (np.ndarray) – input field, last two dimensions (ny, nx)

• dx (int) – grid spacing in x, fld.shape

• cyclic_dim (str, optional) – string ‘x’, ‘y’, ‘xy’, indicating the dimension(s) that are cyclic.

Returns:

gradx of fld with same shape

Return type:

np.ndarray

NEDAS.utils.spatial_operation.grady(fld, dy, cyclic_dim=None)

gradient of input fld in y direction, similar to gradx

NEDAS.utils.spatial_operation.gradx2(fld, dx, cyclic_dim=None)

NEDAS.utils.spatial_operation.grady2(fld, dy, cyclic_dim=None)

NEDAS.utils.spatial_operation.gradxy(fld, dx, dy, cyclic_dim=None)

NEDAS.utils.spatial_operation.laplacian(fld, dx, dy, cyclic_dim=None)

NEDAS.utils.spatial_operation.coarsen(grid, fld, nlevel)

Coarsen the image by downsampling the grid points by factors of 1/2,

Parameters:

• grid (Grid) – the original grid

• fld (np.ndarray) – field with last two dimensions ny,nx

• nlevel (int) – number of resolution levels to go down

Returns:

new grid with lower resolution np.ndarray: new field defined on the new grid

Return type:

Grid

NEDAS.utils.spatial_operation.refine(grid, mask, fld, nlevel)

Refine the image by upsampling the grid points by factors of 2,

Parameters:

• grid (Grid) – the original grid

• fld (np.ndarray) – field with last two dimensions ny,nx

• nlevel (int) – number of resolution levels to go up

Returns:

new grid with higher resolution np.ndarray: new field on new grid

Return type:

Grid

NEDAS.utils.spatial_operation.coarsen_mask(mask, lev1, lev2)

NEDAS.utils.spatial_operation.coarsen_field(fld, lev1, lev2)

NEDAS.utils.spatial_operation.sharpen_field(fld, lev1, lev2)

NEDAS.utils.spatial_operation.warp(grid, fld, u, v)

Warp the image with input vector field

Parameters:

• grid (Grid) – the grid on which the image is defined

• fld (np.ndarray) – input image

• u (np.ndarray) – displacement vector x component, in grid.x units

• v (np.ndarray) – displacement vector y component, in grid.y units

Returns:

the warped image

Return type:

np.ndarray

NEDAS.utils.multiscale module

NEDAS.utils.multiscale.lowpass_response(k2d, k1, k2)

Low-pass spectral response function Input - k2d: array, 2d wavenumber - k1, k2: float, wavenumbers defining the transition zone (from 1 to 0 in response) Return - r: array same dimension as k2d, the response function

NEDAS.utils.multiscale.get_scale_component_spec_bandpass(grid, fld, character_length, s)

NEDAS.utils.multiscale.convolve(grid, fld, rgrid, response)

NEDAS.utils.multiscale.get_scale_component_convolve(grid, fld, character_length, s)

NEDAS.utils.multiscale.get_scale_component(grid, fld, character_length, s)

Get scale component using a bandpass filter in spectral space Input: - grid: Grid object - fld: array, […, ny, nx], the input field - character_length: list of characteristic length for each scale - s: int, scale index Return: - fld: array, the scale component s of input fld

NEDAS.utils.multiscale.get_error_scale_factor(grid, character_length, s)

NEDAS.utils.graphics module

NEDAS.utils.graphics.get_cmap(cmap_name: str)

Get colormap object based on the input name string.

Parameters:

cmap_name (str) – The name of the color map. For cmocean colormaps, the name should be in the format ‘cmocean.<cmap_name>’

Returns:

A colormap object corresponding to the given name.

Return type:

Colormap

Raises:

KeyError – If the colormap name is not found.

NEDAS.utils.graphics.adjust_ax_size(ax, wfac=1.0, hfac=1.0)

Make plot axes a little smaller on right hand side to make room for colorbar.

Even for axes without colorbar in a multi-pane plot, it is still useful to call this function so that the axes will align with those with colorbars.

Parameters:

• ax (matplotlib.axes.Axes) – Matplotlib axes object.

• wfac (float, optional) – Ratio to scale the width of the axes. Defaults to 1.

• hfac (float, optional) – Ratio to scale the height of the axes. Defaults to 1.

NEDAS.utils.graphics.add_colorbar(fig, ax, cmap, vmin, vmax, nlevels=10, fontsize=12, units=None)

Add a colorbar to the right-hand side of an axes.

This function adds a colorbar to the provided axes, using the specified colormap and value range. The number of levels, font size, and unit label can be customized.

Parameters:

• fig (matplotlib.figure.Figure) – Matplotlib figure object.

• ax (matplotlib.axes.Axes) – Matplotlib axes object.

• cmap (matplotlib.colors.Colormap) – Colormap to use for the colorbar.

• vmin (float) – Minimum value for the colorbar.

• vmax (float) – Maximum value for the colorbar.

• nlevels (int, optional) – Number of color levels. Defaults to 10.

• fontsize (int, optional) – Font size for tick labels and unit label. Defaults to 12.

• units (str, optional) – Unit label to display as the colorbar title. Defaults to None.

Returns:

The created colorbar object.

Return type:

matplotlib.colorbar.Colorbar

NEDAS.utils.graphics.arrowhead_xy(x1, x2, y1, y2, hw, hl)

Given a line segment from (x1,y1) to (x2,y2), return the segments that draw an arrow head (for plotting vectors).

Parameters:

• x1 (float) – X-coordinate of the start point of line segment.

• x2 (float) – X-coordinate of the end point of line segment.

• y1 (float) – Y-coordinate of the start point of line segment.

• y2 (float) – Y-coordinate of the end point of line segment.

• hw (float) – Width of the arrow head.

• hl (float) – Length of the arrow head.

Returns:

X-coordinates of the additional line segments forming the arrowhead. list: Y-coordinates of the additional line segments forming the arrowhead.

Return type:

list

NEDAS.utils.graphics.draw_reference_vector_legend(ax, xr, yr, V, L, hw, hl, refcolor, linecolor, ref_units='')

Draw a legend box with reference vector and units string.

Parameters:

• ax (matplotlib.axes.Axes) – Matplotlib axes object.

• xr (float) – X-coordinate of the center of the legend box

• yr (float) – Y-coordinate of the center of the legend box

• V (float) – Velocity scale to be shown as the reference vector.

• L (float) – Length of the reference vector.

• hw (float) – Width of the arrowhead.

• hl (float) – Length of the arrowhead.

• refcolor (str or tuple) – Color of the background in the legend box.

• linecolor (str or tuple) – Color of the reference vector.

• ref_units (str, optional) – Unit label to be shown, default is ‘’.

NEDAS.utils.graphics.draw_line(ax, data, linecolor, linewidth, linestyle, zorder)

Draw line segments.

Parameters:

• ax (matplotlib.axes.Axes) – The axes on which to draw the lines.

• data (dict) – Dictionary containing: - ‘xy’ (list of arrays): Coordinates of points. - ‘parts’ (list of lists): Indices indicating segment divisions.

• linecolor (str or tuple) – Color of the line.

• linewidth (float) – Width of the line.

• linestyle (str) – Style of the line (e.g., ‘-’, ‘–’, ‘:’).

• zorder (int) – Drawing order (higher numbers are drawn on top).

NEDAS.utils.graphics.draw_patch(ax, data, color, zorder)

Draw a filled patch.

Parameters:

• ax (matplotlib.axes.Axes) – The axes on which to draw the lines.

• data (dict) – Dictionary containing: - ‘xy’ (list of arrays): Coordinates of points. - ‘parts’ (list of lists): Indices indicating segment divisions.

• color (str or tuple) – Color of the patch

• zorder (int) – Drawing order (higher numbers are drawn on top).