Skip to content

Commit

Permalink
partially resolve the force problem in extrapolation
Browse files Browse the repository at this point in the history
  • Loading branch information
axsk committed May 15, 2024
1 parent 6068c52 commit 5d68e9a
Show file tree
Hide file tree
Showing 3 changed files with 32 additions and 12 deletions.
10 changes: 7 additions & 3 deletions scripts/villin.jl
Original file line number Diff line number Diff line change
Expand Up @@ -133,10 +133,14 @@ for i in 1:generations

pdblen = ISOKANN.readchemfile(iso.data.sim.pdb) |> length
save_reactive_path(iso, ISOKANN.getcoords(iso.data)[1:pdblen, :] |> cpu;
out="$path/villin_fold_$(simtime)ps.pdb", sigma, maxjump)
out="$path/snapshot/villin_fold_$(simtime)ps.pdb", sigma, maxjump)
#ISOKANN.savecoords("$path/data.pdb", iso)
ISOKANN.Plots.savefig(plot_training(iso), "$path/villin_fold_$(simtime)ps.png")
println("\n status: $path/villin_fold_$(simtime)ps.png \n")
ISOKANN.Plots.savefig(plot_training(iso), "$path/snapshot/villin_fold_$(simtime)ps.png")


symlink("$path/snapshot/villin_fold_$(simtime)ps.pdb", "$path/path.pdb")
symlink("$path/snapshot/villin_fold_$(simtime)ps.png", "$path/training.png")
println("status: $path/ training.png")
ISOKANN.save("$path/iso.jld2", iso)
catch e
@show e
Expand Down
29 changes: 22 additions & 7 deletions src/extrapolate.jl
Original file line number Diff line number Diff line change
Expand Up @@ -79,21 +79,33 @@ end

global trace = []

function energyminimization_chilevel(iso, x0; f_tol=1e-3, alphaguess=1e-5, iterations=20, show_trace=false, skipwater=true)
function energyminimization_chilevel(iso, x0; f_tol=1e-3, alphaguess=1e-5, iterations=20, show_trace=false, skipwater=false, algorithm=Optim.GradientDescent)
sim = iso.data.sim
x = copy(x0) .|> Float64

chi(x) = myonly(chicoords(iso, x))
chilevel = Levelset(chi, chi(x0))
manifold = Levelset(chi, chi(x0))

U(x) = OpenMM.potential(sim, x)
dU(x) = (f = -OpenMM.force(sim, x); (skipwater && zerowater!(sim, f)); f)

global trace = [x0]
U(x) = begin
push!(trace, x)
@show OpenMM.potential(sim, x)
end
dU(x) = begin
push!(trace, x)
f = -OpenMM.force(sim, x)
(skipwater && zerowater!(sim, f))
f
end


linesearch = Optim.LineSearches.HagerZhang(alphamax=alphaguess)
alg = Optim.LBFGS(; linesearch, alphaguess, manifold=chilevel)
alg = algorithm(; linesearch, alphaguess, manifold)


o = Optim.optimize(U, dU, x, alg, 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
return o.minimizer
end

Expand All @@ -112,13 +124,16 @@ struct Levelset{F,T} <: Optim.Manifold
end

function Optim.project_tangent!(M::Levelset, g, x)
replace!(g, NaN => 0)
@assert !any(isnan.(g))
@assert !any(isnan.(x))
#replace!(g, NaN => 0)
u = Zygote.gradient(M.f, x) |> only
u ./= norm(u)
g .-= dot(g, u) * u
end

function Optim.retract!(M::Levelset, x)
@assert !any(isnan.(x))
g = Zygote.withgradient(M.f, x)
u = g.grad |> only
h = M.target - g.val
Expand Down
5 changes: 3 additions & 2 deletions src/simulators/openmm.jl
Original file line number Diff line number Diff line change
Expand Up @@ -221,10 +221,11 @@ function force(sim::OpenMMSimulation, x)
CUDA.reclaim()
setcoords(sim, x)
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'
#m = [sys.getParticleMass(i - 1)._value for i in 1:sys.getNumParticles()]
#f ./= m'
f = vec(f)
@assert(!any(isnan.(f)))
f
Expand Down

0 comments on commit 5d68e9a

Please sign in to comment.