proper defaults for nonbondedForce, enhanced savecoords, allow pass…

…ing of rigidWater and flexibleConstraints

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

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
 

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")

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)

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