Merge pull request #516 from OpenBioSim/feature_somd2_preparefep

Write both SOMD1 and SOMD2 input in prepareFEP

Author: lohedges

nodes/playground/prepareFEP.ipynb

@@ -159,13 +159,6 @@
     "    BSS.Gateway.String(\n",
     "        help=\"The root name for the files describing the perturbation input1->input2.\"\n",
     "    ),\n",
-    ")\n",
-    "node.addInput(\n",
-    "    \"somd2\",\n",
-    "    BSS.Gateway.Boolean(\n",
-    "        help=\"Whether to generate input for SOMD2.\",\n",
-    "        default=False,\n",
-    "    ),\n",
     ")"
    ]
   },
@@ -198,11 +191,9 @@
    "source": [
     "do_mapping = True\n",
     "custom_mapping = node.getInput(\"mapping\")\n",
-    "# print (custom_mapping)\n",
     "if custom_mapping is not None:\n",
     "    do_mapping = False\n",
-    "    mapping = loadMapping(custom_mapping)\n",
-    "    # print (mapping)"
+    "    mapping = loadMapping(custom_mapping)"
    ]
   },
   {
@@ -218,8 +209,7 @@
     "    entries = prematchstring.split(\",\")\n",
     "    for entry in entries:\n",
     "        idxA, idxB = entry.split(\"-\")\n",
-    "        prematch[int(idxA)] = int(idxB)\n",
-    "# print (prematch)"
+    "        prematch[int(idxA)] = int(idxB)"
    ]
   },
   {
@@ -270,21 +260,9 @@
     "        scoring_function=\"RMSDalign\",\n",
     "        timeout=node.getInput(\"timeout\"),\n",
     "    )\n",
+    "    \n",
     "    # We retain the top mapping\n",
-    "    mapping = mappings[0]\n",
-    "    # print (len(mappings))\n",
-    "    # print (mappings)"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "# print (mapping)\n",
-    "# for x in range(0,len(mappings)):\n",
-    "#    print (mappings[x], scores[x])"
+    "    mapping = mappings[0]"
    ]
   },
   {
@@ -293,8 +271,7 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "inverted_mapping = dict([[v, k] for k, v in mapping.items()])\n",
-    "# print (inverted_mapping)"
+    "inverted_mapping = dict([[v, k] for k, v in mapping.items()])"
    ]
   },
   {
@@ -307,6 +284,7 @@
     "# on a root mean squared displacement fit to find the optimal translation vector\n",
     "# (as opposed to merely taking the difference of centroids).\n",
     "lig2 = BSS.Align.rmsdAlign(lig2, lig1, inverted_mapping)\n",
+    "\n",
     "# Merge the two ligands based on the mapping.\n",
     "merged = BSS.Align.merge(\n",
     "    lig1,\n",
@@ -315,9 +293,11 @@
     "    allow_ring_breaking=node.getInput(\"allow_ring_breaking\"),\n",
     "    allow_ring_size_change=node.getInput(\"allow_ring_size_change\"),\n",
     ")\n",
+    "\n",
     "# Create a composite system\n",
     "system1.removeMolecules(lig1)\n",
     "system1.addMolecules(merged)\n",
+    "\n",
     "# Make sure the box vectors are in reduced form.\n",
     "system1.reduceBoxVectors()\n",
     "system1.rotateBoxVectors()"
@@ -331,10 +311,11 @@
    "source": [
     "# Log the mapping used\n",
     "writeLog(lig1, lig2, mapping)\n",
-    "# Are we saving output for SOMD2?\n",
-    "is_somd2 = node.getInput(\"somd2\")\n",
+    "\n",
     "# File root for all output.\n",
     "root = node.getInput(\"output\")\n",
+    "\n",
+    "# Save PDB of merged topology.\n",
     "BSS.IO.saveMolecules(\n",
     "    \"merged_at_lam0.pdb\",\n",
     "    merged,\n",
@@ -345,18 +326,20 @@
     "        \"element\": \"element0\",\n",
     "    },\n",
     ")\n",
-    "if is_somd2:\n",
-    "    BSS.Stream.save(system1, root)\n",
-    "    stream_file = \"%s.bss\" % root\n",
-    "else:\n",
-    "    # Generate package specific input\n",
-    "    protocol = BSS.Protocol.FreeEnergy(\n",
-    "        runtime=2 * BSS.Units.Time.femtosecond, num_lam=3\n",
-    "    )\n",
-    "    process = BSS.Process.Somd(system1, protocol)\n",
-    "    process.getOutput()\n",
-    "    with zipfile.ZipFile(\"somd_output.zip\", \"r\") as zip_hnd:\n",
-    "        zip_hnd.extractall(\".\")"
+    "\n",
+    "\n",
+    "# Generate SOMD1 input.\n",
+    "protocol = BSS.Protocol.FreeEnergy(\n",
+    "    runtime=2 * BSS.Units.Time.femtosecond, num_lam=3\n",
+    ")\n",
+    "process = BSS.Process.Somd(system1, protocol)\n",
+    "process.getOutput()\n",
+    "with zipfile.ZipFile(\"somd_output.zip\", \"r\") as zip_hnd:\n",
+    "    zip_hnd.extractall(\".\")\n",
+    "\n",
+    "# Generate SOMD2 input.\n",
+    "BSS.Stream.save(system1, root)\n",
+    "stream_file = \"%s.bss\" % root"
    ]
   },
   {
@@ -365,12 +348,12 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "if not is_somd2:\n",
-    "    mergedpdb = \"%s.mergeat0.pdb\" % root\n",
-    "    pert = \"%s.pert\" % root\n",
-    "    prm7 = \"%s.prm7\" % root\n",
-    "    rst7 = \"%s.rst7\" % root\n",
-    "    mapping_str = \"%s.mapping\" % root"
+    "# Remap ouput names.\n",
+    "mergedpdb = \"%s.mergeat0.pdb\" % root\n",
+    "pert = \"%s.pert\" % root\n",
+    "prm7 = \"%s.prm7\" % root\n",
+    "rst7 = \"%s.rst7\" % root\n",
+    "mapping_str = \"%s.mapping\" % root"
    ]
   },
   {
@@ -379,19 +362,20 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "if not is_somd2:\n",
-    "    os.replace(\"merged_at_lam0.pdb\", mergedpdb)\n",
-    "    os.replace(\"somd.pert\", pert)\n",
-    "    os.replace(\"somd.prm7\", prm7)\n",
-    "    os.replace(\"somd.rst7\", rst7)\n",
-    "    os.replace(\"somd.mapping\", mapping_str)\n",
-    "    try:\n",
-    "        os.remove(\"somd_output.zip\")\n",
-    "        os.remove(\"somd.cfg\")\n",
-    "        os.remove(\"somd.err\")\n",
-    "        os.remove(\"somd.out\")\n",
-    "    except Exception:\n",
-    "        pass"
+    "# Rename.\n",
+    "os.replace(\"merged_at_lam0.pdb\", mergedpdb)\n",
+    "os.replace(\"somd.pert\", pert)\n",
+    "os.replace(\"somd.prm7\", prm7)\n",
+    "os.replace(\"somd.rst7\", rst7)\n",
+    "os.replace(\"somd.mapping\", mapping_str)\n",
+    "try:\n",
+    "    # Remove redundant files.\n",
+    "    os.remove(\"somd_output.zip\")\n",
+    "    os.remove(\"somd.cfg\")\n",
+    "    os.remove(\"somd.err\")\n",
+    "    os.remove(\"somd.out\")\n",
+    "except Exception:\n",
+    "    pass"
    ]
   },
   {
@@ -400,10 +384,8 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "if is_somd2:\n",
-    "    output = [stream_file]\n",
-    "else:\n",
-    "    output = [mergedpdb, pert, prm7, rst7, mapping_str]\n",
+    "# Set the node output.\n",
+    "output = [mergedpdb, pert, prm7, rst7, mapping_str, stream_file]\n",
     "node.setOutput(\"nodeoutput\", output)"
    ]
   },
@@ -413,6 +395,7 @@
    "metadata": {},
    "outputs": [],
    "source": [
+    "# Validate.\n",
     "node.validate()"
    ]
   }
@@ -433,7 +416,7 @@
    "name": "python",
    "nbconvert_exporter": "python",
    "pygments_lexer": "ipython3",
-   "version": "3.12.2"
+   "version": "3.12.13"
   }
  },
  "nbformat": 4,