From 650303e2ab19748774a7d8cc8dedc5c7451f9c31 Mon Sep 17 00:00:00 2001 From: juaristi22 Date: Fri, 18 Jul 2025 15:06:03 +0200 Subject: [PATCH 1/9] fix build for deployment --- .github/workflows/master.yml | 9 ++------- Makefile | 4 ++++ changelog_entry.yaml | 4 ++++ 3 files changed, 10 insertions(+), 7 deletions(-) diff --git a/.github/workflows/master.yml b/.github/workflows/master.yml index fbe29d22..5e075291 100644 --- a/.github/workflows/master.yml +++ b/.github/workflows/master.yml @@ -74,12 +74,7 @@ jobs: uses: actions/checkout@v4 with: fetch-depth: 0 # Fetch all history for all tags and branches - - name: Set up Python - uses: actions/setup-python@v5 - with: - python-version: ${{ matrix.python-version }} - - - name: Set up Python + - name: Set up Python uses: actions/setup-python@v5 with: python-version: '3.10' @@ -92,7 +87,7 @@ jobs: uv pip install -e ".[dev]" --system - name: Build package run: | - python -m build --wheel --sdist + make build - name: Publish a git tag run: ".github/publish-git-tag.sh || true" - name: Publish to PyPI diff --git a/Makefile b/Makefile index b309000d..e13094cd 100644 --- a/Makefile +++ b/Makefile @@ -19,6 +19,10 @@ test: install: pip install -e ".[dev]" +build: + pip install build + python -m build --wheel --sdist + changelog: build-changelog changelog.yaml --output changelog.yaml --update-last-date --start-from 0.4.5 --append-file changelog_entry.yaml build-changelog changelog.yaml --org PolicyEngine --repo microcalibrate --output CHANGELOG.md --template .github/changelog_template.md diff --git a/changelog_entry.yaml b/changelog_entry.yaml index e69de29b..14f84249 100644 --- a/changelog_entry.yaml +++ b/changelog_entry.yaml @@ -0,0 +1,4 @@ +- bump: patch + changes: + changed: + - Fix build for deployment. From 17f61db9cc99a4f92db18d3152a0b25ea5b7f78f Mon Sep 17 00:00:00 2001 From: juaristi22 Date: Fri, 18 Jul 2025 16:19:15 +0200 Subject: [PATCH 2/9] create astype() methods for microseries and microdataframe --- microdf/generic.py | 42 ++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 42 insertions(+) diff --git a/microdf/generic.py b/microdf/generic.py index 7fe658b1..98c0f9f1 100644 --- a/microdf/generic.py +++ b/microdf/generic.py @@ -512,6 +512,27 @@ def __neg__(self) -> "MicroSeries": def __pos__(self) -> "MicroSeries": return MicroSeries(super().__pos__(), weights=self.weights) + def astype( + self, + dtype, + copy: Optional[bool] = True, + errors: Optional[str] = "raise", + ) -> "MicroSeries": + """Convert MicroSeries to specified data type while preserving weights. + + :param dtype: Data type to convert to. Can be numpy dtype or Python + type. + :param copy: Whether to make a copy of the data (default True). + :param errors: How to handle conversion errors (default "raise"). + :return: New MicroSeries with converted data type and preserved + weights. + """ + converted_series = super().astype(dtype, copy=copy, errors=errors) + return MicroSeries( + converted_series, + weights=self.weights.copy() if copy else self.weights, + ) + def __repr__(self) -> str: return pd.DataFrame( dict(value=self.values, weight=self.weights.values) @@ -928,6 +949,27 @@ def poverty_count( in_poverty = income < threshold return in_poverty.sum() + def astype( + self, + dtype, + copy: Optional[bool] = True, + errors: Optional[str] = "raise", + ) -> "MicroDataFrame": + """Convert MicroDataFrame to specified data type while preserving + weights. + + :param dtype: Data type to convert to. Can be numpy dtype, Python type, + or dict. + :param copy: Whether to make a copy of the data (default True). + :param errors: How to handle conversion errors (default "raise"). + :return: New MicroDataFrame with converted data types and preserved + weights. + """ + converted_df = super().astype(dtype, copy=copy, errors=errors) + return MicroDataFrame( + converted_df, weights=self.weights.copy() if copy else self.weights + ) + def __repr__(self) -> str: df = pd.DataFrame(self) df["weight"] = self.weights From 1882334bb08b62b25134149fcfb9dd8243e1ba57 Mon Sep 17 00:00:00 2001 From: juaristi22 Date: Fri, 18 Jul 2025 17:10:05 +0200 Subject: [PATCH 3/9] write microseries.sqrt() function --- changelog_entry.yaml | 6 +++--- microdf/generic.py | 4 ++++ 2 files changed, 7 insertions(+), 3 deletions(-) diff --git a/changelog_entry.yaml b/changelog_entry.yaml index 14f84249..f1201d77 100644 --- a/changelog_entry.yaml +++ b/changelog_entry.yaml @@ -1,4 +1,4 @@ -- bump: patch +- bump: minor changes: - changed: - - Fix build for deployment. + added: + - Add astype and sqrt methods to MicroSeries and MicroDataFrame. diff --git a/microdf/generic.py b/microdf/generic.py index 98c0f9f1..0f148ae1 100644 --- a/microdf/generic.py +++ b/microdf/generic.py @@ -454,6 +454,10 @@ def __ror__(self, other: Union[int, float, pd.Series]) -> "MicroSeries": def __rxor__(self, other: Union[int, float, pd.Series]) -> "MicroSeries": return MicroSeries(super().__rxor__(other), weights=self.weights) + def sqrt(self) -> "MicroSeries": + sqrt_values = np.sqrt(self.values) + return MicroSeries(sqrt_values, index=self.index, weights=self.weights) + # comparators def __lt__(self, other: Union[int, float, pd.Series]) -> "MicroSeries": From 4171a9479724980551f04e1ac2df8e17a250b524 Mon Sep 17 00:00:00 2001 From: juaristi22 Date: Fri, 18 Jul 2025 17:40:56 +0200 Subject: [PATCH 4/9] make reset_index work in-place --- changelog_entry.yaml | 1 + microdf/generic.py | 60 ++++++++++++++++++++++++++++++++--- microdf/tests/test_generic.py | 56 ++++++++++++++++++++++++++++++++ 3 files changed, 113 insertions(+), 4 deletions(-) diff --git a/changelog_entry.yaml b/changelog_entry.yaml index f1201d77..c69e3f97 100644 --- a/changelog_entry.yaml +++ b/changelog_entry.yaml @@ -2,3 +2,4 @@ changes: added: - Add astype and sqrt methods to MicroSeries and MicroDataFrame. + - Support in-place reset_index. diff --git a/microdf/generic.py b/microdf/generic.py index 0f148ae1..50d2945b 100644 --- a/microdf/generic.py +++ b/microdf/generic.py @@ -810,10 +810,62 @@ def __setattr__(self, key, value) -> None: super().__setattr__(key, value) self.catch_series_relapse() - def reset_index(self) -> "MicroDataFrame": - res = super().reset_index() - res = MicroDataFrame(res, weights=self.weights) - return res + def reset_index( + self, + level: Optional[int] = None, + drop: Optional[bool] = False, + inplace: Optional[bool] = False, + col_level: Optional[int] = 0, + col_fill: Optional[str] = "", + allow_duplicates: Optional[bool] = None, + names: Optional[List[str]] = None, + ) -> Union["MicroDataFrame", None]: + """Reset the index of the MicroDataFrame. + + This method supports all parameters of pandas DataFrame.reset_index(), + including the 'inplace' parameter. + + :param level: Only remove the given levels from the index. Removes all + levels by default. + :param drop: Do not try to insert index into dataframe columns. This + resets the index to the default integer index. + :param inplace: Modify the DataFrame in place (do not create a new + object). + :param col_level: If the columns have multiple levels, determines which + level the labels are inserted into. + :param col_fill: If the columns have multiple levels, determines how + the other levels are named. + :param allow_duplicates: Allow duplicate column labels to be created. + :param names: Using the given string, rename the DataFrame column which + contains the index data. + :return: MicroDataFrame with reset index or None if inplace=True. + """ + if inplace: + weights_backup = self.weights.copy() + # Perform in-place reset on the parent DataFrame + super().reset_index( + level=level, + drop=drop, + inplace=True, + col_level=col_level, + col_fill=col_fill, + allow_duplicates=allow_duplicates, + names=names, + ) + self.weights = weights_backup + self._link_all_weights() + return None + else: + res = super().reset_index( + level=level, + drop=drop, + inplace=False, + col_level=col_level, + col_fill=col_fill, + allow_duplicates=allow_duplicates, + names=names, + ) + return MicroDataFrame(res, weights=self.weights) def copy(self, deep: Optional[bool] = True) -> "MicroDataFrame": res = super().copy(deep) diff --git a/microdf/tests/test_generic.py b/microdf/tests/test_generic.py index 3f6fe5f3..fb7d11e6 100644 --- a/microdf/tests/test_generic.py +++ b/microdf/tests/test_generic.py @@ -230,3 +230,59 @@ def test_additional_ops_return_microseries() -> None: assert isinstance(radd, mdf.MicroSeries) assert isinstance(xor, mdf.MicroSeries) assert isinstance(inv, mdf.MicroSeries) + + +def test_reset_index_inplace() -> None: + df = pd.DataFrame( + {"A": [1, 2, 3, 4], "B": [5, 6, 7, 8]}, index=["a", "b", "c", "d"] + ) + weights = np.array([0.1, 0.2, 0.3, 0.4]) + mdf = MicroDataFrame(df, weights=weights) + + # Test 1: reset_index with inplace=False (default) + mdf_copy = mdf.copy() + result = mdf_copy.reset_index() + assert list(mdf_copy.index) == ["a", "b", "c", "d"] + assert list(result.index) == [0, 1, 2, 3] + assert "index" in result.columns + assert list(result["index"]) == ["a", "b", "c", "d"] + np.testing.assert_array_equal(result.weights.values, weights) + + # Test 2: reset_index with inplace=True + mdf_copy = mdf.copy() + result = mdf_copy.reset_index(inplace=True) + assert result is None + assert list(mdf_copy.index) == [0, 1, 2, 3] + assert "index" in mdf_copy.columns + assert list(mdf_copy["index"]) == ["a", "b", "c", "d"] + np.testing.assert_array_equal(mdf_copy.weights.values, weights) + assert isinstance(mdf_copy["A"], MicroSeries) + assert isinstance(mdf_copy["B"], MicroSeries) + assert isinstance(mdf_copy["index"], MicroSeries) + + # Test 3: reset_index with drop=True + mdf_copy = mdf.copy() + mdf_copy.reset_index(drop=True, inplace=True) + assert list(mdf_copy.index) == [0, 1, 2, 3] + assert "index" not in mdf_copy.columns + assert list(mdf_copy.columns) == ["A", "B"] + np.testing.assert_array_equal(mdf_copy.weights.values, weights) + + # Test 4: Multi-level index + arrays = [["bar", "bar", "baz", "baz"], ["one", "two", "one", "two"]] + multi_index = pd.MultiIndex.from_arrays(arrays, names=["first", "second"]) + df_multi = pd.DataFrame( + {"A": [1, 2, 3, 4], "B": [5, 6, 7, 8]}, index=multi_index + ) + mdf_multi = MicroDataFrame(df_multi, weights=weights) + result = mdf_multi.reset_index(level="first") + assert "first" in result.columns + assert result.index.name == "second" + np.testing.assert_array_equal(result.weights.values, weights) + + # Reset all levels in place + mdf_multi.reset_index(inplace=True) + assert "first" in mdf_multi.columns + assert "second" in mdf_multi.columns + assert list(mdf_multi.index) == [0, 1, 2, 3] + np.testing.assert_array_equal(mdf_multi.weights.values, weights) From 6b907af78d2f29f495dc59d060e363e97cd08bb7 Mon Sep 17 00:00:00 2001 From: juaristi22 Date: Fri, 18 Jul 2025 18:24:18 +0200 Subject: [PATCH 5/9] split generic.py --- changelog_entry.yaml | 1 + microdf/__init__.py | 19 +- microdf/concat.py | 2 +- microdf/microdataframe.py | 433 +++++++++++++++++++++++++ microdf/{generic.py => microseries.py} | 423 ------------------------ microdf/tests/test_generic.py | 3 +- 6 files changed, 443 insertions(+), 438 deletions(-) create mode 100644 microdf/microdataframe.py rename microdf/{generic.py => microseries.py} (57%) diff --git a/changelog_entry.yaml b/changelog_entry.yaml index c69e3f97..c756d312 100644 --- a/changelog_entry.yaml +++ b/changelog_entry.yaml @@ -3,3 +3,4 @@ added: - Add astype and sqrt methods to MicroSeries and MicroDataFrame. - Support in-place reset_index. + - Split generic.py into microdataframe.py and microseries.py. diff --git a/microdf/__init__.py b/microdf/__init__.py index aede573c..ec8e8585 100644 --- a/microdf/__init__.py +++ b/microdf/__init__.py @@ -23,7 +23,6 @@ add_ftt, add_vat, ) -from .generic import MicroDataFrame, MicroSeries from .income_measures import cash_income, market_income, tpc_eci from .inequality import ( bottom_50_pct_share, @@ -37,6 +36,8 @@ top_x_pct_share, ) from .io import read_stata_zip +from .microdataframe import MicroDataFrame, MicroDataFrameGroupBy +from .microseries import MicroSeries, MicroSeriesGroupBy from .poverty import ( deep_poverty_gap, deep_poverty_rate, @@ -79,11 +80,6 @@ "combine_base_reform", "pctchg_base_reform", "agg", - # chart_utils.py - "dollar_format", - "currency_format", - # charts.py - "quantile_pct_chg_plot", # concat.py "concat", # constants.py @@ -130,12 +126,6 @@ "poverty_gap", "squared_poverty_gap", "deep_poverty_gap", - # style.py - "AXIS_COLOR", - "DPI", - "GRID_COLOR", - "TITLE_COLOR", - "set_plot_style", # tax.py "mtr", "tax_from_mtrs", @@ -161,7 +151,10 @@ "weighted_median", "add_weighted_quantiles", "quantile_chg", - # generic.py + # microseries.py "MicroSeries", + "MicroSeriesGroupBy", + # microdataframe.py "MicroDataFrame", + "MicroDataFrameGroupBy", ] diff --git a/microdf/concat.py b/microdf/concat.py index dad6f42a..cc469205 100644 --- a/microdf/concat.py +++ b/microdf/concat.py @@ -3,7 +3,7 @@ import pandas as pd import microdf as mdf -from microdf.generic import MicroDataFrame +from microdf.microdataframe import MicroDataFrame def concat(*args, **kwargs) -> "MicroDataFrame": diff --git a/microdf/microdataframe.py b/microdf/microdataframe.py new file mode 100644 index 00000000..31f4a795 --- /dev/null +++ b/microdf/microdataframe.py @@ -0,0 +1,433 @@ +import copy +import logging +import warnings +from functools import wraps +from typing import Callable, List, Optional, Union + +import numpy as np +import pandas as pd + +from microdf.microseries import MicroSeries, MicroSeriesGroupBy + +logger = logging.getLogger(__name__) + + +class MicroDataFrame(pd.DataFrame): + def __init__(self, *args, weights=None, **kwargs): + """A DataFrame-inheriting class for weighted microdata. Weights can be + provided at initialisation, or using set_weights or set_weight_col. + + :param weights: Array of weights. + :type weights: np.array + """ + super().__init__(*args, **kwargs) + self.weights = None + self.set_weights(weights) + self._link_all_weights() + self.override_df_functions() + + def override_df_functions(self) -> None: + for name in MicroSeries.FUNCTIONS: + + def get_fn(name) -> Callable: + def fn(*args, **kwargs) -> Union[pd.Series, pd.DataFrame]: + is_array = len(args) > 0 and hasattr(args[0], "__len__") + if ( + name in MicroSeries.SCALAR_FUNCTIONS + or name in MicroSeries.AGNOSTIC_FUNCTIONS + and not is_array + ): + results = pd.Series( + [ + getattr(self[col], name)(*args, **kwargs) + for col in self.columns + ] + ) + results.index = self.columns + return results + elif ( + name in MicroSeries.VECTOR_FUNCTIONS + or name in MicroSeries.AGNOSTIC_FUNCTIONS + and is_array + ): + results = pd.DataFrame( + [ + getattr(self[col], name)(*args, **kwargs) + for col in self.columns + ] + ) + results.index = self.columns + return results + + return fn + + setattr(self, name, get_fn(name)) + + def get_args_as_micro_series(*kwarg_names: tuple) -> Callable: + """Decorator for auto-parsing column names into MicroSeries objects. If + given, kwarg_names limits arguments checked to keyword arguments + specified. + + :param arg_names: argument names to restrict to. + :type arg_names: str + """ + + def arg_series_decorator(fn) -> Callable: + @wraps(fn) + def series_function( + self, *args, **kwargs + ) -> Union[pd.Series, pd.DataFrame]: + new_args = [] + new_kwargs = {} + if len(kwarg_names) == 0: + for value in args: + if isinstance(value, str): + if value not in self.columns: + raise Exception("Column not found") + new_args += [self[value]] + else: + new_args += [value] + for name, value in kwargs.items(): + if isinstance(value, str) and ( + len(kwarg_names) == 0 or name in kwarg_names + ): + if value not in self.columns: + raise Exception("Column not found") + new_kwargs[name] = self[value] + else: + new_kwargs[name] = value + return fn(self, *new_args, **new_kwargs) + + return series_function + + return arg_series_decorator + + def __setitem__(self, *args, **kwargs) -> None: + super().__setitem__(*args, **kwargs) + self._link_all_weights() + + def _link_weights(self, column) -> None: + # self[column] = ... triggers __setitem__, which forces pd.Series + # this workaround avoids that + self[column].__class__ = MicroSeries + self[column].set_weights(self.weights) + + def _link_all_weights(self) -> None: + if self.weights is None: + self.set_weights(np.ones((len(self)))) + for column in self.columns: + if column != self.weights_col: + self._link_weights(column) + + def set_weights(self, weights: np.ndarray) -> None: + """Sets the weights for the MicroDataFrame. If a string is received, it + will be assumed to be the column name of the weight column. + + :param weights: Array of weights. + :type weights: np.array + """ + if isinstance(weights, str): + self.weights_col = weights + self.weights = pd.Series(self[weights], dtype=float) + elif weights is not None: + self.weights_col = None + with warnings.catch_warnings(): + warnings.filterwarnings("ignore", category=UserWarning) + self.weights = pd.Series(weights, dtype=float) + self._link_all_weights() + + def set_weight_col(self, column: str) -> None: + """Sets the weights for the MicroDataFrame by specifying the name of + the weight column. + + :param weights: Array of weights. + :type weights: np.array + """ + self.weights = np.array(self[column]) + self.weight_col = column + self._link_all_weights() + + def __getitem__( + self, key: Union[str, List] + ) -> Union[pd.Series, pd.DataFrame]: + result = super().__getitem__(key) + if isinstance(result, pd.DataFrame): + try: + weights = self.weights[key] + except Exception: + weights = self.weights + return MicroDataFrame(result, weights=weights) + return result + + def catch_series_relapse(self) -> None: + for col in self.columns: + if self[col].__class__ == pd.Series: + self._link_weights(col) + + def __setattr__(self, key, value) -> None: + super().__setattr__(key, value) + self.catch_series_relapse() + + def reset_index( + self, + level: Optional[int] = None, + drop: Optional[bool] = False, + inplace: Optional[bool] = False, + col_level: Optional[int] = 0, + col_fill: Optional[str] = "", + allow_duplicates: Optional[bool] = None, + names: Optional[List[str]] = None, + ) -> Union["MicroDataFrame", None]: + """Reset the index of the MicroDataFrame. + + This method supports all parameters of pandas DataFrame.reset_index(), + including the 'inplace' parameter. + + :param level: Only remove the given levels from the index. Removes all + levels by default. + :param drop: Do not try to insert index into dataframe columns. This + resets the index to the default integer index. + :param inplace: Modify the DataFrame in place (do not create a new + object). + :param col_level: If the columns have multiple levels, determines which + level the labels are inserted into. + :param col_fill: If the columns have multiple levels, determines how + the other levels are named. + :param allow_duplicates: Allow duplicate column labels to be created. + :param names: Using the given string, rename the DataFrame column which + contains the index data. + :return: MicroDataFrame with reset index or None if inplace=True. + """ + if inplace: + weights_backup = self.weights.copy() + # Perform in-place reset on the parent DataFrame + super().reset_index( + level=level, + drop=drop, + inplace=True, + col_level=col_level, + col_fill=col_fill, + allow_duplicates=allow_duplicates, + names=names, + ) + self.weights = weights_backup + self._link_all_weights() + return None + else: + res = super().reset_index( + level=level, + drop=drop, + inplace=False, + col_level=col_level, + col_fill=col_fill, + allow_duplicates=allow_duplicates, + names=names, + ) + return MicroDataFrame(res, weights=self.weights) + + def copy(self, deep: Optional[bool] = True) -> "MicroDataFrame": + res = super().copy(deep) + # This changes the original columns to Series. Undo it: + for col in self.columns: + self[col] = MicroSeries(self[col]) + res = MicroDataFrame(res, weights=self.weights.copy(deep)) + return res + + def equals(self, other: "MicroDataFrame") -> bool: + equal_values = super().equals(other) + equal_weights = self.weights.equals(other.weights) + return equal_values and equal_weights + + @get_args_as_micro_series() + def groupby( + self, by: Union[str, List], *args, **kwargs + ) -> "MicroDataFrameGroupBy": + """Returns a GroupBy object with MicroSeriesGroupBy objects for each + column. + + :param by: column to group by + :type by: Union[str, List] + + return: DataFrameGroupBy object with columns using weights + rtype: DataFrameGroupBy + """ + self["__tmp_weights"] = self.weights + gb = super().groupby(by, *args, **kwargs) + weights = copy.deepcopy(gb["__tmp_weights"]) + for col in self.columns: # df.groupby(...)[col]s use weights + res = gb[col] + res.__class__ = MicroSeriesGroupBy + res._init() + res.weights = weights + setattr(gb, col, res) + gb.__class__ = MicroDataFrameGroupBy + gb._init(by) + return gb + + @get_args_as_micro_series() + def poverty_rate(self, income: str, threshold: str) -> float: + """Calculate poverty rate, i.e., the population share with income below + their poverty threshold. + + :param income: Column indicating income. + :type income: str + :param threshold: Column indicating threshold. + :type threshold: str + :return: Poverty rate between zero and one. + :rtype: float + """ + pov = income < threshold + return pov.sum() / pov.count() + + @get_args_as_micro_series() + def deep_poverty_rate(self, income: str, threshold: str) -> float: + """Calculate deep poverty rate, i.e., the population share with income + below half their poverty threshold. + + :param income: Column indicating income. + :type income: str + :param threshold: Column indicating threshold. + :type threshold: str + :return: Deep poverty rate between zero and one. + :rtype: float + """ + pov = income < (threshold / 2) + return pov.sum() / pov.count() + + @get_args_as_micro_series() + def poverty_gap(self, income: str, threshold: str) -> float: + """Calculate poverty gap, i.e., the total gap between income and + poverty thresholds for all people in poverty. + + :param income: Column indicating income. + :type income: str + :param threshold: Column indicating threshold. + :type threshold: str + :return: Poverty gap. + :rtype: float + """ + gaps = (threshold - income)[threshold > income] + return gaps.sum() + + @get_args_as_micro_series() + def deep_poverty_gap(self, income: str, threshold: str) -> float: + """Calculate deep poverty gap, i.e., the total gap between income and + half of poverty thresholds for all people in deep poverty. + + :param income: Column indicating income. + :type income: str + :param threshold: Column indicating threshold. + :type threshold: str + :return: Deep poverty gap. + :rtype: float + """ + deep_threshold = threshold / 2 + gaps = (deep_threshold - income)[deep_threshold > income] + return gaps.sum() + + @get_args_as_micro_series() + def squared_poverty_gap(self, income: str, threshold: str) -> float: + """Calculate squared poverty gap, i.e., the total squared gap between + income and poverty thresholds for all people in poverty. Also known as + the poverty severity index. + + :param income: Column indicating income. + :type income: str + :param threshold: Column indicating threshold. + :type threshold: str + :return: Squared poverty gap. + :rtype: float + """ + gaps = (threshold - income)[threshold > income] + squared_gaps = gaps**2 + return squared_gaps.sum() + + @get_args_as_micro_series() + def poverty_count( + self, + income: Union[MicroSeries, str], + threshold: Union[MicroSeries, str], + ) -> int: + """Calculates the number of entities with income below a poverty + threshold. + + :param income: income array or column name + :type income: Union[MicroSeries, str] + + :param threshold: threshold array or column name + :type threshold: Union[MicroSeries, str] + + return: number of entities in poverty + rtype: int + """ + in_poverty = income < threshold + return in_poverty.sum() + + def astype( + self, + dtype, + copy: Optional[bool] = True, + errors: Optional[str] = "raise", + ) -> "MicroDataFrame": + """Convert MicroDataFrame to specified data type while preserving + weights. + + :param dtype: Data type to convert to. Can be numpy dtype, Python type, + or dict. + :param copy: Whether to make a copy of the data (default True). + :param errors: How to handle conversion errors (default "raise"). + :return: New MicroDataFrame with converted data types and preserved + weights. + """ + converted_df = super().astype(dtype, copy=copy, errors=errors) + return MicroDataFrame( + converted_df, weights=self.weights.copy() if copy else self.weights + ) + + def __repr__(self) -> str: + df = pd.DataFrame(self) + df["weight"] = self.weights + return df[[df.columns[-1]] + list(df.columns[:-1])].__repr__() + + +class MicroDataFrameGroupBy(pd.core.groupby.generic.DataFrameGroupBy): + def _init(self, by: Union[str, List]): + self.columns = list(self.obj.columns) + if isinstance(by, list): + for column in by: + self.columns.remove(column) + elif isinstance(by, str): + self.columns.remove(by) + self.columns.remove("__tmp_weights") + for fn_name in MicroSeries.SCALAR_FUNCTIONS: + + def get_fn(name): + def fn(*args, **kwargs): + return MicroDataFrame( + { + col: getattr(getattr(self, col), name)( + *args, **kwargs + ) + for col in self.columns + } + ) + + return fn + + setattr(self, fn_name, get_fn(fn_name)) + for fn_name in MicroSeries.VECTOR_FUNCTIONS: + + def get_fn(name) -> Callable: + def fn(*args, **kwargs) -> Union[pd.Series, pd.DataFrame]: + return MicroDataFrame( + { + col: getattr(getattr(self, col), name)( + *args, **kwargs + ) + for col in self.columns + } + ) + + return fn + + setattr(self, fn_name, get_fn(fn_name)) diff --git a/microdf/generic.py b/microdf/microseries.py similarity index 57% rename from microdf/generic.py rename to microdf/microseries.py index 50d2945b..89fbd00b 100644 --- a/microdf/generic.py +++ b/microdf/microseries.py @@ -1,6 +1,4 @@ -import copy import logging -import warnings from functools import wraps from typing import Callable, List, Optional, Union @@ -609,424 +607,3 @@ def _weighted_agg_fn( for fn_name in MicroSeries.FUNCTIONS: setattr(self, fn_name, _weighted_agg(fn_name)) - - -class MicroDataFrameGroupBy(pd.core.groupby.generic.DataFrameGroupBy): - def _init(self, by: Union[str, List]): - self.columns = list(self.obj.columns) - if isinstance(by, list): - for column in by: - self.columns.remove(column) - elif isinstance(by, str): - self.columns.remove(by) - self.columns.remove("__tmp_weights") - for fn_name in MicroSeries.SCALAR_FUNCTIONS: - - def get_fn(name): - def fn(*args, **kwargs): - return MicroDataFrame( - { - col: getattr(getattr(self, col), name)( - *args, **kwargs - ) - for col in self.columns - } - ) - - return fn - - setattr(self, fn_name, get_fn(fn_name)) - for fn_name in MicroSeries.VECTOR_FUNCTIONS: - - def get_fn(name) -> Callable: - def fn(*args, **kwargs) -> Union[pd.Series, pd.DataFrame]: - return MicroDataFrame( - { - col: getattr(getattr(self, col), name)( - *args, **kwargs - ) - for col in self.columns - } - ) - - return fn - - setattr(self, fn_name, get_fn(fn_name)) - - -class MicroDataFrame(pd.DataFrame): - def __init__(self, *args, weights=None, **kwargs): - """A DataFrame-inheriting class for weighted microdata. Weights can be - provided at initialisation, or using set_weights or set_weight_col. - - :param weights: Array of weights. - :type weights: np.array - """ - super().__init__(*args, **kwargs) - self.weights = None - self.set_weights(weights) - self._link_all_weights() - self.override_df_functions() - - def override_df_functions(self) -> None: - for name in MicroSeries.FUNCTIONS: - - def get_fn(name) -> Callable: - def fn(*args, **kwargs) -> Union[pd.Series, pd.DataFrame]: - is_array = len(args) > 0 and hasattr(args[0], "__len__") - if ( - name in MicroSeries.SCALAR_FUNCTIONS - or name in MicroSeries.AGNOSTIC_FUNCTIONS - and not is_array - ): - results = pd.Series( - [ - getattr(self[col], name)(*args, **kwargs) - for col in self.columns - ] - ) - results.index = self.columns - return results - elif ( - name in MicroSeries.VECTOR_FUNCTIONS - or name in MicroSeries.AGNOSTIC_FUNCTIONS - and is_array - ): - results = pd.DataFrame( - [ - getattr(self[col], name)(*args, **kwargs) - for col in self.columns - ] - ) - results.index = self.columns - return results - - return fn - - setattr(self, name, get_fn(name)) - - def get_args_as_micro_series(*kwarg_names: tuple) -> Callable: - """Decorator for auto-parsing column names into MicroSeries objects. If - given, kwarg_names limits arguments checked to keyword arguments - specified. - - :param arg_names: argument names to restrict to. - :type arg_names: str - """ - - def arg_series_decorator(fn) -> Callable: - @wraps(fn) - def series_function( - self, *args, **kwargs - ) -> Union[pd.Series, pd.DataFrame]: - new_args = [] - new_kwargs = {} - if len(kwarg_names) == 0: - for value in args: - if isinstance(value, str): - if value not in self.columns: - raise Exception("Column not found") - new_args += [self[value]] - else: - new_args += [value] - for name, value in kwargs.items(): - if isinstance(value, str) and ( - len(kwarg_names) == 0 or name in kwarg_names - ): - if value not in self.columns: - raise Exception("Column not found") - new_kwargs[name] = self[value] - else: - new_kwargs[name] = value - return fn(self, *new_args, **new_kwargs) - - return series_function - - return arg_series_decorator - - def __setitem__(self, *args, **kwargs) -> None: - super().__setitem__(*args, **kwargs) - self._link_all_weights() - - def _link_weights(self, column) -> None: - # self[column] = ... triggers __setitem__, which forces pd.Series - # this workaround avoids that - self[column].__class__ = MicroSeries - self[column].set_weights(self.weights) - - def _link_all_weights(self) -> None: - if self.weights is None: - self.set_weights(np.ones((len(self)))) - for column in self.columns: - if column != self.weights_col: - self._link_weights(column) - - def set_weights(self, weights: np.ndarray) -> None: - """Sets the weights for the MicroDataFrame. If a string is received, it - will be assumed to be the column name of the weight column. - - :param weights: Array of weights. - :type weights: np.array - """ - if isinstance(weights, str): - self.weights_col = weights - self.weights = pd.Series(self[weights], dtype=float) - elif weights is not None: - self.weights_col = None - with warnings.catch_warnings(): - warnings.filterwarnings("ignore", category=UserWarning) - self.weights = pd.Series(weights, dtype=float) - self._link_all_weights() - - def set_weight_col(self, column: str) -> None: - """Sets the weights for the MicroDataFrame by specifying the name of - the weight column. - - :param weights: Array of weights. - :type weights: np.array - """ - self.weights = np.array(self[column]) - self.weight_col = column - self._link_all_weights() - - def __getitem__( - self, key: Union[str, List] - ) -> Union[pd.Series, pd.DataFrame]: - result = super().__getitem__(key) - if isinstance(result, pd.DataFrame): - try: - weights = self.weights[key] - except Exception: - weights = self.weights - return MicroDataFrame(result, weights=weights) - return result - - def catch_series_relapse(self) -> None: - for col in self.columns: - if self[col].__class__ == pd.Series: - self._link_weights(col) - - def __setattr__(self, key, value) -> None: - super().__setattr__(key, value) - self.catch_series_relapse() - - def reset_index( - self, - level: Optional[int] = None, - drop: Optional[bool] = False, - inplace: Optional[bool] = False, - col_level: Optional[int] = 0, - col_fill: Optional[str] = "", - allow_duplicates: Optional[bool] = None, - names: Optional[List[str]] = None, - ) -> Union["MicroDataFrame", None]: - """Reset the index of the MicroDataFrame. - - This method supports all parameters of pandas DataFrame.reset_index(), - including the 'inplace' parameter. - - :param level: Only remove the given levels from the index. Removes all - levels by default. - :param drop: Do not try to insert index into dataframe columns. This - resets the index to the default integer index. - :param inplace: Modify the DataFrame in place (do not create a new - object). - :param col_level: If the columns have multiple levels, determines which - level the labels are inserted into. - :param col_fill: If the columns have multiple levels, determines how - the other levels are named. - :param allow_duplicates: Allow duplicate column labels to be created. - :param names: Using the given string, rename the DataFrame column which - contains the index data. - :return: MicroDataFrame with reset index or None if inplace=True. - """ - if inplace: - weights_backup = self.weights.copy() - # Perform in-place reset on the parent DataFrame - super().reset_index( - level=level, - drop=drop, - inplace=True, - col_level=col_level, - col_fill=col_fill, - allow_duplicates=allow_duplicates, - names=names, - ) - self.weights = weights_backup - self._link_all_weights() - return None - else: - res = super().reset_index( - level=level, - drop=drop, - inplace=False, - col_level=col_level, - col_fill=col_fill, - allow_duplicates=allow_duplicates, - names=names, - ) - return MicroDataFrame(res, weights=self.weights) - - def copy(self, deep: Optional[bool] = True) -> "MicroDataFrame": - res = super().copy(deep) - # This changes the original columns to Series. Undo it: - for col in self.columns: - self[col] = MicroSeries(self[col]) - res = MicroDataFrame(res, weights=self.weights.copy(deep)) - return res - - def equals(self, other: "MicroDataFrame") -> bool: - equal_values = super().equals(other) - equal_weights = self.weights.equals(other.weights) - return equal_values and equal_weights - - @get_args_as_micro_series() - def groupby( - self, by: Union[str, List], *args, **kwargs - ) -> "MicroDataFrameGroupBy": - """Returns a GroupBy object with MicroSeriesGroupBy objects for each - column. - - :param by: column to group by - :type by: Union[str, List] - - return: DataFrameGroupBy object with columns using weights - rtype: DataFrameGroupBy - """ - self["__tmp_weights"] = self.weights - gb = super().groupby(by, *args, **kwargs) - weights = copy.deepcopy(gb["__tmp_weights"]) - for col in self.columns: # df.groupby(...)[col]s use weights - res = gb[col] - res.__class__ = MicroSeriesGroupBy - res._init() - res.weights = weights - setattr(gb, col, res) - gb.__class__ = MicroDataFrameGroupBy - gb._init(by) - return gb - - @get_args_as_micro_series() - def poverty_rate(self, income: str, threshold: str) -> float: - """Calculate poverty rate, i.e., the population share with income below - their poverty threshold. - - :param income: Column indicating income. - :type income: str - :param threshold: Column indicating threshold. - :type threshold: str - :return: Poverty rate between zero and one. - :rtype: float - """ - pov = income < threshold - return pov.sum() / pov.count() - - @get_args_as_micro_series() - def deep_poverty_rate(self, income: str, threshold: str) -> float: - """Calculate deep poverty rate, i.e., the population share with income - below half their poverty threshold. - - :param income: Column indicating income. - :type income: str - :param threshold: Column indicating threshold. - :type threshold: str - :return: Deep poverty rate between zero and one. - :rtype: float - """ - pov = income < (threshold / 2) - return pov.sum() / pov.count() - - @get_args_as_micro_series() - def poverty_gap(self, income: str, threshold: str) -> float: - """Calculate poverty gap, i.e., the total gap between income and - poverty thresholds for all people in poverty. - - :param income: Column indicating income. - :type income: str - :param threshold: Column indicating threshold. - :type threshold: str - :return: Poverty gap. - :rtype: float - """ - gaps = (threshold - income)[threshold > income] - return gaps.sum() - - @get_args_as_micro_series() - def deep_poverty_gap(self, income: str, threshold: str) -> float: - """Calculate deep poverty gap, i.e., the total gap between income and - half of poverty thresholds for all people in deep poverty. - - :param income: Column indicating income. - :type income: str - :param threshold: Column indicating threshold. - :type threshold: str - :return: Deep poverty gap. - :rtype: float - """ - deep_threshold = threshold / 2 - gaps = (deep_threshold - income)[deep_threshold > income] - return gaps.sum() - - @get_args_as_micro_series() - def squared_poverty_gap(self, income: str, threshold: str) -> float: - """Calculate squared poverty gap, i.e., the total squared gap between - income and poverty thresholds for all people in poverty. Also known as - the poverty severity index. - - :param income: Column indicating income. - :type income: str - :param threshold: Column indicating threshold. - :type threshold: str - :return: Squared poverty gap. - :rtype: float - """ - gaps = (threshold - income)[threshold > income] - squared_gaps = gaps**2 - return squared_gaps.sum() - - @get_args_as_micro_series() - def poverty_count( - self, - income: Union[MicroSeries, str], - threshold: Union[MicroSeries, str], - ) -> int: - """Calculates the number of entities with income below a poverty - threshold. - - :param income: income array or column name - :type income: Union[MicroSeries, str] - - :param threshold: threshold array or column name - :type threshold: Union[MicroSeries, str] - - return: number of entities in poverty - rtype: int - """ - in_poverty = income < threshold - return in_poverty.sum() - - def astype( - self, - dtype, - copy: Optional[bool] = True, - errors: Optional[str] = "raise", - ) -> "MicroDataFrame": - """Convert MicroDataFrame to specified data type while preserving - weights. - - :param dtype: Data type to convert to. Can be numpy dtype, Python type, - or dict. - :param copy: Whether to make a copy of the data (default True). - :param errors: How to handle conversion errors (default "raise"). - :return: New MicroDataFrame with converted data types and preserved - weights. - """ - converted_df = super().astype(dtype, copy=copy, errors=errors) - return MicroDataFrame( - converted_df, weights=self.weights.copy() if copy else self.weights - ) - - def __repr__(self) -> str: - df = pd.DataFrame(self) - df["weight"] = self.weights - return df[[df.columns[-1]] + list(df.columns[:-1])].__repr__() diff --git a/microdf/tests/test_generic.py b/microdf/tests/test_generic.py index fb7d11e6..d9648dd2 100644 --- a/microdf/tests/test_generic.py +++ b/microdf/tests/test_generic.py @@ -2,7 +2,8 @@ import pandas as pd import microdf as mdf -from microdf.generic import MicroDataFrame, MicroSeries +from microdf.microdataframe import MicroDataFrame, MicroDataFrameGroupBy +from microdf.microseries import MicroSeries, MicroSeriesGroupBy def test_df_init() -> None: From d08cd5485aa549df27831a2bc77a97802b86eb56 Mon Sep 17 00:00:00 2001 From: juaristi22 Date: Fri, 18 Jul 2025 18:29:04 +0200 Subject: [PATCH 6/9] rename test_generic.py --- .../tests/{test_generic.py => test_microseries_dataframe.py} | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) rename microdf/tests/{test_generic.py => test_microseries_dataframe.py} (98%) diff --git a/microdf/tests/test_generic.py b/microdf/tests/test_microseries_dataframe.py similarity index 98% rename from microdf/tests/test_generic.py rename to microdf/tests/test_microseries_dataframe.py index d9648dd2..8325e49d 100644 --- a/microdf/tests/test_generic.py +++ b/microdf/tests/test_microseries_dataframe.py @@ -2,8 +2,8 @@ import pandas as pd import microdf as mdf -from microdf.microdataframe import MicroDataFrame, MicroDataFrameGroupBy -from microdf.microseries import MicroSeries, MicroSeriesGroupBy +from microdf.microdataframe import MicroDataFrame +from microdf.microseries import MicroSeries def test_df_init() -> None: From 803730f5b202e94a221ee72044f96b84e45acb64 Mon Sep 17 00:00:00 2001 From: juaristi22 Date: Mon, 21 Jul 2025 14:41:11 +0200 Subject: [PATCH 7/9] add option to preserve old weights in set_weights methods --- changelog_entry.yaml | 1 + microdf/microdataframe.py | 22 +++++++++++++++++++--- microdf/microseries.py | 9 ++++++++- 3 files changed, 28 insertions(+), 4 deletions(-) diff --git a/changelog_entry.yaml b/changelog_entry.yaml index c756d312..e656e3cf 100644 --- a/changelog_entry.yaml +++ b/changelog_entry.yaml @@ -4,3 +4,4 @@ - Add astype and sqrt methods to MicroSeries and MicroDataFrame. - Support in-place reset_index. - Split generic.py into microdataframe.py and microseries.py. + - Add optional preserve_old parameter when setting weights. diff --git a/microdf/microdataframe.py b/microdf/microdataframe.py index 31f4a795..9d133b38 100644 --- a/microdf/microdataframe.py +++ b/microdf/microdataframe.py @@ -119,13 +119,22 @@ def _link_all_weights(self) -> None: if column != self.weights_col: self._link_weights(column) - def set_weights(self, weights: np.ndarray) -> None: + def set_weights( + self, + weights: Union[np.ndarray, str], + preserve_old: Optional[bool] = False, + ) -> None: """Sets the weights for the MicroDataFrame. If a string is received, it will be assumed to be the column name of the weight column. :param weights: Array of weights. + :param preserve_old: If True, keeps the old weights as a column when + new weights are provided. :type weights: np.array """ + if preserve_old and self.weights_col is not None: + self["old_" + self.weights_col] = self.weights + if isinstance(weights, str): self.weights_col = weights self.weights = pd.Series(self[weights], dtype=float) @@ -136,15 +145,22 @@ def set_weights(self, weights: np.ndarray) -> None: self.weights = pd.Series(weights, dtype=float) self._link_all_weights() - def set_weight_col(self, column: str) -> None: + def set_weight_col( + self, column: str, preserve_old: Optional[bool] = False + ) -> None: """Sets the weights for the MicroDataFrame by specifying the name of the weight column. :param weights: Array of weights. + :param preserve_old: If True, keeps the old weights as a column when + new weights are provided. :type weights: np.array """ + if preserve_old and self.weights_col is not None: + self["old_" + self.weights_col] = self.weights + self.weights = np.array(self[column]) - self.weight_col = column + self.weights_col = column self._link_all_weights() def __getitem__( diff --git a/microdf/microseries.py b/microdf/microseries.py index 89fbd00b..2cb796c8 100644 --- a/microdf/microseries.py +++ b/microdf/microseries.py @@ -39,15 +39,22 @@ def vector_function(fn: Callable) -> Callable: fn._rtype = pd.Series return fn - def set_weights(self, weights: np.array) -> None: + def set_weights( + self, weights: np.array, preserve_old: Optional[bool] = False + ) -> None: """Sets the weight values. :param weights: Array of weights. + :param preserve_old: If True, keeps the old weights as a column when + new weights are provided. :type weights: np.array. """ if weights is None: self.weights = pd.Series(np.ones_like(self.values), dtype=float) else: + if preserve_old and self.weights is not None: + self["old_weights"] = self.weights + self.weights = pd.Series(weights, dtype=float) @vector_function From f4f7384876891b0c49a97adb84fffce3fdf45e12 Mon Sep 17 00:00:00 2001 From: juaristi22 Date: Tue, 22 Jul 2025 13:32:32 +0200 Subject: [PATCH 8/9] add check to set_weights for matching lengths --- microdf/microdataframe.py | 8 +++++++- microdf/microseries.py | 11 ++++++++++- microdf/tests/test_microseries_dataframe.py | 7 ++++--- microdf/tests/test_weighted.py | 4 ++-- 4 files changed, 23 insertions(+), 7 deletions(-) diff --git a/microdf/microdataframe.py b/microdf/microdataframe.py index 9d133b38..3e8dd591 100644 --- a/microdf/microdataframe.py +++ b/microdf/microdataframe.py @@ -114,7 +114,8 @@ def _link_weights(self, column) -> None: def _link_all_weights(self) -> None: if self.weights is None: - self.set_weights(np.ones((len(self)))) + if len(self) > 0: + self.set_weights(np.ones((len(self)))) for column in self.columns: if column != self.weights_col: self._link_weights(column) @@ -139,6 +140,11 @@ def set_weights( self.weights_col = weights self.weights = pd.Series(self[weights], dtype=float) elif weights is not None: + if len(weights) != len(self): + raise ValueError( + f"Length of weights ({len(weights)}) does not match " + f"length of DataFrame ({len(self)})." + ) self.weights_col = None with warnings.catch_warnings(): warnings.filterwarnings("ignore", category=UserWarning) diff --git a/microdf/microseries.py b/microdf/microseries.py index 2cb796c8..ae73bd2a 100644 --- a/microdf/microseries.py +++ b/microdf/microseries.py @@ -50,8 +50,17 @@ def set_weights( :type weights: np.array. """ if weights is None: - self.weights = pd.Series(np.ones_like(self.values), dtype=float) + if len(self) > 0: + self.weights = pd.Series( + np.ones_like(self.values), dtype=float + ) else: + if len(weights) != len(self): + raise ValueError( + f"Length of weights ({len(weights)}) does not match " + f"length of DataFrame ({len(self)})." + ) + if preserve_old and self.weights is not None: self["old_weights"] = self.weights diff --git a/microdf/tests/test_microseries_dataframe.py b/microdf/tests/test_microseries_dataframe.py index 8325e49d..0cfc4657 100644 --- a/microdf/tests/test_microseries_dataframe.py +++ b/microdf/tests/test_microseries_dataframe.py @@ -80,9 +80,10 @@ def test_mean() -> None: def test_poverty_count() -> None: arr = np.array([10000, 20000, 50000]) w = np.array([1123, 1144, 2211]) - df = MicroDataFrame(weights=w) + df = pd.DataFrame() df["income"] = arr df["threshold"] = 16000 + df = MicroDataFrame(df, weights=w) assert df.poverty_count("income", "threshold") == w[0] @@ -171,10 +172,10 @@ def test_quintile_rank() -> None: assert np.array_equal(s.quintile_rank().values, [5, 3, 4]) -def test_decile_rank_rank() -> None: +def test_decile_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], + weights=[10, 20, 10, 10, 10, 10, 10, 10, 10], ) assert np.array_equal(s.decile_rank().values, [6, 5, 3, 2, 1, 7, 8, 9, 10]) diff --git a/microdf/tests/test_weighted.py b/microdf/tests/test_weighted.py index a6231573..57e4b757 100644 --- a/microdf/tests/test_weighted.py +++ b/microdf/tests/test_weighted.py @@ -15,7 +15,7 @@ df2.y *= 1.5 dfg = pd.concat([df, df2]) dfg["g"] = ["a"] * 3 + ["b"] * 3 -mdg = mdf.MicroDataFrame(dfg[["x", "y", "g"]], weights=W) +mdg = mdf.MicroDataFrame(dfg[["x", "y", "g"]], weights=dfg["w"]) def test_weighted_quantile() -> None: @@ -31,7 +31,7 @@ def test_weighted_median() -> None: def test_weighted_mean() -> None: - # Test umweighted. + # Test unweighted. assert mdf.weighted_mean(df, "x") == 8 / 3 # Test weighted. assert mdf.weighted_mean(df, "x", "w") == 11 / 6 From 227aaa31307383a00b977617e5d3a60fb93f2486 Mon Sep 17 00:00:00 2001 From: juaristi22 Date: Tue, 22 Jul 2025 13:44:27 +0200 Subject: [PATCH 9/9] break down override_df_functions --- microdf/microdataframe.py | 108 +++++++++++++++++++++++++++----------- 1 file changed, 77 insertions(+), 31 deletions(-) diff --git a/microdf/microdataframe.py b/microdf/microdataframe.py index 3e8dd591..e749c8aa 100644 --- a/microdf/microdataframe.py +++ b/microdf/microdataframe.py @@ -27,41 +27,87 @@ def __init__(self, *args, weights=None, **kwargs): self.override_df_functions() def override_df_functions(self) -> None: + """Override DataFrame functions to work with weighted operations.""" for name in MicroSeries.FUNCTIONS: + if name in MicroSeries.SCALAR_FUNCTIONS: + setattr(self, name, self._create_scalar_function(name)) + elif name in MicroSeries.VECTOR_FUNCTIONS: + setattr(self, name, self._create_vector_function(name)) + elif name in MicroSeries.AGNOSTIC_FUNCTIONS: + setattr(self, name, self._create_agnostic_function(name)) + + def _create_scalar_function(self, name: str) -> Callable: + """Create a scalar function that returns a Series of results. + + :param name: Name of the function to create + :return: Function that applies the operation to all columns + """ - def get_fn(name) -> Callable: - def fn(*args, **kwargs) -> Union[pd.Series, pd.DataFrame]: - is_array = len(args) > 0 and hasattr(args[0], "__len__") - if ( - name in MicroSeries.SCALAR_FUNCTIONS - or name in MicroSeries.AGNOSTIC_FUNCTIONS - and not is_array - ): - results = pd.Series( - [ - getattr(self[col], name)(*args, **kwargs) - for col in self.columns - ] - ) - results.index = self.columns - return results - elif ( - name in MicroSeries.VECTOR_FUNCTIONS - or name in MicroSeries.AGNOSTIC_FUNCTIONS - and is_array - ): - results = pd.DataFrame( - [ - getattr(self[col], name)(*args, **kwargs) - for col in self.columns - ] - ) - results.index = self.columns - return results + def fn(*args, **kwargs) -> pd.Series: + results = pd.Series( + [ + getattr(self[col], name)(*args, **kwargs) + for col in self.columns + ] + ) + results.index = self.columns + return results - return fn + return fn + + def _create_vector_function(self, name: str) -> Callable: + """Create a vector function that returns a DataFrame of results. + + :param name: Name of the function to create + :return: Function that applies the operation to all columns + """ + + def fn(*args, **kwargs) -> pd.DataFrame: + results = pd.DataFrame( + [ + getattr(self[col], name)(*args, **kwargs) + for col in self.columns + ] + ) + results.index = self.columns + return results + + return fn + + def _create_agnostic_function(self, name: str) -> Callable: + """Create a function that can be either scalar or vector based on + input. + + :param name: Name of the function to create + :return: Function that applies the operation to all columns + """ + + def fn(*args, **kwargs) -> Union[pd.Series, pd.DataFrame]: + # Check if first argument is array-like + is_array = len(args) > 0 and hasattr(args[0], "__len__") + + if is_array: + # Use vector function behavior + results = pd.DataFrame( + [ + getattr(self[col], name)(*args, **kwargs) + for col in self.columns + ] + ) + results.index = self.columns + return results + else: + # Use scalar function behavior + results = pd.Series( + [ + getattr(self[col], name)(*args, **kwargs) + for col in self.columns + ] + ) + results.index = self.columns + return results - setattr(self, name, get_fn(name)) + return fn def get_args_as_micro_series(*kwarg_names: tuple) -> Callable: """Decorator for auto-parsing column names into MicroSeries objects. If