Ase dataset updates

[lbluque]

This PR includes updates in #630

Updates and fixes in AtomsToGraphs, AseAtomsDataset objects, and LMDBDatabase that make it easier to load and use Atoms with arbitrary properties as ASE style DBs as datasets with the unified OCP trainer.

AtomsToGraphs

1. Add r_stress as config option to include stress in data
2. Add r_data_keys as config option to pass a list of target properties saved in Atoms.info to be included in data (see example below)
3. Edit unit tests for 1 and 2

AseAtomsDataset

1. Add key_mapping and transforms attribute according to changes in Unified OCP Trainer #520
2. Allow using .db, .lmdb, .aselmbd extensions for datasets
3. Allow passing a list of paths for src in config

LMDBDatabase

1. minor bug fixes for handling updating and deleting rows
2. cleanup unit-tests

OCPTrainer

1. cast sid/fid data attributes to list using list() to allow string ids

Here's a minimal example for using:

from pathlib import Path
from ocpmodels.datasets import AseDBDataset

root_dir = Path("/root/path/to/databases")

# paths with ase dbs (ie LMDBDatabase, or dbs from ase: https://wiki.fysik.dtu.dk/ase/ase/db/db.html)
path1 = root_dir / "train_dir_1"
path2 = root_dir / "train_dir_2"
path3 = root_dir / "train_dir_3"

# To load the dataset for S2EFS we set energy, force, stress and magmoms as an example of some other property
# note the magmoms (or any additional property) should be saved in the 'data' attribute of each AtomsRow in the db

dataset = AseDBDataset(
    config={
        "src": [str(path1), str(path2), str(path3)],
        "a2g_args": {
            "r_energy": True,
            "r_forces": True,
            "r_stress": True,
            "r_data_keys": ["magmoms"]
        },
        "key_mapping": {"magmoms": "magnetic_moments"}
    }
)

# note the atoms objects 

# now they should all be part of data objects
print(dataset[0].energy, dataset[0].forces.shape, dataset[0].stress.shape, dataset[0].magnetic_moments.shape)

# you can also directly inspect the corresponding `Atoms` objects
atoms = dataset.get_atoms(0)

# the additional data is saved in the atoms info
print(atoms.info.keys())  # this includes "magmoms"

[mshuaibii]

Looks good! Minor changes.

[codecov]

Codecov Report

Attention: Patch coverage is 79.11392% with 33 lines in your changes are missing coverage. Please review.

Project coverage is 57.23%. Comparing base (fa39a8f) to head (bdbba48).

Files	Patch %	Lines
ocpmodels/datasets/ase_datasets.py	82.81%	11 Missing ⚠️
ocpmodels/datasets/lmdb_database.py	82.97%	8 Missing ⚠️
ocpmodels/trainers/ocp_trainer.py	14.28%	6 Missing ⚠️
ocpmodels/preprocessing/atoms_to_graphs.py	82.35%	3 Missing ⚠️
ocpmodels/common/utils.py	33.33%	2 Missing ⚠️
ocpmodels/datasets/_utils.py	91.66%	1 Missing ⚠️
ocpmodels/datasets/lmdb_dataset.py	66.66%	1 Missing ⚠️
ocpmodels/datasets/oc22_lmdb_dataset.py	50.00%	1 Missing ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##             main     #622      +/-   ##
==========================================
+ Coverage   56.98%   57.23%   +0.25%     
==========================================
  Files         108      109       +1     
  Lines       10262    10287      +25     
==========================================
+ Hits         5848     5888      +40     
+ Misses       4414     4399      -15     

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

[mshuaibii]

Minor comments