Skip to content

Improve the usability of callback functions in ground state optimization #208

@20akshay00

Description

@20akshay00

When playing around with different solvers and new algorithms to compute the ground state of a system, we typically want to track certain properties including but not limited to the energy and galerkin error over the iterations of the solver. Currently there is find_groundstate which takes in an algorithm, and the algorithm takes in an argument for a callback function where the state of the solver is exposed and can be altered. As I am working with FiniteMPS, I have only looked into DMRG and GradientGrassmann so far;

  • GradientGrassmann expects finalize!(x, f, g, numiter) -> x, f, g
  • DMRG expects finalize(iter, state, H, envs) -> state, envs

In order to keep track of some quantities during the optimization, something along these lines has to be done:

function custom_find_groundstate(args...)
  history = []
  function custom_finalizer(iter, state, H, envs) # specific to DMRG
      current_error = maximum(pos -> calc_galerkin(state, pos, envs), 1:length(state))
      current_energy = expectation_value(state, H, envs)
      custom_observable = expectation_value(state, ...)
      push!(history, (; current_error, current_energy, custom_observable, t=Base.time()))
      return state, envs
  end
  
  return find_groundstate(args...; finalize=custom_finalizer), history
end

Alternatively, if some kind of local state is to be preserved, something like this can also be done:

mutable struct IterationData
      energy
      error
      observable
      IterationData() = new([], [], [])
end

function (data::IterationData)(iter, state, H, envs) # specific to DMRG
   push!( data.error, maximum(pos -> calc_galerkin(state, pos, envs), 1:length(state)))
    push!(data.energy, expectation_value(state, H, envs))
    push!(data.observable, expectation_value(state, ...))
    return state, envs
end
  
custom_data = IterationData()
psi, envs, delta = find_groundstate(args...; finalize=custom_data), history

While this is functional, it is a bit clunky and also has to be modified per algorithm as they expect functions with different signatures. There may also be some minor inconsistencies as in #207. Following are possible improvements:

  • Ideally there should be a uniform interface for the signature of finalize, although I am unsure if there is a clean way to do so for these two methods (and any other added later on) given that they operate with different quantities. For example, always having access to at least H (or E = <H>, but consistently either one of them) and psi within finalize makes sense to me.
  • A way to quit the optimization loop in find_groundstate using a custom convergence criteria instead of a hard-coded one which is also solver dependant :/.

My use case for this is really just to quantitatively judge the performance of Gradient methods vs. DMRG for my system in order to decide which one to use hereafter. I think this would be a fairly common use case for anyone working with atypical systems when DMRG is no longer sufficient, so having a convenient way to do this would be quite helpful.

Metadata

Metadata

Assignees

No one assigned

    Labels

    No labels
    No labels

    Type

    No type

    Projects

    No projects

    Milestone

    No milestone

    Relationships

    None yet

    Development

    No branches or pull requests

    Issue actions