Merge pull request #121 from biggraph/feat/bionemo-nims

lynxkite-app/web/package-lock.json

@@ -18,11 +18,13 @@
         "@syncedstore/react": "^0.6.0",
         "@types/node": "^22.10.1",
         "@xyflow/react": "^12.3.5",
+        "3dmol": "^2.4.2",
         "axios": "^1.8.2",
         "daisyui": "^4.12.20",
         "echarts": "^5.5.1",
         "fuse.js": "^7.0.0",
         "json-schema-to-typescript": "^15.0.3",
+        "monaco-editor": "^0.52.2",
         "react": "^18.3.1",
         "react-dom": "^18.3.1",
         "react-markdown": "^9.0.1",
@@ -2366,6 +2368,22 @@
         "d3-zoom": "^3.0.0"
       }
     },
+    "node_modules/3dmol": {
+      "version": "2.4.2",
+      "resolved": "https://registry.npmjs.org/3dmol/-/3dmol-2.4.2.tgz",
+      "integrity": "sha512-7R6uFaF5++s9EpO+R0a3bs6igi5idI+sgrVhtozaZlOGHg6oourYkAlAwYxWRaeBGlNcAd+oEIhbhJYhNxGyZA==",
+      "license": "BSD-3-Clause",
+      "dependencies": {
+        "iobuffer": "^5.3.1",
+        "netcdfjs": "^3.0.0",
+        "pako": "^2.1.0",
+        "upng-js": "^2.1.0"
+      },
+      "engines": {
+        "node": ">=16.16.0",
+        "npm": ">=8.11"
+      }
+    },
     "node_modules/abstract-leveldown": {
       "version": "6.2.3",
       "resolved": "https://registry.npmjs.org/abstract-leveldown/-/abstract-leveldown-6.2.3.tgz",
@@ -4187,6 +4205,12 @@
       "integrity": "sha512-0aO8FkhNZlj/ZIbNi7Lxxr12obT7cL1moPfE4tg1LkX7LlLfC6DeX4l2ZEud1ukP9jNQyNnfzQVqwbwmAATY4Q==",
       "license": "MIT"
     },
+    "node_modules/iobuffer": {
+      "version": "5.4.0",
+      "resolved": "https://registry.npmjs.org/iobuffer/-/iobuffer-5.4.0.tgz",
+      "integrity": "sha512-DRebOWuqDvxunfkNJAlc3IzWIPD5xVxwUNbHr7xKB8E6aLJxIPfNX3CoMJghcFjpv6RWQsrcJbghtEwSPoJqMA==",
+      "license": "MIT"
+    },
     "node_modules/is-alphabetical": {
       "version": "2.0.1",
       "resolved": "https://registry.npmjs.org/is-alphabetical/-/is-alphabetical-2.0.1.tgz",
@@ -5465,8 +5489,7 @@
       "version": "0.52.2",
       "resolved": "https://registry.npmjs.org/monaco-editor/-/monaco-editor-0.52.2.tgz",
       "integrity": "sha512-GEQWEZmfkOGLdd3XK8ryrfWz3AIP8YymVXiPHEdewrUq7mh0qrKrfHLNCXcbB6sTnMLnOZ3ztSiKcciFUkIJwQ==",
-      "license": "MIT",
-      "peer": true
+      "license": "MIT"
     },
     "node_modules/ms": {
       "version": "2.1.3",
@@ -5518,6 +5541,15 @@
       "dev": true,
       "license": "MIT"
     },
