15077298886

from typing import Any, Union
from warnings import warn

import numpy as np
from numpy.typing import NDArray

from .base import Parser
from .utils.file_utils import MultiDict
from .utils.periodic_table import PeriodicTable


class Parameters(Parser):
    """
    Handle files that control parameters for electronic structure calculations.

    If contributing a new parser, please subclass this class, add the new supported file
    type to _supported_files and match statement in this class' __init__, and call the super().__init__ method, include the new file
    type as a kwarg in the super().__init__ call. Optionally include the self.lines line
    in examples.

    ...

    Attributes
    ----------
    _supported_files : dict
        List of supported file types.
    """

    def __init__(self, **kwargs):
        # Parse file information and perform checks
        super().__init__(self._supported_files, **kwargs)

        self._check_binary(False)

    @property
    def _supported_files(self) -> dict:
        # FHI-aims, ...
        return {"control_in": ".in"}

    def __repr__(self) -> str:
        return f"{self.__class__.__name__}({self._format}={self._name})"

    def __init_subclass__(cls, **kwargs):
        # Revert back to the original __init_subclass__ method to avoid checking for
        # required methods in child class of this class too
        return super(Parser, cls).__init_subclass__(**kwargs)


class AimsControl(Parameters):
    """
    FHI-aims control file parser.

    ...

    Attributes
    ----------
    path: str
        path to the aims.out file
    lines: List[str]
        contents of the aims.out file

    Examples
    --------
    >>> ac = AimsControl(control_in="./control.in")
    """

    def __init__(self, **kwargs):
        super().__init__(**kwargs)

    # Use normal methods instead of properties for these methods as we want to specify
    # the setter method using kwargs instead of assigning the value as a dictionary.
    # Then, for consistency, keep get_keywords as a normal function.
    def get_keywords(self) -> MultiDict:
        """
        Get the keywords from the control.in file.

        Returns
        -------
        MultiDict
            Keywords in the control.in file.
        """
        keywords = MultiDict()

        for line in self.lines:
            # Stop at third keyword delimiter if ASE wrote the file
            spl = line.split()
            if len(spl) > 0 and spl[-1] == "(ASE)":
                n_delims = 0
                if line == "#" + ("=" * 79):
                    n_delims += 1
                    if n_delims == 3:
                        # Reached end of keywords
                        break

            elif "#" * 80 in line.strip():
                # Reached the basis set definitions
                break

            if len(spl) > 0 and line[0] != "#":
                keywords[spl[0]] = " ".join(spl[1:])

        return keywords

    def get_species(self) -> list[str]:
        """
        Get the species from a control.in file.

        Returns
        -------
        List[str]
            A list of the species in the control.in file.
        """
        species = []
        for line in self.lines:
            spl = line.split()
            if len(spl) > 0 and spl[0] == "species":
                species.append(line.split()[1])

        return species

    def get_default_basis_funcs(
        self, elements: Union[list[str], None] = None
    ) -> dict[str, str]:
        """
        Get the basis functions.

        Parameters
        ----------
        elements : List[str], optional, default=None
            The elements to parse the basis functions for as chemical symbols.

        Returns
        -------
        Dict[str, str]
            A dictionary of the basis functions for the specified elements.
        """
        # Check that the given elements are valid
        if elements is not None and not set(elements).issubset(
            set(PeriodicTable.element_symbols)
        ):
            raise ValueError("Invalid element(s) given")

        # Warn if the requested elements aren't found in control.in
        if elements is not None and not set(elements).issubset(self.get_species()):
            warn("Could not find all requested elements in control.in", stacklevel=2)

        basis_funcs = {}

        for i, line_1 in enumerate(self.lines):
            spl_1 = line_1.split()
            if "species" in spl_1[0]:
                species = spl_1[1]

                if elements is not None and species not in elements:
                    continue

                for line_2 in self.lines[i + 1 :]:
                    spl = line_2.split()
                    if "species" in spl[0]:
                        break

                    if "#" in spl[0]:
                        continue

                    if "hydro" in line_2:
                        if species in basis_funcs:
                            basis_funcs[species].append(line_2.strip())
                        else:
                            basis_funcs[species] = [line_2.strip()]

        return basis_funcs

    def add_keywords_and_save(self, *args: tuple[str, Any]) -> None:
        """
        Add keywords to the control.in file and write the new control.in to disk.

        Note that files written by ASE or in a format where the keywords are at the top
        of the file followed by the basis sets are the only formats that are supported
        by this function. The keywords need to be added in a Tuple format rather than as
        **kwargs because we need to be able to add multiple of the same keyword.

        Parameters
        ----------
        *args : Tuple[str, Any]
            Keywords to be added to the control.in file.
        """
        # Get the location of the start of the basis sets
        basis_set_start = False

        # if ASE wrote the file, use the 'add' point as the end of keywords delimiter
        # otherwise, use the start of the basis sets as 'add' point
        for i, line_1 in enumerate(self.lines):
            if line_1.strip() == "#" * 80:
                if self.lines[2].split()[-1] == "(ASE)":
                    for j, line_2 in enumerate(reversed(self.lines[:i])):
                        if line_2.strip() == "#" + ("=" * 79):
                            basis_set_start = i - j - 1
                            break
                    break

                # not ASE
                basis_set_start = i
                break

        # Check to make sure basis sets were found
        if not basis_set_start:
            raise IndexError("Could not detect basis sets in control.in")

        # Add the new keywords above the basis sets
        for arg in reversed(args):
            self.lines.insert(basis_set_start, f"{arg[0]:<34} {arg[1]}\n")

        # Write the file
        with open(self.path, "w") as f:
            f.writelines(self.lines)

    def add_cube_cell_and_save(
        self, cell_matrix: NDArray[Any], resolution: int = 100
    ) -> None:
        """
        Add cube output settings to control.in to cover the unit cell specified in
        `cell_matrix` and save to disk.

        Since the default behaviour of FHI-AIMS for generating CUBE files for periodic
        structures with vacuum gives confusing results, this function ensures the cube
        output quantity is calculated for the full unit cell.

        Parameters
        ----------
        cell_matrix : ArrayLike
            2D array defining the unit cell.

        resolution : int
            Number of cube voxels to use for the shortest side of the unit cell.

        """
        if not self.check_periodic():  # Fail for non-periodic structures
            raise TypeError("add_cube_cell doesn't support non-periodic structures")

        shortest_side = min(np.sum(cell_matrix, axis=1))
        resolution = shortest_side / 100.0

        cube_x = (
            2 * int(np.ceil(0.5 * np.linalg.norm(cell_matrix[0, :]) / resolution)) + 1
        )  # Number of cubes along x axis
        x_vector = cell_matrix[0, :] / np.linalg.norm(cell_matrix[0, :]) * resolution
        cube_y = (
            2 * int(np.ceil(0.5 * np.linalg.norm(cell_matrix[1, :]) / resolution)) + 1
        )
        y_vector = cell_matrix[1, :] / np.linalg.norm(cell_matrix[1, :]) * resolution
        cube_z = (
            2 * int(np.ceil(0.5 * np.linalg.norm(cell_matrix[2, :]) / resolution)) + 1
        )
        z_vector = cell_matrix[2, :] / np.linalg.norm(cell_matrix[2, :]) * resolution
        self.add_keywords_and_save(  # Add cube options to control.in
            (
                "cube",
                "origin {} {} {}\n".format(
                    *(np.transpose(cell_matrix @ [0.5, 0.5, 0.5]))
                )
                + "cube edge {} {} {} {}\n".format(cube_x, *x_vector)
                + "cube edge {} {} {} {}\n".format(cube_y, *y_vector)
                + "cube edge {} {} {} {}\n".format(cube_z, *z_vector),
            )
        )
        # print("\tCube voxel resolution is {} Å".format(resolution))

    def remove_keywords_and_save(self, *args: str) -> None:
        """
        Remove keywords from the control.in file and save to disk.

        Note that this will not remove keywords that are commented with a '#'.

        Parameters
        ----------
        *args : str
            Keywords to be removed from the control.in file.
        """
        for keyword in args:
            for i, line in enumerate(self.lines):
                spl = line.split()
                if len(spl) > 0 and spl[0] != "#" and keyword in line:
                    self.lines.pop(i)

        with open(self.path, "w") as f:
            f.writelines(self.lines)

    def check_periodic(self) -> bool:
        """Check if the system is periodic."""
        return "k_grid" in self.get_keywords()

