Source code for directupsampling.utils

from collections import Counter
from collections.abc import Iterable

import numpy as np
from ase.atoms import Atoms
from ase.calculators.calculator import Calculator, all_properties
from ase.calculators.mixing import LinearCombinationCalculator
from ase.calculators.singlepoint import SinglePointCalculator




def ase2mlip_stress(ase_stress: np.ndarray, cell_volume: float):
    mlip_stress = -1.0 * ase_stress * cell_volume
    return mlip_stress





def mlip2ase_stress(mlip_stress: np.ndarray, cell_volume: float):
    ase_stress = -1.0 * mlip_stress / cell_volume
    return ase_stress





def get_calcs_names(*calcs: Calculator):
    names = []
    for i, calc in enumerate(calcs):
        if calc is not None:
            if type(calc) is str:
                name = calc
            else:
                name = calc.name
        else:
            name = i
        names.append(name)
    c = Counter(names)
    for name, occurance in c.most_common():
        if occurance > 1:
            i = 0
            for _ in range(len(names)):
                if names[_] == name:
                    names[_] = name + f"{i}"
                    i += 1
    return tuple(names)





def chunks(iterable: Iterable, n: int, fill_none: bool = False) -> list:
    """Yield n chunks from an iterable."""
    chunked = [[] for _ in range(n)]
    for i, item in enumerate(iterable):
        chunked[i % n].append(item)
    if fill_none:
        max_len = max(len(_) for _ in chunked)
        for chunk in chunked:
            if len(chunk) < max_len:
                chunk.append(None)
    return chunked

    # m = len(lst)//n
    # for i in range(0, len(lst)-m, m):
    #    if i < len(lst)-2*m:
    #        yield lst[i:i + m]
    #    else:
    #        yield lst[i:]




def gather_chunks(chunked: list[list] | None):
    lst = []
    if chunked is not None:
        for _ in range(np.max([len(__) for __ in chunked])):
            for sublst in chunked:
                if len(sublst) > 0:
                    lst.append(sublst.pop(0))
    return lst





def set_volume(atoms: Atoms, volume: float) -> Atoms:
    """Set the volume of an Atoms object, keeping the shape of the cell.

    Parameters
    ----------
    atoms : Atoms
        ASE Atoms object.
    volume : float
        Volume of the simulation cell.

    Returns
    -------
    Atoms

    """
    cell = atoms.get_cell()
    current_volume = atoms.get_volume()
    cell *= (volume / current_volume) ** (1 / 3)
    atoms.set_cell(cell, scale_atoms=True)
    return atoms





def make_single_point_calculator(atoms: Atoms):
    if isinstance(atoms.calc, LinearCombinationCalculator):
        sp_calcs = [
            SinglePointCalculator(atoms, **_.results) for _ in atoms.calc.mixer.calcs
        ]
        for sp_calc, calc in zip(sp_calcs, atoms.calc.mixer.calcs):
            sp_calc.name = calc.name
            sp_calc.implemented_properties = (
                calc.implemented_properties
            )  # needed at least for ASE <=3.23
        new_calc = LinearCombinationCalculator(sp_calcs, atoms.calc.mixer.weights)
    else:
        # Below fixes issue with SinglePointCalculator and some calculators
        # Revise when stable ASE release has stable fix
        results = {k: v for k, v in atoms.calc.results.items() if k in all_properties}
        new_calc = SinglePointCalculator(atoms, **results)
        new_calc.name = atoms.calc.name
    return new_calc