Merge pull request #120 from TheoChem-VU/119-add-a-guess-fragments-me…

…thod-to-adffragmentjob

119 add a guess fragments method to adffragmentjob

examples/job/NaCl.py

@@ -0,0 +1,5 @@
+from tcutility.job import ADFFragmentJob
+
+with ADFFragmentJob() as job:
+	job.molecule('NaCl.xyz')
+

examples/job/NaCl.xyz

@@ -0,0 +1,7 @@
+2
+
+Na 0 0 2 frag=Na
+Cl 0 0 0 frag=Cl
+
+charge_Na = 1
+charge_Cl = -1

src/tcutility/job/adf.py

@@ -1,5 +1,5 @@
 from scm import plams
-from tcutility import log, results, formula, spell_check, data
+from tcutility import log, results, formula, spell_check, data, molecule
 from tcutility.job.ams import AMSJob
 import os
 
@@ -18,7 +18,7 @@ def __init__(self, *args, **kwargs):
         self.single_point()
 
         # by default print the fock matrix
-        self.settings.input.adf.print = 'FMatSFO'
+        self.settings.input.adf.print = 'SFOSiteEnergies'
 
     def __str__(self):
         return f'{self._task}({self._functional}/{self._basis_set}), running in {os.path.join(os.path.abspath(self.rundir), self.name)}'
@@ -230,9 +230,18 @@ class ADFFragmentJob(ADFJob):
     def __init__(self, *args, **kwargs):
         self.childjobs = {}
         super().__init__(*args, **kwargs)
-        self.name = 'complex'
+        self.name = 'EDA'
+
+    def add_fragment(self, mol: plams.Molecule, name: str = None, charge: int = 0, spin_polarization:int = 0):
+        '''
+        Add a fragment to this job. Optionally give the name, charge and spin-polarization of the fragment as well.
 
-    def add_fragment(self, mol, name=None):
+        Args:
+            mol: the molecule corresponding to the fragment.
+            name: the name of the fragment. By default it will be set to ``fragment{N+1}`` if ``N`` is the number of fragments already present.
+            charge: the charge of the fragment to be added.
+            spin_polarization: the spin-polarization of the fragment to be added.
+        '''
         # in case of giving fragments as indices we dont want to add the fragment to the molecule later
         add_frag_to_mol = True  
         # we can be given a list of atoms
@@ -254,6 +263,8 @@ def add_fragment(self, mol, name=None):
         name = name or f'fragment{len(self.childjobs) + 1}'
         self.childjobs[name] = ADFJob(test_mode=self.test_mode)
         self.childjobs[name].molecule(mol)
+        self.childjobs[name].charge(charge)
+        self.childjobs[name].spin_polarization(spin_polarization)
         setattr(self, name, self.childjobs[name])
 
         if not add_frag_to_mol:
@@ -262,9 +273,48 @@ def add_fragment(self, mol, name=None):
         if self._molecule is None:
             self._molecule = self.childjobs[name]._molecule.copy()
         else:
-            self._molecule = self._molecule + self.childjobs[name]._molecule.copy()
+            for atom in self.childjobs[name]._molecule.copy():
+                if any((atom.symbol, atom.coords) == (myatom.symbol, myatom.coords) for myatom in self._molecule):
+                    continue
+                self._molecule.add_atom(atom)
+
+    def guess_fragments(self):
+        '''
+        Guess what the fragments are based on data stored in the molecule provided for this job. 
+        This will automatically set the correct fragment molecules, names, charges and spin-polarizations.
+
+        .. seealso::
+            | :func:`tcutility.molecule.guess_fragments` for an explanation of the xyz-file format required to guess the fragments.
+            | :meth:`ADFFragmentJob.add_fragment` to manually add a fragment.
+
+        .. note::
+            This function will be automatically called if there were no fragments given to this calculation.
+        '''
+        frags = molecule.guess_fragments(self._molecule)
+        if frags is None:
+            log.error('Could not load fragment data for the molecule.')
+            return False
+
+        for fragment_name, fragment in frags.items():
+            charge = fragment.flags.get('charge', 0)
+            spin_polarization = fragment.flags.get('spinpol', 0)
+            self.add_fragment(fragment, fragment_name, charge=charge, spin_polarization=spin_polarization)
+
+        return True
 
     def run(self):
