proteusPy.DisulfideLoader

   1"""
   2This module is part of the proteusPy package, a Python package for
   3the analysis and modeling of protein structures, with an emphasis on disulfide bonds.
   4This work is based on the original C/C++ implementation by Eric G. Suchanek. \n
   5
   6Last revision: 2025-03-26 13:14:09 -egs-
   7"""
   8
   9# Cα N, Cα, Cβ, C', Sγ Å ° ρ
  10
  11# pylint: disable=C0301
  12# pylint: disable=C0302
  13# pylint: disable=W1203
  14# pylint: disable=C0103
  15# pylint: disable=W0612
  16# pylint: disable=R1702
  17
  18# Cα N, Cα, Cβ, C', Sγ Å ° ρ
  19
  20import copy
  21import pickle
  22import time
  23from dataclasses import dataclass, field
  24from pathlib import Path
  25from typing import Any, Dict, List, Optional, Tuple
  26
  27import gdown
  28import pandas as pd
  29import plotly_express as px
  30from pympler import asizeof
  31
  32from proteusPy import __version__
  33from proteusPy.DisulfideBase import Disulfide, DisulfideList
  34from proteusPy.DisulfideClassGenerator import DisulfideClassGenerator
  35from proteusPy.DisulfideClassManager import DisulfideClassManager
  36from proteusPy.DisulfideExceptions import DisulfideParseWarning
  37from proteusPy.DisulfideStats import DisulfideStats
  38from proteusPy.DisulfideVisualization import DisulfideVisualization
  39from proteusPy.logger_config import create_logger
  40from proteusPy.ProteusGlobals import (
  41    DATA_DIR,
  42    LOADER_FNAME,
  43    LOADER_SUBSET_FNAME,
  44    SS_LIST_URL,
  45    SS_PICKLE_FILE,
  46)
  47
  48_logger = create_logger(__name__)
  49
  50try:
  51    # Check if running in Jupyter
  52    shell = get_ipython().__class__.__name__  # type: ignore
  53    if shell == "ZMQInteractiveShell":
  54        from tqdm.notebook import tqdm
  55    else:
  56        from tqdm import tqdm
  57except NameError:
  58    from tqdm import tqdm
  59
  60
  61@dataclass
  62class DisulfideLoader:
  63    """
  64    This class represents the disulfide database itself and is its primary means of accession.
  65    The entirety of the RCSB disulfide database is stored within the class via a
  66    proteusPy.DisulfideList, a ```Pandas``` .csv file, and a ```dict``` of
  67    indices mapping the PDB IDs into their respective list of disulfides. The datastructures allow
  68    simple, direct and flexible access to the disulfide structures contained herein.
  69    This makes it possible to access the disulfides by array index, PDB structure ID, disulfide
  70    name and class ID.
  71
  72    The class also provides methods for plotting distance and angle deviations
  73    as well as torsion statistics for the disulfides in the database.
  74
  75    The class can also render Disulfides overlaid on a common coordinate system to a pyVista
  76    window using the [display_overlay()](#DisulfideLoader.display_overlay) method. See below for examples.
  77
  78    Important note: For typical usage one will access the database via the `Load_PDB_SS()` function.
  79    The difference is that the latter function loads the compressed database from its single
  80    source. The `DisulfideLoader` class is used to build the Disulifde database with a
  81    specific cutoff, or for saving the database to a file.
  82
  83    Cutoff values of -1.0 indicate imposing no cutoffs on the data.
  84
  85    :param verbose: Flag to control output verbosity
  86    :type verbose: bool
  87    :param datadir: Directory containingA data files
  88    :type datadir: str
  89    :param picklefile: Name of the pickle file containing disulfide data
  90    :type picklefile: str
  91    :param quiet: Flag to suppress warnings
  92    :type quiet: bool
  93    :param subset: Flag to load only a subset of data
  94    :type subset: bool
  95    :param cutoff: Distance cutoff, (A) for filtering disulfides. Defaults to -1.0.
  96    :type cutoff: float
  97    :param sg_cutoff: SG distance cutoff, (A) for filtering disulfides. Defaults to -1.0.
  98    :type sg_cutoff: float
  99    :param percentile: Percentile cutoff for filtering disulfides. Must be between 0 and 100.
 100    Filters based on statistical cutoffs derived from the data.
 101    :type percentile: float
 102    :param minimum: Minimum atom distance for filtering disulfides. -1 is no filtering.
 103    :type minimum: float
 104    :param save: Flag to save the Loader to a file
 105    :type save: bool
 106    """
 107
 108    # Fields that serve as both instance attributes and initialization parameters
 109    datadir: str = field(default=DATA_DIR)
 110    picklefile: str = field(default=SS_PICKLE_FILE)
 111    subset: bool = field(default=False)
 112    cutoff: float = field(default=-1.0)
 113    sg_cutoff: float = field(default=-1.0)
 114    verbose: bool = field(default=False)
 115    percentile: float = field(default=-1.0)
 116    quiet: bool = field(default=False)
 117    minimum: float = field(default=-1.0)
 118    saveit: bool = field(default=False)
 119
 120    # Fields that are only used internally and don't need to be initialization parameters
 121    SSList: DisulfideList = field(
 122        default_factory=lambda: DisulfideList([], "ALL_PDB_SS"), init=False
 123    )
 124    SSDict: Dict = field(default_factory=dict, init=False)
 125    TorsionDF: pd.DataFrame = field(default_factory=pd.DataFrame, init=False)
 126    TotalDisulfides: int = field(default=0, init=False)
 127    IDList: List = field(default_factory=list, init=False)
 128    tclass: Optional[DisulfideClassManager] = field(default=None, init=False)
 129    class_generator: Optional[DisulfideClassGenerator] = field(default=None, init=False)
 130    timestamp: float = field(default_factory=time.time, init=False)
 131    version: str = field(default=__version__, init=False)
 132
 133    def __post_init__(self) -> None:
 134        """
 135        Initialize the DisulfideLoader after dataclass initialization.
 136        This method handles loading and processing of the disulfide data.
 137        """
 138
 139        cutoffs = {}
 140        old_length = new_length = 0
 141        full_path = Path(self.datadir) / self.picklefile
 142
 143        if self.verbose and not self.quiet:
 144            _logger.info(
 145                f"Reading disulfides from: {full_path}... ",
 146            )
 147
 148        try:
 149            # Check if the file exists before attempting to open it
 150            if not full_path.exists():
 151                fname = SS_PICKLE_FILE
 152                url = SS_LIST_URL
 153
 154                _fname = Path(DATA_DIR) / fname
 155
 156                if not _fname.exists():
 157                    if self.verbose:
 158                        _logger.info(
 159                            "Master SS list unavailable. Downloading Disulfide Database from Drive..."
 160                        )
 161                    gdown.download(url, str(_fname), quiet=False)
 162
 163            with open(full_path, "rb") as f:
 164                sslist = pickle.load(f)
 165
 166                if self.percentile > 0.0:
 167                    if self.percentile > 100.0:
 168                        raise ValueError("Percentile must be between 0 and 100.")
 169
 170                    cutoffs = DisulfideStats.calculate_cutoff_from_percentile(
 171                        sslist, percentile=self.percentile, verbose=self.verbose
 172                    )
 173
 174                    ca_cutoff = cutoffs["ca_cutoff_percentile"]
 175                    sg_cutoff = cutoffs["sg_cutoff_percentile"]
 176                    self.cutoff = ca_cutoff
 177                    self.sg_cutoff = sg_cutoff
 178
 179                    if self.verbose:
 180                        _logger.info(
 181                            f"Using percentile cutoffs: {ca_cutoff:.2f}, {sg_cutoff:.2f}"
 182                        )
 183
 184                old_length = len(sslist)
 185                filt = sslist.filter_by_distance(
 186                    distance=self.cutoff, distance_type="ca", minimum=-1.0
 187                )
 188                filt = DisulfideList(
 189                    filt,
 190                    f"filtered by Ca cutoff:{self.cutoff:.2f}, Sg cutoff: {self.sg_cutoff:.2f}",
 191                )
 192
 193                new_length = len(filt)
 194
 195                if self.verbose:
 196                    _logger.info(
 197                        "Filtered with Cα cutoff %.2f: old: %d, new: %d",
 198                        self.cutoff,
 199                        old_length,
 200                        new_length,
 201                    )
 202
 203                old_length = new_length
 204                filt = filt.filter_by_distance(
 205                    distance=self.sg_cutoff, distance_type="sg", minimum=-1.0
 206                )
 207                new_length = len(filt)
 208
 209                if self.verbose:
 210                    _logger.info(
 211                        "Filtered with Sγ: cutoff %.2f: old: %d, new: %d",
 212                        self.sg_cutoff,
 213                        old_length,
 214                        new_length,
 215                    )
 216                if self.subset:
 217                    self.SSList = DisulfideList(filt[:5000], "SUBSET_PDB_SS")
 218                else:
 219                    self.SSList = DisulfideList(filt, "ALL_PDB_SS")
 220
 221                self.SSDict = self._create_disulfide_dict()
 222                self.IDList = list(self.SSDict.keys())
 223
 224                self.TorsionDF = self.SSList.torsion_df
 225                self.TotalDisulfides = len(self.SSList)
 226                self.tclass = DisulfideClassManager(self, self.verbose)
 227                self.class_generator = DisulfideClassGenerator(verbose=self.verbose)
 228
 229            if self.verbose:
 230                _logger.info("Loader initialization complete.")
 231                self.describe()
 232
 233        except FileNotFoundError as e:
 234            _logger.error("File not found: %s", full_path)
 235            raise e
 236
 237        except Exception as e:
 238            _logger.error("An error occurred while loading the file: %s", full_path)
 239            raise e
 240        if self.saveit:
 241            self.save(
 242                savepath=DATA_DIR,
 243                verbose=self.verbose,
 244            )
 245
 246    # overload __getitem__ to handle slicing and indexing, and access by name or classid
 247    def __getitem__(self, item: int | slice | str) -> DisulfideList | Disulfide:
 248        """
 249        Implements indexing and slicing to retrieve DisulfideList objects from the
 250        DisulfideLoader. Supports:
 251
 252        - Integer indexing to retrieve a single DisulfideList
 253        - Slicing to retrieve a subset as a DisulfideList
 254        - Lookup by PDB ID to retrieve all Disulfides for that structure
 255        - Lookup by full disulfide name
 256        - Lookup by classid in the format 11111b or 11111o. The last char is the class type.
 257        - Lookup by classid in the format 11111. The base is 8 by default.
 258
 259        :param index: The index or key to retrieve the DisulfideList.
 260        :type index: int, slice, str
 261        :return: A DisulfideList object or a subset of it.
 262        :rtype: DisulfideList
 263        :raises DisulfideException: If the index or name is invalid.
 264        """
 265
 266        res = DisulfideList([], "none")
 267        ind_list = []
 268
 269        if isinstance(item, slice):
 270            indices = range(*item.indices(len(self.SSList)))
 271            ind_list = list(indices)
 272            name = f"pdb_slice[{ind_list[0]}:{ind_list[-1]+1}]"
 273            resolution = self.SSList[0].resolution
 274            sublist = [self.SSList[i] for i in indices]
 275            return DisulfideList(sublist, name, resolution)
 276
 277        if isinstance(item, int):
 278            if item < 0 or item >= self.TotalDisulfides:
 279                _logger.error(
 280                    "DisulfideLoader(): Index %d out of range 0-%d",
 281                    item,
 282                    self.TotalDisulfides - 1,
 283                )
 284                return res
 285
 286            res = self.SSList[item]
 287            return res
 288
 289        # if the item is a string, it could be a PDB ID or a full disulfide name
 290        # or a classid in the format 11111b or 11111o. the last char is the class type
 291
 292        if isinstance(item, str) and len(item) == 6 or len(item) == 5:  # classid
 293            res = self.extract_class(item, verbose=self.verbose)
 294            return res
 295
 296        # PDB_SS['4yys'] return a list of SS
 297        try:
 298            indices = self.SSDict[item]
 299            if indices:
 300                res = DisulfideList([], item)
 301                sslist = self.SSList
 302                for ind in indices:
 303                    res.append(sslist[ind])
 304            else:
 305                # try to find the full disulfide name
 306                res = self.SSList.get_by_name(item)  # full disulfide name
 307
 308        except KeyError as e:
 309            res = self.SSList.get_by_name(item)  # full disulfide name
 310
 311        if not res:
 312            _logger.error("DisulfideLoader(): Cannot find key %s in SSBond DB", item)
 313        return res
 314
 315    def __setitem__(self, index: int, item: Disulfide) -> None:
 316        self.SSList[index] = self._validate_ss(item)
 317
 318    def _validate_ss(self, value: Any) -> Disulfide:
 319        if isinstance(value, Disulfide):
 320            return value
 321        raise TypeError(f"Disulfide object expected, got {type(value).__name__}")
 322
 323    @property
 324    def average_resolution(self) -> float:
 325        """
 326        Return the average structure resolution for the given list.
 327        Result is cached since resolution values don't change after loading.
 328
 329        :return: Average resolution (A)
 330        """
 331        sslist = self.SSList
 332        valid_resolutions = [
 333            ss.resolution
 334            for ss in sslist
 335            if ss.resolution is not None and ss.resolution != -1.0
 336        ]
 337
 338        if not valid_resolutions:
 339            return -1.0
 340
 341        return sum(valid_resolutions) / len(valid_resolutions)
 342
 343    def binary_to_class(self, binary_class: str, base: int = 8) -> list[str]:
 344        """
 345        Convert a binary class string to an octant class string.
 346
 347        :param binary_class: The binary class string to convert.
 348        :param base: The base class to use, 2 or 8.
 349        :return: The octant class list.
 350        """
 351        return self.tclass.binary_to_class(binary_class, base)
 352
 353    def build_ss_from_idlist(self, idlist: List[str]) -> DisulfideList:
 354        """
 355        Return a DisulfideList of Disulfides for a given list of PDBIDs
 356
 357        :param idlist: List of PDBIDs, e.g. ['4yys', '2q7q']
 358        :return: DisulfideList
 359        """
 360        res = DisulfideList([], "RCSB_list")
 361        for pdbid, sslist in self.SSDict.items():
 362            if pdbid in idlist:
 363                for ssid in sslist:
 364                    res.append(self.SSList[ssid])
 365        return res
 366
 367    def _class_indices_from_tors_df(self, class_string: str, base: int = 8) -> pd.Index:
 368        """
 369        Return the row indices of the torsion dataframe that match the class string.
 370
 371        This method is used internally to find the indices of rows in the torsion dataframe
 372        that match the specified class string based on the given base.
 373
 374        :param class_string: The class string to match in the torsion dataframe.
 375        :type class_string: str
 376        :param base: The base class to use for matching, either 2 or 8. Defaults to 8.
 377        :type base: int
 378        :return: The row indices of the torsion dataframe that match the class string.
 379        :rtype: pd.Index
 380        :raises ValueError: If the base is not 2 or 8.
 381        """
 382        tors_df = self.TorsionDF
 383        match base:
 384            case 8:
 385                column = "octant_class_string"
 386            case 2:
 387                column = "binary_class_string"
 388            case _:
 389                raise ValueError(f"Base must be 2 or 8, not {base}")
 390
 391        return tors_df[tors_df[column] == class_string].index
 392
 393    def copy(self) -> "DisulfideLoader":
 394        """
 395        Return a copy of self.
 396
 397        :return: Copy of self
 398        """
 399        return copy.deepcopy(self)
 400
 401    def _create_disulfide_dict(self) -> Dict[str, List[int]]:
 402        """
 403        Create a dictionary from a list of disulfide objects where the key is the pdb_id
 404        and the value is a list of indices of the disulfide objects in the list.
 405
 406        This is an internal method used during initialization.
 407
 408        :param disulfide_list: List of disulfide objects.
 409        :type disulfide_list: list
 410        :return: Dictionary with pdb_id as keys and lists of indices as values.
 411        :rtype: dict
 412        """
 413        disulfide_list = self.SSList
 414
 415        disulfide_dict = {}
 416        for index, disulfide in enumerate(disulfide_list):
 417            if disulfide.pdb_id not in disulfide_dict:
 418                disulfide_dict[disulfide.pdb_id] = []
 419            disulfide_dict[disulfide.pdb_id].append(index)
 420        return disulfide_dict
 421
 422    def get_class_df(self, base: int = 8) -> pd.DataFrame:
 423        """
 424        Return the class incidence dataframe for the input base.
 425        Result is cached since class distributions don't change after loading.
 426
 427        :param base: The base class to use, 2 or 8.
 428        :return: pd.DataFrame
 429        """
 430        return self.tclass.get_class_df(base)
 431
 432    def extract_class(self, clsid: str, verbose: bool = False) -> DisulfideList:
 433        """
 434        Return the list of disulfides corresponding to the input `clsid`.
 435
 436        :param clsid: The class name to extract.
 437        :param verbose: If True, display progress bars, by default False
 438        :return: The list of disulfide bonds from the class.
 439        """
 440
 441        # cls = clsid[:5]
 442        cls = clsid
 443        ss_ids = []
 444        class_disulfides = None
 445
 446        try:
 447            ss_ids = self.tclass[clsid]
 448
 449        except KeyError:
 450            _logger.error("Cannot find key %s in SSBond DB", clsid)
 451            return DisulfideList([], cls, quiet=True)
 452
 453        tot_ss = len(ss_ids)
 454        class_disulfides = DisulfideList([], cls, quiet=True)
 455
 456        _pbar = (
 457            tqdm(range(tot_ss), total=tot_ss, leave=True) if verbose else range(tot_ss)
 458        )
 459
 460        for idx in _pbar:
 461            ssid = ss_ids[idx]
 462            class_disulfides.append(self[ssid])
 463
 464        return class_disulfides
 465
 466    def getlist(self) -> DisulfideList:
 467        """
 468        Return the list of Disulfides contained in the class.
 469
 470        :return: DisulfideList
 471        :rtype: DisulfideList
 472        """
 473        return copy.deepcopy(self.SSList)
 474
 475    def get_by_name(self, name: str = None) -> Optional[Disulfide]:
 476        """
 477        Return the Disulfide with the given name from the list.
 478        Result is cached since disulfide data doesn't change after loading.
 479        """
 480        for ss in self.SSList.data:
 481            if ss.name == name:
 482                return ss  # or ss.copy() !!!
 483        return None
 484
 485    def describe(self, memusg: bool = False) -> None:
 486        """
 487        Reveal key details about the Disulfide database stored in `self`. If `memusg` is True,
 488        the total RAM usage of the object is calculated and displayed — note that this process
 489        may take around 30 seconds on a 2024 MacBook Pro, M3 Max.
 490
 491        :param memusg: Set to True to include the RAM usage of the `DisulfideLoader` object.
 492        :return: None — just the facts!
 493        """
 494        # pylint: disable=E1101
 495        vers = self.version
 496        tot = self.TotalDisulfides
 497        pdbs = len(self.SSDict)
 498        ram = 0
 499        if memusg:
 500            ram = asizeof.asizeof(self) / (1024 * 1024 * 1024)
 501
 502        res = self.average_resolution
 503        cutoff = self.cutoff
 504        sg_cutoff = self.sg_cutoff
 505        percentile = self.percentile
 506        timestr = time.strftime("%Y-%m-%d %H:%M:%S", time.localtime(self.timestamp))
 507        ssMinMax = self.SSList.minmax_energy
 508        ssMin_name: Disulfide = ssMinMax[0].name
 509        ssMax_name: Disulfide = ssMinMax[1].name
 510
 511        print("")
 512        print("    🌟 RCSB Disulfide Database Summary 🌟")
 513        print(f"       🕒 Constructed: {timestr} 🕒")
 514        print(f"PDB IDs Present:               {pdbs}")
 515        print(f"Disulfides Loaded:             {tot}")
 516        print(f"Average Resolution:            {res:.2f} Å")
 517        print(f"Lowest Energy Disulfide:       {ssMin_name}")
 518        print(f"Highest Energy Disulfide:      {ssMax_name}")
 519        print(f"Cα Distance Cutoff:            {cutoff:.2f} Å")
 520        print(f"Sγ Distance Cutoff:            {sg_cutoff:.2f} Å")
 521        print(f"Percentile Cutoff:             {percentile:.2f} %")
 522        if memusg:
 523            print(f"Total RAM Usage:            {ram:.2f} GB")
 524        print(f"     ⚡ proteusPy Version: {vers} ⚡")
 525        print("")
 526
 527        return
 528
 529    def display_overlay(
 530        self, pdbid: str = "", verbose: bool = False, spin: bool = False
 531    ) -> None:
 532        """
 533        Display all disulfides for a given PDB ID overlaid in stick mode against
 534        a common coordinate frame. This allows us to see all of the disulfides
 535        at one time in a single view. Colors vary smoothy between bonds.
 536
 537        :param self: DisulfideLoader object initialized with the database.
 538        :param pdbid: the PDB id string, e.g. 4yys
 539        :param verbose: If True, display progress bars, by default False
 540        :type verbose: bool
 541        :param spin: If True, spin the display, by default False
 542        :type spin: bool
 543        :return: None
 544
 545        Example:
 546        >>> import proteusPy as pp
 547
 548        Instantiate the Loader with the SS database subset.
 549
 550        >>> PDB_SS = pp.Load_PDB_SS(verbose=False, subset=True)
 551
 552        Display the Disulfides from the PDB ID ```4yys```, overlaid onto
 553        a common reference (the proximal disulfides).
 554
 555        >>> PDB_SS.display_overlay('4yys', verbose=False)
 556
 557        You can also slice the loader and display as an overly.
 558        >>> PDB_SS[:8].display_overlay(verbose=False)
 559
 560        """
 561
 562        try:
 563            ssbonds = self[pdbid]
 564        except KeyError:
 565            _logger.error("Cannot find key %s in SSBond DB", pdbid)
 566            return
 567
 568        ssbonds.display_overlay(verbose=verbose, spin=spin)
 569        return
 570
 571    def getTorsions(self, pdbID: Optional[str] = None) -> pd.DataFrame:
 572        """
 573        Return the torsions, distances and energies defined by Torsion_DF_cols
 574
 575        :param pdbID: pdbID, defaults to None, meaning return entire dataset.
 576        :type pdbID: str, optional used to extract for a specific PDB structure. If not specified
 577            then return the entire dataset.
 578        :raises DisulfideParseWarning: Raised if not found
 579        :return: Torsions Dataframe
 580        :rtype: pd.DataFrame
 581
 582        Example:
 583        >>> import proteusPy as pp
 584        >>> PDB_SS = pp.Load_PDB_SS(verbose=False, subset=True)
 585        >>> Tor_DF = PDB_SS.getTorsions()
 586        """
 587        res_df = pd.DataFrame()
 588
 589        if pdbID:
 590            try:
 591                res = self.SSDict[pdbID]
 592                sel = self.TorsionDF["source"] == pdbID
 593                res_df = self.TorsionDF[sel]
 594                return res_df.copy()
 595            except KeyError as e:
 596                mess = f"Cannot find key {pdbID} in SSBond DB"
 597                _logger.error(mess)
 598                raise DisulfideParseWarning(mess) from e
 599        else:
 600            return copy.deepcopy(self.TorsionDF)
 601
 602    def list_binary_classes(self) -> None:
 603        """Enumerate the binary classes"""
 604        for k, v in enumerate(self.tclass.binaryclass_dict):
 605            print(f"Class: |{k}|, |{v}|")
 606
 607    def plot_classes(
 608        self,
 609        base: int = 8,
 610        class_string: Optional[str] = None,
 611        theme: str = "auto",
 612        log: bool = False,
 613        paginated: bool = False,
 614        page_size: int = 200,
 615    ) -> None:
 616        """
 617        Plot the classes for the given base.
 618
 619        :param base: The base class to use, 2 or 8.
 620        :param class_string: The class string to plot.
 621        :param theme: The theme to use for the plot ('auto', 'light', or 'dark').
 622        :param log: Whether to use a log scale for the y-axis.
 623        :param paginated: Whether to paginate the plot.
 624        :param page_size: Number of items per page.
 625        """
 626        # from proteusPy.DisulfideVisualization import DisulfideVisualization
 627
 628        DisulfideVisualization.plot_classes(
 629            self.tclass,
 630            class_string=class_string,
 631            base=base,
 632            theme=theme,
 633            log=log,
 634            page_size=page_size,
 635            paginated=paginated,
 636        )
 637
 638    def plot_classes_vs_cutoff(
 639        self,
 640        cutoff: float,
 641        steps: int = 50,
 642        base: int = 8,
 643        theme: str = "auto",
 644        verbose: bool = False,
 645    ) -> None:
 646        """
 647        Plot the total percentage and number of members for each octant class against the cutoff value.
 648
 649        :param cutoff: Percent cutoff value for filtering the classes.
 650        :type cutoff: float
 651        :param steps: Number of steps to take in the cutoff.
 652        :type steps: int
 653        :param base: The base class to use, 6 or 8.
 654        :type base: int
 655        :param theme: The theme to use for the plot ('auto', 'light', or 'dark'), defaults to 'auto'.
 656        :type theme: str
 657        :param verbose: Whether to display verbose output, defaults to False.
 658        :type verbose: bool
 659        :return: None
 660        :rtype: None
 661        """
 662        # from proteusPy.DisulfideVisualization import DisulfideVisualization
 663
 664        DisulfideVisualization.plot_classes_vs_cutoff(
 665            self.tclass, cutoff, steps, base, theme, verbose
 666        )
 667
 668    def plot_binary_to_eightclass_incidence(
 669        self,
 670        theme: str = "light",
 671        save: bool = False,
 672        savedir: str = ".",
 673        verbose: bool = False,
 674        log: bool = False,
 675    ) -> None:
 676        """Plot the incidence of all octant Disulfide classes for a given binary class.
 677
 678        :param theme: The theme to use for the plot
 679        :type theme: str
 680        :param save: Whether to save the plots
 681        :type save: bool
 682        :param savedir: Directory to save plots to
 683        :type savedir: str
 684        :param verbose: Whether to display verbose output
 685        :type verbose: bool
 686        :param log: Whether to use a log scale for the y-axis
 687        :type log: bool
 688        :return: None
 689        :rtype: None
 690        """
 691
 692        DisulfideVisualization.plot_binary_to_eightclass_incidence(
 693            self.tclass,
 694            theme=theme,
 695            save=save,
 696            savedir=savedir,
 697            verbose=verbose,
 698            log=log,
 699        )
 700
 701    def plot_count_vs_class_df(
 702        self,
 703        class_string: str,
 704        title: str = "title",
 705        theme: str = "auto",
 706        save: bool = False,
 707        savedir: str = ".",
 708        base: int = 8,
 709        verbose: bool = False,
 710        log: bool = False,
 711        sample_size: Optional[int] = None,
 712        page_size: Optional[int] = None,
 713    ) -> None:
 714        """
 715        Plot a line graph of count vs class ID using Plotly for the given disulfide class. The
 716        base selects the class type to plot: 2, 6, or 8, for binary, sextant, or octant classes.
 717
 718        :param class_string: The binary class string to be plotted.
 719        :param title: A string representing the title of the plot (default is 'title').
 720        :param theme: Theme to use for the plot
 721        :param save: Whether to save the plot
 722        :param savedir: Directory to save the plot to
 723        :param base: Base for class IDs (2 or 8)
 724        :param verbose: Whether to display verbose output
 725        :param log: Whether to use log scale for y-axis
 726        :param sample_size: Number of items to sample
 727        :param page_size: Number of items per page
 728        """
 729        # from proteusPy.DisulfideVisualization import DisulfideVisualization
 730        class_list = self.tclass.binary_to_class(class_string, base)
 731        df = self._enumerate_class_fromlist(class_list, base=base)
 732
 733        if sample_size:
 734            DisulfideVisualization.plot_count_vs_class_df_sampled(
 735                df,
 736                title,
 737                theme,
 738                save,
 739                savedir,
 740                base,
 741                verbose,
 742                log,
 743                sample_size,
 744            )
 745        elif page_size:
 746            DisulfideVisualization.plot_count_vs_class_df_paginated(
 747                df, title, theme, save, savedir, base, verbose, log, page_size
 748            )
 749        else:
 750            DisulfideVisualization.plot_count_vs_class_df(
 751                df, title, theme, save, savedir, base, verbose, log
 752            )
 753
 754    def plot_count_vs_classid(
 755        self,
 756        cls: Optional[str] = None,
 757        theme: str = "auto",
 758        base: int = 8,
 759        log: bool = True,
 760    ) -> None:
 761        """
 762        Plot a line graph of count vs class ID using Plotly.
 763
 764        :param cls: Specific class to plot (optional)
 765        :param theme: Theme to use for the plot
 766        :param base: Base for class IDs (2 or 8)
 767        :param log: Whether to use log scale for y-axis
 768        """
 769        # from proteusPy.DisulfideVisualization import DisulfideVisualization
 770
 771        DisulfideVisualization.plot_count_vs_classid(self.tclass, cls, theme, base, log)
 772
 773    def _enumerate_class_fromlist(
 774        self, sslist: List[str], base: int = 8
 775    ) -> pd.DataFrame:
 776        """
 777        Enumerate the classes from a list of class IDs and return a DataFrame with class IDs and their corresponding counts.
 778        Results are cached for improved performance on repeated calls.
 779
 780        :param sslist: A list of class IDs to enumerate.
 781        :param base: The base value for the enumeration, by default 8.
 782        :return: A DataFrame with columns "class_id" and "count" representing the class IDs and their corresponding counts.
 783        """
 784        x = []
 785        y = []
 786
 787        for cls in sslist:
 788            if cls is not None:
 789                _y = self.tclass.sslist_from_classid(cls, base=base)
 790                # it's possible to have 0 SS in a class
 791                if _y is not None:
 792                    # only append if we have both.
 793                    x.append(cls)
 794                    y.append(len(_y))
 795
 796        sslist_df = pd.DataFrame(columns=["class_id", "count"])
 797        sslist_df["class_id"] = x
 798        sslist_df["count"] = y
 799        return sslist_df
 800
 801    def save(
 802        self,
 803        savepath: str = DATA_DIR,
 804        verbose: bool = False,
 805        fname: Optional[str] = None,
 806    ) -> None:
 807        """
 808        Save a copy of the fully instantiated Loader to the specified file.
 809
 810        :param savepath: Path to save the file, defaults to DATA_DIR
 811        :param fname: Filename, defaults to LOADER_FNAME
 812        :param verbose: Verbosity, defaults to False
 813        """
 814        self.version = __version__
 815
 816        fname = None
 817        if self.subset:
 818            fname = LOADER_SUBSET_FNAME
 819        else:
 820            fname = LOADER_FNAME
 821
 822        _fname = Path(savepath) / fname
 823
 824        if verbose:
 825            _logger.info("Writing Disulfide Loader to: %s...", _fname)
 826
 827        with open(str(_fname), "wb+") as f:
 828            pickle.dump(self, f)
 829
 830        if verbose:
 831            _logger.info("Done saving loader.")
 832
 833    def plot_disulfides_vs_pdbid(self, cutoff: int = 1) -> Tuple[List[str], List[int]]:
 834        """
 835        Plots the number of disulfides versus pdbid.
 836
 837        :param cutoff: The minimum number of disulfides a PDB ID must have to be included in the plot.
 838        :type cutoff: int
 839        :return: A tuple containing the list of PDB IDs and the corresponding number of disulfides.
 840        :rtype: tuple
 841        """
 842        pdbids = []
 843        num_disulfides = []
 844
 845        for pdbid, disulfides in self.SSDict.items():
 846            if len(disulfides) > cutoff:
 847                pdbids.append(pdbid)
 848                num_disulfides.append(len(disulfides))
 849
 850        # Create a DataFrame
 851        df = pd.DataFrame({"PDB ID": pdbids, "Number of Disulfides": num_disulfides})
 852        fig = px.bar(
 853            df,
 854            x="PDB ID",
 855            y="Number of Disulfides",
 856            title=f"Disulfides vs PDB ID with cutoff: {cutoff}, {len(pdbids)} PDB IDs",
 857        )
 858        fig.update_layout(
 859            xaxis_title="PDB ID",
 860            yaxis_title="Number of Disulfides",
 861            xaxis_tickangle=-90,
 862        )
 863        fig.show()
 864
 865        return pdbids, num_disulfides
 866
 867    def plot_distances(
 868        self,
 869        distance_type: str = "ca",
 870        cutoff: float = -1,
 871        comparison: str = "less",
 872        theme: str = "auto",
 873        log: bool = True,
 874    ) -> None:
 875        """
 876        Plot the distances for the disulfides in the loader.
 877
 878        :param distance_type: The type of distance to plot ('ca' for Cα-Cα distance, 'sg' for Sγ-Sγ distance)
 879        :param cutoff: The cutoff value for the distance, defaults to -1 (no cutoff)
 880        :param comparison: if 'less' then plot distances less than the cutoff, if 'greater' then plot distances greater than the cutoff
 881        :param theme: The theme to use for the plot ('auto', 'light', or 'dark')
 882        :param log: Whether to use a log scale for the y-axis
 883        """
 884        # from proteusPy.DisulfideVisualization import DisulfideVisualization
 885
 886        # sslist = self.SSList
 887        # distances = sslist.extract_distances(distance_type, comparison, cutoff)
 888
 889        self.SSList.plot_distances(
 890            distance_type=distance_type,
 891            cutoff=cutoff,
 892            comparison=comparison,
 893            theme=theme,
 894            log=log,
 895        )
 896
 897    def plot_deviation_scatterplots(
 898        self, verbose: bool = False, theme: str = "auto"
 899    ) -> None:
 900        """
 901        Plot scatter plots for Bondlength_Deviation, Angle_Deviation Ca_Distance
 902        and SG_Distance.
 903
 904        :param verbose: Whether to display the plot in the notebook. Default is False.
 905        :type verbose: bool
 906        :param theme: One of 'Auto', 'Light', or 'Dark'. Default is 'Auto'.
 907        :type light: str
 908        :return: None
 909        """
 910        self.SSList.plot_deviation_scatterplots(verbose=verbose, theme=theme)
 911
 912    def plot_deviation_histograms(
 913        self, theme: str = "auto", verbose: bool = True
 914    ) -> None:
 915        """
 916        Plot histograms for Bondlength_Deviation, Angle_Deviation, and Ca_Distance.
 917        """
 918        self.SSList.plot_deviation_histograms(theme=theme, verbose=verbose)
 919
 920    def sslist_from_class(
 921        self, class_string: str, base: int = 8, cutoff: float = 0.0
 922    ) -> DisulfideList:
 923        """
 924        Return a DisulfideList containing Disulfides with the given class_string.
 925
 926        :param class_string: The class string to search for.
 927        :param base: The base of the class string. Default is 8.
 928        :param cutoff: The % cutoff value for the class. Default is 0.0.
 929        :return: DisulfideList containing Disulfides with the given class_string.
 930        """
 931        sslist_name = f"{class_string}_{base}_{cutoff:.2f}"
 932        sslist = DisulfideList([], sslist_name)
 933
 934        indices = self._class_indices_from_tors_df(class_string, base=base)
 935
 936        for i in indices:
 937            sslist.append(self[i])
 938
 939        return sslist
 940
 941    def display_torsion_statistics(
 942        self,
 943        class_id: Optional[str] = None,
 944        display: bool = True,
 945        save: bool = False,
 946        fname: str = "ss_torsions.png",
 947        theme: str = "auto",
 948        verbose: bool = False,
 949        dpi: int = 300,
 950        figure_size: tuple[int, int] = (4, 3),
 951    ) -> None:
 952        """
 953        Display torsion and distance statistics for all Disulfides in the loader.
 954        If a class ID is provided, display statistics for that class only.
 955
 956        :param class_id: The class ID to display statistics for. Default is None.
 957        :type class_id: str
 958        :param display: Whether to display the plot in the notebook. Default is True.
 959        :type display: bool
 960        :param save: Whether to save the plot as an image file. Default is False.
 961        :type save: bool
 962        :param fname: The name of the image file to save. Default is 'ss_torsions.png'.
 963        :type fname: str
 964        :param theme: One of 'Auto', 'Light', or 'Dark'. Default is 'Auto'.
 965        :type theme: str
 966        :param verbose: Whether to display verbose output. Default is False.
 967        :type verbose: bool
 968        :param dpi: Dots per inch for the plot. Default is 300.
 969        :type dpi: int
 970        :param figure_size: Size of the figure as a tuple (width, height). Default is (4, 3).
 971        :type figure_size: tuple
 972        :return: None
 973        """
 974        if class_id:
 975            DisulfideVisualization.display_torsion_class_df(
 976                self.TorsionDF,
 977                class_id,
 978                display=display,
 979                save=save,
 980                fname=fname,
 981                theme=theme,
 982                dpi=dpi,
 983                figure_size=figure_size,
 984            )
 985        else:
 986            self.SSList.display_torsion_statistics(
 987                display=display,
 988                save=save,
 989                fname=fname,
 990                theme=theme,
 991                verbose=verbose,
 992                dpi=dpi,
 993                figure_size=figure_size,
 994            )
 995
 996    def classes_vs_cutoff(self, cutoff: float, base: int = 8) -> int:
 997        """
 998        Return number of members for the octant class for a given cutoff value.
 999
1000        :param cutoff: Percent cutoff value for filtering the classes.
1001        :return: None
1002        """
1003
1004        class_df = self.tclass.filter_class_by_percentage(cutoff, base=base)
1005        return class_df.shape[0]
1006
1007    def display_torsion_class_df(
1008        self,
1009        class_id: str,
1010        display: bool = True,
1011        save: bool = False,
1012        fname: str = "ss_torsions.png",
1013        theme: str = "auto",
1014        dpi: int = 300,
1015        figure_size: tuple[int, int] = (4, 3),
1016    ) -> None:
1017        """
1018        Display torsion and distance statistics for a given class ID using the TorsionDF dataframe.
1019
1020        :param class_id: The class ID to display statistics for (e.g. '11111b' for binary or '11111o' for octant)
1021        :param display: Whether to display the plot in the notebook
1022        :param save: Whether to save the plot as an image file
1023        :param fname: The name of the image file to save
1024        :param theme: The theme to use for the plot ('auto', 'light', or 'dark')
1025        :param dpi: DPI (dots per inch) for the saved image, controls the resolution
1026        :param figure_size: Tuple of (width, height) in inches for the figure size
1027        """
1028
1029        DisulfideVisualization.display_torsion_class_df(
1030            self.TorsionDF,
1031            class_id,
1032            display=display,
1033            save=save,
1034            fname=fname,
1035            theme=theme,
1036            dpi=dpi,
1037            figure_size=figure_size,
1038        )
1039
1040    def plot_3d_hexbin_leftright(
1041        self,
1042        width: int = 800,
1043        height: int = 600,
1044        gridsize: int = 80,
1045        tormin: float = -180.0,
1046        tormax: float = 180.0,
1047        scaling: str = "sqrt",
1048        column1: str = "chi2",
1049        column2: str = "chi4",
1050        title: Optional[str] = None,
1051    ) -> None:
1052        """
1053        Create 3D hexbin plots for left and right-handed chi2-chi4 correlations with customizable z-scaling.
1054
1055        :param loader: Loader object to retrieve torsion data
1056        :type loader: proteusPy.PDB_SS
1057        :param width: Window width in pixels
1058        :type width: int, optional
1059        :default width: 800
1060        :param height: Window height in pixels
1061        :type height: int, optional
1062        :default height: 600
1063        :param gridsize: Number of bins for hexbin
1064        :type gridsize: int, optional
1065        :default gridsize: 30
1066        :param tormin: Minimum torsion angle
1067        :type tormin: float, optional
1068        :default tormin: -180.0
1069        :param tormax: Maximum torsion angle
1070        :type tormax: float, optional
1071        :default tormax: 180.0
1072        :param scaling: Scaling method for z-values ('linear', 'sqrt', 'log', 'power')
1073        :type scaling: str, optional
1074        :default scaling: 'sqrt'
1075        :param column1: Name of the first column (x-axis)
1076        :type column1: str, optional
1077        :default column1: 'chi2'
1078        :param column2: Name of the second column (y-axis)
1079        :type column2: str, optional
1080        :default column2: 'chi4'
1081        :param title: Title of the plot
1082        :type title: str, optional
1083        """
1084
1085        DisulfideVisualization.plot_3d_hexbin_leftright(
1086            self,
1087            width=width,
1088            height=height,
1089            gridsize=gridsize,
1090            tormin=tormin,
1091            tormax=tormax,
1092            scaling=scaling,
1093            column1=column1,
1094            column2=column2,
1095            title=title,
1096        )
1097
1098
1099# class ends
1100
1101
1102def Load_PDB_SS(
1103    loadpath: str = DATA_DIR,
1104    verbose: bool = False,
1105    subset: bool = False,
1106    percentile: float = -1.0,
1107    force: bool = False,
1108) -> DisulfideLoader:
1109    """
1110    Load the fully instantiated Disulfide database from the specified file. This function
1111    will load the pre-built database if available, or bootstrap a new loader by downloading
1112    the data from Google Drive if needed. Use the provided parameters to control the loading
1113    behavior, filtering cutoffs, and verbosity.
1114
1115    :param loadpath: Path from which to load the database; defaults to DATA_DIR.
1116    :type loadpath: str
1117    :param verbose: If True, enables verbose logging; defaults to False.
1118    :type verbose: bool
1119    :param subset: If True, loads the subset database; otherwise loads the full database.
1120    :type subset: bool
1121    :param cutoff: Cα distance cutoff used to filter disulfides; defaults to CA_CUTOFF.
1122    :type cutoff: float
1123    :param sg_cutoff: Sγ distance cutoff used to filter disulfides; defaults to SG_CUTOFF.
1124    :type sg_cutoff: float
1125    :param force: If True, forces re-loading from Google Drive even if the file exists; defaults to False.
1126    :type force: bool
1127    :param percentile: Percentile (0-100) to compute cutoffs dynamically; if set to -1.0, the percentile method is not used.
1128    :type percentile: float
1129    :return: An instance of DisulfideLoader containing the loaded disulfide database.
1130    :rtype: DisulfideLoader
1131
1132    Example:
1133        >>> from proteusPy import Load_PDB_SS, create_logger
1134        >>> import logging
1135        >>> _logger = create_logger("testing")
1136        >>> _logger.setLevel(logging.WARNING)
1137        >>> loader = Load_PDB_SS(verbose=False, subset=True)
1138        >>> print(loader[0])
1139        <Disulfide 6dmb_203A_226A, Source: 6dmb, Resolution: 3.0 Å>
1140    """
1141
1142    # normally the .pkl files are local, EXCEPT for the first run from a newly-installed proteusPy
1143    # distribution. In that case we need to download the files for all disulfides and the subset
1144    # from my Google Drive. This is a one-time operation.
1145
1146    _fname_sub = Path(loadpath) / LOADER_SUBSET_FNAME
1147    _fname_all = Path(loadpath) / LOADER_FNAME
1148    _fpath = _fname_sub if subset else _fname_all
1149
1150    sg_cutoff = ca_cutoff = -1.0
1151
1152    if not _fpath.exists() or force is True:
1153        if verbose:
1154            _logger.info(f"Bootstrapping new loader: {str(_fpath)}... ")
1155
1156        loader = Bootstrap_PDB_SS(
1157            loadpath=loadpath,
1158            verbose=verbose,
1159            subset=subset,
1160            force=force,
1161            percentile=percentile,
1162        )
1163        loader.save(
1164            savepath=loadpath,
1165            verbose=verbose,
1166        )
1167        return loader
1168
1169    if verbose:
1170        _logger.info("Reading disulfides from: %s...", _fpath)
1171
1172    with open(_fpath, "rb") as f:
1173        loader = pickle.load(f)
1174    if verbose:
1175        _logger.info("Done reading disulfides from: %s...", _fpath)
1176        loader.describe()
1177
1178    return loader
1179
1180
1181def Bootstrap_PDB_SS(
1182    loadpath: str = DATA_DIR,
1183    verbose: bool = True,
1184    subset: bool = False,
1185    force: bool = False,
1186    fake: bool = False,
1187    percentile: float = -1.0,
1188) -> Optional[DisulfideLoader]:
1189    """
1190    Download and instantiate the disulfide databases from Google Drive.
1191
1192    This function downloads the disulfide master SS list from Google Drive if it doesn't
1193    already exist in the specified load path or if the force flag is set to True.
1194    It then loads the disulfide data from the downloaded file and initializes a
1195    DisulfideLoader instance.
1196
1197    :param loadpath: Path from which to load the data, defaults to DATA_DIR
1198    :type loadpath: str
1199    :param cutoff: Cutoff value for disulfide loading, defaults to -1.0 (no filtering)
1200    :type cutoff: float
1201    :param sg_cutoff: Cutoff value for disulfide loading, defaults to -1.0 (no filtering)
1202    :type sg_cutoff: float
1203    :param verbose: Flag to enable verbose logging, defaults to False
1204    :type verbose: bool
1205    :param subset: Flag to indicate whether to load a subset of the data, defaults to False
1206    :type subset: bool
1207    :param force: Flag to force download even if the file exists, defaults to False
1208    :type force: bool
1209    :return: An instance of DisulfideLoader initialized with the loaded data
1210    :rtype: DisulfideLoader
1211    """
1212
1213    fname = SS_PICKLE_FILE
1214    url = SS_LIST_URL
1215
1216    # _fname = Path(loadpath) / fname
1217    full_path = Path(loadpath) / fname
1218
1219    if not full_path.exists() or force is True:
1220        if verbose:
1221            _logger.warning("Can't find %s. Downloading from Drive...", full_path)
1222
1223        if not fake:
1224            gdown.download(url, str(full_path), quiet=False)
1225        else:
1226            if verbose:
1227                _logger.warning("Fake download: %s", full_path)
1228                return None
1229    if verbose:
1230        _logger.info(
1231            "Building loader from: %s with cutoffs %s s...",
1232            full_path,
1233            percentile,
1234        )
1235
1236    loader = DisulfideLoader(
1237        datadir=DATA_DIR,
1238        subset=subset,
1239        verbose=verbose,
1240        percentile=percentile,
1241    )
1242
1243    if loader.TotalDisulfides == 0:
1244        _logger.error("No disulfides loaded!")
1245        return None
1246
1247    if verbose:
1248        _logger.info("Done building loader.")
1249
1250    return loader
1251
1252
1253if __name__ == "__main__":
1254    import doctest
1255
1256    doctest.testmod()
1257
1258# End of file