Merge remote-tracking branch 'org/master'

Author: avmarchenko

exatomic/algorithms/basis.py

@@ -17,10 +17,46 @@
     _compute_orbitals_nojit)
 
 
+def _setup_orbital(uni, verbose, vector, fps, icoefs, jcoefs=None, irrep=None):
+    """Boilerplate for starting the functions in this module."""
+    t1 = datetime.now()
+    nbf = len(uni.basis_functions)
+    if irrep is not None:
+        nbf = len(uni.basis_set_order.groupby('irrep').get_group(irrep).index)
+    if verbose:
+        p1 = 'Evaluating {} basis functions once.'
+        print(p1.format(nbf))
+    vector = _determine_vector(uni, vector, irrep)
+    fps = _determine_fps(uni, fps, len(vector))
+    x, y, z = numerical_grid_from_field_params(fps)
+    bvs = uni.basis_functions.evaluate(x, y, z, irrep=irrep, verbose=verbose)
+    icoefs = _check_column(uni, 'current_momatrix', icoefs)
+    icoefs = uni.current_momatrix.square(column=icoefs, irrep=irrep).values
+    if jcoefs is not None:
+        jcoefs = _check_column(uni, 'current_momatrix', jcoefs)
+        jcoefs = uni.current_momatrix.square(column=jcoefs).values
+        return t1, vector, fps, x, y, z, bvs, icoefs, jcoefs
+    return t1, vector, fps, x, y, z, bvs, icoefs
+
+
+def _teardown_orbital(uni, verbose, field, t1, inplace, replace, dens=False):
+    """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()))
+    if not inplace: return field
+    if replace and hasattr(uni, '_field'):
+        del uni.__dict__['_field']
+    uni.add_field(field)
+
+
+
 def add_molecular_orbitals(uni, field_params=None, mocoefs=None,
                            vector=None, frame=0, inplace=True,
-                           replace=False, verbose=True):
-    """A universe must contain basis_set, basis_set_order, and
+                           replace=False, verbose=True, irrep=None):
+    """A universe must contain basis_set, [basis_set_order], and
     momatrix attributes to use this function.  Evaluate molecular
     orbitals on a numerical grid.  Attempts to generate reasonable
     defaults if none are provided.  If vector is not provided,
@@ -35,38 +71,24 @@ def add_molecular_orbitals(uni, field_params=None, mocoefs=None,
         vector (int, list, range, np.ndarray): the MO vectors to evaluate
         inplace (bool): if False, return the field obj instead of modifying uni
         replace (bool): if False, do not delete any previous fields
+        irrep (int): if symmetrized, the irrep to which the orbitals belong
 
     Warning:
         If replace is True, removes any fields previously attached to the universe
     """
-    t1 = datetime.now()
-    vector = _determine_vector(uni, vector)
-    bfns = uni.basis_functions
-    fps = _determine_fps(uni, field_params, len(vector))
-    mocoefs = _check_column(uni, 'current_momatrix', mocoefs)
-    if verbose:
-        p1 = 'Evaluating {} basis functions once.'
-        print(p1.format(len(bfns)))
-
-    x, y, z = numerical_grid_from_field_params(fps)
-    bvs = bfns.evaluate(x, y, z)
-    cmat = uni.current_momatrix.square(column=mocoefs).values
-    try: ovs = _compute_orbitals(len(x), bvs, vector, cmat)
-    except: ovs = _compute_orbitals_nojit(len(x), bvs, vector, cmat)
+    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)
     field = _make_field(ovs, fps)
-    t2 = datetime.now()
-    if verbose:
-        p2 = 'Timing: compute orbitals - {:>8.2f}s.'
-        print(p2.format((t2-t1).total_seconds()))
-    if not inplace: return field
-    if replace and hasattr(uni, '_field'):
-        del uni.__dict__['_field']
-    uni.add_field(field)
+    return _teardown_orbital(uni, verbose, field, t1, inplace, replace)
 
 
 def add_density(uni, field_params=None, mocoefs=None, orbocc=None,
                 inplace=True, frame=0, norm='Nd', verbose=True):
-    """A universe must contain basis_set, basis_set_order, and
+    """A universe must contain basis_set, [basis_set_order], and
     momatrix attributes to use this function.  Compute a density
     with C matrix mocoefs and occupation vector orbocc.
 
@@ -77,33 +99,26 @@ def add_density(uni, field_params=None, mocoefs=None, orbocc=None,
         orbocc (str): column in uni.orbital (default 'occupation')
         inplace (bool): if False, return the field obj instead of modifying uni
     """
-    t1 = datetime.now()
-    mocoefs = _check_column(uni, 'current_momatrix', mocoefs)
-    orbocc = mocoefs if orbocc is None and mocoefs != 'coef' else orbocc
+    mocol = mocoefs
+    t1, vector, fps, x, y, z, bvs, mocoefs = \
+        _setup_orbital(uni, verbose, None, field_params, mocoefs)
+    orbocc = mocol if orbocc is None and mocol != 'coef' else orbocc
     orbocc = _check_column(uni, 'orbital', orbocc)
-    bfns = uni.basis_functions
-    orbs = uni.current_momatrix.groupby('orbital')
-    vector = np.array(range(uni.current_momatrix.orbital.max() + 1))
-    fps = _determine_fps(uni, field_params, len(vector))
-
-    x, y, z = numerical_grid_from_field_params(fps)
-    bvs = bfns.evaluate(x, y, z)
-    cmat = uni.current_momatrix.square(column=mocoefs).values
-    try: ovs = _compute_orbitals(len(x), bvs, vector, cmat)
-    except: ovs = _compute_orbitals_nojit(len(x), bvs, vector, cmat)
-    dens = _compute_density(ovs, uni.orbital[orbocc].values)
-    t2 = datetime.now()
-    if verbose:
-        p1 = 'Timing: compute density  - {:>8.2f}s.'
-        print(p1.format((t2-t1).total_seconds()))
-    if not inplace: return _make_field(dens, fps.loc[0])
-    uni.add_field(_make_field(dens, fps.loc[0]))
+    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)
+    field = _make_field(_compute_density(ovs, orbocc), fps.loc[0])
+    return _teardown_orbital(uni, verbose, field, t1, inplace, False)
 
 
 def add_orb_ang_mom(uni, field_params=None, rcoefs=None, icoefs=None,
                     frame=0, orbocc=None, maxes=None, inplace=True,
                     norm='Nd', verbose=True):
-    """A universe must contain basis_set, basis_set_order, and
+    """A universe must contain basis_set, [basis_set_order], and
     momatrix attributes to use this function.  Compute the orbital
     angular momentum.  Requires C matrices from SODIZLDENS.X.X.R,I
     files from Molcas.
@@ -113,43 +128,36 @@ def add_orb_ang_mom(uni, field_params=None, rcoefs=None, icoefs=None,
         field_params (dict): See :func:`~exatomic.algorithms.orbital_util.make_fps`
         rcoefs (str): column in uni.current_momatrix (default 'lreal')
         icoefs (str): column in uni.current_momatrix (default 'limag')
+        maxes (np.ndarray): 3x3 array of magnetic axes (default np.eye(3))
         orbocc (str): column in uni.orbital (default 'lreal')
         inplace (bool): if False, return the field obj instead of modifying uni
     """
     if rcoefs is None or icoefs is None:
         raise Exception("Must specify rcoefs and icoefs")
-    t0 = datetime.now()
-    orbocc = rcoefs if orbocc is None else orbocc
-    rcoefs = _check_column(uni, 'current_momatrix', rcoefs)
-    icoefs = _check_column(uni, 'current_momatrix', icoefs)
+    rcol = rcoefs
+    t1, vector, fps, x, y, z, bvs, rcoefs, icoefs = \
+        _setup_orbital(uni, verbose, None, field_params, rcoefs, jcoefs=icoefs)
+    orbocc = rcol if orbocc is None else orbocc
     if maxes is None:
-        if verbose:
-            print("If magnetic axes are not an identity " \
-                  "matrix, specify maxes.")
         maxes = np.eye(3)
-    t1 = datetime.now()
-    bfns = uni.basis_functions
-    fps = _determine_fps(uni, field_params, 4)
-
-    x, y, z = numerical_grid_from_field_params(fps)
+        if verbose:
+            print("If magnetic axes are not an identity matrix, specify maxes.")
     occvec = uni.orbital[orbocc].values
-    bvs = bfns.evaluate(x, y, z)
-    grx = bfns.evaluate_diff(x, y, z, cart='x')
-    gry = bfns.evaluate_diff(x, y, z, cart='y')
-    grz = bfns.evaluate_diff(x, y, z, cart='z')
+    grx = uni.basis_functions.evaluate_diff(x, y, z, cart='x', verbose=verbose)
+    gry = uni.basis_functions.evaluate_diff(x, y, z, cart='y', verbose=verbose)
+    grz = uni.basis_functions.evaluate_diff(x, y, z, cart='z', verbose=verbose)
     t2 = datetime.now()
     if verbose:
         p1 = 'Timing: grid evaluation  - {:>8.2f}s.'
         print(p1.format((t2-t1).total_seconds()))
-
-    cmatr = uni.current_momatrix.square(column=rcoefs).values
-    cmati = uni.current_momatrix.square(column=icoefs).values
+    print(rcoefs.shape, rcoefs.dtype)
+    print(icoefs.shape, icoefs.dtype)
     curx, cury, curz = _compute_current_density(
-        bvs, grx, gry, grz, cmatr, cmati, occvec, verbose=verbose)
+        bvs, grx, gry, grz, rcoefs, icoefs, occvec, verbose=verbose)
     t3 = datetime.now()
     if verbose:
         p2 = 'Timing: current density  - {:>8.2f}s.'
         print(p2.format((t3-t2).total_seconds()))
