lanery
lanery

Reputation: 5364

Conditional Sum of 2D Numpy Arrays

I have 4 (but in reality an arbitrary amount of) 2D arrays like so:

import numpy as np

c1 = np.ones((75, 100))
c2 = np.ones((75, 100))
c3 = np.ones((75, 100))
c4 = np.ones((75, 100))

c1[22:42, 5:35] = np.random.rand(20, 30) / 2
c2[25:45, 25:55] = np.random.rand(20, 30) / 2
c3[28:48, 45:75] = np.random.rand(20, 30) / 2
c4[31:51, 65:95] = np.random.rand(20, 30) / 2

What I'd like to do is sum the arrays everywhere except for where the arrays overlap. And where there is an overlap, the value ought to be the left array. My instinct is to use np.where, but I cannot think of a clever/neat way of doing so.

Hopefully the below image makes this clear

c_arrays = np.array([c1, c2, c3, c4])
result = c_arrays.sum(axis=0)

fig, ax = plt.subplots()
ax.imshow(result)

enter image description here


Edit: I came up with an awful, recursive solution that at least shows the result I'm looking for. My hope is that someone can offer a much cleaner approach, in particular that isn't recursive

c_arrays_1 = []
for ci, cj in zip(c_arrays, c_arrays[1:]):
    c = np.where(ci + cj < 1, ci, ci + cj - 1)
    c_arrays_1.append(c)

c_arrays_2 = []
for ci, cj in zip(c_arrays_1, c_arrays_1[1:]):
    c = np.where(ci + cj < 1, ci, ci + cj - 1)
    c_arrays_2.append(c)

c_arrays_3 = []
for ci, cj in zip(c_arrays_2, c_arrays_2[1:]):
    c = np.where(ci + cj < 1, ci, ci + cj - 1)
    c_arrays_3.append(c)

fig, ax = plt.subplots()
ax.imshow(c_arrays_3[0])

enter image description here

Upvotes: 0

Views: 934

Answers (1)

Daniel F
Daniel F

Reputation: 14399

Here's a recursive method that I think fits your requirements:

def condsum(*arrs, r = 1):
    if len(arrs) == 1:
        return arrs[0]
    else:
        a = condsum(*arrs[1:], r = r)
        return np.where(a == r, arrs[0], a)

Then you'd just need to do

plt.imshow(condsum(c1, c2, c3, c4))

Upvotes: 1

Related Questions