small fixes, working on debugging minimization. WAIT: why do we divid…

…e the force by the masses?!

scripts/villin.jl

@@ -5,7 +5,7 @@ using PyCall
 
 ## Config
 
-comment = "momenta-long-medlowfriction"
+comment = ""
 
 pdb = "data/villin nowater.pdb"
 steps = 20_000
@@ -17,7 +17,7 @@ minimize = true
 momenta = true
 features = 0.5 # 0 => backbone only
 #forcefields = OpenMM.FORCE_AMBER_IMPLICIT  # TODO: this shouldnb be an option the way we build addwater now
-forcefields = OpenMM.FORCE_AMBER
+forcefields = ISOKANN.OpenMM.FORCE_AMBER
 addwater = true
 padding = 1
 ionicstrength = 0.0
@@ -26,9 +26,9 @@ nx = 20
 nk = 1
 iter = 1000
 generations = 1000
-cutoff = 3000
+cutoff = 2500
 
-kde_padding = 0.05
+kde_padding = 0.02
 extrapolates = 0 # *2
 extrapolate = 0.05
 keepedges = false
@@ -42,6 +42,7 @@ maxjump = 1
 
 path = "out/villin/$(now())"[1:end-4] * comment
 
+
 readdata = nothing
 #readdata = "latest/iso.jld2"
 
@@ -89,18 +90,28 @@ iso = Iso2(data;
     loggers=[])
 
 @time "initial training" run!(iso, iter)
-
+data = nothing
 ## Running
 
 newsim = true
 
 for i in 1:generations
     GC.gc()
+    if length(iso.data) > cutoff
+        iso.data = iso.data[end-cutoff+1:end]
+    end
+    @show varinfo()
+    GC.gc()
 
     @time "extrapolating" ISOKANN.addextrapolates!(iso, extrapolates, stepsize=extrapolate)
     @time "kde sampling" ISOKANN.resample_kde!(iso, nx; padding=kde_padding)
+
+
+
     @time "training" run!(iso, iter)
 
+
+
     simtime = ISOKANN.simulationtime(iso)
 
     @time "saving" begin

src/extrapolate.jl

@@ -28,7 +28,7 @@ extrapolate them by `stepsize` for `steps` steps beyond their extrema,
 resulting in 2n new points.
 If `minimize` is true, the new points are energy minimized.
 """
-function extrapolate(iso, n::Integer, stepsize=0.1, steps=1, minimize=true)
+function extrapolate(iso, n::Integer, stepsize=0.1, steps=1, minimize=true, maxskips=10)
     data = iso.data
     model = iso.model
     coords = flattenlast(data.coords[2])
@@ -40,12 +40,20 @@ function extrapolate(iso, n::Integer, stepsize=0.1, steps=1, minimize=true)
 
     for (p, dir, N) in [(p, -1, n), (reverse(p), 1, 2 * n)]
         for i in p
+            skips > maxskips && break
             try
                 x = extrapolate(iso, coords[:, i], dir * stepsize, steps)
                 minimize && (x = energyminimization_chilevel(iso, x))
+                if data.sim.momenta
+                    x = reshape(x, :, 2)
+                    #x[:, 2] .= 0
+                    x = vec(x)
+                end
+                #&& ISOKANN.OpenMM.set_random_velocities!(data.sim, x)
                 push!(xs, x)
             catch e
-                if isa(e, PyCall.PyError) || isa(e, DomainError)
+                if isa(e, PyCall.PyError) || isa(e, DomainError) || isa(e, AssertionError)
+                    @show e
                     skips += 1
                     continue
                 end
@@ -69,26 +77,42 @@ function extrapolate(iso, x::AbstractVector, step, steps)
     return x
 end
 
-function energyminimization_chilevel(iso, x0; f_tol=1e-1, alphaguess=1e-6, iterations=100, show_trace=true)
+global trace = []
+
+function energyminimization_chilevel(iso, x0; f_tol=1e-3, alphaguess=1e-5, iterations=20, show_trace=false, skipwater=true)
     sim = iso.data.sim
-    x = copy(x0)
+    x = copy(x0) .|> Float64
 
     chi(x) = myonly(chicoords(iso, x))
     chilevel = Levelset(chi, chi(x0))
 
     U(x) = OpenMM.potential(sim, x)
-    dU(x) = -OpenMM.force(sim, x)
+    dU(x) = (f = -OpenMM.force(sim, x); (skipwater && zerowater!(sim, f)); f)
+
+
+    linesearch = Optim.LineSearches.HagerZhang(alphamax=alphaguess)
+    alg = Optim.LBFGS(; linesearch, alphaguess, manifold=chilevel)
 
-    o = Optim.optimize(U, dU, x, Optim.LBFGS(; alphaguess, manifold=chilevel), Optim.Options(; iterations, f_tol, show_trace); inplace=false)
+    o = Optim.optimize(U, dU, x, alg, Optim.Options(; iterations, f_tol, show_trace); inplace=false)
     return o.minimizer
 end
 
+function zerowater!(sim, x)
+    inds = map(sim.pysim.topology.atoms()) do a
+        a.residue.name == "HOH"
+    end
+    x = reshape(x, 3, :)
+    x[:, inds] .= 0
+    vec(x)
+end
+
 struct Levelset{F,T} <: Optim.Manifold
     f::F
     target::T
 end
 
 function Optim.project_tangent!(M::Levelset, g, x)
+    replace!(g, NaN => 0)
     u = Zygote.gradient(M.f, x) |> only
     u ./= norm(u)
     g .-= dot(g, u) * u

src/iso2.jl

@@ -130,21 +130,23 @@ end
 Train iso with adaptive sampling. Sample `nx` new data points followed by `iter` isokann iterations and repeat this `generations` times.
 `cutoff` specifies the maximal data size, after which new data overwrites the oldest data.
 """
