py2ls 0.2.4.5__py3-none-any.whl → 0.2.4.7__py3-none-any.whl

py2ls/ips.py

@@ -60,6 +60,7 @@ except NameError:
 def plt_font(dir_font: str = "/System/Library/Fonts/Hiragino Sans GB.ttc"):
     """
     Add the Chinese (default) font to the font manager
+    show chinese
     Args:
         dir_font (str, optional): _description_. Defaults to "/System/Library/Fonts/Hiragino Sans GB.ttc".
     """
@@ -554,14 +555,28 @@ def shared(*args, strict=True, n_shared=2, verbose=True):
         # Get elements that appear in at least n_shared lists
         shared_elements = [item for item, count in element_count.items() if count >= n_shared]
 
-    shared_elements = flatten(shared_elements, verbose=verbose) 
+    shared_elements = flatten(shared_elements, verbose=verbose)
     if verbose:
         elements2show = shared_elements if len(shared_elements)<10 else shared_elements[:5]
         print(f"{' '*2}{len(shared_elements)} elements shared: {' '*2}{elements2show}")
         print("********* checking shared elements *********")
     return shared_elements
 
-def flatten(nested: Any, unique_list=True, verbose=True):
+def not_shared(*args, strict=True, n_shared=2, verbose=False):
+    """
+    To find the elements in list1 that are not shared with list2 while maintaining the original order of list1
+    usage:
+        list1 = [1, 8, 3, 3, 4, 5]
+        list2 = [4, 5, 6, 7, 8]
+        not_shared(list1,list2)# output [1,3]
+    """
+    _common = shared(*args, strict=strict, n_shared=n_shared, verbose=verbose)
+    list1 = args[0]
+    _not_shared=[item for item in list1 if item not in _common]
+    return flatten(_not_shared, verbose=verbose)
+
+
+def flatten(nested: Any, unique_list=True, verbose=False):
     """
     Recursively flattens a nested structure (lists, tuples, dictionaries, sets) into a single list.
     Parameters:
@@ -589,7 +604,7 @@ def flatten(nested: Any, unique_list=True, verbose=True):
     else:
         return flattened_list
     
-def strcmp(search_term, candidates, ignore_case=True, verbose=False, scorer="WR"):
+def strcmp(search_term, candidates, ignore_case=True,get_rank=False, verbose=False, scorer="WR"):
     """
     Compares a search term with a list of candidate strings and finds the best match based on similarity score.
 
