ModulesAndNodes/Helper_functions.py at main · Sriram2631/ModulesAndNodes · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
import ase
import numpy as np
def _get_structure(
        structure: ase.Atoms,
        cell: np.ndarray,
        indices: np.ndarray,
        positions: np.ndarray | None = None,
        unwrapped_positions: np.ndarray | None = None,
        total_displacements: np.ndarray | None = None,
        *,
        wrap_atoms: bool = True,
    ) -> ase.Atoms:

    from structuretoolkit.common import center_coordinates_in_unit_cell
    from pyiron import ase_to_pyiron
    """Return an updated `Atoms` object based on the provided information.

    Parameters
    ----------
    structure : Atoms
        The reference atomic structure.
    cell : ndarray
        The simulation cell to assign to the new structure.
    indices : ndarray
        Indices of the atoms to include in the new snapshot.
    positions : ndarray, optional
        Wrapped atomic positions.
    unwrapped_positions : ndarray, optional
        Unwrapped atomic positions.
    total_displacements : ndarray, optional
        Total atomic displacements to be added to the initial positions.
    wrap_atoms : bool, optional
        Whether to wrap atoms inside the unit cell (default is True).

    Returns
    -------
    Atoms
        The newly constructed atomic structure with updated positions and cell.

    """
    if indices is not None and len(indices) != len(structure):
        snapshot = Atoms(
            positions=np.zeros((*indices.shape, 3)),
            cell=cell,
            pbc=structure.pbc,
        )
        snapshot.set_array("indices", indices)
    else:
        snapshot = structure.copy()
        if cell is not None:
            snapshot.cell = cell
        if indices is not None:
            snapshot.set_array("indices", indices)

    if wrap_atoms:
        snapshot.positions = positions
        snapshot = center_coordinates_in_unit_cell(snapshot)
    elif unwrapped_positions is not None:
        snapshot.positions = unwrapped_positions
    else:
        snapshot.positions += total_displacements

    snapshot = ase_to_pyiron(snapshot)

    return snapshot