ENH: __array_function__ for np.lib, part 1

np.lib.arraypad through np.lib.nanfunctions

1 files changed, 64 insertions, 0 deletions

diff --git a/numpy/lib/financial.py b/numpy/lib/financial.py
index 06fa1bd92..d1a0cd9c0 100644
--- a/numpy/lib/financial.py
+++ b/numpy/lib/financial.py
@@ -15,6 +15,8 @@ from __future__ import division, absolute_import, print_function
 from decimal import Decimal
 
 import numpy as np
+from numpy.core.overrides import array_function_dispatch
+
 
 __all__ = ['fv', 'pmt', 'nper', 'ipmt', 'ppmt', 'pv', 'rate',
            'irr', 'npv', 'mirr']
@@ -36,6 +38,12 @@ def _convert_when(when):
     except (KeyError, TypeError):
         return [_when_to_num[x] for x in when]
 
+
+def _fv_dispatcher(rate, nper, pmt, pv, when=None):
+    return (rate, nper, pmt, pv)
+
+
+@array_function_dispatch(_fv_dispatcher)
 def fv(rate, nper, pmt, pv, when='end'):
     """
     Compute the future value.
@@ -124,6 +132,12 @@ def fv(rate, nper, pmt, pv, when='end'):
                     (1 + rate*when)*(temp - 1)/rate)
     return -(pv*temp + pmt*fact)
 
+
+def _pmt_dispatcher(rate, nper, pv, fv=None, when=None):
+    return (rate, nper, pv, fv)
+
+
+@array_function_dispatch(_pmt_dispatcher)
 def pmt(rate, nper, pv, fv=0, when='end'):
     """
     Compute the payment against loan principal plus interest.
@@ -216,6 +230,12 @@ def pmt(rate, nper, pv, fv=0, when='end'):
                     (1 + masked_rate*when)*(temp - 1)/masked_rate)
     return -(fv + pv*temp) / fact
 
+
+def _nper_dispatcher(rate, pmt, pv, fv=None, when=None):
+    return (rate, pmt, pv, fv)
+
+
+@array_function_dispatch(_nper_dispatcher)
 def nper(rate, pmt, pv, fv=0, when='end'):
     """
     Compute the number of periodic payments.
@@ -284,6 +304,12 @@ def nper(rate, pmt, pv, fv=0, when='end'):
         B = np.log((-fv+z) / (pv+z))/np.log(1+rate)
         return np.where(rate == 0, A, B)
 
+
+def _ipmt_dispatcher(rate, per, nper, pv, fv=None, when=None):
+    return (rate, per, nper, pv, fv)
+
+
+@array_function_dispatch(_ipmt_dispatcher)
 def ipmt(rate, per, nper, pv, fv=0, when='end'):
     """
     Compute the interest portion of a payment.
@@ -379,6 +405,7 @@ def ipmt(rate, per, nper, pv, fv=0, when='end'):
         pass
     return ipmt
 
+
 def _rbl(rate, per, pmt, pv, when):
     """
     This function is here to simply have a different name for the 'fv'
@@ -388,6 +415,12 @@ def _rbl(rate, per, pmt, pv, when):
     """
     return fv(rate, (per - 1), pmt, pv, when)
 
+
+def _ppmt_dispatcher(rate, per, nper, pv, fv=None, when=None):
+    return (rate, per, nper, pv, fv)
+
+
+@array_function_dispatch(_ppmt_dispatcher)
 def ppmt(rate, per, nper, pv, fv=0, when='end'):
     """
     Compute the payment against loan principal.
@@ -416,6 +449,12 @@ def ppmt(rate, per, nper, pv, fv=0, when='end'):
     total = pmt(rate, nper, pv, fv, when)
     return total - ipmt(rate, per, nper, pv, fv, when)
 
+
+def _pv_dispatcher(rate, nper, pmt, fv=None, when=None):
+    return (rate, nper, nper, pv, fv)
+
+
+@array_function_dispatch(_pv_dispatcher)
 def pv(rate, nper, pmt, fv=0, when='end'):
     """
     Compute the present value.
@@ -520,6 +559,12 @@ def _g_div_gp(r, n, p, x, y, w):
                 (n*t2*x - p*(t1 - 1)*(r*w + 1)/(r**2) + n*p*t2*(r*w + 1)/r +
                  p*(t1 - 1)*w/r))
 
+
+def _rate_dispatcher(nper, pmt, pv, fv, when=None, guess=None, tol=None,
+                     maxiter=None):
+    return (nper, pmt, pv, fv)
+
+
 # Use Newton's iteration until the change is less than 1e-6
 #  for all values or a maximum of 100 iterations is reached.
 #  Newton's rule is
@@ -527,6 +572,7 @@ def _g_div_gp(r, n, p, x, y, w):
 #     where
 #  g(r) is the formula
 #  g'(r) is the derivative with respect to r.
+@array_function_dispatch(_rate_dispatcher)
 def rate(nper, pmt, pv, fv, when='end', guess=None, tol=None, maxiter=100):
     """
     Compute the rate of interest per period.
@@ -598,6 +644,12 @@ def rate(nper, pmt, pv, fv, when='end', guess=None, tol=None, maxiter=100):
     else:
         return rn
 
+
+def _irr_dispatcher(values):
+    return (values,)
+
+
+@array_function_dispatch(_irr_dispatcher)
 def irr(values):
     """
     Return the Internal Rate of Return (IRR).
@@ -677,6 +729,12 @@ def irr(values):
     rate = rate.item(np.argmin(np.abs(rate)))
     return rate
 
+
+def _npv_dispatcher(rate, values):
+    return (values,)
+
+
+@array_function_dispatch(_npv_dispatcher)
 def npv(rate, values):
     """
     Returns the NPV (Net Present Value) of a cash flow series.
@@ -722,6 +780,12 @@ def npv(rate, values):
     values = np.asarray(values)
     return (values / (1+rate)**np.arange(0, len(values))).sum(axis=0)
 
+
+def _mirr_dispatcher(values, finance_rate, reinvest_rate):
+    return (values,)
+
+
+@array_function_dispatch(_mirr_dispatcher)
 def mirr(values, finance_rate, reinvest_rate):
     """
     Modified internal rate of return.