From 7ba16c9c0e48712313d3ba54531d5b0a978aef1e Mon Sep 17 00:00:00 2001
From: Pao <paotorres89@gmail.com>
Date: Fri, 8 Mar 2024 11:07:15 -0300
Subject: [PATCH] added functions to convert format

---
 pyMBE.py | 68 +++++++++++++++++++++++++++++++++++++++++++++++++++++---
 1 file changed, 65 insertions(+), 3 deletions(-)

diff --git a/pyMBE.py b/pyMBE.py
index b6f3246..29f43f8 100644
--- a/pyMBE.py
+++ b/pyMBE.py
@@ -516,6 +516,68 @@ def clean_df_row(self, index, columns_keys_to_clean=["particle_id", "particle_id
             self.add_value_to_df(key=(column_key,''),index=index,new_value=self.np.nan, warning=False)
         return
 
+    def convert_columns_to_original_format (self, df):
+        """
+        Converts the columns of the Dataframe to the original format in pyMBE.
+        
+        Args:
+            df(`DataFrame`): dataframe with pyMBE information as a string  
+        
+        """
+
+        from ast import literal_eval
+
+        columns_with_units = ['diameter', 'epsilon']
+
+        columns_with_list_or_dict = ['sequence', 'residue_list','side_chains', 'parameters_of_the_potential']
+
+        for column_name in columns_with_list_or_dict:
+            df[column_name] = df[column_name].apply(lambda x: literal_eval(x) if self.pd.notnull(x) else x)
+
+        for column_name in columns_with_units:
+
+            df[column_name] = df[column_name].apply(lambda x: self.create_variable_with_units(x) if self.pd.notnull(x) else x)
+
+        df['bond_object'] = df['bond_object'].apply(lambda x: (self.convert_str_to_bond_object(x)) if self.pd.notnull(x) else x)
+
+        return df
+    
+    def convert_str_to_bond_object (self, bond_str):
+        
+        """
+        Convert a row read as a `str` to the corresponding bond object. There are two supported bonds: HarmonicBond and  FeneBond
+
+        Args:
+            bond_str (`strt`): string with the information of a bond object
+
+        Returns:
+            bond_object(`obj`): EsPRESSo bond object 
+        """
+        
+        from ast import literal_eval
+        from espressomd.interactions import HarmonicBond
+        from espressomd.interactions import FeneBond
+
+        supported_bonds = ['HarmonicBond', 'FeneBond']
+
+        for bond in supported_bonds:
+
+            variable = self.re.subn(f'{bond}', '', bond_str)
+
+            if variable[1] == 1: 
+            
+                params = literal_eval(variable[0])
+
+                if bond == 'HarmonicBond':
+
+                    bond_object = HarmonicBond(r_cut =params['r_cut'], k = params['k'], r_0=params['r_0'])
+
+                elif bond == 'FeneBond':
+
+                    bond_object = FeneBond(k = params['k'], d_r_max =params['d_r_max'], r_0=params['r_0'])
+
+        return bond_object   
+
     def copy_df_entry(self, name, column_name, number_of_copies):
         '''
         Creates 'number_of_copies' of a given 'name' in `pymbe.df`.
@@ -1005,7 +1067,7 @@ def create_variable_with_units(self, variable):
         Returns:
             variable_with_units(`obj`): variable with units using the pyMBE UnitRegistry.
         """        
-       if type(variable) is dict: 
+        if type(variable) is dict: 
 
             value=variable.pop('value')
             units=variable.pop('units')
@@ -2021,8 +2083,8 @@ def read_pmb_df (self,filename):
 
         columns_names = self.setup_df()
         df.columns = columns_names
-        self.df=df
-        return df
+        self.df= self.convert_columns_to_original_format(df)
+        return self.df
     
     def read_protein_vtf_in_df (self,filename,unit_length=None):
         """