mlatom 3.4.0b0__tar.gz → 3.6.0__tar.gz

{mlatom-3.4.0b0/src/mlatom.egg-info → mlatom-3.6.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: mlatom
-Version: 3.4.0b0
+Version: 3.6.0
 Summary: A Package for Atomistic Simulations with Machine Learning
 Home-page: http://mlatom.com
 Author: Pavlo O. Dral
@@ -27,7 +27,7 @@ Requires-Dist: tqdm
 # About Program
 MLatom: a Package for Atomistic Simulations with Machine Learning
 
-Version 3.4.0
+Version 3.6.0
 
 **Official website**: http://mlatom.com/    
 **Manual**: http://mlatom.com/manual/  
@@ -60,7 +60,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, Max Pinheiro Jr, Yuming Su, Yiheng Dai, Yangtao Chen, Shuang Zhang, Lina Zhang, Arif Ullah, Quanhao Zhang, Sebastian V. Pios, Yanchi Ou, MLatom: A Package for Atomistic Simulations with Machine Learning,  version 3.4.0, Xiamen University, Xiamen, China, 2013-2024.
+* Pavlo O. Dral, Fuchun Ge, Yi-Fan Hou, Peikun Zheng, Yuxinxin Chen, Bao-Xin Xue, Max Pinheiro Jr, Yuming Su, Yiheng Dai, Yangtao Chen, Shuang Zhang, Lina Zhang, Arif Ullah, Quanhao Zhang, Sebastian V. Pios, Yanchi Ou, MLatom: A Package for Atomistic Simulations with Machine Learning,  version 3.6.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.4.0b0 → mlatom-3.6.0}/README.md

@@ -1,7 +1,7 @@
 # About Program
 MLatom: a Package for Atomistic Simulations with Machine Learning
 
-Version 3.4.0
+Version 3.6.0
 
 **Official website**: http://mlatom.com/    
 **Manual**: http://mlatom.com/manual/  
@@ -34,7 +34,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, Max Pinheiro Jr, Yuming Su, Yiheng Dai, Yangtao Chen, Shuang Zhang, Lina Zhang, Arif Ullah, Quanhao Zhang, Sebastian V. Pios, Yanchi Ou, MLatom: A Package for Atomistic Simulations with Machine Learning,  version 3.4.0, Xiamen University, Xiamen, China, 2013-2024.
+* Pavlo O. Dral, Fuchun Ge, Yi-Fan Hou, Peikun Zheng, Yuxinxin Chen, Bao-Xin Xue, Max Pinheiro Jr, Yuming Su, Yiheng Dai, Yangtao Chen, Shuang Zhang, Lina Zhang, Arif Ullah, Quanhao Zhang, Sebastian V. Pios, Yanchi Ou, MLatom: A Package for Atomistic Simulations with Machine Learning,  version 3.6.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.4.0b0 → mlatom-3.6.0}/setup.py

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

{mlatom-3.4.0b0 → mlatom-3.6.0}/src/mlatom/MLatom.py

@@ -4,7 +4,7 @@
   !---------------------------------------------------------------------------!
   !                                                                           !
   !     MLatom: a Package for Atomistic Simulations with Machine Learning     !
-  !                             MLatom 3.4.0                                  !
+  !                             MLatom 3.6.0                                  !
   !                                   @                                       !
   !                 Xiamen Atomistic Computing Suite (XACS)                   !
   !                                                                           !
@@ -36,7 +36,7 @@
   ! Xue, Max Pinheiro Jr, Yuming Su, Yiheng Dai, Yangtao Chen, Shuang Zhang,  !
   ! Lina Zhang, Arif Ullah, Quanhao Zhang, Sebastian V. Pios, Yanchi Ou,      !
   ! MLatom: A Package for Atomistic Simulations with Machine Learning,        !
-  ! version 3.4.0, Xiamen University, Xiamen, China, 2013-2024.               !
+  ! version 3.6.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.              !

{mlatom-3.4.0b0 → mlatom-3.6.0}/src/mlatom/MLtasks.py

@@ -447,6 +447,7 @@ def freq(args):
         geomopt = simulations.thermochemistry(model=model,
                                         molecule=mol,
                                         **kwargs)
+        mol.dump(f'freq{mol.number}.json',format='json')
         print(' %s ' % ('='*78))
         print(' %s Vibration analysis for molecule %6d' % (' '*20, imol+1))
         print(' %s ' % ('='*78))

{mlatom-3.4.0b0 → mlatom-3.6.0}/src/mlatom/args_class.py

@@ -462,6 +462,9 @@ class mlatom_args(ArgsBase):
         self.set_keyword_alias('ANI-2x', ['ani2x'])
         self.set_keyword_alias('ANI-1x-D4', ['ani1xd4'])
         self.set_keyword_alias('ANI-2x-D4', ['ani2xd4'])
+        self.set_keyword_alias('ANI-1xnr', ['ani1xnr'])
+        self.set_keyword_alias('AIMNet2@B973c', ['aimnet2atb973c'])
+        self.set_keyword_alias('AIMNet2@wb97M-D3', ['aimnet2atwb97md3'])
         self.set_keyword_alias('GFN2-xTB', ['gfn2xtb'])
         self.set_keyword_alias('MNDO/dH', ['mndodh'])
         self.set_keyword_alias('MNDO/H', ['mndoh'])

{mlatom-3.4.0b0 → mlatom-3.6.0}/src/mlatom/data.py

@@ -1486,6 +1486,7 @@ class molecular_trajectory():
             moldb.write_file_with_properties(filename+'.ekin',property_to_write='kinetic_energy')
             moldb.write_file_with_properties(filename+'.epot',property_to_write='energy')
             moldb.write_file_with_properties(filename+'.etot',property_to_write='total_energy')
+            moldb.write_file_with_properties(filename+'.temp',property_to_write='temperature')
             if 'dipole_moment' in moldb.molecules[0].__dict__.keys():
                 with open(filename+'.dp','w') as dpf:
                     for imolecule in moldb.molecules:

{mlatom-3.4.0b0 → mlatom-3.6.0}/src/mlatom/doc.py

@@ -179,6 +179,9 @@ class Doc():
       ANI-2x                   perform ANI-2x      calculations
       ANI-1x-D4                perform ANI-1x-D4   calculations
       ANI-2x-D4                perform ANI-2x-D4   calculations
+      ANI-1xnr                 perform ANI-1xnr    calculations
+      AIMNet2@B973c            perform AIMNet2@B97-3c calculations
+      AIMNet2@wB97M-D3         perform AIMNet2@wB97M-D3 calculations
       ODM2
       ODM2*
       GFN2-xTB

{mlatom-3.4.0b0 → mlatom-3.6.0}/src/mlatom/header.py

@@ -128,6 +128,18 @@ def printHeader(args):
         if args.ani2xd4:
             refItems['D4']  = refs['D4']
             refItems['D4 program']  = refs['D4prog']
+
+    if args.method:
+        if 'DM21' in args.method.upper():
+            refItems['DM21'] = refs['DM21']
+    
+    if args.ani1xnr:
+        refItems['ANI-1xnr'] = refs['ANI-1xnr']
+    if args.aimnet2atb973c:
+        refItems['AIMNet2@B973c'] = refs['AIMNet2']
+    if args.aimnet2atwb97md3:
+        refItems['AIMNet2@wb97M-D3'] = refs['AIMNet2']
+
     
     if args.geomopt:
         if 'optprog=ase' in ' '.join(args.args2pass).lower():

{mlatom-3.4.0b0 → mlatom-3.6.0}/src/mlatom/initial_conditions.py

