From 867791ee0b1e27fdb071e6135f60e4e1108519e6 Mon Sep 17 00:00:00 2001 From: Sikorski Date: Wed, 20 Nov 2024 16:27:56 +0100 Subject: [PATCH] proper defaults for nonbondedForce, enhanced `savecoords`, allow passing of rigidWater and flexibleConstraints --- src/ISOKANN.jl | 1 + src/iso.jl | 20 +++++++++----------- src/plots.jl | 2 +- src/simulators/mopenmm.py | 11 ++++++----- src/simulators/openmm.jl | 11 +++++------ 5 files changed, 22 insertions(+), 23 deletions(-) diff --git a/src/ISOKANN.jl b/src/ISOKANN.jl index eb2ce13..2b402f4 100644 --- a/src/ISOKANN.jl +++ b/src/ISOKANN.jl @@ -77,6 +77,7 @@ export addcoords, addcoords!, resample_kde!, resample_kde export getxs, getys export exit_rates export load_trajectory, save_trajectory +export savecoords export atom_indices export localpdistinds, pdists, restricted_localpdistinds export data_from_trajectory, mergedata diff --git a/src/iso.jl b/src/iso.jl index 853934e..c8aa926 100644 --- a/src/iso.jl +++ b/src/iso.jl @@ -246,20 +246,18 @@ end simulationtime(iso::Iso) = simulationtime(iso.data) """ - savecoords(path::String, iso::Iso, inds=1:numobs(iso.data)) - -Save the coordinates of the specified observation indices from the data of of `iso` to the file `path`. - - savecoords(path::String, iso::Iso, coords::AbstractArray) + savecoords(path::String, iso::Iso, coords::AbstractMatrix=getcoords(iso.data); sorted=true, aligned=true) Save the coordinates of the specified matrix of coordinates to a file, using the molecule in `iso` as a template. +If `sorted` the sort the coordinates by their increasing χ value. If `align` then align each frame to the previous one. """ -function savecoords(path::String, iso::Iso, inds=1:numobs(iso.data)) - coords = getcoords(iso.data)[:,inds] - savecoords(path, iso, coords) -end - -function savecoords(path::String, iso::Iso, coords::AbstractMatrix) +function savecoords(path::String, iso::Iso, coords::AbstractMatrix=getcoords(iso.data); sorted=true, aligned=true) + if sorted + coords = coords[:, cpu(sortperm(dropdims(chicoords(iso, coords), dims=1)))] + end + if aligned + coords = aligntrajectory(coords) + end savecoords(path, iso.data.sim, coords) end diff --git a/src/plots.jl b/src/plots.jl index 10f8e08..c20185a 100644 --- a/src/plots.jl +++ b/src/plots.jl @@ -38,7 +38,7 @@ function plot_training(iso; subdata=nothing) !isnothing(subdata) && (data = subdata) - p1 = plot(losses[1:end], yaxis=:log, title="loss", label="trainloss", xlabel="iter") + p1 = plot(losses, yaxis=:log, title="loss", label="trainloss", xlabel="iter", ylabel="squared loss") for v in filter(x -> isa(x, ValidationLossLogger), iso.loggers) p1 = plot!(v.iters, v.losses, label="validation") diff --git a/src/simulators/mopenmm.py b/src/simulators/mopenmm.py index 945093f..804000e 100644 --- a/src/simulators/mopenmm.py +++ b/src/simulators/mopenmm.py @@ -46,7 +46,7 @@ def trajectory(sim, x0, stepsize, steps, saveevery, mmthreads, withmomenta): return trajectory # from the OpenMM documentation -def defaultsystem(pdb, ligand, forcefields, temp, friction, step, minimize, mmthreads, addwater=False, padding=1, ionicstrength=0, forcefield_kwargs={}): +def defaultsystem(pdb, ligand, forcefields, temp, friction, step, minimize, mmthreads, addwater=False, padding=1, ionicstrength=0, forcefield_kwargs={}, flexibleConstraints = False, rigidWater=True): pdb = PDBFile(pdb) if mmthreads == 'gpu': @@ -56,8 +56,6 @@ def defaultsystem(pdb, ligand, forcefields, temp, friction, step, minimize, mmth platform = Platform.getPlatformByName('CPU') platformargs = {'Threads': str(mmthreads)} - nonbondedMethod = CutoffNonPeriodic - if ligand != "": from openff.toolkit import Molecule from openmmforcefields.generators import SystemGenerator @@ -84,6 +82,9 @@ def defaultsystem(pdb, ligand, forcefields, temp, friction, step, minimize, mmth else: forcefield = ForceField(*forcefields) modeller = Modeller(pdb.topology, pdb.positions) + + nonbondedMethod = CutoffNonPeriodic if pdb.getTopology().getPeriodicBoxVectors() is None else CutoffPeriodic + if addwater: water_force_field = "amber/tip3p_standard.xml" forcefield = ForceField(*forcefields, water_force_field) @@ -96,8 +97,8 @@ def defaultsystem(pdb, ligand, forcefields, temp, friction, step, minimize, mmth system = forcefield.createSystem(modeller.topology, nonbondedMethod=nonbondedMethod, removeCMMotion=False, - #flexibleConstraints=True, - #rigidWater=False, + flexibleConstraints=flexibleConstraints, + rigidWater=rigidWater, **forcefield_kwargs) integrator = LangevinMiddleIntegrator(temp*kelvin, friction/picosecond, step*picoseconds) diff --git a/src/simulators/openmm.jl b/src/simulators/openmm.jl index 5d1e08f..7f1dc4b 100644 --- a/src/simulators/openmm.jl +++ b/src/simulators/openmm.jl @@ -102,16 +102,15 @@ function defaultsystem(; temp=310, # kelvin friction=1, # 1/picosecond step=0.002, # picoseconds - minimize=false, addwater=false, - padding=3, + minimize=addwater, + padding=1, ionicstrength=0.0, mmthreads=CUDA.functional() ? "gpu" : 1, forcefield_kwargs=Dict(), - features=nothing, - #kwargs... + kwargs... ) - @pycall py"defaultsystem"(pdb, ligand, forcefields, temp, friction, step, minimize; addwater, padding, ionicstrength, forcefield_kwargs, mmthreads)::PyObject + @pycall py"defaultsystem"(pdb, ligand, forcefields, temp, friction, step, minimize; addwater, padding, ionicstrength, forcefield_kwargs, mmthreads, kwargs...)::PyObject end @@ -153,7 +152,7 @@ end featurizer(sim::OpenMMSimulation) = featurizer(sim, get(sim.constructor, :features, nothing)) -featurizer(sim, ::Nothing) = natoms(sim) < 100 ? FeaturesAll() : error("No default featurizer specified") +featurizer(sim, ::Nothing) = natoms(sim) < 100 ? FeaturesAll() : error("No default featurizer specified. Specify any of FeaturesAll, FeaturesAtoms, FeaturesCoords, FeaturesPairs") featurizer(sim, atoms::Vector{Int}) = FeaturesAtoms(atoms) featurizer(sim, pairs::Vector{Tuple{Int,Int}}) = FeaturesPairs(pairs) featurizer(sim, features::Function) = features