Add Python support for Prisms (#345)

* Add python Prism wrapper

* Increase tolerance and reuse compare_arrays from tests/mpb.py

* Simplify vector assignment

* Get new scheme data with fixed size blocks

* Revert "Get new scheme data with fixed size blocks"

This reverts commit 0203ab4aa1fc164e9f3b747a84e0b97280bbbd70.

* height is double

* Extend prism out of cell and increase height.

python/geom.py

@@ -341,6 +341,15 @@ def __init__(self, **kwargs):
         super(Ellipsoid, self).__init__(**kwargs)
 
 
+class Prism(GeometricObject):
+
+    def __init__(self, vertices, height, axis=Vector3(z=1), **kwargs):
+        self.vertices = vertices
+        self.height = height
+        self.axis = axis
+        super(Prism, self).__init__(**kwargs)
+
+
 class Matrix(object):
 
     def __init__(self, c1=Vector3(), c2=Vector3(), c3=Vector3()):

python/meep.i

@@ -67,7 +67,7 @@ static int get_attr_int(PyObject *py_obj, int *result, const char *name) {
     }
 
     *result = PyInteger_AsLong(py_attr);
-    Py_DECREF(py_attr);
+    Py_XDECREF(py_attr);
     return 1;
 }
 
@@ -1044,6 +1044,7 @@ void display_geometric_object_info(int indentby, GEOMETRIC_OBJECT o);
         Medium,
         NoisyDrudeSusceptibility,
         NoisyLorentzianSusceptibility,
+        Prism,
         Sphere,
         Susceptibility,
         Vector3,

python/tests/bend_flux.py

@@ -3,6 +3,7 @@
 import unittest
 import numpy as np
 import meep as mp
+from utils import compare_arrays
 
 
 class TestBendFlux(unittest.TestCase):
@@ -15,15 +16,24 @@ def init(self, no_bend=False):
         w = 1
         wvg_ycen = -0.5 * (sy - w - (2 * pad))
         wvg_xcen = 0.5 * (sx - w - (2 * pad))
+        height = 100
 
         if no_bend:
-            geometry = [mp.Block(mp.Vector3(mp.inf, w, mp.inf), center=mp.Vector3(0, wvg_ycen),
-                                 material=mp.Medium(epsilon=12))]
+            no_bend_vertices = [mp.Vector3(-0.5 * sx - 5, wvg_ycen - 0.5 * w),
+                                mp.Vector3(+0.5 * sx + 5, wvg_ycen - 0.5 * w),
+                                mp.Vector3(+0.5 * sx + 5, wvg_ycen + 0.5 * w),
+                                mp.Vector3(-0.5 * sx - 5, wvg_ycen + 0.5 * w)]
+
+            geometry = [mp.Prism(no_bend_vertices, height, material=mp.Medium(epsilon=12))]
         else:
-            geometry = [mp.Block(mp.Vector3(sx - pad, w, mp.inf), center=mp.Vector3(-0.5 * pad, wvg_ycen),
-                                 material=mp.Medium(epsilon=12)),
-                        mp.Block(mp.Vector3(w, sy - pad, mp.inf), center=mp.Vector3(wvg_xcen, 0.5 * pad),
-                                 material=mp.Medium(epsilon=12))]
+            bend_vertices = [mp.Vector3(-0.5 * sx, wvg_ycen - 0.5 * w),
+                             mp.Vector3(wvg_xcen + 0.5 * w, wvg_ycen - 0.5 * w),
+                             mp.Vector3(wvg_xcen + 0.5 * w, 0.5 * sy),
+                             mp.Vector3(wvg_xcen - 0.5 * w, 0.5 * sy),
+                             mp.Vector3(wvg_xcen - 0.5 * w, wvg_ycen + 0.5 * w),
+                             mp.Vector3(-0.5 * sx, wvg_ycen + 0.5 * w)]
+
+            geometry = [mp.Prism(bend_vertices, height, material=mp.Medium(epsilon=12))]
 
         fcen = 0.15
         df = 0.1
@@ -84,12 +94,11 @@ def test_bend_flux(self):
             (0.117171717172, 0.0147547920552, 0.0146151488236),
             (0.118181818182, 0.0194782085272, 0.0192840042241),
             (0.119191919192, 0.0254987474079, 0.0252348211592),
-
         ]
 
         res = list(zip(mp.get_flux_freqs(self.trans), mp.get_fluxes(self.trans), mp.get_fluxes(self.refl)))
 
