From 3df2140a88f5a2c1d57caefa7b9286bbde086768 Mon Sep 17 00:00:00 2001
From: Lindsey <lindseyheagy@gmail.com>
Date: Mon, 1 Jun 2015 14:32:42 -0700
Subject: [PATCH] Fields now take a srcList. A lot of this could be vectorized
 later

---
 simpegEM/FDEM/FieldsFDEM.py | 134 ++++++++++++++++++++----------------
 simpegEM/FDEM/SurveyFDEM.py |   8 ---
 simpegEM/Tests/test_FDEM.py |   7 +-
 3 files changed, 80 insertions(+), 69 deletions(-)

diff --git a/simpegEM/FDEM/FieldsFDEM.py b/simpegEM/FDEM/FieldsFDEM.py
index 0caccafd..e12752b6 100644
--- a/simpegEM/FDEM/FieldsFDEM.py
+++ b/simpegEM/FDEM/FieldsFDEM.py
@@ -21,32 +21,36 @@ class FieldsFDEM_e(FieldsFDEM):
     def startup(self):
         self._edgeCurl = self.survey.prob.mesh.edgeCurl
 
-    def _e(self, e_sol, src):
+    def _e(self, e_sol, srcList):
         e = e_sol
-        e_p = src.e_p(self.survey.prob)
-        if e_p is not None:
-            e += e_p     
+        for i, src in enumerate(srcList):
+            e_p = src.e_p(self.survey.prob)
+            if e_p is not None:
+                e[:,i] += e_p     
         return e
 
-    def _b(self, e_sol, src):
+    def _b(self, e_sol, srcList):
         C = self._edgeCurl
-        b = - 1./(1j*omega(src.freq))*(C * e_sol)
-        S_m, _ = src.eval(self.survey.prob)
-        if S_m is not None:
-            b += 1./(1j*omega(src.freq)) * S_m
+        b = (C * e_sol)
+        for i, src in enumerate(srcList):
+            b[:,i] *= - 1./(1j*omega(src.freq))
+            S_m, _ = src.eval(self.survey.prob)
+            if S_m is not None:
+                b[:,i] += 1./(1j*omega(src.freq)) * S_m
 
-        b_p = src.b_p(self.survey.prob)
-        if b_p is not None:
-            b += b_p 
+            b_p = src.b_p(self.survey.prob)
+            if b_p is not None:
+                b[:,i] += b_p 
 
         return b
 
-    def _bDeriv(self, e, src, v, adjoint=False):
-        S_mDeriv,_ = src.getSourceDeriv(self.survey.prob, v, adjoint)
-        if S_mDeriv is None:
-            return None
-        else:
-            return 1./(1j*omega(src.freq)) * S_mDeriv
+    def _bDeriv(self, e, srcList, v, adjoint=False):
+        raise NotImplementedError('Fields Derivs Not Implemented Yet')
+        # S_mDeriv,_ = src.getSourceDeriv(self.survey.prob, v, adjoint)
+        # if S_mDeriv is None:
+        #     return None
+        # else:
+        #     return 1./(1j*omega(src.freq)) * S_mDeriv
 
 
 class FieldsFDEM_b(FieldsFDEM):
@@ -64,27 +68,32 @@ class FieldsFDEM_b(FieldsFDEM):
         self._MeSigmaI = self.survey.prob.MeSigmaI
         self._MfMui = self.survey.prob.MfMui
 
-    def _b(self, b_sol, src):
+    def _b(self, b_sol, srcList):
         b = b_sol
-        b_p = src.b_p(self.survey.prob)
-        if b_p is not None:
-            b += b_p
+
+        for i, src in enumerate(srcList):
+            b_p = src.b_p(self.survey.prob)
+            if b_p is not None:
+                b[:,i] += b_p
         return b  
 
-    def _e(self, b_sol, src):
+    def _e(self, b_sol, srcList):
         e = self._MeSigmaI * ( self._edgeCurl.T * ( self._MfMui * b_sol))
