LLNL / graphite

A repository for implementing graph network models based on atomic structures.
Other
62 stars 9 forks source link

Errors in denoiser demo notebook #2

Closed nnn911 closed 1 year ago

nnn911 commented 1 year ago

When I tried running the demo notebook I encountered 2 different errors which I hopefully managed to resolve.

My environment:

M1 Mac

sys.version '3.10.12 | packaged by conda-forge | (main, Jun 23 2023, 22:41:52) [Clang 15.0.7 ]' torch.version '1.13.1' torch_geometric.version '2.3.1'

Error 1

When running the 3rd cell I get the following error:

ModuleNotFoundError Traceback (most recent call last) Cell In[7], line 3 1 from torch_geometric.data import Data, Dataset 2 from sklearn.preprocessing import LabelEncoder ----> 3 from graphite.transforms.periodic_radius_graph_ase import PeriodicRadiusGraph_ase 5 class PeriodicStructureDataset(Dataset): 6 def init(self, atoms_list, large_cutoff, duplicate=128):

ModuleNotFoundError: No module named 'graphite.transforms.periodic_radius_graph_ase'

I was able to resolve this error by removing the _ase suffix for both graphite.transforms.periodic_radius_graph_ase and PeriodicRadiusGraph_ase.

Error 2

Running the next cell gives me this error:

KeyError Traceback (most recent call last) File ~/.local/lib/python3.10/site-packages/torch_geometric/data/storage.py:79, in BaseStorage.getattr(self, key) 78 try: ---> 79 return self[key] 80 except KeyError:

File ~/.local/lib/python3.10/site-packages/torch_geometric/data/storage.py:104, in BaseStorage.getitem(self, key) 103 def getitem(self, key: str) -> Any: --> 104 return self._mapping[key]

KeyError: 'box'

This could be resolved by updating the PeriodicRadiusGraph class from data.box to data.cell.

Error 3

After this change, the same cell produces a new error:

RuntimeError Traceback (most recent call last) Cell In[4], line 10 7 LARGE_CUTOFF = 5.0 9 # Ad-hoc Dataset class for applying random perturbation and then downselecting graph edges ---> 10 dataset = PeriodicStructureDataset(ideal_atoms_list, large_cutoff=LARGE_CUTOFF) 12 # Train-validation split 13 num_train = int(0.95 * len(dataset))

Cell In[3], line 20, in PeriodicStructureDataset.init(self, atoms_list, large_cutoff, duplicate) 12 x = LabelEncoder().fit_transform(atoms.numbers) 13 data = Data( 14 x = torch.tensor(x).long(), 15 pos = torch.tensor(atoms.positions).float(), (...) 18 numbers = torch.tensor(atoms.numbers).long(), 19 ) ---> 20 data = self.radiusgraph(data) 21 self.dataset.append(data) 22 self.dataset = [d.clone() for d in self.dataset for in range(duplicate)]

File ~/Work/ovito_plugins/graphite/src/graphite/transforms/periodic_radius_graph.py:21, in PeriodicRadiusGraph.call(self, data) 19 def call(self, data): 20 pos, box = data.pos, data.cell ---> 21 edge_index, edge_vec = periodic_radius_graph(pos, box, self.cutoff) ... ---> 77 coords = torch.div(pos, cell_dims, rounding_mode='floor').to(torch.long) 78 shifts = cell_shift(coords, shape) 80 # If any atom is outside the box, wrap the positions and the cell coordinates back into the box.

RuntimeError: The size of tensor a (1024) must match the size of tensor b (3) at non-singleton dimension 0

Which can be traced to the same code section and fixed replacing data.cell with data.cell.diag().

Conclusion

After these 2 minor changes I was able to use the denoiser as expected.

tim-hsu commented 1 year ago

@nnn911, thanks so much for pointing it out. The source code constantly gets updated so sometimes it's hard to keep track of whether the notebooks still work.

I merged your fix. I have also added an alternative ASE implementation defined in the notebook.

To provide a bit more context: The ASE neighbor list (graph construction) implementation is quite feature-rich and robust. I recommend using it for production. But ultimately I don't want ase to be a hard install dependency of graphite, so I plan to not have ASE-related codes in the graphite source code.