+        '''
+        Run the ``ADFFragmentJob``. This involves setting up the calculations for each fragment as well as the parent job. 
+        It will also submit a calculation with the SCF iterations set to 0 in order to facilitate investigation of the field effects using PyOrb.
+        '''
+        # check if the user defined fragments for this job
+        if len(self.childjobs) == 0:
+            log.warn('Fragments were not specified yet, trying to read them from the xyz file ...')
+            # if they did not define the fragments, try to guess them using the xyz-file
+            if not self.guess_fragments():
+                log.error('Please specify the fragments using ADFFragmentJob.add_fragment or specify them in the xyz file.')
+                raise 
+
         mol_str = " + ".join([formula.molecule(child._molecule) for child in self.childjobs.values()])
         log.flow(f'ADFFragmentJob [{mol_str}]', ['start'])
         # obtain some system wide properties of the molecules
@@ -304,7 +354,7 @@ def run(self):
             log.flow(f'Spin-Polarization: {child.settings.input.adf.SpinPolarization or 0}', ['straight', 'straight'])
             # the child name will be prepended with SP showing that it is the singlepoint calculation
             child.name = f'frag_{childname}'
-            child.rundir = self.rundir
+            child.rundir = j(self.rundir, self.name)
 
             # add the path to the child adf.rkf file as a dependency to the parent job
             self.settings.input.adf.fragments[childname] = j(child.workdir, 'adf.rkf')
@@ -320,7 +370,7 @@ def run(self):
             child.run()
             self.dependency(child)
 
-            log.flow(f'SlurmID:  {child.slurm_job_id}', ['straight', 'skip', 'end'])
+            log.flow(f'SlurmID:  {child.slurm_job_id}', ['straight', 'skip', 'straight'])
             log.flow(f'Work dir: {child.workdir}', ['straight', 'skip', 'end'])
             log.flow()
 
@@ -340,17 +390,19 @@ def run(self):
         # set the _molecule to None, otherwise it will overwrite the atoms block
         self._molecule = None
         # run this job
+        self.rundir = j(self.rundir, self.name)
+        self.name = 'complex'
         log.flow(log.Emojis.good + ' Submitting parent job', ['split'])
         super().run()
         log.flow(f'SlurmID: {self.slurm_job_id}', ['straight', 'end'])
         log.flow()
 
         # also do the calculation with SCF cycles set to 1
-        self.settings.input.adf.SCF.Iterations = 1
+        self.settings.input.adf.SCF.Iterations = 0
         self.settings.input.adf.print = 'FMatSFO'  # by default print the fock matrix for each SCF cycle
         self.settings.input.adf.AllPoints = 'Yes'
         self.settings.input.adf.FullFock = 'Yes'
-        self.name = self.name + '_SCF1'
+        self.name = 'complex_SCF0'
         log.flow(log.Emojis.good + ' Submitting extra job with 1 SCF cycle', ['split'])
 
         super().run()

src/tcutility/job/generic.py

@@ -1,4 +1,4 @@
-from tcutility import log, results, slurm
+from tcutility import log, results, slurm, molecule
 import subprocess as sp
 import os
 import stat
