def test_shell_dof():
assert (get_shell_dof(-3) == 7)
assert (get_shell_dof(-2) == 5)
assert (get_shell_dof(-1) == 4)
assert (get_shell_dof(0) == 1)
assert (get_shell_dof(1) == 3)
assert (get_shell_dof(2) == 6)
assert (get_shell_dof(3) == 10)
def test_shell_dof():
assert(get_shell_dof(-3)==7)
assert(get_shell_dof(-2)==5)
assert(get_shell_dof(-1)==4)
assert(get_shell_dof( 0)==1)
assert(get_shell_dof( 1)==3)
assert(get_shell_dof( 2)==6)
assert(get_shell_dof( 3)==10)
def get_poly_conversion(shell_type):
# conversion from normalised Cartesian to normalized pure
assert shell_type < -1
alpha = Symbol("alpha")
x = Symbol("x")
y = Symbol("y")
z = Symbol("z")
xyz = (x,y,z)
px = Wild("px")
py = Wild("py")
pz = Wild("pz")
c = Wild("c")
num_dof_in = get_shell_dof(-shell_type)
num_dof_out = get_shell_dof(shell_type)
lcs = numpy.zeros((num_dof_in, num_dof_out), dtype=object)
for i_out, (poly, pure_wfn_norm) in enumerate(get_polys(shell_type, alpha, xyz)):
poly = poly.expand()
if isinstance(poly, C.Add):
terms = poly.args
else:
terms = [poly]
coeffs = {}
for term in terms:
d = term.evalf(20).match(c*x**px*y**py*z**pz)
key = (int(d[px]), int(d[py]), int(d[pz]))
coeffs[key] = d[c]
for i_in, key in enumerate(iter_cartesian_powers(abs(shell_type))):
cart_wfn_norm = get_cartesian_wfn_norm(alpha, key[0], key[1], key[2])
norm_ratio = mypowsimp((cart_wfn_norm/pure_wfn_norm).evalf(20))
lc = float(coeffs.get(key, 0)*norm_ratio)
if abs(lc-int(lc)) < 1e-12:
lc = int(lc)
lcs[i_in,i_out] = lc
return lcs
def write(self, max_shell=3):
# C source code
f_c = open("%s_auto.c" % self.name, "w")
f_h = open("%s_auto.h" % self.name, "w")
f_pyf = open("%s_auto.pyf.inc" % self.name, "w")
f_c_header = open("../../HEADER.c")
f_c.write(f_c_header.read())
f_c_header.seek(0)
f_h.write(f_c_header.read())
f_c_header.close()
f_f_header = open("../../HEADER.f")
f_pyf.write(f_f_header.read())
f_f_header.close()
fn_names = []
# top lines i C file
print >> f_c, "#include <math.h>"
print >> f_c, "#include \"%s_auto.h\"" % self.name
for include in self.includes:
print >> f_c, "#include \"%s\"" % include
print >> f_c
# top lines in H file
guard_name = "%s_AUTO_H" % self.name.upper()
print >> f_h, "#ifndef %s" % guard_name
print >> f_h, "#define %s" % guard_name
print >> f_h
print >> f_h, "#define MAX_SHELL %i" % max_shell
print >> f_h, "#define NUM_SHELL_TYPES %i" % (2*max_shell+1)
print >> f_h, "#define MAX_SHELL_DOF %i" % max(get_shell_dof(shell_type) for shell_type in xrange(-max_shell, max_shell+1))
print >> f_h, "#define CHECK_ALLOC(pointer) if (pointer==NULL) {result = -1; goto EXIT; }"
print >> f_h, "#define CHECK_SHELL(shell_type) if (abs(shell_type) > MAX_SHELL) { result = -2; goto EXIT; }"
print >> f_h, "#define GET_SHELL_DOF(shell_type) ((shell_type<-1)?(-2*shell_type+1):((shell_type==-1)?(4):(((shell_type+1)*(shell_type+2))/2)))"
print >> f_h
# add all compute routines
for st_row in self.iter_shell_types(max_shell):
print st_row
fn_name = self.write_routine(f_pyf, f_c, f_h, st_row)
fn_names.append(fn_name)
# collect all compute routines in an array
arg_c_types = []
for ag in self.arg_groups:
for arg in ag.args:
arg_c_types.append(arg.get_c_type())
print >> f_c, "typedef void (*fntype)(%s, double*);" % (", ".join(arg_c_types))
print >> f_c, "static const fntype fns[%i] = {%s};" % (len(fn_names), ", ".join(fn_names))
print >> f_c
# dispatch routine that takes the right function from the array
all_c_types_names = []
switches = []
for ag in self.arg_groups:
switch = ag.get_switch_name()
if switch is not None:
all_c_types_names.append("int %s" % switch)
switches.append(switch)
all_c_types_names.append(ag.get_c_types_names())
fn_name = "%s_dispatch" % self.name
c_names = self.get_c_names()
print >> f_h, "void %s_dispatch(%s, double* out);" % (self.name, ", ".join(all_c_types_names))
print >> f_c, "void %s_dispatch(%s, double* out)" % (self.name, ", ".join(all_c_types_names))
print >> f_c, "{"
factor = 1
offsets = []
for switch in switches:
if factor == 1:
offsets.append("%i+%s" % (max_shell, switch))
else:
offsets.append("%i*(%i+%s)" % (factor, max_shell, switch))
factor *= 2*max_shell + 1
print >> f_c, " fns[%s](%s, out);" % ("+".join(offsets), c_names)
print >> f_c, "}"
# close the guard
print >> f_h
print >> f_h, "#endif"
# close the files
f_c.close()
f_h.close()
f_pyf.close()
