8325072180

Committed 18 Mar 2024 10:32AM UTC coverage: 91.525% (-0.04%) from 91.566%

Build # 8325072180

Build Type

Pull #70

github

Committed by

web-flow

Commit Message

Merge 8d74312d2 into 7d9ea5872

Pull Request Pull Request #70: Add geomopt to CLI

Run Details

51 of 58 new or added lines in 4 files covered. (87.93%)

2 existing lines in 1 file now uncovered.

270 of 295 relevant lines covered (91.53%)

3.66 hits per line

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

93.27

/janus_core/single_point.py

"""Prepare and perform single point calculations."""

from collections.abc import Collection
from pathlib import Path
from typing import Any, Optional

from ase import Atoms
from ase.io import read, write
from numpy import isfinite, ndarray

from janus_core.janus_types import (
    Architectures,
    ASEReadArgs,
    ASEWriteArgs,
    CalcResults,
    Devices,
    MaybeList,
    MaybeSequence,
)
from janus_core.log import config_logger
from janus_core.mlip_calculators import choose_calculator


class SinglePoint:
    """
    Prepare and perform single point calculations.

    Parameters
    ----------
    struct : Optional[MaybeSequence[Atoms]]
        ASE Atoms structure(s) to simulate. Required if `struct_path` is None.
        Default is None.
    struct_path : Optional[str]
        Path of structure to simulate. Required if `struct` is None.
        Default is None.
    struct_name : Optional[str]
        Name of structure. Default is inferred from chemical formula if `struct`
        is specified, else inferred from `struct_path`.
    architecture : Literal[architectures]
        MLIP architecture to use for single point calculations.
        Default is "mace_mp".
    device : Devices
        Device to run model on. Default is "cpu".
    read_kwargs : ASEReadArgs
        Keyword arguments to pass to ase.io.read. Default is {}.
    calc_kwargs : Optional[dict[str, Any]]
        Keyword arguments to pass to the selected calculator. Default is {}.
    log_kwargs : Optional[dict[str, Any]]
            Keyword arguments to pass to `config_logger`. Default is {}.

    Attributes
    ----------
    architecture : Architectures
        MLIP architecture to use for single point calculations.
    struct : MaybeSequence[Atoms]
        ASE Atoms structure(s) to simulate.
    device : Devices
        Device to run MLIP model on.
    struct_path : Optional[str]
        Path of structure to simulate.
    struct_name : Optional[str]
        Name of structure.
    logger : logging.Logger
        Logger if log file has been specified.

    Methods
    -------
    read_structure(**kwargs)
        Read structure and structure name.
    set_calculator(**kwargs)
        Configure calculator and attach to structure.
    run_single_point(properties=None)
        Run single point calculations.
    """

    def __init__(
        self,
        struct: Optional[MaybeSequence[Atoms]] = None,
        struct_path: Optional[str] = None,
        struct_name: Optional[str] = None,
        architecture: Architectures = "mace_mp",
        device: Devices = "cpu",
        read_kwargs: Optional[ASEReadArgs] = None,
        calc_kwargs: Optional[dict[str, Any]] = None,
        log_kwargs: Optional[dict[str, Any]] = None,
    ) -> None:
        """
        Read the structure being simulated and attach an MLIP calculator.

        Parameters
        ----------
        struct : Optional[MaybeSequence[Atoms]]
            ASE Atoms structure(s) to simulate. Required if `struct_path`
            is None. Default is None.
        struct_path : Optional[str]
            Path of structure to simulate. Required if `struct` is None.
            Default is None.
        struct_name : Optional[str]
            Name of structure. Default is inferred from chemical formula if `struct`
            is specified, else inferred from `struct_path`.
        architecture : Architectures
            MLIP architecture to use for single point calculations.
            Default is "mace_mp".
        device : Devices
            Device to run MLIP model on. Default is "cpu".
        read_kwargs : Optional[ASEReadArgs]
            Keyword arguments to pass to ase.io.read. Default is {}.
        calc_kwargs : Optional[dict[str, Any]]
            Keyword arguments to pass to the selected calculator. Default is {}.
        log_kwargs : Optional[dict[str, Any]]
            Keyword arguments to pass to `config_logger`. Default is {}.
        """
        if struct and struct_path:
            raise ValueError(
                "You cannot specify both the ASE Atoms structure (`struct`) "
                "and a path to the structure file (`struct_path`)"
            )

        if not struct and not struct_path:
            raise ValueError(
                "Please specify either the ASE Atoms structure (`struct`) "
                "or a path to the structure file (`struct_path`)"
            )

        read_kwargs = read_kwargs if read_kwargs else {}
        calc_kwargs = calc_kwargs if calc_kwargs else {}
        log_kwargs = log_kwargs if log_kwargs else {}

        if log_kwargs and "filename" not in log_kwargs:
            raise ValueError("'filename' must be included in `log_kwargs`")

        log_kwargs.setdefault("name", __name__)
        self.logger = config_logger(**log_kwargs)

        self.architecture = architecture
        self.device = device
        self.struct_path = struct_path
        self.struct_name = struct_name

        # Read structure if given as path
        if self.struct_path:
            self.read_structure(**read_kwargs)
        else:
            self.struct = struct
            if not self.struct_name:
                self.struct_name = self.struct.get_chemical_formula()

        # Configure calculator
        self.set_calculator(**calc_kwargs)

        if self.logger:
            self.logger.info("Single point calculator configured")

    def read_structure(self, **kwargs) -> None:
        """
        Read structure and structure name.

        If the file contains multiple structures, only the last configuration
        will be read by default.

        Parameters
        ----------
        **kwargs
            Keyword arguments passed to ase.io.read.
        """
        if self.struct_path:
            self.struct = read(self.struct_path, **kwargs)
            if not self.struct_name:
                self.struct_name = Path(self.struct_path).stem
        else:
            raise ValueError("`struct_path` must be defined")

    def set_calculator(
        self, read_kwargs: Optional[ASEReadArgs] = None, **kwargs
    ) -> None:
        """
        Configure calculator and attach to structure.

        Parameters
        ----------
        read_kwargs : Optional[ASEReadArgs]
            Keyword arguments to pass to ase.io.read. Default is {}.
        **kwargs
            Additional keyword arguments passed to the selected calculator.
        """
        calculator = choose_calculator(
            architecture=self.architecture,
            device=self.device,
            **kwargs,
        )
        if self.struct is None:
            read_kwargs = read_kwargs if read_kwargs else {}
            self.read_structure(**read_kwargs)

        if isinstance(self.struct, list):
            for struct in self.struct:
                struct.calc = calculator
        else:
            self.struct.calc = calculator

    def _get_potential_energy(self) -> MaybeList[float]:
        """
        Calculate potential energy using MLIP.

        Returns
        -------
        MaybeList[float]
            Potential energy of structure(s).
        """
        if isinstance(self.struct, list):
            return [struct.get_potential_energy() for struct in self.struct]

        return self.struct.get_potential_energy()

    def _get_forces(self) -> MaybeList[ndarray]:
        """
        Calculate forces using MLIP.

        Returns
        -------
        MaybeList[ndarray]
            Forces of structure(s).
        """
        if isinstance(self.struct, list):
            return [struct.get_forces() for struct in self.struct]

        return self.struct.get_forces()

    def _get_stress(self) -> MaybeList[ndarray]:
        """
        Calculate stress using MLIP.

        Returns
        -------
        MaybeList[ndarray]
            Stress of structure(s).
        """
        if isinstance(self.struct, list):
            return [struct.get_stress() for struct in self.struct]

        return self.struct.get_stress()

    @staticmethod
    def _remove_invalid_props(
        struct: Atoms,
        results: CalcResults = None,
        properties: Collection[str] = (),
    ) -> None:
        """
        Remove any invalid properties from calculated results.

        Parameters
        ----------
        struct : Atoms
            ASE Atoms structure with attached calculator results.
        results : CalcResults
            Dictionary of calculated results. Default is {}.
        properties : Collection[str]
            Physical properties requested to be calculated. Default is ().
        """
        results = results if results else {}

        # Find any properties with non-finite values
        rm_keys = [
            prop
            for prop in struct.calc.results
            if not isfinite(struct.calc.results[prop]).all()
        ]

        # Raise error if property was explicitly requested, otherwise remove
        for prop in rm_keys:
            if prop in properties:
                raise ValueError(
                    f"'{prop}' contains non-finite values for this structure."
                )
            del struct.calc.results[prop]
            if prop in results:
                del results[prop]

    def _clean_results(
        self, results: CalcResults = None, properties: Collection[str] = ()
    ) -> None:
        """
        Remove NaN and inf values from results and calc.results dictionaries.

        Parameters
        ----------
        results : CalcResults
            Dictionary of calculated results. Default is {}.
        properties : Collection[str]
            Physical properties requested to be calculated. Default is ().
        """
        results = results if results else {}

        if isinstance(self.struct, list):
            for image in self.struct:
                self._remove_invalid_props(image, results, properties)
        else:
            self._remove_invalid_props(self.struct, results, properties)

    def run_single_point(
        self,
        properties: MaybeSequence[str] = (),
        write_results: bool = False,
        write_kwargs: Optional[ASEWriteArgs] = None,
    ) -> CalcResults:
        """
        Run single point calculations.

        Parameters
        ----------
        properties : MaybeSequence[str]
            Physical properties to calculate. If not specified, "energy",
            "forces", and "stress" will be returned.
        write_results : bool
            True to write out structure with results of calculations. Default is False.
        write_kwargs : Optional[ASEWriteArgs],
            Keyword arguments to pass to ase.io.write if saving structure with
            results of calculations. Default is {}.

        Returns
        -------
        CalcResults
            Dictionary of calculated results.
        """
        results: CalcResults = {}
        if isinstance(properties, str):
            properties = [properties]

        for prop in properties:
            if prop not in ["energy", "forces", "stress"]:
                raise NotImplementedError(
                    f"Property '{prop}' cannot currently be calculated."
                )

        write_kwargs = write_kwargs if write_kwargs else {}
        if write_kwargs and "filename" not in write_kwargs:
            raise ValueError("'filename' must be included in write_kwargs")

        if self.logger:
            self.logger.info("Starting single point calculation")

        if "energy" in properties or len(properties) == 0:
            results["energy"] = self._get_potential_energy()
        if "forces" in properties or len(properties) == 0:
            results["forces"] = self._get_forces()
        if "stress" in properties or len(properties) == 0:
            results["stress"] = self._get_stress()

        # Remove meaningless values from results e.g. stress for non-periodic systems
        self._clean_results(results, properties=properties)

        if self.logger:
            self.logger.info("Single point calculation complete")

        if write_results:
            if "filename" not in write_kwargs:
                filename = f"{self.struct_name}-results.xyz"
                write_kwargs["filename"] = Path(".").absolute() / filename
            write(images=self.struct, **write_kwargs)

        return results

