mlatom 3.3.0__tar.gz → 3.4.0__tar.gz

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: mlatom
-Version: 3.3.0
+Version: 3.4.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.3.0
+Version 3.4.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.3.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.4.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.3.0 → mlatom-3.4.0}/README.md

@@ -1,7 +1,7 @@
 # About Program
 MLatom: a Package for Atomistic Simulations with Machine Learning
 
-Version 3.3.0
+Version 3.4.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.3.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.4.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.3.0 → mlatom-3.4.0}/setup.py

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

{mlatom-3.3.0 → mlatom-3.4.0}/src/mlatom/MLatom.py

@@ -4,7 +4,7 @@
   !---------------------------------------------------------------------------!
   !                                                                           !
   !     MLatom: a Package for Atomistic Simulations with Machine Learning     !
-  !                             MLatom 3.3.0                                  !
+  !                             MLatom 3.4.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.3.0, Xiamen University, Xiamen, China, 2013-2024.               !
+  ! version 3.4.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.3.0 → mlatom-3.4.0}/src/mlatom/MLtasks.py

@@ -367,7 +367,10 @@ def geomopt(args):
     molDB = loading_data(XYZfile=args.XYZfile, charges=args.charges, multiplicities=args.multiplicities)
     model = loading_model(args)
     fname = args.optXYZ
-    if os.path.exists(fname): stopper.stopMLatom(f'File {fname} already exists; please delete or rename it')
+    # if os.path.exists(fname): stopper.stopMLatom(f'File {fname} already exists; please delete or rename it')
+    if os.path.exists(fname): 
+        os.remove(fname)
+        print(f' * Warning * File {fname} already exists; MLatom will overwrite {fname}')
     db_opt = data.molecular_database()
     kwargs = {}
     if args.optProg:        kwargs['program'] = args.optprog
@@ -375,16 +378,28 @@ def geomopt(args):
     if args.ase.fmax:       kwargs['convergence_criterion_for_forces'] = float(args.ase.fmax)
     if args.ase.steps:      kwargs['maximum_number_of_steps'] = int(args.ase.steps)
     if args.ase.optimizer:  kwargs['optimization_algorithm'] = args.ase.optimizer
+    if len(molDB.molecules)<=10: 
+        kwargs['print_properties'] = 'all'
+    else:
+        kwargs['print_properties'] = 'min'
+    if args.printall:       kwargs['print_properties'] = 'all'
+    if args.printmin:       kwargs['print_properties'] = None
+    if args.dumpopttrajs:
+        kwargs['dump_trajectory_interval'] = 1
+    else:
+        kwargs['dump_trajectory_interval'] = None
+    kwargs['format'] = 'json'
     for imol, mol in enumerate(molDB):
         mol.number = imol+1
+        print(' %s ' % ('='*78))
+        print(' Optimization of molecule %d' % (imol+1))
+        print(' %s \n' % ('='*78))
+        kwargs['filename'] = f'opttraj{imol+1}.json'
         geomopt = simulations.optimize_geometry(model=model,
                                         initial_molecule=mol,
                                         **kwargs)
         db_opt.molecules.append(geomopt.optimized_molecule)
-        print(' %s ' % ('='*78))
-        print(' Optimization of molecule %d' % (imol+1))
-        print(' %s \n' % ('='*78))
-        print(f'   {"Iteration":^10s}    {"Energy (Hartree)":^25s}') # units?
+        print(f'\n\n   {"Iteration":^10s}    {"Energy (Hartree)":^25s}')
         for step in geomopt.optimization_trajectory.steps:
             print('   %10d    %25.13f' % (step.step+1, step.molecule.energy))
         if args.AIQM1 or args.AIQM1DFT or args.AIQM1DFTstar:
@@ -522,6 +537,8 @@ def MLTPA(args):
 # Reusable functions below. Name with -ing form
 def loading_data(XYZfile, Yfile=None, YgradXYZfile=None, charges=None, multiplicities=None):
     assert XYZfile, 'please provide data file(s) needed.'
