breaks.h

Go to the documentation of this file.

/******************************************************************************
 *
 *
 * Project: sgs
 * Purpose: C++ implementation of raster stratification using breaks
 * Author: Joseph Meyer
 * Date: September, 2025
 *
 ******************************************************************************/

 
#include <iostream>
 
#include <boost/asio/thread_pool.hpp>
#include <boost/asio/post.hpp>
 
#include "utils/raster.h"
#include "utils/helper.h"
 
namespace sgs {
namespace breaks {

inline void processMapPixel(
    size_t index,
    helper::RasterBandMetaData& dataBand,
    void *p_dataBuffer,
    helper::RasterBandMetaData& stratBand,
    void *p_stratBuffer,
    std::vector<double>& bandBreaks,
    size_t multiplier,
    bool& mapNan,
    size_t& mapStrat)
{
    double val = helper::getPixelValueDependingOnType<double>(dataBand.type, p_dataBuffer, index);
    bool isNan = std::isnan(val) || (double)val == dataBand.nan;
    mapNan |= isNan;
        
    size_t strat = 0;
    if (!isNan) {
        auto it = std::lower_bound(bandBreaks.begin(), bandBreaks.end(), val);
        strat = (it == bandBreaks.end()) ? bandBreaks.size() : std::distance(bandBreaks.begin(), it);
    }
 
    helper::setStrataPixelDependingOnType(stratBand.type, p_stratBuffer, index, isNan, strat);
 
    if (!mapNan) {
        mapStrat += strat * multiplier;
    }
}   

inline void
processPixel(
    size_t index,
    void *p_data,
    helper::RasterBandMetaData *p_dataBand,
    void *p_strat,
    helper::RasterBandMetaData *p_stratBand,
    std::vector<double>& bandBreaks)
{
    double val = helper::getPixelValueDependingOnType<double>(p_dataBand->type, p_data, index);
    bool isNan = std::isnan(val) || val == p_dataBand->nan;
 
    size_t strat = 0;
    if (!isNan) {
        auto it = std::lower_bound(bandBreaks.begin(), bandBreaks.end(), val);
        strat = (it == bandBreaks.end()) ? bandBreaks.size() : std::distance(bandBreaks.begin(), it);
    }
 
    helper::setStrataPixelDependingOnType(p_stratBand->type, p_strat, index, isNan, strat);
}

raster::GDALRasterWrapper *breaks(
    raster::GDALRasterWrapper *p_raster,
    std::map<int, std::vector<double>> breaks,
    bool map,
    std::string filename,
    bool largeRaster,
    int threads,
    std::string tempFolder,
    std::map<std::string, std::string> driverOptions)
{
    GDALAllRegister();
 
    int height = p_raster->getHeight();
    int width = p_raster->getWidth();
    double *geotransform = p_raster->getGeotransform();
    std::string projection = std::string(p_raster->getDataset()->GetProjectionRef());
 
    GDALDataset *p_dataset = nullptr;
 
    size_t bandCount = breaks.size();
    std::vector<std::vector<double>> bandBreaks;
    std::vector<std::string> bandNames = p_raster->getBands();
 
    std::vector<helper::RasterBandMetaData> dataBands(bandCount);
    std::vector<helper::RasterBandMetaData> stratBands(bandCount + map);
    std::vector<helper::VRTBandDatasetInfo> VRTBandInfo;
 
    bool isMEMDataset = !largeRaster && filename == "";
    bool isVRTDataset = largeRaster && filename == "";
 
    std::mutex dataBandMutex;
    std::mutex stratBandMutex;
    //VRT bands are each their own dataset so they can each have their own mutex
    std::vector<std::mutex> stratBandMutexes(isVRTDataset * (bandCount + map));
 
    std::string driver;
    if (isMEMDataset || isVRTDataset) {
        driver = isMEMDataset ? "MEM" : "VRT";
        p_dataset = helper::createVirtualDataset(driver, width, height, geotransform, projection);
    }
    else {
        std::filesystem::path filepath = filename;
        std::string extension = filepath.extension().string();
 
        if (extension == ".tif") {
            driver = "Gtiff";
        }
        else {
            throw std::runtime_error("sgs only supports .tif files right now");
        }
    }
 
    //allocate, read, and initialize raster data and breaks information
    GDALDataType stratPixelType = GDT_Int8;
    size_t stratPixelSize = 1;
    size_t band = 0;
    for (auto const& [key, val] : breaks) {
        helper::RasterBandMetaData *p_dataBand = &dataBands[band];
        helper::RasterBandMetaData *p_stratBand = &stratBands[band];
 
        //get and store metadata from input raster band
        GDALRasterBand *p_band = p_raster->getRasterBand(key);
        p_dataBand->p_band = p_band;
        p_dataBand->type = p_raster->getRasterBandType(key);
        p_dataBand->size = p_raster->getRasterBandTypeSize(key);
        p_dataBand->p_buffer = largeRaster ? nullptr : p_raster->getRasterBandBuffer(key);
        p_dataBand->nan = p_band->GetNoDataValue();
        p_dataBand->p_mutex = &dataBandMutex;
        p_band->GetBlockSize(&p_dataBand->xBlockSize, &p_dataBand->yBlockSize);
 
        //sort and add band breaks vector
        std::vector<double> valCopy = val; //have to create copy to alter the band breaks in iteration loop
        std::sort(valCopy.begin(), valCopy.end());
        bandBreaks.push_back(valCopy);
        
        //update metadata of new strat raster
        size_t maxStrata = val.size() + 1;
        helper::setStratBandTypeAndSize(maxStrata, &p_stratBand->type, &p_stratBand->size);
        p_stratBand->name = "strat_" + bandNames[key];
        p_stratBand->xBlockSize = map ? dataBands[0].xBlockSize : p_dataBand->xBlockSize;
        p_stratBand->yBlockSize = map ? dataBands[0].yBlockSize : p_dataBand->yBlockSize;
        p_stratBand->p_mutex = &stratBandMutex; //overwritten if VRT dataset
 
        //update dataset with new band information
        if (isMEMDataset) {
            helper::addBandToMEMDataset(p_dataset, *p_stratBand);
        }
        else if (isVRTDataset) {
            helper::createVRTBandDataset(p_dataset, *p_stratBand, tempFolder, std::to_string(key), VRTBandInfo, driverOptions); 
            p_stratBand->p_mutex = &stratBandMutexes[band];
        }
        else { //non-virtual dataset driver
            if (stratPixelSize < p_stratBand->size) {
                stratPixelSize = p_stratBand->size;
                stratPixelType = p_stratBand->type;
            }
        }
 
        band++;
    }
 
    //set multipliers if mapped stratification  
    std::vector<size_t> multipliers(breaks.size(), 1);
    if (map) {
        //determine the stratification band index multipliers of the mapped band
        for (size_t i = 0; i < bandCount - 1; i++) {
            multipliers[i + 1] = multipliers[i] * (bandBreaks[i].size() + 1);
        }
 
        //update info of new strat raster map band
        helper::RasterBandMetaData *p_stratBand = &stratBands.back();
        size_t maxStrata = multipliers.back() * (bandBreaks.back().size() + 1);
        helper::setStratBandTypeAndSize(maxStrata, &p_stratBand->type, &p_stratBand->size);
        p_stratBand->name = "strat_map";
        p_stratBand->xBlockSize = dataBands[0].xBlockSize;
        p_stratBand->yBlockSize = dataBands[0].yBlockSize;
        p_stratBand->p_mutex = &stratBandMutex; //overwritten if VRT dataset
 
        //update dataset with new band information
        if (isMEMDataset) {
            helper::addBandToMEMDataset(p_dataset, *p_stratBand);            
        }
        else if (isVRTDataset) {
            helper::createVRTBandDataset(
                p_dataset, 
                *p_stratBand, 
                tempFolder, 
                "map", 
                VRTBandInfo, 
                driverOptions
            ); 
            p_stratBand->p_mutex = &stratBandMutexes.back();
        }
        else { //non-virtual dataset driver
            if (stratPixelSize < p_stratBand->size) {
                stratPixelSize = p_stratBand->size;
                stratPixelType = p_stratBand->type;
            }
        }
    }   
 
    //create full non-virtual dataset now that we have all required band information
    if (!isMEMDataset && !isVRTDataset) {
        bool useTiles = stratBands[0].xBlockSize != width &&
                stratBands[0].yBlockSize != height;
 
        for (size_t band = 0; band < stratBands.size(); band++) {
            stratBands[band].size = stratPixelSize;
            stratBands[band].type = stratPixelType;
            stratBands[band].p_buffer = !largeRaster ? VSIMalloc3(height, width, stratPixelSize) : nullptr;
        }
 
        p_dataset = helper::createDataset(
            filename,
            driver,
            width,
            height,
            geotransform,
            projection,
            stratBands.data(),
            stratBands.size(),
            useTiles,
            driverOptions
        );
    }
 
    //iterate through all pixels and update the stratified raster bands
    if (largeRaster) {
        pybind11::gil_scoped_acquire acquire;
        boost::asio::thread_pool pool(threads);
 
        if (map) {
            //use the first raster band to determine block size
            int xBlockSize = dataBands[0].xBlockSize;
            int yBlockSize = dataBands[0].yBlockSize;
 
            int xBlocks = (p_raster->getWidth() + xBlockSize - 1) / xBlockSize;
            int yBlocks = (p_raster->getHeight() + yBlockSize - 1) / yBlockSize;
            int chunkSize = yBlocks / threads;
 
            for (int yBlockStart = 0; yBlockStart < yBlocks; yBlockStart += chunkSize) {
                int yBlockEnd = std::min(yBlockStart + chunkSize, yBlocks);
            
                boost::asio::post(pool, [
                    bandCount,
                    xBlockSize, 
                    yBlockSize, 
                    yBlockStart, 
                    yBlockEnd, 
                    xBlocks, 
                    &dataBands, 
                    &stratBands, 
                    &bandBreaks,
                    &multipliers
                ] {
                    std::vector<void *> dataBuffers(dataBands.size());
                    std::vector<void *> stratBuffers(stratBands.size());
                    for (size_t band = 0; band < dataBuffers.size(); band++) {
                        dataBuffers[band] = VSIMalloc3(xBlockSize, yBlockSize, dataBands[band].size);
                        stratBuffers[band] = VSIMalloc3(xBlockSize, yBlockSize, stratBands[band].size);
                    }
                    stratBuffers.back() = VSIMalloc3(xBlockSize, yBlockSize, stratBands.back().size);
 
                    for (int yBlock = yBlockStart; yBlock < yBlockEnd; yBlock++) {
                        for (int xBlock = 0; xBlock < xBlocks; xBlock++) {
                            int xValid, yValid;
                            dataBands[0].p_mutex->lock();
                            dataBands[0].p_band->GetActualBlockSize(xBlock, yBlock, &xValid, &yValid);
                            dataBands[0].p_mutex->unlock();
 
                            //read raster band data into band buffers
                            for (size_t band = 0; band < bandCount; band++) {
                                helper::rasterBandIO(
                                    dataBands[band], 
                                    dataBuffers[band], 
                                    xBlockSize, 
                                    yBlockSize, 
                                    xBlock, 
                                    yBlock, 
                                    xValid, 
                                    yValid, 
                                    true //read = true
                                );      
                            }
 
                            //process blocked band data
                            for (int y = 0; y < yValid; y++) {
                                size_t index = static_cast<size_t>(y * xBlockSize);
                                for (int x = 0; x < xValid; x++) {
                                    bool mapNan = false;
                                    size_t mapStrat = 0;
                                    
                                    for (size_t band = 0; band < bandCount; band++) {
                                        processMapPixel(
                                            index, 
                                            dataBands[band], 
                                            dataBuffers[band], 
                                            stratBands[band], 
                                            stratBuffers[band], 
                                            bandBreaks[band],
                                            multipliers[band],
                                            mapNan,
                                            mapStrat
                                        );
                                    }
                                
                                    helper::setStrataPixelDependingOnType(
                                        stratBands.back().type,
                                        stratBuffers.back(),
                                        index,
                                        mapNan,
                                        mapStrat
                                    );
 
                                    index++;
                                }
                            }
                    
                            //write strat band data
                            for (size_t band = 0; band <= bandCount; band++) {
                                helper::rasterBandIO(
                                    stratBands[band],
                                    stratBuffers[band],
                                    xBlockSize,
                                    yBlockSize,
                                    xBlock,
                                    yBlock,
                                    xValid,
                                    yValid,
                                    false //read = false
                                );
                            }
                        }
                    }
 
                    for (size_t band = 0; band < dataBuffers.size(); band++) {
                        VSIFree(dataBuffers[band]);
                        VSIFree(stratBuffers[band]);
                    }
                    VSIFree(stratBuffers.back());
                });
            }   
        }
        else {  
            for (size_t band = 0; band < bandCount; band++) {
                helper::RasterBandMetaData* p_dataBand = &dataBands[band];
                helper::RasterBandMetaData* p_stratBand = &stratBands[band];
 
                int xBlockSize = p_dataBand->xBlockSize;
                int yBlockSize = p_dataBand->yBlockSize;
                    
                int xBlocks = (p_raster->getWidth() + xBlockSize - 1) / xBlockSize;
                int yBlocks = (p_raster->getHeight() + yBlockSize - 1) / yBlockSize;           
                int chunkSize = yBlocks / threads;
                
                for (int yBlockStart = 0; yBlockStart < yBlocks; yBlockStart += chunkSize) {
                    int yBlockEnd = std::min(yBlocks, yBlockStart + chunkSize);
                    std::vector<double> *p_breaks = &bandBreaks[band];
                    
                    boost::asio::post(pool, [
                        xBlockSize, 
                        yBlockSize, 
                        yBlockStart, 
                        yBlockEnd, 
                        xBlocks, 
                        p_dataBand, 
                        p_stratBand, 
                        p_breaks
                    ] {
                        void *p_data = VSIMalloc3(xBlockSize, yBlockSize, p_dataBand->size);
                        void *p_strat = VSIMalloc3(xBlockSize, yBlockSize, p_stratBand->size);
 
                        int xValid, yValid;
                        for (int yBlock = yBlockStart; yBlock < yBlockEnd; yBlock++) {
                            for (int xBlock = 0; xBlock < xBlocks; xBlock++) {
                                p_dataBand->p_mutex->lock();
                                p_dataBand->p_band->GetActualBlockSize(xBlock, yBlock, &xValid, &yValid);
                                p_dataBand->p_mutex->unlock();
                                
                                //read block into memory
                                helper::rasterBandIO(
                                    *p_dataBand,
                                    p_data,
                                    xBlockSize,
                                    yBlockSize,
                                    xBlock,
                                    yBlock,
                                    xValid,
                                    yValid,
                                    true //read = true
                                );
 
                                //process block
                                for (int y = 0; y < yValid; y++) {
                                    size_t index = static_cast<size_t>(y * xBlockSize);
                                    for (int x = 0; x < xValid; x++) {
                                        processPixel(index, p_data, p_dataBand, p_strat, p_stratBand, *p_breaks);
                                        index++;
                                    }
                                }
                                
                                //write resulting stratifications to disk
                                helper::rasterBandIO(
                                    *p_stratBand,
                                    p_strat,
                                    xBlockSize,
                                    yBlockSize,
                                    xBlock,
                                    yBlock,
                                    xValid,
                                    yValid,
                                    false //read = false
                                );
                            }
                        }
                        VSIFree(p_data);
                        VSIFree(p_strat);
                    });
                }
            }
        }
        pool.join();
        pybind11::gil_scoped_release release;
    }
    else {
        size_t pixelCount = static_cast<size_t>(p_raster->getWidth()) * static_cast<size_t>(p_raster->getHeight());
        if (map) {
            for (size_t index = 0; index < pixelCount; index++) {
                bool mapNan = false;
                size_t mapStrat = 0;
                                    
                for (size_t band = 0; band < bandCount; band++) {
                    processMapPixel(
                        index, 
                        dataBands[band], 
                        dataBands[band].p_buffer, 
                        stratBands[band], 
                        stratBands[band].p_buffer, 
                        bandBreaks[band],
                        multipliers[band],
                        mapNan,
                        mapStrat
                    );
                }
 
                helper::setStrataPixelDependingOnType(
                    stratBands.back().type,
                    stratBands.back().p_buffer,
                    index,
                    mapNan,
                    mapStrat
                );
            }
        }
        else {
            for (size_t band = 0; band < bandCount; band++) {
                for (size_t i = 0; i < pixelCount; i++) {
                    processPixel(
                        i,
                            dataBands[band].p_buffer,   
                        &dataBands[band], 
                        stratBands[band].p_buffer,
                        &stratBands[band],
                        bandBreaks[band]
                    );
                }
            }
        }
 
        //if non-virtual dataset then write in-memory output bands to disk
        if (!isVRTDataset && !isMEMDataset) {
            CPLErr err;
            for (size_t band = 0; band < stratBands.size(); band++) {
                err = stratBands[band].p_band->RasterIO(
                    GF_Write,
                    0,
                    0,
                    width,
                    height,
                    stratBands[band].p_buffer,
                    width,
                    height, 
                    stratBands[band].type,
                    0,
                    0
                );
                if (err) {
                    throw std::runtime_error("error writing band to file.");
                }
            }
        }
    }
 
    //close and add all of the VRT sub datasets as bands
    if (isVRTDataset) {
        for (size_t band = 0; band < VRTBandInfo.size(); band++) {
            GDALClose(VRTBandInfo[band].p_dataset);
            helper::addBandToVRTDataset(p_dataset, stratBands[band], VRTBandInfo[band]);
        }
    }
 
    //if the bands are in memory, populate a vector of just bands to use in GDALRasterWrapper creation
    std::vector<void *> buffers(stratBands.size());
    if (!largeRaster) {
        for (size_t band = 0; band < stratBands.size(); band++) {
            buffers[band] = stratBands[band].p_buffer;
        }
    }
 
    return largeRaster ?
            new raster::GDALRasterWrapper(p_dataset) :
        new raster::GDALRasterWrapper(p_dataset, buffers);
 
}
 
} //namespace breaks
} //namespace sgs