From 74f470504889963ee692aa89904d0da1230b801b Mon Sep 17 00:00:00 2001
From: Lindsey Heagy <lheagy@eos.ubc.ca>
Date: Mon, 14 Mar 2016 13:25:09 -0700
Subject: [PATCH] cleanup of TDEM example

---
 SimPEG/EM/TDEM/SurveyTDEM.py            | 32 +++++++++++++++----------
 SimPEG/Examples/EM_TDEM_1D_Inversion.py | 12 +++++-----
 2 files changed, 25 insertions(+), 19 deletions(-)

diff --git a/SimPEG/EM/TDEM/SurveyTDEM.py b/SimPEG/EM/TDEM/SurveyTDEM.py
index 07783f29..acd154ba 100644
--- a/SimPEG/EM/TDEM/SurveyTDEM.py
+++ b/SimPEG/EM/TDEM/SurveyTDEM.py
@@ -131,14 +131,18 @@ class BaseSrc(SimPEG.Survey.BaseSrc):
         return getattr(self, '_waveform', None)
     @waveform.setter
     def waveform(self, val):
-        # if self.waveform is None:
-        val._assertMatchesPair(self.waveformPair)
-        self._waveform = val
+        if self.waveform is None:
+            print val
+            val._assertMatchesPair(self.waveformPair)
+            self._mapping = val
+        else:
+            self._mapping = self.PropMap(val)
 
 
-    def __init__(self, rxList, waveform = None ):
-        self.waveform = waveform or StepOffWaveform()
-        SimPEG.Survey.BaseSrc.__init__(self, rxList)
+    def __init__(self, rxList, waveform = StepOffWaveform(), **kwargs):
+        self.waveform = waveform
+        SimPEG.Survey.BaseSrc.__init__(self, rxList, **kwargs)
+
 
 
     def bInitial(self, prob):
@@ -172,15 +176,17 @@ class BaseSrc(SimPEG.Survey.BaseSrc):
 
 
 class MagDipole(BaseSrc):
-    def __init__(self, rxList, waveform=None, loc=None, orientation='Z', moment=1., mu=mu_0):
 
-        self.loc = loc
-        self.orientation = orientation
-        assert orientation in ['X','Y','Z'], "Orientation (right now) doesn't actually do anything! The methods in SrcUtils should take care of this..."
-        self.moment = moment
-        self.mu = mu
+    waveform = None
+    loc = None
+    orientation = 'Z'
+    moment = 1.
+    mu = mu_0
+
+    def __init__(self, rxList, **kwargs):
+        assert self.orientation in ['X','Y','Z'], "Orientation (right now) doesn't actually do anything! The methods in SrcUtils should take care of this..."
         self.integrate = False
-        BaseSrc.__init__(self, rxList, waveform)
+        BaseSrc.__init__(self, rxList, **kwargs)
 
     def _bfromVectorPotential(self, prob):
         if prob._eqLocs is 'FE':
diff --git a/SimPEG/Examples/EM_TDEM_1D_Inversion.py b/SimPEG/Examples/EM_TDEM_1D_Inversion.py
index 629811ca..c92b972d 100644
--- a/SimPEG/Examples/EM_TDEM_1D_Inversion.py
+++ b/SimPEG/Examples/EM_TDEM_1D_Inversion.py
@@ -42,10 +42,10 @@ def run(plotIt=True):
 
 
     rxOffset=1e-3
-    rx = EM.TDEM.RxTDEM(np.array([[rxOffset, 0., 30]]), np.logspace(-5,-3, 31), 'bz')
-    src = EM.TDEM.SrcTDEM_VMD_MVP([rx], np.array([0., 0., 80]))
-    survey = EM.TDEM.SurveyTDEM([src])
-    prb = EM.TDEM.ProblemTDEM_b(mesh, mapping=mapping)
+    rx = EM.TDEM.Rx(np.array([[rxOffset, 0., 30]]), np.logspace(-5,-3, 31), 'bz')
+    src = EM.TDEM.SurveyTDEM.MagDipole([rx], loc=np.array([0., 0., 80]))
+    survey = EM.TDEM.Survey([src])
+    prb = EM.TDEM.Problem_b(mesh, mapping=mapping)
 
     prb.Solver = SolverLU
     prb.timeSteps = [(1e-06, 20),(1e-05, 20), (0.0001, 20)]
@@ -53,9 +53,9 @@ def run(plotIt=True):
 
     # create observed data
     std = 0.05
-    
+
     survey.dobs = survey.makeSyntheticData(mtrue,std)
-    survey.std = std 
+    survey.std = std
     survey.eps = 1e-5*np.linalg.norm(survey.dobs)
 
     if plotIt: