update modify_residue_number and add batch_modify_residues

This commit is contained in:
2024-02-02 11:30:33 +08:00
parent cacef78dfc
commit ac2c0ad23b

View File

@@ -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