diff --git a/README.md b/README.md index 2d67ffd..f94cb7f 100644 --- a/README.md +++ b/README.md @@ -36,6 +36,10 @@ - **Fourier transforms** — package [`fft`](fft) is `numpy.fft` on arrays (`FFT`/`RFFT`/`FFT2`/`FFTN`/… with `n`/`s`, `axes` and `norm`, single precision kept), on [go-fft](https://github.com/go-fft/fft). +- **Notebooks and storage** — `Repr` prints like NumPy (summarised above + 1000 elements), `Display` gives a notebook a text and an HTML table, and + `MarshalBinary`/`UnmarshalBinary` encode an array the same way on every + architecture (so gob and the go-notebook kernel keep arrays). - **Memory reuse** — `Workspace`, an arena for loops: bind the inputs with `ws.Use`, compute as usual, `ws.Reset()` at the end of each pass. Results stop costing the garbage collector anything (see below). diff --git a/display.go b/display.go new file mode 100644 index 0000000..7b8054c --- /dev/null +++ b/display.go @@ -0,0 +1,260 @@ +package ndarray + +import ( + "html" + "math" + "strconv" + "strings" +) + +// Printing, laid out the way NumPy prints: nested brackets, every element +// right-aligned to one width, lines broken at 75 columns, and, above 1000 +// elements, only the first and last three items of each axis with "..." +// between them. Elements are written by the package's own formatting (the +// shortest form that reads back exactly), so the digits may differ from +// NumPy's, which pads every element to a common precision. + +const ( + printThreshold = 1000 // above this many elements, summarise + printEdge = 3 // items kept at each end of a summarised axis + printWidth = 75 // line width +) + +// Repr returns the array as NumPy's repr shows it: array([...]), with the +// dtype when it is not the default for its kind, and the shape when the +// elements are summarised. +func (a *Array) Repr() string { + const open = "array(" + body := a.format(len(open)) + var b strings.Builder + b.WriteString(open) + b.WriteString(body) + if a.Size() > printThreshold { + b.WriteString(", shape=(") + for i, d := range a.shape { + if i > 0 { + b.WriteString(", ") + } + b.WriteString(strconv.Itoa(d)) + } + if len(a.shape) == 1 { + b.WriteByte(',') + } + b.WriteByte(')') + } + switch a.dtype { + case Float64, Int64, Bool, Complex128: + default: + b.WriteString(", dtype=" + a.dtype.String()) + } + b.WriteByte(')') + return b.String() +} + +// Display shows the array in a notebook: its Repr as text, and as an HTML +// table for one and two dimensions. +func (a *Array) Display() map[string][]byte { + out := map[string][]byte{"text/plain": []byte(a.Repr())} + if len(a.shape) == 1 || len(a.shape) == 2 { + out["text/html"] = []byte(a.htmlTable()) + } + return out +} + +// format lays out the elements, the first line starting at column indent. +func (a *Array) format(indent int) string { + if len(a.shape) == 0 { + return a.element(a.contiguousStore(), 0) + } + flat := a.contiguousStore() + summarise := a.Size() > printThreshold + width := 0 + a.visit(summarise, func(i int) { + width = max(width, len(a.element(flat, i))) + }) + var b strings.Builder + a.formatAxis(&b, flat, 0, 0, indent, width, summarise) + return b.String() +} + +// visit calls f with the flat index of every element that is printed. +func (a *Array) visit(summarise bool, f func(int)) { + var walk func(axis, base int) + walk = func(axis, base int) { + if axis == len(a.shape) { + f(base) + return + } + stride := prod(a.shape[axis+1:]) + for _, i := range shownIndices(a.shape[axis], summarise) { + if i >= 0 { + walk(axis+1, base+i*stride) + } + } + } + walk(0, 0) +} + +// shownIndices lists the indices of an axis of length n that are printed, +// with -1 where the "..." goes. +func shownIndices(n int, summarise bool) []int { + if !summarise || n <= 2*printEdge { + out := make([]int, n) + for i := range out { + out[i] = i + } + return out + } + out := []int{0, 1, 2, -1} + for i := n - printEdge; i < n; i++ { + out = append(out, i) + } + return out +} + +// formatAxis writes the sub-array at flat offset base starting at axis. +func (a *Array) formatAxis(b *strings.Builder, flat any, axis, base, indent, width int, summarise bool) { + b.WriteByte('[') + stride := prod(a.shape[axis+1:]) + idx := shownIndices(a.shape[axis], summarise) + last := axis == len(a.shape)-1 + col := indent + 1 // the column after the bracket + for k, i := range idx { + if last { + s := "..." + if i >= 0 { + s = strings.Repeat(" ", width-len(a.element(flat, base+i))) + a.element(flat, base+i) + } + if k > 0 { + b.WriteByte(',') + col++ + if col+1+len(s)+1 > printWidth { + b.WriteString("\n" + strings.Repeat(" ", indent+1)) + col = indent + 1 + } else { + b.WriteByte(' ') + col++ + } + } + b.WriteString(s) + col += len(s) + continue + } + if k > 0 { + b.WriteString("," + strings.Repeat("\n", len(a.shape)-axis-1) + strings.Repeat(" ", indent+1)) + } + if i < 0 { + b.WriteString("...") + continue + } + a.formatAxis(b, flat, axis+1, base+i*stride, indent+1, width, summarise) + } + b.WriteByte(']') +} + +// element writes element i of a contiguous storage slice for printing: +// integral floats end in a point ("2."), as in NumPy. +func (a *Array) element(flat any, i int) string { + s := elemString(flat, i) + switch v := flat.(type) { + case []float64: + return pointed(s, v[i]) + case []float32: + return pointed(s, float64(v[i])) + case []bool: + if v[i] { + return "True" + } + return "False" + case []complex64: + return complexElement(complex128(v[i]), 32) + case []complex128: + return complexElement(v[i], 64) + } + return s +} + +// complexElement writes a complex element as NumPy does inside an array: +// 1.+2.j, no parentheses. +func complexElement(z complex128, bits int) string { + re, im := real(z), imag(z) + sign := "+" + if im < 0 || (im == 0 && math.Signbit(im)) { + sign, im = "-", -im + } + return pointed(floatString(re, bits), re) + sign + pointed(floatString(im, bits), im) + "j" +} + +// pointed adds NumPy's trailing point to an integral float written without +// an exponent. +func pointed(s string, v float64) string { + if v == math.Trunc(v) && !math.IsInf(v, 0) && !strings.ContainsAny(s, "e.") { + return s + "." + } + return s +} + +// htmlTable writes a one- or two-dimensional array as a table, summarised +// like the text form. +func (a *Array) htmlTable() string { + rows, cols := 1, a.shape[0] + if len(a.shape) == 2 { + rows, cols = a.shape[0], a.shape[1] + } + flat := a.contiguousStore() + summarise := a.Size() > printThreshold + var b strings.Builder + b.WriteString(`") + ci := shownIndices(cols, summarise) + if len(a.shape) == 2 { + b.WriteString("") + for _, j := range ci { + b.WriteString("") + } + b.WriteString("") + } + b.WriteString("") + for _, i := range shownIndices(rows, summarise && len(a.shape) == 2) { + b.WriteString("") + if len(a.shape) == 2 { + b.WriteString("") + } + for _, j := range ci { + if i < 0 || j < 0 { + b.WriteString("") + continue + } + b.WriteString("") + } + b.WriteString("") + } + b.WriteString("
`) + b.WriteString(html.EscapeString(shapeString(a.shape) + " " + a.dtype.String())) + b.WriteString("
" + index(j) + "
" + index(i) + "⋯" + html.EscapeString(a.element(flat, i*cols+j)) + "
") + return b.String() +} + +// index writes a row or column header, "⋯" for the summarised gap. +func index(i int) string { + if i < 0 { + return "⋯" + } + return strconv.Itoa(i) +} + +// shapeString writes a shape as NumPy does: (3, 4), (5,), (). +func shapeString(shape []int) string { + var b strings.Builder + b.WriteByte('(') + for i, d := range shape { + if i > 0 { + b.WriteString(", ") + } + b.WriteString(strconv.Itoa(d)) + } + if len(shape) == 1 { + b.WriteByte(',') + } + b.WriteByte(')') + return b.String() +} diff --git a/marshal.go b/marshal.go new file mode 100644 index 0000000..397fa62 --- /dev/null +++ b/marshal.go @@ -0,0 +1,79 @@ +package ndarray + +import ( + "encoding/binary" + "errors" + "fmt" +) + +// The binary form of an array, as MarshalBinary writes it: +// +// "ND" 0x01 magic and version +// dtype one byte, the DType +// ndim uvarint +// shape[0..ndim) uvarints +// elements row-major, little-endian, DType.Size bytes each +// (a bool is one byte, 0 or 1; a complex number is its +// real part, then its imaginary part) +// +// The same bytes come out on every architecture. + +var marshalMagic = [3]byte{'N', 'D', 1} + +// ErrFormat is returned by UnmarshalBinary for bytes that are not an array. +var ErrFormat = errors.New("ndarray: not an encoded array") + +// MarshalBinary encodes the array: its dtype, its shape and its elements in +// row-major order. A view is encoded as the array it shows. It implements +// encoding.BinaryMarshaler, so an array can be kept by gob and by a +// notebook kernel. +func (a *Array) MarshalBinary() ([]byte, error) { + b := make([]byte, 0, 16+a.Size()*a.dtype.Size()) + b = append(b, marshalMagic[:]...) + b = append(b, byte(a.dtype)) + b = binary.AppendUvarint(b, uint64(len(a.shape))) + for _, d := range a.shape { + b = binary.AppendUvarint(b, uint64(d)) + } + return binary.Append(b, binary.LittleEndian, a.contiguousStore()) +} + +// UnmarshalBinary decodes what MarshalBinary wrote into a, replacing its +// contents. The bytes may come from anywhere: a shape the length of the data +// does not account for is an error before anything is allocated. +func (a *Array) UnmarshalBinary(b []byte) error { + if len(b) < 4 || [3]byte(b[:3]) != marshalMagic { + return ErrFormat + } + dt := DType(b[3]) + if int(dt) >= numDTypes { + return fmt.Errorf("%w: dtype %d", ErrFormat, b[3]) + } + b = b[4:] + ndim, n := binary.Uvarint(b) + if n <= 0 || ndim > 64 { + return fmt.Errorf("%w: bad rank", ErrFormat) + } + b = b[n:] + shape := make([]int, ndim) + for i := range shape { + d, n := binary.Uvarint(b) + if n <= 0 || d > maxSize { + return fmt.Errorf("%w: bad dimension", ErrFormat) + } + shape[i], b = int(d), b[n:] + } + if err := validateShape(shape); err != nil { + return fmt.Errorf("%w: %v", ErrFormat, err) + } + // Divide rather than multiply: size*16 can overflow for a complex128 + // shape that validateShape (which counts 8 bytes an element) allows. + if size := prod(shape); len(b)%dt.Size() != 0 || len(b)/dt.Size() != size { + return fmt.Errorf("%w: %d bytes for %d elements of %v", ErrFormat, len(b), size, dt) + } + s := makeStore(dt, prod(shape)) + // The length was checked above, so Decode cannot run short. + _, _ = binary.Decode(b, binary.LittleEndian, s) + *a = *fromStore(s, shape) + return nil +} diff --git a/marshal_test.go b/marshal_test.go new file mode 100644 index 0000000..a775782 --- /dev/null +++ b/marshal_test.go @@ -0,0 +1,152 @@ +package ndarray + +import ( + "bytes" + "encoding/binary" + "encoding/gob" + "errors" + "strings" + "testing" +) + +func TestMarshalRoundTrip(t *testing.T) { + for _, dt := range allDTypes() { + a := must(Arange(0, 12, 1)).AsType(dt) + a = must(a.Reshape(3, 4)) + if dt == Bool { + a = must(FromSlice([]bool{true, false, true, true, false, false, true, false, true, true, true, false}, 3, 4)) + } + if dt.IsComplex() { + a = must(a.Add(Scalar(-2.5i))) + } + // A strided view is encoded as the array it shows. + v := must(a.Slice(All(), Step(2))) + for _, x := range []*Array{a, v, must(ZerosOf(dt)), must(ZerosOf(dt, 0, 5))} { + b, err := x.MarshalBinary() + if err != nil { + t.Fatal(err) + } + var y Array + if err := y.UnmarshalBinary(b); err != nil { + t.Fatalf("%v %v: %v", dt, x.Shape(), err) + } + if y.DType() != dt || y.String() != x.String() { + t.Errorf("%v round trip: %v, want %v", dt, y.String(), x.String()) + } + } + } +} + +// TestMarshalLayout pins the bytes, so that every architecture (s390x is +// big-endian) writes the same ones. +func TestMarshalLayout(t *testing.T) { + a := must(FromSlice([]int16{1, -2}, 2)) + b, _ := a.MarshalBinary() + want := []byte{'N', 'D', 1, byte(Int16), 1, 2, 1, 0, 0xfe, 0xff} + if !bytes.Equal(b, want) { + t.Errorf("bytes = %v, want %v", b, want) + } +} + +func TestUnmarshalRejects(t *testing.T) { + good, _ := must(FromSlice([]float64{1, 2, 3}, 3)).MarshalBinary() + head := func(dt DType, dims ...uint64) []byte { + b := []byte{'N', 'D', 1, byte(dt)} + b = binary.AppendUvarint(b, uint64(len(dims))) + for _, d := range dims { + b = binary.AppendUvarint(b, d) + } + return b + } + cases := map[string][]byte{ + "empty": nil, + "magic": append([]byte("XD"), good[2:]...), + "version": append([]byte{'N', 'D', 2}, good[3:]...), + "dtype": head(DType(99)), + "no rank": []byte{'N', 'D', 1, 0}, + "rank": head(Float64, make([]uint64, 65)...), + "cut dim": append(head(Float64), 1), + "huge dim": head(Float64, 1<<62), + "overflow": head(Float64, 1<<40, 1<<40), + "short data": append(head(Float64, 3), 1, 2, 3), + "ragged data": append(head(Complex128, 1), make([]byte, 17)...), + "long data": append(head(Float64, 1), make([]byte, 16)...), + "lying shape": head(Complex128, 1<<58), + } + for name, b := range cases { + var a Array + if err := a.UnmarshalBinary(b); !errors.Is(err, ErrFormat) { + t.Errorf("%s: err = %v", name, err) + } + } +} + +func TestGobKeepsArrays(t *testing.T) { + type holder struct{ A *Array } + h := holder{A: must(FromSlice([]complex64{1 + 2i}, 1))} + var buf bytes.Buffer + if err := gob.NewEncoder(&buf).Encode(h); err != nil { + t.Fatal(err) + } + var h2 holder + if err := gob.NewDecoder(&buf).Decode(&h2); err != nil || Item[complex64](h2.A, 0) != 1+2i { + t.Errorf("gob: %v %v", h2.A, err) + } +} + +func TestRepr(t *testing.T) { + big := must(Arange(0, 2000, 1)) + cases := []struct { + a *Array + want string + }{ + {must(FromSlice([]int64{0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11}, 3, 4)), + "array([[ 0, 1, 2, 3],\n [ 4, 5, 6, 7],\n [ 8, 9, 10, 11]])"}, + {must(FromSlice([]float32{1.5, 2, 3}, 3)), "array([1.5, 2., 3.], dtype=float32)"}, + {must(FromSlice([]complex128{1 + 2i, complex(0, -1)}, 2)), "array([1.+2.j, 0.-1.j])"}, + {must(FromSlice([]complex64{complex(0, float32(negZero()))}, 1)), "array([0.-0.j], dtype=complex64)"}, + {must(FromSlice([]bool{true, false}, 2)), "array([ True, False])"}, + {must(FromSlice([]uint8{7}, []int{}...)), "array(7, dtype=uint8)"}, + {must(FromSlice([]float64{1e20, 0.5}, 2)), "array([1e+20, 0.5])"}, + {big, "array([ 0., 1., 2., ...,"}, + {must(big.Reshape(40, 50)), " ...,\n"}, + {must(big.Reshape(40, 50)), "shape=(40, 50))"}, + {big, "shape=(2000,))"}, + {must(must(Arange(0, 24, 1)).Reshape(2, 3, 4)), "11.]],\n\n [[12."}, + {must(Arange(0, 40, 1)), "12.,\n 13."}, + } + for _, c := range cases { + if got := c.a.Repr(); !strings.Contains(got, c.want) { + t.Errorf("Repr =\n%s\nwant it to contain\n%s", got, c.want) + } + } +} + +func negZero() float64 { z := 0.0; return -z } + +func TestDisplay(t *testing.T) { + m := must(FromSlice([]int64{1, 2, 3, 4}, 2, 2)) + d := m.Display() + if string(d["text/plain"]) != m.Repr() { + t.Error("text/plain is not the Repr") + } + h := string(d["text/html"]) + for _, want := range []string{`(2, 2) int64`, "1", "4"} { + if !strings.Contains(h, want) { + t.Errorf("html lacks %q:\n%s", want, h) + } + } + if h := string(must(FromSlice([]bool{true}, 1)).Display()["text/html"]); !strings.Contains(h, "True") || !strings.Contains(h, "(1,) bool") || strings.Contains(h, "") { + t.Errorf("1-D html:\n%s", h) + } + big := must(must(Arange(0, 4000, 1)).Reshape(40, 100)) + if h := string(big.Display()["text/html"]); !strings.Contains(h, "⋯") || !strings.Contains(h, "⋯") || !strings.Contains(h, "3999.") { + t.Errorf("summarised html") + } + if _, ok := must(Zeros(2, 2, 2)).Display()["text/html"]; ok { + t.Error("3-D html") + } + if shapeString([]int{}) != "()" { + t.Error("0-d shape") + } +}