Classes

class Pocket (file_path: str = '',
block_type: str = '',
block_content: str = '',
color='red',
name=None,
index: int | None = 0,
props: dict | None = None)

Expand source code

class Pocket:
    """
    A class representing a molecular pocket structure with various properties and methods for manipulation.
    This class handles loading, visualization, and analysis of molecular pocket structures,
    primarily dealing with PDB files. It supports initialization from either a file path
    or direct content block.

    Attributes:
        color (str): Color representation of the pocket (default: "red").
        index (int): Index identifier for the pocket (default: 0).
        props (dict): Dictionary containing pocket properties.
        name (str): Name identifier for the pocket.
        file_path (Path): Path to the pocket file.
        structure (AtomArray): Biotite structure object representing the pocket.
        coordinates (ndarray): Numpy array of atomic coordinates.
        block_type (str): Type of structure block (e.g., "pdb").
        block_content (str): Content of the structure block.

    Args:
        file_path (str, optional): Path to the pocket structure file.
        block_type (str, optional): Type of structure block.
        block_content (str, optional): Content of the structure block.
        color (str, optional): Color for visualization (default: "red").
        name (str, optional): Name identifier for the pocket.
        index (int, optional): Index identifier (default: 0).
        props (dict, optional): Dictionary of pocket properties.

    Raises:
        ValueError: If the structure cannot be loaded.
        ValueError: If neither or both file_path and block_content are provided.
        ValueError: If the block_type is not supported (only "pdb" is supported).
        FileNotFoundError: If the specified file_path doesn't exist.

    Examples:
        # Create from file
        pocket = Pocket(file_path="path/to/pocket.pdb", name="Pocket1")

        # Create from content
        pocket = Pocket(block_content="...", block_type="pdb", name="Pocket2")

        # Visualize pocket
        pocket.visualize()

        # Get pocket center
        center = pocket.get_center()
    """

    def __init__(
        self,
        file_path: str = "",
        block_type: str = "",
        block_content: str = "",
        color="red",
        name=None,
        index: Optional[int] = 0,
        props: Optional[dict] = None,
    ):
        self.color = color
        self.index = index
        self.props = props

        self.name = name
        self.file_path = None
        self.structure = None
        self.coordinates = None

        file_path_obj = Path(file_path).absolute() if file_path else None
        extension = file_path_obj.suffix.lower() if file_path_obj else ""
        if not block_type and extension:
            block_type = extension.lstrip(".")  # Remove the leading dot
        self.block_type = block_type.lower()
        self.block_content = block_content

        sources_provided = sum(bool(x) for x in [file_path, block_content])
        if sources_provided != 1:
            raise ValueError("Please provide exactly one of file_path or block_content.")

        from_block = False
        try:
            if file_path:
                self.file_path = Path(file_path).absolute()
                if not self.file_path.exists():
                    raise FileNotFoundError(f"The file {self.file_path} does not exist.")

                if not self.block_type:
                    self.block_type = self.file_path.suffix.lstrip(".").lower()

                self.block_content = self.file_path.read_text()

                pocket_file_dir = self.get_directory()
                if str(pocket_file_dir) != str(self.file_path.parent):
                    try:
                        destination = Path(pocket_file_dir) / self.file_path.name
                        shutil.copy2(self.file_path, destination)
                        self.file_path = destination
                    except Exception as e:
                        DEFAULT_LOGGER.log_error(f"Failed to copy file to destination: {str(e)}")
                        raise
            elif block_content:
                self.block_content = block_content
                if not self.block_type:
                    raise ValueError("block_type must be provided when initializing with block_content.")
                from_block = True
                pocket_file_dir = self.get_directory()

            if self.block_content:
                if self.block_type not in ["pdb"]:
                    raise ValueError(f"Only pdb file formats are supported (given {self.block_type})")
                self.structure = self.load_structure_from_block(self.block_content, self.block_type)

            if self.structure is None:
                raise ValueError("Structure could not be loaded.")

            # Type checking for AtomArrayStack
            if isinstance(self.structure, AtomArrayStack):
                self.structure = self.structure[0]

            DEFAULT_LOGGER.log_info(
                f"Loaded structure from {self.file_path if self.file_path else 'block content'}. Selected structure index: {0}"
            )

            if self.name is None:
                if self.file_path:
                    self.name = self.file_path.stem
                else:
                    self.name = "Unknown_Pocket"
                    directory = Path(pocket_file_dir)
                    num = len(list(directory.glob(f"{self.name}*")))
                    self.name = f"{self.name}_{num + 1}"

            self.coordinates = self.structure.coord

            if from_block:
                directory = Path(pocket_file_dir)
                self.file_path = directory / f"{self.name}.{self.block_type}"
                self.write_to_file(self.file_path)

        except Exception as e:
            DEFAULT_LOGGER.log_error(f"Failed to initialize pocket: {str(e)}")
            raise

    def load_structure_from_block(self, block_content: str, block_type: str):
        """
        Load molecular structure from a text block content.

        This method creates a Structure object from a text block containing structural data
        in a supported format.

        Args:
            block_content (str): Text content containing the structure data
            block_type (str): Format of the structure data (currently only "pdb" supported)

        Returns:
            Structure: A Structure object representing the molecular structure

        Raises:
            ValueError: If block_type is not supported
        """
        if block_type == "pdb":
            pdb_file = PDBFile.read(io.StringIO(block_content))
            structure = pdb_file.get_structure()
        else:
            raise ValueError(f"Unsupported block type: {block_type}")
        return structure

    @staticmethod
    def load_structure(structure_file_path: str):
        """
        Load a protein structure from a PDB file.

        Args:
            structure_file_path (str): Path to the PDB file containing the protein structure.

        Returns:
            Structure: The loaded protein structure object.

        Raises:
            FileNotFoundError: If the specified PDB file does not exist.
            ValueError: If the PDB file is invalid or cannot be parsed.
        """

        structure_file = PDBFile.read(structure_file_path)
        structure = structure_file.get_structure()
        return structure

    def write_to_file(self, output_path: str, output_format: str = "pdb"):
        """
        Write the current structure to a file in the specified format.
        This method writes the current structure to a file, with support for different output formats.
        If the output format is not PDB, it first writes to a temporary PDB file and then converts to
        the desired format.

        Args:
            output_path (str): Path where the structure file should be written.
            output_format (str, optional): Format of the output file. Defaults to "pdb".

        Raises:
            Exception: If writing to file fails, the error is logged via DEFAULT_LOGGER.

        Example:
            >>> pocket.write_to_file("structure.pdb")
            >>> pocket.write_to_file("structure.mol2", output_format="mol2")
        """

        def write_to_pdb_file(structure, output_path):
            pdb_file = PDBFile()
            pdb_file.set_structure(structure)
            pdb_file.write(output_path)

        try:
            path = Path(output_path)
            if not path.parent.exists():
                path.parent.mkdir(parents=True, exist_ok=True)

            if path.suffix.lower() != ".pdb":
                with tempfile.NamedTemporaryFile(delete=True) as temp:
                    write_to_pdb_file(self.structure, temp.name)
                    convert_file("pdb", temp.name, output_format, output_path)
            else:
                write_to_pdb_file(self.structure, output_path)
            DEFAULT_LOGGER.log_info(f"Current structure written to {output_path}.")

        except Exception as e:
            DEFAULT_LOGGER.log_error(f"Failed to write structure to file {output_path}: {str(e)}")

    @jupyter_visualization
    def visualize(self):
        """
        Visualizes the protein pocket in 3D using ProteinViewer.

        This method creates a 3D visualization of the pocket structure using the ProteinViewer class.
        It configures the visualization with appropriate colors and displays the pocket name and properties.

        Returns:
            The rendered 3D visualization of the protein pocket

        Notes:
            - Uses the pocket's file path and name for visualization
            - Applies color coding based on the pocket's index
            - Falls back to index 0 if the pocket index exceeds available colors
            - Displays pocket name and properties in the visualization legend
        """
        pocket_paths = [str(self.file_path)]
        pocket_names = ["Name: " + self.name + " | " + self.pocket_props()]

        viewer = ProteinViewer("", format="pdb")

        pocket_config = viewer.get_pocket_visualization_config()
        # Ensure self.index is within the bounds of surface_colors
        if self.index >= len(pocket_config.surface_colors):
            # Log a warning or adjust index appropriately
            DEFAULT_LOGGER.log_warning(f"Index {self.index} is out of bounds for surface_colors. Resetting to 0.")
            self.index = 0  # Default to the first color if out of bounds

        pocket_config.surface_colors = [pocket_config.surface_colors[self.index]]
        print(
            "\n|\n".join(
                colored("■", pocket_config.surface_colors[0]) + " " + pocket_name for pocket_name in pocket_names
            )
        )

        return viewer.render_protein_with_pockets(pocket_paths=pocket_paths, pocket_config=pocket_config)

    def pocket_props(self):
        """
        Formats the properties of a protein pocket into a single string line.

        Returns:
            str: A formatted string containing pocket properties in the format:
                 'Volume: {value}ų | Drugability score: {value}'
                 Returns empty string if no properties are available.
        """
        properties_line = ""
        if self.props:
            properties_line = (
                f"Volume: {self.props.get('volume', 'N/A')}ų | "
                f"Drugability score: {self.props.get('drugability_score', 'N/A')}"
            )
        return properties_line

    def __repr__(self):
        properties_line = ""
        if self.props:
            properties_line = (
                f"  Volume: {self.props.get('volume', 'N/A')}ų, "
                f"Total SASA: {self.props.get('total_SASA', 'N/A')} "
                f"Polar SASA: {self.props.get('polar_SASA', 'N/A')} "
                f"Polar/Apolar SASA ratio: {self.props.get('polar_apolar_SASA_ratio', 'N/A')} "
                f"Hydrophobicity: {self.props.get('hydrophobicity', 'N/A')} "
                f"Polarity: {self.props.get('polarity', 'N/A')} "
                f"Drugability score: {self.props.get('drugability_score', 'N/A')}"
            )

        return (
            f"Pocket:\n  Name: {self.name}\n{properties_line}  Block type: {self.block_type}\n"
            "Available Fields: {block_type, block_content, file_path, name, coordinates}"
        )

    def __str__(self):
        properties_line = ""
        if self.props:
            properties_line = (
                f"  Volume: {self.props.get('volume', 'N/A')}ų, "
                f"Total SASA: {self.props.get('total_SASA', 'N/A')}, "
                f"Polar SASA: {self.props.get('polar_SASA', 'N/A')}, "
                f"Polar/Apolar SASA ratio: {self.props.get('polar_apolar_SASA_ratio', 'N/A')}, "
                f"Hydrophobicity: {self.props.get('hydrophobicity', 'N/A')}, "
                f"Polarity: {self.props.get('polarity', 'N/A')}, "
                f"Drugability score: {self.props.get('drugability_score', 'N/A')}"
            )

        return (
            f"Pocket:\n  Name: {self.name}\n{properties_line}  Block type: {self.block_type}\n"
            "Available Fields: {block_type, block_content, file_path, name, coordinates}"
        )

    def get_center(self) -> Optional[List[float]]:
        """
        Calculate and return the center coordinates of the pocket.

        This method computes the arithmetic mean of all coordinates in the pocket
        to determine its center point.

        Returns:
            Optional[List[float]]: A list containing the x, y, z coordinates of the pocket's center.
                                  Returns None if coordinates are not available.
        """
        if self.coordinates is None:
            DEFAULT_LOGGER.log_warning("Coordinates are not available for this Pocket.")
            return None
        center = self.coordinates.mean(axis=0)
        DEFAULT_LOGGER.log_info(f"Calculated center coordinates: {center.tolist()}")
        return [float(x) for x in center.tolist()]

    @staticmethod
    def get_directory() -> str:
        """
        Returns the base directory path for storing pocket-related data.

        This method creates (if not exists) and returns the path to a 'pockets' directory
        under the working directory. The directory is created with parent directories if needed.

        Returns:
            str: Absolute path to the pockets base directory as a string
        """
        pockets_base_dir = Path(WORKING_DIR) / "pockets"
        pockets_base_dir.mkdir(parents=True, exist_ok=True)

        return str(pockets_base_dir)

    def update_coordinates(self, coords: np.ndarray):
        """
        Updates the coordinates of the pocket structure.

        Args:
            coords (np.ndarray): New coordinates to update the pocket structure with.
                                Should be a numpy array containing the coordinate data.

        Updates:
            - self.structure.coord: Updates the coordinates in the structure object
            - self.coordinates: Updates the local coordinates attribute

        Note:
            This method performs an in-place update of the coordinates and logs the action.
        """
        self.structure.coord = coords
        self.coordinates = coords
        DEFAULT_LOGGER.log_info("Pocket coordinates has been inplaced updated.")

