include the dashboard into the package, convenience trajectorydata ge…

…nerators

src/ISOKANN.jl

@@ -99,4 +99,7 @@ include("reactionpath2.jl")
 include("IsoMu/IsoMu.jl")
 include("vgv/vgv.jl")
 
+include("makie.jl")
+include("bonito.jl")
+
 end

src/bonito.jl

@@ -1,11 +1,5 @@
-using Revise
 using Bonito
 using WGLMakie
-using ISOKANN
-
-includet("makie.jl")
-
-#iso = Iso2(OpenMMSimulation(steps=30, pdb="data/vgv.pdb", forcefields=ISOKANN.OpenMM.FORCE_AMBER_IMPLICIT), loggers=[], opt=ISOKANN.NesterovRegularized(1e-3, 0), nx=10, nk=1, minibatch=64)
 
 USEGPU = true
 ISO = nothing
@@ -111,10 +105,13 @@ function isocreator()
         ); width="300px",), isoo
 end
 
-app = App(title="ISOKANN Dashboard") do session
-    return content(session)
-end
 
+function serve()
+    app = App(title="ISOKANN Dashboard") do session
+        return content(session)
+    end
 
-server = Bonito.get_server()
-route!(server, "/" => app)
+    server = Bonito.get_server()
+    route!(server, "/" => app)
+    return app
+end

src/data.jl

@@ -140,16 +140,6 @@ end
 
 @deprecate shuffledata(data) shuffleobs(data)
 
-# TODO: this does not belong here!
-"""  trajectory(sim, nx)
-generate a trajectory of length `nx` from the simulation `sim`"""
-function trajectory(sim, nx)
-    siml = deepcopy(sim)
-    logevery = round(Int, sim.T / sim.dt)
-    siml.T = sim.T * nx
-    xs = solve(siml; logevery=logevery)
-    return xs
-end
 
 
 ### Data I/O

src/makie.jl

@@ -1,5 +1,4 @@
 using WGLMakie
-using ISOKANN
 using WGLMakie.Observables: throttle
 using ThreadPools
 
@@ -71,7 +70,7 @@ end
 
 
 
-function dashboard(iso::Iso2, session)
+function dashboard(iso::Iso2, session=nothing)
     coords = Observable(iso.data.coords[1] |> cpu)
     chis = Observable(ISOKANN.chis(iso) |> vec |> cpu)
     icur = Observable(1)
@@ -144,7 +143,7 @@ function dashboard(iso::Iso2, session)
         ThreadPools.@tspawnat 1 begin
             try
                 last = time()
-                while isready(session) && run.active[]
+                while (isnothing(session) || isready(session)) && run.active[]
 
                     run!(iso)
                     adaptivesampling.active[] && ISOKANN.resample_kde!(iso, 1; padding=0.01, bandwidth=0.1)

src/models.jl

@@ -52,11 +52,12 @@ end
 `features` allows to pass a featurizer as preprocessor,
 `activation` determines the activation function for each but the last layer
 `lastactivation` can be used to modify the last layers activation function """
-function pairnet(n::Int=22; layers=3, features=identity, activation=Flux.sigmoid, lastactivation=identity, nout=1)
+function pairnet(n::Int=22; layers=3, features=identity, activation=Flux.sigmoid, lastactivation=identity, nout=1, layernorm=false)
     float32(x) = Float32.(x)
     nn = Flux.Chain(
         #float32,
         features,
+        layernorm ? Flux.LayerNorm(n) : identity,
         [Flux.Dense(
             round(Int, n^(l / layers)),
             round(Int, n^((l - 1) / layers)),

src/simulation.jl

@@ -28,6 +28,9 @@ function randx0(sim::IsoSimulation, nx)
     return xs
 end
 
+trajectory(sim::IsoSimulation, steps) = error("not implemented")
+laggedtrajectory(sim::IsoSimulation, nx) = error("not implemented")
+
 ###
 
 """
@@ -191,3 +194,14 @@ end
 function datasize((xs, ys)::Tuple)
     return size(xs), size(ys)
 end
+
+function trajectorydata(sim::IsoSimulation, steps; reverse=false, kwargs...)
+    xs = laggedtrajectory(sim, steps)
+    SimulationData(sim, data_from_trajectory(xs; reverse), kwargs...)
+end
+
+function trajectoryburstdata(sim, steps, nk; kwargs)
+    xs = laggedtrajectory(sim, steps)
+    ys = propagate(sim, xs, nk)
+    SimulationData(sim, ys, kwargs...)
+end

src/simulators/langevin.jl

@@ -29,6 +29,16 @@ function propagate(l::AbstractLangevin, x0::AbstractMatrix, ny)
     return ys
 end
 
+"""  trajectory(sim, nx)
+generate a trajectory of length `nx` from the simulation `sim`"""
+function trajectory(sim, nx)
+    siml = deepcopy(sim)
+    logevery = round(Int, sim.T / sim.dt)
+    siml.T = sim.T * nx
+    xs = solve(siml; logevery=logevery)
+    return xs
+end
+
 function solve_end(l::AbstractLangevin; u0)
     StochasticDiffEq.solve(SDEProblem(l, u0))[:, end]
 end

src/simulators/mopenmm.py

@@ -34,15 +34,8 @@ def trajectory(sim, x0, stepsize, steps, saveevery, mmthreads, withmomenta):
   trajectory[0] = x0
 
   c = newcontext(sim.context, mmthreads)
-
-  if withmomenta:
-    n = len(x0) // 2
-    c.setPositions(x0[:n])
-    c.setVelocities(x0[n:])
-  else:
-    c.setPositions(x0)
-    c.setVelocitiesToTemperature(sim.integrator.getTemperature())
-
+
+  set_numpy_state(c, x0)
   c.getIntegrator().setStepSize(stepsize)
 
   for n in range(1,n_states):

src/simulators/openmm.jl

@@ -8,7 +8,7 @@ import ISOKANN: ISOKANN, IsoSimulation,
     propagate, dim, randx0,
     featurizer, defaultmodel,
     savecoords, getcoords, force, pdb,
-    force, potential, lagtime
+    force, potential, lagtime, trajectory
 
 export OpenMMSimulation, FORCE_AMBER, FORCE_AMBER_IMPLICIT
 
@@ -193,6 +193,12 @@ function trajectory(s::OpenMMSimulation, x0::AbstractVector{T}=getcoords(s), ste
     return xs
 end
 
+function ISOKANN.laggedtrajectory(s::OpenMMSimulation, nlags)
+    steps = s.steps * nlags
+    saveevery = s.steps
+    trajectory(s, getcoords(s), steps, saveevery)
+end
+
 getcoords(sim::OpenMMSimulation) = getcoords(sim.pysim, sim.momenta)#::Vector
 setcoords(sim::OpenMMSimulation, coords) = setcoords(sim.pysim, coords, sim.momenta)