Skip to content

Commit

Permalink
Merge branch 'main' into agoose77/feat-shape-item-none
Browse files Browse the repository at this point in the history
  • Loading branch information
@agoose77
agoose77 authored Jan 30, 2023
2 parents 1107e25 + 82536ce commit 0682a21
Show file tree
Hide file tree
Showing 5 changed files with 96 additions and 81 deletions.
28 changes: 21 additions & 7 deletions src/awkward/operations/ak_mean.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -186,16 +186,16 @@ def _impl(x, weight, axis, keepdims, mask_identity):
sumw = ak.operations.ak_count._impl(
x,
axis,
keepdims,
mask_identity,
keepdims=True,
mask_identity=True,
highlevel=True,
behavior=None,
)
sumwx = ak.operations.ak_sum._impl(
x,
axis,
keepdims,
mask_identity,
keepdims=True,
mask_identity=True,
highlevel=True,
behavior=None,
)
Expand All @@ -211,12 +211,26 @@ def _impl(x, weight, axis, keepdims, mask_identity):
sumwx = ak.operations.ak_sum._impl(
x * weight,
axis,
keepdims,
mask_identity,
keepdims=True,
mask_identity=True,
highlevel=True,
behavior=None,
)
return sumwx / sumw

out = sumwx / sumw

if not mask_identity:
out = ak.highlevel.Array(ak.operations.fill_none(out, np.nan, axis=-1))

if axis is None:
if not keepdims:
out = out[(0,) * out.ndim]
else:
if not keepdims:
posaxis = ak._util.maybe_posaxis(out.layout, axis, 1)
out = out[(slice(None, None),) * posaxis + (0,)]

return out


@ak._connect.numpy.implements("mean")
Expand Down
87 changes: 32 additions & 55 deletions src/awkward/operations/ak_ptp.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -79,65 +79,42 @@ def ptp(array, axis=None, *, keepdims=False, mask_identity=True, flatten_records

def _impl(array, axis, keepdims, mask_identity):
behavior = ak._util.behavior_of(array)
array = ak.highlevel.Array(
ak.operations.to_layout(array, allow_record=False, allow_other=False),
behavior=behavior,
)
layout = ak.operations.to_layout(array, allow_record=False, allow_other=False)

with np.errstate(invalid="ignore", divide="ignore"):
if axis is None:
out = ak.operations.ak_max._impl(
array,
axis,
keepdims,
None,
mask_identity,
highlevel=True,
behavior=None,
) - ak.operations.ak_min._impl(
array,
axis,
keepdims,
None,
mask_identity,
highlevel=True,
behavior=None,
)
if not mask_identity and out is None:
out = 0
maxi = ak.operations.ak_max._impl(
layout,
axis,
True,
None,
mask_identity,
highlevel=True,
behavior=behavior,
)
mini = ak.operations.ak_min._impl(
layout,
axis,
True,
None,
True,
highlevel=True,
behavior=behavior,
)
out = maxi - mini

# Check that removed code was not needed!
assert maxi is not None and mini is not None

if not mask_identity:
out = ak.highlevel.Array(ak.operations.fill_none(out, 0, axis=-1))

if axis is None:
if not keepdims:
out = out[(0,) * out.ndim]
else:
maxi = ak.operations.ak_max._impl(
array,
axis,
True,
None,
mask_identity,
highlevel=True,
behavior=None,
)
mini = ak.operations.ak_min._impl(
array,
axis,
True,
None,
True,
highlevel=True,
behavior=None,
)

if maxi is None or mini is None:
out = None

else:
out = maxi - mini

if not mask_identity:
out = ak.highlevel.Array(ak.operations.fill_none(out, 0, axis=-1))

if not keepdims:
posaxis = ak._util.maybe_posaxis(out.layout, axis, 1)
out = out[(slice(None, None),) * posaxis + (0,)]
if not keepdims:
posaxis = ak._util.maybe_posaxis(out.layout, axis, 1)
out = out[(slice(None, None),) * posaxis + (0,)]

return out

Expand Down
19 changes: 16 additions & 3 deletions src/awkward/operations/ak_std.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -171,17 +171,30 @@ def _impl(x, weight, ddof, axis, keepdims, mask_identity):
)

