Skip to content

Commit

Permalink
replace chemfiles with mdtraj for loading and saving trajectories
Browse files Browse the repository at this point in the history
  • Loading branch information
axsk committed Aug 21, 2024
1 parent 93a15cf commit daf14b1
Show file tree
Hide file tree
Showing 6 changed files with 36 additions and 75 deletions.
32 changes: 1 addition & 31 deletions Manifest.toml
Original file line number Diff line number Diff line change
Expand Up @@ -2,7 +2,7 @@

julia_version = "1.10.4"
manifest_format = "2.0"
project_hash = "50147434ae4e72454b4d3631ab3c93b006b9387f"
project_hash = "faf4567aaf643eff1653f6ca013dc56f535da286"

[[deps.ADTypes]]
git-tree-sha1 = "99a6f5d0ce1c7c6afdb759daa30226f71c54f6b0"
Expand Down Expand Up @@ -154,12 +154,6 @@ git-tree-sha1 = "c06a868224ecba914baa6942988e2f2aade419be"
uuid = "a9b6321e-bd34-4604-b9c9-b65b8de01458"
version = "0.1.0"

[[deps.AtomsBase]]
deps = ["LinearAlgebra", "PeriodicTable", "Printf", "Requires", "StaticArrays", "Unitful", "UnitfulAtomic"]
git-tree-sha1 = "995c2b6b17840cd87b722ce9c6cdd72f47bab545"
uuid = "a963bdd2-2df7-4f54-a1ee-49d51e6be12a"
version = "0.3.5"

[[deps.Automa]]
deps = ["PrecompileTools", "TranscodingStreams"]
git-tree-sha1 = "014bc22d6c400a7703c0f5dc1fdc302440cf88be"
Expand Down Expand Up @@ -356,18 +350,6 @@ weakdeps = ["SparseArrays"]
[deps.ChainRulesCore.extensions]
ChainRulesCoreSparseArraysExt = "SparseArrays"

[[deps.Chemfiles]]
deps = ["AtomsBase", "Chemfiles_jll", "DocStringExtensions", "PeriodicTable", "Unitful", "UnitfulAtomic"]
git-tree-sha1 = "82fe5e341c793cb51149d993307da9543824b206"
uuid = "46823bd8-5fb3-5f92-9aa0-96921f3dd015"
version = "0.10.41"

[[deps.Chemfiles_jll]]
deps = ["Artifacts", "JLLWrappers", "Libdl"]
git-tree-sha1 = "f3743181e30d87c23d9c8ebd493b77f43d8f1890"
uuid = "78a364fa-1a3c-552a-b4bb-8fa0f9c1fcca"
version = "0.10.4+0"

[[deps.CloseOpenIntervals]]
deps = ["Static", "StaticArrayInterface"]
git-tree-sha1 = "05ba0d07cd4fd8b7a39541e31a7b0254704ea581"
Expand Down Expand Up @@ -1992,12 +1974,6 @@ git-tree-sha1 = "8489905bcdbcfac64d1daa51ca07c0d8f0283821"
uuid = "69de0a69-1ddd-5017-9359-2bf0b02dc9f0"
version = "2.8.1"

[[deps.PeriodicTable]]
deps = ["Base64", "Unitful"]
git-tree-sha1 = "238aa6298007565529f911b734e18addd56985e1"
uuid = "7b2266bf-644c-5ea3-82d8-af4bbd25a884"
version = "1.2.1"

[[deps.Pipe]]
git-tree-sha1 = "6842804e7867b115ca9de748a0cf6b364523c16d"
uuid = "b98c9c47-44ae-5843-9183-064241ee97a0"
Expand Down Expand Up @@ -2786,12 +2762,6 @@ weakdeps = ["ConstructionBase", "InverseFunctions"]
ConstructionBaseUnitfulExt = "ConstructionBase"
InverseFunctionsUnitfulExt = "InverseFunctions"

[[deps.UnitfulAtomic]]
deps = ["Unitful"]
git-tree-sha1 = "903be579194534af1c4b4778d1ace676ca042238"
uuid = "a7773ee8-282e-5fa2-be4e-bd808c38a91a"
version = "1.0.0"

[[deps.UnitfulLatexify]]
deps = ["LaTeXStrings", "Latexify", "Unitful"]
git-tree-sha1 = "975c354fcd5f7e1ddcc1f1a23e6e091d99e99bc8"
Expand Down
2 changes: 0 additions & 2 deletions Project.toml
Original file line number Diff line number Diff line change
Expand Up @@ -8,7 +8,6 @@ BioStructures = "de9282ab-8554-53be-b2d6-f6c222edabfc"
Bonito = "824d6782-a2ef-11e9-3a09-e5662e0c26f8"
CUDA = "052768ef-5323-5732-b1bb-66c8b64840ba"
ChainRulesCore = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4"
Chemfiles = "46823bd8-5fb3-5f92-9aa0-96921f3dd015"
Combinatorics = "861a8166-3701-5b0c-9a16-15d98fcdc6aa"
DataFrames = "a93c6f00-e57d-5684-b7b6-d8193f3e46c0"
Distances = "b4f34e82-e78d-54a5-968a-f98e89d6e8f7"
Expand Down Expand Up @@ -54,7 +53,6 @@ MollyExt = "Molly"
BioStructures = "4.0.0"
CUDA = "5.2.0"
ChainRulesCore = "1.23.0"
Chemfiles = "0.10.41"
Combinatorics = "1.0.2"
DataFrames = "1.6.1"
Distances = "0.10.11"
Expand Down
3 changes: 1 addition & 2 deletions src/ISOKANN.jl
Original file line number Diff line number Diff line change
Expand Up @@ -24,7 +24,6 @@ using StaticArrays: @SVector
using StochasticDiffEq: StochasticDiffEq
using LinearAlgebra: pinv, norm, I, schur

import Chemfiles
import ProgressMeter
import ChainRulesCore
import Flux
Expand Down Expand Up @@ -68,7 +67,7 @@ export SimulationData
export getxs, getys
export exit_rates

export reactionpath_minimum, reactionpath_ode, writechemfile
export reactionpath_minimum, reactionpath_ode

include("subsample.jl") # adaptive sampling
include("pairdists.jl") # pair distances
Expand Down
3 changes: 2 additions & 1 deletion src/data.jl
Original file line number Diff line number Diff line change
Expand Up @@ -175,5 +175,6 @@ function exportsorted(iso, path="out/sorted.pdb")
p = iso.model(xs) |> vec |> sortperm
xs = ISOKANN.getcoords(iso.data)
println("saving sorted data to $path")
ISOKANN.writechemfile(path, ISOKANN.aligntrajectory(xs[:, p] |> cpu); source=pdb(iso.data))
traj = ISOKANN.aligntrajectory(xs[:, p] |> cpu)
save_trajectory(path, traj, top=pdb(iso.data))
end
69 changes: 31 additions & 38 deletions src/molutils.jl
Original file line number Diff line number Diff line change
Expand Up @@ -65,44 +65,6 @@ standardform(x::AbstractArray, rotationhandles=(2, 11, 19)) =

standardform(x::AbstractArray, sim::IsoSimulation) = standardform(x, rotationhandles(sim))


## Chemfiles to read/write trajectory data

function readchemfile(source::String, steps=nothing)
traj = Chemfiles.Trajectory(source, 'r')
readchemfile(traj, steps)
end

# todo: should we return this wih flattened coords?
function readchemfile(traj::Chemfiles.Trajectory, steps=nothing)
isnothing(steps) && (steps = 1:length(traj))
frame = Chemfiles.read_step(traj, 0)
xs = Array{Float32}(undef, size(Chemfiles.positions(frame))..., length(steps))
for (i, s) in enumerate(steps)
Chemfiles.read_step!(traj, s - 1, frame)
xs[:, :, i] .= Chemfiles.positions(frame)
end
return xs ./ 10 # convert from Angstom to nm
end

function writechemfile(filename, data::Array{<:Any,3}; source)
trajectory = Chemfiles.Trajectory(source, 'r')
frame = Chemfiles.read(trajectory)
trajectory = Chemfiles.Trajectory(filename, 'w', uppercase(split(filename, ".")[end]))
for i in 1:size(data, 3)
Chemfiles.positions(frame) .= data[:, :, i] .* 10 # convert from nm to Angstrom
write(trajectory, frame)
end
close(trajectory)
end

function writechemfile(filename, data::Array{<:Any,2}; source)
n = size(data, 2)
data = reshape(data, 3, :, n)
writechemfile(filename, data; source)
end


### alignment of pointclouds / trajectories using procrustes alignment
function aligntrajectory(traj::AbstractVector)
aligned = [centermean(traj[1])]
Expand Down Expand Up @@ -146,3 +108,34 @@ function split_first_dimension(A, d)
s1, s2... = size(A)
reshape(A, (d, div(s1, d), s2...))
end

"""
load_trajectory(filename; top=nothing, kwargs...)
wrapper around Python's `mdtraj.load()`.
Returns a (3 * natom, nframes) shaped array.
"""
function load_trajectory(filename; top::Union{Nothing,String}=nothing, stride=nothing, atom_indices=nothing)
mdtraj = pyimport_conda("mdtraj", "mdtraj", "conda-forge")

if filename[end-2:end] != "pdb" && isnothing(top)
error("must supply topology file (.pdb) to the top argument")
end

traj = MDTRAJ.load(filename; top, stride, atom_indices)
xs = permutedims(PyArray(py"$traj.xyz"o), (3, 2, 1))
xs = reshape(xs, :, size(xs, 3))::Matrix{Float32}
end

"""
save_trajectory(filename, coords::AbstractMatrix; top::String)
save the trajectory given in `coords` to `filename` with the topology provided by the file `top`
"""
function save_trajectory(filename, coords::AbstractMatrix; top::String)
mdtraj = pyimport_conda("mdtraj", "mdtraj", "conda-forge")
traj = mdtraj.load(filename, stride=-1)
xyz = reshape(coords, 3, :, size(coords, 2))
traj.xyz = PyReverseDims(xyz)
traj.save(filename)
end
2 changes: 1 addition & 1 deletion src/reactionpath2.jl
Original file line number Diff line number Diff line change
Expand Up @@ -151,6 +151,6 @@ function save_reactive_path(iso::Iso2, coords::AbstractMatrix=getcoords(iso.data
path = centercoords(path)
println("saving reactive path of length $(length(ids)) to $out")
mkpath(dirname(out))
writechemfile(out, path; source)
save_trajectory(out, path, top=source)
return ids
end

0 comments on commit daf14b1

Please sign in to comment.