Merge pull request #116 from tjduigna/master

Molcas symmetrized orbital generation

Author: tjduigna

exatomic/algorithms/basis.py

@@ -347,12 +347,8 @@ def _radial(self, x, y, z, alphas, cs, rs=None, pre=None):
     def _angular(self, shl, x, y, z, *ang):
         """Generates the symbolic angular portion of a basis function.
         Substitutes symbolic (_i) -> (_i - iA) for i in [x, y, z]."""
-        if len(ang) == 3:
-            l, m, n = ang
-            sym = _x ** l * _y ** m * _z ** n
-        else:
-            L, ml = ang
-            sym = self._sh[L][ml]
+        if len(ang) == 3: sym = _x ** ang[0] * _y ** ang[1] * _z ** ang[2]
+        else: sym = self._sh[ang[0]][ang[1]]
         return sym.subs({_x: _x - x, _y: _y - y, _z: _z - z})
 
 
@@ -379,13 +375,11 @@ def _evaluate_gau_bso_sym(self, xs, ys, zs, irrep=None):
         p['L'] = [self._shells[i].L for i in p['shldx']]
         shls = p.groupby(['center', 'L'])
         norms = [shl.norm_contract() for shl in self._shells]
-        ocens = [(col, col.replace('ocen', 'sign'))
-                 for col in bso.columns if col.startswith('ocen')]
-        ocens = [c for t in ocens for c in t]
-        for cen, L, ml, irrep, args in zip(bso['center'],
-                                           bso['L'], bso['ml'],
-                                           bso['irrep'],
-                                           *(bso[col] for col in ocens)):
+        ocens = [c for col in bso.columns if col.startswith('ocen')
+                 for c in (col, col.replace('ocen', 'sign'))]
+        for i, (cen, L, ml, irrep) in enumerate(zip(bso['center'],
+                                                    bso['L'], bso['ml'],
+                                                    bso['irrep'])):
             ax, ay, az = self._xyzs[cen]
             shldx = shls.get_group((cen, L)).shldx.values[0]
             ishl = self._shells[shldx]
@@ -394,7 +388,9 @@ def _evaluate_gau_bso_sym(self, xs, ys, zs, irrep=None):
             r = self._radial(ax, ay, az, ishl.alphas,
                              norm[:,cache[irrep][cen][L][ml]])
             oterm = None
-            for ocen, sign in zip(args[::2], args[1::2]):
+            for oc, si in zip(ocens[::2], ocens[1::2]):
+                ocen = bso[oc].iloc[i]
+                sign = bso[si].iloc[i]
                 if ocen >= 0:
                     ox, oy, oz = self._xyzs[ocen]
                     ao = self._angular(ishl, ox, oy, oz, L, ml)
@@ -416,7 +412,7 @@ def _evaluate_gau_bso_sym(self, xs, ys, zs, irrep=None):
     def _evaluate_gau_bso(self, xs, ys, zs, irrep=None):
         """Evaluates a Gaussian basis set according to the order specified
         by the :class:`~exatomic.core.basis.BasisSetOrder` and returns a
-        numpy array."""
+        numpy array of numerical basis function values."""
         cnt = 0
         if xs is not None: flds = np.empty((len(self), len(xs)))
         else: flds = Series([None for _ in range(len(self))])
@@ -523,7 +519,7 @@ def __repr__(self):
     def __init__(self, uni, frame=0, cartp=True):
         # Attach relevant uni attributes
         self._meta = uni.meta
-        self._bso = uni.basis_set_order
+        self._bso = uni.current_basis_set_order
         ptrs, xyzs, shells = uni.enumerate_shells()
         self._ptrs = ptrs
         self._xyzs = xyzs
@@ -547,4 +543,35 @@ def __init__(self, uni, frame=0, cartp=True):
         self._expnt = _r ** 2
         if not self._meta['gaussian']:
             self._expnt = _r
-            self._pre = uni.basis_set_order['prefac']
+            self._pre = uni.current_basis_set_order['prefac']
+
+
+def compute_uncontracted_basis_set_order(uni):
+    bso = uni.basis_set_order
+    prims = uni.basis_set.primitives_by_shell()
+    df = []
+    pcen, pL, pml = -1, -1, -1
+    if uni.meta['program'] == 'molcas':
+        for cen, L, ml in zip(bso['center'], bso['L'], bso['ml']):
+            if not all(p == c for p, c in zip((pcen, pL, pml), (cen, L, ml))):
+                seht = uni.atom.loc[cen]['set']
+                nprim = prims.loc[(seht, L)]
+                for i in range(nprim):
+                    df.append((cen, L, ml, i, 0))
+            pcen = cen
+            pL = L
+            pml = ml
+    elif uni.meta['program'] == 'nwchem':
+        if not uni.meta['spherical']:
+            raise NotImplementedError('Sorry.')
+        from exatomic.nwchem.basis import spherical_ordering_function as func
+        prims = prims.groupby('set')
+        for i, seht in enumerate(uni.atom['set']):
+            shls = prims.get_group(seht)
+            for (seht, L), nprim in shls.items():
+                for j in range(nprim):
+                    for ml in func(L):
+                        df.append([i, L, ml, j, 0])
+    cols = ('center', 'L', 'ml', 'shell', 'frame')
+    uni.meta['uncontracted'] = True
+    return pd.DataFrame(df, columns=cols)

exatomic/algorithms/numerical.py

@@ -7,6 +7,7 @@
 Everything in this module is implemented in numba.
 """
 import numpy as np
