10391705242

Committed 14 Aug 2024 04:52PM UTC coverage: 95.275% (-0.06%) from 95.337%

Build # 10391705242

Build Type

push

github

Committed by

web-flow

Commit Message

Handle blending of stress calculations (#52)

* fix tuple type hint

* torchani stress voigt form

* test blending with stresses

* robust stress shape handling in blending

* Revert "torchani stress voigt form"

This reverts commit 2a5691996.

* clean docstring;flake8

Run Details

12 of 13 new or added lines in 2 files covered. (92.31%)

746 of 783 relevant lines covered (95.27%)

0.95 hits per line

Source File
Press 'n' to go to next uncovered line, 'b' for previous

82.35

/cascade/utils.py

from io import StringIO
from copy import deepcopy

import numpy as np
import ase
from ase import Atoms
from ase.calculators.calculator import Calculator
from ase.calculators.singlepoint import SinglePointCalculator


# Taken from ExaMol
# Taken from jitterbug
def write_to_string(atoms: Atoms, fmt: str, **kwargs) -> str:
    """Write an ASE atoms object to string

    Args:
        atoms: Structure to write
        fmt: Target format
        kwargs: Passed to the write function
    Returns:
        Structure written in target format
    """

    out = StringIO()
    atoms.write(out, fmt, **kwargs)
    return out.getvalue()


def read_from_string(atoms_msg: str, fmt: str) -> Atoms:
    """Read an ASE atoms object from a string

    Args:
        atoms_msg: String format of the object to read
        fmt: Format (cannot be autodetected)
    Returns:
        Parsed atoms object
    """

    out = StringIO(str(atoms_msg))  # str() ensures that Proxies are resolved
    return ase.io.read(out, format=fmt)


def canonicalize(atoms: Atoms) -> Atoms:
    """A convenience function to standardize the format of an ase.Atoms object

    The main motivation is to freeze the Atoms.calc attribute into an immutable
    singlepoint calculator

    Args:
        atoms: Structure to write
    Returns:
        Atoms object that has been serialized and deserialized
    """
    # TODO (wardlt): Make it so the single-point calculator can hold unknown properties? (see discussion https://gitlab.com/ase/ase/-/issues/782)
    out_atoms = atoms.copy()
    if atoms.calc is not None:
        old_calc = atoms.calc
        out_atoms.calc = SinglePointCalculator(atoms)
        out_atoms.calc.results = old_calc.results.copy()
    return out_atoms


def apply_calculator(
        calc: Calculator,
        traj: list[Atoms]) -> list[Atoms]:
    """Run a calculator on every atoms object in a list, returning a new list

    Args:
        calc: the calculator to be applied
        traj: the list of atoms to have the calculator applied to
    Returns:
        list of atoms that have been canonicalized s.t. their results are stored in a SinglePointCalculator
    """
    traj = deepcopy(traj)
    for i, atoms in enumerate(traj):
        atoms.calc = calc
        atoms.get_forces()
        traj[i] = canonicalize(atoms)
    return traj


def to_voigt(stress: np.ndarray) -> np.ndarray:
    """Converts a (3,3) stress tensor to voigt form, or do nothing if its already in this form

    Args:
        stress: a stess tensor of shape (3,3) or (6,), in which case this function is a no-op
    Returns:
        stress: a (6,) stress tensor in voigt form
    """
    if stress.shape == (3, 3):
        stress = np.array([stress[0, 0], stress[1, 1], stress[2, 2],
                           stress[1, 2], stress[0, 2], stress[0, 1]])
    elif stress.shape != (6, ):
        raise ValueError(f"Stress tensor must either be of shape (3,3) or (6,), got {stress.shape}")
    return stress

1	from io import StringIO	1✔
2	from copy import deepcopy	1✔
3
4	import numpy as np	1✔
5	import ase	1✔
6	from ase import Atoms	1✔
7	from ase.calculators.calculator import Calculator	1✔
8	from ase.calculators.singlepoint import SinglePointCalculator	1✔
9
10
11	# Taken from ExaMol
12	# Taken from jitterbug
13	def write_to_string(atoms: Atoms, fmt: str, **kwargs) -> str:	1✔
14	"""Write an ASE atoms object to string
15
16	Args:
17	atoms: Structure to write
18	fmt: Target format
19	kwargs: Passed to the write function
20	Returns:
21	Structure written in target format
22	"""
23
24	out = StringIO()	×
25	atoms.write(out, fmt, **kwargs)	×
26	return out.getvalue()	×
27
28
29	def read_from_string(atoms_msg: str, fmt: str) -> Atoms:	1✔
30	"""Read an ASE atoms object from a string
31
32	Args:
33	atoms_msg: String format of the object to read
34	fmt: Format (cannot be autodetected)
35	Returns:
36	Parsed atoms object
37	"""
38
39	out = StringIO(str(atoms_msg)) # str() ensures that Proxies are resolved	×
40	return ase.io.read(out, format=fmt)	×
41
42
43	def canonicalize(atoms: Atoms) -> Atoms:	1✔
44	"""A convenience function to standardize the format of an ase.Atoms object
45
46	The main motivation is to freeze the Atoms.calc attribute into an immutable
47	singlepoint calculator
48
49	Args:
50	atoms: Structure to write
51	Returns:
52	Atoms object that has been serialized and deserialized
53	"""
54	# TODO (wardlt): Make it so the single-point calculator can hold unknown properties? (see discussion https://gitlab.com/ase/ase/-/issues/782)
55	out_atoms = atoms.copy()	1✔
56	if atoms.calc is not None:	1✔
57	old_calc = atoms.calc	1✔
58	out_atoms.calc = SinglePointCalculator(atoms)	1✔
59	out_atoms.calc.results = old_calc.results.copy()	1✔
60	return out_atoms	1✔
61
62
63	def apply_calculator(	1✔
64	calc: Calculator,
65	traj: list[Atoms]) -> list[Atoms]:
66	"""Run a calculator on every atoms object in a list, returning a new list
67
68	Args:
69	calc: the calculator to be applied
70	traj: the list of atoms to have the calculator applied to
71	Returns:
72	list of atoms that have been canonicalized s.t. their results are stored in a SinglePointCalculator
73	"""
74	traj = deepcopy(traj)	1✔
75	for i, atoms in enumerate(traj):	1✔
76	atoms.calc = calc	1✔
77	atoms.get_forces()	1✔
78	traj[i] = canonicalize(atoms)	1✔
79	return traj	1✔
80
81
82	def to_voigt(stress: np.ndarray) -> np.ndarray:	1✔
83	"""Converts a (3,3) stress tensor to voigt form, or do nothing if its already in this form
84
85	Args:
86	stress: a stess tensor of shape (3,3) or (6,), in which case this function is a no-op
87	Returns:
88	stress: a (6,) stress tensor in voigt form
89	"""
90	if stress.shape == (3, 3):	1✔
91	stress = np.array([stress[0, 0], stress[1, 1], stress[2, 2],	1✔
92	stress[1, 2], stress[0, 2], stress[0, 1]])
93	elif stress.shape != (6, ):	1✔
NEW 94	raise ValueError(f"Stress tensor must either be of shape (3,3) or (6,), got {stress.shape}")	×
95	return stress	1✔

globus-labs / cascade / 10391705242

Source File Press 'n' to go to next uncovered line, 'b' for previous

Source File
Press 'n' to go to next uncovered line, 'b' for previous