diff --git a/src/py4vasp/_config.py b/src/py4vasp/_config.py index 3d55cae7..2c3404e2 100644 --- a/src/py4vasp/_config.py +++ b/src/py4vasp/_config.py @@ -1,4 +1,4 @@ # Copyright © VASP Software GmbH, # Licensed under the Apache License 2.0 (http://www.apache.org/licenses/LICENSE-2.0) -VASP_COLORS = ("#4C265F", "#2FB5AB", "#2C68FC", "#A82C35", "#808080") -VASP_PURPLE, VASP_CYAN, VASP_BLUE, VASP_RED, VASP_GRAY = VASP_COLORS +VASP_COLORS = ("#4C265F", "#2FB5AB", "#2C68FC", "#A82C35", "#808080", "#212529") +VASP_PURPLE, VASP_CYAN, VASP_BLUE, VASP_RED, VASP_GRAY, VASP_DARK = VASP_COLORS diff --git a/src/py4vasp/_third_party/graph/__init__.py b/src/py4vasp/_third_party/graph/__init__.py index f1972421..5db9f5ad 100644 --- a/src/py4vasp/_third_party/graph/__init__.py +++ b/src/py4vasp/_third_party/graph/__init__.py @@ -1,8 +1,11 @@ # Copyright © VASP Software GmbH, # Licensed under the Apache License 2.0 (http://www.apache.org/licenses/LICENSE-2.0) -from py4vasp._config import VASP_COLORS +import copy + +from py4vasp._config import VASP_BLUE, VASP_COLORS, VASP_GRAY, VASP_RED from py4vasp._util import import_ +from .contour import Contour from .graph import Graph from .mixin import Mixin from .plot import plot @@ -13,6 +16,10 @@ if import_.is_imported(go) and import_.is_imported(pio): axis_format = {"showexponent": "all", "exponentformat": "power"} + contour = copy.copy(pio.templates["ggplot2"].data.contour[0]) + contour.colorscale = [[0, VASP_BLUE], [0.5, VASP_GRAY], [1, VASP_RED]] + data = {"contour": (contour,)} layout = {"colorway": VASP_COLORS, "xaxis": axis_format, "yaxis": axis_format} - pio.templates["vasp"] = go.layout.Template(layout=layout) + pio.templates["vasp"] = go.layout.Template(data=data, layout=layout) + pio.templates["ggplot2"].layout.shapedefaults = {} pio.templates.default = "ggplot2+vasp" diff --git a/src/py4vasp/_third_party/graph/contour.py b/src/py4vasp/_third_party/graph/contour.py new file mode 100644 index 00000000..3f34a931 --- /dev/null +++ b/src/py4vasp/_third_party/graph/contour.py @@ -0,0 +1,178 @@ +# Copyright © VASP Software GmbH, +# Licensed under the Apache License 2.0 (http://www.apache.org/licenses/LICENSE-2.0) +from __future__ import annotations + +import dataclasses +import itertools + +import numpy as np + +from py4vasp import _config +from py4vasp._third_party.graph import trace +from py4vasp._util import import_ + +ff = import_.optional("plotly.figure_factory") +go = import_.optional("plotly.graph_objects") +interpolate = import_.optional("scipy.interpolate") + + +@dataclasses.dataclass +class Contour(trace.Trace): + """Represents data on a 2d slice through the unit cell. + + This class creates a visualization of the data within the unit cell based on its + configuration. Currently it supports the creation of heatmaps and quiver plots. + For heatmaps each data point corresponds to one point on the grid. For quiver plots + each data point should be a 2d vector within the plane. + """ + + _interpolation_factor = 2 + """If the lattice does not align with the cartesian axes, the data is interpolated + to by approximately this factor along each line.""" + _shift_label_pixels = 10 + "Shift the labels by this many pixels to avoid overlap." + + data: np.array + """2d or 3d grid data in the plane spanned by the lattice vectors. If the data is + the dimensions should be the ones of the grid, if the data is 3d the first dimension + should be a 2 for a vector in the plane of the grid and the other two dimensions + should be the grid.""" + lattice: Lattice + """2 vectors spanning the plane in which the data is represented. Each vector should + have two components, so remove any element normal to the plane.""" + label: str + "Assign a label to the visualization that may be used to identify one among multiple plots." + isolevels: bool = False + "Defines whether isolevels should be added or a heatmap is used." + supercell: np.array = (1, 1) + "Multiple of each lattice vector to be drawn." + show_cell: bool = True + "Show the unit cell in the resulting visualization." + + def to_plotly(self): + lattice_supercell = np.diag(self.supercell) @ self.lattice.vectors + # swap a and b axes because that is the way plotly expects the data + data = np.tile(self.data, self.supercell).T + if self._is_contour(): + yield self._make_contour(lattice_supercell, data), self._options() + elif self._is_heatmap(): + yield self._make_heatmap(lattice_supercell, data), self._options() + else: + yield self._make_quiver(lattice_supercell, data), self._options() + + def _is_contour(self): + return self.data.ndim == 2 and self.isolevels + + def _is_heatmap(self): + return self.data.ndim == 2 and not self.isolevels + + def _make_contour(self, lattice, data): + x, y, z = self._interpolate_data_if_necessary(lattice, data) + return go.Contour(x=x, y=y, z=z, name=self.label, autocontour=True) + + def _make_heatmap(self, lattice, data): + x, y, z = self._interpolate_data_if_necessary(lattice, data) + return go.Heatmap(x=x, y=y, z=z, name=self.label, colorscale="turbid_r") + + def _make_quiver(self, lattice, data): + u = data[:, :, 0].flatten() + v = data[:, :, 1].flatten() + meshes = [ + np.linspace(np.zeros(2), vector, num_points, endpoint=False) + for vector, num_points in zip(reversed(lattice), data.shape) + # remember that b and a axis are swapped + ] + x, y = np.array([sum(points) for points in itertools.product(*meshes)]).T + fig = ff.create_quiver(x, y, u, v, scale=1) + return fig.data[0] + + def _interpolate_data_if_necessary(self, lattice, data): + if self._interpolation_required(): + x, y, z = self._interpolate_data(lattice, data) + else: + x, y, z = self._use_data_without_interpolation(lattice, data) + return x, y, z + + def _interpolation_required(self): + y_position_first_vector = self.lattice.vectors[0, 1] + x_position_second_vector = self.lattice.vectors[1, 0] + return not np.allclose((y_position_first_vector, x_position_second_vector), 0) + + def _interpolate_data(self, lattice, data): + area_cell = abs(np.cross(lattice[0], lattice[1])) + points_per_area = data.size / area_cell + points_per_line = np.sqrt(points_per_area) * self._interpolation_factor + lengths = np.sum(np.abs(lattice), axis=0) + shape = np.ceil(points_per_line * lengths).astype(int) + line_mesh_a = self._make_mesh(lattice, data.shape[1], 0) + line_mesh_b = self._make_mesh(lattice, data.shape[0], 1) + x_in, y_in = (line_mesh_a[:, np.newaxis] + line_mesh_b[np.newaxis, :]).T + x_in = x_in.flatten() + y_in = y_in.flatten() + z_in = data.flatten() + x_out, y_out = np.meshgrid( + np.linspace(x_in.min(), x_in.max(), shape[0]), + np.linspace(y_in.min(), y_in.max(), shape[1]), + ) + z_out = interpolate.griddata((x_in, y_in), z_in, (x_out, y_out), method="cubic") + return x_out[0], y_out[:, 0], z_out + + def _use_data_without_interpolation(self, lattice, data): + x = self._make_mesh(lattice, data.shape[1], 0) + y = self._make_mesh(lattice, data.shape[0], 1) + return x, y, data + + def _make_mesh(self, lattice, num_point, index): + vector = index if self._interpolation_required() else (index, index) + return ( + np.linspace(0, lattice[vector], num_point, endpoint=False) + + 0.5 * lattice[vector] / num_point + ) + + def _options(self): + return { + "shapes": self._create_unit_cell(), + "annotations": self._label_unit_cell_vectors(), + } + + def _create_unit_cell(self): + if not self.show_cell: + return () + pos_to_str = lambda pos: f"{pos[0]} {pos[1]}" + vectors = self.lattice.vectors + corners = (vectors[0], vectors[0] + vectors[1], vectors[1]) + to_corners = (f"L {pos_to_str(corner)}" for corner in corners) + path = f"M 0 0 {' '.join(to_corners)} Z" + unit_cell = {"type": "path", "line": {"color": _config.VASP_DARK}, "path": path} + return (unit_cell,) + + def _label_unit_cell_vectors(self): + if self.lattice.cut is None: + return [] + vectors = self.lattice.vectors + labels = tuple("abc".replace(self.lattice.cut, "")) + return [ + { + "text": label, + "showarrow": False, + "x": 0.5 * vectors[i, 0], + "y": 0.5 * vectors[i, 1], + **self._shift_label(vectors[i], vectors[1 - i]), + } + for i, label in enumerate(labels) + ] + + def _shift_label(self, current_vector, other_vector): + invert = np.cross(current_vector, other_vector) < 0 + norm = np.linalg.norm(current_vector) + shifts = self._shift_label_pixels * current_vector[::-1] / norm + if invert: + return { + "xshift": -shifts[0], + "yshift": shifts[1], + } + else: + return { + "xshift": shifts[0], + "yshift": -shifts[1], + } diff --git a/src/py4vasp/_third_party/graph/graph.py b/src/py4vasp/_third_party/graph/graph.py index 6824658e..8d44152c 100644 --- a/src/py4vasp/_third_party/graph/graph.py +++ b/src/py4vasp/_third_party/graph/graph.py @@ -9,7 +9,9 @@ from py4vasp import exception from py4vasp._config import VASP_COLORS +from py4vasp._third_party.graph.contour import Contour from py4vasp._third_party.graph.series import Series +from py4vasp._third_party.graph.trace import Trace from py4vasp._util import import_ go = import_.optional("plotly.graph_objects") @@ -25,7 +27,7 @@ class Graph(Sequence): parameters set in this class. """ - series: Series or Sequence[Series] + series: Trace or Sequence[Trace] "One or more series shown in the graph." xlabel: str = None "Label for the x axis." @@ -71,11 +73,14 @@ def to_plotly(self): "Convert the graph to a plotly figure." figure = self._make_plotly_figure() for trace, options in self._generate_plotly_traces(): - if options["row"] is None: + if options.get("row") is None: figure.add_trace(trace) else: figure.add_trace(trace, row=options["row"], col=1) - + for shape in options.get("shapes", ()): + figure.add_shape(**shape) + for annotation in options.get("annotations", ()): + figure.add_annotation(**annotation) return figure def show(self): @@ -107,9 +112,8 @@ def _ipython_display_(self): def _generate_plotly_traces(self): colors = itertools.cycle(VASP_COLORS) for series in self: - if not series.color: - series = replace(series, color=next(colors)) - yield from series._generate_traces() + series = _set_color_if_not_present(series, colors) + yield from series.to_plotly() def _make_plotly_figure(self): figure = self._figure_with_one_or_two_y_axes() @@ -120,12 +124,13 @@ def _make_plotly_figure(self): return figure def _figure_with_one_or_two_y_axes(self): + has_secondary_y_axis = lambda series: isinstance(series, Series) and series.y2 if self._subplot_on: max_row = max(series.subplot for series in self) figure = subplots.make_subplots(rows=max_row, cols=1) figure.update_layout(showlegend=False) return figure - elif any(series.y2 for series in self): + elif any(has_secondary_y_axis(series) for series in self): return subplots.make_subplots(specs=[[{"secondary_y": True}]]) else: return go.Figure() @@ -142,6 +147,8 @@ def _set_xaxis_options(self, figure): figure.layout.xaxis.tickmode = "array" figure.layout.xaxis.tickvals = tuple(self.xticks.keys()) figure.layout.xaxis.ticktext = self._xtick_labels() + if self._all_are_contour(): + figure.layout.xaxis.visible = False def _xtick_labels(self): # empty labels will be overwritten by plotly so we put a single space in them @@ -156,6 +163,17 @@ def _set_yaxis_options(self, figure): figure.layout.yaxis.title.text = self.ylabel if self.y2label: figure.layout.yaxis2.title.text = self.y2label + if self._all_are_contour(): + figure.layout.yaxis.visible = False + if self._any_are_contour(): + figure.layout.yaxis.scaleanchor = "x" + figure.layout.height = 500 + + def _all_are_contour(self): + return all(isinstance(series, Contour) for series in self) + + def _any_are_contour(self): + return any(isinstance(series, Contour) for series in self) def to_frame(self): """Convert graph to a pandas dataframe. @@ -213,7 +231,16 @@ def _name_column(self, series, suffix, idx=None): @property def _subplot_on(self): - return any(series.subplot for series in self) + has_subplot = lambda series: isinstance(series, Series) and series.subplot + return any(has_subplot(series) for series in self) + + +def _set_color_if_not_present(series, color_iterator): + if isinstance(series, Contour): + return series + if not series.color: + series = replace(series, color=next(color_iterator)) + return series Graph._fields = tuple(field.name for field in fields(Graph)) diff --git a/src/py4vasp/_third_party/graph/series.py b/src/py4vasp/_third_party/graph/series.py index 802af191..16a9e890 100644 --- a/src/py4vasp/_third_party/graph/series.py +++ b/src/py4vasp/_third_party/graph/series.py @@ -5,13 +5,14 @@ import numpy as np from py4vasp import exception +from py4vasp._third_party.graph import trace from py4vasp._util import import_ go = import_.optional("plotly.graph_objects") @dataclass -class Series: +class Series(trace.Trace): """Represents a single series in a graph. Typically this corresponds to a single line of x-y data with an optional name used @@ -55,7 +56,7 @@ def __setattr__(self, key, value): assert not self._frozen or hasattr(self, key) super().__setattr__(key, value) - def _generate_traces(self): + def to_plotly(self): first_trace = True for item in enumerate(np.atleast_2d(np.array(self.y))): yield self._make_trace(*item, first_trace), {"row": self.subplot} @@ -123,5 +124,8 @@ def _common_options(self, first_trace): "yaxis": "y2" if self.y2 else "y", } + def _generate_shapes(self): + return () + Series._fields = tuple(field.name for field in fields(Series)) diff --git a/src/py4vasp/_third_party/graph/trace.py b/src/py4vasp/_third_party/graph/trace.py new file mode 100644 index 00000000..503d0cff --- /dev/null +++ b/src/py4vasp/_third_party/graph/trace.py @@ -0,0 +1,14 @@ +# Copyright © VASP Software GmbH, +# Licensed under the Apache License 2.0 (http://www.apache.org/licenses/LICENSE-2.0) +from abc import ABC, abstractmethod + + +class Trace(ABC): + """Defines a base class with all methods that need to be implemented for Graph to + work as intended""" + + @abstractmethod + def to_plotly(self): + """Use yield to generate one or more plotly traces. Each returned element should + be a tuple (trace, dict) where the trace can be used as data for plotly and the + options modify the generation of the figure.""" diff --git a/src/py4vasp/_util/slicing.py b/src/py4vasp/_util/slicing.py new file mode 100644 index 00000000..4b797a6d --- /dev/null +++ b/src/py4vasp/_util/slicing.py @@ -0,0 +1,238 @@ +# Copyright © VASP Software GmbH, +# Licensed under the Apache License 2.0 (http://www.apache.org/licenses/LICENSE-2.0) +import dataclasses + +import numpy as np + +from py4vasp import exception + +INDICES = {"a": 0, "b": 1, "c": 2} +AXIS = ("x", "y", "z") + + +@dataclasses.dataclass +class Plane: + """Defines a plane in 2d produced from cutting the 3d unit cell by removing one + lattice vector.""" + + vectors: np.array + "2d vectors spanning the plane." + cell: np.array = None + "Lattice vectors of the unit cell." + cut: str = None + "Lattice vector cut to get the plane, if not set, no labels will be added" + + +def grid_scalar(data, plane, fraction): + """Takes a 2d slice of a 3d grid data. + + Often, we want to generate a 2d slice of a 3d data on a grid. One example would be + to visualize a slice through a plane as a contour plot. This routine facilitates + this task implementing taking the cut of the 3d data. + + Parameters + ---------- + data : np.ndarray + 3d data on a grid from which a 2d plane is extracted. + plane : Plane + Defines the 2d plane to which the data is reduced. + fraction : float + Determines which plane of the grid data is used. Periodic boundaries are assumed. + + Returns + ------- + np.ndarray + A 2d array where the dimension selected by cut has been removed by selecting + the plane according to the specificied fraction. + """ + _raise_error_if_cut_unknown(plane.cut) + index = INDICES[plane.cut] + length = data.shape[index] + slice_ = [slice(None), slice(None), slice(None)] + slice_[index] = np.round(length * fraction).astype(np.int_) % length + return data[tuple(slice_)] + + +def grid_vector(data, plane, fraction): + """Takes a 2d slice of grid data where every datapoint is a 3d vector. + + Often, we want to generate a 2d slice of data on a grid. One example would be to + visualize a slice through a plane as a quiver plot. This routine facilitates this + task implementing taking the cut of the 3d data. + + Parameters + ---------- + data : np.ndarray + Data on a grid where every point is a vector. The dimensions should be + (3d vector, grid_a, grid_b, grid_c) and will be reduced to a 2d vector on a + 2d grid. + plane : Plane + Defines the 2d plane to which the data is reduced. + fraction : float + Determines which plane of the grid data is used. Periodic boundaries are assumed. + + Returns + ------- + np.ndarray + A 3d array where the dimension selected by cut has been removed by selecting + the plane according to the specificied fraction. The vector is projected onto + the plane for visualization. + """ + _raise_error_if_cut_unknown(plane.cut) + index = INDICES[plane.cut] + length = data.shape[index + 1] # add 1 to account for the vector dimension + slice_ = [slice(None), slice(None), slice(None), slice(None)] + slice_[index + 1] = np.round(length * fraction).astype(np.int_) % length + return _project_vectors_to_plane(plane, data[tuple(slice_)]) + + +def _project_vectors_to_plane(plane, data): + # We want to want to project the vector r onto the plane spanned by the vectors + # u and v. Let the result be s = a u + b v. We can obtain the projected vector by + # minimizing the length |r - s| which leads to two conditions + # (r - a u - b v).u = 0 + # (r - a u - b v).v = 0 + # solving for a and b yields + # a = (v^2 u.r - u.v v.r) / N + # b = (u^2 v.r - u.v u.r) / N + # N = u^2 v^2 - (u.v)^2 + u2 = np.dot(plane.vectors[0], plane.vectors[0]) + v2 = np.dot(plane.vectors[1], plane.vectors[1]) + uv = np.dot(plane.vectors[0], plane.vectors[1]) + vectors = np.delete(plane.cell, INDICES[plane.cut], axis=0) + ur = np.tensordot(vectors[0], data, axes=(0, 0)) + vr = np.tensordot(vectors[1], data, axes=(0, 0)) + N = u2 * v2 - uv**2 + a = (v2 * ur - uv * vr) / N + b = (u2 * vr - uv * ur) / N + au = np.multiply.outer(plane.vectors[0], a) + bv = np.multiply.outer(plane.vectors[1], b) + return au + bv + + +def plane(cell, cut, normal="auto"): + """Takes a 2d slice of a 3d cell and projects it onto 2d coordinates. + + For simplicity in the documentation, we will assume that the cut is in the plane of + the a and b lattice vector. We use Greek letter α and β to refer to the vectors in + the 2d coordinate system. This routine computes the transformation Q such that + + .. math:: + + \alpha &= Q a \\ + \beta &= Q b + + Here, the matrix Q has a 2 × 3 shape so that the vectors α and β have only 2 + dimensions. Furthermore, we want that the matrix Q fulfills certain properties, + namely that the lengths of the vectors and their angle is not changed. If the + vectors lie within e.g. the x-y plane, α and β should be the same as a and b just + with the zero for the z component removed. However, to avoid irregular behavior + where small changes to the unit cell cause large changes to the selected vectors + this mode only works if the normal vector is close (~15°) to a Cartesian axis. + Otherwise the user needs to provide the axis with the *normal* argument. + Alternatively, one can align the first remaining lattice vector with the x axis by + passing None for *normal*. + + Parameters + ---------- + cell : np.ndarray + A 3 × 3 array defining the three lattice vectors of the unit cell. + cut : str + Either "a", "b", or "c" to define which lattice vector is removed to get the slice. + normal : str + Set the Cartesian direction "x", "y", or "z" parallel to which the normal of + the plane is rotated. Alteratively, set it to "auto" to rotate to the closest + Cartesian axis. If you set it to None, the normal will not be considered and + the first remaining lattice vector will be aligned with the x axis instead. + + Returns + ------- + Plane + A 2d representation of the plane with some information to transform data to it. + """ + _raise_error_if_cut_unknown(cut) + vectors = np.delete(cell, INDICES[cut], axis=0) + if normal is not None: + return Plane(_rotate_normal_to_cartesian_axis(vectors, normal), cell, cut) + else: + return Plane(_rotate_first_vector_to_x_axis(vectors), cell, cut) + + +def _rotate_first_vector_to_x_axis(vectors): + u, v = np.linalg.norm(vectors, axis=1) + x = np.dot(*vectors) / (u * v) + return np.array([[u, 0], [v * x, v * np.sqrt(1 - x**2)]]) + + +def _rotate_normal_to_cartesian_axis(vectors, normal): + old_normal = _get_old_normal(vectors) + index_axis = _get_index_axis(old_normal, normal) + new_normal = _get_new_normal_from_cartesian_axis(old_normal, index_axis) + rotation_matrix = _get_rotation_matrix((old_normal, new_normal)) + new_vectors = vectors @ rotation_matrix.T + return np.delete(new_vectors, index_axis, axis=1) + + +def _get_old_normal(vectors): + old_normal = np.cross(*vectors).astype(np.float_) + return old_normal / np.linalg.norm(old_normal) + + +def _get_index_axis(old_normal, normal): + if normal in AXIS: + return AXIS.index(normal) + elif normal == "auto": + index = np.argmax(np.abs(old_normal)) + _raise_error_if_direction_is_not_obvious(np.abs(old_normal[index])) + return index + else: + _raise_unknown_normal_error(normal) + + +def _get_new_normal_from_cartesian_axis(old_normal, index): + new_normal = np.zeros(3) + # do not use sign function because it is 0 if old_normal[index] == 0 + new_normal[index] = 1 if old_normal[index] >= 0 else -1 + return new_normal + + +def _get_rotation_matrix(vectors): + cos_angle = np.dot(*vectors) + v = np.cross(*vectors) + if np.linalg.norm(v) < 1e-10: + return np.eye(3) + V = np.cross(np.eye(3), v) + return np.eye(3) + V + V @ V / (1 + cos_angle) + + +def _get_slice(shape, cut, fraction): + return slice_ + + +def _raise_error_if_cut_unknown(cut): + if cut in INDICES: + return + message = """\ +The selected choice {cut} is invalid. Please select one of the three lattice vectors +with passing *cut* ("a", "b", or "c").""" + raise exception.IncorrectUsage(message.format(cut=cut)) + + +def _raise_error_if_direction_is_not_obvious(largest_element): + if largest_element > 0.95: + return + message = """\ +You did not specify the Cartesian direction to which the normal of the cut plane will +be rotated to. py4vasp tries to determine the axis automatically but in this case no +axis is close to the normal of the plane. Please pass the additional argument *normal* +("x", "y", or "z") to specify to which axis py4vasp should use as normal vector for the +plane.""" + raise exception.IncorrectUsage(message) + + +def _raise_unknown_normal_error(normal): + message = """\ +The selected normal {normal} is invalid. Please select one of the Cartesian axis ("x", +"y", or "z") or "auto" to specify to which axis the normal is rotated. "auto" will use +a closeby Cartesian axis if possible.""" + raise exception.IncorrectUsage(message.format(normal=normal)) diff --git a/src/py4vasp/calculation/_density.py b/src/py4vasp/calculation/_density.py index 40377a93..19fb5c46 100644 --- a/src/py4vasp/calculation/_density.py +++ b/src/py4vasp/calculation/_density.py @@ -3,8 +3,8 @@ import numpy as np from py4vasp import _config, calculation, exception -from py4vasp._third_party import view -from py4vasp._util import documentation, import_, index, select +from py4vasp._third_party import graph, view +from py4vasp._util import documentation, import_, index, select, slicing from py4vasp.calculation import _base, _structure pretty = import_.optional("IPython.lib.pretty") @@ -20,6 +20,45 @@ } _MAGNETIZATION = ("magnetization", "mag", "m") +_PLANE = """\ +You need to specify a plane defined by two of the lattice vectors by selecting +a *cut* along the third one. You must only select a single cut and the value +should correspond to the fractional length along this third lattice vector. +py4vasp will then create a plane from the other two lattice vectors and +generate a contour plot within this plane. + +Usually, the first remaining lattice vector is aligned with the x-axis and the +second one such that the angle between the vectors is preserved. You can +overwrite this choice by defining a normal direction. Then py4vasp will rotate +the normal vector of the plane to align with the specified direction. This is +particularly useful if the lattice vector you cut is aligned with a Cartesian +direction. +""" + +_COMMON_PARAMETERS = """\ +a, b, c : float + You must select exactly one of these to specify which of the three lattice + vectors you want to remove to form a plane. The assigned value represents + the fractional length along this lattice vector, so `a = 0.3` will remove + the first lattice vector and then take the grid points at 30% of the length + of the first vector in the b-c plane. The fractional height uses periodic + boundary conditions. + +supercell : int or np.ndarray + Replicate the contour plot periodically a given number of times. If you + provide two different numbers, the resulting cell will be the two remaining + lattice vectors multiplied by the specific number. + +normal : str or None + If not set, py4vasp will align the first remaining lattice vector with the + x-axis and the second one such that the angle between the lattice vectors + is preserved. You can set it to "x", "y", or "z"; then py4vasp will rotate + the plane in such a way that the normal direction aligns with the specified + Cartesian axis. This may look better if the normal direction is close to a + Cartesian axis. You may also set it to "auto" so that py4vasp chooses a + close Cartesian axis if it can find any. +""" + def _join_with_emphasis(data): emph_data = [f"*{x}*" for x in filter(lambda key: key != _INTERNAL, data)] @@ -289,6 +328,128 @@ def _use_symmetric_isosurface(self, component): _raise_is_collinear_error() return component > 0 + @_base.data_access + @documentation.format(plane=_PLANE, common_parameters=_COMMON_PARAMETERS) + def to_contour( + self, selection=None, *, a=None, b=None, c=None, supercell=None, normal=None + ): + """Generate a contour plot of the selected component of the density. + + {plane} + + Parameters + ---------- + selection : str + Select which component of the density you want to visualize. Please use the + `selections` method to get all available choices. + + {common_parameters} + + Returns + ------- + graph + A contour plot in the plane spanned by the 2 remaining lattice vectors. + + + Examples + -------- + + Cut a plane through the magnetization density at the origin of the third lattice + vector. + + >>> calc.density.to_contour("3", c=0) + + Replicate a plane in the middle of the second lattice vector 2 times in each + direction. + + >>> calc.density.to_contour(b=0.5, supercell=2) + + Take a slice of the kinetic energy density along the first lattice vector and + rotate it such that the normal of the plane aligns with the x axis. + + >>> calc.density.to_contour("tau", a=0.3, normal="x") + """ + cut, fraction = self._get_cut(a, b, c) + plane = slicing.plane(self._structure.lattice_vectors(), cut, normal) + 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) + contours = [ + self._contour(selector, selection, plane, fraction, supercell) + for selection in selections + ] + return graph.Graph(contours) + + def _contour(self, selector, selection, plane, fraction, supercell): + density = selector[selection].T + data = slicing.grid_scalar(density, plane, fraction) + label = self._label(selector.label(selection)) or "charge" + contour = graph.Contour(data, plane, label, isolevels=True) + if supercell is not None: + contour.supercell = np.ones(2, dtype=np.int_) * supercell + return contour + + @_base.data_access + @documentation.format(plane=_PLANE, common_parameters=_COMMON_PARAMETERS) + def to_quiver(self, *, a=None, b=None, c=None, supercell=None, normal=None): + """Generate a quiver plot of magnetization density. + + {plane} + + For a collinear calculation, the magnetization density will be aligned with the + y axis of the plane. For noncollinear calculations, the magnetization density + is projected into the plane. + + Parameters + ---------- + {common_parameters} + + Returns + ------- + graph + A quiver plot in the plane spanned by the 2 remaining lattice vectors. + + + Examples + -------- + + Cut a plane at the origin of the third lattice vector. + + >>> calc.density.to_quiver(c=0) + + Replicate a plane in the middle of the second lattice vector 2 times in each + direction. + + >>> calc.density.to_quiver(b=0.5, supercell=2) + + Take a slice of the spin components of the kinetic energy density along the + first lattice vector and rotate it such that the normal of the plane aligns with + the x axis. + + >>> calc.density.to_quiver("tau", a=0.3, normal="x") + """ + cut, fraction = self._get_cut(a, b, c) + plane = slicing.plane(self._structure.lattice_vectors(), cut, normal) + if self.is_collinear(): + data = slicing.grid_scalar(self._raw_data.charge[1].T, plane, fraction) + data = np.array((np.zeros_like(data), data)) + else: + data = slicing.grid_vector(self.to_numpy()[1:], plane, fraction) + label = self._selection or "magnetization" + quiver_plot = graph.Contour(5.0 * data, plane, label) + if supercell is not None: + quiver_plot.supercell = np.ones(2, dtype=np.int_) * supercell + return graph.Graph([quiver_plot]) + + def _get_cut(self, a, b, c): + _raise_error_cut_selection_incorrect(a, b, c) + if a is not None: + return "a", a + if b is not None: + return "b", b + return "c", c + @_base.data_access def is_nonpolarized(self): "Returns whether the density is not spin polarized." @@ -348,3 +509,20 @@ 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_cut_selection_incorrect(*selections): + # only a single element may be selected + selected_elements = sum(selection is not None for selection in selections) + if selected_elements == 0: + raise exception.IncorrectUsage( + "You have not selected a lattice vector along which the slice should be " + "constructed. Please set exactly one of the keyword arguments (a, b, c) " + "to a real number that specifies at which fraction of the lattice vector " + "the plane is." + ) + if selected_elements > 1: + raise exception.IncorrectUsage( + "You have selected more than a single element. Please use only one of " + "(a, b, c) and not multiple choices." + ) diff --git a/src/py4vasp/calculation/_structure.py b/src/py4vasp/calculation/_structure.py index ecffaa28..81510c3d 100644 --- a/src/py4vasp/calculation/_structure.py +++ b/src/py4vasp/calculation/_structure.py @@ -144,8 +144,8 @@ def to_dict(self): {examples} """ return { - "lattice_vectors": self._lattice_vectors(), - "positions": self._positions(), + "lattice_vectors": self.lattice_vectors(), + "positions": self.positions(), "elements": self._topology().elements(), "names": self._topology().names(), } @@ -171,8 +171,8 @@ def to_view(self, supercell=None): make_3d = lambda array: array if array.ndim == 3 else array[np.newaxis] return view.View( elements=np.atleast_2d(self._topology().elements()), - lattice_vectors=make_3d(self._lattice_vectors()), - positions=make_3d(self._positions()), + lattice_vectors=make_3d(self.lattice_vectors()), + positions=make_3d(self.positions()), supercell=self._parse_supercell(supercell), ) @@ -261,6 +261,32 @@ def to_POSCAR(self): message = "Converting multiple structures to a POSCAR is currently not implemented." raise exception.NotImplemented(message) + @_base.data_access + def lattice_vectors(self): + """Return the lattice vectors spanning the unit cell + + Returns + ------- + np.ndarray + Lattice vectors of the unit cell in Å. + """ + lattice_vectors = _LatticeVectors(self._raw_data.cell.lattice_vectors) + return self._scale() * lattice_vectors[self._get_steps()] + + @_base.data_access + def positions(self): + """Return the direct coordinates of all ions in the unit cell. + + Direct or fractional coordinates measure the position of the ions in terms of + the lattice vectors. Hence they are dimensionless quantities. + + Returns + ------- + np.ndarray + Positions of all ions in terms of the lattice vectors. + """ + return self._raw_data.positions[self._get_steps()] + @_base.data_access @documentation.format(examples=_slice.examples("structure", "cartesian_positions")) def cartesian_positions(self): @@ -273,7 +299,7 @@ def cartesian_positions(self): {examples} """ - return self._positions() @ self._lattice_vectors() + return self.positions() @ self.lattice_vectors() @_base.data_access @documentation.format(examples=_slice.examples("structure", "volume")) @@ -287,7 +313,7 @@ def volume(self): {examples} """ - return np.abs(np.linalg.det(self._lattice_vectors())) + return np.abs(np.linalg.det(self.lattice_vectors())) @_base.data_access def number_atoms(self): @@ -331,10 +357,6 @@ def _parse_supercell(self, supercell): def _topology(self): return calculation.topology.from_data(self._raw_data.topology) - def _lattice_vectors(self): - lattice_vectors = _LatticeVectors(self._raw_data.cell.lattice_vectors) - return self._scale() * lattice_vectors[self._get_steps()] - def _scale(self): if isinstance(self._raw_data.cell.scale, np.float_): return self._raw_data.cell.scale @@ -343,9 +365,6 @@ def _scale(self): else: return 1.0 - def _positions(self): - return self._raw_data.positions[self._get_steps()] - def _get_steps(self): return self._steps if self._is_trajectory else () diff --git a/tests/calculation/test_density.py b/tests/calculation/test_density.py index fef5cd34..e2586f06 100644 --- a/tests/calculation/test_density.py +++ b/tests/calculation/test_density.py @@ -255,6 +255,198 @@ def test_plotting_supercell(supercell, reference_density, Assert): check_view(reference_density, expected, Assert, supercell=supercell) +@pytest.mark.parametrize( + "kwargs, index, position", + (({"a": 0.1}, 0, 1), ({"b": 0.7}, 1, 8), ({"c": 1.3}, 2, 4)), +) +def test_contour_of_charge(nonpolarized_density, kwargs, index, position, Assert): + graph = nonpolarized_density.to_contour(**kwargs) + slice_ = [slice(None), slice(None), slice(None)] + slice_[index] = position + data = nonpolarized_density.ref.output["charge"][tuple(slice_)] + assert len(graph) == 1 + series = graph.series[0] + Assert.allclose(series.data, data) + lattice_vectors = nonpolarized_density.ref.structure.lattice_vectors() + lattice_vectors = np.delete(lattice_vectors, index, axis=0) + expected_products = lattice_vectors @ lattice_vectors.T + actual_products = series.lattice.vectors @ series.lattice.vectors.T + Assert.allclose(actual_products, expected_products) + assert series.label == "charge" + + +def test_incorrect_slice_raises_error(nonpolarized_density): + with pytest.raises(exception.IncorrectUsage): + nonpolarized_density.to_contour() + with pytest.raises(exception.IncorrectUsage): + nonpolarized_density.to_contour(a=1, b=2) + with pytest.raises(exception.IncorrectUsage): + nonpolarized_density.to_contour(3) + + +@pytest.mark.parametrize( + "selection", ["3", "sigma_z", "z", "sigma_3", "magnetization", "mag", "m"] +) +def test_collinear_to_contour(selection, collinear_density, Assert): + source = collinear_density.ref.source + if source == "charge": + expected_label = selection + expected_data = collinear_density.ref.output["magnetization"][:, :, 7] + else: + expected_label = f"{source}({selection})" + expected_data = collinear_density.ref.output[source][1, :, :, 7] + if selection in ("magnetization", "mag", "m"): + # magnetization not allowed for tau + return + expected_lattice = collinear_density.ref.structure.lattice_vectors()[:2, :2] + graph = collinear_density.to_contour(selection, c=-0.5) + assert len(graph) == 1 + series = graph.series[0] + Assert.allclose(series.data, expected_data) + Assert.allclose(series.lattice.vectors, expected_lattice) + assert series.label == expected_label + assert series.isolevels + + +@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_noncollinear_to_contour(noncollinear_density, selections, Assert): + source = noncollinear_density.ref.source + if source == "charge": + if "(" in selections[0]: # magnetization filtered from selections + expected_labels = ("1", "2", "3") + else: + expected_labels = selections + expected_data = noncollinear_density.ref.output["magnetization"][:, :, 5, :] + else: + expected_labels = (f"{source}({selection})" for selection in selections) + expected_data = noncollinear_density.ref.output[source][1:, :, 5, :] + if "(" in selections[0]: # magnetization not allowed for tau + return + graph = noncollinear_density.to_contour(" ".join(selections), b=0.4) + expected_lattice = noncollinear_density.ref.structure.lattice_vectors()[::2, ::2] + assert len(graph) == len(expected_data) + for density, label, series in zip(expected_data, expected_labels, graph.series): + Assert.allclose(series.data, density) + Assert.allclose(series.lattice.vectors, expected_lattice) + assert series.label == label + + +def test_to_contour_supercell(nonpolarized_density, Assert): + graph = nonpolarized_density.to_contour(a=0, supercell=2) + Assert.allclose(graph.series[0].supercell, (2, 2)) + graph = nonpolarized_density.to_contour(a=0, supercell=(2, 1)) + Assert.allclose(graph.series[0].supercell, (2, 1)) + + +def test_raise_error_if_normal_not_reasonable(nonpolarized_density): + with pytest.raises(exception.IncorrectUsage): + # cell is not close to cartesian vectors + nonpolarized_density.to_contour(a=0, normal="auto") + with pytest.raises(exception.IncorrectUsage): + nonpolarized_density.to_contour(a=0, normal="unknown choice") + + +@pytest.mark.parametrize( + "normal, rotation", + [ + ("auto", np.eye(2)), + ("x", np.array([[0, -1], [1, 0]])), + ("y", np.diag((1, -1))), + ("z", np.eye(2)), + ], +) +def test_to_contour_normal_vector(collinear_density, normal, rotation, Assert): + graph = collinear_density.to_contour(c=0.5, normal=normal) + lattice_vectors = collinear_density.ref.structure.lattice_vectors() + expected_lattice = lattice_vectors[:2, :2] @ rotation + Assert.allclose(graph.series[0].lattice.vectors, expected_lattice) + + +def test_to_contour_operation(collinear_density, Assert): + graph = collinear_density.to_contour("0 - 3", c=0) + source = collinear_density.ref.source + output = collinear_density.ref.output + if source == "charge": + expected_data = output["charge"][:, :, 0] - output["magnetization"][:, :, 0] + else: + expected_data = output[source][0, :, :, 0] - output[source][1, :, :, 0] + Assert.allclose(graph.series[0].data, expected_data) + + +def test_collinear_to_quiver(collinear_density, Assert): + source = collinear_density.ref.source + if source == "charge": + expected_label = "magnetization" + expected_data = collinear_density.ref.output["magnetization"][8] + else: + expected_label = source + expected_data = collinear_density.ref.output[source][1, 8] + expected_data = np.array((np.zeros_like(expected_data), expected_data)) + lattice_vectors = collinear_density.ref.structure.lattice_vectors()[1:] + expected_lattice = np.diag(np.linalg.norm(lattice_vectors, axis=1)) + graph = collinear_density.to_quiver(a=-0.2) + assert len(graph) == 1 + series = graph.series[0] + Assert.allclose(series.data, 5.0 * expected_data) + Assert.allclose(series.lattice.vectors, expected_lattice) + assert series.label == expected_label + + +def test_noncollinear_to_quiver(noncollinear_density, Assert): + source = noncollinear_density.ref.source + if source == "charge": + expected_label = "magnetization" + expected_data = noncollinear_density.ref.output["magnetization"][:2, :, :, 6] + else: + expected_label = source + expected_data = noncollinear_density.ref.output[source][1:3, :, :, 6] + expected_lattice = noncollinear_density.ref.structure.lattice_vectors()[:2, :2] + graph = noncollinear_density.to_quiver(c=1.4) + assert len(graph) == 1 + series = graph.series[0] + Assert.allclose(series.data, 5.0 * expected_data) + Assert.allclose(series.lattice.vectors, expected_lattice) + assert series.label == expected_label + + +def test_to_quiver_supercell(collinear_density, Assert): + graph = collinear_density.to_quiver(a=0, supercell=2) + Assert.allclose(graph.series[0].supercell, (2, 2)) + graph = collinear_density.to_quiver(a=0, supercell=(2, 1)) + Assert.allclose(graph.series[0].supercell, (2, 1)) + + +@pytest.mark.parametrize( + "normal, rotation", + [ + ("auto", np.eye(2)), + ("x", np.array([[0, -1], [1, 0]])), + ("y", np.diag((1, -1))), + ("z", np.eye(2)), + ], +) +def test_to_quiver_normal_vector(noncollinear_density, normal, rotation, Assert): + unrotated_graph = noncollinear_density.to_quiver(c=0.5) + rotated_graph = noncollinear_density.to_quiver(c=0.5, normal=normal) + expected_lattice = unrotated_graph.series[0].lattice.vectors @ rotation + Assert.allclose(rotated_graph.series[0].lattice.vectors, expected_lattice) + expected_data = (unrotated_graph.series[0].data.T @ rotation).T + Assert.allclose(rotated_graph.series[0].data, expected_data) + + def test_to_numpy(reference_density, Assert): source = reference_density.ref.source if source == "charge": @@ -307,4 +499,5 @@ def test_print(reference_density, format_): def test_factory_methods(raw_data, check_factory_methods): data = raw_data.density("Fe3O4 collinear") - check_factory_methods(calculation.density, data) + parameters = {"to_contour": {"a": 0.3}} + check_factory_methods(calculation.density, data, parameters) diff --git a/tests/calculation/test_structure.py b/tests/calculation/test_structure.py index 08d99f1e..2ee5bbe7 100644 --- a/tests/calculation/test_structure.py +++ b/tests/calculation/test_structure.py @@ -264,6 +264,18 @@ def test_supercell_wrong_size(Sr2TiO4, not_a_supercell, not_core): Sr2TiO4.to_ase(not_a_supercell) +@pytest.mark.parametrize("steps", [-1, 0, slice(1, 3), slice(None)]) +def test_lattice_vectors(Sr2TiO4, steps, Assert): + structure = Sr2TiO4 if steps == -1 else Sr2TiO4[steps] + Assert.allclose(structure.lattice_vectors(), Sr2TiO4.ref.lattice_vectors[steps]) + + +@pytest.mark.parametrize("steps", [-1, 0, slice(1, 3), slice(None)]) +def test_positions(Sr2TiO4, steps, Assert): + structure = Sr2TiO4 if steps == -1 else Sr2TiO4[steps] + Assert.allclose(structure.positions(), Sr2TiO4.ref.positions[steps]) + + def test_cartesian_positions(Sr2TiO4, Fe3O4, Ca3AsBr3, Assert, not_core): check_cartesian_positions(Sr2TiO4, Assert) check_cartesian_positions(Fe3O4, Assert) diff --git a/tests/conftest.py b/tests/conftest.py index 67e10984..3f63cae2 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -47,7 +47,7 @@ def not_core(is_core): class _Assert: @staticmethod - def allclose(actual, desired): + def allclose(actual, desired, tolerance=1): if _is_none(actual): assert _is_none(desired) elif getattr(actual, "dtype", None) == np.bool_: @@ -58,7 +58,7 @@ def allclose(actual, desired): actual, mask_actual = _finite_subset(actual) desired, mask_desired = _finite_subset(desired) assert np.all(mask_actual == mask_desired) - assert_array_almost_equal_nulp(actual, desired, 30) + assert_array_almost_equal_nulp(actual, desired, tolerance * 30) @staticmethod def same_structure(actual, desired): diff --git a/tests/third_party/graph/test_graph.py b/tests/third_party/graph/test_graph.py index 26d05829..1fd8b09f 100644 --- a/tests/third_party/graph/test_graph.py +++ b/tests/third_party/graph/test_graph.py @@ -6,8 +6,11 @@ import numpy as np import pytest -from py4vasp import exception -from py4vasp._third_party.graph import Graph, Series +from py4vasp import _config, exception +from py4vasp._third_party.graph import Contour, Graph, Series +from py4vasp._util import import_, slicing + +px = import_.optional("plotly.express") @pytest.fixture @@ -59,6 +62,46 @@ def subplot(): return Series(x1, y1, subplot=1), Series(x2, y2, subplot=2) +@pytest.fixture +def rectangle_contour(): + return Contour( + data=np.linspace(0, 10, 20 * 18).reshape((20, 18)), + lattice=slicing.Plane(np.diag([4.0, 3.6]), cut="c"), + label="rectangle contour", + isolevels=True, + ) + + +@pytest.fixture +def tilted_contour(): + return Contour( + data=np.linspace(0, 5, 16 * 20).reshape((16, 20)), + lattice=slicing.Plane(np.array([[2, 3], [2, -3]]), cut="b"), + label="tilted contour", + supercell=(2, 1), + show_cell=False, + ) + + +@pytest.fixture +def simple_quiver(): + return Contour( + data=np.array([[(y, x) for x in range(3)] for y in range(5)]).T, + lattice=slicing.Plane(np.diag((3, 5)), cut="a"), + label="quiver plot", + ) + + +@pytest.fixture +def complex_quiver(): + return Contour( + data=np.linspace(-3, 3, 2 * 12 * 10).reshape((2, 12, 10)), + lattice=slicing.Plane([[3, 2], [-3, 2]]), # cut not set + label="quiver plot", + supercell=(3, 2), + ) + + def test_basic_graph(parabola, Assert, not_core): graph = Graph(parabola) fig = graph.to_plotly() @@ -373,3 +416,183 @@ def test_nonexisting_attribute_raises_error(parabola): graph = Graph(parabola) with pytest.raises(AssertionError): graph.nonexisting = "not possible" + + +def test_normal_series_does_not_set_contour_layout(parabola, not_core): + graph = Graph(parabola) + fig = graph.to_plotly() + assert fig.layout.xaxis.visible is None + assert fig.layout.yaxis.visible is None + assert fig.layout.yaxis.scaleanchor is None + + +def test_contour(rectangle_contour, Assert, not_core): + graph = Graph(rectangle_contour) + fig = graph.to_plotly() + assert len(fig.data) == 1 + # plotly expects y-x order + Assert.allclose(fig.data[0].z, rectangle_contour.data.T) + # shift because the points define the centers of the rectangles + Assert.allclose(fig.data[0].x, np.linspace(0, 4, 20, endpoint=False) + 0.1) + Assert.allclose(fig.data[0].y, np.linspace(0, 3.6, 18, endpoint=False) + 0.1) + assert fig.data[0].name == rectangle_contour.label + assert fig.data[0].autocontour + # text explicitly that it is False to prevent None passing the test + assert fig.layout.xaxis.visible == False + assert fig.layout.yaxis.visible == False + assert len(fig.layout.shapes) == 1 + check_unit_cell(fig.layout.shapes[0], x="4.0", y="3.6", zero="0.0") + check_annotations(rectangle_contour.lattice, fig.layout.annotations, Assert) + assert fig.layout.yaxis.scaleanchor == "x" + + +def test_contour_supercell(rectangle_contour, Assert, not_core): + supercell = np.asarray((3, 5)) + rectangle_contour.supercell = supercell + graph = Graph(rectangle_contour) + fig = graph.to_plotly() + assert len(fig.data) == 1 + # plotly expects y-x order + assert all(fig.data[0].z.T.shape == supercell * rectangle_contour.data.shape) + assert len(fig.data[0].x) == 60 + assert len(fig.data[0].y) == 90 + assert len(fig.layout.shapes) == 1 + check_unit_cell(fig.layout.shapes[0], x="4.0", y="3.6", zero="0.0") + check_annotations(rectangle_contour.lattice, fig.layout.annotations, Assert) + + +def check_unit_cell(unit_cell, x, y, zero): + assert unit_cell.type == "path" + assert unit_cell.path == f"M 0 0 L {x} {zero} L {x} {y} L {zero} {y} Z" + assert unit_cell.line.color == _config.VASP_DARK + + +def check_annotations(lattice, annotations, Assert): + assert len(lattice.vectors) == len(annotations) + sign = np.sign(np.cross(*lattice.vectors)) + labels = "abc".replace(lattice.cut, "") + for vector, label, annotation in zip(lattice.vectors, labels, annotations): + assert annotation.showarrow == False + assert annotation.text == label + Assert.allclose((annotation.x, annotation.y), 0.5 * vector) + expected_shift = sign * 10 * vector[::-1] / np.linalg.norm(vector) + Assert.allclose((annotation.xshift, -annotation.yshift), expected_shift) + sign *= -1 + + +def test_contour_interpolate(tilted_contour, Assert, not_core): + graph = Graph(tilted_contour) + fig = graph.to_plotly() + area_cell = 12.0 + points_per_area = tilted_contour.data.size / area_cell + points_per_line = np.sqrt(points_per_area) * tilted_contour._interpolation_factor + lengths = np.array([6, 9]) # this accounts for the 2 x 1 supercell + expected_shape = np.ceil(points_per_line * lengths).astype(int) + expected_average = np.average(tilted_contour.data) + assert len(fig.data) == 1 + # plotly expects y-x order + assert all(fig.data[0].z.T.shape == expected_shape) + assert fig.data[0].x.size == expected_shape[0] + assert fig.data[0].y.size == expected_shape[1] + finite = np.isfinite(fig.data[0].z) + assert np.isclose(np.average(fig.data[0].z[finite]), expected_average) + assert len(fig.layout.shapes) == 0 + expected_colorscale = px.colors.get_colorscale("turbid_r") + assert len(fig.data[0].colorscale) == len(expected_colorscale) + for actual, expected in zip(fig.data[0].colorscale, expected_colorscale): + Assert.allclose(actual[0], expected[0]) + assert actual[1] == expected[1] + check_annotations(tilted_contour.lattice, fig.layout.annotations, Assert) + + +def test_mix_contour_and_series(two_lines, rectangle_contour, not_core): + graph = Graph([rectangle_contour, two_lines]) + fig = graph.to_plotly() + assert len(fig.data) == 3 + assert fig.layout.xaxis.visible is None + assert fig.layout.yaxis.visible is None + assert fig.layout.yaxis.scaleanchor == "x" + + +def test_simple_quiver(simple_quiver, Assert, not_core): + graph = Graph(simple_quiver) + fig = graph.to_plotly() + data_size = simple_quiver.data.size // 2 + assert len(fig.data) == 1 + actual = split_data(fig.data[0], data_size, Assert) + arrows = actual.tips - actual.positions + for (x, y), (u, v) in zip(actual.positions, arrows): + Assert.allclose(x, v) + Assert.allclose(y, u) + assert len(fig.layout.shapes) == 1 + check_unit_cell(fig.layout.shapes[0], x="3", y="5", zero="0") + check_annotations(simple_quiver.lattice, fig.layout.annotations, Assert) + assert fig.layout.yaxis.scaleanchor == "x" + + +def test_complex_quiver(complex_quiver, Assert, not_core): + graph = Graph(complex_quiver) + fig = graph.to_plotly() + data_size = np.prod(complex_quiver.supercell) * complex_quiver.data.size // 2 + vectors = np.array(complex_quiver.lattice.vectors) + step_a = vectors[0] / complex_quiver.data.shape[1] + mesh_a = np.arange(complex_quiver.supercell[0] * complex_quiver.data.shape[1]) + step_b = vectors[1] / complex_quiver.data.shape[2] + mesh_b = np.arange(complex_quiver.supercell[1] * complex_quiver.data.shape[2]) + expected_positions = np.array( + [a * step_a + b * step_b for b in mesh_b for a in mesh_a] + ) + work = complex_quiver.data + work = np.block([[work, work], [work, work], [work, work]]).T + expected_tips = expected_positions + work.reshape(expected_positions.shape) + expected_barb_length = 0.3 * np.linalg.norm(work, axis=-1).flatten() + # + assert len(fig.data) == 1 + actual = split_data(fig.data[0], data_size, Assert) + Assert.allclose(actual.positions, expected_positions, tolerance=10) + Assert.allclose(actual.tips, expected_tips, tolerance=10) + Assert.allclose(actual.barb_length, expected_barb_length) + assert len(fig.layout.annotations) == 0 + + +@dataclasses.dataclass +class ContourData: + positions: np.ndarray = None + first_tips: np.ndarray = None + second_tips: np.ndarray = None + first_barb_length: np.ndarray = None + second_barb_length: np.ndarray = None + + +def split_data(data, data_size, Assert): + actual = ContourData() + assert data.mode == "lines" + # The data contains first a line between each grid point and the tip of the arrow; + # each of the lines are separated by a None. The second part is the tip of the arrow + # consisting of two lines; again tips are separated by None. + assert len(data.x) == len(data.y) == 7 * data_size + # first element contains positions + slice_ = slice(0, 3 * data_size, 3) + actual.positions = np.array((data.x[slice_], data.y[slice_])).T + # second element of both parts contain the tip of the arrows + slice_ = slice(1, 3 * data_size, 3) + actual.tips = np.array((data.x[slice_], data.y[slice_])).T + slice_ = slice(3 * data_size + 1, None, 4) + other_tips = np.array((data.x[slice_], data.y[slice_])).T + Assert.allclose(other_tips, actual.tips) + # third element of first part and fourth element of second part should be None (=separator) + slice_ = slice(2, 3 * data_size, 3) + assert all(element is None for element in data.x[slice_]) + assert all(element is None for element in data.y[slice_]) + slice_ = slice(3 * data_size + 3, None, 4) + assert all(element is None for element in data.x[slice_]) + assert all(element is None for element in data.y[slice_]) + # the first and third element of the second part contain the barb of the arrow + slice_ = slice(3 * data_size, None, 4) + barb = np.array((data.x[slice_], data.y[slice_])).T + actual.barb_length = np.linalg.norm(barb - actual.tips, axis=-1) + slice_ = slice(3 * data_size + 2, None, 4) + other_barb = np.array((data.x[slice_], data.y[slice_])).T + other_barb_length = np.linalg.norm(other_barb - actual.tips, axis=-1) + Assert.allclose(other_barb_length, actual.barb_length, tolerance=10) + return actual diff --git a/tests/util/test_convert.py b/tests/util/test_convert.py index eade333b..390acee5 100644 --- a/tests/util/test_convert.py +++ b/tests/util/test_convert.py @@ -42,6 +42,7 @@ def test_hex_to_rgb(Assert): [44, 104, 252], [168, 44, 53], [128, 128, 128], + [33, 37, 41], ] expected = np.array(colors) / 255 actual = np.array([to_rgb(color) for color in VASP_COLORS]) diff --git a/tests/util/test_slicing.py b/tests/util/test_slicing.py new file mode 100644 index 00000000..9eb1034d --- /dev/null +++ b/tests/util/test_slicing.py @@ -0,0 +1,182 @@ +# Copyright © VASP Software GmbH, +# Licensed under the Apache License 2.0 (http://www.apache.org/licenses/LICENSE-2.0) +import numpy as np +import pytest + +from py4vasp import exception +from py4vasp._util import slicing + + +@pytest.mark.parametrize("cut, index", [("a", 0), ("b", 1), ("c", 2)]) +def test_orthorhombic(cut, index, Assert): + cell = np.diag((3, 4, 5)) + expected_plane = np.delete(np.delete(cell, index, axis=0), index, axis=1) + actual_plane = slicing.plane(cell, cut) + Assert.allclose(actual_plane.vectors, expected_plane) + assert actual_plane.cut == cut + + +@pytest.mark.parametrize("cut, indices", [("a", (0, 2)), ("b", (1, 0)), ("c", (2, 1))]) +def test_unusual_orthorhombic(cut, indices, Assert): + cell = np.roll(np.diag((3, 4, 5)), shift=1, axis=0) + expected_plane = np.delete(np.delete(cell, indices[0], axis=0), indices[1], axis=1) + actual_plane = slicing.plane(cell, cut) + Assert.allclose(actual_plane.vectors, expected_plane) + assert actual_plane.cut == cut + + +@pytest.mark.parametrize("cut, index", [("a", 0), ("b", 1), ("c", 2)]) +def test_nearly_orthorhombic(cut, index, Assert): + cell = 0.01 * np.ones((3, 3)) + np.diag((3, 4, 5)) + expected_plane = np.delete(cell, index, axis=0) + # the expected plane is in 3d and the actual plane in 2d should have same angles + # and lengths + approximate_plane = np.delete(np.delete(cell, index, axis=0), index, axis=1) + actual_plane = slicing.plane(cell, cut).vectors + Assert.allclose(actual_plane @ actual_plane.T, expected_plane @ expected_plane.T) + assert np.max(np.abs(approximate_plane - actual_plane)) < 0.1 + + +XX = np.sqrt(1 - np.sqrt(0.75)) + + +@pytest.mark.parametrize( + "normal, expected_plane", + [ + ("x", [[XX, 1 + XX], [1 + XX, XX]]), + ("y", [[1, 1], [1 + XX, -XX]]), + ("z", [[XX + 1, -XX], [1, 1]]), + ], +) +def test_rotate_to_user_defined_axis(normal, expected_plane, Assert): + cell = np.array([[0, 1, 1], [1, 0, 1], [1, 1, 0]]) + actual_plane = slicing.plane(cell, "a", normal) + Assert.allclose(actual_plane.vectors, expected_plane) + Assert.allclose(actual_plane.cell, cell) + + +@pytest.mark.parametrize( + "normal, rotation", + (("x", [[1, 0], [0, 1]]), ("y", [[-1, 0], [0, 1]]), ("z", [[0, 1], [-1, 0]])), +) +def test_normal_with_orthonormal_cell(normal, rotation, Assert): + cell = np.diag((3, 4, 5)) + expected_plane = np.dot(np.delete(np.delete(cell, 0, axis=0), 0, axis=1), rotation) + actual_plane = slicing.plane(cell, "a", normal) + Assert.allclose(actual_plane.vectors, expected_plane) + Assert.allclose(actual_plane.cell, cell) + assert actual_plane.cut == "a" + + +@pytest.mark.parametrize("cut, diagonal", [("a", [4, 5]), ("b", [3, 5]), ("c", [3, 4])]) +def test_no_rotation_orthorhombic_cell(cut, diagonal, Assert): + cell = np.diag((3, 4, 5)) + expected_plane = np.diag(diagonal) + actual_plane = slicing.plane(cell, cut, normal=None) + Assert.allclose(actual_plane.vectors, expected_plane) + + +@pytest.mark.parametrize("cut, index", [("a", 0), ("b", 1), ("c", 2)]) +def test_no_rotation_nontrivial_cell(cut, index, Assert): + cell = np.array([[0, 1, 1.1], [0.9, 0, 1.1], [0.9, 1, 0]]) + vectors = np.delete(cell, index, axis=0) + expected_products = vectors @ vectors.T + actual_plane = slicing.plane(cell, cut, normal=None).vectors + Assert.allclose(actual_plane[0, 1], 0) + actual_products = actual_plane @ actual_plane.T + Assert.allclose(actual_products, expected_products) + + +def test_raise_error_if_normal_is_not_obvious(): + with pytest.raises(exception.IncorrectUsage): + slicing.plane(np.array([[0, 1, 1], [1, 0, 1], [1, 1, 0]]), "a") + + +def test_raise_error_for_unknown_choices(): + with pytest.raises(exception.IncorrectUsage): + slicing.plane(np.eye(3), "unknown") + with pytest.raises(exception.IncorrectUsage): + slicing.plane(np.eye(3), "a", normal="unknown") + + +@pytest.mark.parametrize("cut", ("a", "b", "c")) +@pytest.mark.parametrize("fraction", (-0.4, 0, 0.4, 0.8, 1.2)) +def test_slice_grid_scalar(cut, fraction, Assert): + grid_scalar = np.random.random((10, 12, 14)) + 0.1 + if cut == "a": + index = np.round(fraction * 10).astype(np.int_) % 10 + expected_data = grid_scalar[index, :, :] + elif cut == "b": + index = np.round(fraction * 12).astype(np.int_) % 12 + expected_data = grid_scalar[:, index, :] + else: + index = np.round(fraction * 14).astype(np.int_) % 14 + expected_data = grid_scalar[:, :, index] + plane = slicing.Plane( + vectors=None, cut=cut + ) # vectors & cell should not be necessary + actual_data = slicing.grid_scalar(grid_scalar, plane, fraction) + Assert.allclose(actual_data, expected_data) + + +@pytest.mark.parametrize("cut", ("a", "b", "c")) +@pytest.mark.parametrize("fraction", (-0.4, 0, 0.4, 0.8, 1.2)) +def test_slice_grid_vector(cut, fraction, Assert): + grid_vector = np.random.random((3, 10, 12, 14)) + 0.1 + ignored = 99 + if cut == "a": + index = np.round(fraction * 10).astype(np.int_) % 10 + cell = np.diag((ignored, 2, 3)) + expected_data = grid_vector[1:, index, :, :] + elif cut == "b": + index = np.round(fraction * 12).astype(np.int_) % 12 + expected_data = grid_vector[::2, :, index, :] + cell = np.diag((2, ignored, 3)) + else: + index = np.round(fraction * 14).astype(np.int_) % 14 + expected_data = grid_vector[:2, :, :, index] + cell = np.diag((2, 3, ignored)) + plane = slicing.Plane(vectors=[[2, 0], [0, 3]], cell=cell, cut=cut) + actual_data = slicing.grid_vector(grid_vector, plane, fraction) + Assert.allclose(actual_data, expected_data) + + +@pytest.mark.parametrize("cut, index", [("a", 0), ("b", 1), ("c", 2)]) +def test_slice_grid_vector_nontrivial_cell(cut, index, Assert): + fraction = 0 + cell = np.array([[0, 1, 1.1], [0.9, 0, 1.1], [0.9, 1, 0]]) + vectors = np.delete(cell, index, axis=0) + normal = np.cross(*vectors) + plane = slicing.plane(cell, cut, normal=None) + alpha = np.linspace(-1, 2, 10) + beta = np.linspace(-2, 1, 12) + gamma = np.linspace(-1.5, 1.5, 14) + grid_vector = np.array( + [ + [ + [a * vectors[0] + b * vectors[1] + c * normal for a in alpha] + for b in beta + ] + for c in gamma + ] + ).T + if cut == "a": + expected_data = np.array( + [ + [alpha[0] * plane.vectors[0] + b * plane.vectors[1] for b in beta] + for c in gamma + ] + ).T + elif cut == "b": + expected_data = np.array( + [ + [a * plane.vectors[0] + beta[0] * plane.vectors[1] for a in alpha] + for c in gamma + ] + ).T + else: + expected_data = np.array( + [[a * plane.vectors[0] + b * plane.vectors[1] for a in alpha] for b in beta] + ).T + actual_data = slicing.grid_vector(grid_vector, plane, fraction) + Assert.allclose(actual_data, expected_data, tolerance=3)