+import pandas as pd
 from numba import (jit, jitclass, deferred_type,
                    optional, int64, float64, boolean)
 from exatomic.base import nbche
@@ -159,6 +160,78 @@ def momatrix_as_square(movec):
     return square
 
 
+############################################
+# Reordering matrix elements can be useful #
+############################################
+
+@jit(nopython=True, nogil=True, cache=nbche)
+def _index_map(old, new):
+    """
+    Basis functions are uniquely defined by 4 indices;
+    atomic center, L value, ml value and so-called "shell" index.
+    shell is defined here as corresponding to a column index in an instance
+    of a :class:`~exatomic.algorithms.numerical.Shell`. This function
+    simply finds the mapping between the `old` basis set ordering scheme
+    and the new one.
+
+    Args:
+        old (np.ndarray): order [center, L, ml, shell]
+        new (np.ndarray): order [center, L, ml, shell]
+
+    Returns:
+        mappr (np.ndarray): old -> new indices
+    """
+    mappr = np.empty(len(new), dtype=np.int64)
+    for ni in range(len(new)):
+        for oi in range(len(old)):
+            if (new[ni] == old[oi]).all():
+                mappr[ni] = oi
+                break
+    return mappr
+
+@jit(nopython=True, nogil=True, cache=nbche)
+def _reorder_matrix(old, new, values):
+    """
+    Reorders matrix elements according to an old and new basis set order.
+
+    Args:
+        old (np.ndarray): order [center, L, ml, shell]
+        new (np.ndarray): order [center, L, ml, shell]
+
+    Returns:
+        nvals (np.ndarray): reordered matrix
+    """
+    nvals = np.empty(values.shape)
+    mappr = _index_map(old, new)
+    for ni, oi in enumerate(mappr):
+        for nj, oj in enumerate(mappr):
+            nvals[ni, nj] = values[oi, oj]
+    return nvals
+
+
+def reorder_matrix(uni_to_reorder, ordered_uni, attr='momatrix', mocoefs='coef'):
+    """
+    Reorders matrix elements in a uni_to_reorder by the basis set order
+    defined in ordered_uni.
+
+    Args:
+        uni_to_reorder (:class:`~exatomic.core.universe.Universe`): uni to reorder
+        ordered_uni (:class:`~exatomic.core.universe.Universe`): ordered uni
+        attr (str): specify if non-standard matrices (default "momatrix")
+        mocoefs (str): column name in gettattr(uni_to_reorder, attr)
+
+    Returns:
+        reordered (pd.DataFrame): reordered matrix with labeled columns and indices
+    """
+    cols = ['center', 'L', 'ml', 'shell']
+    old = uni_to_reorder.current_basis_set_order[cols].values.astype(np.int64)
+    new = ordered_uni.current_basis_set_order[cols].values.astype(np.int64)
+    val = getattr(uni_to_reorder, attr).square(column=mocoefs).values
+    cmat = _reorder_matrix(old, new, val)
+    idxs = pd.Index(range(val.shape[0]), name='chi')
+    cols = pd.Index(range(val.shape[0]), name='orbital')
+    return pd.DataFrame(cmat, columns=cols, index=idxs)
+
 #######################
 # Basis set expansion #
 #######################
