SteinUnitAverager

class unit_averaging.averager.SteinUnitAverager(focus_function, ind_estimates, ind_covar_ests)[source]

Bases: OptimalUnitAverager

Unit averaging scheme that shrinks the target estimator to the overall mean.

This class implements an MSE-optimal unit averaging scheme which treats the target unit as unrestricted and restricts all other units (in the sense of OptimalUnitAverager). The resulting estimator is an example of a Stein-like shrinkage scheme: the target unit’s estimates is optimally shrank towars the average of all other units.

Parameters:
  • focus_function (BaseFocusFunction) – Focus function expressing the transformation of interest.

  • ind_estimates (np.ndarray | list | dict[str | int, np.ndarray | list]) – Individual unit estimates. Can be a list, numpy array, or dictionary. Each unit-specific estimate should be a NumPy array or list. The first dimension of ind_estimates indexes units (rows or dictionary entries).

  • ind_covar_ests (np.ndarray | list | dict[str | int, np.ndarray | list]) – Individual unit covariance estimates. Can be a list, numpy array, or dictionary. Each unit-specific covariance estimate should be a NumPy array or list of lists. The first dimension of ind_covar_ests indexes units (rows or dictionary entries).

Attributes:
  • ind_estimates (np.ndarray) – Array of individual unit estimates.

  • ind_covar_ests (np.ndarray) – Array of individual unit covariance estimates.

  • unrestricted_units_bool (np.ndarray) – Boolean array indicating which units are unrestricted. Unlike for OptimalUnitAverager, these are automatically set during fit(). Stein average scheme leaves the target unit unrestricted and restricts all other units.

  • keys (np.ndarray) – Array of keys corresponding to the units. The individual estimates are converted to numpy arrays internally. If ind_estimates is a dictionary, the keys are preserved in the keys attribute. If ind_estimates is a list or array, keys defaults to numeric indices (0, 1, 2, …).

  • weights (np.ndarray) – The computed weights for each unit.

  • estimate (float) – The computed unit averaging estimate.

  • focus_function (BaseFocusFunction) – Focus function expressing the transformation of interest.

  • target_id (int | str) – The ID of the target unit. Initialized as None, set by calling fit().

Example

>>> from unit_averaging import SteinUnitAverager+, InlineFocusFunction
>>> 1

Methods

average(focus_function=None)[source]

Perform unit averaging with the fitted weights.

This method computes the unit averaging estimate using the fitted weights. It can accept a different focus function and reuse the fitted weights.

Parameters:

focus_function (BaseFocusFunction | None) – Focus function to use in computing the averaging estimator. Expresses the parameter of interest. If None, defaults to the focus function used in fitting.

Returns:

The unit averaging estimate.

Return type:

float

Raises:

TypeError – If weights have not been fitted yet by calling fit()

fit(target_id)[source]

Compute the unit averaging weights and the averaging estimator.

Parameters:

target_id (int | str) – ID of the target unit. This is specified in terms of the keys attribute, which are either numeric indices (if ind_estimates was an array or list) or dictionary keys (if ind_estimates was a dictionary)

Raises:

ValueError – If the target unit is not found in the keys.