Merge pull request #86 from jcapriot/complex_mu_kernel

Allows complex mu values in rTE kernel

Author: jcapriot

geoana/kernels/_extensions/_rTE.cpp

@@ -9,19 +9,20 @@ void funcs::rTE(
     double * frequencies,
     double * lambdas,
     complex_t * sigmas,
-    double * mus,
+    complex_t * mus,
     double * h,
     size_t n_frequency,
     size_t n_filter,
     size_t n_layers
-){
+) noexcept(true){
     double mu_0 = 4.0E-7*M_PI;
     complex_t j(0.0, 1.0);
 
     complex_t k2, u, Yh, Y, tanhuh, Y0;
     complex_t *sigi;
-    double *mui;
-    double omega, mu_c, l2;
+    complex_t *mui;
+    double omega, l2;
+    complex_t mu_c;
 
     for(size_t i_filt=0, i=0; i_filt<n_filter; ++i_filt){
         l2 = lambdas[i_filt] * lambdas[i_filt];
@@ -55,12 +56,12 @@ void funcs::rTEgrad(
     double * frequencies,
     double * lambdas,
     complex_t * sigmas,
-    double * mus,
+    complex_t * mus,
     double * h,
     size_t n_frequency,
     size_t n_filter,
     size_t n_layers
-){
+) noexcept(true){
     double mu_0 = 4.0E-7*M_PI;
     complex_t j(0.0, 1.0);
 
@@ -74,7 +75,7 @@ void funcs::rTEgrad(
     vec_complex_t k2s(n_layers), us(n_layers), Yhs(n_layers);
     vec_complex_t Ys(n_layers-1), tanhs(n_layers-1);
 
-    double *mui;
+    complex_t *mui;
     complex_t *sigi, *TE_dhi, *TE_dmui, *TE_dsigmai;
     complex_t gyh0, bot, gy, gtanh, Y0;
     complex_t gu, gmu, gk2;

geoana/kernels/_extensions/_rTE.h

@@ -14,12 +14,12 @@ namespace funcs {
         double *frequencies,
         double *lambdas,
         complex_t *sigmas,
-        double *mus,
+        complex_t *mus,
         double *depths,
         size_t n_frequency,
         size_t n_filter,
         size_t n_layers
-    );
+    ) noexcept(true);
 
     void rTEgrad(
         complex_t * TE_dsigma,
@@ -28,12 +28,12 @@ namespace funcs {
         double * frequencies,
         double * lambdas,
         complex_t * sigmas,
-        double * mus,
+        complex_t * mus,
         double * h,
         size_t n_frequency,
         size_t n_filter,
         size_t n_layers
-    );
+    ) noexcept(true);
 }
 
 #endif

geoana/kernels/_extensions/rTE.pyx

@@ -19,12 +19,12 @@ cdef extern from "_rTE.h" namespace "funcs":
         double *frequencies,
         double *lambdas,
         complex_t *sigmas,
-        double *mus,
+        complex_t *mus,
         double *thicks,
         SIZE_t n_frequency,
         SIZE_t n_filter,
         SIZE_t n_layers
-        ) nogil
+        ) noexcept nogil
 
     void rTEgrad(
         complex_t * TE_dsigma,
@@ -33,12 +33,12 @@ cdef extern from "_rTE.h" namespace "funcs":
         double * frequencies,
         double * lambdas,
         complex_t * sigmas,
-        double * mus,
+        complex_t * mus,
         double * h,
         SIZE_t n_frequency,
         SIZE_t n_filter,
         SIZE_t n_layers
-    ) nogil
+    ) noexcept nogil
 
 def rTE_forward(frequencies, lamb, sigma, mu, thicknesses):
     """Compute reflection coefficients for Transverse Electric (TE) mode.
@@ -67,7 +67,7 @@ def rTE_forward(frequencies, lamb, sigma, mu, thicknesses):
     """
     # Sigma and mu must be fortran contiguous
     sigma = np.require(sigma, dtype=np.complex128, requirements="F")
-    mu = np.require(mu, dtype=np.float64, requirements="F")
+    mu = np.require(mu, dtype=np.complex128, requirements="F")
 
     # These just must be contiguous
     lamb = np.require(lamb, dtype=np.float64, requirements="C")
@@ -92,7 +92,7 @@ def rTE_forward(frequencies, lamb, sigma, mu, thicknesses):
         REAL_t[:] f = frequencies
         REAL_t[:] lam = lamb
         COMPLEX_t[:, :] sig = sigma
-        REAL_t[:, :] c_mu = mu
+        COMPLEX_t[:, :] c_mu = mu
         REAL_t[:] hs = thicknesses
 
     cdef:
@@ -109,12 +109,13 @@ def rTE_forward(frequencies, lamb, sigma, mu, thicknesses):
         # Types are same size and both pairs of numbers (r, i), so can cast one to the other
         complex_t *out_p = <complex_t *> &out[0, 0]
         complex_t *sig_p = <complex_t *> &sig[0, 0]
