support for single-precision floating point for fields array functions (

NanoComp#1833)

* switch dft-related functions to using realnum from double

* more fixes

* more type conversions from double to realnum

* adjust check tolerance of tests/integrate.cpp based on floating-point precision

* more fixes

* rebase from master from fix merge conflicts

* slight adjustment to tolerances in unit tests and update docs

* remove type check in test_adjoint_solver.py

* revert return types of integration functions to double

* revert return type of process_dft_component to double

* cleanup

doc/docs/Python_User_Interface.md

@@ -3602,7 +3602,9 @@ current simulation time.
 + `cmplx`: `boolean`; if `True`, return complex-valued data otherwise return
   real-valued data (default).
 
-+ `arr`: optional field to pass a pre-allocated NumPy array of the correct size,
++ `arr`: optional parameter to pass a pre-allocated NumPy array of the correct size and
+  type (either `numpy.float32` or `numpy.float64` depending on the [floating-point precision
+  of the fields and materials](Build_From_Source.md#floating-point-precision-of-the-fields-and-materials-arrays))
   which will be overwritten with the field/material data instead of allocating a
   new array.  Normally, this will be the array returned from a previous call to
   `get_array` for a similar slice, allowing one to re-use `arr` (e.g., when
@@ -3655,7 +3657,8 @@ def get_dft_array(self, dft_obj, component, num_freq):
 
 <div class="method_docstring" markdown="1">
 
-Returns the Fourier-transformed fields as a NumPy array.
+Returns the Fourier-transformed fields as a NumPy array. The type is either `numpy.complex64`
+or `numpy.complex128` depending on the [floating-point precision of the fields](Build_From_Source.md#floating-point-precision-of-the-fields-and-materials-arrays).
 
 **Parameters:**
 

python/adjoint/optimization_problem.py

@@ -248,7 +248,7 @@ def adjoint_run(self):
 
             # Store adjoint fields for each design set of design variables
             self.D_a.append(utils.gather_design_region_fields(self.sim,self.design_region_monitors,self.frequencies))
-        
+
         # reset the m number
         if utils._check_if_cylindrical(self.sim):
             self.sim.m = -self.sim.m

python/adjoint/utils.py

@@ -37,14 +37,14 @@ def __init__(
 
     def update_design_parameters(self, design_parameters):
         self.design_parameters.update_weights(design_parameters)
-    
+
     def update_beta(self,beta):
         self.design_parameters.beta=beta
 
     def get_gradient(self, sim, fields_a, fields_f, frequencies, finite_difference_step):
         num_freqs = onp.array(frequencies).size
         shapes = []
-        '''We have the option to linear scale the gradients up front
+        '''We have the option to linearly scale the gradients up front
         using the scalegrad parameter (leftover from MPB API). Not
         currently needed for any existing feature (but available for
         future use)'''
@@ -54,16 +54,18 @@ def get_gradient(self, sim, fields_a, fields_f, frequencies, finite_difference_s
             returns a singleton element for the forward and adjoint fields.
             This only occurs when we are in 2D and only working with a particular
             polarization (as the other fields are never stored). For example, the
-            2D in-plane polarization consists of a single scalar Ez field 
+            2D in-plane polarization consists of a single scalar Ez field
             (technically, meep doesn't store anything for these cases, but get_dft_array
             still returns 0).
-            
-            Our get_gradient algorithm, however, requires we pass an array of 
+
+            Our get_gradient algorithm, however, requires we pass an array of
             zeros with the proper shape as the design_region.'''
             spatial_shape = sim.get_array_slice_dimensions(component, vol=self.volume)[0]
             if (fields_a[component_idx][0,...].size == 1):
-                fields_a[component_idx] = onp.zeros(onp.insert(spatial_shape,0,num_freqs))
-                fields_f[component_idx] = onp.zeros(onp.insert(spatial_shape,0,num_freqs))
+                fields_a[component_idx] = onp.zeros(onp.insert(spatial_shape,0,num_freqs),
+                                                    dtype=onp.float32 if mp.is_single_precision() else onp.float64)
+                fields_f[component_idx] = onp.zeros(onp.insert(spatial_shape,0,num_freqs),
+                                                    dtype=onp.float32 if mp.is_single_precision() else onp.float64)
             if _check_if_cylindrical(sim):
                 '''For some reason, get_dft_array returns the field
                 components in a different order than the convention used
@@ -78,15 +80,16 @@ def get_gradient(self, sim, fields_a, fields_f, frequencies, finite_difference_s
         shapes = onp.asarray(shapes).flatten(order='C')
         fields_a = onp.concatenate(fields_a)
         fields_f = onp.concatenate(fields_f)
-        
+
         grad = onp.zeros((num_freqs, self.num_design_params))  # preallocate
         geom_list = sim.geometry
         f = sim.fields
         vol = sim._fit_volume_to_simulation(self.volume)
-        
+
         # compute the gradient
         mp._get_gradient(grad,scalegrad,fields_a,fields_f,
-        sim.gv,vol.swigobj,onp.array(frequencies),sim.geps,shapes,finite_difference_step)
+                         sim.gv,vol.swigobj,onp.array(frequencies),
+                         sim.geps,shapes,finite_difference_step)
         return onp.squeeze(grad).T
 
 def _check_if_cylindrical(sim):
@@ -203,7 +206,7 @@ def gather_design_region_fields(
       fairly awkward to inspect directly.  Their primary use case is supporting
       gradient calculations.
     """
-    fwd_fields = []
+    design_region_fields = []
     for monitor in design_region_monitors:
         fields_by_component = []
         for component in _compute_components(simulation):
@@ -212,8 +215,8 @@ def gather_design_region_fields(
                 fields = simulation.get_dft_array(monitor, component, freq_idx)
                 fields_by_freq.append(_make_at_least_nd(fields))
             fields_by_component.append(onp.stack(fields_by_freq))
-        fwd_fields.append(fields_by_component)
-    return fwd_fields
+        design_region_fields.append(fields_by_component)
+    return design_region_fields
 
 
 def validate_and_update_design(
@@ -258,7 +261,7 @@ def create_adjoint_sources(
         monitors: Iterable[ObjectiveQuantity],
         monitor_values_grad: onp.ndarray) -> List[mp.Source]:
     monitor_values_grad = onp.asarray(monitor_values_grad,
-                                      dtype=onp.complex128)
+                                      dtype=onp.complex64 if mp.is_single_precision() else onp.complex128)
     if not onp.any(monitor_values_grad):
         raise RuntimeError(
             'The gradient of all monitor values is zero, which '
@@ -273,7 +276,7 @@ def create_adjoint_sources(
     for monitor_idx, monitor in enumerate(monitors):
         # `dj` for each monitor will have a shape of (num frequencies,)
         dj = onp.asarray(monitor_values_grad[monitor_idx],
-                         dtype=onp.complex128)
+                         dtype=onp.complex64 if mp.is_single_precision() else onp.complex128)
         if onp.any(dj):
             adjoint_sources += monitor.place_adjoint_source(dj)
     assert adjoint_sources

python/adjoint/wrapper.py

@@ -229,7 +229,9 @@ def _simulate_fwd(design_variables):
         def _simulate_rev(res, monitor_values_grad):
             """Runs adjoint simulation, returning VJP of design wrt monitor values."""
             fwd_fields = jax.tree_map(
-                lambda x: onp.asarray(x, dtype=onp.complex128), res[0])
+                lambda x: onp.asarray(x,
+                                      dtype=onp.complex64 if mp.is_single_precision() else onp.complex128),
+                res[0])
             design_variable_shapes = res[1]
             adj_fields = self._run_adjoint_simulation(monitor_values_grad)
             vjps = self._calculate_vjps(fwd_fields, adj_fields,

python/meep.i

@@ -150,7 +150,7 @@ static std::complex<double> py_amp_func_wrap(const meep::vec &v) {
     return ret;
 }
 
-static std::complex<double> py_field_func_wrap(const std::complex<double> *fields,
+static std::complex<double> py_field_func_wrap(const std::complex<meep::realnum> *fields,
                                                const meep::vec &loc,
                                                void *data_) {
     SWIG_PYTHON_THREAD_SCOPED_BLOCK;
@@ -441,25 +441,25 @@ PyObject *_get_dft_array(meep::fields *f, dft_type dft, meep::component c, int n
     // Return value: New reference
     int rank;
     size_t dims[3];
-    std::complex<double> *dft_arr = f->get_dft_array(dft, c, num_freq, &rank, dims);
+    std::complex<meep::realnum> *dft_arr = f->get_dft_array(dft, c, num_freq, &rank, dims);
 
-    if (dft_arr==NULL){ // this can happen e.g. if component c vanishes by symmetry
-     std::complex<double> d[1] = {std::complex<double>(0,0)};
-     return PyArray_SimpleNewFromData(0, 0, NPY_CDOUBLE, d);
+    if (dft_arr == NULL) { // this can happen e.g. if component c vanishes by symmetry
+      std::complex<double> d[1] = {std::complex<double>(0,0)};
+      return PyArray_SimpleNewFromData(0, 0, sizeof(meep::realnum) == sizeof(float) ? NPY_CFLOAT : NPY_CDOUBLE, d);
     }
 
     if (rank == 0) // singleton results
-     return PyArray_SimpleNewFromData(0, 0, NPY_CDOUBLE, dft_arr);
+      return PyArray_SimpleNewFromData(0, 0, sizeof(meep::realnum) == sizeof(float) ? NPY_CFLOAT : NPY_CDOUBLE, dft_arr);
 
     size_t length = 1;
     npy_intp *arr_dims = new npy_intp[rank];
     for (int i = 0; i < rank; ++i) {
-        arr_dims[i] = dims[i];       // implicit size_t -> int cast, presumed safe for individual array dimensions
-        length *= dims[i];
+         arr_dims[i] = dims[i];       // implicit size_t -> int cast, presumed safe for individual array dimensions
+         length *= dims[i];
     }
 
-    PyObject *py_arr = PyArray_SimpleNew(rank, arr_dims, NPY_CDOUBLE);
-    memcpy(PyArray_DATA((PyArrayObject*)py_arr), dft_arr, sizeof(std::complex<double>) * length);
+    PyObject *py_arr = PyArray_SimpleNew(rank, arr_dims, sizeof(meep::realnum) == sizeof(float) ? NPY_CFLOAT : NPY_CDOUBLE);
+    memcpy(PyArray_DATA((PyArrayObject*)py_arr), dft_arr, sizeof(std::complex<meep::realnum>) * length);
     delete[] dft_arr;
     if (arr_dims) delete[] arr_dims;
 
@@ -844,8 +844,8 @@ meep::volume_list *make_volume_list(const meep::volume &v, int c,
 
 %inline %{
 void _get_gradient(PyObject *grad, double scalegrad, PyObject *fields_a, PyObject *fields_f, 
-    meep::grid_volume *grid_volume, meep::volume *where, PyObject *frequencies, 
-    meep_geom::geom_epsilon *geps, PyObject *fields_shapes, double fd_step) {
+                   meep::grid_volume *grid_volume, meep::volume *where, PyObject *frequencies,
+                   meep_geom::geom_epsilon *geps, PyObject *fields_shapes, double fd_step) {
     // clean the gradient array
     PyArrayObject *pao_grad = (PyArrayObject *)grad;
     if (!PyArray_Check(pao_grad)) meep::abort("grad parameter must be numpy array.");
@@ -859,14 +859,14 @@ void _get_gradient(PyObject *grad, double scalegrad, PyObject *fields_a, PyObjec
     if (!PyArray_Check(pao_fields_a)) meep::abort("adjoint fields parameter must be numpy array.");
     if (!PyArray_ISCARRAY(pao_fields_a)) meep::abort("Numpy adjoint fields array must be C-style contiguous.");
     if (PyArray_NDIM(pao_fields_a) !=1) {meep::abort("Numpy adjoint fields array must have 1 dimension.");}
-    std::complex<double> *fields_a_c = (std::complex<double> *)PyArray_DATA(pao_fields_a);
+    std::complex<meep::realnum> *fields_a_c = (std::complex<meep::realnum> *)PyArray_DATA(pao_fields_a);
 
     // clean the forward fields array
     PyArrayObject *pao_fields_f = (PyArrayObject *)fields_f;
     if (!PyArray_Check(pao_fields_f)) meep::abort("forward fields parameter must be numpy array.");
     if (!PyArray_ISCARRAY(pao_fields_f)) meep::abort("Numpy forward fields array must be C-style contiguous.");
     if (PyArray_NDIM(pao_fields_f) !=1) {meep::abort("Numpy forward fields array must have 1 dimension.");}
-    std::complex<double> *fields_f_c = (std::complex<double> *)PyArray_DATA(pao_fields_f);
+    std::complex<meep::realnum> *fields_f_c = (std::complex<meep::realnum> *)PyArray_DATA(pao_fields_f);
 
     // clean shapes array
     PyArrayObject *pao_fields_shapes = (PyArrayObject *)fields_shapes;
@@ -1025,20 +1025,20 @@ void _get_gradient(PyObject *grad, double scalegrad, PyObject *fields_a, PyObjec
     $1 = (size_t *)array_data($input);
 }
 
-%typecheck(SWIG_TYPECHECK_POINTER, fragment="NumPy_Fragments") double* slice {
+%typecheck(SWIG_TYPECHECK_POINTER, fragment="NumPy_Fragments") meep::realnum* slice {
     $1 = is_array($input);
 }
 
-%typemap(in, fragment="NumPy_Macros") double* slice {
-    $1 = (double *)array_data($input);
+%typemap(in, fragment="NumPy_Macros") meep::realnum* slice {
+    $1 = (meep::realnum *)array_data($input);
 }
 
-%typecheck(SWIG_TYPECHECK_POINTER, fragment="NumPy_Fragments") std::complex<double>* slice {
+%typecheck(SWIG_TYPECHECK_POINTER, fragment="NumPy_Fragments") std::complex<meep::realnum>* slice {
     $1 = is_array($input);
 }
 
-%typemap(in) std::complex<double>* slice {
-    $1 = (std::complex<double> *)array_data($input);
+%typemap(in) std::complex<meep::realnum>* slice {
+    $1 = (std::complex<meep::realnum> *)array_data($input);
 }
 
 %typecheck(SWIG_TYPECHECK_POINTER) meep::component {

python/simulation.py

@@ -3337,7 +3337,9 @@ def get_array(self, component=None, vol=None, center=None, size=None, cmplx=None
         + `cmplx`: `boolean`; if `True`, return complex-valued data otherwise return
           real-valued data (default).
 
-        + `arr`: optional field to pass a pre-allocated NumPy array of the correct size,
+        + `arr`: optional parameter to pass a pre-allocated NumPy array of the correct size and
+          type (either `numpy.float32` or `numpy.float64` depending on the [floating-point precision
+          of the fields and materials](Build_From_Source.md#floating-point-precision-of-the-fields-and-materials-arrays))
           which will be overwritten with the field/material data instead of allocating a
           new array.  Normally, this will be the array returned from a previous call to
           `get_array` for a similar slice, allowing one to re-use `arr` (e.g., when
@@ -3407,7 +3409,10 @@ def get_array(self, component=None, vol=None, center=None, size=None, cmplx=None
             arr = np.require(arr, requirements=['C', 'W'])
 
         else:
-            arr = np.zeros(dims, dtype=np.complex128 if cmplx else np.float64)
+            if mp.is_single_precision():
+                arr = np.zeros(dims, dtype=np.complex64 if cmplx else np.float32)
+            else:
+                arr = np.zeros(dims, dtype=np.complex128 if cmplx else np.float64)
 
         if np.iscomplexobj(arr):
             self.fields.get_complex_array_slice(v, component, arr, frequency, snap)
@@ -3418,7 +3423,8 @@ def get_array(self, component=None, vol=None, center=None, size=None, cmplx=None
 
     def get_dft_array(self, dft_obj, component, num_freq):
         """
-        Returns the Fourier-transformed fields as a NumPy array.
+        Returns the Fourier-transformed fields as a NumPy array. The type is either `numpy.complex64`
+        or `numpy.complex128` depending on the [floating-point precision of the fields](Build_From_Source.md#floating-point-precision-of-the-fields-and-materials-arrays).
 
         **Parameters:**
 
@@ -3464,7 +3470,7 @@ def get_source(self, component, vol=None, center=None, size=None):
         dim_sizes = np.zeros(3, dtype=np.uintp)
         mp._get_array_slice_dimensions(self.fields, v, dim_sizes, True, False)
         dims = [s for s in dim_sizes if s != 0]
-        arr = np.zeros(dims, dtype=np.complex128)
+        arr = np.zeros(dims, dtype=np.complex64 if mp.is_single_precision() else np.complex128)
         self.fields.get_source_slice(v, component, arr)
         return arr
 

python/tests/test_adjoint_jax.py

@@ -175,7 +175,8 @@ def test_design_region_monitor_helpers(self):
         for value in design_region_fields[0]:
             self.assertIsInstance(value, onp.ndarray)
             self.assertEqual(value.ndim, 4)  # dims: freq, x, y, pad
-            self.assertEqual(value.dtype, onp.complex128)
+            self.assertEqual(value.dtype,
+                             onp.complex64 if mp.is_single_precision() else onp.complex128)
 
 
 class WrapperTest(ApproxComparisonTestCase):

python/tests/test_adjoint_solver.py

@@ -15,7 +15,8 @@
 resolution = 30
 
 silicon = mp.Medium(epsilon=12)
-sapphire = mp.Medium(epsilon_diag=(10.225,10.225,9.95),epsilon_offdiag=(-0.825,-0.55*np.sqrt(3/2),0.55*np.sqrt(3/2)))
+sapphire = mp.Medium(epsilon_diag=(10.225,10.225,9.95),
+                     epsilon_offdiag=(-0.825,-0.55*np.sqrt(3/2),0.55*np.sqrt(3/2)))
 
 sxy = 5.0
 cell_size = mp.Vector3(sxy,sxy,0)
@@ -469,7 +470,7 @@ def test_complex_fields(self):
 
             ## compare objective results
             print("Ez2 -- adjoint solver: {}, traditional simulation: {}".format(adjsol_obj,Ez2_unperturbed))
-            self.assertClose(adjsol_obj,Ez2_unperturbed,epsilon=1e-8)
+            self.assertClose(adjsol_obj,Ez2_unperturbed,epsilon=1e-6)
 
             ## compute perturbed |Ez|^2
             Ez2_perturbed = forward_simulation_complex_fields(p+dp, frequencies)
@@ -533,7 +534,7 @@ def test_offdiagonal(self):
             adj_scale = (dp[None,:]@adjsol_grad).flatten()
             fd_grad = S12_perturbed-S12_unperturbed
             print("Directional derivative -- adjoint solver: {}, FD: {}".format(adj_scale,fd_grad))
-            tol = 0.04
+            tol = 0.05 if mp.is_single_precision() else 0.04
             self.assertClose(adj_scale,fd_grad,epsilon=tol)
 if __name__ == '__main__':
     unittest.main()

python/tests/test_cavity_arrayslice.py

@@ -70,15 +70,15 @@ def test_2d_slice(self):
 
     def test_1d_slice_user_array(self):
         self.sim.run(until_after_sources=0)
-        arr = np.zeros(126, dtype=np.float64)
+        arr = np.zeros(126, dtype=np.float32 if mp.is_single_precision() else np.float64)
         vol = mp.Volume(center=self.center_1d, size=self.size_1d)
         self.sim.get_array(mp.Hz, vol, arr=arr)
         tol = 1e-5 if mp.is_single_precision() else 1e-8
         self.assertClose(self.expected_1d, arr, epsilon=tol)
 
     def test_2d_slice_user_array(self):
         self.sim.run(until_after_sources=0)
-        arr = np.zeros((126, 38), dtype=np.float64)
+        arr = np.zeros((126, 38), dtype=np.float32 if mp.is_single_precision() else np.float64)
         vol = mp.Volume(center=self.center_2d, size=self.size_2d)
         self.sim.get_array(mp.Hz, vol, arr=arr)
         tol = 1e-5 if mp.is_single_precision() else 1e-8
@@ -106,14 +106,14 @@ def test_1d_complex_slice(self):
         self.sim.run(until_after_sources=0)
         vol = mp.Volume(center=self.center_1d, size=self.size_1d)
         hl_slice1d = self.sim.get_array(mp.Hz, vol, cmplx=True)
-        self.assertTrue(hl_slice1d.dtype == np.complex128)
+        self.assertTrue(hl_slice1d.dtype == np.complex64 if mp.is_single_precision() else np.complex128)
         self.assertTrue(hl_slice1d.shape[0] == 126)
 
     def test_2d_complex_slice(self):
         self.sim.run(until_after_sources=0)
         vol = mp.Volume(center=self.center_2d, size=self.size_2d)
         hl_slice2d = self.sim.get_array(mp.Hz, vol, cmplx=True)
-        self.assertTrue(hl_slice2d.dtype == np.complex128)
+        self.assertTrue(hl_slice2d.dtype == np.complex64 if mp.is_single_precision() else np.complex128)
         self.assertTrue(hl_slice2d.shape[0] == 126 and hl_slice2d.shape[1] == 38)
 
 

scheme/meep.i

@@ -33,9 +33,9 @@ static inline std::complex<double> my_complex_func2(double t, void *f) {
 }
 
 typedef struct { SCM func; int nf; } my_field_func_data;
-static inline std::complex<double> my_field_func(const std::complex<double> *fields,
-					    const meep::vec &loc,
-					    void *data_) {
+static inline std::complex<double> my_field_func(const std::complex<meep::realnum> *fields,
+                                                 const meep::vec &loc,
+                                                 void *data_) {
   my_field_func_data *data = (my_field_func_data *) data_;
   int num_items = data->nf;
   cnumber *items = new cnumber[num_items];