Skip to content

Commit

Permalink
add option to not create 1D index in expand_dims
Browse files Browse the repository at this point in the history
  • Loading branch information
TomNicholas committed Apr 19, 2024
1 parent 3e848eb commit 95d453c
Showing 1 changed file with 17 additions and 5 deletions.
22 changes: 17 additions & 5 deletions xarray/core/dataset.py
Original file line number Diff line number Diff line change
Expand Up @@ -4479,6 +4479,7 @@ def expand_dims(
self,
dim: None | Hashable | Sequence[Hashable] | Mapping[Any, Any] = None,
axis: None | int | Sequence[int] = None,
create_1d_index: bool = True,
**dim_kwargs: Any,
) -> Self:
"""Return a new object with an additional axis (or axes) inserted at
Expand All @@ -4488,6 +4489,8 @@ def expand_dims(
If dim is already a scalar coordinate, it will be promoted to a 1D
coordinate consisting of a single value.
The
Parameters
----------
dim : hashable, sequence of hashable, mapping, or None
Expand All @@ -4503,6 +4506,8 @@ def expand_dims(
multiple axes are inserted. In this case, dim arguments should be
same length list. If axis=None is passed, all the axes will be
inserted to the start of the result array.
create_1d_index : bool, default is True
Whether to create new PandasIndex objects for new 1D coordinate variables.
**dim_kwargs : int or sequence or ndarray
The keywords are arbitrary dimensions being inserted and the values
are either the lengths of the new dims (if int is given), or their
Expand Down Expand Up @@ -4622,6 +4627,8 @@ def expand_dims(
# save the coordinates to the variables dict, and set the
# value within the dim dict to the length of the iterable
# for later use.

# TODO should we have an option to not create a variable here?
index = PandasIndex(v, k)
indexes[k] = index
variables.update(index.create_variables())
Expand Down Expand Up @@ -4660,11 +4667,16 @@ def expand_dims(
variables[k] = v.set_dims(dict(all_dims))
else:
if k not in variables:
# If dims includes a label of a non-dimension coordinate,
# it will be promoted to a 1D coordinate with a single value.
index, index_vars = create_default_index_implicit(v.set_dims(k))
indexes[k] = index
variables.update(index_vars)
if create_1d_index:
# If dims includes a label of a non-dimension coordinate,
# it will be promoted to a 1D coordinate with a single value.
index, index_vars = create_default_index_implicit(v.set_dims(k))
indexes[k] = index
variables.update(index_vars)
else:
# create 1D variable without creating a new index
new_1d_var = v.set_dims(k)
variables.update({k: new_1d_var})

return self._replace_with_new_dims(
variables, coord_names=coord_names, indexes=indexes
Expand Down

0 comments on commit 95d453c

Please sign in to comment.