From 05d4fc7f7b52752558f7351b4ea273470d1f6768 Mon Sep 17 00:00:00 2001 From: Sudarshan Vijay Date: Thu, 22 Feb 2024 11:57:59 +0100 Subject: [PATCH] Only one grid scalar for a trajectory is currently implemented --- src/py4vasp/_third_party/view/view.py | 46 +++++++++++++++++---------- tests/third_party/view/test_view.py | 37 ++++++++++++++++----- 2 files changed, 59 insertions(+), 24 deletions(-) diff --git a/src/py4vasp/_third_party/view/view.py b/src/py4vasp/_third_party/view/view.py index a91a495c..93605d65 100644 --- a/src/py4vasp/_third_party/view/view.py +++ b/src/py4vasp/_third_party/view/view.py @@ -9,6 +9,7 @@ import numpy as np import numpy.typing as npt +from py4vasp import exception from py4vasp._util import convert, import_ ase = import_.optional("ase") @@ -112,9 +113,21 @@ def _create_atoms(self, idx_traj): atoms = atoms.repeat(self.supercell) return atoms + def _iterate_trajectory_frames(self): + if self.grid_scalars: + for grid_scalar in self.grid_scalars: + if len(self.positions) > 1 and len(grid_scalar.quantity) > 1: + raise exception.NotImplemented( + """\ +Currently isosurfaces are implemented only for cases where there is only one frame in +the trajectory. Make sure that either only one frame for the positions attribute is +supplied with its corresponding grid scalar.""" + ) + return range(len(self.positions)) + def to_ngl(self): trajectory = [] - for idx_traj in range(len(self.lattice_vectors)): + for idx_traj in self._iterate_trajectory_frames(): atoms = self._create_atoms(idx_traj) trajectory.append(atoms) ngl_trajectory = nglview.ASETrajectory(trajectory) @@ -136,25 +149,26 @@ def to_ngl(self): return widget def show_isosurface(self, widget): - iter_traj = list(range(len(self.lattice_vectors))) + iter_traj = self._iterate_trajectory_frames() for grid_scalar, idx_traj in itertools.product(self.grid_scalars, iter_traj): atoms = self._create_atoms(idx_traj) - data = grid_scalar.quantity[idx_traj] - with tempfile.TemporaryDirectory() as tmp: - filename = os.path.join(tmp, CUBE_FILENAME) - ase_cube.write_cube(open(filename, "w"), atoms=atoms, data=data) - widget.add_component(filename) - if grid_scalar.isosurfaces: - for isosurface in grid_scalar.isosurfaces: - isosurface_options = { - "isolevel": isosurface.isolevel, - "color": isosurface.color, - "opacity": isosurface.opacity, - } - widget.add_surface(**isosurface_options) + if idx_traj == 0: + data = grid_scalar.quantity[idx_traj] + with tempfile.TemporaryDirectory() as tmp: + filename = os.path.join(tmp, CUBE_FILENAME) + ase_cube.write_cube(open(filename, "w"), atoms=atoms, data=data) + widget.add_component(filename) + if grid_scalar.isosurfaces: + for isosurface in grid_scalar.isosurfaces: + isosurface_options = { + "isolevel": isosurface.isolevel, + "color": isosurface.color, + "opacity": isosurface.opacity, + } + widget.add_surface(**isosurface_options) def show_arrows_at_atoms(self, widget): - iter_traj = list(range(len(self.lattice_vectors))) + iter_traj = self._iterate_trajectory_frames() for _arrows, idx_traj in itertools.product(self.ion_arrows, iter_traj): atoms = self._create_atoms(idx_traj) _, transformation = atoms.cell.standard_form() diff --git a/tests/third_party/view/test_view.py b/tests/third_party/view/test_view.py index 92c1266f..05ed9f4a 100644 --- a/tests/third_party/view/test_view.py +++ b/tests/third_party/view/test_view.py @@ -92,8 +92,8 @@ def view3d(request, not_core): charge_grid_scalar.isosurfaces = [isosurface] grid_scalars = [charge_grid_scalar] else: - charge_grid_scalar = GridQuantity(np.random.rand(2, 12, 10, 8), "charge") - potential_grid_scalar = GridQuantity(np.random.rand(2, 12, 10, 8), "potential") + charge_grid_scalar = GridQuantity(np.random.rand(1, 12, 10, 8), "charge") + potential_grid_scalar = GridQuantity(np.random.rand(1, 12, 10, 8), "potential") potential_grid_scalar.isosurfaces = [isosurface] grid_scalars = [charge_grid_scalar, potential_grid_scalar] view = View(grid_scalars=grid_scalars, **inputs) @@ -102,6 +102,18 @@ def view3d(request, not_core): return view +@pytest.fixture +def view3d_fail(not_core): + inputs = base_input_view(is_structure=False) + isosurface = Isosurface(isolevel=0.1, color="#2FB5AB", opacity=0.6) + charge_grid_scalar = GridQuantity(np.random.rand(2, 12, 10, 8), "charge") + potential_grid_scalar = GridQuantity(np.random.rand(2, 12, 10, 8), "potential") + potential_grid_scalar.isosurfaces = [isosurface] + grid_scalars = [charge_grid_scalar, potential_grid_scalar] + view = View(grid_scalars=grid_scalars, **inputs) + return view + + @pytest.fixture(params=[True, False]) def view_arrow(request, not_core): is_structure = request.param @@ -164,12 +176,21 @@ def test_isosurface(view3d): assert widget.get_state()["_ngl_msg_archive"][1]["args"][0]["binary"] == False for idx in range(len(view3d.lattice_vectors)): for grid_scalar in view3d.ref.grid_scalars: - expected_data = grid_scalar.quantity[idx] - state = widget.get_state() - output_cube = state["_ngl_msg_archive"][idx + 1]["args"][0]["data"] - output_data = ase_cube.read_cube(io.StringIO(output_cube))["data"] - assert expected_data.shape == output_data.shape - np.allclose(expected_data, output_data) + # If you pass in a grid scalar into a trajectory, I presume that you want to view + # the isosurface only for the first index of the trajectory. If you have more than one + # grid scalar in your data file then you should get an error. + if idx == 0: + expected_data = grid_scalar.quantity[idx] + state = widget.get_state() + output_cube = state["_ngl_msg_archive"][idx + 1]["args"][0]["data"] + output_data = ase_cube.read_cube(io.StringIO(output_cube))["data"] + assert expected_data.shape == output_data.shape + np.allclose(expected_data, output_data) + + +@pytest.mark.xfail +def test_fail_isosurface(view3d_fail): + widget = view3d_fail.to_ngl() def test_ion_arrows(view_arrow):