ZMPY3D_demo_zm.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# This notebook demonstrates how to calculate 3D Zernike moments using ZMPY3D with TensorFlow.\n",
    "#\n",
    "# This notebook primarily consists of the following steps: \n",
    "#     1. Install ZMPY3D_TF.\n",
    "#     2. Define necessary parameters.\n",
    "#     3. Load precalculated cache.\n",
    "#     4. Download example PDB data with coordinates.\n",
    "#     5. Convert coordinate data into a voxel.\n",
    "#     6. Create a callable function for generating Zernike moments and normalization.\n",
    "#     7. Obtain the results.\n",
    "#     8. A command line interface (CLI) example"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Install ZMPY3D versions for TensorFlow.\n",
    "! pip install ZMPY3D-TF\n",
    "print(f\"It is recommended to restart the Python kernel for the IPython notebook.\")\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Download example data from GitHub using curl\n",
    "! curl -OJL https://github.com/tawssie/ZMPY3D/raw/main/1WAC_A.txt\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "import ZMPY3D_TF as z\n",
    "import tensorflow as tf\n",
    "import pickle\n",
    "\n",
    "MaxOrder = 20 # Set the maximum order to 20\n",
    "GridWidth= 1.00 # Set the grid width to 1.00 (equivalent to 1 angstrom)\n",
    "Param=z.get_global_parameter() # Retrieve the necessary parameters\n",
    "\n",
    "# Find the cache_data directory based on the site package location of ZMPY3D.\n",
    "LogCacheFilePath=z.__file__.replace('__init__.py', 'cache_data') + '/LogG_CLMCache_MaxOrder{:02d}.pkl'.format(MaxOrder)\n",
    "\n",
    "with open(LogCacheFilePath, 'rb') as file:\n",
    "    CachePKL = pickle.load(file)\n",
    "\n",
    "# Convert the precalculated cache data into TensorFlow tensor objects\n",
    "GCache_pqr_linear= tf.convert_to_tensor(CachePKL['GCache_pqr_linear'])\n",
    "GCache_complex= tf.convert_to_tensor(CachePKL['GCache_complex'])\n",
    "GCache_complex_index= tf.convert_to_tensor(CachePKL['GCache_complex_index'])\n",
    "CLMCache3D= tf.convert_to_tensor(CachePKL['CLMCache3D'],dtype=tf.complex128)\n",
    "CLMCache= tf.convert_to_tensor(CachePKL['CLMCache'], dtype=tf.float64)\n",
    "\n",
    "\n",
    "print(f\"Now using the MaxOrder of {MaxOrder} and the GridWidth of {GridWidth}.\")\n",
    "print(f\"Pre-calculated parameters have been loaded successfully.\")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%time\n",
    "PDBFileName='./1WAC_A.txt'\n",
    "\n",
    "\n",
    "# Convert structure data into coordinates\n",
    "[XYZ,AA_NameList]=z.get_pdb_xyz_ca(PDBFileName)\n",
    "\n",
    "# Convert coordinates into voxels using precalculated Gaussian densities\n",
    "ResidueBox=z.get_residue_gaussian_density_cache(Param)\n",
    "[Voxel3D,Corner]=z.fill_voxel_by_weight_density(XYZ,AA_NameList,Param['residue_weight_map'],GridWidth,ResidueBox[GridWidth])\n",
    "\n",
    "# Convert the voxel data into a tensor object\n",
    "Voxel3D=tf.convert_to_tensor(Voxel3D,dtype=tf.float64)\n",
    "\n",
    "print(f\"Converting PDB to 3D voxel grid with NumPy on CPU, then transferring to GPU memory as Tensorflow objects.\")\n",
    "print(f\"Time elapsed is as follows:\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%time\n",
    "\n",
    "Dimension_BBox_scaled=tf.shape(Voxel3D)\n",
    "\n",
    "MaxOrder=tf.convert_to_tensor(MaxOrder,dtype=tf.int64)\n",
    "\n",
    "X_sample = tf.range(Dimension_BBox_scaled[0] + 1,dtype=tf.float64)\n",
    "Y_sample = tf.range(Dimension_BBox_scaled[1] + 1,dtype=tf.float64)\n",
    "Z_sample = tf.range(Dimension_BBox_scaled[2] + 1,dtype=tf.float64)\n",
    "\n",
    "# Calculate the volume mass and the center of mass\n",
    "[VolumeMass,Center,_]=z.calculate_bbox_moment(Voxel3D,1,X_sample,Y_sample,Z_sample)\n",
    "\n",
    "# Calculate the weights for sphere sampling\n",
    "[AverageVoxelDist2Center,_]=z.calculate_molecular_radius(Voxel3D,Center,VolumeMass,tf.convert_to_tensor(Param['default_radius_multiplier'], dtype=tf.float64))\n",
    "Sphere_X_sample, Sphere_Y_sample, Sphere_Z_sample=z.get_bbox_moment_xyz_sample(Center,AverageVoxelDist2Center,Dimension_BBox_scaled)\n",
    "\n",
    "# Apply weights to the geometric moments\n",
    "_,_,SphereBBoxMoment=z.calculate_bbox_moment(Voxel3D\n",
    "                                  ,MaxOrder\n",
    "                                  ,Sphere_X_sample\n",
    "                                  ,Sphere_Y_sample\n",
    "                                  ,Sphere_Z_sample)\n",
    "\n",
    "# Convert to scaled 3D Zernike moments\n",
    "ZMoment_scaled,_=z.calculate_bbox_moment_2_zm(MaxOrder\n",
    "                                   , GCache_complex\n",
    "                                   , GCache_pqr_linear\n",
    "                                   , GCache_complex_index\n",
    "                                   , CLMCache3D\n",
    "                                   , SphereBBoxMoment)\n",
    "\n",
    "# Convert the scaled 3D Zernike moments into 3DZD-based descriptors\n",
    "ZM_3DZD_invariant=z.get_3dzd_121_descriptor(ZMoment_scaled)\n",
    "ZM_3DZD_invariant_121=tf.reshape(tf.boolean_mask(ZM_3DZD_invariant, ~tf.math.is_nan(ZM_3DZD_invariant)), [-1])\n",
    "\n",
    "\n",
    "print(f\"The dimensions of the voxel being used are {Voxel3D.shape}.\")\n",
    "print(f\"Transforming the 3D voxel into Zernike moments with global normalization, 3DZD style, yields 121 descriptors.\")\n",
    "print(f\"Time elapsed is as follows:\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "executionInfo": {
     "elapsed": 26905,
     "status": "ok",
     "timestamp": 1711572606231,
     "user": {
      "displayName": "Sean Lai",
      "userId": "14374858823610391148"
     },
     "user_tz": 420
    },
    "id": "Fl7FrGrptnLd",
    "outputId": "344edd85-0107-4785-de89-c640fa06184b",
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "@tf.function\n",
    "def OneTimeConversion_TF(Voxel3D,MaxOrder):\n",
    "    Dimension_BBox_scaled=tf.shape(Voxel3D)\n",
    "    MaxOrder=tf.convert_to_tensor(MaxOrder,dtype=tf.int64)\n",
    "    \n",
    "    X_sample = tf.range(Dimension_BBox_scaled[0] + 1,dtype=tf.float64)\n",
    "    Y_sample = tf.range(Dimension_BBox_scaled[1] + 1,dtype=tf.float64)\n",
    "    Z_sample = tf.range(Dimension_BBox_scaled[2] + 1,dtype=tf.float64)\n",
    "    \n",
    "    [VolumeMass,Center,_]=z.calculate_bbox_moment(Voxel3D,1,X_sample,Y_sample,Z_sample)\n",
    "    [AverageVoxelDist2Center,_]=z.calculate_molecular_radius(Voxel3D,Center,VolumeMass,tf.convert_to_tensor(Param['default_radius_multiplier'], dtype=tf.float64))\n",
    "    Sphere_X_sample, Sphere_Y_sample, Sphere_Z_sample=z.get_bbox_moment_xyz_sample(Center,AverageVoxelDist2Center,Dimension_BBox_scaled)\n",
    "    \n",
    "    _,_,SphereBBoxMoment=z.calculate_bbox_moment(Voxel3D\n",
    "                                      ,MaxOrder\n",
    "                                      ,Sphere_X_sample\n",
    "                                      ,Sphere_Y_sample\n",
    "                                      ,Sphere_Z_sample)\n",
    "    \n",
    "    ZMoment_scaled,_=z.calculate_bbox_moment_2_zm(MaxOrder\n",
    "                                       , GCache_complex\n",
    "                                       , GCache_pqr_linear\n",
    "                                       , GCache_complex_index\n",
    "                                       , CLMCache3D\n",
    "                                       , SphereBBoxMoment)\n",
    "    \n",
    "    ZM_3DZD_invariant=z.get_3dzd_121_descriptor(ZMoment_scaled)\n",
    "    \n",
    "    ZM_3DZD_invariant_121=tf.reshape(tf.boolean_mask(ZM_3DZD_invariant, ~tf.math.is_nan(ZM_3DZD_invariant)), [-1])\n",
    "    return ZM_3DZD_invariant_121\n",
    "\n",
    "print(f\"Merge all steps into a single callable Tensorfow function, OneTimeConversion_TF, decorated with @tf.function for optimized graph compilation.\")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "executionInfo": {
     "elapsed": 977,
     "status": "ok",
     "timestamp": 1711572616966,
     "user": {
      "displayName": "Sean Lai",
      "userId": "14374858823610391148"
     },
     "user_tz": 420
    },
    "id": "sZXn24b_OBcv",
    "outputId": "85c614fa-a7c7-46ec-fc93-eb740091c212"
   },
   "outputs": [],
   "source": [
    "%%time\n",
    "# Use OneTimeConversion_TF with @tf.function to calculate and normalize 3D Zernike moments\n",
    "print(OneTimeConversion_TF(Voxel3D,MaxOrder))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "executionInfo": {
     "elapsed": 977,
     "status": "ok",
     "timestamp": 1711572616966,
     "user": {
      "displayName": "Sean Lai",
      "userId": "14374858823610391148"
     },
     "user_tz": 420
    },
    "id": "sZXn24b_OBcv",
    "outputId": "85c614fa-a7c7-46ec-fc93-eb740091c212"
   },
   "outputs": [],
   "source": [
    "\n",
    "# Alternatively, use a system call to compute results via CLI\n",
    "# ./ZMPY3D_TF_CLI_ZM PDBFile GridWidth MaximumOrder NormOrder Mode\n",
    "! ZMPY3D_TF_CLI_ZM \"./1WAC_A.txt\" 1.0 20 2 1\n",
    "\n"
   ]
  }
 ],
 "metadata": {
  "accelerator": "GPU",
  "colab": {
   "authorship_tag": "ABX9TyMENoVGMAAZ51Qh0Uw6LnV6",
   "collapsed_sections": [
    "3Hn12aJzpKR_",
    "DyMa0OG-pe8T",
    "9DbKZZY_pxLk",
    "QZKvszjqp58S",
    "Qp5hA6bBqHqO",
    "UFVne5plq0jK",
    "fVasOaTdrEy2"
   ],
   "gpuType": "T4",
   "provenance": [
    {
     "file_id": "1KRMJipBkbYroHJ6SeXHaPaX5-4WBBy6Z",
     "timestamp": 1711574522738
    },
    {
     "file_id": "1Gy7cTTaL2NPdz7T6Tk4jZDSVM5TwT_Bf",
     "timestamp": 1711323812481
    }
   ]
  },
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "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.16"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}