Create Boltzmann velocity distributions from Simulations directly

Project.toml

@@ -1,7 +1,9 @@
 name = "NQCDynamics"
 uuid = "36248dfb-79eb-4f4d-ab9c-e29ea5f33e14"
 authors = ["James <james.gardner1421@gmail.com>"]
-version = "0.13.7"
+
+version = "0.13.8"
+
 
 [deps]
 AdvancedHMC = "0bf59076-c3b1-5ca4-86bd-e02cd72cde3d"

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

@@ -0,0 +1,35 @@
+"""
+Add-on functions for NQCDistributions to make generating distributions more convenient.
+
+"""
+
+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.
+"""
+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

src/NQCDynamics.jl

@@ -20,6 +20,8 @@ export Simulation,
        masses
 
 @reexport using NQCDistributions
+# Simulation-aware version of nuclear Boltzmann distribution. 
+include("NQCDistributions-convenience.jl")
 
 include("DynamicsUtils/DynamicsUtils.jl")
 @reexport using .DynamicsUtils: get_positions, get_velocities

test/Core/simulations.jl

@@ -1,6 +1,8 @@
 using Test
 using NQCDynamics
 using Unitful
+using Distributions
+using PyCall
 
 atoms = Atoms([:C, :H])
 cell = InfiniteCell()
@@ -25,7 +27,7 @@ model = NQCModels.Free()
 @testset "get_temperature" 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
@@ -34,7 +36,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
@@ -65,3 +67,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