+    if not os.path.exists(XYZfile):
+        stopper.stopMLatom(f'xyz file {XYZfile} is not found!')
     molecular_database = data.molecular_database.from_xyz_file(XYZfile)
     if Yfile: 
         molecular_database.add_scalar_properties_from_file(Yfile)

{mlatom-3.3.0 → mlatom-3.4.0}/src/mlatom/aiqm1.py

@@ -80,8 +80,9 @@ class aiqm1(models.torchani_model):
         
         for atom in molecule.atoms:
             if not atom.atomic_number in [1, 6, 7, 8]:
-                print(' * Warning * Molecule contains elements other than CHNO, no calculations performed')
-                return
+                errmsg = ' * Warning * Molecule contains elements other than CHNO, no calculations performed'
+                # print(errmsg)
+                raise ValueError(errmsg)
         
         if nstates >1:
             mol_copy = molecule.copy()

{mlatom-3.3.0 → mlatom-3.4.0}/src/mlatom/args_class.py

@@ -202,53 +202,53 @@ class ArgsBase():
     @staticmethod
     def _args_extractor(string: str) -> List[str]:
         pair_dict = {'(': ')', '[': ']', '{': '}', '<': '>', "'": "'", '"': '"'}
-        pair_level = 0
         pair_right: List[str] = []
         tmp = ''
-        prev = ''
+        comment = False
         out_string_list: List[str] = []
         for chr in string:
             pair_right_last = pair_right[-1] if pair_right else ''
+            if chr == '#':
+                comment = True
+            if chr == '\n':
+                comment = False
+            if comment:
+                continue
             if chr in pair_dict.keys() and chr != pair_right_last:
-                pair_level += 1
                 pair_right.append(pair_dict[chr])
             elif chr == pair_right_last:
-                pair_level -= 1
                 pair_right.pop()
-            elif chr == ' ' and pair_level == 0 and prev.strip():  # end
-                out_string_list.append(tmp)
-                tmp = ''
-                prev = chr
-                continue
-            if '"' in pair_right or "'" in pair_right:
+            if pair_right:
                 tmp += chr
+            elif chr in [' ', '\n'] and tmp:  # end
+                out_string_list.append(tmp.strip())
+                tmp = ''
             else:
                 tmp += chr.strip()
-            prev = chr
-        if pair_level:
+        if pair_right:
             stopper.stopMLatom(f'pair character unmatched args: "{tmp}"')
         else:
             if tmp: out_string_list.append(tmp)
+        
         return out_string_list
     
     def parse_input_file(self, file: str):
-        try:
-            with open(file) as f:
-                content = f.read().splitlines()
-        except:
-            self._raise_error(f'can not open file {file}! exit!'); exit()
+        with open(file) as f:
+            content = self._args_extractor(f.read())
         self.parse_input_content(content)
     
     def parse_input_content(self, content: Union[List[str], str]):
-        help = False
         if type(content) is str:
-            content = content.splitlines()
-        for c in self._clean_input_content(content):
+            content = [content]
+        _content = []
+        for c in content:
+            _content.extend(self._args_extractor(c) if '\n' not in c else [c])
+        for c in _content:
             splitted = c.split('=')
             if len(splitted) == 1:
                 key = splitted[0]
                 if key.lower() in ['help', '-help', '-h', '--help']:
-                    help = True
+                    self._print_doc()
                 elif key.lower() in self.ignore_list:
                     pass
                 else:
@@ -265,8 +265,6 @@ class ArgsBase():
                     self.data[key] = value
             else:
                 self._raise_error(f'error happends at your input file, error content:\n    {c}')
-        if help:
-            self._print_doc()
     
     @classmethod
     def _multi_level_dict(cls, key: str, value: Any) -> Tuple[str, Union[AttributeDict, Dict[str, Union[str, Dict[str, Any]]]]]:
@@ -283,19 +281,7 @@ class ArgsBase():
         try: Doc.printDoc(self.args_dict)
         except: self._warning_print('Doc not imported')
         stopper.stopMLatom('')