A class representing a molecular pocket structure with various properties and methods for manipulation. This class handles loading, visualization, and analysis of molecular pocket structures, primarily dealing with PDB files. It supports initialization from either a file path or direct content block.

Attributes

color : str

Color representation of the pocket (default: "red").

index : int

Index identifier for the pocket (default: 0).

props : dict

Dictionary containing pocket properties.

name : str

Name identifier for the pocket.

file_path : Path

Path to the pocket file.

structure : AtomArray

Biotite structure object representing the pocket.

coordinates : ndarray

Numpy array of atomic coordinates.

block_type : str

Type of structure block (e.g., "pdb").

block_content : str

Content of the structure block.

Args

file_path : str, optional

Path to the pocket structure file.

block_type : str, optional

Type of structure block.

block_content : str, optional

Content of the structure block.

color : str, optional

Color for visualization (default: "red").

name : str, optional

Name identifier for the pocket.

index : int, optional

Index identifier (default: 0).

props : dict, optional

Dictionary of pocket properties.

Raises

ValueError

If the structure cannot be loaded.

ValueError

If neither or both file_path and block_content are provided.

ValueError

If the block_type is not supported (only "pdb" is supported).

FileNotFoundError

If the specified file_path doesn't exist.

Examples

Create from file

pocket = Pocket(file_path="path/to/pocket.pdb", name="Pocket1")

