Drop unexpected columsn when writing PDB dataframes #124

BioPandas · Apr 21, 2023 · 7674e37 · 7674e37
1 parent 1ff101c
commit 7674e37
Showing 1 changed file with 3 additions and 0 deletions.
diff --git a/biopandas/pdb/pandas_pdb.py b/biopandas/pdb/pandas_pdb.py
@@ -610,6 +610,9 @@ def to_pdb(self, path, records=None, gz=False, append_newline=True):
 
         dfs = {r: self.df[r].copy() for r in records if not self.df[r].empty}
 
+        # Drop unexpected columns
+        dfs = {k: v[[pdb_df_columns]] for k, v in dfs.items() if k in {"ATOM", "HETATM"}}
+
         for r in dfs:
             for col in pdb_records[r]:
                 dfs[r][col["id"]] = dfs[r][col["id"]].apply(col["strf"])