Convert one-bracket rate scales without UnboundLocalError

changelog.d/fix-singleton-tax-scale-conversion.fixed.md

@@ -0,0 +1 @@
+Converting a one-bracket rate scale no longer raises `UnboundLocalError`, zero-origin marginal singletons preserve their tax calculations, and converting an empty linear-average scale gives an empty marginal scale.

policyengine_core/taxscales/linear_average_rate_tax_scale.py

@@ -59,10 +59,15 @@ def to_marginal(self) -> MarginalRateTaxScale:
             unit=self.unit,
         )
 
+        if len(self.rates) == 0:
+            return marginal_tax_scale
+
         previous_i = 0
         previous_threshold = 0
+        terminal_rate = self.rates[0]
 
         for threshold, rate in zip(self.thresholds[1:], self.rates[1:]):
+            terminal_rate = rate
             if threshold != float("Inf"):
                 i = rate * threshold
                 marginal_tax_scale.add_bracket(
@@ -72,6 +77,6 @@ def to_marginal(self) -> MarginalRateTaxScale:
                 previous_i = i
                 previous_threshold = threshold
 
-        marginal_tax_scale.add_bracket(previous_threshold, rate)
+        marginal_tax_scale.add_bracket(previous_threshold, terminal_rate)
 
         return marginal_tax_scale

policyengine_core/taxscales/marginal_rate_tax_scale.py

@@ -267,7 +267,12 @@ def to_average(self) -> taxscales.LinearAverageRateTaxScale:
             unit=self.unit,
         )
 
-        average_tax_scale.add_bracket(0, 0)
+        if len(self.thresholds) == len(self.rates) == 1 and self.thresholds[0] == 0:
+            # A zero-origin singleton has a constant average rate. Starting
+            # at a zero rate and interpolating toward infinity would levy no tax.
+            average_tax_scale.add_bracket(0, self.rates[0])
+        else:
+            average_tax_scale.add_bracket(0, 0)
 
         if self.thresholds:
             i = 0
@@ -284,6 +289,7 @@ def to_average(self) -> taxscales.LinearAverageRateTaxScale:
                 previous_threshold = threshold
                 previous_rate = rate
 
-            average_tax_scale.add_bracket(float("Inf"), rate)
+            # Far above the last threshold the average rate approaches the top rate.
+            average_tax_scale.add_bracket(float("Inf"), previous_rate)
 
         return average_tax_scale

tests/core/tax_scales/test_linear_average_rate_tax_scale.py

@@ -88,3 +88,55 @@ def test_to_marginal():
         [0.1, 0.25, 0.4, 0.5],
         absolute_error_margin=1e-10,
     )
+
+
+@pytest.mark.parametrize("rate", [0, 0.37])
+def test_to_marginal_of_a_single_bracket(rate):
+    # One bracket used to raise UnboundLocalError: the top rate was read from
+    # a loop variable that only the second and later brackets set.
+    tax_base = numpy.array([0, 1, 1.5, 1_000])
+    tax_scale = taxscales.LinearAverageRateTaxScale()
+    tax_scale.add_bracket(0, rate)
+
+    result = tax_scale.to_marginal()
+
+    assert result.thresholds == [0]
+    assert result.rates == [rate]
+    tools.assert_near(
+        result.calc(tax_base), tax_scale.calc(tax_base), absolute_error_margin=1e-10
+    )
+
+
+@pytest.mark.parametrize("rates_type", [list, numpy.array])
+def test_to_marginal_uses_last_paired_rate(rates_type):
+    tax_scale = taxscales.LinearAverageRateTaxScale()
+    tax_scale.thresholds = [0, 100]
+    tax_scale.rates = rates_type([0.1, 0.2, 0.8])
+
+    result = tax_scale.to_marginal()
+
+    assert result.thresholds == [0, 100]
+    # Tax at 100 is 100 * 0.2 = 20, so the first marginal rate is 20 / 100.
+    # The terminal rate is the last paired rate 0.2, ignoring surplus 0.8.
+    assert result.rates == [0.2, 0.2]
+
+
+@pytest.mark.parametrize("rates_type", [list, numpy.array])
+def test_to_marginal_uses_last_paired_rate_at_infinity(rates_type):
+    tax_scale = taxscales.LinearAverageRateTaxScale()
+    tax_scale.thresholds = [0, 100, numpy.inf]
+    tax_scale.rates = rates_type([0.1, 0.2, 0.3, 0.8])
+
+    result = tax_scale.to_marginal()
+
+    assert result.thresholds == [0, 100]
+    # The finite knot gives 100 * 0.2 / 100 = 0.2; the infinity knot
+    # supplies the terminal rate 0.3 even though its tax is not integrated.
+    assert result.rates == [0.2, 0.3]
+
+
+def test_to_marginal_of_no_brackets():
+    result = taxscales.LinearAverageRateTaxScale().to_marginal()
+
+    assert result.thresholds == []
+    assert result.rates == []