Create from content

pocket = Pocket(block_content="…", block_type="pdb", name="Pocket2")

Visualize pocket

pocket.visualize()

Get pocket center

center = pocket.get_center()

Static methods

def get_directory() ‑> str

Expand source code

@staticmethod
def get_directory() -> str:
    """
    Returns the base directory path for storing pocket-related data.

    This method creates (if not exists) and returns the path to a 'pockets' directory
    under the working directory. The directory is created with parent directories if needed.

    Returns:
        str: Absolute path to the pockets base directory as a string
    """
    pockets_base_dir = Path(WORKING_DIR) / "pockets"
    pockets_base_dir.mkdir(parents=True, exist_ok=True)

    return str(pockets_base_dir)

Returns the base directory path for storing pocket-related data.

This method creates (if not exists) and returns the path to a 'pockets' directory under the working directory. The directory is created with parent directories if needed.

Returns

str

Absolute path to the pockets base directory as a string

def load_structure(structure_file_path: str)

Expand source code

@staticmethod
def load_structure(structure_file_path: str):
    """
    Load a protein structure from a PDB file.

    Args:
        structure_file_path (str): Path to the PDB file containing the protein structure.

    Returns:
        Structure: The loaded protein structure object.

    Raises:
        FileNotFoundError: If the specified PDB file does not exist.
        ValueError: If the PDB file is invalid or cannot be parsed.
    """

    structure_file = PDBFile.read(structure_file_path)
    structure = structure_file.get_structure()
    return structure

