REF/ENH: robust norms dtype, plus pep-8, refactor for consistency

Author: josef-pkt

statsmodels/robust/norms.py

@@ -3,6 +3,16 @@
 # TODO: add plots to weighting functions for online docs.
 
 
+def _cabs(x):
+    """absolute value function that changes complex sign based on real sign
+
+    This could be useful for complex step derivatives of functions that
+    need abs. Not yet used.
+    """
+    sign = (x.real >= 0) * 2 - 1
+    return sign * x
+
+
 class RobustNorm:
     """
     The parent class for the norms used for robust regression.
@@ -627,23 +637,25 @@ def rho(self, z):
 
             rho(z) = a*(b + c - a)                  for \|z\| > c
         """
+        a, b, c = self.a, self.b, self.c
 
-        z = np.abs(z)
         z_isscalar = np.isscalar(z)
         z = np.atleast_1d(z)
-        a = self.a; b = self.b; c = self.c
+
         t1, t2, t3 = self._subset(z)
-        t34 = ~(t1 | t2 )
-        v = np.zeros(z.shape, float)
+        t34 = ~(t1 | t2)
+        dt = np.promote_types(z.dtype, "float")
+        v = np.zeros(z.shape, dtype=dt)
+        z = np.abs(z)
         v[t1] = z[t1]**2 * 0.5
         v[t2] = (a * (z[t2] - a) + a**2 * 0.5)
-        v[t3] = a * (c - z[t3])**2  / (c - b) * (-0.5)
+        v[t3] = a * (c - z[t3])**2 / (c - b) * (-0.5)
         v[t34] += a * (b + c - a) * 0.5
 
         if z_isscalar:
-            return v[0]
-        else:
-            return v
+            v = v[0]
+
+        return v
 
     def psi(self, z):
         r"""
@@ -667,14 +679,26 @@ def psi(self, z):
 
             psi(z) = 0                            for \|z\| > c
         """
-        z = np.asarray(z)
-        a = self.a; b = self.b; c = self.c
+        a, b, c = self.a, self.b, self.c
+
+        z_isscalar = np.isscalar(z)
+        z = np.atleast_1d(z)
+
         t1, t2, t3 = self._subset(z)
+        dt = np.promote_types(z.dtype, "float")
+        v = np.zeros(z.shape, dtype=dt)
         s = np.sign(z)
         z = np.abs(z)
-        v =  (t1 * z*s +
-                 t2 * a*s +
-                 t3 * a*s * (c - z) / (c - b))
+
+        v[t1] = z[t1] * s[t1]
+        v[t2] = a * s[t2]
+        v[t3] = a * s[t3] * (c - z[t3]) / (c - b)
+        # v = (t1 * z*s +
+        #          t2 * a*s +
+        #          t3 * a*s * (c - z) / (c - b))
+
+        if z_isscalar:
+            v = v[0]
         return v
 
     def weights(self, z):
