BENCH: add scalar and array pmt benchmarks

benchmarks/benchmarks.py

@@ -3,6 +3,27 @@
 import numpy_financial as npf
 
 
+class Pmt:
+    """Compare scalar overhead with array throughput for loan payments."""
+
+    param_names = ["size", "rate", "when"]
+    params = [["scalar", 1, 10, 1000], [0.0, 0.05], ["end", "begin"]]
+
+    def setup(self, size, rate, when):
+        self.rate = rate
+        self.nper = 120
+        self.pv = 1_000_000.0
+        self.fv = -100_000.0
+        if size != "scalar":
+            self.rate = np.full(size, rate)
+            self.nper = np.full(size, self.nper)
+            self.pv = np.full(size, self.pv)
+        self.when = when
+
+    def time_pmt(self, size, rate, when):
+        npf.pmt(self.rate, self.nper, self.pv, self.fv, self.when)
+
+
 class Npv2D:
 
     param_names = ["n_cashflows", "cashflow_lengths", "rates_lengths"]

numpy_financial/_financial.py

@@ -13,7 +13,17 @@
 
 from collections.abc import Iterable, Mapping, Sequence
 from decimal import Decimal
-from typing import Any, Callable, Final, Literal, Protocol, TypeAlias, TypeVar, overload
+from typing import (
+    Any,
+    Callable,
+    Final,
+    Literal,
+    Protocol,
+    TypeAlias,
+    TypeVar,
+    cast,
+    overload,
+)
 
 import numpy as np
 import numpy.typing as npt
@@ -379,12 +389,30 @@ def pmt(rate, nper, pv, fv: Any = 0, when: _When = 'end'):
 
     """
     when = _convert_when(when)
-    (rate, nper, pv, fv, when) = map(np.array, [rate, nper, pv, fv, when])
+    scalar = (type(rate) in (float, np.float64)
+            and type(when) is int
+            and all(type(value) in (int, float, np.float64)
+                    and (type(value) is not int or -2**63 <= value < 2**63)
+                    for value in (nper, pv, fv)))
+    if scalar:
+        # Keep NumPy floating-point error handling, without allocating arrays
+        # for ordinary scalar loans. Other dtypes retain the general path.
+        rate, nper, pv, fv = map(np.float64, (rate, nper, pv, fv))
+        when = cast(int, when)
+    else:
+        (rate, nper, pv, fv, when) = map(np.array, [rate, nper, pv, fv, when])
     temp = (1 + rate) ** nper
     mask = (rate == 0)
-    masked_rate = np.where(mask, 1, rate)
-    fact = np.where(mask != 0, nper,
-                    (1 + masked_rate * when) * (temp - 1) / masked_rate)
+    if scalar:
+        masked_rate = np.float64(1) if mask else rate
+    else:
+        masked_rate = np.where(mask, 1, rate)
+    fact = (1 + masked_rate * when) * (temp - 1) / masked_rate
+    if scalar:
+        if mask:
+            fact = nper
+    else:
+        fact = np.where(mask, nper, fact)
     return -(fv + pv * temp) / fact
 
 @overload

numpy_financial/tests/test_pmt_scalar.py

@@ -0,0 +1,34 @@
+import numpy as np
+import pytest
+from numpy.testing import assert_allclose
+
+import numpy_financial as npf
+
+
+@pytest.mark.parametrize("rate", [0.0, 0.05, -0.05, -1.0, -1.5, 1e-12,
+                                  np.inf, np.nan])
+@pytest.mark.parametrize("nper", [0, -3, 0.5, 120, 10000])
+@pytest.mark.parametrize("when", ["begin", "end"])
+def test_scalar_matches_array_path(rate, nper, when):
+    # A zero-dimensional rate array exercises the general implementation.
+    with np.errstate(all="ignore"):
+        expected = npf.pmt(np.array(rate), nper, 1000, -100, when)
+        actual = npf.pmt(rate, nper, 1000, -100, when)
+    assert type(actual) is type(expected)
+    assert_allclose(actual, expected, rtol=1e-13)
+
+
+@pytest.mark.parametrize("args", [(0.0, 0, 100), (-1.5, 0.5, 100),
+                                  (10.0, 10000, 100)])
+def test_scalar_respects_seterr(args):
+    with np.errstate(all="raise"):
+        with pytest.raises(FloatingPointError):
+            npf.pmt(*args)
+
+
+@pytest.mark.parametrize("value", [np.float32(100), np.int64(100), 2**64])
+def test_scalar_other_types(value):
+    expected = npf.pmt(np.array(0.05), 10, value)
+    actual = npf.pmt(0.05, 10, value)
+    assert type(actual) is type(expected)
+    assert_allclose(actual, expected)