minor changes, add openmmtrajectory.jl example

Manifest.toml

@@ -2,7 +2,7 @@
 
 julia_version = "1.10.4"
 manifest_format = "2.0"
-project_hash = "faf4567aaf643eff1653f6ca013dc56f535da286"
+project_hash = "82551888b97af7274c9467a3e0d8fa7fcbd8a068"
 
 [[deps.ADTypes]]
 git-tree-sha1 = "99a6f5d0ce1c7c6afdb759daa30226f71c54f6b0"
@@ -154,6 +154,12 @@ git-tree-sha1 = "c06a868224ecba914baa6942988e2f2aade419be"
 uuid = "a9b6321e-bd34-4604-b9c9-b65b8de01458"
 version = "0.1.0"
 
+[[deps.AtomsBase]]
+deps = ["LinearAlgebra", "PeriodicTable", "Printf", "Requires", "StaticArrays", "Unitful", "UnitfulAtomic"]
+git-tree-sha1 = "995c2b6b17840cd87b722ce9c6cdd72f47bab545"
+uuid = "a963bdd2-2df7-4f54-a1ee-49d51e6be12a"
+version = "0.3.5"
+
 [[deps.Automa]]
 deps = ["PrecompileTools", "TranscodingStreams"]
 git-tree-sha1 = "014bc22d6c400a7703c0f5dc1fdc302440cf88be"
@@ -350,6 +356,18 @@ weakdeps = ["SparseArrays"]
     [deps.ChainRulesCore.extensions]
     ChainRulesCoreSparseArraysExt = "SparseArrays"
 
+[[deps.Chemfiles]]
+deps = ["AtomsBase", "Chemfiles_jll", "DocStringExtensions", "PeriodicTable", "Unitful", "UnitfulAtomic"]
+git-tree-sha1 = "82fe5e341c793cb51149d993307da9543824b206"
+uuid = "46823bd8-5fb3-5f92-9aa0-96921f3dd015"
+version = "0.10.41"
+
+[[deps.Chemfiles_jll]]
+deps = ["Artifacts", "JLLWrappers", "Libdl"]
+git-tree-sha1 = "f3743181e30d87c23d9c8ebd493b77f43d8f1890"
+uuid = "78a364fa-1a3c-552a-b4bb-8fa0f9c1fcca"
+version = "0.10.4+0"
+
 [[deps.CloseOpenIntervals]]
 deps = ["Static", "StaticArrayInterface"]
 git-tree-sha1 = "05ba0d07cd4fd8b7a39541e31a7b0254704ea581"
@@ -1974,6 +1992,12 @@ git-tree-sha1 = "8489905bcdbcfac64d1daa51ca07c0d8f0283821"
 uuid = "69de0a69-1ddd-5017-9359-2bf0b02dc9f0"
 version = "2.8.1"
 
+[[deps.PeriodicTable]]
+deps = ["Base64", "Unitful"]
+git-tree-sha1 = "238aa6298007565529f911b734e18addd56985e1"
+uuid = "7b2266bf-644c-5ea3-82d8-af4bbd25a884"
+version = "1.2.1"
+
 [[deps.Pipe]]
 git-tree-sha1 = "6842804e7867b115ca9de748a0cf6b364523c16d"
 uuid = "b98c9c47-44ae-5843-9183-064241ee97a0"
@@ -2762,6 +2786,12 @@ weakdeps = ["ConstructionBase", "InverseFunctions"]
     ConstructionBaseUnitfulExt = "ConstructionBase"
     InverseFunctionsUnitfulExt = "InverseFunctions"
 
+[[deps.UnitfulAtomic]]
+deps = ["Unitful"]
+git-tree-sha1 = "903be579194534af1c4b4778d1ace676ca042238"
+uuid = "a7773ee8-282e-5fa2-be4e-bd808c38a91a"
+version = "1.0.0"
+
 [[deps.UnitfulLatexify]]
 deps = ["LaTeXStrings", "Latexify", "Unitful"]
 git-tree-sha1 = "975c354fcd5f7e1ddcc1f1a23e6e091d99e99bc8"

Project.toml

@@ -8,6 +8,7 @@ BioStructures = "de9282ab-8554-53be-b2d6-f6c222edabfc"
 Bonito = "824d6782-a2ef-11e9-3a09-e5662e0c26f8"
 CUDA = "052768ef-5323-5732-b1bb-66c8b64840ba"
 ChainRulesCore = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4"
+Chemfiles = "46823bd8-5fb3-5f92-9aa0-96921f3dd015"
 Combinatorics = "861a8166-3701-5b0c-9a16-15d98fcdc6aa"
 DataFrames = "a93c6f00-e57d-5684-b7b6-d8193f3e46c0"
 Distances = "b4f34e82-e78d-54a5-968a-f98e89d6e8f7"

scripts/multitraj.jl

@@ -5,7 +5,8 @@ using ISOKANN:
     pdists,
     atom_indices,
     data_from_trajectory,
