implement numpy.save

code/micropython.mk

@@ -19,6 +19,7 @@ SRC_USERMOD += $(USERMODULES_DIR)/numpy/create.c
 SRC_USERMOD += $(USERMODULES_DIR)/numpy/fft/fft.c
 SRC_USERMOD += $(USERMODULES_DIR)/numpy/fft/fft_tools.c
 SRC_USERMOD += $(USERMODULES_DIR)/numpy/filter.c
+SRC_USERMOD += $(USERMODULES_DIR)/numpy/io/io.c
 SRC_USERMOD += $(USERMODULES_DIR)/numpy/linalg/linalg.c
 SRC_USERMOD += $(USERMODULES_DIR)/numpy/linalg/linalg_tools.c
 SRC_USERMOD += $(USERMODULES_DIR)/numpy/numerical.c

code/numpy/io/io.c

@@ -0,0 +1,163 @@
+/*
+ * This file is part of the micropython-ulab project,
+ *
+ * https://github.com/v923z/micropython-ulab
+ *
+ * The MIT License (MIT)
+ *
+ * Copyright (c) 2022 Zoltán Vörös
+*/
+
+#include <string.h>
+
+#include "py/builtin.h"
+#include "py/obj.h"
+#include "py/runtime.h"
+#include "py/stream.h"
+
+#include "../../ndarray.h"
+#include "io.h"
+
+#define ULAB_IO_BUFFER_SIZE         128
+
+
+static mp_obj_t io_save(mp_obj_t fname, mp_obj_t ndarray_) {
+    ndarray_obj_t *ndarray = MP_OBJ_TO_PTR(ndarray_);
+    int error;
+    char *buffer = m_new(char, ULAB_IO_BUFFER_SIZE);
+    uint8_t offset = 0;
+
+    // test for endianness
+    uint16_t x = 1;
+    int8_t endian = (x >> 8) == 1 ? '>' : '<';
+
+    mp_obj_t open_args[2] = {
+        fname,
+        MP_OBJ_NEW_QSTR(MP_QSTR_wb)
+    };
+
+    mp_obj_t npy = mp_builtin_open(2, open_args, (mp_map_t *)&mp_const_empty_map);
+    const mp_stream_p_t *fout = mp_get_stream(npy);
+
+    // write header;
+    // magic string + header length, which is always 128 - 10 = 118, represented as a little endian uint16 (0x76, 0x00)
+    // + beginning of the dictionary describing the array
+    memcpy(buffer, "\x93NUMPY\x01\x00\x76\x00{'descr': '", 21);
+    offset += 21;
+
+    buffer[offset] = endian;
+    if((ndarray->dtype == NDARRAY_UINT8) || (ndarray->dtype == NDARRAY_INT8)) {
+        // for single-byte data, the endianness doesn't matter
+        buffer[offset] = '|';
+    }
+    offset++;
+    switch(ndarray->dtype) {
+        case NDARRAY_UINT8:
+            memcpy(buffer+offset, "u1", 2);
+            break;
+        case NDARRAY_INT8:
+            memcpy(buffer+offset, "i1", 2);
+            break;
+        case NDARRAY_UINT16:
+            memcpy(buffer+offset, "u2", 2);
+            break;
+        case NDARRAY_INT16:
+            memcpy(buffer+offset, "i2", 2);
+            break;
+        case NDARRAY_FLOAT:
+            #if MICROPY_FLOAT_IMPL == MICROPY_FLOAT_IMPL_FLOAT
+            memcpy(buffer+offset, "f4", 2);
+            #else
+            memcpy(buffer+offset, "f8", 2);
+            #endif
+            break;
+        #if ULAB_SUPPORTS_COMPLEX
+        case NDARRAY_COMPLEX:
+            #if MICROPY_FLOAT_IMPL == MICROPY_FLOAT_IMPL_FLOAT
+            memcpy(buffer+offset, "c8", 2);
+            #else
+            memcpy(buffer+offset, "c16", 2);
+            #endif
+            offset++;
+            break;
+        #endif
+    }
+
+    offset += 2;
+    memcpy(buffer+offset, "', 'fortran_order': False, 'shape': (", 37);
+    offset += 37;
+
+    if(ndarray->ndim == 1) {
+        offset += sprintf(buffer+offset, "%ld,", ndarray->shape[ULAB_MAX_DIMS - 1]);
+    } else {
+        for(uint8_t i = 0; i < ndarray->ndim - 1; i++) {
+            offset += sprintf(buffer+offset, "%ld, ", ndarray->shape[ULAB_MAX_DIMS - i - 1]);
+        }
+        offset += sprintf(buffer+offset, "%ld", ndarray->shape[ULAB_MAX_DIMS - 1]);
+    }
+    memcpy(buffer+offset, "), }", 4);
+    offset += 4;
+    // pad with space till the very end
+    memset(buffer+offset, 32, ULAB_IO_BUFFER_SIZE - offset - 1);
+    buffer[ULAB_IO_BUFFER_SIZE - 1] = '\n';
+    fout->write(npy, buffer, ULAB_IO_BUFFER_SIZE, &error);
+
+    // write the array data
+    uint8_t sz = ndarray->itemsize;
+    offset = 0;
+
+    uint8_t *array = (uint8_t *)ndarray->array;
+
+    // TODO: if flatiter is available, we can save the loop expansion
+    #if ULAB_MAX_DIMS > 3
+    size_t i = 0;
+    do {
+    #endif
+        #if ULAB_MAX_DIMS > 2
+        size_t j = 0;
+        do {
+        #endif
+            #if ULAB_MAX_DIMS > 1
+            size_t k = 0;
+            do {
+            #endif
+                size_t l = 0;
+                do {
+                    memcpy(buffer+offset, array, sz);
+                    offset += sz;
+                    if(offset == ULAB_IO_BUFFER_SIZE) {
+                        fout->write(npy, buffer, offset, &error);
+                        offset = 0;
+                    }
+                    array += ndarray->strides[ULAB_MAX_DIMS - 1];
+                    l++;
+                } while(l <  ndarray->shape[ULAB_MAX_DIMS - 1]);
+            #if ULAB_MAX_DIMS > 1
+                array -= ndarray->strides[ULAB_MAX_DIMS - 1] * ndarray->shape[ULAB_MAX_DIMS-1];
+                array += ndarray->strides[ULAB_MAX_DIMS - 2];
+                k++;
+            } while(k <  ndarray->shape[ULAB_MAX_DIMS - 2]);
+            #endif
+        #if ULAB_MAX_DIMS > 2
+            array -= ndarray->strides[ULAB_MAX_DIMS - 2] * ndarray->shape[ULAB_MAX_DIMS-2];
+            array += ndarray->strides[ULAB_MAX_DIMS - 3];
+            j++;
+        } while(j <  ndarray->shape[ULAB_MAX_DIMS - 3]);
+        #endif
+    #if ULAB_MAX_DIMS > 3
+        array -= ndarray->strides[ULAB_MAX_DIMS - 3] * ndarray->shape[ULAB_MAX_DIMS-3];
+        array += ndarray->strides[ULAB_MAX_DIMS - 4];
+        i++;
+    } while(i <  ndarray->shape[ULAB_MAX_DIMS - 4]);
+    #endif
+
+    if(offset != 0) {
+        fout->write(npy, buffer, offset, &error);
+    }
+
+    m_del(char, buffer, ULAB_IO_BUFFER_SIZE);
+    return mp_const_none;
+}
+
+MP_DEFINE_CONST_FUN_OBJ_2(io_save_obj, io_save);
+

