Drop unexpected columns when writing PDB dataframes #124

biopandas/pdb/pandas_pdb.py

@@ -702,7 +702,7 @@ def to_pdb(self, path, records=None, gz=False, append_newline=True):
                     for idx in range(dfs[r][c].values.shape[0]):
                         if len(dfs[r][c].values[idx]) > 8:
                             dfs[r][c].values[idx] = str(dfs[r][c].values[idx]).strip()
-                if c in {"line_idx", "OUT"}:
+                if c in {"line_idx", "OUT", "model_id"}:
                     pass
                 elif r in {"ATOM", "HETATM"} and c not in pdb_df_columns:
                     warn(
@@ -712,19 +712,8 @@ def to_pdb(self, path, records=None, gz=False, append_newline=True):
                 else:
                     dfs[r]["OUT"] = dfs[r]["OUT"] + dfs[r][c]
 
-        if pd_version < LooseVersion("0.17.0"):
-            warn(
-                "You are using an old pandas version (< 0.17)"
-                " that relies on the old sorting syntax."
-                " Please consider updating your pandas"
-                " installation to a more recent version.",
-                DeprecationWarning,
-            )
-            dfs.sort(columns="line_idx", inplace=True)
-
-        elif pd_version < LooseVersion("0.23.0"):
+        if pd_version < LooseVersion("0.23.0"):
             df = pd.concat(dfs)
-
         else:
             df = pd.concat(dfs, sort=False)
 

biopandas/pdb/tests/test_write_pdb.py

@@ -8,6 +8,7 @@
 import warnings
 
 import pandas as pd
+
 from biopandas.pdb import PandasPdb
 
 TESTDATA_FILENAME = os.path.join(os.path.dirname(__file__), "data", "3eiy.pdb")
@@ -77,6 +78,18 @@ def test_anisou():
     os.remove(OUTFILE)
     assert f1 == four_eiy
 
+
+def test_write_with_model_id():
+    """Test writing a dataframe with a model ID column added."""
+    ppdb = PandasPdb()
+    ppdb.read_pdb(TESTDATA_FILENAME)
+    df.label_models()
+    ppdb.to_pdb(path=OUTFILE, records=None)
+    with open(OUTFILE, "r") as f:
+        f1 = f.read()
+    os.remove(OUTFILE)
+    assert f1 == f2
+
 
 def test_add_remark():
     """Test adding a REMARK entry."""