mlatom 3.16.2__tar.gz → 3.17.0__tar.gz

{mlatom-3.16.2/src/mlatom.egg-info → mlatom-3.17.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: mlatom
-Version: 3.16.2
+Version: 3.17.0
 Summary: A Package for AI-enhanced computational chemistry
 Home-page: http://mlatom.com
 Author: Pavlo O. Dral
@@ -14,20 +14,11 @@ Classifier: Environment :: Console
 Requires-Python: >=3.8
 Description-Content-Type: text/markdown
 License-File: LICENSE.txt
-Requires-Dist: numpy
-Requires-Dist: scipy
-Requires-Dist: h5py
-Requires-Dist: pyh5md
-Requires-Dist: torch
-Requires-Dist: torchani
-Requires-Dist: matplotlib
-Requires-Dist: statsmodels
-Requires-Dist: tqdm
 
 # About Program
 MLatom: a Package for Atomistic Simulations with Machine Learning
 
-Version 3.16.2
+Version 3.17.0
 
 **Official website**: http://mlatom.com/    
 **Documentation**: http://mlatom.com/docs/                     
@@ -56,7 +47,7 @@ When this Software or its derivatives are used in scientific publications, it sh
 * Pavlo O. Dral, Fuchun Ge, Yi-Fan Hou, Peikun Zheng, Yuxinxin Chen, Mario Barbatti, Olexandr Isayev, Cheng Wang, Bao-Xin Xue, Max Pinheiro Jr, Yuming Su, Yiheng Dai, Yangtao Chen, Lina Zhang, Shuang Zhang, Arif Ullah, Quanhao Zhang, Yanchi Ou. [MLatom 3: A Platform for Machine Learning-enhanced Computational Chemistry Simulations and Workflows](https://doi.org/10.1021/acs.jctc.3c01203).   *J. Chem. Theory Comput.* **2024**, *20*, 1193--1213.                                                                                               
 * Pavlo O. Dral, Fuchun Ge, Bao-Xin Xue, Yi-Fan Hou, Max Pinheiro Jr, Jianxing Huang, Mario Barbatti. [MLatom 2: An Integrative Platform for Atomistic Machine Learning](http://doi.org/10.1007/s41061-021-00339-5). *Top. Curr. Chem.* **2021**, *379*, 27.       
 * Pavlo O. Dral. [MLatom: A Program Package for Quantum Chemical Research Assisted by Machine Learning](http://dx.doi.org/10.1002/jcc.26004). *J. Comput. Chem.* **2019**, *40*, 2339--2347.                                                                                      
-* Pavlo O. Dral, Fuchun Ge, Yi-Fan Hou, Peikun Zheng, Yuxinxin Chen, Bao-Xin Xue, Mikolaj Martyka, Max Pinheiro Jr, Yuming Su, Yiheng Dai, Yangtao Chen, Shuang Zhang, Lina Zhang, Arif Ullah, Quanhao Zhang, Sebastian V. Pios, Yanchi Ou, Matheus O. Bispo, Vignesh B. Kumar, Xin-Yu Tong, MLatom: A Package for Atomistic Simulations with Machine Learning,  version 3.16.2, Xiamen University, Xiamen, China, 2013-2024.
+* Pavlo O. Dral, Fuchun Ge, Yi-Fan Hou, Peikun Zheng, Yuxinxin Chen, Bao-Xin Xue, Mikolaj Martyka, Max Pinheiro Jr, Yuming Su, Yiheng Dai, Yangtao Chen, Shuang Zhang, Lina Zhang, Arif Ullah, Quanhao Zhang, Sebastian V. Pios, Yanchi Ou, Matheus O. Bispo, Vignesh B. Kumar, Xin-Yu Tong, MLatom: A Package for Atomistic Simulations with Machine Learning,  version 3.17.0, Xiamen University, Xiamen, China, 2013-2024.
   [MLatom.com](http://mlatom.com).
                                                                       
 The citations for MLatom's interfaces and features shall be eventually

{mlatom-3.16.2 → mlatom-3.17.0}/README.md

@@ -1,7 +1,7 @@
 # About Program
 MLatom: a Package for Atomistic Simulations with Machine Learning
 
-Version 3.16.2
+Version 3.17.0
 
 **Official website**: http://mlatom.com/    
 **Documentation**: http://mlatom.com/docs/                     
@@ -30,7 +30,7 @@ When this Software or its derivatives are used in scientific publications, it sh
 * Pavlo O. Dral, Fuchun Ge, Yi-Fan Hou, Peikun Zheng, Yuxinxin Chen, Mario Barbatti, Olexandr Isayev, Cheng Wang, Bao-Xin Xue, Max Pinheiro Jr, Yuming Su, Yiheng Dai, Yangtao Chen, Lina Zhang, Shuang Zhang, Arif Ullah, Quanhao Zhang, Yanchi Ou. [MLatom 3: A Platform for Machine Learning-enhanced Computational Chemistry Simulations and Workflows](https://doi.org/10.1021/acs.jctc.3c01203).   *J. Chem. Theory Comput.* **2024**, *20*, 1193--1213.                                                                                               
 * Pavlo O. Dral, Fuchun Ge, Bao-Xin Xue, Yi-Fan Hou, Max Pinheiro Jr, Jianxing Huang, Mario Barbatti. [MLatom 2: An Integrative Platform for Atomistic Machine Learning](http://doi.org/10.1007/s41061-021-00339-5). *Top. Curr. Chem.* **2021**, *379*, 27.       
 * Pavlo O. Dral. [MLatom: A Program Package for Quantum Chemical Research Assisted by Machine Learning](http://dx.doi.org/10.1002/jcc.26004). *J. Comput. Chem.* **2019**, *40*, 2339--2347.                                                                                      
-* Pavlo O. Dral, Fuchun Ge, Yi-Fan Hou, Peikun Zheng, Yuxinxin Chen, Bao-Xin Xue, Mikolaj Martyka, Max Pinheiro Jr, Yuming Su, Yiheng Dai, Yangtao Chen, Shuang Zhang, Lina Zhang, Arif Ullah, Quanhao Zhang, Sebastian V. Pios, Yanchi Ou, Matheus O. Bispo, Vignesh B. Kumar, Xin-Yu Tong, MLatom: A Package for Atomistic Simulations with Machine Learning,  version 3.16.2, Xiamen University, Xiamen, China, 2013-2024.
+* Pavlo O. Dral, Fuchun Ge, Yi-Fan Hou, Peikun Zheng, Yuxinxin Chen, Bao-Xin Xue, Mikolaj Martyka, Max Pinheiro Jr, Yuming Su, Yiheng Dai, Yangtao Chen, Shuang Zhang, Lina Zhang, Arif Ullah, Quanhao Zhang, Sebastian V. Pios, Yanchi Ou, Matheus O. Bispo, Vignesh B. Kumar, Xin-Yu Tong, MLatom: A Package for Atomistic Simulations with Machine Learning,  version 3.17.0, Xiamen University, Xiamen, China, 2013-2024.
   [MLatom.com](http://mlatom.com).
                                                                       
 The citations for MLatom's interfaces and features shall be eventually

{mlatom-3.16.2 → mlatom-3.17.0}/setup.py

@@ -43,7 +43,7 @@ with open(os.path.join(here, 'README.md')) as f:
 
 setup(
     name = pkg_name,
-    version = "3.16.2",
+    version = "3.17.0",
     author = "Pavlo O. Dral",
     author_email = "admin@mlatom.com",
     license = 'MIT (modified)',

{mlatom-3.16.2 → mlatom-3.17.0}/src/mlatom/MLatom.py

@@ -4,7 +4,7 @@
   !---------------------------------------------------------------------------!
   !                                                                           !
   !     MLatom: a Package for Atomistic Simulations with Machine Learning     !
-  !                             MLatom 3.16.2                                 !
+  !                             MLatom 3.17.0                                 !
   !                                   @                                       !
   !                 Xiamen Atomistic Computing Suite (XACS)                   !
   !                                                                           !
@@ -38,7 +38,7 @@
   ! Sebastian V. Pios, Yanchi Ou, Matheus O. Bispo, Vignesh B. Kumar,         !
   ! Xin-Yu Tong,                                                              !
   ! MLatom: A Package for Atomistic Simulations with Machine Learning,        !
-  ! version 3.16.2, Xiamen University, Xiamen, China, 2013-2024.              !
+  ! version 3.17.0, Xiamen University, Xiamen, China, 2013-2024.              !
   !                                                                           !
   ! The citations for MLatom's interfaces and features shall be eventually    !
   ! included too. See header.py, ref.json and http://mlatom.com.              !
@@ -56,9 +56,24 @@
 '''
 
 import os, sys, time
-from mlatom import header
-from mlatom.MLtasks import CLItasks
-from mlatom.args_class import mlatom_args
+# import header
+# from mlatom.MLtasks import CLItasks
+# from mlatom.args_class import mlatom_args
+import importlib.util
+# ~POD, 2025.03.23
+# the complicated import below is required to load the same instance of mlatom,
+# where this script is located.
+
+dir_path = os.path.dirname(os.path.abspath(__file__))
+path2init = os.path.join(dir_path, '__init__.py')
+dirname = os.path.basename(dir_path)
+spec = importlib.util.spec_from_file_location(dirname, path2init)
+mlatom_with_this_file = importlib.util.module_from_spec(spec)
+sys.modules[dirname] = mlatom_with_this_file
+spec.loader.exec_module(mlatom_with_this_file)
+header = mlatom_with_this_file.header
+CLItasks = mlatom_with_this_file.MLtasks.CLItasks
+mlatom_args = mlatom_with_this_file.args_class.mlatom_args
 
 def run(argv = []):
     starttime = time.time()

{mlatom-3.16.2 → mlatom-3.17.0}/src/mlatom/MLtasks.py

@@ -23,7 +23,7 @@ class CLItasks(object):
     def __init__(self, args):
         self.args = args
     def run(self):
-        if self.args.MLprog.lower() == 'mlatomf' and self.args.task in self.MLatomF_tasks and not self.args.method and not self.args.hyperparameter_optimization['hyperparameters']:
+        if self.args.MLprog.lower() == 'mlatomf' and self.args.task in self.MLatomF_tasks and not self.args.method and not self.args.hyperparameter_optimization['hyperparameters'] and not self.args.uaiqm:
             return run_with_mlatomF(self.args)
         else: 
             return globals()[self.args.task](self.args)
@@ -39,10 +39,12 @@ def run_with_mlatomF(args):
         post_delta_learning(args)
     return results
     
-# Functions for tasks bleow. Name with corresponding task name in args_class.mLatom_args._task_list
+# Functions for tasks below. Name with corresponding task name in args_class.mLatom_args._task_list
 def useMLmodel(args):
     if args.deltaLearn:
         pre_delta_learning(args)
+    if args.uaiqm:
+        get_uaiqm_args(args)
     model = loading_model(args)
     molecular_database = loading_data(args.XYZfile, args.Yfile, args.YgradXYZfile, charges=args.charges, multiplicities=args.multiplicities)
     t_predict = predicting(model=model, molecular_database=molecular_database, value=bool(args.YestFile), gradient=bool(args.YgradXYZestFile), hessian=bool(args.hessianEstFile), ismethod=args.method)
@@ -60,6 +62,13 @@ def useMLmodel(args):
                     printing_aiqm2_results(molecule=mol)
                 elif args.ani1ccx or args.ani1x or args.ani2x or args.ani1xd4 or args.ani2xd4 or args.ani1ccxgelu or args.ani1xgelu or args.ani1ccxgelud4 or args.ani1xgelud4:
                     printing_animethod_results(methodname=model.method, molecule=mol)
+        elif 'uaiqm' in args.method:
+            nmol = 0
+            for mol in molecular_database:
+                nmol += 1
+                print('\n\n Properties of molecule %d\n' % nmol)
+                uversion = None if not args.uversion else args.uversion
+                printing_uaiqm_results(molecule=mol, method=args.method, uversion=args.uversion)
         else:
             for imol in range(len(molecular_database.molecules)): print(' Energy of molecule %6d: %25.13f Hartree' % (imol+1, molecular_database.molecules[imol].energy))
         print('')
@@ -374,6 +383,8 @@ def tsfreq(args):
 
 def geomopt(args, return_moldb=False):
     from . import simulations
+    if args.uaiqm:
+        get_uaiqm_args(args)
     molDB = loading_data(XYZfile=args.XYZfile, charges=args.charges, multiplicities=args.multiplicities)
     model = loading_model(args)
     fname = args.optXYZ
@@ -420,6 +431,11 @@ def geomopt(args, return_moldb=False):
             print('\n\n Final properties of molecule %d\n' % (imol+1))
             printing_aiqm2_results(molecule=geomopt.optimized_molecule)
             print('\n')
+        elif args.method and 'uaiqm' in args.method:
+            print('\n\n Final properties of molecule %d\n' % (imol+1))
+            uversion = None if not args.uversion else args.uversion
+            printing_uaiqm_results(molecule=geomopt.optimized_molecule, method=args.method, uversion=uversion)
+            print('\n')
         elif args.ani1ccx or args.ani1x or args.ani2x or args.ani1xd4 or args.ani2xd4 or args.ani1ccxgelu or args.ani1xgelu or args.ani1ccxgelud4 or args.ani1xgelud4:
             print('\n\n Final properties of molecule %d\n' % (imol+1))
             printing_animethod_results(methodname=model.method, molecule=geomopt.optimized_molecule)
@@ -441,6 +457,8 @@ def freq(args, molDB=None):
     from . import simulations
     if molDB is None:
         molDB = loading_data(args.XYZfile, charges=args.charges, multiplicities=args.multiplicities)
+    if args.uaiqm:
+        get_uaiqm_args(args)
     model = loading_model(args)
     kwargs = {}
     if args.freqProg:       kwargs['program'] = args.freqProg
@@ -518,6 +536,9 @@ def freq(args, molDB=None):
                 #     scaling_factor = 0.962
                 
             if scaling:
+                mol_copy = mol.copy()
+                mol_copy.frequencies = np.array(mol_copy.frequencies) * scaling_factor
+                mol_copy.dump(f'scaled_freq_gaussian{mol.number}.log',format='gaussian')
                 print(' Scale frequencies linearly')
                 print(f' Scaling factor: {scaling_factor}')
                 print('   Mode     Frequencies     Reduced masses     Force Constants       IR intensities')
@@ -536,6 +557,10 @@ def freq(args, molDB=None):
         if args.AIQM2 or args.AIQM2DFT:
             printing_aiqm2_results(molecule=mol)
             print('')
+        elif args.method and 'uaiqm' in args.method:
+            uversion = None if not args.uversion else args.uversion
+            printing_uaiqm_results(molecule=mol, method=args.method, uversion=args.uversion)
+            print('')
         elif args.ani1ccx or args.ani1x or args.ani2x or args.ani1xd4 or args.ani2xd4 or args.ani1ccxgelu or args.ani1xgelu or args.ani1ccxgelud4 or args.ani1xgelud4:
             printing_animethod_results(methodname=model.method, molecule=mol)
             print('')
@@ -657,6 +682,10 @@ def loading_method(args):
             kwargs['qm_program'] = args.qmprog
         else:
             kwargs['program'] = args.qmprog 
+    if args.uversion:
+        kwargs['version'] = args.uversion
+    if 'uwarning' in args.args_dict:
+        kwargs['warning'] = args.uwarning
     method = models.methods(**kwargs)
     method.set_num_threads(args.nthreads)
     return method
@@ -1050,6 +1079,21 @@ def post_delta_learning(args):
                         strtmp += '%25.13f   ' % (ygradxyzb[imol][iatom][idim] + corr[imol][iatom][idim])
                     fyestt.writelines('%s\n' % strtmp)
 
+def get_uaiqm_args(args):
+    molecular_database = loading_data(args.XYZfile, args.Yfile, args.YgradXYZfile, charges=args.charges, multiplicities=args.multiplicities)
+    model = models.uaiqm(method='uaiqm_optimal')
+    if not args.ncpus:
+        args.ncpus = 8
+    if not args.time_budget:
+        args.time_budget = None
+    model.select_optimal(
+        molecule=molecular_database[0],
+        nCPUs=args.ncpus,
+        time_budget=args.time_budget
+    )
+    args.method = model.method
+    args.uversion = model.version  
+
 def printing_aiqm1_results(aiqm1=True, molecule=None):
     fmt = ' %-41s: %15.8f Hartree'
     # find aiqm1 keyword
@@ -1078,6 +1122,22 @@ def printing_aiqm2_results(molecule=None):
     print(fmt % ('GFN2-xTB* contribution', molecule.gfn2xtbstar.energy), file=sys.stdout)
     print(fmt % ('D4 contribution', molecule.d4wb97x.energy))
     print(fmt % ('Total energy', molecule.energy))
+
+def printing_uaiqm_results(molecule=None, method=None, uversion=None):
+    print(f' Selected UAIQM method: {method}')
+    if uversion:
+        print(f' Selected version: {uversion}\n')
+    else:
+        print(f' Selected version: latest\n')
+    fmt = ' %-41s: %15.8f Hartree'
+    print(fmt % ('Standard deviation of ML contribution', molecule.energy_standard_deviation), end='')
+    print(' %15.5f kcal/mol' % (molecule.energy_standard_deviation * constants.Hartree2kcalpermol))
+    if molecule.baseline_energy:
+        print(fmt % ('Baseline contribution', molecule.baseline_energy))
+    print(fmt % ('NN contribution', molecule.MLs_energy))
+    if molecule.dispersion_energy:
+        print(fmt % ('D4 contribution', molecule.dispersion_energy), file=sys.stdout)
+    print(fmt % ('Total energy', molecule.energy))
     
 def printing_animethod_results(methodname=None, molecule=None):
     if 'D4'.casefold() in methodname.casefold():

mlatom-3.17.0/src/mlatom/__init__.py

@@ -0,0 +1,11 @@
+#!/usr/bin/env python3
+from . import args_class, MLtasks, header, data, models, plot, simulations, stats, xyz, namd, constants, spectra
+from .simulations import optimize_geometry, irc, freq, thermochemistry, md, generate_initial_conditions, vibrational_spectrum, md_parallel
+from .data import atom, molecule, molecular_database
+from .models import kreg, ani, msani
+from .models import methods, aiqm1, aiqm2, dens, ani_methods, aimnet2_methods, gaussian_methods, pyscf_methods, orca_methods, turbomole_methods, mndo_methods, sparrow_methods, xtb_methods, dftbplus_methods, columbus_methods, dftd3_methods, dftd4_methods
+from . import al_utils
+from .al import al
+from .gap_md import gap_md, gap_model
+# add-ons
+from .models import uaiqm

mlatom-3.17.0/src/mlatom/aimnet2.py

@@ -0,0 +1,110 @@
+'''
+.. code-block::
+
+  !---------------------------------------------------------------------------! 
+  ! aimnet2: Universal AIMnet2 models                                         ! 
+  ! Implementations by: Yuxinxin Chen                                         !
+  !---------------------------------------------------------------------------! 
+'''
+from .model_cls import torchani_model, method_model
+from .decorators import doc_inherit
+
+class aimnet2_methods(torchani_model, method_model):
+    '''
+    Universal ML methods with AIMNet2: https://doi.org/10.26434/chemrxiv-2023-296ch
+
+    Arguments:
+        method (str): A string that specifies the method. Available choices: ``'AIMNet2@b973c'`` and ``'AIMNet2@wb97m-d3'``.
+        device (str, optional): Indicate which device the calculation will be run on, i.e. 'cpu' for CPU, 'cuda' for Nvidia GPUs. When not speficied, it will try to use CUDA if there exists valid ``CUDA_VISIBLE_DEVICES`` in the environ of system.
+
+    '''
+    
+    supported_methods = ['AIMNet2@b973c', 'AIMNet2@wb97m-d3']
+    element_symbols_available = ['H', 'B', 'C', 'N', 'O', 'F', 'Si', 'P', 'S', 'Cl', 'As', 'Se', 'Br', 'I']
+
+    def __init__(self, method: str = 'AIMNet2@b973c', device = None):
+        import torch
+        if device is None:
+            if torch.cuda.is_available():
+                self.device = torch.device('cuda')
+            else:
+                self.device = torch.device('cpu')
+        else:
+            self.device = torch.device(device)
+        self.model_path = self.parse_aimnet2_resources(method)
+        self.model = torch.jit.load(self.model_path)
+        
+    @doc_inherit
+    def predict(
+            self, 
+            molecular_database = None, 
+            molecule = None,
+            calculate_energy: bool = False,
+            calculate_energy_gradients: bool = False, 
+            calculate_hessian: bool = False,
+            batch_size: int = 2**16,
+        ) -> None:
+        import numpy as np
+
+        molDB = super().predict(molecular_database=molecular_database, molecule=molecule)
+
+        for element_symbol in np.unique(np.concatenate(molDB.element_symbols)):
+            if element_symbol not in self.element_symbols_available:
+                print(f' * Warning * Molecule contains elements \'{element_symbol}\', which is not supported by method \'{self.method}\' that only supports {self.element_symbols_available}, no calculations performed')
+                return
+
+        for mol in molDB:
+            self.predict_for_molecule(
+                molecule=mol, 
+                calculate_energy=calculate_energy, 
+                calculate_energy_gradients=calculate_energy_gradients, 
+                calculate_hessian=calculate_hessian)
+
+    def predict_for_molecule(
+        self,
+        molecule, 
+        calculate_energy: bool = False, 
+        calculate_energy_gradients: bool = False, 
+        calculate_hessian: bool = False):
+        
+        import torch
+
+        coord = torch.as_tensor(molecule.xyz_coordinates).to(torch.float).to(self.device).unsqueeze(0)
+        numbers = torch.as_tensor(molecule.atomic_numbers).to(torch.long).to(self.device).unsqueeze(0)
+        charge = torch.tensor([molecule.charge], dtype=torch.float, device=self.device)
+        nninput = dict(coord=coord, numbers=numbers, charge=charge)
+        prev = torch.is_grad_enabled()
+        torch._C._set_grad_enabled(calculate_energy_gradients)
+        if calculate_energy_gradients:
+            nninput['coord'].requires_grad_(True)
+        nnoutput = self.model(nninput)
+        if calculate_energy:
+            molecule.energy = nnoutput['energy'].item()
+        if calculate_energy_gradients:
+            if 'forces' in nnoutput:
+                f = nnoutput['forces'][0]
+            else:
+                f = - torch.autograd.grad(nnoutput['energy'], nninput['coord'])[0][0]
+            forces = f.detach().cpu().numpy()
+            molecule.add_xyz_vectorial_property(forces, 'energy_gradients')
+
+        if calculate_hessian:
+            print('Hessian not available yet')
+            molecule.hessian = np.zeros((len(molecule),)*2)
+        
+        torch._C._set_grad_enabled(prev)
+
+    @staticmethod
+    def parse_aimnet2_resources(method):
+        import requests, os
+        local_dir = os.path.expanduser('~/.local/AIMNet2/')
+        repo_name = "AIMNet2"
+        tag_name = method.lower().replace('@', '_')
+        url = "https://github.com/isayevlab/{}/raw/old/models/{}_ens.jpt".format(repo_name, tag_name)
+        if not os.path.exists(local_dir+f'{tag_name}_ens.jpt'):
+            os.makedirs(local_dir, exist_ok=True)
+            print(f'Downloading {method} model parameters ...')
+            resource_res = requests.get(url)
+            with open(local_dir+f'{tag_name}_ens.jpt', 'wb') as f:
+                f.write(resource_res.content)
+        return local_dir + f'{tag_name}_ens.jpt'

{mlatom-3.16.2 → mlatom-3.17.0}/src/mlatom/aiqm1.py

@@ -9,13 +9,10 @@
 '''
 import numpy as np
 import os
-from . import data, models, stopper
+from . import data
+from .model_cls import model, torchani_model, method_model, model_tree_node
 
-import torch
-import torchani
-from torchani.utils import ChemicalSymbolsToInts
-
-class aiqm1(models.torchani_model, metaclass=models.meta_method):
+class aiqm1(torchani_model, method_model):
     """
     The Artificial intelligence-quantum mechanical method as in the `AIQM1 paper`_.
 
@@ -33,7 +30,7 @@ class aiqm1(models.torchani_model, metaclass=models.meta_method):
         mol = ml.data.molecule()
         mol.read_from_xyz_file(filename='ethanol.xyz')
         # Run AIQM1 calculation
-        aiqm1 = ml.models.methods(method='AIQM1', qm_program='MNDO')
+        aiqm1 = ml.methods(method='AIQM1', qm_program='MNDO')
         aiqm1.predict(molecule=mol, calculate_energy=True, calculate_energy_gradients=True)
         # Get energy, gradient, and prediction uncertainty of AIQM1 
         energy = mol.energy
@@ -42,27 +39,30 @@ class aiqm1(models.torchani_model, metaclass=models.meta_method):
 
 
     """
+    
+    supported_methods = ['AIQM1', 'AIQM1@DFT', 'AIQM1@DFT*']
     atomic_energies = {'AIQM1': {1:-0.50088038, 6:-37.79221710, 7:-54.53360298, 8:-75.00986203},
                        'AIQM1@DFT': {1:-0.50139362, 6:-37.84623117, 7:-54.59175573, 8:-75.07674376}}
     atomic_energies['AIQM1@DFT*'] = atomic_energies['AIQM1@DFT']
     
     def __init__(self, method='AIQM1', qm_program=None, qm_program_kwargs={}, dftd4_kwargs={}, **kwargs):
+        from .models import methods
         self.method = method.upper()
         self.qm_program = qm_program
         self.qm_program_kwargs = qm_program_kwargs
         modelname = self.method.lower().replace('*','star').replace('@','at')
         ani_nn_children = []
         for ii in range(8):
-            nn_i = models.model_tree_node(name=f'{modelname}_nn{ii}', operator='predict', model=ani_nns_in_aiqm1(method=self.method, model_index=ii))
+            nn_i = model_tree_node(name=f'{modelname}_nn{ii}', operator='predict', model=ani_nns_in_aiqm1(method=self.method, model_index=ii))
             ani_nn_children.append(nn_i)
-        ani_nns = models.model_tree_node(name=f'{modelname}_nn', children=ani_nn_children, operator='average')
-        shift = models.model_tree_node(name=f'{modelname}_atomic_energy_shift', operator='predict', model=atomic_energy_shift(method=self.method))
-        odm2star = models.model_tree_node(name='odm2star', operator='predict', model=models.methods(method='ODM2*', program=qm_program, **qm_program_kwargs))
+        ani_nns = model_tree_node(name=f'{modelname}_nn', children=ani_nn_children, operator='average')
+        shift = model_tree_node(name=f'{modelname}_atomic_energy_shift', operator='predict', model=atomic_energy_shift(method=self.method))
+        odm2star = model_tree_node(name='odm2star', operator='predict', model=methods(method='ODM2*', program=qm_program, **qm_program_kwargs))
         aiqm1_children = [ani_nns, shift, odm2star]
         if self.method != 'AIQM1@DFT*':
-            d4 = models.model_tree_node(name='d4_wb97x', operator='predict', model=models.methods(method='D4', functional='wb97x', **dftd4_kwargs))
+            d4 = model_tree_node(name='d4_wb97x', operator='predict', model=methods(method='D4', functional='wb97x', **dftd4_kwargs))
             aiqm1_children.append(d4)
-        self.aiqm1_model = models.model_tree_node(name=modelname, children=aiqm1_children, operator='sum')
+        self.aiqm1_model = model_tree_node(name=modelname, children=aiqm1_children, operator='sum')
     
     def predict(self, molecular_database=None, molecule=None,
                 calculate_energy=True, calculate_energy_gradients=False, calculate_hessian=False, calculate_dipole_derivatives=False, nstates=1, current_state=0, **kwargs):
@@ -96,7 +96,8 @@ class aiqm1(models.torchani_model, metaclass=models.meta_method):
             molecule.energy = self.atomic_energies[self.method][molecule.atoms[0].atomic_number]
             standard_atom = data.atom(atomic_number=molecule.atoms[0].atomic_number)
             if molecule.charge != 0 or molecule.multiplicity != standard_atom.multiplicity:
-                odm2model = models.methods(method='ODM2*', program=self.qm_program)
+                from .models import methods
+                odm2model = methods(method='ODM2*', program=self.qm_program)
                 mol_odm2 = molecule.copy()
                 odm2model.predict(molecule=mol_odm2, nstates=nstates, **kwargs)
                 mol_standard_odm2 = molecule.copy() ; mol_standard_odm2.charge = 0; mol_standard_odm2.multiplicity=standard_atom.multiplicity
@@ -129,7 +130,7 @@ class aiqm1(models.torchani_model, metaclass=models.meta_method):
                 molecule.__dict__[f'{modelname}_nn'].standard_deviation(properties=properties+atomic_properties)
             
 
-class atomic_energy_shift(models.model):
+class atomic_energy_shift(model):
     atomic_energy_shifts = {'AIQM1': {1: -4.29365862e-02, 6: -3.34329586e+01, 7: -4.69301173e+01, 8: -6.29634763e+01},
                             'AIQM1@DFT': {1: -4.27888067e-02, 6: -3.34869833e+01, 7: -4.69896148e+01, 8: -6.30294433e+01}}
     atomic_energy_shifts['AIQM1@DFT*'] = atomic_energy_shifts['AIQM1@DFT']
@@ -166,10 +167,11 @@ class atomic_energy_shift(models.model):
                     ndim = len(mol.atoms) * 3
                     mol.hessian = np.zeros(ndim*ndim).reshape(ndim,ndim)
 
-class ani_nns_in_aiqm1(models.torchani_model):
+class ani_nns_in_aiqm1(torchani_model):
     species_order = [1, 6, 7, 8]
     
     def __init__(self, method='AIQM1', model_index = 0):
+        import torch
         if method == 'AIQM1':
             self.level = 'cc'
         elif method in ['AIQM1@DFT', 'AIQM1@DFT*']:
@@ -180,6 +182,7 @@ class ani_nns_in_aiqm1(models.torchani_model):
         self.load_model()
     
     def define_aev(self):
+        import torch
         Rcr = 5.2000e+00
         Rca = 4.0000e+00
         EtaR = torch.tensor([1.6000000e+01], device=self.device)
@@ -189,19 +192,23 @@ class ani_nns_in_aiqm1(models.torchani_model):
         EtaA = torch.tensor([8.0000000e+00], device=self.device)
         ShfA = torch.tensor([9.0000000e-01, 1.6750000e+00,  2.4499998e+00, 3.2250000e+00], device=self.device)
         num_species = len(self.species_order)
+        import torchani
         aev_computer = torchani.AEVComputer(Rcr, Rca, EtaR, ShfR, EtaA, Zeta, ShfA, ShfZ, num_species)
         self.aev_computer = aev_computer
 
     def load_model(self):
+        import torch
         mlatomdir=os.path.dirname(__file__)
         dirname = os.path.join(mlatomdir, 'aiqm1_model')
         method = 'aiqm1_' + self.level
         self.define_nn()
         checkpoint = torch.load(os.path.join(dirname, f'{method}_cv{self.model_index}.pt'), map_location=self.device)
         self.nn.load_state_dict(checkpoint['nn'])
+        import torchani
         self.model  = torchani.nn.Sequential(self.aev_computer, self.nn).to(self.device).double()
 
     def define_nn(self):
+        import torch
         aev_dim = self.aev_computer.aev_length
         H_network = torch.nn.Sequential(
             torch.nn.Linear(aev_dim, 160),
@@ -243,6 +250,7 @@ class ani_nns_in_aiqm1(models.torchani_model):
             torch.nn.Linear(96, 1)
         )
         
+        import torchani
         nn = torchani.ANIModel([H_network, C_network, N_network, O_network])
         self.nn = nn
     
@@ -250,6 +258,8 @@ class ani_nns_in_aiqm1(models.torchani_model):
                 calculate_energy=True, calculate_energy_gradients=False, calculate_hessian=False, nstates=1, **kwargs):
         molDB = super().predict(molecular_database=molecular_database, molecule=molecule)
         
+        from torchani.utils import ChemicalSymbolsToInts
+        
         calculate_energy_gradients = bool(np.array(calculate_energy_gradients).any())
         calculate_hessian = bool(np.array(calculate_hessian).any())
         species_to_tensor = ChemicalSymbolsToInts(self.species_order)
@@ -264,6 +274,7 @@ class ani_nns_in_aiqm1(models.torchani_model):
                     mol.electronic_states.append(mol_copy.copy())
                 molecules = mol.electronic_states
             
+            import torch
             for mol in molecules:
                 atomic_numbers = np.array([atom.atomic_number for atom in mol.atoms])
                 xyz_coordinates = torch.tensor(np.array(mol.xyz_coordinates).astype('float')).to(self.device).requires_grad_(calculate_energy_gradients or calculate_hessian)
@@ -278,6 +289,7 @@ class ani_nns_in_aiqm1(models.torchani_model):
                         for iatom in range(len(mol.atoms)):
                             mol.atoms[iatom].energy_gradients = grads[iatom]
                 if calculate_hessian:
+                    import torchani
                     ANI_NN_hessian = torchani.utils.hessian(xyz_coordinates, energies=ANI_NN_energy)
                     mol.hessian = ANI_NN_hessian[0].detach().cpu().numpy()
 

{mlatom-3.16.2 → mlatom-3.17.0}/src/mlatom/aiqm2.py

@@ -1,9 +1,8 @@
+import os
 from . import data, models, constants
-from multiprocessing import cpu_count
-from .interfaces.torchani_interface import ani
-import os 
+from .model_cls import method_model, model_tree_node
 
-class aiqm2(models.model, metaclass=models.meta_method):
+class aiqm2(method_model):
 
     """ 
     GFN2-xTB based artificial intelligence quantum-mechanical method 2 (AIQM2)
@@ -20,7 +19,7 @@ class aiqm2(models.model, metaclass=models.meta_method):
         mol = ml.data.molecule()
         mol.read_from_xyz_file(filename='ethanol.xyz')
         # Run AIQM2 calculation
-        aiqm2 = ml.models.methods(method='aiqm2')
+        aiqm2 = ml.methods(method='aiqm2')
         aiqm2.predict(molecule=mol, calculate_energy=True, calculate_energy_gradients=True, calculate_hessian=True)
         # Get energy, gradient, and uncertainty of AIQM2 
         energy = mol.energy
@@ -29,7 +28,8 @@ class aiqm2(models.model, metaclass=models.meta_method):
         std = mol.aiqm2_model.energy_standard_deviation
 
     """ 
-
+    
+    supported_methods = ['AIQM2', 'AIQM2@DFT', 'AIQM2@DFT*']
     mlatomdir=os.path.dirname(__file__)
     dirname = os.path.join(mlatomdir, 'aiqm2_model')
 
@@ -37,7 +37,8 @@ class aiqm2(models.model, metaclass=models.meta_method):
         self,
         method: str = 'AIQM2',
         working_directory: str = '.',
-        qm_program_kwargs: dict = {}
+        qm_program_kwargs: dict = {},
+        nthreads: int = 1
     ):
 
         self.method = method.lower()
@@ -45,6 +46,16 @@ class aiqm2(models.model, metaclass=models.meta_method):
         self.working_directory = working_directory
         self.qm_program_kwargs = qm_program_kwargs
         self.load()
+        self.nthreads = nthreads
+
+    @property
+    def nthreads(self):
+        return self._nthreads
+
+    @nthreads.setter
+    def nthreads(self, value):
+        self._nthreads = value
+        self.aiqm2_model.nthreads = self._nthreads
 
     def predict(
         self, 
@@ -85,27 +96,37 @@ class aiqm2(models.model, metaclass=models.meta_method):
         molecule.__dict__[f'{self.model_name}_nn'].standard_deviation(properties=['energy'])
 
     def load(self):
+        from .models import methods
         if not os.path.exists(os.path.join(self.dirname, f'{self.model_name}_cv0.pt')):
             self.download_models()
         
         model_paths = [os.path.join(self.dirname, f'{self.model_name}_cv{ii}.pt') for ii in range(8)]
 
-        baseline = models.model_tree_node(
+        baseline = model_tree_node(
             name='gfn2xtbstar',
-            model=models.methods(method='GFN2-xTB*', **self.qm_program_kwargs),
+            model=methods(method='GFN2-xTB*', **self.qm_program_kwargs),
             operator='predict'
         )
-        d4 = models.model_tree_node(
+        d4 = model_tree_node(
             name='d4wb97x',
-            model=models.methods(method='D4', functional='wb97x', working_directory=self.working_directory),
+            model=methods(method='D4', functional='wb97x', working_directory=self.working_directory),
             operator='predict'
         )
-        nn = models.model_tree_node(
+        from .interfaces.torchani_interface import ani
+        class ani_wrapper(ani):
+            def __init__(self,**kwargs):
+                super().__init__(**kwargs)
+
+            def predict(self,**kwargs):
+                if 'calculate_dipole_derivatives' in kwargs.keys():
+                    del kwargs['calculate_dipole_derivatives']
+                super().predict(**kwargs)
+        nn = model_tree_node(
             name=f'{self.model_name}_nn',
             children=[
-                models.model_tree_node(
+                model_tree_node(
                     name=f'{self.model_name}_nn_{ii}',
-                    # model=models.ani(
+                    # model=ani(
                     #     model_file=model_paths[ii],
                     #     verbose=0),
                     model=ani_wrapper(
@@ -117,7 +138,7 @@ class aiqm2(models.model, metaclass=models.meta_method):
             operator='average'
         )
 
-        self.aiqm2_model = models.model_tree_node(
+        self.aiqm2_model = model_tree_node(
             name=self.model_name,
             children=[baseline, nn, d4],
             operator='sum'
@@ -134,11 +155,4 @@ class aiqm2(models.model, metaclass=models.meta_method):
             with open(f'{self.dirname}/{self.model_name}_cv{ii}.pt','wb') as f:
                 f.write(resource_res.content)
 
-class ani_wrapper(ani):
-    def __init__(self,**kwargs):
-        super().__init__(**kwargs)
 
-    def predict(self,**kwargs):
-        if 'calculate_dipole_derivatives' in kwargs.keys():
-            del kwargs['calculate_dipole_derivatives']
-        super().predict(**kwargs)