Merge pull request #10 from JuDFTteam/develop

Major improvements

aiida_spirit/parsers.py

@@ -11,7 +11,7 @@
 from aiida.common import exceptions
 from aiida.orm import Dict, ArrayData
 from masci_tools.io.common_functions import search_string
-from .calculations import _RETLIST, _SPIRIT_STDOUT
+from .calculations import _RETLIST, _SPIRIT_STDOUT, _ATOM_TYPES
 
 SpiritCalculation = CalculationFactory('spirit')
 
@@ -52,12 +52,13 @@ def parse(self, **kwargs):
             return self.exit_codes.ERROR_MISSING_OUTPUT_FILES
 
         # parse information from output file (number of iterations, convergence info, ...)
-        output_node, mag, energ = self.parse_retrieved()
-        self.out('output_parameters', output_node)
-        self.out('magnetization', mag)
-        self.out('energies', energ)
+        retrieved_dict = self.parse_retrieved()
+
+        for key, value in retrieved_dict.items():
+            self.out(key, value)
 
         # check consistency of spirit_version_info with the inputs
+        output_node = retrieved_dict['output_parameters']
         if 'pinning' in self.node.inputs:
             version_info = output_node['spirit_version_info']
             if not 'enabled' in version_info['Pinning']:
@@ -69,7 +70,19 @@ def parse(self, **kwargs):
 
         return ExitCode(0)
 
-    def parse_retrieved(self):
+    def _retrieve_if_found(self, filename, *args, **kwargs):
+        """Retrieves a file and loads it with `np.loadtxt`.
+        The `*args` and `**kwargs` are passed to `np.loadtxt`.
+        If the file is not found it returns None."""
+        retrieved = self.retrieved
+        if filename in retrieved.list_object_names():
+            with retrieved.open(filename, 'r') as _f:
+                return np.loadtxt(_f, *args, **kwargs)
+        else:
+            self.logger.info('{} not found!'.format(filename))
+            return None
+
+    def parse_retrieved(self):  # pylint: disable=too-many-locals
         """Parse the output from the retrieved and create aiida nodes"""
 
         retrieved = self.retrieved
@@ -84,32 +97,82 @@ def parse_retrieved(self):
 
         # parse output files
         self.logger.info('Parsing energy archive')
-        with retrieved.open('spirit_Image-00_Energy-archive.txt') as _f:
-            energ = np.loadtxt(_f, skiprows=1)
+        energ = self._retrieve_if_found('spirit_Image-00_Energy-archive.txt',
+                                        skiprows=1)
+
         self.logger.info('Parsing initial magnetization')
-        with retrieved.open('spirit_Image-00_Spins-initial.ovf') as _f:
-            m_init = np.loadtxt(_f)
+        m_init = self._retrieve_if_found('spirit_Image-00_Spins-initial.ovf')
+
         self.logger.info('Parsing final magnetization')
-        with retrieved.open('spirit_Image-00_Spins-final.ovf') as _f:
-            m_final = np.loadtxt(_f)
+        m_final = self._retrieve_if_found('spirit_Image-00_Spins-final.ovf')
+
+        self.logger.info('Parsing atom types')
+        atyp = self._retrieve_if_found(_ATOM_TYPES)
+
+        self.logger.info('Parsing MC output')
+        out_mc = self._retrieve_if_found('output_mc.txt')
+
+        # Write dictionary of retrieved quantities
+        _retrieved_dict = {'output_parameters': output_node}
 
         # collect arrays in ArrayData
