Node-based pyiron-base

[JNmpi]

Summary

Detailed Description

Further Information, Files, and Links

[JNmpi]

Fundamental concepts

Thanks @liamhuber for all the enhancements. At the IPAM workshop at UCLA, I had the chance to discuss with Jan concepts for bringing together various ideas from pyiron, ironflow and other workflow modules. These ideas are aimed to develop a more modular and easier-to-maintain pyiron. The basic building block will be nodes. The framework must therefore provide all tools to create, manage, and upscale them easily. In the following a brief summary of the concept is given:

        @pyiron.node
        def multiply(x=1, y=2):
               return x * y 

The main idea is that a node can be represented by a function with well-defined input, output and execution body. Rather than having to define explicitly the input and output when using a class this structure appears more pythonic and is used e.g. by dask. A main advantage would be the low entry barrier - all a new user would have to do is to write a common python function and decorate it with the pyiron.node. This approach can be also easily extended to provide type/ontology information via typing:

        from pyiron import onto
        
        @pyiron.node(register(onto.atomistic), log_input=True)
        def multiply(x: onto.types.atomistic.BulkModule=1 , y: int=2): -> float:
            return x * y 

The idea would be to provide ontological types like you introduced in the latest versions of ironflow. The nodes can be individually used and connected similar like in dask without having to explicitly define a workflow. A decorated node would be delayed and executed only once the highest level node is executed (similar like the delayed mode in dask):

          c = multiply(2, 3)    # no execution, only delayed object created
          d = multiply(c, 4)    # no execution, only delayed object created
          d.run() # only now c and d will be evaluated

Pyiron-related concepts

While the node-based pyiron could be used without explicitly defining a workflow, a common simulation should probably have one:

from pyiron import Workflow
        
        wf = Workflow(name='test', onto=onto.atomistic)
        
        Al = wf.create.structure.bulk('Al')
        job = wf.create.job.Lammps(name='Lammps_job', structure='Al')

The new syntax would be very similar to the existing one, so users should very easily adapt to it. A couple of things are to note:

• Workflow provides arguments/functionality that is not in Project (i.e. definition of its core ontology/application)
• The entire input to Lammps can be provided by function arguments (e.g. structure='Al'). For convenience and consistency, we should still leave our existing input notation (i.e. job.input['para'] = para) but the preferred style should be the new functional one.
• The Workflow object provides on one hand autocompletion to easily access and browse through all available nodes, on the other hand this syntax appends the nodes to the workflow.

Serial execution

The following example is an extreme case, and may be a good test case for checking the concept and code:

        structure = wf.create.structure.bulk('Al')
        job_old = None
        for i_step in range(10):   # cannot be captured by workflow (write as macrocode)
              job = wf.create.job.Lammps(name=‘Lammps_job',
                                       structure=job.get_structure(-1))
              
              if job_old is not None:  # write decorated function to capture it in workflows
                   if np.abs(job.output.energy_tot - job_old.output.energy_tot) < eps:
                          break

Such a construct would fail in dask delay.

VASP example

Register nodes

In contrast to having a complex module like VASP in a single node we should have smaller and more flexible ones. This is sketched below together with ideas regarding node registration:

        @pyiron.node(register(onto.atomistic.code.vasp.exe))
        def VASP_exe(incar:FilePath, poscar:FilePath, potcar:FilePath, kpoints:FilePath):
            work = WorkingDirectory(path='.', 
                                    files=[incar, poscar, potcar, kpoints]
                                   )
            work.run('vasp.exe -f my_mode')
            return work

        @pyiron.node(register(onto.atomistic.code.vasp.parser_outcar))
        def VASP_parser_outcar(outcar:FilePath, select=[], exclude=[]):
            out_dict = my_parser(outcar, select)
            return out_dict  # or iodata object


        @pyiron.node(register(onto.atomistic.code.vasp.parser_incar))
        def VASP_parser_incar(incar:DataIO):
            incar_str = my_parser(DataIO)
            return incar_str


        @pyiron.node(register(onto.atomistic.code.vasp.parser_poscar))
        def VASP_parser_poscar(structure: onto.atomistic.structure):
            poscar_str = my_parser(structure)
            return poscar_str