@@ -169,12 +242,6 @@ def _enum_cartesian(L):
     # for z in range(L + 1):
     #     for y in range(L + 1 - z):
     #         yield (L - y - z, y, z)
-    # Naive combinations with replacements
-    # for i in range(L, -1, -1):
-    #     for j in range(L, -1, -1):
-    #         for k in range(L, -1, -1):
-    #             if i + j + k == L:
-    #                 yield (i, j, k)
     # Double loop CWR
     for x in range(L, -1, -1):
         for z in range(L + 1 - x):

exatomic/algorithms/orbital.py

@@ -8,13 +8,14 @@
 set of operations that are provided by this module and wrapped into a clean API.
 """
 import numpy as np
+from numba import TypingError
 from datetime import datetime
 from exatomic.base import sym2z
 from .orbital_util import (
     numerical_grid_from_field_params, _determine_fps,
     _determine_vector, _compute_orb_ang_mom, _compute_current_density,
-    _compute_orbitals, _compute_density, _check_column, _make_field,
-    _compute_orbitals_nojit)
+    _compute_density, _check_column, _make_field,
+    _compute_orbitals_numba, _compute_orbitals_numpy)
 
 
 def _setup_orbital(uni, verbose, vector, fps, icoefs, jcoefs=None, irrep=None):
@@ -38,17 +39,20 @@ def _setup_orbital(uni, verbose, vector, fps, icoefs, jcoefs=None, irrep=None):
         return t1, vector, fps, x, y, z, bvs, icoefs, jcoefs
     return t1, vector, fps, x, y, z, bvs, icoefs
 
+def _compute_orbital(verbose, npts, bvs, vector, cmat):
+    try: ovs = _compute_orbitals_numba(npts, bvs, vector, cmat)
+    except (ValueError, IndexError, AssertionError, TypingError) as e:
+        if verbose: print('numba eval failed, falling back to numpy')
+        ovs = _compute_orbitals_numpy(npts, bvs, vector, cmat)
+    return ovs
 
-def _teardown_orbital(uni, verbose, field, t1, inplace, replace, dens=False):
+def _teardown_orbital(uni, verbose, field, t1, inplace, name='orbitals'):
     """Boilerplate for finishing the functions in this module."""
     if verbose:
         t2 = datetime.now()
-        kind = 'density ' if dens else 'orbitals'
-        p2 = 'Timing: compute {} - {:>8.2f}s.'
-        print(p2.format(kind, (t2-t1).total_seconds()))
+        p2 = 'Timing: compute {:<8} - {:>8.2f}s.'
+        print(p2.format(name, (t2-t1).total_seconds()))
     if not inplace: return field
-    if replace and hasattr(uni, '_field'):
-        del uni.__dict__['_field']
     uni.add_field(field)
 
 
@@ -76,14 +80,12 @@ def add_molecular_orbitals(uni, field_params=None, mocoefs=None,
     Warning:
         If replace is True, removes any fields previously attached to the universe
     """
+    if replace and hasattr(uni, '_field'): del uni.__dict__['_field']
     t1, vector, fps, x, y, z, bvs, mocoefs = \
         _setup_orbital(uni, verbose, vector, field_params, mocoefs, irrep=irrep)
-    try: ovs = _compute_orbitals(len(x), bvs, vector, mocoefs)
-    except (ValueError, IndexError, AssertionError) as e:
-        if verbose: print('Falling back to numpy orbital evaluation.')
-        ovs = _compute_orbitals_nojit(len(x), bvs, vector, mocoefs)
+    ovs = _compute_orbital(verbose, len(x), bvs, vector, mocoefs)
     field = _make_field(ovs, fps)