-        mag = ArrayData()
-        mag.set_array(
-            'initial',
-            np.nan_to_num(m_init))  # nan_to_num is needed with defects
-        mag.set_array('final', np.nan_to_num(m_final))
-        mag.extras['description'] = {
-            'initial': 'initial directions of the magnetization vectors',
-            'final': 'final directions of the magnetization vectors',
-        }
-        energies = ArrayData()
-        energies.set_array('energies', energ)
-        energies.extras['description'] = {
-            'energies': 'energy convergence',
-        }
-
-        return output_node, mag, energies
+        if m_init is not None and m_final is not None:
+            mag = ArrayData()
+            mag.set_array(
+                'initial',
+                np.nan_to_num(m_init))  # nan_to_num is needed with defects
+            mag.set_array('final', np.nan_to_num(m_final))
+            mag.extras['description'] = {
+                'initial': 'initial directions of the magnetization vectors',
+                'final': 'final directions of the magnetization vectors',
+            }
+            _retrieved_dict.update({'magnetization': mag})
+
+        if energ is not None:
+            energies = ArrayData()
+            energies.set_array('energies', energ)
+            energies.extras['description'] = {
+                'energies': 'Energy convergence with iterations.',
+            }
+            _retrieved_dict.update({'energies': energies})
+
+        if atyp is not None:
+            atypes = ArrayData()
+            atypes.set_array('atom_types', atyp)
+            atypes.extras['description'] = {
+                'atom_types': 'list of atom types for all positions',
+            }
+            _retrieved_dict.update({'atom_types': atypes})
+
+        # Only add mc if it is found
+        if out_mc is not None:
+            output_mc = ArrayData()
+
+            # Associante the columns of out_mc with individual arrays
+            array_names = [
+                'temperature', 'energy', 'magnetization', 'susceptibility',
+                'specific_heat', 'binder_cumulant'
+            ]
+            for i, name in enumerate(array_names):
+                output_mc.set_array(name, out_mc[:, i])
+
+            output_mc.extras['description'] = {
+                'temperature':
+                'The temperature at which the sampling was performed',
+                'energy':
+                'The average energy at the sampled temperature',
+                'magnetization':
+                'The average spin direction at the sample temperature',
+                'susceptibility':
+                'The susceptibility at the sampled temperature',
+                'specific_heat':
+                'The specific heat at the sampled temperature',
+                'binder_cumulant':
+                'The binder_cumulant at the sampled temperature',
+            }
+            _retrieved_dict.update({'monte_carlo': output_mc})
+
+        return _retrieved_dict
 
 
 def parse_outfile(txt):

aiida_spirit/tools/_vfr.py

@@ -15,7 +15,7 @@
 # pylint: disable=line-too-long,global-statement
 
 
-def setup(vfr_frame_id=''):
+def setup(vfr_frame_id='', height_px=600, width_percent=100):
     """Create a frame in the jupyter ntoebook to into which the spins are shown."""
     global _vfr_frame_id_mapping, _next_frame_id, _frame_id_suffix
     if vfr_frame_id not in _vfr_frame_id_mapping:
@@ -58,7 +58,9 @@ def setup(vfr_frame_id=''):
                 '''");
         var iframe_url = "https://judftteam.github.io/aiida-spirit/vfr_notebook_view/?origin=" + window.location.origin + "&frame_id='''
                 + vfr_frame_id + '''";
-        frame_wrapper.innerHTML = '<iframe src="' + iframe_url + '" style="width: 100%; height:600px; border: 1px solid black;"></iframe>';
+        frame_wrapper.innerHTML = '<iframe src="' + iframe_url + '" style="width: '''
+                + str(width_percent) + '''%; height:''' + str(height_px) +
+                '''px; border: 1px solid black;"></iframe>';
     </script>
     ''')
 

aiida_spirit/tools/plotting.py

@@ -8,18 +8,19 @@
 from .get_from_remote import list_remote_files, get_file_content_from_remote
 
 
-def init_spinview(vfr_frame_id=''):
+def init_spinview(vfr_frame_id='', height_px=600, width_percent=100):
     """
     Initialize the vfrendering HTML object.
 
     Needs to be called before the show_spins function can set the spins
     :param vfr_frame_id: a string that controls if multiple windows
     should be openend. Use the same string in show_spins to update this.
-    This is not fully implemented yet and does not work in this version.
+    :param height_px: height of the window in pixels
+    :param width_percent: window width in percent of the current width size
     """
 
     # initialize vfrendering HTML object
-    view = setup(vfr_frame_id)
+    view = setup(vfr_frame_id, height_px, width_percent)
 
     return view
 