code/numpy/io/io.h

@@ -0,0 +1,16 @@
+/*
+ * This file is part of the micropython-ulab project,
+ *
+ * https://github.com/v923z/micropython-ulab
+ *
+ * The MIT License (MIT)
+ *
+ * Copyright (c) 2022 Zoltán Vörös
+*/
+
+#ifndef _ULAB_IO_
+#define _ULAB_IO_
+
+MP_DECLARE_CONST_FUN_OBJ_2(io_save_obj);
+
+#endif

code/numpy/numpy.c

@@ -23,6 +23,7 @@
 #include "create.h"
 #include "fft/fft.h"
 #include "filter.h"
+#include "io/io.h"
 #include "linalg/linalg.h"
 #include "numerical.h"
 #include "stats.h"
@@ -265,6 +266,9 @@ static const mp_rom_map_elem_t ulab_numpy_globals_table[] = {
     #if ULAB_NUMPY_HAS_ROLL
         { MP_OBJ_NEW_QSTR(MP_QSTR_roll), (mp_obj_t)&numerical_roll_obj },
     #endif
+    #if ULAB_NUMPY_HAS_SAVE
+        { MP_OBJ_NEW_QSTR(MP_QSTR_save), (mp_obj_t)&io_save_obj },
+    #endif
     #if ULAB_NUMPY_HAS_SORT
         { MP_OBJ_NEW_QSTR(MP_QSTR_sort), (mp_obj_t)&numerical_sort_obj },
     #endif

code/ulab.c

@@ -33,7 +33,7 @@
 #include "user/user.h"
 #include "utils/utils.h"
 
-#define ULAB_VERSION 4.0.0
+#define ULAB_VERSION 4.1.0
 #define xstr(s) str(s)
 #define str(s) #s
 

code/ulab.h

@@ -478,6 +478,10 @@
 #define ULAB_NUMPY_HAS_ROLL             (1)
 #endif
 
+#ifndef ULAB_NUMPY_HAS_SAVE
+#define ULAB_NUMPY_HAS_SAVE             (1)
+#endif
+
 #ifndef ULAB_NUMPY_HAS_SORT
 #define ULAB_NUMPY_HAS_SORT             (1)
 #endif

docs/ulab-change-log.md

@@ -1,3 +1,9 @@
+Sat, 8 Jan 2022
+
+    version 4.1.0
+
+    implement numpy.save
+
 Fri, 3 Dec 2021
 
 version 3.3.8