Fix pfam_id column missing on empty dataframes

Merck · Apr 1, 2022 · 3b6edb0 · 3b6edb0
1 parent e3952d7
commit 3b6edb0
Showing 1 changed file with 4 additions and 0 deletions.
diff --git a/deepbgc/features.py b/deepbgc/features.py
@@ -200,6 +200,10 @@ def __init__(self, column='pfam_id', sequence_as_vector=True):
         self.sequence_as_vector = sequence_as_vector
 
     def transform(self, X):
+        if X.empty:
+            X = pd.DataFrame({self.column: []})
+        if self.column not in X.columns:
+            raise ValueError('Unexpected pfam dataframe, column "{}" not found in: {}'.format(self.column, X.columns))
         # Turn each pfam ID into a vector
         values = pd.get_dummies(X[self.column]).reindex(columns=self.unique_values, fill_value=0)
         if self.sequence_as_vector: