From 4d4ccb7ee497d79317bdf9f52fb3de9d30b8d298 Mon Sep 17 00:00:00 2001
From: Ardavan Oskooi <ardavan.oskooi@gmail.com>
Date: Mon, 4 Jan 2021 08:55:51 -0800
Subject: [PATCH] simplify element-wise array copy in get_source_slice and add
 unit test for get_source (#1462)

---
 python/tests/simulation.py | 16 +++++++++++++---
 src/array_slice.cpp        | 13 +++++++------
 2 files changed, 20 insertions(+), 9 deletions(-)

diff --git a/python/tests/simulation.py b/python/tests/simulation.py
index b874b3eae..b897f955b 100644
--- a/python/tests/simulation.py
+++ b/python/tests/simulation.py
@@ -577,9 +577,19 @@ def print_field(sim):
 
     def test_source_slice(self):
         sim = self.init_simple_simulation()
-        sim.run(until=5)
-        slice = sim.get_source(mp.Ez)
-        print(slice)
+        sim.run(until=1)
+
+        vol1d = mp.Volume(center=mp.Vector3(0.1234,0), size=mp.Vector3(0,5.07))
+        source_slice = sim.get_source(mp.Ez, vol=vol1d)
+        x,y,z,w = sim.get_array_metadata(vol=vol1d)
+        self.assertEqual(source_slice.shape, w.shape)
+        self.assertEqual(np.sum(source_slice), 0)
+
+        vol2d = mp.Volume(center=mp.Vector3(-0.541,0.791), size=mp.Vector3(3.5,2.8))
+        source_slice = sim.get_source(mp.Ez, vol=vol2d)
+        x,y,z,w = sim.get_array_metadata(vol=vol2d)
+        self.assertEqual(source_slice.shape, w.shape)
+        self.assertNotEqual(np.sum(source_slice), 0)
 
     def test_has_mu(self):
 
diff --git a/src/array_slice.cpp b/src/array_slice.cpp
index 4d562ecde..d2ffb51aa 100644
--- a/src/array_slice.cpp
+++ b/src/array_slice.cpp
@@ -714,15 +714,16 @@ cdouble *fields::get_source_slice(const volume &where, component source_slice_co
   rank = get_array_slice_dimensions(where, dims, dirs, true, false);
   slice_size = dims[0] * (rank >= 2 ? dims[1] : 1) * (rank == 3 ? dims[2] : 1);
 
-  if (slice != 0) {
-    for (size_t i = 0; i < slice_size; ++i)
-      slice[i] = slice_collapsed[i];
-    slice = array_to_all(slice, slice_size);
+  if (slice) {
+    memcpy(slice, slice_collapsed, 2 * slice_size * sizeof(double));
+    delete[] (cdouble*) slice_collapsed;
   }
+  else
+    slice = slice_collapsed;
 
-  slice_collapsed = array_to_all(slice_collapsed, slice_size);
+  array_to_all(slice, slice_size);
 
-  return slice_collapsed;
+  return slice;
 }
 
 /***************************************************************/