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_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")