cannot convert vasp MLAIMD traj to extxyz ASE format

[Asif-Iqbal-Bhatti]

Hello,

I would like to ask, is there a way to read all the trajectories from MLAIMD and convert that to an ASE extxyz file? I have this script, but that is way slow, and I am batching to fit in memory.

'''
calc = Calculation.from_file("vaspout.h5")
print(dir(calc.energy[:]))
print(dir(calc))

def MLMD_AIMD_to_EXTXYZ(xyzpath, calc, buffer_size=5000):
#struct_tomdtraj = calc.structure[:].to_mdtraj()
#print(dir(struct_tomdtraj))
#struct_tomdtraj.save_xyz('mdtraj_data.xyz')

md_steps = calc.structure[:].number_steps()
name = calc.topology
buffer = []

for i in tqdm(range(md_steps), desc="Processing MLAIMD frames"):
    atoms = calc.structure[i].to_ase()
    energy = calc.energy[i].read()['total energy   ETOTAL']
    forces = calc.force[i].read()['forces']
    stress = calc.stress[i].read()['stress'] # in kbar with -ve being opposite to ASE

    #stress_ASE = np.array(
    #    [
    #        stress[0][0],
    #        stress[1][1],
    #        stress[2][2],
    #        stress[1][2],
    #        stress[2][0],
    #        stress[0][1],
    #    ]
    #)
    ## internal stress
    #stress_eVA3 = -1 * stress_ASE / 1602.1766208  # from kbar to eV/Angstrom^3
                
    atoms.calc = SinglePointCalculator(atoms, energy=energy, forces=forces, stress=stress)
    atoms.info = {"System": f"{name}"}
    buffer.append(atoms)

    if len(buffer) >= buffer_size:
        write(f'{xyzpath}.extxyz', buffer, append=True)
        buffer.clear()

if buffer:
    write(f'{xyzpath}', buffer, append=True)

######################################################

MLMD_AIMD_to_EXTXYZ(EXTXYZ_NAME, calc)
'''

Is there a way to do much quicker using numpy vector.

Regards,
Asif

[martin-schlipf]

I fear the speed may actually be limited by writing the extxyz file. If that is not the case and it is really the reading of the data that is the limit you could implement your own to_extxyz function

def to_extxyz(filename, calc, slice_):
    # read data with py4vasp, use slice_ to read many data at the same time while
    # also limiting the amount of data read to avoid memory issues
    structure_data = calc.structure[slice_].read()
    forces = calc.force[slice_].read()["forces"]
    stress = calc.stress[slice_].read()["stress"]

    # write data to extxyz file
    num_step, num_atom, _ = forces.shape
    with open(filename, "a") as xyz_file:
        for step in range(num_step):
            # write header for current step
            for atom in range(num_atom):
                # write the content of structure step to xyz file

Possibly you could eliminate the atom loop by a clever use of np.array2string but I'm not convinced that would be faster than your own loop. You could also convert first to a text in memory and then write a longer text at once but that may be happening under the hood anyway.

[Asif-Iqbal-Bhatti]

So with your suggestion I added the following lines, and there is a considerable speedup, and memory issues are limited. I did not have the time to profile it, but what I noticed reading 50000 steps from ASE read and write was approximately the same. My previous code took 6 hours. I am sharing the code; if it can be made more efficient, someone can add more to it, and then that would be useful. I was thinking of doing that in C++, but reading .h5 is very complicated. The issue is ASE and pymatgen do not parse MLAIMD steps, only DFT. That is why I have to go through this path to feed this data to MACE, SevenNet, ..., and Fairchem Potentials.

'''
'''
def vaspouth5toextxyz():
calc = Calculation.from_file("vaspout.h5")

def to_extxyz(filename, calc, slice_):
    structure_data = calc.structure[slice_].read()
    forces = calc.force[slice_].read()["forces"]
    stress = calc.stress[slice_].read()["stress"]

    num_step, num_atom, _ = forces.shape

    species = [name.split("_")[0] for name in structure_data["names"]]

    with open(filename, "a") as xyz_file:
        for local_step in range(num_step):
            lattice = structure_data["lattice_vectors"][local_step]
            frac_positions = structure_data["positions"][local_step]
            positions = frac_positions @ lattice # MAT MUL in numpy
            step_forces = forces[local_step]
            step_stress = stress[local_step]

            xyz_file.write(f"{num_atom}\n")

            stress_flat = " ".join(f"{x:.6f}" for x in step_stress.flatten())
            lattice_flat = " ".join(f"{x:.6f}" for x in lattice.flatten())
            xyz_file.write(
                f'Lattice="{lattice_flat}" '
                f'Properties=species:S:1:pos:R:3:forces:R:3 '
                f'pbc="T T T" '
                f'Stress="{stress_flat}"\n'
            )

            for atom in range(num_atom):
                sp = species[atom]
                x, y, z = positions[atom]
                fx, fy, fz = step_forces[atom]
                xyz_file.write(f"{sp} {x:.6f} {y:.6f} {z:.6f} {fx:.6f} {fy:.6f} {fz:.6f}\n")
        
# --- MAIN ---
n_steps = calc.structure[:].number_steps()
batch = 5000
for start in range(0, n_steps, batch):
    stop = min(start + batch, n_steps)
    print(f"Writing steps {start}:{stop}")
    to_extxyz(EXTXYZ_NAME, calc, slice(start, stop))

'''
'''

[martin-schlipf]

Glad to see that you have a faster version now. You can time the reading and writing separately. I think the writing should dominate now so moving to reading the hdf5 with C++ would not help.

You could do one minor modification. The following two lines of code should be equivalent. But that is only for readability and
not for performance.

species = [name.split("_")[0] for name in structure_data["names"]]
species = structure_data["elements"]

You could do the transformation to Cartesian units outside of the local_step loop. py4vasp even provides a routine to do that

calc.structure[slice_].cartesian_positions()