-        np.testing.assert_allclose(expected, res[:20])
+        compare_arrays(self, np.array(expected), np.array(res[:20]), tol=1e-7)
 
         # Real run
         self.sim = None
@@ -124,7 +133,7 @@ def test_bend_flux(self):
 
         res = list(zip(mp.get_flux_freqs(self.trans), mp.get_fluxes(self.trans), mp.get_fluxes(self.refl)))
 
-        np.testing.assert_allclose(expected, res[:20])
+        compare_arrays(self, np.array(expected), np.array(res[:20]), tol=1e-3)
 
 
 if __name__ == '__main__':

python/tests/geom.py

@@ -245,6 +245,15 @@ def test_contains_point(self):
         self.assertIn(zeros(), e)
 
 
+class TestPrism(unittest.TestCase):
+
+    def test_contains_point(self):
+        vertices = [gm.Vector3(-1, 1), gm.Vector3(1, 1), gm.Vector3(1, -1), gm.Vector3(-1, -1)]
+        p = gm.Prism(vertices, height=1)
+        self.assertIn(zeros(), p)
+        self.assertNotIn(gm.Vector3(2, 2), p)
+
+
 class TestMedium(unittest.TestCase):
 
     def test_D_conductivity(self):

python/tests/mpb.py

@@ -14,6 +14,7 @@
 from scipy.optimize import ridder
 import meep as mp
 from meep import mpb
+from utils import compare_arrays
 
 
 class TestModeSolver(unittest.TestCase):
@@ -178,20 +179,7 @@ def check_fields_against_h5(self, ref_path, field, suffix=''):
             ref_arr[1::3] = ref_y.ravel()
             ref_arr[2::3] = ref_z.ravel()
 
-            self.compare_arrays(ref_arr, field)
-
-    def compare_arrays(self, exp, res, tol=1e-3):
-        exp_1d = exp.ravel()
-        res_1d = res.ravel()
-
-        norm_exp = np.linalg.norm(exp_1d)
-        norm_res = np.linalg.norm(res_1d)
-
-        if norm_exp == 0:
-            self.assertEqual(norm_res, 0)
-        else:
-            diff = np.linalg.norm(res_1d - exp_1d) / norm_exp
-            self.assertLess(diff, tol)
+            compare_arrays(self, ref_arr, field)
 
     def compare_h5_files(self, ref_path, res_path, tol=1e-3):
         with h5py.File(ref_path) as ref:
@@ -200,7 +188,7 @@ def compare_h5_files(self, ref_path, res_path, tol=1e-3):
                     if k == 'description':
                         self.assertEqual(ref[k].value, res[k].value)
                     else:
-                        self.compare_arrays(ref[k].value, res[k].value, tol=tol)
+                        compare_arrays(self, ref[k].value, res[k].value, tol=tol)
 
     def test_update_band_range_data(self):
         brd = []
@@ -407,7 +395,7 @@ def test_compute_field_energy(self):
 
         expected_energy_in_dielectric = 0.6990769686037558
 
-        self.compare_arrays(np.array(expected_energy), np.array(energy))
+        compare_arrays(self, np.array(expected_energy), np.array(energy))
         self.assertAlmostEqual(expected_energy_in_dielectric, energy_in_dielectric, places=3)
 
     def test_compute_group_velocity(self):
@@ -693,7 +681,7 @@ def get_dpwr(ms, band):
         with h5py.File(ref_path, 'r') as f:
             expected = f['data-new'].value
 
-        self.compare_arrays(expected, converted_dpwr[-1])
+        compare_arrays(self, expected, converted_dpwr[-1])
 
     def test_hole_slab(self):
         from mpb_hole_slab import ms
@@ -958,7 +946,7 @@ def test_tri_rods(self):
             expected_re = f['z.r-new'].value
             expected_im = f['z.i-new'].value
             expected = np.vectorize(complex)(expected_re, expected_im)
-            self.compare_arrays(expected, new_efield)
+            compare_arrays(self, expected, new_efield)
 
         ms.run_te()
 
@@ -992,7 +980,7 @@ def test_tri_rods(self):
 
         with h5py.File(ref_path, 'r') as f:
             ref = f['data-new'].value
-            self.compare_arrays(ref, new_eps, tol=1e-3)
+            compare_arrays(self, ref, new_eps, tol=1e-3)
 
     def test_subpixel_averaging(self):
         ms = self.init_solver()
@@ -1065,7 +1053,7 @@ def test_run_te_with_mu_material(self):
         mu = ms.get_mu()
         with h5py.File(data_path, 'r') as f:
             data = f['data'].value
-            self.compare_arrays(data, mu)
+            compare_arrays(self, data, mu)
 
     def test_output_tot_pwr(self):
         ms = self.init_solver()
@@ -1083,7 +1071,7 @@ def test_output_tot_pwr(self):
         with h5py.File(ref_path, 'r') as f:
             expected = f['data'].value
 
-        self.compare_arrays(expected, arr)
+        compare_arrays(self, expected, arr)
 
     def test_get_eigenvectors(self):
         ms = self.init_solver()
@@ -1196,7 +1184,7 @@ def test_epsilon_input_file(self):
         ]
 
         self.check_band_range_data(expected_brd, ms.band_range_data)
-        self.compare_arrays(expected_freqs, ms.all_freqs[-1])
+        compare_arrays(self, expected_freqs, ms.all_freqs[-1])
 
         self.check_gap_list(expected_gap_list, ms.gap_list)
 
@@ -1340,7 +1328,7 @@ def test_multiply_bloch_in_map_data(self):
         md = mpb.MPBData(rectify=True, resolution=32, periods=3)
         result2 = md.convert(efield_no_bloch)
 
-        self.compare_arrays(result1, result2, tol=1e-7)
+        compare_arrays(self, result1, result2, tol=1e-7)
 
     def test_poynting(self):
         ms = self.init_solver()

python/tests/utils.py

@@ -0,0 +1,15 @@
+import numpy as np
+
+
+def compare_arrays(test_instance, exp, res, tol=1e-3):
+    exp_1d = exp.ravel()
+    res_1d = res.ravel()
+
+    norm_exp = np.linalg.norm(exp_1d)
+    norm_res = np.linalg.norm(res_1d)
+
+    if norm_exp == 0:
+        test_instance.assertEqual(norm_res, 0)
+    else:
+        diff = np.linalg.norm(res_1d - exp_1d) / norm_exp
+        test_instance.assertLess(diff, tol)

python/typemap_utils.cpp

@@ -541,6 +541,49 @@ static int pyellipsoid_to_ellipsoid(PyObject *py_ell, geometric_object *e) {
     return 1;
 }
 
+static int pyprism_to_prism(PyObject *py_prism, geometric_object *p) {
+    material_type material;
+    double height;
+    vector3 axis;
+
+    if (!get_attr_material(py_prism, &material) ||
+        !get_attr_dbl(py_prism, &height, "height") ||
+        !get_attr_v3(py_prism, &axis, "axis")) {
+
+        return 0;
+    }
+
+    PyObject *py_vert_list = PyObject_GetAttrString(py_prism, "vertices");
+
+    if (!py_vert_list) {
+        PyErr_PrintEx(0);
+        return 0;
+    }
+
+    if (!PyList_Check(py_vert_list)) {
+        PyErr_SetString(PyExc_TypeError, "Expected Prism.vertices to be a list\n");
+        return 0;
+    }
+
+    int num_vertices = PyList_Size(py_vert_list);
+    vector3 *vertices = new vector3[num_vertices];
+
+    for (Py_ssize_t i = 0; i < num_vertices; ++i) {
+        vector3 v3;
+        if (!pyv3_to_v3(PyList_GetItem(py_vert_list, i), &v3)) {
+            return 0;
+        }
+        vertices[i] = v3;
+    }
+
+    *p = make_prism(material, vertices, num_vertices, height, axis);
+
+    delete [] vertices;
+    Py_DECREF(py_vert_list);
+
+    return 1;
+}
+
 static int py_gobj_to_gobj(PyObject *po, geometric_object *o) {
     int success = 0;
     std::string go_type = py_class_name_as_string(po);
@@ -563,6 +606,9 @@ static int py_gobj_to_gobj(PyObject *po, geometric_object *o) {
     else if (go_type == "Ellipsoid") {
         success = pyellipsoid_to_ellipsoid(po, o);
     }
+    else if (go_type == "Prism") {
+        success = pyprism_to_prism(po, o);
+    }
     else {
         PyErr_Format(PyExc_TypeError, "Error: %s is not a valid GeometricObject type\n", go_type.c_str());
         return 0;