How can I optimize the processing time of the zonal_stats module using multiprocessing in Python?

Question

I am carrying out a spectral feature extraction in Python 3 starting from an 8-band image and a shapefile (consisting of about a million polygons) superimposed on it. I calculate the mean and standard deviation of the pixels for each polygon and join them to the polygon as additional attributes. The most intuitive way was to use the zonal_stats module which automatically saves the result in geojason (maintaining geographic information). The rest of the code allows merging the individual geojasons into one. The process, in particular the calculation of the spectral features, takes 21 hours. Given that I am a beginner with Python and with the world of programming, how could I implement multiprocessing in order to reduce processing times?

P.S .: My goal is to have a shp or geojason file with the spectral features calculated for each polygon and to drastically reduce the processing time. I use a workstation in which an Intel Xeon processor with 12 physical cores and 24 virtual cores.

from rasterstats import zonal_stats

polig = '...\\poligons.shp'
img_Berg = '....\\8_band_image.tif'    

col = ['mean', 'std']

#the most demanding part of code starts from here ....

def zstats(b):
    zs = zonal_stats(polig, img_Berg,  band=b, stats=['mean', 'std'], geojson_out=True)
    return zs

stats_blu = zstats(4)
stats_green = zstats(7)
stats_red = zstats(6)
stats_rededge = zstats(8)
stats_nir1 = zstats(5)
stats_nir2 = zstats(3)
stats_gndvi = zstats(1)
stats_dsm = zstats(2)

#.....to here

def change_name(d):                      
    for i in range(len(d)):
        a = d[i]['properties']
        if d == stats_blu:
            a['mean_blu'] = a.pop('mean')
            a[' std_blu'] = a.pop('std')
        elif d == stats_green:
            a['mean_green'] = a.pop('mean')
            a[' std_green'] = a.pop('std')
        elif d == stats_red:
            a['mean_red'] = a.pop('mean')
            a[' std_red'] = a.pop('std')
        elif d == stats_rededge:
            a['mean_rededg'] = a.pop('mean')
            a[' std_rededg'] = a.pop('std')
        elif d == stats_nir1:
            a['mean_nir1'] = a.pop('mean')
            a[' std_nir1'] = a.pop('std')
        elif d == stats_nir2:
            a['mean_nir2'] = a.pop('mean')
            a[' std_nir2'] = a.pop('std')
        elif d == stats_gndvi:
            a['mean_gndvi'] = a.pop('mean')
            a[' std_gndvi'] = a.pop('std')
        elif d == stats_dsm:
            a['mean_dsm'] = a.pop('mean')
            a[' std_dsm'] = a.pop('std')
    return(d)


a = change_name(stats_blu) 
b = change_name(stats_green) 
c = change_name(stats_red) 
d = change_name(stats_rededge) 
e = change_name(stats_nir1) 
f = change_name(stats_nir2) 
g = change_name(stats_gndvi) 
h = change_name(stats_dsm) 

def merge(a, b, c, d, e, f, g, h):
    l = [a, b, c, d, e, f, g]
    for i in range(len(a)):
        for j in l:
            z = j[i]['properties'].copy()
            h[i]['properties'].update(z)

merge(a, b, c, d, e, f, g, h)

Output_geoj = h.copy()

Answer 1

A quick way to test if a speedup is possible is to only process some of the polygons in one run and save the result as a geojson-file. If you change your script to take a batch number and total number of batches as input, you can run it from the command line. And if you start multiple runs at the same time you test if a speedup is possible. I have written some untested code

import json
from math import ceil
import sys

import fiona
from rasterstats import zonal_stats

# batch_num is from 1 to total_batches
total_batches = int(sys.argv[1])
batch_num = int(sys.argv[2])

if batch_num > total_batches:
    raise ValueError("Batch number is greater than total batches")

polig = 'poligons.shp'
img_Berg = '8_band_image.tif'    

with fiona.open(polig) as shp:
    features = list(shp)

batch_size = ceil(len(features) / total_batches)
batch_features = features[(batch_num-1)*batch_size:(batch_num*batch_size)]

print(batch_size, len(batch_features))

def zstats(band, postfix):
    zs = zonal_stats(batch_features, img_Berg,  band=band, stats=['mean', 'std'], geojson_out=True)
    for f in zs:
        props = f['properties']
        props['mean'+postfix] = props.pop('mean')
        props[' std'+postfix] = props.pop('std')
    return zs

to_process = [
    (4, 'blu'),
    (7, 'green'),
    (6, 'red'),
    (8, 'rededg'),
    (5, 'nir1'),
    (3, 'nir2'),
    (1, 'gndvi'),
    (2, 'dsm')
]

# The processing happens here
stats = [zstats(*args) for args in to_process]

# Then merging properties into a copy of the last geojson object
merged_stats = stats[-1].copy()
for s in stats:
    for i in range(len(merged_stats)):
        props = s[i]['properties'].copy()
        merged_stats[i]['properties'].update(props)

output_file = 'polig-stats-{}-of-{}.geojson'.format(batch_num, total_batches)

with open(output_file, 'w') as f:
    json.dump(merged_stats, f)

To test the code you save it in a file e.g. batch-zonal-stats.py and call it from the terminal like this

$ python batch-zonal-stats.py 10000 1

This will run the first of 10000 batches and save a file polig-stats-1-of-10000.geojson in the same folder as your shapefile. This should run pretty fast since it only process 100 polygons and you can check the output file. If everything is ok you can open maybe 8 terminals and then run this in the first terminal

$ python batch-zonal-stats.py 8 1

and this in the second terminal

$ python batch-zonal-stats.py 8 2

and so on until all 8 terminals runs one batch. Check to see if the cpu usage is closer to 100% now. If everything works write here again, and then we can try to rewrite it to a single script using multiprocessing :)

EDIT: I didn't see your comment about using 20-30% ram. That might be a problem. If the problem is loading the many features from the shapefile, my solution might not work. If that is the case I think you should write another script that splits the shapefile into one file for each batch, and then let each batch open a separate shapefile. The solution above open the same large shapefile 8 times which wastes a lot of ram I guess :)

EDIT 2: Small changes to code