from ..tbl.table import Table
from ..fun.misc import reduce_list as rl
[docs]
class StepAtLeast(Table):
"""Select trajectories by the step number.
Args:
reqs[0] (Table): Trajectory data.
param["step"] (int): Step number for selection. If you set step=2,
trajectories containing at least three points (frames) are
selected.
param["group_depth"] (int): Depth number of trajectory number column.
param["split_depth"] (int): File split depth number.
Returns:
Table: Selected trajectory Table
"""
[docs]
def set_info(self, param={}):
"""Copy info from reqs[0] and add params.
"""
self.info.copy_req(0)
self.info.add_param(
"step", param["step"], "int32",
"Step number for trajectory selection")
self.info.set_group_depth(param["group_depth"])
self.info.set_split_depth(param["split_depth"])
[docs]
@staticmethod
def process(reqs, param):
"""Select trajectories by the step number.
Args:
reqs[0] (pandas.DataFrame): Trajectory data.
param["index_cols"] (list of str): Index column names for
:meth:`pandas.DataFrame.groupby`. List should include the
trajectory number column.
Returns:
pandas.DataFrame: Selected trajectory table
"""
df = reqs[0].copy()
grouped = df.groupby(rl(param["index_cols"]))
df = grouped.filter(lambda x: len(x) > param["step"])
return df.reset_index(drop=True)
[docs]
class StepRange(Table):
"""Select trajectories by the step number range.
Args:
reqs[0] (Table): Trajectory data.
param["step_range"] (list of int): Minimum and maximum numbers of
trajectory step.
param["group_depth"] (int): Data split depth number.
param["split_depth"] (int): File split depth number.
Returns:
Table: Selected trajectory Table
"""
[docs]
def set_info(self, param={}):
"""Copy info from reqs[0] and add params.
"""
self.info.copy_req(0)
self.info.add_param(
"step_range", param["step_range"], "list of int",
"Minimum and maximum numbers of the trajectory step")
self.info.set_group_depth(param["group_depth"])
self.info.set_split_depth(param["split_depth"])
[docs]
@staticmethod
def process(reqs, param):
"""Select trajectories by the step number range.
Args:
reqs[0] (pandas.DataFrame): Trajectory data.
param["step_range"] (list of int): Minimum and maximum numbers
of trajectory step.
param["index_cols"] (list of str): Index column names for
:meth:`pandas.DataFrame.groupby`. List should include the
trajectory number column.
Returns:
pandas.DataFrame: Selected trajectory table
"""
df = reqs[0].copy()
grouped = df.groupby(rl(param["index_cols"]))
df = grouped.filter(lambda x: len(x) > param["step_range"][0])
grouped = df.groupby(rl(param["index_cols"]))
df = grouped.filter(lambda x: len(x) <= param["step_range"][1] + 1)
return df.reset_index(drop=True)