Create VASP (macro node)

        @pyiron.node(register(onto.atomistic.code.vasp))
        def VASP(incar: DataIO, 
                 structure: onto.atomistic.structure,
                 calculator: onto.node.calculator
                ):
            
            onto_vasp = onto.atomistic.code.vasp
        
            vasp = onto_vasp.VASP_exe(incar=onto_vasp.VASP_parser_incar(incar=incar), 
                                      structure=onto_vasp.VASP_parser_structure(structure=structure)
                                     )
            out_dict = onto_vasp.VASP_parser_outcar(vasp.outcar, select=['energy_tot'], exclude=[])
            
            return out_dict

This is only a sketch of first ideas. Comments and suggestions are very welcome.

[JNmpi]

A few more ideas and pseudo code regarding node based pyiron. The examples demonstrate typical workflows we use.

Example workflows

Single

    wf = Workflow(name='test', 
                  onto=onto.atomistic,
                  sql=None, # True - use default, db_connector, ...
                  log_level = 'low', # 'debug', ...
                  store_working_directory = True,
                 )

Note that we could add many options in the workflow creation. Examples are whether to store the nodes in a database, whether and which input and output data to put into hdf5, etc.

    job = wf.create.code.VASP(name='Lammps_job',
                              structure=wf.create.structure.bulk(‘Al’),
                              calc=wf.create.calculator.MD(T=300)
                            )  

Note that the new nodes allow to provide all input via function parameters.

Parallel

Run over large number of structure stored in a structure container:

    jobs = wf.create.code.VASP(name='Lammps_job',
                               structure=wf.create.structure_container(‘Al_data’),
                               calc=wf.create.calculator.MD(T=300)
                               )

or run over a list of temperatures:

    structures = wf.create.structure_container(‘Al_data’)
    calculators = wf.create.calculator.MD(T=np.arange(100, 1000, 100))
    STRUC, CALC = wf.np.meshgrid(strctures, calculators)
    
    jobs = wf.create.code.VASP(name='Lammps_job',
                               structure=STRUC,
                               calc=CALC
                               )

Serial

    @pyiron.node
    def iter_energy(energy_convergence):
        job = wf.create.code.VASP(name='Lammps_job',
                                  structure=wf.create.structure.bulk(‘Al’),
                                  calc=wf.create.calculator.MD(T=300)
                                 )    
        
        for i in range(10):        
            job_new = wf.create.code.VASP(name='Lammps_job',
                                      structure=job.get_structure(-1))   
            if np.abs(job.output.energy_tot - job_new.output.energy_tot) < energy_convergence:
                break
                
            job = job_new    
            
        return job.get_structure(-1)   

and apply nodes:

    wf.iter_energy(energy_convergence=1e-4)
    wf.run()

[liamhuber]

Hi Jörg,

Super, yep, I'm on board with this. I'll be on vacation through Wednesday so I won't be able to look at this in depth, but a bunch of the earlier stuff you suggest is already implemented over in contrib, minus the syntactic sugar of doing it with a decorator (which should be easy).

One thing that is notably missing in your examples is labels for outputs. IMO these are absolutely critical for constructing complex graphs where inputs and outputs can be connected with any complexity -- the only way to get around this is to restrict these functions to only return a single output value, which I think is too harsh.

Concretely, look at your example setting a kwarg "structure=wf.create.structure.bulk('Al')" -- in this case "bulk" is a node, and in principal may return multiple values (although in this case it only returns one). So I really feel extremely strongly that we need something like "structure = wf.create.structure.bulk('Al').structure", and to require specifying labels for the output, eg in the decorator like "@pyiron.node('structure')" or "@pyiron.node(('energy_pot', 'forces')".

If you and Jan haven't yet, please run the contrib workflow notebook(s?) to see what parts of your plans the existing infrastructure already already covers. All the starting stuff I think is just a matter of adding decorator syntactic sugar to existing functionality. The later stuff with specifying parallelization, database connection, etc all still needs to be done.

