Chimie Théorique » QMX

"""Read and write ASE2's netCDF trajectory files."""

from ase.io.pupynere import NetCDFFile

from ase.atoms import Atoms

from ase.calculators.singlepoint import SinglePointCalculator

def read_netcdf(filename, index=-1):

    nc = NetCDFFile(filename)

    dims = nc.dimensions

    vars = nc.variables

    positions = vars['CartesianPositions']

    numbers = vars['AtomicNumbers'][:]

    pbc = vars['BoundaryConditions'][:]

    cell = vars['UnitCell']

    tags = vars['Tags'][:]

    if not tags.any():

        tags = None

    magmoms = vars['MagneticMoments'][:]

    if not magmoms.any():

        magmoms = None

    nimages = positions.shape[0]

    attach_calculator = False

    if 'PotentialEnergy' in vars:

        energy = vars['PotentialEnergy']

        attach_calculator = True

    else:

        energy = nimages * [None]

    if 'CartesianForces' in vars:

        forces = vars['CartesianForces']

        attach_calculator = True

    else:

        forces = nimages * [None]

    if 'Stress' in vars:

        stress = vars['Stress']

        attach_calculator = True

    else:

        stress = nimages * [None]

    if isinstance(index, int):

        indices = [index]

    else:

        indices = range(nimages)[index]

    images = []

    for i in indices:

        atoms = Atoms(positions=positions[i],

                      numbers=numbers,

                      cell=cell[i],

                      pbc=pbc,

                      tags=tags, magmoms=magmoms)

        if attach_calculator:

            calc = SinglePointCalculator(energy[i], forces[i], stress[i],

                                         None, atoms) ### Fixme magmoms

            atoms.set_calculator(calc)

        images.append(atoms)

    if isinstance(index, int):

        return images[0]

    else:

        return images

class LOA:

    def __init__(self, images):

        self.set_atoms(images[0])

    def __len__(self):

        return len(self.atoms)

    def set_atoms(self, atoms):

        self.atoms = atoms

    def GetPotentialEnergy(self):

        return self.atoms.get_potential_energy()

    def GetCartesianForces(self):

        return self.atoms.get_forces()

    def GetUnitCell(self):

        return self.atoms.get_cell()

    def GetAtomicNumbers(self):

        return self.atoms.get_atomic_numbers()

    def GetCartesianPositions(self):

        return self.atoms.get_positions()

    def GetBoundaryConditions(self):

        return self.atoms.get_pbc()

def write_netcdf(filename, images):

    from ASE.Trajectories.NetCDFTrajectory import NetCDFTrajectory

    if not isinstance(images, (list, tuple)):

        images = [images]

    loa = LOA(images)

    traj = NetCDFTrajectory(filename, loa)

    for atoms in images:

        loa.set_atoms(atoms)

        traj.Update()

    traj.Close()