@@ -699,32 +723,43 @@ def weights(self, z):
 
             weights(z) = 0                            for \|z\| > c
         """
-        z = np.asarray(z)
-        a = self.a
-        b = self.b
-        c = self.c
+        a, b, c = self.a, self.b, self.c
+
+        z_isscalar = np.isscalar(z)
+        z = np.atleast_1d(z)
+
         t1, t2, t3 = self._subset(z)
 
-        v = np.zeros_like(z)
+        dt = np.promote_types(z.dtype, "float")
+        v = np.zeros(z.shape, dtype=dt)
         v[t1] = 1.0
         abs_z = np.abs(z)
         v[t2] = a / abs_z[t2]
         abs_zt3 = abs_z[t3]
         v[t3] = a * (c - abs_zt3) / (abs_zt3 * (c - b))
-        v[np.where(np.isnan(v))] = 1.  # TODO: for some reason 0 returns a nan?
+
+        if z_isscalar:
+            v = v[0]
         return v
 
     def psi_deriv(self, z):
         """Derivative of psi function, second derivative of rho function.
         """
-        t1, _, t3 = self._subset(z)
         a, b, c = self.a, self.b, self.c
+
+        z_isscalar = np.isscalar(z)
         z = np.atleast_1d(z)
-        # default is t1
-        d = np.zeros_like(z)
+
+        t1, _, t3 = self._subset(z)
+
+        dt = np.promote_types(z.dtype, "float")
+        d = np.zeros(z.shape, dtype=dt)
         d[t1] = 1.0
         zt3 = z[t3]
         d[t3] = -(a * np.sign(zt3) * zt3) / (np.abs(zt3) * (c - b))
+
+        if z_isscalar:
+            d = d[0]
         return d
 
 

statsmodels/robust/tests/results/results_norms.py

@@ -11,4 +11,4 @@
                -0.3333333333333333, 0.0],
     weights=[0.0, 0.1111111111111111, 0.75, 1.0, 1.0, 1.0, 0.75,
              0.1111111111111111, 0.0],
-    )
+    )

statsmodels/robust/tests/test_norms.py

@@ -4,19 +4,54 @@
 from numpy.testing import assert_allclose
 
 from statsmodels.robust import norms
+from statsmodels.tools.numdiff import (
+    _approx_fprime_scalar,
+    # _approx_fprime_cs_scalar,  # not yet
+    )
 from .results import results_norms as res_r
 
 cases = [
     (norms.Hampel, (1.5, 3.5, 8.), res_r.res_hampel)
     ]
 
+dtypes = ["int", np.float64, np.complex128]
+
+
+@pytest.mark.parametrize("dtype", dtypes)
 @pytest.mark.parametrize("case", cases)
-def test_norm(case):
+def test_norm(case, dtype):
     ncls, args, res = case
     norm = ncls(*args)
-    x = np.array([-9., -6, -2, -1, 0, 1, 2, 6, 9])
+    x = np.array([-9, -6, -2, -1, 0, 1, 2, 6, 9], dtype=dtype)
+
+    weights = norm.weights(x)
+    rho = norm.rho(x)
+    psi = norm.psi(x)
+    psi_deriv = norm.psi_deriv(x)
+    assert_allclose(weights, res.weights, rtol=1e-12, atol=1e-20)
+    assert_allclose(rho, res.rho, rtol=1e-12, atol=1e-20)
+    assert_allclose(psi, res.psi, rtol=1e-12, atol=1e-20)
+    assert_allclose(psi_deriv, res.psi_deriv, rtol=1e-12, atol=1e-20)
+
+    dtype2 = np.promote_types(dtype, "float")
+    assert weights.dtype == dtype2
+    assert rho.dtype == dtype2
+    assert psi.dtype == dtype2
+    assert psi_deriv.dtype == dtype2
+
+    psid = _approx_fprime_scalar(x, norm.rho)
+    assert_allclose(psid, res.psi, rtol=1e-6, atol=1e-8)
+    psidd = _approx_fprime_scalar(x, norm.psi)
+    assert_allclose(psidd, res.psi_deriv, rtol=1e-6, atol=1e-8)
+
+    # complex step derivatives are not yet supported if method uses np.abs
+    # psid = _approx_fprime_cs_scalar(x, norm.rho)
+    # assert_allclose(psid, res.psi, rtol=1e-12, atol=1e-20)
+    # psidd = _approx_fprime_cs_scalar(x, norm.psi)
+    # assert_allclose(psidd, res.psi_deriv, rtol=1e-12, atol=1e-20)
 
-    assert_allclose(norm.weights(x), res.weights, rtol=1e-12, atol=1e-20)
-    assert_allclose(norm.rho(x), res.rho, rtol=1e-12, atol=1e-20)
-    assert_allclose(norm.psi(x), res.psi, rtol=1e-12, atol=1e-20)
-    assert_allclose(norm.psi_deriv(x), res.psi_deriv, rtol=1e-12, atol=1e-20)
+    # check scalar value
+    methods = ["weights", "rho", "psi", "psi_deriv"]
+    for meth in methods:
+        resm = [getattr(norm, meth)(xi) for xi in x]
+        assert_allclose(resm, getattr(res, meth))