Source code for pyscf.dft.xc_deriv

#!/usr/bin/env python
# -*- coding: utf-8 -*-
# Copyright 2022 The PySCF Developers. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
#

'''
Transform XC functional derivatives between different representations
'''

import math
import itertools
from functools import lru_cache
import ctypes
import numpy as np
from pyscf import lib

libdft = lib.load_library('libdft')


[docs] def transform_vxc(rho, vxc, xctype, spin=0): r''' Transform libxc functional derivatives to the derivative tensor of parameters in rho: [[density_a, (nabla_x)_a, (nabla_y)_a, (nabla_z)_a, tau_a], [density_b, (nabla_x)_b, (nabla_y)_b, (nabla_z)_b, tau_b]]. The output tensor has the shape: * spin polarized LDA : [2,1,N] GGA : [2,4,N] MGGA: [2,5,N] * spin unpolarized LDA : [1,N] GGA : [4,N] MGGA: [5,N] ''' rho = np.asarray(rho, order='C') if xctype == 'GGA': order = 1 nvar = 4 fr = vxc[0].T fg = vxc[1].T elif xctype == 'MGGA': order = 2 nvar = 5 fr = vxc[0].T fg = vxc[1].T ft = vxc[3].T else: # LDA order = 0 nvar = 1 fr = vxc[0].T ngrids = rho.shape[-1] if spin == 1: if order == 0: vp = fr.reshape(2, nvar, ngrids) else: vp = np.empty((2,nvar, ngrids)) vp[:,0] = fr #:vp[:,1:4] = np.einsum('abg,bxg->axg', _stack_fg(fg), rho[:,1:4]) vp[:,1:4] = _stack_fg(fg, rho=rho) if order > 1: vp[:,4] = ft else: if order == 0: vp = fr.reshape(nvar, ngrids) else: vp = np.empty((nvar, ngrids)) vp[0] = fr vp[1:4] = 2 * fg * rho[1:4] if order > 1: vp[4] = ft return vp
[docs] def transform_fxc(rho, vxc, fxc, xctype, spin=0): r''' Transform libxc functional derivatives to the derivative tensor of parameters in rho: [[density_a, (nabla_x)_a, (nabla_y)_a, (nabla_z)_a, tau_a], [density_b, (nabla_x)_b, (nabla_y)_b, (nabla_z)_b, tau_b]]. The output tensor has the shape: * spin polarized LDA : [2,1,2,1,N] GGA : [2,4,2,4,N] MGGA: [2,5,2,5,N] * spin unpolarized LDA : [1,1,N] GGA : [4,4,N] MGGA: [5,5,N] ''' rho = np.asarray(rho, order='C') if xctype == 'GGA': order = 1 nvar = 4 fg = vxc[1].T frr = fxc[0].T frg = fxc[1].T fgg = fxc[2].T elif xctype == 'MGGA': order = 2 nvar = 5 fg = vxc[1].T frr, frg, fgg, ftt, frt, fgt = [fxc[i].T for i in [0, 1, 2, 4, 6, 9]] else: # LDA order = 0 nvar = 1 frr = fxc[0].T ngrids = rho.shape[-1] if spin == 1: if order == 0: vp = _stack_frr(frr).reshape(2,nvar, 2,nvar, ngrids).transpose(1,3,0,2,4) else: vp = np.empty((2,nvar, 2,nvar, ngrids)).transpose(1,3,0,2,4) vp[0,0] = _stack_frr(frr) i3 = np.arange(3) #:qgg = _stack_fgg(fgg) #:qgg = np.einsum('abcdg,axg->xbcdg', qgg, rho[:,1:4]) #:qgg = np.einsum('xbcdg,cyg->xybdg', qgg, rho[:,1:4]) qgg = _stack_fgg(fgg, rho=rho).transpose(1,3,0,2,4) qgg[i3,i3] += _stack_fg(fg) vp[1:4,1:4] = qgg frg = frg.reshape(2,3,ngrids) #:qrg = _stack_fg(frg, axis=1) #:qrg = np.einsum('rabg,axg->xrbg', qrg, rho[:,1:4]) qrg = _stack_fg(frg, axis=1, rho=rho).transpose(2,0,1,3) vp[0,1:4] = qrg vp[1:4,0] = qrg.transpose(0,2,1,3) if order > 1: fgt = fgt.reshape(3,2,ngrids) #:qgt = _stack_fg(fgt, axis=0) #:qgt = np.einsum('abrg,axg->xbrg', qgt, rho[:,1:4]) qgt = _stack_fg(fgt, axis=0, rho=rho).transpose(1,0,2,3) vp[1:4,4] = qgt vp[4,1:4] = qgt.transpose(0,2,1,3) qrt = frt.reshape(2,2,ngrids) vp[0,4] = qrt vp[4,0] = qrt.transpose(1,0,2) vp[4,4] = _stack_frr(ftt) vp = vp.transpose(2,0,3,1,4) else: if order == 0: vp = frr.reshape(nvar, nvar, ngrids) else: vp = np.empty((nvar, nvar, ngrids)) vp[0,0] = frr i3 = np.arange(3) qgg = 4 * fgg * rho[1:4] * rho[1:4,None] qgg[i3,i3] += fg * 2 vp[1:4,1:4] = qgg vp[0,1:4] = vp[1:4,0] = 2 * frg * rho[1:4] if order > 1: vp[4,1:4] = vp[1:4,4] = 2 * fgt * rho[1:4] vp[0,4] = frt vp[4,0] = frt vp[4,4] = ftt return vp
[docs] def transform_kxc(rho, fxc, kxc, xctype, spin=0): r''' Transform libxc functional derivatives to the derivative tensor of parameters in rho: [[density_a, (nabla_x)_a, (nabla_y)_a, (nabla_z)_a, tau_a], [density_b, (nabla_x)_b, (nabla_y)_b, (nabla_z)_b, tau_b]]. The output tensor has the shape: * spin polarized LDA : [2,1,2,1,2,1,N] GGA : [2,4,2,4,2,4,N] MGGA: [2,5,2,5,2,5,N] * spin unpolarized LDA : [1,1,1,N] GGA : [4,4,4,N] MGGA: [5,5,5,N] ''' rho = np.asarray(rho, order='C') if xctype == 'GGA': order = 1 nvar = 4 frg = fxc[1].T fgg = fxc[2].T frrr, frrg, frgg, fggg = [x.T for x in kxc[:4]] elif xctype == 'MGGA': order = 2 nvar = 5 frg = fxc[1].T fgg = fxc[2].T fgt = fxc[9].T frrr, frrg, frgg, fggg, frrt, frgt, frtt, fggt, fgtt, fttt = \ [kxc[i].T for i in [0, 1, 2, 3, 5, 7, 10, 12, 15, 19]] else: # LDA order = 0 nvar = 1 frrr = kxc[0].T ngrids = rho.shape[-1] if spin == 1: if order == 0: vp = _stack_frrr(frrr).reshape(2,nvar, 2,nvar, 2,nvar, ngrids).transpose(1,3,5,0,2,4,6) else: vp = np.empty((2,nvar, 2,nvar, 2,nvar, ngrids)).transpose(1,3,5,0,2,4,6) vp[0,0,0] = _stack_frrr(frrr) i3 = np.arange(3) #:qggg = _stack_fggg(fggg) #:qggg = np.einsum('abcdefg,axg->xbcdefg', qggg, rho[:,1:4]) #:qggg = np.einsum('xbcdefg,cyg->xybdefg', qggg, rho[:,1:4]) #:qggg = np.einsum('xybdefg,ezg->xyzbdfg', qggg, rho[:,1:4]) qggg = _stack_fggg(fggg, rho=rho).transpose(1,3,5,0,2,4,6) qgg = _stack_fgg(fgg) qgg = np.einsum('abcdg,axg->xbcdg', qgg, rho[:,1:4]) for i in range(3): qggg[:,i,i] += qgg qggg[i,:,i] += qgg.transpose(0,2,1,3,4) qggg[i,i,:] += qgg.transpose(0,2,3,1,4) vp[1:4,1:4,1:4] = qggg frgg = frgg.reshape(2,6,ngrids) #:qrgg = _stack_fgg(frgg, axis=1) #:qrgg = np.einsum('rabcdg,axg->xrbcdg', qrgg, rho[:,1:4]) #:qrgg = np.einsum('xrbcdg,cyg->xyrbdg', qrgg, rho[:,1:4]) qrgg = _stack_fgg(frgg, axis=1, rho=rho).transpose(2,4,0,1,3,5) qrg = _stack_fg(frg.reshape(2,3,ngrids), axis=1) qrgg[i3,i3] += qrg vp[0,1:4,1:4] = qrgg vp[1:4,0,1:4] = qrgg.transpose(0,1,3,2,4,5) vp[1:4,1:4,0] = qrgg.transpose(0,1,3,4,2,5) frrg = frrg.reshape(3,3,ngrids) qrrg = _stack_frr(frrg, axis=0) #:qrrg = _stack_fg(qrrg, axis=2) #:qrrg = np.einsum('rsabg,axg->rsxbg', qrrg, rho[:,1:4]) qrrg = _stack_fg(qrrg, axis=2, rho=rho).transpose(3,0,1,2,4) vp[0,0,1:4] = qrrg vp[0,1:4,0] = qrrg.transpose(0,1,3,2,4) vp[1:4,0,0] = qrrg.transpose(0,3,1,2,4) if order > 1: fggt = fggt.reshape(6,2,ngrids) #:qggt = _stack_fgg(fggt, axis=0) #:qggt = np.einsum('abcdrg,axg->xbcdrg', qggt, rho[:,1:4]) #:qggt = np.einsum('xbcdrg,cyg->xybdrg', qggt, rho[:,1:4]) qggt = _stack_fgg(fggt, axis=0, rho=rho).transpose(1,3,0,2,4,5) qgt = _stack_fg(fgt.reshape(3,2,ngrids), axis=0) i3 = np.arange(3) qggt[i3,i3] += qgt vp[1:4,1:4,4] = qggt vp[1:4,4,1:4] = qggt.transpose(0,1,2,4,3,5) vp[4,1:4,1:4] = qggt.transpose(0,1,4,2,3,5) qgtt = _stack_frr(fgtt.reshape(3,3,ngrids), axis=1) #:qgtt = _stack_fg(qgtt, axis=0) #:qgtt = np.einsum('abrsg,axg->xbrsg', qgtt, rho[:,1:4]) qgtt = _stack_fg(qgtt, axis=0, rho=rho).transpose(1,0,2,3,4) vp[1:4,4,4] = qgtt vp[4,1:4,4] = qgtt.transpose(0,2,1,3,4) vp[4,4,1:4] = qgtt.transpose(0,2,3,1,4) frgt = frgt.reshape(2,3,2,ngrids) #:qrgt = _stack_fg(frgt, axis=1) #:qrgt = np.einsum('rabsg,axg->xrbsg', qrgt, rho[:,1:4]) qrgt = _stack_fg(frgt, axis=1, rho=rho).transpose(2,0,1,3,4) vp[0,1:4,4] = qrgt vp[0,4,1:4] = qrgt.transpose(0,1,3,2,4) vp[1:4,0,4] = qrgt.transpose(0,2,1,3,4) vp[4,0,1:4] = qrgt.transpose(0,3,1,2,4) vp[1:4,4,0] = qrgt.transpose(0,2,3,1,4) vp[4,1:4,0] = qrgt.transpose(0,3,2,1,4) qrrt = _stack_frr(frrt.reshape(3,2,ngrids), axis=0) vp[0,0,4] = qrrt vp[0,4,0] = qrrt.transpose(0,2,1,3) vp[4,0,0] = qrrt.transpose(2,0,1,3) qrtt = _stack_frr(frtt.reshape(2,3,ngrids), axis=1) vp[0,4,4] = qrtt vp[4,0,4] = qrtt.transpose(1,0,2,3) vp[4,4,0] = qrtt.transpose(1,2,0,3) vp[4,4,4] = _stack_frrr(fttt, axis=0) vp = vp.transpose(3,0,4,1,5,2,6) else: if order == 0: vp = frrr.reshape(nvar, nvar, nvar, ngrids) else: vp = np.empty((nvar, nvar, nvar, ngrids)) vp[0,0,0] = frrr i3 = np.arange(3) qggg = 8 * fggg * rho[1:4] * rho[1:4,None] * rho[1:4,None,None] qgg = 4 * fgg * rho[1:4] for i in range(3): qggg[i,i,:] += qgg qggg[i,:,i] += qgg qggg[:,i,i] += qgg vp[1:4,1:4,1:4] = qggg qrgg = 4 * frgg * rho[1:4] * rho[1:4,None] qrgg[i3,i3] += frg * 2 vp[0,1:4,1:4] = qrgg vp[1:4,0,1:4] = qrgg vp[1:4,1:4,0] = qrgg qrrg = 2 * frrg * rho[1:4] vp[0,0,1:4] = qrrg vp[0,1:4,0] = qrrg vp[1:4,0,0] = qrrg if order > 1: qggt = 4 * fggt * rho[1:4] * rho[1:4,None] qggt[i3,i3] += fgt * 2 vp[1:4,1:4,4] = qggt vp[1:4,4,1:4] = qggt vp[4,1:4,1:4] = qggt qgtt = 2 * fgtt * rho[1:4] vp[1:4,4,4] = qgtt vp[4,1:4,4] = qgtt vp[4,4,1:4] = qgtt qrgt = 2 * frgt * rho[1:4] vp[0,1:4,4] = qrgt vp[0,4,1:4] = qrgt vp[1:4,0,4] = qrgt vp[4,0,1:4] = qrgt vp[1:4,4,0] = qrgt vp[4,1:4,0] = qrgt vp[0,0,4] = frrt vp[0,4,0] = frrt vp[4,0,0] = frrt vp[0,4,4] = frtt vp[4,0,4] = frtt vp[4,4,0] = frtt vp[4,4,4] = fttt return vp
[docs] def transform_lxc(rho, fxc, kxc, lxc, xctype, spin=0): r''' Transform libxc vxc functional output to the derivative tensor of parameters in rho: [density, nabla_x, nabla_y, nabla_z, tau]. The output tensor has the shape: * spin polarized LDA : [2,1,2,1,2,1,2,1,N] GGA : [2,4,2,4,2,4,2,4,N] MGGA: [2,5,2,5,2,5,2,5,N] * spin unpolarized LDA : [1,1,1,1,N] GGA : [4,4,4,4,N] MGGA: [5,5,5,5,N] ''' raise NotImplementedError
def _stack_fg(fg, axis=0, rho=None, out=None): '''fg [uu, ud, dd] -> [[uu*2, ud], [du, dd*2]]''' if rho is None: qg = _stack_frr(fg, axis) if axis == 0: qg[0,0] *= 2 qg[1,1] *= 2 elif axis == 1: qg[:,0,0] *= 2 qg[:,1,1] *= 2 elif axis == 2: qg[:,:,0,0] *= 2 qg[:,:,1,1] *= 2 else: raise NotImplementedError return qg else: rho = np.asarray(rho, order='C') fg = np.asarray(fg, order='C') if fg.shape[axis] != 3: fg = fg.reshape(fg.shape[:axis] + (3, -1) + fg.shape[axis+1:]) nvar, ngrids = rho.shape[1:] ncounts = int(np.prod(fg.shape[:axis])) mcounts = int(np.prod(fg.shape[axis+1:-1])) qg = np.empty(fg.shape[:axis] + (2, 3) + fg.shape[axis+1:]) rho = libdft.VXCfg_to_direct_deriv( qg.ctypes.data_as(ctypes.c_void_p), fg.ctypes.data_as(ctypes.c_void_p), rho.ctypes.data_as(ctypes.c_void_p), ctypes.c_int(ncounts), ctypes.c_int(nvar), ctypes.c_int(mcounts), ctypes.c_int(ngrids)) return qg def _stack_frr(frr, axis=0): '''frr [u_u, u_d, d_d] -> [[u_u, u_d], [d_u, d_d]]''' if frr.shape[axis] != 3: frr = frr.reshape(frr.shape[:axis] + (3, -1) + frr.shape[axis+1:]) slices = [slice(None)] * frr.ndim slices[axis] = [[0,1],[1,2]] return frr[tuple(slices)] def _stack_fgg(fgg, axis=0, rho=None): ''' fgg [uu_uu, uu_ud, uu_dd, ud_ud, ud_dd, dd_dd] -> [[uu_uu, ud_ud, ud_dd], [ud_uu, ud_ud, ud_dd], [dd_uu, dd_ud, dd_dd]] -> tensor with shape [2,2, 2,2, ...] ''' if fgg.shape[axis] != 6: fgg = fgg.reshape(fgg.shape[:axis] + (6, -1) + fgg.shape[axis+1:]) slices = [slice(None)] * fgg.ndim slices[axis] = [[0, 1, 2], [1, 3, 4], [2, 4, 5]] fgg = fgg[tuple(slices)] return _stack_fg(_stack_fg(fgg, axis=axis+1, rho=rho), axis=axis, rho=rho) def _stack_frrr(frrr, axis=0): ''' frrr [u_u_u, u_u_d, u_d_d, d_d_d] -> tensor with shape [2, 2, 2, ...] ''' if frrr.shape[axis] != 4: frrr = frrr.reshape(frrr.shape[:axis] + (4, -1) + frrr.shape[axis+1:]) slices = [slice(None)] * frrr.ndim slices[axis] = [[[0, 1], [1, 2]], [[1, 2], [2, 3]]] return frrr[tuple(slices)] def _stack_fggg(fggg, axis=0, rho=None): ''' fggg [uu_uu_uu, uu_uu_ud, uu_uu_dd, uu_ud_ud, uu_ud_dd, uu_dd_dd, ud_ud_ud, ud_ud_dd, ud_dd_dd, dd_dd_dd] -> tensor with shape [2,2, 2,2, 2,2, ...] ''' if fggg.shape[axis] != 10: fggg = fggg.reshape(fggg.shape[:axis] + (10, 2) + fggg.shape[axis+1:]) slices = [slice(None)] * fggg.ndim slices[axis] = [[[0, 1, 2], [1, 3, 4], [2, 4, 5]], [[1, 3, 4], [3, 6, 7], [4, 7, 8]], [[2, 4, 5], [4, 7, 8], [5, 8, 9]]] fggg = fggg[tuple(slices)] fggg = _stack_fg(fggg, axis=axis+2, rho=rho) fggg = _stack_fg(fggg, axis=axis+1, rho=rho) return _stack_fg(fggg, axis=axis, rho=rho) @lru_cache(100) def _product_uniq_indices(nvars, order): ''' Indexing the symmetry unique elements in cartesian product ''' # n = 0 # for i range(nvars): # for j in range(i, nvars): # for k in range(k, nvars): # ... # prod[(i,j,k,...)] = n # n += 1 prod = np.ones([nvars] * order, dtype=int) uniq_idx = list(itertools.combinations_with_replacement(range(nvars), order)) prod[tuple(np.array(uniq_idx).T)] = np.arange(len(uniq_idx)) idx = np.where(prod) sidx = np.sort(idx, axis=0) prod[idx] = prod[tuple(sidx)] return prod def _pair_combinations(lst): n = len(lst) if n <= 2: return [[tuple(lst)]] a = lst[0] results = [] for i in range(1, n): pair = (a, lst[i]) rests = _pair_combinations(lst[1:i]+lst[i+1:]) for rest in rests: rest.append(pair) results.extend(rests) return results def _unfold_gga(rho, xc_val, spin, order, nvar, xlen, reserve=0): assert nvar >= 4 ngrids = rho.shape[-1] n_transform = order - reserve if spin == 0: nvar2 = nvar drv = libdft.VXCunfold_sigma_spin0 else: nvar2 = nvar * 2 drv = libdft.VXCunfold_sigma_spin1 # xc_val[idx] expands unique xc elements to n-order tensors idx = _product_uniq_indices(xlen, order) xc_tensor = np.empty((xlen**reserve * nvar2**n_transform, ngrids)) xc_tensor[:idx.size] = xc_val[idx.ravel()] buf = np.empty_like(xc_tensor) for i in range(n_transform): xc_tensor, buf = buf, xc_tensor ncounts = xlen**(order-1-i) * nvar2**i drv(xc_tensor.ctypes, buf.ctypes, rho.ctypes, ctypes.c_int(ncounts), ctypes.c_int(nvar), ctypes.c_int(ngrids)) if spin == 0: xc_tensor = xc_tensor.reshape([xlen]*reserve + [nvar]*n_transform + [ngrids]) else: xc_tensor = xc_tensor.reshape([xlen]*reserve + [2,nvar]*n_transform + [ngrids]) return xc_tensor def _diagonal_indices(idx, order): assert order > 0 n = len(idx) diag_idx = [np.asarray(idx)] for i in range(1, order): last_dim = diag_idx[-1] diag_idx = [np.repeat(idx, n) for x in diag_idx] diag_idx.append(np.tile(last_dim, n)) # repeat(diag_idx, 2) return tuple(x for x in diag_idx for i in range(2)) _XC_NVAR = { ('HF', 0): (1, 1), ('HF', 1): (1, 2), ('LDA', 0): (1, 1), ('LDA', 1): (1, 2), ('GGA', 0): (4, 2), ('GGA', 1): (4, 5), ('MGGA', 0): (5, 3), ('MGGA', 1): (5, 7), }
[docs] def transform_xc(rho, xc_val, xctype, spin, order): '''General transformation to construct XC derivative tensor''' xc_val = np.asarray(xc_val, order='C') if order == 0: return xc_val[0] nvar, xlen = _XC_NVAR[xctype, spin] ngrids = xc_val.shape[-1] offsets = [0] + [count_combinations(xlen, o) for o in range(order+1)] p0, p1 = offsets[order:order+2] if xctype == 'LDA' or xctype == 'HF': xc_out = xc_val[p0:p1] if spin == 0: return xc_out.reshape([1]*order + [ngrids]) else: idx = _product_uniq_indices(xlen, order) return xc_out[idx].reshape([2,1]*order + [ngrids]) rho = np.asarray(rho, order='C') assert rho.size == (spin+1) * nvar * ngrids xc_tensor = _unfold_gga(rho, xc_val[p0:p1], spin, order, nvar, xlen) nabla_idx = np.arange(1, 4) if spin == 0: for n_pairs in range(1, order//2+1): p0, p1 = offsets[order-n_pairs:order-n_pairs+2] xc_sub = _unfold_gga(rho, xc_val[p0:p1], spin, order-n_pairs, nvar, xlen, n_pairs) # Just the sigma components xc_sub = xc_sub[(1,)*n_pairs] * 2**n_pairs # low_sigmas indicates the index for the extra sigma terms low_sigmas = itertools.combinations(range(order), n_pairs*2) pair_combs = [list(itertools.chain(*p[::-1])) for p in _pair_combinations(list(range(n_pairs*2)))] diag_idx = _diagonal_indices(nabla_idx, n_pairs) for dim_lst in low_sigmas: rest_dims = [i for i in range(xc_tensor.ndim) if i not in dim_lst] for pair_comb in pair_combs: xc_tensor_1 = xc_tensor.transpose( [dim_lst[i] for i in pair_comb] + rest_dims) xc_tensor_1[diag_idx] += xc_sub else: i3to2x2 = _product_uniq_indices(2, 2) for n_pairs in range(1, order//2+1): p0, p1 = offsets[order-n_pairs:order-n_pairs+2] xc_sub = _unfold_gga(rho, xc_val[p0:p1], spin, order-n_pairs, nvar, xlen, n_pairs) # Just the sigma components xc_sub = xc_sub[(slice(2,5),) * n_pairs] for i in range(n_pairs): xc_sub[(slice(None),)*i+(0,)] *= 2 xc_sub[(slice(None),)*i+(2,)] *= 2 sigma_idx = (i3to2x2[(slice(None),)*2 + (np.newaxis,)*(i*2)] for i in reversed(range(n_pairs))) xc_sub = xc_sub[tuple(sigma_idx)] low_sigmas = itertools.combinations(range(order), n_pairs*2) pair_combs = [list(itertools.chain(*p[::-1])) for p in _pair_combinations(list(range(n_pairs*2)))] diag_idx = _diagonal_indices(nabla_idx, n_pairs) for dim_lst in low_sigmas: dim_lst2 = np.array(dim_lst)*2 + np.array([1, 0])[:,None] rest_dims = [i for i in range(xc_tensor.ndim) if i not in dim_lst2] for pair_comb in pair_combs: leading_dims = list(dim_lst2[:,pair_comb].ravel()) xc_tensor_1 = xc_tensor.transpose(leading_dims + rest_dims) xc_tensor_1[diag_idx] += xc_sub return xc_tensor
[docs] def count_combinations(nvar, order): '''sum(len(combinations_with_replacement(range(nvar), o) for o in range(order))''' return lib.comb(nvar+order, order)
[docs] def ud2ts(v_ud): '''XC derivatives spin-up/spin-down representations to total-density/spin-density representations''' v_ud = np.asarray(v_ud, order='C') v_ts = np.empty_like(v_ud) nvar, ngrids = v_ts.shape[-2:] nvar2 = nvar * 2 order = v_ud.ndim // 2 drv = libdft.VXCud2ts v_ts, v_ud = v_ud, v_ts for i in range(order): v_ts, v_ud = v_ud, v_ts n = nvar2**i nvg = nvar * nvar2**(order-1-i) * ngrids #:c = np.array([[.5, .5], # vrho = (va + vb) / 2 #: [.5, -.5]]) # vs = (va - vb) / 2 #:v_ts = np.einsum('ra,...ax...g->...rx...g', c, v_ud) drv(v_ts.ctypes, v_ud.ctypes, ctypes.c_int(n), ctypes.c_int(nvg)) return v_ts
[docs] def ts2ud(v_ts): '''XC derivatives total-density/spin-density representations to spin-up/spin-down representations''' v_ts = np.asarray(v_ts, order='C') v_ud = np.empty_like(v_ts) nvar, ngrids = v_ts.shape[-2:] nvar2 = nvar * 2 order = v_ud.ndim // 2 drv = libdft.VXCts2ud v_ts, v_ud = v_ud, v_ts for i in range(order): v_ts, v_ud = v_ud, v_ts n = nvar2**i nvg = nvar * nvar2**(order-1-i) * ngrids #:c = np.array([[1, 1], # va = vrho + vs #: [1, -1]]) # vb = vrho - vs #:v_ud = np.einsum('ra,...ax...g->...rx...g', c, v_ts) drv(v_ud.ctypes, v_ts.ctypes, ctypes.c_int(n), ctypes.c_int(nvg)) return v_ud