@@ -623,6 +638,11 @@ def strcmp(search_term, candidates, ignore_case=True, verbose=False, scorer="WR"
             similarity_scores = [fuzz.ratio(str1_, word) for word in str2_]
         else:
             similarity_scores = [fuzz.WRatio(str1_, word) for word in str2_]
+        if get_rank:
+            idx = [similarity_scores.index(i) for i in sorted(similarity_scores,reverse=True)]
+            if verbose:
+                display([candidates[ii] for ii in idx])
+            return [candidates[ii] for ii in idx]
         best_match_index = similarity_scores.index(max(similarity_scores))
         best_match_score = similarity_scores[best_match_index]
     else:
@@ -1554,13 +1574,26 @@ def unzip(dir_path, output_dir=None):
             tar_ref.extractall(output_dir)
         return output_dir
     # Handle .gz files
-    if dir_path.endswith(".gz"):
+    if dir_path.endswith(".gz") or dir_path.endswith(".gzip"):
         import gzip
 
         output_file = os.path.splitext(dir_path)[0]  # remove the .gz extension
-        with gzip.open(dir_path, "rb") as gz_file:
-            with open(output_file, "wb") as out_file:
-                shutil.copyfileobj(gz_file, out_file)
+        try:
+            with gzip.open(dir_path, "rb") as gz_file:
+                with open(output_file, "wb") as out_file:
+                    shutil.copyfileobj(gz_file, out_file)
+            print(f"unzipped '{dir_path}' to '{output_file}'")
+        except FileNotFoundError:
+            print(f"Error: The file '{dir_path}' was not found.")
+        except PermissionError:
+            print(f"Error: Permission denied when accessing '{dir_path}' or writing to '{output_file}'.")
+        except Exception as e:
+            try:
+                import tarfile
+                with tarfile.open(dir_path, 'r:gz') as tar:
+                    tar.extractall(path=output_file)
+            except Exception as final_e:
+                print(f"An final unexpected error occurred: {final_e}")
         return output_file
 
     # Handle .zip files
@@ -1642,9 +1675,12 @@ def is_df_abnormal(df: pd.DataFrame, verbose=False) -> bool:
     False: normal
 
     """
+    if not isinstance(df, pd.DataFrame):
+        return False
+    df.columns = df.columns.astype(str)# 把它变成str, 这样就可以进行counts运算了
     # Initialize a list to hold messages about abnormalities
     messages = []
-    is_abnormal = False
+    is_abnormal = True
     # Check the shape of the DataFrame
     actual_shape = df.shape
     messages.append(f"Shape of DataFrame: {actual_shape}")
@@ -1739,10 +1775,12 @@ def fload(fpath, kind=None, **kwargs):
             content = file.read()
         return content
 
-    def load_html(fpath):
-        with open(fpath, "r") as file:
-            content = file.read()
-        return content
+    # def load_html(fpath):
+    #     with open(fpath, "r") as file:
+    #         content = file.read()
+    #     return content
+    def load_html(fpath,**kwargs):
+        return pd.read_html(fpath,**kwargs)
 
     def load_json(fpath, **kwargs):
         output=kwargs.pop("output","json")
@@ -1956,8 +1994,8 @@ def fload(fpath, kind=None, **kwargs):
                                 # display(df.head(2))
                                 # print(f"is_df_abnormal:{is_df_abnormal(df, verbose=0)}")
                                 if not is_df_abnormal(df, verbose=0):
-                                    display(df.head(2))
-                                    print(f"shape: {df.shape}")
+                                    display(df.head(2)) if isinstance(df, pd.DataFrame) else display("it is not a DataFrame")
+                                    print(f"shape: {df.shape}") if isinstance(df, pd.DataFrame) else display("it is not a DataFrame")
                                     return df
                             except EmptyDataError as e:
                                 continue
@@ -1981,6 +2019,42 @@ def fload(fpath, kind=None, **kwargs):
             pass
         return df
 
+
+    def load_parquet(fpath, **kwargs):
+        """
+        Load a Parquet file into a Pandas DataFrame with advanced options.
+
+        Parameters:
+        - fpath (str): The file path to the Parquet file.
+        - engine (str): The engine to use for reading the Parquet file (default is 'pyarrow').
+        - columns (list): List of columns to load. If None, loads all columns.
+        - verbose (bool): If True, prints additional information about the loading process.
+        - filters (list): List of filter conditions for predicate pushdown.
+        - **kwargs: Additional keyword arguments for `pd.read_parquet`.
+
+        Returns:
+        - df (DataFrame): The loaded DataFrame.
+        """
+        
+        engine = kwargs.get("engine", "pyarrow")
+        verbose = kwargs.pop("verbose", False)
+        
+        if verbose:
+            use_pd("read_parquet", verbose=verbose)
+        try:
+            df = pd.read_parquet(fpath, engine=engine, **kwargs)
+            if verbose:
+                if 'columns' in kwargs:
+                    print(f"Loaded columns: {kwargs['columns']}")
+                else:
+                    print("Loaded all columns.")
+            print(f"shape: {df.shape}")
+        except Exception as e:
+            print(f"An error occurred while loading the Parquet file: {e}")
+            df = None
+
+        return df 
+
     def load_ipynb(fpath, **kwargs):
         as_version = kwargs.get("as_version", 4)
         with open(fpath, "r") as file:
@@ -2049,51 +2123,21 @@ def fload(fpath, kind=None, **kwargs):
         kind = kind.lower()
     kind = kind.lstrip(".").lower()
     img_types = [
-        "bmp",
-        "eps",
-        "gif",
-        "icns",
-        "ico",
-        "im",
-        "jpg",
-        "jpeg",
-        "jpeg2000",
-        "msp",
-        "pcx",
-        "png",
-        "ppm",
-        "sgi",
-        "spider",
-        "tga",
-        "tiff",
-        "tif",
-        "webp",
-        "json",
+        "bmp","eps","gif","png","jpg","jpeg","jpeg2000","tiff","tif",
+        "icns","ico","im","msp","pcx","ppm","sgi","spider","tga","webp",
     ]
     doc_types = [
-        "docx",
-        "txt",
-        "md",
-        "html",
-        "json",
-        "yaml",
-        "xml",
-        "csv",
-        "xlsx",
-        "pdf",
+        "docx","pdf",
+        "txt","csv","xlsx","tsv","parquet","snappy",
+        "md","html",
+        "json","yaml","xml",
         "ipynb",
+        "mtx"
     ]
     zip_types = [
-        "gz",
-        "zip",
-        "7z",
-        "tar",
-        "tar.gz",
-        "tar.bz2",
-        "bz2",
-        "xz",
-        "rar",
-        "tgz",
+        "gz","zip","7z","rar","tgz",
+        "tar","tar.gz","tar.bz2",
+        "bz2","xz","gzip"
     ]
     other_types = ["fcs"]
     supported_types = [*doc_types, *img_types, *zip_types, *other_types]
@@ -2122,14 +2166,14 @@ def fload(fpath, kind=None, **kwargs):
     elif kind == "txt" or kind == "md":
         return load_txt_md(fpath)
     elif kind == "html":
-        return load_html(fpath)
+        return load_html(fpath, **kwargs)
     elif kind == "json":
-        return load_json(fpath)
+        return load_json(fpath, **kwargs)
     elif kind == "yaml":
         return load_yaml(fpath)
     elif kind == "xml":
         return load_xml(fpath)
-    elif kind == "csv":
+    elif kind in ["csv","tsv"]:
         content = load_csv(fpath, **kwargs)
         return content
     elif kind in ["ods", "ods", "odt"]:
@@ -2140,14 +2184,25 @@ def fload(fpath, kind=None, **kwargs):
         engine = kwargs.get("engine", "xlrd")
         kwargs.pop("engine", None)
         content = load_excel(fpath, engine=engine, **kwargs)
+        print(f"shape: {content.shape}")
         display(content.head(3))
         return content
     elif kind == "xlsx":
         content = load_excel(fpath, **kwargs)
         display(content.head(3))
+        print(f"shape: {content.shape}")
+        return content
+    elif kind=='mtx':
+        from scipy.io import mmread
+        dat_mtx=mmread(fpath)
+        content=pd.DataFrame.sparse.from_spmatrix(dat_mtx,**kwargs)
+        display(content.head(3))
+        print(f"shape: {content.shape}")
         return content
     elif kind == "ipynb":
         return load_ipynb(fpath, **kwargs)
+    elif kind in ['parquet','snappy']:
+        return load_parquet(fpath,**kwargs)
     elif kind == "pdf":
         # print('usage:load_pdf(fpath, page="all", verbose=False)')
         return load_pdf(fpath, **kwargs)
@@ -2193,9 +2248,7 @@ def fload(fpath, kind=None, **kwargs):
         return meta, data
 
     else:
-        # try:
-        #     content = load_csv(fpath, **kwargs)
-        # except:
+        print("direct reading...")
         try:
             try:
                 with open(fpath, "r", encoding="utf-8") as f:
@@ -2495,6 +2548,25 @@ def fsave(
         tree = etree.ElementTree(root)
         tree.write(fpath, pretty_print=True, xml_declaration=True, encoding="UTF-8")
 
+    def save_parquet(fpath:str, data:pd.DataFrame, **kwargs): 
+        engine = kwargs.pop("engine","auto") # auto先试pyarrow, 不行就转为fastparquet, {‘auto’, ‘pyarrow’, ‘fastparquet’}
+        compression=kwargs.pop("compression",None) # Use None for no compression. Supported options: ‘snappy’, ‘gzip’, ‘brotli’, ‘lz4’, ‘zstd’
+        try:
+            # Attempt to save with "pyarrow" if engine is set to "auto"
+                data.to_parquet(fpath, engine=engine, compression=compression, **kwargs)
+                print(f"DataFrame successfully saved to {fpath} with engine '{engine}' and {compression} compression.") 
+        except Exception as e:
+            print(f"Error using with engine '{engine}' and {compression} compression: {e}")
+            if "Sparse" in str(e):
+                try:
+                    # Handle sparse data by converting columns to dense
+                    print("Attempting to convert sparse columns to dense format...")
+                    data = data.apply(lambda x: x.sparse.to_dense() if pd.api.types.is_sparse(x) else x)
+                    save_parquet(fpath, data=data,**kwargs)
+                except Exception as last_e:
+                    print(f"After converted sparse columns to dense format, Error using with engine '{engine}' and {compression} compression: {last_e}")
+                    
+
     if kind is None:
         _, kind = os.path.splitext(fpath)
         kind = kind.lower()
@@ -2540,6 +2612,15 @@ def fsave(
         save_yaml(fpath, content, **kwargs)
     elif kind == "ipynb":
         save_ipynb(fpath, content, **kwargs)
+    elif kind.lower() in ["parquet","pq","big","par"]:
+        compression=kwargs.pop("compression",None) # Use None for no compression. Supported options: ‘snappy’, ‘gzip’, ‘brotli’, ‘lz4’, ‘zstd’
+        # fix the fpath ends
+        if not '.parquet' in fpath:
+            fpath=fpath.replace(kind, 'parquet')
+        if compression is not None:
+            if not fpath.endswith(compression):
+                fpath=fpath+f".{compression}"
+        save_parquet(fpath=fpath, data=content,compression=compression,**kwargs)
     else:
         try:
             netfinder.downloader(url=content, dir_save=dirname(fpath), kind=kind)
@@ -3058,8 +3139,11 @@ def figsave(*args, dpi=300):
 
 def is_str_color(s):
     # Regular expression pattern for hexadecimal color codes
-    color_code_pattern = r"^#([A-Fa-f0-9]{6}|[A-Fa-f0-9]{8})$"
-    return re.match(color_code_pattern, s) is not None
+    if isinstance(s,str):
+        color_code_pattern = r"^#([A-Fa-f0-9]{6}|[A-Fa-f0-9]{8})$"
+        return re.match(color_code_pattern, s) is not None
+    else:
+        return True
 
 
 def is_num(s):
@@ -5509,7 +5593,21 @@ def df_reducer(
 ) -> pd.DataFrame:
     """
     Reduces the dimensionality of the selected DataFrame using PCA or UMAP.
-
+    method: 
+        1. 'umap': 
+            - big dataset and global structure, often preferred in large-scale datasets for 
+            visualization and dimensionality reduction, balancing speed and quality of visualization.
+            - t-SNE excels at preserving local structure (i.e., clusters), but it often loses global 
+            relationships, causing clusters to appear in arbitrary proximities to each other.
+        2. 'pca': 
+            - t-SNE excels at preserving local structure (i.e., clusters), but it often loses global 
+                relationships, causing clusters to appear in arbitrary proximities to each other.
+            - useful as a preprocessing step and in datasets where linear relationships dominate.
+        3. 't-SNE': 
+            a. t-SNE excels at preserving local structure (i.e., clusters), but it often loses global 
+                relationships, causing clusters to appear in arbitrary proximities to each other.
+            b. often preferred in large-scale datasets for visualization and dimensionality 
+                reduction, balancing speed and quality of visualization.
     Parameters:
     -----------
     data : pd.DataFrame

py2ls/mol.py

@@ -0,0 +1,289 @@
+import os
+import subprocess
+from rdkit import Chem
+from rdkit.Chem import AllChem,Draw
+from openbabel import openbabel
+import matplotlib.pyplot as plt
+# import pymol2  # 使用 PyMOL API 进行分子展示
+
+from typing import Any, Dict, Union, List
+
+def load_mol(fpath: str) -> Union[Dict[str, Any], None]:
+    """
+    Master function to read various molecular structure files and return a consistent molecule dictionary.
+    Supports formats: .pdb, .mol, .sdf, .xyz, .gro, and others through RDKit, Pybel, MDAnalysis, and ASE.
+    
+    Parameters:
+    - fpath (str): Path to the molecular file
+    
+    Returns:
+    - mol_dict (Dict[str, Any]): Dictionary with molecule information:
+        - 'atoms': List of atom information dictionaries
+        - 'bonds': List of bond information dictionaries
+        - 'metadata': Metadata for molecule (e.g., file name)
+    """
+    ext = os.path.splitext(fpath)[-1].lower()  # Get the file extension
+
+    def create_atom_dict(atom) -> Dict[str, Any]:
+        """Helper to create a consistent atom dictionary."""
+        return {
+            'element': atom.atomic_symbol,
+            'coords': atom.coords,
+            'index': atom.idx,
+            'charge': atom.formalcharge
+        }
+    
+    def create_bond_dict(bond) -> Dict[str, Any]:
+        """Helper to create a consistent bond dictionary."""
+        return {
+            'start_atom_idx': bond.GetBeginAtomIdx(),
+            'end_atom_idx': bond.GetEndAtomIdx(),
+            'bond_type': bond.GetBondTypeAsDouble()
+        }
+
+    mol_dict = {
+        "atoms": [],
+        "bonds": [],
+        "metadata": {
+            "file_name": os.path.basename(fpath),
+            "format": ext
+        }
+    }
+
+    try:
+        # Handling with RDKit (for .mol and .sdf)
+        if ext in ['.mol', '.sdf']:
+            from rdkit import Chem
+            if ext == '.mol':
+                mol = Chem.MolFromMolFile(fpath)
+                if mol is None:
+                    raise ValueError("RDKit failed to parse the .mol file.")
+                atoms = mol.GetAtoms()
+                bonds = mol.GetBonds()
+            elif ext == '.sdf':
+                supplier = Chem.SDMolSupplier(fpath)
+                mol = next(supplier, None)
+                if mol is None:
+                    raise ValueError("RDKit failed to parse the .sdf file.")
+                atoms = mol.GetAtoms()
+                bonds = mol.GetBonds()
+
+            # Populate atom and bond data
+            mol_dict["atoms"] = [
+                {
+                    "element": atom.GetSymbol(),
+                    "coords": atom.GetOwningMol().GetConformer().GetAtomPosition(atom.GetIdx()),
+                    "index": atom.GetIdx(),
+                    "charge": atom.GetFormalCharge()
+                }
+                for atom in atoms
+            ]
+            mol_dict["bonds"] = [
+                create_bond_dict(bond)
+                for bond in bonds
+            ]
+
+        # Handling with Pybel (supports multiple formats: .pdb, .mol, .xyz, etc.)
+        elif ext in ['.pdb', '.mol', '.xyz', '.sdf']:
+            from openbabel import pybel
+
+            mol = next(pybel.readfile(ext[1:], fpath), None)
+            if mol is None:
+                raise ValueError("Pybel failed to parse the file.")
+            # Populate atom and bond data
+            mol_dict["atoms"] = [
+                {
+                    "element": atom.type,
+                    "coords": atom.coords,
+                    "index": atom.idx,
+                    "charge": atom.partialcharge
+                }
+                for atom in mol.atoms
+            ]
+            mol_dict["bonds"] = [
+                {
+                    "start_atom_idx": bond.GetBeginAtomIdx(),
+                    "end_atom_idx": bond.GetEndAtomIdx(),
+                    "bond_type": bond.GetBondOrder()
+                }
+                for bond in openbabel.OBMolBondIter(mol.OBMol)
+            ]
+
+        # Handling with MDAnalysis (for .pdb, .gro, and trajectory files)
+        elif ext in ['.pdb', '.gro', '.xyz', '.xtc', '.dcd', '.trr']:
+            import MDAnalysis as mda
+            u = mda.Universe(fpath)
+            atoms = u.atoms
+            mol_dict["atoms"] = [
+                {
+                    "element": atom.name,
+                    "coords": atom.position,
+                    "index": atom.id,
+                    "charge": atom.charge if hasattr(atom, 'charge') else None
+                }
+                for atom in atoms
+            ]
+            mol_dict["bonds"] = [
+                {"start_atom_idx": bond[0], "end_atom_idx": bond[1], "bond_type": 1}
+                for bond in u.bonds.indices
+            ]
+
+        # Handling with ASE (for .xyz, .pdb, and other atomic structure formats)
+        elif ext in ['.xyz', '.pdb', '.vasp', '.cif']:
+            from ase.io import read as ase_read
+            atoms = ase_read(fpath)
+            mol_dict["atoms"] = [
+                {
+                    "element": atom.symbol,
+                    "coords": atom.position,
+                    "index": i,
+                    "charge": None
+                }
+                for i, atom in enumerate(atoms)
+            ]
+            # ASE does not explicitly support bonds by default, so bonds are not populated here.
+
+        else:
+            raise ValueError(f"Unsupported file extension: {ext}")
+
+    except Exception as e:
+        print(f"Error loading molecule from {fpath}: {e}")
+        return None
+
+    return mol_dict
+ 
+class DockingConfig:
+    def __init__(self, receptor_file, ligand_smiles_list, center=(0, 0, 0), size=(20, 20, 20), output_dir="docking_results"):
+        self.receptor_file = receptor_file
+        self.ligand_smiles_list = ligand_smiles_list
+        self.center = center
+        self.size = size
+        self.output_dir = output_dir
+        os.makedirs(output_dir, exist_ok=True)
+
+def mol_to_pdbqt(mol, output_file):
+    """Converts an RDKit Mol object to PDBQT format."""
+    obConversion = openbabel.OBConversion()
+    obConversion.SetInAndOutFormats("mol", "pdbqt")
+    obMol = openbabel.OBMol()
+    obConversion.ReadString(obMol, Chem.MolToMolBlock(mol))
+    obConversion.WriteFile(obMol, output_file)
+
+def prepare_ligand(smiles, ligand_id):
+    """Prepare the ligand file in PDBQT format."""
+    mol = Chem.MolFromSmiles(smiles)
+    mol = Chem.AddHs(mol)
+    AllChem.EmbedMolecule(mol)
+    AllChem.UFFOptimizeMolecule(mol)
+    ligand_file = f"ligand_{ligand_id}.pdbqt"
+    mol_to_pdbqt(mol, ligand_file)
+    return ligand_file
+
+def run_docking(receptor_file, ligand_file, output_file, center, size):
+    """Runs Vina docking using the receptor and ligand files."""
+    vina_command = [
+        "vina",
+        "--receptor", receptor_file,
+        "--ligand", ligand_file,
+        "--center_x", str(center[0]),
+        "--center_y", str(center[1]),
+        "--center_z", str(center[2]),
+        "--size_x", str(size[0]),
+        "--size_y", str(size[1]),
+        "--size_z", str(size[2]),
+        "--out", output_file,
+        "--log", output_file.replace(".pdbqt", ".log")
+    ]
+    subprocess.run(vina_command, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
+
+def parse_vina_output(output_file):
+    """Parses Vina output log file to extract docking scores."""
+    scores = []
+    with open(output_file.replace(".pdbqt", ".log"), 'r') as f:
+        for line in f:
+            if line.startswith("REMARK VINA RESULT"):
+                score = float(line.split()[3])
+                scores.append(score)
+    return scores
+
+def docking_master_function(config: DockingConfig):
+    """Master function to run molecular docking for multiple ligands."""
+    receptor_pdbqt = config.receptor_file
+    results = {}
+
+    for i, smiles in enumerate(config.ligand_smiles_list):
+        ligand_file = prepare_ligand(smiles, ligand_id=i)
+        output_file = os.path.join(config.output_dir, f"docked_ligand_{i}.pdbqt")
+        
+        # Run docking for each ligand
+        run_docking(
+            receptor_file=receptor_pdbqt,
+            ligand_file=ligand_file,
+            output_file=output_file,
+            center=config.center,
+            size=config.size
+        )
+
+        # Parse docking results and store them
+        scores = parse_vina_output(output_file)
+        results[smiles] = scores
+        print(f"Ligand {i} (SMILES: {smiles}) docking scores: {scores}")
+
+        # Visualize individual docking result
+        visualize_docking(config.receptor_file, output_file, f"{config.output_dir}/ligand_{i}_visualization.png")
+
+        # Clean up intermediate files
+        os.remove(ligand_file)
+
+    # Plot binding affinity distribution
+    plot_binding_affinities(results, f"{config.output_dir}/binding_affinities.png")
+    return results
+
+def visualize_docking(receptor_file, ligand_file, dir_save):
+    """Generates a 2D visualization of the docking result using RDKit and Matplotlib."""
+    # Load the receptor and ligand molecules
+    receptor = Chem.MolFromPDBFile(receptor_file, removeHs=False)
+    ligand = Chem.MolFromPDBFile(ligand_file, removeHs=False)
+
+    # Draw the receptor and ligand
+    img = Draw.MolToImage(receptor, size=(300, 300))
+    img_ligand = Draw.MolToImage(ligand, size=(300, 300))
+
+    # Save images
+    img.save(dir_save.replace('.png', '_receptor.png'))
+    img_ligand.save(dir_save.replace('.png', '_ligand.png'))
+    
+    print(f"Saved 2D visualizations to {dir_save.replace('.png', '_receptor.png')} and {dir_save.replace('.png', '_ligand.png')}")
+
+
+def plot_binding_affinities(results, dir_save):
+    """Plots binding affinities for all ligands."""
+    ligands = list(results.keys())
+    affinities = [min(scores) for scores in results.values()]  # Minimum binding affinity per ligand
+
+    plt.figure(figsize=(10, 6))
+    plt.barh(ligands, affinities, color="skyblue")
+    plt.xlabel("Binding Affinity (kcal/mol)")
+    plt.ylabel("Ligands (SMILES)")
+    plt.title("Binding Affinities of Different Ligands")
+    plt.gca().invert_yaxis()
+    plt.tight_layout()
+    plt.savefig(dir_save)
+    plt.show()
+    print(f"Saved binding affinity plot to {dir_save}")
+    
+# 示例使用
+if __name__ == "__main__":
+    # 配置
+    receptor_file = "receptor.pdbqt"
+    ligand_smiles_list = ["CCO", "CCC", "CCN"]  # 示例的配体SMILES列表
+    docking_config = DockingConfig(
+        receptor_file=receptor_file,
+        ligand_smiles_list=ligand_smiles_list,
+        center=(10, 10, 10),  # 假设对接中心
+        size=(20, 20, 20)     # 假设对接区域大小
+    )
+
+    # 运行master function
+    docking_results = docking_master_function(docking_config)
+    print("Final docking results:", docking_results)