-        _,S_e = src.eval(self.survey.prob)
-        if S_e is not None:
-            e += -self._MeSigmaI*S_e
 
-        e_p = src.e_p(self.survey.prob)
-        if e_p is not None:
-            e += e_p
+        for i,src in enumerate(srcList): 
+            _,S_e = src.eval(self.survey.prob)
+            if S_e is not None:
+                e += -self._MeSigmaI*S_e
+
+            e_p = src.e_p(self.survey.prob)
+            if e_p is not None:
+                e[:,i] += e_p
 
         return e
 
-    def _eDeriv(self, b_sol, src, v, adjoint=False):
-        _,S_eDeriv = src.getSourceDeriv(self.survey.prob, v, adjoint)
+    def _eDeriv(self, b_sol, srcList, v, adjoint=False):
+        raise NotImplementedError('Fields Derivs Not Implemented Yet')
+        _,S_eDeriv = src.evalDeriv(self.survey.prob, v, adjoint)
 
         if S_eDeriv is None:
             return None
@@ -108,30 +117,35 @@ class FieldsFDEM_j(FieldsFDEM):
         self._MfSigmai = self.survey.prob.MfSigmai
         self._curModel = self.survey.prob.curModel
 
-    def _j(self, j_sol, src):
+    def _j(self, j_sol, srcList):
         j = j_sol
-        j_p = src.j_p(self.survey.prob) 
-        if j_p is not None:
-            j += j_p
+        for i, src in enumerate(srcList):
+            j_p = src.j_p(self.survey.prob) 
+            if j_p is not None:
+                j[:,i] += j_p
         return j
 
-    def _h(self, j_sol, src): 
+    def _h(self, j_sol, srcList): 
         MeMuI = self._MeMuI
         C = self._edgeCurl
         MfSigmai = self._MfSigmai
 
-        h = - 1./(1j*omega(src.freq)) * MeMuI * (C.T * (MfSigmai * j_sol) ) 
-        S_m,_ = src.eval(self.survey.prob)
-        if S_m is not None:
-            h += 1./(1j*omega(src.freq)) * MeMuI * S_m
+        h =  MeMuI * (C.T * (MfSigmai * j_sol) ) 
+
+        for i, src in enumerate(srcList):
+            h[:,i] *= -1./(1j*omega(src.freq))
+            S_m,_ = src.eval(self.survey.prob)
+            if S_m is not None:
+                h[:,i] += 1./(1j*omega(src.freq)) * MeMuI * S_m
+
+            h_p = src.h_p(self.survey.prob)
+            if h_p is not None:
+                h[:,i] += h_p 
 
-        h_p = src.h_p(self.survey.prob)
-        if h_p is not None:
-            h += h_p 
         return h
 
-    def _hDeriv(self, j_sol, src, v, adjoint=False):
-
+    def _hDeriv(self, j_sol, srcList, v, adjoint=False):
+        raise NotImplementedError('Fields Derivs Not Implemented Yet')
         sig = self._curModel.transform
         sigi = 1/sig
         dsig_dm = self._curModel.transformDeriv
@@ -165,24 +179,28 @@ class FieldsFDEM_h(FieldsFDEM):
         self._MeMuI = self.survey.prob.MeMuI
         self._MfSigmai = self.survey.prob.MfSigmai
 
-    def _h(self, h_sol, src):
+    def _h(self, h_sol, srcList):
         h = h_sol
-        h_p = src.h_p(self.survey.prob)
-        if h_p is not None:
-            h += h_p
-        return h
+        for i, src in enumerate(srcList):
+            h_p = src.h_p(self.survey.prob)
+            if h_p is not None:
+                h[:,i] += h_p
+            return h
 
-    def _j(self, h_sol, src):
+    def _j(self, h_sol, srcList):
         j = self._edgeCurl*h_sol
