Integrating Out Dimension from MultiDimensional Array using Parallel Processing

You can use Numba so to speed up this code by a large margin. Numba is a JIT compiler that is able to compile Numpy-based code to fast native codes (so loops are not a issue with it, in fact this is a good idea to use loops in Numba).

The first thing to do is to pre-compute np.sin(PA) once so to avoid repeated computations. Then, dir_flux * np.sin(PA) can be computed using a for loop and the result can be stored in a pre-allocated array so not to perform millions of expensive small array allocations. The outer loop can be executed using multiple threads using prange and the Numba flag parallel=True. It can be further accelerated using the flag fastmath=True assuming the input values are not special (like NaN or Inf or very very small: see subnormal numbers).

While this should theoretically enough got get a fast code, the current implementation of np.trapz is not efficient as it performs expensive allocations. One can easily rewrite the function so not to allocated any additional arrays.

Here are the resulting code:

import numpy as np
import numba as nb

@nb.njit('(float64[::1], float64[::1])')
def trapz(y, x):
    s = 0.0
    for i in range(x.size-1):
        dx = x[i+1] - x[i]
        dy = y[i] + y[i+1]
        s += dx * dy
    return s * 0.5

@nb.njit('(float64[:,:,:,:], float64[:])', parallel=True)
def compute(flux, PA):
    sl, si, sj, sk = flux.shape
    assert sl == PA.size
    data = np.empty((si, sj, sk))
    flattenPA = np.ascontiguousarray(PA)
    sinPA = np.sin(flattenPA)
    for i in nb.prange(si):
        tmp = np.empty(sl)
        for j in range(sj):
            for k in range(sk):
                dir_flux = flux[:, i, j, k]
                for l in range(sl):
                    tmp[l] = dir_flux[l] * sinPA[l]
                omni_flux = trapz(tmp, flattenPA)
                data[i, j, k] = omni_flux
    return data

for idx,filename in enumerate(filenames):
    # Read NetCDF Data File as 'raw_data'
    flux=raw_data['FluxHydrogen'][:]   #This is size (72,35,25,20)
    PA=raw_data['PitchAngleGrid'][:]   #This is size (72)
    data[idx] = compute(flux, PA)

Note flux and PA must be Numpy arrays. Also note that trapz is accurate as long as len(PA) is relatively small and np.std(PA) is not huge. Otherwise a pair-wise summation or even a (paranoid) Kahan summation should help (note Numpy use a pair-wise summation). In practice, results are the same on random normal numbers.

Further optimizations

The code can be made even faster by making flux accesses more contiguous. An efficient transposition can be used to do that (the one of Numpy is not efficient). However, this is not simple to do on 4D arrays. Another solution is to compute the trapz operation on whole lines of the k dimension. This makes the computation very efficient and nearly memory-bound on my machine. Here is the code:

@nb.njit('(float64[:,:,:,:], float64[:])', fastmath=True, parallel=True)
def compute(flux, PA):
    sl, si, sj, sk = flux.shape
    assert sl == PA.size
    data = np.empty((si, sj, sk))
    sinPA = np.sin(PA)
    premultPA = PA * 0.5
    for i in nb.prange(si):
        for j in range(sj):
            dir_flux = flux[:, i, j, :]
            data[i, j, :].fill(0.0)
            for l in range(sl-1):
                dx = premultPA[l+1] - premultPA[l]
                fact1 = dx * sinPA[l]
                fact2 = dx * sinPA[l+1]
                for k in range(sk):
                    data[i, j, k] += fact1 * dir_flux[l, k] + fact2 * dir_flux[l+1, k]
    return data

Note the premultiplication make the computation slightly less precise.

Results

Here are results on random numbers (like @DominikStańczak used) on my 6-core machine (i5-9600KF processor):

Initial sequential solution:                       193.14 ms  (±  1.8 ms)
DominikStańczak sequential vectorized solution:      8.68 ms  (± 48.1 µs)
Numba parallel solution without fastmath:            0.48 ms  (±  6.7 µs)
Numba parallel solution without fastmath:            0.38 ms  (±  9.5 µs)
Best Numba solution (with fastmath):                 0.32 ms  (±  5.2 µs)
Optimal lower-bound execution:                       0.24 ms  (RAM bandwidth saturation)

Thus, the fastest Numba version is 27 times faster than the (sequential) version of @DominikStańczak and 604 times faster than the initial one. It is nearly optimal.

As a first step, let's vectorize the code itself. I'm just going to deal with doing this on a per-file basis for now, to show you how to get rid of the nested for loop:

shape = (72, 35, 25, 20)

flux = np.random.normal(size=shape)
PA = np.random.normal(size=shape[0])

Now, timing your implementation, rewritten a little:

%%timeit
data = np.empty(shape[1:])
for i in range(shape[1]):
    for j in range(shape[2]):
        for k in range(shape[3]):
            dir_flux=flux[:,i,j,k]
            omni_flux=np.trapz(dir_flux*np.sin(PA),PA)
            data[i,j,k]=omni_flux
# 211 ms ± 4.86 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

My first idea was to pull sin out of the for loop, because there's no need to recalculate it each time, but that got me 10ms tops. However, if instead of the for loop, we used plain numpy vectorization via broadcasting, turning sin_PA into a (72, 1, 1, 1)-shaped array:

%%timeit
sin_PA = np.sin(PA).reshape(-1, 1, 1, 1)
data = np.trapz(flux * sin_PA, x=PA, axis=0)
# 9.03 ms ± 554 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)

That's a 20 times speed up, nothing to scoff at. I estimate it'd take about three minutes for all your files. You can also use np.allclose to verify the results agree up to floating point error.

If you still need to parallelize this afterwards, I would use dask.array In fact, if you've got your data in netcdf4 files, I would use xarray (which is helpful for multidimensional data anyway) to read those, and then run the trapz computations on that with Dask enabled in the backend. I think that's the simplest way to achieve easy multiprocessing in this case. Here's a quick sketch:

import xarray
from Dask.distributed import Client
client = Client()

file_data = xarray.open_mfdataset(filenames, parallel=True)

# massage the data a little, probably

flux = file_data["FluxHydrogen"]
PA = file_data["PitchAngleGrid"]
integrand = flux * np.sin(PA)   # most element-wise numpy operations work on xarray ones or Dask based ones without a hitch
data = integrand.integrate(coord="PitchAngle")   # or some such name for the dimension you're integrating out