Merge pull request #316 from v923z/stats

simplify array contraction

Author: v923z

code/micropython.mk

@@ -18,6 +18,7 @@ SRC_USERMOD += $(USERMODULES_DIR)/numpy/linalg/linalg.c
 SRC_USERMOD += $(USERMODULES_DIR)/numpy/linalg/linalg_tools.c
 SRC_USERMOD += $(USERMODULES_DIR)/numpy/numerical/numerical.c
 SRC_USERMOD += $(USERMODULES_DIR)/numpy/poly/poly.c
+SRC_USERMOD += $(USERMODULES_DIR)/numpy/stats/stats.c
 SRC_USERMOD += $(USERMODULES_DIR)/numpy/vector/vector.c
 SRC_USERMOD += $(USERMODULES_DIR)/user/user.c
 

code/ndarray.c

@@ -598,7 +598,6 @@ bool ndarray_is_dense(ndarray_obj_t *ndarray) {
 
 ndarray_obj_t *ndarray_new_ndarray(uint8_t ndim, size_t *shape, int32_t *strides, uint8_t dtype) {
     // Creates the base ndarray with shape, and initialises the values to straight 0s
-    // the function should work in the general n-dimensional case
     ndarray_obj_t *ndarray = m_new_obj(ndarray_obj_t);
     ndarray->base.type = &ulab_ndarray_type;
     ndarray->dtype = dtype == NDARRAY_BOOL ? NDARRAY_UINT8 : dtype;
@@ -618,10 +617,12 @@ ndarray_obj_t *ndarray_new_ndarray(uint8_t ndim, size_t *shape, int32_t *strides
         ndarray->len *= shape[i-1];
     }
 
-    uint8_t *array = m_new(byte, ndarray->itemsize * ndarray->len);
+    // if the length is 0, still allocate a single item, so that contractions can be handled
+    size_t len = ndarray->itemsize * MAX(1, ndarray->len);
+    uint8_t *array = m_new(byte, len);
     // this should set all elements to 0, irrespective of the of the dtype (all bits are zero)
     // we could, perhaps, leave this step out, and initialise the array only, when needed
-    memset(array, 0, ndarray->len * ndarray->itemsize);
+    memset(array, 0, len);
     ndarray->array = array;
     return ndarray;
 }

code/numpy/linalg/linalg.c

@@ -385,14 +385,8 @@ static mp_obj_t linalg_norm(size_t n_args, const mp_obj_t *pos_args, mp_map_t *k
         // always get a float, so that we don't have to resolve the dtype later
         mp_float_t (*func)(void *) = ndarray_get_float_function(ndarray->dtype);
         shape_strides _shape_strides = tools_reduce_axes(ndarray, axis);
-        mp_float_t *rarray = NULL;
-        ndarray_obj_t *results = NULL;
-        if((axis != mp_const_none) && (ndarray->ndim > 1)) {
-            results = ndarray_new_dense_ndarray(MAX(1, ndarray->ndim-1), _shape_strides.shape, NDARRAY_FLOAT);
-            rarray = results->array;
-        } else {
-            rarray = m_new(mp_float_t, 1);
-        }
+        ndarray_obj_t *results = ndarray_new_dense_ndarray(_shape_strides.ndim, _shape_strides.shape, NDARRAY_FLOAT);
+        mp_float_t *rarray = (mp_float_t *)results->array;
 
         #if ULAB_MAX_DIMS > 3
         size_t i = 0;
@@ -418,28 +412,26 @@ static mp_obj_t linalg_norm(size_t n_args, const mp_obj_t *pos_args, mp_map_t *k
                         l++;
                     } while(l < _shape_strides.shape[0]);
                     *rarray = MICROPY_FLOAT_C_FUN(sqrt)(dot * (count - 1));
-                    if(results != NULL) {
-                        rarray++;
-                    }
                 #if ULAB_MAX_DIMS > 1
+                    rarray += _shape_strides.increment;
                     array -= _shape_strides.strides[0] * _shape_strides.shape[0];
                     array += _shape_strides.strides[ULAB_MAX_DIMS - 1];
                     k++;
                 } while(k < _shape_strides.shape[ULAB_MAX_DIMS - 1]);
                 #endif
             #if ULAB_MAX_DIMS > 2
-                array -= _shape_strides.strides[ULAB_MAX_DIMS - 1] * _shape_strides.shape[ULAB_MAX_DIMS-1];
+                array -= _shape_strides.strides[ULAB_MAX_DIMS - 1] * _shape_strides.shape[ULAB_MAX_DIMS - 1];
                 array += _shape_strides.strides[ULAB_MAX_DIMS - 2];
                 j++;
             } while(j < _shape_strides.shape[ULAB_MAX_DIMS - 2]);
             #endif
         #if ULAB_MAX_DIMS > 3
-            array -= _shape_strides.strides[ULAB_MAX_DIMS - 2] * _shape_strides.shape[ULAB_MAX_DIMS-2];
+            array -= _shape_strides.strides[ULAB_MAX_DIMS - 2] * _shape_strides.shape[ULAB_MAX_DIMS - 2];
             array += _shape_strides.strides[ULAB_MAX_DIMS - 3];
             i++;
         } while(i < _shape_strides.shape[ULAB_MAX_DIMS - 3]);
         #endif
-        if(results == NULL) {
+        if(results->ndim == 0) {
             return mp_obj_new_float(*rarray);
         }
         return results;