numpy · charris · Mar 23, 2017 · Apr 25, 2014 · charris · Mar 23, 2017
diff --git a/doc/release/1.13.0-notes.rst b/doc/release/1.13.0-notes.rst
@@ -121,6 +121,16 @@ This means that, e.g., it is now possible to do the following:
     [array([[ 1. ,  1. , -0.5], [ 1. ,  1. , -0.5]]), 
      array([[ 2. ,  2. ,  2. ], [ 2. ,  1.7,  0.5]])]
 
+``np.heaviside`` computes the Heaviside function
+------------------------------------------------
+The new function ``np.heaviside(x, h0)`` (a ufunc) computes the Heaviside
+function::
+
+                       { 0   if x < 0,
+    heaviside(x, h0) = { h0  if x == 0,
+                       { 1   if x > 0.
+
+
 Improvements
 ============
 

diff --git a/numpy/core/code_generators/generate_umath.py b/numpy/core/code_generators/generate_umath.py
@@ -550,6 +550,12 @@ def english_upper(s):
           TD(ints, simd=[('avx2', ints)]),
           TD(O, f='PyNumber_Rshift'),
           ),
+'heaviside':
+    Ufunc(2, 1, None,
+          docstrings.get('numpy.core.umath.heaviside'),
+          None,
+          TD(flts, f='heavisi
8000
de', astype={'e':'f'}),
+          ),
 'degrees':
     Ufunc(1, 1, None,
           docstrings.get('numpy.core.umath.degrees'),

diff --git a/numpy/core/code_generators/ufunc_docstrings.py b/numpy/core/code_generators/ufunc_docstrings.py
@@ -915,6 +915,50 @@ def add_newdoc(place, name, doc):
 
     """)
 
+add_newdoc('numpy.core.umath', 'heaviside',
+    """
+    Compute the Heaviside step function.
+
+    The Heaviside step function is defined as::
+
+                             0   if x < 0
+        heaviside(x, h0) =  h0   if x == 0
+                             1   if x > 0
+
+    where `h0` is often taken to be 0.5, but 0 and 1 are also sometimes used.
+
+    Parameters
+    ----------
+    x : array_like
+        Input values.
+    h0 : array_like
+        The value of the function at x = 0.
+    out : ndarray, optional
+            Array into which the output is placed. Its type is preserved and it
+        must be of the right shape to hold the output. See doc.ufuncs.
+
+    Returns
+    -------
+    out : ndarray
+        The output array, element-wise Heaviside step function of `x`.
+
+    Notes
+    -----
+    .. versionadded:: 1.13.0
+
+    References
+    ----------
+    .. Wikipedia, "Heaviside step function",
+       https://en.wikipedia.org/wiki/Heaviside_step_function
+
+    Examples
+    --------
+    >>> np.heaviside([-1.5, 0, 2.0], 0.5)
+    array([ 0. ,  0.5,  1. ])
+    >>> np.heaviside([-1.5, 0, 2.0], 1)
+    array([ 0.,  1.,  1.])
+    """)
+
 add_newdoc('numpy.core.umath', 'divide',
     """
     Divide arguments element-wise.

diff --git a/numpy/core/include/numpy/npy_math.h b/numpy/core/include/numpy/npy_math.h
@@ -316,19 +316,22 @@ NPY_INPLACE double npy_rad2deg(double x);
 NPY_INPLACE double npy_logaddexp(double x, double y);
 NPY_INPLACE double npy_logaddexp2(double x, double y);
 NPY_INPLACE double npy_divmod(double x, double y, double *modulus);
+NPY_INPLACE double npy_heaviside(double x, double h0);
 
 NPY_INPLACE float npy_deg2radf(float x);
 NPY_INPLACE float npy_rad2degf(float x);
 NPY_INPLACE float npy_logaddexpf(float x, float y);
 NPY_INPLACE float npy_logaddexp2f(float x, float y);
 NPY_INPLACE float npy_divmodf(float x, float y, float *modulus);
+NPY_INPLACE float npy_heavisidef(float x, float h0);
 
 NPY_INPLACE npy_longdouble npy_deg2radl(npy_longdouble x);
 NPY_INPLACE npy_longdouble npy_rad2degl(npy_longdouble x);
 NPY_INPLACE npy_longdouble npy_logaddexpl(npy_longdouble x, npy_longdouble y);
 NPY_INPLACE npy_longdouble npy_logaddexp2l(npy_longdouble x, npy_longdouble y);
 NPY_INPLACE npy_longdouble npy_divmodl(npy_longdouble x, npy_longdouble y,
                            npy_longdouble *modulus);
+NPY_INPLACE npy_longdouble npy_heavisidel(npy_longdouble x, npy_longdouble h0);
 
 #define npy_degrees npy_rad2deg
 #define npy_degreesf npy_rad2degf

diff --git a/numpy/core/src/npymath/npy_math_internal.h.src b/numpy/core/src/npymath/npy_math_internal.h.src
@@ -541,6 +541,22 @@ NPY_INPLACE @type@ npy_cbrt@c@(@type@ x)
  * #C = F, ,L#
  */
 
+@type@ npy_heaviside@c@(@type@ x, @type@ h0)
+{
+    if (npy_isnan(x)) {
+        return (@type@) NPY_NAN;
+    }
+    else if (x == 0) {
+        return h0;
+    }
+    else if (x < 0) {
+        return (@type@) 0.0;
+    }
+    else {
+        return (@type@) 1.0;
+    }
+}
+
 #define LOGE2    NPY_LOGE2@c@
 #define LOG2E    NPY_LOG2E@c@
 #define RAD2DEG  (180.0@c@/NPY_PI@c@)

diff --git a/numpy/core/tests/test_umath.py b/numpy/core/tests/test_umath.py
@@ -1197,6 +1197,28 @@ def test_radians(self):
         assert_almost_equal(ncu.radians(-90.0), -0.5*np.pi)
 
 
+class TestHeavside(TestCase):
+    def test_heaviside(self):
+        x = np.array([[-30.0, -0.1, 0.0, 0.2], [7.5, np.nan, np.inf, -np.inf]])
+        expectedhalf = np.array([[0.0, 0.0, 0.5, 1.0], [1.0, np.nan, 1.0, 0.0]])
+        expected1 = expectedhalf.copy()
+        expected1[0, 2] = 1
+
+        h = ncu.heaviside(x, 0.5)
+        assert_equal(h, expectedhalf)
+
+        h = ncu.heaviside(x, 1.0)
+        assert_equal(h, expected1)
+
+        x = x.astype(np.float32)
+
+        h = ncu.heaviside(x, np.float32(0.5))
+        assert_equal(h, expectedhalf.astype(np.float32))
+
+        h = ncu.heaviside(x, np.float32(1.0))
+        assert_equal(h, expected1.astype(np.float32))
+
+
 class TestSign(TestCase):
     def test_sign(self):
         a = np.array([np.inf, -np.inf, np.nan, 0.0, 3.0, -3.0])