diff --git a/internal/vm/ndarray.go b/internal/vm/ndarray.go index 3d657b4c..9d4afe65 100644 --- a/internal/vm/ndarray.go +++ b/internal/vm/ndarray.go @@ -19,8 +19,8 @@ import ( // The Ruby surface is the one the go-ndarray playground's code generator emits // (Go | Ruby | NumPy): creation (zeros/ones/full/arange/linspace/eye/from), // views (reshape/transpose/flatten/squeeze/expand_dims and NumPy basic indexing -// through #[]), elementwise ufuncs and **, reductions with an axis and -// keepdims:, products and joins. Semantics are go-ndarray's, which are NumPy's, +// through #[]), elementwise ufuncs and **, reductions with an axis (positional, +// a.sum(0), or NumPy's keyword, a.sum(axis: 0)) and keepdims:, products and joins. Semantics are go-ndarray's, which are NumPy's, // in float64. // // Every method checks its own arguments. A wrong count raises ArgumentError, a @@ -167,6 +167,23 @@ func ndKeep(kw map[string]object.Value) bool { return ok && v.Truthy() } +// ndOptAxis reads the optional axis of a reduction, given positionally (a.sum(0)) +// or as NumPy's keyword (a.sum(axis: 0)). Absent or nil means no axis; giving +// both forms is an ArgumentError, as NumPy's duplicate argument is a TypeError. +func (vm *VM) ndOptAxis(args []object.Value, kw map[string]object.Value) (int, bool) { + v, ok := kw["axis"] + if len(args) == 1 { + if ok { + raise("ArgumentError", "axis given both positionally and as a keyword") + } + v, ok = args[0], true + } + if !ok || object.IsNil(v) { + return 0, false + } + return vm.ndInt(v), true +} + // ndOnes is the all-ones shape of the given rank: what a whole-array reduction // with keepdims: true leaves. func ndOnes(ndim int) []int { @@ -549,11 +566,11 @@ func (vm *VM) registerNDArray() { // the reduced axes as length 1 (all of them when there is no axis). reduce := func(name string, whole func(*nd.Array) (float64, error), along func(*nd.Array, int, bool) (*nd.Array, error)) { d(name, func(vm *VM, v object.Value, args []object.Value, _ *Proc) object.Value { - args, kw := ndKw(args, "keepdims") + args, kw := ndKw(args, "axis", "keepdims") ndArgc(args, 0, 1) a, keep := self(v), ndKeep(kw) - if len(args) == 1 && !object.IsNil(args[0]) { - return mustArray(along(a, vm.ndInt(args[0]), keep)) + if axis, ok := vm.ndOptAxis(args, kw); ok { + return mustArray(along(a, axis, keep)) } r, err := whole(a) if err != nil { @@ -578,11 +595,11 @@ func (vm *VM) registerNDArray() { // indices along the axis. argReduce := func(name string, whole func(*nd.Array) (int, error), along func(*nd.Array, int, bool) (*nd.Array, error)) { d(name, func(vm *VM, v object.Value, args []object.Value, _ *Proc) object.Value { - args, kw := ndKw(args, "keepdims") + args, kw := ndKw(args, "axis", "keepdims") ndArgc(args, 0, 1) a, keep := self(v), ndKeep(kw) - if len(args) == 1 && !object.IsNil(args[0]) { - return mustArray(along(a, vm.ndInt(args[0]), keep)) + if axis, ok := vm.ndOptAxis(args, kw); ok { + return mustArray(along(a, axis, keep)) } i, err := whole(a) if err != nil { @@ -599,9 +616,10 @@ func (vm *VM) registerNDArray() { // cumsum: flattened with no axis, else along the axis. d("cumsum", func(vm *VM, v object.Value, args []object.Value, _ *Proc) object.Value { + args, kw := ndKw(args, "axis") ndArgc(args, 0, 1) - if len(args) == 1 && !object.IsNil(args[0]) { - return mustArray(self(v).CumSum(vm.ndInt(args[0]))) + if axis, ok := vm.ndOptAxis(args, kw); ok { + return mustArray(self(v).CumSum(axis)) } return &NDArray{a: self(v).CumSumFlat()} }) diff --git a/internal/vm/ndarray_test.go b/internal/vm/ndarray_test.go index 56098f02..c0bd8356 100644 --- a/internal/vm/ndarray_test.go +++ b/internal/vm/ndarray_test.go @@ -548,7 +548,13 @@ func TestNDArrayArgumentChecks(t *testing.T) { {`NDArray.zeros(2, 0).max(1)`, "ArgumentError", "zero-length axis"}, {a + `a.argmax(1, 2)`, "ArgumentError", "(given 2, expected 0..1)"}, {a + `a.argmax("x")`, "TypeError", "String into Integer"}, - {a + `a.argmin(axis: 1)`, "ArgumentError", "unknown keyword: :axis"}, + {a + `a.argmin(axes: 1)`, "ArgumentError", "unknown keyword: :axes"}, + {a + `a.sum(0, axis: 1)`, "ArgumentError", "both positionally and as a keyword"}, + {a + `a.argmax(0, axis: 1)`, "ArgumentError", "both positionally and as a keyword"}, + {a + `a.cumsum(0, axis: 1)`, "ArgumentError", "both positionally and as a keyword"}, + {a + `a.sum(axis: "x")`, "TypeError", "String into Integer"}, + {a + `a.sum(axis: 2)`, "IndexError", "axis out of range"}, + {a + `a.cumsum(axes: 1)`, "ArgumentError", "unknown keyword: :axes"}, {a + `a.argmax(4)`, "IndexError", "axis out of range"}, {`NDArray.zeros(0).argmin(keepdims: true)`, "ArgumentError", "empty"}, {a + `a.cumsum(1, 2)`, "ArgumentError", "(given 2, expected 0..1)"}, @@ -585,3 +591,45 @@ func TestNDArrayArgumentChecks(t *testing.T) { } } } + +// TestNDArrayAxisKeyword: NumPy spells the axis as a keyword (a.sum(axis=0)), so +// a Ruby user writes a.sum(axis: 0). Before, only keepdims: was accepted and +// axis: raised "unknown keyword: :axis". Values are NumPy's, by hand, for +// a = arange(6).reshape(2, 3) = [[0,1,2],[3,4,5]]. The positional form stays. +func TestNDArrayAxisKeyword(t *testing.T) { + const a = `a = NDArray.arange(6).reshape(2, 3); ` + for _, c := range []struct{ src, want string }{ + // a.sum(axis=0) == [3, 5, 7]; a.sum(axis=1) == [3, 12]. + {a + `p a.sum(axis: 0).to_a`, "[3.0, 5.0, 7.0]\n"}, + {a + `p a.sum(axis: 1).to_a`, "[3.0, 12.0]\n"}, + {a + `p a.sum(axis: -1).to_a`, "[3.0, 12.0]\n"}, + // Control: the positional form gives the same answers. + {a + `p a.sum(0).to_a`, "[3.0, 5.0, 7.0]\n"}, + {a + `p a.sum(1).to_a`, "[3.0, 12.0]\n"}, + // a.sum(axis=1, keepdims=True).shape == (2, 1), in either order of keywords. + {a + `p a.sum(axis: 1, keepdims: true).shape`, "[2, 1]\n"}, + {a + `p a.sum(keepdims: true, axis: 1).shape`, "[2, 1]\n"}, + {a + `p a.sum(1, keepdims: true).shape`, "[2, 1]\n"}, + // axis=None (nil) is the whole array. + {a + `p a.sum(axis: nil)`, "15.0\n"}, + {a + `p a.sum(axis: nil, keepdims: true).shape`, "[1, 1]\n"}, + {a + `p a.sum`, "15.0\n"}, + // The other reductions. + {a + `p a.mean(axis: 0).to_a`, "[1.5, 2.5, 3.5]\n"}, + {a + `p a.prod(axis: 1).to_a`, "[0.0, 60.0]\n"}, + {a + `p a.max(axis: 1).to_a`, "[2.0, 5.0]\n"}, + {a + `p a.min(axis: 0).to_a`, "[0.0, 1.0, 2.0]\n"}, + {a + `p a.argmax(axis: 1).to_a`, "[2.0, 2.0]\n"}, + {a + `p a.argmin(axis: 0).to_a`, "[0.0, 0.0, 0.0]\n"}, + {a + `p a.argmax(1).to_a`, "[2.0, 2.0]\n"}, + {a + `p a.argmax(axis: nil)`, "5\n"}, + // cumsum(axis=0) == [[0,1,2],[3,5,7]]; no axis flattens. + {a + `p a.cumsum(axis: 0).to_a`, "[0.0, 1.0, 2.0, 3.0, 5.0, 7.0]\n"}, + {a + `p a.cumsum(0).to_a`, "[0.0, 1.0, 2.0, 3.0, 5.0, 7.0]\n"}, + {a + `p a.cumsum.to_a`, "[0.0, 1.0, 3.0, 6.0, 10.0, 15.0]\n"}, + } { + if got := eval(t, c.src); got != c.want { + t.Errorf("src=%q got=%q want=%q", c.src, got, c.want) + } + } +}