Skip to content

Commit

Permalink
Only one grid scalar for a trajectory is currently implemented
Browse files Browse the repository at this point in the history
  • Loading branch information
sudarshanv01 committed Feb 22, 2024
1 parent b96c38a commit 05d4fc7
Show file tree
Hide file tree
Showing 2 changed files with 59 additions and 24 deletions.
46 changes: 30 additions & 16 deletions src/py4vasp/_third_party/view/view.py
Original file line number Diff line number Diff line change
Expand Up @@ -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")
Expand Down Expand Up @@ -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)
Expand All @@ -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()
Expand Down
37 changes: 29 additions & 8 deletions tests/third_party/view/test_view.py
Original file line number Diff line number Diff line change
Expand Up @@ -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)
Expand All @@ -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
Expand Down Expand Up @@ -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):
Expand Down

0 comments on commit 05d4fc7

Please sign in to comment.