PwParser: Fix trajectory verification for fixed_coords (#794)

For `relax` and `vc-relax` calculations, the `PwParser` does an extra
check on the `TrajectoryData` output to see if the structure has been
properly converged. However, when comparing the final forces with the
threshold, it does not consider that some site coordinates may have been
fixed with the `fixed_coords` setting.

Here we extract the `fixed_coords` setting from the `PwCalculation`
inputs, and pass it to the `verify_convergence_trajectory` function.
This has been adapted to ignore forces for fixed coordinates by setting
the corresponding forces to zero before comparing them with the force
threshold.

src/aiida_quantumespresso/calculations/__init__.py

@@ -2,17 +2,17 @@
 """Base `CalcJob` for implementations for pw.x and cp.x of Quantum ESPRESSO."""
 import abc
 import copy
-from functools import partial
 import numbers
 import os
 from types import MappingProxyType
 import warnings
 
 from aiida import orm
-from aiida.common import datastructures, exceptions
+from aiida.common import AttributeDict, datastructures, exceptions
 from aiida.common.lang import classproperty
 from aiida.common.warnings import AiidaDeprecationWarning
 from aiida.plugins import DataFactory
+import numpy
 from qe_tools.converters import get_parameters_from_cell
 
 from aiida_quantumespresso.data.hubbard_structure import HubbardStructureData
@@ -163,9 +163,31 @@ def validate_inputs(cls, value, port_namespace):
         if structure_kinds != pseudo_kinds:
             return f'The `pseudos` specified and structure kinds do not match: {pseudo_kinds} vs {structure_kinds}'
 
+        if 'settings' in value:
+            settings = _uppercase_dict(value['settings'].get_dict(), dict_name='settings')
+
+            # Validate the FIXED_COORDS setting
+            fixed_coords = settings.get('FIXED_COORDS', None)
+
+            if fixed_coords is not None:
+
+                fixed_coords = numpy.array(fixed_coords)
+
+                if len(fixed_coords.shape) != 2 or fixed_coords.shape[1] != 3:
+                    return 'The `fixed_coords` setting must be a list of lists with length 3.'
+
+                if fixed_coords.dtype != bool:
+                    return 'All elements in the `fixed_coords` setting lists must be either `True` or `False`.'
+
+                if 'structure' in value:
+                    nsites = len(value['structure'].sites)
+
+                    if len(fixed_coords) != nsites:
+                        return f'Input structure has {nsites} sites, but fixed_coords has length {len(fixed_coords)}'
+
     @classmethod
     def validate_parallelization(cls, value, _):
-        """Validate the ``parallelization`` port."""
+        """Validate the ``parallelization`` input."""
         if value:
             value_dict = value.get_dict()
             unknown_flags = set(value_dict.keys()) - set(cls._ENABLED_PARALLELIZATION_FLAGS)
@@ -501,47 +523,36 @@ def _generate_PWCPinputdata(cls, parameters, settings, pseudos, structure, kpoin
         del atomic_species_card_list
 
         # ------------ ATOMIC_POSITIONS -----------
-        # Check on validity of FIXED_COORDS
-        fixed_coords_strings = []
-        fixed_coords = settings.pop('FIXED_COORDS', None)
-        if fixed_coords is None:
-            # No fixed_coords specified: I store a list of empty strings
-            fixed_coords_strings = [''] * len(structure.sites)
+        coordinates = [site.position for site in structure.sites]
+        if use_fractional:
+            atomic_positions_card_header = 'ATOMIC_POSITIONS crystal\n'
+            coordinates = numpy.dot(coordinates, numpy.linalg.inv(numpy.array(structure.cell))).tolist()
         else:
-            if len(fixed_coords) != len(structure.sites):
-                raise exceptions.InputValidationError(
-                    f'Input structure contains {len(structure.sites)} sites, but fixed_coords has length '
-                    f'{len(fixed_coords)}'
-                )
+            atomic_positions_card_header = 'ATOMIC_POSITIONS angstrom\n'
 
-            for i, this_atom_fix in enumerate(fixed_coords):
-                if len(this_atom_fix) != 3:
-                    raise exceptions.InputValidationError(f'fixed_coords({i + 1:d}) has not length three')
-                for fixed_c in this_atom_fix:
-                    if not isinstance(fixed_c, bool):
-                        raise exceptions.InputValidationError(f'fixed_coords({i + 1:d}) has non-boolean elements')
-
-                if_pos_values = [cls._if_pos(_) for _ in this_atom_fix]
-                fixed_coords_strings.append('  {:d} {:d} {:d}'.format(*if_pos_values))  # pylint: disable=consider-using-f-string
+        atomic_positions_card_list = [
+            '{0} {1:18.10f} {2:18.10f} {3:18.10f}'.format(site.kind_name.ljust(6), *site_coords)  # pylint: disable=consider-using-f-string
+            for site, site_coords in zip(structure.sites, coordinates)
+        ]
 
-        abs_pos = [_.position for _ in structure.sites]
-        if use_fractional:
-            import numpy as np
-            atomic_positions_card_list = ['ATOMIC_POSITIONS crystal\n']
-            coordinates = np.dot(np.array(abs_pos), np.linalg.inv(np.array(structure.cell)))
-        else:
-            atomic_positions_card_list = ['ATOMIC_POSITIONS angstrom\n']
-            coordinates = abs_pos
+        fixed_coords = settings.pop('FIXED_COORDS', None)
 
-        for site, site_coords, fixed_coords_string in zip(structure.sites, coordinates, fixed_coords_strings):
-            atomic_positions_card_list.append(
-                '{0} {1:18.10f} {2:18.10f} {3:18.10f} {4}\n'.format(  # pylint: disable=consider-using-f-string
-                    site.kind_name.ljust(6), site_coords[0], site_coords[1], site_coords[2], fixed_coords_string
+        if fixed_coords is not None:
+            # Note that this check is also in the validation of the top-level namespace, but this is only checked in
+            # in case the structure is provided
+            if len(fixed_coords) != len(structure.sites):
+                raise exceptions.InputValidationError(
+                    f'Input structure has {len(structure.sites)} sites, but fixed_coords has length {len(fixed_coords)}'
                 )
-            )
+            fixed_coords_strings = [
+                ' {:d} {:d} {:d}'.format(*row) for row in numpy.int32(numpy.invert(fixed_coords)).tolist()  # pylint: disable=consider-using-f-string
+            ]
+            atomic_positions_card_list = [
+                atomic_pos_str + fixed_coords_str
+                for atomic_pos_str, fixed_coords_str in zip(atomic_positions_card_list, fixed_coords_strings)
+            ]
 
-        atomic_positions_card = ''.join(atomic_positions_card_list)
-        del atomic_positions_card_list
+        atomic_positions_card = atomic_positions_card_header + '\n'.join(atomic_positions_card_list) + '\n'
 
         # Optional ATOMIC_FORCES card
         atomic_forces = settings.pop('ATOMIC_FORCES', None)
@@ -750,17 +761,6 @@ def _generate_PWCPinputdata(cls, parameters, settings, pseudos, structure, kpoin
 
         return inputfile, local_copy_list_to_append
 
-    @staticmethod
-    def _if_pos(fixed):
-        """Return 0 if fixed is True, 1 otherwise.
-
-        Useful to convert from the boolean value of fixed_coords to the value required by Quantum Espresso as if_pos.
-        """
-        if fixed:
-            return 0
-
-        return 1
-
 
 def _lowercase_dict(dictionary, dict_name):
     return _case_transform_dict(dictionary, dict_name, '_lowercase_dict', str.lower)

src/aiida_quantumespresso/parsers/pw.py

@@ -227,6 +227,7 @@ def is_ionically_converged(self, trajectory, except_final_scf=False):
         :param trajectory: the output trajectory data
         :param except_final_scf: if True will return whether the calculation is converged except for the final scf.
         """
+        from aiida_quantumespresso.calculations import _uppercase_dict
         from aiida_quantumespresso.utils.defaults.calculation import pw
         from aiida_quantumespresso.utils.validation.trajectory import verify_convergence_trajectory
 
@@ -236,6 +237,11 @@ def is_ionically_converged(self, trajectory, except_final_scf=False):
         threshold_stress = parameters.get('CELL', {}).get('press_conv_thr', pw.press_conv_thr)
         external_pressure = parameters.get('CELL', {}).get('press', 0)
 
+        if 'settings' in self.node.inputs:
+            settings = _uppercase_dict(self.node.inputs.settings.get_dict(), dict_name='settings')
+
+        fixed_coords = settings.get('FIXED_COORDS', None)
+
         # Through the `cell_dofree` the degrees of freedom of the cell can be constrained, which makes the threshold on
         # the stress hard to interpret. Therefore, unless the `cell_dofree` is set to the default `all` where the cell
         # is fully unconstrained, the stress is ignored even if an explicit `press_conv_thr` is specified in the inputs.
@@ -245,10 +251,15 @@ def is_ionically_converged(self, trajectory, except_final_scf=False):
             threshold_stress = None
 
         if relax_type == 'relax':
-            return verify_convergence_trajectory(trajectory, -1, *[threshold_forces, None])
+            return verify_convergence_trajectory(
+                trajectory=trajectory,
+                index=-1,
+                threshold_forces=threshold_forces,
+                fixed_coords=fixed_coords
+            )
 
         if relax_type == 'vc-relax':
-            values = [threshold_forces, threshold_stress, external_pressure]
+            values = [threshold_forces, threshold_stress, external_pressure, fixed_coords]
             converged_relax = verify_convergence_trajectory(trajectory, -2, *values)
             converged_final = verify_convergence_trajectory(trajectory, -1, *values)
 

src/aiida_quantumespresso/utils/validation/trajectory.py

@@ -1,31 +1,42 @@
 # -*- coding: utf-8 -*-
 """Utilities for the validation of `TrajectoryData` content."""
+from typing import List, Optional
+
+from aiida.orm import TrajectoryData
 import numpy
 
 
 def verify_convergence_trajectory(
-    trajectory, index=-1, threshold_forces=None, threshold_stress=None, reference_pressure=0.
-):
-    """Verify that the data of the given `TrajectoryData` is converged with respect to the given thresholds.
-
-    There are up to three properties controlled for convergence: total energy, forces and stress.
-    If no threshold is specified the check is skipped. If any of the checks fails, either because the corresponding
-    array, or index within that array, does not exist in the `TrajectoryData`, None is returned. If all thresholds are
-    successfully met, `True` is returned, if at least one of them fails, `False` is returned.
+    trajectory: TrajectoryData,
+    index: int = -1,
+    threshold_forces: Optional[float] = None,
+    threshold_stress: Optional[float] = None,
+    reference_pressure: float = 0,
+    fixed_coords: Optional[List[List[bool]]] = None
+) -> bool:
+    """Verify that the data of the given ``TrajectoryData`` is converged with respect to the given thresholds.
+
+    Two are properties checked for convergence: forces and stress. If no threshold is specified for a property, the
+    check is skipped. If all thresholds are successfully met, `True` is returned. If at least one of them fails, `False`
+    is returned.
 
     The stress is converged if the absolute difference of its corresponding pressure and the reference pressure is
     within the given threshold.
 
+    For calculations with fixed coordinates using the ``fixed_coords`` setting, the fixed coordinates can be provided.
+    The forces that correspond to a fixed coordinate are ignored when comparing with the threshold.
+
     :param trajectory: the `TrajectoryData`
     :param index: the frame index of the trajectory data to check, default is `-1` meaning the last frame
     :param threshold_forces: the force threshold in Ry / bohr
     :param threshold_stress: the stress threshold in kbar
     :param reference_pressure: reference pressure in kbar
+    :param fixed_coords: list of fixed coordinates in the calculation for each site
     :return: `True` if all thresholds are valid, `False` otherwise
     :raises ValueError: if any of the arrays or indices don't exist
     """
     if threshold_forces is not None:
-        converged_forces = verify_convergence_forces(trajectory, index, threshold_forces)
+        converged_forces = verify_convergence_forces(trajectory, index, threshold_forces, fixed_coords)
     else:
         converged_forces = True
 
@@ -37,12 +48,21 @@ def verify_convergence_trajectory(
     return converged_forces and converged_stress
 
 
-def verify_convergence_forces(trajectory, index=-1, threshold=None):
+def verify_convergence_forces(
+    trajectory: TrajectoryData,
+    index: int = -1,
+    threshold: Optional[float] = None,
+    fixed_coords: Optional[List[List[bool]]] = None
+) -> bool:
     """Verify that the `forces` of the given `TrajectoryData` are converged with respect to given threshold.
 
+    For calculations with fixed coordinates using the ``fixed_coords`` setting, the fixed coordinates can be provided.
+    The forces that correspond to a fixed coordinate are ignored when comparing with the threshold.
+
     :param trajectory: the `TrajectoryData`
     :param index: the frame index of the trajectory data to check, default is `-1` meaning the last frame
     :param threshold: the force threshold in Ry / bohr
+    :param fixed_coords: list of fixed coordinates in the calculation for each site
     :return: `True` if threshold is valid, `False` otherwise
     :raises ValueError: if the `forces` array or given index does not exist
     """
@@ -54,14 +74,26 @@ def verify_convergence_forces(trajectory, index=-1, threshold=None):
     threshold *= CONSTANTS.ry_to_ev / CONSTANTS.bohr_to_ang  # Convert to eV / Å
 
     try:
-        forces = trajectory.get_array('forces')[index]
+        abs_forces = numpy.abs(trajectory.get_array('forces')[index])
     except (KeyError, IndexError) as exception:
         raise ValueError('the `forces` array does not exist or the given index exceeds the length.') from exception
 
-    return numpy.max(abs(forces)) < threshold
+    if fixed_coords is not None:
+        fixed_coords = numpy.array(fixed_coords)
+        # Set the forces corresponding to fixed coordinates to zero, as they should not be checked versus threshold
+        # Since `fixed_coords` is a list of lists of booleans, where `True` indicates that a coordinate is fixed, we can
+        # invert this array and multiply it with the forces to set all fixed coordinates to zero.
+        abs_forces *= numpy.invert(fixed_coords)
+
+    return numpy.all(abs_forces < threshold)
 
 
-def verify_convergence_stress(trajectory, index=-1, threshold=None, reference_pressure=0.):
+def verify_convergence_stress(
+    trajectory: TrajectoryData,
+    index: int = -1,
+    threshold: Optional[float] = None,
+    reference_pressure: float = 0,
+) -> bool:
     """Verify that the `stress` of the given `TrajectoryData` are converged with respect to given threshold.
 
     The stress is converged if the absolute difference of its corresponding pressure and the reference pressure is

tests/calculations/test_cp/test_cp_autopilot_False_.in

@@ -21,8 +21,8 @@
 ATOMIC_SPECIES
 Si     28.0855 Si.upf
 ATOMIC_POSITIONS angstrom
-Si           0.0000000000       0.0000000000       0.0000000000 
-Si           1.3575000000       1.3575000000       1.3575000000 
+Si           0.0000000000       0.0000000000       0.0000000000
+Si           1.3575000000       1.3575000000       1.3575000000
 CELL_PARAMETERS angstrom
       2.7150000000       2.7150000000       0.0000000000
       2.7150000000       0.0000000000       2.7150000000

tests/calculations/test_cp/test_cp_autopilot_True_.in

@@ -21,8 +21,8 @@
 ATOMIC_SPECIES
 Si     28.0855 Si.upf
 ATOMIC_POSITIONS angstrom
-Si           0.0000000000       0.0000000000       0.0000000000 
-Si           1.3575000000       1.3575000000       1.3575000000 
+Si           0.0000000000       0.0000000000       0.0000000000
+Si           1.3575000000       1.3575000000       1.3575000000
 CELL_PARAMETERS angstrom
       2.7150000000       2.7150000000       0.0000000000
       2.7150000000       0.0000000000       2.7150000000

tests/calculations/test_pw.py

@@ -342,3 +342,35 @@ def test_pw_validate_inputs_restart_nscf(
     inputs['parameters'] = orm.Dict(parameters)
     with pytest.warns(Warning, match='`restart_mode` should be set to `from_scratch` for a `.*`.'):
         generate_calc_job(fixture_sandbox, entry_point_name, inputs)
+
+
+def test_fixed_coords(fixture_sandbox, generate_calc_job, generate_inputs_pw, file_regression):
+    """Test a ``PwCalculation`` where the ``fixed_coords`` setting was provided."""
+    entry_point_name = 'quantumespresso.pw'
+
+    inputs = generate_inputs_pw()
+    inputs['settings'] = orm.Dict(dict={'FIXED_COORDS': [[True, True, False], [False, True, False]]})
+    generate_calc_job(fixture_sandbox, entry_point_name, inputs)
+
+    with fixture_sandbox.open('aiida.in') as handle:
+        input_written = handle.read()
+
+    file_regression.check(input_written, encoding='utf-8', extension='.in')
+
+
+@pytest.mark.parametrize(['fixed_coords', 'error_message'], [
+    ([[True, True], [False, True]], 'The `fixed_coords` setting must be a list of lists with length 3.'),
+    ([[True, True, 1], [False, True, False]
+      ], 'All elements in the `fixed_coords` setting lists must be either `True` or `False`.'),
+    ([
+        [True, True, False],
+    ], 'Input structure has 2 sites, but fixed_coords has length 1'),
+])
+def test_fixed_coords_validation(fixture_sandbox, generate_calc_job, generate_inputs_pw, fixed_coords, error_message):
+    """Test the validation of the ``fixed_coords`` setting for the ``PwCalculation``."""
+    entry_point_name = 'quantumespresso.pw'
+
+    inputs = generate_inputs_pw()
+    inputs['settings'] = orm.Dict(dict={'FIXED_COORDS': fixed_coords})
+    with pytest.raises(ValueError, match=error_message):
+        generate_calc_job(fixture_sandbox, entry_point_name, inputs)

tests/calculations/test_pw/test_fixed_coords.in

@@ -0,0 +1,28 @@
+&CONTROL
+  calculation = 'scf'
+  outdir = './out/'
+  prefix = 'aiida'
+  pseudo_dir = './pseudo/'
+  verbosity = 'high'
+/
+&SYSTEM
+  ecutrho =   2.4000000000d+02
+  ecutwfc =   3.0000000000d+01
+  ibrav = 0
+  nat = 2
+  ntyp = 1
+/
+&ELECTRONS
+  electron_maxstep = 60
+/
+ATOMIC_SPECIES
+Si     28.0855 Si.upf
+ATOMIC_POSITIONS angstrom
+Si           0.0000000000       0.0000000000       0.0000000000 0 0 1
+Si           1.3575000000       1.3575000000       1.3575000000 1 0 1
+K_POINTS automatic
+2 2 2 0 0 0
+CELL_PARAMETERS angstrom
+      2.7150000000       2.7150000000       0.0000000000
+      2.7150000000       0.0000000000       2.7150000000
+      0.0000000000       2.7150000000       2.7150000000

tests/calculations/test_pw/test_pw_default.in

@@ -18,8 +18,8 @@
 ATOMIC_SPECIES
 Si     28.0855 Si.upf
 ATOMIC_POSITIONS angstrom
-Si           0.0000000000       0.0000000000       0.0000000000 
-Si           1.3575000000       1.3575000000       1.3575000000 
+Si           0.0000000000       0.0000000000       0.0000000000
+Si           1.3575000000       1.3575000000       1.3575000000
 K_POINTS automatic
 2 2 2 0 0 0
 CELL_PARAMETERS angstrom

tests/calculations/test_pw/test_pw_ibrav.in

@@ -18,7 +18,7 @@
 ATOMIC_SPECIES
 Si     28.0855 Si.upf
 ATOMIC_POSITIONS angstrom
-Si           0.0000000000       0.0000000000       0.0000000000 
-Si           1.3575000000       1.3575000000       1.3575000000 
+Si           0.0000000000       0.0000000000       0.0000000000
+Si           1.3575000000       1.3575000000       1.3575000000
 K_POINTS automatic
 2 2 2 0 0 0