"""
Utility classes and functions that the models share.
"""
from __future__ import annotations
from collections.abc import (
Callable,
ItemsView,
Iterable,
Iterator,
KeysView,
Mapping,
MutableMapping,
Sequence,
ValuesView,
)
from typing import Any, Protocol, TypeVar, cast
import numpy as np
from pandas import DataFrame, Index, MultiIndex, Series
from linearmodels.typing import AnyArray, Label
_KT = TypeVar("_KT")
_VT_co = TypeVar("_VT_co", covariant=True)
class SupportsKeysAndGetItem(Protocol[_KT, _VT_co]):
"""Protocol for objects that have ``keys`` and can be indexed by their keys."""
def keys(self) -> Iterable[_KT]:
"""Return the keys."""
...
def __getitem__(self, __k: _KT, /) -> _VT_co: ...
def _new_attr_dict_(*args: Iterable[tuple[Any, Any]]) -> AttrDict:
attr_dict = AttrDict()
for k, v in args:
attr_dict[k] = v
return attr_dict
[docs]
class AttrDict(MutableMapping):
"""
Ordered dictionary-like object that exposes keys as attributes.
Parameters
----------
*args : mapping or sequence of tuple
A mapping, or a sequence of (key, value) pairs, with the initial items.
The arguments are those of ``dict``.
**kwargs
Items to add, with the names of the keyword arguments as keys.
"""
[docs]
def update(
self,
*args: SupportsKeysAndGetItem[Any, Any] | Iterable[tuple[Any, Any]],
**kwargs: Any,
) -> None:
"""
Update AD from dictionary or iterable E and F.
If E is present and has a .keys() method, then does: for k in E: AD[k] = E[k]
If E is present and lacks a .keys() method, then does: for k, v in E: AD[k] = v
In either case, this is followed by: for k in F: AD[k] = F[k]
Parameters
----------
*args : mapping or iterable of tuple
At most one mapping, or iterable of (key, value) pairs, which is E.
**kwargs
The items of F, with the names of the keyword arguments as keys.
"""
self.__private_dict__.update(*args, **kwargs)
[docs]
def clear(self) -> None:
"""Remove all items from AD."""
self.__private_dict__.clear()
[docs]
def copy(self) -> AttrDict:
"""
Create a shallow copy of AD.
Returns
-------
AttrDict
A copy that has the same keys and values.
"""
ad = AttrDict()
for key in self.__private_dict__.keys():
ad[key] = self.__private_dict__[key]
return ad
[docs]
def keys(self) -> KeysView[Any]:
"""
Return an ordered list-like object providing a view on AD's keys.
Returns
-------
KeysView
View of the keys.
"""
return self.__private_dict__.keys()
[docs]
def items(self) -> ItemsView[Any, Any]:
"""
Return an ordered list-like object providing a view on AD's items.
Returns
-------
ItemsView
View of the (key, value) pairs.
"""
return self.__private_dict__.items()
[docs]
def values(self) -> ValuesView[Any]:
"""
Return an ordered list-like object providing a view on AD's values.
Returns
-------
ValuesView
View of the values.
"""
return self.__private_dict__.values()
[docs]
def pop(self, key: Label, default: Any = None) -> Any:
"""
Remove specified key and return the corresponding value.
Parameters
----------
key : hashable
The key to remove.
default : object, optional
The value to return if the key is not found. The default is None.
Returns
-------
object
The value of the key, or ``default`` if the key is not found.
"""
return self.__private_dict__.pop(key, default)
def __reduce__(
self,
) -> tuple[
Callable[[Iterable[tuple[Any, Any]]], AttrDict], tuple[tuple[Any, Any], ...]
]:
return _new_attr_dict_, tuple((k, v) for k, v in self.items())
def __len__(self) -> int:
return self.__private_dict__.__len__()
def __repr__(self) -> str:
out = self.__private_dict__.__str__()
return "AttrDict" + out
def __str__(self) -> str:
return self.__repr__()
def __init__(
self, *args: Mapping[Any, Any] | Sequence[tuple[Any, Any]], **kwargs: Any
) -> None:
self.__dict__["__private_dict__"] = dict(*args, **kwargs)
def __contains__(self, item: Label) -> bool:
return self.__private_dict__.__contains__(item)
def __getitem__(self, item: Label) -> Any:
return self.__private_dict__[item]
def __setitem__(self, key: Label, value: Any) -> None:
if key == "__private_dict__":
raise KeyError("__private_dict__ is reserved and cannot be set.")
self.__private_dict__[key] = value
def __delitem__(self, key: Label) -> None:
del self.__private_dict__[key]
def __getattr__(self, key: Label) -> Any:
if key not in self.__private_dict__:
raise AttributeError
return self.__private_dict__[key]
def __setattr__(self, key: Label, value: Any) -> None:
if key == "__private_dict__":
raise AttributeError("__private_dict__ is invalid")
self.__private_dict__[key] = value
def __delattr__(self, key: Label) -> None:
del self.__private_dict__[key]
def __dir__(self) -> Iterable[str]:
out = [str(key) for key in self.__private_dict__.keys()]
out += list(super().__dir__())
filtered = [key for key in out if key.isidentifier()]
return sorted(set(filtered))
def __iter__(self) -> Iterator[Label]:
return self.__private_dict__.__iter__()
def ensure_unique_column(col_name: str, df: DataFrame, addition: str = "_") -> str:
"""
Create a column name that is not a column of a DataFrame.
Parameters
----------
col_name : str
The preferred name.
df : DataFrame
The frame that must not have a column with the name.
addition : str, optional
The string that is added before and after the name until it is not the
name of a column of ``df``. The default is ``"_"``.
Returns
-------
str
``col_name`` if it is not a column of ``df``, and otherwise the name with
``addition`` added at both ends as often as needed.
"""
while col_name in df:
col_name = addition + col_name + addition
return col_name
def panel_to_frame(
x: AnyArray | None,
items: Sequence[Label] | Index,
major_axis: Sequence[Label] | Index,
minor_axis: Sequence[Label] | Index,
swap: bool = False,
) -> DataFrame:
"""
Construct a multiindex DataFrame using Panel-like arguments.
Parameters
----------
x : ndarray or None
3-d array with shape (nitem, nmajor, nminor). If None, the frame has no
values.
items : sequence or Index
Labels of the items, which become the columns.
major_axis : sequence or Index
Labels of the major axis.
minor_axis : sequence or Index
Labels of the minor axis.
swap : bool, optional
Whether to swap the major and minor axes. The default is False.
Returns
-------
DataFrame
The frame, with a two-level index named ``"major"`` and ``"minor"``.
Notes
-----
This function is equivalent to
Panel(x, items, major_axis, minor_axis).to_frame()
if ``swap`` is True, it is equivalent to
Panel(x, items, major_axis, minor_axis).swapaxes(1,2).to_frame()
"""
nmajor = np.arange(len(major_axis)).tolist()
nminor = np.arange(len(minor_axis)).tolist()
final_levels = [major_axis, minor_axis]
mi = MultiIndex.from_product([nmajor, nminor])
if x is not None:
shape = x.shape
x = x.reshape((shape[0], shape[1] * shape[2])).T
df = DataFrame(x, columns=Index(items), index=mi)
if swap:
df.index = mi.swaplevel()
df = df.sort_index()
final_levels = [minor_axis, major_axis]
mi_index = cast("MultiIndex", df.index)
df.index = mi_index.set_levels(levels=final_levels, level=[0, 1])
df.index.names = ["major", "minor"]
return df
class DataFrameWrapper:
"""
Wrapper around a pandas DataFrame deferring frame's construction.
Parameters
----------
values : ndarray
The data to use in the DataFrame's constructor.
columns : list of str, optional
The column names of the frame. The default is the integers from 0.
index : Index or list of str, optional
The index to use. The default is the integers from 0.
"""
def __init__(
self,
values: AnyArray,
*,
columns: list[str] | None = None,
index: Index | list[str] | None = None,
) -> None:
self._values = values
self._columns = columns
if isinstance(index, list):
index = Index(index)
self._index: Index | None = index
def __call__(self) -> DataFrame:
return DataFrame(self._values, columns=self._columns, index=self._index)
class SeriesWrapper:
"""
Wrapper around a pandas Series deferring series' construction.
Parameters
----------
values : ndarray
The data to use in the Series' constructor.
name : str, optional
The name of the series. The default is no name.
index : Index or list of str, optional
The index to use. The default is the integers from 0.
"""
def __init__(
self,
values: AnyArray,
*,
name: str | None = None,
index: Index | list[str] | None = None,
) -> None:
self._values = values
self._name = name
if isinstance(index, list):
index = Index(index)
self._index: Index | None = index
def __call__(self) -> Series:
return Series(self._values, name=self._name, index=self._index)