style: refactor tpt (#161)

docs/src/examples/duffing_TPT.md

@@ -12,10 +12,12 @@ using OrdinaryDiffEq, DelaunayTriangulation, Contour
 ## System
 
 The Duffing oscillator is a simple model for a nonlinear oscillator with a double-well potential. The equation of motion for the Duffing oscillator under additive Gaussian noise is given by:
-\begin{equation}
+
+```math
 \dot{x} = p, \\
 \dot{p} = -\gamma p - \nabla U + \sqrt{\frac{2\gamma}{\beta}} \dot{W},
-\end{equation}
+```
+
 with the potential energy $U(x) = \frac{1}{4}x^4 - \frac{1}{2}x^2$ and the kinetic energy $K(p) = p^2/2$. The parameters $\gamma$ and $\beta=1/k_b T$ control the strength of the dissipation and noise, respectively. $W$ is a Wiener process, and the noise term is scaled by $\sqrt{2\gamma/\beta}$ to ensure the correct temperature scaling for a Langevin type system defined by the Hamiltonian $H$.
 
 ```@example TPT
@@ -47,7 +49,7 @@ langevin_sys = Langevin(Hamiltonian, divfree, KE, gamma, beta)
 
 ## Phase space mesh
 
-We can easily evaluate and visualise the Hamiltonian and equally spaced grid in phase space.
+We can easily evaluate and visualize the Hamiltonian and equally spaced grid in phase space.
 
 ```@example TPT
 nx, ny = 41, 41
@@ -73,7 +75,7 @@ streamplot!(v, -2..2, -2..2, linewidth = 0.5, colormap = [:black, :black], grids
 fig
 ```
 
-The overdamped undriven duffing oscillator, is autonomous and respect detailed balance. As such the maximum likelihood path is the path that is parallel to drift and can be computed with the string method. If one know the saddle point, one can easily compute the MLP by solving for the (reverse) flow/drift from the saddle point to the minima. As such, the maximum likelihood transition path from (-1,0) to (1,0) gives:
+The undriven duffing oscillator, is autonomous and respect detailed balance. As such the maximum likelihood path is the path that is parallel to drift and can be computed with the string method. If one know the saddle point, one can easily compute the MLP by solving for the (reverse) flow/drift from the saddle point to the minima. As such, the maximum likelihood transition path from (-1,0) to (1,0) gives:
 
 ```@example TPT
 # Generate and plot the maximum likelihood transition path from (-1,0) to (1,0)
@@ -137,7 +139,7 @@ scatter!(pts_outer[:,1], pts_outer[:,2], label="Outer points")
 fig
 ```
 
-We would like to compute the committor, the reactive current, and the reaction rate for the overdamped Duffing oscillator with additive Gaussian noise. We compute these quantities on a triangular mesh between the before computed boundaries.
+We would like to compute the committor, the reactive current, and the reaction rate for the Duffing oscillator with additive Gaussian noise. We compute these quantities on a triangular mesh between the before computed boundaries.
 
 ```@example TPT
 box = [xmin, xmax, ymin, ymax]
@@ -170,9 +172,11 @@ fig
 ## Committor functions
 
 A committor measures the probability that a system, starting at a given point in phase space, will reach one designated region before another. Formally, for two disjoint sets A and B, the forward committor $q_+(x, p)$ from A to B gives the likelihood that a trajectory initiated at x will reach B before A under the system’s dynamics. The committor boundary-value problem for a Langevin system is given by:
-\begin{equation}
+
+```math
  p \frac{\mathrm{d}q}{\mathrm{d}x} - U'(x) \frac{\mathrm{d}q}{\mathrm{d}p} + \gamma [-p \frac{\mathrm{d}q}{\mathrm{d}p} + \beta^{-1} \frac{\mathrm{d}^2 q}{\mathrm{d}p^2}] = 0,
-\end{equation}
+```
+
 for $(x,p) \in (A\cup B)^c$, with boundary conditions $q(\partial A) = 0$, $q(\partial B) = 1$, and $\nabla \nabla q = 0$ on the outer boundary ${(x,p) : H(x,p) = \mathrm{Hbdry}}$. The homogeneous Neumann boundary condition $\nabla \nabla q = 0$ means that the trajectory reflects from the outer boundary whenever it reaches it. We can compute the committor function for the Duffing oscillator using the `committor` function.
 
 ```@example TPT
@@ -187,6 +191,7 @@ tricontourf(Triangulation(mesh.pts', mesh.tri'), q)
 ```
 
 We can also compute the backward committor $q_{-}(x, p)$ from A to B, which is the probability that a trajectory initiated at x will reach A before B under the system’s dynamics. Hence, we must reverse the drift function in the Langevin system and swap the boundaries A and B in the committor function
+
 ```@example TPT
 function divfree1(x,y)
     f1,f2 = divfree(x,y)
@@ -202,18 +207,22 @@ qminus = committor(langevin_sys_reverse, mesh, Bind, Aind)
 tricontourf(Triangulation(mesh.pts', mesh.tri'), qminus)
 ```
 
-For non-equilibrium processes, such as the transitions in the double-well of the Duffing, we have that the $q_{-}\neq 1-q_+$. In particular, for langevin systems of the form in the system above time reversal involves a momentum flip such that $q_{-}(x, p)= 1-q_+(x, -p)$.
+For non-equilibrium processes, such as the transitions in the double-well of the Duffing, we have that the $q_{-}\neq 1-q_+$. In particular, for Langevin systems of the form in the system above time reversal involves a momentum flip such that $q_{-}(x, p)= 1-q_+(x, -p)$.
 
 ## Probability Density of Reactive Trajectories
 
-In general, we are interested in reactive trajectories that start in A and end in B without going back to B. The probability density of finding a reactive trajectory at a point in phase space is given by:
+In general, we are interested in reactive trajectories that start in A and ends in B without going back to A. The probability density of finding a reactive trajectory at a point in phase space is given by:
+
 ```math
 \rho_R(x, p) = \rho(x, p) q₊(x, p) q₋(x, p),
 ```
-where $\rho(x, p)$ is the probability density of finding a trajectory at $(x,p)$, $\rho(x, p)$ is also called the invariant probability density of the system. For a overdamped langevin system the invariant probability density:
+
+where $\rho(x, p)$ is the probability density of finding a trajectory at $(x,p)$, $\rho(x, p)$ is also called the invariant probability density of the system. For an overdamped Langevin system the invariant probability density:
+
 ```math
 \rho(x, p) \approx exp(-\beta H(x,p))/Z
 ```
+
 with $Z=\int exp(-\beta H(x,p)) \mathrm{d}x \mathrm{d}p$ the normalization. We can compute the integrated invariant probability density `Z` for the mesh using the `invariant_pdf` function.
 
 ```@example TPT

ext/DistMesh2D.jl

@@ -9,7 +9,7 @@ using CriticalTransitions: Mesh
 
 export distmesh2D, dellipse, ddiff
 export get_ellipse, reparametrization
-export find_boundary, huniform, dunion
+export huniform, dunion
 
 """
     huniform(p::AbstractMatrix)
@@ -49,14 +49,6 @@ function triarea(pts, tri)
     return d12[:, 1] .* d13[:, 2] - d12[:, 2] .* d13[:, 1]
 end
 
-function CriticalTransitions.find_boundary(pts, a, radii, h0)
-    xc, yc = a
-    rx, ry = radii
-    circle = @. sqrt((pts[:, 1] - xc)^2 / rx^2 + (pts[:, 2] - yc)^2 / ry^2)
-    ind = findall(circle .- 1.0 .< h0 * 1e-2)
-    return length(ind), vec(ind)
-end
-
 """
     fixmesh(pts::AbstractMatrix, tri::AbstractMatrix)
 
@@ -283,4 +275,33 @@ function CriticalTransitions.reparametrization(path, h)
     path_y = itp_y.(g1)
     return [path_x path_y]
 end
+
+function CriticalTransitions.TransitionPathMesh(
+    mesh::Mesh, point_a, point_b, radii, density
+)
+    Na = round(Int, π * sum(radii) / density) # the number of points on the A-circle
+    Nb = Na
+
+    ptsA = get_ellipse(point_a, radii, Na)
+    ptsB = get_ellipse(point_b, radii, Na)
+
+    function dfuncA(p)
+        return dellipse(p, point_a, radii)
+    end
+
+    function dfuncB(p)
+        return dellipse(p, point_b, radii)
+    end
+
+    xa, ya = point_a
+    xb, yb = point_b
+    rx, ry = radii
+
+    bboxA = [xa - rx, xa + rx, ya - ry, ya + ry]
+    Amesh = distmesh2D(dfuncA, huniform, density, bboxA, ptsA)
+    bboxB = [xb - rx, xb + rx, yb - ry, yb + ry]
+    Bmesh = distmesh2D(dfuncB, huniform, density, bboxB, ptsB)
+    return TransitionPathMesh(mesh, Amesh, Bmesh, point_a, point_b, radii, density)
+end
+
 end # module

src/CriticalTransitions.jl

@@ -46,6 +46,7 @@ include("trajectories/TransitionEnsemble.jl")
 include("trajectories/simulation.jl")
 include("trajectories/transition.jl")
 
+include("transition_path_theory/TransitionPathMesh.jl")
 include("transition_path_theory/langevin.jl")
 include("transition_path_theory/committor.jl")
 include("transition_path_theory/invariant_pdf.jl")
@@ -81,6 +82,7 @@ export basins, basinboundary, basboundary
 export intervals_to_box
 
 export distmesh2D, dellipse, ddiff
+export TransitionPathMesh, Committor
 export get_ellipse, reparametrization
 export find_boundary, huniform, dunion
 

src/extension_functions.jl

@@ -13,11 +13,6 @@ function distmesh2D end
 function get_ellipse end
 function reparametrization end
 
-struct Mesh
-    pts::Matrix{Float64}
-    tri::Matrix{Int}
-end
-
 # ## Method error handling
 # We also inject a method error handler, which
 # prints a suggestion if the Proj extension is not loaded.

src/transition_path_theory/TransitionPathMesh.jl

@@ -0,0 +1,43 @@
+struct Mesh{T,I}
+    pts::Matrix{T}
+    tri::Matrix{I}
+end
+
+struct TransitionPathMesh{T,I}
+    mesh::Mesh{T,I}
+    Amesh::Mesh{T,I}
+    Bmesh::Mesh{T,I}
+    point_a::Tuple{T,T}
+    point_b::Tuple{T,T}
+    radii::Tuple{T,T}
+    density::T
+end
+
+function find_boundary(pts, a, radii, h0)
+    xc, yc = a
+    rx, ry = radii
+    circle = @. sqrt((pts[:, 1] - xc)^2 / rx^2 + (pts[:, 2] - yc)^2 / ry^2)
+    ind = findall(circle .- 1.0 .< h0 * 1e-2)
+    return length(ind), vec(ind)
+end
+
+function find_boundary_A(TPmesh::TransitionPathMesh; set=:A)
+    return find_boundary(
+        TPmesh.mesh.pts,
+        set == :A ? TPmesh.point_a : TPmesh.point_b,
+        TPmesh.radii,
+        TPmesh.density,
+    )
+end
+
+function reverse_AB(TPmesh::TransitionPathMesh)
+    return TransitionPathMesh(
+        TPmesh.mesh,
+        TPmesh.Bmesh,
+        TPmesh.Amesh,
+        TPmesh.point_b,
+        TPmesh.point_a,
+        TPmesh.radii,
+        TPmesh.density,
+    )
+end

src/transition_path_theory/committor.jl

@@ -1,3 +1,4 @@
+
 """
 $(TYPEDSIGNATURES)
 
@@ -112,3 +113,34 @@ function committor(sys::Langevin, mesh::Mesh, Aind, Bind)
     q[free_nodes] = A[free_nodes, free_nodes] \ b[free_nodes]
     return vec(q)
 end
+
+function _committor(sys::Langevin, TPmesh::TransitionPathMesh)
+    KE, divfree, beta, gamma = sys.kinetic, sys.driftfree, sys.beta, sys.gamma
+    pts, tri = TPmesh.mesh.pts, TPmesh.mesh.tri
+
+    _, Aind = find_boundary_A(TPmesh; set=:A)
+    _, Bind = find_boundary_A(TPmesh; set=:B)
+
+    return committor(sys, TPmesh.mesh, Aind, Bind)
+end
+
+struct Committor
+    mesh::TransitionPathMesh
+    forward::Vector{Float64}
+    backward::Vector{Float64}
+    Z::Float64
+
+    function Committor(sys::Langevin, mesh::TransitionPathMesh)
+        forwardsys = sys
+        backwardsys = remake(sys; driftfree=(x, p) -> -1 .* sys.driftfree(x, p))
+        forwardmesh = mesh
+        backwardmesh = reverse_AB(mesh)
+
+        forward = _committor(forwardsys, forwardmesh)
+        backward = _committor(backwardsys, backwardmesh)
+        return new(mesh, forward, backward, partition_function(sys, mesh))
+    end
+end
+function committor(sys::Langevin, mesh::TransitionPathMesh)
+    return Committor(sys, mesh)
+end

src/transition_path_theory/invariant_pdf.jl

@@ -70,3 +70,7 @@ function invariant_pdf(sys::Langevin, mesh::Mesh, Amesh::Mesh, Bmesh::Mesh)
 
     return Z
 end
+
+function partition_function(sys::Langevin, mesh::TransitionPathMesh)
+    return invariant_pdf(sys, mesh.mesh, mesh.Amesh, mesh.Bmesh)
+end

src/transition_path_theory/langevin.jl

@@ -22,3 +22,14 @@ struct Langevin{H,D,KE,T}
     "Inverse temperature parameter (β = 1/kT) that sets the noise intensity."
     beta::T
 end
+
+function SciMLBase.remake(
+    sys::Langevin;
+    Hamiltonian=sys.Hamiltonian,
+    driftfree=sys.driftfree,
+    kinetic=sys.kinetic,
+    gamma=sys.gamma,
+    beta=sys.beta,
+)
+    return Langevin(Hamiltonian, driftfree, kinetic, gamma, beta)
+end

src/transition_path_theory/probability.jl

@@ -38,6 +38,10 @@ function probability_reactive(sys::Langevin, mesh::Mesh, q, qminus, Z)
     return prob
 end
 
+function probability_reactive(sys::Langevin, q::Committor)
+    return probability_reactive(sys, q.mesh.mesh, q.forward, q.backward, q.Z)
+end
+
 """
 $(TYPEDSIGNATURES)
 
@@ -87,3 +91,7 @@ function probability_last_A(sys::Langevin, mesh::Mesh, Ames::Mesh, qminus, Z)
     prob /= Z
     return prob
 end
+
+function probability_last_A(sys::Langevin, q::Committor)
+    return probability_last_A(sys, q.mesh.mesh, q.mesh.Amesh, q.backward, q.Z)
+end

src/transition_path_theory/reactive_current.jl

@@ -79,3 +79,48 @@ function reactive_current(sys::Langevin, mesh::Mesh, q, qminus, Z)
 
     return Rcurrent_verts, Rrate
 end
+
+function reactive_current(sys::Langevin, q::Committor)
+    return reactive_current(sys::Langevin, q.mesh.mesh::Mesh, q.forward, q.backward, q.Z)
+end
+
+function reactive_rate(sys::Langevin, mesh::Mesh, q, Z)
+    ham, divfree, beta, gamma = sys.Hamiltonian, sys.driftfree, sys.beta, sys.gamma
+    pts, tri = mesh.pts, mesh.tri
+    Npts = size(pts, 1)
+    Ntri = size(tri, 1)
+    # find the reactive current and the transition rate
+    # reactive current at the centers of mesh triangles
+    Rcurrent = zeros(Float64, (Ntri, 2))
+    Rrate = 0.0
+
+    for j in 1:Ntri
+        ind = tri[j, :]
+        verts = pts[ind, :]
+        qtri = q[ind]
+
+        a = [
+            verts[2, 1]-verts[1, 1] verts[2, 2]-verts[1, 2]
+            verts[3, 1]-verts[1, 1] verts[3, 2]-verts[1, 2]
+        ]
+        b = [qtri[2] - qtri[1], qtri[3] - qtri[1]]
+
+        g = a \ b
+
+        Aux = ones(3, 3)
+        Aux[2:3, :] .= transpose(verts)
+        tri_area = 0.5 * abs(det(Aux))
+
+        vmid = reshape(sum(verts; dims=1) / 3, 1, 2) # midpoint of mesh triangle
+        mu = exp(-beta * ham(vmid[:, 1], vmid[:, 2])[1])
+        Rrate += g[1]^2 * mu * tri_area
+    end
+
+    Rrate = Rrate * gamma / (Z * beta)
+
+    return Rrate
+end
+
+function reactive_rate(sys::Langevin, q::Committor)
+    return reactive_rate(sys::Langevin, q.mesh.mesh::Mesh, q.forward, q.Z)
+end