Predict at atom level

[ma763]

Hello,

Thank you so much for this project, very interesting!

I am trying to use this package to do a regression prediction to the atom level instead of predicting a property for the whole crystal. I am trying to input the atom labels but struggling to change the data.py script accordingly... what would be the easiest way to input the atom labels?

And I guess I need to output the atom prediction before the pooling? Does this correspond to return out of def forward(): in model.py?

Thanks!
Marta

[txie-93]

I think you will need to get the atom features before pooling, as you suggested. The atom features are atom_fea at the following line.

cgcnn/cgcnn/model.py

Line 154 in d612a69

atom_fea = conv_func(atom_fea, nbr_fea, nbr_fea_idx)

[ma763]

Okay, that's helpful, will try this!

However, I am still struggling to input the atom label data for training (instead of giving a label for the whole crystal structure). Is it even possible?
Or if not, do you recommend that I try another package for what I am trying to do?

I really appreciate your help, thank you very much.

[txie-93]

You will need to change the data loader in that case.

Check out the dataset class here.

cgcnn/cgcnn/data.py

Line 350 in d612a69

return (atom_fea, nbr_fea, nbr_fea_idx), target, cif_id

You may also need to change the collate function.

cgcnn/cgcnn/data.py

Line 96 in d612a69

def collate_pool(dataset_list):