mlatom 2.0.1__py3.7.egg

MLatom/interfaces/TorchANI/interface_TorchANI.py

@@ -0,0 +1,323 @@
+#!/usr/bin/python3
+'''
+  !---------------------------------------------------------------------------! 
+  ! Interface_TorchANI: Interface between TorchANI and MLatom                 ! 
+  ! Implementations by: Fuchun Ge and Max Pinheiro Jr                         ! 
+  !---------------------------------------------------------------------------! 
+'''
+import numpy as np
+import os, sys, subprocess, time, shutil, re, math, random
+import stopper
+from args_class import ArgsBase
+
+filedir = os.path.dirname(__file__)
+
+class Args(ArgsBase):
+    def __init__(self):
+        super().__init__()
+        self.add_default_dict_args([
+            'learningCurve'
+            ],
+            bool
+        )
+        self.add_default_dict_args([
+            'xyzfile', 'yfile', 'ygradxyzfile','itrainin','itestin','isubtrainin','ivalidatein','mlmodelin','setname'
+            ],
+            ""
+        )        
+        self.add_dict_args({
+            'mlmodeltype': 'ANI',
+            'mlmodelout': "ANIbestmodel.pt",
+            'sampling': "random",
+            'yestfile': "enest.dat",
+            'ygradxyzestfile': "gradest.dat",
+            'lcNtrains': [],
+            'natom': 0,
+            'atype': []
+        })
+        self.parse_input_content([
+            'ani.batch_size=8',
+            'ani.max_epochs=10000000',
+            'ani.early_stopping_learning_rate=0.00001',
+            'ani.force_coefficient=0.1',
+            'ani.batch_sizes=0',
+            'ani.patience=100',
+            'ani.lrfactor=0.5',
+            'ani.Rcr=5.2',
+            'ani.Rca=3.5',
+            'ani.EtaR=16',
+            'ani.ShfR=0.9,1.16875,1.4375,1.70625,1.975,2.24375,2.5125,2.78125,3.05,3.31875,3.5875,3.85625,4.125,4.39375,4.6625,4.93125',
+            'ani.Zeta=32',
+            'ani.ShfZ=0.19634954,0.58904862,0.9817477,1.3744468,1.7671459,2.1598449,2.552544,2.9452431',
+            'ani.EtaA=8',
+            'ani.ShfA=0.90,1.55,2.20,2.85',
+            'ani.Neuron_l1=160',
+            'ani.Neuron_l2=128',
+            'ani.Neuron_l3=96',
+            'ani.AF1=CELU',
+            'ani.AF2=CELU',
+            'ani.AF3=CELU'
+            ])
+
+    def parse(self, argsraw):
+        self.parse_input_content(argsraw)
+        self.argProcess()
+
+    def argProcess(self):
+        for k, v in self.ani.data.items():
+            if k.lower() in ['shfr','etar','zeta','shfz','shfz','etaa','shfa']:
+                self.ani.data[k] = torch.tensor([float(i) for i in str(v).split(',')])
+            elif k.lower() in ['neuron_l1','neuron_l2','neuron_l3']:
+                self.ani.data[k] =[int(i) for i in str(v).split(',')]
+            elif k.lower() in ['af1','af2','af3']:
+                self.ani.data[k] = v.split(',')
+
+
+        with open(self.xyzfile,'r') as f:
+            self.natom = int(f.readline())
+            exec('f.readline()')
+            self.atype = [f.readline().split()[0] for i in range(self.natom)]
+        
+        if not self.mlmodelin:
+            self.mlmodelin = self.mlmodelout
+        
+        species_order = sorted(set(self.atype), key=lambda x: self.atype.index(x))
+        if len(self.ani.Neuron_l1) == 1: self.ani.Neuron_l1 = self.ani.Neuron_l1*len(species_order)
+        if len(self.ani.Neuron_l2) == 1: self.ani.Neuron_l2 = self.ani.Neuron_l2*len(species_order)
+        if len(self.ani.Neuron_l3) == 1: self.ani.Neuron_l3 = self.ani.Neuron_l3*len(species_order)
+        # if self.lcyonly: self.ygradxyzfile=''
+        
+        if not self.ani.batch_size:
+            if self.ntrain: self.ani.batch_size = int(math.sqrt(self.ntrain))
+            else: self.ani.batch_size=8
+
+        if self.learningcurve:
+            self.lcNtrains = [int(i) for i in str(self.lcNtrains).split(',')]
+            if self.ani.batch_sizes:
+                self.ani.batch_sizes = [int(i) for i in self.lcNtrains.split(',')]
+                self.ani.batch_size = self.ani.batch_sizes[self.lcNtrains.index(self.ntrain)]
+
+class ANICls(object):
+    dataConverted = False
+    @classmethod
+    def load(cls):
+        loaded=False
+        if not loaded:
+            
+            available = True
+            try: 
+                import h5py
+                import torch
+                from . import TorchANI_train
+                from . import TorchANI_predict
+            except: 
+                available = False
+            
+
+            globals()['available'] = available
+            globals()['h5py'] = h5py
+            globals()['torch'] = torch
+            globals()['TorchANI_train'] = TorchANI_train
+            globals()['TorchANI_predict'] = TorchANI_predict
+
+            cls.batch_size = 8
+            cls.max_epochs = 10000000
+            cls.early_stopping_learning_rate = 1.0E-5
+            cls.force_coefficient = 0.1
+            cls.batch_sizes = []
+            cls.patience = 100
+            cls.lrfactor = 0.5
+            cls.Rcr = 5.2000e+00
+            cls.Rca = 3.5000e+00
+            cls.EtaR = torch.tensor([1.6000000e+01])
+            cls.ShfR = torch.tensor([9.0000000e-01, 1.1687500e+00, 1.4375000e+00, 1.7062500e+00, 1.9750000e+00, 2.2437500e+00, 2.5125000e+00, 2.7812500e+00, 3.0500000e+00, 3.3187500e+00, 3.5875000e+00, 3.8562500e+00, 4.1250000e+00, 4.3937500e+00, 4.6625000e+00, 4.9312500e+00])
+            cls.Zeta = torch.tensor([3.2000000e+01])
+            cls.ShfZ = torch.tensor([1.9634954e-01, 5.8904862e-01, 9.8174770e-01, 1.3744468e+00, 1.7671459e+00, 2.1598449e+00, 2.5525440e+00, 2.9452431e+00])
+            cls.EtaA = torch.tensor([8.0000000e+00])
+            cls.ShfA = torch.tensor([9.0000000e-01, 1.5500000e+00, 2.2000000e+00, 2.8500000e+00])
+            cls.Neuron_l1 = [160]
+            cls.Neuron_l2 = [128]
+            cls.Neuron_l3 = [96]
+            cls.AF1 = ['CELU']
+            cls.AF2 = ['CELU']
+            cls.AF3 = ['CELU']
+
+            loaded=True
+
+    def __init__(self, argsANI = sys.argv[1:]):
+        print(' ___________________________________________________________\n\n%s' % __doc__)
+    
+    @classmethod
+    def convertdata(cls, argsANI, subdatasets):
+        cls.load()
+        def convert(fileout, setname, yorgrad=False):
+            # convert data to .h5 format
+            prefix = ''
+            if args.learningcurve: prefix = '../'
+            if setname: 
+                coordfile = prefix+'xyz.dat_'+setname
+                yfile = prefix+'y.dat_'+setname
+                gradfile = prefix+'grad.dat_'+setname
+            else: 
+                coordfile = args.xyzfile
+                if args.yfile and yorgrad: yfile = args.yfile
+                if args.ygradxyzfile and yorgrad: gradfile = args.ygradxyzfile
+            
+            with open(coordfile ,'r') as fxyz:
+                natom_max=0
+                for line in fxyz:
+                    natom = int(line)
+                    natom_max = max(natom_max, natom)
+                    for i in range(natom+1):
+                        fxyz.readline()
+                fxyz.seek(0)
+                data={}
+                # data['smiles'] = np.array([args.smiles.encode('ascii')])
+                data['species'] = np.array([i.encode('ascii') for i in args.atype])
+                data['coordinates'] = np.array([])
+                if args.yfile and yorgrad: 
+                    fy =  open(yfile,'r')
+                    data['energies'] = np.array([])
+                if args.ygradxyzfile and yorgrad: 
+                    fgrad = open(gradfile,'r')
+                    data['forces'] = np.array([])
+                for line in fxyz:
+                    natom = int(line)
+                    if args.ygradxyzfile and yorgrad: fgrad.readline()
+
+                    if args.yfile and yorgrad: data['energies'] =np.append(data['energies'], float(fy.readline()))
+                    fxyz.readline()
+                    if args.ygradxyzfile and yorgrad: fgrad.readline()
+
+                    for i in range(natom):
+                        data['coordinates'] = np.append(data['coordinates'], [float(i) for i in fxyz.readline().split()[-3:]])
+                        if args.ygradxyzfile and yorgrad: 
+                            data['forces'] = np.append(data['forces'], [-1*float(i) for i in fgrad.readline().split()[-3:]]) 
+                    for i in range(natom_max-natom):
+                        data['coordinates'] = np.append(data['coordinates'], np.zeros(3))
+                        if args.ygradxyzfile and yorgrad: 
+                            data['forces'] = np.append(data['forces'], np.zeros(3)) 
+            
+                if args.yfile and yorgrad: fy.close()
+                if args.ygradxyzfile and yorgrad: fgrad.close()
+                data['coordinates'] = data['coordinates'].reshape(-1,natom_max,3)
+                if args.ygradxyzfile and yorgrad: data['forces'] = data['forces'].reshape(-1,natom_max,3)
+
+                hf = h5py.File(fileout, 'w')
+                grp = hf.create_group('dataset')
+                subgrp = grp.create_group('molecule')
+                for k, v in data.items():
+                    subgrp[k] = v
+
+        if not available:
+            stopper.stopMLatom('Please install all Python module required for TorchANI')
+            
+        args = Args()
+        args.parse(argsANI)
+        print('\n Converting data...\n\n')
+        if subdatasets:
+            if 'Subtrain'  in subdatasets:
+                for i in subdatasets:
+                    if i in ['Subtrain', 'Validate']: convert('ANI_'+i.lower()+'.h5', i.lower(),True)
+                    else: convert('ANI_'+i.lower()+'.h5', i.lower())
+            else:
+                for i in subdatasets:
+                    if i == 'Train': convert('ANI_'+i.lower()+'.h5', i.lower(),True)
+                    else: convert('ANI_'+i.lower()+'.h5', i.lower())
+        else:   convert('ANI.h5', '')
+
+    @classmethod
+    def createMLmodel(cls, argsANI, subdatasets):
+        cls.load()
+        args = Args()
+        args.parse(argsANI)
+
+        if not cls.dataConverted or args.learningcurve:
+            cls.convertdata(argsANI, subdatasets)
+            cls.dataConverted = True
+        
+        starttime = time.time()
+
+        # if args.mlmodeltype.lower() in ['ani1x','ani1ccx','ani2x']:
+        #     print('ML model already built in TorchANI')
+        if False:
+            pass
+        else:
+            TorchANI_train.train(args)
+
+        endtime = time.time()
+        wallclock = endtime - starttime
+        return wallclock
+
+    @classmethod
+    def useMLmodel(cls, argsANI, subdatasets):
+        cls.load()
+        args = Args()
+        args.parse(argsANI)
+
+        if not cls.dataConverted: 
+            cls.convertdata(argsANI, subdatasets)
+            cls.dataConverted = True
+
+        starttime = time.time()
+
+        TorchANI_predict.predict(args)
+        
+        endtime = time.time()
+        wallclock = endtime - starttime
+        return wallclock
+
+def printHelp():
+    helpText = __doc__ + '''
+  To use Interface_ANI, please install TorchANI and its dependencies
+
+  Arguments with their default values:
+    MLprog=TorchANI            enables this interface
+    MLmodelType=ANI            requests ANI model
+
+    ani.batch_size=8           batch size
+    ani.max_epochs=10000000    max epochs
+    
+    ani.early_stopping_learning_rate=0.00001
+                               learning rate that triggers early-stopping
+    
+    ani.force_coefficient=0.1  weight for force
+    ani.Rcr=5.2                radial cutoff radius
+    ani.Rca=3.5                angular cutoff radius
+    ani.EtaR=1.6               radial smoothness in radial part
+    
+    ani.ShfR=0.9,1.16875,      radial shifts in radial part
+    1.4375,1.70625,1.975,
+    2.24375,2.5125,2.78125,
+    3.05,3.31875,3.5875,
+    3.85625,4.125,4.9375,
+    4.6625,4.93125
+    
+    ani.Zeta=32                angular smoothness
+    
+    ani.ShfZ=0.19634954,       angular shifts
+    0.58904862,0.9817477,
+    1.3744468,1.7671459,
+    2.1598449,2.552544,
+    2.9452431
+    
+    ani.EtaA=8                 radial smoothness in angular part
+    ani.ShfA=0.9,1.55,2.2,2.85 radial shifts in angular part
+    ani.Neuron_l1=160          number of neurons in layer 1
+    ani.Neuron_l2=128          number of neurons in layer 2
+    ani.Neuron_l3=96           number of neurons in layer 3
+    ani.AF1='CELU'             acitivation function for layer 1
+    ani.AF2='CELU'             acitivation function for layer 2
+    ani.AF3='CELU'             acitivation function for layer 3
+
+  Cite TorchANI:
+    X. Gao, F. Ramezanghorbani, O. Isayev, J. S. Smith, A. E. Roitberg,
+    J. Chem. Inf. Model. 2020, 60, 3408
+    
+  Cite ANI model:
+    J. S. Smith, O. Isayev, A. E. Roitberg, Chem. Sci. 2017, 8, 3192
+'''
+    print(helpText)
+
+if __name__ == '__main__':
+    ANICls()

