rework the langevin interface

src/ISOKANN.jl

@@ -79,7 +79,6 @@ include("simulators/langevin.jl")  # for the simulators
 include("isotarget.jl")
 include("iso2.jl")
 include("plots.jl")           # visualizations
-include("simulators/potentials.jl")
 
 include("simulators/openmm.jl")
 

src/simulation.jl

@@ -195,12 +195,12 @@ function datasize((xs, ys)::Tuple)
     return size(xs), size(ys)
 end
 
-function trajectorydata(sim::IsoSimulation, steps; reverse=false, kwargs...)
+function trajectorydata_linear(sim::IsoSimulation, steps; reverse=false, kwargs...)
     xs = laggedtrajectory(sim, steps)
     SimulationData(sim, data_from_trajectory(xs; reverse), kwargs...)
 end
 
-function trajectoryburstdata(sim, steps, nk; kwargs)
+function trajectorydata_bursts(sim, steps, nk; kwargs)
     xs = laggedtrajectory(sim, steps)
     ys = propagate(sim, xs, nk)
     SimulationData(sim, ys, kwargs...)

src/simulators/langevin.jl

@@ -1,14 +1,17 @@
-using Parameters
-import StochasticDiffEq  # does this incur invalidations?
+#using Parameters
+using StaticArrays: @SVector
+using StochasticDiffEq: StochasticDiffEq
 import ForwardDiff
 
+# Abstract type defining the Overdamped Langevin dynamics
+# mandatory interface methods: potential, sigma, dim, lagtime, dt
 abstract type AbstractLangevin <: IsoSimulation end
-# interface methods: potential(l), sigma(l), dim(l)
 
 featurizer(::AbstractLangevin) = identity
-lagtime(l::AbstractLangevin) = tmax(l)
+integrator(::AbstractLangevin) = StochasticDiffEq.EM()
+defaultmodel(l::AbstractLangevin; nout, kwargs...) = smallnet(dim(l), nout)
 
-function SDEProblem(l::AbstractLangevin, x0=randx0(l), T=tmax(l); dt=dt(l), alg=integrator(l), kwargs...)
+function SDEProblem(l::AbstractLangevin, x0=randx0(l), T=lagtime(l); dt=dt(l), alg=integrator(l), kwargs...)
     drift(x,p,t) = force(l, x)
     noise(x,p,t) = sigma(l, x)
     StochasticDiffEq.SDEProblem(drift, noise, x0, T, alg=alg, dt=dt; kwargs...)
@@ -18,71 +21,115 @@ function force(l::AbstractLangevin, x)
     - ForwardDiff.gradient(x->potential(l, x), x)
 end
 
-propagate(l::AbstractLangevin, x0::CuArray, ny) = gpu(propagate(l, cpu(x0), ny))
-
 function propagate(l::AbstractLangevin, x0::AbstractMatrix, ny)
     dim, nx = size(x0)
     ys = zeros(dim, ny, nx)
     Threads.@threads for (i, j) in [(i, j) for j in 1:ny, i in 1:nx]
-        ys[:, j, i] = solve_end(l; u0=x0[:, i])
+        ys[:, j, i] = trajectory(l; x0, saveat=lagtime(l))[:, end]
     end
     return ys
 end
 
-"""  trajectory(sim, nx)
-generate a trajectory of length `nx` from the simulation `sim`"""
-function trajectory(sim, nx)
-    siml = deepcopy(sim)
-    logevery = round(Int, sim.T / sim.dt)
-    siml.T = sim.T * nx
-    xs = solve(siml; logevery=logevery)
-    return xs
+propagate(l::AbstractLangevin, x0::CuArray, ny) = gpu(propagate(l, cpu(x0), ny))
+
+"""  trajectory(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. """
+function trajectory(l::AbstractLangevin; T=lagtime(l), x0=randx0(l), save_start=false, saveat=dt(l), dt=dt(l))
+    sde = SDEProblem(l, x0, T)
+    sol = StochasticDiffEq.solve(sde; saveat, save_start)
+    xs = reduce(hcat, sol.u)
+    return xs::Matrix
 end
 
-function solve_end(l::AbstractLangevin; u0)
-    StochasticDiffEq.solve(SDEProblem(l, u0))[:, end]
+laggedtrajectory(l::AbstractLangevin, steps) = trajectory(l, T=lagtime(l) * steps, saveat=lagtime(l))
+
+# helper functions to generate random initial data
+
+randx0(l::AbstractLangevin) = randx0(l, 1) |> vec
+
+function randx0(l::AbstractLangevin, n)
+    s = support(l)
+    x0 = rand(size(s, 1), n) .* (s[:, 2] .- s[:, 1]) .+ s[:, 1]
+    return x0
 end
 
+supportbox(l, x::Number) = supportbox(l, [-x, x])
+supportbox(l, x::Vector) = repeat(x', outer=dim(l))
+supportbox(l, x::Matrix) = x
+
+# dispatch potential for higher dimensional arrays
+potential(l::AbstractLangevin, x::AbstractArray) = mapslices(x -> potential(l, x), x; dims=1)
+
+### Concrete implementations of the AbstractLangevin type
+
 ##  Generic Diffusion in a potential
 @kwdef struct Diffusion{T} <: AbstractLangevin
     potential::T
     dim::Int64=1
-    sigma::Vector{Float64} = [1.0]
+    sigma::Union{Number,Vector,Matrix} = 1.0
+    dt::Float64 = 0.01
+    lagtime::Float64 = 1.0
+    support::Union{Number,Vector,Matrix} = 1.0
 end
 
-tmax(d::AbstractLangevin) = tmax(d.potential)
-tmax(d) = 1.0
-
-integrator(d::AbstractLangevin) = StochasticDiffEq.SROCK2()
+potential(d::Diffusion, x::Vector) = d.potential(x)
+dim(l::Diffusion) = l.dim
+sigma(l::Diffusion, x) = l.sigma
+dt(d::Diffusion) = d.dt
+lagtime(d::Diffusion) = d.lagtime
+support(l::Diffusion) = supportbox(l, l.support)
 
-dt(d::AbstractLangevin) = dt(d.potential)
-dt(d) = 0.01
+## Double- and Triplewell potentials implemented using the generic `Diffusion`
 
-potential(d::Diffusion, x) = d.potential(x)
-sigma(l::Diffusion, x) = l.sigma
-dim(l::Diffusion) = l.dim
-support(l::Diffusion) = repeat([-1.5 1.5], outer=[dim(l)]) :: Matrix{Float64}
+Doublewell(; kwargs...) = Diffusion(;
+    potential=doublewell,
+    support=1.5,
+    kwargs...)
 
-randx0(l::Diffusion, n) = reduce(hcat, [randx0(l) for i in 1:n])
-function randx0(l::Diffusion)
-    s = support(l)
-    x0 = rand(size(s, 1)) .* (s[:,2] .- s[:,1]) .+ s[:,1]
-    return x0
-end
+doublewell(x) = ((x[1])^2 - 1)^2
 
-doublewell(x) = ((x[1])^2 - 1) ^ 2
 
-Doublewell(;kwargs...) = Diffusion(;potential=doublewell, kwargs...)
+Triplewell(; kwargs...) = Diffusion(;
+    potential=triplewell,
+    dim=2,
+    sigma=[1.0, 1.0],
+    support=[-2 2; -1.5 2.5],
+    kwargs...)
 
+triplewell(x) = triplewell(x...)
 triplewell(x,y) = (3/4 * exp(-x^2 - (y-1/3)^2)
             - 3/4 * exp(-x^2 - (y-5/3)^2)
             - 5/4 * exp(-(x-1)^2 - y^2)
             - 5/4 * exp(-(x+1)^2 - y^2)
             + 1/20 * x^4 + 1/20 * (y-1/3)^4)
 
-triplewell(x) = triplewell(x...)
 
+MuellerBrown(; kwargs...) = Diffusion(;
+    potential=mueller_brown,
+    dim=2,
+    sigma=1.0,
+    support=[-1.5 1; -0.5 2],
+    dt=0.001,
+    kwargs...)
+
+mueller_brown(x::AbstractVector) = mueller_brown(x...)
+function mueller_brown(x, y)
+    -200 * exp(-1 * (x - 1)^2 + -10 * (y)^2) -
+    100 * exp(-1 * (x)^2 + -10 * (y - 0.5)^2) -
+    170 * exp(-6.5 * (x + 0.5)^2 + 11 * (x + 0.5) * (y - 1.5) + -6.5 * (y - 1.5)^2) +
+    15 * exp(0.7 * (x + 1)^2 + 0.6 * (x + 1) * (y - 1) + 0.7 * (y - 1)^2)
+end
+
+#=
+## Example of a custom type for a specific process
 
-Triplewell() = Diffusion(; potential=triplewell, dim=2, sigma=[1.0, 1.0])
+struct CustomMuellerBrown <: AbstractLangevin end
 
-defaultmodel(sim::AbstractLangevin; nout, kwargs...) = smallnet(dim(sim), nout)
+integrator(m::MuellerBrown) = StochasticDiffEq.EM()
+lagtime(m::MuellerBrown) = 0.01
+dt(m::MuellerBrown) = 0.0001
+dim(::MuellerBrown) = 2
+potential(::MuellerBrown, x) = mueller_brown(x)
+sigma(l::MuellerBrown, x) = 2.0
+support(l::MuellerBrown) = [-1.5 1; -0.5 2]
+=#