ENH: added vector operators: divergence, curl and laplacian

numpy/lib/function_base.py

@@ -1320,6 +1320,258 @@ def gradient(f, *varargs, **kwargs):
     else:
         return outvals
 
+def divergence(v,*varargs,**kwargs):
+    """
+    Return the divergence of an N-dimensional vector field of N 
+    components, each of dimension N.
+
+    The divergence is computed using second order accurate central
+    differences in the interior and either first differences or second
+    order accurate one-sides (forward or backwards) differences at the
+    boundaries. The returned gradient hence has the same shape as the
+    input array.
+
+    Parameters
+    ----------  
+    v : list of numpy arrays
+        Each of these array is the N-dimensional projection of the vector
+        field on the corresponding axis. 
+
+    varargs : scalar or list of scalar, optional
+        N scalars specifying the sample distances for each dimension,
+        i.e. `dx`, `dy`, `dz`, ... Default distance: 1.
+        single scalar specifies sample distance for all dimensions.
+        if `axis` is given, the number of varargs must equal the number of axes.
+    edge_order : {1, 2}, optional
+        Gradient is calculated using N\ :sup:`th` order accurate differences
+        at the boundaries. Default: 1.
+
+
+    Returns
+    -------
+    divergence : numpy array
+        The output corresponds to the divergence of the input vector field. 
+        This means that the output array has the form 
+        dAx/dx + dAy/dy + dAz/dz + ... for an input vector field 
+        A = (Ax, Ay, Az, ...) up to N dimensions.
+
+    Example
+    -------
+    This example shows how the calculated divergence of the 2-D field 
+    (0.5*x**2, -y*x) (whose divergence should be 0 everywhere) returns a 0 array
+
+    >>> import numpy as np
+    >>> X,Y=np.mgrid[0:2000,0:2000]
+    >>> a1=0.5*X**2
+    >>> a2=-Y*X
+    >>> c=[a1,a2]
+    >>> d=np.divergence(c)
+    >>> print d
+    [[ 0.  0.  0. ...,  0.  0.  0.]
+     [ 0.  0.  0. ...,  0.  0.  0.]
+     [ 0.  0.  0. ...,  0.  0.  0.]
+     ..., 
+     [ 0.  0.  0. ...,  0.  0.  0.]
+     [ 0.  0.  0. ...,  0.  0.  0.]
+     [ 0.  0.  0. ...,  0.  0.  0.]
+    """
+    N = [0]*len(v)
+    for i, v_c in enumerate(v):
+        v_c = np.asanyarray(v_c)
+        N[i] = len(v_c.shape) # Get number of dimensions from every component
+    if False in [N[0] == N[i] for i in range(len(v))]: 
+        #Check if all components are the same size
+        raise ValueError("Not all components of input"
+	                 " have the same number of dimensions")
+    else:
+        if False in [np.shape(v[0]) == np.shape(v[i]) for i in range(len(v))]:
+            raise ValueError("Not all components of input are the same size")
+        else:
+	    N = N[0] 
+	# If all vector field components are the same shape,
+	#N becomes the number of dimensions
+
+    axes = tuple(range(N))
+    outvals=np.zeros(np.shape(v_c)) #Initialize output
+
+    for i,ax in enumerate(axes):
+        outvals += np.gradient(v[ax], *varargs, axis=ax,edge_order=0)
+
+    return outvals
+
+
+def curl(v,*varargs,**kwargs):
+    """
+    Return the curl of a 3-D vector field.
+
+    The curl is computed using second order accurate central differences
+    in the interior and either first differences or second order accurate
+    one-sides (forward or backwards) differences at the boundaries. The
+    returned gradient hence has the same shape as the input array.
+
+    Parameters
+    ----------
+    v : list of numpy arrays
+        Each of these array is the N-dimensional projection 
+        of the vector field on the corresponding axis. 
+
+    varargs : scalar or list of scalar, optional
+        N scalars specifying the sample distances for each dimension,
+        i.e. `dx`, `dy`, `dz`, ... Default distance: 1.
+        single scalar specifies sample distance for all dimensions.
+        if `axis` is given, the number of varargs must equal the number of axes.
+    edge_order : {1, 2}, optional
+        Gradient is calculated using N\ :sup:`th` order accurate differences
+        at the boundaries. Default: 1.
+
+    axis : None or int or tuple of ints, optional
+        Gradient is calculated only along the given axis or axes
+        The default (axis = None) is to calculate the gradient for 
+        all the axes of the input array.
+        axis may be negative, in which case it counts from the last to the first axis.
+
+    Returns
+    -------
+    curl : list of numpy arrays
+        The output corresponds to the curl of the input 3-D vector field.
+        This means that the output has the form 
+        (dAz/dy-dAy/dz, dAx/dz-dAz/dx, dAy/dx-dAx/dy) of an input vector 
+        field A = (Ax, Ay, Ax).
+
+    Example
+    -------
+    The following example shows in matplotlib how the vector field (0,0,x*y) has the 
+    correct value, nonzero only in the third component of the resulting vector field
+    after applying curl
+
+    >>> import matplotlib.pylab as plt
+    >>> import numpy as np
+    >>> X,Y,Z=np.mgrid[0:200,0:200,0:200]
+    >>> a0=X*Y
+    >>> a1=np.zeros(np.shape(X))
+    >>> a2=np.zeros(np.shape(X))
+    >>> a=[a0,a1,a2]
+    >>> [cx,cy,cz]=np.curl(a)
+    >>> plt.imshow(cx[:,:,1])
+    >>> plt.colorbar()
+    >>> plt.show()
+    >>> plt.imshow(cy[:,:,1])
+    >>> plt.colorbar()
+    >>> plt.show()
+    >>> plt.imshow(cz[:,:,1])
+    >>> plt.colorbar()
+    >>> plt.show()
+    """
+    N = [0]*len(v)
+    for i, v_c in enumerate(v):
+        v_c = np.asanyarray(v_c)
+	N[i] = len(v_c.shape) 
+        # Extract the number of dimensions from every component v_c
+    if False in [N[i] == 3 for i in range(len(v))]: 
+        #Check if all components are the same size
+        raise ValueError("Not all components of input vector field"
+                         " have three dimensions")
+    else:
+        if False in [np.shape(v[0]) == np.shape(v[i]) for i in range(len(v))]:
+            raise ValueError("Not all components of input are the same size")
+        else:
+            N = N[0] 
+        # If all vector field components are the same shape,
+        #N becomes the number of dimensions
+
+    outvals=[np.zeros(np.shape(v_c))] * 3 #Initialize output vector field
+
+    lst_axes = [1,2,2,0,0,1] # ordered list of dimension of the derivatives
+    lst_args = [2,1,0,2,1,0] # ordered list of arguments for the derivatives
+
+    for i in range(6):
+        # select the appropriate derivative and argument for each of the 6 
+        # required computations 
+        ax = lst_axes[i]
+        arg = lst_args[i]
+
+        out = np.gradient(v[arg], *varargs, axis=ax, edge_order=0)
+
+        if i == 0:
+            out0 = out
+        elif i == 1:
+            out1 = out
+        elif i == 2:
+            out2 = out
+        elif i == 3:
+            out3 = out
+        elif i == 4:
+            out4 = out
+        else:
+            out5 = out
+
+    outvals[0] = out0 - out1
+    outvals[1] = out2 - out3
+    outvals[2] = out4 - out5
+    return outvals
+
+
+def laplace(f):
+    """
+    Return the laplacian of input.
+
+    Computes the laplacian of an N-dimensional numpy array (scalar
+    field), or of a list of N N-dimensional numpy arrays (vector
+    field) as the composition of the numpy functions gradient and
+    divergence. In the case of an input vector field, the result
+    is the vector laplacian, in which a list of the laplacian of
+    each component is returned.
+    For more information on the computations, type help(gradient) or 
+    help(divergence) 
+
+    Parameters
+    ----------
+    f : array_like
+        Input array
+
+    Returns
+    -------
+    l : array_like
+        Laplacian of input. If the input is a scalar field f
+        (single numpy array) the output is (d^2(f)/dx_0^2 +
+        d^2(f)/dx_1^2 + d^2(f)/dx_2^2 + ...), where d^2(f)/dx_i^2
+        refers to the second derivative of f with respect to x_i.
+        If the input is a scalar field the output is a list of 
+        numpy arrays, being each of them the laplacian operator
+        applied to each of the input vector field components f_i
+
+    See Also
+    --------
+    gradient, divergence
+
+    Example
+    -------
+    This example returns an array that grows linearly in the X axis
+    after applying the laplacian function to the array X**3+Y**2+Z**2:
+
+    >>> X,Y,Z=np.mgrid[0:200,0:200,0:200]
+    >>> f=X**3+Y**2+Z**2
+    >>> d=laplace(f)
+    >>> plt.imshow(d[:,:,1])
+    >>> plt.colorbar()
+    >>> plt.show()
+    """
+    if type(f) == np.ndarray:
+        l = np.divergence(np.gradient(f, *varargs, edge_order=0))
+    elif type(f) == list:
+        if False in [np.shape(f[0]) == np.shape(f[i]) for i in range(len(f))]:
+            raise TypeError("All components of the input vector field "
+			    "must be the same shape")
+        else:
+            l = []
+            for i in range(len(f)):
+                laplace_comp = np.divergence(np.gradient(f[i], *varargs, edge_order=0))
+                l.append(laplace_comp)
+    else:
+        raise TypeError("Please, enter a numpy array or a list of"
+			" numpy arrays of the same shape")
+    return l
+
 
 def diff(a, n=1, axis=-1):
     """