1	from typing import Any, Union	1✔
2	from warnings import warn	1✔
3
4	import numpy as np	1✔
5	from numpy.typing import NDArray	1✔
6
7	from .base import Parser	1✔
8	from .utils.file_utils import MultiDict	1✔
9	from .utils.periodic_table import PeriodicTable	1✔
10
11
12	class Parameters(Parser):	1✔
13	"""
14	Handle files that control parameters for electronic structure calculations.
15
16	If contributing a new parser, please subclass this class, add the new supported file
17	type to _supported_files and match statement in this class' __init__, and call the super().__init__ method, include the new file
18	type as a kwarg in the super().__init__ call. Optionally include the self.lines line
19	in examples.
20
21	...
22
23	Attributes
24	----------
25	_supported_files : dict
26	List of supported file types.
27	"""
28
29	def __init__(self, **kwargs):	1✔
30	# Parse file information and perform checks
31	super().__init__(self._supported_files, **kwargs)	1✔
32
33	self._check_binary(False)	1✔
34
35	@property	1✔
36	def _supported_files(self) -> dict:	1✔
37	# FHI-aims, ...
38	return {"control_in": ".in"}	1✔
39
40	def __repr__(self) -> str:	1✔
41	return f"{self.__class__.__name__}({self._format}={self._name})"	×
42
43	def __init_subclass__(cls, **kwargs):	1✔
44	# Revert back to the original __init_subclass__ method to avoid checking for
45	# required methods in child class of this class too
46	return super(Parser, cls).__init_subclass__(**kwargs)	1✔
47
48
49	class AimsControl(Parameters):	1✔
50	"""
51	FHI-aims control file parser.
52
53	...
54
55	Attributes
56	----------
57	path: str
58	path to the aims.out file
59	lines: List[str]
60	contents of the aims.out file
61
62	Examples
63	--------
64	>>> ac = AimsControl(control_in="./control.in")
65	"""
66
67	def __init__(self, **kwargs):	1✔
68	super().__init__(**kwargs)	1✔
69
70	# Use normal methods instead of properties for these methods as we want to specify
71	# the setter method using kwargs instead of assigning the value as a dictionary.
72	# Then, for consistency, keep get_keywords as a normal function.
73	def get_keywords(self) -> MultiDict:	1✔
74	"""
75	Get the keywords from the control.in file.
76
77	Returns
78	-------
79	MultiDict
80	Keywords in the control.in file.
81	"""
82	keywords = MultiDict()	1✔
83
84	for line in self.lines:	1✔
85	# Stop at third keyword delimiter if ASE wrote the file
86	spl = line.split()	1✔
87	if len(spl) > 0 and spl[-1] == "(ASE)":	1✔
88	n_delims = 0	1✔
89	if line == "#" + ("=" * 79):	1✔
90	n_delims += 1	×
91	if n_delims == 3:	×
92	# Reached end of keywords
93	break	×
94
95	elif "#" * 80 in line.strip():	1✔
96	# Reached the basis set definitions
97	break	1✔
98
99	if len(spl) > 0 and line[0] != "#":	1✔
100	keywords[spl[0]] = " ".join(spl[1:])	1✔
101
102	return keywords	1✔
103
104	def get_species(self) -> list[str]:	1✔
105	"""
106	Get the species from a control.in file.
107
108	Returns
109	-------
110	List[str]
111	A list of the species in the control.in file.
112	"""
113	species = []	1✔
114	for line in self.lines:	1✔
115	spl = line.split()	1✔
116	if len(spl) > 0 and spl[0] == "species":	1✔
117	species.append(line.split()[1])	1✔
118
119	return species	1✔
120
121	def get_default_basis_funcs(	1✔
122	self, elements: Union[list[str], None] = None
123	) -> dict[str, str]:
124	"""
125	Get the basis functions.
126
127	Parameters
128	----------
129	elements : List[str], optional, default=None
130	The elements to parse the basis functions for as chemical symbols.
131
132	Returns
133	-------
134	Dict[str, str]
135	A dictionary of the basis functions for the specified elements.
136	"""
137	# Check that the given elements are valid
138	if elements is not None and not set(elements).issubset(	1✔
139	set(PeriodicTable.element_symbols)
140	):
141	raise ValueError("Invalid element(s) given")	×
142
143	# Warn if the requested elements aren't found in control.in
144	if elements is not None and not set(elements).issubset(self.get_species()):	1✔
145	warn("Could not find all requested elements in control.in", stacklevel=2)	×
146
147	basis_funcs = {}	1✔
148
149	for i, line_1 in enumerate(self.lines):	1✔
150	spl_1 = line_1.split()	1✔
151	if "species" in spl_1[0]:	1✔
152	species = spl_1[1]	1✔
153
154	if elements is not None and species not in elements:	1✔
155	continue	×
156
157	for line_2 in self.lines[i + 1 :]:	1✔
158	spl = line_2.split()	1✔
159	if "species" in spl[0]:	1✔
160	break	1✔
161
162	if "#" in spl[0]:	1✔
163	continue	1✔
164
165	if "hydro" in line_2:	1✔
166	if species in basis_funcs:	1✔
167	basis_funcs[species].append(line_2.strip())	1✔
168	else:
169	basis_funcs[species] = [line_2.strip()]	1✔
170
171	return basis_funcs	1✔
172
173	def add_keywords_and_save(self, *args: tuple[str, Any]) -> None:	1✔
174	"""
175	Add keywords to the control.in file and write the new control.in to disk.
176
177	Note that files written by ASE or in a format where the keywords are at the top
178	of the file followed by the basis sets are the only formats that are supported
179	by this function. The keywords need to be added in a Tuple format rather than as
180	**kwargs because we need to be able to add multiple of the same keyword.
181
182	Parameters
183	----------
184	*args : Tuple[str, Any]
185	Keywords to be added to the control.in file.
186	"""
187	# Get the location of the start of the basis sets
188	basis_set_start = False	1✔
189
190	# if ASE wrote the file, use the 'add' point as the end of keywords delimiter
191	# otherwise, use the start of the basis sets as 'add' point
192	for i, line_1 in enumerate(self.lines):	1✔
193	if line_1.strip() == "#" * 80:	1✔
194	if self.lines[2].split()[-1] == "(ASE)":	1✔
195	for j, line_2 in enumerate(reversed(self.lines[:i])):	1✔
196	if line_2.strip() == "#" + ("=" * 79):	1✔
197	basis_set_start = i - j - 1	1✔
198	break	1✔
199	break	1✔
200
201	# not ASE
202	basis_set_start = i	1✔
203	break	1✔
204
205	# Check to make sure basis sets were found
206	if not basis_set_start:	1✔
207	raise IndexError("Could not detect basis sets in control.in")	×
208
209	# Add the new keywords above the basis sets
210	for arg in reversed(args):	1✔
211	self.lines.insert(basis_set_start, f"{arg[0]:<34} {arg[1]}\n")	1✔
212
213	# Write the file
214	with open(self.path, "w") as f:	1✔
215	f.writelines(self.lines)	1✔
216
217	def add_cube_cell_and_save(	1✔
218	self, cell_matrix: NDArray[Any], resolution: int = 100
219	) -> None:
220	"""
221	Add cube output settings to control.in to cover the unit cell specified in
222	`cell_matrix` and save to disk.
223
224	Since the default behaviour of FHI-AIMS for generating CUBE files for periodic
225	structures with vacuum gives confusing results, this function ensures the cube
226	output quantity is calculated for the full unit cell.
227
228	Parameters
229	----------
230	cell_matrix : ArrayLike
231	2D array defining the unit cell.
232
233	resolution : int
234	Number of cube voxels to use for the shortest side of the unit cell.
235
236	"""
237	if not self.check_periodic(): # Fail for non-periodic structures	1✔
238	raise TypeError("add_cube_cell doesn't support non-periodic structures")	1✔
239
240	shortest_side = min(np.sum(cell_matrix, axis=1))	1✔
241	resolution = shortest_side / 100.0	1✔
242
243	cube_x = (	1✔
244	2 * int(np.ceil(0.5 * np.linalg.norm(cell_matrix[0, :]) / resolution)) + 1
245	) # Number of cubes along x axis
246	x_vector = cell_matrix[0, :] / np.linalg.norm(cell_matrix[0, :]) * resolution	1✔
247	cube_y = (	1✔
248	2 * int(np.ceil(0.5 * np.linalg.norm(cell_matrix[1, :]) / resolution)) + 1
249	)
250	y_vector = cell_matrix[1, :] / np.linalg.norm(cell_matrix[1, :]) * resolution	1✔
251	cube_z = (	1✔
252	2 * int(np.ceil(0.5 * np.linalg.norm(cell_matrix[2, :]) / resolution)) + 1
253	)
254	z_vector = cell_matrix[2, :] / np.linalg.norm(cell_matrix[2, :]) * resolution	1✔
255	self.add_keywords_and_save( # Add cube options to control.in	1✔
256	(
257	"cube",
258	"origin {} {} {}\n".format(
259	*(np.transpose(cell_matrix @ [0.5, 0.5, 0.5]))
260	)
261	+ "cube edge {} {} {} {}\n".format(cube_x, *x_vector)
262	+ "cube edge {} {} {} {}\n".format(cube_y, *y_vector)
263	+ "cube edge {} {} {} {}\n".format(cube_z, *z_vector),
264	)
265	)
266	# print("\tCube voxel resolution is {} Å".format(resolution))
267
268	def remove_keywords_and_save(self, *args: str) -> None:	1✔
269	"""
270	Remove keywords from the control.in file and save to disk.
271
272	Note that this will not remove keywords that are commented with a '#'.
273
274	Parameters
275	----------
276	*args : str
277	Keywords to be removed from the control.in file.
278	"""
279	for keyword in args:	1✔
280	for i, line in enumerate(self.lines):	1✔
281	spl = line.split()	1✔
282	if len(spl) > 0 and spl[0] != "#" and keyword in line:	1✔
283	self.lines.pop(i)	1✔
284
285	with open(self.path, "w") as f:	1✔
286	f.writelines(self.lines)	1✔
287
288	def check_periodic(self) -> bool:	1✔
289	"""Check if the system is periodic."""
290	return "k_grid" in self.get_keywords()	1✔

maurergroup / dfttoolkit / 15077298886

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