SAMOS density calculation

get_density_from_diffusion.py

from typing import Iterable

import numpy as np

from ase.io import read
from ase.atoms import Atoms
from samos.trajectory import Trajectory
from samos.analysis.get_gaussian_density import get_gaussian_density

def calculate_gaussian_density(
    ase_trajectory: Iterable[Atoms],
    out_file: str,
    element: str,
):
    t = Trajectory()

    # This sets the unit cell from the first frame. Not sure what to do
    # in case unit cell is allowed to vary during MD...
    t.set_atoms(next(iter(ase_trajectory)))


    positions = np.array([a.positions for a in ase_trajectory])
    t.set_array(t._POSITIONS_KEY, positions)
    get_gaussian_density(t, element=element, outputfile=out_file)


if __name__ == "__main__":
    from argparse import ArgumentParser
    parser = ArgumentParser()
    parser.add_argument("-i", "--input-file", type=str, required=True)
    parser.add_argument("-o", "--output-file", type=str, default="out.xsf")
    parser.add_argument("-e", "--element-of-interest", type=str, default="Li")
    args = parser.parse_args()
    print("Called with args:")
    for k, v in vars(args).items():
        print(f"{k:15} = {v}")

    ase_traj = read(
        args.input_file,
        index=':',
    )
    assert isinstance(ase_traj, list)
    print("Trajectory length =", len(ase_traj))
    calculate_gaussian_density(
        ase_trajectory=ase_traj,
        out_file=args.output_file,
        element=args.element_of_interest,
    )