Using the "trajectory" kwarg of OTF-MD engines leads to exceptions later on

[smheidrich]

According to the documentation, trajectory is a valid (even required, albeit "usually set to None") key for the md_params dictionary that's a parameter of the otf_md() function. So I would expect to be able to let ASE's own MD machinery record a trajectory as usual by setting it to anything other than None, e.g. a filename.

Here is a minimal reproducible example script in which I try to do just this (a bit long because FLARE requires a lot of boilerplate as far as I can tell):

flare_ase_trajectory_bug.py

#!/usr/bin/env python3
import ase
from ase.calculators.emt import EMT
import numpy as np
from ase import units
from ase.md.velocitydistribution import (MaxwellBoltzmannDistribution,
                                         Stationary, ZeroRotation)
from flare.ase.otf_md import otf_md
from flare.ase.logger import OTFLogger

# Dummy system
a = ase.Atoms("H3")
a.positions[1:] += [[2,0,0], [2,2,0]]
a.set_cell(np.diag([10,10,10]))
a.set_pbc(True)

# Copied from https://flare.readthedocs.io/en/latest/tutorials/ase.html )
from flare.gp import GaussianProcess
from flare.ase.calculator import FLARE_Calculator
import flare.mc_simple as mc_simple

kernel = mc_simple.two_plus_three_body_mc
kernel_grad = mc_simple.two_plus_three_body_mc_grad
energy_force_kernel = mc_simple.two_plus_three_mc_force_en
energy_kernel = mc_simple.two_plus_three_mc_en

hyps = np.array([0.1, 1., 0.001, 1, 0.03])
two_cut = 4.0
three_cut = 4.0
cutoffs = np.array([two_cut, three_cut])
hyp_labels = ['sig2', 'ls2', 'sig3', 'ls3', 'noise']
opt_algorithm = 'LBFGS'

gp_model = GaussianProcess(kernel, kernel_grad, hyps, cutoffs,
                           energy_kernel=energy_kernel,
                           energy_force_kernel=energy_force_kernel,
                           hyp_labels=hyp_labels,
                           opt_algorithm=opt_algorithm, par=True)

flare_calc = FLARE_Calculator(gp_model, par=True, use_mapping=False)

# For minimal example, just use ASE's EMT calculator as dummy instead of DFT
dft_calc=EMT()

# MD and FLARE OTF stuff
a.set_calculator(flare_calc)

# set up OTF MD engine
md_params = {'timestep': 1 * units.fs,
             'trajectory': 'otf_md.traj', # XXX everything fine if this is None
             'dt': None}

otf_params = {'dft_calc': dft_calc,
              'init_atoms': [0],
              'std_tolerance_factor': 0.5,
              'max_atoms_added' : len(a.positions),
              'freeze_hyps': 0,
              'restart_from': None,
              'use_mapping': a.calc.use_mapping}

# intialize velocity
temperature = 600
MaxwellBoltzmannDistribution(a, temperature * units.kB)
Stationary(a) # zero linear momentum
ZeroRotation(a) # zero angular momentum

test_otf = otf_md('VelocityVerlet', a, md_params, otf_params)

# set up logger
test_otf.attach(OTFLogger(test_otf, a,
    logfile='otf_run.log', mode="w", data_in_logfile=True),
    interval=1)

# rescale temperature to 1200K at step 11
rescale_temp = [1200]
rescale_steps = [11]

# run otf
number_of_steps = 10
test_otf.otf_run(number_of_steps, rescale_temp, rescale_steps)

When I execute it, a regular ASE-compatible trajectory does in fact get written to disk, but after some more progress, this exception is raised:

Traceback (most recent call last):
  File "./flare_ase_trajectory_bug.py", line 80, in <module>
    test_otf.otf_run(number_of_steps, rescale_temp, rescale_steps)
  File "/home/smheidrich/flare/flare/ase/otf.py", line 135, in otf_run
    self.observers[0][0].add_atom_info(atom, self.stds[atom])
AttributeError: 'function' object has no attribute 'add_atom_info'

Putting None for trajectory instead makes it run fine.

I guess the issue is that ASE's own trajectory logger gets attached as the first observer to the OTF object, but FLARE relies on having its own logger as the first observer.

Initially I thought this would be a quick fix, but making this part of the code tolerate ASE's trajectory logger only leads to further problems in other places... so I thought I'd ask here before I waste much more time on this - perhaps it's an easy fix for you guys.

Everything was tested with the current revision on master (7afbd4b).