tests/core/tax_scales/test_marginal_rate_tax_scale.py

@@ -222,6 +222,42 @@ def test_to_average():
     )
 
 
+@pytest.mark.parametrize("rate", [0, 0.37])
+def test_to_average_of_a_single_bracket(rate):
+    # One bracket used to raise UnboundLocalError: the top rate was read from
+    # a loop variable that only the second and later brackets set.
+    tax_scale = taxscales.MarginalRateTaxScale()
+    tax_scale.add_bracket(0, rate)
+
+    result = tax_scale.to_average()
+
+    # For rate 0.37: 1 * 0.37 = 0.37, 1.5 * 0.37 = 0.555,
+    # and 1_000 * 0.37 = 370. Negative bases owe no marginal tax.
+    tax_base = numpy.array([-1, 0, 1, 1.5, 1_000])
+    tools.assert_near(
+        result.calc(tax_base),
+        [0, 0, rate, 1.5 * rate, 1_000 * rate],
+        absolute_error_margin=1e-10,
+    )
+    assert result.thresholds == [0, numpy.inf]
+    assert result.rates == [rate, rate]
+    assert result.to_marginal().rates == [rate]
+
+
+@pytest.mark.parametrize("rates_type", [list, numpy.array])
+def test_to_average_uses_last_paired_rate(rates_type):
+    tax_scale = taxscales.MarginalRateTaxScale()
+    tax_scale.thresholds = [0, 100]
+    tax_scale.rates = rates_type([0.1, 0.2, 0.8])
+
+    result = tax_scale.to_average()
+
+    assert result.thresholds == [0, 100, numpy.inf]
+    # Tax at 100 is 100 * 0.1 = 10, so its average rate is 10 / 100 = 0.1.
+    # The terminal marginal rate paired with threshold 100 is 0.2, not 0.8.
+    assert result.rates == [0, 0.1, 0.2]
+
+
 def test_rate_from_bracket_indice():
     tax_base = numpy.array([0, 1_000, 1_500, 50_000])
     tax_scale = taxscales.MarginalRateTaxScale()

tests/core/tax_scales/test_rate_scale_conversion_property.py

