Add dft_energy Python support (NanoComp#747)

* Add dft_energy Python support

* Add new functions to meep package

* Add explanation of DftObj class

* Add test

python/Makefile.am

@@ -77,6 +77,7 @@ TESTS =                                   \
     $(TEST_DIR)/cavity_farfield.py        \
     $(TEST_DIR)/chunks.py                 \
     $(TEST_DIR)/cyl_ellipsoid.py          \
+    $(TEST_DIR)/dft_energy.py             \
     $(TEST_DIR)/dft_fields.py             \
     $(TEST_DIR)/field_functions.py        \
     $(TEST_DIR)/force.py                  \
@@ -91,7 +92,7 @@ TESTS =                                   \
     $(MODE_COEFFS_TEST)                   \
     $(MODE_DECOMPOSITION_TEST)            \
     $(TEST_DIR)/multilevel_atom.py        \
-    $(TEST_DIR)/oblique_source.py        \
+    $(TEST_DIR)/oblique_source.py         \
     $(TEST_DIR)/physical.py               \
     $(TEST_DIR)/pw_source.py              \
     $(TEST_DIR)/refl_angular.py           \

python/meep.i

@@ -1169,6 +1169,12 @@ meep::volume_list *make_volume_list(const meep::volume &v, int c,
     }
 }
 
+%apply double *flux {
+    double *electric,
+    double *magnetic,
+    double *total
+};
+
 // Tells Python to take ownership of the h5file* this function returns so that
 // it gets garbage collected and the file gets closed.
 %newobject meep::fields::open_h5file;