MLatom/interfaces/sGDML/interface_sGDML.py

@@ -0,0 +1,276 @@
+#!/usr/bin/python3
+'''
+  !---------------------------------------------------------------------------! 
+  ! Interface_sGDML: Interface between sGDML and MLatom                       ! 
+  ! Implementations by: Fuchun Ge                                             ! 
+  !---------------------------------------------------------------------------! 
+'''
+from typing import Tuple
+import numpy as np
+import os, sys, subprocess, time, shutil, re, math, random
+import stopper
+from args_class import ArgsBase
+
+filedir = os.path.dirname(__file__)
+
+class Args(ArgsBase):
+    def __init__(self):
+        super().__init__()
+        self.add_default_dict_args([
+            'learningCurve',
+            ],
+            bool
+        )
+        self.add_default_dict_args([
+            'mlmodeltype'
+            ],
+            'sgdml'
+        )
+        self.add_default_dict_args([
+            'xyzfile', 'yfile', 'ygradxyzfile','itrainin','itestin','isubtrainin','ivalidatein','mlmodelin','setname'
+            ],
+            ""
+        )        
+        self.add_dict_args({
+            'nthreads': 1,
+            'mlmodelout': "sGDML",
+            'sampling': "random",
+            'yestfile': "enest.dat",
+            'ygradxyzestfile': "gradest.dat"
+        })
+        self.parse_input_content([
+            'sGDML.torch=false ',
+            'sGDML.gdml=false ',
+            'sGDML.cprsn=false ',
+            'sGDML.no_E=false ',
+            'sGDML.s=false'
+        ])
+
+    def parse(self, argsraw):
+        self.parse_input_content(argsraw)
+        self.argProcess()
+
+    def argProcess(self):
+        with open(self.xyzfile,'r') as f:
+            sGDMLCls.natom = int(f.readline())
+            exec('f.readline()')
+            sGDMLCls.atype = [f.readline().split()[0] for i in range(sGDMLCls.natom)]
+        
+        # for estAccMLmodel that MLmodelIn may not be provided
+        if self.mlmodelout[-4:] != '.npz': self.mlmodelout += '.npz'
+        if not self.mlmodelin:
+            self.mlmodelin = self.mlmodelout
+
+        if self.sGDML.s:
+            sigmas = self.sGDML.s.split(',')
+            self.sGDML.s = []
+            for i in range(len(sigmas)):
+                self.sGDML.s.append(sigmas[i])
+            self.sGDML.s.reverse()
+        if self.mlmodeltype.lower()=='gdml':
+            self.sGDML.gdml = True
+
+            
+
+class sGDMLCls(object):
+    setname = ''
+    atype = ''
+    natom = 0
+    nsubtrain = 0
+    nvalidate = 0
+    dataConverted = False
+    @classmethod
+    def load(cls):
+        loaded=False
+        if not loaded:
+            try:
+                sGDMLbin = os.environ['sGDML']
+                from . import sgdml_dataset_from_extxyz
+                from sgdml.predict import GDMLPredict
+                from sgdml.utils import io
+                globals()['sGDMLbin'] = sGDMLbin
+                globals()['sgdml_dataset_from_extxyz'] = sgdml_dataset_from_extxyz
+                globals()['GDMLPredict'] = GDMLPredict
+                globals()['io'] = io
+            except:
+                stopper.stopMLatom('Please specify sGDML bin in $sGDML')
+            loaded=True
+
+    def __init__(self, argssGDML = sys.argv[1:]):
+            print(' ___________________________________________________________\n\n%s' % __doc__)
+    
+    @classmethod
+    def convertdata(cls, argssGDML, subdatasets):
+        cls.load()
+        def convert(fileout, setname, yorgrad=False):
+            # convert data to .exyz format
+            prefix = ''
+            if args.learningcurve: prefix = '../'
+            if setname: 
+                coordfile = prefix+'xyz.dat_'+setname
+                yfile = prefix+'y.dat_'+setname
+                gradfile = prefix+'grad.dat_'+setname
+            else: 
+                coordfile = args.xyzfile
+                if args.yfile and yorgrad: yfile = args.yfile
+                if args.ygradxyzfile and yorgrad: gradfile = args.ygradxyzfile
+            with open(fileout, 'w') as fo, open(coordfile ,'r') as fxyz:
+                if args.yfile and yorgrad: fy =  open(yfile,'r')
+                if args.ygradxyzfile and yorgrad: fgrad = open(gradfile,'r')
+                for line in fxyz:
+                    if setname == 'subtrain': cls.nsubtrain += 1
+                    if setname == 'validate': cls.nvalidate += 1
+                    natom = int(line)
+                    fo.write(line)
+                    if args.ygradxyzfile and yorgrad: fgrad.readline()
+
+                    if args.yfile and yorgrad: fo.write('%s '% fy.readline()[:-1])
+                    fo.write('\n')
+                    fxyz.readline()
+                    if args.ygradxyzfile and yorgrad: fgrad.readline()
+
+                    for i in range(natom):
+                        fo.write(fxyz.readline()[:-1])
+                        if args.ygradxyzfile and yorgrad: 
+                            fo.write('%24.15f %24.15f %24.15f' % tuple(-1*float(i) for i in fgrad.readline().split()[-3:]))
+                        fo.write('\n')
+                if args.yfile and yorgrad: fy.close()
+                if args.ygradxyzfile and yorgrad: fgrad.close()
+                
+            if yorgrad: sgdml_dataset_from_extxyz.convert(fileout)
+
+        args = Args()
+        args.parse(argssGDML)
+        print('\n Converting data...\n\n')
+        if subdatasets:
+            if 'Subtrain' in subdatasets:
+                for i in subdatasets:
+                    if i in ['Subtrain', 'Validate']: convert('sGDML_'+i.lower()+'.xyz', i.lower(),True)
+                    else: convert('sGDML_'+i.lower()+'.xyz', i.lower())
+            else:
+                for i in subdatasets:
+                    if i == 'Train': convert('sGDML_'+i.lower()+'.xyz', i.lower(),True)
+                    else: convert('sGDML_'+i.lower()+'.xyz', i.lower())
+        else:   convert('sGDML.xyz', '')
+
+    @classmethod
+    def createMLmodel(cls, argssGDML, subdatasets):
+        cls.load()
+        args = Args()
+        args.parse(argssGDML)
+        
+        if not cls.dataConverted or args.learningcurve:
+            cls.convertdata(argssGDML, subdatasets)
+            cls.dataConverted = True
+
+        cmds = [
+            [sGDMLbin, 'create', '-p', str(args.nthreads), '-v', 'sGDML_validate.npz', '--task_dir', 'sGDMLtask', 'sGDML_subtrain.npz', str(cls.nsubtrain), str(cls.nvalidate)],
+            [sGDMLbin, 'train', '-p', str(args.nthreads), 'sGDMLtask'],
+            [sGDMLbin, 'validate', '-p', str(args.nthreads), 'sGDMLtask', 'sGDML_validate.npz'], 
+            [sGDMLbin, 'select', '-p', str(args.nthreads), '--model_file', args.mlmodelout, 'sGDMLtask'], 
+            [sGDMLbin, 'show', '-o', '-p', str(args.nthreads), args.mlmodelout]
+        ]
+        
+        if args.sGDML.E_cstr: cmds[0].insert(6, '--E_cstr')
+        if args.sGDML.no_E: cmds[0].insert(6, '--no_E')
+        if args.sGDML.cprsn: cmds[0].insert(6, '--cprsn')
+        if args.sGDML.gdml: cmds[0].insert(6, '--gdml')
+        if args.sGDML.s: 
+            for i in args.sGDML.s:
+                cmds[0].insert(6, i)
+            cmds[0].insert(6, '-s')
+        if args.sGDML.torch: 
+            cmds[0].insert(4, '--torch')
+            cmds[1].insert(4, '--torch')
+            cmds[2].insert(4, '--torch')
+            cmds[3].insert(4, '--torch')
+        
+        cls.nsubtrain = 0
+        cls.nvalidate = 0
+
+        FNULL = open(os.devnull, 'w')
+        
+        starttime = time.time()
+
+        for cmd in cmds:
+            print('> sgdml '+' '.join(cmd[1:]))
+            sys.stdout.flush()
+            proc = subprocess.Popen(cmd, stdout=subprocess.PIPE ,stderr=subprocess.PIPE)
+            for line in iter(proc.stdout.readline, b''):
+                print(line.decode('ascii').replace('\n',''))
+                sys.stdout.flush()
+            proc.stdout.close()  
+        FNULL.close()     
+
+        endtime = time.time()
+        wallclock = endtime - starttime
+        return wallclock
+
+
+    @classmethod
+    def useMLmodel(cls, argssGDML, subdatasets):
+        cls.load()
+        args = Args()
+        args.parse(argssGDML)
+
+        if not cls.dataConverted: 
+            cls.convertdata(argssGDML, subdatasets)
+            cls.dataConverted = True
+
+        if args.setname: args.setname='_'+args.setname
+
+        starttime = time.time()
+
+        model = np.load(args.mlmodelin)
+        gdml = GDMLPredict(model)
+
+        r,_ = io.read_xyz('sGDML'+args.setname+'.xyz') # 9 atoms
+        e,f = gdml.predict(r)
+
+        line="%d\n\n" % int(cls.natom)
+        with open(args.yestfile,'wb') as ff:
+            np.savetxt(ff, e, fmt='%12.6f', delimiter=" ")
+        with open(args.ygradxyzestfile,'wb') as ff:
+            for force in f:
+                ff.write(line.encode('utf-8'))
+                np.savetxt(ff, -1*force.reshape(-1,3), fmt='%12.6f', delimiter=" ")
+                
+        endtime = time.time()
+        wallclock = endtime - starttime
+        return wallclock
+def printHelp():
+    helpText = __doc__ + '''
+  To use Interface_sGDML, please set enviromental variable $sGDML
+  to the path of the sGDML executable
+
+  Arguments with their default values:
+    MLprog=sGDML               enables this interface
+    MLmodelType=S              requests model S
+      sGDML                    [defaut]
+      GDML
+    sgdml.gdml=False           use GDML instead of sGDML
+    sgdml.cprsn=False          compress kernel matrix along symmetric 
+                               degrees of freedom
+    sgdml.no_E=False           do not predict energies
+    sgdml.E_cstr=False         include the energy constraints in the
+                               kernel
+    sgdml.s=<s1>[,<s2>[,...]]  set hyperparameter sigma
+           =<start>:[<step>:]<stop>     
+ 
+  Cite sGDML program:
+    S. Chmiela, H. E. Sauceda, I. Poltavsky, K.-R. Müller, A. Tkatchenko,
+    ‎Comput. Phys. Commun. 2019, 240, 38
+
+  Cite GDML method, if you use it:
+    S. Chmiela, A. Tkatchenko, H. E. Sauceda, I. Poltavsky, K. T. Schütt,
+    K.-R. Müller, Sci. Adv. 2017, 3, e1603015
+    
+  Cite sGDML method, if you use it:
+    S. Chmiela, H. E. Sauceda, K.-R. Müller, A. Tkatchenko,
+    Nat. Commun. 2018, 9, 3887
+    
+'''
+    print(helpText)
+
+if __name__ == '__main__':
+    sGDMLCls()