@@ -0,0 +1,165 @@
+"""Properties of converting between marginal and linear-average rate scales.
+
+For equal-length scales with 1 to 6 brackets, rates in [0, 1], first threshold
+0, and strictly increasing positive integer upper thresholds up to 1,000,000:
+
+* ``to_average`` and then ``to_marginal`` give back the same thresholds and
+  rates (one bracket included; it used to raise ``UnboundLocalError``);
+* the average scale's last bracket starts at infinity with the top rate;
+* a singleton marginal conversion preserves tax at finite bases within
+  [-1,000,000, 1,000,000].
+
+For such a linear-average rate scale, ``to_marginal``
+gives a marginal scale that levies threshold x rate at every threshold (the
+average rate there), and the top rate on the whole base above the last one.
+Examples are in ``test_marginal_rate_tax_scale.py`` and
+``test_linear_average_rate_tax_scale.py``.
+
+Appending surplus rates to a scale with at least two brackets leaves both
+conversion results unchanged, preserving the last paired terminal rate.
+"""
+
+from __future__ import annotations
+
+import numpy
+import pytest
+
+# The smoke job installs Core without the dev extra but collects every module.
+pytest.importorskip("hypothesis")
+
+from hypothesis import given, settings  # noqa: E402
+from hypothesis import strategies as st  # noqa: E402
+
+from policyengine_core import taxscales  # noqa: E402
+
+
+@st.composite
+def _brackets(draw, min_upper_thresholds=0):
+    upper_thresholds = draw(
+        st.lists(
+            st.integers(min_value=1, max_value=1_000_000),
+            min_size=min_upper_thresholds,
+            max_size=5,
+            unique=True,
+        )
+    )
+    thresholds = [0] + sorted(upper_thresholds)
+    rates = draw(
+        st.lists(
+            st.integers(min_value=0, max_value=1_000).map(lambda r: r / 1_000),
+            min_size=len(thresholds),
+            max_size=len(thresholds),
+        )
+    )
+    return thresholds, rates
+
+
+@settings(max_examples=300, deadline=None)
+@given(_brackets())
+def test_marginal_to_average_round_trips(brackets):
+    thresholds, rates = brackets
+    marginal = taxscales.MarginalRateTaxScale()
+    for threshold, rate in zip(thresholds, rates):
+        marginal.add_bracket(threshold, rate)
+
+    average = marginal.to_average()
+    round_trip = average.to_marginal()
+
+    assert average.thresholds[-1] == numpy.inf
+    assert average.rates[-1] == rates[-1]
+    assert round_trip.thresholds == thresholds
+    numpy.testing.assert_allclose(round_trip.rates, rates, rtol=0, atol=1e-9)
+
+
+@settings(max_examples=300, deadline=None)
+@given(
+    rate=st.integers(min_value=0, max_value=1_000).map(lambda r: r / 1_000),
+    bases=st.lists(
+        st.floats(
+            min_value=-1_000_000,
+            max_value=1_000_000,
+            allow_nan=False,
+            allow_infinity=False,
+        ),
+        min_size=1,
+        max_size=10,
+    ),
+)
+def test_singleton_marginal_to_average_preserves_calculation(rate, bases):
+    marginal = taxscales.MarginalRateTaxScale()
+    marginal.add_bracket(0, rate)
+    tax_base = numpy.array([-1_000, -1, 0, 1, 1.5, 1_000, *bases])
+
+    average = marginal.to_average()
+
+    # One zero-origin marginal bracket taxes max(base, 0) * rate.
+    expected = numpy.maximum(tax_base, 0) * rate
+    numpy.testing.assert_allclose(average.calc(tax_base), expected, rtol=1e-12)
+    numpy.testing.assert_allclose(
+        average.calc(tax_base), marginal.calc(tax_base), rtol=1e-12
+    )
+
+
+@settings(max_examples=300, deadline=None)
+@given(_brackets())
+def test_average_to_marginal_levies_the_average_rate(brackets):
+    thresholds, rates = brackets
+    average = taxscales.LinearAverageRateTaxScale()
+    for threshold, rate in zip(thresholds, rates):
+        average.add_bracket(threshold, rate)
+
+    marginal = average.to_marginal()
+
+    # A linear-average scale's rate at each threshold is the average rate
+    # there, so the marginal scale levies threshold x rate at every
+    # threshold, and the top rate on the whole base from the last one up.
+    tax_base = numpy.array(thresholds, dtype=float)
+    above = 2.0 * thresholds[-1] + 1
+    # ``MarginalRateTaxScale.calc`` shifts each threshold t by t x eps (it
+    # scales them by 1 + eps), which moves the tax by about that much times
+    # the bracket's rate. Close thresholds near a million can make that rate
+    # large, so bound the rounding by the inputs rather than a fixed 1e-6.
+    eps = numpy.finfo(numpy.float64).eps
+    marginal_rates = numpy.abs(numpy.diff(tax_base * numpy.array(rates))) / (
+        numpy.diff(tax_base)
+    )
+    rounding = 8 * eps * ((marginal_rates * tax_base[1:]).sum() + above)
+    numpy.testing.assert_allclose(
+        marginal.calc(numpy.append(tax_base, above)),
+        numpy.append(tax_base * numpy.array(rates), above * rates[-1]),
+        rtol=1e-9,
+        atol=1e-6 + rounding,
+    )
+
+
+@pytest.mark.parametrize(
+    "scale_type, conversion",
+    [
+        (taxscales.MarginalRateTaxScale, "to_average"),
+        (taxscales.LinearAverageRateTaxScale, "to_marginal"),
+    ],
+)
+@settings(max_examples=300, deadline=None)
+@given(
+    brackets=_brackets(min_upper_thresholds=1),
+    surplus_rates=st.lists(
+        st.integers(min_value=0, max_value=1_000).map(lambda r: r / 1_000),
+        min_size=1,
+        max_size=3,
+    ),
+)
+def test_conversions_ignore_surplus_rates(
+    scale_type, conversion, brackets, surplus_rates
+):
+    thresholds, rates = brackets
+    paired = scale_type()
+    for threshold, rate in zip(thresholds, rates):
+        paired.add_bracket(threshold, rate)
+    surplus = paired.copy()
+    surplus.rates.extend(surplus_rates)
+
+    expected = getattr(paired, conversion)()
+    result = getattr(surplus, conversion)()
+
+    assert result.thresholds == expected.thresholds
+    assert result.rates == expected.rates