-function runadaptive!(iso; generations=1, nx=10, iter=100, cutoff=Inf, keepedges=false, extrapolates=0, extrapolation=0.01)
+function runadaptive!(iso; generations=1, nx=10, iter=100, cutoff=Inf, keepedges=false, extrapolates=0, extrapolation=0.01, kde=0, kde_padding=0)
     p = ProgressMeter.Progress(generations)
     t_sim = 0.
+    t_kde = 0.0
     t_train = 0.
     t_extra = 0.0
     for g in 1:generations
         GC.gc()
         t_sim += @elapsed adddata!(iso, nx; keepedges)
+        t_kde += @elapsed ISOKANN.resample_kde!(iso, kde; padding=kde_padding)
+
+        t_extra += @elapsed ISOKANN.addextrapolates!(iso, extrapolates, stepsize=extrapolation)
 
         if length(iso.data) > cutoff
             iso.data = iso.data[end-cutoff+1:end]
         end
 
-        t_extra += @elapsed ISOKANN.addextrapolates!(iso, extrapolates, stepsize=extrapolation)
-
         t_train += @elapsed run!(iso, iter, showprogress=false)
 
         ProgressMeter.next!(p;
@@ -153,7 +155,7 @@ function runadaptive!(iso; generations=1, nx=10, iter=100, cutoff=Inf, keepedges
                 (:loss, iso.losses[end]),
                 (:iterations, length(iso.losses)),
                 (:data, size(getys(iso.data))),
-                ("t_sim, t_train, t_extra", (t_sim, t_train, t_extra)),
+                ("t_sim, t_train, t_extra, t_kde", (t_sim, t_train, t_extra, t_kde)),
                 ("simulated time", "$(simulationtime(iso))"), #TODO: doesnt work with cutoff
                 (:macrorates, exit_rates(iso))],
             ignore_predictor=true)

src/models.jl

@@ -19,7 +19,7 @@ inputdim(model::Flux.Dense) = size(model.weight, 2)
 outputdim(model::Flux.Chain) = outputdim(model.layers[end])
 outputdim(model::Flux.Dense) = size(model.weight, 1)
 
-iscuda(m::Flux.Chain) = Flux.params(m)[1] isa CuArray
+iscuda(m::Flux.Chain) = m[2].weight isa CuArray
 
 
 

src/pairdists.jl

@@ -96,7 +96,7 @@ function pdists(x::AbstractArray, inds::Vector{Tuple{T,T}}) where {T}
   d, s... = size(x)
   b = reshape(x, d, :)
   p = pdists(b, inds)
-  return reshape(p, :, s...)
+    return reshape(p, length(inds), s...)
 end
 
 # convenience wrapper

src/simulation.jl

@@ -73,6 +73,9 @@ function SimulationData(sim::IsoSimulation, xs::AbstractMatrix, nk::Int; kwargs.
             xs = xs[:, e.select]
             ys = e.result[:, :, e.select]
             @warn "SimulationData: discarded $(nx-sum(e.select))/$nx starting points due to simulation errors."
+            if sum(e.select) == 0
+                return sim
+            end
         else
             rethrow(e)
         end

src/simulators/mopenmm.py

@@ -126,7 +126,14 @@ def set_numpy_state(context, x, withmomenta):
         context.setVelocities(x[n:])
     else:
         context.setPositions(x)
-        context.setVelocitiesToTemperature(sim.integrator.getTemperature())
+        context.setVelocitiesToTemperature(context.getIntegrator().getTemperature())
+
+def set_random_velocities(context,  mmthreads):
+    context.setVelocitiesToTemperature(context.getIntegrator().getTemperature())
+    v = context.getState(getVelocities=True).getVelocities(asNumpy=True).value_in_unit(nanometer/picosecond).flatten()
+    return v
+
+
 
 def newcontext(context, mmthreads):
   if mmthreads == 'gpu':

src/simulators/openmm.jl

@@ -187,6 +187,11 @@ setcoords(sim::OpenMMSimulation, coords) = setcoords(sim.pysim, coords, sim.mome
 
 getcoords(sim::PyObject, momenta) = py"get_numpy_state($sim.context, $momenta).flatten()"
 
+function minimize(sim::OpenMMSimulation, coords, iter=100)
+    setcoords(sim, coords)
+    sim.pysim.minimizeEnergy(maxIterations=iter)
+    return getcoords(sim)
+end
 
 function setcoords(sim::PyObject, coords::AbstractVector{T}, momenta) where {T}
     t = reinterpret(Tuple{T,T,T}, Array(coords))
@@ -200,26 +205,38 @@ function setcoords(sim::PyObject, coords::AbstractVector{T}, momenta) where {T}
     end
 end
 
+""" mutates the state in sim """
+function set_random_velocities!(sim, x)
+    v = py"set_random_velocities($(sim.pysim.context), $(sim.mmthreads))"
+    n = length(x) ÷ 2
+    x[n+1:end] = v
+
+    return x
+end
+
 force(sim::OpenMMSimulation, x::CuArray) = force(sim, Array(x)) |> cu
 potential(sim::OpenMMSimulation, x::CuArray) = potential(sim, Array(x))
 
 function force(sim::OpenMMSimulation, x)
     CUDA.reclaim()
-    sim = sim.pysim
     setcoords(sim, x)
-    f = sim.context.getState(getForces=true).getForces(asNumpy=true)
-    f = f.value_in_unit(f.unit) |> permutedims |> vec
+    sys = sim.pysim.system
+    m = [sys.getParticleMass(i - 1)._value for i in 1:sys.getNumParticles()]
+    f = sim.pysim.context.getState(getForces=true).getForces(asNumpy=true)
+    f = f.value_in_unit(f.unit) |> permutedims
+    f ./= m'
+    f = vec(f)
+    @assert(!any(isnan.(f)))
+    f
 end
 
 function potential(sim::OpenMMSimulation, x)
     CUDA.reclaim()
-    sim = sim.pysim
     setcoords(sim, x)
-    v = sim.context.getState(getEnergy=true).getPotentialEnergy()
+    v = sim.pysim.context.getState(getEnergy=true).getPotentialEnergy()
     v = v.value_in_unit(v.unit)
 end
 
-
 ### PYTHON CODE