meshgrid (#740) (#746)

* feat: `meshgrid` (#740)

• Default implementation equivalent to NumPy's with `copy=True` and `sparse=False`
• Add 2D, 3D, 4D test cases
• Update docs

* docs: update version to 6.12.0

Author: iyassou

code/numpy/create.c

@@ -747,6 +747,143 @@ mp_obj_t create_logspace(size_t n_args, const mp_obj_t *pos_args, mp_map_t *kw_a
 MP_DEFINE_CONST_FUN_OBJ_KW(create_logspace_obj, 2, create_logspace);
 #endif
 
+#if ULAB_NUMPY_HAS_MESHGRID
+//| def meshgrid(*xi, indexing="xy"):
+//|     """
+//|     .. param: xi
+//|        1-D arrays representing the coordinates of a grid
+//|     .. param: indexing
+//|        "xy" (Cartesian) or "ij" (matrix) indexing of output
+//| 
+//|     Return a tuple of coordinate matrices from coordinate vectors.
+//|     Implements the NumPy equivalent with copy=True and sparse=False."""
+//|     ...
+
+mp_obj_t create_meshgrid(size_t n_args, const mp_obj_t *pos_args, mp_map_t *kw_args) {
+    static const mp_arg_t allowed_args[] = {
+        { MP_QSTR_indexing, MP_ARG_KW_ONLY | MP_ARG_OBJ, { .u_rom_obj = MP_ROM_QSTR(MP_QSTR_xy) } },
+    };
+
+    mp_arg_val_t args[MP_ARRAY_SIZE(allowed_args)];
+    mp_arg_parse_all(0, pos_args, kw_args, MP_ARRAY_SIZE(allowed_args), allowed_args, args);
+
+    if(!mp_obj_is_str(args[0].u_obj)) {
+        mp_raise_TypeError(MP_ERROR_TEXT("indexing should be 'xy' or 'ij'"));
+    }
+    bool cartesian = true;
+    const char *_indexing = mp_obj_str_get_str(args[0].u_obj);
+    if(strcmp(_indexing, "ij") == 0) {
+        cartesian = false;
+    } else if (strcmp(_indexing, "xy") != 0) {
+        mp_raise_ValueError(MP_ERROR_TEXT("indexing must be 'xy' or 'ij'"));
+    }
+
+    if(n_args == 0) {
+        return mp_const_empty_tuple;
+    }
+    if(n_args > ULAB_MAX_DIMS) {
+        mp_raise_ValueError(MP_ERROR_TEXT("too many input arrays"));
+    }
+    
+    size_t *shape = m_new0(size_t, ULAB_MAX_DIMS);
+    size_t shape_offset = ULAB_MAX_DIMS - n_args;
+    for(size_t i = 0; i < n_args; i++) {
+        if(!mp_obj_is_type(pos_args[i], &ulab_ndarray_type)) {
+            mp_raise_TypeError(MP_ERROR_TEXT("arguments must be ndarrays"));
+        }
+        ndarray_obj_t *in = MP_OBJ_TO_PTR(pos_args[i]);
+        if(in->ndim != 1) {
+            mp_raise_ValueError(MP_ERROR_TEXT("arguments must be 1D arrays"));
+        }
+        shape[shape_offset + i] = in->shape[ULAB_MAX_DIMS - 1];
+    }
+    if(cartesian && n_args >= 2) {
+        SWAP(size_t, shape[shape_offset], shape[shape_offset + 1]);
+    }
+
+    mp_obj_tuple_t *tuple = MP_OBJ_TO_PTR(mp_obj_new_tuple(n_args, NULL));
+    for(size_t p = 0; p < n_args; p++) {
+        ndarray_obj_t *source = MP_OBJ_TO_PTR(pos_args[p]);
+        ndarray_obj_t *target = ndarray_new_dense_ndarray(n_args, shape, source->dtype);
+        uint8_t *tarray = (uint8_t *)target->array;
+        uint8_t *sarray = (uint8_t *)source->array;
+        size_t itemsize = source->itemsize;
+        int32_t source_stride = source->strides[ULAB_MAX_DIMS - 1];
+
+        size_t active_axis_index = p;
+        if(cartesian && n_args >= 2) {
+            if(p == 0) {
+                active_axis_index = 1;
+            } else if(p == 1) {
+                active_axis_index = 0;
+            }
+        }
+        size_t active_dim_ulab = shape_offset + active_axis_index;
+
+        #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 {
+                        // Determine the index in the source array based on
+                        // which loop corresponds to the active dimension
+                        size_t source_index = 0;
+                        #if ULAB_MAX_DIMS > 3
+                        if(active_dim_ulab == ULAB_MAX_DIMS - 4) {
+                            source_index = i;
+                        }
+                        #endif
+                        #if ULAB_MAX_DIMS > 2
+                        if(active_dim_ulab == ULAB_MAX_DIMS - 3) {
+                            source_index = j;
+                        }
+                        #endif
+                        #if ULAB_MAX_DIMS > 1
+                        if(active_dim_ulab == ULAB_MAX_DIMS - 2) {
+                            source_index = k;
+                        }
+                        #endif
+                        if(active_dim_ulab == ULAB_MAX_DIMS - 1) {
+                            source_index = l;
+                        }
+
+                        memcpy(tarray, sarray + (source_index * source_stride), itemsize);
+                        tarray += itemsize;
+                        
+                        l++;
+                    } while(l < target->shape[ULAB_MAX_DIMS - 1]);
+                #if ULAB_MAX_DIMS > 1
+                    k++;
+                } while(k < target->shape[ULAB_MAX_DIMS - 2]);
+                #endif
+            #if ULAB_MAX_DIMS > 2
+                j++;
+            } while(j < target->shape[ULAB_MAX_DIMS - 3]);
+            #endif
+        #if ULAB_MAX_DIMS > 3
+            i++;
+        } while(i < target->shape[ULAB_MAX_DIMS - 4]);
+        #endif
+
+        tuple->items[p] = MP_OBJ_FROM_PTR(target);
+    }
+
+    m_del(size_t, shape, ULAB_MAX_DIMS);
+    return MP_OBJ_FROM_PTR(tuple);
+}
+
+MP_DEFINE_CONST_FUN_OBJ_KW(create_meshgrid_obj, 0, create_meshgrid);
+#endif
+
 #if ULAB_NUMPY_HAS_ONES
 //| def ones(shape: Union[int, Tuple[int, ...]], *, dtype: _DType = ulab.numpy.float) -> ulab.numpy.ndarray:
 //|    """

