Fix bug when feature is not in data

im-ethz · Mar 15, 2022 · 9aafe9e · 9aafe9e
1 parent 81b1bc6
commit 9aafe9e
Showing 1 changed file with 11 additions and 10 deletions.
diff --git a/preprocess_trainingpeaks.py b/preprocess_trainingpeaks.py
@@ -545,16 +545,17 @@ def features(i, verbose=0):
 	# TODO: combined_pedal_smoothness
 
 def drop_extremes(df, feature, comp, value):
-	if comp == 'lower':
-		mask = df[feature] <= value
-	elif comp == 'higher':
-		mask = df[feature] >= value
-	elif comp == 'equal':
-		mask = df[feature] == value
-
-	if mask.sum() > 0:
-		print(f"DROP: {mask.sum()} with {feature} {comp} than {value}")
-	df.loc[mask, feature] = np.nan
+	if feature in df:
+		if comp == 'lower':
+			mask = df[feature] <= value
+		elif comp == 'higher':
+			mask = df[feature] >= value
+		elif comp == 'equal':
+			mask = df[feature] == value
+
+		if mask.sum() > 0:
+			print(f"DROP: {mask.sum()} with {feature} {comp} than {value}")
+		df.loc[mask, feature] = np.nan
 	return df
 
 def extremes(i):