On the dynamic change of workchain process class inputs

Hello everyone,
I have the following use-case. I’d like to write a WorkChain that performs a molecular dynamics (say, using aiida-lammps), and afterwards computes some quantity out of the trajectory. Since the part of the dynamics is always the same, i am wondering if there is a simple solution to change “dynamically” the workchain (in the expose_inputs) for the calculation of a specific quantity.

An example of what i would do is as follows:

class ExtractQuantityWorkChain(WorkChain):

    _quantity_calculator_process_class = CalculateSomethingWorkChain

    @classmethod
    def set_quantity_calculator_workchain(cls, entry_point_name: str = None, workchain: WorkChain = None):
        """Set the quantity predictor WorkChain to use."""
        if entry_point_name is not None:
           cls._quantity_calculator_process_class = WorkflowFactory(entry_point_name)
        elif workchain is not None:
            cls._quantity_calculator_process_class = workchain
        else:
            raise ValueError('At least one input is required, but none were provided.')

    @classmethod
    def define(cls, spec):
        """Define the process specification."""
        # yapf: disable
        super().define(spec)

        spec.expose_inputs(cls._quantity_calculator_process_class, namespace='quantity_workchain', exclude=('trajectory'))
        spec.expose_inputs(LammpsBaseWorkChain, namespace='lammps_workchain')

[...]

I am wondering though if this is a safe approach. I am worried that by having the set_quantity_calculator_workchain one also changes internally the class variable at runtime of other running workchains.

Any pointer is highly appreciated.

When would this set_quantity_calculator_workchain be called? That is going to be the main question. You can change the class attribute, but that change would only exist in that Python interpreter. It is not persisted outside of it. That means that if the interpreter shuts down and the class is reloaded, the cls._quantity_calculator_process_class is reset to the default, i.e, CalculateSomethingWorkChain. This would mean that it could pollute the class for other running workchains as long as it is in the same interpreter, but that would be the case for multiple workchains being run in the same daemon worker.

I think what you would really need here is templating, i.e., through a metaclass. For each different CalculateSomethingWorkChain type, you construct dynamically a separate ExtractQuantityWorkChain class that fixes the _quantity_calculator_process_class . I have never done this and not sure if there are any assumptions in AiiDA’s code that would prevent this approach from being used. It is also a bit laborious and perhaps unnecessarily complex.

You could always just make the output namespace dynamic so you can add any outputs. You don’t have to expose explicit ports. Of course you lose some validation but this may be acceptable still

Thanks a lot @sphuber for the quick reply!

This would mean that it could pollute the class for other running workchains as long as it is in the same interpreter, but that would be the case for multiple workchains being run in the same daemon worker.

Indeed, this was my fear. This may be likely, and thus I would discard this option.

You could always just make the output namespace dynamic so you can add any outputs. You don’t have to expose explicit ports. Of course you lose some validation but this may be acceptable still

Right, this is actually a much simpler solution and I definetely think it’s the best trade-off here! Thanks for the pointer!