Load a protein structure from a PDB file.

Args

structure_file_path : str

Path to the PDB file containing the protein structure.

Returns

Structure

The loaded protein structure object.

Raises

FileNotFoundError

If the specified PDB file does not exist.

ValueError

If the PDB file is invalid or cannot be parsed.

Methods

def get_center(self) ‑> List[float] | None

Expand source code

def get_center(self) -> Optional[List[float]]:
    """
    Calculate and return the center coordinates of the pocket.

    This method computes the arithmetic mean of all coordinates in the pocket
    to determine its center point.

    Returns:
        Optional[List[float]]: A list containing the x, y, z coordinates of the pocket's center.
                              Returns None if coordinates are not available.
    """
    if self.coordinates is None:
        DEFAULT_LOGGER.log_warning("Coordinates are not available for this Pocket.")
        return None
    center = self.coordinates.mean(axis=0)
    DEFAULT_LOGGER.log_info(f"Calculated center coordinates: {center.tolist()}")
    return [float(x) for x in center.tolist()]

Calculate and return the center coordinates of the pocket.

This method computes the arithmetic mean of all coordinates in the pocket to determine its center point.

Returns

Optional[List[float]]

A list containing the x, y, z coordinates of the pocket's center. Returns None if coordinates are not available.

def load_structure_from_block(self, block_content: str, block_type: str)

Expand source code

def load_structure_from_block(self, block_content: str, block_type: str):
    """
    Load molecular structure from a text block content.

    This method creates a Structure object from a text block containing structural data
    in a supported format.

    Args:
        block_content (str): Text content containing the structure data
        block_type (str): Format of the structure data (currently only "pdb" supported)

    Returns:
        Structure: A Structure object representing the molecular structure

    Raises:
        ValueError: If block_type is not supported
    """
    if block_type == "pdb":
        pdb_file = PDBFile.read(io.StringIO(block_content))
        structure = pdb_file.get_structure()
    else:
        raise ValueError(f"Unsupported block type: {block_type}")
    return structure

