Skip to content

Commit

Permalink
Merge pull request #462 from davidhassell/dask-field-halo
Browse files Browse the repository at this point in the history 
dask: `Field.halo`
  • Loading branch information
@davidhassell
davidhassell authored Oct 26, 2022
2 parents 8908b9f + 5c51308 commit a843944
Show file tree
Hide file tree
Showing 4 changed files with 127 additions and 75 deletions.
11 changes: 10 additions & 1 deletion cf/data/data.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -7299,7 +7299,6 @@ def insert_dimension(self, position=0, inplace=False):
return d

@daskified(_DASKIFIED_VERBOSE)
@_deprecated_kwarg_check("size")
@_inplace_enabled(default=False)
@_manage_log_level_via_verbosity
def halo(
Expand Down Expand Up @@ -7487,6 +7486,16 @@ def halo(
"""
from dask.array.core import concatenate

if size is not None:
_DEPRECATION_ERROR_KWARGS(
self,
"halo",
{"size": None},
message="Use the 'depth' parameter instead.",
version="TODODASKVER",
removed_at="5.0.0",
) # pragma: no cover

d = _inplace_enabled_define_and_cleanup(self)

ndim = d.ndim
Expand Down
99 changes: 58 additions & 41 deletions cf/field.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -13193,12 +13193,13 @@ def grad_xy(self, x_wrap=None, one_sided_at_boundary=False, radius=None):
@_manage_log_level_via_verbosity
def halo(
self,
size,
depth,
axes=None,
tripolar=None,
fold_index=-1,
inplace=False,
verbose=None,
size=None,
):
"""Expand the field construct by adding a halo to its data.

Expand Down Expand Up @@ -13226,42 +13227,43 @@ def halo(

:Parameters:

size: `int` or `dict`
depth: `int` or `dict`
Specify the size of the halo for each axis.

If *size* is a non-negative `int` then this is the halo
size that is applied to all of the axes defined by the
*axes* parameter.
If *depth* is a non-negative `int` then this is the
halo size that is applied to all of the axes defined
by the *axes* parameter.

Alternatively, halo sizes may be assigned to axes
individually by providing a `dict` for which a key
specifies an axis (by passing the axis description to a
call of the field construct's `domain_axis` method. For
example, for a value of ``'X'``, the domain axis construct
returned by ``f.domain_axis('X')``) with a corresponding
value of the halo size for that axis. Axes not specified
by the dictionary are not expanded, and the *axes*
parameter must not also be set.
specifies an axis (by passing the axis description to
a call of the field construct's `domain_axis`
method. For example, for a value of ``'X'``, the
domain axis construct returned by
``f.domain_axis('X')``) with a corresponding value of
the halo size for that axis. Axes not specified by the
dictionary are not expanded, and the *axes* parameter
must not also be set.

*Parameter example:*
Specify a halo size of 1 for all otherwise selected
axes: ``size=1``
axes: ``1``

*Parameter example:*
Specify a halo size of zero ``size=0``. This results in
Specify a halo size of zero: ``0``. This results in
no change to the data shape.

*Parameter example:*
For data with three dimensions, specify a halo size of 3
for the first dimension and 1 for the second dimension:
``size={0: 3, 1: 1}``. This is equivalent to ``size={0:
3, 1: 1, 2: 0}``
For data with three dimensions, specify a halo size
of 3 for the first dimension and 1 for the second
dimension: ``{0: 3, 1: 1}``. This is equivalent to
``{0: 3, 1: 1, 2: 0}``

*Parameter example:*
Specify a halo size of 2 for the "longitude" and
"latitude" axes: ``size=2, axes=['latutude',
'longitude']``, or equivalently ``size={'latutude': 2,
'longitude': 2}``.
"latitude" axes: ``depth=2, axes=['latutude',
'longitude']``, or equivalently ``depth={'latitude':
2, 'longitude': 2}``.

axes: (sequence of) `str` or `int`, optional
Select the domain axes to be expanded, defined by the
Expand Down Expand Up @@ -13295,27 +13297,28 @@ def halo(

tripolar: `dict`, optional
A dictionary defining the "X" and "Y" axes of a global
tripolar domain. This is necessary because in the global
tripolar case the "X" and "Y" axes need special treatment,
as described above. It must have keys ``'X'`` and ``'Y'``,
whose values identify the corresponding domain axis
construct by passing the value to a call of the field
construct's `domain_axis` method. For example, for a value
of ``'ncdim%i'``, the domain axis construct returned by
tripolar domain. This is necessary because in the
global tripolar case the "X" and "Y" axes need special
treatment, as described above. It must have keys
``'X'`` and ``'Y'``, whose values identify the
corresponding domain axis construct by passing the
value to a call of the field construct's `domain_axis`
method. For example, for a value of ``'ncdim%i'``, the
domain axis construct returned by
``f.domain_axis('ncdim%i')``.

The "X" and "Y" axes must be a subset of those identified
by the *size* or *axes* parameter.
The "X" and "Y" axes must be a subset of those
identified by the *depth* or *axes* parameter.

See the *fold_index* parameter.

*Parameter example:*
Define the "X" and Y" axes by their netCDF dimension
names: ``tripolar={'X': 'ncdim%i', 'Y': 'ncdim%j'}``
names: ``{'X': 'ncdim%i', 'Y': 'ncdim%j'}``

*Parameter example:*
Define the "X" and Y" axes by positions 2 and 1
respectively of the data: ``tripolar={'X': 2, 'Y': 1}``
respectively of the data: ``{'X': 2, 'Y': 1}``

fold_index: `int`, optional
Identify which index of the "Y" axis corresponds to the
Expand All @@ -13328,6 +13331,9 @@ def halo(

{{verbose: `int` or `str` or `None`, optional}}

size: deprecated at version TODODASKVER
Use the *depth* parameter instead.

:Returns:

`Field` or `None`
Expand Down Expand Up @@ -13418,21 +13424,32 @@ def halo(
"""
f = _inplace_enabled_define_and_cleanup(self)

# Set the halo size for each axis.
# TODODASKAPI
if size is not None:
_DEPRECATION_ERROR_KWARGS(
self,
"halo",
{"size": None},
message="Use the 'depth' parameter instead.",
version="TODODASKVER",
removed_at="5.0.0",
) # pragma: no cover

# Set the halo depth for each axis.
data_axes = f.get_data_axes(default=())
if isinstance(size, dict):
if isinstance(depth, dict):
if axes is not None:
raise ValueError(
"Can't set existing axes when size is a dict."
"Can't set existing axes when depth is a dict."
)

axis_halo = {
self.domain_axis(k, key=True): v for k, v in size.items()
self.domain_axis(k, key=True): v for k, v in depth.items()
}

if not set(data_axes).issuperset(axis_halo):
raise ValueError(
f"Can't apply halo: Bad axis specification: {size!r}"
f"Can't apply halo: Bad axis specification: {depth!r}"
)
else:
if axes is None:
Expand All @@ -13441,7 +13458,7 @@ def halo(
if isinstance(axes, (str, int)):
axes = (axes,)

axis_halo = {self.domain_axis(k, key=True): size for k in axes}
axis_halo = {self.domain_axis(k, key=True): depth for k in axes}

if tripolar:
# Find the X and Y axes of a tripolar grid
Expand Down Expand Up @@ -13474,10 +13491,10 @@ def halo(
tripolar_axes = {X: "X", Y: "Y"}

# Add halos to the field construct's data
size = {data_axes.index(axis): h for axis, h, in axis_halo.items()}
depth = {data_axes.index(axis): h for axis, h, in axis_halo.items()}

f.data.halo(
size=size,
depth,
tripolar=tripolar,
fold_index=fold_index,
inplace=True,
Expand Down Expand Up @@ -13510,7 +13527,7 @@ def halo(
}

c.halo(
size=construct_size,
construct_size,
tripolar=construct_tripolar,
fold_index=fold_index,
inplace=True,
Expand Down
46 changes: 28 additions & 18 deletions cf/mixin/propertiesdata.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -16,6 +16,7 @@
)
from ..functions import (
_DEPRECATION_ERROR_ATTRIBUTE,
_DEPRECATION_ERROR_KWARGS,
_DEPRECATION_ERROR_METHOD,
default_netCDF_fillvals,
)
Expand Down Expand Up @@ -4899,12 +4900,13 @@ def get_filenames(self):
@_manage_log_level_via_verbosity
def halo(
self,
size,
depth,
axes=None,
tripolar=None,
fold_index=-1,
inplace=False,
verbose=None,
size=None,
):
"""Expand the data by adding a halo.
Expand All @@ -4929,10 +4931,10 @@ def halo(
:Parameters:
size: `int` or `dict`
depth: `int` or `dict`
Specify the size of the halo for each axis.
If *size* is a non-negative `int` then this is the
If *depth* is a non-negative `int` then this is the
halo size that is applied to all of the axes defined
by the *axes* parameter.
Expand All @@ -4946,22 +4948,22 @@ def halo(
*Parameter example:*
Specify a halo size of 1 for all otherwise selected
axes: ``size=1``
axes: ``1``
*Parameter example:*
Specify a halo size of zero ``size=0``. This results
in no change to the data shape.
Specify a halo size of zero: ``0``. This results in
no change to the data shape.
*Parameter example:*
For data with three dimensions, specify a halo size
of 3 for the first dimension and 1 for the second
dimension: ``size={0: 3, 1: 1}``. This is equivalent
to ``size={0: 3, 1: 1, 2: 0}``
dimension: ``{0: 3, 1: 1}``. This is equivalent to
``{0: 3, 1: 1, 2: 0}``
*Parameter example:*
Specify a halo size of 2 for the first and last
dimensions `size=2, axes=[0, -1]`` or equivalently
``size={0: 2, -1: 2}``.
dimensions `depth=2, axes=[0, -1]`` or equivalently
``depth={0: 2, -1: 2}``.
axes: (sequence of) `int`
Select the domain axes to be expanded, defined by
Expand All @@ -4979,14 +4981,13 @@ def halo(
positions in the data.
The "X" and "Y" axes must be a subset of those
identified by the *size* or *axes* parameter.
identified by the *depth* or *axes* parameter.
See the *fold_index* parameter.
*Parameter example:*
Define the "X" and Y" axes by positions 2 and 1
respectively of the data: ``tripolar={'X': 2, 'Y':
1}``
respectively of the data: ``{'X': 2, 'Y': 1}``
fold_index: `int`, optional
Identify which index of the "Y" axis corresponds to
Expand All @@ -4999,6 +5000,9 @@ def halo(
{{verbose: `int` or `str` or `None`, optional}}
size: deprecated at version TODODASKVER
Use the *depth* parameter instead.
:Returns:
`{{class}}` or `None`
Expand All @@ -5007,19 +5011,25 @@ def halo(
**Examples**
TODO
TODO
"""
_kwargs = [f"{k}={v!r}" for k, v in locals().items()]
_ = f"{self.__class__.__name__}.halo("
logger.info("{}{}".format(_, (",\n" + " " * len(_)).join(_kwargs)))
if size is not None:
_DEPRECATION_ERROR_KWARGS(
self,
"halo",
{"size": None},
message="Use the 'depth' parameter instead.",
version="TODODASKVER",
removed_at="5.0.0",
) # pragma: no cover

v = _inplace_enabled_define_and_cleanup(self)

data = v.get_data(None)
if data is not None:
data.halo(
size=size,
depth,
axes=axes,
tripolar=tripolar,
fold_index=fold_index,
Expand Down
Loading

0 comments on commit a843944

Please sign in to comment.