+    "node_modules/netcdfjs": {
+      "version": "3.0.0",
+      "resolved": "https://registry.npmjs.org/netcdfjs/-/netcdfjs-3.0.0.tgz",
+      "integrity": "sha512-LOvT8KkC308qtpUkcBPiCMBtii7ZQCN6LxcVheWgyUeZ6DQWcpSRFV9dcVXLj/2eHZ/bre9tV5HTH4Sf93vrFw==",
+      "license": "MIT",
+      "dependencies": {
+        "iobuffer": "^5.3.2"
+      }
+    },
     "node_modules/no-case": {
       "version": "3.0.4",
       "resolved": "https://registry.npmjs.org/no-case/-/no-case-3.0.4.tgz",
@@ -5649,6 +5681,12 @@
       "integrity": "sha512-g4+387DXDKlZzHkP+9FLt8yKj8+/3tOkPv7DVTJGGRm00RkEWgqbFstX1mXJ4M0VDYhUqsTOiISqNOJnhAu3PQ==",
       "license": "MIT"
     },
+    "node_modules/pako": {
+      "version": "2.1.0",
+      "resolved": "https://registry.npmjs.org/pako/-/pako-2.1.0.tgz",
+      "integrity": "sha512-w+eufiZ1WuJYgPXbV/PO3NCMEc3xqylkKHzp8bxp1uW4qaSNQUkwmLLEc3kKsfz8lpV1F8Ht3U1Cm+9Srog2ug==",
+      "license": "(MIT AND Zlib)"
+    },
     "node_modules/parent-module": {
       "version": "1.0.1",
       "resolved": "https://registry.npmjs.org/parent-module/-/parent-module-1.0.1.tgz",
@@ -7107,6 +7145,21 @@
         "browserslist": ">= 4.21.0"
       }
     },
+    "node_modules/upng-js": {
+      "version": "2.1.0",
+      "resolved": "https://registry.npmjs.org/upng-js/-/upng-js-2.1.0.tgz",
+      "integrity": "sha512-d3xzZzpMP64YkjP5pr8gNyvBt7dLk/uGI67EctzDuVp4lCZyVMo0aJO6l/VDlgbInJYDY6cnClLoBp29eKWI6g==",
+      "license": "MIT",
+      "dependencies": {
+        "pako": "^1.0.5"
+      }
+    },
+    "node_modules/upng-js/node_modules/pako": {
+      "version": "1.0.11",
+      "resolved": "https://registry.npmjs.org/pako/-/pako-1.0.11.tgz",
+      "integrity": "sha512-4hLB8Py4zZce5s4yd9XzopqwVv/yGNhV1Bl8NTmCq1763HeK2+EwVTv+leGeL13Dnh2wfbqowVPXCIO0z4taYw==",
+      "license": "(MIT AND Zlib)"
+    },
     "node_modules/uri-js": {
       "version": "4.4.1",
       "resolved": "https://registry.npmjs.org/uri-js/-/uri-js-4.4.1.tgz",

lynxkite-app/web/package.json

@@ -21,11 +21,13 @@
     "@syncedstore/react": "^0.6.0",
     "@types/node": "^22.10.1",
     "@xyflow/react": "^12.3.5",
+    "3dmol": "^2.4.2",
     "axios": "^1.8.2",
     "daisyui": "^4.12.20",
     "echarts": "^5.5.1",
     "fuse.js": "^7.0.0",
     "json-schema-to-typescript": "^15.0.3",
+    "monaco-editor": "^0.52.2",
     "react": "^18.3.1",
     "react-dom": "^18.3.1",
     "react-markdown": "^9.0.1",

lynxkite-app/web/src/workspace/Workspace.tsx

@@ -37,6 +37,7 @@ import { LynxKiteState } from "./LynxKiteState";
 import NodeSearch, { type OpsOp, type Catalog, type Catalogs } from "./NodeSearch.tsx";
 import NodeWithGraphCreationView from "./nodes/GraphCreationNode.tsx";
 import NodeWithImage from "./nodes/NodeWithImage.tsx";
+import NodeWithMolecule from "./nodes/NodeWithMolecule.tsx";
 import NodeWithParams from "./nodes/NodeWithParams";
 import NodeWithTableView from "./nodes/NodeWithTableView.tsx";
 import NodeWithVisualization from "./nodes/NodeWithVisualization.tsx";
@@ -175,6 +176,7 @@ function LynxKiteFlow() {
       image: NodeWithImage,
       table_view: NodeWithTableView,
       graph_creation_view: NodeWithGraphCreationView,
+      molecule: NodeWithMolecule,
     }),
     [],
   );

