Subject: Issue with ase_interface.py in SchNetPack and Atomref-based Model Predictions

[Dale911]

Dear SchNetPack Developers,

I am encountering an issue when using ase_interface.py for prediction and molecular dynamics (MD) simulations with a model that was trained using atom references.

Specifically, my model was trained with atomref and RemoveOffsets/AddOffsets transformations as follows:

trn.RemoveOffsets( property=ENERGY_KEY, remove_mean=True, remove_atomrefs=True, atomrefs=atomrefs_tensor, # Passed as a tensor property_mean=means[ENERGY_KEY], is_extensive=False ) The MAE on the validation set is approximately 2 kcal/mol, which is reasonable. However, when I use the model to predict energies on the training set via the ASE interface, the predictions are highly inaccurate. For example, a raw energy of ~1056 Hartree is predicted as only ~856 Hartree. (MAE -HUGH)

It seems the atom references or offsets might not be correctly handled during inference with the ASE interface. Could you please advise on how to properly use ase_interface.py with models trained using atomref and RemoveOffsets?

Any insights or suggestions would be greatly appreciated.

Best regards, Dale

[stefaanhessmann]

Hi @Dale911,

it looks like postprocessing is not working properly for you. Could you check if it is turned on and if the correct modules are defined as postprocessors:

print(model.do_postprocessing)
print(model.postprocessors)

Best, Stefaan

[Dale911]

I check my model is print(model.do_postprocessing) ...: print(model.postprocessors) True ModuleList( (0): CastTo64() (1): AddOffsets() )

[stefaanhessmann]

Could you please double check if all settings in RemoveOffsets and AddOffsets are equal? Especially the provided your custom atomrefs-Tensor.

And does the issue also happen when you use the model outside of the ase_interface?

If this does not solve the issue, there might be some bug and we will have a deeper look.

[Dale911]

I used the same of settting for RemoveOffsets and AddOffsets as my training script:

data = spk.data.AtomsDataModule(
    datapath=DATABASE_PATH,
    batch_size=BATCH_SIZE_PER_GPU,
    num_train=num_train_actual,        # Use calculated number
    num_val=num_val_actual,            # Use calculated number
    transforms=[
        trn.ASENeighborList(cutoff=CUTOFF),
        # Only remove mean, NO atomrefs
        trn.RemoveOffsets(
            property=ENERGY_KEY, 
            remove_mean=True, 
            remove_atomrefs=False,      # DISABLED atomrefs
            property_mean=means[ENERGY_KEY],
            is_extensive=True
        ),
        trn.CastTo32()
    ],
    property_units={
        ENERGY_KEY: ENERGY_UNITS, 
        FORCES_KEY: FORCES_UNITS
    },
    num_workers=NUM_WORKERS,
    pin_memory=PIN_MEMORY
)
   
nnpot = spk.model.NeuralNetworkPotential(
        representation=schnet,
        input_modules=[pairwise_distance],
        output_modules=[pred_energy, pred_forces],
        postprocessors=[
            trn.CastTo64(),
            # Only add mean back, NO atomrefs
            trn.AddOffsets(
                ENERGY_KEY, 
                add_mean=True, 
                add_atomrefs=False,         # DISABLED atomrefs
                property_mean=means[ENERGY_KEY],
                is_extensive=True
            )
        ]
    )

i also train another model withou atomref (mean no setting up add/remove-offsets), and then no problem with the spk's prediction,

Thanks for your suppport

Bests Dale

[sundusaijaz]

Dear Dale,

Thanks for reporting this issue. We took a closer look and found out that the RemoveOffsets transform wasn’t handling intensive properties (like energy per atom) quite right—the atom references were being subtracted as a total sum for whole batch, rather than averaged per atom.

We’ve now fixed this by making sure the atom reference bias is divided by the number of atoms when is_extensive=False. To make sure everything works as expected, we also added comprehensive test cases (see tests/atomistic/test_offset.py) to cover both extensive and intensive properties and to double-check consistency between training and inference.

You can find all the details and the complete test suite along with the fix in PR #732

Best, Sundus