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(`
`)
+ b.WriteString(html.EscapeString(shapeString(a.shape) + " " + a.dtype.String()))
+ b.WriteString("")
+ ci := shownIndices(cols, summarise)
+ if len(a.shape) == 2 {
+ b.WriteString(" | ")
+ for _, j := range ci {
+ b.WriteString("" + index(j) + " | ")
+ }
+ b.WriteString("
")
+ }
+ b.WriteString("")
+ for _, i := range shownIndices(rows, summarise && len(a.shape) == 2) {
+ b.WriteString("")
+ if len(a.shape) == 2 {
+ b.WriteString("| " + index(i) + " | ")
+ }
+ for _, j := range ci {
+ if i < 0 || j < 0 {
+ b.WriteString("⋯ | ")
+ continue
+ }
+ b.WriteString("" + html.EscapeString(a.element(flat, i*cols+j)) + " | ")
+ }
+ b.WriteString("
")
+ }
+ b.WriteString("
")
+ 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")
+ }
+}