lynxkite-app/web/src/workspace/nodes/NodeWithMolecule.tsx

@@ -0,0 +1,73 @@
+import React, { useEffect, type CSSProperties } from "react";
+import NodeWithParams from "./NodeWithParams";
+
+const NodeWithMolecule = (props: any) => {
+  const containerRef = React.useRef<HTMLDivElement>(null);
+  const viewerRef = React.useRef<any>(null);
+
+  useEffect(() => {
+    const config = props.data?.display?.value;
+    if (!config || !containerRef.current) return;
+    async function run() {
+      const $3Dmol = await import("3dmol");
+
+      try {
+        // Initialize viewer only once
+        if (!viewerRef.current) {
+          viewerRef.current = $3Dmol.createViewer(containerRef.current, {
+            backgroundColor: "white",
+          });
+        }
+
+        const viewer = viewerRef.current;
+
+        // Clear previous models
+        viewer.clear();
+
+        // Add new model and style it
+        viewer.addModel(config.data, config.format);
+        viewer.setStyle({}, { stick: {} });
+        viewer.zoomTo();
+        viewer.render();
+      } catch (error) {
+        console.error("Error rendering 3D molecule:", error);
+      }
+    }
+    run();
+    const resizeObserver = new ResizeObserver(() => {
+      viewerRef.current?.resize();
+    });
+
+    const observed = containerRef.current;
+    resizeObserver.observe(observed);
+    return () => {
+      resizeObserver.unobserve(observed);
+      if (viewerRef.current) {
+        viewerRef.current.clear();
+      }
+    };
+  }, [props.data?.display?.value]);
+
+  const nodeStyle: CSSProperties = {
+    display: "flex",
+    flexDirection: "column",
+    height: "100%",
+  };
+
+  const vizStyle: CSSProperties = {
+    flex: 1,
+    minHeight: "300px",
+    border: "1px solid #ddd",
+    borderRadius: "4px",
+    overflow: "hidden",
+    position: "relative",
+  };
+
+  return (
+    <NodeWithParams nodeStyle={nodeStyle} collapsed {...props}>
+      <div style={vizStyle} ref={containerRef} />
+    </NodeWithParams>
+  );
+};
+
+export default NodeWithMolecule;

lynxkite-bio/src/lynxkite_bio/nims.py

@@ -1,107 +1,143 @@
 """Wrappers for BioNeMo NIMs."""
 
-from enum import Enum
 from lynxkite_graph_analytics import Bundle
 from lynxkite.core import ops
 import joblib
+import httpx
+import pandas as pd
 import os
 
-NIM_URLS = os.environ.get("NIM_URLS", "http://localhost:8000").split(",")
 
 mem = joblib.Memory(".joblib-cache")
 ENV = "LynxKite Graph Analytics"
 op = ops.op_registration(ENV)
 
+key = os.getenv("NVCF_RUN_KEY")
 
-class MSASearchTypes(Enum):
-    ALPHAFOLD2 = "ALPHAFOLD2"
-    ESM2 = "ESM2"
 
-
-class AlignmentFormats(Enum):
-    FASTA = "fasta"
-    A3M = "a3m"
-    STOCKHOLM = "stockholm"
-    CLUSTAL = "clustal"
-    PDB = "pdb"
-    PIR = "pir"
-    MSF = "msf"
-    TSV = "tsv"
+async def query_bionemo_nim(
+    url: str,
+    payload: dict,
+):
+    headers = {
+        "Authorization": f"Bearer {key}",
+        "NVCF-POLL-SECONDS": "500",
+        "Content-Type": "application/json",
+    }
+    try:
+        print(f"Sending request to {url}")
+        async with httpx.AsyncClient(timeout=600) as client:
+            response = await client.post(url, json=payload, headers=headers)
+        print(f"Received response from {url}", response.status_code)
+        response.raise_for_status()
+        return response.json()
+    except httpx.RequestError as e:
+        raise ValueError(f"Query failed: {e}")
 
 
 @op("MSA-search")
 @mem.cache
