From ac2c0ad23b5a7bb3e8f150f72e1bcb1d87dc3a51 Mon Sep 17 00:00:00 2001
From: lingyuzeng <pylyzeng@gmail.com>
Date: Fri, 2 Feb 2024 11:30:33 +0800
Subject: [PATCH] update modify_residue_number and add batch_modify_residues

---
 analysis_pdb.py | 24 ++++++++++++++++++------
 1 file changed, 18 insertions(+), 6 deletions(-)

diff --git a/analysis_pdb.py b/analysis_pdb.py
index 1109c37..3460f59 100755
--- a/analysis_pdb.py
+++ b/analysis_pdb.py
@@ -146,27 +146,39 @@ class PDBAnalyzer:
         # Return a new instance of PDBAnalyzer pointing to the cleaned file
         return cls(out_file)
     
-    def modify_residue_number(self, chain_id: str, original_res_num: int, new_res_num: int) -> PandasPdb:
+    def modify_residue_number(self, chain_id: str, original_res_num: int, new_res_num: int) -> None:
         """
         Modify the residue number for a specific chain in a PDB file.
-        
+
         Args:
         chain_id (str): The chain identifier in the PDB file.
         original_res_num (int): The original residue number to find.
         new_res_num (int): The new residue number to replace with.
-        
-        Returns:
-        PandasPdb: The modified PDB structure as a PandasPdb object.
         """
         # Check if the residue number to be replaced exists
         mask = (self.biodata.df['ATOM']['chain_id'] == chain_id) & \
                (self.biodata.df['ATOM']['residue_number'] == original_res_num)
         if not mask.any():
             raise ValueError(f"Residue number {original_res_num} in chain {chain_id} does not exist in the PDB file.")
-        
+
         # Update the residue number
         self.biodata.df['ATOM'].loc[mask, 'residue_number'] = new_res_num
 
+    def batch_modify_residues(self, modifications: dict) -> PandasPdb:
+        """
+        Batch modify multiple residues in a PDB file based on a dictionary of modifications.
+
+        Args:
+        modifications (dict): A dictionary with chain identifiers as keys and a list of
+                              tuples (original_res_num, new_res_num) as values.
+
+        Returns:
+        PandasPdb: The modified PDB structure as a PandasPdb object.
+        """
+        for chain_id, changes in modifications.items():
+            for original_res_num, new_res_num in changes:
+                self.modify_residue_number(chain_id, original_res_num, new_res_num)
+
         # Return the modified PandasPdb object
         return self.biodata