Introducing state

[CedricTravelletti]

This introduces state updating in the calculator.

See this PR in GeometryOptimization.jl for a detailed discussion of why this is needed.

[tjjarvinen]

There is a need for this kind of feature.

• Do not create new types to add features
• Update algorithm can be problematic - how is it intended to be used? This need more explaining on how to do this
• It would be good to add some tests. Like add it to some of the already existing test functions. (These need a general update, so we could do this on another PR too).

[tjjarvinen]

Do not add AbstractCalculator type. It is not needed and it's presence will cause a lot of problems.

[tjjarvinen]

DummyState is not needed. We just use Missing to do the same

[tjjarvinen]

I would use

calculator_internal_state(::Any) = missing

update_calculator_state(
    update_algorithm::Any,
    initial_internal_state::Missing,
    initial_state::Any,
    new_state::Any
) = missing    # return new internal state

I would also add documentation on why to do this.

Also, update_algorithm is potentially very problematic. This stems from where are you going to define it. If it is defined for individual packages (like DFTK) then it should not be in the interface, or it will cause problems.

You will potentially call the command with different calculator types (e.g. when comparing different methods) and this causes errors, if algorithm in not defined for both of them. You could ignore algorithm all together to solve this. But this then goes on to the question, why it is there in the first place?

[CedricTravelletti]

Thanks a lot for your comment @tjjarvinen . This was helpful and made us come with a new point of view.

Let me briefly sketch why this whole state business was needed in the first place:

• There are circumstances where a calculator is part of a whole computation pipeline (think geometry optimisation).
  By that I mean that the calculator is used in a series of calculations on systems that share common properties.
• In such a situation, it makes sense to store intermediate results of previous computations and use them to speed up future computations by warming-up the calculator in a clever way.

This is where the idea of having a calculator state that can be extracted and provided as a starting point for computations originated in the first place.

[CedricTravelletti]

Now on to our new point of view and proposition, which tries to include your comments.

What we need:

1. Extract intermediate computational results from the calculator, which users can then externally process to provide sensible starting points for other computations down the line.
2. Warm up calculator with a sensible state of starting parameters (if we have them) that would speed up the current calculation.

Implementation Suggestion

1. This can be solved using your suggestion: calculator_internal_state(::Any) = missing
   Then individual calculators can define methods for calculator_internal_state and choose what they want to expose.
   For DFT calculators this would be density, orbitals, ...; for empirical potentials calculators this could be a neighbors list, or other stuff.
2. Warming up the calculator can by done via the kwargs.
   We specify in the interace that there is a kwarg warm_up=nothing. This is sensible, since for a lot of different types of calculators, there is some warming up that can be done. It is thus meaningful to include this default kwarg in the interface.

What a typical pipeline (geometry optimisation) would then look like

function warm_up_DFTK!(new_system, calculator::DFTKCalculator, compute_cache)
    compute_cache.push!(calculator_internal_state(calculator))
    interpolate_starting_point(new_system, compute_cache)
end

end


compute_cache = ... # can be a REDIS cache, or something on disk, or whatever.

for k in ...
    AtomsCalculators.potential_energy(system, calculator;
                                      warm_up = warm_up_DFTK!(system, calculator, compute_cache))
end

And DFTKCalculator would be responsible to choose what to do with the warm_up:

AtomsCalculators.@generate_interface function AtomsCalculators.potential_energy(
        system::AbstractSystem, calculator::DFTKCalculator; warm_up=nothing,
        kwargs...)
    if warm_up != nothing
        # prepare some stuff.
    end
end

[CedricTravelletti]

eager to know what you think.

[tjjarvinen]

Thank you for the clarification!

My own idea was something like this

int_states = [ ]

for sys in some_systems
    # do some work with my_calculator
    if something_interesting
        push!(int_states,  (get_internal_state(my_calculator), sys)  ) 
    end 
end

for (int_state, sys) in int_states
   set_internal_state!(my_calculator, sys, int_state)
   # do something else like
   AtomsCalculators.potential_energy(sys, my_calculator) 
end

But, I also like your idea of using a keyword. Meaning that we would have keywords as part of the interface definition.

However, I can see a potential issue coming from this. Let's say our calculator is a combination of several calculators, e.g. a QM/MM (or QM/ML) calculator using DFTK for QM part. If we use keyword in the interface this causes some issues. Because both QM and MM calculators get the same keywords and, if both have warm_up implemented. Then MM calculator will fail, if you use QM calculator specific warm_up.

With set_internal_state! command, we use dispatch to get over this problem, if we define

function set_internal_state!(calc, sys, int_states::AbstractArray)
   foreach( int_states ) do state
       set_internal_state!(cals, sys, state)
   end
end

We can then use an array like [dftk_warm_up, mm_warm_up] as an input for set_internal_state! and everything would work.

Then there is the naming issue as well. We are using now calculator_internal_state and warm_up. Are those good ones? I wrote get_internal_state and set_internal_state! in this post. But I am not sure about them either. Because they could be linked to AtomsBase system structure or calculator structure. So maybe get_calculator_internal_state and set_calculator_internal_state would be better. But that is also a bit clunky?

We need to add also

• Documentation entry that lists all keywords that are part of the interface. This could be a new page in online documentation.
• Testing functionality. This would be added to test_potential_energy, test_forces and test_virial functions.
• Documentation entry explaining how to use this feature. This could also be a new page.

[cortner]

I didn't have time to read through the details, but can I just ask : why do we not simply adopt the Lux model for parameters and states? It is very lightweight, and very general. What is there that it cannot do?

[CedricTravelletti]

@cortner

I didn't have time to read through the details, but can I just ask : why do we not simply adopt the Lux model for parameters and states? It is very lightweight, and very general. What is there that it cannot do?

I had a look at Lux. I agree that it would be a good model for us. There is one caveat tough: adopting the Lux model would require calculators to return results and state upon each call. This would require changing the interface so that calls to (say) AtomsCalculators.potential_energy retrun both a value and a state.

One possible way around it would be to add a keyword return_state to all function calls, that, if provided by the user, will be mutated to store the computational state of the calculator at the end of the computation.

Action

It seems to me that in order to use the Lux model, we need to decide between the following:

1. Change the interface and have all functions return both a value and a state.
2. Add a return_state argument that can be mutated to save the state of the calculator at the end of the computation.

Personally, I think that at this point of development (using AtomsCalculators for geometry optimization), this question of state cannot be avoided anymore, so we should make a decision.

Happy to hear what you guys think.

[CedricTravelletti]

For reference: (from Lux documentation)

Design Principles

• Layers must be immutable – cannot store any parameter/state but rather store the information to construct them

• Layers are pure functions

• Layers return a Tuple containing the result and the updated state

• Given same inputs the outputs must be same

[cortner]

1. One doesn't have to adopt all aspects of Lux.
2. One can write invisible wrappers for calculators that take care of the boilerplate and can be adapted to different scenarios.

This is all being discussed right now on Zulip. Please follow the discussion there and try to make the calls with CEMIX

[rkurchin]

Trying to catch up on things I've fallen behind on here and haven't read this in tremendous detail, but is it correct that this is superseded by #12? If so, unless we think there's a chance we'll revert back to this approach, should we close this to keep things clear/tidy?

[cortner]

you are probably right, @CedricTravelletti can confirm?

[CedricTravelletti]

@rkurchin and @cortner You are both right, this is superseeded by #12 . Closing this one now.