Load molecular structure from a text block content.

This method creates a Structure object from a text block containing structural data in a supported format.

Args

block_content : str

Text content containing the structure data

block_type : str

Format of the structure data (currently only "pdb" supported)

Returns

Structure

A Structure object representing the molecular structure

Raises

ValueError

If block_type is not supported

def pocket_props(self)

Expand source code

def pocket_props(self):
    """
    Formats the properties of a protein pocket into a single string line.

    Returns:
        str: A formatted string containing pocket properties in the format:
             'Volume: {value}ų | Drugability score: {value}'
             Returns empty string if no properties are available.
    """
    properties_line = ""
    if self.props:
        properties_line = (
            f"Volume: {self.props.get('volume', 'N/A')}ų | "
            f"Drugability score: {self.props.get('drugability_score', 'N/A')}"
        )
    return properties_line

Formats the properties of a protein pocket into a single string line.

Returns

str

A formatted string containing pocket properties in the format: 'Volume: {value}ų | Drugability score: {value}' Returns empty string if no properties are available.

def update_coordinates(self, coords: numpy.ndarray)

Expand source code

def update_coordinates(self, coords: np.ndarray):
    """
    Updates the coordinates of the pocket structure.

    Args:
        coords (np.ndarray): New coordinates to update the pocket structure with.
                            Should be a numpy array containing the coordinate data.

    Updates:
        - self.structure.coord: Updates the coordinates in the structure object
        - self.coordinates: Updates the local coordinates attribute

    Note:
        This method performs an in-place update of the coordinates and logs the action.
    """
    self.structure.coord = coords
    self.coordinates = coords
    DEFAULT_LOGGER.log_info("Pocket coordinates has been inplaced updated.")

Updates the coordinates of the pocket structure.

Args

coords : np.ndarray

New coordinates to update the pocket structure with. Should be a numpy array containing the coordinate data.

Updates

• self.structure.coord: Updates the coordinates in the structure object
• self.coordinates: Updates the local coordinates attribute

Note

This method performs an in-place update of the coordinates and logs the action.

def visualize(*args, **kwargs)

Expand source code

def wrapper(*args, **kwargs):
    html_visualization = func(*args, **kwargs)
    return JupyterViewer.visualize(html_visualization)

def write_to_file(self, output_path: str, output_format: str = 'pdb')

Expand source code

def write_to_file(self, output_path: str, output_format: str = "pdb"):
    """
    Write the current structure to a file in the specified format.
    This method writes the current structure to a file, with support for different output formats.
    If the output format is not PDB, it first writes to a temporary PDB file and then converts to
    the desired format.

    Args:
        output_path (str): Path where the structure file should be written.
        output_format (str, optional): Format of the output file. Defaults to "pdb".

    Raises:
        Exception: If writing to file fails, the error is logged via DEFAULT_LOGGER.

    Example:
        >>> pocket.write_to_file("structure.pdb")
        >>> pocket.write_to_file("structure.mol2", output_format="mol2")
    """

    def write_to_pdb_file(structure, output_path):
        pdb_file = PDBFile()
        pdb_file.set_structure(structure)
        pdb_file.write(output_path)

    try:
        path = Path(output_path)
        if not path.parent.exists():
            path.parent.mkdir(parents=True, exist_ok=True)

        if path.suffix.lower() != ".pdb":
            with tempfile.NamedTemporaryFile(delete=True) as temp:
                write_to_pdb_file(self.structure, temp.name)
                convert_file("pdb", temp.name, output_format, output_path)
        else:
            write_to_pdb_file(self.structure, output_path)
        DEFAULT_LOGGER.log_info(f"Current structure written to {output_path}.")

    except Exception as e:
        DEFAULT_LOGGER.log_error(f"Failed to write structure to file {output_path}: {str(e)}")

Write the current structure to a file in the specified format. This method writes the current structure to a file, with support for different output formats. If the output format is not PDB, it first writes to a temporary PDB file and then converts to the desired format.

Args

output_path : str

Path where the structure file should be written.

output_format : str, optional

Format of the output file. Defaults to "pdb".

Raises

Exception

If writing to file fails, the error is logged via DEFAULT_LOGGER.

Example

>>> pocket.write_to_file("structure.pdb")
>>> pocket.write_to_file("structure.mol2", output_format="mol2")