Search code examples
pythonnumpynumerical-methods

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?


Solution

  • 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?