Chimie Théorique » QMX

# VelocityDistributions.py -- set up a velocity distribution

"""Module for setting up e.g. Maxwell-Boltzmann velocity distributions.

Currently, only one function is defined, MaxwellBoltzmannDistribution,

which sets the momenta of a list of atoms according to a

Maxwell-Boltzmann distribution at a given temperature.

"""

import numpy as np

import sys

# For parallel GPAW simulations, the random velocities should be distributed.

# Uses gpaw world communicator as default, but allow option of

# specifying other communicator (for ensemble runs)

if '_gpaw' in sys.modules:

    # http://wiki.fysik.dtu.dk/gpaw

    from gpaw.mpi import world as gpaw_world

else:

    gpaw_world = None

def _maxwellboltzmanndistribution(masses, temp, communicator=gpaw_world):

    xi = np.random.standard_normal((len(masses),3))

    if communicator is not None:

       communicator.broadcast(xi, 0)

    momenta = xi * np.sqrt(masses * temp)[:,np.newaxis]

    return momenta

def MaxwellBoltzmannDistribution(atoms, temp, communicator=gpaw_world):

    """Sets the momenta to a Maxwell-Boltzmann distribution."""

    momenta = _maxwellboltzmanndistribution(atoms.get_masses(), temp, communicator)

    atoms.set_momenta(momenta)

def Stationary(atoms):

    "Sets the center-of-mass momentum to zero."

    p = atoms.get_momenta()

    p0 = np.sum(p, 0)

    # We should add a constant velocity, not momentum, to the atoms

    m = atoms.get_masses()

    mtot = np.sum(m)

    v0 = p0/mtot

    p -= v0*m[:,np.newaxis]

    atoms.set_momenta(p)