Generate a dataframe with conditional combination of range of numbers

here is a cleaned up version of your first solution that I've verified returns the same results (with dummy definition for cal_P) that should be relatively performant while avoiding memory issues:

def cal_P(*x):
    return np.sum(np.array(x)**2)

%%timeit
x1=10
y1=8
m=7
P_List = []
m_range = np.arange(0, m+1)
for m in m_range:
    records = []
    for h in np.arange(0, x1+1):
        for i in np.arange(0, max(0, (x1-h+1))):
            j = m-i
            if j>=0:
                for k in np.arange(0, max(0, y1-j)+1):
                    arguments = { "h":h, "i":i, "j":j, "k":k}
                    records.append(arguments)
    df = pd.DataFrame(records)
    df["P"]= df.apply(cal_P, axis=1)
    P_List.append(df.P.sum())

P_List (both versions): [5709, 9403, 11840, 13556, 14881, 15955, 16744, 17056].

• original (1st): 2.29 s ± 147 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

• improved: 63.8 ms ± 219 µs per loop (mean ± std. dev. of 7 runs, 10 loops each)

I suspect df.append is what was slow. In general, single row operations, esp. when modifying the df size, will be expensive.

The best way to speed up your code is to profile it (e.g. cProfile) to see what's slow - it's often something easy to change. If you really need high performance, maybe iterate over one more index before invoking meshgrid.

I converted your second code sample to generate df into the following function:

def f1(x1, y1, m):
    m_range = np.arange(0, m  + 1)
    h_range = np.arange(0, x1 + 1)
    i_range = np.arange(0, x1 + 1)
    k_range = np.arange(0, y1 + x1 + 2)
    j_range = np.arange(0, x1 + m  + 1)
    data = np.array(np.meshgrid(m_range, h_range, i_range,
        j_range, k_range)).T.reshape(-1, 5)
    df = pd.DataFrame(data=data, columns=['m', 'h', 'i', 'j', 'k'])
    i_cond = df.i <= (x1-df.h)
    k_cond = df.k <= (y1-(df.m-df.i))
    j_cond = df.j == (df.m-df.i)
    rowNo1 = df.index.size
    df = df.drop(df[~(i_cond & k_cond & j_cond)].index)
    rowNo2 = df.index.size
    return df, rowNo1, rowNo2

Note that it returns 3 results:

• df - the DataFrame itself,
• rowNo1 - the number of rows before drop,
• rowNo2 - the number of rows after drop.

I executed it on your data:

df, rN1, rN2 = f1(10, 8, 7)
print(f'{rN1}, {rN2}')

getting:

348480, 2010

Note that rowNo1 is quite big.

Note also that inclusion of j_range in the meshgrid can be omitted and then the row number is substantially smaller.

Column named j can be computed after df has been generated, using df.m - df.i formula.

Then "unnecessary" rows can be eliminated, this time using df.m >= df.i formula.

So my proposition how to generate the source DataFrame is:

def f2(x1, y1, m):
    m_range = np.arange(0, m  + 1)
    h_range = np.arange(0, x1 + 1) 
    i_range = np.arange(0, x1 + 1)
    k_range = np.arange(0, y1 + x1 + 2)
    data = np.array(np.meshgrid(m_range, h_range, i_range,
        k_range)).T.reshape(-1, 4)
    df = pd.DataFrame(data=data, columns=["m", "h", "i", "k"])
    df.insert(loc=3, column='j', value=df.m - df.i)
    i_cond = df.i <= (x1-df.h) 
    k_cond = df.k <= (y1-(df.m-df.i))
    j_cond = df.m >= df.i
    rowNo1 = df.index.size
    df = df.drop(df[~(i_cond & k_cond & j_cond)].index)
    rowNo2 = df.index.size
    return df, rowNo1, rowNo2

For comparison, I still compute row numbers before and after drop.

Then I executed:

df2, rN1, rN2 = f2(10, 8, 7)
print(f'{rN1}, {rN2}')

getting:

19360, 2010

Note that this time rowNo1 is substantially smaller, so there is a better chance that this function will not exhaust available memory on greater values of x1, y1 and m.

To check that both results (df and df2) are the same, I ran:

df.sort_values(['m', 'h', 'i', 'j', 'k'], ignore_index=True).equals(
    df2.sort_values(['m', 'h', 'i', 'j', 'k'], ignore_index=True))

getting True.

And the last important factor is the execution speed. Using %timeit I got:

• 120 ms for your code (f1),
• 7.69 ms for my code (f2) - over 15 times faster.

Try my function, without returning both row numbers, on your real data.

I noticed that my first solution still leads to a Memory exhausted exception. So I came up with other, "2-step" solution, more resistant to memory exhaustion:

• First h and i columns are generated, based on np.triu_indices function.
• Then, for each group with particular i value, there are added j and k columns.

The code is as follows:

1. A function to generate h and i columns:

   def getHI(x1):
       ''' Generate "h" and "i" columns '''
       h, i = np.triu_indices(x1 + 1)
       return pd.DataFrame(np.vstack((h, x1 - i)).T, columns=['h', 'i'])
   

2. A function to add j and k columns, for the current group of rows (for some i):

   def addJK(grp, m, y1):
       ''' Add "j" and "k" columns to the current group '''
       i = grp.iloc[0,1]    # "i" column from the first row
       j = m - i
       if (y1 - j < 0) or (j < 0): return None
       kRng = y1 - j + 1
       res = pd.DataFrame(np.hstack((np.tile(grp.values, (kRng, 1)),
           np.repeat(np.arange(kRng), grp.index.size)[:, np.newaxis])),
           columns=['h', 'i', 'k'])
       res.insert(2, 'j', j)  # Add "j" column
       return res
   

3. A new version of your cal_P:

   def cal_P2(hijk):
       return (hijk.h + hijk.i + hijk.j + hijk.k).astype(float).sum()
   

   The most important change is that the computation is vectorized, (the source object is a whole DataFrame, not individual row). This way you avoid apply, which operates much slower.

   Another change is that the result is of float type, because in your target version the result will be just float.

   Of course, your "true" cal_P2 function must be also vectorized, otherwise the execution time will be much longer.

4. A function generating the sum of probabilities, for given x1, y1 and m values:

   def getP(x1, y1, m):
       df_HI = getHI(x1)
       df_HIJK = df_HI.groupby('i', group_keys=False).apply(addJK, m=m, y1=y1)
       return cal_P2(df_HIJK)
   

5. A function generating a list of sum of probabilities, for each m from 0 to mm:

   def f2(x1, y1, mm):
       return [ getP(x1, y1, m) for m in range(mm + 1) ]
   

Note that because m is used inside the loop, I named the top value of the loop as mm, so the code is more tidy.

I tested this function on your test data:

f2(10, 8, 7)  

getting:

[891.0, 1691.0, 2380.0, 2940.0, 3355.0, 3611.0, 3696.0, 3600.0]

Compare it with the result generated by your code. Compare also the execution times.

Then check the execution time for gradually increasing values of x1, y1 and mm. I'm have no idea whether it is the best possible solution, but at least you can treat it as one more option to try and compare with other options you had so far.

An alternative to consider is to compute getP (the "individual" probability, for particular x1, y1 and m values), at least during testing the execution time.