@@ -202,7 +202,7 @@ def generate_initial_conditions(molecule=None, generation_method=None, number_of
             mass = mass_.reshape(Natoms,1)
             total_mass = np.sum(mass_)
             for irepeat in range(number_of_initial_conditions):
-                if eliminate_angular_momentum:
+                if eliminate_angular_momentum and not molecule.is_it_linear():
                     velocities_all[irepeat] = getridofang(coordinates_all[irepeat],velocities_all[irepeat],mass_)
                     v_cm = sum(velocities_all[irepeat]*mass)/total_mass
                     velocities_all[irepeat] -= v_cm
@@ -212,7 +212,9 @@ def generate_initial_conditions(molecule=None, generation_method=None, number_of
                     new_molecule.atoms[iatom].xyz_coordinates = coordinates_all[irepeat][iatom]
                     new_molecule.atoms[iatom].xyz_velocities = velocities_all[irepeat][iatom]
                 init_cond_db.molecules.append(new_molecule)
-        
+        elif generation_method.casefold() == 'harmonic-quantum-boltzmann'.casefold():
+            init_cond_db = harmonic_quantum_Boltzmann_sampling.sample(npoints=number_of_initial_conditions,molecule=molecule,temperature=initial_temperature,use_hessian=use_hessian)
+
         if filter_by_energy_window:
             if iteration == 1:
                 result = excitation_energy_window_filter(init_cond_db,**window_filter_kwargs)
@@ -247,6 +249,10 @@ def generate_initial_conditions(molecule=None, generation_method=None, number_of
     if len(init_cond_db) > target_number_of_initial_conditions:
        init_cond_db = init_cond_db[:target_number_of_initial_conditions]
 
+    # Change the random seed so that the user-defined one only affects initial conditions sampling
+    if not random_seed is None:
+        np.random.seed()
+
     return init_cond_db
 
 def read_velocities_from_file(filename):
@@ -573,6 +579,168 @@ class wigner_sampling():
                     q[iatom][icoord] += cn[i][iatom][icoord]*cint[i]*fac 
                     v[iatom][icoord] += cn[i][iatom][icoord]*dcint[i]*fac 
         return q,v # Unit: a.u.
+
+class harmonic_quantum_Boltzmann_sampling():
+    #
+    # The harmonic quantum Boltzmann sampling in this class follows what is shown in the following paper (it is called thermal sampling in VENUS manual):
+    # J. Phys. Chem. A 1998, 102, 3648-3658
+    #
+    # The quanta ni of each normal mode is first sampled from the harmonic quantum Boltzman distribution function:
+    #
+    #   p(ni) = exp(-ni * h * vi / kB / T) * (1 - (exp(-h * vi / kB / T))) 
+    #
+    # where vi is the frequency of normal mode i, h is the Planck constant, kB is the Boltzmann constant and T is the temperature
+    # 
+    # The energy of normal mode i is calculated as Ei = (ni + 0.5) * h * vi
+    #
+    # The mass weighted normal mode coordinates Qi and momenta Pi are
+    #
+    #   Qi = Ai * cos(2 * pi * Ri)
+    #   Pi = -Ai * wi * sin(2 * pi * Ri)
+    #
+    # where wi = 2 * pi * vi, Ai = sqrt(2 * Ei) / wi and Ri is a uniform random number on [0,1]
+    #
+    # The mass weighted momentum Prc of reaction coordinate is chosed from a thermal distribution 
+    # 
+    #   Prc = ± sqrt(-2 * kB * T * ln(1 - R))
+    #
+    # where R is a uniform random number on [0,1]
+    #
+
+    def __init__(self):
+        pass 
+
+    @classmethod 
+    def sample(cls,npoints,molecule,temperature,use_hessian):
+        if not use_hessian:
+            if not ('frequencies' in molecule.__dict__ and 'normal_modes' in molecule[0].__dict__):
+                print('Frequencies and normal modes not found, try to calculate them from Hessian matrix')
+                if not 'hessian' in molecule.__dict__.keys():
+                    stopper.stopMLatom('Hessian matrix not found -- cannot do wigner sampling')
+                simulations.freq.freq_modified_from_TorchANI(molecule=molecule,normal_mode_normalization='mass deweighted normalized')
+        else:
+            if not 'hessian' in molecule.__dict__.keys():
+                stopper.stopMLatom('Hessian matrix not found -- cannot do wigner sampling')
+            simulations.freq.freq_modified_from_TorchANI(molecule=molecule,normal_mode_normalization='mass deweighted normalized')
+        freq = molecule.frequencies 
+        nm_masses = molecule.reduced_masses
+        Natoms = len(molecule)
+        
+        nnegative = 0 
+        while freq[nnegative] < 0:
+            nnegative += 1 
+        # freq = freq[nnegative:]
+        nm = np.zeros((len(freq),Natoms,3))
+
+        for imode in range(len(freq)):
+            for iatom in range(Natoms):
+                nm[imode][iatom] = molecule[iatom].normal_modes[imode]
+            nm[imode] /= np.sqrt(np.sum(nm[imode]**2))
+
+        q_list = np.array([molecule.xyz_coordinates]*npoints)
+        v_list = np.zeros((npoints,Natoms,3))
+
+        for imode in range(nnegative,len(freq)):
+            qq,vv = cls.get_vq(temperature,freq[imode],nm_masses[imode],npoints)
+            for isample in range(npoints):
+                # print(isample)
+                q_list[isample] += nm[imode]*qq[isample]
+                # print(nm[imode]*vv[isample])
+                v_list[isample] += nm[imode]*vv[isample]
+
+        # Deal with reaction coordinate
+        for ii in range(nnegative):
+            for isample in range(npoints):
+                scale = np.random.choice([-1,1])*np.sqrt(-2*constants.kB_in_Hartree*temperature*np.log(1-np.random.random())) / np.sqrt(nm_masses[ii] * constants.ram2au)/ constants.au2fs * constants.Bohr2Angstrom
+                v_list[isample] += scale * nm[ii]
+
+        init_cond_db = data.molecular_database() 
+        for isample in range(npoints):
+            new_molecule = molecule.copy(atomic_labels=[],molecular_labels=[])
+            for iatom in range(Natoms):
+                new_molecule.atoms[iatom].xyz_coordinates = q_list[isample][iatom]
+                new_molecule.atoms[iatom].xyz_velocities = v_list[isample][iatom]
+            init_cond_db.molecules.append(new_molecule)
+
+        return init_cond_db
+    
+    @classmethod 
+    def sample_quanta(cls,temperature,freq,nsample):
+        hvkT = constants.planck_constant*freq*constants.speed_of_light*100/constants.kB/temperature
+        
+        nn_len = 100
+        nn = np.array([ii for ii in range(nn_len)])
+        pp = np.array([np.exp(-ii*hvkT)*(1-np.exp(-hvkT)) for ii in nn])
+        while 1.0-np.sum(pp) > 1e-10:
+            nn_len += 50
+            nn = np.array([ii for ii in range(nn_len)])
+            pp = np.array([np.exp(-ii*hvkT)*(1-np.exp(-hvkT)) for ii in nn])
+        rand = np.random.choice(nn,size=nsample,p=pp)
+        return rand
+    
+    @classmethod 
+    def get_energy(cls,rand,freq):
+        return (rand+0.5)*constants.planck_constant*freq*constants.speed_of_light*100/1000*constants.Avogadro_constant*constants.kJpermol2Hartree # Hartree
+    
+    @classmethod 
+    def get_vq(cls,temperature,freq,mass,nsample):
+        energy = cls.get_energy(cls.sample_quanta(temperature,freq,nsample),freq)
+        rand = np.random.random(len(energy))
+        omega = freq * constants.speed_of_light*100*2*np.pi / 1.0E15 / constants.fs2au # au
+        AA = np.sqrt(2*energy)/omega # au
+        qq = AA * np.cos(2*np.pi*rand) / np.sqrt(mass*constants.ram2au) * constants.Bohr2Angstrom # Angstrom
+        vv = -omega * AA * np.sin(2*np.pi*rand) / np.sqrt(mass*constants.ram2au) * constants.Bohr2Angstrom / constants.au2fs # Angstrom/fs 
+        return qq,vv
+
+# def wignersample(npoints,molecule):
+#     qlist = [] 
+#     vlist = []
+#     geomEq = np.array([each.xyz_coordinates for each in molecule.atoms])
+#     mass = molecule.get_nuclear_masses() 
+#     mass = mass.reshape(1,len(mass))
+#     # Calculate normal modes from Hessian matrix 
+#     #nm,freq,ele,linear_int = readGaussianNM(nmfile)
+#     if not 'hessian' in molecule.__dict__.keys():
+#         stopper.stopMLatom('Hessian matrix not found -- cannot do wigner sampling')
+
+#     linear = molecule.is_it_linear()
+#     Natoms = len(molecule.atoms)
+#     if linear:
+#         ntriv = 5 
+#         linear_int = 1
+#     else:
+#         ntriv = 6
+#         linear_int = 0
+
+#     # freq,nm,_,_ = vibrational_analysis(mass,molecule.hessian,mode_type='MDU')
+#     simulations.freq.freq_modified_from_TorchANI(molecule=molecule,normal_mode_normalization='mass deweighted unnormalized')
+#     freq = molecule.frequencies 
+#     nm = np.zeros(3*Natoms,Natoms,3)
+#     for itriv in range(ntriv):
+#         nm[itriv] = 1.0 / np.sqrt(3*Natoms)
+#     for imode in range(ntriv,3*Natoms):
+#         for iatom in range(Natoms):
+#             nm[ntriv+imode][iatom] = molecule[iatom].normal_modes[imode]
+
+#     numcoo = len(freq)
+#     #print(len(nm),len(freq),len(ele),linear)
+#     atom_mass = [each.nuclear_mass for each in molecule.atoms]
+#     cart_nms, w_nmode = nm2cart(nm,atom_mass)
+
+#     geomEq_au = np.array(geomEq) / constants.Bohr2Angstrom
+
+#     anq, amp, freq_au, atmau = rdmol(atom_mass,linear_int,freq) 
+
+#     for ipoint in range(npoints):
+#         q,v = inqp(geomEq_au,w_nmode[ntriv:],freq_au,anq,amp,numcoo,atmau)
+#         qlist.append(q)
+#         vlist.append(v)
+    
+#     # Transform from a.u. to Angstrom & fs
+#     qlist = np.array(qlist) * constants.Bohr2Angstrom
+#     vlist = np.array(vlist) * constants.Bohr2Angstrom / constants.au2fs
+
+#     return qlist, vlist, linear_int
     
 # generate random velocities without angular momentum for linear molecule
 def generate_random_velocities_for_linear_molecule(molecule):

{mlatom-3.4.0b0 → mlatom-3.6.0}/src/mlatom/interfaces/__init__.py

@@ -10,6 +10,10 @@ def ani(**kwargs):
     from .torchani_interface import ani_methods as interface
     return interface(**kwargs)
 
+def aimnet2(**kwargs):
+    from .torchani_interface import aimnet2_methods as interface 
+    return interface(**kwargs)
+
 def mndo(**kwargs):
     from .mndo_interface import mndo_methods as interface
     return interface(**kwargs)

{mlatom-3.4.0b0 → mlatom-3.6.0}/src/mlatom/interfaces/gaussian_interface.py

@@ -186,13 +186,18 @@ def run_gaussian_job(**kwargs):
         memory = f"%mem={kwargs['memory']}\n"
     gaussian_keywords = f'{memory}%nproc={nthreads}\n' + gaussian_keywords 
     
+    if 'cwd' in kwargs:
+        cwd = kwargs['cwd']
+    else:
+        cwd='.'
+        
     if 'model_predict_kwargs' in kwargs:
         model_predict_kwargs_str = str(kwargs['model_predict_kwargs'])
     else:
         model_predict_kwargs_str = "{}"
     
     model_predict_kwargs_str_file = 'model_predict_kwargs'
-    with open(model_predict_kwargs_str_file, 'w') as f:
+    with open(os.path.join(cwd, model_predict_kwargs_str_file), 'w') as f:
         f.write(model_predict_kwargs_str)
     
     if 'external_task' in kwargs:
@@ -227,11 +232,6 @@ def run_gaussian_job(**kwargs):
                 kwargs['method'] = 'p ' + kwargs['method']
             gaussian_keywords += '# '+kwargs['method']+'\n'
 
-    if 'cwd' in kwargs:
-        cwd = kwargs['cwd']
-    else:
-        cwd='.'
-
     if 'writechk' in kwargs:
         writechk = kwargs['writechk']
     else:

{mlatom-3.4.0b0 → mlatom-3.6.0}/src/mlatom/interfaces/pyscf_interface.py

@@ -72,6 +72,10 @@ class pyscf_methods(OMP_pyscf):
         pyscf_mol.unit = 'Ang'
         pyscf_mol.build()
 
+        # DM21
+        if 'DM21' in self.method.upper():
+            self.predict_for_molecule_DM21(molecule=molecule, pyscf_mol=pyscf_mol, calculate_energy=calculate_energy, calculate_energy_gradients=calculate_energy_gradients, calculate_hessian=calculate_hessian, **kwargs)
+            return
 
         # HF
         if 'HF' == self.method.upper():
@@ -193,8 +197,50 @@ class pyscf_methods(OMP_pyscf):
                     h[ii*3:(ii+1)*3, jj*3:(jj+1)*3] = hess[ii][jj]
             molecule.hessian = h / constants.Bohr2Angstrom**2
 
-       
-    
+    def predict_for_molecule_DM21(self, molecule=None, pyscf_mol=None, calculate_energy=True, calculate_energy_gradients=False, calculate_hessian=False, **kwargs):
+        # reference: https://github.com/google-deepmind/deepmind-research/tree/f5de0ede8430809180254ee957abf36ed62579ef/density_functional_approximation_dm21
+        # METHODS AVAILABE:
+        # DM21 - trained on molecules dataset, and fractional charge, and fractional spin constraints.
+        # DM21m - trained on molecules dataset.
+        # DM21mc - trained on molecules dataset, and fractional charge constraints.
+        # DM21mu - trained on molecules dataset, and electron gas constraints.
+
+        from pyscf import dft
+        try:
+            import density_functional_approximation_dm21 as dm21
+        except:
+            errmsg = 'Please install required packages for DM21. For more details, please refer to https://github.com/google-deepmind/deepmind-research/tree/master/density_functional_approximation_dm21.'
+            raise ModuleNotFoundError(errmsg)
+        
+        if pyscf_mol.spin == 0:
+            pyscf_method = dft.RKS(pyscf_mol)
+        else:
+            pyscf_method = dft.UKS(pyscf_mol)
+
+        pyscf_method.xc = 'B3LYP'
+        pyscf_method.run()
+        dm0 = pyscf_method.make_rdm1()
+        pyscf_method._numint = dm21.NeuralNumInt(dm21.Functional.__dict__[self.method])
+        # relax convergence tolerances to increase success to convergence
+        pyscf_method.conv_tol = 1E-6
+        pyscf_method.conv_tol_grad = 1E-3
+        try:
+            pyscf_method.kernel(dm0=dm0)
+        except:
+            errmsg = 'DM21 cannot converge properly'
+            stopper.stopMLatom(errmsg)
+                
+        if calculate_energy:
+            molecule.energy = pyscf_method.e_tot
+
+        if calculate_energy_gradients:           
+            errmsg = 'DM21 by pyscf does not support gradients calculation'
+            raise ValueError(errmsg)
+
+        if calculate_hessian:
+            errmsg = 'DM21 by pyscf does not support hessian calculation'
+            raise ValueError(errmsg)
+
     @doc_inherit
     def predict(self, molecule=None, molecular_database=None, calculate_energy=True, calculate_energy_gradients=False, calculate_hessian=False, **kwargs):
         '''

{mlatom-3.4.0b0 → mlatom-3.6.0}/src/mlatom/interfaces/torchani_interface.py

@@ -15,8 +15,11 @@ import numpy as np
 import tqdm
 from collections import OrderedDict
 import torch
+from torch import nn, Tensor
 import torchani
 from torchani.data import TransformableIterable, IterableAdapter
+from typing import Dict, List
+
 
 from .. import data
 from .. import models
@@ -649,7 +652,9 @@ class ani_methods(models.torchani_model):
 
     def model_setup(self, method):
         self.method = method
-        if 'ANI-1x'.casefold() in method.casefold():
+        if 'ANI-1xnr'.casefold() in method.casefold():
+            self.model = load_ani1xnr_model()
+        elif 'ANI-1x'.casefold() in method.casefold():
             self.model = torchani.models.ANI1x(periodic_table_index=True).to(self.device).double()
         elif 'ANI-1ccx'.casefold() in method.casefold():
             self.model = torchani.models.ANI1ccx(periodic_table_index=True).to(self.device).double()
@@ -726,6 +731,132 @@ class ani_methods(models.torchani_model):
             modelname = self.method.lower().replace('-','')
         molecule.__dict__[f'{modelname}'].standard_deviation(properties=properties+atomic_properties)
 
+def load_ani1xnr_model():
+    # ANI-1xnr https://github.com/atomistic-ml/ani-1xnr/
+    # Universal reactive ML potential ANI-1xnr: https://doi.org/10.1038/s41557-023-01427-3
+    species = ['H', 'C', 'N', 'O']
+    def parse_ani1xnr_resources():
+        import requests
+        import zipfile
+        import io
+        local_dir = os.path.expanduser('~/.local/ANI1xnr/')
+        url = "https://github.com/atomistic-ml/ani-1xnr/archive/refs/heads/main.zip"
+        if not os.path.exists(local_dir+'ani-1xnr-main'):
+            os.makedirs(local_dir, exist_ok=True)
+            print(f'Downloading ANI-1xnr model parameters ...')
+            resource_res = requests.get(url)
+            resource_zip = zipfile.ZipFile(io.BytesIO(resource_res.content))
+            resource_zip.extractall(local_dir)
+        return local_dir
+    model_prefix = parse_ani1xnr_resources() + 'ani-1xnr-main/model/ani-1xnr/'
+    # const_file, sae_file, ensemble_prefix, ensemble_size = parse_ani1xnr_resources()
+    const_file = model_prefix + 'rHCNO-5.2R_32-3.5A_a8-4.params'
+    sae_file = model_prefix + 'sae_linfit.dat'
+    consts = torchani.neurochem.Constants(const_file)
+    aev_computer = torchani.AEVComputer(**consts)
+    energy_shifter = torchani.neurochem.load_sae(sae_file)
+    species_converter = torchani.nn.SpeciesConverter(species)
+
+    neural_networks = []
+    for ii in range(8):
+        neural_network = torchani.neurochem.load_model(consts.species, model_prefix + f'train{ii}/networks')
+        neural_networks.append(torchani.nn.Sequential(species_converter, aev_computer, neural_network, energy_shifter))
+
+    model_ensemble = torchani.nn.Ensemble(modules=neural_networks)
+    model_ensemble.species = species
+
+    return model_ensemble
+
+class aimnet2_methods(models.torchani_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.
+
+    '''
+    
+    available_methods = models.methods.methods_map['aimnet2']
+    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: str = 'cuda' if torch.cuda.is_available() else 'cpu', **kwargs):
+        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: data.molecular_database = None, 
+            molecule: data.molecule = None,
+            calculate_energy: bool = False,
+            calculate_energy_gradients: bool = False, 
+            calculate_hessian: bool = False,
+            batch_size: int = 2**16,
+        ) -> None:
+
+        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):
+
+        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
+        local_dir = os.path.expanduser('~/.local/AIMNet2/')
+        repo_name = "AIMNet2"
+        tag_name = method.lower().replace('@', '_')
+        url = "https://github.com/isayevlab/{}/raw/main/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'
+
 def printHelp():
     helpText = __doc__.replace('.. code-block::\n\n', '') + '''
   To use Interface_ANI, please install TorchANI and its dependencies

{mlatom-3.4.0b0 → mlatom-3.6.0}/src/mlatom/md.py

@@ -141,6 +141,12 @@ class md():
         
         self.stop_function = stop_function
         self.stop_function_kwargs = stop_function_kwargs
+
+        self.linearity = self.molecule_with_initial_conditions.is_it_linear()
+        if self.linearity:
+            self.degrees_of_freedom = 3 * self.Natoms - 5
+        else:
+            self.degrees_of_freedom = 3 * self.Natoms - 6
         self.propagate()
         
     def propagate(self):
@@ -200,6 +206,7 @@ class md():
             velocity = np.copy(molecule.get_xyz_vectorial_properties('xyz_velocities'))
             
             molecule.total_energy = molecule.energy + molecule.kinetic_energy
+            molecule.temperature = molecule.kinetic_energy / (constants.kB_in_Hartree*self.degrees_of_freedom/2)
             trajectory_step.step = istep 
             trajectory_step.time = istep * self.time_step
             trajectory_step.molecule = molecule