-        _,S_e = src.eval(self.survey.prob)
-        if S_e is not None:
-            j += -S_e
-        j_p = src.j_p(self.survey.prob)
-        if j_p is not None:
-            j += j_p 
+        for i, src in enumerate(srcList):
+            _,S_e = src.eval(self.survey.prob)
+            if S_e is not None:
+                j[:,i] += -S_e
+
+            j_p = src.j_p(self.survey.prob)
+            if j_p is not None:
+                j[:,i] += j_p 
         return j
 
-    def _jDeriv(self, h_sol, src, v, adjoint=False):
+    def _jDeriv(self, h_sol, srcList, v, adjoint=False):
+        raise NotImplementedError('Fields Derivs Not Implemented Yet')
         _,S_eDeriv = src.getSourceDeriv(self.survey.prob, v, adjoint)
         if S_eDeriv is None:
             return None
diff --git a/simpegEM/FDEM/SurveyFDEM.py b/simpegEM/FDEM/SurveyFDEM.py
index 35f7ec97..ff1f50fe 100644
--- a/simpegEM/FDEM/SurveyFDEM.py
+++ b/simpegEM/FDEM/SurveyFDEM.py
@@ -98,10 +98,6 @@ class SrcFDEM(Survey.BaseSrc):
     def eval(self, prob):
         S_m = self._getS_m(prob)
         S_e = self._getS_e(prob)
-
-        if S_m is not None and S_m.ndim == 1: S_m = Utils.mkvc(S_m,2)
-        if S_e is not None and S_e.ndim == 1: S_e = Utils.mkvc(S_e,2)
-
         return S_m, S_e 
 
     def evalDeriv(self, prob, v, adjoint=None):
@@ -109,22 +105,18 @@ class SrcFDEM(Survey.BaseSrc):
 
     def b_p(self,prob):
         b_p = self._getb_p(prob)
-        if b_p is not None and b_p.ndim == 1: b_p = Utils.mkvc(b_p,2)
         return b_p 
 
     def h_p(self,prob):
         h_p = self._geth_p(prob)
-        if h_p is not None and h_p.ndim == 1: h_p = Utils.mkvc(h_p,2)
         return h_p
 
     def e_p(self,prob):
         e_p = self._gete_p(prob)
-        if e_p is not None and e_p.ndim == 1: e_p = Utils.mkvc(e_p,2)
         return e_p
 
     def j_p(self,prob):
         j_p = self._getj_p(prob)
-        if j_p is not None and j_p.ndim == 1: j_p = Utils.mkvc(j_p,2)
         return j_p
 
     def _getb_p(self,prob):
diff --git a/simpegEM/Tests/test_FDEM.py b/simpegEM/Tests/test_FDEM.py
index 5458f42a..f6f71982 100644
--- a/simpegEM/Tests/test_FDEM.py
+++ b/simpegEM/Tests/test_FDEM.py
@@ -17,7 +17,7 @@ TOL = 1e-4
 FLR = 1e-20 # "zero", so if residual below this --> pass regardless of order
 CONDUCTIVITY = 1e1
 MU = mu_0
-freq = 1e-1
+freq = [1e-1, 2e-1]
 addrandoms = True
 
 
@@ -35,9 +35,10 @@ def getProblem(fdemType, comp):
     x = np.array([np.linspace(-30,-15,3),np.linspace(15,30,3)]) #don't sample right by the source
     XYZ = Utils.ndgrid(x,x,np.r_[0.])
     Rx0 = EM.FDEM.RxFDEM(XYZ, comp)
-    Src0 = EM.FDEM.SrcFDEM_MagDipole([Rx0], loc=np.r_[0.,0.,0.], freq=freq)
+    Src0 = EM.FDEM.SrcFDEM_MagDipole([Rx0], loc=np.r_[0.,0.,0.], freq=freq[0])
+    Src1 = EM.FDEM.SrcFDEM_MagDipole([Rx0], loc=np.r_[0.,0.,0.], freq=freq[1])
 
-    survey = EM.FDEM.SurveyFDEM([Src0])
+    survey = EM.FDEM.SurveyFDEM([Src0, Src1])
 
 
     if verbose: