microdf-python 0.4.4__py3-none-any.whl → 0.4.6__py3-none-any.whl

microdf/tests/test_generic.py

@@ -0,0 +1,232 @@
+import numpy as np
+import pandas as pd
+
+import microdf as mdf
+from microdf.generic import MicroDataFrame, MicroSeries
+
+
+def test_df_init() -> None:
+    arr = np.array([0, 1, 1])
+    w = np.array([3, 0, 9])
+    df = mdf.MicroDataFrame({"a": arr}, weights=w)
+    assert df.a.mean() == np.average(arr, weights=w)
+
+    df = mdf.MicroDataFrame()
+    df["a"] = arr
+    df.set_weights(w)
+    assert df.a.mean() == np.average(arr, weights=w)
+
+    df = mdf.MicroDataFrame()
+    df["a"] = arr
+    df["w"] = w
+    df.set_weight_col("w")
+    assert df.a.mean() == np.average(arr, weights=w)
+
+
+def test_series_getitem() -> None:
+    arr = np.array([0, 1, 1])
+    w = np.array([3, 0, 9])
+    s = mdf.MicroSeries(arr, weights=w)
+    assert s[[1, 2]].sum() == np.sum(arr[[1, 2]] * w[[1, 2]])
+
+    assert s[1:3].sum() == np.sum(arr[1:3] * w[1:3])
+
+
+def test_sum() -> None:
+    arr = np.array([0, 1, 1])
+    w = np.array([3, 0, 9])
+    series = mdf.MicroSeries(arr, weights=w)
+    assert series.sum() == (arr * w).sum()
+
+    arr = np.linspace(-20, 100, 100)
+    w = np.linspace(1, 3, 100)
+    series = mdf.MicroSeries(arr)
+    series.set_weights(w)
+    assert series.sum() == (arr * w).sum()
+
+    # Verify that an error is thrown when passing weights of different size
+    # from the values.
+    w = np.linspace(1, 3, 101)
+    series = mdf.MicroSeries(arr)
+    try:
+        series.set_weights(w)
+        assert False
+    except Exception:
+        pass
+
+
+def test_mean() -> None:
+    arr = np.array([3, 0, 2])
+    w = np.array([4, 1, 1])
+    series = mdf.MicroSeries(arr, weights=w)
+    assert series.mean() == np.average(arr, weights=w)
+
+    arr = np.linspace(-20, 100, 100)
+    w = np.linspace(1, 3, 100)
+    series = mdf.MicroSeries(arr)
+    series.set_weights(w)
+    assert series.mean() == np.average(arr, weights=w)
+
+    w = np.linspace(1, 3, 101)
+    series = mdf.MicroSeries(arr)
+    try:
+        series.set_weights(w)
+        assert False
+    except Exception:
+        pass
+
+
+def test_poverty_count() -> None:
+    arr = np.array([10000, 20000, 50000])
+    w = np.array([1123, 1144, 2211])
+    df = MicroDataFrame(weights=w)
+    df["income"] = arr
+    df["threshold"] = 16000
+    assert df.poverty_count("income", "threshold") == w[0]
+
+
+def test_median() -> None:
+    # 1, 2, 3, 4, *4*, 4, 5, 5, 5
+    arr = np.array([1, 2, 3, 4, 5])
+    w = np.array([1, 1, 1, 3, 3])
+    series = mdf.MicroSeries(arr, weights=w)
+    assert series.median() == 4
+
+
+def test_unweighted_groupby() -> None:
+    df = mdf.MicroDataFrame({"x": [1, 2], "y": [3, 4], "z": [5, 6]})
+    assert (df.groupby("x").z.sum().values == np.array([5.0, 6.0])).all()
+
+
+def test_multiple_groupby() -> None:
+    df = mdf.MicroDataFrame({"x": [1, 2], "y": [3, 4], "z": [5, 6]})
+    assert (df.groupby(["x", "y"]).z.sum() == np.array([5, 6])).all()
+
+
+def test_concat() -> None:
+    df1 = mdf.MicroDataFrame({"x": [1, 2]}, weights=[3, 4])
+    df2 = mdf.MicroDataFrame({"y": [5, 6]}, weights=[7, 8])
+    # Verify that pd.concat returns DataFrame (probably no way to fix this).
+    pd_long = pd.concat([df1, df2])
+    assert isinstance(pd_long, pd.DataFrame)
+    assert not isinstance(pd_long, mdf.MicroDataFrame)
+    # Verify that mdf.concat works.
+    mdf_long = mdf.concat([df1, df2])
+    assert isinstance(mdf_long, mdf.MicroDataFrame)
+    # Weights should be preserved.
+    assert mdf_long.weights.equals(pd.concat([df1.weights, df2.weights]))
+    # Verify it works horizontally too (take the first set of weights).
+    mdf_wide = mdf.concat([df1, df2], axis=1)
+    assert isinstance(mdf_wide, mdf.MicroDataFrame)
+    assert mdf_wide.weights.equals(df1.weights)
+
+
+def test_set_index() -> None:
+    d = mdf.MicroDataFrame(dict(x=[1, 2, 3]), weights=[4, 5, 6])
+    assert d.x.__class__ == MicroSeries
+    d.index = [1, 2, 3]
+    assert d.x.__class__ == MicroSeries
+
+
+def test_reset_index() -> None:
+    d = mdf.MicroDataFrame(dict(x=[1, 2, 3]), weights=[4, 5, 6])
+    assert d.reset_index().__class__ == MicroDataFrame
+
+
+def test_cumsum() -> None:
+    s = mdf.MicroSeries([1, 2, 3], weights=[4, 5, 6])
+    assert np.array_equal(s.cumsum().values, [4, 14, 32])
+
+    s = mdf.MicroSeries([2, 1, 3], weights=[5, 4, 6])
+    assert np.array_equal(s.cumsum().values, [10, 14, 32])
+
+    s = mdf.MicroSeries([3, 1, 2], weights=[6, 4, 5])
+    assert np.array_equal(s.cumsum().values, [18, 22, 32])
+
+
+def test_rank() -> None:
+    s = mdf.MicroSeries([1, 2, 3], weights=[4, 5, 6])
+    assert np.array_equal(s.rank().values, [4, 9, 15])
+
+    s = mdf.MicroSeries([3, 1, 2], weights=[6, 4, 5])
+    assert np.array_equal(s.rank().values, [15, 4, 9])
+
+    s = mdf.MicroSeries([2, 1, 3], weights=[5, 4, 6])
+    assert np.array_equal(s.rank().values, [9, 4, 15])
+
+
+def test_percentile_rank() -> None:
+    s = mdf.MicroSeries([4, 2, 3, 1], weights=[20, 40, 20, 20])
+    assert np.array_equal(s.percentile_rank().values, [100, 60, 80, 20])
+
+
+def test_quartile_rank() -> None:
+    s = mdf.MicroSeries([4, 2, 3], weights=[25, 50, 25])
+    assert np.array_equal(s.quartile_rank().values, [4, 2, 3])
+
+
+def test_quintile_rank() -> None:
+    s = mdf.MicroSeries([4, 2, 3], weights=[20, 60, 20])
+    assert np.array_equal(s.quintile_rank().values, [5, 3, 4])
+
+
+def test_decile_rank_rank() -> None:
+    s = mdf.MicroSeries(
+        [5, 4, 3, 2, 1, 6, 7, 8, 9],
+        weights=[10, 20, 10, 10, 10, 10, 10, 10, 10, 10],
+    )
+    assert np.array_equal(s.decile_rank().values, [6, 5, 3, 2, 1, 7, 8, 9, 10])
+
+
+def test_copy_equals() -> None:
+    d = mdf.MicroDataFrame(
+        {"x": [1, 2], "y": [3, 4], "z": [5, 6]}, weights=[7, 8]
+    )
+    d_copy = d.copy()
+    d_copy_diff_weights = d_copy.copy()
+    d_copy_diff_weights.weights *= 2
+    assert d.equals(d_copy)
+    assert not d.equals(d_copy_diff_weights)
+    # Same for a MicroSeries.
+    assert d.x.equals(d_copy.x)
+    assert not d.x.equals(d_copy_diff_weights.x)
+
+
+def test_subset() -> None:
+    df = mdf.MicroDataFrame(
+        {"x": [1, 2], "y": [3, 4], "z": [5, 6]}, weights=[7, 8]
+    )
+    df_no_z = mdf.MicroDataFrame({"x": [1, 2], "y": [3, 4]}, weights=[7, 8])
+    assert df[["x", "y"]].equals(df_no_z)
+    df_no_z_diff_weights = df_no_z.copy()
+    df_no_z_diff_weights.weights += 1
+    assert not df[["x", "y"]].equals(df_no_z_diff_weights)
+
+
+def test_value_subset() -> None:
+    d = mdf.MicroDataFrame({"x": [1, 2, 3], "y": [1, 2, 2]}, weights=[4, 5, 6])
+    d2 = d[d.y > 1]
+    assert d2.y.shape == d2.weights.shape
+
+
+def test_bitwise_ops_return_microseries() -> None:
+    s1 = mdf.MicroSeries([True, False, True], weights=[1, 2, 3])
+    s2 = mdf.MicroSeries([False, False, True], weights=[1, 2, 3])
+    and_result = s1 & s2
+    or_result = s1 | s2
+    assert isinstance(and_result, mdf.MicroSeries)
+    assert isinstance(or_result, mdf.MicroSeries)
+    expected_and = mdf.MicroSeries([False, False, True], weights=[1, 2, 3])
+    expected_or = mdf.MicroSeries([True, False, True], weights=[1, 2, 3])
+    assert and_result.equals(expected_and)
+    assert or_result.equals(expected_or)
+
+
+def test_additional_ops_return_microseries() -> None:
+    s = mdf.MicroSeries([1, 2, 3], weights=[4, 5, 6])
+    radd = 1 + s
+    xor = s ^ mdf.MicroSeries([0, 1, 0], weights=[4, 5, 6])
+    inv = ~mdf.MicroSeries([True, False], weights=[1, 1])
+    assert isinstance(radd, mdf.MicroSeries)
+    assert isinstance(xor, mdf.MicroSeries)
+    assert isinstance(inv, mdf.MicroSeries)

microdf/tests/test_inequality.py

@@ -0,0 +1,20 @@
+import pandas as pd
+
+import microdf as mdf
+
+
+def test_top_pct() -> None:
+    x = list(range(1, 11))  # 1 to 10. Sum = 10 * 11 / 2 = 55.
+    df = pd.DataFrame({"x": x})
+    ms = mdf.MicroSeries(x)
+    RES = 10 / 55
+    assert mdf.top_10_pct_share(df, "x") == RES
+    assert ms.top_10_pct_share() == RES
+    x = list(range(1, 4))
+    df = pd.DataFrame({"x": x, "w": x})
+    ms = mdf.MicroSeries(x, weights=x)
+    # This is equivalent to [1, 2, 2, 3, 3, 3]
+    # Sum = 14, top half is 9.
+    RES = 9 / 14
+    assert mdf.top_50_pct_share(df, "x", "w") == RES
+    assert ms.top_50_pct_share() == RES

microdf/tests/test_io.py

@@ -0,0 +1,10 @@
+import microdf as mdf
+
+
+def test_read_stata_zip() -> None:
+    """"""
+    SCF2016 = "https://www.federalreserve.gov/econres/files/scfp2016s.zip"
+    COLS = ["wgt", "networth"]
+    df = mdf.read_stata_zip(SCF2016, columns=COLS)
+    assert df.columns.tolist() == COLS
+    assert df.shape[0] > 0

