Create Boltzmann velocity distributions from Simulations directly

Project.toml

@@ -2,6 +2,7 @@ name = "NQCDynamics"
 uuid = "36248dfb-79eb-4f4d-ab9c-e29ea5f33e14"
 authors = ["James <james.gardner1421@gmail.com>"]
 
+
 version = "0.15.0"
 
 

docs/src/NQCDistributions/overview.md

@@ -84,7 +84,14 @@ using Unitful
 velocity = VelocityBoltzmann(300u"K", rand(10), (3, 10))
 rand(velocity)
 ```
-This can be handed directly to the [`DynamicalDistribution`](@ref) when Boltzmann
+
+[`VelocityBoltzmann`](@ref) can also be called with a [`Simulation`](@ref), since this contains 
+both the atomic masses and the desired dimensions of the system.
+
+If the `Simulation` was set up with a `Model` which implements `NQCModels.mobileatoms`, *immobile* atoms 
+are correctly initialised with zero velocity. 
+
+The resulting distribution can be handed directly to the [`DynamicalDistribution`](@ref) when Boltzmann
 velocities are required.
 ```@example boltzmannvelocity
 distribution = DynamicalDistribution(velocity, 1, (3, 10))

src/NQCDistributions-convenience.jl

@@ -1,4 +1,5 @@
 """
+
 Add-on functions for NQCDistributions to make generating distributions more convenient. 
 
 """
@@ -8,21 +9,30 @@ import Distributions
 """
     VelocityBoltzmann(temperature, sim::AbstractSimulation; center = zeros(size(sim)))
 
-Generates a VelocityBoltzmann distribution covering an entire Simulation. 
-If `NQCModels.mobileatoms` is modified, velocities for immobile atoms will always be zero. 
+
+Generates a VelocityBoltzmann distribution covering an entire Simulation.
+If `NQCModels.mobileatoms` is modified, velocities for immobile atoms will always be zero.
 """
-function NQCDistributions.VelocityBoltzmann(temperature, sim::AbstractSimulation; center = zeros(Float64, size(sim)))
-	if length(NQCModels.mobileatoms(sim))==size(sim)[2]
-		# Simple case: All atoms should be moving, so give them all Boltzmann distributions. 
-		return NQCDistributions.VelocityBoltzmann(temperature, sim.atoms.masses, size(sim); center=center)
-	else
-		# Not all atoms should be moving, so assign Boltzmann distribution only to mobile atoms. 
-		sampleable_array=convert(Matrix{Distributions.UnivariateDistribution}, hcat([[Distributions.Dirac(center[i,j]) for i in 1:size(sim)[1]] for j in 1:size(sim)[2]]...))
-		for i in NQCModels.mobileatoms(sim)
-			for j in NQCModels.dofs(sim)
-				sampleable_array[j,i]=NQCDistributions.VelocityBoltzmann(temperature, sim.atoms.masses[i]; center=center[j,i])
-			end
-		end
-		return NQCDistributions.UnivariateArray(sampleable_array)
-	end
+function NQCDistributions.VelocityBoltzmann(temperature, sim::AbstractSimulation; center=zeros(Float64, size(sim)))
+    if length(NQCModels.mobileatoms(sim)) == size(sim)[2]
+        # Simple case: All atoms should be moving, so give them all Boltzmann distributions.
+        return NQCDistributions.VelocityBoltzmann(temperature, sim.atoms.masses, size(sim); center=center)
+    else
+        # Not all atoms should be moving, so assign Boltzmann distribution only to mobile atoms.
+        sampleable_array = convert(Matrix{Distributions.UnivariateDistribution}, hcat([[Distributions.Dirac(center[i, j]) for i in 1:size(sim)[1]] for j in 1:size(sim)[2]]...))
+        for i in NQCModels.mobileatoms(sim)
+            for j in NQCModels.dofs(sim)
+                sampleable_array[j, i] = NQCDistributions.VelocityBoltzmann(temperature, sim.atoms.masses[i]; center=center[j, i])
+            end
+        end
+        return NQCDistributions.UnivariateArray(sampleable_array)
+    end
+end
+
+function NQCDistributions.DynamicalDistribution(velocities, positions, simulation::Simulation; classical=Int[])
+    mobile_atoms = NQCModels.mobileatoms(simulation.calculator.model, 1)
+    simulation_dims = size(simulation)
+    frozen_indices = symdiff(1:simulation_dims[2], mobile_atoms)
+    return NQCDistributions.DynamicalDistribution(velocities, positions, simulation_dims; frozen_atoms=frozen_indices, classical=classical)
+
 end

test/Core/simulations.jl

@@ -1,6 +1,8 @@
 using Test
 using NQCDynamics
-using Unitful, UnitfulAtomic
+using Unitful
+using Distributions
+using PyCall
 
 atoms = Atoms([:C, :H])
 cell = InfiniteCell()
@@ -43,7 +45,7 @@ end
 @testset "get_temperature(Simulation)" begin
     sim = Simulation(atoms, model; temperature=10u"K")
     @test NQCDynamics.get_temperature(sim) isa Real
-    sim = Simulation(atoms, model; temperature=x->10u"K")
+    sim = Simulation(atoms, model; temperature=x -> 10u"K")
     @test NQCDynamics.get_temperature(sim) isa Real
     sim = Simulation(atoms, model; temperature=10)
     @test NQCDynamics.get_temperature(sim, 1) isa Real
@@ -52,7 +54,7 @@ end
 @testset "get_ring_polymer_temperature" begin
     sim = RingPolymerSimulation(atoms, model, 10; temperature=10u"K")
     @test NQCDynamics.get_ring_polymer_temperature(sim) isa Real
-    sim = RingPolymerSimulation(atoms, model, 10; temperature=x->10u"K")
+    sim = RingPolymerSimulation(atoms, model, 10; temperature=x -> 10u"K")
     @test NQCDynamics.get_ring_polymer_temperature(sim) isa Real
     sim = RingPolymerSimulation(atoms, model, 10; temperature=10)
     @test NQCDynamics.get_ring_polymer_temperature(sim, 1) isa Real
@@ -83,3 +85,32 @@ end
     @test masses(sim, 2) == atoms.masses[2]
     @test masses(sim, CartesianIndex(1, 2, 3)) == atoms.masses[2]
 end
+
+@testset "mobileatoms and distribution generation" begin
+    sim = Simulation(atoms, model)
+    @test NQCModels.mobileatoms(sim) == 1:2
+    @test NQCDistributions.VelocityBoltzmann(10, sim).sampleable == VelocityBoltzmann(10, atoms.masses, size(sim)).sampleable
+
+    # Generate a structure with constraints from ASE as a typical example
+    ase_constraints = pyimport("ase.constraints")
+    ase_io = pyimport("ase.io")
+    structure = ase_io.read("artifacts/desorption_test.xyz", index=1)
+    structure.set_constraint(ase_constraints.FixAtoms(indices=collect(0:17)))
+    nqcd_atoms, nqcd_positions, nqcd_cell = NQCDynamics.convert_from_ase_atoms(structure)
+    # Test the frozen model
+    frozen_model = AdiabaticASEModel(structure)
+    frozen_sim = Simulation(nqcd_atoms, frozen_model; cell=cell, temperature=10.0)
+
+    # Pass: Atoms 1-18 are frozen, 19-56 are mobile
+    @test NQCModels.mobileatoms(frozen_sim) == collect(19:56)
+    test_dist = NQCDistributions.VelocityBoltzmann(10, frozen_sim)
+    # Pass: First DOF of first atom should be frozen, so Dirac distribution
+    findfirst(x -> isa(x, Distributions.Dirac), test_dist.sampleable) == CartesianIndex(1, 1)
+    # Pass: First DOF of Atom 19 should be mobile, so Normal distribution
+    findfirst(x -> isa(x, Distributions.Normal), test_dist.sampleable) == CartesianIndex(1, 19)
+
+    # Test full distribution generation
+    test_dist = NQCDistributions.DynamicalDistribution(fill(3.0, size(nqcd_positions)), nqcd_positions, frozen_sim)
+    dist_random = rand(test_dist)
+    @test all(dist_random.v[:, 1:18] .== 0.0) # Velocities should be zero for frozen atoms.
+end