@@ -207,7 +207,7 @@ def molecule(self, mol: Union[str, plams.Molecule, plams.Atom, list[plams.Atom]]
             self._molecule = mol
 
         elif isinstance(mol, str) and os.path.exists(mol):
-            self._molecule = plams.Molecule(mol)
+            self._molecule = molecule.load(mol)
 
         elif isinstance(mol, str):
             self._molecule_path = os.path.abspath(mol)

src/tcutility/molecule.py

@@ -1,5 +1,6 @@
 from scm import plams
 from tcutility.results import result
+from tcutility import ensure_list
 from typing import Dict
 
 def parse_str(s: str):
@@ -128,7 +129,9 @@ def parse_flags(args):
 
 def guess_fragments(mol: plams.Molecule) -> Dict[str, plams.Molecule]:
     '''
-    Guess fragments based on data from the xyz file. Two methods are supported, see the tabs below.
+    Guess fragments based on data from the xyz file. Two methods are currently supported, see the tabs below. 
+    We also support reading of charges and spin-polarizations for the fragments. 
+    They should be given as ``charge_{fragment_name}`` and ``spinpol_{fragment_name}`` respectively.
     
     
     .. tabs::
@@ -150,6 +153,8 @@ def guess_fragments(mol: plams.Molecule) -> Dict[str, plams.Molecule]:
 
                 frag_Donor = 1, 3-5
                 frag_Acceptor = 2, 6-8
+                charge_Donor = -1
+                spinpol_Acceptor = 2
             
             In this case, fragment atom indices must be provided below the coordinates. 
             The fragment name must be prefixed with ``frag_``. Indices can be given as integers or as ranges using ``-``.
@@ -168,14 +173,19 @@ def guess_fragments(mol: plams.Molecule) -> Dict[str, plams.Molecule]:
                 H      -0.58149793       1.00718395       1.13712667 frag=Acceptor
                 H      -0.58149793      -1.00718395       1.13712667 frag=Acceptor
                 H       1.16299585      -0.00000000       1.13712667 frag=Acceptor
+
+                charge_Donor = -1
+                spinpol_Acceptor = 2
             
             In this case, fragment atoms are marked with the `frag` flag which gives the name of the fragment the atom belongs to.
 
     Args:
-        mol: the molecule that is to be split into fragments. It should have defined either method shown above. If it does not define these methods this function returns ``None``.
+        mol: the molecule that is to be split into fragments. It should have defined either method shown above. 
+             If it does not define these methods this function returns ``None``.
 
     Returns:
-        A dictionary containing fragment names as keys and :class:`plams.Molecule <scm.plams.mol.molecule.Molecule>` objects as values. Atoms that were not included by either method will be placed in the molecule object with key ``None``.
+        A dictionary containing fragment names as keys and :class:`plams.Molecule <scm.plams.mol.molecule.Molecule>` objects as values. 
+        Atoms that were not included by either method will be placed in the molecule object with key ``None``.
 
     '''
 
@@ -185,7 +195,7 @@ def guess_fragments(mol: plams.Molecule) -> Dict[str, plams.Molecule]:
         fragment_mols = {frag.removeprefix('frag_'): plams.Molecule() for frag in fragment_flags}
         for frag in fragment_flags:
             indices = []
-            index_line = mol.flags[frag]
+            index_line = ensure_list(mol.flags[frag])
             for indx in index_line:
                 if isinstance(indx, int):
                     indices.append(indx)
@@ -195,6 +205,12 @@ def guess_fragments(mol: plams.Molecule) -> Dict[str, plams.Molecule]:
                     raise ValueError(f'Fragment index {indx} could not be parsed.')
 
             [fragment_mols[frag.removeprefix('frag_')].add_atom(mol[i]) for i in indices]
+            fragment_mols[frag.removeprefix('frag_')].flags = {'tags': set()}
+            if f'charge_{frag}' in mol.flags:
+                fragment_mols[frag.removeprefix('frag_')].flags['charge'] = mol.flags[f'charge_{frag}']
+            if f'spinpol_{frag}' in mol.flags:
+                fragment_mols[frag.removeprefix('frag_')].flags['spinpol'] = mol.flags[f'spinpol_{frag}']
+
         return fragment_mols
 
     # second method, check if the atoms have a frag= flag defined
@@ -205,6 +221,13 @@ def guess_fragments(mol: plams.Molecule) -> Dict[str, plams.Molecule]:
             # get the fragment the atom belongs to and add it to the list
             fragment_mols[atom.flags.get('frag')].add_atom(atom)
 
+        for frag in fragment_names:
+            fragment_mols[frag].flags = {'tags': set()}
+            if f'charge_{frag}' in mol.flags:
+                fragment_mols[frag].flags['charge'] = mol.flags[f'charge_{frag}']
+            if f'spinpol_{frag}' in mol.flags:
+                fragment_mols[frag].flags['spinpol'] = mol.flags[f'spinpol_{frag}']
+
         return fragment_mols
 
 

test/fixtures/xyz/NaCl.xyz

@@ -0,0 +1,7 @@
+2
+
+Na 0 0 2 frag=Na
+Cl 0 0 0 frag=Cl
+
+charge_Na = 1
+charge_Cl = -1

test/fixtures/xyz/NaCl_homolytic.xyz

@@ -0,0 +1,7 @@
+2
+
+Na 0 0 2 frag=Na
+Cl 0 0 0 frag=Cl
+
+spinpol_Na = -1
+spinpol_Cl = 1

test/test_molecules.py

@@ -62,6 +62,27 @@ def test_mol_copy() -> None:
     mol.copy()
 
 
+def test_guess_fragments() -> None:
+    xyzfile = j(os.path.split(__file__)[0], "fixtures", "xyz", "NaCl.xyz")
+    mol = molecule.load(xyzfile)
+    frags = molecule.guess_fragments(mol)
+    assert frags['Na'][1].symbol == 'Na'
+
+
+def test_guess_fragments2() -> None:
+    xyzfile = j(os.path.split(__file__)[0], "fixtures", "xyz", "NaCl.xyz")
+    mol = molecule.load(xyzfile)
+    frags = molecule.guess_fragments(mol)
+    assert frags['Na'].flags['charge'] == 1
+
+
+def test_guess_fragments3() -> None:
+    xyzfile = j(os.path.split(__file__)[0], "fixtures", "xyz", "NaCl_homolytic.xyz")
+    mol = molecule.load(xyzfile)
+    frags = molecule.guess_fragments(mol)
+    assert frags['Cl'].flags['spinpol'] == 1
+
+
 if __name__ == "__main__":
     import pytest