-    
-    @classmethod
-    def _clean_input_content(cls, content: List[str]) -> List[str]:
-        to_return: List[str] = []
-        for c in content:
-            usable = c.split('#')[0].strip()
-            while usable.find(' =') != -1:
-                usable = usable.replace(' =', '=')
-            while usable.find('= ') != -1:
-                usable = usable.replace('= ', '=')
-            to_return.extend(cls._args_extractor(usable))
-        return to_return
-    
+
     @property
     def args_dict(self) -> Dict[str, Any]:
         return AttributeDict.normal_dict(self.data)
@@ -616,6 +602,14 @@ class mlatom_args(ArgsBase):
             'DOF':-6,
             'linear':0,
         })
+        # geomopt output
+        self.add_default_dict_args(
+            [
+                'printall',                # print out all information in geometry optimization
+                'printmin',                # print out minimal infomation in geometry optimization        
+                'dumpopttrajs'             # whether to dump optimization trajectory
+            ], ''
+        )
     
         self.defualt_args2pass = self.args_string_list(['', None])
 
@@ -628,23 +622,35 @@ class mlatom_args(ArgsBase):
         if len(argsraw) == 0:
             Doc.printDoc({})
             stopper.stopMLatom('At least one option should be provided')
-        elif len(argsraw) == 1:
-            if os.path.exists(argsraw[0]):
-                # self.parse_input_file(argsraw[0])
-                self.argsraw = []
-                with open(argsraw[0]) as f:
-                    content = f.read().splitlines()
-                    for c in self._clean_input_content(content):
-                        self.argsraw += [c]
-                self.parse_input_content(self.argsraw)
-            else:
-                self.parse_input_content(argsraw[0])
-        elif len(argsraw) >= 2:
+        elif len(argsraw) == 1 and os.path.exists(argsraw[0]):
+            self.parse_input_file(argsraw[0])
+        else:
             self.parse_input_content(argsraw)
         
         self._post_operations()
+            
+    def _post_operations(self):
+        if not self.MLprog:
+            if self.MLmodelType :
+                try: self.MLprog = default_MLprog[self.MLmodelType.lower()]
+                except: stopper.stopMLatom('Unkown MLmodelType')
+            else:
+                self.MLprog = 'MLatomF'
         self._checkArgs()
         self._check_hyperopt()