On a first read through, the absence of output naming is really only thing that worries me here, the rest of it looks brilliant so far.

Very practically, I want to first prioritize some performance enhancements for ironflow (decoupling port status model logic from the draw call should be sufficient), but once that's done I am excited to go over to contrib and start integrating the ideas here into the graph infrastructure 👍👍👍

[liamhuber]

Actually I have a second concern: the use of a global variable ('wf') in the final example. I think this is a closely connected concern to my worries about Marvin's contrib work where he allows entire nodes to be passed as input to other nodes. I completely agree that we need this type of macro functionality, I just think we'll need to be a little out careful about the implementation.

[JNmpi]

Hi Liam,

Thanks for your super-quick and very positive reply. I am glad that we both see the advantages and potential of these formulations and that so much is already realized and implemented in ironflow. Once you are back from your vacation it would be good to have Zoom meeting to discuss the next steps in more detail.

With respect to your questions/topics some first thoughts below:

• ... only return a single output value: This could be resolved by returning several output values and typing them. A simple example is sketched below:

      @pyiron.node()
      def myfunc(A: a, B: b) -> (O1: o1, O2: o2):
            o1, o2 = a, b
            return o1, o2
  

The types O1, O2 etc. can be ontologically enriched using concepts of the typing module link. In particular the Annotation function looks promising to provide metadata.

