update docs

docs/src/introduction.md

@@ -0,0 +1,32 @@
+# Introduction
+
+This package provides the core ISOKANN algorithm as well as some wrappers and convenience routines to work with different kind of simulations and data.
+
+The core ISOKANN algorithm is accessed by the `Iso2` type,
+which holds the neural network, optimizer, ISOKANN parameters and training data.
+
+You can construct it by passing a tuple of `(xs, ys)` of arrays as input data. Here `xs` is a matrix where the columns are starting points of trajectories and `ys` is a 3 dimensional array where `ys[d,k,n]` is the `d`-th coordinate of the `k`-th Koopman-replica of the `n`-th trajectory.
+
+To start training the neural network simply call the `run!` function passing the `Iso2` object and the number of ISOKANN iterations.
+The resulting \chi values can be obtained via the `chis` method
+
+```julia
+iso=Iso2((rand(3,100), rand(3,10,100)))
+run!(iso)
+chis(iso)
+```
+
+We also supply some basic simulations which can generate the data, e.g. [`Doublewell`](@ref), [`MuellerBrown`], [`Diffusion`], [`MollySimulation`] and [`OpenMMSimulation`].
+You can use the [`isodata`] function to sample data for ISOKANN.
+
+```julia
+sim = Doublewell()
+data = isodata(sim, 100, 20)
+iso = Iso2(data)
+```
+
+We also provide different type of wrappers to load simulations [`vgv`] or generate data from trajectories [`IsoMu`].
+
+Experimental support for adaptive sampling is provided by [`run(iso, sim; ny)`].
+
+

src/ISOKANN.jl

@@ -48,7 +48,6 @@ export AdamRegularized, pairnet#, Adam
 export plot_training, scatter_ramachandran
 export reactionpath
 export cpu, gpu
-export iso2
 export Iso2
 export Doublewell, Triplewell, MuellerBrown
 export chis

src/IsoMu/IsoMu.jl

@@ -5,9 +5,11 @@ using Plots
 
 using DataFrames: DataFrame
 using ISOKANN
-import ISOKANN: reactive_path, save_reactive_path, aligntrajectory
+using ISOKANN: plot_reactive_path, writechemfile, aligntrajectory
 using Distances: pairwise, Euclidean
 
+import ISOKANN: reactive_path, save_reactive_path
+
 #using FileIO
 import BioStructures
 import Flux
@@ -16,7 +18,7 @@ import Optimisers
 import Graphs
 import Molly
 
-export DataLink, train!, save_reactive_path, isokann, meanvelocity, adjust!, gpu!, paperplot, paperplot
+export DataLink, isokann, train!, save_reactive_path, paperplot, adjust!, gpu!
 
 include("datalink.jl")
 include("isomu.jl")

src/IsoMu/isomu.jl

@@ -6,7 +6,7 @@
 Combine ISOKANN with a DataLink
 
 # Fields
-- `data::DataLink`: reference to the data 
+- `data::DataLink`: reference to the data
 - `learnrate::Float64=0.001`: learnrate for ADAM
 - `regularization::Float64=0.0001`: regularization of the NN
 - `hidden::Vector{Int}=[300, 100, 10]`: Size of hidden layers
@@ -65,7 +65,7 @@ pdbfile(mu::IMu) = pdbfile(mu.data)
 numobs(mu::IMu) = size(xs(mu), 2)
 
 
-""" 
+"""
     reactive_path(mu::IMu; kwargs...) where {T} -> (Vector{Int}, Matrix)
 
 Compute the shortest Onsager-Machlup path through the data with the χ value as time.
@@ -77,14 +77,15 @@ Compute the shortest Onsager-Machlup path through the data with the χ value as
 - `sigma::Float64=0.1`: the temperature for the allowed trajectory
 - `window::Interval`: which coordinates to consider for the path
 - `method`: Either one of
-            `QuantilePath(quant)` for a top-bot quartile path 
+            `QuantilePath(quant)` for a top-bot quartile path
             `FromToPath(from::Int, to::Int)` specific frame numbers
             `FullPath()` from start to end
 
 # Returns
-- `inds::Vector`: indices of the selected path 
+- `inds::Vector`: indices of the selected path
 - `path::Matrix`: aligned coordinates of the path
 """
+# TODO: this should dispatch to ISOKANN.reactive_path
 function reactive_path(
     mu::IMu;
     sigma=0.1,
@@ -117,7 +118,7 @@ isincreasing(x) = sum(diff(x) .> 0) > length(x) / 2
     save_reactive_path(mu::IMu; out=joinpath(outdir(mu.data), "reactive_path.pdb"),
     kwargs...) -> Vector{Int}
 
-Compute the reactive path for the IMu object and save it as pdb 
+Compute the reactive path for the IMu object and save it as pdb
 (defaulting to the outdir of mu) see also `reactive_path` for further arguments
 
 # Arguments
@@ -130,6 +131,7 @@ Compute the reactive path for the IMu object and save it as pdb
 # Returns
 - `inds::Vector{Int}`: the indices of the selected path
 """
+# TODO: this should dispatch to ISOKANN.reactive_path
 function save_reactive_path(mu::IMu;
     out=joinpath(outdir(mu.data), "reactive_path.pdb"),
     kwargs...)
@@ -158,7 +160,7 @@ function meanvelocity(mu::IMu)
     title!(mu.data.dir)
 end
 
-function adjust!(mu::IMu, args...; kwargs...)
+function Flux.adjust!(mu::IMu, args...; kwargs...)
     mu.iso.opt isa NamedTuple || error("have to train! model once before tuning parameters")
     Flux.adjust!(mu.iso.opt, args...; kwargs...)
 end

src/reactionpath2.jl

@@ -97,6 +97,7 @@ end
 
 
 # TODO: check how this works with IsoMu
+# TODO: scaling by 10 shouldnt belong here
 function save_reactive_path(iso::Iso2, coords::AbstractMatrix;
     sigma=1, out="out/reactive_path.pdb", source)
     chi = chis(iso) |> vec |> cpu

test/runtests.jl

@@ -5,7 +5,7 @@ using Test
 @testset "ISOKANN.jl" verbose = true begin
 
     @testset "IsoRun CPU" begin
-        iso = ISOKANN.IsoRun(nd=10) |> run!
+        iso = IsoRun(nd=10) |> run!
         @test true
     end
 
@@ -14,7 +14,7 @@ using Test
         using CUDA
         if CUDA.functional()
             CUDA.allowscalar(false)
-            iso = ISOKANN.IsoRun(nd=10) |> gpu |> run!
+            iso = IsoRun(nd=10) |> gpu |> run!
             @test true
         else
             @info "No functional GPU found. Marking GPU test as broken."
@@ -39,8 +39,8 @@ using Test
 
         @testset "Iso2 Dialanine with adaptive sampling" begin
             sim = MollyLangevin(sys=PDB_ACEMD())
-            model = ISOKANN.pairnet(484, features=ISOKANN.flatpairdists, nout=3)
-            opt = Flux.setup(Flux.AdamW(1e-3, (0.9, 0.999), 1e-4), model)
+            model = pairnet(484, features=ISOKANN.flatpairdists, nout=3)
+            opt = Flux.AdamW(1e-3, (0.9, 0.999), 1e-4)
 
             iso = Iso2(sim; nx=100, nk=10, nd=1, model, opt)
             iso = run!(iso, sim, 10, 100, ny=10)