-    return _teardown_orbital(uni, verbose, field, t1, inplace, replace)
+    return _teardown_orbital(uni, verbose, field, t1, inplace)
 
 
 def add_density(uni, field_params=None, mocoefs=None, orbocc=None,
@@ -106,13 +108,9 @@ def add_density(uni, field_params=None, mocoefs=None, orbocc=None,
     orbocc = _check_column(uni, 'orbital', orbocc)
     vector = uni.orbital[~np.isclose(uni.orbital[orbocc], 0)].index.values
     orbocc = uni.orbital.loc[vector][orbocc].values
-    try: ovs = _compute_orbitals(len(x), bvs, vector, mocoefs)
-    except (ValueError, IndexError, AssertionError) as e:
-        #if verbose:
-        print('Falling back to numpy orbital evaluation.')
-        ovs = _compute_orbitals_nojit(len(x), bvs, vector, mocoefs)
+    ovs = _compute_orbital(verbose, len(x), bvs, vector, mocoefs)
     field = _make_field(_compute_density(ovs, orbocc), fps.loc[0])
-    return _teardown_orbital(uni, verbose, field, t1, inplace, False)
+    return _teardown_orbital(uni, verbose, field, t1, inplace, name='density')
 
 
 def add_orb_ang_mom(uni, field_params=None, rcoefs=None, icoefs=None,
@@ -150,8 +148,6 @@ def add_orb_ang_mom(uni, field_params=None, rcoefs=None, icoefs=None,
     if verbose:
         p1 = 'Timing: grid evaluation  - {:>8.2f}s.'
         print(p1.format((t2-t1).total_seconds()))
-    print(rcoefs.shape, rcoefs.dtype)
-    print(icoefs.shape, icoefs.dtype)
     curx, cury, curz = _compute_current_density(
         bvs, grx, gry, grz, rcoefs, icoefs, occvec, verbose=verbose)
     t3 = datetime.now()
@@ -160,4 +156,4 @@ def add_orb_ang_mom(uni, field_params=None, rcoefs=None, icoefs=None,
         print(p2.format((t3-t2).total_seconds()))
     field = _make_field(_compute_orb_ang_mom(
         x, y, z, curx, cury, curz, maxes), fps)
-    return _teardown_orbital(uni, False, field, t1, inplace, False)
+    return _teardown_orbital(uni, verbose, field, t1, inplace, name='angmom')

exatomic/algorithms/orbital_util.py

@@ -220,10 +220,7 @@ def _determine_vector(uni, vector, irrep=None):
                                  'could not be found in uni.orbital.')
         else:
             ihomo = iorb[iorb['occupation'] < 1.98]
-            print(ihomo)
             ihomo = ihomo.vector.values[0]
-            print(max(0, ihomo-5))
-            print(min(ihomo + 7, len(iorb.index)))
             return np.array(range(max(0, ihomo-5),
                                   min(ihomo + 7, len(iorb.index))))
     # If specified, carry on
@@ -294,15 +291,14 @@ def _check_column(uni, df, key):
 
 
 @jit(nopython=True, nogil=True, parallel=nbpll)
-def _compute_orbitals(npts, bvs, vecs, cmat):
+def _compute_orbitals_numba(npts, bvs, vecs, cmat):
     """Compute orbitals from numerical basis functions."""
     ovs = np.empty((len(vecs), npts), dtype=np.float64)
     for i, vec in enumerate(vecs):
         ovs[i] = np.dot(cmat[:, vec], bvs)
     return ovs
 
-
-def _compute_orbitals_nojit(npts, bvs, vecs, cmat):
+def _compute_orbitals_numpy(npts, bvs, vecs, cmat):
     """Compute orbitals from numerical basis functions."""
     ovs = np.empty((len(vecs), npts), dtype=np.float64)
     for i, vec in enumerate(vecs):

exatomic/core/editor.py

@@ -8,6 +8,7 @@
 found file formats directly into :class:`~exatomic.container.Universe` objects.
 """
 import six
+import pandas as pd
 from exa import Editor as _Editor
 from exa import TypedMeta
 from .universe import Universe
@@ -56,11 +57,14 @@ def to_universe(self, **kws):
                 meta.update(self.meta)
             else:
                 meta = self.meta
-        kwargs = {'name': name, 'description': description,
-                  'meta': meta}
+        kwargs = {'name': name, 'meta': meta,
+                  'description': description}
         attrs = [attr.replace('parse_', '')
                  for attr in vars(self.__class__).keys()
                  if attr.startswith('parse_')]
+        extras = {key: val for key, val in vars(self).items()
+                  if isinstance(val, pd.DataFrame)
+                  and key[1:] not in attrs}
         for attr in attrs:
             result = None
             try:
@@ -72,4 +76,5 @@ def to_universe(self, **kws):
             if result is not None:
                 kwargs[attr] = result
         kwargs.update(kws)
+        kwargs.update(extras)
         return Universe(**kwargs)