From 6a613350f35a629bdcff2a96427583cb22b43226 Mon Sep 17 00:00:00 2001
From: Sikorski <sikorski@zib.de>
Date: Mon, 22 Jan 2024 16:07:01 +0100
Subject: [PATCH] prototype openmm interface

---
 Manifest.toml | 19 +++++++++++-
 Project.toml  |  1 +
 src/openmm.jl | 84 +++++++++++++++++++++++++++++++++++++++++++++++++++
 3 files changed, 103 insertions(+), 1 deletion(-)
 create mode 100644 src/openmm.jl

diff --git a/Manifest.toml b/Manifest.toml
index 7cf261c..2b97f16 100644
--- a/Manifest.toml
+++ b/Manifest.toml
@@ -2,7 +2,7 @@
 
 julia_version = "1.10.0"
 manifest_format = "2.0"
-project_hash = "13e03dcf8c31e0ac86aedcdfc539300d2edf1af2"
+project_hash = "ad146ecae413bfb5492b94014259778ed1274664"
 
 [[deps.ADTypes]]
 git-tree-sha1 = "41c37aa88889c171f1300ceac1313c06e891d245"
@@ -391,6 +391,12 @@ git-tree-sha1 = "8cfa272e8bdedfa88b6aefbbca7c19f1befac519"
 uuid = "f0e56b4a-5159-44fe-b623-3e5288b988bb"
 version = "2.3.0"
 
+[[deps.Conda]]
+deps = ["Downloads", "JSON", "VersionParsing"]
+git-tree-sha1 = "51cab8e982c5b598eea9c8ceaced4b58d9dd37c9"
+uuid = "8f4d0f93-b110-5947-807f-2305c1781a2d"
+version = "1.10.0"
+
 [[deps.ConstructionBase]]
 deps = ["LinearAlgebra"]
 git-tree-sha1 = "c53fc348ca4d40d7b371e71fd52251839080cbc9"
@@ -1874,6 +1880,12 @@ git-tree-sha1 = "00099623ffee15972c16111bcf84c58a0051257c"
 uuid = "92933f4c-e287-5a05-a399-4b506db050ca"
 version = "1.9.0"
 
+[[deps.PyCall]]
+deps = ["Conda", "Dates", "Libdl", "LinearAlgebra", "MacroTools", "Serialization", "VersionParsing"]
+git-tree-sha1 = "9816a3826b0ebf49ab4926e2b18842ad8b5c8f04"
+uuid = "438e738f-606a-5dbb-bf0a-cddfbfd45ab0"
+version = "1.96.4"
+
 [[deps.Qt6Base_jll]]
 deps = ["Artifacts", "CompilerSupportLibraries_jll", "Fontconfig_jll", "Glib_jll", "JLLWrappers", "Libdl", "Libglvnd_jll", "OpenSSL_jll", "Vulkan_Loader_jll", "Xorg_libSM_jll", "Xorg_libXext_jll", "Xorg_libXrender_jll", "Xorg_libxcb_jll", "Xorg_xcb_util_cursor_jll", "Xorg_xcb_util_image_jll", "Xorg_xcb_util_keysyms_jll", "Xorg_xcb_util_renderutil_jll", "Xorg_xcb_util_wm_jll", "Zlib_jll", "libinput_jll", "xkbcommon_jll"]
 git-tree-sha1 = "37b7bb7aabf9a085e0044307e1717436117f2b3b"
@@ -2482,6 +2494,11 @@ git-tree-sha1 = "7209df901e6ed7489fe9b7aa3e46fb788e15db85"
 uuid = "3d5dd08c-fd9d-11e8-17fa-ed2836048c2f"
 version = "0.21.65"
 
+[[deps.VersionParsing]]
+git-tree-sha1 = "58d6e80b4ee071f5efd07fda82cb9fbe17200868"
+uuid = "81def892-9a0e-5fdd-b105-ffc91e053289"
+version = "1.3.0"
+
 [[deps.VertexSafeGraphs]]
 deps = ["Graphs"]
 git-tree-sha1 = "8351f8d73d7e880bfc042a8b6922684ebeafb35c"