microdf/tests/test_optional_dependency.py

@@ -0,0 +1,54 @@
+import sys
+import types
+
+import pytest
+
+from microdf._optional import VERSIONS, import_optional_dependency
+
+
+def test_import_optional() -> None:
+    """"""
+    match = "Missing .*notapackage.* pip .* conda .* notapackage"
+    with pytest.raises(ImportError, match=match):
+        import_optional_dependency("notapackage")
+
+    result = import_optional_dependency("notapackage", raise_on_missing=False)
+    assert result is None
+
+
+def test_xlrd_version_fallback() -> None:
+    """"""
+    pytest.importorskip("xlrd")
+    import_optional_dependency("xlrd")
+
+
+def test_bad_version() -> None:
+    """"""
+    name = "fakemodule"
+    module = types.ModuleType(name)
+    module.__version__ = "0.9.0"
+    sys.modules[name] = module
+    VERSIONS[name] = "1.0.0"
+
+    match = "microdf requires .*1.0.0.* of .fakemodule.*'0.9.0'"
+    with pytest.raises(ImportError, match=match):
+        import_optional_dependency("fakemodule")
+
+    with pytest.warns(UserWarning):
+        result = import_optional_dependency("fakemodule", on_version="warn")
+    assert result is None
+
+    module.__version__ = "1.0.0"  # exact match is OK
+    result = import_optional_dependency("fakemodule")
+    assert result is module
+
+
+def test_no_version_raises() -> None:
+    """"""
+    name = "fakemodule"
+    module = types.ModuleType(name)
+    sys.modules[name] = module
+    VERSIONS[name] = "1.0.0"
+
+    with pytest.raises(ImportError, match="Can't determine .* fakemodule"):
+        import_optional_dependency(name)

microdf/tests/test_poverty.py

@@ -0,0 +1,73 @@
+import numpy as np
+import pandas as pd
+
+import microdf as mdf
+
+df = pd.DataFrame(
+    {
+        "income": [-10, 0, 10, 20],
+        "threshold": [15, 10, 15, 10],
+        "weight": [1, 2, 3, 4],
+    }
+)
+md = mdf.MicroDataFrame(df[["income", "threshold"]], weights=df.weight)
+
+
+def test_poverty_rate() -> None:
+    # Unweighted
+    assert np.allclose(mdf.poverty_rate(df, "income", "threshold"), 3 / 4)
+    # Weighted
+    assert np.allclose(
+        mdf.poverty_rate(df, "income", "threshold", "weight"), 6 / 10
+    )
+    assert np.allclose(md.poverty_rate("income", "threshold"), 6 / 10)
+
+
+def test_deep_poverty_rate() -> None:
+    # Unweighted
+    assert np.allclose(mdf.deep_poverty_rate(df, "income", "threshold"), 2 / 4)
+    # Weighted
+    assert np.allclose(
+        mdf.deep_poverty_rate(df, "income", "threshold", "weight"), 3 / 10
+    )
+    assert np.allclose(md.deep_poverty_rate("income", "threshold"), 3 / 10)
+
+
+def test_poverty_gap() -> None:
+    # Unweighted
+    assert np.allclose(mdf.poverty_gap(df, "income", "threshold"), 25 + 10 + 5)
+    # Weighted
+    RES = 25 * 1 + 10 * 2 + 5 * 3
+    assert np.allclose(
+        mdf.poverty_gap(df, "income", "threshold", "weight"), RES
+    )
+    assert np.allclose(md.poverty_gap("income", "threshold"), RES)
+
+
+def test_squared_poverty_gap() -> None:
+    # Unweighted
+    assert np.allclose(
+        mdf.squared_poverty_gap(df, "income", "threshold"),
+        25**2 + 10**2 + 5**2,
+    )
+    # Weighted
+    RES = 1 * (25**2) + 2 * (10**2) + 3 * (5**2)
+    assert np.allclose(
+        mdf.squared_poverty_gap(df, "income", "threshold", "weight"),
+        RES,
+    )
+    assert np.allclose(md.squared_poverty_gap("income", "threshold"), RES)
+
+
+def test_deep_poverty_gap() -> None:
+    # Unweighted
+    assert np.allclose(
+        mdf.deep_poverty_gap(df, "income", "threshold"), 17.5 + 5 + 0 + 0
+    )
+    # Weighted
+    RES = 17.5 * 1 + 5 * 2 + 0 * 3 + 0 * 4
+    assert np.allclose(
+        mdf.deep_poverty_gap(df, "income", "threshold", "weight"), RES
+    )
+    # Same in MicroDataFrame.
+    assert np.allclose(md.deep_poverty_gap("income", "threshold"), RES)

microdf/tests/test_quantile_chg.py

@@ -0,0 +1,18 @@
+import pandas as pd
+
+import microdf as mdf
+
+V1 = [1, 2, 3]
+V2 = [4, 5, 6]
+W1 = [7, 8, 9]
+W2 = [10, 11, 12]
+DF1 = pd.DataFrame({"v": V1, "w": W1})
+DF2 = pd.DataFrame({"v": V2, "w": W2})
+
+
+def test_quantile_chg() -> None:
+    mdf.quantile_chg(DF1, DF2, "v", "w", "v", "w")
+
+
+def test_quantile_pct_chg_plot() -> None:
+    mdf.quantile_pct_chg_plot(DF1, DF2, "v", "w", "v", "w")

microdf/tests/test_tax.py

@@ -0,0 +1,54 @@
+import numpy as np
+import pandas as pd
+import pytest
+
+import microdf as mdf
+
+
+def test_tax() -> None:
+    """"""
+    # Consider a MTR schedule of 0% up to 10,000, then 10% after that.
+    BRACKETS = [0, 10e3]
+    RATES = [0, 0.1]
+    INCOME = [0, 5e3, 10e3, 10e3 + 1, 20e3]
+    EXPECTED = [0, 0, 0, 0.1, 1e3]
+    res = mdf.tax_from_mtrs(INCOME, BRACKETS, RATES)
+    pd.testing.assert_series_equal(res, pd.Series(EXPECTED))
+    # Try with 10% avoidance.
+    EXPECTED_10PCT_AVOIDANCE = [0, 0, 0, 0, 800.0]
+    res_10pct_avoidance = mdf.tax_from_mtrs(INCOME, BRACKETS, RATES, 0.1)
+    pd.testing.assert_series_equal(
+        res_10pct_avoidance, pd.Series(EXPECTED_10PCT_AVOIDANCE)
+    )
+    # Try with avoidance elasticity of 2.
+    EXPECTED_E2_AVOIDANCE = [
+        0,
+        0,
+        0,
+        0,  # Taxable base becomes (10e3 + 1) * (1 - 2 * 0.1)
+        # Taxable base becomes 20e3 * (exp(-2 * 0.1)).
+        0.1 * (20e3 * np.exp(-0.2) - 10e3),
+    ]
+    res_e2_avoidance = mdf.tax_from_mtrs(
+        INCOME, BRACKETS, RATES, avoidance_elasticity=2
+    )
+    pd.testing.assert_series_equal(
+        res_e2_avoidance, pd.Series(EXPECTED_E2_AVOIDANCE)
+    )
+    # Try with flat avoidance elasticity of 2.
+    EXPECTED_E2_AVOIDANCE_FLAT = [
+        0,
+        0,
+        0,
+        0,  # Taxable base becomes (10e3 + 1) * (1 - 2 * 0.1)
+        600.0,
+    ]  # Taxable base becomes 20e3 * (1 - 2 * 0.1) = 16e3.
+    res_e2_avoidance_flat = mdf.tax_from_mtrs(
+        INCOME, BRACKETS, RATES, avoidance_elasticity_flat=2
+    )
+    pd.testing.assert_series_equal(
+        res_e2_avoidance_flat, pd.Series(EXPECTED_E2_AVOIDANCE_FLAT)
+    )
+    # Ensure error when passing both rate and elasticity.
+    with pytest.raises(Exception):
+        mdf.tax_from_mtrs(INCOME, BRACKETS, RATES, 0.1, 2)

microdf/tests/test_taxcalc.py

@@ -0,0 +1,25 @@
+import pytest
+
+import microdf as mdf
+
+try:
+    import taxcalc as tc
+
+    _HAVE_TAXCALC = True
+except ImportError:
+    _HAVE_TAXCALC = False
+
+
+def test_calc_df() -> None:
+    """"""
+    if not _HAVE_TAXCALC:
+        pytest.skip("taxcalc is not installed")
+    mdf.calc_df()
+
+
+def test_static_baseline_calc() -> None:
+    """"""
+    if not _HAVE_TAXCALC:
+        pytest.skip("taxcalc is not installed")
+    recs = tc.Records.cps_constructor()
+    mdf.static_baseline_calc(recs, 2020)

microdf/tests/test_utils.py

@@ -0,0 +1,39 @@
+import pandas as pd
+
+import microdf as mdf
+
+
+def test_cartesian_product() -> None:
+    """"""
+    res = mdf.cartesian_product(
+        {"a": [1, 2, 3], "b": ["val1", "val2"], "c": [100, 101]}
+    )
+    EXPECTED = pd.DataFrame(
+        {
+            "a": [1, 1, 1, 1, 2, 2, 2, 2, 3, 3, 3, 3],
+            "b": [
+                "val1",
+                "val1",
+                "val2",
+                "val2",
+                "val1",
+                "val1",
+                "val2",
+                "val2",
+                "val1",
+                "val1",
+                "val2",
+                "val2",
+            ],
+            "c": [100, 101, 100, 101, 100, 101, 100, 101, 100, 101, 100, 101],
+        }
+    )
+    pd.testing.assert_frame_equal(res, EXPECTED)
+
+
+def test_flatten() -> None:
+    """"""
+    L = [[[1, 2, 3], [4, 5]], 6]
+    res = list(mdf.flatten(L))
+    EXPECTED = [1, 2, 3, 4, 5, 6]
+    assert res == EXPECTED

microdf/tests/test_weighted.py

@@ -0,0 +1,74 @@
+import pandas as pd
+import pytest
+
+import microdf as mdf
+
+X = [1, 5, 2]
+Y = [0, -6, 3]
+W = [4, 1, 1]
+df = pd.DataFrame({"x": X, "y": Y, "w": W})
+ms = mdf.MicroSeries(X, weights=W)
+md = mdf.MicroDataFrame(df[["x", "y"]], weights=W)
+# Also make a version with groups.
+df2 = df.copy(deep=True)
+df2.x *= 2
+df2.y *= 1.5
+dfg = pd.concat([df, df2])
+dfg["g"] = ["a"] * 3 + ["b"] * 3
+mdg = mdf.MicroDataFrame(dfg[["x", "y", "g"]], weights=W)
+
+
+def test_weighted_quantile() -> None:
+    Q = [0, 0.5, 1]
+    mdf.weighted_quantile(df, "x", "w", Q).tolist()
+
+
+def test_weighted_median() -> None:
+    assert mdf.weighted_median(df, "x") == 2
+    mdf.weighted_median(df, "x", "w")
+    # Test with groups.
+    mdf.weighted_median(dfg, "x", "w", "g")
+
+
+def test_weighted_mean() -> None:
+    # Test umweighted.
+    assert mdf.weighted_mean(df, "x") == 8 / 3
+    # Test weighted.
+    assert mdf.weighted_mean(df, "x", "w") == 11 / 6
+    # Test weighted with multiple columns.
+    assert mdf.weighted_mean(df, ["x", "y"], "w").tolist() == [11 / 6, -3 / 6]
+    # Test grouped.
+    mdf.weighted_mean(dfg, "x", "w", "g")
+    mdf.weighted_mean(dfg, ["x", "y"], "w", "g")
+
+
+def test_weighted_sum() -> None:
+    # Test unweighted.
+    assert mdf.weighted_sum(df, "x") == 8
+    # Test weighted.
+    assert mdf.weighted_sum(df, "x", "w") == 11
+    # Test weighted with multiple columns.
+    assert mdf.weighted_sum(df, ["x", "y"], "w").tolist() == [11, -3]
+    # Test grouped.
+    mdf.weighted_sum(dfg, "x", "w", "g")
+    mdf.weighted_sum(dfg, ["x", "y"], "w", "g")
+
+
+def test_gini() -> None:
+    # Test nothing breaks.
+    ms.gini()
+    # Unweighted.
+    mdf.gini(df, "x")
+    # Weighted
+    mdf.gini(df, "x", "w")
+    # Unweighted, grouped
+    mdf.gini(dfg, "x", groupby="g")
+    # Weighted, grouped
+    mdf.gini(dfg, "x", "w", groupby="g")
+    # Test old and new match.
+    assert ms.gini() == mdf.gini(df, "x", "w")
+
+
+def test_add_weighted_quantiles() -> None:
+    with pytest.deprecated_call():
+        mdf.add_weighted_quantiles(df, "x", "w")

microdf/ubi.py

@@ -1,32 +1,34 @@
+from typing import List, Optional, Union
+
 import numpy as np
+import pandas as pd
 
 import microdf as mdf
 
 
 def ubi_or_bens(
-    df,
-    ben_cols,
-    max_ubi="max_ubi",
-    ubi="ubi",
-    bens="bens",
-    update_income_measures=None,
-):
-    """Calculates whether a tax unit will take UBI or benefits,
-       and adjusts values accordingly.
+    df: pd.DataFrame,
+    ben_cols: Union[str, List[str]],
+    max_ubi: str = "max_ubi",
+    ubi: str = "ubi",
+    bens: str = "bens",
+    update_income_measures: Optional[List[str]] = None,
+) -> None:
+    """Calculates whether a tax unit will take UBI or benefits, and adjusts
+    values accordingly.
 
     :param df: DataFrame.
     :param ben_cols: List of columns for benefits.
-    :param max_ubi: Column name of the maximum UBI, before accounting
-            for benefits. Defaults to 'max_ubi'.
+    :param max_ubi: Column name of the maximum UBI, before accounting for
+        benefits. Defaults to 'max_ubi'.
     :param ubi: Column name to add representing the UBI. Defaults to 'ubi'.
     :param bens: Column name to add representing total benefits (after
-            adjustment). Defaults to 'bens'.
-    :param update_income_measures: List of income measures to update.
-            Defaults to ['expanded_income', 'aftertax_income'].
+        adjustment). Defaults to 'bens'.
+    :param update_income_measures: List of income measures to update. Defaults
+        to ['expanded_income', 'aftertax_income'].
     :returns: Nothing. Benefits in ben_cols are adjusted, ubi and bens columns
         are added, and expanded_income and aftertax_income are updated
         according to the net difference.
-
     """
     if update_income_measures is None:
         update_income_measures = ["expanded_income", "aftertax_income"]