forked from orest-d/p4vasp
-
Notifications
You must be signed in to change notification settings - Fork 0
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
* Refactored structure viewer into its own object * Add unit tests for structure viewer - structure and arrows
- Loading branch information
Showing
9 changed files
with
1,106 additions
and
272 deletions.
There are no files selected for viewing
Large diffs are not rendered by default.
Oops, something went wrong.
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,86 @@ | ||
from py4vasp.data import _util | ||
from py4vasp.exceptions import RefinementException | ||
import functools | ||
import numpy as np | ||
|
||
|
||
class StructureViewer: | ||
"""Collection of data and elements to be displayed in a structure viewer""" | ||
|
||
def __init__(self, structure, show_cell=True, supercell=None, show_axes=False, axes_length=3, arrows=None): | ||
self.structure = structure | ||
self.show_cell = show_cell | ||
self.supercell = supercell | ||
self.show_axes = show_axes | ||
self.axes_length = axes_length | ||
self.arrows = arrows | ||
|
||
def with_arrows(self, arrows): | ||
self.arrows = arrows | ||
return self | ||
|
||
def show(self): | ||
import nglview | ||
from nglview.shape import Shape | ||
|
||
structure = self.structure.to_pymatgen() | ||
if self.supercell is not None: | ||
structure.make_supercell(self.supercell) | ||
|
||
view = nglview.show_pymatgen(structure) | ||
if self.show_cell: | ||
view.add_representation(repr_type="unitcell") | ||
if self.show_axes or self.arrows is not None: | ||
shape = Shape(view=view) | ||
if self.show_axes: | ||
shape.add_arrow( | ||
[0, 0, 0], [self.axes_length, 0, 0], [1, 0, 0], 0.2) | ||
shape.add_arrow( | ||
[0, 0, 0], [0, self.axes_length, 0], [0, 1, 0], 0.2) | ||
shape.add_arrow( | ||
[0, 0, 0], [0, 0, self.axes_length], [0, 0, 1], 0.2) | ||
if self.arrows is not None: | ||
for (x, y, z), (vx, vy, vz) in zip(structure.cart_coords, self.arrows): | ||
shape.add_arrow( | ||
[x, y, z], [x+vx, y+vy, z+vz], [0.1, 0.1, 0.8], 0.2) | ||
|
||
return view | ||
|
||
|
||
class Structure: | ||
def __init__(self, raw_structure): | ||
self._raw = raw_structure | ||
self.structure_viewer = None | ||
|
||
def read(self): | ||
return self.to_dict() | ||
|
||
def to_dict(self): | ||
return { | ||
"cell": self._raw.cell.lattice_vectors[:], | ||
"cartesian_positions": self._raw.cartesian_positions[:], | ||
"species": list(self._raw.species), | ||
} | ||
|
||
def __len__(self): | ||
return len(self._raw.cartesian_positions) | ||
|
||
def to_pymatgen(self): | ||
import pymatgen as mg | ||
return mg.Structure( | ||
lattice=mg.Lattice(self._raw.cell.lattice_vectors), | ||
species=[specie.decode("ascii") for specie in self._raw.species], | ||
coords=self._raw.cartesian_positions, | ||
coords_are_cartesian=True | ||
) | ||
|
||
def plot(self, show_cell=True, supercell=None, show_axes=False, axes_length=3): | ||
self.structure_viewer = StructureViewer( | ||
self, show_cell=show_cell, supercell=supercell, show_axes=show_axes, axes_length=axes_length) | ||
return self.structure_viewer.show() | ||
|
||
def plot_arrows(self, arrows): | ||
if self.structure_viewer is None: | ||
self.plot() | ||
self.structure_viewer = self.structure_viewer.with_arrows(arrows) | ||
return self.structure_viewer.show() |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,75 @@ | ||
from py4vasp.data import Structure | ||
from py4vasp.exceptions import RefinementException | ||
import py4vasp.raw as raw | ||
import pytest | ||
import numpy as np | ||
|
||
|
||
@pytest.fixture | ||
def raw_structure(): | ||
number_atoms = 20 | ||
shape = (number_atoms, 3) | ||
return raw.Structure( | ||
cell=raw.Cell(scale=1.0, lattice_vectors=np.eye(3)), | ||
cartesian_positions=np.arange(np.prod(shape)).reshape(shape), | ||
species=np.array(["C"]*number_atoms, dtype="S2"), | ||
) | ||
|
||
|
||
def test_read(raw_structure, Assert): | ||
actual = Structure(raw_structure).read() | ||
assert (actual["cell"] == raw_structure.cell.lattice_vectors).all() | ||
Assert.allclose(actual["cartesian_positions"], | ||
raw_structure.cartesian_positions) | ||
assert (actual["species"] == raw_structure.species).all() | ||
|
||
|
||
def test_to_pymatgen(raw_structure, Assert): | ||
structure = Structure(raw_structure) | ||
mg_structure = structure.to_pymatgen() | ||
a, b, c = mg_structure.lattice.as_dict()["matrix"] | ||
assert a == [1, 0, 0] | ||
assert b == [0, 1, 0] | ||
assert c == [0, 0, 1] | ||
|
||
|
||
def test_plot(raw_structure, Assert): | ||
structure = Structure(raw_structure) | ||
assert structure.structure_viewer is None | ||
view = structure.plot() | ||
assert structure.structure_viewer is not None | ||
|
||
view = structure.plot(show_cell=False) | ||
assert [(msg["methodName"], msg["args"]) | ||
for msg in view.get_state()["_ngl_msg_archive"]][1:] == [] | ||
|
||
view = structure.plot(show_cell=True) | ||
assert [(msg["methodName"], msg["args"]) for msg in view.get_state()[ | ||
"_ngl_msg_archive"]][1:] == [('addRepresentation', ['unitcell'])] | ||
|
||
view = structure.plot(show_cell=False, show_axes=False) | ||
assert [(msg["methodName"], msg["args"]) | ||
for msg in view.get_state()["_ngl_msg_archive"]][1:] == [] | ||
|
||
view = structure.plot(show_cell=False, show_axes=True, axes_length=5) | ||
assert [(msg["methodName"], msg["args"]) for msg in view.get_state()["_ngl_msg_archive"]][1:] == [ | ||
('addShape', [ | ||
'shape', [('arrow', [0, 0, 0], [5, 0, 0], [1, 0, 0], 0.2)]]), | ||
('addShape', [ | ||
'shape', [('arrow', [0, 0, 0], [0, 5, 0], [0, 1, 0], 0.2)]]), | ||
('addShape', [ | ||
'shape', [('arrow', [0, 0, 0], [0, 0, 5], [0, 0, 1], 0.2)]]), | ||
] | ||
|
||
|
||
def test_plot_arrows(raw_structure, Assert): | ||
structure = Structure(raw_structure) | ||
assert structure.structure_viewer is None | ||
view = structure.plot_arrows([(0, 0, 1) for i in range(len(structure))]) | ||
assert structure.structure_viewer is not None | ||
assert [(msg["methodName"], msg["args"]) for msg in view.get_state()[ | ||
"_ngl_msg_archive"]][1] == ('addRepresentation', ['unitcell']) | ||
assert [(msg["methodName"], msg["args"]) for msg in view.get_state()[ | ||
"_ngl_msg_archive"]][2] == ('addShape', ['shape', [('arrow', [0.0, 1.0, 2.0], [0.0, 1.0, 3.0], [0.1, 0.1, 0.8], 0.2)]]) | ||
assert sum(msg["methodName"] == 'addShape' for msg in view.get_state()[ | ||
"_ngl_msg_archive"]) == len(structure) |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters