index | modules | gitlab

# Source code for ase.md.npt

    '''Constant pressure/stress and temperature dynamics.

    Combined Nose-Hoover and Parrinello-Rahman dynamics, creating an NPT
    (or N,stress,T) ensemble.

    The method is the one proposed by Melchionna et al. [1] and later
    modified by Melchionna [2].  The differential equations are integrated
    using a centered difference method [3].

     1. S. Melchionna, G. Ciccotti and B. L. Holian, "Hoover NPT dynamics
        for systems varying in shape and size", Molecular Physics 78, p. 533
        (1993).

     2. S. Melchionna, "Constrained systems and statistical distribution",
        Physical Review E, 61, p. 6165 (2000).

     3. B. L. Holian, A. J. De Groot, W. G. Hoover, and C. G. Hoover,
        "Time-reversible equilibrium and nonequilibrium isothermal-isobaric
        simulations with centered-difference Stoermer algorithms.", Physical
        Review A, 41, p. 4552 (1990).
    '''
    import sys
    import weakref
    from typing import IO, Optional, Tuple, Union

    import numpy as np

    from ase import Atoms, units
    from ase.md.md import MolecularDynamics

    linalg = np.linalg

    # Delayed imports:  If the trajectory object is reading a special ASAP version
    # of HooverNPT, that class is imported from Asap.Dynamics.NPTDynamics.

    class NPT(MolecularDynamics):

        classname = "NPT"  # Used by the trajectory.
        _npt_version = 2   # Version number, used for Asap compatibility.

        def __init__(
            self,
            atoms: Atoms,
            timestep: float,
            temperature: Optional[float] = None,
            externalstress: Optional[float] = None,
            ttime: Optional[float] = None,
            pfactor: Optional[float] = None,
            *,
            temperature_K: Optional[float] = None,
            mask: Optional[Union[Tuple[int], np.ndarray]] = None,
            trajectory: Optional[str] = None,
            logfile: Optional[Union[IO, str]] = None,
            loginterval: int = 1,
            append_trajectory: bool = False,
        ):
            '''Constant pressure/stress and temperature dynamics.

            Combined Nose-Hoover and Parrinello-Rahman dynamics, creating an
            NPT (or N,stress,T) ensemble.

            The method is the one proposed by Melchionna et al. [1] and later
            modified by Melchionna [2].  The differential equations are integrated
            using a centered difference method [3].  See also NPTdynamics.tex

            The dynamics object is called with the following parameters:

            atoms: Atoms object
                The list of atoms.

            timestep: float
                The timestep in units matching eV, Å, u.

            temperature: float (deprecated)
                The desired temperature in eV.

            temperature_K: float
                The desired temperature in K.

            externalstress: float or nparray
                The external stress in eV/A^3.  Either a symmetric
                3x3 tensor, a 6-vector representing the same, or a
                scalar representing the pressure.  Note that the
                stress is positive in tension whereas the pressure is
                positive in compression: giving a scalar p is
                equivalent to giving the tensor (-p, -p, -p, 0, 0, 0).

            ttime: float
                Characteristic timescale of the thermostat, in ASE internal units
                Set to None to disable the thermostat.

                WARNING: Not specifying ttime sets it to None, disabling the
                thermostat.

            pfactor: float
                A constant in the barostat differential equation.  If
                a characteristic barostat timescale of ptime is
                desired, set pfactor to ptime^2 * B
                (where ptime is in units matching
                eV, Å, u; and B is the Bulk Modulus, given in eV/Å^3).
                Set to None to disable the barostat.
                Typical metallic bulk moduli are of the order of
                100 GPa or 0.6 eV/A^3.

                WARNING: Not specifying pfactor sets it to None, disabling the
                barostat.

            mask: None or 3-tuple or 3x3 nparray (optional)
                Optional argument.  A tuple of three integers (0 or 1),
                indicating if the system can change size along the
                three Cartesian axes.  Set to (1,1,1) or None to allow
                a fully flexible computational box.  Set to (1,1,0)
                to disallow elongations along the z-axis etc.
                mask may also be specified as a symmetric 3x3 array
                indicating which strain values may change.

            Useful parameter values:

            * The same timestep can be used as in Verlet dynamics, i.e. 5 fs is fine
              for bulk copper.

            * The ttime and pfactor are quite critical[4], too small values may
              cause instabilites and/or wrong fluctuations in T / p.  Too
              large values cause an oscillation which is slow to die.  Good
              values for the characteristic times seem to be 25 fs for ttime,
              and 75 fs for ptime (used to calculate pfactor), at least for
              bulk copper with 15000-200000 atoms.  But this is not well
              tested, it is IMPORTANT to monitor the temperature and
              stress/pressure fluctuations.

            References:

            1) S. Melchionna, G. Ciccotti and B. L. Holian, Molecular
               Physics 78, p. 533 (1993).

            2) S. Melchionna, Physical
               Review E 61, p. 6165 (2000).

            3) B. L. Holian, A. J. De Groot, W. G. Hoover, and C. G. Hoover,
               Physical Review A 41, p. 4552 (1990).

            4) F. D. Di Tolla and M. Ronchetti, Physical
               Review E 48, p. 1726 (1993).

            '''

            MolecularDynamics.__init__(self, atoms, timestep, trajectory,
                                       logfile, loginterval,
                                       append_trajectory=append_trajectory)
            # self.atoms = atoms
            # self.timestep = timestep
            if externalstress is None and pfactor is not None:
                raise TypeError("Missing 'externalstress' argument.")
            self.zero_center_of_mass_momentum(verbose=1)
            self.temperature = units.kB * self._process_temperature(
                temperature, temperature_K, 'eV')
            if externalstress is not None:
                self.set_stress(externalstress)
            self.set_mask(mask)
            self.eta = np.zeros((3, 3), float)
            self.zeta = 0.0
            self.zeta_integrated = 0.0
            self.initialized = 0
            self.ttime = ttime
            self.pfactor_given = pfactor
            self._calculateconstants()
            self.timeelapsed = 0.0
            self.frac_traceless = 1

        def set_temperature(self, temperature=None, *, temperature_K=None):
            """Set the temperature.

            Parameters:

            temperature: float (deprecated)
                The new temperature in eV.  Deprecated, use ``temperature_K``.

            temperature_K: float (keyword-only argument)
                The new temperature, in K.
            """
            self.temperature = units.kB * self._process_temperature(
                temperature, temperature_K, 'eV')
            self._calculateconstants()

        def set_stress(self, stress):
            """Set the applied stress.

            Must be a symmetric 3x3 tensor, a 6-vector representing a symmetric
            3x3 tensor, or a number representing the pressure.

            Use with care, it is better to set the correct stress when creating
            the object.
            """

            if np.isscalar(stress):
                stress = np.array([-stress, -stress, -stress, 0.0, 0.0, 0.0])
            else:
                stress = np.array(stress)
                if stress.shape == (3, 3):
                    if not self._issymmetric(stress):
                        raise ValueError(
                            "The external stress must be a symmetric tensor.")
                    stress = np.array((stress[0, 0], stress[1, 1],
                                       stress[2, 2], stress[1, 2],
                                       stress[0, 2], stress[0, 1]))
                elif stress.shape != (6,):
                    raise ValueError("The external stress has the wrong shape.")
            self.externalstress = stress

        def set_mask(self, mask):
            """Set the mask indicating dynamic elements of the computational box.

            If set to None, all elements may change.  If set to a 3-vector
            of ones and zeros, elements which are zero specify directions
            along which the size of the computational box cannot change.
            For example, if mask = (1,1,0) the length of the system along
            the z-axis cannot change, although xz and yz shear is still
            possible.  May also be specified as a symmetric 3x3 array indicating
            which strain values may change.

            Use with care, as you may "freeze in" a fluctuation in the strain rate.
            """
            if mask is None:
                mask = np.ones((3,))
            if not hasattr(mask, "shape"):
                mask = np.array(mask)
            if mask.shape != (3,) and mask.shape != (3, 3):
                raise RuntimeError('The mask has the wrong shape ' +
                                   '(must be a 3-vector or 3x3 matrix)')
            else:
                mask = np.not_equal(mask, 0)  # Make sure it is 0/1

            if mask.shape == (3,):
                self.mask = np.outer(mask, mask)
            else:
                self.mask = mask

        def set_fraction_traceless(self, fracTraceless):
            """set what fraction of the traceless part of the force
            on eta is kept.

            By setting this to zero, the volume may change but the shape may not.
            """
            self.frac_traceless = fracTraceless

        def get_strain_rate(self):
            """Get the strain rate as an upper-triangular 3x3 matrix.

            This includes the fluctuations in the shape of the computational box.

            """
            return np.array(self.eta, copy=1)

        def set_strain_rate(self, rate):
            """Set the strain rate.  Must be an upper triangular 3x3 matrix.

            If you set a strain rate along a direction that is "masked out"
            (see ``set_mask``), the strain rate along that direction will be
            maintained constantly.
            """
            if not (rate.shape == (3, 3) and self._isuppertriangular(rate)):
                raise ValueError("Strain rate must be an upper triangular matrix.")
            self.eta = rate
            if self.initialized:
                # Recalculate h_past and eta_past so they match the current value.
                self._initialize_eta_h()

        def get_time(self):
            "Get the elapsed time."
            return self.timeelapsed

        def run(self, steps):
            """Perform a number of time steps."""
            if not self.initialized:
                self.initialize()
            else:
                if self.have_the_atoms_been_changed():
                    raise NotImplementedError(
                        "You have modified the atoms since the last timestep.")

            for _ in range(steps):
                self.step()
                self.nsteps += 1
                self.call_observers()

        def have_the_atoms_been_changed(self):
            "Checks if the user has modified the positions or momenta of the atoms"
            limit = 1e-10
            h = self._getbox()
            if max(abs((h - self.h).ravel())) > limit:
                self._warning("The computational box has been modified.")
                return 1
            expected_r = np.dot(self.q + 0.5, h)
            err = max(abs((expected_r - self.atoms.get_positions()).ravel()))
            if err > limit:
                self._warning("The atomic positions have been modified: " +
                              str(err))
                return 1
            return 0

        def step(self):
            """Perform a single time step.

            Assumes that the forces and stresses are up to date, and that
            the positions and momenta have not been changed since last
            timestep.
            """

            # Assumes the following variables are OK
            # q_past, q, q_future, p, eta, eta_past, zeta, zeta_past, h, h_past
            #
            # q corresponds to the current positions
            # p must be equal to self.atoms.GetCartesianMomenta()
            # h must be equal to self.atoms.GetUnitCell()
            #
            # print "Making a timestep"
            dt = self.dt
            h_future = self.h_past + 2 * dt * np.dot(self.h, self.eta)
            if self.pfactor_given is None:
                deltaeta = np.zeros(6, float)
            else:
                stress = self.stresscalculator()
                deltaeta = -2 * dt * (self.pfact * linalg.det(self.h) *
                                      (stress - self.externalstress))

            if self.frac_traceless == 1:
                eta_future = self.eta_past + self.mask * \
                    self._makeuppertriangular(deltaeta)
            else:
                trace_part, traceless_part = self._separatetrace(
                    self._makeuppertriangular(deltaeta))
                eta_future = (self.eta_past + trace_part +
                              self.frac_traceless * traceless_part)

            deltazeta = 2 * dt * self.tfact * (self.atoms.get_kinetic_energy() -
                                               self.desiredEkin)
            zeta_future = self.zeta_past + deltazeta
            # Advance time
            self.timeelapsed += dt
            self.h_past = self.h
            self.h = h_future
            self.inv_h = linalg.inv(self.h)
            self.q_past = self.q
            self.q = self.q_future
            self._setbox_and_positions(self.h, self.q)
            self.eta_past = self.eta
            self.eta = eta_future
            self.zeta_past = self.zeta
            self.zeta = zeta_future
            self._synchronize()  # for parallel simulations.
            self.zeta_integrated += dt * self.zeta
            force = self.forcecalculator()
            self._calculate_q_future(force)
            self.atoms.set_momenta(
                np.dot(self.q_future - self.q_past, self.h / (2 * dt)) *
                self._getmasses())
            # self.stresscalculator()

        def forcecalculator(self):
            return self.atoms.get_forces(md=True)

        def stresscalculator(self):
            return self.atoms.get_stress(include_ideal_gas=True)

        def initialize(self):
            """Initialize the dynamics.

            The dynamics requires positions etc for the two last times to
            do a timestep, so the algorithm is not self-starting.  This
            method performs a 'backwards' timestep to generate a
            configuration before the current.

            This is called automatically the first time ``run()`` is called.
            """
            # print "Initializing the NPT dynamics."
            dt = self.dt
            atoms = self.atoms
            self.h = self._getbox()
            if not self._isuppertriangular(self.h):
                print("I am", self)
                print("self.h:")
                print(self.h)
                print("Min:", min((self.h[1, 0], self.h[2, 0], self.h[2, 1])))
                print("Max:", max((self.h[1, 0], self.h[2, 0], self.h[2, 1])))
                raise NotImplementedError(
                    "Can (so far) only operate on lists of atoms where the "
                    "computational box is an upper triangular matrix.")
            self.inv_h = linalg.inv(self.h)
            # The contents of the q arrays should migrate in parallel simulations.
            # self._make_special_q_arrays()
            self.q = np.dot(self.atoms.get_positions(), self.inv_h) - 0.5
            # zeta and eta were set in __init__
            self._initialize_eta_h()
            deltazeta = dt * self.tfact * (atoms.get_kinetic_energy() -
                                           self.desiredEkin)
            self.zeta_past = self.zeta - deltazeta
            self._calculate_q_past_and_future()
            self.initialized = 1

        def get_gibbs_free_energy(self):
            """Return the Gibb's free energy, which is supposed to be conserved.

            Requires that the energies of the atoms are up to date.

            This is mainly intended as a diagnostic tool.  If called before the
            first timestep, Initialize will be called.
            """
            if not self.initialized:
                self.initialize()
            n = self._getnatoms()
            # tretaTeta = sum(diagonal(matrixmultiply(transpose(self.eta),
            #                                        self.eta)))
            contractedeta = np.sum((self.eta * self.eta).ravel())
            gibbs = (self.atoms.get_potential_energy() +
                     self.atoms.get_kinetic_energy()
                     - np.sum(self.externalstress[0:3]) * linalg.det(self.h) / 3.0)
            if self.ttime is not None:
                gibbs += (1.5 * n * self.temperature *
                          (self.ttime * self.zeta)**2 +
                          3 * self.temperature * (n - 1) * self.zeta_integrated)
            else:
                assert self.zeta == 0.0
            if self.pfactor_given is not None:
                gibbs += 0.5 / self.pfact * contractedeta
            else:
                assert contractedeta == 0.0
            return gibbs

        def get_center_of_mass_momentum(self):
            "Get the center of mass momentum."
            return self.atoms.get_momenta().sum(0)

        def zero_center_of_mass_momentum(self, verbose=0):
            "Set the center of mass momentum to zero."
            cm = self.get_center_of_mass_momentum()
            abscm = np.sqrt(np.sum(cm * cm))
            if verbose and abscm > 1e-4:
                self._warning(
                    self.classname +
                    ": Setting the center-of-mass momentum to zero "
                    "(was %.6g %.6g %.6g)" % tuple(cm))
            self.atoms.set_momenta(self.atoms.get_momenta() -
                                   cm / self._getnatoms())

        def attach_atoms(self, atoms):
            """Assign atoms to a restored dynamics object.

            This function must be called to set the atoms immediately after the
            dynamics object has been read from a trajectory.
            """
            try:
                self.atoms
            except AttributeError:
                pass
            else:
                raise RuntimeError("Cannot call attach_atoms on a dynamics "
                                   "which already has atoms.")
            MolecularDynamics.__init__(self, atoms, self.dt)
            limit = 1e-6
            h = self._getbox()
            if max(abs((h - self.h).ravel())) > limit:
                raise RuntimeError(
                    "The unit cell of the atoms does not match "
                    "the unit cell stored in the file.")
            self.inv_h = linalg.inv(self.h)
            self.q = np.dot(self.atoms.get_positions(), self.inv_h) - 0.5
            self._calculate_q_past_and_future()
            self.initialized = 1

        def attach(self, function, interval=1, *args, **kwargs):
            """Attach callback function or trajectory.

            At every *interval* steps, call *function* with arguments
            *args* and keyword arguments *kwargs*.

            If *function* is a trajectory object, its write() method is
            attached, but if *function* is a BundleTrajectory (or another
            trajectory supporting set_extra_data(), said method is first
            used to instruct the trajectory to also save internal
            data from the NPT dynamics object.
            """
            if hasattr(function, "set_extra_data"):
                # We are attaching a BundleTrajectory or similar
                function.set_extra_data("npt_init",
                                        WeakMethodWrapper(self, "get_init_data"),
                                        once=True)
                function.set_extra_data("npt_dynamics",
                                        WeakMethodWrapper(self, "get_data"))
            MolecularDynamics.attach(self, function, interval, *args, **kwargs)

        def get_init_data(self):
            "Return the data needed to initialize a new NPT dynamics."
            return {'dt': self.dt,
                    'temperature': self.temperature,
                    'desiredEkin': self.desiredEkin,
                    'externalstress': self.externalstress,
                    'mask': self.mask,
                    'ttime': self.ttime,
                    'tfact': self.tfact,
                    'pfactor_given': self.pfactor_given,
                    'pfact': self.pfact,
                    'frac_traceless': self.frac_traceless}

        def get_data(self):
            "Return data needed to restore the state."
            return {'eta': self.eta,
                    'eta_past': self.eta_past,
                    'zeta': self.zeta,
                    'zeta_past': self.zeta_past,
                    'zeta_integrated': self.zeta_integrated,
                    'h': self.h,
                    'h_past': self.h_past,
                    'timeelapsed': self.timeelapsed}

        @classmethod
        def read_from_trajectory(cls, trajectory, frame=-1, atoms=None):
            """Read dynamics and atoms from trajectory (Class method).

            Simultaneously reads the atoms and the dynamics from a BundleTrajectory,
            including the internal data of the NPT dynamics object (automatically
            saved when attaching a BundleTrajectory to an NPT object).

            Arguments:

            trajectory
                The filename or an open BundleTrajectory object.

            frame (optional)
                Which frame to read.  Default: the last.

            atoms (optional, internal use only)
                Pre-read atoms.  Do not use.
            """
            if isinstance(trajectory, str):
                if trajectory.endswith('/'):
                    trajectory = trajectory[:-1]
                if trajectory.endswith('.bundle'):
                    from ase.io.bundletrajectory import BundleTrajectory
                    trajectory = BundleTrajectory(trajectory)
                else:
                    raise ValueError(
                        f"Cannot open '{trajectory}': unsupported file format")
            # trajectory is now a BundleTrajectory object (or compatible)
            if atoms is None:
                atoms = trajectory[frame]
            init_data = trajectory.read_extra_data('npt_init', 0)
            frame_data = trajectory.read_extra_data('npt_dynamics', frame)
            dyn = cls(atoms, timestep=init_data['dt'],
                      temperature=init_data['temperature'],
                      externalstress=init_data['externalstress'],
                      ttime=init_data['ttime'],
                      pfactor=init_data['pfactor_given'],
                      mask=init_data['mask'])
            dyn.desiredEkin = init_data['desiredEkin']
            dyn.tfact = init_data['tfact']
            dyn.pfact = init_data['pfact']
            dyn.frac_traceless = init_data['frac_traceless']
            for k, v in frame_data.items():
                setattr(dyn, k, v)
            return (dyn, atoms)

        def _getbox(self):
            "Get the computational box."
            return self.atoms.get_cell()

        def _getmasses(self):
            "Get the masses as an Nx1 array."
            return np.reshape(self.atoms.get_masses(), (-1, 1))

        def _separatetrace(self, mat):
            """return two matrices, one proportional to the identity
            the other traceless, which sum to the given matrix
            """
            tracePart = ((mat[0][0] + mat[1][1] + mat[2][2]) / 3.) * np.identity(3)
            return tracePart, mat - tracePart

        # A number of convenient helper methods
        def _warning(self, text):
            "Emit a warning."
            sys.stderr.write("WARNING: " + text + "\n")
            sys.stderr.flush()

        def _calculate_q_future(self, force):
            "Calculate future q.  Needed in Timestep and Initialization."
            dt = self.dt
            id3 = np.identity(3)
            alpha = (dt * dt) * np.dot(force / self._getmasses(),
                                       self.inv_h)
            beta = dt * np.dot(self.h, np.dot(self.eta + 0.5 * self.zeta * id3,
                                              self.inv_h))
            inv_b = linalg.inv(beta + id3)
            self.q_future = np.dot(2 * self.q +
                                   np.dot(self.q_past, beta - id3) + alpha,
                                   inv_b)

        def _calculate_q_past_and_future(self):
            def ekin(p, m=self.atoms.get_masses()):
                p2 = np.sum(p * p, -1)
                return 0.5 * np.sum(p2 / m) / len(m)
            p0 = self.atoms.get_momenta()
            m = self._getmasses()
            p = np.array(p0, copy=1)
            dt = self.dt
            for i in range(2):
                self.q_past = self.q - dt * np.dot(p / m, self.inv_h)
                self._calculate_q_future(self.atoms.get_forces(md=True))
                p = np.dot(self.q_future - self.q_past, self.h / (2 * dt)) * m
                e = ekin(p)
                if e < 1e-5:
                    # The kinetic energy and momenta are virtually zero
                    return
                p = (p0 - p) + p0

        def _initialize_eta_h(self):
            self.h_past = self.h - self.dt * np.dot(self.h, self.eta)
            if self.pfactor_given is None:
                deltaeta = np.zeros(6, float)
            else:
                deltaeta = (-self.dt * self.pfact * linalg.det(self.h)
                            * (self.stresscalculator() - self.externalstress))
            if self.frac_traceless == 1:
                self.eta_past = self.eta - self.mask * \
                    self._makeuppertriangular(deltaeta)
            else:
                trace_part, traceless_part = self._separatetrace(
                    self._makeuppertriangular(deltaeta))
                self.eta_past = (self.eta - trace_part -
                                 self.frac_traceless * traceless_part)

        def _makeuppertriangular(self, sixvector):
            "Make an upper triangular matrix from a 6-vector."
            return np.array(((sixvector[0], sixvector[5], sixvector[4]),
                             (0, sixvector[1], sixvector[3]),
                             (0, 0, sixvector[2])))

        @staticmethod
        def _isuppertriangular(m) -> bool:
            "Check that a matrix is on upper triangular form."
            return m[1, 0] == m[2, 0] == m[2, 1] == 0.0

        def _calculateconstants(self):
            """(Re)calculate some constants when pfactor,
            ttime or temperature have been changed."""

            n = self._getnatoms()
            if self.ttime is None:
                self.tfact = 0.0
            else:
                self.tfact = 2.0 / (3 * n * self.temperature *
                                    self.ttime * self.ttime)
            if self.pfactor_given is None:
                self.pfact = 0.0
            else:
                self.pfact = 1.0 / (self.pfactor_given * linalg.det(self._getbox()))
                # self.pfact = 1.0/(n * self.temperature * self.ptime * self.ptime)
            self.desiredEkin = 1.5 * (n - 1) * self.temperature

        def _setbox_and_positions(self, h, q):
            """Set the computational box and the positions."""
            self.atoms.set_cell(h)
            r = np.dot(q + 0.5, h)
            self.atoms.set_positions(r)

        # A few helper methods, which have been placed in separate methods
        # so they can be replaced in the parallel version.
        def _synchronize(self):
            """Synchronize eta, h and zeta on all processors.

            In a parallel simulation, eta, h and zeta are communicated
            from the master to all slaves, to prevent numerical noise from
            causing them to diverge.

            In a serial simulation, do nothing.
            """
            # This is a serial simulation object.  Do nothing.

        def _getnatoms(self):
            """Get the number of atoms.

            In a parallel simulation, this is the total number of atoms on all
            processors.
            """
            return len(self.atoms)

        def _make_special_q_arrays(self):
            """Make the arrays used to store data about the atoms.

            In a parallel simulation, these are migrating arrays.  In a
            serial simulation they are ordinary Numeric arrays.
            """
            natoms = len(self.atoms)
            self.q = np.zeros((natoms, 3), float)
            self.q_past = np.zeros((natoms, 3), float)
            self.q_future = np.zeros((natoms, 3), float)

    class WeakMethodWrapper:
        """A weak reference to a method.

        Create an object storing a weak reference to an instance and
        the name of the method to call.  When called, calls the method.

        Just storing a weak reference to a bound method would not work,
        as the bound method object would go away immediately.
        """

        def __init__(self, obj, method):
            self.obj = weakref.proxy(obj)
            self.method = method

        def __call__(self, *args, **kwargs):
            m = getattr(self.obj, self.method)
            return m(*args, **kwargs)