from collections import Counter
from logging import getLogger
from typing import Optional, Tuple
import pandas
import numpy
from .dataframe_helpers import dataframe_shuffle
logger = getLogger("pandas-streaming")
[docs]
class ImbalancedSplitException(Exception):
"""
Raised when an imbalanced split is detected.
"""
[docs]
def train_test_split_weights(
df,
weights=None,
test_size=0.25,
train_size=None,
shuffle=True,
fail_imbalanced=0.05,
random_state=None,
):
"""
Splits a database in train/test given, every row
can have a different weight.
:param df: :class:`pandas.DataFrame` or see
:class:`StreamingDataFrame <pandas_streaming.df.dataframe.StreamingDataFrame>`
:param weights: None or weights or weights column name
:param test_size: ratio for the test partition
(if *train_size* is not specified)
:param train_size: ratio for the train partition
:param shuffle: shuffles before the split
:param fail_imbalanced: raises an exception if relative weights
difference is higher than this value
:param random_state: seed for random generators
:return: train and test :class:`pandas.DataFrame`
If the dataframe is not shuffled first, the function
will produce two datasets which are unlikely to be randomized
as the function tries to keep equal weights among both paths
without using randomness.
"""
if hasattr(df, "iter_creation"):
raise NotImplementedError( # pragma: no cover
"Not implemented yet for StreamingDataFrame."
)
if isinstance(df, numpy.ndarray):
raise NotImplementedError( # pragma: no cover
"Not implemented on numpy arrays."
)
if shuffle:
df = dataframe_shuffle(df, random_state=random_state)
if weights is None:
if test_size == 0 or train_size == 0:
raise ValueError(
f"test_size={test_size} or train_size={train_size} cannot be null (1)."
)
from sklearn.model_selection import train_test_split
return train_test_split(
df, test_size=test_size, train_size=train_size, random_state=random_state
)
if isinstance(weights, pandas.Series):
weights = list(weights)
elif isinstance(weights, str):
weights = list(df[weights])
if len(weights) != df.shape[0]:
raise ValueError(
"Dimension mismatch between weights and dataframe " # noqa: UP030
"{0} != {1}".format(df.shape[0], len(weights))
)
p = (1 - test_size) if test_size else None
if train_size is not None:
p = train_size
test_size = 1 - p
if p is None or min(test_size, p) <= 0:
raise ValueError(
f"test_size={test_size} or train_size={train_size} cannot be null (2)."
)
ratio = test_size / p
if random_state is None:
randint = numpy.random.randint
else:
state = numpy.random.RandomState(random_state)
randint = state.randint
balance = 0
train_ids = []
test_ids = []
test_weights = 0
train_weights = 0
for i in range(df.shape[0]):
w = weights[i]
if balance == 0:
h = randint(0, 1)
totest = h == 0
else:
totest = balance < 0
if totest:
test_ids.append(i)
balance += w
test_weights += w
else:
train_ids.append(i)
balance -= w * ratio
train_weights += w * ratio
r = abs(train_weights - test_weights) / (1.0 * (train_weights + test_weights))
if r >= fail_imbalanced:
raise ImbalancedSplitException( # pragma: no cover
"Split is imbalanced: train_weights={0} test_weights={1} r={2}." # noqa: UP030
"".format(train_weights, test_weights, r)
)
return df.iloc[train_ids, :], df.iloc[test_ids, :]
[docs]
def train_test_connex_split(
df,
groups,
test_size=0.25,
train_size=None,
stratify=None,
hash_size=9,
unique_rows=False,
shuffle=True,
fail_imbalanced=0.05,
keep_balance=None,
stop_if_bigger=None,
return_cnx=False,
must_groups=None,
random_state=None,
verbose=0,
):
"""
This split is for a specific case where data is linked
in many ways. Let's assume we have three ids as we have
for online sales: *(product id, user id, card id)*.
As we may need to compute aggregated features,
we need every id not to be present in both train and
test set. The function computes the connected components
and breaks each of them in two parts for train and test.
:param df: :epkg:`pandas:DataFrame`
:param groups: columns name for the ids
:param test_size: ratio for the test partition
(if *train_size* is not specified)
:param train_size: ratio for the train partition
:param stratify: column holding the stratification
:param hash_size: size of the hash to cache information about partition
:param unique_rows: ensures that rows are unique
:param shuffle: shuffles before the split
:param fail_imbalanced: raises an exception if relative weights difference
is higher than this value
:param stop_if_bigger: (float) stops a connected components from being
bigger than this ratio of elements, this should not be used
unless a big components emerges, the algorithm stops merging
but does not guarantee it returns the best cut,
the value should be close to 0
:param keep_balance: (float), if not None, does not merge connected components
if their relative sizes are too different,
the value should be close to 1
:param return_cnx: returns connected components as a third results
:param must_groups: column name for ids which must not be shared by
train/test partitions
:param random_state: seed for random generator
:param verbose: verbosity (uses logging)
:return: Two see :class:`StreamingDataFrame
<pandas_streaming.df.dataframe.StreamingDataFrame>`, one
for train, one for test.
The list of ids must hold in memory.
There is no streaming implementation for the ids.
.. exref::
:title: Splits a dataframe, keep ids in separate partitions
:tag: dataframe
In some data science problems, rows are not independant
and share common value, most of the time ids. In some
specific case, multiple ids from different columns are
connected and must appear in the same partition.
Testing that each id column is evenly split and do not
appear in both sets in not enough. Connected components
are needed.
.. runpython::
:showcode:
from pandas import DataFrame
from pandas_streaming.df import train_test_connex_split
df = DataFrame([dict(user="UA", prod="PAA", card="C1"),
dict(user="UA", prod="PB", card="C1"),
dict(user="UB", prod="PC", card="C2"),
dict(user="UB", prod="PD", card="C2"),
dict(user="UC", prod="PAA", card="C3"),
dict(user="UC", prod="PF", card="C4"),
dict(user="UD", prod="PG", card="C5"),
])
train, test = train_test_connex_split(
df, test_size=0.5, groups=['user', 'prod', 'card'],
fail_imbalanced=0.6)
print(train)
print(test)
If *return_cnx* is True, the third results contains:
* connected components for each id
* the dataframe with connected components as a new column
.. runpython::
:showcode:
from pandas import DataFrame
from pandas_streaming.df import train_test_connex_split
df = DataFrame([dict(user="UA", prod="PAA", card="C1"),
dict(user="UA", prod="PB", card="C1"),
dict(user="UB", prod="PC", card="C2"),
dict(user="UB", prod="PD", card="C2"),
dict(user="UC", prod="PAA", card="C3"),
dict(user="UC", prod="PF", card="C4"),
dict(user="UD", prod="PG", card="C5"),
])
train, test, cnx = train_test_connex_split(
df, test_size=0.5, groups=['user', 'prod', 'card'],
fail_imbalanced=0.6, return_cnx=True)
print(cnx[0])
print(cnx[1])
"""
if stratify is not None:
raise NotImplementedError( # pragma: no cover
"Option stratify is not implemented."
)
if groups is None or len(groups) == 0:
raise ValueError( # pragma: no cover
"groups is empty. Use regular train_test_split."
)
if hasattr(df, "iter_creation"):
raise NotImplementedError( # pragma: no cover
"Not implemented yet for StreamingDataFrame."
)
if isinstance(df, numpy.ndarray):
raise NotImplementedError( # pragma: no cover
"Not implemented on numpy arrays."
)
if shuffle:
df = dataframe_shuffle(df, random_state=random_state)
dfids = df[groups].copy()
if must_groups is not None:
dfids_must = df[must_groups].copy()
name = "connex"
while name in dfids.columns:
name += "_"
one = "weight"
while one in dfids.columns:
one += "_"
# Connected components.
elements = list(range(dfids.shape[0]))
counts_cnx = {i: {i} for i in elements}
connex = {}
avoids_merge = {}
def do_connex_components(dfrows, local_groups, kb, sib):
"run connected components algorithms"
itern = 0
modif = 1
while modif > 0 and itern < len(elements):
if df.shape[0] > 10000:
logger.info(
"[train_test_connex_split] iteration=%d-#nb connect=%d - "
"modif=%s",
itern,
len(set(elements)),
modif,
)
modif = 0
itern += 1
for i, row in enumerate(dfrows.itertuples(index=False, name=None)):
vals = [
val
for val in zip(local_groups, row)
if not isinstance(val[1], float) or not numpy.isnan(val[1])
]
c = elements[i]
for val in vals:
if val not in connex:
connex[val] = c
modif += 1
set_c = set(connex[val] for val in vals)
set_c.add(c)
new_c = min(set_c)
add_pair_c = []
for c in set_c:
if c == new_c or (new_c, c) in avoids_merge:
continue
if kb is not None:
maxi = min(len(counts_cnx[new_c]), len(counts_cnx[c]))
if maxi > 5:
diff = len(counts_cnx[new_c]) + len(counts_cnx[c]) - maxi
r = diff / float(maxi)
if r > kb:
if verbose: # pragma: no cover
logger.info(
"[train_test_connex_split] balance "
"r=%1.4f>%1.2f, #[%d]=%d, #[%d]=%d",
r,
kb,
new_c,
len(counts_cnx[new_c]),
c,
len(counts_cnx[c]),
)
continue
if sib is not None:
r = (len(counts_cnx[new_c]) + len(counts_cnx[c])) / float(
len(elements)
)
if r > sib:
logger.info(
"[train_test_connex_split] "
"no merge r=%1.4f>%1.2f, #[%d]=%d, #[%d]=%d",
r,
sib,
new_c,
len(counts_cnx[new_c]),
c,
len(counts_cnx[c]),
)
avoids_merge[new_c, c] = i
continue
add_pair_c.append(c)
if len(add_pair_c) > 0:
for c in add_pair_c:
modif += len(counts_cnx[c])
for ii in counts_cnx[c]:
elements[ii] = new_c
counts_cnx[new_c] = counts_cnx[new_c].union(counts_cnx[c])
counts_cnx[c] = set()
keys = list(vals)
for val in keys:
if connex[val] == c:
connex[val] = new_c
modif += 1
if must_groups:
do_connex_components(dfids_must, must_groups, None, None)
do_connex_components(dfids, groups, keep_balance, stop_if_bigger)
# final
dfids[name] = elements
dfids[one] = 1
grsum = dfids[[name, one]].groupby(name, as_index=False).sum()
for g in groups:
logger.info("[train_test_connex_split] #nb in '%d':", len(set(dfids[g])))
logger.info(
"[train_test_connex_split] #connex %d/%d", grsum.shape[0], dfids.shape[0]
)
if grsum.shape[0] <= 1:
raise ValueError( # pragma: no cover
"Every element is in the same connected components."
)
# Statistics: top connected components
if verbose:
# Global statistics
counts = Counter(elements)
cl = [(v, k) for k, v in counts.items()]
cum = 0
maxc = None
logger.info(
"[train_test_connex_split] number of connected components: %d",
len(set(elements)),
)
for i, (v, k) in enumerate(sorted(cl, reverse=True)):
if i == 0:
maxc = k, v
if i >= 10:
break
cum += v
logger.info(
"[train_test_connex_split] c=%s #elements=%s cumulated=%d/%d",
k,
v,
cum,
len(elements),
)
# Most important component
logger.info(
"[train_test_connex_split] first row of the biggest component %d", maxc
)
tdf = dfids[dfids[name] == maxc[0]]
logger.info("[train_test_connex_split] % s", tdf.head(n=10))
# Splits.
train, test = train_test_split_weights(
grsum,
weights=one,
test_size=test_size,
train_size=train_size,
shuffle=shuffle,
fail_imbalanced=fail_imbalanced,
random_state=random_state,
)
train.drop(one, inplace=True, axis=1)
test.drop(one, inplace=True, axis=1)
# We compute the final dataframe.
def double_merge(d):
"merge twice"
merge1 = dfids.merge(d, left_on=name, right_on=name)
merge2 = df.merge(merge1, left_on=groups, right_on=groups)
return merge2
train_f = double_merge(train)
test_f = double_merge(test)
if return_cnx:
return train_f, test_f, (connex, dfids)
else:
return train_f, test_f
[docs]
def train_test_apart_stratify(
df: pandas.DataFrame,
group,
test_size: Optional[float] = 0.25,
train_size: Optional[float] = None,
stratify: Optional[str] = None,
force: bool = False,
random_state: Optional[int] = None,
sorted_indices: bool = False,
) -> Tuple["StreamingDataFrame", "StreamingDataFrame"]: # noqa: F821
"""
This split is for a specific case where data is linked
in one way. Let's assume we have two ids as we have
for online sales: *(product id, category id)*.
A product can have multiple categories. We need to have
distinct products on train and test but common categories
on both sides.
:param df: :epkg:`pandas:DataFrame`
:param group: columns name for the ids
:param test_size: ratio for the test partition
(if *train_size* is not specified)
:param train_size: ratio for the train partition
:param stratify: column holding the stratification
:param force: if True, tries to get at least one example on the test side
for each value of the column *stratify*
:param random_state: seed for random generators
:param sorted_indices: sort index first,
see issue `41 <https://github.com/sdpython/pandas-streaming/issues/41>`
:return: Two see :class:`StreamingDataFrame
<pandas_streaming.df.dataframe.StreamingDataFrame>`, one
for train, one for test.
The list of ids must hold in memory.
There is no streaming implementation for the ids.
This split was implemented for a case of a multi-label
classification. A category (*stratify*) is not exclusive
and an observation can be assigned to multiple
categories. In that particular case, the method
:func:`sklearn.model_selection.train_test_split`
can not directly be used.
.. runpython::
:showcode:
import pandas
from pandas_streaming.df import train_test_apart_stratify
df = pandas.DataFrame([dict(a=1, b="e"),
dict(a=1, b="f"),
dict(a=2, b="e"),
dict(a=2, b="f")])
train, test = train_test_apart_stratify(
df, group="a", stratify="b", test_size=0.5)
print(train)
print('-----------')
print(test)
"""
if stratify is None:
raise ValueError("stratify must be specified.") # pragma: no cover
if group is None:
raise ValueError("group must be specified.") # pragma: no cover
if hasattr(df, "iter_creation"):
raise NotImplementedError("Not implemented yet for StreamingDataFrame.")
if isinstance(df, numpy.ndarray):
raise NotImplementedError("Not implemented on numpy arrays.")
p = (1 - test_size) if test_size else None
if train_size is not None:
p = train_size
test_size = 1 - p
if p is None or min(test_size, p) <= 0:
raise ValueError( # pragma: no cover
f"test_size={test_size} or train_size={train_size} cannot be null"
)
couples = df[[group, stratify]].itertuples(name=None, index=False)
hist = Counter(df[stratify])
sorted_hist = [(v, k) for k, v in hist.items()]
sorted_hist.sort()
ids = {c: set() for c in hist}
for g, s in couples:
ids[s].add(g)
if random_state is None:
permutation = numpy.random.permutation
else:
state = numpy.random.RandomState(random_state)
permutation = state.permutation
split = {}
for _, k in sorted_hist:
indices = sorted(ids[k]) if sorted_indices else ids[k]
not_assigned, assigned = [], []
for c in indices:
if c in split:
assigned.append(c)
else:
not_assigned.append(c)
if len(not_assigned) == 0:
continue
nb_test = sum(split[c] for c in assigned)
expected = min(len(ids[k]), int(test_size * len(ids[k]) + 0.5)) - nb_test
if force and expected == 0 and nb_test == 0:
nb_train = len(assigned) - nb_test
if nb_train > 0 or len(not_assigned) > 1:
expected = min(1, len(not_assigned))
if expected > 0:
permutation(not_assigned)
for e in not_assigned[:expected]:
split[e] = 1
for e in not_assigned[expected:]:
split[e] = 0
else:
for c in not_assigned:
split[c] = 0
train_set = set(k for k, v in split.items() if v == 0)
test_set = set(k for k, v in split.items() if v == 1)
train_df = df[df[group].isin(train_set)]
test_df = df[df[group].isin(test_set)]
return train_df, test_df