+        self._multi_lines_to_file()
+
+    def _multi_lines_to_file(self):
+        import hashlib
+        for arg in self.args2pass:
+            if '\n' in arg:
+                key, value = arg.split('=')
+                tmpfile = f"{key}_{hashlib.md5(value.encode('utf-8')).hexdigest()[:6]}"
+                if 'xyz' in key.lower():
+                    tmpfile += '.xyz'
+                with open(tmpfile, 'w') as f:
+                    f.write(value.strip("'").strip('"').strip() + '\n')
+                self.parse_input_content([f'{key}={tmpfile}'])
 
     def _check_hyperopt(self):
         self.hyperparameter_optimization = {
@@ -660,14 +666,6 @@ class mlatom_args(ArgsBase):
                 self._hyperopt_str_dict[key.split('.')[-1]] = value
         if self._hyperopt_str_dict:
             self.hyperparameter_optimization['optimization_algorithm'] = self.hyperopt.algorithm
-            
-    def _post_operations(self):
-        if not self.MLprog:
-            if self.MLmodelType :
-                try: self.MLprog = default_MLprog[self.MLmodelType.lower()]
-                except: stopper.stopMLatom('Unkown MLmodelType')
-            else:
-                self.MLprog = 'MLatomF'
 
     def _checkArgs(self):
         if not self._task:

{mlatom-3.3.0 → mlatom-3.4.0}/src/mlatom/conversions.py

@@ -2,7 +2,7 @@ import numpy as np
 
 def smi2xyz(smi_strings):
     from openbabel import pybel as pb
-    if type(smi_strings) == str:
+    if isinstance(smi_strings, str):
         smi_strings =  smi_strings.strip().split('\n')
     xyz_strings = []
     for smi_string in smi_strings:
@@ -13,7 +13,7 @@ def smi2xyz(smi_strings):
 
 def xyz2smi(xyz_strings):
     from openbabel import pybel as pb
-    if type(xyz_strings) == str:
+    if isinstance(xyz_strings, str):
         xyz_strings = split_xyz_string(xyz_strings)
     smi_strings = []
     for string in xyz_strings:

{mlatom-3.3.0 → mlatom-3.4.0}/src/mlatom/data.py

@@ -585,6 +585,37 @@ class molecule:
 
     def align(self, ref, pivot='CoM'):
         pass
+    
+    def info(self, properties = 'all', return_string=False):
+        printstrs = []
+        printstrs += [f" Molecule with {len(self.get_element_symbols())} atom(s): {', '.join(self.get_element_symbols())}", '']
+        if (properties == 'all' or 'xyz_coordinates' in properties) and 'xyz_coordinates' in self.atoms[0].__dict__:
+            printstrs += [f' XYZ coordinates, Angstrom\n']
+            iatom = 0
+            for atom in self.atoms:
+                iatom += 1
+                printstrs += [' %-4d %-3s %18.6f %18.6f %18.6f' % (iatom, atom.element_symbol,
+        atom.xyz_coordinates[0], atom.xyz_coordinates[1], atom.xyz_coordinates[2])]
+            printstrs += ['']
+        if (properties == 'all' or 'distance_matrix' in properties) and 'xyz_coordinates' in self.atoms[0].__dict__:
+            printstrs += [f' Interatomic distance matrix, Angstrom\n']
+            dist_mat = self.get_internuclear_distance_matrix()
+            printstrs += [f'{dist_mat}']
+            printstrs += ['']
+        if (properties == 'all' or 'energy' in properties) and 'energy' in self.__dict__:
+            printstrs += [' Energy: %18.6f Hartree\n' % self.energy]
+        if (properties == 'all' or 'energy_gradients' in properties) and 'energy_gradients' in self.atoms[0].__dict__:
+            printstrs += [f' Energy gradients, Hartree/Angstrom\n']
+            iatom = 0
+            for atom in self.atoms:
+                iatom += 1
+                printstrs += [' %-4d %-3s %18.6f %18.6f %18.6f' % (iatom, atom.element_symbol,
+        atom.energy_gradients[0], atom.energy_gradients[1], atom.energy_gradients[2])]
+            printstrs += ['']
+            printstrs += [' Energy gradients norm: %18.6f Hartree/Angstrom\n' % np.linalg.norm(self.energy_gradients)]
+        printstr = '\n'.join(printstrs)
+        if return_string: return printstr
+        else: print(printstr)
 
     def __add__(self, obj):
         if isinstance(obj, molecular_database):
@@ -593,7 +624,8 @@ class molecule:
             return molecular_database([self] + [obj])
         
     def __str__(self):
-        return f"molecule with {len(self.get_element_symbols())} atom(s): {', '.join(self.get_element_symbols())}"
+        printstr = self.info(properties = 'all', return_string=True)
+        return printstr
    
     def __iter__(self):
         for atom in self.atoms:

{mlatom-3.3.0 → mlatom-3.4.0}/src/mlatom/doc.py

@@ -250,6 +250,7 @@ class Doc():
       optProg=scipy            use scipy package [default]
       optProg=gaussian         use Gaussian program
       optProg=ASE              use ASE
+      optProg=geometric        use geomeTRIC
       optxyz=S                 save optimized geometries in file S [default: optgeoms.xyz]
     The following options only used for ASE program:
         ase.fmax=R                    threshold of maximum force (in eV/A)

{mlatom-3.3.0 → mlatom-3.4.0}/src/mlatom/header.py

@@ -132,6 +132,8 @@ def printHeader(args):
     if args.geomopt:
         if 'optprog=ase' in ' '.join(args.args2pass).lower():
             refItems['Atomic simulation environment (ASE)']  = refs['ASE']
+        elif 'optprog=geometric' in ' '.join(args.args2pass).lower():
+            refItems['geomeTRIC program']  = refs['geometric']
         elif 'optprog=gaussian' in ' '.join(args.args2pass).lower():
             refItems['Gaussian program']  = refs['Gaussian']
         elif args.freq and 'optprog=pyscf' in ' '.join(args.args2pass).lower():
@@ -174,6 +176,8 @@ def printHeader(args):
     if args.ts:
         if 'optprog=ase' in ' '.join(args.args2pass).lower():
             refItems['Atomic simulation environment (ASE)']  = refs['ASE']
+        if 'optprog=geometric' in ' '.join(args.args2pass).lower():
+            refItems['geomeTRIC program']  = refs['geometric']
         elif 'optprog=gaussian' in ' '.join(args.args2pass).lower():
             refItems['Gaussian program']  = refs['Gaussian']
         elif "GAUSS_EXEDIR" in os.environ:

{mlatom-3.3.0 → mlatom-3.4.0}/src/mlatom/initial_conditions.py

@@ -6,7 +6,7 @@ from . import constants
 
   !---------------------------------------------------------------------------! 
   ! initial_conditions: Module for generating initial conditions              ! 
-  ! Implementations by: Yi-Fan Hou & Pavlo O. Dral                            ! 
+  ! Implementations by: Yi-Fan Hou, Lina Zhang, Fuchun Ge, Pavlo O. Dral      ! 
   !---------------------------------------------------------------------------! 
 '''
 import numpy as np 
@@ -14,7 +14,13 @@ from . import data
 from . import stopper
 from . import constants
 
-def excitation_energy_window_filter(molecular_database=None, model=None, model_predict_kwargs={}, target_excitation_energy=1, window_half_width=0.1, f_max=None, random_seed=None):
+def excitation_energy_window_filter(molecular_database=None,
+                                    model=None,
+                                    model_predict_kwargs={},
+                                    target_excitation_energy=1,
+                                    window_half_width=0.1,
+                                    f_max=None,
+                                    random_seed=None):
     molDB = molecular_database
     if not random_seed is None:
         np.random.seed(random_seed)
@@ -66,16 +72,18 @@ def generate_initial_conditions(molecule=None, generation_method=None, number_of
     Generate initial conditions
 
     Arguments:
-        molecule (:class:`data.molecule`): Molecule with necessary information
-        generation_method (str): Initial condition generation method 
-        number_of_initial_conditions (int): Number of initial conditions to generate, 1 by default
-        file_with_initial_xyz_coordinates (str): File with initial xyz coordinates, only valid for ``generation_method='user-defined'``
-        file_with_initial_xyz_velocities (str): File with initial xyz velocities, only valid for ``generation_method='user-defined'``
-        eliminate_angular_momentum (bool): Remove angular momentum from velocities, valid for ``generation_method='random'`` and ``generation_method='wigner'``
-        degrees_of_freedom (int): Degrees of freedom of the molecule, by default remove translational and rotational degrees of freedom. It can be a negative value, which means that some value is subtracted from 3*Natoms
-        initial_temperature (float): Initial temperature in Kelvin, control random initial velocities
-        initial_kinetic_energy (float): Initial kinetic energy in Hartree, control random initial velocities
-        random_seed (int): Random seed for numpy random number generator
+        molecule (:class:`data.molecule`): molecule with necessary information
+        generation_method (str): initial condition generation method, see below the table
+        number_of_initial_conditions (int): number of initial conditions to generate, 1 by default
+        file_with_initial_xyz_coordinates (str): file with initial xyz coordinates, only valid for ``generation_method='user-defined'``
+        file_with_initial_xyz_velocities (str): file with initial xyz velocities, only valid for ``generation_method='user-defined'``
+        eliminate_angular_momentum (bool): remove angular momentum from velocities, valid for ``generation_method='random'`` and ``generation_method='wigner'``
+        degrees_of_freedom (int): degrees of freedom of the molecule, by default remove translational and rotational degrees of freedom. It can be a negative value, which means that some value is subtracted from 3*Natoms
+        initial_temperature (float): initial temperature in Kelvin, control random initial velocities
+        initial_kinetic_energy (float): initial kinetic energy in Hartree, control random initial velocities
+        random_seed (int): random seed for numpy random number generator (do not use unless you want to obtain the same results every time)
+        filter_by_energy_window (bool): filter by excitation energy window
+        window_filter_kwargs (dict): keyword arguments for filtering the energy window, see below the table
 
     .. table::
         :align: center
@@ -83,13 +91,26 @@ def generate_initial_conditions(molecule=None, generation_method=None, number_of
         =============================  =============================================
         generation_method              description
         =============================  =============================================
-        ``'user-defined'`` (default)   Use user-defined initial conditions
-        ``'random'``                   Generate random velocities
-        ``'maxwell-boltzmann'``        Randomly generate initial velocities from Maxwell-Boltzmann distribution
-        ``'wigner'``                   Use Wigner sampling as implemented in `Newton-X <https://doi.org/10.1021/acs.jctc.2c00804>`__
+        ``'user-defined'`` (default)   use user-defined initial conditions
+        ``'random'``                   generate random velocities
+        ``'maxwell-boltzmann'``        randomly generate initial velocities from Maxwell-Boltzmann distribution
+        ``'wigner'``                   use Wigner sampling as implemented in `Newton-X <https://doi.org/10.1021/acs.jctc.2c00804>`__
         =============================  =============================================
 
         
+    .. table::
+        :align: center
+        
+        ================================  ================================================================================================================
+        window_filter_kwargs              description
+        ================================  ================================================================================================================
+        model                             model or method that can calculate excitation energies and oscillator strengths
+        model_predict_kwargs              keyword arguments for above model, typically ``nstates`` specifying how many states to calculate
+        target_excitation_energy (float)  in eV
+        window_half_width (float)         in eV
+        random_seed (int)                 random seed for numpy random number generator (do not use unless you want to obtain the same results every time)
+        ================================  ================================================================================================================
+        
     Returns:
         A molecular database (:class:`ml.data.molecular_database`) with ``number_of_initial_conditions`` initial conditions
 
@@ -113,6 +134,22 @@ def generate_initial_conditions(molecule=None, generation_method=None, number_of
                                                       generation_method = 'wigner',
                                                       number_of_initial_conditions = 1)
 
+        # Sample with filtering by excitation energy window. Requires the model for calculating excitation energies and oscillator strengths.
+        model = ml.models.methods(method='AIQM1') 
+        model_predict_kwargs={'nstates':9} # requests calculation of 9 electronic states
+        window_filter_kwargs={'model':model,
+                              'model_predict_kwargs':model_predict_kwargs, 
+                              'target_excitation_energy':5.7, # eV
+                              'window_half_width':0.1, # eV}
+        init_cond_db = ml.generate_initial_conditions(molecule=mol,
+                                                    generation_method='wigner',
+                                                    number_of_initial_conditions=5,
+                                                    initial_temperature=0,
+                                                    random_seed=0,
+                                                    use_hessian=False,
+                                                    filter_by_energy_window=True,
+                                                    window_filter_kwargs=window_filter_kwargs)
+        
     .. note::
 
         Wigner sampling needs Hessian matrix. You can use ``ml.models.methods.predict(molecule=mol,calculate_hessian=True)`` to get Hessian matrix.

{mlatom-3.3.0 → mlatom-3.4.0}/src/mlatom/interfaces/ase_interface.py

@@ -24,6 +24,11 @@ def optimize_geometry(initial_molecule, model, convergence_criterion_for_forces,
     if optimization_algorithm == None:
         optimization_algorithm = 'LBFGS'
     
+    if 'model_predict_kwargs' in kwargs:
+        model_predict_kwargs = kwargs['model_predict_kwargs']
+    else:
+        model_predict_kwargs = {}
+    
     optimization_trajectory = data.molecular_trajectory()
     globals()['initial_molecule'] = initial_molecule
     globals()['optimization_trajectory'] = optimization_trajectory
@@ -64,7 +69,7 @@ def optimize_geometry(initial_molecule, model, convergence_criterion_for_forces,
             
     # atoms.set_calculator(MLatomCalculator(model=model, save_optimization_trajectory=True))
     # atoms.set_calculator() is deprecated
-    atoms.calc = MLatomCalculator(model=model, save_optimization_trajectory=True)
+    atoms.calc = MLatomCalculator(model=model, save_optimization_trajectory=True, model_predict_kwargs=model_predict_kwargs)
     
     from ase import optimize
     opt = optimize.__dict__[optimization_algorithm](atoms)
@@ -86,16 +91,26 @@ def transition_state(initial_molecule, model,
                      **kwargs):
     if optimization_algorithm == None:
         optimization_algorithm = 'dimer'
+
+    if 'model_predict_kwargs' in kwargs:
+        # model_predict_kwargs = kwargs['model_predict_kwargs']
+        model_predict_kwargs = kwargs.pop('model_predict_kwargs')
+    else:
+        model_predict_kwargs = {}
+
     if optimization_algorithm.casefold() == 'dimer'.casefold():
         return dimer_method(initial_molecule, model, 
+                            model_predict_kwargs,
                             convergence_criterion_for_forces,
                             maximum_number_of_steps, **kwargs)
     elif optimization_algorithm.casefold() == 'NEB'.casefold():
         return nudged_elastic_band(initial_molecule, kwargs.pop('final_molecule'), model, 
-                            convergence_criterion_for_forces,
-                            maximum_number_of_steps,  **kwargs)
+                                   model_predict_kwargs,
+                                   convergence_criterion_for_forces,
+                                   maximum_number_of_steps,  **kwargs)
 
 def dimer_method(initial_molecule, model, 
+                 model_predict_kwargs,
                  convergence_criterion_for_forces,
                  maximum_number_of_steps,  **kwargs):
     optimization_trajectory = data.molecular_trajectory()
@@ -107,7 +122,7 @@ def dimer_method(initial_molecule, model,
         initial_molecule.write_file_with_xyz_coordinates(filename=xyzfilename)
         atoms = io.read(xyzfilename, index=':', format='xyz')[0]
 
-    atoms.calc = MLatomCalculator(model=model, save_optimization_trajectory=True)
+    atoms.calc = MLatomCalculator(model=model,  model_predict_kwargs= model_predict_kwargs, save_optimization_trajectory=True)
 
     from ase.dimer import DimerControl, MinModeAtoms, MinModeTranslate
 
@@ -121,6 +136,7 @@ def dimer_method(initial_molecule, model,
     return optimization_trajectory
 
 def nudged_elastic_band(initial_molecule, final_molecule, model, 
+                        model_predict_kwargs,
                         convergence_criterion_for_forces,
                         maximum_number_of_steps,
                         number_of_middle_images=3, **kwargs):
@@ -144,7 +160,7 @@ def nudged_elastic_band(initial_molecule, final_molecule, model,
     neb = NEB(images, **kwargs)
     neb.interpolate()
     for image in images[1:number_of_middle_images+1]:
-        image.calc = MLatomCalculator(model=model, save_optimization_trajectory=True)
+        image.calc = MLatomCalculator(model=model,  model_predict_kwargs= model_predict_kwargs, save_optimization_trajectory=True)
     optimizer = MDMin(neb, trajectory='A2B.traj')
     optimizer.run(fmax=convergence_criterion_for_forces,
                   steps=maximum_number_of_steps)
@@ -153,9 +169,10 @@ def nudged_elastic_band(initial_molecule, final_molecule, model,
 
 class MLatomCalculator(Calculator):
     implemented_properties = ['energy', 'forces']
-    def __init__(self, model, save_optimization_trajectory = False):
+    def __init__(self, model,  model_predict_kwargs, save_optimization_trajectory = False):
         super(MLatomCalculator, self).__init__()
         self.model = model
+        self.model_predict_kwargs =  model_predict_kwargs
         self.save_optimization_trajectory = save_optimization_trajectory
 
     def calculate(self, atoms=None, properties=['energy'], system_changes=all_changes):
@@ -172,7 +189,7 @@ class MLatomCalculator(Calculator):
             coordinates = mol_from_file.xyz_coordinates
             current_molecule.xyz_coordinates = coordinates
             
-            self.model.predict(molecule=current_molecule, calculate_energy=True, calculate_energy_gradients=True)
+            self.model._predict_geomopt(molecule=current_molecule, calculate_energy=True, calculate_energy_gradients=True, **self.model_predict_kwargs)
             if not 'energy' in current_molecule.__dict__:
                 raise ValueError('model did not return any energy')