AssertionError: "coordinates shape is not align atoms" when using custom dict input for UniMolRepr.get_repr()

[srdjan125]

Describe the bug

I’m trying to use the UniMolRepr class with a custom molecule input in the form of a dict containing 'atoms' and 'coordinates', following the documentation. However, I get an AssertionError deep inside the library stating:

AssertionError: coordinates shape is not align atoms

Here is my minimal example code:

from unimol_tools import UniMolRepr

data = {
'atoms': ['C', 'C', 'O', 'N', 'H', 'H'],
'coordinates': [
[0.0, 0.0, 0.0],
[1.2, 0.0, 0.0],
[2.4, 0.1, 0.0],
[0.0, 1.3, 0.0],
[-0.5, -0.8, 0.0],
[1.7, -0.9, 0.0]
]
}

clf = UniMolRepr(data_type='molecule', remove_hs=True)
result = clf.get_repr(data, return_atomic_reprs=True)

print("cls_repr shape:", result['cls_repr'].shape)
print("atomic_reprs shape:", result['atomic_reprs'].shape)

The atoms list and coordinates list have matching lengths (6 atoms, 6 coordinate sets). The coordinates are well-formed 3D lists.

Full error traceback snippet:

AssertionError: coordinates shape is not align atoms

I’ve verified that the lengths and shapes are correct, but it still fails inside the inner_coords function in conformer.py. Could you please clarify the expected input format for the data dict or let me know if this is a bug?

Thank you!

unimol_tools Version

latet one

Expected behavior

I provide data as given in your documenations, I don't understand why do I get this erroe

To Reproduce

No response

Environment

No response

Additional Context

No response

[darkkoo]

Currently you need to change the input to
data = {
'atoms': [['C', 'C', 'O', 'N', 'H', 'H']],
'coordinates': [[
[0.0, 0.0, 0.0],
[1.2, 0.0, 0.0],
[2.4, 0.1, 0.0],
[0.0, 1.3, 0.0],
[-0.5, -0.8, 0.0],
[1.7, -0.9, 0.0]
]]
}
Although there is only one molecule, it still needs to be wrapped as a 2D list. This problem will be solved in the next version. Thank you!

[srdjan125]

Hi @darkkoo thanx, I haved figured it out today by finding example on your another page: https://deepwiki.com/deepmodeling/Uni-Mol/5.4-representation-system-(unimolrepr) .
It would be good if the example of expected format was also given here: https://unimol.readthedocs.io/en/yn-docs/train.html#uni-mol-representation.
And yesterday I found another page which stated that data should be: atoms = ['C', 'C', 'C']
coordinates = [[0,0,0], [0,0,1], [0,0,2]] which was throwing error (I cannot recall which page exactly).

[emotionor]

@srdjan125 Thank you for your advice and reminders, we will improve the document and check the pages.