Magnetism (orest-d#17)

* Implement logic to extract magnetism
* Implement basic magnetism reading from file
* Implement reading the data into a dictionary
* Make charges and moments accessible individually
* Add magnetism to structure.
* Implement plotting the moments along the z axis
* Implement noncollinear magnetism
* Implement charge-only case
* Add documentation for new magnetism functionality

src/py4vasp/data/__init__.py

@@ -21,6 +21,7 @@
 from .dos import Dos
 from .energy import Energy
 from .kpoints import Kpoints
+from .magnetism import Magnetism
 from .projectors import Projectors
 from .topology import Topology
 from .trajectory import Trajectory

src/py4vasp/data/magnetism.py

@@ -0,0 +1,142 @@
+from py4vasp.data import _util
+import numpy as np
+import functools
+
+_step_parameter = """Parameters
+----------
+steps : int or range
+    If present specifies which steps of the simulation should be extracted.
+    By default the data of all steps is read.
+"""
+
+_index_note = """Notes
+-----
+The index order is different compared to the raw data when noncollinear calculations
+are used. This routine returns the magnetic moments as (steps, atoms, orbitals,
+directions)."""
+
+_to_dict_doc = f""" Read the charges and magnetization data into a dictionary.
+
+{_step_parameter}
+
+Returns
+-------
+dict
+    Contains the charges and magnetic moments generated by Vasp projected
+    on atoms and orbitals.
+
+{_index_note}
+"""
+
+_charges_doc = f""" Read the charges of the selected steps.
+
+{_step_parameter}
+
+Returns
+-------
+np.ndarray
+    Contains the charges for the selected steps projected on atoms and orbitals.
+"""
+
+_total_charges_doc = f""" Read the total charges of the selected steps.
+
+{_step_parameter}
+
+Returns
+-------
+np.ndarray
+    Contains the total charges for the selected steps projected on atoms. This
+    corresponds to the charges summed over the orbitals.
+"""
+
+_moments_doc = f""" Read the magnetic moments of the selected steps.
+
+{_step_parameter}
+
+Returns
+-------
+np.ndarray
+    Contains the magnetic moments for the selected steps projected on atoms and
+    orbitals.
+
+{_index_note}
+"""
+
+_total_moments_doc = f""" Read the total magnetic moments of the selected steps.
+
+{_step_parameter}
+
+Returns
+-------
+np.ndarray
+    Contains the total magnetic moments for the selected steps projected on atoms.
+    This corresponds to the magnetic moments summed over the orbitals.
+"""
+
+
+class Magnetism:
+    """ The evolution of the magnetization over the simulation.
+
+    This class gives access to the magnetic moments and charges projected on the
+    different orbitals on every atom.
+
+    Parameters
+    ----------
+    raw_magnetism
+        Dataclass containing the charges and magnetic moments read from Vasp.
+    """
+
+    def __init__(self, raw_magnetism):
+        self._raw = raw_magnetism
+
+    @classmethod
+    @_util.add_doc(_util.from_file_doc("orbital-resolved magnetic moments"))
+    def from_file(cls, file=None):
+        return _util.from_file(cls, file, "magnetism")
+
+    @_util.add_doc(_to_dict_doc)
+    def to_dict(self, steps=None):
+        return {"charges": self.charges(steps), "moments": self.moments(steps)}
+
+    @functools.wraps(to_dict)
+    def read(self, *args):
+        return self.to_dict(*args)
+
+    @_util.add_doc(_charges_doc)
+    def charges(self, steps=None):
+        steps = self._default_steps_if_none(steps)
+        return self._raw.moments[steps, 0, :, :]
+
+    @_util.add_doc(_moments_doc)
+    def moments(self, steps=None):
+        steps = self._default_steps_if_none(steps)
+        if self._raw.moments.shape[1] == 1:
+            return None
+        elif self._raw.moments.shape[1] == 2:
+            return self._raw.moments[steps, 1, :, :]
+        else:
+            moments = self._raw.moments[steps, 1:, :, :]
+            direction_axis = 1 if moments.ndim == 4 else 0
+            return np.moveaxis(moments, direction_axis, -1)
+
+    @_util.add_doc(_total_charges_doc)
+    def total_charges(self, steps=None):
+        return self._sum_over_orbitals(self.charges(steps))
+
+    @_util.add_doc(_total_moments_doc)
+    def total_moments(self, steps=None):
+        if self._raw.moments.shape[1] == 1:
+            return None
+        elif self._raw.moments.shape[1] == 2:
+            return self._sum_over_orbitals(self.moments(steps))
+        else:
+            steps = self._default_steps_if_none(steps)
+            total_moments = self._sum_over_orbitals(self._raw.moments[steps, 1:, :, :])
+            direction_axis = 1 if total_moments.ndim == 3 else 0
+            return np.moveaxis(total_moments, direction_axis, -1)
+
+    def _default_steps_if_none(self, steps):
+        return steps if steps is not None else range(len(self._raw.moments))
+
+    def _sum_over_orbitals(self, quantity):
+        return np.sum(quantity, axis=-1)

src/py4vasp/data/structure.py

@@ -1,4 +1,4 @@
-from py4vasp.data import _util, Viewer3d, Topology
+from py4vasp.data import _util, Viewer3d, Topology, Magnetism
 import ase
 import numpy as np
 import functools
@@ -17,6 +17,9 @@ class Structure:
         Dataclass containing the raw data defining the structure.
     """
 
+    length_moments = 1.5
+    "Length in Å how a magnetic moment is displayed relative to the largest moment."
+
     def __init__(self, raw_structure):
         self._raw = raw_structure
 
@@ -39,12 +42,20 @@ def to_dict(self):
             "cell": self._raw.cell.scale * self._raw.cell.lattice_vectors[:],
             "positions": self._raw.positions[:],
             "elements": Topology(self._raw.topology).elements(),
+            **self._read_magnetic_moments(),
         }
 
     @functools.wraps(to_dict)
     def read(self):
         return self.to_dict()
 
+    def _read_magnetic_moments(self):
+        if self._raw.magnetism is not None:
+            magnetism = Magnetism(self._raw.magnetism)
+            return {"magnetic_moments": magnetism.total_moments(-1)}
+        else:
+            return {}
+
     def __len__(self):
         return len(self._raw.positions)
 
@@ -69,6 +80,8 @@ def to_ase(self, supercell=None):
             scaled_positions=data["positions"],
             pbc=True,
         )
+        if "magnetic_moments" in data:
+            structure.set_initial_magnetic_moments(data["magnetic_moments"])
         if supercell is not None:
             structure *= supercell
         return structure
@@ -85,12 +98,32 @@ def to_viewer3d(self, supercell=None):
         Returns
         -------
         Viewer3d
-            Visualize the structure as a 3d figure.
+            Visualize the structure as a 3d figure, adding the magnetic momements
+            of the atoms if present.
         """
         viewer = Viewer3d.from_structure(self, supercell=supercell)
         viewer.show_cell()
+        moments = self._prepare_magnetic_moments_for_plotting()
+        if moments is not None:
+            viewer.show_arrows_at_atoms(moments)
         return viewer
 
     @functools.wraps(to_viewer3d)
     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"]
+        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
+
+    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))

