The `trajectory` class variable is not accessible when using `Sella`

[Andrew-S-Rosen]

Sella right now is not storing the trajectory class variable like the other ASE optimizers.

Here is an example, using the master branch of ASE:

from ase.build import bulk
from ase.calculators.emt import EMT
from ase.optimize import BFGS
from sella.optimize import Sella

atoms = bulk("Cu") * (2, 1, 1)
atoms.calc = EMT()
dyn = BFGS(atoms,trajectory="test.traj")
dyn.run()
assert dyn.trajectory is not None # <---- this passes (note: using the `master` branch of ASE!)

atoms = bulk("Cu") * (2, 1, 1)
atoms.calc = EMT()
dyn = Sella(atoms,trajectory="test.traj")
dyn.run()
assert dyn.trajectory is not None # <---- this fails

Running the above snippet will yield an assertion error for Sella but not for BFGS. For those developing packages around Sella, it would be very convenient to be able to access the same properties as the other optimizers. Tagging @samblau.

[ehermes]

In general, Sella does not guarantee field compatibility with built-in ASE Optimizer objects due to the additional complexity of the code. However, in this instance, it's rather straightforward to implement a fix, see #24. Let me know if there are any other fields that you would like to have exposed in a similar way to ASE's built-in Optimizers. Depending on the parameter, I may or may not be able to deliver a fix.

[Andrew-S-Rosen]

@ehermes: Thanks! Good to know. The trajectory is the most important one for me! Then we can close this issue 👍

[ehermes]

I'm going to try rolling this change into an unrelated bugfix that was reported to me on Slack, then make a 2.3.1 release.

[ehermes]

When I run your test, I get the following output:

      Step     Time          Energy         fmax
BFGS:    0 17:42:00       -0.011363        0.0000
Traceback (most recent call last):
  File "/home/eric/test/sella/test.py", line 10, in <module>
    assert dyn.trajectory is not None
           ^^^^^^^^^^^^^^
AttributeError: 'BFGS' object has no attribute 'trajectory'

That is, the BFGS object doesn't have a trajectory attribute... However, when I comment out the BFGS block, it runs without error:

No GPU/TPU found, falling back to CPU. (Set TF_CPP_MIN_LOG_LEVEL=0 and rerun for more info.)
     Step     Time          Energy         fmax         cmax       rtrust          rho
Sella   0 17:42:35       -0.011363       0.0000       0.0000       0.1000       1.0000

Are you sure the problem isn't with BFGS? I think you shouldn't be relying on the trajectory field of the optimizer, anyway. Remember that you can create an instance of the Trajectory object yourself, and pass that to the Optimizer, instead of relying on the Optimizer to create the object for you.

[Andrew-S-Rosen]

@ehermes: Are you using the master branch of ASE here? I think it's likely the discrepancy.

[ehermes]

Yes, I am using the master branch.

[ehermes]

In any case I suggest using a workflow along the lines of the following, which will avoid this issue entirely:

from ase.io.trajectory import Trajectory
from sella import Sella

trajectory = Trajectory()
dyn = Sella(atoms, trajectory=trajectory)

This way, you are in control of the creation of the Trajectory object, so you will always have a reference to it, and you don't need to prod into the unexposed inner fields of the Optimizer object (which is not guaranteed to be a stable API).

[Andrew-S-Rosen]

Yes, I am using the master branch.

Thanks for checking. That's a bit odd. Should have been patched upstream in ASE a little while ago.

I think you shouldn't be relying on the trajectory field of the optimizer, anyway. Remember that you can create an instance of the Trajectory object yourself, and pass that to the Optimizer, instead of relying on the Optimizer to create the object for you.

Yeah, that sounds like the smarter way to go in the end. I think that's compatible with what we have in mind here. Thanks for the suggestion! I'll try it out.

[Andrew-S-Rosen]

@ehermes:

I tried out your alternate suggestion but ran into issues there as well.

from ase.build import bulk
from ase.calculators.emt import EMT
from ase.optimize import BFGS
from sella.optimize import Sella
from ase.io.trajectory import Trajectory

atoms = bulk("Cu") * (2, 1, 1)
atoms.calc = EMT()
dyn = BFGS(atoms,trajectory=Trajectory("test.traj", "w"))
dyn.run() # <--- this runs

atoms = bulk("Cu") * (2, 1, 1)
atoms.calc = EMT()
dyn = Sella(atoms,trajectory=Trajectory("test.traj", "w"))
dyn.run() # <--- this does not

Traceback:

Cell In[3], line 4
      2 atoms.calc = EMT()
      3 dyn = Sella(atoms,trajectory=Trajectory("test.traj", "w"))
----> 4 dyn.run()

File ~/software/miniconda/envs/quacc/lib/python3.10/site-packages/ase/optimize/optimize.py:275, in Optimizer.run(self, fmax, steps)
    273 if steps:
    274     self.max_steps = steps
--> 275 return Dynamics.run(self)

File ~/software/miniconda/envs/quacc/lib/python3.10/site-packages/ase/optimize/optimize.py:158, in Dynamics.run(self)
    151 def run(self):
    152     """Run dynamics algorithm.
    153
    154     This method will return when the forces on all individual
    155     atoms are less than *fmax* or when the number of steps exceeds
    156     *steps*."""
--> 158     for converged in Dynamics.irun(self):
    159         pass
    160     return converged

File ~/software/miniconda/envs/quacc/lib/python3.10/site-packages/ase/optimize/optimize.py:130, in Dynamics.irun(self)
    127 yield False
    129 if self.nsteps == 0:
--> 130     self.log()
    131     self.call_observers()
    133 # run the algorithm until converged or max_steps reached

File ~/software/miniconda/envs/quacc/lib/python3.10/site-packages/sella/optimize/optimize.py:293, in Sella.log(self, forces)
    291 if self.logfile is None:
    292     return
--> 293 _, fmax, cmax = self.pes.converged(self.fmax)
    294 e = self.pes.get_f()
    295 T = strftime("%H:%M:%S", localtime())

File ~/software/miniconda/envs/quacc/lib/python3.10/site-packages/sella/peswrapper.py:298, in PES.converged(self, fmax, cmax)
    297 def converged(self, fmax, cmax=1e-5):
--> 298     fmax1 = np.linalg.norm(self.get_projected_forces(), axis=1).max()
    299     cmax1 = np.linalg.norm(self.get_res())
    300     conv = (fmax1 < fmax) and (cmax1 < cmax)

File ~/software/miniconda/envs/quacc/lib/python3.10/site-packages/sella/peswrapper.py:293, in PES.get_projected_forces(self)
    291 def get_projected_forces(self):
    292     """Returns Nx3 array of atomic forces orthogonal to constraints."""
--> 293     g = self.get_g()
    294     Ufree = self.get_Ufree()
    295     return -((Ufree @ Ufree.T) @ g).reshape((-1, 3))

File ~/software/miniconda/envs/quacc/lib/python3.10/site-packages/sella/peswrapper.py:230, in PES.get_g(self)
    229 def get_g(self):
--> 230     self._update()
    231     return self.curr['g'].copy()

File ~/software/miniconda/envs/quacc/lib/python3.10/site-packages/sella/peswrapper.py:193, in PES._update(self, feval)
    190 drdx, Ucons, Unred, Ufree = self._calc_basis()
    192 if feval:
--> 193     f, g = self.eval()
    194 else:
    195     f = None

File ~/software/miniconda/envs/quacc/lib/python3.10/site-packages/sella/peswrapper.py:159, in PES.eval(self)
    157 f = self.atoms.get_potential_energy()
    158 g = -self.atoms.get_forces().ravel()
--> 159 self.write_traj()
    160 return f, g

File ~/software/miniconda/envs/quacc/lib/python3.10/site-packages/sella/peswrapper.py:153, in PES.write_traj(self)
    151 def write_traj(self):
    152     if self.traj is not None:
--> 153         self.traj.write()

File ~/software/miniconda/envs/quacc/lib/python3.10/site-packages/ase/io/trajectory.py:134, in TrajectoryWriter.write(self, atoms, **kwargs)
    131 if atoms is None:
    132     atoms = self.atoms
--> 134 for image in atoms.iterimages():
    135     self._write_atoms(image, **kwargs)

AttributeError: 'NoneType' object has no attribute 'iterimages'

[ehermes]

You need to set the atoms kwarg to Trajectory. Try the following:

from ase.build import bulk
from ase.calculators.emt import EMT
from sella.optimize import Sella
from ase.io.trajectory import Trajectory

atoms = bulk("Cu") * (2, 1, 1)
atoms.calc = EMT()
dyn = Sella(atoms,trajectory=Trajectory("test.traj", "w", atoms=atoms))
dyn.run()

Note that the example you have provided fails for me on BFGS as well:

      Step     Time          Energy         fmax
BFGS:    0 18:17:28       -0.011363        0.0000
Traceback (most recent call last):
  File "/home/eric/test/sella/traj_test.py", line 10, in <module>
    dyn.run() # <--- this runs
    ^^^^^^^^^
  File "/home/eric/.local/lib/python3.11/site-packages/ase/optimize/optimize.py", line 269, in run
    return Dynamics.run(self)
           ^^^^^^^^^^^^^^^^^^
  File "/home/eric/.local/lib/python3.11/site-packages/ase/optimize/optimize.py", line 156, in run
    for converged in Dynamics.irun(self):
  File "/home/eric/.local/lib/python3.11/site-packages/ase/optimize/optimize.py", line 129, in irun
    self.call_observers()
  File "/home/eric/.local/lib/python3.11/site-packages/ase/optimize/optimize.py", line 108, in call_observers
    function(*args, **kwargs)
  File "/home/eric/.local/lib/python3.11/site-packages/ase/io/trajectory.py", line 131, in write
    for image in atoms.iterimages():
                 ^^^^^^^^^^^^^^^^
AttributeError: 'NoneType' object has no attribute 'iterimages'

[ehermes]

Also, just to be clear, the purpose of manually constructing the Trajectory object outside of Sella is so that you have a reference to the trajectory object in your main script... so you probably want to bind it to a variable:

from ase.build import bulk
from ase.calculators.emt import EMT
from sella.optimize import Sella
from ase.io.trajectory import Trajectory

atoms = bulk("Cu") * (2, 1, 1)
atoms.calc = EMT()
trajectory = Trajectory('test.traj', 'w', atoms=atoms)
dyn = Sella(atoms,trajectory=trajectory)
dyn.run()

Otherwise you'll still be forced to poke at the dyn.trajectory attribute, which as I've mentioned, is not necessarily a stable API.

[Andrew-S-Rosen]

@ehermes: Thank you this is absolutely brilliant. You are 100% right that this is a much better way to go. Thank you for reading my mind and providing a much more sound alternative despite me giving you basically no information to go on. I'm going to close this issue as largely being unnecessary.