diff --git a/scripts/villin.jl b/scripts/villin.jl index 62a0851..b874bd9 100644 --- a/scripts/villin.jl +++ b/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 diff --git a/src/extrapolate.jl b/src/extrapolate.jl index d026ef4..5b1b534 100644 --- a/src/extrapolate.jl +++ b/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 diff --git a/src/iso2.jl b/src/iso2.jl index 5cf697b..3fb51c0 100644 --- a/src/iso2.jl +++ b/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) diff --git a/src/models.jl b/src/models.jl index ccd8c07..d5861e9 100644 --- a/src/models.jl +++ b/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 diff --git a/src/pairdists.jl b/src/pairdists.jl index 681c157..b09cc0f 100644 --- a/src/pairdists.jl +++ b/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 diff --git a/src/simulation.jl b/src/simulation.jl index 01d5f19..79e4a59 100644 --- a/src/simulation.jl +++ b/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 diff --git a/src/simulators/mopenmm.py b/src/simulators/mopenmm.py index a30f3b0..b8c0be1 100644 --- a/src/simulators/mopenmm.py +++ b/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': diff --git a/src/simulators/openmm.jl b/src/simulators/openmm.jl index c01fbde..fe9b185 100644 --- a/src/simulators/openmm.jl +++ b/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