Updated codes, fixed bug in dataMT and moved notebooks to a folder.

simpegMT/ProblemMT.py

@@ -1,8 +1,8 @@
 from SimPEG import Survey, Problem, Utils, Models, np, sp, SolverLU as SimpegSolver
 from simpegEM.FDEM import BaseFDEMProblem
-from simpegEM.Utils import omega
+from simpegEM.Utils.EMUtils import omega
 from scipy.constants import mu_0
-from SurveyMT import SurveyMT, FieldsMT
+from SurveyMT import SurveyMT, FieldsMT, DataMT
 import multiprocessing, sys, time
 
 
@@ -14,7 +14,7 @@ class BaseMTProblem(BaseFDEMProblem):
 
 
     surveyPair = SurveyMT
-    dataPair = Survey.Data
+    dataPair = DataMT
 
     Solver = SimpegSolver
     solverOpts = {}
@@ -25,14 +25,6 @@ class BaseMTProblem(BaseFDEMProblem):
     # Might need to add more stuff here.
 
 
-    @property
-    def MeSigmaBG(self):
-        #TODO: hardcoded to sigma as the model
-        if getattr(self, '_MeSigmaBG', None) is None:
-            sigmaBG = self.backModel
-            self._MeSigmaBG = self.mesh.getEdgeInnerProduct(sigmaBG)
-        return self._MeSigmaBG
-
     
 class ProblemMT_eForm_ps(BaseMTProblem):
     """ 
@@ -44,6 +36,7 @@ class ProblemMT_eForm_ps(BaseMTProblem):
     _fieldType = 'e'
     _eqLocs    = 'FE'
     fieldsPair = FieldsMT
+
     # Set new properties 
     # Background model
     @property
@@ -62,6 +55,15 @@ class ProblemMT_eForm_ps(BaseMTProblem):
             if hasattr(self, prop):
                 delattr(self, prop)
 
+    @property
+    def MeDeltaSigma(self):
+        #TODO: hardcoded to sigma as the model
+        if getattr(self, '_MeDeltaSigma', None) is None:
+            sigma = self.curModel
+            sigmaBG = self.backModel
+            self._MeDeltaSigma = self.mesh.getEdgeInnerProduct(sigma - sigmaBG)
+        return self._MeDeltaSigma
+
     def __init__(self, mesh, **kwargs):
         BaseMTProblem.__init__(self, mesh, **kwargs)
 
@@ -104,8 +106,136 @@ class ProblemMT_eForm_ps(BaseMTProblem):
         # Get the background electric fields
         from simpegMT.Sources import homo1DModelSource
         eBG_bp = homo1DModelSource(self.mesh,freq,backSigma)
-        deltM = self.curModel - self.backModel
-        Abg = -1j*omega(freq)*deltM*eBG_bp
+        deltM = self.MeDeltaSigma
+        Abg = -1j*omega(freq)*deltM
+
+        return Abg*eBG_bp, eBG_bp
+
+    def getRHSderiv(self, freq, backSigma, u, v, adjoint=False):
+        raise NotImplementedError('getRHSDeriv not implemented yet')
+        return None
+
+    def fields(self, m, m_back):
+        '''
+        Function to calculate all the fields for the model m.
+
+        :param np.ndarray (nC,) m: Conductivity model
+        :param np.ndarray (nC,) m_back: Background conductivity model
+        '''
+        self.curModel = m
+        self.backModel = m_back
+        # RHS, CalcFields = self.getRHS(freq,m_back), self.calcFields
+
+        F = FieldsMT(self.mesh, self.survey)
+        for freq in self.survey.freqs:
+            if self.verbose:
+                startTime = time.time()
+                print 'Starting work for {:.3e}'.format(freq)
+                sys.stdout.flush()
+            A = self.getA(freq)
+            rhs, e_p = self.getRHS(freq,m_back)
+            Ainv = self.Solver(A, **self.solverOpts)
+            e_s = Ainv * rhs 
+            e = e_p + e_s
+            # Store the fields
+            Src = self.survey.getSources(freq)
+            # Store the fieldss
+            F[Src, 'e_px'] = e[:,0]
+            F[Src, 'e_py'] = e[:,1]
+            # Note curl e = -iwb so b = -curl/iw
+            b = -( self.mesh.edgeCurl * e )/( 1j*omega(freq) )
+            F[Src, 'b_px'] = b[:,0]
+            F[Src, 'b_py'] = b[:,1]
+            if self.verbose:
+                print 'Ran for {:f} seconds'.format(time.time()-startTime)
+                sys.stdout.flush()
+        return F
+        
+class ProblemMT_eForm_Tp(BaseMTProblem):
+    """ 
+    A MT problem solving a e formulation and a primary/secondary fields decompostion.
+
+    Solves the equation
+    """
+
+    _fieldType = 'e'
+    _eqLocs    = 'FE'
+    fieldsPair = FieldsMT
+
+    # Set new properties 
+    # Background model
+    @property
+    def backModel(self):
+        """
+            Sets the model, and removes dependent mass matrices.
+        """
+        return getattr(self, '_backModel', None)
+
+    @backModel.setter
+    def backModel(self, value):
+        if value is self.backModel:
+            return # it is the same!
+        self._backModel = Models.Model(value, self.mapping)
+        for prop in self.deleteTheseOnModelUpdate:
+            if hasattr(self, prop):
+                delattr(self, prop)
+
+    @property
+    def MeSigmaBack(self):
+        #TODO: hardcoded to sigma as the model
+        if getattr(self, '_MeSigmaBack', None) is None:
+            sigma = self.curModel
+            sigmaBG = self.backModel
+            self._MeSigmaBack = self.mesh.getEdgeInnerProduct(sigmaBG)
+        return self._MeSigmaBack
+
+    def __init__(self, mesh, **kwargs):
+        BaseMTProblem.__init__(self, mesh, **kwargs)
+
+    def getA(self, freq):
+        """
+            Function to get the A matrix.
+
+            :param float freq: Frequency
+            :rtype: scipy.sparse.csr_matrix
+            :return: A
+        """
+        mui = self.MfMui
+        sig = self.MeSigma
+        C = self.mesh.edgeCurl
+
+        return C.T*mui*C + 1j*omega(freq)*sig
+
+    def getADeriv(self, freq, u, v, adjoint=False):
+        sig = self.curTModel
+        dsig_dm = self.curTModelDeriv
+        dMe_dsig = self.mesh.getEdgeInnerProductDeriv(sig, v=u)
+
+        if adjoint:
+            return 1j * omega(freq) * ( dsig_dm.T * ( dMe_dsig.T * v ) )
+
+        return 1j * omega(freq) * ( dMe_dsig * ( dsig_dm * v ) )
+
+    def getRHS(self, freq, backSigma):
+        """
+            Function to return the right hand side for the system.
+            :param float freq: Frequency
+            :param numpy.ndarray (nC,) backSigma: Background conductivity model
+            :rtype: numpy.ndarray (nE, 2)
+            :return: one RHS for both polarizations
+        """
+        # Get sources for the frequency
+        src = self.survey.getSources(freq)
+        # Make sure that there is 2 polarizations. 
+        # assert len()
+        # Get the background electric fields
+        from simpegMT.Sources import homo1DModelSource
+        eBG_bp = homo1DModelSource(self.mesh,freq,backSigma)
+        MeBack = self.MeSigmaBack
+        # Set up the A system
+        mui = self.MfMui
+        C = self.mesh.edgeCurl
+        Abg = C.T*mui*C + 1j*omega(freq)*MeBack
 
         return Abg*eBG_bp, eBG_bp
     def getRHSderiv(self, freq, backSigma, u, v, adjoint=False):
@@ -124,28 +254,27 @@ class ProblemMT_eForm_ps(BaseMTProblem):
         # RHS, CalcFields = self.getRHS(freq,m_back), self.calcFields
 
         F = FieldsMT(self.mesh, self.survey)
-
-        if verbose:
-            startTime = time.time()
-            print 'Starting work for {:.3e}'.format(freq)
-            sys.stdout.flush()
-        A = self.getA(freq)
-        rhs, e_p = self.getRHS(freq,m_back)
-        Ainv = self.Solver(A, **self.solverOpts)
-        e_s = Ainv * rhs 
-        e = e_p + e_s
-        # Store the fields
-        Src = self.survey.getSources(freq)
-        # Store the fieldss
-        F[Src, 'e_px'] = e[:,0]
-        F[Src, 'e_py'] = e[:,1]
-        # Note curl e = -iwb so b = -curl/iw
-        b = -( self.mesh.edgeCurl * e )/( 1j*omega(freq) )
-        F[Src, 'b_px'] = b[:,0]
-        F[Src, 'b_py'] = b[:,1]
-        if verbose:
-            print 'Ran for {:f} seconds'.format(time.time()-startTime)
-            sys.stdout.flush()
+        for freq in self.survey.freqs:
+            if self.verbose:
+                startTime = time.time()
+                print 'Starting work for {:.3e}'.format(freq)
+                sys.stdout.flush()
+            A = self.getA(freq)
+            rhs, e_p = self.getRHS(freq,m_back)
+            Ainv = self.Solver(A, **self.solverOpts)
+            e_s = Ainv * rhs 
+            e = e_p + e_s
+            # Store the fields
+            Src = self.survey.getSources(freq)
+            # Store the fieldss
+            F[Src, 'e_px'] = e[:,0]
+            F[Src, 'e_py'] = e[:,1]
+            # Note curl e = -iwb so b = -curl/iw
+            b = -( self.mesh.edgeCurl * e )/( 1j*omega(freq) )
+            F[Src, 'b_px'] = b[:,0]
+            F[Src, 'b_py'] = b[:,1]
+            if self.verbose:
+                print 'Ran for {:f} seconds'.format(time.time()-startTime)
+                sys.stdout.flush()
         return F
         
-

simpegMT/SurveyMT.py

@@ -244,10 +244,10 @@ class DataMT(Survey.Data):
             tArrRec = np.concatenate((src.freq*np.ones((locs.shape[0],1)),locs,np.nan*np.ones((locs.shape[0],8))),axis=1).view(dtRI)
             # np.array([(src.freq,rx.locs[0,0],rx.locs[0,1],rx.locs[0,2],np.nan ,np.nan ,np.nan ,np.nan ,np.nan ,np.nan ,np.nan ,np.nan ) for rx in src.rxList],dtype=dtRI)
             # Get the type and the value for the DataMT object as a list
-            typeList = [[rx.rxType,self[src,rx][0]] for rx in src.rxList]
+            typeList = [[rx.rxType,self[src,rx]] for rx in src.rxList]
             # Insert the values to the temp array
             for nr,(key,val) in enumerate(typeList):
-                tArrRec[key] = val
+                tArrRec[key] = mkvc(val,2)
             # Masked array 
             mArrRec = np.ma.MaskedArray(rec2ndarr(tArrRec),mask=np.isnan(rec2ndarr(tArrRec))).view(dtype=tArrRec.dtype)
             # Unique freq and loc of the masked array