def assert_array_equal(arg1: tp.ArrayLike, arg2: tp.ArrayLike) -> None:
"""Raise exception if the first argument and the second argument have different metadata or values."""
arg1 = _to_any_array(arg1)
arg2 = _to_any_array(arg2)
assert_meta_equal(arg1, arg2)
if is_pandas(arg1) and is_pandas(arg2):
if arg1.equals(arg2):
return
elif not is_pandas(arg1) and not is_pandas(arg2):
if np.array_equal(arg1, arg2):
return
raise AssertionError(f"Arrays {arg1} and {arg2} do not match")
def to_2d(arg: tp.ArrayLike,
raw: bool = False,
expand_axis: int = 1) -> tp.AnyArray2d:
"""Reshape argument to two dimensions.
If `raw` is True, returns NumPy array.
If 1-dim, will expand along axis 1 (i.e., Series to DataFrame with one column)."""
arg = to_any_array(arg, raw=raw)
if arg.ndim == 2:
return arg
elif arg.ndim == 1:
if checks.is_series(arg):
if expand_axis == 0:
return pd.DataFrame(arg.values[None, :], columns=arg.index)
elif expand_axis == 1:
return arg.to_frame()
return np.expand_dims(arg, expand_axis)
elif arg.ndim == 0:
return arg.reshape((1, 1))
raise ValueError(
f"Cannot reshape a {arg.ndim}-dimensional array to 2 dimensions")
def soft_to_ndim(arg: tp.ArrayLike,
ndim: int,
raw: bool = False) -> tp.AnyArray:
"""Try to softly bring `arg` to the specified number of dimensions `ndim` (max 2)."""
arg = to_any_array(arg, raw=raw)
if ndim == 1:
if arg.ndim == 2:
if arg.shape[1] == 1:
if checks.is_frame(arg):
return arg.iloc[:, 0]
return arg[:, 0] # downgrade
if ndim == 2:
if arg.ndim == 1:
if checks.is_series(arg):
return arg.to_frame()
return arg[:, None] # upgrade
return arg # do nothing
def to_1d(arg: tp.ArrayLike, raw: bool = False) -> tp.AnyArray1d:
"""Reshape argument to one dimension.
If `raw` is True, returns NumPy array.
If 2-dim, will collapse along axis 1 (i.e., DataFrame with one column to Series)."""
arg = to_any_array(arg, raw=raw)
if arg.ndim == 2:
if arg.shape[1] == 1:
if checks.is_frame(arg):
return arg.iloc[:, 0]
return arg[:, 0]
if arg.ndim == 1:
return arg
elif arg.ndim == 0:
return arg.reshape((1, ))
raise ValueError(
f"Cannot reshape a {arg.ndim}-dimensional array to 1 dimension")
def to_dict(arg: tp.ArrayLike, orient: str = 'dict') -> dict:
"""Convert object to dict."""
arg = to_pd_array(arg)
if orient == 'index_series':
return {arg.index[i]: arg.iloc[i] for i in range(len(arg.index))}
return arg.to_dict(orient)