From e319dad9225d21dd5574cef43369f920b9a07de9 Mon Sep 17 00:00:00 2001
From: "Steven G. Johnson" <stevenj@mit.edu>
Date: Tue, 29 Dec 2020 15:27:23 -0500
Subject: [PATCH] return numpy arrays from std::vectors (#1458)

---
 python/meep.i | 18 ++++++++++++++++++
 1 file changed, 18 insertions(+)

diff --git a/python/meep.i b/python/meep.i
index 58f5b10c3..0364ceee3 100644
--- a/python/meep.i
+++ b/python/meep.i
@@ -1471,6 +1471,18 @@ void _get_gradient(PyObject *grad, PyObject *fields_a, PyObject *fields_f, PyObj
 %template(IntVector) std::vector<int>;
 %template(DoubleVector) std::vector<double>;
 
+// use NumPy arrays for returning common std::vector types:
+%typemap(out) std::vector<double> {
+    npy_intp vec_len = (npy_intp) $1.size();
+    $result = PyArray_SimpleNew(1, &vec_len, NPY_DOUBLE);
+    memcpy(PyArray_DATA((PyArrayObject*) $result), &$1[0], vec_len * sizeof(double));
+}
+%typemap(out) std::vector<int> {
+    npy_intp vec_len = (npy_intp) $1.size();
+    $result = PyArray_SimpleNew(1, &vec_len, NPY_INT);
+    memcpy(PyArray_DATA((PyArrayObject*) $result), &$1[0], vec_len * sizeof(int));
+}
+
 %include "vec.i"
 %include "meep.hpp"
 %include "meep/mympi.hpp"
@@ -1482,6 +1494,12 @@ void _get_gradient(PyObject *grad, PyObject *fields_a, PyObject *fields_f, PyObj
 %include "std_complex.i"
 %template(ComplexVector) std::vector<std::complex<double> >;
 
+%typemap(out) std::vector<std::complex<double> > {
+    npy_intp vec_len = (npy_intp) $1.size();
+    $result = PyArray_SimpleNew(1, &vec_len, NPY_COMPLEX128);
+    memcpy(PyArray_DATA((PyArrayObject*) $result), &$1[0], vec_len * sizeof(double) * 2);
+}
+
 std::vector<struct meep::sourcedata> meep::dft_near2far::near_sourcedata(const meep::vec &x_0, double* farpt_list, size_t nfar_pts, std::complex<double>* dJ);
 
 void meep::fields::add_srcdata(struct meep::sourcedata cur_data, meep::src_time *src, size_t n, std::complex<double>* amp_arr);