src/py4vasp/data/viewer3d.py

@@ -10,6 +10,9 @@ class _Arrow3d(NamedTuple):
     color: np.ndarray
     radius: float = 0.2
 
+    def to_serializable(self):
+        return list(self.tail), list(self.tip), list(self.color), float(self.radius)
+
 
 _x_axis = _Arrow3d(tail=np.zeros(3), tip=np.array((3, 0, 0)), color=[1, 0, 0])
 _y_axis = _Arrow3d(tail=np.zeros(3), tip=np.array((0, 3, 0)), color=[0, 1, 0])
@@ -117,4 +120,4 @@ def hide_arrows_at_atoms(self):
         self._arrows = []
 
     def _make_arrow(self, arrow):
-        return self._ngl.shape.add_arrow(*arrow)
+        return self._ngl.shape.add_arrow(*(arrow.to_serializable()))

src/py4vasp/raw/file.py

@@ -159,6 +159,21 @@ def cell(self):
             lattice_vectors=self._h5f["results/positions/lattice_vectors"],
         )
 
+    def magnetism(self):
+        """ Read the magnetisation data of the crystal.
+
+        Returns
+        -------
+        raw.Magnetism
+            The magnetic moments and charges on every atom in orbital resolved
+            representation.
+        """
+        self._assert_not_closed()
+        key = "intermediate/history/magnetism/moments"
+        if key not in self._h5f:
+            return None
+        return raw.Magnetism(moments=self._h5f[key])
+
     def structure(self):
         """ Read the structure information.
 
@@ -171,6 +186,7 @@ def structure(self):
             topology=self.topology(),
             cell=self.cell(),
             positions=self._h5f["results/positions/position_ions"],
+            magnetism=self.magnetism(),
         )
 
     def energy(self):

src/py4vasp/raw/rawdata.py

@@ -64,6 +64,14 @@ class Cell:
     __eq__ = _dataclass_equal
 
 
+@dataclass
+class Magnetism:
+    "Data about the magnetism in the system."
+    moments: np.ndarray
+    "Contains the charge and magnetic moments atom and orbital resolved."
+    __eq__ = _dataclass_equal
+
+
 @dataclass
 class Structure:
     "Structural information of the system."
@@ -73,6 +81,8 @@ class Structure:
     "Unit cell of the crystal or simulation cell for molecules."
     positions: np.ndarray
     "Position of all atoms in the unit cell in units of the lattice vectors."
+    magnetism: Magnetism = None
+    "Magnetization of every atom in the unit cell."
     __eq__ = _dataclass_equal
 
 

tests/data/test_magnetism.py

@@ -0,0 +1,93 @@
+from py4vasp.data import Magnetism
+from .test_topology import raw_topology
+import py4vasp.raw as raw
+import numpy as np
+import pytest
+
+
+@pytest.fixture
+def raw_magnetism(raw_topology):
+    number_steps = 4
+    number_atoms = len(raw_topology.elements)
+    lmax = 3
+    number_components = 2
+    shape = (number_steps, number_components, number_atoms, lmax)
+    magnetism = raw.Magnetism(moments=np.arange(np.prod(shape)).reshape(shape))
+    magnetism.charges = magnetism.moments[:, 0, :, :]
+    magnetism.total_charges = np.sum(magnetism.charges, axis=2)
+    magnetism.magnetic_moments = magnetism.moments[:, 1, :, :]
+    magnetism.total_moments = np.sum(magnetism.magnetic_moments, axis=2)
+    return magnetism
+
+
+@pytest.fixture
+def noncollinear_magnetism(raw_magnetism):
+    shape = raw_magnetism.moments.shape
+    shape = (shape[0] // 2, shape[1] * 2, 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)
+
+
+def test_read(raw_magnetism, Assert):
+    actual = Magnetism(raw_magnetism).read()
+    Assert.allclose(actual["charges"], raw_magnetism.charges)
+    Assert.allclose(actual["moments"], raw_magnetism.magnetic_moments)
+    actual = Magnetism(raw_magnetism).read(-1)
+    Assert.allclose(actual["charges"], raw_magnetism.charges[-1])
+    Assert.allclose(actual["moments"], raw_magnetism.magnetic_moments[-1])
+
+
+def test_charges(raw_magnetism, Assert):
+    actual = Magnetism(raw_magnetism).charges()
+    Assert.allclose(actual, raw_magnetism.charges)
+
+
+def test_moments(raw_magnetism, Assert):
+    actual = Magnetism(raw_magnetism).moments()
+    Assert.allclose(actual, raw_magnetism.magnetic_moments)
+
+
+def test_total_charges(raw_magnetism, Assert):
+    actual = Magnetism(raw_magnetism).total_charges()
+    Assert.allclose(actual, raw_magnetism.total_charges)
+    actual = Magnetism(raw_magnetism).total_charges(range(2))
+    Assert.allclose(actual, raw_magnetism.total_charges[0:2])
+
+
+def test_total_moments(raw_magnetism, Assert):
+    actual = Magnetism(raw_magnetism).total_moments()
+    Assert.allclose(actual, raw_magnetism.total_moments)
+    actual = Magnetism(raw_magnetism).total_moments(3)
+    Assert.allclose(actual, raw_magnetism.total_moments[3])
+
+
+def test_noncollinear(noncollinear_magnetism, Assert):
+    actual = Magnetism(noncollinear_magnetism)
+    Assert.allclose(actual.charges(), noncollinear_magnetism.moments[:, 0])
+    step = 0
+    moments = actual.moments(step)
+    for new_order in np.ndindex(moments.shape):
+        atom, orbital, component = new_order
+        old_order = (step, component + 1, atom, orbital)  # 0 component is charge
+        Assert.allclose(moments[new_order], noncollinear_magnetism.moments[old_order])
+    total_moments = actual.total_moments()
+    for new_order in np.ndindex(total_moments.shape):
+        step, atom, component = new_order
+        old_order = (step, component + 1, atom)  # 0 component is charge
+        expected_total_moment = np.sum(noncollinear_magnetism.moments[old_order])
+        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])
+    assert actual.moments() is None
+    assert actual.total_moments() is None