1	"""Prepare and perform single point calculations."""
2
3	from collections.abc import Collection	4✔
4	from pathlib import Path	4✔
5	from typing import Any, Optional	4✔
6
7	from ase import Atoms	4✔
8	from ase.io import read, write	4✔
9	from numpy import isfinite, ndarray	4✔
10
11	from janus_core.janus_types import (	4✔
12	Architectures,
13	ASEReadArgs,
14	ASEWriteArgs,
15	CalcResults,
16	Devices,
17	MaybeList,
18	MaybeSequence,
19	)
20	from janus_core.log import config_logger	4✔
21	from janus_core.mlip_calculators import choose_calculator	4✔
22
23
24	class SinglePoint:	4✔
25	"""
26	Prepare and perform single point calculations.
27
28	Parameters
29	----------
30	struct : Optional[MaybeSequence[Atoms]]
31	ASE Atoms structure(s) to simulate. Required if `struct_path` is None.
32	Default is None.
33	struct_path : Optional[str]
34	Path of structure to simulate. Required if `struct` is None.
35	Default is None.
36	struct_name : Optional[str]
37	Name of structure. Default is inferred from chemical formula if `struct`
38	is specified, else inferred from `struct_path`.
39	architecture : Literal[architectures]
40	MLIP architecture to use for single point calculations.
41	Default is "mace_mp".
42	device : Devices
43	Device to run model on. Default is "cpu".
44	read_kwargs : ASEReadArgs
45	Keyword arguments to pass to ase.io.read. Default is {}.
46	calc_kwargs : Optional[dict[str, Any]]
47	Keyword arguments to pass to the selected calculator. Default is {}.
48	log_kwargs : Optional[dict[str, Any]]
49	Keyword arguments to pass to `config_logger`. Default is {}.
50
51	Attributes
52	----------
53	architecture : Architectures
54	MLIP architecture to use for single point calculations.
55	struct : MaybeSequence[Atoms]
56	ASE Atoms structure(s) to simulate.
57	device : Devices
58	Device to run MLIP model on.
59	struct_path : Optional[str]
60	Path of structure to simulate.
61	struct_name : Optional[str]
62	Name of structure.
63	logger : logging.Logger
64	Logger if log file has been specified.
65
66	Methods
67	-------
68	read_structure(**kwargs)
69	Read structure and structure name.
70	set_calculator(**kwargs)
71	Configure calculator and attach to structure.
72	run_single_point(properties=None)
73	Run single point calculations.
74	"""
75
76	def __init__(	4✔
77	self,
78	struct: Optional[MaybeSequence[Atoms]] = None,
79	struct_path: Optional[str] = None,
80	struct_name: Optional[str] = None,
81	architecture: Architectures = "mace_mp",
82	device: Devices = "cpu",
83	read_kwargs: Optional[ASEReadArgs] = None,
84	calc_kwargs: Optional[dict[str, Any]] = None,
85	log_kwargs: Optional[dict[str, Any]] = None,
86	) -> None:
87	"""
88	Read the structure being simulated and attach an MLIP calculator.
89
90	Parameters
91	----------
92	struct : Optional[MaybeSequence[Atoms]]
93	ASE Atoms structure(s) to simulate. Required if `struct_path`
94	is None. Default is None.
95	struct_path : Optional[str]
96	Path of structure to simulate. Required if `struct` is None.
97	Default is None.
98	struct_name : Optional[str]
99	Name of structure. Default is inferred from chemical formula if `struct`
100	is specified, else inferred from `struct_path`.
101	architecture : Architectures
102	MLIP architecture to use for single point calculations.
103	Default is "mace_mp".
104	device : Devices
105	Device to run MLIP model on. Default is "cpu".
106	read_kwargs : Optional[ASEReadArgs]
107	Keyword arguments to pass to ase.io.read. Default is {}.
108	calc_kwargs : Optional[dict[str, Any]]
109	Keyword arguments to pass to the selected calculator. Default is {}.
110	log_kwargs : Optional[dict[str, Any]]
111	Keyword arguments to pass to `config_logger`. Default is {}.
112	"""
113	if struct and struct_path:	4✔
114	raise ValueError(	4✔
115	"You cannot specify both the ASE Atoms structure (`struct`) "
116	"and a path to the structure file (`struct_path`)"
117	)
118
119	if not struct and not struct_path:	4✔
120	raise ValueError(	4✔
121	"Please specify either the ASE Atoms structure (`struct`) "
122	"or a path to the structure file (`struct_path`)"
123	)
124
125	read_kwargs = read_kwargs if read_kwargs else {}	4✔
126	calc_kwargs = calc_kwargs if calc_kwargs else {}	4✔
127	log_kwargs = log_kwargs if log_kwargs else {}	4✔
128
129	if log_kwargs and "filename" not in log_kwargs:	4✔
NEW 130	raise ValueError("'filename' must be included in `log_kwargs`")	×
131
132	log_kwargs.setdefault("name", __name__)	4✔
133	self.logger = config_logger(**log_kwargs)	4✔
134
135	self.architecture = architecture	4✔
136	self.device = device	4✔
137	self.struct_path = struct_path	4✔
138	self.struct_name = struct_name	4✔
139
140	# Read structure if given as path
141	if self.struct_path:	4✔
142	self.read_structure(**read_kwargs)	4✔
143	else:
144	self.struct = struct	4✔
145	if not self.struct_name:	4✔
146	self.struct_name = self.struct.get_chemical_formula()	4✔
147
148	# Configure calculator
149	self.set_calculator(**calc_kwargs)	4✔
150
151	if self.logger:	4✔
152	self.logger.info("Single point calculator configured")	4✔
153
154	def read_structure(self, **kwargs) -> None:	4✔
155	"""
156	Read structure and structure name.
157
158	If the file contains multiple structures, only the last configuration
159	will be read by default.
160
161	Parameters
162	----------
163	**kwargs
164	Keyword arguments passed to ase.io.read.
165	"""
166	if self.struct_path:	4✔
167	self.struct = read(self.struct_path, **kwargs)	4✔
168	if not self.struct_name:	4✔
169	self.struct_name = Path(self.struct_path).stem	4✔
170	else:
171	raise ValueError("`struct_path` must be defined")	×
172
173	def set_calculator(	4✔
174	self, read_kwargs: Optional[ASEReadArgs] = None, **kwargs
175	) -> None:
176	"""
177	Configure calculator and attach to structure.
178
179	Parameters
180	----------
181	read_kwargs : Optional[ASEReadArgs]
182	Keyword arguments to pass to ase.io.read. Default is {}.
183	**kwargs
184	Additional keyword arguments passed to the selected calculator.
185	"""
186	calculator = choose_calculator(	4✔
187	architecture=self.architecture,
188	device=self.device,
189	**kwargs,
190	)
191	if self.struct is None:	4✔
192	read_kwargs = read_kwargs if read_kwargs else {}	×
193	self.read_structure(**read_kwargs)	×
194
195	if isinstance(self.struct, list):	4✔
196	for struct in self.struct:	4✔
197	struct.calc = calculator	4✔
198	else:
199	self.struct.calc = calculator	4✔
200
201	def _get_potential_energy(self) -> MaybeList[float]:	4✔
202	"""
203	Calculate potential energy using MLIP.
204
205	Returns
206	-------
207	MaybeList[float]
208	Potential energy of structure(s).
209	"""
210	if isinstance(self.struct, list):	4✔
211	return [struct.get_potential_energy() for struct in self.struct]	4✔
212
213	return self.struct.get_potential_energy()	4✔
214
215	def _get_forces(self) -> MaybeList[ndarray]:	4✔
216	"""
217	Calculate forces using MLIP.
218
219	Returns
220	-------
221	MaybeList[ndarray]
222	Forces of structure(s).
223	"""
224	if isinstance(self.struct, list):	4✔
225	return [struct.get_forces() for struct in self.struct]	×
226
227	return self.struct.get_forces()	4✔
228
229	def _get_stress(self) -> MaybeList[ndarray]:	4✔
230	"""
231	Calculate stress using MLIP.
232
233	Returns
234	-------
235	MaybeList[ndarray]
236	Stress of structure(s).
237	"""
238	if isinstance(self.struct, list):	4✔
239	return [struct.get_stress() for struct in self.struct]	×
240
241	return self.struct.get_stress()	4✔
242
243	@staticmethod	4✔
244	def _remove_invalid_props(	4✔
245	struct: Atoms,
246	results: CalcResults = None,
247	properties: Collection[str] = (),
248	) -> None:
249	"""
250	Remove any invalid properties from calculated results.
251
252	Parameters
253	----------
254	struct : Atoms
255	ASE Atoms structure with attached calculator results.
256	results : CalcResults
257	Dictionary of calculated results. Default is {}.
258	properties : Collection[str]
259	Physical properties requested to be calculated. Default is ().
260	"""
261	results = results if results else {}	4✔
262
263	# Find any properties with non-finite values
264	rm_keys = [	4✔
265	prop
266	for prop in struct.calc.results
267	if not isfinite(struct.calc.results[prop]).all()
268	]
269
270	# Raise error if property was explicitly requested, otherwise remove
271	for prop in rm_keys:	4✔
272	if prop in properties:	4✔
273	raise ValueError(	4✔
274	f"'{prop}' contains non-finite values for this structure."
275	)
276	del struct.calc.results[prop]	4✔
277	if prop in results:	4✔
278	del results[prop]	4✔
279
280	def _clean_results(	4✔
281	self, results: CalcResults = None, properties: Collection[str] = ()
282	) -> None:
283	"""
284	Remove NaN and inf values from results and calc.results dictionaries.
285
286	Parameters
287	----------
288	results : CalcResults
289	Dictionary of calculated results. Default is {}.
290	properties : Collection[str]
291	Physical properties requested to be calculated. Default is ().
292	"""
293	results = results if results else {}	4✔
294
295	if isinstance(self.struct, list):	4✔
296	for image in self.struct:	4✔
297	self._remove_invalid_props(image, results, properties)	4✔
298	else:
299	self._remove_invalid_props(self.struct, results, properties)	4✔
300
301	def run_single_point(	4✔
302	self,
303	properties: MaybeSequence[str] = (),
304	write_results: bool = False,
305	write_kwargs: Optional[ASEWriteArgs] = None,
306	) -> CalcResults:
307	"""
308	Run single point calculations.
309
310	Parameters
311	----------
312	properties : MaybeSequence[str]
313	Physical properties to calculate. If not specified, "energy",
314	"forces", and "stress" will be returned.
315	write_results : bool
316	True to write out structure with results of calculations. Default is False.
317	write_kwargs : Optional[ASEWriteArgs],
318	Keyword arguments to pass to ase.io.write if saving structure with
319	results of calculations. Default is {}.
320
321	Returns
322	-------
323	CalcResults
324	Dictionary of calculated results.
325	"""
326	results: CalcResults = {}	4✔
327	if isinstance(properties, str):	4✔
328	properties = [properties]	4✔
329
330	for prop in properties:	4✔
331	if prop not in ["energy", "forces", "stress"]:	4✔
332	raise NotImplementedError(	4✔
333	f"Property '{prop}' cannot currently be calculated."
334	)
335
336	write_kwargs = write_kwargs if write_kwargs else {}	4✔
337	if write_kwargs and "filename" not in write_kwargs:	4✔
338	raise ValueError("'filename' must be included in write_kwargs")	×
339
340	if self.logger:	4✔
341	self.logger.info("Starting single point calculation")	4✔
342
343	if "energy" in properties or len(properties) == 0:	4✔
344	results["energy"] = self._get_potential_energy()	4✔
345	if "forces" in properties or len(properties) == 0:	4✔
346	results["forces"] = self._get_forces()	4✔
347	if "stress" in properties or len(properties) == 0:	4✔
348	results["stress"] = self._get_stress()	4✔
349
350	# Remove meaningless values from results e.g. stress for non-periodic systems
351	self._clean_results(results, properties=properties)	4✔
352
353	if self.logger:	4✔
354	self.logger.info("Single point calculation complete")	4✔
355
356	if write_results:	4✔
357	if "filename" not in write_kwargs:	4✔
358	filename = f"{self.struct_name}-results.xyz"	4✔
359	write_kwargs["filename"] = Path(".").absolute() / filename	4✔
360	write(images=self.struct, **write_kwargs)	4✔
361
362	return results	4✔

stfc / janus-core / 8325072180

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