FAIR-Chem/fairchem

Run-relaxations from LMDB / ASE databases

Closed this issue · 3 comments

I’m trying to run main.py --mode run-relaxations directly from an LMDB and are getting an error relating to the FixAtoms constraint and the mask format. 

I’ve tried multiple different formats of .fixed attributes for the graphs objects (bool, tensor, indices) but run into similar issues.

Traceback (most recent call last):
  File "/lustre/hpc/kemi/clausen/fairchem/src/fairchem/../../main.py", line 8, in <module>
    main()
  File "/lustre/hpc/kemi/clausen/fairchem/src/fairchem/core/_cli.py", line 108, in main
    runner_wrapper(args.distributed, config)
  File "/lustre/hpc/kemi/clausen/fairchem/src/fairchem/core/_cli.py", line 55, in runner_wrapper
    Runner(distributed=distributed)(config)
  File "/lustre/hpc/kemi/clausen/fairchem/src/fairchem/core/_cli.py", line 42, in __call__
    self.task.run()
  File "/lustre/hpc/kemi/clausen/fairchem/src/fairchem/core/tasks/task.py", line 98, in run
    self.trainer.run_relaxations()
  File "/lustre/hpc/kemi/clausen/fairchem/src/fairchem/core/trainers/ocp_trainer.py", line 570, in run_relaxations
    relaxed_batch = ml_relax(
                    ^^^^^^^^^
  File "/lustre/hpc/kemi/clausen/fairchem/src/fairchem/core/common/relaxation/ml_relaxation.py", line 75, in ml_relax
    relaxed_batch = optimizer.run(fmax=fmax, steps=steps)
                    ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/lustre/hpc/kemi/clausen/fairchem/src/fairchem/core/common/relaxation/optimizers/lbfgs_torch.py", line 127, in run
    self.write(energy, forces, update_mask)
  File "/lustre/hpc/kemi/clausen/fairchem/src/fairchem/core/common/relaxation/optimizers/lbfgs_torch.py", line 215, in write
    atoms_objects = batch_to_atoms(self.batch)
                    ^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/lustre/hpc/kemi/clausen/fairchem/src/fairchem/core/common/relaxation/ase_utils.py", line 54, in batch_to_atoms
    constraint=FixAtoms(mask=fixed[idx].tolist()),
               ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/groups/kemi/clausen/miniconda3/envs/fairchem/lib/python3.11/site-packages/ase/constraints.py", line 137, in __init__
    raise ValueError(
ValueError: The indices array contains duplicates. Perhaps you want to specify a mask instead, but forgot the mask= keyword.

Ideally, I would also like to run the relaxations from an ASE database but it seems that this is not supported as the BaseTrainer does not recognize task.relax_dataset.format?

  1. Can you provide a minimum example (maybe writing a single toy ASE atoms object to an LMDB or ASE db and trying the relaxation?). Might be nice to add this as a test case for one of our models and/or add it to one of the tutorial docs.
  2. We moved the config for the relax_dataset format from task.relax_data.format to dataset.relax.format to keep it close to the dataset configs https://github.com/FAIR-Chem/fairchem/blob/main/src/fairchem/core/trainers/base_trainer.py#L366

Can you also share the ase version you are using. Because the error you are hitting is very explicitly addressed in the code by setting a mask - constraint=FixAtoms(mask=fixed[idx].tolist())

After your answers I consulted the updated example config and that cleared things up a bit.

Also, even though I had a version of the main branch from last week, I pulled a new version and this has apparently solved the constraints issue. I'll let you know if I experience something similar again but this can be closed.