with np.errstate(invalid="ignore", divide="ignore"):
return ufuncs.sqrt(
out = ufuncs.sqrt(
ak.operations.ak_var._impl(
x,
weight,
ddof,
axis,
keepdims,
mask_identity,
keepdims=True,
mask_identity=True,
)
)

if not mask_identity:
out = ak.highlevel.Array(ak.operations.fill_none(out, np.nan, axis=-1))

if axis is None:
if not keepdims:
out = out[(0,) * out.ndim]
else:
if not keepdims:
posaxis = ak._util.maybe_posaxis(out.layout, axis, 1)
out = out[(slice(None, None),) * posaxis + (0,)]

return out


@ak._connect.numpy.implements("std")
def _nep_18_impl_std(
Expand Down
37 changes: 26 additions & 11 deletions src/awkward/operations/ak_var.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -176,45 +176,60 @@ def _impl(x, weight, ddof, axis, keepdims, mask_identity):
)

with np.errstate(invalid="ignore", divide="ignore"):
xmean = ak.operations.ak_mean._impl(x, weight, axis, False, mask_identity)
xmean = ak.operations.ak_mean._impl(
x, weight, axis, keepdims=True, mask_identity=True
)
if weight is None:
sumw = ak.operations.ak_count._impl(
x,
axis,
keepdims,
mask_identity,
keepdims=True,
mask_identity=True,
highlevel=True,
behavior=None,
)
sumwxx = ak.operations.ak_sum._impl(
(x - xmean) ** 2,
axis,
keepdims,
mask_identity,
keepdims=True,
mask_identity=True,
highlevel=True,
behavior=None,
)
else:
sumw = ak.operations.ak_sum._impl(
x * 0 + weight,
axis,
keepdims,
mask_identity,
keepdims=True,
mask_identity=True,
highlevel=True,
behavior=None,
)
sumwxx = ak.operations.ak_sum._impl(
(x - xmean) ** 2 * weight,
axis,
keepdims,
mask_identity,
keepdims=True,
mask_identity=True,
highlevel=True,
behavior=None,
)
if ddof != 0:
return (sumwxx / sumw) * (sumw / (sumw - ddof))
out = (sumwxx / sumw) * (sumw / (sumw - ddof))
else:
return sumwxx / sumw
out = sumwxx / sumw

if not mask_identity:
out = ak.highlevel.Array(ak.operations.fill_none(out, np.nan, axis=-1))

if axis is None:
if not keepdims:
out = out[(0,) * out.ndim]
else:
if not keepdims:
posaxis = ak._util.maybe_posaxis(out.layout, axis, 1)
out = out[(slice(None, None),) * posaxis + (0,)]

return out


@ak._connect.numpy.implements("var")
Expand Down
6 changes: 1 addition & 5 deletions tests/test_2020_reduce_axis_none.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -128,7 +128,6 @@ def test_std():
assert np.isnan(ak.std(array[2], axis=None, mask_identity=False))


@pytest.mark.xfail(reason="fix mask_identity=False")
def test_std_no_mask_axis_none():
assert ak._util.arrays_approx_equal(
ak.std(array[-1:], axis=None, keepdims=True, mask_identity=True),
Expand All @@ -150,7 +149,6 @@ def test_var():
assert np.isnan(ak.var(array[2], axis=None, mask_identity=False))


@pytest.mark.xfail(reason="fix mask_identity=False")
def test_var_no_mask_axis_none():
assert ak._util.arrays_approx_equal(
ak.var(array[-1:], axis=None, keepdims=True, mask_identity=True),
Expand All @@ -172,7 +170,6 @@ def test_mean():
assert np.isnan(ak.mean(array[2], axis=None, mask_identity=False))


@pytest.mark.xfail(reason="fix mask_identity=False")
def test_mean_no_mask_axis_none():
assert ak._util.arrays_approx_equal(
ak.mean(array[-1:], axis=None, keepdims=True, mask_identity=True),
Expand All @@ -191,10 +188,9 @@ def test_ptp():
ak.ptp(array, axis=None, keepdims=True, mask_identity=True),
ak.Array([10.0]).mask[[True]],
)
assert np.isinf(ak.ptp(array[2], axis=None, mask_identity=False))
assert ak.ptp(array[2], axis=None, mask_identity=False) == pytest.approx(0.0)


@pytest.mark.xfail(reason="fix mask_identity=False")
def test_ptp_no_mask_axis_none():
assert ak._util.arrays_approx_equal(
ak.ptp(array[-1:], axis=None, keepdims=True, mask_identity=True),
Expand Down

0 comments on commit 0682a21

Please sign in to comment.