• the use of a global variable ('wf') in the final example. This can be likely prevented by the following slight extension:

        @pyiron.node
        def iter_energy(wf, energy_convergence):
            job = wf.create.code.VASP(name='Lammps_job',    
            ...  
  

  The application of the node would be unchanged, i.e., wf behaves like self in a class:

        wf.create.node.iter_energy(energy_convergence=1e-4)  # wf is not in the function call but provided by the attached class object 
        wf.run()
  

One final question. In your comment you mentioned 'contrib workflow notebook'. Are these the notebooks in the ironflow repository (or the one in the pyiron_ontology)?

[liamhuber]

Hi Joerg,

A live chat sounds great!

Re annotations, I think that works nicely for adding onto typing in addition to data typing, although annotations currently do not support kwargs (like "o type=onto.foo"), so we would need to force a fixed ordering in annotations. For this reason I lean towards using a dict, eg as a kwarg in the decorator with keys matching the variable names, to provide this data. The big problem I see using annotations to label output data is that I want output labels to be absolutely mandatory, and forcing people to learn about annotations and use them seems harder to understand than adding a positional arg to the node decorator. Also we could probably enforce it as a requirement, but we'd need to add extra rail guards, where adding it as an arg to the decorator just right away makes it clear.

Re the demo notebook, this is actually not at all integrated with ironflow yet, it's over in contrib in notebooks/workflow_example.ipynb

[liamhuber]

Missed the global var discussion because of indentation.

It may indeed be possible to provide it by the scope of the class, but then you'd still need something like 'self.wf', and self will (at a minimum) look very strange in the context of a function definition (even if the decorator means we return a class instance). I bet we can find a solution, it will just take some thinking to get it both functional and intuitive.

[JNmpi]

Let's try to schedule a meeting on Thursday or Friday. I am presently at the DPG meeting in Dresden and will be back in Düsseldorf on Friday.

It would be definitely good to explore possible option to provide the labels for the output variables. Decorators may be a good option. We should only make sure that our solution is as close to standard python, so that the barrier for users is as low as possible.

Thanks also for the latest development on ironflow. The ontologic features are really great and it is exciting to play with it. The only issue is the sluggish behavior, which makes it often hard to know whether a click did not work or something is still happening. I am therefore looking very much forward to the next development on speeding things up. Then one can fully enjoy the really cool features and the great concepts that you have already implemented. Really great work!

[JNmpi]

Hi Liam, I had now a look at your workflow notebook in pyiron_contrib. Really very nice! I see also the strong links to my thoughts. An important task of the decorator would be to make the following statement more intuitive and python-like:

      def plus_minus_one(x: int | float = 1) -> tuple[int | float, int | float]:
          return x+1, x-1
      
      node = Node(plus_minus_one, ("p1", "m1"))

With the new formulation this could read like

      @pyiron.node(repository)   # repository where the node would be registered (could be local or global)
      def plus_minus_one(x: int | float = 1) -> tuple[int | float, int | float]:
          return x+1, x-1
      
      plus_minus.inputs.x.type_hint  

The last line is only an example to show that all constructions in your notebook should work, i.e., the decorator converted the function into a node object.

[JNmpi]

One more thought regarding your notebook. For code applications, it may be helpful to offer a lazy mode, i.e., the following statement should just build the workflow but not run it:

    wf.structure = nodes.BulkStructure(repeat=3, cubic=True, element="Al")
    wf.engine = nodes.Lammps(structure=wf.structure.outputs.structure)
    wf.calc = nodes.CalcMD(job=wf.engine.outputs.job)
    wf.plot = nodes.Scatter(
        x=wf.calc.outputs.steps, 
        y=wf.calc.outputs.temperature
    )

To actually run it one would have to call the following line:

  wf.run()

Here one could also specify where to run it, i.e., the queue, the number of cores etc. It would be also nice to have an option to convert the code into a graph and vice versa:

    wf.visualize() 
    wf.to_graph()   # alternative

[jan-janssen]

While I am currently at the IPAM workshop, I would like to join the meeting to synchronize the discussions, so just keep me in the loop.

[JNmpi]

Hi Jan, great that you will join. We have not yet set up a meeting but a good choice may be Friday afternoon when I will be back at home.

[jan-janssen]

I have a couple of recent developments from the IPAM workshop, which might also be helpful for this discussion:

• structuretoolkit - move the structure analysis to a separate module so they can be used with ASE Atoms directly. The idea is that if people like this stand-alone module then they might also be more likely to give pyiron a try. The pull request to merge these changes back into pyiron_atomistics is available Use Structuretoolkit pyiron_atomistics#994
• pyiron_lammps - the pure python interface to LAMMPS now supports multiple ASE structures as well as multiple pandas DataFrames for interatomic potentials. Just like numpy if only a single ASE structure and a single potential are provided only a single set of elastic constants is calculated. If either multiple ASE structures or multiple potential dataframes are provided then multiple sets of elastic constants are calculated. Finally, if a list of ASE structures and a list of interatomic potential dataframes are provided with the same length - e.g. as they are generated from np.meshgrid then again a list of elastic constants is calculated. Implement new numpy like syntax lammpsparser#12
• pympipool - now subtasks can use multiple MPI ranks. For running LAMMPS simulation this gives the user the full flexibility, either a single LAMMPS calculation is using all MPI ranks, or each LAMMPS simulation uses one MPI rank each or any other split of MPI ranks to LAMMPS simulations is possible, as long as all LAMMPS simulation use the same number of MPI ranks and all MPI ranks are used. Enable MPI parallel sub tasks executorlib#28

While these developments focus more on the scalability of pyiron, I could see them being beneficial in simplifying the development of complex workflows and hopefully provide the ideal test bed for the developments discussed above.

[liamhuber]

I'm just on mobile so my responses are pretty limited in depth, sorry.

Re meeting: Friday sounds good. At present I can be free any time. We should keep @pmrv in the loop here too in case he wants to attend; there is both synergy and some conflict between the graph stuff and the tinybase stuff.

Re standard python/decorators/link: indeed, I think getting the existing graph stuff working with decorators should be super fast, then adding the fancier bits on top can be more iterative. I'm super excited about this direction.

Re lazy evaluation: there is some support for this! My existing node stuff has init flags for turning on/off the initial run and running automatically on update. Definitely not as smooth as your example, but the groundwork exists at least.

Re structuretoolkit pyiron_lammps, etc: I am super excited to get this, as the coupling between current pyiron jobs and nodes is a huge pain to manage! Personally I would be happy to only support Lammps forever, but we will need to think carefully about data storage and making sure we can still accommodate more expensive codes like vasp. But for now getting it all on-the-fly as facilitated by pyiron_lammps is super exciting