ISOKANN.jl
Documentation for ISOKANN.jl
Main entry points
ISOKANN.SimulationData
— Typestruct SimulationData{S,D,C,F}
A struct combining a simulation with the simulated coordinates and corresponding ISOKANN trainingsdata
Fields
sim::S
: The simulation object.data::D
: The ISOKANN trainings data.coords::C
: The orginal coordinates of the simulations.featurizer::F
: A function mapping coordinates to ISOKANN features.
ISOKANN.OpenMM.OpenMMSimulation
— TypeOpenMMSimulation(; pdb, steps, ...)
OpenMMSimulation(; py, steps)
Constructs an OpenMM simulation object. Either use OpenMMSimulation(;py, steps)
where py
is the location of a .py python script creating a OpenMM simulation object or supply a .pdb file via
pdb` and the following parameters (see also defaultsystem):
Arguments
pdb::String
: Path to the PDB file.ligand::String
: Path to ligand file.forcefields::Vector{String}
: List of force field XML files.temp::Float64
: Temperature in Kelvin.friction::Float64
: Friction coefficient in 1/picosecond.step::Float64
: Integration step size in picoseconds.steps::Int
: Number of simulation steps.features
: Which features to use for learning the chi function. - A vector ofInt
denotes the indices of all atoms to compute the pairwise distances from. - A vector of CartesianIndex{2} computes the specific distances between the atom pairs. - A number denotes the radius below which all pairs of atoms will be used (computed only on the starting configuration) - Ifnothing
all pairwise distances are used.minimize::Bool
: Whether to perform energy minimization on first state.nthreads
: The number of threads to use for parallelization of multiple simulations.mmthreads
: The number of threads to use for each OpenMM simulation. Set to "gpu" to use the GPU platform.
Returns
OpenMMSimulation
: An OpenMMSimulation object.
ISOKANN.propagate
— Functionpropagate(sim::OpenMMSimulation, x0::AbstractMatrix, nk)
Propagates nk
replicas of the OpenMMSimulation sim
from the inintial states x0
.
Arguments
sim
: An OpenMMSimulation object.x0
: Matrix containing the initial states as columnsnk
: The number of replicas to create.
ISOKANN.Iso
— TypeIso(data; opt=NesterovRegularized(), model=defaultmodel(data), gpu=false, kwargs...)
Iso(sim::IsoSimulation; nx=100, nk=10, nd=1, kwargs...)
Convenience constructor which generates the SimulationData
from the simulation sim
and constructs the Iso object. See also Iso(data; kwargs...)
Arguments
sim::IsoSimulation
: TheIsoSimulation
object.nx::Int
: The number of starting points.nk::Int
: The number of koopman samples.nout::Int
: Dimension of the χ function.
ISOKANN.run!
— Functionrun!(iso::Iso, n=1, epochs=1)
Run the training process for the Iso model.
Arguments
iso::Iso
: The Iso model to train.n::Int
: The number of (outer) Koopman iterations.epochs::Int
: The number of (inner) epochs to train the model for each Koopman evaluation.
Public API
ISOKANN.Iso
— MethodIso(data; opt=NesterovRegularized(), model=defaultmodel(data), gpu=false, kwargs...)
ISOKANN.Iso
— MethodIso(sim::IsoSimulation; nx=100, nk=10, nd=1, kwargs...)
Convenience constructor which generates the SimulationData
from the simulation sim
and constructs the Iso object. See also Iso(data; kwargs...)
Arguments
sim::IsoSimulation
: TheIsoSimulation
object.nx::Int
: The number of starting points.nk::Int
: The number of koopman samples.nout::Int
: Dimension of the χ function.
ISOKANN.SimulationData
— MethodSimulationData(sim::IsoSimulation, nx::Int, nk::Int; ...)
SimulationData(sim::IsoSimulation, xs::AbstractMatrix, nk::Int; ...)
SimulationData(sim::IsoSimulation, (xs,ys); ...)
Generates SimulationData from a simulation with either
nx
initial points andnk
Koopman samplesxs
as initial points andnk
Koopman samplexs
as inintial points andys
as Koopman samples
ISOKANN.AdamRegularized
— FunctionAdam with L2 regularization. Note that this is different from AdamW (Adam+WeightDecay) (c.f. Decay vs L2 Reg.)
ISOKANN.data_from_trajectory
— Methoddata_from_trajectory(xs::AbstractMatrix; reverse=false) :: DataTuple
Generate the lag-1 data from the trajectory xs
. If reverse
is true, also take the time-reversed lag-1 data.
ISOKANN.laggedtrajectory
— Methodlaggedtrajectory(data::SimulationData, n) = laggedtrajectory(data.sim, n, x0=data.coords[1][:, end])
Simulate a trajectory comprising of n
simulations from the last point in data
ISOKANN.load_trajectory
— Methodload_trajectory(filename; top=nothing, kwargs...)
wrapper around Python's mdtraj.load()
. Returns a (3 * natom, nframes) shaped array.
ISOKANN.localpdistinds
— Methodlocalpdistinds(coords::AbstractMatrix, radius)
Given coords
of shape ( 3n x frames ) return the pairs of indices whose minimal distance along all frames is at least once lower then radius
ISOKANN.mergedata
— Methodmergedata(d1::SimulationData, d2::SimulationData)
Merge the data and features of d1
and d2
, keeping the simulation and features of d1
. Note that there is no check if simulation features agree.
ISOKANN.pairnet
— MethodFully connected neural network with layers
layers from n
to nout
dimensions. features
allows to pass a featurizer as preprocessor, activation
determines the activation function for each but the last layer lastactivation
can be used to modify the last layers activation function
ISOKANN.pdists
— Methodpdists(coords::AbstractArray, inds::Vector{<:Tuple})
Compute the pairwise distances between the particles specified by the tuples inds
over all frames in traj
. Assumes a column contains all 3n coordinates.
ISOKANN.reactionpath_minimum
— Functionreactionpath_minimum(iso::Iso, x0; steps=100)
Compute the reaction path by integrating ∇χ with orthogonal energy minimization.
Arguments
iso::Iso
: The isomer for which the reaction path minimum is to be computed.x0
: The starting point for the reaction path computation.steps=100
: The number of steps to take along the reaction path.
ISOKANN.reactionpath_ode
— Methodreactionpath_ode(iso, x0; steps=101, extrapolate=0, orth=0.01, solver=OrdinaryDiffEq.Tsit5(), dt=1e-3, kwargs...)
Compute the reaction path by integrating ∇χ as well as orth
* F orthogonal to ∇χ where F is the original force field.
Arguments
iso::Iso
: The isomer for which the reaction path minimum is to be computed.x0
: The starting point for the reaction path computation.steps=100
: The number of steps to take along the reaction path.minimize=false
: Whether to minimize the orthogonal to ∇χ before integration.extrapolate=0
: How fast to extrapolate beyond χ 0 and 1.orth=0.01
: The weight of the orthogonal force field.solver=OrdinaryDiffEq.Tsit5()
: The solver to use for the ODE integration.dt=1e-3
: The initial time step for the ODE integration.
ISOKANN.reactive_path
— Methodreactive_path(xi::AbstractVector, coords::AbstractMatrix; sigma, maxjump=1, method=QuantilePath(0.05), normalize=false, sortincreasing=true)
Find the maximum likelihood path (under the model of brownion motion with noise sigma
) through coords
with times xi
. Supports either CPU or GPU arrays.
Arguments
coords
: (ndim x npoints) matrix of coordinates.xi
: time coordinate of the npoints pointssigma
: spatial noise strength of the model.maxjump
: upper bound to the jump in timexi
along the path.method
: eitherFromToPath
,QuantilePath
,FullPath
orMaxPath
, specifying the end points of the pathnormalize
: whether to normalize allcoords
firstsortincreasing
: return the path from lower to higherxi
values
ISOKANN.restricted_localpdistinds
— Methodrestricted_localpdistinds(coords, radius, atoms)
Like localdists
, but consider only the atoms with index in atoms
ISOKANN.run!
— Functionrun!(iso::Iso, n=1, epochs=1)
Run the training process for the Iso model.
Arguments
iso::Iso
: The Iso model to train.n::Int
: The number of (outer) Koopman iterations.epochs::Int
: The number of (inner) epochs to train the model for each Koopman evaluation.
ISOKANN.runadaptive!
— Methodrunadaptive!(iso; generations=1, nx=10, iter=100, cutoff=Inf)
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.
ISOKANN.save_reactive_path
— Functionsave_reactive_path(iso::Iso,
coords::AbstractMatrix=getcoords(iso.data) |> cpu;
sigma=1,
maxjump=1,
out="out/reactive_path.pdb",
source=pdbfile(iso.data),
kwargs...)
Extract and save the reactive path of a given iso
.
Computes the maximum likelihood path with parameter sigma
along the given data points, aligns it and saves it to the out
path.
See also reactive_path
.
Arguments
iso::Iso
: The Iso for which the reactive path is computed.out="out/reactive_path.pdb"
: The output file path for saving the reactive path.source
: The source .pdb file providing the topologykwargs...
: additional parameters passed toreactive_path
.
Returns
ids
: The IDs of the reactive path.
ISOKANN.save_trajectory
— Methodsave_trajectory(filename, coords::AbstractMatrix; top::String)
save the trajectory given in coords
to filename
with the topology provided by the file top
using mdtraj.
ISOKANN.trajectorydata_bursts
— Methodtrajectorydata_bursts(sim::IsoSimulation, steps, nk; kwargs...)
Simulate a single long trajectory of steps
times the lagtime and start nk
burst trajectories at each step for the Koopman samples.
x0–-x––x–- / | / | y y y y
ISOKANN.trajectorydata_linear
— Methodtrajectorydata_linear(sim::IsoSimulation, steps; reverse=false, kwargs...)
Simulate a single long trajectory of steps
times the lagtime and use this "chain" to generate the corresponding ISOKANN data. If reverse
is true, also add the time-reversed transitions
x (<)–> x (<)–> x
Internal API
ISOKANN.DataTuple
— TypeDataTuple = Tuple{Matrix{T},Array{T,3}} where {T<:Number}
We represent data as a tuple of xs and ys.
xs is a matrix of size (d, n) where d is the dimension of the system and n the number of samples. ys is a tensor of size (d, k, n) where k is the number of koopman samples.
ISOKANN.IsoSimulation
— Typeabstract type IsoSimulation
Abstract type representing an IsoSimulation. Should implement the methods getcoords
, propagate
, dim
ISOKANN.Stabilize2
— TypeTransformStabilize(transform, last=nothing)
Wraps another transform and permutes its target to match the previous target
Currently we also have the stablilization (wrt to the model though) inside each Transform. TODO: Decide which to keep
ISOKANN.TransformISA
— TypeTransformISA(permute)
Compute the target via the inner simplex algorithm (without feasiblization routine). permute
specifies whether to apply the stabilizing permutation
ISOKANN.TransformPseudoInv
— TypeTransformPseudoInv(normalize, direct, eigenvecs, permute)
Compute the target by approximately inverting the action of K with the Moore-Penrose pseudoinverse.
If direct==true
solve chi * pinv(K(chi))
, otherwise inv(K(chi) * pinv(chi)))
. eigenvecs
specifies whether to use the eigenvectors of the schur matrix. normalize
specifies whether to renormalize the resulting target vectors. permute
specifies whether to permute the target for stability.
ISOKANN.TransformShiftscale
— TypeTransformShiftscale()
Classical 1D shift-scale (ISOKANN 1)
ISOKANN._pickclosestloop
— Methodscales with n=length(hs)
ISOKANN.adddata
— Methodadddata(data::D, model, sim, ny, lastn=1_000_000)::D
Generate new data for ISOKANN by adaptive subsampling using the chi-stratified/-uniform method.
- Adaptively subsample
ny
points fromdata
uniformly along theirmodel
values. - propagate according to the simulation
model
. - return the newly obtained data concatenated to the input data
The subsamples are taken only from the lastn
last datapoints in data
.
Examples
julia> (xs, ys) = adddata((xs,ys), chi, mollysim)
ISOKANN.adddata
— Methodadddata(d::SimulationData, model, n)
χ-stratified subsampling. Select n samples amongst the provided ys/koopman points of d
such that their χ-value according to model
is approximately uniformly distributed and propagate them. Returns a new SimulationData
which has the new data appended.
ISOKANN.addextrapolates!
— Methodaddextrapolates!(iso, n, stepsize=0.01, steps=10)
Sample new data starting points obtained by extrapolating the chi function beyond the current extrema and attach it to the iso
objects data.
Samples n
points at the lower and upper end each, resulting in 2n new points. step
is the magnitude of chi-value-change per step and
steps`` is the number of steps to take. E.g. 10 steps of stepsize 0.01 result in a change in chi of about 0.1.
The obtained data is filtered such that unstable simulations should be removed, which may result in less then 2n points being added.
ISOKANN.bootstrap
— Methodbootstrap(sim, nx, ny) :: DataTuple
compute initial data by propagating the molecules initial state to obtain the xs and propagating them further for the ys
ISOKANN.centercoords
— Methodcentercoords any given states by shifting their individual 3d mean to the origin
ISOKANN.chi_exit_rate
— Methodcompute the chi exit rate as per Ernst, Weber (2017), chap. 3.3
ISOKANN.energyminimization_chilevel
— Methodenergyminimization_chilevel(iso, x0; f_tol=1e-3, alphaguess=1e-5, iterations=20, show_trace=false, skipwater=false, algorithm=Optim.GradientDescent, xtol=nothing)
Local energy minimization on the current levelset of the chi function
ISOKANN.exportdata
— Functionexportdata(data::AbstractArray, model, sys, path="out/data.pdb")
Export data to a PDB file.
This function takes an AbstractArray data
, sorts it according to the model
evaluation, removes duplicates, transforms it to standard form and saves it as a PDB file to path
.
ISOKANN.extrapolate
— Functionextrapolate(iso, n, stepsize=0.1, steps=1, minimize=true)
Take the n
most extreme points of the chi-function of the iso
object and 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.
ISOKANN.fixperm
— Methodfixperm(new, old)
Permutes the rows of new
such as to minimize L1 distance to old
.
Arguments
new
: The data to match to the reference data.old
: The reference data.
ISOKANN.flatpairdists
— Functionflatpairdists(x)
Assumes each col of x to be a flattened representation of multiple 3d coords. Returns the flattened pairwise distances as columns.
ISOKANN.flattenfirst
— Methodcollapse the first and second dimension of the array A
into the first dimension
ISOKANN.growmodel
— MethodGiven a model and return a copy with its last layer replaced with given output dimension n
ISOKANN.inputdim
— Methodobtain the input dimension of a Flux model
ISOKANN.kabsch
— Methodcompute R such that R*p is closest to q
ISOKANN.koopman
— Methodevluation of koopman by shiftscale(mean(model(data))) on the data
ISOKANN.load
— Methodload(path::String, iso::Iso)
Load the Iso object from a JLD2 file Note that it will be loaded to the CPU, even if it was saved on the GPU. An OpenMMSimulation will be reconstructed anew from the saved pdb file.
ISOKANN.model_with_opt
— Functionconvenience wrapper returning the provided model with the default AdamW optimiser
ISOKANN.outputdim
— MethodObtain the output dimension of a Flux model
ISOKANN.pairdistfeatures
— Methodpairdistfeatures(inds::AbstractVector)
Returns a featurizer function which computes the pairwise distances between the particles specified by inds
ISOKANN.pickclosest_sort
— Methodpickclosest(haystack, needles)
Return the indices into haystack which lie closest to needles
without duplicates by removing haystack candidates after a match. Note that this is not invariant under pertubations of needles
scales with n log(n) m where n=length(haystack), m=length(needles)
ISOKANN.plotatoms!
— Functionscatter plot of all first "O" atoms of the starting points xs
as well as the "O" atoms from the koopman samples to the first point from ys
ISOKANN.plotlossdata
— Functioncombined plot of loss graph and current atoms
ISOKANN.reactionforce
— Functionreactionforce(iso, sim, x, direction, orth=1)
Compute the vector f
with colinear component to dchi/dx such that dchi/dx * f = 1 and orth*forcefield in the orthogonal space
ISOKANN.resample_kde
— Methodresample_kde(xs, ys, n; kwargs...)
Return n
indices of ys
such that the corresponding points "fill the gaps" in the KDE of xs
. For possible kwargs
see kde_needles
.
ISOKANN.save
— Methodsave(path::String, iso::Iso)
Save the complete Iso object to a JLD2 file
ISOKANN.savecoords
— Functionsavecoords(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)
Save the coordinates of the specified matrix of coordinates to a file, using the molecule in iso
as a template.
ISOKANN.saveextrema
— Methodsaveextrema(path::String, iso::Iso)
Save the two extermal configurations (metastabilities) to the file path
.
ISOKANN.scatter_chifix
— Methodfixed point plot, i.e. x vs model(x)
ISOKANN.selectrows
— Methodgiven an array of arbitrary shape, select the rows inds
in the first dimension
ISOKANN.shiftscale
— Methodempirical shift-scale operation
ISOKANN.simulationtime
— Methodsimulationtime(iso::Iso)
print and return the total simulation time contained in the data of iso
in nanoseconds.
ISOKANN.sqpairdist
— Methodsqpairdist(x::AbstractArray)
Compute the squared pairwise distances between the columns of x
. If x
has 3 dimensions, the computation is batched along the 3rd dimension.
ISOKANN.subsample
— Methodsubsample(model, data::Array, n) :: Matrix
subsample(model, data::Tuple, n) :: Tuple
Subsample n
points of data
uniformly in model
. If model
returns multiple values per sample, subsample along each dimension.
ISOKANN.subsample
— Methodsubsample(data, nx)
return a random subsample of nx
points from data
ISOKANN.subsample_inds
— Methodsubsample_inds(model, xs, n) :: Vector{Int}
Returns n
indices of xs
such that model(xs[inds])
is approximately uniformly distributed.
ISOKANN.subsample_uniformgrid
— Methodsubsbample_uniformgrid(ys, n) -> is
given a list of values ys
, return n
indices
issuch that
ys[is]` are approximately uniform by picking the closest points to a randomly perturbed grid in [0,1].
ISOKANN.trajectory
— Methodtrajectory(l::AbstractLangevin; T=lagtime(l), x0=randx0(l), save_start=false, saveat=lagtime(l), dt=dt(l)) generate a trajectory of length T
, starting at x0 with stepsize dt
, saving the output every saveat
time.
ISOKANN.writechemfile
— Methodwritechemfile(filename, data::Array{<:Any,2}; source)
Save the coordinates in data
to filename
with source
as template using the Chemfiles library
ISOKANN.OpenMM.integrate_langevin
— Functionintegrate_langevin(sim::OpenMMSimulation, x0=getcoords(sim); steps=steps(sim), F_ext::Union{Function,Nothing}=nothing, saveevery::Union{Int, nothing}=nothing)
Integrate the Langevin equations with a Euler-Maruyama scheme, allowing for external forces.
- Fext: An additional force perturbation. It is expected to have the form Fext(F, x) and mutating the provided force F.
- saveevery: If
nothing
, returns just the last point, otherwise returns an array saving everysaveevery
frame.
ISOKANN.OpenMM.set_random_velocities!
— Methodmutates the state in sim
ISOKANN.propagate
— Methodpropagate(sim::OpenMMSimulation, x0::AbstractMatrix, nk)
Propagates nk
replicas of the OpenMMSimulation sim
from the inintial states x0
.
Arguments
sim
: An OpenMMSimulation object.x0
: Matrix containing the initial states as columnsnk
: The number of replicas to create.
ISOKANN.randx0
— Methodgenerate n
random inintial points for the simulation mm
ISOKANN.savecoords
— Methodsavecoords(path, sim::OpenMMSimulation, coords::AbstractArray{T})
Save the given coordinates
in a .pdb file using OpenMM