notebooks/quickstart.ipynb

{
 "cells": [
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "upper-filing",
   "metadata": {},
   "source": [
    "# rcsbsearchapi quickstart\n",
    "\n",
    "This notebook contains examples from the rcsbsearchapi [quickstart](https://rcsbsearchapi.readthedocs.io/en/latest/quickstart.html)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "african-monthly",
   "metadata": {},
   "outputs": [],
   "source": [
    "from rcsbsearchapi.search import TextQuery, AttributeQuery, SequenceQuery, SeqMotifQuery, StructSimilarityQuery, Attr\n",
    "from rcsbsearchapi import rcsb_attributes as attrs\n",
    "from rcsbsearchapi.const import STRUCTURE_ATTRIBUTE_SEARCH_SERVICE\n",
    "from rcsbsearchapi.const import CHEMICAL_ATTRIBUTE_SEARCH_SERVICE"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "sublime-karen",
   "metadata": {},
   "source": [
    "## Operator syntax\n",
    "\n",
    "Here is an example from the [RCSB PDB Search API](http://search.rcsb.org/#search-example-1) page, using the operator syntax. This query finds symmetric dimers having a twofold rotation with the DNA-binding domain of a heat-shock transcription factor. Full text and attribute search are used.\n",
    "\n",
    "Note the use of standard comparison operators (`==`, `>` etc) for rcsb attributes and set operators for combining queries."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "enabling-america",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Create terminals for each query\n",
    "q1 = TextQuery(\"heat-shock transcription factor\")\n",
    "q2 = attrs.rcsb_struct_symmetry.symbol == \"C2\"\n",
    "q3 = attrs.rcsb_struct_symmetry.kind == \"Global Symmetry\"\n",
    "q4 = attrs.rcsb_entry_info.polymer_entity_count_DNA >= 1\n",
    "\n",
    "# combined using bitwise operators (&, |, ~, etc)\n",
    "query = q1 & (q2 & q3 & q4) # AND of all queries\n",
    "\n",
    "# Call the query to execute it\n",
    "for assemblyid in query(\"assembly\"): # return type specified as \"assembly\"\n",
    "    print(assemblyid)\n"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "accomplished-passion",
   "metadata": {},
   "source": [
    "Attribute names can be found in the [RCSB PDB schema](http://search.rcsb.org/rcsbsearch/v2/metadata/schema). They can also be found via tab completion, or by iterating:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "supported-observer",
   "metadata": {},
   "outputs": [],
   "source": [
    "[a.attribute for a in attrs if \"authors\" in a.attribute]"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "other-grant",
   "metadata": {},
   "source": [
    "## Fluent syntax\n",
    "\n",
    "Here is the same example using the [fluent](https://en.wikipedia.org/wiki/Fluent_interface) syntax:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "polish-indonesia",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Start with a Attr or TextQuery, then add terms\n",
    "results = TextQuery(\"heat-shock transcription factor\").and_(\n",
    "    # Add attribute node as fully-formed AttributeQuery\n",
    "    AttributeQuery(attribute=\"rcsb_struct_symmetry.symbol\", operator=\"exact_match\", value=\"C2\") \\\n",
    "    # Add attribute node as Attr with chained operations\n",
    "    .and_(Attr(\"rcsb_struct_symmetry.kind\")).exact_match(\"Global Symmetry\") \\\n",
    "    # Add attribute node by name (converted to Attr) with chained operations\n",
    "    .and_(\"rcsb_entry_info.polymer_entity_count_DNA\").greater_or_equal(1)\n",
    "    ).exec(\"assembly\")\n",
    "# Exec produces an iterator of IDs\n",
    "\n",
    "for assemblyid in results:\n",
    "    print(assemblyid)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "032e01e6",
   "metadata": {},
   "source": [
    "## Structural Attribute Search and Chemical Search Combination\n",
    "\n",
    "Grouping of a Structural Attribute query and Chemical Attribute query is permitted as long as grouping is done correctly and search services are specified accordingly. More details on attributes that are available for text searches can be found on the [RCSB PDB Search API](https://search.rcsb.org/#search-attributes) page."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fafc71bf",
   "metadata": {},
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "eef8a33f",
   "metadata": {},
   "outputs": [],
   "source": [
    "# By default, service is set to \"text\" for structural attribute search\n",
    "q1 = AttributeQuery(\"exptl.method\", \"exact_match\", \"electron microscopy\",\n",
    "                    STRUCTURE_ATTRIBUTE_SEARCH_SERVICE # this constant specifies \"text\" service\n",
    "                    )\n",
    "\n",
    "# Need to specify chemical attribute search service - \"text_chem\"\n",
    "q2 = AttributeQuery(\"drugbank_info.brand_names\", \"contains_phrase\", \"tylenol\",\n",
    "                    CHEMICAL_ATTRIBUTE_SEARCH_SERVICE # this constant specifies \"text_chem\" service\n",
    "                    )\n",
    "\n",
    "query = q1 & q2 # combining queries\n",
    "\n",
    "list(query())\n"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "4b9a442c",
   "metadata": {},
   "source": [
    "## Computed Structure Models\n",
    "\n",
    "The [RCSB PDB Search API](https://search.rcsb.org/#results_content_type) page provides information on how to include Computed Structure Models (CSMs) into a search query. Here is a code example below.\n",
    "\n",
    "This query returns IDs for experimental and computed structure models associated with \"hemoglobin\". Queries for *only* computed models or *only* experimental models can also be made (default)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "85b707a5",
   "metadata": {},
   "outputs": [],
   "source": [
    "q1 = TextQuery(\"hemoglobin\")\n",
    "\n",
    "# add parameter as a list with either \"computational\" or \"experimental\" or both as list values\n",
    "q2 = q1(return_content_type=[\"computational\", \"experimental\"])\n",
    "\n",
    "list(q2)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "2c4f3120",
   "metadata": {},
   "source": [
    "## Return Types and Attribute Search\n",
    "\n",
    "A search query can return different result types when a return type is specified. Below are examples on specifying return types Polymer Entities,\n",
    "\n",
    "Non-polymer Entities, Polymer Instances, and Molecular Definitions, using a Structure Attribute query. More information on return types can be found in the [RCSB PDB Search API](https://search.rcsb.org/#building-search-request) page."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "508ba82a",
   "metadata": {},
   "outputs": [],
   "source": [
    "q1 = AttributeQuery(\"rcsb_entry_container_identifiers.entry_id\", \"in\", [\"4HHB\"]) # query for 4HHB deoxyhemoglobin\n",
    "\n",
    "print(\"Polymer Entities:\")\n",
    "for poly in q1(\"polymer_entity\"): # include return type as a string parameter for query object\n",
    "    print(poly)\n",
    "\n",
    "print(\"Non-polymer Entities:\")\n",
    "for nonPoly in q1(\"non_polymer_entity\"):\n",
    "    print(nonPoly)\n",
    "\n",
    "print(\"Polymer Instances:\")\n",
    "for polyInst in q1(\"polymer_instance\"):\n",
    "    print(polyInst)\n",
    "\n",
    "print(\"Molecular Definitions:\")\n",
    "for mol in q1(\"mol_definition\"):\n",
    "    print(mol)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "669974e0",
   "metadata": {},
   "source": [
    "## Counting Results\n",
    "\n",
    "If only the number of results is desired, the count function can be used. This query returns the number of experimental models associated with \"hemoglobin\"."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8ef9edac",
   "metadata": {},
   "outputs": [],
   "source": [
    "q1 = TextQuery(\"hemoglobin\")\n",
    "\n",
    "# N.B., Just as shown above for `query()`, `return_type` and `return_content_type` can also be specified as parameters to `count()`\n",
    "q1.count()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4e335e45",
   "metadata": {},
   "source": [
    "## Obtaining Scores for Results\n",
    "\n",
    "Results can be returned alongside additional metadata, including result scores. To return this metadata, set the `results_verbosity` parameter to \"verbose\" (all metadata), \"minimal\" (scores only), or \"compact\" (default, no metadata). If set to \"verbose\" or \"minimal\", results will be returned as a list of dictionaries. For example, here we get all experimental models associated with \"hemoglobin\", along with their scores."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "05321bba",
   "metadata": {},
   "outputs": [],
   "source": [
    "q1 = TextQuery(\"hemoglobin\")\n",
    "for idscore in list(q1(results_verbosity=\"minimal\")):\n",
    "    print(idscore)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "e9c4bb0f",
   "metadata": {},
   "source": [
    "## Sequence Query\n",
    "\n",
    "Below is an example from the [RCSB PDB Search API](https://search.rcsb.org/#search-example-3) page. Queries can be made using DNA, RNA, and protein sequences when specified using the SearchQuery class. In this example, we are finding macromolecular PDB entities that share 90% sequence identity with GTPase HRas protein from *Gallus gallus* (*Chicken)*."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6be93cf2",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Use SequenceQuery class and add parameters\n",
    "results = SequenceQuery(\"MTEYKLVVVGAGGVGKSALTIQLIQNHFVDEYDPTIEDSYRKQVVIDGET\" +\n",
    "                        \"CLLDILDTAGQEEYSAMRDQYMRTGEGFLCVFAINNTKSFEDIHQYREQI\" +\n",
    "                        \"KRVKDSDDVPMVLVGNKCDLPARTVETRQAQDLARSYGIPYIETSAKTRQ\" +\n",
    "                        \"GVEDAFYTLVREIRQHKLRKLNPPDESGPGCMNCKCVIS\", 1, 0.9)\n",
    "\n",
    "# results(\"polymer_entity\") produces an iterator of IDs with return type - polymer entities\n",
    "for polyid in results(\"polymer_entity\"):\n",
    "    print(polyid)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "adopted-gnome",
   "metadata": {},
   "source": [
    "## Sequence Motif Query\n",
    "\n",
    "Below is an example from the [RCSB PDB Search API](https://search.rcsb.org/#search-example-6) page, using the sequence motif search function. \n",
    "This query retrives occurences of the His2/Cys2 Zinc Finger DNA-binding domain as represented by its PROSITE signature."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8a96fd62",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Use SeqMotifQuery class and add parameters\n",
    "results = SeqMotifQuery(\"C-x(2,4)-C-x(3)-[LIVMFYWC]-x(8)-H-x(3,5)-H.\",\n",
    "                        pattern_type=\"prosite\",\n",
    "                        sequence_type=\"protein\")\n",
    "\n",
    "# results(\"polymer_entity\") produces an iterator of IDs with return type - polymer entities\n",
    "for polyid in results(\"polymer_entity\"):\n",
    "    print(polyid)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0eaa8e55",
   "metadata": {},
   "source": [
    "Below is an example query for the zinc finger motif that binds Zn in a DNA-binding domain:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8363d7e2",
   "metadata": {},
   "outputs": [],
   "source": [
    "results = SeqMotifQuery(\"C.{2,4}C.{12}H.{3,5}H\", pattern_type=\"regex\", sequence_type=\"protein\")\n",
    "\n",
    "for polyid in results(\"polymer_entity\"):\n",
    "    print(polyid)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4d41bc89",
   "metadata": {},
   "source": [
    "Below is an example query for SH3 domains:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "22d92d94",
   "metadata": {},
   "outputs": [],
   "source": [
    "# By default, the pattern_type argument is \"simple\" and the sequence_type argument is \"protein\".\n",
    "results = SeqMotifQuery(\"XPPXP\")  # X is used as a \"variable residue\" and can be any amino acid. \n",
    "\n",
    "for polyid in results(\"polymer_entity\"):\n",
    "    print(polyid)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "baf150f9",
   "metadata": {},
   "source": [
    "All 3 of these pattern types can be used to search for DNA and RNA sequences as well.\n",
    "Demonstrated are 2 queries, one DNA and one RNA, using the simple pattern type:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "91b035e4",
   "metadata": {},
   "outputs": [],
   "source": [
    "from rcsbsearchapi.search import SeqMotifQuery\n",
    "\n",
    "# DNA query: this is a query for a T-Box.\n",
    "dna = SeqMotifQuery(\"TCACACCT\", sequence_type=\"dna\")\n",
    "\n",
    "print(\"DNA results:\")\n",
    "for polyid in dna(\"polymer_entity\"):\n",
    "    print(polyid)\n",
    "\n",
    "# RNA query: 6C RNA motif\n",
    "rna = SeqMotifQuery(\"CCCCCC\", sequence_type=\"rna\")\n",
    "print(\"RNA results:\")\n",
    "for polyid in rna(\"polymer_entity\"):\n",
    "    print(polyid)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ad25a8a0",
   "metadata": {},
   "source": [
    "## Structure Similarity Query\n",
    "\n",
    "The PDB archive can be queried using the 3D shape of a protein structure. To perform this query, 3D protein structure data must be provided as an input or parameter, A chain ID or assembly ID must be specified, whether the input structure data should be compared to Assemblies or Polymer Entity Instance (Chains) is required, and defining the search type as either strict or relaxed is required. More information on how Structure Similarity Queries work can be found on the [RCSB PDB Structure Similarity Search](https://www.rcsb.org/docs/search-and-browse/advanced-search/structure-similarity-search) page."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "13ab72d1",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Basic query: querying using entry ID and default values assembly ID \"1\", operator \"strict\", and target search space \"Assemblies\"\n",
    "q1 = StructSimilarityQuery(entry_id=\"4HHB\")\n",
    "\n",
    "# Same example but with parameters explicitly specified\n",
    "q1 = StructSimilarityQuery(structure_search_type=\"entry_id\",\n",
    "                           entry_id=\"4HHB\",## Structure Similarity Query\n",
    "                           structure_input_type=\"assembly_id\",\n",
    "                           assembly_id=\"1\",\n",
    "                           operator=\"strict_shape_match\",\n",
    "                           target_search_space=\"assembly\"\n",
    "                           )\n",
    "for id in q1(\"assembly\"):\n",
    "    print(id)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "38ad5f07",
   "metadata": {},
   "source": [
    "Below is a more complex example that utilizes chain ID, relaxed search operator, and polymer entity instance or target search space. Specifying whether the input structure type is chain id or assembly id is very important. For example, specifying chain ID as the input structure type but inputting an assembly ID can lead to an error."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "72d7cec5",
   "metadata": {},
   "outputs": [],
   "source": [
    "# More complex query with entry ID value \"4HHB\", chain ID \"B\", operator \"relaxed\", and target search space \"Chains\"\n",
    "q2 = StructSimilarityQuery(structure_search_type=\"entry_id\",\n",
    "                                   entry_id=\"4HHB\",\n",
    "                                   structure_input_type=\"chain_id\",\n",
    "                                   chain_id=\"B\",\n",
    "                                   operator=\"relaxed_shape_match\",\n",
    "                                   target_search_space=\"polymer_entity_instance\")\n",
    "list(q2())"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1144d0fb",
   "metadata": {},
   "source": [
    "Structure similarity queries also allow users to upload a file from their local computer or input a file url from the website to query the PDB archive for similar proteins. The file represents a target protein structure in the file formats \"cif\", \"bcif\", \"pdb\", \"cif.gz\", or \"pdb.gz\". If a user wants to use a file url for queries, the user must specify the structure search type, the value (being the url), and the file format of the file. This is also the same case for file upload, except the value is the absolute path leading to the file that is in the local machine. An example for file url is below for 4HHB (hemoglobin)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2d83cd13",
   "metadata": {},
   "outputs": [],
   "source": [
    "q3 = StructSimilarityQuery(structure_search_type=\"file_url\",\n",
    "                           file_url=\"https://files.rcsb.org/view/4HHB.cif\",\n",
    "                           file_format=\"cif\")\n",
    "list(q3())\n",
    "\n",
    "# If you want to upload your own structure file for similarity search, you can do so by using the `file_path` parameter:\n",
    "q4 = StructSimilarityQuery(structure_search_type=\"file_upload\",\n",
    "                           file_path=\"/PATH/TO/FILE.cif\",  # specify local model file path\n",
    "                           file_format=\"cif\")\n",
    "list(q4())"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0f894e48",
   "metadata": {},
   "source": [
    "### Structure Motif Query Examples\n",
    "\n",
    "The PDB Archive can also be queried by using a \"motif\" found in these 3D structures. To perform this type of query, an entry_id or a file URL/path must be provided, along with residues (which are parts of 3D structures.) This is the bare minimum needed to make a search, but there are lots of other parameters that can be added to a Structure Motif Query (see [full search schema](https://search.rcsb.org/redoc/index.html)).\n",
    "\n",
    "To make a Structure Motif Query, you must first define anywhere from 2-10 \"residues\" that will be used in the query. Each individual residue has a Chain ID, Operator, Residue Number, and Exchanges (optional) that can be declared in that order using positonal arguments, or using the \"chain_id\", \"struct_oper_id\", and \"label_seq_id\" to define what parameter you are passing through. All 3 of the required parameters must be included, or the package will throw an AssertionError. \n",
    "\n",
    "Each residue can have a maximum of 4 Exchanges, and each query can only have 16 exchanges total. Violating any of these rules will cause the package to throw an AssertionError. \n",
    "\n",
    "Examples of how to instantiate Residues can be found below. These can then be put into a list and passed through to a Structure Motif Query."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "842be053",
   "metadata": {},
   "outputs": [],
   "source": [
    "from rcsbsearchapi.search import StructureMotifResidue\n",
    "\n",
    "# construct a Residue with a Chain ID of A, an operator of 1, a residue \n",
    "# number of 192, and Exchanges of \"LYS\" and \"HIS\"\n",
    "Res1 = StructureMotifResidue(\"A\", \"1\", 192, [\"LYS\", \"HIS\"])\n",
    "# as for what is a valid \"Exchange\", the package provides these as a literal,\n",
    "# and they should be type checked. \n",
    "\n",
    "# you can also specify the arguments:\n",
    "# this query is the same as above. \n",
    "Res2 = StructureMotifResidue(struct_oper_id=\"1\", chain_id=\"A\", exchanges=[\"LYS\", \"HIS\"], label_seq_id=192)\n",
    "\n",
    "# after delcaring a minimum of 2 and as many as 10 residues, they can be passed into a list for use in the query itself:\n",
    "Res3 = StructureMotifResidue(\"A\", \"1\", 162)  # exchanges are optional\n",
    "\n",
    "ResList = [Res1, Res3]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c35487df",
   "metadata": {},
   "source": [
    "From there, these Residues can be used in a query. As stated before, you can only include 2 - 10 residues in a query. If you fail to provide residues for a query, or provide the wrong amount, the package will throw a ValueError. \n",
    "\n",
    "For a Structure Motif Query using an entry_id, the only other necessary value that must be passed into the query is the residue list. The default type of query is an entry_id query. \n",
    "\n",
    "As this type of query has a lot of optional parameters, do *not* use positional arguments as more than likely an error will occur. \n",
    "\n",
    "Below is an example of a basic entry_id Structure Motif Query, with the residues declared earlier:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "04698fd0",
   "metadata": {},
   "outputs": [],
   "source": [
    "from rcsbsearchapi.search import StructMotifQuery\n",
    "\n",
    "q1 = StructMotifQuery(entry_id=\"2MNR\", residue_ids=ResList)\n",
    "q1_res = list(q1(return_type='polymer_entity'))\n",
    "print(len(q1_res))\n",
    "print(q1.count(return_type='polymer_entity'))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6c9d0e94",
   "metadata": {},
   "source": [
    "Like with Structure Similarity Queries, a file url or filepath can also be provided to the program. These can take the place of an entry_id. \n",
    "\n",
    "For a file url query, you *must* provide both a valid file URL (a string), and the file's file extension (also as a string). Failure to provide these elements correctly will cause the package to throw an AssertionError. \n",
    "\n",
    "Below is an example of the same query as above, only this time providing a file url:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "14bd65a5",
   "metadata": {},
   "outputs": [],
   "source": [
    "link = \"https://files.rcsb.org/view/2MNR.cif\"\n",
    "q2 = StructMotifQuery(structure_search_type=\"file_url\", url=link, file_extension=\"cif\", residue_ids=ResList)\n",
    "# structure_search_type MUST be provided. A mismatched query type will cause an error. \n",
    "list(q2())"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "06e2eeb0",
   "metadata": {},
   "source": [
    "Like with Structure Similarity Queries, a filepath to a file may also be provided. This file must be a valid file accepted by the search API. A file extension must also be provided with the file upload. \n",
    "\n",
    "The query would look something like this. Note that this is abstracted for the purpose of notebook portability.\n",
    "```python\n",
    "filepath = \"/absolute/path/to/file.cif\"\n",
    "q3 = StructMotifQuery(structure_search_type=\"file_upload\", file_path=filepath, file_extension=\"cif\", residue_ids=ResList)\n",
    "\n",
    "list(q3())\n",
    "```\n",
    "There are many additional parameters that Structure Motif Query supports. These include a variety of features such as backbone distance tolerance, side chain distance tolerance, angle tolerance, RMSD cutoff, limits (stop searching after this many hits), atom pairing schemes, motif pruning strategy, allowed structures, and excluded structures. These can be mixed and matched as needed to make accurate and useful queries. All of these have some default value which is used when a parameter isn't provided. These parameters conform to the defaults used by the Search API. \n",
    "\n",
    "Below will demonstrate how to define these parameters using non-positional arguments:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d38e7b32",
   "metadata": {},
   "outputs": [],
   "source": [
    "# specifying backbone distance tolerance: 0-3, default is 1\n",
    "# allowed backbone distance tolerance in Angstrom. \n",
    "backbone = StructMotifQuery(entry_id=\"2MNR\", backbone_distance_tolerance=2, residue_ids=ResList)\n",
    "list(backbone())\n",
    "\n",
    "# specifying sidechain distance tolerance: 0-3, default is 1\n",
    "# allowed side-chain distance tolerance in Angstrom.\n",
    "sidechain = StructMotifQuery(entry_id=\"2MNR\", side_chain_distance_tolerance=2, residue_ids=ResList)\n",
    "list(sidechain())\n",
    "\n",
    "# specifying angle tolerance: 0-3, default is 1\n",
    "# allowed angle tolerance in multiples of 20 degrees. \n",
    "angle = StructMotifQuery(entry_id=\"2MNR\", angle_tolerance=2, residue_ids=ResList)\n",
    "list(angle())\n",
    "\n",
    "# specifying RMSD cutoff: >=0, default is 2\n",
    "# Threshold above which hits will be filtered by RMSD\n",
    "rmsd = StructMotifQuery(entry_id=\"2MNR\", rmsd_cutoff=1, residue_ids=ResList)\n",
    "list(rmsd())\n",
    "\n",
    "# specifying limit: >=0, default excluded\n",
    "# Stop accepting results after this many hits. \n",
    "limit = StructMotifQuery(entry_id=\"2MNR\", limit=100, residue_ids=ResList)\n",
    "list(limit())\n",
    "\n",
    "# specifying atom pairing scheme, default = \"SIDE_CHAIN\"\n",
    "# ENUM: \"ALL\", \"BACKBONE\", \"SIDE_CHAIN\", \"PSUEDO_ATOMS\"\n",
    "# this is typechecked by a literal. \n",
    "# Which atoms to consider to compute RMSD scores and transformations. \n",
    "atom = StructMotifQuery(entry_id=\"2MNR\", atom_pairing_scheme=\"ALL\", residue_ids=ResList)\n",
    "list(atom())\n",
    "\n",
    "# specifying motif pruning strategy, default = \"KRUSKAL\"\n",
    "# ENUM: \"NONE\", \"KRUSKAL\"\n",
    "# this is typechecked by a literal in the package. \n",
    "# Specifies how many query motifs are \"pruned\". KRUSKAL leads to less stringent queries, and faster results.\n",
    "pruning = StructMotifQuery(entry_id=\"2MNR\", motif_pruning_strategy=\"NONE\", residue_ids=ResList)\n",
    "list(pruning())\n",
    "\n",
    "# specifying allowed structures, default excluded\n",
    "# specify the structures you wish to allow in the return result. As an example,\n",
    "# we could only allow the results from the limited query we ran earlier. \n",
    "allowed = StructMotifQuery(entry_id=\"2MNR\", allowed_structures=list(limit()), residue_ids=ResList)\n",
    "list(allowed())\n",
    "\n",
    "# specifying structures to exclude, default excluded\n",
    "# specify structures to exclude from a query. We could, for example,\n",
    "# exclude the results of the previous allowed query. \n",
    "excluded = StructMotifQuery(entry_id=\"2MNR\", excluded_structures=list(allowed()), residue_ids=ResList)\n",
    "list(excluded())"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cce26406",
   "metadata": {},
   "source": [
    "The Structure Motif Query can be used to make some very specific queries. Below is an example of a query that retrives occurances of the enolase superfamily, a group of proteins diverse in sequence and structure that are all capable of abstracting a proton from a carboxylic acid. Position-specific exchanges are crucial to represent this superfamily accurately."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5d9c9a2b",
   "metadata": {},
   "outputs": [],
   "source": [
    "Res1 = StructureMotifResidue(\"A\", \"1\", 162, [\"LYS\", \"HIS\"])\n",
    "Res2 = StructureMotifResidue(\"A\", \"1\", 193)\n",
    "Res3 = StructureMotifResidue(\"A\", \"1\", 219)\n",
    "Res4 = StructureMotifResidue(\"A\", \"1\", 245, [\"GLU\", \"ASP\", \"ASN\"])\n",
    "Res5 = StructureMotifResidue(\"A\", \"1\", 295, [\"HIS\", \"LYS\"])\n",
    "\n",
    "ResList = [Res1, Res2, Res3, Res4, Res5]\n",
    "\n",
    "query = StructMotifQuery(entry_id=\"2MNR\", residue_ids=ResList)\n",
    "\n",
    "\n",
    "print(len(list(query())))\n",
    "print(query.count())"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2c3fed55",
   "metadata": {},
   "source": [
    "## Faceted Queries\n",
    "\n",
    "In order to group and perform calculations and statistics on PDB data by using a simple search query, you can use a faceted query (or facets). Facets arrange search results into categories (buckets) based on the requested field values. More information on Faceted Queries can be found [here](https://search.rcsb.org/#using-facets). All facets should be provided with `name`, `aggregation_type`, and `attribute` values. Depending on the aggregation type, other parameters must also be specified. The `facets()` function runs the query `q` using the specified facet(s), and returns a list of dictionaries:\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "51ccfd5b",
   "metadata": {},
   "outputs": [],
   "source": [
    "from rcsbsearchapi.search import Facet\n",
    "\n",
    "q = AttributeQuery(\"rcsb_accession_info.initial_release_date\", operator=\"greater\", value=\"2019-08-20\")\n",
    "q.facets(facets=Facet(name=\"Methods\", aggregation_type=\"terms\", attribute=\"exptl.method\"))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9dd33bc2",
   "metadata": {},
   "source": [
    "#### Term Facets\n",
    "Terms faceting is a multi-bucket aggregation where buckets are dynamically built - one per unique value. We can specify the minimum count (`>= 0`) for a bucket to be returned using the parameter `min_interval_population` (default value `1`). We can also control the number of buckets returned (`<= 65336`) using the parameter `max_num_intervals` (default value `65336`)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ff6a9653",
   "metadata": {},
   "outputs": [],
   "source": [
    "# This is the default query, used by the RCSB Search API when no query is explicitly specified.\n",
    "# This default query will be used for most of the examples found below for faceted queries.\n",
    "base_q = AttributeQuery(\"rcsb_entry_info.structure_determination_methodology\", operator=\"exact_match\", value=\"experimental\") \n",
    "\n",
    "base_q.facets(facets=Facet(name=\"Journals\", aggregation_type=\"terms\", attribute=\"rcsb_primary_citation.rcsb_journal_abbrev\", min_interval_population=1000))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5b90bdcf",
   "metadata": {},
   "source": [
    "#### Histogram Facets\n",
    "Histogram facets build fixed-sized buckets (intervals) over numeric values. The size of the intervals must be specified in the parameter `interval`. We can also specify `min_interval_population` if desired."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1110f97d",
   "metadata": {},
   "outputs": [],
   "source": [
    "base_q.facets(return_type=\"polymer_entity\", facets=Facet(name=\"Formula Weight\", aggregation_type=\"histogram\", attribute=\"rcsb_polymer_entity.formula_weight\", interval=50, min_interval_population=1))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2ad3dbc5",
   "metadata": {},
   "source": [
    "#### Date Histogram Facets\n",
    "Similar to histogram facets, date histogram facetes build buckets over date values. For date histogram aggregations, we must specify `interval=\"year\"`. Again, we may also specify `min_interval_population`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e826a1f2",
   "metadata": {},
   "outputs": [],
   "source": [
    "base_q.facets(facets=Facet(name=\"Release Date\", aggregation_type=\"date_histogram\", attribute=\"rcsb_accession_info.initial_release_date\", interval=\"year\", min_interval_population=1))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2f562ab6",
   "metadata": {},
   "source": [
    "#### Range Facets\n",
    "We can define the buckets ourselves by using range facets. In order to specify the ranges, we use the `Range` class. Note that the range includes the `start` value and excludes the `end` value (`include_lower` and `include_upper` should not be specified). If the `start` or `end` is omitted, the minimum or maximum boundaries will be used by default. The buckets should be provided as a list of `Range` objects to the `ranges` parameter. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c56c0c89",
   "metadata": {},
   "outputs": [],
   "source": [
    "from rcsbsearchapi.search import Range\n",
    "\n",
    "base_q.facets(facets=Facet(name=\"Resolution Combined\", aggregation_type=\"range\", attribute=\"rcsb_entry_info.resolution_combined\", ranges=[Range(start=None,end=2), Range(start=2, end=2.2), Range(start=2.2, end=2.4), Range(start=4.6, end=None)]))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "212b7b6f",
   "metadata": {},
   "source": [
    "#### Date Range Facets\n",
    "Date range facets allow us to specify date values as bucket ranges, using [date math expressions](https://search.rcsb.org/#date-math-expressions)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5492f1c9",
   "metadata": {},
   "outputs": [],
   "source": [
    "base_q.facets(facets=Facet(name=\"Release Date\", aggregation_type=\"date_range\", attribute=\"rcsb_accession_info.initial_release_date\", ranges=[Range(start=None,end=\"2020-06-01||-12M\"), Range(start=\"2020-06-01\", end=\"2020-06-01||+12M\"), Range(start=\"2020-06-01||+12M\", end=None)]))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a6b64dec",
   "metadata": {},
   "source": [
    "#### Cardinality Facets \n",
    "Cardinality facets return a single value: the count of distinct values returned for a given field. A `precision_threshold` (`<= 40000`, default value `40000`) may be specified."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3c800a0e",
   "metadata": {},
   "outputs": [],
   "source": [
    "base_q.facets(facets=Facet(name=\"Organism Names Count\", aggregation_type=\"cardinality\", attribute=\"rcsb_entity_source_organism.ncbi_scientific_name\"))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "484c6c61",
   "metadata": {},
   "source": [
    "#### Multidimensional Facets\n",
    "Complex, multi-dimensional aggregations are possible by specifying additional facets in the `nested_facets` parameter, as in the example below:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8bff1718",
   "metadata": {},
   "outputs": [],
   "source": [
    "f1 = Facet(name=\"Polymer Entity Types\", aggregation_type=\"terms\", attribute=\"rcsb_entry_info.selected_polymer_entity_types\")\n",
    "f2 = Facet(name=\"Release Date\", aggregation_type=\"date_histogram\", attribute=\"rcsb_accession_info.initial_release_date\", interval=\"year\")\n",
    "base_q.facets(facets=Facet(name=\"Experimental Method\", aggregation_type=\"terms\", attribute=\"rcsb_entry_info.experimental_method\", nested_facets=[f1, f2]))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "346eb145",
   "metadata": {},
   "source": [
    "#### Filter Facets\n",
    "Filters allow us to filter documents that contribute to bucket count. Similar to queries, we can group several `TerminalFilter`s into a single `GroupFilter`. We can combine a filter with a facet using the `FilterFacet` class. Terminal filters should specify an `attribute` and `operator`, as well as possible a `value` and whether or not it should be a `negation` and/or `case_sensitive`. Group filters should specify a `logical_operator` (which should be either `\"and\"` or `\"or\"`) and a list of filters (`nodes`) that should be combined. Finally, the `FilterFacet` should be provided with a filter and a (list of) facet(s). Here are some examples:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7aced389",
   "metadata": {},
   "outputs": [],
   "source": [
    "from rcsbsearchapi.search import TerminalFilter, GroupFilter, FilterFacet\n",
    "tf1 = TerminalFilter(attribute=\"rcsb_polymer_instance_annotation.type\", operator=\"exact_match\", value=\"CATH\")\n",
    "tf2 = TerminalFilter(attribute=\"rcsb_polymer_instance_annotation.annotation_lineage.id\", operator=\"in\", value=[\"2.140.10.30\", \"2.120.10.80\"])\n",
    "ff2 = FilterFacet(filters=tf2, facets=Facet(\"CATH Domains\", \"terms\", \"rcsb_polymer_instance_annotation.annotation_lineage.id\", min_interval_population=1))\n",
    "ff1 = FilterFacet(filters=tf1, facets=ff2)\n",
    "base_q.facets(\"polymer_instance\", ff1)\n",
    "\n",
    "tf1 = TerminalFilter(attribute=\"rcsb_struct_symmetry.kind\", operator=\"exact_match\", value=\"Global Symmetry\", negation=False)\n",
    "f2 = Facet(name=\"ec_terms\", aggregation_type=\"terms\", attribute=\"rcsb_polymer_entity.rcsb_ec_lineage.id\")\n",
    "f1 = Facet(name=\"sym_symbol_terms\", aggregation_type=\"terms\", attribute=\"rcsb_struct_symmetry.symbol\", nested_facets=f2)\n",
    "ff = FilterFacet(filters=tf1, facets=f1)\n",
    "q1 = AttributeQuery(\"rcsb_assembly_info.polymer_entity_count\", operator=\"equals\", value=1)\n",
    "q2 = AttributeQuery(\"rcsb_assembly_info.polymer_entity_instance_count\", operator=\"greater\", value=1)\n",
    "q = q1 & q2\n",
    "q.facets(\"assembly\", ff)\n",
    "\n",
    "tf1 = TerminalFilter(attribute=\"rcsb_polymer_entity_group_membership.aggregation_method\", operator=\"exact_match\", value=\"sequence_identity\")\n",
    "tf2 = TerminalFilter(attribute=\"rcsb_polymer_entity_group_membership.similarity_cutoff\", operator=\"equals\", value=100)\n",
    "gf = GroupFilter(logical_operator=\"and\", nodes=[tf1, tf2])\n",
    "ff = FilterFacet(filters=gf, facets=Facet(\"Distinct Protein Sequence Count\", \"cardinality\", \"rcsb_polymer_entity_group_membership.group_id\"))\n",
    "base_q.facets(\"polymer_entity\", ff)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5c306677",
   "metadata": {},
   "source": [
    "For a more practical example, see the [Covid-19 notebook](covid.ipynb)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.9.7"
  },
  "toc": {
   "base_numbering": 1,
   "nav_menu": {},
   "number_sections": true,
   "sideBar": true,
   "skip_h1_title": false,
   "title_cell": "Table of Contents",
   "title_sidebar": "Contents",
   "toc_cell": false,
   "toc_position": {},
   "toc_section_display": true,
   "toc_window_display": false
  },
  "varInspector": {
   "cols": {
    "lenName": 16,
    "lenType": 16,
    "lenVar": 40
   },
   "kernels_config": {
    "python": {
     "delete_cmd_postfix": "",
     "delete_cmd_prefix": "del ",
     "library": "var_list.py",
     "varRefreshCmd": "print(var_dic_list())"
    },
    "r": {
     "delete_cmd_postfix": ") ",
     "delete_cmd_prefix": "rm(",
     "library": "var_list.r",
     "varRefreshCmd": "cat(var_dic_list()) "
    }
   },
   "types_to_exclude": [
    "module",
    "function",
    "builtin_function_or_method",
    "instance",
    "_Feature"
   ],
   "window_display": false
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}