CODEWITHSUNDEEP
pythonnumpynumerical-methods
GlassSeaHorse
GlassSeaHorse

Reputation: 459

How to make numpy.cumsum start after the first value

I have:

import numpy as np

position = np.array([4, 4.34, 4.69, 5.02, 5.3, 5.7, ..., 4])
x = (B/position**2)*dt

A = np.cumsum(x)
assert A[0] == 0  # I want this to be true.

Where B and dt are scalar constants. This is for a numerical integration problem with initial condition of A[0] = 0. Is there a way to set A[0] = 0 and then do a cumsum for everything else?

Upvotes: 7

Views: 8456

Answers (5)

Nellì
Nellì

Reputation: 11

Another way is to simply add a 0 to the list on which you are performing the cumsum:

import numpy as np

position = np.array([4, 4.34, 4.69, 5.02, 5.3, 5.7, ..., 4])
x = (B/position**2)*dt

# A = np.cumsum([0] + x) # This works only if x is a list
A = np.concatenate([[0], np.cumsum(x)]) # Both with lists and np arrays

This works both if x is a list or a np.ndarray.

Upvotes: 1

Piotr Dabkowski
Piotr Dabkowski

Reputation: 5939

You can use roll (shift right by 1) and then set the first entry to zero.

Upvotes: 0

Mateen Ulhaq
Mateen Ulhaq

Reputation: 27231

1D cumsum

A wrapper around np.cumsum that sets first element to 0:

def cumsum(pmf):
    cdf = np.empty(len(pmf) + 1, dtype=pmf.dtype)
    cdf[0] = 0
    np.cumsum(pmf, out=cdf[1:])
    return cdf

Example usage:

>>> np.arange(1, 11)
array([ 1,  2,  3,  4,  5,  6,  7,  8,  9, 10])

>>> cumsum(np.arange(1, 11))
array([ 0,  1,  3,  6, 10, 15, 21, 28, 36, 45, 55])

N-D cumsum

A wrapper around np.cumsum that sets first element to 0, and works with N-D arrays:

def cumsum(pmf, axis=None, dtype=None):
    if axis is None:
        pmf = pmf.reshape(-1)
        axis = 0

    if dtype is None:
        dtype = pmf.dtype

    idx = [slice(None)] * pmf.ndim

    # Create array with extra element along cumsummed axis.
    shape = list(pmf.shape)
    shape[axis] += 1
    cdf = np.empty(shape, dtype)

    # Set first element to 0.
    idx[axis] = 0
    cdf[tuple(idx)] = 0

    # Perform cumsum on remaining elements.
    idx[axis] = slice(1, None)
    np.cumsum(pmf, axis=axis, dtype=dtype, out=cdf[tuple(idx)])

    return cdf

Example usage:

>>> np.arange(1, 11).reshape(2, 5)
array([[ 1,  2,  3,  4,  5],
       [ 6,  7,  8,  9, 10]])

>>> cumsum(np.arange(1, 11).reshape(2, 5), axis=-1)
array([[ 0,  1,  3,  6, 10, 15],
       [ 0,  6, 13, 21, 30, 40]])

Upvotes: 4

jkalden
jkalden

Reputation: 1588

I don't understand what exactly your problem is, but here are some things you can do to have A[0] = 0.

You can create A to be longer by one index to have the zero as the first entry:

# initialize example data
import numpy as np
B = 1
dt = 1
position =  np.array([4, 4.34, 4.69, 5.02, 5.3, 5.7])

# do calculation
A = np.zeros(len(position) + 1)
A[1:] = np.cumsum((B/position**2)*dt)

Result:

A = [ 0.          0.0625      0.11559096  0.16105356  0.20073547  0.23633533 0.26711403]
len(A) == len(position) + 1

Alternatively, you can manipulate the calculation to substract the first entry of the result:

# initialize example data
import numpy as np
B = 1
dt = 1
position =  np.array([4, 4.34, 4.69, 5.02, 5.3, 5.7])

# do calculation
A = np.cumsum((B/position**2)*dt)
A = A - A[0]

Result:

[ 0.          0.05309096  0.09855356  0.13823547  0.17383533  0.20461403]
len(A) == len(position)

As you see, the results have different lengths. Is one of them what you expect?

Upvotes: 10

Farzad Abdolhosseini
Farzad Abdolhosseini

Reputation: 360

I totally understand your pain, I wonder why Numpy doesn't allow this with np.cumsum. Anyway, though I'm really late and there's already another good answer, I prefer this one a bit more:

np.cumsum(np.pad(array, (1, 0), "constant"))

where array in your case is (B/position**2)*dt. You can change the order of np.pad and np.cumsum as well. I'm just adding a zero to the start of the array and calling np.cumsum.

Upvotes: 5

Related Questions