-    ang_mom = _compute_orb_ang_mom(x, y, z, curx, cury, curz, maxes)
-    if not inplace: return _make_field(ang_mom, fps)
-    uni.add_field(_make_field(ang_mom, fps))
+    field = _make_field(_compute_orb_ang_mom(
+        x, y, z, curx, cury, curz, maxes), fps)
+    return _teardown_orbital(uni, False, field, t1, inplace, False)

exatomic/algorithms/orbital.py

@@ -9,6 +9,7 @@
 These are their stories.
 '''
 from __future__ import division
+import six
 import numpy as np
 import pandas as pd
 from numba import jit
@@ -173,6 +174,7 @@ def _compute_current_density(bvs, gvx, gvy, gvz, cmatr, cmati, occvec, verbose=T
     curx = np.zeros(npts, dtype=np.float64)
     cury = np.zeros(npts, dtype=np.float64)
     curz = np.zeros(npts, dtype=np.float64)
+    cval = np.zeros(nbas, dtype=np.float64)
     if verbose:
         fp = FloatProgress(description='Computing:')
         display(fp)
@@ -192,20 +194,43 @@ def _compute_current_density(bvs, gvx, gvy, gvz, cmatr, cmati, occvec, verbose=T
             gvxnu = gvx[nu]
             gvynu = gvy[nu]
             gvznu = gvz[nu]
-            cval = evaluate('-0.5 * (occvec * (crmu * cinu - cimu * crnu))').sum()
-            curx = evaluate('curx + cval * (bvmu * gvxnu - gvxmu * bvnu)')
-            cury = evaluate('cury + cval * (bvmu * gvynu - gvymu * bvnu)')
-            curz = evaluate('curz + cval * (bvmu * gvznu - gvzmu * bvnu)')
+            cval = evaluate('-0.5 * (occvec * (crmu * cinu - cimu * crnu))', out=cval)
+            csum = cval.sum()
+            evaluate('curx + csum * (bvmu * gvxnu - gvxmu * bvnu)', out=curx)
+            evaluate('cury + csum * (bvmu * gvynu - gvymu * bvnu)', out=cury)
+            evaluate('curz + csum * (bvmu * gvznu - gvzmu * bvnu)', out=curz)
     if verbose:
         fp.close()
     return curx, cury, curz
 
 
-def _determine_vector(uni, vector):
+def _determine_vector(uni, vector, irrep=None):
     """Find some orbital indices in a universe."""
+    if irrep is not None: # Symmetry is fun
+        iorb = uni.orbital.groupby('irrep').get_group(irrep)
+        if vector is not None: # Check if vectors are in irrep
+            # Input vectors appropriately indexed by irrep
+            if all((i in iorb.vector.values for i in vector)):
+                return np.array(vector)
+            # Input vectors indexed in terms of total vectors
+            elif all((i in iorb.index.values for i in vector)):
+                return iorb.loc[vector]['vector'].values
+            else:
+                raise ValueError('One or more specified vectors '
+                                 '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
     if isinstance(vector, int): return np.array([vector])
-    typs = (list, tuple, range, np.ndarray)
+    typs = (list, tuple, six.moves.range, np.ndarray)
     if isinstance(vector, typs): return np.array(vector)
+    # Try to find some reasonable default
     norb = len(uni.basis_set_order.index)
     if vector is None:
         if norb < 10:

exatomic/algorithms/orbital_util.py

@@ -2,8 +2,8 @@
 # Copyright (c) 2015-2018, Exa Analytics Development Team
 # Distributed under the terms of the Apache License 2.0
 """
-Tests for symbolic basis functions
-######################################
+Tests for symbolic basis function utilities
+#############################################
 """
 from __future__ import absolute_import
 from __future__ import print_function
@@ -14,7 +14,7 @@
 from exatomic.base import resource
 from exatomic import nwchem, molcas
 from ..basis import (cart_lml_count, spher_lml_count, solid_harmonics,
-                     enum_cartesian, car2sph, Symbolic, BasisFunctions)
+                     enum_cartesian, car2sph, BasisFunctions)
 
 
 class TestCartesianToSpherical(TestCase):
@@ -45,11 +45,11 @@ def setUp(self):
         self.mo = mo.to_universe()
 
     def test_basis_functions(self):
-        nwfns = self.nw.basis_functions.list()
-        mofns = self.mo.basis_functions.list()
-        nw = sorted(list(nwfns[0]._expr.expand().as_coefficients_dict().values()))
-        mo = sorted(list(mofns[0]._expr.expand().as_coefficients_dict().values()))
+        nwfns = self.nw.basis_functions.evaluate()
+        mofns = self.mo.basis_functions.evaluate()
+        nw = sorted(list(nwfns[0].expand().as_coefficients_dict().values()))
+        mo = sorted(list(mofns[0].expand().as_coefficients_dict().values()))
         for a, b in zip(nw, mo):
             self.assertTrue(np.isclose(np.float64(a), np.float64(b)))
-        self.assertFalse(len(nwfns[11]._expr.expand().as_coefficients_dict()) ==
-                         len(mofns[11]._expr.expand().as_coefficients_dict()))
+        self.assertFalse(len(nwfns[11].expand().as_coefficients_dict()) ==
+                         len(mofns[11].expand().as_coefficients_dict()))

exatomic/algorithms/tests/test_basis.py

@@ -29,7 +29,7 @@ def setUp(self):
 
     def test_compare_fields(self):
         res = compare_fields(self.uni, self.chk, verbose=False)
-        self.assertTrue(np.isclose(len(res), sum(res)))
+        self.assertTrue(np.isclose(len(res), sum(res), rtol=5e-4))
 
 
 

exatomic/algorithms/tests/test_orbital.py

@@ -192,15 +192,11 @@ class BasisSetOrder(DataFrame):
     | prefac            | float    | prefactor (optional - for STOs)           |
     +-------------------+----------+-------------------------------------------+
     """
-    _columns = ['center', 'L', 'shell']
+    _columns = ['center', 'L']
     _index = 'chi'
     _cardinal = ('frame', np.int64)
     _categories = {'L': np.int64}
 
-    #@property
-    #def _constructor(self):
-    #    return BasisSetOrder
-
 
 class Overlap(DataFrame):
     """
@@ -231,12 +227,18 @@ class Overlap(DataFrame):
     #def _constructor(self):
     #    return Overlap
 
-    def square(self, frame=0, column='coef'):
-        """Return a 'square' matrix DataFrame of the Overlap."""
-        sq = pd.DataFrame(_square(self[column].values))
-        sq.index.name = 'chi0'
-        sq.columns.name = 'chi1'
-        return sq
+    def square(self, frame=0, column='coef', mocoefs=None):
+        """Return a 'square' matrix DataFrame of the Overlap.
+
+        Args:
+            column (str): column of coefficients to reshape
+            mocoefs (str): alias for `column`
+            frame (int): default 0
+        """
+        if mocoefs is not None: column = mocoefs
+        sq = _square(self[column].values)
+        return pd.DataFrame(sq, index=pd.Index(range(sq.shape[0]), name='chi0'),
+                            columns=pd.Index(range(sq.shape[1]), name='chi1'))
 
     @classmethod
     def from_column(cls, source):
@@ -246,11 +248,8 @@ def from_column(cls, source):
         if isinstance(source, np.ndarray):
             vals = source
         elif isinstance(source, six.string_types):
-            if os.sep in source:
-                vals = pd.read_csv(source, header=None).values.flatten()
-            else:
-            # except FileNotFoundError:
-                vals = pd.read_csv(StringIO(source), header=None).values.flatten()
+            if os.sep not in source: source = StringIO(source)
+            vals = pd.read_csv(source, header=None).values.flatten()
         else:
             # Without a catchall, _tri_indices may through UnboundLocalError
             raise TypeError("Invalid type for source: {}".format(type(source)))