code/numpy/create.h

@@ -57,6 +57,11 @@ mp_obj_t create_logspace(size_t , const mp_obj_t *, mp_map_t *);
 MP_DECLARE_CONST_FUN_OBJ_KW(create_logspace_obj);
 #endif
 
+#if ULAB_NUMPY_HAS_MESHGRID
+mp_obj_t create_meshgrid(size_t , const mp_obj_t *, mp_map_t *);
+MP_DECLARE_CONST_FUN_OBJ_KW(create_meshgrid_obj);
+#endif
+
 #if ULAB_NUMPY_HAS_ONES
 mp_obj_t create_ones(size_t , const mp_obj_t *, mp_map_t *);
 MP_DECLARE_CONST_FUN_OBJ_KW(create_ones_obj);

code/numpy/numpy.c

@@ -163,6 +163,9 @@ static const mp_rom_map_elem_t ulab_numpy_globals_table[] = {
     #if ULAB_NUMPY_HAS_LOGSPACE
         { MP_ROM_QSTR(MP_QSTR_logspace), MP_ROM_PTR(&create_logspace_obj) },
     #endif
+    #if ULAB_NUMPY_HAS_MESHGRID
+        { MP_ROM_QSTR(MP_QSTR_meshgrid), MP_ROM_PTR(&create_meshgrid_obj) },
+    #endif
     #if ULAB_NUMPY_HAS_ONES
         { MP_ROM_QSTR(MP_QSTR_ones), MP_ROM_PTR(&create_ones_obj) },
     #endif

code/ulab.c

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

code/ulab.h

@@ -365,6 +365,10 @@
 #define ULAB_NUMPY_HAS_LOGSPACE         (1)
 #endif
 
+#ifndef ULAB_NUMPY_HAS_MESHGRID
+#define ULAB_NUMPY_HAS_MESHGRID         (1)
+#endif
+
 #ifndef ULAB_NUMPY_HAS_ONES
 #define ULAB_NUMPY_HAS_ONES             (1)
 #endif

docs/manual/source/conf.py

@@ -27,7 +27,7 @@
 author = 'Zoltán Vörös'
 
 # The full version, including alpha/beta/rc tags
-release = '6.11.0'
+release = '6.12.0'
 
 
 # -- General configuration ---------------------------------------------------

docs/ulab-convert.ipynb

@@ -14,7 +14,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 1,
+   "execution_count": null,
    "metadata": {
     "ExecuteTime": {
      "end_time": "2022-02-09T06:27:15.118699Z",
@@ -61,7 +61,7 @@
     "author = 'Zoltán Vörös'\n",
     "\n",
     "# The full version, including alpha/beta/rc tags\n",
-    "release = '6.11.0'\n",
+    "release = '6.12.0'\n",
     "\n",
     "\n",
     "# -- General configuration ---------------------------------------------------\n",

docs/ulab-ndarray.ipynb

@@ -507,7 +507,7 @@
    "source": [
     "# Array initialisation functions\n",
     "\n",
-    "There are nine functions that can be used for initialising an array. Starred functions accept `complex` as the value of the `dtype`, if the firmware was compiled with complex support.\n",
+    "There are several functions that can be used for initialising an array. Starred functions accept `complex` as the value of the `dtype`, if the firmware was compiled with complex support.\n",
     "\n",
     "1. [numpy.arange](#arange)\n",
     "1. [numpy.concatenate](#concatenate)\n",
@@ -518,6 +518,7 @@
     "1. [numpy.full*](#full)\n",
     "1. [numpy.linspace*](#linspace)\n",
     "1. [numpy.logspace](#logspace)\n",
+    "1. [numpy.meshgrid](#meshgrid)\n",
     "1. [numpy.ones*](#ones)\n",
     "1. [numpy.zeros*](#zeros)"
    ]
@@ -1152,6 +1153,59 @@
     "print('num=5:\\t\\t\\t', np.logspace(1, 10, num=5, endpoint=False, base=2))"
    ]
   },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## meshgrid\n",
+    "\n",
+    "`numpy`: https://numpy.org/doc/stable/reference/generated/numpy.meshgrid.html\n",
+    "\n",
+    "The `meshgrid` function returns coordinate matrices from coordinate vectors. Given *N* one-dimensional arrays, it returns a tuple of *N* N-dimensional arrays.\n",
+    "\n",
+    "Make sure *N* is not larger than the maximum dimension of your firmware.\n",
+    "\n",
+    "The `indexing` keyword argument determines the arrangement of the output. The default is `'xy'` (Cartesian indexing), where the first two dimensions are swapped. This is commonly used for generating grids for 2D plotting. If `indexing` is set to `'ij'` (matrix indexing), the dimensions of the output arrays correspond exactly to the order of the input arrays.\n",
+    "\n",
+    "Note that unlike `numpy`, `ulab`'s implementation always behaves as if `sparse=False` and `copy=True` i.e. it always returns fully allocated dense arrays."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 12,
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "X:\n",
+      "array([[1, 2, 3],\n",
+      "       [1, 2, 3]], dtype=uint8)\n",
+      "Y:\n",
+      "array([[10, 10, 10],\n",
+      "       [20, 20, 20]], dtype=uint8)\n",
+      "\n",
+      "\n"
+     ]
+    }
+   ],
+   "source": [
+    "%%micropython -unix 1\n",
+    "\n",
+    "from ulab import numpy as np\n",
+    "\n",
+    "x = np.array([1, 2, 3], dtype=np.uint8)\n",
+    "y = np.array([10, 20], dtype=np.uint8)\n",
+    "\n",
+    "X, Y = np.meshgrid(x, y)\n",
+    "\n",
+    "print(\"X:\")\n",
+    "print(X)\n",
+    "print(\"Y:\")\n",
+    "print(Y)"
+   ]
+  },
   {
    "cell_type": "markdown",
    "metadata": {},

tests/2d/numpy/meshgrid.py

@@ -0,0 +1,39 @@
+try:
+    from ulab import numpy as np
+except:
+    import numpy as np
+
+# 1. Standard 2D Case - XY Indexing (Default)
+# Expected Output: x repeated in rows, y repeated in cols. Shape (len(y), len(x))
+print("--- 2D xy ---")
+x = np.array([1, 2, 3], dtype=np.uint8)
+y = np.array([10, 20], dtype=np.uint8)
+
+xv, yv = np.meshgrid(x, y, indexing='xy')
+print(xv)
+print(yv)
+
+# 2. Standard 2D Case - IJ Indexing
+# Expected Output: x repeated in cols, y repeated in rows. Shape (len(x), len(y))
+print("\n--- 2D ij ---")
+xi, yi = np.meshgrid(x, y, indexing='ij')
+print(xi)
+print(yi)
+
+print("\n=== 2D Strides ===")
+# x: [0, 1, 2, 3] -> [::2] -> [0, 2] (len 2)
+x = np.array([0, 1, 2, 3], dtype=np.uint8)[::2]
+# y: [10, 20, 30] -> [::-1] -> [30, 20, 10] (len 3)
+y = np.array([10, 20, 30], dtype=np.uint8)[::-1]
+
+print("--- XY ---")
+# Shape: (3, 2)
+mx, my = np.meshgrid(x, y, indexing='xy')
+print(mx)
+print(my)
+
+print("--- IJ ---")
+# Shape: (2, 3)
+mx, my = np.meshgrid(x, y, indexing='ij')
+print(mx)
+print(my)

tests/2d/numpy/meshgrid.py.exp

@@ -0,0 +1,27 @@
+--- 2D xy ---
+array([[1, 2, 3],
+       [1, 2, 3]], dtype=uint8)
+array([[10, 10, 10],
+       [20, 20, 20]], dtype=uint8)
+
+--- 2D ij ---
+array([[1, 1],
+       [2, 2],
+       [3, 3]], dtype=uint8)
+array([[10, 20],
+       [10, 20],
+       [10, 20]], dtype=uint8)
+
+=== 2D Strides ===
+--- XY ---
+array([[0, 2],
+       [0, 2],
+       [0, 2]], dtype=uint8)
+array([[30, 30],
+       [20, 20],
+       [10, 10]], dtype=uint8)
+--- IJ ---
+array([[0, 0, 0],
+       [2, 2, 2]], dtype=uint8)
+array([[30, 20, 10],
+       [30, 20, 10]], dtype=uint8)