Manipulating raw eigenvectors (#221)

* Getting eigenvectors

* get_eigenvectors working

* set_eigenvectors working

* load and save eigenvectors working

* Call fix_hfield_phase in get_eigenvectors test

libpympb/pympb.cpp

@@ -1345,13 +1345,6 @@ void mode_solver::get_epsilon_tensor(int c1, int c2, int imag, int inv) {
   }
 }
 
-void mode_solver::load_eigenvectors(char *filename) {
-  meep::master_printf("Loading eigenvectors from \"%s\"...\n", filename);
-  // TODO: Write in python
-  // evectmatrixio_readall_raw(filename, H);
-  curfield_reset();
-}
-
 std::vector<mpb_real> mode_solver::get_freqs() {
   return freqs;
 }
@@ -2074,6 +2067,36 @@ void mode_solver::get_lattice(double data[3][3]) {
   matrix3x3_to_arr(data, Rm);
 }
 
+std::vector<int> mode_solver::get_eigenvectors_slice_dims(int num_bands) {
+  std::vector<int> res(3);
+  res[0] = H.localN;
+  res[1] = H.c;
+  res[2] = num_bands;
+
+  return res;
+}
+
+void mode_solver::get_eigenvectors(int p_start, int p, std::complex<mpb_real> *cdata, int size) {
+
+  for (int i = 0, j = p_start; i < size; i += p, j += H.p) {
+    for (int k = 0; k < p; ++k) {
+      cdata[i + k] = std::complex<mpb_real>(H.data[j + k].re, H.data[j + k].im);
+    }
+  }
+}
+
+void mode_solver::set_eigenvectors(int b_start, std::complex<mpb_real> *cdata, int size) {
+  int columns = size / H.n;
+
+  for (int i = 0, j = b_start; i < size; i += columns, j += H.p) {
+    for (int k = 0; k < columns; ++k) {
+      H.data[j + k].re = cdata[i + k].real();
+      H.data[j + k].im = cdata[i + k].imag();
+    }
+  }
+  curfield_reset();
+}
+
 double mode_solver::get_eigensolver_flops() {
   return eigensolver_flops;
 }

libpympb/pympb.hpp

@@ -105,7 +105,6 @@ struct mode_solver {
   bool has_mu();
   bool material_has_mu(void *mt);
   void curfield_reset();
-  void load_eigenvectors(char *filename);
 
   size_t get_field_size();
 
@@ -124,6 +123,9 @@ struct mode_solver {
   void set_curfield_cmplx(std::complex<mpb_real> *cdata, int size);
 
   void get_lattice(double data[3][3]);
+  void get_eigenvectors(int p_start, int p, std::complex<mpb_real> *cdata, int size);
+  std::vector<int> get_eigenvectors_slice_dims(int num_bands);
+  void set_eigenvectors(int b_start, std::complex<mpb_real> *cdata, int size);
 
   std::vector<mpb_real> compute_field_energy();
   double compute_energy_in_objects(geometric_object_list objects);

python/solver.py

@@ -222,6 +222,26 @@ def get_tot_pwr(self, which_band):
 
         return tot_pwr
 
+    def get_eigenvectors(self, first_band, num_bands):
+        dims = self.mode_solver.get_eigenvectors_slice_dims(num_bands)
+        ev = np.zeros(np.prod(dims), dtype=np.complex128)
+        self.mode_solver.get_eigenvectors(first_band - 1, num_bands, ev)
+        return ev.reshape(dims)
+
+    def set_eigenvectors(self, ev, first_band):
+        self.mode_solver.set_eigenvectors(first_band - 1, ev.flatten())
+
+    def save_eigenvectors(self, filename):
+        with h5py.File(filename, 'w') as f:
+            ev = self.get_eigenvectors(1, self.num_bands)
+            f['rawdata'] = ev
+
+    def load_eigenvectors(self, filename):
+        with h5py.File(filename, 'r') as f:
+            ev = f['rawdata'].value
+            self.set_eigenvectors(ev, 1)
+            self.mode_solver.curfield_reset()
+
     # The band-range-data is a list of tuples, each consisting of a (min, k-point)
     # tuple and a (max, k-point) tuple, with each min/max pair describing the
     # frequency range of a band and the k-points where it achieves its minimum/maximum.
@@ -590,7 +610,7 @@ def run_parity(self, p, reset_fields, *band_functions):
         )
 
         if isinstance(reset_fields, basestring):
-            self.mode_solver.load_eigenvectors(reset_fields)
+            self.load_eigenvectors(reset_fields)
 
         print("{} k-points".format(len(self.k_points)))
 

python/tests/mpb.py

@@ -953,7 +953,6 @@ def test_tri_rods(self):
 
     def test_subpixel_averaging(self):
         ms = self.init_solver()
-        ms.tolerance = 1e-12
         ms.run_te()
 
         expected_brd = [
@@ -1026,5 +1025,57 @@ def test_output_tot_pwr(self):
 
         self.compare_h5_files(ref_path, res_path)
 
+    def test_get_eigenvectors(self):
+        ms = self.init_solver()
+        ms.run_te(mpb.fix_hfield_phase)
+
+        def compare_eigenvectors(ref_fn, start, cols):
+            with h5py.File(os.path.join(self.data_dir, ref_fn), 'r') as f:
+                expected = f['rawdata'].value
+                # Reshape the last dimension of 2 reals into one complex
+                expected = np.vectorize(complex)(expected[..., 0], expected[..., 1])
+                ev = ms.get_eigenvectors(start, cols)
+                np.testing.assert_allclose(expected, ev, rtol=1e-3)
+
+        # Get all columns
+        compare_eigenvectors('tutorial-te-eigenvectors.h5', 1, 8)
+        # Get last column
+        compare_eigenvectors('tutorial-te-eigenvectors-8-1.h5', 8, 1)
+        # Get columns 3,4, and 5
+        compare_eigenvectors('tutorial-te-eigenvectors-3-3.h5', 3, 3)
+
+    def test_set_eigenvectors(self):
+        ms = self.init_solver()
+
+        def set_H_to_zero_and_check(start, num_bands):
+            ev = ms.get_eigenvectors(start, num_bands)
+            self.assertNotEqual(np.count_nonzero(ev), 0)
+            zeros = np.zeros(ev.shape, dtype=np.complex128)
+            ms.set_eigenvectors(zeros, start)
+            new_ev = ms.get_eigenvectors(start, num_bands)
+            self.assertEqual(np.count_nonzero(new_ev), 0)
+
+        ms.run_te()
+        set_H_to_zero_and_check(8, 1)
+        ms.run_te()
+        set_H_to_zero_and_check(1, 8)
+        ms.run_te()
+        set_H_to_zero_and_check(3, 3)
+
+    def test_load_and_save_eigenvectors(self):
+        ms = self.init_solver()
+        ms.run_te()
+        fn = self.filename_prefix + '.h5'
+
+        ev = ms.get_eigenvectors(8, 1)
+        zeros = np.zeros(ev.shape, dtype=np.complex128)
+        ms.set_eigenvectors(zeros, 8)
+        ms.save_eigenvectors(fn)
+
+        ms.run_te()
+        ms.load_eigenvectors(fn)
+        new_ev = ms.get_eigenvectors(8, 1)
+        self.assertEqual(np.count_nonzero(new_ev), 0)
+
 if __name__ == '__main__':
     unittest.main()