Skip to content

Using mocked VASP for workchain tests

Bonan Zhu edited this page Jan 15, 2022 · 2 revisions

Introduction

Constructing workchain tests is a pain, usually the actual calculations have to be run to test the correctness of workflow. This cannot be done on typicaly online CI tests due to licensing issues. Even if it is possible, such tests would be very slow to complete.

Instead, this plugin has implemented a mock code that simulates actual VASP code that runs the calculations. The mock-vasp command will try to match input files with existing output files living in a repository, based on the hashes generated from the contents of INCAR, KPOINTS, POSCAR files. For backward compatibility, it can also match existing results based on the name of the system provided in the POSCAR (disabled in mock-vasp-strict executable). However, this should not be used except for the very simply cases.

The central machinary of the mock code is the MockRegistry class, which manages a repository of existing calculations. To generate data for running a workchain test, one must actually run the workchain using a real VASP executable, and the then deposit the results into the registry.

Generating data for the MockRegistry

One way to temporary switch from mock code to real code is to point the REAL_VASP_PATH environmental variable to the actual VASP executable. This swaps the mock executable with a real one. In addition, the actual POTCAR data must be uploaded prior to the calculations. Note that the POTCAR is not used for matching outputs. This can be done by calling the uplaod_real_pseudopotentials function inside the test function.

def upload_real_pseudopotentials(path):
    """
    Upload real pseudopotentials for workchain test mock deposition


    This function should be called once before the REAL vasp calculation is launch to setup the
    correct POTCARs
    """
    global POTCAR_FAMILY_NAME  # pylint: disable=global-statement
    POTCAR_FAMILY_NAME = 'TEMP'
    potcar_data_cls = DataFactory('vasp.potcar')
    potcar_data_cls.upload_potcar_family(path, 'TEMP', 'TEMP-REALPOTCARS', stop_if_existing=False, dry_run=False)

There is also a function for depositing the associated CalcJobNode run by a given WorkChainNode:

def upload_real_workchain(node, name):
    """
    Upload the workchain to the repository to make it work with mocking

    This function should be called once after the REAL vasp calculation is run during the test
    """
    reg = MockRegistry()
    print(reg.base_path)
    reg.upload_aiida_work(node, name)

The process of generating mock data is best illustrated by an example. Let's consider this test function in test_vasp_wc.py:

def test_vasp_wc_ionic_magmom_carry(fresh_aiida_env, potentials, mock_vasp_strict):
    """Test with mocked vasp code for handling ionic convergence issues"""
    from aiida.orm import Code
    from aiida.plugins import WorkflowFactory
    from aiida.engine import run


    workchain = WorkflowFactory('vasp.vasp')


    mock_vasp_strict.store()
    create_authinfo(computer=mock_vasp_strict.computer, store=True)


    incar = dict(INCAR_IONIC_CONV)
    incar['ispin'] = 2
    incar['lorbit'] = 10
    incar['nupdown'] = 2
    inputs = setup_vasp_workchain(si_structure(), incar, 8)
    inputs.verbose = get_data_node('bool', True)


    # The test calculation contain NELM breaches during the relaxation - set to ignore it.
    inputs.handler_overrides = get_data_node('dict', dict={'ignore_nelm_breach_relax': True})
    inputs.settings = get_data_node('dict', dict={'parser_settings': {
        'add_structure': True,
        'add_site_magnetization': True,
    }})
    inputs.max_iterations = get_data_node('int', 2)


    _, node = run.get_node(workchain, **inputs)
    assert node.exit_status == 0


    called_nodes = list(node.called)
    called_nodes.sort(key=lambda x: x.ctime)
    # Check that the second node takes the magnetization of the first node
    assert called_nodes[1].inputs.parameters['magmom'] == [0.646]

To generate the mock data, two simple modification should be carried out:

def test_vasp_wc_ionic_magmom_carry(fresh_aiida_env, potentials, mock_vasp_strict):
    """Test with mocked vasp code for handling ionic convergence issues"""
    from aiida.orm import Code
    from aiida.plugins import WorkflowFactory
    from aiida.engine import run


    workchain = WorkflowFactory('vasp.vasp')
    upload_real_pseudopotentials("<path_to_POTCAR_library>")   # <----- Upload real POTCARs for the real VASP calculations

    mock_vasp_strict.store()
    create_authinfo(computer=mock_vasp_strict.computer, store=True)


    incar = dict(INCAR_IONIC_CONV)
    incar['ispin'] = 2
    incar['lorbit'] = 10
    incar['nupdown'] = 2
    inputs = setup_vasp_workchain(si_structure(), incar, 8)
    inputs.verbose = get_data_node('bool', True)


    # The test calculation contain NELM breaches during the relaxation - set to ignore it.
    inputs.handler_overrides = get_data_node('dict', dict={'ignore_nelm_breach_relax': True})
    inputs.settings = get_data_node('dict', dict={'parser_settings': {
        'add_structure': True,
        'add_site_magnetization': True,
    }})
    inputs.max_iterations = get_data_node('int', 2)


    _, node = run.get_node(workchain, **inputs)
    assert node.exit_status == 0


    called_nodes = list(node.called)
    called_nodes.sort(key=lambda x: x.ctime)
    # Check that the second node takes the magnetization of the first node
    assert called_nodes[1].inputs.parameters['magmom'] == [0.646]
    # <----- Save the results under a sub folder named `vasp-wc-magmon-carry`.
    # <----- The name is just for convenience, since it is input content that will be matched with
    upload_real_workchain(node, "vasp-wc-test-magmom-carry")   

After the edit, run the test with

export REAL_VASP_PATH="<path_to_real_vasp_executable>"
pytest test_vasp_wc.py::test_vasp_wc_ionic_magmom_carry
unset REAL_VASP_PATH

It is very important that pytest is instructed to run the target test ONLY, otherwise other mock tests will be run with actual VASP executable as well.

After this, unset the REAL_VASP_PATH environmental variable and remove the upload_real_* functional in the test. The test should be runnable nevertheless, as now the output files are extracted from the repository directly.

Finally, don't forget to add new additions to the registry to git with git add.