working on the water experiments

data/systems/vgvapg.pdb

@@ -1,3 +1,4 @@
+CRYST1   30.000   30.000   30.000  90.00  90.00  90.00 P 1           1
 MODEL        0
 ATOM      1  N   VAL A   0      -0.428  38.346  25.000  1.00  0.00           N  
 ATOM      2  H   VAL A   0      -0.441  38.991  24.202  1.00  0.00           H  

scripts/waterimpexp.jl

@@ -4,10 +4,10 @@ using Plots
 
 VAL_NX = 1000 # number of validation samples
 VAL_STEPS = 10_000 # steps between validation samples
-VAL_NK = 10 # validation burst size
+VAL_NK = 100 # validation burst size
 
 NX = 100_000 # number of training samples
-NITER = 10_000 # number of training iterations
+NITER = 1000 # number of training iterations
 
 REVERSE = true # also count reverse transitions
 PDB = "data/systems/alanine dipeptide.pdb"
@@ -24,7 +24,11 @@ valsim = OpenMMSimulation(addwater=true, padding=1; features)
 valtraj = laggedtrajectory(valsim, VAL_NX, steps=VAL_STEPS)
 valdata = SimulationData(valsim, valtraj, VAL_NK) |> gpu
 
-isos = map(enumerate(sims)) do (i, sim)
+scatter_ramachandran(valtraj)
+
+isos = []
+for (i, sim) in enumerate(sims)
+
     @info "generating data $i"
     xs = laggedtrajectory(sim, NX)
     scatter_ramachandran(xs) |> display
@@ -33,9 +37,16 @@ isos = map(enumerate(sims)) do (i, sim)
 
     vallog = ISOKANN.ValidationLossLogger(data=valdata)
     @info "training model $i"
-    iso = Iso(data, loggers=Any[vallog], autoplot=30)
+    iso = Iso(data, loggers=Any[vallog], minibatch=1000)
+
+    push!(isos, iso)
 
     run!(iso, NITER)
     plot_training(iso)
     savefig("water trained $i.png")
 end
+
+
+function deltadchi(iso1=isos[1], iso2=isos[2], coords=valtraj)
+    delta = reduce(hcat, [ISOKANN.dchidx(iso2, x) - ISOKANN.dchidx(iso1, x) for x in eachcol(coords |> gpu)])
+end

src/makie.jl

@@ -69,13 +69,13 @@ function plotmol(c, pysim, color=1; grad=nothing, kwargs...)
     if !isnothing(grad)
         grad = @lift($(observe(grad))[:, $i])
         plotgrad!(ax, @lift(reshape($col, 3, :)), @lift(reshape($grad, 3, :)),
-            arrowsize=0.01, lengthscale=0.2, linecolor=:red, linewidth=0.005)
+            arrowsize=0.01, lengthscale=0.2, linecolor=:red, linewidth=0.001)
     end
 
     return fig
 end
 
-function plotgrad!(ax, c::Observable{T}, dc::Observable{T}; kwargs...) where {T<:AbstractMatrix}
+function plotgrad!(ax, c::Observable{<:AbstractMatrix}, dc::Observable{<:AbstractMatrix}; kwargs...)
 
     x = @lift vec($c[1, :])
     y = @lift vec($c[2, :])
@@ -84,6 +84,9 @@ function plotgrad!(ax, c::Observable{T}, dc::Observable{T}; kwargs...) where {T<
     v = @lift vec($dc[2, :])
     w = @lift vec($dc[3, :])
 
+    @show x[] |> size
+    @show u[] |> size
+
     arrows!(ax, x, y, z, u, v, w; kwargs...)
 end
 

src/plots.jl

@@ -64,7 +64,7 @@ function plot_training(iso; subdata=nothing)
             push!(ps, l.plot())
         end
     end
-    plot(ps..., layout=(length(ps), 1), size=(400, 300 * length(ps)))
+    plot(ps..., layout=(length(ps), 1), size=(400, 300 * length(ps)), fmt=:png)
 end
 
 function plot_chi(iso; target=true)

src/simulators/openmm.jl

@@ -214,9 +214,44 @@ function propagate(sim::OpenMMSimulation, x0::AbstractMatrix, nk)
     return ys
 end
 
+"""
+    laggedtrajectory(sim::OpenMMSimulation, lags; steps=steps(sim), resample_velocities=true, kwargs...)
+
+Generate a lagged trajectory for a given OpenMMSimulation.
+E.g. x0--x--x--x  for `lags=3` and `steps=2`
+
+# Arguments
+- `sim::OpenMMSimulation`: The simulation object.
+- `lags`: The number of steps.
+- `steps`: The lagtime, i.e. number of steps to take in the simulation.
+- `resample_velocities`: Whether to resample velocities according to Maxwell-Boltzman at each step.
+- `kwargs...`: Additional keyword arguments to pass to the `trajectory` function.
+
+# Returns
+- A matrix of `lags` samples which each have `steps` simulation-steps inbetween them.
+"""
 laggedtrajectory(sim::OpenMMSimulation, lags; steps=steps(sim), resample_velocities=true, kwargs...) =
     trajectory(sim, lags * steps; saveevery=steps, resample_velocities, kwargs...)
 
+
+"""
+    trajectory(sim::OpenMMSimulation{Nothing}, steps=steps(sim); saveevery=1, x0=getcoords(sim), resample_velocities=false, throw=false, showprogress=true, reclaim=true)
+
+Simulates the trajectory of an OpenMM simulation.
+
+# Arguments
+- `sim::OpenMMSimulation{Nothing}`: The OpenMM simulation object.
+- `steps`: The number of steps to simulate. Defaults to the number of steps defined in the simulation object.
+- `saveevery`: Interval at which to save the trajectory. Defaults to 1.
+- `x0`: Initial coordinates for the simulation. Defaults to the current coordinates of the simulation object.
+- `resample_velocities`: Whether to resample velocities at the start of the simulation. Defaults to `false`.
+- `throw`: Whether to throw an error if the simulation fails. If false it returns the trajectory computed so far. Defaults to `false`.
+- `showprogress`: Whether to display a progress bar during the simulation. Defaults to `true`.
+- `reclaim`: Whether to reclaim CUDA memory before the simulation. Defaults to `true`.
+
+# Returns
+- The trajectory of the simulation as a matrix of coordinates.
+"""
 function trajectory(sim::OpenMMSimulation{Nothing}, steps=steps(sim); saveevery=1, x0=getcoords(sim), resample_velocities=false, throw=false, showprogress=true, reclaim=true)
     reclaim && claim_memory(sim)
     n = div(steps, saveevery)