Merge pull request #26 from person142/ipmt-cleanups

MAINT: small cleanups to `ipmt`

Author: Kai-Striega

numpy_financial/_financial.py

@@ -417,21 +417,26 @@ def ipmt(rate, per, nper, pv, fv=0, when='end'):
                                                         pv, fv, when)
 
     total_pmt = pmt(rate, nper, pv, fv, when)
-    ipmt = np.array(_rbl(rate, per, total_pmt, pv, when) * rate)
+    ipmt_array = np.array(_rbl(rate, per, total_pmt, pv, when) * rate)
 
     # Payments start at the first period, so payments before that
     # don't make any sense.
-    ipmt[per < 1] = _value_like(ipmt, np.nan)
+    ipmt_array[per < 1] = _value_like(ipmt_array, np.nan)
     # If payments occur at the beginning of a period and this is the
     # first period, then no interest has accrued.
     per1_and_begin = (when == 1) & (per == 1)
-    ipmt[per1_and_begin] = _value_like(ipmt, 0)
+    ipmt_array[per1_and_begin] = _value_like(ipmt_array, 0)
     # If paying at the beginning we need to discount by one period.
     per_gt_1_and_begin = (when == 1) & (per > 1)
-    ipmt[per_gt_1_and_begin] = (
-        ipmt[per_gt_1_and_begin] / (1 + rate[per_gt_1_and_begin])
+    ipmt_array[per_gt_1_and_begin] = (
+        ipmt_array[per_gt_1_and_begin] / (1 + rate[per_gt_1_and_begin])
     )
-    return ipmt
+
+    if np.ndim(ipmt_array) == 0:
+        # Follow the ufunc convention of returning scalars for scalar
+        # and 0d array inputs.
+        return ipmt_array.item(0)
+    return ipmt_array
 
 
 def _rbl(rate, per, pmt, pv, when):

numpy_financial/tests/test_financial.py

@@ -265,27 +265,6 @@ def test_decimal_with_when(self):
                               Decimal('60'), Decimal('55000'), Decimal('0'),
                               'end'))
 
-        # begin
-        assert_equal(npf.ipmt(Decimal('0.1') / Decimal('12'), Decimal('1'),
-                              Decimal('24'), Decimal('2000'),
-                              Decimal('0'), Decimal('1')).flat[0],
-                     npf.ipmt(Decimal('0.1') / Decimal('12'), Decimal('1'),
-                              Decimal('24'), Decimal('2000'),
-                              Decimal('0'), 'begin').flat[0])
-        # end
-        assert_equal(npf.ipmt(Decimal('0.1') / Decimal('12'), Decimal('1'),
-                              Decimal('24'), Decimal('2000'),
-                              Decimal('0')).flat[0],
-                     npf.ipmt(Decimal('0.1') / Decimal('12'), Decimal('1'),
-                              Decimal('24'), Decimal('2000'),
-                              Decimal('0'), 'end').flat[0])
-        assert_equal(npf.ipmt(Decimal('0.1') / Decimal('12'), Decimal('1'),
-                              Decimal('24'), Decimal('2000'),
-                              Decimal('0'), Decimal('0')).flat[0],
-                     npf.ipmt(Decimal('0.1') / Decimal('12'), Decimal('1'),
-                              Decimal('24'), Decimal('2000'),
-                              Decimal('0'), 'end').flat[0])
-
     def test_broadcast(self):
         assert_almost_equal(npf.nper(0.075, -2000, 0, 100000., [0, 1]),
                             [21.5449442, 20.76156441], 4)
@@ -374,6 +353,33 @@ def test_when_is_end(self, when):
             result = npf.ipmt(0.1 / 12, 1, 24, 2000, 0, when)
         assert_allclose(result, -16.666667, rtol=1e-6)
 
+
+    @pytest.mark.parametrize('when', [Decimal('1'), 'begin'])
+    def test_when_is_begin_decimal(self, when):
+        result = npf.ipmt(
+            Decimal('0.1') / Decimal('12'),
+            Decimal('1'),
+            Decimal('24'),
+            Decimal('2000'),
+            Decimal('0'),
+            when,
+        )
+        assert result == 0
+
+    @pytest.mark.parametrize('when', [None, Decimal('0'), 'end'])
+    def test_when_is_end_decimal(self, when):
+        # Computed using Google Sheet's IPMT
+        desired = Decimal('-16.666667')
+        args = (
+            Decimal('0.1') / Decimal('12'),
+            Decimal('1'),
+            Decimal('24'),
+            Decimal('2000'),
+            Decimal('0')
+        )
+        result = npf.ipmt(*args) if when is None else npf.ipmt(*args, when)
+        assert_almost_equal(result, desired, decimal=5)
+
     @pytest.mark.parametrize('per, desired', [
         (0, numpy.nan),
         (1, 0),
@@ -418,6 +424,14 @@ def test_decimal_broadcasting(self):
         )
         assert_almost_equal(result, desired, decimal=4)
 
+    def test_0d_inputs(self):
+        args = (0.1 / 12, 1, 24, 2000)
+        # Scalar inputs should return a scalar.
+        assert numpy.isscalar(npf.ipmt(*args))
+        args = (numpy.array(args[0]),) + args[1:]
+        # 0d array inputs should return a scalar.
+        assert numpy.isscalar(npf.ipmt(*args))
+
 
 class TestFv:
     def test_float(self):