diff --git a/Project.toml b/Project.toml
index 800aa1b..52221d6 100644
--- a/Project.toml
+++ b/Project.toml
@@ -27,6 +27,7 @@ PlotlyBase = "a03496cd-edff-5a9b-9e67-9cda94a718b5"
 PlotlyKaleido = "f2990250-8cf9-495f-b13a-cce12b45703c"
 Plots = "91a5bcdd-55d7-5caf-9e0b-520d859cae80"
 ProgressMeter = "92933f4c-e287-5a05-a399-4b506db050ca"
+PyCall = "438e738f-606a-5dbb-bf0a-cddfbfd45ab0"
 Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
 SlurmClusterManager = "c82cd089-7bf7-41d7-976b-6b5d413cbe0a"
 SnoopPrecompile = "66db9d55-30c0-4569-8b51-7e840670fc0c"
diff --git a/src/openmm.jl b/src/openmm.jl
new file mode 100644
index 0000000..fe4665c
--- /dev/null
+++ b/src/openmm.jl
@@ -0,0 +1,84 @@
+
+using PyCall
+import ISOKANN: propagate
+
+py"""
+from openmm.app import *
+from openmm import *
+from openmm.unit import *
+from sys import stdout
+"""
+
+nanometer = py"nanometer"
+
+struct OpenMMSimulation
+    pysim::PyObject
+    steps::Int
+end
+
+###
+
+function propagate(sim::OpenMMSimulation, x0)
+    steps = sim.steps
+    pysim = sim.pysim
+    setcoords(pysim, x0)
+    pysim.step(steps)
+    getcoords(pysim)
+end
+
+function propagate(sim::OpenMMSimulation, x0::AbstractMatrix, ny)
+    dim, nx = size(x0)
+    ys = zeros(dim, nx, ny)
+    for i in 1:nx, j in 1:ny
+        ys[:, i, j] = propagate(sim, x0[:, i])
+    end
+    return ys
+end
+
+function getcoords(sim::PyObject)
+    x = sim.context.getState(getPositions=true).getPositions().value_in_unit(nanometer)
+    reinterpret(Float64, x)
+end
+
+function setcoords(sim::PyObject, coords)
+    x = reinterpret(Tuple{Float64,Float64,Float64}, coords) * nanometer
+    sim.context.setPositions(x)
+end
+
+###
+
+function openmm_examplesys(;
+    temp=300,
+    friction=1,
+    step=0.004,
+    pdb="/home/htc/bzfsikor/.julia/conda/3/share/openmm/examples/input.pdb",
+    forcefields=["amber14-all.xml", "amber14/tip3pfb.xml"],
+    steps=100)
+
+    py"""
+    pdb = PDBFile($pdb)
+    forcefield = ForceField(*$forcefields)
+    system = forcefield.createSystem(pdb.topology, nonbondedMethod=PME,
+            nonbondedCutoff=1*nanometer, constraints=HBonds)
+    integrator = LangevinMiddleIntegrator($temp*kelvin, $friction/picosecond, $step*picoseconds)
+    simulation = Simulation(pdb.topology, system, integrator)
+    simulation.context.setPositions(pdb.positions)
+    """
+    # simulation.minimizeEnergy()
+    # simulation.reporters.append(PDBReporter('output.pdb', 1000))
+    # simulation.reporters.append(StateDataReporter(stdout, 1000, step=True,
+    #        potentialEnergy=True, temperature=True))
+    # simulation.step(10000)
+
+    sim = py"simulation"
+    return OpenMMSimulation(sim, steps)
+end
+
+###
+
+function test_getsetcoords(sim)
+    x = getcoords(sim)
+    bak = copy(x)
+    setcoords(sim, x)
+    @assert getcoords(sim) == bak
+end