Merge pull request #33 from ONSBigData/i28/revisit-index-calculator

I28/revisit index calculator

Author: mitches-got-glitches

precon/__init__.py

@@ -10,8 +10,8 @@
     impute_base_prices,
     get_base_prices,
     get_quality_adjusted_prices,
-    base_price_fill_shift,
-    
+    ffill_shift,
+
 )
 from precon.chaining import chain, unchain
 from precon.contributions import contributions, contributions_with_double_update

precon/base_prices.py

@@ -12,8 +12,9 @@
 from precon._validation import _handle_axis, _list_convert
 from precon.index_methods import calculate_index
 from precon.helpers import (
-    flip, 
+    axis_slice,
     axis_vals_as_frame,
+    flip,
     subset_shared_axis,
 )
 from precon.weights import reindex_weights_to_indices
@@ -69,13 +70,13 @@ 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.    
+    # 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.
@@ -84,7 +85,12 @@ def impute_base_prices(
         weights = weights.mask(to_impute, 0)
 
     # Get the base prices to start with from given base period.
-    start_prices = get_base_prices(prices, base_period, axis=axis, ffill_shift=False)
+    start_prices = get_base_prices(
+        prices,
+        base_period,
+        axis=axis,
+        fill_shift=False,
+    )
     base_prices = start_prices.copy()
 
     if not shift_imputed_values:
@@ -153,27 +159,22 @@ def impute_base_prices(
 
 
 def get_base_prices(
-        prices: pd.DataFrame,
-        base_period: Union[int, Sequence[int]] = 1,
-        axis: pd._typing.Axis = 0,
-        ffill_shift: bool = True,
-        ) -> pd.DataFrame:
-    """Return prices at base month with optional ffill and shift.
-
-    Default behaviour is to fill forward values within the year and
-    shift one period, since base prices usually start being used in
-    the following period up to the next base period. Will return
-    NaNs in non-base month if ffill=False.
+    prices: pd.DataFrame,
+    base_period: Union[int, Sequence[int]] = 1,
+    axis: pd._typing.Axis = 0,
+    fill_shift: bool = True,
+) -> pd.DataFrame:
+    """Return base prices with optional fill and shift.
 
     Parameters
     ----------
     prices : DataFrame
-    base_period : int, or list of ints
-        The base periods to select base prices from.
+    base_period : int, or list of ints, defaults to 1
+        Base period/s to select base prices from.
     axis : {0, 1} int, defaults to 0
         Fill and shift direction.
-    ffill_shift : bool, defaults to True
-        Switch to forward fill values within the year and shift by one
+    fill_shift : bool, defaults to True
+        Switch to forward fill base prices within year and shift one
         period.
 
     Returns
@@ -183,35 +184,37 @@ def get_base_prices(
 
     Notes
     -----
-    The base prices are forward filled within each year so that base
-    prices are not filled when prices have stopped being collected.
-    When shifting, base prices are shifted by one period. So for a base
-    period of Jan (int=1) base prices are shifted on to the Feb-Jan+1
-    time delta in which they apply. A base price is needed for the Jan
-    period at the start of the series, so the function fills the
-    shifted values with the unshifted values to achieve this
-    
-    TODO: Make this work for any base period.
-    
+    When using fill_shift, the base prices are forward filled within
+    each year so that base prices are not filled when prices have
+    stopped being collected. Base prices are also shifted by one period,
+    so for a base period of Jan (int=1) base prices are shifted on to
+    the Feb-Jan+1 time delta in which they apply.
+
     """
     base_period = _list_convert(base_period)
-    
+
     # Only prices in the base periods are not NaN.
     months = axis_vals_as_frame(prices, axis, converter=lambda x: x.month)
     base_prices = prices.where(months.isin(base_period))
 
-    if ffill_shift:
-        # Fill base prices forward within the year and shift one.
-        return base_price_fill_shift(base_prices, axis)
-    
-    return base_prices
+    # Ensure the prices in the first period are taken as base prices
+    # even if not a period given by base_period parameter.
+    first_period = axis_slice(0, axis)
 
+    if not base_prices.iloc[first_period].isna().all():
+        base_prices.iloc[first_period] = prices.iloc[first_period]
 
-def base_price_fill_shift(
+    if fill_shift:
+        return ffill_shift(base_prices, axis)
+    else:
+        return base_prices
+
+
+def ffill_shift(
     base_prices: pd.DataFrame,
     axis: int = 0
 ) -> pd.DataFrame:
-    """Fill forward base prices and shift one period.
+    """Fill forward base prices within year and shift one period.
 
     Parameters
     ----------
@@ -233,8 +236,6 @@ def base_price_fill_shift(
     period at the start of the series, so the function fills the
     shifted values with the unshifted values to achieve this.
 
-    TODO: Make this work for any base period.
-
     """
     return (
         base_prices.groupby(lambda x: x.year, axis=axis)

precon/pipelines.py

@@ -1,57 +1,65 @@
 """A set of common pipeline functions to create National Statistics."""
-from typing import Optional
+from typing import Optional, Union, Sequence
 
 import pandas as pd
 from pandas._typing import Axis
 
 from precon._validation import _handle_axis
-from precon.base_prices import impute_base_prices, get_base_prices
+from precon.base_prices import (
+    impute_base_prices,
+    get_base_prices,
+    ffill_shift,
+)
 from precon.index_methods import calculate_index
 from precon.helpers import flip
 
 
 def index_calculator(
-        prices: pd.DataFrame,
-        index_method: str,
-        shift_imputed_values: bool = False,
-        to_impute: Optional[pd.DataFrame] = None,
-        weights: Optional[pd.DataFrame] = None,
-        adjustments: Optional[pd.DataFrame] = None,
-        exclusions: Optional[pd.DataFrame] = None,
-        base_period: int = 1,
-        axis: Axis = 1,
-        ) -> pd.Series:
+    prices: pd.DataFrame,
+    index_method: str,
+    shift_imputed_values: bool = False,
+    to_impute: Optional[pd.DataFrame] = None,
+    weights: Optional[pd.DataFrame] = None,
+    adjustments: Optional[pd.DataFrame] = None,
+    exclusions: Optional[pd.DataFrame] = None,
+    base_prices: Optional[pd.DataFrame] = None,
+    base_period: Union[int, Sequence[int]] = 1,
+    axis: Axis = 1,
+) -> pd.Series:
     """Calculates an index given prices and an index method, with
     optional arguments for base price imputation.
 
     Parameters
     ----------
-    prices: DataFrame
+    prices : DataFrame
         The prices with which to calculate the index.
-    method: {'jevons', 'dutot', 'carli', 'laspeyres', 'geometric_laspeyres'}
+    method : {'jevons', 'dutot', 'carli', 'laspeyres', 'geometric_laspeyres'}
         Method to calculate the index.
     shift_imputed_values: bool, defaults to True
         True if imputed values are shifted onto the following period.
-    to_impute: DataFrame, optional
+    to_impute : DataFrame, optional
         A boolean mask of where to impute.
-    weights: DataFrame, optional
+    weights : DataFrame, optional
         The weights to use if the index method requires it.
-    adjustments: DataFrame, optional
-        Adjustment values to apply to prices for quality adjustment. If
-        there is no adjustment for a price then the adjustment value
-        should be zero.
-    exclusions: DataFrame, optional
+    adjustments : DataFrame, optional
+        Adjustment factors to multiply by base prices for quality
+        adjustments. A factor of 1 means no adjustment takes place.
+    exclusions : DataFrame, optional
         A boolean mask of prices to exclude from the final index
         calculation.
-    base_period: int, defaults to 1
-        Base period to select initial base prices from.
+    base_prices : DataFrame, optional
+        Base prices to be forward filled and shifted before used in
+        index calculation.
+    base_period : int, or list of ints, defaults to 1
+        Base period/s to select initial base prices from.
     axis : {0 or 'index', 1 or 'columns'}, defaults to 0
         The axis that holds the time series values.
 
     Returns
     -------
     Series
         The index.
+
     """
     axis = _handle_axis(axis)
 
@@ -60,8 +68,7 @@ def index_calculator(
         # the final index calculation
         weights = weights.mask(exclusions, 0)
 
-    # Impute the base prices if necessary, if not just take the prices
-    # in the base period and fill forward
+    # Impute the base prices if necessary.
     if to_impute is not None:
         base_prices = impute_base_prices(
             prices,
@@ -74,7 +81,14 @@ def index_calculator(
             adjustments=adjustments,
         )
     else:
-        base_prices = get_base_prices(prices, base_period, axis)
+        if base_prices is not None:
+            base_prices = ffill_shift(base_prices, axis)
+        else:
+            # Get base prices from prices.
+            base_prices = get_base_prices(prices, base_period, axis)
+
+        if adjustments is not None:
+            base_prices *= adjustments
 
     return calculate_index(
         prices,