From 2e0aadf1132d05cc7b897678de8cb5271d32d47a Mon Sep 17 00:00:00 2001
From: rowanc1 <rowan@visiblegeology.com>
Date: Fri, 4 Jul 2014 14:46:58 -0700
Subject: [PATCH] Generalize sources a little bit to make them easier to use

---
 simpegEM/Sources/CircularLoop.py   | 21 ++++++++++++++++-----
 simpegEM/Sources/magneticDipole.py | 22 +++++++++++++++++-----
 2 files changed, 33 insertions(+), 10 deletions(-)

diff --git a/simpegEM/Sources/CircularLoop.py b/simpegEM/Sources/CircularLoop.py
index 338bc732..6f5f2233 100644
--- a/simpegEM/Sources/CircularLoop.py
+++ b/simpegEM/Sources/CircularLoop.py
@@ -7,14 +7,25 @@ def MagneticLoopVectorPotential(txLoc, obsLoc, component, radius):
         at given locations
 
         :param numpy.ndarray txLoc: Location of the transmitter(s) (x, y, z)
-        :param numpy.ndarray obsLoc: Where the potentials will be calculated (x, y, z)
-        :param str component: The component to calculate - 'x', 'y', or 'z'
+        :param numpy.ndarray,SimPEG.Mesh obsLoc: Where the potentials will be calculated (x, y, z) or a SimPEG Mesh
+        :param str,list component: The component to calculate - 'x', 'y', or 'z' if an array, or grid type if mesh, can be a list
         :param numpy.ndarray I: Input current of the loop
         :param numpy.ndarray radius: radius of the loop
         :rtype: numpy.ndarray
         :return: The vector potential each dipole at each observation location
     """
 
+    if type(component) in [list, tuple]:
+        out = range(len(component))
+        for i, comp in enumerate(component):
+            out[i] = MagneticLoopVectorPotential(txLoc, obsLoc, comp, radius)
+        return np.concatenate(out)
+
+    if isinstance(obsLoc, Mesh.BaseMesh):
+        mesh = obsLoc
+        assert component in ['Ex','Ey','Ez','Fx','Fy','Fz'], "Components must be in: ['Ex','Ey','Ez','Fx','Fy','Fz']"
+        return MagneticLoopVectorPotential(txLoc, getattr(mesh,'grid'+component), component[1], radius)
+
     txLoc = np.atleast_2d(txLoc)
     obsLoc = np.atleast_2d(obsLoc)
 
@@ -24,7 +35,7 @@ def MagneticLoopVectorPotential(txLoc, obsLoc, component, radius):
     if component=='z':
         A = np.zeros((n, nTx))
         if nTx ==1:
-            return A.flatten()        
+            return A.flatten()
         return A
 
     else:
@@ -51,10 +62,10 @@ def MagneticLoopVectorPotential(txLoc, obsLoc, component, radius):
             elif component == 'y':
                 A[ind, i] = Aphi[ind] * ( x[ind] / r[ind] )
             else:
-                raise ValueError('Invalid component')               
+                raise ValueError('Invalid component')
 
         if nTx == 1:
-            return A.flatten()                
+            return A.flatten()
         return A
 
 if __name__ == '__main__':
diff --git a/simpegEM/Sources/magneticDipole.py b/simpegEM/Sources/magneticDipole.py
index 162af089..9f6166c8 100644
--- a/simpegEM/Sources/magneticDipole.py
+++ b/simpegEM/Sources/magneticDipole.py
@@ -1,5 +1,6 @@
 import numpy as np
 from scipy.constants import mu_0, pi
+from SimPEG import Mesh
 
 def MagneticDipoleVectorPotential(txLoc, obsLoc, component, dipoleMoment=(0., 0., 1.)):
     """
@@ -7,18 +8,29 @@ def MagneticDipoleVectorPotential(txLoc, obsLoc, component, dipoleMoment=(0., 0.
         at given locations 'ref. <http://en.wikipedia.org/wiki/Dipole#Magnetic_vector_potential>'
 
         :param numpy.ndarray txLoc: Location of the transmitter(s) (x, y, z)
-        :param numpy.ndarray obsLoc: Where the potentials will be calculated (x, y, z)
-        :param str component: The component to calculate - 'x', 'y', or 'z'
+        :param numpy.ndarray,SimPEG.Mesh obsLoc: Where the potentials will be calculated (x, y, z) or a SimPEG Mesh
+        :param str,list component: The component to calculate - 'x', 'y', or 'z' if an array, or grid type if mesh, can be a list
         :param numpy.ndarray dipoleMoment: The vector dipole moment
         :rtype: numpy.ndarray
         :return: The vector potential each dipole at each observation location
     """
 
-    if component=='x':
+    if type(component) in [list, tuple]:
+        out = range(len(component))
+        for i, comp in enumerate(component):
+            out[i] = MagneticDipoleVectorPotential(txLoc, obsLoc, comp, dipoleMoment=dipoleMoment)
+        return np.concatenate(out)
+
+    if isinstance(obsLoc, Mesh.BaseMesh):
+        mesh = obsLoc
+        assert component in ['Ex','Ey','Ez','Fx','Fy','Fz'], "Components must be in: ['Ex','Ey','Ez','Fx','Fy','Fz']"
+        return MagneticDipoleVectorPotential(txLoc, getattr(mesh,'grid'+component), component[1], dipoleMoment=dipoleMoment)
+
+    if component == 'x':
         dimInd = 0
-    elif component=='y':
+    elif component == 'y':
         dimInd = 1
-    elif component=='z':
+    elif component == 'z':
         dimInd = 2
     else:
         raise ValueError('Invalid component')