How to save MPS data during the DMRG calculation

Hello,
I’m doing a DMRG calculation using Google Colab, but it sometimes fails due to the errors such as out of memory or a sudden disconnection. If the MPS data can be saved during the DMRG calculation, the calculation can be continued from that time. But I’m at a loss as to how to rewrite the code. Is there any good way to implement this?

Thank you.

If you don’t want to just pass in a small number of sweeps outside of DMRG and take care of it that way, I imagine you could use a custom observer to modify the measure! method which gives you access to the current MPS along with where in the sweep it is.

1 Like

In addition to Ryan’s answer (I agree a custom observer is the way to go, if you want to save an MPS during a calculation), you may find the following helpful:

1 Like

I do exactly that, here’s my code for reference (it does some other stuff such as saving specific attributes as well).

function saveITensorObject(obj, g::HDF5.Group, name::String, attrs::Dict)
  if haskey(g, name)
    delete_object(g, name)
  end

  write(g, name, obj)

  writeAttributes(g, attrs)

  return nothing
end


function saveITensorObject(obj, filePath::String, groupPath::String, name::String, attrs::Dict)
  f = h5open(filePath, "cw")
  g = create_or_get_group(f, groupPath)
  saveITensorObject(obj, g, name, attrs)
  close(f)

  return nothing
end


struct DMRGSaver <: AbstractObserver
  hdf5Path::String
  group::String
end


function ITensors.checkdone!(o::DMRGSaver; kwargs...)
  energy = kwargs[:energy]
  psi = kwargs[:psi]
  println("Energy per site = $(real(energy) / length(siteinds(psi)))")
  
  attrs = Dict([("energy", energy), ("currentSweep", kwargs[:sweep]), ("linkDim", maxlinkdim(psi))])
  saveITensorObject(psi, o.hdf5Path, o.group, "MPS", attrs)

  return false
end
2 Likes

If I can make a small suggestion, the

h5open(filePath, "cw") do f
  g = create_or_get_group(f, groupPath)
  saveITensorObject(obj, g, name, attrs)
end

method of opening the file is a bit safer in case of a crash (e.g. see here)

2 Likes