How to migrate your plugin to work with the `atomistic` package

In the following we present guidelines on how to migrate your existing plugin to use the new AiiDA data types delivered in the aiida-atomistic code. This mainly has to do with the properties attribute of new atomistic StructureData nodes.

Main changes from the orm to the atomistic StructureData¶

There are three main differences which you should take into account when migrating:

• properties attribute (a subclass of the pydantic BaseModel)
• immutability of StructureData
• site-based approach to define the properties

The properties attribute in the new atomistic StructureData allows you to store additional information associated with each site/atom in the structure. To migrate your plugin, you need to update your code to access and manipulate the properties attribute instead of the previous methods. This is needed for all the properties, even the old ones, now under the properties attribute.

• pbc: structure.properties.pbc
• cell: structure.properties.cell
• sites: structure.properties.sites

... and so on. A full list of properties can be accessed via the classmethod get_supported_properties.

Additionally, the StructureData nodes in the atomistic module are now immutable even before being stored in the database. This means that you cannot modify the structure once it is created. If you need to make changes, you should create a new StructureDataMutable which allows you to modify the properties. This can be done via the get_value method. Once finished, you can initialize a new StructureData object using StructureData.from_mutable(mutable_object). For further details, please have a look at the dedicated section.

Site-based approach: Kind class is dropped¶

All the properties are defined in site-base approach, i.e. we do no more support the kind-based mapping. This is less compact description but allows a code-agnostic definition of properties. The mapping to kinds can always be done afterwards (i.e. in the aiida-quantumespresso plugin). Now the kinds are just a property of the structure, as well as all the others. Kinds can be generated on demand using the get_kinds or other methods (see the dedicated section).

Backward compatiblity¶

For some backward compatibility support, refer to the documentation. The support for the old orm.StructureData is meant to be dropped soon.

How a plugin should behave: parsing the structure properties for input file generation¶

When using the StructureData in your plugin, the idea is that each defined properties should be used in the calculation. The rationale is that, in this way, we have no ambiguity about if a property was used or not in the job. For example, if we use a structure with defined magmoms, the related DFT simulation submitted should be a magnetic one. It sounds confusing if a non-magnetic calculation was done using a magnetic structure.

To check on the supported properties, the plugin should define somewhere a list of supported ones, and control against the list of defined properties for a given StructureData. The properties to be checked are the one effectively extend the structure: “magmoms”, “charges”, “hubbard” and so on.

To easily check the plugin compatibility with the defined properties, we provide the method check_plugin_support:

This will check both the defined properties of the structure and the supported ones in the given plugin. If some property stored in the structure is not supported in the plugin, the plugin_check variable will be a non-empty set containing all the unsupported properties, and then some handler should be triggered if its length is larger than zero. See the next section for more details.

Dealing with unsupported properties¶

What if your plugin does not support (for example) magnetic calculations, but the structure contains magmoms in its properties? We recommend two solutions:

• except with some error (for example NotImplementedError), maybe even returning an AiiDA exit_code (depending on when you perform the check);
• raise a Warning and using a calcfunction to generate a new StructureData without the given unsupported set of properties;

In the first case, the plugin will stop and except. Let’s take the example of the aiida-quantumespresso PwCalculation. We provide a check in the super class BasePwCpInputGenerator. We suggest to put the properties validation in a validate_inputs classmethod. We define the list of _supported_properties as class attribute of the BasePwCpInputGenerator, and we include the following code snippet in the validate_inputs method (which is invoked in the corresponding method of the subclass):

class BasePwCpInputGenerator(CalcJob):
    """`CalcJob` implementation for the pw.x code of Quantum ESPRESSO."""
    ...
    _supported_properties = ["magmoms","hubbard"]
    ...
    @classmethod
    def validate_inputs(cls, value, port_namespace):
        """Validate the entire inputs namespace.
        
        Check the StructureData to contains only supported properties. The supported properties are the ones that are 
        defined in the _supported_properties class attribute of the super class BasePwCpInputGenerator.
        """
    ...
        """Check if the structure if the atomistic `StructureData` and 
        if contains unsupported properties
        """
        if isinstance(value['structure'], LegacyStructureData):
            raise exceptions.InputValidationError('LegacyStructureData (orm.StructureData) is no more supported, \
                use StructureData instead. You can convert a LegacyStructureData to a \
                StructureData using the `to_atomistic` method of the legacy.')
        else:
            plugin_check = value['structure'].check_plugin_support(cls._supported_properties)
            if len(plugin_check)>0:
                raise NotImplementedError(f'The input structure contains one or more unsupported properties \
                    for this process: {plugin_check}') 
    ...

In this way, each subclass of BasePwCpInputGenerator will perform the check on the same set of supported properties. You can define for each subclass the StructureData check, in case they support different properties.

In the second case we just raise a Warning and produce a new StructureData without the unsupported properties. This approach is better suited for a WorkChain, in which we may want to perform again the check and stripe on the fly the properties we don’t want. Using as example the PwBaseWorkChain:

class PwBaseWorkChain(ProtocolMixin, BaseRestartWorkChain):
    ...
    def validate_structure(self,):
        from aiida_atomistic.data.structure.utils import generate_striped_structure
        plugin_check = self.inputs.pw.structure.check_plugin_support(PwCalculation._supported_properties)
        if len(plugin_check)>0:
            # Generate a new StructureData without the unsupported properties.
            self.inputs.pw.structure = generate_striped_structure(self.inputs.pw.structure, orm.List(list(plugin_check)))
    ...

The generate_striped_structure is a calcfunction, so provenance is still preserved here.

Custom properties should be taken into account as well (see the dedicated section).

Custom properties¶

These properties are the ones not officially supported in the atomistic package. They should be defined only if the given calcjob/plugin supports them, meaning that they will be used in the calculation.

If instead one just wants to attach more information to a StructureData node, but not use it in any of the simulations, it is possible to set the extras:

structure.base.extras.set("property_1", [1,2,3,4])