Added subsetting to aggregate and impute_base_prices

The subsetting using the function subset_shared_axis added to helpers.py
is useful in the event of the functions being applied in a groupby.
Because the arguments are not also chunked, into groups, this extra step
does that so that large DataFrame's aren't being computed each time.

Author: mitches-got-glitches

precon/aggregation.py

@@ -12,7 +12,7 @@
 )
 
 from precon.weights import get_weight_shares, reindex_weights_to_indices
-from precon.helpers import flip
+from precon.helpers import flip, subset_shared_axis
 from precon._validation import _handle_axis
 
 
@@ -64,6 +64,10 @@ def aggregate(
         'geomean': _geo_mean_aggregate,
     }
     agg_method = methods_lib.get(method)
+    
+    # Subset the metadata axis to match those of indices, for quicker
+    # handling of function when applied by groupby.
+    weights = subset_shared_axis(weights, indices, axis)
 
     # Make sure that the indices and weights have the same time series
     # axis before aggregating.

precon/helpers.py

@@ -1,6 +1,8 @@
 # -*- coding: utf-8 -*-
-import pandas as pd
+from typing import Sequence, Union, Optional
 
+import pandas as pd
+from pandas._typing import Level
 
 def reindex_and_fill(df, other, first='ffill', axis=0):
     """Reindex and fill the DataFrame or Series by given index and axis.
@@ -156,3 +158,48 @@ def reduce_to_only_differing_periods(df, axis):
     series to only the periods where the values have changed.
     """
     return df[df.ne(df.shift(1, axis=axis))].dropna()
+
+
+def subset_shared_axis(
+        df: pd.DataFrame,
+        other: pd.DataFrame,
+        axis: int = 0,
+        droplevel: Optional[Union[Level, Sequence[Level]]] = None,
+        ) -> pd.DataFrame:
+    """Subsets a DataFrame by it's shared axis with the other.
+
+    Optional behaviour to drop levels of the other axis before
+    looking for the shared axis.
+
+    Parameters
+    ----------
+    df : DataFrame
+    other : Object of the same data type
+        Its indices on the given axis are used to define the subset of
+        indices for this object.
+    axis : {0, 1} int, default 0
+        The axis to subset the shared index.
+    droplevel : int, str, or list-like
+        If a string is given, must be the name of a level. If
+        list-like, elements must be names or positional indexes of
+        levels.
+
+    Returns
+    -------
+    DataFrame
+        The subsetted frame.
+
+    """
+    if not df.axes[axis].eq(other.axes[axis]):
+        
+        other_axis = other.axes[axis]
+        
+        if droplevel:
+            other_axis = other_axis.droplevel(droplevel)
+        
+        shared_axis = df.axes[axis].isin(other_axis)
+        return df.loc[axis_slice(shared_axis, axis)]
+    
+    else:
+        return df
+

precon/imputation.py

@@ -11,7 +11,7 @@
 
 from precon._validation import _handle_axis
 from precon.index_methods import calculate_index
-from precon.helpers import flip, axis_slice
+from precon.helpers import flip, axis_slice, subset_shared_axis
 from precon.weights import reindex_weights_to_indices
 
 
@@ -64,6 +64,14 @@ def impute_base_prices(
     """
     axis = _handle_axis(axis)
 
+    # Subset the metadata axis to match those of indices, for quicker
+    # handling of function when applied by groupby.    
+    to_impute = subset_shared_axis(to_impute, prices, flip(axis))
+    if weights is not None:
+        weights = subset_shared_axis(weights, prices, flip(axis))
+    if adjustments is not None:
+        adjustments = subset_shared_axis(adjustments, prices, flip(axis))
+    
     # Ensure the weights are in the same shape as the prices and
     # exclude the prices to impute from the imputation index
     # calculation by setting weights to zero.