From 1f3d4a8dcf799003b64fd442023e4fd5fe29d2ce Mon Sep 17 00:00:00 2001 From: Jorge Morgado Vega Date: Thu, 8 Jun 2023 13:52:59 -0400 Subject: [PATCH] Update example 4 Co-authored-by: gvieralopez --- examples/example_04.py | 12 ++++++------ 1 file changed, 6 insertions(+), 6 deletions(-) diff --git a/examples/example_04.py b/examples/example_04.py index 2c0c6d1..28904d7 100644 --- a/examples/example_04.py +++ b/examples/example_04.py @@ -14,23 +14,23 @@ Dataset.uci_characters(), Dataset.uci_movement_libras(), ] - + +# Select the desired features to be extracted from the trajectories featurizer = featurizers.UniversalFeaturizer() for dataset in datasets: print(f"\nDataset: {dataset.name}\n") - # Split the dataset into train and test + # Split the dataset into train and test and filter out short trajectories train, test = dataset.filter( - lambda traj, _: len(traj) >= 5 and traj.r.delta.norm.sum() > 0 + lambda traj, label: len(traj) >= 5 + and traj.r.delta.norm.sum() > 0 + and dataset.label_counts[label] > 5 ).split( train_size=0.7, random_state=SEED, ) - # Select the desired features to be extracted from the trajectories - featurizer = featurizers.UniversalFeaturizer() - # Define the model model = XGBoostModel(featurizer=featurizer)