Merge pull request PolusAI#248 from jfennick/move_mm_viewer

move molecular modeling viewer

Author: jfennick

README.md

@@ -4,3 +4,17 @@ This repository contains workflows for various molecular modeling tasks. The wor
 
 ## Quick Start
 Follow the [installation instructions](https://github.com/PolusAI/workflow-inference-compiler#quick-start) here.
+
+## Jupyter notebook visualization
+
+You can view 3D molecular structures in the Jupyter notebook `src/vis/viewer.ipynb`. The visualization currently needs to be in its own conda environment.
+
+```
+install/install_conda.sh
+conda create --name vis
+conda activate vis
+install/install_nglview.sh
+pip install -e ".[all]"
+```
+
+![Plots](docs/tree_viewer.png)

docs/dev/api.rst

@@ -6,4 +6,12 @@ mm_workflows.schemas.biobb
 
 mm_workflows.schemas.gromacs_mdp
 ------------------------------------
-.. automodule:: mm_workflows.schemas.gromacs_mdp
+.. automodule:: mm_workflows.schemas.gromacs_mdp
+
+vis.filewatcher
+------------------------------------
+.. automodule:: vis.filewatcher
+
+vis.viewer
+------------------------------------
+.. automodule:: vis.viewer

install/install_nglview.sh

@@ -0,0 +1,24 @@
+#!/bin/bash -e
+# NOTE: mamba is a drop-in replacement for conda, just much faster.
+# (i.e. You can replace mamba with conda below.)
+# See https://github.com/conda-forge/miniforge#mambaforge-pypy3
+CONDA=conda
+if [ "$(which mamba)" ]; then
+    CONDA=mamba
+fi
+
+# mdtraj needs binary build dependencies (specifically cython) so do not install via pip
+$CONDA install -y -c conda-forge mdtraj
+
+# mdtraj and pypy CANNOT currently be installed into the same environment:
+# $CONDA install -y -c conda-forge mdtraj pypy
+
+$CONDA install -y pip
+
+pip install jupyterlab ipytree
+
+# See release notes about jupyterlab
+# https://github.com/nglviewer/nglview/releases/tag/v3.0.4
+pip install nglview
+
+jupyter-nbextension enable nglview --py --sys-prefix

pyproject.toml

@@ -147,7 +147,7 @@ ignore_errors = false
 # versioneer to the exclude=regex above, but this works.
 
 [[tool.mypy.overrides]]
-module = ["BioSimSpace.*", "MDAnalysis.*", "pymol.*", "rdkit.*", "nmrformd.*", "setuptools.*", "parmed.*", "pdbfixer.*", "openmm.*", "mdtraj.*"]
+module = ["BioSimSpace.*", "MDAnalysis.*", "pymol.*", "rdkit.*", "nmrformd.*", "setuptools.*", "parmed.*", "pdbfixer.*", "openmm.*", "mdtraj.*", "nglview.*"]
 ignore_missing_imports = true
 
 

src/vis/__init__.py

@@ -0,0 +1,42 @@
+import glob
+import os
+from pathlib import Path
+from typing import Dict, List
+
+
+# NOTE: You should be very careful when using file watchers! Most libraries
+# (watchdog, watchfiles, etc) will use operating system / platform-specific
+# APIs to check for changes (for performance reasons). However, this can cause
+# problems in some cases, specifically for network drives.
+# See https://stackoverflow.com/questions/45441623/using-watchdog-of-python-to-monitoring-afp-shared-folder-from-linux
+# I'm 99% sure the same problem happens with Docker containers. Either way, the
+# solution is to use polling. However, for unknown reasons, simply replacing
+# Observer with PollingObserver doesn't seem to work! So we are forced to write
+# our own basic file watcher using glob.
+
+
+def file_watcher_glob(cachedir_path: Path, patterns: List[str], prev_files: Dict[str, float]) -> Dict[str, float]:
+    """Determines whether files (specified by the given glob patterns) have been either recently created or modified.\n
+    Note that this is a workaround due to an issue with using standard file-watching libraries.
+
+    Args:
+        cachedir_path (Path): The --cachedir directory of the main workflow.
+        patterns (List[str]): The glob patterns which specifies the files to be watched.
+        prev_files (Dict[str, float]): This should be the return value from the previous function call.
+
+    Returns:
+        Dict[str, float]: A dictionary containing the filepaths and last modification times.
+    """
+    changed_files = {}
+    file_patterns = [str(cachedir_path / f'**/{pattern}') for pattern in patterns]
+    file_paths_all = [glob.glob(fp, recursive=True) for fp in file_patterns]
+    file_paths_flat = [path for fps in file_paths_all for path in fps]
+    for file in file_paths_flat:
+        mtime = os.path.getmtime(file)
+        if file not in prev_files:
+            # created
+            changed_files[file] = mtime
+        elif mtime > prev_files[file]:
+            # modified
+            changed_files[file] = mtime
+    return changed_files

src/vis/filewatcher.py

@@ -0,0 +1,55 @@
+{
+ "cells": [
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from vis import viewer\n",
+    "# Once the workflow has finished executing, you can view the output files.\n",
+    "viewer.tree_viewer()"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from vis import viewer\n",
+    "# Alternatively, you can view the files in realtime while the workflow is executing.\n",
+    "# This runs for about 10 mins before printing 'done'.\n",
+    "# You can always re-run the cell and/or increase num_iterations.\n",
+    "num_iterations = 60 * 10\n",
+    "viewer.realtime_viewer(num_iterations)"
+   ]
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "Python 3.10.6",
+   "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.10.6"
+  },
+  "vscode": {
+   "interpreter": {
+    "hash": "612ab4556636a7986aff5653e5d5ed5cb1e0c640daa830f0dae50b030d37dcd4"
+   }
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 2
+}