-def msa_search(
+async def msa_search(
     bundle: Bundle,
     *,
     protein_table: str,
     protein_column: str,
     e_value: float = 0.0001,
     iterations: int = 1,
-    search_type: MSASearchTypes = MSASearchTypes.ALPHAFOLD2,
-    output_alignment_formats: list[AlignmentFormats] = [
-        AlignmentFormats.FASTA,
-        AlignmentFormats.A3M,
-    ],
-    databases: str = '["Uniref30_2302", "colabfold_envdb_202108", "PDB70_220313"]',
+    search_type: str = "alphafold2",
+    output_alignment_formats: str = "a3m",
+    databases: str = "Uniref30_2302,colabfold_envdb_202108",
 ):
     bundle = bundle.copy()
+    bundle.dfs[protein_table]["alignments"] = None
+
+    formats = [format.strip() for format in output_alignment_formats.split(",")]
+    dbs = [db.strip() for db in databases.split(",")]
+
+    for idx, protein_sequence in enumerate(bundle.dfs[protein_table][protein_column]):
+        print(f"Processing protein {idx + 1}/{len(bundle.dfs[protein_table])}")
+
+        response = await query_bionemo_nim(
+            url="https://health.api.nvidia.com/v1/biology/colabfold/msa-search/predict",
+            payload={
+                "sequence": protein_sequence,
+                "e_value": e_value,
+                "iterations": iterations,
+                "search_type": search_type,
+                "output_alignment_formats": formats,
+                "databases": dbs,
+            },
+        )
+
+        bundle.dfs[protein_table].at[idx, "alignments"] = response["alignments"]
+
     return bundle
 
 
 @op("Query OpenFold2")
 @mem.cache
-def query_openfold2(
+async def query_openfold2(
     bundle: Bundle,
     *,
     protein_table: str,
     protein_column: str,
     alignment_table: str,
     alignment_column: str,
-    use_templates: bool = False,
-    relaxed_prediction: bool = False,
-    databases: str = '["Uniref30_2302", "colabfold_envdb_202108", "PDB70_220313"]',
+    selected_models: str = "1,2",
+    relax_prediction: bool = False,
 ):
     bundle = bundle.copy()
+
+    bundle.dfs[protein_table]["folded_protein"] = None
+    selected_models_list = [int(model) for model in selected_models.split(",")]
+
+    for idx in range(len(bundle.dfs[protein_table])):
+        print(f"Processing protein {idx + 1}/{len(bundle.dfs[protein_table])}")
+
+        protein = bundle.dfs[protein_table][protein_column].iloc[idx]
+        alignments = bundle.dfs[alignment_table][alignment_column].iloc[idx]
+
+        response = await query_bionemo_nim(
+            url="https://health.api.nvidia.com/v1/biology/openfold/openfold2/predict-structure-from-msa-and-template",
+            payload={
+                "sequence": protein,
+                "alignments": alignments,
+                "selected_models": selected_models_list,
+                "relax_prediction": relax_prediction,
+            },
+        )
+
+        folded_protein = response["structures_in_ranked_order"].pop(0)["structure"]
+        bundle.dfs[protein_table].at[idx, "folded_protein"] = folded_protein
+
+    bundle.dfs["openfold"] = pd.DataFrame()
+
     return bundle
 
 
-@op("View molecules", view="visualization")
-def view_molecules(
+@op("View molecule", view="molecule")
+def view_molecule(
     bundle: Bundle,
     *,
     molecule_table: str,
     molecule_column: str,
-    color="spectrum",
+    row_index: int = 0,
 ):
+    molecule_data = bundle.dfs[molecule_table][molecule_column].iloc[row_index]
+
     return {
-        "series": [
-            {
-                "type": "pie",
-                "radius": ["40%", "70%"],
-                "itemStyle": {
-                    "borderRadius": 10,
-                    "borderColor": "#fff",
-                    "borderWidth": 2,
-                },
-                "data": [
-                    {"value": 2, "name": "Hydrogen"},
-                    {"value": 1, "name": "Sulfur"},
-                    {"value": 4, "name": "Oxygen"},
-                ],
-            }
-        ]
+        "data": molecule_data,
+        "format": "pdb"
+        if molecule_data.startswith("ATOM")
+        else "sdf"
+        if molecule_data.startswith("CTfile")
+        else "smiles",
     }
 
 
-@op("Known drug")
-def known_drug(*, drug_name: str):
-    return Bundle()
-
-
 @op("Query GenMol")
 @mem.cache
-def query_genmol(
+async def query_genmol(
     bundle: Bundle,
     *,
     molecule_table: str,
@@ -113,12 +149,26 @@ def query_genmol(
     scoring: str = "QED",
 ):
     bundle = bundle.copy()
+
+    response = await query_bionemo_nim(
+        url="https://health.api.nvidia.com/v1/biology/nvidia/genmol/generate",
+        payload={
+            "smiles": bundle.dfs[molecule_table][molecule_column].iloc[0],
+            "num_molecules": num_molecules,
+            "temperature": temperature,
+            "noise": noise,
+            "step_size": step_size,
+            "scoring": scoring,
+        },
+    )
+    generated_ligands = "\n".join([v["smiles"] for v in response["molecules"]])
+    bundle.dfs[molecule_table]["ligands"] = generated_ligands
     return bundle
 
 
 @op("Query DiffDock")
 @mem.cache
-def query_diffdock(
+async def query_diffdock(
     proteins: Bundle,
     ligands: Bundle,
     *,
@@ -126,8 +176,29 @@ def query_diffdock(
     protein_column: str,
     ligand_table: str,
     ligand_column: str,
+    ligand_file_type: str = "txt",
     num_poses=10,
     time_divisions=20,
     num_steps=18,
 ):
-    return proteins
+    response = await query_bionemo_nim(
+        url="https://health.api.nvidia.com/v1/biology/mit/diffdock",
+        payload={
+            "protein": proteins.dfs[protein_table][protein_column].iloc[0],
+            "ligand": ligands.dfs[ligand_table][ligand_column].iloc[0],
+            "ligand_file_type": ligand_file_type,
+            "num_poses": num_poses,
+            "time_divisions": time_divisions,
+            "num_steps": num_steps,
+        },
+    )
+    bundle = Bundle()
+    bundle.dfs["diffdock_table"] = pd.DataFrame()
+    bundle.dfs["diffdock_table"]["protein"] = [response["protein"]] * len(response["status"])
+    bundle.dfs["diffdock_table"]["ligand"] = [response["ligand"]] * len(response["status"])
+    bundle.dfs["diffdock_table"]["trajectory"] = response["trajectory"]
+    bundle.dfs["diffdock_table"]["ligand_positions"] = response["ligand_positions"]
+    bundle.dfs["diffdock_table"]["position_confidence"] = response["position_confidence"]
+    bundle.dfs["diffdock_table"]["status"] = response["status"]
+
+    return bundle

lynxkite-core/src/lynxkite/core/ops.py

@@ -180,6 +180,7 @@ class Op(BaseConfig):
                 "table_view",
                 "graph_creation_view",
                 "image",
+                "molecule",
             ]:
                 # If the operation is some kind of visualization, we use the output as the
                 # value to display by default.