-    lastcat
+    lastcat,
+    *
 
 DATADIR = "/data/numerik/ag_cmd/trajectory_transfers_for_isokann/data/8EF5_500ns_pka_7.4_no_capping_310.10C"
 MAXRADIUS = 0.5 # angstrom
@@ -15,7 +16,7 @@ trajfiles = ["$DATADIR/traj.dcd" for i in 1:2]
 pdbfile = "$DATADIR/struct.pdb"
 
 
-molecule = load_trajectory(pdbfile, top=pdbfile)
+molecule = load_trajectory(pdbfile)
 pdist_inds = restricted_localpdistinds(molecule, MAXRADIUS, atom_indices(pdbfile, "not water and name==CA"))
 
 

scripts/openmmtrajectory.jl

@@ -0,0 +1,26 @@
+using ISOKANN
+
+pdb = "data/openmm8ef5-processed.pdb"
+radius = 1.0
+nk = 2
+samples = 100
+trainsteps = 100
+
+###
+
+sim = OpenMMSimulation(; pdb, minimize=true, gpu=true)
+
+x0 = getcoords(sim)
+cainds = atom_indices(pdb, "name==CA")
+featinds = restricted_localpdistinds(x0, radius, cainds)
+
+traj = laggedtrajectory(sim, samples)
+data = SimulationData(sim, traj, nk, featurizer=coords -> pdists(coords, featinds))
+
+iso = Iso2(data, opt=NesterovRegularized(), gpu=true)
+
+run!(iso, trainsteps)
+
+iso.data = addcoords(iso.data, laggedtrajectory(sim, samples, x0=iso.data.coords[1][:, end]))
+
+run!(iso, trainsteps)

src/ISOKANN.jl

@@ -55,6 +55,7 @@ export pairnet
 #export MollyLangevin, propagate, solve#, MollySDE
 
 export propagate
+export laggedtrajectory
 
 export run!, runadaptive!
 export AdamRegularized, NesterovRegularized
@@ -65,10 +66,13 @@ export Iso2
 export Doublewell, Triplewell, MuellerBrown
 export chis
 export SimulationData
+export addcoords
 export getxs, getys
 export exit_rates
-export atom_indices
 export load_trajectory, save_trajectory
+export atom_indices
+export localpdistinds, pdists, restricted_localpdistinds
+
 
 export reactionpath_minimum, reactionpath_ode
 

src/bonito.jl

@@ -20,7 +20,7 @@ end
 logrange(x1, x2, n) = round.((10^y for y in range(log10(x1), log10(x2), length=n)), sigdigits=1)
 
 function isocreator()
-    pdb = Dropdown(["data/alanine-dipeptide-nowater.pdb", "data/vgv.pdb", "data/villin nowater.pdb"])
+    pdb = Dropdown(["data/alanine-dipeptide-nowater.pdb", "data/vgv.pdb", "data/villin nowater.pdb", "temp/8eftstepped.pdb"])
     pdbid = TextField("")
     steps = StylableSlider(1:1000, value=10)
     temperature = StylableSlider(-10:70, value=30)

src/pairdists.jl

@@ -79,7 +79,6 @@ Like `localdists`, but consider only the atoms with index in `atoms`
 """
 function restricted_localpdistinds(coords, radius, atoms)
     rc = reshape(reshape(coords, 3, :, size(coords, 2))[:, atoms, :], :, size(coords, 2))
-    @show size(rc)
     pairs = localpdistinds(rc, radius)
     map(pairs) do (a, b)
         (atoms[a], atoms[b])

src/simulation.jl

@@ -137,13 +137,18 @@ Merge the data and features of `d1` and `d2`, keeping the simulation and feature
 Note that there is no check if simulation features agree.
 """
 function Base.merge(d1::SimulationData, d2::SimulationData)
-    features = lastcat.(d1.features, d2.features)
     coords = lastcat.(d1.coords, d2.coords)
+    d2f = if d1.featurizer == d2.featurizer
+        d2.features
+    else
+        d1.featurizer.(d2.coords)
+    end
+    features = lastcat.(d1.features, d2f)
     return SimulationData(d1.sim, features, coords, d1.featurizer)
 end
 
 function addcoords(d::SimulationData, coords::AbstractMatrix)
-    merge(d, SimulationData(d.sim, coords, nk(d)))
+    merge(d, SimulationData(d.sim, coords, nk(d), featurizer=d.featurizer))
 end
 
 

src/simulators/openmm.jl

@@ -195,10 +195,10 @@ function trajectory(s::OpenMMSimulation, x0::AbstractVector{T}=getcoords(s), ste
     return xs
 end
 
-function ISOKANN.laggedtrajectory(s::OpenMMSimulation, n_lags, steps_per_lag = s.steps)
+function ISOKANN.laggedtrajectory(s::OpenMMSimulation, n_lags, steps_per_lag=s.steps; x0=getcoords(s))
     steps = steps_per_lag * n_lags
     saveevery = steps_per_lag
-    trajectory(s, getcoords(s), steps, saveevery)
+    trajectory(s, x0, steps, saveevery)
 end
 
 getcoords(sim::OpenMMSimulation) = getcoords(sim.pysim, sim.momenta)#::Vector