Singleton DFT fix (#1333)

* fix bug that caused segfault for singleton elements * exit earlier * revert initialization Co-authored-by: Alec Hammond <[email protected]>
NanoComp · Sep 9, 2020 · b092291 · b092291
1 parent a5e6237
commit b092291
Show file tree

Hide file tree

Showing 2 changed files with 10 additions and 6 deletions.
diff --git a/python/meep.i b/python/meep.i
@@ -438,11 +438,14 @@ PyObject *_get_dft_array(meep::fields *f, dft_type dft, meep::component c, int n
     size_t dims[3];
     std::complex<double> *dft_arr = f->get_dft_array(dft, c, num_freq, &rank, dims);
 
-    if (rank==0 || dft_arr==NULL){ // this can happen e.g. if component c vanishes by symmetry
+    if (dft_arr==NULL){ // this can happen e.g. if component c vanishes by symmetry
      std::complex<double> d[1] = {std::complex<double>(0,0)};
      return PyArray_SimpleNewFromData(0, 0, NPY_CDOUBLE, d);
     }
 
+    if (rank == 0) // singleton results
+     return PyArray_SimpleNewFromData(0, 0, NPY_CDOUBLE, dft_arr);
+
     size_t length = 1;
     npy_intp *arr_dims = new npy_intp[rank];
     for (int i = 0; i < rank; ++i) {

diff --git a/src/array_slice.cpp b/src/array_slice.cpp
@@ -662,6 +662,10 @@ double *collapse_array(double *array, int *rank, size_t dims[3], direction dirs[
       reduced_dims[reduced_rank++] = dims[r];
     }
   }
+  if (reduced_rank==0) {
+    *rank = 0;
+    return array; // return array as is for singleton use case
+  }
   if (reduced_rank == full_rank) return array; // nothing to collapse
 
   /*--------------------------------------------------------------*/
@@ -673,11 +677,8 @@ double *collapse_array(double *array, int *rank, size_t dims[3], direction dirs[
   else if (full_rank == 3) {
     stride[0] = dims[1] * dims[2];
     stride[1] = dims[2];
-    if (reduced_stride[0] != 0)
-      reduced_stride[0] = reduced_dims[1];
-    else if (reduced_stride[1] != 0)
-      reduced_stride[1] = reduced_dims[1];
-    // else: two degenerate dimensions->reduced array is 1-diml, no strides needed
+    if (reduced_rank == 2)
+      reduced_stride[reduced_stride[0] != 0 ? 0 : 1] = reduced_dims[1];
   }
 
   /*--------------------------------------------------------------*/