Initial high level interface to get_array_slice

python/Makefile.am

@@ -33,6 +33,7 @@ TESTS =                                 \
     $(TEST_DIR)/3rd_harm_1d.py          \
     $(TEST_DIR)/antenna_radiation.py    \
     $(TEST_DIR)/bend_flux.py            \
+    $(TEST_DIR)/cavity_arrayslice.py    \
     $(TEST_DIR)/cyl_ellipsoid.py        \
     $(TEST_DIR)/geom.py                 \
     $(TEST_DIR)/holey_wvg_bands.py      \

python/examples/cavity_arrayslice.py

@@ -1,11 +1,8 @@
-import unittest
 import meep as mp
 import numpy as np
 import matplotlib.pyplot as plt
 
-##################################################
 # set up the geometry
-##################################################
 eps = 13
 w = 1.2
 r = 0.36
@@ -33,60 +30,44 @@
     geometry.append(mp.Cylinder(r, center=mp.Vector3(d / -2 - i)))
 
 sim = mp.Simulation(cell_size=cell,
-                         geometry=geometry,
-                         sources=[],
-                         boundary_layers=[mp.PML(dpml)],
-                         resolution=20)
+                    geometry=geometry,
+                    sources=[],
+                    boundary_layers=[mp.PML(dpml)],
+                    resolution=20)
 
-##################################################
-# add sources 
-##################################################
+# add sources
 sim.sources = [mp.Source(mp.GaussianSource(fcen, fwidth=df),
-                    mp.Hz, mp.Vector3())]
+               mp.Hz, mp.Vector3())]
 
-#sim.symmetries = [mp.Mirror(mp.Y, phase=-1), mp.Mirror(mp.X, phase=-1)]
-
-##################################################
 # run until sources are finished (and no later)
-##################################################
-sim._run_sources_until(0,[])
+sim._run_sources_until(0, [])
 
-##################################################
 # get 1D and 2D array slices
-##################################################
-xMin = -0.25*sx;
-xMax = +0.25*sx;
-yMin = -0.15*sy;
-yMax = +0.15*sy;
+xMin = -0.25 * sx
+xMax = +0.25 * sx
+yMin = -0.15 * sy
+yMax = +0.15 * sy
 
 # 1D slice of Hz data
-dims1d=np.zeros(3,dtype=np.int32)
-v1d=mp.volume( mp.vec(xMin, 0.0), mp.vec(xMax, 0.0) )
-rank1d  = sim.fields.get_array_slice_dimensions(v1d, dims1d);
-NX1 = dims1d[0];
-slice1d = np.zeros(NX1, dtype=np.double);
-sim.fields.get_array_slice(v1d, mp.Hz, slice1d);
+size_1d = mp.Vector3(xMax - xMin)
+center_1d = mp.Vector3((xMin + xMax) / 2)
+slice1d = sim.get_array(center=center_1d, size=size_1d, component=mp.Hz)
 
 # 2D slice of Hz data
-dims2d=np.zeros(3,dtype=np.int32)
-v2d=mp.volume( mp.vec(xMin, yMin), mp.vec(xMax, yMax) )
-rank2d  = sim.fields.get_array_slice_dimensions(v2d, dims2d);
-NX2 = dims2d[0];
-NY2 = dims2d[1];
-N2 = NX2*NY2;
-slice2d = np.zeros(N2, dtype=np.double);
-sim.fields.get_array_slice(v2d, mp.Hz, slice2d);
+size_2d = mp.Vector3(xMax - xMin, yMax - yMin)
+center_2d = mp.Vector3((xMin + xMax) / 2, (yMin + yMax) / 2)
+slice2d = sim.get_array(center=center_2d, size=size_2d, component=mp.Hz)
 
 # plot 1D slice
-plt.subplot(1,2,1);
-x1d=np.linspace(xMin, xMax, dims1d[0]);
-plt.plot(x1d, slice1d);
+plt.subplot(1, 2, 1)
+x1d = np.linspace(xMin, xMax, len(slice1d))
+plt.plot(x1d, slice1d)
 
 # plot 2D slice
