Implemented atomic representation for a single element type and valence electron only representation

qstack/spahm/rho/atom.py

@@ -19,7 +19,7 @@ def check_file(mol_file):
 def get_repr(mol, elements, charge, spin,
              open_mod=defaults.omod, dm=None,
              guess=defaults.guess, model=defaults.model, xc=defaults.xc,
-             auxbasis=defaults.auxbasis):
+             auxbasis=defaults.auxbasis, only_z=None):
 
     # User-defined options
     elements = sorted(list(set(elements)))
@@ -31,11 +31,16 @@ def get_repr(mol, elements, charge, spin,
     # Compute density matrices
     if dm is None:
         dm = spahm.compute_spahm.get_guess_dm(mol, guess, openshell=spin, xc=xc)
+# if only the representations for a given atom type are to be computed restricts the considered atomic indices
+    if only_z is not None:
+        only_i = [i for i,z in enumerate(mol.elements) if z in only_z]
+    else:
+        only_i = range(mol.natm) #otherwise consider the full list of atomic indices
 
     rep = []
     for omod in open_mod:
         DM      = utils.dm_open_mod(dm, omod) if spin is not None else dm
-        c_df    = df_wrapper(mol, DM, auxbasis)
+        c_df    = df_wrapper(mol, DM, auxbasis, only_i=only_i)
         vectors = sym_wrapper(c_df, mol, idx, ao, ao_len, M, elements)
         if spin is None:
             rep = vectors
@@ -45,9 +50,7 @@ def get_repr(mol, elements, charge, spin,
     if spin is not None:
         rep = np.hstack(rep)
 
-    mrep = []
-    for q, v in zip(mol.elements, rep):
-        mrep.append(np.array((q, v)))
+    mrep = [np.array((q,v), dtype=object) for q,v in zip(np.array(mol.elements)[only_i], rep)]
     return np.array(mrep)
 
 
@@ -61,6 +64,7 @@ def main():
     parser.add_argument('--model',     dest='model',     default=defaults.model,               type=str, help=f"the model to use when creating the representation (default: {defaults.model})")
     parser.add_argument('--dm',        dest='dm',        default=None,                         type=str, help="a density matrix to load instead of computing the guess")
     parser.add_argument('--species',   dest='elements',  default=defaults.elements, nargs='+', type=str, help="the elements contained in the database")
+    parser.add_argument('--only',   dest='only_z', default=None, nargs='+', type=str, help="The restricted list of elements for which you want to generate the representation")
     parser.add_argument('--charge',    dest='charge',    default=0,                            type=int, help='total charge of the system (default: 0)')
     parser.add_argument('--spin',      dest='spin',      default=None,                         type=int, help='number of unpaired electrons (default: None) (use 0 to treat a closed-shell system in a UHF manner)')
     parser.add_argument('--xc',        dest='xc',        default=defaults.xc,                  type=str, help=f'DFT functional for the SAD guess (default: {defaults.xc})')
@@ -72,10 +76,10 @@ def main():
     mol = compound.xyz_to_mol(check_file(args.mol), args.basis, charge=args.charge, spin=args.spin, unit=args.units)
     dm = None if args.dm is None else np.load(args.dm)
 
-    representations = get_repr(mol, elements, args.charge, args.spin,
+    representations = get_repr(mol, args.elements, args.charge, args.spin,
                                open_mod=args.omod,
                                dm=dm, guess=args.guess, model=args.model,
-                               xc=args.xc, auxbasis=args.auxbasis)
+                               xc=args.xc, auxbasis=args.auxbasis, only_z=args.only_z)
 
     # output dir
     cwd = os.getcwd()

qstack/spahm/rho/atomic_density.py

@@ -3,11 +3,15 @@
 from . import lowdin
 
 
-def fit(mol, dm, aux_basis, short=False, w_slicing=True):
+def fit(mol, dm, aux_basis, short=False, w_slicing=True, only_i=None):
 
     L = lowdin.Lowdin_split(mol, dm)
 
-    dm_slices = mol.aoslice_nr_by_atom()[:,2:]
+    if only_i is not None:
+        dm_slices = mol.aoslice_nr_by_atom()[only_i,2:]
+    else:
+        dm_slices = mol.aoslice_nr_by_atom()[:,2:]
+
     auxmol = compound.make_auxmol(mol, aux_basis)
     eri2c, eri3c = fields.decomposition.get_integrals(mol, auxmol)[1:]
 

qstack/spahm/rho/dmb_rep_atom.py

@@ -29,17 +29,17 @@ def df_sad_diff(mol, dm, auxbasis):
         dm = dm - dm_sad
         return fields.decomposition.decompose(mol, dm, auxbasis)[1]
 
-    def df_lowdin_long(mol, dm, auxbasis):
-        return atomic_density.fit(mol, dm, auxbasis)
+    def df_lowdin_long(mol, dm, auxbasis, only_i=[]):
+        return atomic_density.fit(mol, dm, auxbasis, only_i=only_i)
 
-    def df_lowdin_short(mol, dm, auxbasis):
-        return atomic_density.fit(mol, dm, auxbasis, short=True)
+    def df_lowdin_short(mol, dm, auxbasis, only_i=[]):
+        return atomic_density.fit(mol, dm, auxbasis, short=True, only_i=only_i)
 
-    def df_lowdin_long_x(mol, dm, auxbasis):
-        return atomic_density.fit(mol, dm, auxbasis, w_slicing=False)
+    def df_lowdin_long_x(mol, dm, auxbasis, only_i=[]):
+        return atomic_density.fit(mol, dm, auxbasis, w_slicing=False, only_i=only_i)
 
-    def df_lowdin_short_x(mol, dm, auxbasis):
-        return atomic_density.fit(mol, dm, auxbasis, short=True, w_slicing=False)
+    def df_lowdin_short_x(mol, dm, auxbasis, only_i=[]):
+        return atomic_density.fit(mol, dm, auxbasis, short=True, w_slicing=False, only_i=only_i)
 
     def df_occup(mol, dm, auxbasis):
         L = lowdin.Lowdin_split(mol, dm)

tests/test_spahm_a.py

@@ -43,9 +43,24 @@ def test_water_SAD_guess_close_shell():
         assert(a[0] == a_true[0])                        # atom type
         assert(np.linalg.norm(a[1]-a_true[1]) < 1e-08)   # atom representations
 
+def test_water_single_element():
+    path = os.path.dirname(os.path.realpath(__file__))
+    mol = compound.xyz_to_mol(path+'/data/H2O.xyz', 'minao', charge=0, spin=None)
+
+    X = atom.get_repr(mol, ["H", "O"], 0, None, dm=None,
+                      guess='LB', model='lowdin-long-x', auxbasis='ccpvdzjkfit', only_z=['O']) #requesting reps for O-atom only
+
+    X_true = np.load(path+'/data/SPAHM_a_H2O/X_H2O.npy', allow_pickle=True)
+    a = X[0]
+    assert(X.shape == np.array(X_true[0], ndmin=2).shape)
+    for a_true in X_true:
+        if a[0] == a_true[0]:                       # atom type
+            assert(np.linalg.norm(a[1]-a_true[1]) < 1e-08)   # atom representations
+
 
 
 if __name__ == '__main__':
     test_water()
     test_water_SAD_guess_close_shell()
     test_water_SAD_guess_open_shell()
+    test_water_single_element()