diff --git a/iodata/basis.py b/iodata/basis.py index 5dea1a13..510eec80 100644 --- a/iodata/basis.py +++ b/iodata/basis.py @@ -28,8 +28,9 @@ from numbers import Integral from typing import Union -import attr +import attrs import numpy as np +from numpy.typing import NDArray from .attrutils import validate_shape @@ -100,7 +101,7 @@ def angmom_its(angmom: Union[int, list[int]]) -> Union[str, list[str]]: return ANGMOM_CHARS[angmom] -@attr.s(auto_attribs=True, slots=True, on_setattr=[attr.setters.validate, attr.setters.convert]) +@attrs.define class Shell: """A shell in a molecular basis representing (generalized) contractions with the same exponents. @@ -126,11 +127,11 @@ class Shell: """ - icenter: int - angmoms: list[int] = attr.ib(validator=validate_shape(("coeffs", 1))) - kinds: list[str] = attr.ib(validator=validate_shape(("coeffs", 1))) - exponents: np.ndarray = attr.ib(validator=validate_shape(("coeffs", 0))) - coeffs: np.ndarray = attr.ib(validator=validate_shape(("exponents", 0), ("kinds", 0))) + icenter: int = attrs.field() + angmoms: list[int] = attrs.field(validator=validate_shape(("coeffs", 1))) + kinds: list[str] = attrs.field(validator=validate_shape(("coeffs", 1))) + exponents: NDArray = attrs.field(validator=validate_shape(("coeffs", 0))) + coeffs: NDArray = attrs.field(validator=validate_shape(("exponents", 0), ("kinds", 0))) @property def nbasis(self) -> int: @@ -156,7 +157,7 @@ def ncon(self) -> int: return len(self.angmoms) -@attr.s(auto_attribs=True, slots=True, on_setattr=[attr.setters.validate, attr.setters.convert]) +@attrs.define class MolecularBasis: """A complete molecular orbital or density basis set. @@ -205,9 +206,9 @@ class MolecularBasis: """ - shells: list[Shell] - conventions: dict[str, str] - primitive_normalization: str + shells: list[Shell] = attrs.field() + conventions: dict[str, str] = attrs.field() + primitive_normalization: str = attrs.field() @property def nbasis(self) -> int: @@ -222,12 +223,12 @@ def get_segmented(self): shells.append( Shell(shell.icenter, [angmom], [kind], shell.exponents, coeffs.reshape(-1, 1)) ) - return attr.evolve(self, shells=shells) + return attrs.evolve(self, shells=shells) def convert_convention_shell( conv1: list[str], conv2: list[str], reverse=False -) -> tuple[np.ndarray, np.ndarray]: +) -> tuple[NDArray, NDArray]: """Return a permutation vector and sign changes to convert from 1 to 2. The transformation from convention 1 to convention 2 can be done applying @@ -289,7 +290,7 @@ def convert_convention_shell( def convert_conventions( molbasis: MolecularBasis, new_conventions: dict[str, list[str]], reverse=False -) -> tuple[np.ndarray, np.ndarray]: +) -> tuple[NDArray, NDArray]: """Return a permutation vector and sign changes to convert from 1 to 2. The transformation from molbasis.convention to the new convention can be done @@ -339,7 +340,7 @@ def convert_conventions( return np.array(permutation), np.array(signs) -def iter_cart_alphabet(n: int) -> np.ndarray: +def iter_cart_alphabet(n: int) -> NDArray: """Loop over powers of Cartesian basis functions in alphabetical order. See https://theochem.github.io/horton/2.1.1/tech_ref_gaussian_basis.html diff --git a/iodata/formats/chgcar.py b/iodata/formats/chgcar.py index 40fc1800..619e1c53 100644 --- a/iodata/formats/chgcar.py +++ b/iodata/formats/chgcar.py @@ -26,6 +26,7 @@ """ import numpy as np +from numpy.typing import NDArray from ..docstrings import document_load_one from ..periodic import sym2num @@ -37,7 +38,7 @@ PATTERNS = ["CHGCAR*", "AECCAR*"] -def _load_vasp_header(lit: LineIterator) -> tuple[str, np.ndarray, np.ndarray, np.ndarray]: +def _load_vasp_header(lit: LineIterator) -> tuple[str, NDArray, NDArray, NDArray]: """Load the cell and atoms from a VASP file format. Parameters diff --git a/iodata/formats/cp2klog.py b/iodata/formats/cp2klog.py index 2fec908e..77d74a0b 100644 --- a/iodata/formats/cp2klog.py +++ b/iodata/formats/cp2klog.py @@ -21,6 +21,7 @@ from typing import Union import numpy as np +from numpy.typing import NDArray from scipy.special import factorialk from ..basis import HORTON2_CONVENTIONS, MolecularBasis, Shell, angmom_sti @@ -42,9 +43,7 @@ } -def _get_cp2k_norm_corrections( - ell: int, alphas: Union[float, np.ndarray] -) -> Union[float, np.ndarray]: +def _get_cp2k_norm_corrections(ell: int, alphas: Union[float, NDArray]) -> Union[float, NDArray]: """Compute the corrections for the normalization of the basis functions. This correction is needed because the CP2K atom code works with a different @@ -236,7 +235,7 @@ def _read_cp2k_occupations_energies( def _read_cp2k_orbital_coeffs( lit: LineIterator, oe: list[tuple[int, int, float, float]] -) -> dict[tuple[int, int], np.ndarray]: +) -> dict[tuple[int, int], NDArray]: """Read the expansion coefficients of the orbital from an open CP2K ATOM output. Parameters @@ -294,11 +293,11 @@ def _get_norb_nel(oe: list[tuple[int, int, float, float]]) -> tuple[int, float]: def _fill_orbitals( - orb_coeffs: np.ndarray, - orb_energies: np.ndarray, - orb_occupations: np.ndarray, + orb_coeffs: NDArray, + orb_energies: NDArray, + orb_occupations: NDArray, oe: list[tuple[int, int, float, float]], - coeffs: dict[tuple[int, int], np.ndarray], + coeffs: dict[tuple[int, int], NDArray], obasis: MolecularBasis, restricted: bool, ): diff --git a/iodata/formats/cube.py b/iodata/formats/cube.py index 0a9c3863..df190970 100644 --- a/iodata/formats/cube.py +++ b/iodata/formats/cube.py @@ -29,6 +29,7 @@ from typing import TextIO import numpy as np +from numpy.typing import NDArray from ..docstrings import document_dump_one, document_load_one from ..iodata import IOData @@ -42,7 +43,7 @@ def _read_cube_header( lit: LineIterator, -) -> tuple[str, np.ndarray, np.ndarray, np.ndarray, dict[str, np.ndarray], np.ndarray]: +) -> tuple[str, NDArray, NDArray, NDArray, dict[str, NDArray], NDArray]: """Load header data from a CUBE file object. Parameters @@ -62,7 +63,7 @@ def _read_cube_header( # skip the second line next(lit) - def read_grid_line(line: str) -> tuple[int, np.ndarray]: + def read_grid_line(line: str) -> tuple[int, NDArray]: """Read a grid line from the cube file.""" words = line.split() return ( @@ -83,7 +84,7 @@ def read_grid_line(line: str) -> tuple[int, np.ndarray]: cellvecs = axes * shape.reshape(-1, 1) cube = {"origin": origin, "axes": axes, "shape": shape} - def read_atom_line(line: str) -> tuple[int, float, np.ndarray]: + def read_atom_line(line: str) -> tuple[int, float, NDArray]: """Read an atomic number and coordinate from the cube file.""" words = line.split() return ( @@ -106,7 +107,7 @@ def read_atom_line(line: str) -> tuple[int, float, np.ndarray]: return title, atcoords, atnums, cellvecs, cube, atcorenums -def _read_cube_data(lit: LineIterator, cube: dict[str, np.ndarray]): +def _read_cube_data(lit: LineIterator, cube: dict[str, NDArray]): """Load cube data from a CUBE file object. Parameters @@ -150,10 +151,10 @@ def load_one(lit: LineIterator) -> dict: def _write_cube_header( f: TextIO, title: str, - atcoords: np.ndarray, - atnums: np.ndarray, - cube: dict[str, np.ndarray], - atcorenums: np.ndarray, + atcoords: NDArray, + atnums: NDArray, + cube: dict[str, NDArray], + atcorenums: NDArray, ): print(title, file=f) print("OUTER LOOP: X, MIDDLE LOOP: Y, INNER LOOP: Z", file=f) @@ -169,7 +170,7 @@ def _write_cube_header( print(f"{atnums[i]:5d} {q: 11.6f} {x: 11.6f} {y: 11.6f} {z: 11.6f}", file=f) -def _write_cube_data(f: TextIO, cube_data: np.ndarray, block_size: int): +def _write_cube_data(f: TextIO, cube_data: NDArray, block_size: int): counter = 0 for value in cube_data.flat: f.write(f" {value: 12.5E}") diff --git a/iodata/formats/fchk.py b/iodata/formats/fchk.py index 1d86c09c..62f62b00 100644 --- a/iodata/formats/fchk.py +++ b/iodata/formats/fchk.py @@ -23,6 +23,7 @@ from typing import Optional, TextIO import numpy as np +from numpy.typing import NDArray from ..basis import HORTON2_CONVENTIONS, MolecularBasis, Shell, convert_conventions from ..docstrings import document_dump_one, document_load_many, document_load_one @@ -473,7 +474,7 @@ def _load_dm(label: str, fchk: dict, result: dict, key: str): result[key] = _triangle_to_dense(fchk[label]) -def _triangle_to_dense(triangle: np.ndarray) -> np.ndarray: +def _triangle_to_dense(triangle: NDArray) -> NDArray: """Convert a symmetric matrix in triangular storage to a dense square matrix. Parameters @@ -512,7 +513,7 @@ def _dump_real_scalars(name: str, val: float, f: TextIO): print(f"{name:40} R {float(val): 16.8E}", file=f) -def _dump_integer_arrays(name: str, val: np.ndarray, f: TextIO): +def _dump_integer_arrays(name: str, val: NDArray, f: TextIO): """Dumper for a array of integers.""" nval = val.size if nval != 0: @@ -527,7 +528,7 @@ def _dump_integer_arrays(name: str, val: np.ndarray, f: TextIO): k = 0 -def _dump_real_arrays(name: str, val: np.ndarray, f: TextIO): +def _dump_real_arrays(name: str, val: NDArray, f: TextIO): """Dumper for a array of float.""" nval = val.size if nval != 0: diff --git a/iodata/formats/gamess.py b/iodata/formats/gamess.py index 0428420d..608bc916 100644 --- a/iodata/formats/gamess.py +++ b/iodata/formats/gamess.py @@ -19,6 +19,7 @@ """GAMESS punch file format.""" import numpy as np +from numpy.typing import NDArray from ..docstrings import document_load_one from ..utils import LineIterator, angstrom @@ -29,7 +30,7 @@ PATTERNS = ["*.dat"] -def _read_data(lit: LineIterator) -> tuple: +def _read_data(lit: LineIterator) -> tuple[str, str, list[str]]: """Extract ``title``, ``symmetry`` and ``symbols`` from the punch file.""" title = next(lit).strip() symmetry = next(lit).split()[0] @@ -46,7 +47,7 @@ def _read_data(lit: LineIterator) -> tuple: return title, symmetry, symbols -def _read_coordinates(lit: LineIterator, result: dict) -> tuple: +def _read_coordinates(lit: LineIterator, result: dict[str]) -> tuple[NDArray, NDArray]: """Extract ``numbers`` and ``coordinates`` from the punch file.""" for _ in range(2): next(lit) @@ -67,7 +68,7 @@ def _read_coordinates(lit: LineIterator, result: dict) -> tuple: return numbers, coordinates -def _read_energy(lit: LineIterator, result: dict) -> tuple: +def _read_energy(lit: LineIterator, result: dict[str]) -> tuple[float, NDArray]: """Extract ``energy`` and ``gradient`` from the punch file.""" energy = float(next(lit).split()[1]) natom = len(result["symbols"]) @@ -81,7 +82,7 @@ def _read_energy(lit: LineIterator, result: dict) -> tuple: return energy, gradient -def _read_hessian(lit: LineIterator, result: dict) -> np.ndarray: +def _read_hessian(lit: LineIterator, result: dict[str]) -> NDArray: """Extract ``hessian`` from the punch file.""" # check that $HESS is not already parsed if "athessian" in result: @@ -102,7 +103,7 @@ def _read_hessian(lit: LineIterator, result: dict) -> np.ndarray: return hessian -def _read_masses(lit: LineIterator, result: dict) -> np.ndarray: +def _read_masses(lit: LineIterator, result: dict[str]) -> NDArray: """Extract ``masses`` from the punch file.""" natom = len(result["symbols"]) masses = np.zeros(natom, float) @@ -119,7 +120,7 @@ def _read_masses(lit: LineIterator, result: dict) -> np.ndarray: "PUNCH", ["title", "energy", "grot", "atgradient", "athessian", "atmasses", "atnums", "atcoords"], ) -def load_one(lit: LineIterator) -> dict: +def load_one(lit: LineIterator) -> dict[str]: """Do not edit this docstring. It will be overwritten.""" result = {} while True: diff --git a/iodata/formats/gaussianlog.py b/iodata/formats/gaussianlog.py index b9341550..6a2815fb 100644 --- a/iodata/formats/gaussianlog.py +++ b/iodata/formats/gaussianlog.py @@ -28,6 +28,7 @@ """ import numpy as np +from numpy.typing import NDArray from ..docstrings import document_load_one from ..utils import LineIterator, set_four_index_element @@ -73,7 +74,7 @@ def load_one(lit: LineIterator) -> dict: return result -def _load_twoindex_g09(lit: LineIterator, nbasis: int) -> np.ndarray: +def _load_twoindex_g09(lit: LineIterator, nbasis: int) -> NDArray: """Load a two-index operator from a GAUSSIAN LOG file format. Parameters @@ -106,7 +107,7 @@ def _load_twoindex_g09(lit: LineIterator, nbasis: int) -> np.ndarray: return result -def _load_fourindex_g09(lit: LineIterator, nbasis: int) -> np.ndarray: +def _load_fourindex_g09(lit: LineIterator, nbasis: int) -> NDArray: """Load a four-index operator from a GAUSSIAN LOG file. Parameters diff --git a/iodata/formats/mol2.py b/iodata/formats/mol2.py index d28d5dcf..7e481bff 100644 --- a/iodata/formats/mol2.py +++ b/iodata/formats/mol2.py @@ -26,6 +26,7 @@ from typing import TextIO import numpy as np +from numpy.typing import NDArray from ..docstrings import ( document_dump_many, @@ -83,9 +84,7 @@ def load_one(lit: LineIterator) -> dict: return result -def _load_helper_atoms( - lit: LineIterator, natoms: int -) -> tuple[np.ndarray, np.ndarray, np.ndarray, tuple]: +def _load_helper_atoms(lit: LineIterator, natoms: int) -> tuple[NDArray, NDArray, NDArray, tuple]: """Load element numbers, coordinates and atomic charges.""" atnums = np.empty(natoms) atcoords = np.empty((natoms, 3)) @@ -112,7 +111,7 @@ def _load_helper_atoms( return atnums, atcoords, atchgs, attypes -def _load_helper_bonds(lit: LineIterator, nbonds: int) -> tuple[np.ndarray]: +def _load_helper_bonds(lit: LineIterator, nbonds: int) -> NDArray: """Load bond information. Each line in a bond definition has the following structure diff --git a/iodata/formats/molden.py b/iodata/formats/molden.py index 964752eb..9353f836 100644 --- a/iodata/formats/molden.py +++ b/iodata/formats/molden.py @@ -28,8 +28,9 @@ import copy from typing import TextIO, Union -import attr +import attrs import numpy as np +from numpy.typing import NDArray from ..basis import ( HORTON2_CONVENTIONS, @@ -225,9 +226,7 @@ def _load_low(lit: LineIterator) -> dict: return result -def _load_helper_atoms( - lit: LineIterator, cunit: float -) -> tuple[np.ndarray, np.ndarray, np.ndarray]: +def _load_helper_atoms(lit: LineIterator, cunit: float) -> tuple[NDArray, NDArray, NDArray]: """Load element numbers and coordinates.""" atnums = [] atcorenums = [] @@ -357,9 +356,9 @@ def _load_helper_coeffs(lit: LineIterator) -> tuple: def _is_normalized_properly( obasis: MolecularBasis, - atcoords: np.ndarray, - orb_alpha: np.ndarray, - orb_beta: np.ndarray, + atcoords: NDArray, + orb_alpha: NDArray, + orb_beta: NDArray, norm_threshold: float = 1e-4, ) -> bool: """Test the normalization of the occupied and virtual orbitals. @@ -551,7 +550,7 @@ def _fix_obasis_normalize_contractions(obasis: MolecularBasis) -> MolecularBasis fixed_shells = [] for shell in obasis.shells: shell_obasis = MolecularBasis( - [attr.evolve(shell, icenter=0)], obasis.conventions, obasis.primitive_normalization + [attrs.evolve(shell, icenter=0)], obasis.conventions, obasis.primitive_normalization ) # 2) Get the first diagonal element of the overlap matrix olpdiag = compute_overlap(shell_obasis, np.zeros((1, 3), float))[0, 0] diff --git a/iodata/formats/molekel.py b/iodata/formats/molekel.py index 96489e1f..299a6c0f 100644 --- a/iodata/formats/molekel.py +++ b/iodata/formats/molekel.py @@ -26,6 +26,7 @@ from typing import TextIO import numpy as np +from numpy.typing import NDArray from ..basis import MolecularBasis, Shell, angmom_its, angmom_sti, convert_conventions from ..docstrings import document_dump_one, document_load_one @@ -56,7 +57,7 @@ def _load_helper_charges(lit: LineIterator) -> dict: return {"mulliken": np.array(atcharges)} -def _load_helper_atoms(lit: LineIterator) -> tuple[np.ndarray, np.ndarray]: +def _load_helper_atoms(lit: LineIterator) -> tuple[NDArray, NDArray]: atnums = [] atcoords = [] for line in lit: @@ -105,7 +106,7 @@ def _load_helper_obasis(lit: LineIterator) -> MolecularBasis: return MolecularBasis(shells, CONVENTIONS, "L2") -def _load_helper_coeffs(lit: LineIterator, nbasis: int) -> tuple[np.ndarray, np.ndarray]: +def _load_helper_coeffs(lit: LineIterator, nbasis: int) -> tuple[NDArray, NDArray]: coeffs = [] energies = [] irreps = [] @@ -143,7 +144,7 @@ def _load_helper_coeffs(lit: LineIterator, nbasis: int) -> tuple[np.ndarray, np. return np.hstack(coeffs), np.array(energies), irreps -def _load_helper_occ(lit: LineIterator) -> np.ndarray: +def _load_helper_occ(lit: LineIterator) -> NDArray: occs = [] for line in lit: if line.strip() == "$END": diff --git a/iodata/formats/mwfn.py b/iodata/formats/mwfn.py index f625a24f..6ede364f 100644 --- a/iodata/formats/mwfn.py +++ b/iodata/formats/mwfn.py @@ -19,6 +19,7 @@ """Multiwfn MWFN file format.""" import numpy as np +from numpy.typing import NDArray from ..basis import HORTON2_CONVENTIONS, MolecularBasis, Shell from ..docstrings import document_load_one @@ -175,7 +176,7 @@ def _load_helper_shells(lit: LineIterator, nshell: int) -> dict: def _load_helper_section( lit: LineIterator, nprim: int, start: str, skip: int, dtype: np.dtype -) -> np.ndarray: +) -> NDArray: """Read single or multiple line(s) section.""" section = [] while len(section) < nprim: diff --git a/iodata/formats/orcalog.py b/iodata/formats/orcalog.py index e5619f0e..d6c38513 100644 --- a/iodata/formats/orcalog.py +++ b/iodata/formats/orcalog.py @@ -21,6 +21,7 @@ from typing import TextIO import numpy as np +from numpy.typing import NDArray from ..docstrings import document_load_one from ..utils import LineIterator @@ -87,7 +88,7 @@ def _helper_number_atoms(lit: LineIterator) -> int: return natom -def _helper_geometry(lit: TextIO, natom: int) -> tuple[np.ndarray, np.ndarray]: +def _helper_geometry(lit: TextIO, natom: int) -> tuple[NDArray, NDArray]: """Load coordinates form a ORCA output file format. Parameters @@ -119,7 +120,7 @@ def _helper_geometry(lit: TextIO, natom: int) -> tuple[np.ndarray, np.ndarray]: return atnums, atcoords -def _helper_scf_energies(lit: TextIO) -> tuple[np.ndarray, np.ndarray]: +def _helper_scf_energies(lit: TextIO) -> tuple[NDArray, NDArray]: """Load energies from each SCF cycle from a ORCA output file format. Parameters diff --git a/iodata/formats/qchemlog.py b/iodata/formats/qchemlog.py index 719d480d..320f9d46 100644 --- a/iodata/formats/qchemlog.py +++ b/iodata/formats/qchemlog.py @@ -22,6 +22,7 @@ """ import numpy as np +from numpy.typing import NDArray from ..docstrings import document_load_one from ..orbitals import MolecularOrbitals @@ -307,7 +308,7 @@ def _helper_orbital_energies_unrestricted(lit: LineIterator) -> tuple: return subdata -def _helper_section(start: str, end: str, lit: LineIterator, backward: bool = False) -> np.ndarray: +def _helper_section(start: str, end: str, lit: LineIterator, backward: bool = False) -> NDArray: """Load data between starting and ending strings.""" data = [] for line in lit: @@ -324,7 +325,7 @@ def _helper_section(start: str, end: str, lit: LineIterator, backward: bool = Fa return np.array(data, dtype=float) -def _helper_mulliken(lit: LineIterator) -> np.ndarray: +def _helper_mulliken(lit: LineIterator) -> NDArray: """Load mulliken net atomic charges.""" # skip line between 'Ground-State Mulliken Net Atomic Charges' line & atomic charge entries while True: @@ -358,7 +359,7 @@ def _helper_dipole_moments(lit: LineIterator) -> tuple: return dipole, quadrupole, dipole_tol -def _helper_polar(lit: LineIterator) -> np.ndarray: +def _helper_polar(lit: LineIterator) -> NDArray: """Load polarizability matrix.""" next(lit) polarizability_tensor = [] @@ -369,7 +370,7 @@ def _helper_polar(lit: LineIterator) -> np.ndarray: return np.array(polarizability_tensor, dtype=float) -def _helper_hessian(lit: LineIterator, natom: int) -> np.ndarray: +def _helper_hessian(lit: LineIterator, natom: int) -> NDArray: """Load hessian matrix.""" # hessian in Cartesian coordinates, shape(3 * natom, 3 * natom) col_idx = [int(i) for i in next(lit).split()] diff --git a/iodata/formats/wfn.py b/iodata/formats/wfn.py index dc97b915..aada13da 100644 --- a/iodata/formats/wfn.py +++ b/iodata/formats/wfn.py @@ -30,6 +30,7 @@ from typing import TextIO import numpy as np +from numpy.typing import NDArray from ..basis import MolecularBasis, Shell, convert_conventions from ..docstrings import document_dump_one, document_load_one @@ -137,7 +138,7 @@ def _load_helper_num(lit: LineIterator) -> list[int]: return num_mo, nprim, num_atoms -def _load_helper_atoms(lit: LineIterator, num_atoms: int) -> tuple[np.ndarray, np.ndarray]: +def _load_helper_atoms(lit: LineIterator, num_atoms: int) -> tuple[NDArray, NDArray]: """Read the coordinates of the atoms.""" atnums = np.empty(num_atoms, int) atcoords = np.empty((num_atoms, 3), float) @@ -158,7 +159,7 @@ def _load_helper_atoms(lit: LineIterator, num_atoms: int) -> tuple[np.ndarray, n def _load_helper_section( lit: LineIterator, n: int, start: str, skip: int, step: int, dtype: np.dtype -) -> np.ndarray: +) -> NDArray: """Read CENTRE ASSIGNMENTS, TYPE ASSIGNMENTS, and EXPONENTS sections.""" section = [] while len(section) < n: @@ -174,7 +175,7 @@ def _load_helper_section( return np.array(section, dtype=dtype) -def _load_helper_mo(lit: LineIterator, nprim: int) -> tuple[int, float, float, np.ndarray]: +def _load_helper_mo(lit: LineIterator, nprim: int) -> tuple[int, float, float, NDArray]: """Read one section of MO information.""" line = next(lit) if not line.startswith("MO"): @@ -201,7 +202,7 @@ def _load_helper_energy(lit: LineIterator) -> float: return energy, virial -def _load_helper_multiwfn(lit: LineIterator, num_mo: int) -> np.ndarray: +def _load_helper_multiwfn(lit: LineIterator, num_mo: int) -> NDArray: """Read MO spin information from MULTIWFN extension.""" for line in lit: if "$MOSPIN $END" in line: @@ -258,8 +259,8 @@ def load_wfn_low(lit: LineIterator) -> tuple: def build_obasis( - icenters: np.ndarray, type_assignments: np.ndarray, exponents: np.ndarray, lit: LineIterator -) -> tuple[MolecularBasis, np.ndarray]: + icenters: NDArray, type_assignments: NDArray, exponents: NDArray, lit: LineIterator +) -> tuple[MolecularBasis, NDArray]: """Construct a basis set using the arrays read from a WFN or WFX file. Parameters @@ -351,7 +352,7 @@ def build_obasis( return obasis, permutation -def get_mocoeff_scales(obasis: MolecularBasis) -> np.ndarray: +def get_mocoeff_scales(obasis: MolecularBasis) -> NDArray: """Get the L2-normalization of the un-normalized Cartesian basis functions. Parameters @@ -461,7 +462,7 @@ def _format_helper_section(header: str, skip: int, spec: str, nline: int) -> tup return f"{header[:skip].ljust(skip)}{spec * nline}", len(spec) -def _dump_helper_section(f: TextIO, data: np.ndarray, fmt: str, skip: int, step: int, nline: int): +def _dump_helper_section(f: TextIO, data: NDArray, fmt: str, skip: int, step: int, nline: int): """Write a CENTRE_ASSIGNMENTS, TYPE_ASSIGNMENTS, or EXPONENTS section to file ``f``.""" while len(data) > 0: chunk = data[:nline] diff --git a/iodata/inputs/common.py b/iodata/inputs/common.py index b95e3947..2955e956 100644 --- a/iodata/inputs/common.py +++ b/iodata/inputs/common.py @@ -18,7 +18,7 @@ # -- """Utilities for writing input files.""" -import attr +import attrs import numpy as np from ..iodata import IOData @@ -30,7 +30,7 @@ def populate_fields(data: IOData) -> dict: """Generate a dictionary with fields to replace in the template.""" # load IOData dict using attr.asdict because the IOData class uses __slots__ - fields = attr.asdict(data, recurse=False) + fields = attrs.asdict(data, recurse=False) # store atomic coordinates in angstrom fields["atcoords"] = data.atcoords / angstrom # set general defaults diff --git a/iodata/iodata.py b/iodata/iodata.py index 3edabb6e..4a89d03c 100644 --- a/iodata/iodata.py +++ b/iodata/iodata.py @@ -18,8 +18,11 @@ # -- """Module for handling input/output from different file formats.""" -import attr +from typing import Optional + +import attrs import numpy as np +from numpy.typing import NDArray from .attrutils import convert_array_to, validate_shape from .basis import MolecularBasis @@ -29,7 +32,7 @@ __all__ = ["IOData"] -@attr.s(auto_attribs=True, slots=True, on_setattr=[attr.setters.validate, attr.setters.convert]) +@attrs.define class IOData: """A container class for data loaded from (or to be written to) a file. @@ -172,78 +175,78 @@ class IOData: """ - atcharges: dict = attr.ib(factory=dict) - atcoords: np.ndarray = attr.ib( + atcharges: dict = attrs.field(factory=dict) + atcoords: Optional[NDArray] = attrs.field( default=None, converter=convert_array_to(float), - validator=attr.validators.optional(validate_shape("natom", 3)), + validator=attrs.validators.optional(validate_shape("natom", 3)), ) - _atcorenums: np.ndarray = attr.ib( + _atcorenums: Optional[NDArray] = attrs.field( default=None, converter=convert_array_to(float), - validator=attr.validators.optional(validate_shape("natom")), + validator=attrs.validators.optional(validate_shape("natom")), ) - atffparams: dict = attr.ib(factory=dict) - atfrozen: np.ndarray = attr.ib( + atffparams: dict = attrs.field(factory=dict) + atfrozen: Optional[NDArray] = attrs.field( default=None, converter=convert_array_to(bool), - validator=attr.validators.optional(validate_shape("natom")), + validator=attrs.validators.optional(validate_shape("natom")), ) - atgradient: np.ndarray = attr.ib( + atgradient: Optional[NDArray] = attrs.field( default=None, converter=convert_array_to(float), - validator=attr.validators.optional(validate_shape("natom", 3)), + validator=attrs.validators.optional(validate_shape("natom", 3)), ) - athessian: np.ndarray = attr.ib( + athessian: Optional[NDArray] = attrs.field( default=None, converter=convert_array_to(float), - validator=attr.validators.optional(validate_shape(None, None)), + validator=attrs.validators.optional(validate_shape(None, None)), ) - atmasses: np.ndarray = attr.ib( + atmasses: Optional[NDArray] = attrs.field( default=None, converter=convert_array_to(float), - validator=attr.validators.optional(validate_shape("natom")), + validator=attrs.validators.optional(validate_shape("natom")), ) - atnums: np.ndarray = attr.ib( + atnums: Optional[NDArray] = attrs.field( default=None, converter=convert_array_to(int), - validator=attr.validators.optional(validate_shape("natom")), + validator=attrs.validators.optional(validate_shape("natom")), ) - basisdef: str = None - bonds: np.ndarray = attr.ib( + basisdef: Optional[str] = attrs.field(default=None) + bonds: Optional[NDArray] = attrs.field( default=None, converter=convert_array_to(int), - validator=attr.validators.optional(validate_shape(None, 3)), + validator=attrs.validators.optional(validate_shape(None, 3)), ) - cellvecs: np.ndarray = attr.ib( + cellvecs: Optional[NDArray] = attrs.field( default=None, converter=convert_array_to(float), - validator=attr.validators.optional(validate_shape(None, 3)), + validator=attrs.validators.optional(validate_shape(None, 3)), ) - _charge: float = None - core_energy: float = None - cube: Cube = None - energy: float = None - extcharges: np.ndarray = attr.ib( + _charge: Optional[float] = attrs.field(default=None) + core_energy: Optional[float] = attrs.field(default=None) + cube: Optional[Cube] = attrs.field(default=None) + energy: Optional[float] = attrs.field(default=None) + extcharges: NDArray = attrs.field( default=None, converter=convert_array_to(float), - validator=attr.validators.optional(validate_shape(None, 4)), + validator=attrs.validators.optional(validate_shape(None, 4)), ) - extra: dict = attr.ib(factory=dict) - g_rot: float = None - lot: str = None - mo: MolecularOrbitals = None - moments: dict = attr.ib(factory=dict) - _nelec: float = None - obasis: MolecularBasis = None - obasis_name: str = None - one_ints: dict = attr.ib(factory=dict) - one_rdms: dict = attr.ib(factory=dict) - run_type: str = None - _spinpol: float = None - title: str = None - two_ints: dict = attr.ib(factory=dict) - two_rdms: dict = attr.ib(factory=dict) + extra: dict = attrs.field(factory=dict) + g_rot: Optional[float] = attrs.field(default=None) + lot: Optional[str] = attrs.field(default=None) + mo: Optional[MolecularOrbitals] = attrs.field(default=None) + moments: dict = attrs.field(factory=dict) + _nelec: Optional[float] = attrs.field(default=None) + obasis: Optional[MolecularBasis] = attrs.field(default=None) + obasis_name: Optional[str] = attrs.field(default=None) + one_ints: dict = attrs.field(factory=dict) + one_rdms: dict = attrs.field(factory=dict) + run_type: Optional[str] = attrs.field(default=None) + _spinpol: Optional[float] = attrs.field(default=None) + title: Optional[str] = attrs.field(default=None) + two_ints: dict = attrs.field(factory=dict) + two_rdms: dict = attrs.field(factory=dict) def __attrs_post_init__(self): # Trigger setter to acchieve consistency in properties @@ -262,7 +265,7 @@ def __attrs_post_init__(self): # Public interfaces to private attributes @property - def atcorenums(self) -> np.ndarray: + def atcorenums(self) -> NDArray: """Return effective core charges.""" if self._atcorenums is None and self.atnums is not None: self.atcorenums = self.atnums.astype(float) diff --git a/iodata/orbitals.py b/iodata/orbitals.py index bca69d26..9a862370 100644 --- a/iodata/orbitals.py +++ b/iodata/orbitals.py @@ -18,8 +18,11 @@ # -- """Data structure for molecular orbitals.""" -import attr +from typing import Optional + +import attrs import numpy as np +from numpy.typing import NDArray from .attrutils import convert_array_to, validate_shape @@ -55,7 +58,7 @@ def validate_norbab(mo, attribute, value): raise ValueError("In case of restricted orbitals, norba must be equal to norbb.") -@attr.s(auto_attribs=True, slots=True, on_setattr=[attr.setters.validate, attr.setters.convert]) +@attrs.define class MolecularOrbitals: """Class of Orthonormal Molecular Orbitals. @@ -92,28 +95,28 @@ class MolecularOrbitals: """ - kind: str = attr.ib( - validator=attr.validators.in_(["restricted", "unrestricted", "generalized"]) + kind: str = attrs.field( + validator=attrs.validators.in_(["restricted", "unrestricted", "generalized"]) ) - norba: int = attr.ib(validator=validate_norbab) - norbb: int = attr.ib(validator=validate_norbab) - occs: np.ndarray = attr.ib( + norba: int = attrs.field(validator=validate_norbab) + norbb: int = attrs.field(validator=validate_norbab) + occs: Optional[NDArray] = attrs.field( default=None, converter=convert_array_to(float), - validator=attr.validators.optional(validate_shape("norb")), + validator=attrs.validators.optional(validate_shape("norb")), ) - coeffs: np.ndarray = attr.ib( + coeffs: Optional[NDArray] = attrs.field( default=None, converter=convert_array_to(float), - validator=attr.validators.optional(validate_shape(None, "norb")), + validator=attrs.validators.optional(validate_shape(None, "norb")), ) - energies: np.ndarray = attr.ib( + energies: Optional[NDArray] = attrs.field( default=None, converter=convert_array_to(float), - validator=attr.validators.optional(validate_shape("norb")), + validator=attrs.validators.optional(validate_shape("norb")), ) - irreps: np.ndarray = attr.ib( - default=None, validator=attr.validators.optional(validate_shape("norb")) + irreps: Optional[NDArray] = attrs.field( + default=None, validator=attrs.validators.optional(validate_shape("norb")) ) @property diff --git a/iodata/overlap.py b/iodata/overlap.py index 85a9d8af..cca52633 100644 --- a/iodata/overlap.py +++ b/iodata/overlap.py @@ -20,7 +20,7 @@ from typing import Optional, Union -import attr +import attrs import numpy as np import scipy.special from numpy.typing import NDArray @@ -65,10 +65,10 @@ def factorial2(n: Union[int, NDArray[int]]) -> Union[int, NDArray[int]]: def compute_overlap( obasis0: MolecularBasis, - atcoords0: np.ndarray, + atcoords0: NDArray, obasis1: Optional[MolecularBasis] = None, - atcoords1: Optional[np.ndarray] = None, -) -> np.ndarray: + atcoords1: Optional[NDArray] = None, +) -> NDArray: r"""Compute overlap matrix for the given molecular basis set(s). .. math:: @@ -265,7 +265,7 @@ def compute_overlap_gaussian_1d(self, x1, x2, n1, n2, two_at): return value -def _compute_cart_shell_normalizations(shell: Shell) -> np.ndarray: +def _compute_cart_shell_normalizations(shell: Shell) -> NDArray: """Return normalization constants for the primitives in a given shell. Parameters @@ -275,12 +275,12 @@ def _compute_cart_shell_normalizations(shell: Shell) -> np.ndarray: Returns ------- - np.ndarray + NDArray The normalization constants, always for Cartesian functions, even when shell is pure. """ - shell = attr.evolve(shell, kinds=["c"] * shell.ncon) + shell = attrs.evolve(shell, kinds=["c"] * shell.ncon) result = [] for angmom in shell.angmoms: for exponent in shell.exponents: @@ -289,7 +289,7 @@ def _compute_cart_shell_normalizations(shell: Shell) -> np.ndarray: return np.array(result) -def gob_cart_normalization(alpha: np.ndarray, n: np.ndarray) -> np.ndarray: +def gob_cart_normalization(alpha: NDArray, n: NDArray) -> NDArray: """Compute normalization of exponent. Parameters @@ -301,7 +301,7 @@ def gob_cart_normalization(alpha: np.ndarray, n: np.ndarray) -> np.ndarray: Returns ------- - np.ndarray + NDArray The normalization constant for the gaussian cartesian basis. """ diff --git a/iodata/test/common.py b/iodata/test/common.py index 12b4f9f3..68080495 100644 --- a/iodata/test/common.py +++ b/iodata/test/common.py @@ -155,9 +155,9 @@ def check_orthonormal(mo_coeffs, ao_overlap, atol=1e-5): Parameters ---------- - mo_coeffs : np.ndarray, shape=(nbasis, mo_count) + mo_coeffs : NDArray, shape=(nbasis, mo_count) Molecular orbital coefficients. - ao_overlap : np.ndarray, shape=(nbasis, nbasis) + ao_overlap : NDArray, shape=(nbasis, nbasis) Atomic orbital overlap matrix. atol : float Absolute tolerance in deviation from identity matrix. diff --git a/iodata/test/test_attrutils.py b/iodata/test/test_attrutils.py index cdf5cdba..d4e99935 100644 --- a/iodata/test/test_attrutils.py +++ b/iodata/test/test_attrutils.py @@ -18,19 +18,20 @@ # -- """Unit tests for iodata.attrutils.""" -import attr +import attrs import numpy as np import pytest from numpy.testing import assert_allclose +from numpy.typing import NDArray from ..attrutils import convert_array_to, validate_shape -@attr.s(auto_attribs=True, slots=True, on_setattr=attr.setters.convert) +@attrs.define class FooBar: """Just a silly class for testing convert_array_to.""" - spam: np.ndarray = attr.ib(converter=convert_array_to(float)) + spam: NDArray = attrs.field(converter=convert_array_to(float)) def test_convert_array_to_init(): @@ -59,22 +60,22 @@ def test_convert_array_to_assign(): assert fb.spam is None -@attr.s(auto_attribs=True, slots=True, on_setattr=attr.setters.validate) +@attrs.define class Spam: """Just a silly class for testing validate_shape.""" - egg0: np.ndarray = attr.ib(validator=validate_shape(1, None, None)) - egg1: np.ndarray = attr.ib(validator=validate_shape(("egg0", 2), ("egg2", 1))) - egg2: np.ndarray = attr.ib(validator=validate_shape(2, ("egg1", 1))) - egg3: np.ndarray = attr.ib(validator=validate_shape(("leg", 0))) - leg: str = attr.ib(validator=validate_shape(("egg3", 0))) + egg0: NDArray = attrs.field(validator=validate_shape(1, None, None)) + egg1: NDArray = attrs.field(validator=validate_shape(("egg0", 2), ("egg2", 1))) + egg2: NDArray = attrs.field(validator=validate_shape(2, ("egg1", 1))) + egg3: NDArray = attrs.field(validator=validate_shape(("leg", 0))) + leg: str = attrs.field(validator=validate_shape(("egg3", 0))) def test_validate_shape_init(): # Construct a Spam instance with valid arguments. This should just work spam = Spam(np.zeros((1, 7, 4)), np.zeros((4, 3)), np.zeros((2, 3)), np.zeros(5), "abcde") # Double check - attr.validate(spam) + attrs.validate(spam) # Call constructor with invalid arguments with pytest.raises(TypeError): _ = Spam( @@ -131,7 +132,7 @@ def test_validate_shape_assign(): # Construct a Spam instance with valid arguments. This should just work spam = Spam(np.zeros((1, 7, 4)), np.zeros((4, 3)), np.zeros((2, 3)), np.zeros(5), "abcde") # Double check - attr.validate(spam) + attrs.validate(spam) # assign invalid attributes with pytest.raises(TypeError): spam.egg0 = np.zeros((2, 7, 4)) @@ -145,33 +146,33 @@ def test_validate_shape_assign(): spam.leg = "abcd" -@attr.s(slots=True) +@attrs.define class NoName0: """Test exception in validate_shape: unsupported item in shape_requirements.""" - xxx: str = attr.ib(validator=validate_shape(["asdfsa", 3])) + xxx: str = attrs.field(validator=validate_shape(["asdfsa", 3])) -@attr.s(slots=True) +@attrs.define class NoName1: """Test exception in validate_shape: unsupported item in shape_requirements.""" - xxx: str = attr.ib(validator=validate_shape(("asdfsa",))) + xxx: str = attrs.field(validator=validate_shape(("asdfsa",))) -@attr.s(slots=True) +@attrs.define class NoName2: """Test exception in validate_shape: other doest not exist.""" - xxx: str = attr.ib(validator=validate_shape("other")) + xxx: str = attrs.field(validator=validate_shape("other")) -@attr.s(slots=True) +@attrs.define class NoName3: """Test exception in validate_shape: other is not an array.""" - xxx: str = attr.ib(validator=validate_shape(("other", 1))) - other = attr.ib() + xxx: str = attrs.field(validator=validate_shape(("other", 1))) + other = attrs.field() def test_validate_shape_exceptions(): diff --git a/iodata/test/test_basis.py b/iodata/test/test_basis.py index 2fe6dc24..e076e77c 100644 --- a/iodata/test/test_basis.py +++ b/iodata/test/test_basis.py @@ -18,7 +18,7 @@ # -- """Unit tests for iodata.obasis.""" -import attr +import attrs import numpy as np import pytest from numpy.testing import assert_equal @@ -121,7 +121,7 @@ def test_shell_validators(): # It should not raise a TypeError. shell = Shell(0, [0, 0], ["c", "c"], np.zeros(6), np.zeros((6, 2))) # Rerun the validators as a double check. - attr.validate(shell) + attrs.validate(shell) # Tests with invalid constructor arguments. with pytest.raises(TypeError): Shell(0, [0, 0], ["c", "c"], np.zeros(6), np.zeros((6, 2, 2))) diff --git a/iodata/test/test_json.py b/iodata/test/test_json.py index 5e70f3be..4b205c20 100644 --- a/iodata/test/test_json.py +++ b/iodata/test/test_json.py @@ -34,7 +34,7 @@ # Tests for qcschema_molecule -# GEOMS: dict of str: np.ndarray(N, 3) +# GEOMS: dict of str: NDArray(N, 3) GEOMS = { "LiCl": np.array([[0.000000, 0.000000, -1.631761], [0.000000, 0.000000, 0.287958]]), "OHr": np.array([[0.0, 0.0, -0.12947694], [0.0, -1.49418734, 1.02744651]]), diff --git a/iodata/test/test_molden.py b/iodata/test/test_molden.py index a4ffe842..bab0e8df 100644 --- a/iodata/test/test_molden.py +++ b/iodata/test/test_molden.py @@ -21,7 +21,7 @@ import os import warnings -import attr +import attrs import numpy as np import pytest from numpy.testing import assert_allclose, assert_equal @@ -557,7 +557,7 @@ def test_load_dump_consistency(tmpdir, fn, match): mol1.obasis = mol1.obasis.get_segmented() # - Set default irreps in mol1, if not present. if mol1.mo.irreps is None: - mol1.mo = attr.evolve(mol1.mo, irreps=["1a"] * mol1.mo.norb) + mol1.mo = attrs.evolve(mol1.mo, irreps=["1a"] * mol1.mo.norb) # - Remove the one_rdms from mol1. mol1.one_rdms = {} compare_mols(mol1, mol2) diff --git a/iodata/test/test_overlap.py b/iodata/test/test_overlap.py index 9a2e5bc4..6ad705a2 100644 --- a/iodata/test/test_overlap.py +++ b/iodata/test/test_overlap.py @@ -20,7 +20,7 @@ import itertools -import attr +import attrs import numpy as np import pytest from numpy.testing import assert_allclose @@ -102,7 +102,7 @@ def test_load_fchk_o2_cc_pvtz_cart_num(): ref = np.load(str(fn_npy)) with as_file(files("iodata.test.data").joinpath("o2_cc_pvtz_cart.fchk")) as fn_fchk: data = load_one(fn_fchk) - obasis = attr.evolve(data.obasis, conventions=OVERLAP_CONVENTIONS) + obasis = attrs.evolve(data.obasis, conventions=OVERLAP_CONVENTIONS) assert_allclose(ref, compute_overlap(obasis, data.atcoords), rtol=1.0e-5, atol=1.0e-8) @@ -154,7 +154,7 @@ def test_overlap_two_basis_different(fn0, fn1): # overlap matrix. atcoords = np.concatenate([mol0.atcoords, mol1.atcoords]) shells = mol0.obasis.shells + [ - attr.evolve(shell, icenter=shell.icenter + mol0.natom) for shell in mol1.obasis.shells + attrs.evolve(shell, icenter=shell.icenter + mol0.natom) for shell in mol1.obasis.shells ] obasis = MolecularBasis(shells, OVERLAP_CONVENTIONS, "L2") olp_big = compute_overlap(obasis, atcoords) diff --git a/iodata/utils.py b/iodata/utils.py index 724abc06..8cf238e8 100644 --- a/iodata/utils.py +++ b/iodata/utils.py @@ -20,9 +20,10 @@ import warnings -import attr +import attrs import numpy as np import scipy.constants as spc +from numpy.typing import NDArray from scipy.linalg import eigh from .attrutils import validate_shape @@ -119,7 +120,7 @@ def back(self, line): self.lineno -= 1 -@attr.s(auto_attribs=True, slots=True, on_setattr=[attr.setters.validate, attr.setters.convert]) +@attrs.define class Cube: """The volumetric data from a cube (or similar) file. @@ -136,9 +137,9 @@ class Cube: """ - origin: np.ndarray = attr.ib(validator=validate_shape(3)) - axes: np.ndarray = attr.ib(validator=validate_shape(3, 3)) - data: np.ndarray = attr.ib(validator=validate_shape(None, None, None)) + origin: NDArray = attrs.field(validator=validate_shape(3)) + axes: NDArray = attrs.field(validator=validate_shape(3, 3)) + data: NDArray = attrs.field(validator=validate_shape(None, None, None)) @property def shape(self): @@ -147,7 +148,7 @@ def shape(self): def set_four_index_element( - four_index_object: np.ndarray, i0: int, i1: int, i2: int, i3: int, value: float + four_index_object: NDArray, i0: int, i1: int, i2: int, i3: int, value: float ): """Assign values to a four index object, account for 8-fold index symmetry. @@ -174,7 +175,7 @@ def set_four_index_element( four_index_object[i3, i0, i1, i2] = value -def volume(cellvecs: np.ndarray) -> float: +def volume(cellvecs: NDArray) -> float: """Calculate the (generalized) cell volume. Parameters @@ -200,7 +201,7 @@ def volume(cellvecs: np.ndarray) -> float: raise ValueError("Argument cellvecs should be of shape (x, 3), where x is in {1, 2, 3}") -def derive_naturals(dm: np.ndarray, overlap: np.ndarray) -> tuple[np.ndarray, np.ndarray]: +def derive_naturals(dm: NDArray, overlap: NDArray) -> tuple[NDArray, NDArray]: """Derive natural orbitals from a given density matrix. Parameters @@ -232,7 +233,7 @@ def derive_naturals(dm: np.ndarray, overlap: np.ndarray) -> tuple[np.ndarray, np return coeffs, occs -def check_dm(dm: np.ndarray, overlap: np.ndarray, eps: float = 1e-4, occ_max: float = 1.0): +def check_dm(dm: NDArray, overlap: NDArray, eps: float = 1e-4, occ_max: float = 1.0): """Check if the density matrix has eigenvalues in the proper range. Parameters