-plt.subplot(1,2,2);
-dy = (yMax-yMin) / dims2d[1];
-dx = (xMax-xMin) / dims2d[0];
-(x2d,y2d)=np.mgrid[ slice(xMin, xMax, dx), slice(yMin, yMax, dy)];
-plt.contourf( x2d, y2d, np.reshape(slice2d, (dims2d[0], dims2d[1])) )
-plt.colorbar();
+plt.subplot(1, 2, 2)
+dy = (yMax - yMin) / slice2d.shape[1]
+dx = (xMax - xMin) / slice2d.shape[0]
+(x2d, y2d) = np.mgrid[slice(xMin, xMax, dx), slice(yMin, yMax, dy)]
+plt.contourf(x2d, y2d, slice2d)
+plt.colorbar()
 plt.show()

python/meep.i

@@ -332,8 +332,39 @@ meep::volume_list *make_volume_list(const meep::volume &v, int c,
 }
 
 // Typemap suite for array_slice
-// TODO: add (cdouble *, int) version
-%apply (double* INPLACE_ARRAY1, int DIM1) {(double *slice, int slice_length)};
+
+%typecheck(SWIG_TYPECHECK_POINTER, fragment="NumPy_Fragments") double* slice {
+    $1 = is_array($input);
+}
+
+%typemap(in, fragment="NumPy_Macros") double* slice {
+    $1 = (double *)array_data($input);
+}
+
+%typecheck(SWIG_TYPECHECK_POINTER, fragment="NumPy_Fragments") std::complex<double>* slice {
+    $1 = is_array($input);
+}
+
+%typemap(in) std::complex<double>* slice {
+    $1 = (std::complex<double> *)array_data($input);
+}
+
+%typecheck(SWIG_TYPECHECK_POINTER) meep::component {
+    $1 = PyInteger_Check($input) && PyInteger_AsLong($input) < 100;
+}
+
+%typemap(in) meep::component {
+    $1 = static_cast<meep::component>(PyInteger_AsLong($input));
+}
+
+%typecheck(SWIG_TYPECHECK_POINTER) meep::derived_component {
+    $1 = PyInteger_Check($input) && PyInteger_AsLong($input) >= 100;
+}
+
+%typemap(in) meep::derived_component {
+    $1 = static_cast<meep::derived_component>(PyInteger_AsLong($input));
+}
+
 %apply int INPLACE_ARRAY1[ANY] { int [3] };
 
 // Rename python builtins

python/simulation.py

@@ -652,13 +652,44 @@ def output_components(self, fname, *components):
         if self.output_append_h5 is None:
             self.output_h5_hook(self.fields.h5file_name(fname, self._get_filename_prefix(), True))
 
+    def get_array(self, center, size, component=mp.Ez, cmplx=None, arr=None):
+        dim_sizes = np.zeros(3, dtype=np.int32)
+        vol = Volume(center, size=size, dims=self.dimensions, is_cylindrical=self.is_cylindrical)
+        self.fields.get_array_slice_dimensions(vol.swigobj, dim_sizes)
+
+        dims = [s for s in dim_sizes if s != 0]
+
+        if cmplx is None:
+            cmplx = component < mp.Dielectric and not self.fields.is_real
+
+        if arr is not None:
+            if cmplx and not np.iscomplexobj(arr):
+                raise ValueError("Requested a complex slice, but provided array of type {}.".format(arr.dtype))
+
+            for a, b in zip(arr.shape, dims):
+                if a != b:
+                    fmt = "Expected dimensions {}, but got {}"
+                    raise ValueError(fmt.format(dims, arr.shape))
+
+            arr = np.require(arr, requirements=['C', 'W'])
+
+        else:
+            arr = np.zeros(dims, dtype=np.complex128 if cmplx else np.float64)
+
+        if np.iscomplexobj(arr):
+            self.fields.get_complex_array_slice(vol.swigobj, component, arr)
+        else:
+            self.fields.get_array_slice(vol.swigobj, component, arr)
+
+        return arr
+
     def change_k_point(self, k):
         self.k_point = k
 
         if self.fields:
-            cond1 = not (not self.k_point or self.k_point == mp.Vector3())
+            needs_complex_fields = not (not self.k_point or self.k_point == mp.Vector3())
 
-            if cond1 and self.fields.is_real != 0:
+            if needs_complex_fields and self.fields.is_real:
                 self.fields = None
                 self._init_fields()
             else:
@@ -680,11 +711,11 @@ def run(self, *step_funcs, **kwargs):
         if kwargs:
             raise ValueError("Unrecognized keyword arguments: {}".format(kwargs.keys()))
 
-        if until_after_sources:
+        if until_after_sources is not None:
             self._run_sources_until(until_after_sources, step_funcs)
-        elif k_points:
+        elif k_points is not None:
             return self._run_k_points(k_points, *step_funcs)
-        elif until:
+        elif until is not None:
             self._run_until(until, step_funcs)
         else:
             raise ValueError("Invalid run configuration")

python/tests/cavity_arrayslice.py

@@ -0,0 +1,114 @@
+from __future__ import division
+
+import os
+import unittest
+
+import meep as mp
+import numpy as np
+
+
+class TestCavityArraySlice(unittest.TestCase):
+
+    data_dir = os.path.abspath(os.path.realpath(os.path.join(__file__, '..', 'data')))
+    expected_1d = np.load(os.path.join(data_dir, 'cavity_arrayslice_1d.npy'))
+    expected_2d = np.load(os.path.join(data_dir, 'cavity_arrayslice_2d.npy'))
+
+    def setUp(self):
+
+        r = 0.36
+        d = 1.4
+        sy = 6
+        pad = 2
+        dpml = 1
+        sx = (2 * (pad + dpml + 3)) + d - 1
+
+        cell = mp.Vector3(sx, sy, 0)
+
+        blk = mp.Block(size=mp.Vector3(1e20, 1.2, 1e20),
+                       material=mp.Medium(epsilon=13))
+
+        geometry = [blk]
+
+        for i in range(3):
+            geometry.append(mp.Cylinder(r, center=mp.Vector3(d / 2 + i)))
+
+        for i in range(3):
+            geometry.append(mp.Cylinder(r, center=mp.Vector3(d / -2 - i)))
+
+        sources = [mp.Source(mp.GaussianSource(0.25, fwidth=0.2), mp.Hz, mp.Vector3())]
+
+        self.sim = mp.Simulation(
+            cell_size=cell,
+            geometry=geometry,
+            sources=sources,
+            boundary_layers=[mp.PML(dpml)],
+            resolution=20
+        )
+
+        self.x_min = -0.25 * sx
+        self.x_max = +0.25 * sx
+        self.y_min = -0.15 * sy
+        self.y_max = +0.15 * sy
+
+        self.size_1d = mp.Vector3(self.x_max - self.x_min)
+        self.center_1d = mp.Vector3((self.x_min + self.x_max) / 2)
+
+        self.size_2d = mp.Vector3(self.x_max - self.x_min, self.y_max - self.y_min)
+        self.center_2d = mp.Vector3((self.x_min + self.x_max) / 2, (self.y_min + self.y_max) / 2)
+
+    def test_1d_slice(self):
+        self.sim.run(until_after_sources=0)
+        hl_slice1d = self.sim.get_array(center=self.center_1d, size=self.size_1d, component=mp.Hz)
+        np.testing.assert_allclose(self.expected_1d, hl_slice1d)
+
+    def test_2d_slice(self):
+        self.sim.run(until_after_sources=0)
+        hl_slice2d = self.sim.get_array(center=self.center_2d, size=self.size_2d, component=mp.Hz)
+        np.testing.assert_allclose(self.expected_2d, hl_slice2d)
+
+    def test_1d_slice_user_array(self):
+        self.sim.run(until_after_sources=0)
+        arr = np.zeros(126, dtype=np.float64)
+        self.sim.get_array(center=self.center_1d, size=self.size_1d, component=mp.Hz, arr=arr)
+        np.testing.assert_allclose(self.expected_1d, arr)
+
+    def test_2d_slice_user_array(self):
+        self.sim.run(until_after_sources=0)
+        arr = np.zeros((126, 38), dtype=np.float64)
+        self.sim.get_array(center=self.center_2d, size=self.size_2d, component=mp.Hz, arr=arr)
+        np.testing.assert_allclose(self.expected_2d, arr)
+
+    def test_illegal_user_array(self):
+        self.sim.run(until_after_sources=0)
+
+        with self.assertRaises(ValueError):
+            arr = np.zeros(128)
+            self.sim.get_array(center=self.center_1d, size=self.size_1d,
+                               component=mp.Hz, arr=arr)
+
+        with self.assertRaises(ValueError):
+            arr = np.zeros((126, 39))
+            self.sim.get_array(center=self.center_2d, size=self.size_2d,
+                               component=mp.Hz, arr=arr)
+
+        with self.assertRaises(ValueError):
+            arr = np.zeros((126, 38))
+            self.sim.get_array(center=self.center_2d, size=self.size_2d,
+                               component=mp.Hz, cmplx=True, arr=arr)
+
+    def test_1d_complex_slice(self):
+        self.sim.run(until_after_sources=0)
+        hl_slice1d = self.sim.get_array(center=self.center_1d, size=self.size_1d,
+                                        component=mp.Hz, cmplx=True)
+        self.assertTrue(hl_slice1d.dtype == np.complex128)
+        self.assertTrue(hl_slice1d.shape[0] == 126)
+
+    def test_2d_complex_slice(self):
+        self.sim.run(until_after_sources=0)
+        hl_slice2d = self.sim.get_array(center=self.center_2d, size=self.size_2d,
+                                        component=mp.Hz, cmplx=True)
+        self.assertTrue(hl_slice2d.dtype == np.complex128)
+        self.assertTrue(hl_slice2d.shape[0] == 126 and hl_slice2d.shape[1] == 38)
+
+if __name__ == '__main__':
+    unittest.main()