@@ -37,20 +38,22 @@ def _plot_spins_vfr(  # pylint: disable=too-many-arguments
 
     # set positions and direciton of the spins
     positions = []
-    for iz in range(n_basis_cells[0]):
+    for iz in range(n_basis_cells[2]):
         for iy in range(n_basis_cells[1]):
-            for ix in range(n_basis_cells[2]):
+            for ix in range(n_basis_cells[0]):
                 for p in pos_cell:
                     # set position
                     positions.append(p + ix * cell[0] + iy * cell[1] +
                                      iz * cell[2])
-    # make flattened array
+    # make flattened arrays
     positions = np.array(positions).reshape(-1, 3)
+    directions = np.array(spin_directions).reshape(-1, 3)
 
     # normalize directions
-    directions = np.array(spin_directions).reshape(-1, 3)
-    for ixyz in range(len(directions)):
-        directions[ixyz, :] /= np.linalg.norm(directions[ixyz, :])
+    norm = np.linalg.norm(directions, axis=1)
+    norm[norm == 0] = 1  # prevent divide by zero
+    for ixyz in range(3):
+        directions[:, ixyz] /= norm
 
     # scaling factor for the directions
     directions *= scale_spins
@@ -71,12 +74,13 @@ def _plot_spins_vfr(  # pylint: disable=too-many-arguments
         )
 
 
-def show_spins(  # pylint: disable=inconsistent-return-statements,too-many-arguments
+def show_spins(  # pylint: disable=inconsistent-return-statements,too-many-arguments,too-many-locals, too-many-branches
         spirit_calc,
         show_final_structure=True,
         scale_spins=1.0,
         list_spin_files_on_remote=False,
         use_remote_spins_id=None,
+        mask=None,
         vfr_frame_id=''):
     """
     Update the vfrendering spin view plot with the final or initial spin structure.
@@ -89,6 +93,8 @@ def show_spins(  # pylint: disable=inconsistent-return-statements,too-many-argum
     :param list_spin_files_on_remote: print a list of the available spin image files on the remote folder.
     :param use_remote_spins_id: show neither final nor initial spin structure but show the structure of
         a certain checkpoint (see list_spin_files_on_remote=True output for available checkpoints).
+    :param mask: mask which can be used to hide or rescale spins (mask is multiplied to the spins,
+                 i.e. mask==0 hides a spin, mask>1 enhaces its size).
     :param vfr_frame_id: if given this allows to control into which spinview frame the spins are shown.
         Should be the same as in the init_spinview. This is not fully implemented yet and does not work in this version.
     """
@@ -104,11 +110,40 @@ def show_spins(  # pylint: disable=inconsistent-return-statements,too-many-argum
     pos_cell = np.array([i.position for i in struc.sites])
 
     # get initial or final spin directions
+    if 'magnetization' not in spirit_calc.outputs:
+        if not spirit_calc.is_finished_ok and spirit_calc.is_terminated:
+            raise ValueError('SpiritCalculation did not finish ok.')
+        print(
+            'SpiritCalculation does not have magnetization output node. Cannot plot anything'
+        )
+        return
+
     m = spirit_calc.outputs.magnetization
+    minit = m.get_array('initial')
+    mfinal = m.get_array('final')
     if show_final_structure:
-        m = m.get_array('final')
+        m = mfinal
     else:
-        m = m.get_array('initial')
+        m = minit
+
+    # apply a scaling mask
+    if mask is not None:
+        if m[:, 0].shape != mask.shape:
+            raise ValueError(
+                f'mask array is not of the right shape. Got {mask.shape} but expected {m[:,0].shape}.'
+            )
+        for ixyz in range(3):
+            m[:, ixyz] *= mask
+
+    if 'defects' in spirit_calc.inputs:
+        if 'atom_types' in spirit_calc.outputs:
+            # hide defects
+            atom_types = spirit_calc.outputs.atom_types.get_array('atom_types')
+            m[atom_types < 0] = 0
+        else:
+            # fallback if atom_types are not there
+            # these are the positions where the initial and final spins are the same (hide those if we have defects)
+            m[(minit == mfinal).all(axis=1)] = 0
 
     # print a list of files that are still on the remote and which can be plotted
     if list_spin_files_on_remote or use_remote_spins_id is not None: