diff --git a/.github/dependabot.yml b/.github/dependabot.yml new file mode 100644 index 0000000..8b73709 --- /dev/null +++ b/.github/dependabot.yml @@ -0,0 +1,13 @@ +version: 2 +updates: +- package-ecosystem: pip + directory: "/" + schedule: + interval: weekly + open-pull-requests-limit: 2 + allow: + - dependency-type: direct + - dependency-type: indirect + ignore: + - dependency-name: mistune + - dependency-name: numpy diff --git a/.github/workflows/testing.yml b/.github/workflows/testing.yml index f7f76a3..4896efe 100644 --- a/.github/workflows/testing.yml +++ b/.github/workflows/testing.yml @@ -2,46 +2,22 @@ on: push: branches: - main - tags: - 'v*' - pull_request: branches: - main - - jobs: - - paper: - runs-on: ubuntu-latest - name: Paper Draft - steps: - - name: Checkout - uses: actions/checkout@v2 - - name: Build draft PDF - uses: openjournals/openjournals-draft-action@master - with: - journal: joss - # This should be the path to the paper within your repo. - paper-path: paper/paper.md - - name: Upload - uses: actions/upload-artifact@v1 - with: - name: paper - # This is the output path where Pandoc will write the compiled - # PDF. Note, this should be the same directory as the input - # paper.md - path: paper/paper.pdf lint: runs-on: ubuntu-latest steps: - - uses: actions/checkout@v2 - - - uses: actions/setup-python@v2 + - uses: actions/checkout@v3 + - uses: actions/setup-python@v4 with: - python-version: 3.8 + python-version: "3.10" + cache: pip + cache-dependency-path: pyproject.toml - name: Install dependencies run: | @@ -58,7 +34,7 @@ jobs: runs-on: ubuntu-latest strategy: matrix: - python-version: [ 3.8, 3.9 ] + python-version: [ "3.9", "3.10", "3.11" ] steps: - uses: actions/checkout@v2.3.4 @@ -101,10 +77,11 @@ jobs: - uses: actions/setup-python@v3 with: - python-version: 3.9 + python-version: 3.11 - name: Install dependencies run: | + python -m pip install --upgrade pip pip install -e .[strict] pip install -e .[docs] - name: Build diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index 0137b3b..13a2ec5 100644 --- a/.pre-commit-config.yaml +++ b/.pre-commit-config.yaml @@ -1,63 +1,36 @@ default_language_version: python: python3 -exclude: '^src/atomate2/vasp/schemas/calc_types/' repos: -- repo: https://github.com/charliermarsh/ruff-pre-commit - rev: v0.0.250 - hooks: - - id: ruff - args: [--fix] -- repo: https://github.com/pre-commit/pre-commit-hooks - rev: v4.4.0 - hooks: - - id: check-yaml - - id: fix-encoding-pragma - args: [--remove] - - id: end-of-file-fixer - - id: trailing-whitespace -- repo: https://github.com/psf/black - rev: 22.12.0 - hooks: - - id: black -- repo: https://github.com/asottile/blacken-docs - rev: v1.12.1 - hooks: - - id: blacken-docs - additional_dependencies: [black] - exclude: README.md -- repo: https://github.com/pycqa/flake8 - rev: 6.0.0 - hooks: - - id: flake8 - entry: pflake8 - files: ^src/ - additional_dependencies: - - pyproject-flake8==6.0.0 - - flake8-bugbear==22.12.6 - - flake8-typing-imports==1.14.0 - - flake8-docstrings==1.6.0 - - flake8-rst-docstrings==0.3.0 - - flake8-rst==0.8.0 -- repo: https://github.com/pre-commit/pygrep-hooks - rev: v1.10.0 - hooks: - - id: python-use-type-annotations - - id: rst-backticks - - id: rst-directive-colons - - id: rst-inline-touching-normal +- repo: https://github.com/astral-sh/ruff-pre-commit + # Ruff version. + rev: v0.1.8 + hooks: + # Run the linter. + - id: ruff + args: [ --fix ] + # Run the formatter. + - id: ruff-format - repo: https://github.com/pre-commit/mirrors-mypy - rev: v0.991 + rev: v1.5.1 hooks: - id: mypy files: ^src/ + args: + - --namespace-packages + - --explicit-package-bases additional_dependencies: - tokenize-rt==4.1.0 - types-pkg_resources==0.1.2 - - types-paramiko - repo: https://github.com/codespell-project/codespell - rev: v2.2.2 + rev: v2.2.6 hooks: - id: codespell - stages: [commit, commit-msg] - args: [--ignore-words-list, 'titel,statics,ba,nd,te'] - types_or: [python, rst, markdown] + name: codespell + description: Checks for common misspellings in text files. + entry: codespell + language: python + types: [text] + args: [ + --ignore-words-list, 'titel,statics,ba,nd,te,mater,commun,vise,dscribe', + --skip, "*.ipynb,./tests,*paper*", + ] \ No newline at end of file diff --git a/paper/fig1.png b/paper/fig1.png deleted file mode 100644 index 37fbda0..0000000 Binary files a/paper/fig1.png and /dev/null differ diff --git a/paper/paper.bib b/paper/paper.bib deleted file mode 100644 index 73bc52e..0000000 --- a/paper/paper.bib +++ /dev/null @@ -1,40 +0,0 @@ -@article{Kresse1993Jan, - author = {Kresse, G. and Hafner, J.}, - title = {{Ab initio molecular dynamics for liquid metals}}, - journal = {Phys. Rev. B}, - volume = {47}, - number = {1}, - pages = {558--561(R)}, - year = {1993}, - month = jan, - issn = {2469-9969}, - publisher = {American Physical Society}, - doi = {10.1103/PhysRevB.47.558} -} - -@article{Giannozzi2009Sep, - author = {Giannozzi, P. and Baroni, S. and Bonini, N. and Calandra, M. and Car, R. and Cavazzoni, C. and Ceresoli, D. and Chiarotti, G. L. and Cococcioni, M. and Dabo, I. and de Gironcoli, S. and Fabris, S. and Fratesi, G. and Gebauer, R. and Gerstmann, U. and Gougoussis, C. and Kokalj, A. and Lazzeri, M. and Martin-Samos, L. and Marzari, N. and Mauri, F. and Mazzarello, R. and Paolini, S. and Pasquarello, A. and Paulatto, L. and Sbraccia, C. and Seitsonen, A. P. J. and Smogunov, A. and Umari, P. and Vanderbilt, D.}, - title = {{QUANTUM ESPRESSO: A modular and open-source software project for quantum simulations of materials}}, - journal = {J. Phys.: Condens. Matter}, - volume = {21}, - number = {39}, - pages = {395502}, - year = {2009}, - month = sep, - issn = {0953-8984}, - publisher = {IOP Publishing}, - doi = {10.1088/0953-8984/21/39/395502} -} - -@article{Ong2013Feb, - author = {Ong, Shyue Ping and Richards, William Davidson and Jain, Anubhav and Hautier, Geoffroy and Kocher, Michael and Cholia, Shreyas and Gunter, Dan and Chevrier, Vincent L. and Persson, Kristin A. and Ceder, Gerbrand}, - title = {{Python Materials Genomics (pymatgen): A robust, open-source python library for materials analysis}}, - journal = {Comput. Mater. Sci.}, - volume = {68}, - pages = {314--319}, - year = {2013}, - month = feb, - issn = {0927-0256}, - publisher = {Elsevier}, - doi = {10.1016/j.commatsci.2012.10.028} -} diff --git a/paper/paper.md b/paper/paper.md deleted file mode 100644 index 2feb21d..0000000 --- a/paper/paper.md +++ /dev/null @@ -1,70 +0,0 @@ ---- -title: 'mp-PyRho: Regridding periodic charge densities for machine-learning applications' - -tags: - - python - - materials science - - machine learning - - periodic data -authors: - - name: Jimmy-Xuan Shen - orcid: 0000-0002-2743-7531 - affiliation: "1, 2" # (Multiple affiliations must be quoted) - - name: Kristin A. Persson - orcid: 0000-0002-8410-4639 - affiliation: "1, 3" -affiliations: - - name: Department of Materials Science and Engineering, University of California, Berkeley, Berkeley, California 94720, United States - index: 1 - - name: Lawrence Livermore National Laboratory, Livermore, California 94550, United States - index: 2 - - name: Energy Sciences Area, Lawrence Berkeley National Laboratory, Berkeley, California 94720, United States - index: 3 -date: October 2022 -bibliography: paper.bib ---- - -# Summary - -The electronic charge density is a central quantity in the field of computational material science. -Since most materials simulations codes, such as the Vienna Ab initio Simulation Package (VASP)[@Kresse1993Jan] and Quantum ESPRESSO[@Giannozzi2009Sep], assume periodic boundary conditions, the calculations usually stores the charge densities on a three-dimensional grid that is regular in the directions of the lattice vectors. -This makes the calculations, especially performing FFT's on these charge densities straightforward. -However, these non-orthogonal, and more importantly material-specific grids, means that we cannot directly compare the charge densities from different simulations. -Despite how data-rich the charge densities are, using periodic charge densities for machine learning is difficult because: - -1. Different charge densities are usually stored on non-commensurate grids, meaning that one deep-learning Convolutional Neural Network (CNN) cannot be trained to process different crystal structures. - -2. Only a single charge density is generated for each simulation, this is not ideal for Deep Learning since we want to train on data sets that represent many copies of the same system. This is like trying to train a CNN to recognize a cat by only showing it one picture of a cat. - -The `mp-pyRho` package aims to give users full control over the how the charge density is represented. -The input data can be thought of as an infinite periodic field, and the goal of this package is to allow the users to crop out a parallelepiped of arbitrary shape, position, orientation, and grid density. -If the output grids are cubic with the same dimensions, then they can be used as inputs for deep learning applications. - -`mp-pyRho` is a Python package that allows users to control the representation of periodic data in arbitrary dimensions. -The mathematics of the re-gridding relies on `scipy.interpolate` and is general for all dimensions. -However, special consideration was given to the case of 3D data, where we have integrated with the `VolumetricData` objects from the `pymatgen` package [Ong2013Feb]. - -The code is responsible for: -1. Up-scaling the periodic data to a higher resolution grid using Fourier interpolation. - -2. Re-gridding the data on a new grid, with arbitrary shape, position, and orientation. - -3. Generate arbitrary crops of the periodic data. - -All three of these capabilities can be seen in Figure 1. - -![Figure 1. (a) Demonstration of Fourier interpolation of a 2D periodic field. (b) Demonstration of regridding on a 2D periodic field. (c) Demonstration of cropping of a 3D periodic field.](fig1.png) - -As charge density data become more availibe through modern quantum chemistry databases, The cropping capabilities of `mp-pyRho` can process these data sets and make them more accessible for machine learning applications. - - -# Statement of need - -Grid-interpolation and grid-restructuring are common tasks in computational geometry, and efficient algorithms exist to perform these tasks. -However, since 3D periodic data is extremely common in computational materials science and not common elsewhere, there is a lack of software that is specifically designed to handle this type of data. -`mp-pyRho` is designed to be a general tool for regridding periodic data, and we have integrated it with the `VolumetricData` objects from the `pymatgen` package [@Ong2013Feb] to make it easy to use for materials science applications. - -# Acknowledgements - -We acknowledge contributions from Jason M. Munro, Matthew K. Horton, and Shyam Dwaraknath for their helpful discussions during the development of this software and Kristin A. Persson for supervising during the development of this software. -This work was supported by the US Department of Energy, Office of Science, Office of Basic Energy Sciences, Materials Sciences and Engineering Division under contract no. DE-AC02-05-CH11231 (Materials Project program KC23MP). diff --git a/pyproject.toml b/pyproject.toml index 262162b..2c03c44 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -16,7 +16,7 @@ classifiers = [ ] authors = [{name = "Jimmy-Xuan Shen", email = "jmmshn@gmail.com"}] dependencies = [ - 'pymatgen>=2022.9.21', + 'pymatgen>=2023.9.21', ] description = "Tools for re-griding periodic volumetric quantum chemistry data for machine-learning purposes." @@ -27,18 +27,20 @@ name = "mp-pyrho" readme = "README.md" requires-python = ">=3.8" +[tool.setuptools_scm] + [project.optional-dependencies] dev = ["pre-commit==2.17.0"] docs = [ - "numpydoc==1.3.1", - "sphinx==5.0.2", - "furo==2022.6.21", - "m2r2==0.3.2", + "numpydoc==1.6.0", + "sphinx==7.2.6", + "furo==2024.1.29", + "m2r2==v0.3.3.post2", "ipython==8.4.0", "nbsphinx==0.8.9", "nbsphinx-link==1.3.0", - "sphinx-autodoc-typehints==1.18.3", - "sphinx-autoapi==1.8.4", + "sphinx-autodoc-typehints==2.0.0", + "sphinx-autoapi==3.0.0", ] tests = [ "pytest==7.1.2", @@ -46,7 +48,7 @@ tests = [ "hypothesis==6.40.1", ] strict = [ - "pymatgen==2022.9.21", + "pymatgen==2024.2.20", ] [project.urls] @@ -54,60 +56,48 @@ homepage = "https://materialsproject.github.io/pyrho/" repository = "https://materialsproject.github.io/pyrho" [tool.ruff] -# Enable pycodestyle (`E`) and Pyflakes (`F`) codes by default. -select = ["E", "F"] -ignore = [] - -# Allow autofix for all enabled rules (when `--fix`) is provided. -fixable = ["E", "F"] -unfixable = [] - -# Exclude a variety of commonly ignored directories. -exclude = [ - ".bzr", - ".direnv", - ".eggs", - ".git", - ".hg", - ".mypy_cache", - ".nox", - ".pants.d", - ".pytype", - ".ruff_cache", - ".svn", - ".tox", - ".venv", - "__pypackages__", - "_build", - "buck-out", - "build", - "dist", - "node_modules", - "venv", +src = ["src",] +extend-exclude = ["docs", "tests"] +line-length = 88 +indent-width = 4 + +[tool.ruff.lint] +ignore = ["E203", "E501", "F401", "SLF001"] +select = [ + "A001", + "A002", + "B018", + "D", + "E", + "F", + "I001", + "PLE", + "PLW", + # "ERA001", + "RUF", + "D409", + "TCH", + "TID251", + "T20", + "UP032", + "W605", ] -# Same as Black. -line-length = 120 - -[tool.versioningit.vcs] -default-tag = "0.0.1" -method = "git" +[tool.ruff.format] +# Like Black, use double quotes for strings. +quote-style = "double" -[tool.isort] -profile = "black" +# Like Black, indent with spaces, rather than tabs. +indent-style = "space" -[tool.flake8] -max-line-length = 120 -max-doc-length = 120 -select = "C, E, F, W, B, B950" -extend-ignore = "E203, W503, E501, F401, RST21" -min-python-version = "3.8.0" -docstring-convention = "numpy" -rst-roles = "class, func, ref, obj" +# Like Black, respect magic trailing commas. +skip-magic-trailing-comma = false -[tool.pydocstyle] -ignore = "D100,D203,D213,D405,D415" +# Like Black, automatically detect the appropriate line ending. +line-ending = "auto" +[tool.ruff.lint.pydocstyle] +convention = "numpy" [tool.mypy] ignore_missing_imports = true diff --git a/src/pyrho/charge_density.py b/src/pyrho/charge_density.py index 354754f..06a7a83 100644 --- a/src/pyrho/charge_density.py +++ b/src/pyrho/charge_density.py @@ -1,23 +1,26 @@ """Chang Density Objects: Periodic Grid + Lattice / Atoms.""" - from __future__ import annotations import math import warnings from dataclasses import dataclass -from typing import Dict, List, Tuple, Union +from typing import TYPE_CHECKING import numpy as np import numpy.typing as npt from monty.dev import deprecated from monty.json import MSONable +from pymatgen.analysis.structure_matcher import ElementComparator, StructureMatcher from pymatgen.core.lattice import Lattice -from pymatgen.core.structure import Structure from pymatgen.io.vasp import Chgcar, Poscar, VolumetricData from pyrho.pgrid import PGrid from pyrho.utils import get_sc_interp +if TYPE_CHECKING: + from pymatgen.core.structure import Structure + + __all__ = ["ChargeDensity"] @@ -40,7 +43,7 @@ class ChargeDensity(MSONable): """ - pgrids: Dict[str, PGrid] + pgrids: dict[str, PGrid] structure: Structure normalization: str | None = "vasp" @@ -53,7 +56,7 @@ def __post_init__(self): """ lattices = [self.pgrids[key].lattice for key in self.pgrids.keys()] if not all( - np.allclose(self.structure.lattice._matrix, lattice) for lattice in lattices + np.allclose(self.structure.lattice.matrix, lattice) for lattice in lattices ): raise ValueError("Lattices are not identical") @@ -80,7 +83,7 @@ def normalized_data(self) -> dict[str, npt.NDArray]: } @property - def grid_shape(self) -> Tuple[int, int, int]: + def grid_shape(self) -> tuple[int, int, int]: """Return the shape of the charge density.""" return self.pgrids["total"].grid_shape @@ -138,7 +141,7 @@ def from_pmg( """ pgrids = { - k: PGrid(v, vdata.structure.lattice._matrix) for k, v in vdata.data.items() + k: PGrid(v, vdata.structure.lattice.matrix) for k, v in vdata.data.items() } return cls( pgrids=pgrids, structure=vdata.structure, normalization=normalization @@ -151,7 +154,7 @@ def reorient_axis(self) -> None: ``c`` is in the positive-z halve of space """ - args: Tuple[float, float, float, float, float, float] = ( + args: tuple[float, float, float, float, float, float] = ( self.structure.lattice.abc + self.structure.lattice.angles ) self.structure.lattice = Lattice.from_parameters(*args, vesta=True) @@ -185,7 +188,7 @@ def get_data_in_cube(self, s: float, ngrid: int, key: str = "total") -> npt.NDAr def get_transformed( self, sc_mat: npt.NDArray, - grid_out: Union[List[int], int], + grid_out: list[int] | int, origin: npt.ArrayLike = (0, 0, 0), up_sample: int = 1, ) -> "ChargeDensity": @@ -230,8 +233,6 @@ def get_transformed( ngrid = grid_out / new_structure.volume mult = (np.prod(lengths) / ngrid) ** (1 / 3) grid_out = [int(math.floor(max(l_ / mult, 1))) for l_ in lengths] - else: - grid_out = grid_out pgrids = {} for k, pgrid in self.normalized_pgrids.items(): @@ -260,14 +261,29 @@ def to_Chgcar(self) -> Chgcar: Chgcar: The charge density object + """ + return self.to_VolumetricData(cls=Chgcar, normalization="vasp") + + def to_VolumetricData( + self, cls=VolumetricData, normalization: str = "vasp" + ) -> VolumetricData: + """Convert the charge density to a ``pymatgen.io.vasp.outputs.VolumetricData`` object. + + Scale and convert each key in the pgrids dictionary and create a ``VolumetricData`` object + + Returns + ------- + VolumetricData: + The charge density object + """ struct = self.structure.copy() data_dict = {} for k, v in self.normalized_data.items(): data_dict[k] = _scaled_data( - v, lattice=self.structure.lattice, normalization="vasp" + v, lattice=self.structure.lattice, normalization=normalization ) - return Chgcar(Poscar(struct), data=data_dict) + return cls(Poscar(structure=struct), data_dict) @classmethod def from_file( @@ -311,85 +327,83 @@ def from_hdf5( """ return cls.from_pmg(pmg_obj.from_hdf5(filename)) - # - # _, new_rho = get_sc_interp(self.rho, sc_mat, grid_sizes=grid_out) - # new_rho = new_rho.reshape(grid_out) - # - # grid_shifts = [ - # int(t * g) for t, g in zip(translation - np.round(translation), grid_out) - # ] - # - # new_rho = roll_array(new_rho, grid_shifts) - # return self.__class__.from_rho(new_rho, new_structure) - - -# class SpinChargeDensity(MSONable, ChargeABC): -# def __init__(self, chargeden_dict: Dict, aug_charge: Dict = None): -# """ -# Wrapper class that parses multiple sets of grid data on the same lattice - -# Args: -# chargeden_dict: A dictionary containing multiple charge density objects -# typically in the format {'total' : ChargeDen1, 'diff' : ChargeDen2} -# """ -# self.chargeden_dict = chargeden_dict -# self.aug_charge = aug_charge -# self._tmp_key = next( -# iter(self.chargeden_dict) -# ) # get one key in the dictionary to make writing the subsequent code easier - -# @classmethod -# def from_pmg_volumetric_data( -# cls, vdata: VolumetricData, data_keys=("total", "diff") -# ): -# chargeden_dict = {} -# data_aug = getattr(vdata, "data_aug", None) -# for k in data_keys: -# chargeden_dict[k] = ChargeDensity.from_pmg(vdata, data_key=k) -# return cls(chargeden_dict, aug_charge=data_aug) - -# @property -# def lattice(self) -> Lattice: -# return self.chargeden_dict[self._tmp_key].lattice - -# def to_Chgcar(self) -> Chgcar: -# struct = self.chargeden_dict[self._tmp_key].structure -# data_ = {k: v.renormalized_data for k, v in self.chargeden_dict.items()} -# return Chgcar(Poscar(struct), data_, data_aug=self.aug_charge) - -# def to_VolumetricData(self) -> VolumetricData: -# key_ = next(iter(self.chargeden_dict)) -# struct = self.chargeden_dict[key_].structure -# data_ = {k: v.renormalized_data for k, v in self.chargeden_dict.items()} -# return VolumetricData(struct, data_) - -# def get_reshaped( -# self, -# sc_mat: npt.ArrayLike, -# grid_out: Union[List, int], -# origin: npt.ArrayLike = (0.0, 0.0, 0.0), -# up_sample: int = 1, -# ) -> "SpinChargeDensity": -# new_spin_charge = {} -# for k, v in self.chargeden_dict.items(): -# new_spin_charge[k] = v.get_reshaped_cell(sc_mat, frac_shift, grid_out) -# factor = int( -# new_spin_charge[self._tmp_key].structure.num_sites -# / self.chargeden_dict[self._tmp_key].structure.num_sites -# ) -# new_aug = {} -# if self.aug_charge is not None: -# for k, v in self.aug_charge.items(): -# new_aug[k] = multiply_aug(v, factor) -# return self.__class__(new_spin_charge, new_aug) - -# def reorient_axis(self) -> None: -# for k, v in self.chargeden_dict: -# v.reorient_axis() + +def get_matched_structure_mapping( + uc_struct: Structure, sc_struct: Structure, sm: StructureMatcher | None = None +) -> tuple[npt.NDArray, npt.ArrayLike] | None: + """Get the mapping of the supercell to the unit cell. + + Get the mapping from the supercell structure onto the base structure, + Note: this only works for structures that are exactly matched. + + Parameters + ---------- + uc_struct: host structure, smaller cell + sc_struct: bigger cell + sm: StructureMatcher instance + + Returns + ------- + sc_m : supercell matrix to apply to s1 to get s2 + total_t : translation to apply on s1 * sc_m to get s2 + """ + if sm is None: + sm = StructureMatcher( + primitive_cell=False, comparator=ElementComparator(), attempt_supercell=True + ) + s1, s2 = sm._process_species([sc_struct.copy(), uc_struct.copy()]) + trans = sm.get_transformation(s1, s2) + if trans is None: + return None + sc, t, mapping = trans + temp = s2.copy().make_supercell(sc) + ii, jj = 0, mapping[0] + vec = np.round(sc_struct[ii].frac_coords - temp[jj].frac_coords) + return sc, t + vec + + +def get_volumetric_like_sc( + vd: VolumetricData, + sc_struct: Structure, + grid_out: npt.ArrayLike, + up_sample: int = 1, + sm: StructureMatcher | None = None, + normalization: str | None = "vasp", +): + """Get the volumetric data in the supercell. + + Parameters + ---------- + vd: VolumeData instance + sc_struct: supercell structure. + grid_out: grid size to output the volumetric data. + up_sample: up sampling factor. + sm: StructureMatcher instance + normalization: normalization method for the volumetric data. + default is "vasp" which assumes the normalization is the + same as VASP's CHGCAR file. If None, no normalization is + done. + + Returns + ------- + VolumetricData: volumetric data in the supercell + """ + trans = get_matched_structure_mapping(vd.structure, sc_struct=sc_struct, sm=sm) + if trans is None: + raise ValueError("Could not find a supercell mapping") + sc_mat, total_t = trans + cden = ChargeDensity.from_pmg(vd, normalization=normalization) + orig = np.dot(total_t, sc_mat) + cden_transformed = cden.get_transformed( + sc_mat=sc_mat, origin=-orig, grid_out=grid_out, up_sample=up_sample + ) + return cden_transformed.to_VolumetricData( + cls=vd.__class__, normalization=normalization + ) @deprecated -def multiply_aug(data_aug: List[str], factor: int) -> List[str]: +def multiply_aug(data_aug: list[str], factor: int) -> list[str]: """Update the data in the augmentation charge. The original idea here was to use to to speed up some vasp calculations for @@ -414,8 +428,8 @@ def multiply_aug(data_aug: List[str], factor: int) -> List[str]: Each line of the augmentation data. """ - res: List[str] = [] - cur_block: List[str] = [] + res: list[str] = [] + cur_block: list[str] = [] cnt = 0 for ll in data_aug: if "augmentation" in ll: @@ -429,13 +443,12 @@ def multiply_aug(data_aug: List[str], factor: int) -> List[str]: cur_block = [ll] else: cur_block.append(ll) - else: - for _ in range(factor): - cnt += 1 - cur_block[ - 0 - ] = f"augmentation occupancies{cnt:>4}{cur_block[0].split()[-1]:>4}\n" - res.extend(cur_block) + for _ in range(factor): + cnt += 1 + cur_block[ + 0 + ] = f"augmentation occupancies{cnt:>4}{cur_block[0].split()[-1]:>4}\n" + res.extend(cur_block) return res diff --git a/src/pyrho/pgrid.py b/src/pyrho/pgrid.py index 154fd5e..a8f4882 100644 --- a/src/pyrho/pgrid.py +++ b/src/pyrho/pgrid.py @@ -1,14 +1,17 @@ """Python class for ND grid data volumetric data.""" + from __future__ import annotations -from typing import List, Union +from typing import TYPE_CHECKING import numpy as np -import numpy.typing as npt from monty.json import MSONable from pyrho.utils import gaussian_smear, get_sc_interp, interpolate_fourier +if TYPE_CHECKING: + import numpy.typing as npt + class PGrid(MSONable): """Class representing of _periodic_ grid data. @@ -38,7 +41,7 @@ def __init__(self, grid_data: npt.NDArray, lattice: npt.NDArray): def _transform_data( self, sc_mat: npt.ArrayLike, - grid_out: List[int], + grid_out: npt.ArrayLike, origin: npt.ArrayLike | None = None, up_sample: int = 1, ) -> npt.NDArray: @@ -72,7 +75,9 @@ def _transform_data( arr_in=self.grid_data, shape=[g_dim_ * up_sample for g_dim_ in self.grid_data.shape], ) - _, new_data = get_sc_interp(interp_grid_data, sc_mat, grid_sizes=grid_out, origin=origin) # type: ignore + _, new_data = get_sc_interp( + interp_grid_data, sc_mat, grid_sizes=grid_out, origin=origin + ) # type: ignore new_data = new_data.reshape(grid_out) return new_data @@ -110,8 +115,8 @@ def __truediv__(self, factor: float) -> PGrid: def get_transformed( self, - sc_mat: Union[List[List[int]], npt.NDArray], - grid_out: List[int], + sc_mat: list[list[int]] | npt.NDArray, + grid_out: list[int], origin: npt.NDArray | None = None, up_sample: int = 1, ) -> PGrid: @@ -142,7 +147,7 @@ def get_transformed( return PGrid(grid_data=new_data, lattice=new_lattice) def lossy_smooth_compression( - self, grid_out: List, smear_std: float = 0.2 + self, grid_out: list, smear_std: float = 0.2 ) -> npt.NDArray: """Perform Fourier interpolation then Gaussian smoothing. diff --git a/src/pyrho/utils.py b/src/pyrho/utils.py index 032241b..f8f4d3e 100644 --- a/src/pyrho/utils.py +++ b/src/pyrho/utils.py @@ -2,13 +2,15 @@ from __future__ import annotations from itertools import combinations -from typing import Iterable, List, Tuple, Union +from typing import TYPE_CHECKING, Iterable, List, Tuple, Union import numpy as np -from numpy.typing import ArrayLike, NDArray from scipy.interpolate import RegularGridInterpolator from scipy.ndimage import convolve +if TYPE_CHECKING: + from numpy.typing import ArrayLike, NDArray + __all__ = [ "pad_arr", "interpolate_fourier", @@ -55,11 +57,7 @@ def pad_arr(arr_in: NDArray, shape: List[int]) -> NDArray: padded data """ - # for _, isize in enumerate(shape): - # if isize < arr_in.shape[_]: - # raise Warning( - # "Some dimension of output array is smaller than the same dimension of input array." - # ) + def get_slice(idig, idim, bound_pairs): if idig == "0": return slice(0, bound_pairs[idim][0]) diff --git a/src/pyrho/vis/scatter.py b/src/pyrho/vis/scatter.py index fd1d07c..08b5266 100644 --- a/src/pyrho/vis/scatter.py +++ b/src/pyrho/vis/scatter.py @@ -1,10 +1,15 @@ """Helper functions to visualize the data in plotly.""" from __future__ import annotations +from typing import TYPE_CHECKING + import numpy as np import plotly.graph_objs as go from matplotlib import pyplot as plt -from matplotlib.axes import Axes + +if TYPE_CHECKING: + from matplotlib.axes import Axes + """Visualization functions do the scatter plots in plotly since it seems to be more efficient.""" diff --git a/tests/conftest.py b/tests/conftest.py index e5101e5..3e3bc45 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -31,11 +31,13 @@ def get_xy(self, lat_mat, grids, origin=(0, 0)): """Get the x and y coordinates for a given pair of lattice vectors and grid size. Args: + ---- lat_mat: lattice vectors grids: grid size origin_cart: origin of the lattice vectors in cartesian coordinates Returns: + ------- XX: x coordinates for the grid in the shape of the grid YY: y coordinates for the grid in the shape of the grid diff --git a/tests/test_charge_density.py b/tests/test_charge_density.py index 2ef1b29..402f81b 100644 --- a/tests/test_charge_density.py +++ b/tests/test_charge_density.py @@ -2,7 +2,7 @@ import pytest from pymatgen.io.vasp import Chgcar -from pyrho.charge_density import ChargeDensity +from pyrho.charge_density import ChargeDensity, get_volumetric_like_sc def test_charge_density(test_dir): @@ -33,3 +33,10 @@ def test_charge_density(test_dir): np.testing.assert_allclose( chgcar_sc.data["total"], chgcar_transformed.data["total"], atol=0.1, rtol=0.01 ) # since the chgcar is scaled, the absolute tolerance is 0.1 + +def test_mapping(test_dir): + chgcar_uc = Chgcar.from_file(test_dir / "CHGCAR.uc.vasp") + chgcar_sc = Chgcar.from_file(test_dir / "CHGCAR.sc1.vasp") + grid_out = chgcar_sc.dim + chgcar_out = get_volumetric_like_sc(chgcar_uc, chgcar_sc.structure, grid_out=grid_out, up_sample=2) + np.testing.assert_allclose(chgcar_out.data["total"], chgcar_sc.data["total"], atol=0.1, rtol=0.01) \ No newline at end of file diff --git a/tests/vis/test_scatter.py b/tests/vis/test_scatter.py index 273a9cf..9c9a87d 100644 --- a/tests/vis/test_scatter.py +++ b/tests/vis/test_scatter.py @@ -4,7 +4,6 @@ def test_get_scatter_plot(): - data = [ [ [1, 1, 1, 1, 1, 1, 1],