diff --git a/gsrp_smart_util.py b/gsrp_smart_util.py index 7146625dc54389d5c8a09b59fcd6880f562b2547..1096884311ad2ebc2da1e9310d1b4d72a4406651 100644 --- a/gsrp_smart_util.py +++ b/gsrp_smart_util.py @@ -111,19 +111,19 @@ def add(mem1, mem2, cc, t_max, id1, id2, n_ind, mem_limit=np.infty): # assume len(id2) == 1 if (2 + len(id1)) * mem1[0].size * mem2[0].size * mem1[0].itemsize > mem_limit: return - idx1, idx2 = np.meshgrid(np.arange(len(mem1[0])), np.arange(len(mem2[0]))) - idx1, idx2 = idx1.ravel(), idx2.ravel() + idx2, idx1 = np.nonzero((mem1[1].max(0) - mem2[1][0][:, None] <= t_max) & + (mem2[1][0][:, None] - mem1[1].min(0) <= t_max)) out_tij = np.empty((len(id1) + 1, len(idx1)), mem1[1].dtype) np.take(mem1[1], idx1, axis=1, out=out_tij[:-1]) np.take(mem2[1], idx2, axis=1, out=out_tij[-1:]) - mask = reduce_all(np.abs(out_tij[:-1] - out_tij[-1:]) <= t_max) # Faster than numpy for large array - out_tij = np.compress(mask, out_tij, axis=1) - idx1, idx2 = idx1[mask], idx2[mask] out_val = mem1[0][idx1] + mem2[0][idx2] tij_dep = out_tij[:-1] - out_tij[-1:] tij_dep *= np.array([1 if i > id2 else -1 for i in id1])[:, np.newaxis] ch_dep = np.array([num_ind_dep(i, id2, n_ind) if i > id2 else num_ind_dep(id2, i, n_ind) for i in id1]) - out_val += cc[ch_dep[:, np.newaxis], tij_dep].sum(0) + if len(ch_dep) == 1: + out_val += cc[ch_dep, tij_dep[0]] + else: + out_val += np.sum(cc[c, t] for c, t in zip(ch_dep, tij_dep)) return out_val, out_tij