From b63e86ca11743b005dd0db7ede9e69fa8b052f55 Mon Sep 17 00:00:00 2001 From: cvmbb Date: Wed, 29 Jan 2025 17:57:48 +0100 Subject: [PATCH 1/3] update --- aimnet2calc/aimnet2ase.py | 21 ++++++++++++---- aimnet2calc/nblist.py | 8 +++---- test/test_ase.py | 50 +++++++++++++++++++++++++++++++++++++++ test/test_calculator.py | 5 ++-- 4 files changed, 73 insertions(+), 11 deletions(-) create mode 100644 test/test_ase.py diff --git a/aimnet2calc/aimnet2ase.py b/aimnet2calc/aimnet2ase.py index 1155cb0..6ee121f 100644 --- a/aimnet2calc/aimnet2ase.py +++ b/aimnet2calc/aimnet2ase.py @@ -1,12 +1,14 @@ -from ase.calculators.calculator import Calculator, all_changes -from aimnet2calc import AIMNet2Calculator from typing import Union -import torch + import numpy as np +import torch +from ase.calculators.calculator import Calculator, all_changes + +from aimnet2calc.calculator import AIMNet2Calculator class AIMNet2ASE(Calculator): - implemented_properties = ['energy', 'forces', 'free_energy', 'charges', 'stress'] + implemented_properties = ['energy', 'forces', 'free_energy', 'charges', 'stress', 'dipole_moment'] def __init__(self, base_calc: Union[AIMNet2Calculator, str] = 'aimnet2', charge=0, mult=1): super().__init__() if isinstance(base_calc, str): @@ -56,6 +58,12 @@ def update_tensors(self): if self._t_mol_idx is None: self.mol_idx = torch.zeros(len(self.atoms), dtype=torch.int64, device=self.base_calc.device) + def get_dipole_moment(self,atoms): + charges = self.get_charges()[:, np.newaxis] + positions = atoms.get_positions() + return np.sum(charges * positions, axis=0) + + def calculate(self, atoms=None, properties=['energy'], system_changes=all_changes): super().calculate(atoms, properties, system_changes) self.update_tensors() @@ -77,8 +85,11 @@ def calculate(self, atoms=None, properties=['energy'], system_changes=all_change for k, v in results.items(): results[k] = v.detach().cpu().numpy() - self.results['energy'] = results['energy'] + # ase expects us to return a scalar + self.results['energy'] = results['energy'][0] self.results['charges'] = results['charges'] + self.results['dipole_moment'] = self.get_dipole_moment(self.atoms) + if 'forces' in properties: self.results['forces'] = results['forces'] if 'stress' in properties: diff --git a/aimnet2calc/nblist.py b/aimnet2calc/nblist.py index 623ff6e..f34ce54 100644 --- a/aimnet2calc/nblist.py +++ b/aimnet2calc/nblist.py @@ -134,12 +134,12 @@ def _nblist_pbc_cuda(conn_mat, shifts): return idx_j, mat_pad, shifts[S_idx] -def _nblist_pbc_cpu(conn_mat, shifts, device): +def _nblist_pbc_cpu(conn_mat, shifts): conn_mat = conn_mat.cpu().numpy() mat_idxj, mat_pad, mat_S_idx = _cpu_dense_nb_mat_sft(conn_mat) - mat_idxj = torch.from_numpy(mat_idxj).to(device) - mat_pad = torch.from_numpy(mat_pad).to(device) - mat_S_idx = torch.from_numpy(mat_S_idx).to(device) + mat_idxj = torch.from_numpy(mat_idxj).cpu() + mat_pad = torch.from_numpy(mat_pad).cpu() + mat_S_idx = torch.from_numpy(mat_S_idx).cpu() mat_S = shifts[mat_S_idx] return mat_idxj, mat_pad, mat_S diff --git a/test/test_ase.py b/test/test_ase.py new file mode 100644 index 0000000..faa8876 --- /dev/null +++ b/test/test_ase.py @@ -0,0 +1,50 @@ +import os + +import ase +import numpy as np + +from aimnet2calc.aimnet2ase import AIMNet2ASE + +MODELS = ('aimnet2', 'aimnet2_b973c') +DIR = os.path.dirname(__file__) + + +def _struct_pbc(): + filename = os.path.join(DIR, '1008775.cif') + return ase.io.read(filename) + + +def _struct_list(): + filename = os.path.join(DIR, 'mols_size_var.xyz') + return ase.io.read(filename, index=':') + + +def _stuct_batch(): + filename = os.path.join(DIR, 'mols_size_36.xyz') + return ase.io.read(filename, index=':') + + +def _test_dipole(calc, atoms): + atoms.calc = calc + e =atoms.get_potential_energy() + assert e.shape == () + assert e < 0.0 + + dm = atoms.get_dipole_moment() + assert dm.shape == (3,) + assert np.linalg.norm(dm) > 0.0 + + +def test_calculator(): + for model in MODELS: + print('Testing model:', model) + calculator = AIMNet2ASE(model) + for atoms_list, runtype in zip((_stuct_batch(), _struct_list(), _struct_pbc()), ('batch', 'list', 'pbc')): + if runtype == 'batch' or runtype == 'list': + for atoms in atoms_list: + _test_dipole(calculator, atoms) + + else: + _test_dipole(calculator, atoms_list) + + diff --git a/test/test_calculator.py b/test/test_calculator.py index 38ec8b7..9711a73 100644 --- a/test/test_calculator.py +++ b/test/test_calculator.py @@ -1,8 +1,9 @@ -import ase.io -from aimnet2calc.calculator import AIMNet2Calculator import os + +import ase.io import numpy as np +from aimnet2calc.calculator import AIMNet2Calculator MODELS = ('aimnet2', 'aimnet2_b973c') DIR = os.path.dirname(__file__) From 4809e9374ee0ce42a71826e0baddc1a99167aa71 Mon Sep 17 00:00:00 2001 From: cvmbb Date: Wed, 29 Jan 2025 20:56:32 +0100 Subject: [PATCH 2/3] minor changes to get_dipole_moment implementatiion --- aimnet2calc/nblist.py | 6 +++--- test/test_ase.py | 9 ++++++--- 2 files changed, 9 insertions(+), 6 deletions(-) diff --git a/aimnet2calc/nblist.py b/aimnet2calc/nblist.py index f34ce54..aba073b 100644 --- a/aimnet2calc/nblist.py +++ b/aimnet2calc/nblist.py @@ -137,9 +137,9 @@ def _nblist_pbc_cuda(conn_mat, shifts): def _nblist_pbc_cpu(conn_mat, shifts): conn_mat = conn_mat.cpu().numpy() mat_idxj, mat_pad, mat_S_idx = _cpu_dense_nb_mat_sft(conn_mat) - mat_idxj = torch.from_numpy(mat_idxj).cpu() - mat_pad = torch.from_numpy(mat_pad).cpu() - mat_S_idx = torch.from_numpy(mat_S_idx).cpu() + mat_idxj = torch.from_numpy(mat_idxj) + mat_pad = torch.from_numpy(mat_pad) + mat_S_idx = torch.from_numpy(mat_S_idx) mat_S = shifts[mat_S_idx] return mat_idxj, mat_pad, mat_S diff --git a/test/test_ase.py b/test/test_ase.py index faa8876..5a0c4e3 100644 --- a/test/test_ase.py +++ b/test/test_ase.py @@ -28,12 +28,15 @@ def _test_dipole(calc, atoms): atoms.calc = calc e =atoms.get_potential_energy() assert e.shape == () - assert e < 0.0 + assert hasattr(atoms, 'get_charges') + q = atoms.get_charges() + assert q.shape == (len(atoms),) + + assert hasattr(atoms, 'get_dipole_moment') dm = atoms.get_dipole_moment() assert dm.shape == (3,) - assert np.linalg.norm(dm) > 0.0 - + def test_calculator(): for model in MODELS: From c09168cd40f50886849fb588737304465082f134 Mon Sep 17 00:00:00 2001 From: cvmbb Date: Thu, 30 Jan 2025 17:07:39 +0100 Subject: [PATCH 3/3] using item nstead of accessing index --- aimnet2calc/aimnet2ase.py | 3 +-- test/test_ase.py | 2 +- 2 files changed, 2 insertions(+), 3 deletions(-) diff --git a/aimnet2calc/aimnet2ase.py b/aimnet2calc/aimnet2ase.py index 6ee121f..67fe0ca 100644 --- a/aimnet2calc/aimnet2ase.py +++ b/aimnet2calc/aimnet2ase.py @@ -85,8 +85,7 @@ def calculate(self, atoms=None, properties=['energy'], system_changes=all_change for k, v in results.items(): results[k] = v.detach().cpu().numpy() - # ase expects us to return a scalar - self.results['energy'] = results['energy'][0] + self.results['energy'] = results['energy'].item() self.results['charges'] = results['charges'] self.results['dipole_moment'] = self.get_dipole_moment(self.atoms) diff --git a/test/test_ase.py b/test/test_ase.py index 5a0c4e3..4143814 100644 --- a/test/test_ase.py +++ b/test/test_ase.py @@ -27,7 +27,7 @@ def _stuct_batch(): def _test_dipole(calc, atoms): atoms.calc = calc e =atoms.get_potential_energy() - assert e.shape == () + assert isinstance(e, float) assert hasattr(atoms, 'get_charges') q = atoms.get_charges()