python/typemap_utils.cpp

@@ -19,8 +19,12 @@
 
 #if PY_MAJOR_VERSION >= 3
     #define PyObject_ToCharPtr(n) PyUnicode_AsUTF8(n)
+    #define PyInteger_Check(n) PyLong_Check(n)
+    #define PyInteger_AsLong(n) PyLong_AsLong(n)
 #else
     #define PyObject_ToCharPtr(n) PyString_AsString(n)
+    #define PyInteger_Check(n) PyInt_Check(n)
+    #define PyInteger_AsLong(n) PyInt_AsLong(n)
 #endif
 
 static PyObject *py_geometric_object() {

src/array_slice.cpp

@@ -461,9 +461,8 @@ cdouble *fields::get_complex_array_slice(const volume &where,
 }
 
 double *fields::get_array_slice(const volume &where, component c,
-                                double *slice, int slice_length)
+                                double *slice)
 {
-  (void) slice_length;
   std::vector<component> components(1);
   components[0]=c;
   return (double *)do_get_array_slice(where, components,
@@ -473,9 +472,8 @@ double *fields::get_array_slice(const volume &where, component c,
 
 double *fields::get_array_slice(const volume &where,
                                 derived_component c,
-                                double *slice, int slice_length)
+                                double *slice)
 {
-  (void) slice_length;
   int nfields;
   component carray[12];
   field_rfunction rfun = derived_component_func(c, gv, nfields, carray);
@@ -486,9 +484,8 @@ double *fields::get_array_slice(const volume &where,
 }
 
 cdouble *fields::get_complex_array_slice(const volume &where, component c,
-                                         cdouble *slice, int slice_length)
+                                         cdouble *slice)
 {
-  (void) slice_length;
   std::vector<component> components(1);
   components[0]=c;
   return (cdouble *)do_get_array_slice(where, components,

src/meep.hpp

@@ -1314,14 +1314,12 @@ class fields {
 
   // alternative entry points for when you have no field
   // function, i.e. you want just a single component or 
-  // derived component.) (The slice_length parameter is only
-  // present to facilitate SWIG-generated python wrappers;
-  // it is ignored by the C code).
-  double *get_array_slice(const volume &where, component c, double *slice=0, int slice_length=0);
-  double *get_array_slice(const volume &where, derived_component c, double *slice=0, int slice_length=0);
+  // derived component.)
+  double *get_array_slice(const volume &where, component c, double *slice=0);
+  double *get_array_slice(const volume &where, derived_component c, double *slice=0);
   std::complex<double> *get_complex_array_slice(const volume &where,
                                                 component c,
-                                                std::complex<double> *slice=0, int slice_length=0);
+                                                std::complex<double> *slice=0);
 
   // master routine for all above entry points
   void *do_get_array_slice(const volume &where,