CODEWITHSUNDEEP
pythonarraysnumpy
dubbbdan
dubbbdan

Reputation: 2740

Numpy: apply along axis with input arrays with different indicies

I have a problem where I have two spatial datasets, 1) temperature (temp), and 2) depth(sd) that need to be combined in a calculation. The metric I am calculating cfgi metric for the first day (i.e. t=0) is to be initialized as all zeros. Each subsequent day (i.e. t=1) is calculated as:

cfgi = cfgi_t=0 - temp_t=1 * np.exp(-0.2*sd_t=1)

where t=0 or t=1 should be interpreted as a subscript

Lets say I have these data, where axis=0 represents daily surfaces:

temp = np.array([-1., -1., -1., -1., -1., -1., -1., -1., -1., -1., -1., -1., -1.,
       -1., -1., -1., -1., -1., -1., -1., -1., -1., -1., -1., -1., -2.,
       -2., -2., -2., -2., -2., -2., -2., -2., -2., -2., -2., -2., -2.,
       -2., -2., -2., -2., -2., -2., -2., -2., -2., -2., -2., -3., -3.,
       -3., -3., -3., -3., -3., -3., -3., -3., -3., -3., -3., -3., -3.,
       -3., -3., -3., -3., -3., -3., -3., -3., -3., -3., -4., -4., -4.,
       -4., -4., -4., -4., -4., -4., -4., -4., -4., -4., -4., -4., -4.,
       -4., -4., -4., -4., -4., -4., -4., -4., -4., -5., -5., -5., -5.,
       -5., -5., -5., -5., -5., -5., -5., -5., -5., -5., -5., -5., -5.,
       -5., -5., -5., -5., -5., -5., -5., -5.]).reshape((5,5,5))

sd = np.array([1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1.,
       1., 1., 1., 1., 1., 1., 1., 1., 2., 2., 2., 2., 2., 2., 2., 2., 2.,
       2., 2., 2., 2., 2., 2., 2., 2., 2., 2., 2., 2., 2., 2., 2., 2., 3.,
       3., 3., 3., 3., 3., 3., 3., 3., 3., 3., 3., 3., 3., 3., 3., 3., 3.,
       3., 3., 3., 3., 3., 3., 3., 4., 4., 4., 4., 4., 4., 4., 4., 4., 4.,
       4., 4., 4., 4., 4., 4., 4., 4., 4., 4., 4., 4., 4., 4., 4., 5., 5.,
       5., 5., 5., 5., 5., 5., 5., 5., 5., 5., 5., 5., 5., 5., 5., 5., 5.,
       5., 5., 5., 5., 5., 5.]).reshape(5,5,5)

cfgi = np.zeros(sd.shape)

My first idea was to use np.apply_along_axis with this function:

def calculate_CFGI(a):
    return a[0] - a[1]* np.exp(-0.2*a[2])

This works, but I am hoping someone could help me achieve my desired result without a for-loop.

for i in range(sd.shape[0]-1):
    c = cfgi[i,:,:]
    t = temp[i+1,:,]
    s = sd[i+1,:,]
    combined = np.stack([c,t,s])
    cfgi[i+1,:,:] = np.apply_along_axis(calculate_CFGI, 0, combined)

Desired Result:

array([[[0.        , 0.        , 0.        , 0.        , 0.        ],
        [0.        , 0.        , 0.        , 0.        , 0.        ],
        [0.        , 0.        , 0.        , 0.        , 0.        ],
        [0.        , 0.        , 0.        , 0.        , 0.        ],
        [0.        , 0.        , 0.        , 0.        , 0.        ]],

       [[1.34064009, 1.34064009, 1.34064009, 1.34064009, 1.34064009],
        [1.34064009, 1.34064009, 1.34064009, 1.34064009, 1.34064009],
        [1.34064009, 1.34064009, 1.34064009, 1.34064009, 1.34064009],
        [1.34064009, 1.34064009, 1.34064009, 1.34064009, 1.34064009],
        [1.34064009, 1.34064009, 1.34064009, 1.34064009, 1.34064009]],

       [[2.987075  , 2.987075  , 2.987075  , 2.987075  , 2.987075  ],
        [2.987075  , 2.987075  , 2.987075  , 2.987075  , 2.987075  ],
        [2.987075  , 2.987075  , 2.987075  , 2.987075  , 2.987075  ],
        [2.987075  , 2.987075  , 2.987075  , 2.987075  , 2.987075  ],
        [2.987075  , 2.987075  , 2.987075  , 2.987075  , 2.987075  ]],

       [[4.78439086, 4.78439086, 4.78439086, 4.78439086, 4.78439086],
        [4.78439086, 4.78439086, 4.78439086, 4.78439086, 4.78439086],
        [4.78439086, 4.78439086, 4.78439086, 4.78439086, 4.78439086],
        [4.78439086, 4.78439086, 4.78439086, 4.78439086, 4.78439086],
        [4.78439086, 4.78439086, 4.78439086, 4.78439086, 4.78439086]],

       [[6.62378806, 6.62378806, 6.62378806, 6.62378806, 6.62378806],
        [6.62378806, 6.62378806, 6.62378806, 6.62378806, 6.62378806],
        [6.62378806, 6.62378806, 6.62378806, 6.62378806, 6.62378806],
        [6.62378806, 6.62378806, 6.62378806, 6.62378806, 6.62378806],
        [6.62378806, 6.62378806, 6.62378806, 6.62378806, 6.62378806]]])

Upvotes: 1

Views: 34

Answers (1)

Mad Physicist
Mad Physicist

Reputation: 114320

Let's look at your relation carefully:

 cfgi[0, ...] = 0
 cfgi[1, ...] = 0 - temp[1, ...] * np.exp(-0.2 * sd[1, ...])
 cfgi[2, ...] = 0 - temp[1, ...] * np.exp(-0.2 * sd[1, ...]) - temp[2, ...] * np.exp(-0.2 * sd[2, ...])
 ...

You can say x = temp * np.exp(0.2 * sd), and write:

cfgi[n, ...] = 0 - x[1, ...] - x[2, ...] - ... - x[n, ...]

Hopefully you can see the pattern here:

cfgi = np.zeros_like(temp)
cfgi[1:, ...] = -np.cumsum(temp[1:, ...] * np.exp(-0.2 * sd[1:, ...]), axis=0)

Upvotes: 1

Related Questions