Fix magnetic moments (orest-d#18)

* Fix structure class when only charges are contained in Magnetism
* Fix vanishing magnetic moments
* Fix position of arrows for nonstandard cells

src/py4vasp/data/magnetism.py

@@ -96,7 +96,9 @@ def from_file(cls, file=None):
 
     @_util.add_doc(_to_dict_doc)
     def to_dict(self, steps=None):
-        return {"charges": self.charges(steps), "moments": self.moments(steps)}
+        moments = self.moments(steps)
+        moments = {"moments": moments} if moments is not None else {}
+        return {"charges": self.charges(steps), **moments}
 
     @functools.wraps(to_dict)
     def read(self, *args):

src/py4vasp/data/structure.py

@@ -38,11 +38,12 @@ def to_dict(self):
             all the atoms in units of the lattice vectors and the elements of
             the atoms.
         """
+        moments = self._read_magnetic_moments()
         return {
             "cell": self._raw.cell.scale * self._raw.cell.lattice_vectors[:],
             "positions": self._raw.positions[:],
             "elements": Topology(self._raw.topology).elements(),
-            **self._read_magnetic_moments(),
+            **({"magnetic_moments": moments} if moments is not None else {}),
         }
 
     @functools.wraps(to_dict)
@@ -51,10 +52,9 @@ def read(self):
 
     def _read_magnetic_moments(self):
         if self._raw.magnetism is not None:
-            magnetism = Magnetism(self._raw.magnetism)
-            return {"magnetic_moments": magnetism.total_moments(-1)}
+            return Magnetism(self._raw.magnetism).total_moments(-1)
         else:
-            return {}
+            return None
 
     def __len__(self):
         return len(self._raw.positions)
@@ -112,18 +112,26 @@ def to_viewer3d(self, supercell=None):
     def plot(self, *args):
         return self.to_viewer3d(*args)
 
-    def _prepare_magnetic_moments_for_plotting(self,):
-        if self._raw.magnetism is None:
-            return None
-        moments = self._read_magnetic_moments()["magnetic_moments"]
+    def _prepare_magnetic_moments_for_plotting(self):
+        moments = self._read_magnetic_moments()
         moments = self._convert_to_moment_to_3d_vector(moments)
-        rescale_moments = self.length_moments / np.max(np.linalg.norm(moments, axis=1))
-        return rescale_moments * moments
+        max_length_moments = self._max_length_moments(moments)
+        if max_length_moments > 1e-15:
+            rescale_moments = self.length_moments / max_length_moments
+            return rescale_moments * moments
+        else:
+            return None
 
     def _convert_to_moment_to_3d_vector(self, moments):
-        if moments.ndim == 2:
-            return moments
-        moments = moments.reshape((len(moments), 1))
-        no_new_moments = (0, 0)
-        add_zero_for_xy_axis = (2, 0)
-        return np.pad(moments, (no_new_moments, add_zero_for_xy_axis))
+        if moments is not None and moments.ndim == 1:
+            moments = moments.reshape((len(moments), 1))
+            no_new_moments = (0, 0)
+            add_zero_for_xy_axis = (2, 0)
+            moments = np.pad(moments, (no_new_moments, add_zero_for_xy_axis))
+        return moments
+
+    def _max_length_moments(self, moments):
+        if moments is not None:
+            return np.max(np.linalg.norm(moments, axis=1))
+        else:
+            return 0.0

src/py4vasp/data/viewer3d.py

@@ -49,6 +49,9 @@ 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)
+        # ngl works with the standard form, so we need to store the positions in the same format
+        standard_cell, _ = ase.cell.standard_form()
+        ase.set_cell(standard_cell, scale_atoms=True)
         res = cls(nglview.show_ase(ase))
         res._positions = ase.positions
         if supercell is not None:

tests/data/test_magnetism.py

@@ -28,6 +28,14 @@ def noncollinear_magnetism(raw_magnetism):
     return raw_magnetism
 
 
+@pytest.fixture
+def charge_only(raw_magnetism):
+    shape = raw_magnetism.moments.shape
+    shape = (shape[0] * shape[1], 1, shape[2], shape[3])
+    raw_magnetism.moments = raw_magnetism.moments.reshape(shape)
+    return raw_magnetism
+
+
 def test_from_file(raw_magnetism, mock_file, check_read):
     with mock_file("magnetism", raw_magnetism) as mocks:
         check_read(Magnetism, mocks, raw_magnetism)
@@ -83,11 +91,9 @@ def test_noncollinear(noncollinear_magnetism, Assert):
         Assert.allclose(total_moments[new_order], expected_total_moment)
 
 
-def test_charge_only(raw_magnetism, Assert):
-    shape = raw_magnetism.moments.shape
-    shape = (shape[0] * shape[1], 1, shape[2], shape[3])
-    raw_magnetism.moments = raw_magnetism.moments.reshape(shape)
-    actual = Magnetism(raw_magnetism)
-    Assert.allclose(actual.charges(), raw_magnetism.moments[:, 0])
+def test_charge_only(charge_only, Assert):
+    actual = Magnetism(charge_only)
+    Assert.allclose(actual.charges(), charge_only.moments[:, 0])
     assert actual.moments() is None
     assert actual.total_moments() is None
+    assert "moments" not in actual.read()

tests/data/test_structure.py

@@ -1,7 +1,7 @@
 from unittest.mock import patch
 from py4vasp.data import Structure, Magnetism
 from .test_topology import raw_topology
-from .test_magnetism import raw_magnetism, noncollinear_magnetism
+from .test_magnetism import raw_magnetism, noncollinear_magnetism, charge_only
 import py4vasp.data as data
 import py4vasp.raw as raw
 import pytest
@@ -15,7 +15,8 @@ def raw_structure(raw_topology):
     structure = raw.Structure(
         topology=raw_topology,
         cell=raw.Cell(scale=2.0, lattice_vectors=np.eye(3)),
-        positions=np.arange(np.prod(shape)).reshape(shape) / np.prod(shape),
+        # shift the positions of 0 to avoid relative comparison between tiny numbers
+        positions=(0.5 + np.arange(np.prod(shape)).reshape(shape)) / np.prod(shape),
     )
     structure.actual_cell = structure.cell.scale * structure.cell.lattice_vectors
     return structure
@@ -122,3 +123,27 @@ def test_noncollinear_magnetism(noncollinear_magnetism, raw_structure, Assert):
     rescale_moments = Structure.length_moments / largest_moment
     for actual, expected in zip(actual_moments, expected_moments):
         Assert.allclose(actual, expected * rescale_moments)
+
+
+def test_charge_only(charge_only, raw_structure, Assert):
+    raw_structure.magnetism = charge_only
+    structure = Structure(raw_structure)
+    assert "magnetic_moments" not in structure.read()
+    cm_init = patch.object(data.Viewer3d, "__init__", autospec=True, return_value=None)
+    cm_cell = patch.object(data.Viewer3d, "show_cell")
+    cm_arrows = patch.object(data.Viewer3d, "show_arrows_at_atoms")
+    with cm_init, cm_cell, cm_arrows as arrows:
+        structure.plot()
+        arrows.assert_not_called()
+
+
+def test_vanishing_moments(raw_magnetism, raw_structure, Assert):
+    raw_magnetism.moments = np.zeros_like(raw_magnetism.moments)
+    raw_structure.magnetism = raw_magnetism
+    structure = Structure(raw_structure)
+    cm_init = patch.object(data.Viewer3d, "__init__", autospec=True, return_value=None)
+    cm_cell = patch.object(data.Viewer3d, "show_cell")
+    cm_arrows = patch.object(data.Viewer3d, "show_arrows_at_atoms")
+    with cm_init, cm_cell, cm_arrows as arrows:
+        structure.plot()
+        arrows.assert_not_called()

tests/data/test_viewer3d.py

@@ -123,3 +123,11 @@ def test_serializable():
     arrow = _Arrow3d(np.zeros(3), np.ones(3), np.ones(1))
     for element in arrow.to_serializable():
         json.json_clean(element)
+
+
+def test_standard_form(raw_structure, Assert):
+    x = np.sqrt(0.5)
+    raw_structure.cell.lattice_vectors = np.array([[x, x, 0], [-x, x, 0], [0, 0, 1]])
+    viewer = make_viewer(raw_structure)
+    expected_positions = raw_structure.cell.scale * raw_structure.positions
+    Assert.allclose(viewer._positions, expected_positions)