@@ -1378,6 +1384,7 @@ PyObject *_get_array_slice_dimensions(meep::fields *f, const meep::volume &where
     from .simulation import (
         Absorber,
         Ldos,
+        EnergyRegion,
         FluxRegion,
         ForceRegion,
         Harminv,
@@ -1405,12 +1412,16 @@ PyObject *_get_array_slice_dimensions(meep::fields *f, const meep::volume &where
         during_sources,
         GDSII_vol,
         get_center_and_size,
+        get_eigenmode_freqs,
+        get_electric_energy,
+        get_energy_freqs,
         get_flux_freqs,
         get_fluxes,
-        get_eigenmode_freqs,
         get_force_freqs,
         get_forces,
+        get_magnetic_energy,
         get_near2far_freqs,
+        get_total_energy,
         in_point,
         in_volume,
         interpolate,
@@ -1457,6 +1468,7 @@ PyObject *_get_array_slice_dimensions(meep::fields *f, const meep::volume &where
         output_sfield_r,
         output_sfield_p,
         py_v3_to_vec,
+        scale_energy_fields,
         scale_flux_fields,
         scale_force_fields,
         scale_near2far_fields,

python/simulation.py

@@ -175,6 +175,7 @@ def __init__(self, center=None, size=Vector3(), direction=mp.AUTOMATIC, weight=1
 ModeRegion = FluxRegion
 Near2FarRegion = FluxRegion
 ForceRegion = FluxRegion
+EnergyRegion = FluxRegion
 
 
 class FieldsRegion(object):
@@ -191,7 +192,29 @@ def __init__(self, where=None, center=None, size=None):
 
 
 class DftObj(object):
-
+    """Wrapper around dft objects that allows delayed initialization of the structure.
+
+    When splitting the structure into chunks for parallel simulations, we want to
+    know all of the details of the simulation in order to ensure that each processor
+    gets a similar amount of work. The problem with DFTs is that the 'add_flux' style
+    methods immediately initialize the structure and fields. So, if the user adds
+    multiple DFT objects to the simulation, the load balancing code only knows about
+    the first one and can't split the work up nicely. To circumvent this, we delay
+    the execution of the 'add_flux' methods as late as possible. When 'add_flux' (or
+    add_near2far, etc.) is called, we
+
+    1. Create an instance of the appropriate subclass of DftObj (DftForce, DftFlux,
+    etc.). Set its args property to the list of arguments passed to add_flux, and
+    set its func property to the 'real' add_flux, which is prefixed by an underscore.
+
+    2. Add this DftObj to the list Simulation.dft_objects. When we actually run the
+    simulation, we call Simulation._evaluate_dft_objects, which calls dft.func(*args)
+    for each dft in the list.
+
+    If the user tries to access a property or call a function on the DftObj before
+    Simulation._evaluate_dft_objects is called, then we initialize the C++ object
+    through swigobj_attr and return the property they requested.
+    """
     def __init__(self, func, args):
         self.func = func
         self.args = args
@@ -324,6 +347,28 @@ def mu(self):
     def flux(self, direction, where, resolution):
         return self.swigobj_attr('flux')(direction, where.swigobj, resolution)
 
+
+class DftEnergy(DftObj):
+
+    def __init__(self, func, args):
+        super(DftEnergy, self).__init__(func, args)
+        self.nfreqs = args[2]
+        self.regions = args[3]
+        self.num_components = 12
+
+    @property
+    def electric(self):
+        return self.swigobj_attr('electric')
+
+    @property
+    def magnetic(self):
+        return self.swigobj_attr('magnetic')
+
+    @property
+    def total(self):
+        return self.swigobj_attr('total')
+
+
 class DftFields(DftObj):
 
     def __init__(self, func, args):
@@ -1336,6 +1381,44 @@ def _add_near2far(self, fcen, df, nfreq, near2fars):
             self.init_sim()
         return self._add_fluxish_stuff(self.fields.add_dft_near2far, fcen, df, nfreq, near2fars)
 
+    def add_energy(self, fcen, df, nfreq, *energys):
+        en = DftEnergy(self._add_energy, [fcen, df, nfreq, energys])
+        self.dft_objects.append(en)
+        return en
+
+    def _add_energy(self, fcen, df, nfreq, energys):
+        if self.fields is None:
+            self.init_sim()
+        return self._add_fluxish_stuff(self.fields.add_dft_energy, fcen, df, nfreq, energys)
+
+    def _display_energy(self, name, func, energys):
+        if energys:
+            freqs = get_energy_freqs(energys[0])
+            display_csv(self, "{}-energy".format(name), zip(freqs, *[func(f) for f in energys]))
+
+    def display_electric_energy(self, *energys):
+        self._display_energy('electric', get_electric_energy, energys)
+
+    def display_magnetic_energy(self, *energys):
+        self._display_energy('magnetic', get_magnetic_energy, energys)
+
+    def display_total_energy(self, *energys):
+        self._display_energy('total', get_total_energy, energys)
+
+    def load_energy(self, fname, energy):
+        if self.fields is None:
+            self.init_sim()
+        energy.load_hdf5(self.fields, fname, '', self.get_filename_prefix())
+
+    def save_energy(self, fname, energy):
+        if self.fields is None:
+            self.init_sim()
+        energy.save_hdf5(self.fields, fname, '', self.get_filename_prefix())
+
+    def load_minus_energy(self, fname, energy):
+        self.load_energy(fname, energy)
+        energy.scale_dfts(-1.0)
+
     def get_farfield(self, f, v):
         return mp._get_farfield(f.swigobj, py_v3_to_vec(self.dimensions, v, is_cylindrical=self.is_cylindrical))
 
@@ -2620,6 +2703,26 @@ def get_near2far_freqs(f):
     return np.linspace(f.freq_min, f.freq_min + f.dfreq * f.Nfreq, num=f.Nfreq, endpoint=False).tolist()
 
 
+def scale_energy_fields(s, ef):
+    df.scale_dfts(s)
+
+
+def get_energy_freqs(f):
+    return np.linspace(f.freq_min, f.freq_min + f.dfreq * f.Nfreq, num=f.Nfreq, endpoint=False).tolist()
+
+
+def get_electric_energy(f):
+    return f.electric()
+
+
+def get_magnetic_energy(f):
+    return f.magnetic()
+
+
+def get_total_energy(f):
+    return f.total()
+
+
 def interpolate(n, nums):
     res = []
     if isinstance(nums[0], mp.Vector3):

python/tests/dft_energy.py

@@ -0,0 +1,44 @@
+from __future__ import division
+
+import unittest
+import meep as mp
+
+# compute group velocity of a waveguide mode using two different methods
+# (1) ratio of Poynting flux to energy density
+# (2) via MPB from get-eigenmode-coefficients
+
+
+class TestDftEnergy(unittest.TestCase):
+
+    def test_dft_energy(self):
+        resolution = 20
+        cell = mp.Vector3(10, 5)
+        geom = [mp.Block(size=mp.Vector3(mp.inf, 1, mp.inf), material=mp.Medium(epsilon=12))]
+        pml = [mp.PML(1)]
+        fsrc  = 0.15
+        sources = [mp.EigenModeSource(src=mp.GaussianSource(frequency=fsrc, fwidth=0.2*fsrc),
+                                      center=mp.Vector3(-3), size=mp.Vector3(y=5),
+                                      eig_band=1, eig_parity=mp.ODD_Z+mp.EVEN_Y,
+                                      eig_match_freq=True)]
+
+        sim = mp.Simulation(resolution=resolution, cell_size=cell, geometry=geom,
+                            boundary_layers=pml, sources=sources)
+
+        flux = sim.add_flux(fsrc, 0, 1, mp.FluxRegion(center=mp.Vector3(3), size=mp.Vector3(y=5)))
+        energy = sim.add_energy(fsrc, 0, 1, mp.EnergyRegion(center=mp.Vector3(3), size=mp.Vector3(y=5)))
+        sim.run(until_after_sources=100)
+
+        res = sim.get_eigenmode_coefficients(flux, [1], eig_parity=mp.ODD_Z+mp.EVEN_Y)
+        mode_vg = res.vgrp[0]
+        poynting_flux  = mp.get_fluxes(flux)[0]
+        e_energy = mp.get_electric_energy(energy)[0]
+        ratio_vg = (0.5 * poynting_flux) / e_energy
+        m_energy = mp.get_magnetic_energy(energy)[0]
+        t_energy = mp.get_total_energy(energy)[0]
+
+        self.assertAlmostEqual(m_energy + e_energy, t_energy)
+        self.assertAlmostEqual(ratio_vg, mode_vg, places=3)
+
+
+if __name__ == '__main__':
+    unittest.main()