From 32f3ef301aaa54e98589e07507c87fb665b73cda Mon Sep 17 00:00:00 2001
From: seogi_macbook <sgkang09@github.com>
Date: Tue, 12 Jan 2016 18:52:39 -0800
Subject: [PATCH] Modifications for general waveform

---
 SimPEG/EM/TDEM/BaseTDEM.py   |  4 ++--
 SimPEG/EM/TDEM/SurveyTDEM.py | 30 ++++++++++++++++++++++++------
 SimPEG/Problem.py            |  8 ++++++++
 3 files changed, 34 insertions(+), 8 deletions(-)

diff --git a/SimPEG/EM/TDEM/BaseTDEM.py b/SimPEG/EM/TDEM/BaseTDEM.py
index e36d76b8..4d4b39a7 100644
--- a/SimPEG/EM/TDEM/BaseTDEM.py
+++ b/SimPEG/EM/TDEM/BaseTDEM.py
@@ -31,7 +31,7 @@ class FieldsTDEM(Problem.TimeFields):
 
 
 class BaseTDEMProblem(BaseTimeProblem, BaseEMProblem):
-    """docstring for ProblemTDEM1D"""
+    """docstring for BaseTDEMProblem"""
     def __init__(self, mesh, mapping=None, **kwargs):
         BaseTimeProblem.__init__(self, mesh, mapping=mapping, **kwargs)
 
@@ -43,7 +43,7 @@ class BaseTDEMProblem(BaseTimeProblem, BaseEMProblem):
         # Create a fields storage object
         F = self._FieldsForward_pair(self.mesh, self.survey)
         for src in self.survey.srcList:
-            # Set the initial conditions
+            # Set the initial conditions            
             F[src,:,0] = src.getInitialFields(self.mesh)
         F = self.forward(m, self.getRHS, F=F)
         if self.verbose: print '%s\nDone calculating fields(m)\n%s'%('*'*50,'*'*50)
diff --git a/SimPEG/EM/TDEM/SurveyTDEM.py b/SimPEG/EM/TDEM/SurveyTDEM.py
index 7f6e5c04..6a65c0fe 100644
--- a/SimPEG/EM/TDEM/SurveyTDEM.py
+++ b/SimPEG/EM/TDEM/SurveyTDEM.py
@@ -99,14 +99,33 @@ class SrcTDEM_VMD_MVP(SrcTDEM):
 
 
 class SrcTDEM_CircularLoop_MVP(SrcTDEM):
-
-    def __init__(self,rxList,loc,radius):
+    def __init__(self,rxList,loc,radius,waveformType):
         self.loc = loc
         self.radius = radius
-        SrcTDEM.__init__(self,rxList)
+        self.waveformType = waveformType
+        SrcTDEM.__init__(self,rxList)        
 
     def getInitialFields(self, mesh):
         """Circular Loop, magnetic vector potential"""
+        if self.waveformType == "STEPOFF":
+            print ">> Step waveform: Non-zero initial condition"
+            if mesh._meshType is 'CYL':
+                if mesh.isSymmetric:
+                    MVP = MagneticLoopVectorPotential(self.loc, mesh, 'Ey', self.radius)
+                else:
+                    raise NotImplementedError('Non-symmetric cyl mesh not implemented yet!')
+            elif mesh._meshType is 'TENSOR':
+                MVP = MagneticLoopVectorPotential(self.loc, mesh, ['Ex','Ey','Ez'], self.radius)
+            else:
+                raise Exception('Unknown mesh for CircularLoop')
+            return {"b": mesh.edgeCurl*MVP}
+        elif self.waveformType == "GENERAL":
+            print ">> General waveform: Zero initial condition"
+            return {"b": np.zeros(mesh.nF)}
+        else:
+            raise NotImplementedError("Only use STEPOFF or GENERAL")
+
+    def getMeS(self, mesh, MfMui):
         if mesh._meshType is 'CYL':
             if mesh.isSymmetric:
                 MVP = MagneticLoopVectorPotential(self.loc, mesh, 'Ey', self.radius)
@@ -115,9 +134,8 @@ class SrcTDEM_CircularLoop_MVP(SrcTDEM):
         elif mesh._meshType is 'TENSOR':
             MVP = MagneticLoopVectorPotential(self.loc, mesh, ['Ex','Ey','Ez'], self.radius)
         else:
-            raise Exception('Unknown mesh for CircularLoop')
-
-        return {"b": mesh.edgeCurl*MVP}
+            raise Exception('Unknown mesh for CircularLoop')            
+        return mesh.edgeCurl.T*MfMui*mesh.edgeCurl*MVP
 
 
 class SurveyTDEM(Survey.BaseSurvey):
diff --git a/SimPEG/Problem.py b/SimPEG/Problem.py
index 607cff5b..be29ff01 100644
--- a/SimPEG/Problem.py
+++ b/SimPEG/Problem.py
@@ -158,6 +158,9 @@ class BaseProblem(object):
 
 class BaseTimeProblem(BaseProblem):
     """Sets up that basic needs of a time domain problem."""
+    
+    waveformType = "STEPOFF"    
+    current = None
 
     @property
     def timeSteps(self):
@@ -184,6 +187,11 @@ class BaseTimeProblem(BaseProblem):
         self._timeSteps = Utils.meshTensor(value)
         del self.timeMesh
 
+    def currentwaveform(self, wave):
+        self._timeSteps = np.diff(wave[:,0])
+        self.current = wave[:,1]
+        self.waveformType = "GENERAL"
+
     @property
     def nT(self):
         "Number of time steps."