add tests & proper support for negative named axes

src/awkward/_namedaxis.py

@@ -73,6 +73,34 @@ def _prettify(ax: AxisName) -> str:
  )
 
 
+def _neg2pos_axis(
+ axis: int,
+ total: int,
+) -> int:
+ """
+ Converts a negative axis index to a positive one.
+
+ This function takes a negative axis index and the total number of axes and returns the corresponding positive axis index.
+ If the input axis index is already positive, it is returned as is.
+
+ Args:
+ axis (int): The axis index to convert. Can be negative.
+ total (int): The total number of axes.
+
+ Returns:
+ int: The positive axis index corresponding to the input axis index.
+
+ Examples:
+ >>> _neg2pos_axis(-1, 3)
+ 2
+ >>> _neg2pos_axis(1, 3)
+ 1
+ """
+ if axis < 0:
+ return total + axis
+ return axis
+
+
 def _get_named_axis(
  ctx: MaybeSupportsNamedAxis | AttrsNamedAxisMapping | tp.Mapping | tp.Any,
 ) -> AxisMapping:
@@ -365,6 +393,7 @@ def _keep_named_axis(
 def _keep_named_axis_up_to(
  named_axis: AxisMapping,
  axis: int,
+ total: int,
 ) -> AxisMapping:
  """
  Determines the new named axis after keeping all axes up to the specified axis. This function is useful when an operation
@@ -373,26 +402,20 @@ def _keep_named_axis_up_to(
  Args:
  named_axis (AxisMapping): The current named axis.
  axis (int): The index of the axis up to which to keep.
+ total (int): The total number of axes.
 
  Returns:
  AxisMapping: The new named axis after keeping all axes up to the specified axis.
 
  Examples:
- >>> _keep_named_axis_up_to({"x": 0, "y": 2}, 0)
- {"x": 0}
- >>> _keep_named_axis_up_to({"x": 0, "y": 2}, 1)
- {"x": 0}
- >>> _keep_named_axis_up_to({"x": 0, "y": 2}, 2)
- {"x": 0, "y": 2}
- >>> _keep_named_axis_up_to({"x": 0, "y": -2}, 0)
+ >>> _keep_named_axis_up_to({"x": 0, "y": 1, "z": 2}, 1, 3)
+ {"x": 0, "y": 1}
+ >>> _keep_named_axis_up_to({"x": 0, "y": 1, "z": 2}, -1, 3)
+ {"x": 0, "y": 1, "z": 2}
+ >>> _keep_named_axis_up_to({"x": 0, "y": 1, "z": 2}, 0, 3)
  {"x": 0}
- >>> _keep_named_axis_up_to({"x": 0, "y": -2}, 1)
- {"x": 0, "y": -2}
- >>> _keep_named_axis_up_to({"x": 0, "y": -2}, 2)
- {"x": 0, "y": -2}
  """
- if axis < 0:
- raise ValueError("The axis must be a positive integer.")
+ axis = _neg2pos_axis(axis, total)
  out = {}
  for k, v in named_axis.items():
  if v >= 0 and v <= axis:
@@ -458,7 +481,11 @@ def _remove_named_axis(
  total = len(named_axis)
 
  # remove the specified axis
- out = {ax: pos for ax, pos in named_axis.items() if pos != axis}
+ out = {
+ ax: pos
+ for ax, pos in named_axis.items()
+ if _neg2pos_axis(pos, total) != _neg2pos_axis(axis, total)
+ }
 
  return _adjust_pos_axis(out, axis, total, direction=-1)
 
@@ -504,14 +531,14 @@ def _adjust(pos: int, axis: int, direction: int) -> int:
  # -> change position by direction
  if pos >= axis:
  return pos + direction
- # positive axis and position smaller than the removed/added (positive) axis, but greater than 0
- # -> keep position
- elif pos >= 0:
- return pos
  # positive axis and negative position
  # -> change position by direction
+ elif pos < 0 and pos + total < axis:
+ return pos - direction
+ # positive axis and position smaller than the removed/added (positive) axis, but greater than 0
+ # -> keep position
  else:
- return _adjust(pos, axis - total, direction)
+ return pos
  # negative axis
  else:
  # negative axis and position smaller than the removed/added (negative) axis
@@ -520,7 +547,7 @@ def _adjust(pos: int, axis: int, direction: int) -> int:
  return pos - direction
  # negative axis and positive position
  # -> change position by inverse direction
- elif pos > axis + total:
+ elif pos > 0 and pos > axis + total:
  return pos + direction
  # negative axis and position greater than the removed/added (negative) axis, but smaller than 0
  # -> keep position

src/awkward/contents/content.py

@@ -16,7 +16,7 @@
 )
 from awkward._behavior import get_array_class, get_record_class
 from awkward._kernels import KernelError
-from awkward._layout import wrap_layout
+from awkward._layout import maybe_posaxis, wrap_layout
 from awkward._meta.meta import Meta
 from awkward._namedaxis import (
  NamedAxis,
@@ -591,7 +591,7 @@ def _getitem(self, where, named_axis: Type[NamedAxis] = NamedAxis):
  for x in nextwhere:
  if x is np.newaxis or x is None:
  n_newaxis += 1
- n_total = self.purelist_depth
+ n_total = self.minmax_depth[1]
  n_slice_none = n_total - (len(_nextwhere) - n_newaxis - 1)
  # expand `[...]` to `[:]*n_slice_none`
  _nextwhere = (
@@ -602,26 +602,33 @@ def _getitem(self, where, named_axis: Type[NamedAxis] = NamedAxis):
 
  # now propagate named axis
  _named_axis = _keep_named_axis(named_axis.mapping, None)
- iter_named_axis = iter(dict(_named_axis).items())
- _adjust_dims = 0
+ _adjust_dim = 0
  # this loop does the following:
  # - remove a named axis for integer indices, e.g. `a[1, 2]`
  # - add a named axis for None (or np.newaxis) indices, e.g. `a[..., None]`
  # - keep named axis for any other index, e.g. `a[:]`, `a[0:1]`, or `a[a>0]`
  # (these may only remove elements, but not dimensions)
- for i, nw in enumerate(_nextwhere):
- i += _adjust_dims
- for _, pos in iter_named_axis:
- if pos == i:
+ for dim, nw in enumerate(_nextwhere):
+ dim_adjusted = dim + _adjust_dim
+ total_adjusted = self.minmax_depth[1] + _adjust_dim
+ for _, pos in _named_axis.items():
+ if maybe_posaxis(self, pos, 0) == dim_adjusted:
  break
+
  if is_integer(nw) or (is_array_like(nw) and nw.ndim == 0):
  _named_axis = _remove_named_axis(
- _named_axis, i, self.purelist_depth
+ named_axis=_named_axis,
+ axis=dim_adjusted,
+ total=total_adjusted,
  )
- _adjust_dims -= 1
+ _adjust_dim -= 1
  elif nw is None:
- _named_axis = _add_named_axis(_named_axis, i, self.purelist_depth)
- _adjust_dims += 1
+ _named_axis = _add_named_axis(
+ named_axis=_named_axis,
+ axis=dim_adjusted,
+ total=total_adjusted,
+ )
+ _adjust_dim += 1
 
  # set propagated named axis
  named_axis.mapping = _named_axis

src/awkward/highlevel.py

@@ -1109,12 +1109,12 @@ def __getitem__(self, where):
  )
 
  if NamedAxis.mapping:
- out = ak.operations.ak_with_named_axis._impl(
+ return ak.operations.ak_with_named_axis._impl(
  out,
  named_axis=NamedAxis.mapping,
  highlevel=True,
- behavior=self._behavior,
- attrs=self._attrs,
+ behavior=None,
+ attrs=None,
  )
  return out
 

src/awkward/operations/ak_cartesian.py

@@ -256,6 +256,7 @@ def _impl(arrays, axis, nested, parameters, with_name, highlevel, behavior, attr
  # Handle named axis
  # Step 1: Normalize named axis to positional axis
  axis = _named_axis_to_positional_axis(out_named_axis, axis)
+ max_ndim = max([layout.minmax_depth[1] for layout in layouts])
 
  if with_name is not None:
  if parameters is None:
@@ -284,7 +285,7 @@ def _impl(arrays, axis, nested, parameters, with_name, highlevel, behavior, attr
  if nested is None or nested is False:
  nested = []
  elif nested is True:
- out_named_axis = _add_named_axis(out_named_axis, 0)
+ out_named_axis = _add_named_axis(out_named_axis, 0, max_ndim)
  if fields is not None:
  nested = list(fields)[:-1]
  else:
@@ -311,7 +312,7 @@ def _impl(arrays, axis, nested, parameters, with_name, highlevel, behavior, attr
  "[0, len(arrays) - 1) for an iterable of arrays"
  )
  for n in nested:
- out_named_axis = _add_named_axis(out_named_axis, n)
+ out_named_axis = _add_named_axis(out_named_axis, n, max_ndim)
 
  backend = next((layout.backend for layout in layouts), cpu)
  if posaxis == 0:

src/awkward/operations/ak_is_none.py

@@ -61,9 +61,7 @@ def _impl(array, axis, highlevel, behavior, attrs):
 
  # Step 2: propagate named axis from input to output,
  # use strategy "keep up to" (see: awkward._namedaxis)
- out_named_axis = _keep_named_axis_up_to(
- named_axis, axis + layout.minmax_depth[1] if axis < 0 else axis
- )
+ out_named_axis = _keep_named_axis_up_to(named_axis, axis, layout.minmax_depth[1])
 
  def action(layout, depth, backend, lateral_context, **kwargs):
  posaxis = maybe_posaxis(layout, axis, depth)

src/awkward/operations/ak_local_index.py

@@ -108,9 +108,7 @@ def _impl(array, axis, highlevel, behavior, attrs):
 
  # Step 2: propagate named axis from input to output,
  # use strategy "keep up to" (see: awkward._namedaxis)
- out_named_axis = _keep_named_axis_up_to(
- named_axis, axis + layout.minmax_depth[1] if axis < 0 else axis
- )
+ out_named_axis = _keep_named_axis_up_to(named_axis, axis, layout.minmax_depth[1])
 
  out = ak._do.local_index(layout, axis)
 

src/awkward/operations/ak_singletons.py

@@ -76,7 +76,9 @@ def _impl(array, axis, highlevel, behavior, attrs):
 
  # Step 2: propagate named axis from input to output,
  # use strategy "add one" (see: awkward._namedaxis)
- out_named_axis = _add_named_axis(named_axis, axis + 1)
+ out_named_axis = _add_named_axis(
+ named_axis, (axis + 1) if axis >= 0 else axis, layout.minmax_depth[1]
+ )
 
  def action(layout, depth, backend, **kwargs):
  posaxis = maybe_posaxis(layout, axis, depth)