+        complex_t *mu_p = <complex_t *> &c_mu[0, 0]
         REAL_t *h_p = NULL
     if n_layers > 1:
         h_p = &hs[0]
 
     with nogil:
-        rTE(out_p, &f[0], &lam[0], sig_p, &c_mu[0, 0], h_p,
+        rTE(out_p, &f[0], &lam[0], sig_p, mu_p, h_p,
             n_frequency, n_filter, n_layers)
 
     return np.array(out)
@@ -153,7 +154,7 @@ def rTE_gradient(frequencies, lamb, sigma, mu, thicknesses):
     """
     # Require that they are all the same ordering as sigma
     sigma = np.require(sigma, dtype=np.complex128, requirements="F")
-    mu = np.require(mu, dtype=np.float64, requirements="F")
+    mu = np.require(mu, dtype=np.complex128, requirements="F")
     lamb = np.require(lamb, dtype=np.float64, requirements="C")
     frequencies = np.require(frequencies, dtype=np.float64, requirements="C")
     thicknesses = np.require(thicknesses, dtype=np.float64, requirements="C")
@@ -176,7 +177,7 @@ def rTE_gradient(frequencies, lamb, sigma, mu, thicknesses):
         REAL_t[:] f = frequencies
         REAL_t[:] lam = lamb
         COMPLEX_t[:, :] sig = sigma
-        REAL_t[:, :] c_mu = mu
+        COMPLEX_t[:, :] c_mu = mu
         REAL_t[:] hs = thicknesses
 
     cdef:
@@ -194,6 +195,7 @@ def rTE_gradient(frequencies, lamb, sigma, mu, thicknesses):
     cdef:
         # Types are same size and both pairs of numbers (r, i), so can cast one to the other
         complex_t *sig_p = <complex_t *> &sig[0, 0]
+        complex_t *mu_p = <complex_t *> &c_mu[0, 0]
         complex_t *gsig_p = <complex_t *> &gsig[0, 0, 0]
         complex_t *gmu_p = <complex_t *> &gmu[0, 0, 0]
         complex_t *gh_p = NULL
@@ -203,7 +205,7 @@ def rTE_gradient(frequencies, lamb, sigma, mu, thicknesses):
         gh_p = <complex_t *> &gh[0, 0, 0]
 
     with nogil:
-        rTEgrad(gsig_p, gmu_p, gh_p, &f[0], &lam[0], sig_p, &c_mu[0, 0], h_p,
+        rTEgrad(gsig_p, gmu_p, gh_p, &f[0], &lam[0], sig_p, mu_p, h_p,
             n_frequency, n_filter, n_layers)
 
     return np.array(gsig), np.array(gh), np.array(gmu)

tests/em/fdem/test_em_fdem_layered.py

@@ -294,9 +294,9 @@ def test_rTE_jacobian(self):
             thicknesses = np.ones(n_layer-1)
             lamb = np.logspace(0, 3, n_lambda)
             np.random.seed(0)
-            sigma = 1E-1*(1 + 0.1 * np.random.rand(n_layer, n_frequency))
-            mu = mu_0 * (1 + 0.1 * np.random.rand(n_layer, n_frequency))
-            dmu = mu_0 * 0.1 * np.random.rand(n_layer, n_frequency)
+            sigma = 1E-1*(1 + 1E-4j + 0.1 * np.random.rand(n_layer, n_frequency))
+            mu = mu_0 * (1 + 1.0E-4j + 0.1 * np.random.rand(n_layer, n_frequency))
+            dmu = mu_0 * (0.1 + 1.0E-5j) * np.random.rand(n_layer, n_frequency)
 
             def rte_sigma(x):
                 sigma = x.reshape(n_layer, n_frequency)
@@ -354,8 +354,8 @@ def setUp(self):
         thicknesses = np.ones(n_layer-1)
         lamb = np.logspace(-5, -3, n_lambda)
         np.random.seed(123)
-        sigma = 1E-1 * (1 + 1.0/(n_layer*n_frequency) * np.arange(n_layer*n_frequency).reshape(n_layer, n_frequency))
-        mu = mu_0 * (1 + 1.0/(n_layer*n_frequency) * np.arange(n_layer*n_frequency).reshape(n_layer, n_frequency))
+        sigma = 1E-1 * (1 + 1E-4j + 1.0/(n_layer*n_frequency) * np.arange(n_layer*n_frequency).reshape(n_layer, n_frequency))
+        mu = mu_0 * (1 + 1.0E-4j + 1.0/(n_layer*n_frequency) * np.arange(n_layer*n_frequency).reshape(n_layer, n_frequency))
 
         self.rTE1 = rTE_forward(frequencies, lamb, sigma, mu, thicknesses)
         self.rTE2 = _rTE_forward(frequencies, lamb, sigma, mu, thicknesses)