'Script for making contracted basis set in python

I'm trying to make a script which generates a contracted basis set that I need in my scientific calculations. The output file contains basically 2D grid of numbers. In the first column I have exponents and the second column and columns after that contain contraction coefficients. In the example below I have four functions (that correspond to 4 exponents) and in this contraction scheme these four functions are contracted into two functions in such a way that in the first function I have 3 exponents and hence 3 nonzero coefficients. The second function consist only of 1 function and hence 1 nonzero coefficient. So, 4 functions contracted into two as 3 + 1 it would look like this

exponent1    coefficient1 0

exponent2    coefficient2 0

exponent3    coefficient3 0

exponent4    0 coefficient4

so that the non-zero numbers in the column are contracted together. In my script I of course try to make a general scheme so that if I have n exponents I can make a contraction scheme where I can have any kind of scheme between 1 and n functions.

The script is as follows now

#!/usr/bin/python3

def contract(e, i, c, n):
        """
        e = list of exponents == number of functions.
        Functions e are contracted to i functions whose coefficients
        are determined by c and n is a list which tells how the contraction
        is done (e.g. n = [3, 1] --> functions contracted into two as 3 + 1.)
        """

        num = len(e)
        grid = [[0 for i in range(i + 1)] for x in range(num)]

        for num1, row1 in enumerate(grid):
                row1[0] = e[num1] #add exponents

        for g in grid:
                print(g)

i = 2
e = [0, 1, 2, 3]
c = [4, 5, 6, 7]
n = [3, 1]

contract(e, i, c, n)

This works so far and the output this produces now is:

[0, 0, 0]
[1, 0, 0]
[2, 0, 0]
[3, 0, 0]

The output should be:

[0, 4, 0]
[1, 5, 0]
[2, 6, 0]
[3, 0, 7]

The problem is that I don't know how I could do the rest of the code. So how could I get the coefficients to correct places? I will then somehow print the numbers in the list so the final output file should in this case look like:

0 4 0
1 5 0
2 6 0
3 0 7

Does anyone have an idea how I could do this?



Sources

This article follows the attribution requirements of Stack Overflow and is licensed under CC BY-SA 3.0.

Source: Stack Overflow

Solution Source