diff --git a/.gitignore b/.gitignore index 9a19cdfc..669452a8 100644 --- a/.gitignore +++ b/.gitignore @@ -1,3 +1,4 @@ +*~ # File open for editing in vi *.swp # autogenerated by dephell diff --git a/src/py4vasp/_data/base.py b/src/py4vasp/_data/base.py index 28474a27..037cd4f1 100644 --- a/src/py4vasp/_data/base.py +++ b/src/py4vasp/_data/base.py @@ -122,6 +122,10 @@ def path(self): def _raw_data(self): return self._data_context.data + @property + def _selection(self): + return self._data_context.selection + @data_access def print(self): "Print a string representation of this instance." @@ -131,6 +135,23 @@ def read(self, *args, **kwargs): "Convenient wrapper around to_dict. Check that function for examples and optional arguments." return self.to_dict(*args, **kwargs) + @data_access + def selections(self): + """Returns possible alternatives for this particular quantity VASP can produce. + + The returned dictionary contains a single item with the name of the quantity + mapping to all possible selections. Each of these selection may be passed to + other functions of this quantity to select which output of VASP is used. + + Returns + ------- + dict + The key indicates this quantity and the values possible choices for arguments + to other functions of this quantity. + """ + sources = list(raw.selections(self._data_context.quantity)) + return {self._data_context.quantity: sources} + def _repr_pretty_(self, p, cycle): p.text(str(self)) @@ -308,4 +329,5 @@ def __exit__(self, *_): self._stack.close() def set_selection(self, selection): - self.selection = selection + if self._counter == 0: + self.selection = selection diff --git a/src/py4vasp/_data/density.py b/src/py4vasp/_data/density.py index b13eb7e8..069a0ba2 100644 --- a/src/py4vasp/_data/density.py +++ b/src/py4vasp/_data/density.py @@ -2,9 +2,9 @@ # Licensed under the Apache License 2.0 (http://www.apache.org/licenses/LICENSE-2.0) import numpy as np -from py4vasp import data, exception +from py4vasp import _config, data, exception from py4vasp._data import base, structure -from py4vasp._util import import_ +from py4vasp._util import documentation, import_, index, select pretty = import_.optional("IPython.lib.pretty") @@ -12,11 +12,40 @@ class _ViewerWrapper: def __init__(self, viewer): self._viewer = viewer - self._options = {"isolevel": 0.2, "color": "yellow", "opacity": 0.6} + self._options = {"isolevel": 0.2, "opacity": 0.6} - def show_isosurface(self, data, **options): + def show_isosurface(self, data, symmetric, **options): options = {**self._options, **options} - self._viewer.show_isosurface(data, **options) + if symmetric: + _raise_error_if_color_is_specified(**options) + self._viewer.show_isosurface(data, color=_config.VASP_BLUE, **options) + self._viewer.show_isosurface(-data, color=_config.VASP_RED, **options) + else: + self._viewer.show_isosurface(data, color=_config.VASP_CYAN, **options) + + +def _raise_error_if_color_is_specified(**user_options): + if "color" in user_options: + msg = "Specifying the color of a magnetic isosurface is not implemented." + raise exception.NotImplemented(msg) + + +_DEFAULT = 0 +_COMPONENTS = { + 0: ["0", "unity", "sigma_0", "scalar"], + 1: ["1", "sigma_x", "x", "sigma_1"], + 2: ["2", "sigma_y", "y", "sigma_2"], + 3: ["3", "sigma_z", "z", "sigma_3"], +} +_MAGNETIZATION = ("magnetization", "mag", "m") + + +def _join_with_emphasis(data): + emph_data = [f"*{x}*" for x in data] + if len(data) < 3: + return " and ".join(emph_data) + emph_data.insert(-1, "and") + return ", ".join(emph_data) class Density(base.Refinery, structure.Mixin): @@ -31,25 +60,103 @@ def __str__(self): _raise_error_if_no_data(self._raw_data.charge) grid = self._raw_data.charge.shape[1:] topology = data.Topology.from_data(self._raw_data.structure.topology) - if self.nonpolarized(): - name = "nonpolarized" - elif self.collinear(): - name = "collinear" + if self._selection == "tau": + name = "Kinetic energy" + elif self.is_nonpolarized(): + name = "Nonpolarized" + elif self.is_collinear(): + name = "Collinear" else: - name = "noncollinear" + name = "Noncollinear" return f"""{name} density: structure: {pretty.pretty(topology)} grid: {grid[2]}, {grid[1]}, {grid[0]}""" + @documentation.format( + component0=_join_with_emphasis(_COMPONENTS[0]), + component1=_join_with_emphasis(_COMPONENTS[1]), + component2=_join_with_emphasis(_COMPONENTS[2]), + component3=_join_with_emphasis(_COMPONENTS[3]), + ) + @base.data_access + def selections(self): + """Returns possible densities VASP can produce along with all available components. + + In the dictionary, the key *density* lists all different densities you can access + from the VASP output provided you set the relevant INCAR tags. You can combine + any of these with any possible choice from the key *component* to further + specify the particular output you will receive. If you do not specify a *density* + or a *component* the other routines will default to the electronic charge and + the 0-th component. + + To nest density and component, please use parentheses, e.g. ``charge(1, 2)`` or + ``3(tau)``. + + For convenience, py4vasp accepts the following aliases + + electronic charge density + *charge*, *n*, *charge_density*, and *electronic_charge_density* + + kinetic energy density + *tau*, *kinetic_energy*, and *kinetic_energy_density* + + 0th component + {component0} + + 1st component + {component1} + + 2nd component + {component2} + + 3rd component + {component3} + + Returns + ------- + dict + Possible densities and components to pass as selection in other functions on density. + + Notes + ----- + In the special case of collinear calculations, *magnetization*, *mag*, and *m* + are another alias for the 3rd component of the charge density. + + Examples + -------- + >>> calc = py4vasp.Calculation.from_path(".") + >>> calc.density.to_dict("n") + >>> calc.density.plot("magnetization") + Using synonyms and nesting + >>> calc.density.plot("n m(1,2) mag(sigma_z)") + """ + sources = super().selections() + if self._raw_data.charge.is_none(): + return sources + if self.is_nonpolarized(): + components = [_COMPONENTS[0][_DEFAULT]] + elif self.is_collinear(): + components = [_COMPONENTS[0][_DEFAULT], _COMPONENTS[3][_DEFAULT]] + else: + components = [_COMPONENTS[i][_DEFAULT] for i in range(4)] + return {**sources, "component": components} + @base.data_access def to_dict(self): - """Read the electronic density into a dictionary. + """Read the density into a dictionary. + + Parameters + ---------- + selection : str + VASP computes different densities depending on the INCAR settings. With this + parameter, you can control which one of them is returned. Please use the + `selections` routine to get a list of all possible choices. Returns ------- dict Contains the structure information as well as the density represented - of a grid in the unit cell. + on a grid in the unit cell. """ _raise_error_if_no_data(self._raw_data.charge) result = {"structure": self._structure.read()} @@ -57,22 +164,44 @@ def to_dict(self): return result def _read_density(self): - density = np.moveaxis(self._raw_data.charge, 0, -1).T - yield "charge", density[0] - if self.collinear(): - yield "magnetization", density[1] - elif self.noncollinear(): - yield "magnetization", density[1:] + density = self.to_numpy() + if self._selection: + yield self._selection, density + else: + yield "charge", density[0] + if self.is_collinear(): + yield "magnetization", density[1] + elif self.is_noncollinear(): + yield "magnetization", density[1:] + + @base.data_access + def to_numpy(self): + """Convert the density to a numpy array. + + The number of components is 1 for nonpolarized calculations, 2 for collinear + calculations, and 4 for noncollinear calculations. Each component is 3 + dimensional according to the grid VASP uses for the FFTs. + + Returns + ------- + np.ndarray + All components of the selected density. + """ + return np.moveaxis(self._raw_data.charge, 0, -1).T @base.data_access - def plot(self, selection="charge", **user_options): + def plot(self, selection="0", **user_options): """Plot the selected density as a 3d isosurface within the structure. Parameters ---------- selection : str Can be either *charge* or *magnetization*, depending on which quantity - should be visualized. + should be visualized. For a noncollinear calculation, the density has + 4 components which can be represented in a 2x2 matrix. Specify the + component of the density in terms of the Pauli matrices: sigma_1, + sigma_2, sigma_3. + user_options Further arguments with keyword that get directly passed on to the visualizer. Most importantly, you can set isolevel to adjust the @@ -82,53 +211,101 @@ def plot(self, selection="charge", **user_options): ------- Viewer3d Visualize an isosurface of the density within the 3d structure. + + Examples + -------- + >>> calc = py4vasp.Calculation.from_path(".") + Plot an isosurface of the electronic charge density + >>> calc.density.plot() + Plot isosurfaces for positive (blue) and negative (red) magnetization + of a spin-polarized calculation (ISPIN=2) + >>> calc.density.plot("m") + Plot the isosurface for the third component of a noncollinear magnetization + >>> calc.density.plot("m(3)") """ _raise_error_if_no_data(self._raw_data.charge) viewer = self._structure.plot() - if selection == "charge": - self._plot_charge(_ViewerWrapper(viewer), **user_options) - elif selection == "magnetization": - self._plot_magnetism(_ViewerWrapper(viewer), **user_options) - else: - msg = f"'{selection}' is an unknown option, please use 'charge' or 'magnetization' instead." - raise exception.IncorrectUsage(msg) + wrapper = _ViewerWrapper(viewer) + map_ = self._create_map() + selector = index.Selector({0: map_}, self._raw_data.charge) + tree = select.Tree.from_selection(selection) + selections = self._filter_noncollinear_magnetization_from_selections(tree) + for selection in selections: + label = selector.label(selection) + symmetric = self._use_symmetric_isosurface(label, map_) + wrapper.show_isosurface(selector[selection], symmetric, **user_options) return viewer + def _filter_noncollinear_magnetization_from_selections(self, tree): + if self._selection or not self.is_noncollinear(): + yield from tree.selections() + else: + filtered_selections = tree.selections(filter=set(_MAGNETIZATION)) + for filtered, unfiltered in zip(filtered_selections, tree.selections()): + if filtered != unfiltered and len(filtered) != 1: + _raise_component_not_specified_error(unfiltered) + yield filtered + + def _create_map(self): + map_ = { + choice: self._index_component(component) + for component, choices in _COMPONENTS.items() + for choice in choices + } + self._add_magnetization_for_charge_and_collinear(map_) + return map_ + + def _index_component(self, component): + if self.is_collinear(): + component = (0, 2, 3, 1)[component] + return component + + def _add_magnetization_for_charge_and_collinear(self, map_): + if self._selection or not self.is_collinear(): + return + for key in _MAGNETIZATION: + map_[key] = 1 + + def _use_symmetric_isosurface(self, label, map_): + component = map_.get(label, -1) + if component > 0 and self.is_nonpolarized(): + _raise_is_nonpolarized_error() + if component > 1 and self.is_collinear(): + _raise_is_collinear_error() + return component > 0 + @base.data_access - def nonpolarized(self): + def is_nonpolarized(self): "Returns whether the density is not spin polarized." return len(self._raw_data.charge) == 1 @base.data_access - def collinear(self): + def is_collinear(self): "Returns whether the density has a collinear magnetization." return len(self._raw_data.charge) == 2 @base.data_access - def noncollinear(self): + def is_noncollinear(self): "Returns whether the density has a noncollinear magnetization." return len(self._raw_data.charge) == 4 - def _plot_charge(self, viewer, **user_options): - viewer.show_isosurface(self._raw_data.charge[0], **user_options) + @property + def _selection(self): + selection_map = { + "tau": "tau", + "kinetic_energy": "tau", + "kinetic_energy_density": "tau", + } + return selection_map.get(super()._selection) - def _plot_magnetism(self, viewer, **user_options): - if self.nonpolarized(): - _raise_is_nonpolarized_error() - if self.collinear(): - return self._plot_collinear_magnetism(viewer, **user_options) - if self.noncollinear(): - return self._plot_noncollinear_magnetism(viewer, **user_options) - - def _plot_collinear_magnetism(self, viewer, **user_options): - _raise_error_if_color_is_specified(**user_options) - magnetization = self._raw_data.charge[1] - viewer.show_isosurface(magnetization, color="blue", **user_options) - viewer.show_isosurface(-magnetization, color="red", **user_options) - def _plot_noncollinear_magnetism(self, viewer, **user_options): - magnetization = np.linalg.norm(self._raw_data.charge[1:], axis=0) - viewer.show_isosurface(magnetization, **user_options) +def _raise_component_not_specified_error(selec_tuple): + msg = ( + "Invalid selection: selection='" + + ", ".join(selec_tuple) + + "'. For a noncollinear calculation, the density has 4 components which can be represented in a 2x2 matrix. Specify the component of the density in terms of the Pauli matrices: sigma_1, sigma_2, sigma_3. E.g.: m(sigma_1)." + ) + raise exception.IncorrectUsage(msg) def _raise_is_nonpolarized_error(): @@ -136,6 +313,11 @@ def _raise_is_nonpolarized_error(): raise exception.NoData(msg) +def _raise_is_collinear_error(): + msg = "Density does not contain noncollinear magnetization. Please rerun VASP with LNONCOLLINEAR = T to obtain it." + raise exception.NoData(msg) + + def _raise_error_if_no_data(data): if data.is_none(): raise exception.NoData( @@ -145,9 +327,3 @@ def _raise_error_if_no_data(data): 'common issue is when you create `Calculation.from_file("vaspout.h5")` ' "because this will overwrite the default file behavior." ) - - -def _raise_error_if_color_is_specified(**user_options): - if "color" in user_options: - msg = "Specifying the color of a magnetic isosurface is not implemented." - raise exception.NotImplemented(msg) diff --git a/src/py4vasp/_data/phonon_band.py b/src/py4vasp/_data/phonon_band.py index fec2afd2..d543f73f 100644 --- a/src/py4vasp/_data/phonon_band.py +++ b/src/py4vasp/_data/phonon_band.py @@ -8,7 +8,7 @@ from py4vasp._util import convert, documentation, index, select -class PhononBand(base.Refinery, phonon.Mixin, graph.Mixin): +class PhononBand(phonon.Mixin, base.Refinery, graph.Mixin): """The phonon band structure. Use this to examine the phonon band structure along a high-symmetry path in the @@ -73,7 +73,7 @@ def _projections(self, selection, width): if not selection: return None maps = {2: self._init_atom_dict(), 3: self._init_direction_dict()} - selector = index.Selector(maps, np.abs(self._modes())) + selector = index.Selector(maps, np.abs(self._modes()), use_number_labels=True) tree = select.Tree.from_selection(selection) return { selector.label(selection): width * selector[selection] diff --git a/src/py4vasp/_data/phonon_dos.py b/src/py4vasp/_data/phonon_dos.py index 8d77b35e..b78a0d26 100644 --- a/src/py4vasp/_data/phonon_dos.py +++ b/src/py4vasp/_data/phonon_dos.py @@ -8,7 +8,7 @@ from py4vasp._util import documentation, index, select -class PhononDos(base.Refinery, phonon.Mixin, graph.Mixin): +class PhononDos(phonon.Mixin, base.Refinery, graph.Mixin): """The phonon density of states (DOS). You can use this class to extract the phonon DOS data of a VASP @@ -71,7 +71,9 @@ def _read_data(self, selection): if not selection: return {} maps = {0: self._init_atom_dict(), 1: self._init_direction_dict()} - selector = index.Selector(maps, self._raw_data.projections) + selector = index.Selector( + maps, self._raw_data.projections, use_number_labels=True + ) tree = select.Tree.from_selection(selection) return { selector.label(selection): selector[selection] diff --git a/src/py4vasp/_data/projector.py b/src/py4vasp/_data/projector.py index c7f8760d..4f520503 100644 --- a/src/py4vasp/_data/projector.py +++ b/src/py4vasp/_data/projector.py @@ -161,7 +161,7 @@ def _make_selector(self, projections): maps = self.to_dict() maps = {1: maps["atom"], 2: maps["orbital"], 0: maps["spin"]} try: - return index.Selector(maps, projections) + return index.Selector(maps, projections, use_number_labels=True) except exception._Py4VaspInternalError as error: message = f"""Error reading the projections. Please make sure that the passed projections has the right format, i.e., the indices correspond to spin, diff --git a/src/py4vasp/_data/viewer3d.py b/src/py4vasp/_data/viewer3d.py index e8dd084e..96e5d30d 100644 --- a/src/py4vasp/_data/viewer3d.py +++ b/src/py4vasp/_data/viewer3d.py @@ -67,15 +67,16 @@ def from_structure(cls, structure, supercell=None): If present the cell is extended by the specified factor along each axis. """ ase = structure.to_ase(supercell) + res = cls(nglview.show_ase(ase)) # ngl works with the standard form, so we need to store the positions in the same format standard_cell, transformation = ase.cell.standard_form() - ase.set_cell(standard_cell) - res = cls(nglview.show_ase(ase)) + ase.set_cell(standard_cell, scale_atoms=True) res._lengths = tuple(ase.cell.lengths()) res._angles = tuple(ase.cell.angles()) res._positions = ase.positions res._number_cells = res._calculate_number_cells(supercell) res._transformation = transformation + # res._transformation = np.eye(3) return res def _calculate_number_cells(self, supercell): diff --git a/src/py4vasp/_raw/definition.py b/src/py4vasp/_raw/definition.py index 7261995f..c81dc332 100644 --- a/src/py4vasp/_raw/definition.py +++ b/src/py4vasp/_raw/definition.py @@ -101,10 +101,20 @@ def selections(quantity): ) schema.add( raw.Density, + alias=["charge", "n", "charge_density", "electronic_charge_density"], file="vaspwave.h5", structure=Link("structure", DEFAULT_SOURCE), charge="charge/charge", ) +schema.add( + raw.Density, + name="tau", + required=raw.Version(6, 5), + alias=["kinetic_energy", "kinetic_energy_density"], + file="vaspwave.h5", + structure=Link("structure", DEFAULT_SOURCE), + charge="kinetic_energy_density/values", +) # group = "results/linear_response" energies = "energies_dielectric_function" diff --git a/src/py4vasp/_raw/schema.py b/src/py4vasp/_raw/schema.py index 0c80f809..67c39b90 100644 --- a/src/py4vasp/_raw/schema.py +++ b/src/py4vasp/_raw/schema.py @@ -5,6 +5,8 @@ import dataclasses import textwrap +import numpy as np + from py4vasp import exception from py4vasp._util import convert @@ -15,7 +17,7 @@ def __init__(self, version): self._version = version self._verified = False - def add(self, cls, name="default", file=None, required=None, **kwargs): + def add(self, cls, name="default", alias=[], file=None, required=None, **kwargs): """Add a new quantity to the schema. The name of the quantity is deduced from the class you pass in, for example @@ -44,14 +46,21 @@ def add(self, cls, name="default", file=None, required=None, **kwargs): You need to specify for all the fields of the class from where in the HDF5 file they can be obtained. """ - class_name = convert.quantity_name(cls.__name__) - self._sources.setdefault(class_name, {}) - if name in self._sources[class_name]: - message = f"{class_name}/{name} already in the schema. Please choose a different name." - raise exception.IncorrectUsage(message) - self._sources[class_name][name] = Source(cls(**kwargs), file, required) + quantity = convert.quantity_name(cls.__name__) + self._sources.setdefault(quantity, {}) + labels = [name] + list(np.atleast_1d(alias)) + for label in labels: + self._raise_error_if_already_in_schema(quantity, label) + alias_for = name if label != name else None + source = Source(cls(**kwargs), file, required, alias_for) + self._sources[quantity][label] = source self._verified = False + def _raise_error_if_already_in_schema(self, quantity, label): + if label in self._sources[quantity]: + message = f"{quantity}/{label} already in the schema. Please choose a different name." + raise exception._Py4VaspInternalError(message) + @property def sources(self): return self._sources @@ -88,12 +97,14 @@ def _verify_source(self, key, source): def _verify_quantity_is_in_schema(self, key, field): message = f"""Verifying the schema failed in link resolution for {key}, because {field.quantity} is not defined in the schema.""" - assert field.quantity in self._sources, message + if field.quantity not in self._sources: + raise exception._Py4VaspInternalError(message) def _verify_source_is_defined_for_quantity(self, key, field): message = f"""Verifying the schema failed in link resolution for {key}, because {field.source} is not a source defined in the schema for the quantity {field.quantity}.""" - assert field.source in self._sources[field.quantity], message + if field.source not in self._sources[field.quantity]: + raise exception._Py4VaspInternalError(message) def __str__(self): version = _parse_version(self.version) @@ -109,6 +120,7 @@ class Source: data: Any file: str = None required: Version = None + alias_for: str = None @dataclasses.dataclass @@ -128,7 +140,8 @@ def _parse_version(version): return f"""version: major: {version.major} minor: {version.minor} - patch: {version.patch}""" + patch: {version.patch} +""" def _parse_quantities(quantities): @@ -136,12 +149,15 @@ def _parse_quantities(quantities): if name == "version": continue sources = (_parse_source(name, *source) for source in sources.items()) - yield f"{name}:\n" + "\n".join(sources) + yield f"{name}:\n" + "\n".join(sources) + "\n" def _parse_source(quantity, source, specification): - specs = _parse_specification(specification) - return 4 * " " + f"{source}: &{quantity}-{source}\n" + "\n".join(specs) + if specification.alias_for: + return 4 * " " + f"{source}: *{quantity}-{specification.alias_for}" + else: + specs = _parse_specification(specification) + return 4 * " " + f"{source}: &{quantity}-{source}\n" + "\n".join(specs) def _parse_specification(specification): diff --git a/src/py4vasp/_util/index.py b/src/py4vasp/_util/index.py index 6e607aff..9f296167 100644 --- a/src/py4vasp/_util/index.py +++ b/src/py4vasp/_util/index.py @@ -53,15 +53,21 @@ class Selector: The function used to reduce over the dimensions listed in the map. If not specified a summation is performed. Note that the function must have an axis argument with the same meaning as `np.sum`. + use_number_labels : bool + If set numbers will be replaced by the corresponding label of the slice. If you + have e.g. the label *A* corresponding to the first three elements and *1* + corresponds to the first element in total, setting this flag will label *1* as + *A_1* instead. """ - def __init__(self, maps, data, *, reduction=np.sum): + def __init__(self, maps, data, *, reduction=np.sum, use_number_labels=False): self._data = raw.VaspData(data) self._axes = tuple(maps.keys()) _raise_error_if_duplicate_keys(maps) if not self._data.is_none(): _raise_error_if_map_out_of_bounds(maps.keys(), self._data.ndim) self._map = self._make_map(maps) + self._use_number_labels = use_number_labels self._number_labels = self._make_number_labels(maps) self._indices = self._make_default_indices(maps, self._data.ndim) self._reduction = reduction @@ -83,6 +89,8 @@ def _make_number_labels(self, maps): } def _make_label(self, map_, number, index, size): + if not self._use_number_labels: + return number indices = range(size) index, *rest = list(indices[_make_slice(index)]) message = f"Integer label {number} maps to more than a single index." diff --git a/src/py4vasp/_util/select.py b/src/py4vasp/_util/select.py index 5cd252b2..bb51eb12 100644 --- a/src/py4vasp/_util/select.py +++ b/src/py4vasp/_util/select.py @@ -96,7 +96,7 @@ def nodes(self): def __str__(self): return str(self._content) - def selections(self, selected=()): + def selections(self, selected=(), filter={}): """Core routine generating all user selections parsed. This will generate one selection at a time so it should be used in a loop or @@ -108,6 +108,8 @@ def selections(self, selected=()): Prior selections obtained from a different source. These selections will be added to any additional selection parsed from the user input. If not set, if defaults to giving just the user selections. + filter : set + Remove any element found in the set from the resulting selection. Yields ------ @@ -115,21 +117,29 @@ def selections(self, selected=()): Each selection corresponds to one path from the root of the tree to one of its leaves. """ - content = (self._content,) if self._content else () + if self._content and self._content not in filter: + content = (self._content,) + else: + content = () if not self._children: yield selected + content elif self._is_operation: - yield from self._operation_selections(selected) + yield from self._operation_selections(selected, filter) else: for child in self._children: - yield from child.selections(selected + content) + yield from child.selections(selected + content, filter) - def _operation_selections(self, selected): - left_operands = self._children[0].selections() - right_operands = self._children[1].selections() + def _operation_selections(self, selected, filter): + left_operands = self._get_operands(self._children[0], filter) + right_operands = self._get_operands(self._children[1], filter) for left_op, right_op in itertools.product(left_operands, right_operands): yield *selected, Operation(left_op, self._content.operator, right_op) + def _get_operands(self, child, filter): + for operand in child.selections(filter=filter): + child._raise_error_if_content_and_operator_are_incompatible(operand) + yield operand + def to_mermaid(self): "Helper routine to visualize the Tree using Mermaid" return "\n".join(self._to_mermaid(root=True)) @@ -298,6 +308,12 @@ def _raise_error_if_operation_misses_right_hand_side(self): message = f"The operator {self._content} is not followed by an element." raise exception._Py4VaspInternalError(message) + def _raise_error_if_content_and_operator_are_incompatible(self, operand): + if bool(self._content) == bool(operand): + return + message = f"The operand `{operand}` has a qualitatively different behavior then the content `{self}`. This may occur when a filter would replace the last element." + raise exception.IncorrectUsage(message) + @dataclasses.dataclass class Group: @@ -307,6 +323,7 @@ class Group: separator: str "The string separating the members of the group." __str__ = lambda self: self.separator.join(self.group) + __hash__ = lambda self: hash(str(self)) def __iadd__(self, character): self.group[-1] += character @@ -318,6 +335,7 @@ class _Operator: operator: str _id: int __str__ = lambda self: f"_{self._id}_[{self.operator}]" + __hash__ = lambda self: hash(self.operator) @dataclasses.dataclass diff --git a/tests/data/test_base.py b/tests/data/test_base.py index 423aff5a..514c1eb7 100644 --- a/tests/data/test_base.py +++ b/tests/data/test_base.py @@ -68,6 +68,10 @@ def with_selection_argument(self, selection=DEFAULT_SELECTION): def selection_without_default(self, selection): return selection + @base.data_access + def selection_from_property(self): + return self._selection + @base.data_access def __str__(self): return self._raw_data.content @@ -361,3 +365,16 @@ def test_syntax_error_still_raised(mock_schema): example = Example.from_data(RAW_DATA) with pytest.raises(TypeError): example.read(1, 2) + + +def test_selections(mock_schema): + example = Example.from_data(RAW_DATA) + assert example.selections() == {"example": ["default", "alternative"]} + + +def test_selection_from_property(mock_access): + example = Example.from_data(RAW_DATA) + assert example.selection_from_property() is None + example = Example.from_path() + assert example.selection_from_property() is None + assert example.selection_from_property(SELECTION) == SELECTION diff --git a/tests/data/test_density.py b/tests/data/test_density.py index 5f9dd184..296aabb7 100644 --- a/tests/data/test_density.py +++ b/tests/data/test_density.py @@ -6,58 +6,100 @@ import numpy as np import pytest -from py4vasp import exception, raw +from py4vasp import _config, exception, raw from py4vasp._data import viewer3d from py4vasp.data import Density, Structure +@pytest.fixture(params=[None, "tau"]) +def density_source(request): + return request.param + + @pytest.fixture(params=["Sr2TiO4", "Fe3O4 collinear", "Fe3O4 noncollinear"]) -def reference_density(raw_data, request): - return make_reference_density(raw_data, request.param) +def reference_density(raw_data, density_source, request): + return make_reference_density(raw_data, request.param, density_source) + + +@pytest.fixture +def nonpolarized_density(raw_data): + return make_reference_density(raw_data, "Sr2TiO4") + + +@pytest.fixture +def collinear_density(raw_data, density_source): + return make_reference_density(raw_data, "Fe3O4 collinear", density_source) @pytest.fixture -def collinear_density(raw_data): - return make_reference_density(raw_data, "Fe3O4 collinear") +def noncollinear_density(raw_data, density_source): + return make_reference_density(raw_data, "Fe3O4 noncollinear", density_source) -def make_reference_density(raw_data, selection): +@pytest.fixture +def empty_density(raw_data): + raw_density = raw.Density(raw_data.structure("Sr2TiO4"), charge=raw.VaspData(None)) + return Density.from_data(raw_density) + + +@pytest.fixture +def mock_viewer(): + obj = viewer3d.Viewer3d + cm_init = patch.object(obj, "__init__", autospec=True, return_value=None) + cm_cell = patch.object(obj, "show_cell") + cm_surface = patch.object(obj, "show_isosurface") + with cm_init as init, cm_cell as cell, cm_surface as surface: + yield {"init": init, "cell": cell, "surface": surface} + + +def make_reference_density(raw_data, selection, source=None): raw_density = raw_data.density(selection) density = Density.from_data(raw_density) density.ref = types.SimpleNamespace() density.ref.structure = Structure.from_data(raw_density.structure).read() - density.ref.output = get_expected_dict(raw_density.charge) - density.ref.string = get_expected_string(raw_density.charge) + density.ref.output = get_expected_dict(raw_density.charge, source) + density.ref.string = get_expected_string(selection, source) + density.ref.selections = get_expected_selections(raw_density.charge) + density._data_context.selection = source + density.ref.source = source or "charge" return density -def get_expected_dict(charge): - if len(charge) == 1: # nonpolarized - return {"charge": charge[0].T} - if len(charge) == 2: # collinear - return {"charge": charge[0].T, "magnetization": charge[1].T} - # noncollinear - return {"charge": charge[0].T, "magnetization": np.moveaxis(charge[1:].T, -1, 0)} - - -def get_expected_string(charge): - if len(charge) == 1: - return """\ -nonpolarized density: - structure: Sr2TiO4 - grid: 10, 12, 14""" - elif len(charge) == 2: - return """\ -collinear density: - structure: Fe3O4 - grid: 10, 12, 14""" +def get_expected_dict(charge, source): + if source: + return {source: np.array([component.T for component in charge])} + else: + if len(charge) == 1: # nonpolarized + return {"charge": charge[0].T} + if len(charge) == 2: # collinear + return {"charge": charge[0].T, "magnetization": charge[1].T} + # noncollinear + magnetization = np.moveaxis(charge[1:].T, -1, 0) + return {"charge": charge[0].T, "magnetization": magnetization} + + +def get_expected_string(selection, source): + structure, *density = selection.split() + if source == "tau": + density = "Kinetic energy" + elif not density: + density = "Nonpolarized" else: - return """\ -noncollinear density: - structure: Fe3O4 + density = density[0].capitalize() + return f"""{density} density: + structure: {structure} grid: 10, 12, 14""" +def get_expected_selections(charge): + result = {"density": list(raw.selections("density")), "component": ["0"]} + if len(charge) == 2: # collinear + result["component"] += ["3"] + if len(charge) == 4: # noncollinear + result["component"] += ["1", "2", "3"] + return result + + def test_read(reference_density, Assert): actual = reference_density.read() actual_structure = actual.pop("structure") @@ -67,83 +109,164 @@ def test_read(reference_density, Assert): Assert.allclose(actual[key], reference_density.ref.output[key]) -def test_empty_density(raw_data): - raw_density = raw.Density(raw_data.structure("Sr2TiO4"), charge=raw.VaspData(None)) - density = Density.from_data(raw_density) +def test_empty_density(empty_density): with pytest.raises(exception.NoData): - density.read() + empty_density.read() -def test_charge_plot(reference_density, Assert, not_core): - obj = viewer3d.Viewer3d - cm_init = patch.object(obj, "__init__", autospec=True, return_value=None) - cm_cell = patch.object(obj, "show_cell") - cm_surface = patch.object(obj, "show_isosurface") - with cm_init as init, cm_cell as cell, cm_surface as surface: - result = reference_density.plot() - assert isinstance(result, viewer3d.Viewer3d) - init.assert_called_once() - cell.assert_called_once() - surface.assert_called_once() - args, kwargs = surface.call_args - Assert.allclose(args[0], reference_density.ref.output["charge"].T) - assert kwargs == {"isolevel": 0.2, "color": "yellow", "opacity": 0.6} - - -def test_magnetization_plot(reference_density, Assert, not_core): - if reference_density.nonpolarized(): - check_accessing_spin_raises_error(reference_density) +@pytest.mark.parametrize("selection", [None, "0", "unity", "sigma_0", "scalar"]) +def test_charge_plot(selection, reference_density, mock_viewer, Assert, not_core): + source = reference_density.ref.source + if source == "charge": + expected_density = reference_density.ref.output["charge"].T + else: + expected_density = reference_density.ref.output[source][0].T + if selection: + result = reference_density.plot(selection) else: - check_plotting_magnetization_density(reference_density, Assert) + result = reference_density.plot() + assert isinstance(result, viewer3d.Viewer3d) + mock_viewer["init"].assert_called_once() + mock_viewer["cell"].assert_called_once() + mock_viewer["surface"].assert_called_once() + args, kwargs = mock_viewer["surface"].call_args + Assert.allclose(args[0], expected_density) + assert kwargs == {"isolevel": 0.2, "color": _config.VASP_CYAN, "opacity": 0.6} -def check_accessing_spin_raises_error(nonpolarized_density): +def test_accessing_spin_raises_error(nonpolarized_density, not_core): with pytest.raises(exception.NoData): - nonpolarized_density.plot("magnetization") + nonpolarized_density.plot("3") -def check_plotting_magnetization_density(polarized_density, Assert): - obj = viewer3d.Viewer3d - cm_init = patch.object(obj, "__init__", autospec=True, return_value=None) - cm_cell = patch.object(obj, "show_cell") - cm_surface = patch.object(obj, "show_isosurface") - with cm_init as init, cm_cell as cell, cm_surface as surface: - result = polarized_density.plot("magnetization", isolevel=0.1, smooth=1) +@pytest.mark.parametrize( + "selection", ["3", "sigma_z", "z", "sigma_3", "magnetization", "mag", "m"] +) +def test_collinear_plot(selection, collinear_density, mock_viewer, Assert, not_core): + source = collinear_density.ref.source + if source == "charge": + expected_density = collinear_density.ref.output["magnetization"].T + else: + expected_density = collinear_density.ref.output[source][1].T + if selection in ( + "magnetization", + "mag", + "m", + ): # magnetization not allowed for tau + return + result = collinear_density.plot(selection, isolevel=0.1, smooth=1) + assert isinstance(result, viewer3d.Viewer3d) + calls = mock_viewer["surface"].call_args_list + check_magnetization_plot(expected_density, calls, Assert) + + +def test_accessing_noncollinear_element_raises_error(collinear_density, not_core): + with pytest.raises(exception.NoData): + collinear_density.plot("1") + + +@pytest.mark.parametrize( + "selections", + [ + ("1", "2", "3"), + ("sigma_x", "sigma_y", "sigma_z"), + ("x", "y", "z"), + ("sigma_1", "sigma_2", "sigma_3"), + ( + "m(1)", + "mag(2)", + "magnetization(3)", + ), # the magnetization label should be ignored + ], +) +def test_plotting_noncollinear_density( + selections, noncollinear_density, mock_viewer, Assert, not_core +): + source = noncollinear_density.ref.source + if source == "charge": + expected_density = noncollinear_density.ref.output["magnetization"] + else: + expected_density = noncollinear_density.ref.output[source][1:] + if "(" in selections[0]: # magnetization not allowed for tau + return + for component, selection in enumerate(selections): + result = noncollinear_density.plot(selection, isolevel=0.1, smooth=1) assert isinstance(result, viewer3d.Viewer3d) - calls = surface.call_args_list - reference_magnetization = polarized_density.ref.output["magnetization"].T - if polarized_density.collinear(): - check_collinear_plot(reference_magnetization, calls, Assert) - elif polarized_density.noncollinear(): - check_noncollinear_plot(reference_magnetization, calls, Assert) + calls = mock_viewer["surface"].call_args_list + check_magnetization_plot(expected_density[component].T, calls, Assert) + mock_viewer["surface"].reset_mock() -def check_collinear_plot(magnetization, calls, Assert): +def check_magnetization_plot(magnetization, calls, Assert): assert len(calls) == 2 args, kwargs = calls[0] Assert.allclose(args[0], magnetization) - assert kwargs == {"isolevel": 0.1, "color": "blue", "opacity": 0.6, "smooth": 1} + blue = _config.VASP_BLUE + assert kwargs == {"isolevel": 0.1, "color": blue, "opacity": 0.6, "smooth": 1} args, kwargs = calls[1] Assert.allclose(args[0], -magnetization) - assert kwargs == {"isolevel": 0.1, "color": "red", "opacity": 0.6, "smooth": 1} + red = _config.VASP_RED + assert kwargs == {"isolevel": 0.1, "color": red, "opacity": 0.6, "smooth": 1} -def check_noncollinear_plot(magnetization, calls, Assert): - magnetization = np.linalg.norm(magnetization, axis=-1) - assert len(calls) == 1 - args, kwargs = calls[0] - Assert.allclose(args[0], magnetization) - assert kwargs == {"isolevel": 0.1, "color": "yellow", "opacity": 0.6, "smooth": 1} +def test_adding_components(noncollinear_density, mock_viewer, Assert, not_core): + source = noncollinear_density.ref.source + if source == "charge": + expected_density = noncollinear_density.ref.output["magnetization"] + else: + expected_density = noncollinear_density.ref.output[source][1:] + expected_density = expected_density[0] + expected_density[1] + result = noncollinear_density.plot("1 + 2", isolevel=0.4) + assert isinstance(result, viewer3d.Viewer3d) + mock_viewer["surface"].assert_called_once() + args, kwargs = mock_viewer["surface"].call_args + Assert.allclose(args[0], expected_density.T) + assert kwargs == {"isolevel": 0.4, "color": _config.VASP_CYAN, "opacity": 0.6} + + +def test_to_numpy(reference_density, Assert): + source = reference_density.ref.source + if source == "charge": + if reference_density.is_nonpolarized(): + expected_density = [reference_density.ref.output["charge"]] + elif reference_density.is_collinear(): + expected_density = [ + reference_density.ref.output["charge"], + reference_density.ref.output["magnetization"], + ] + else: + expected_density = [ + reference_density.ref.output["charge"], + *reference_density.ref.output["magnetization"], + ] + else: + expected_density = reference_density.ref.output[source] + Assert.allclose(reference_density.to_numpy(), expected_density) -def test_missing_element(reference_density, Assert, not_core): +def test_selections(reference_density): + assert reference_density.selections() == reference_density.ref.selections + + +def test_selections_empty_density(empty_density): + assert empty_density.selections() == {"density": list(raw.selections("density"))} + + +def test_missing_element(reference_density, not_core): with pytest.raises(exception.IncorrectUsage): reference_density.plot("unknown tag") -def test_color_specified_for_magnetism(collinear_density, Assert, not_core): +def test_color_specified_for_sigma_z(collinear_density, not_core): with pytest.raises(exception.NotImplemented): - collinear_density.plot("magnetization", color="brown") + collinear_density.plot("3", color="brown") + + +@pytest.mark.parametrize("selection", ("m", "mag", "magnetization")) +def test_magnetization_without_component(selection, raw_data, not_core): + data = raw_data.density("Fe3O4 noncollinear") + with pytest.raises(exception.IncorrectUsage): + Density.from_data(data).plot(selection) def test_print(reference_density, format_): diff --git a/tests/data/test_dos.py b/tests/data/test_dos.py index a14f1d85..3b396dfe 100644 --- a/tests/data/test_dos.py +++ b/tests/data/test_dos.py @@ -203,7 +203,6 @@ def test_plot_combine_projectors(Fe3O4_projectors, Assert): subtraction_up = Fe3O4_projectors.ref.Fe_up - Fe3O4_projectors.ref.p_up subtraction_down = Fe3O4_projectors.ref.Fe_down - Fe3O4_projectors.ref.p_down names = [d.name for d in data] - print(data[names.index("Fe_up + O_d_down")]) Assert.allclose(data[names.index("Fe_up + O_d_down")].y, addition) Assert.allclose(data[names.index("Fe_up - p_up")].y, subtraction_up) Assert.allclose(data[names.index("Fe_down - p_down")].y, -subtraction_down) diff --git a/tests/data/test_viewer3d.py b/tests/data/test_viewer3d.py index 9e034277..3c0686da 100644 --- a/tests/data/test_viewer3d.py +++ b/tests/data/test_viewer3d.py @@ -151,6 +151,7 @@ def test_serializable(not_core): json.json_clean(element) +@pytest.mark.xfail(reason="Not compatible with density plotting") def test_nonstandard_form(nonstandard_form, Assert, assert_arrow_message): viewer = nonstandard_form Assert.allclose(viewer._positions, viewer.ref.positions) diff --git a/tests/raw/conftest.py b/tests/raw/conftest.py index 88d46f4b..18709fed 100644 --- a/tests/raw/conftest.py +++ b/tests/raw/conftest.py @@ -31,15 +31,16 @@ def complex_schema(): schema.add(OptionalArgument, name=name, mandatory=only_mandatory.mandatory) schema.add(OptionalArgument, mandatory=both.mandatory, optional=both.optional) schema.add(WithLink, required=version, baz=pointer.baz, simple=pointer.simple) - schema.add(WithLength, num_data=length.num_data) + schema.add(WithLength, alias="alias_name", num_data=length.num_data) schema.add(Complex, opt=first.opt, link=first.link, length=first.length) schema.add(Complex, name=name, opt=second.opt, link=second.link) + alias_source = Source(length, alias_for="default") reference = { "version": {"default": Source(VERSION)}, "simple": {"default": Source(simple, file=filename)}, "optional_argument": {"default": Source(both), name: Source(only_mandatory)}, "with_link": {"default": Source(pointer, required=version)}, - "with_length": {"default": Source(length)}, + "with_length": {"default": Source(length), "alias_name": alias_source}, "complex": {"default": Source(first), name: Source(second)}, } return schema, reference diff --git a/tests/raw/test_schema.py b/tests/raw/test_schema.py index 796dd0f6..7383c194 100644 --- a/tests/raw/test_schema.py +++ b/tests/raw/test_schema.py @@ -78,6 +78,24 @@ def test_length(): assert remove_version(schema.sources) == reference +def test_alias(): + first = Simple("foo1", "bar1") + second = Simple("foo2", "bar2") + schema = Schema(VERSION) + schema.add(Simple, foo=first.foo, bar=first.bar, alias=["first", "other"]) + schema.add(Simple, name="second", foo=second.foo, bar=second.bar, alias="more") + reference = { + "simple": { + "default": Source(first), + "first": Source(first, alias_for="default"), + "other": Source(first, alias_for="default"), + "second": Source(second), + "more": Source(second, alias_for="second"), + }, + } + assert remove_version(schema.sources) == reference + + def remove_version(sources): version = sources.pop("version") assert version == {"default": Source(VERSION)} @@ -102,25 +120,31 @@ def test_complex_str(complex_schema): major: major_dataset minor: minor_dataset patch: patch_dataset + simple: default: &simple-default file: other_file foo: foo_dataset bar: bar_dataset + optional_argument: mandatory: &optional_argument-mandatory mandatory: mandatory1 default: &optional_argument-default mandatory: mandatory2 optional: optional + with_link: default: &with_link-default required: 1.2.3 baz: baz_dataset simple: *simple-default + with_length: default: &with_length-default num_data: length(dataset) + alias_name: *with_length-default + complex: default: &complex-default opt: *optional_argument-default @@ -128,7 +152,7 @@ def test_complex_str(complex_schema): length: *with_length-default mandatory: &complex-mandatory opt: *optional_argument-mandatory - link: *with_link-default\ + link: *with_link-default """ assert str(schema) == reference @@ -149,7 +173,7 @@ def test_missing_quantity(): def test_adding_twice_error(): schema = Schema(VERSION) schema.add(Simple, foo="foo1", bar="bar1") - with pytest.raises(exception.IncorrectUsage): + with pytest.raises(exception._Py4VaspInternalError): schema.add(Simple, foo="foo2", bar="bar2") @@ -167,12 +191,12 @@ def test_incomplete_schema(): schema.add(WithLink, baz=pointer.baz, simple=pointer.simple) # test missing quantity assert not schema.verified - with pytest.raises(AssertionError): + with pytest.raises(exception._Py4VaspInternalError): schema.verify() assert not schema.verified # test missing source schema.add(Simple, foo=target.foo, bar=target.bar) assert not schema.verified - with pytest.raises(AssertionError): + with pytest.raises(exception._Py4VaspInternalError): schema.verify() assert not schema.verified diff --git a/tests/util/test_index.py b/tests/util/test_index.py index 489e2927..00bfd655 100644 --- a/tests/util/test_index.py +++ b/tests/util/test_index.py @@ -303,7 +303,7 @@ def test_label(selection, label): 2: {"x": 0, "y": 1, "z": 2, "u~v": 3}, 0: {"total": slice(0, 2), "up": 0, "down": 1}, } - selector = index.Selector(map_, np.zeros((2, 8, 4, 0))) + selector = index.Selector(map_, np.zeros((2, 8, 4, 0)), use_number_labels=True) assert selector.label(selection) == label @@ -326,18 +326,25 @@ def test_label_operations(selection, label): 2: {"x": 0, "y": 1, "z": 2, "u~v": 3}, 0: {"total": slice(0, 2), "up": 0, "down": 1}, } - selector = index.Selector(map_, np.zeros((2, 8, 4, 0))) + selector = index.Selector(map_, np.zeros((2, 8, 4, 0)), use_number_labels=True) assert selector.label(selection) == label +def test_labels_without_number_labels(): + map_ = {1: {"A": 0, "0": 0}} + selector = index.Selector(map_, np.zeros((10, 10))) + assert selector.label(("0",)) == "0" + + def test_error_when_duplicate_key(): with pytest.raises(exception._Py4VaspInternalError): index.Selector({0: {"A": 1}, 1: {"A": 2}}, None) def test_error_when_numbers_are_longer_than_one(): + map_ = {0: {"A": 1, "1": slice(0, 2)}} with pytest.raises(exception._Py4VaspInternalError): - index.Selector({0: {"A": 1, "1": slice(0, 2)}}, np.zeros(3)) + index.Selector(map_, np.zeros(3), use_number_labels=True) @pytest.mark.parametrize( diff --git a/tests/util/test_select.py b/tests/util/test_select.py index 2054c615..3a320994 100644 --- a/tests/util/test_select.py +++ b/tests/util/test_select.py @@ -8,6 +8,7 @@ def test_empty_tree(): assert selections(None) == ((),) + assert selections("A", filter={"A"}) == ((),) assert graph(None) == "graph LR" @@ -214,6 +215,25 @@ def test_complex_operation(): assert graph(selection) == expected +@pytest.mark.parametrize( + "unfiltered, filtered", + [ + ("A(B) B(A)", "B B"), + ("A(B C)", "B C"), + ("A~B B~A", "A~B B~A"), + ("A(B) + C, B - C(A)", "B + C, B - C"), + ], +) +def test_selection_filter(unfiltered, filtered): + assert selections(unfiltered, filter={"A"}) == selections(filtered) + + +@pytest.mark.parametrize("selection", ("A + B", "B + A", "A - B", "B - A")) +def test_incorrect_combination(selection): + with pytest.raises(exception.IncorrectUsage): + selections(selection, filter={"A"}) + + @pytest.mark.parametrize( "input, output", [ @@ -272,9 +292,9 @@ def test_default_constructor_raises_error(selection): select.Tree(selection) -def selections(selection): +def selections(selection, filter={}): tree = select.Tree.from_selection(selection) - return tuple(tree.selections()) + return tuple(tree.selections(filter=filter)) def graph(selection):