How to calculate raster layer based on another raster layer using gdal in python?

Question

I have a raster A .tif layer and a raster B .tif layer. I want to identify all pixels in raster B that are < 2 and set those corresponding pixels in raster A to 0, given that these rasters pixel grids are aligned. I understand that QGIS has conditional functions I can use for this within raster calculator, but since I need to loop through many raster files, I will need to rely and python and gdal. Trying to work this out, I would think code would look something like:

Calc("A where B < 2, 0, A", A=raster_A.tif, B=raster_B.tif, outfile=new_raster)

Though I am quite confused about using the proper logic and syntax here. How can I accomplish this conditional raster calculation using python and gdal?

Answer 1

I know that you are specifically requesting GDAL but I think rasterio could be easier to manipulate in your case and more pythonic. So if your environment authorize you to run pip install rasterio then here is a suggestion using pure rasterio + numpy trick:

import rasterio as rio
import numpy as np 

with rio.open("raster_A.tif") as ra, rio.open("raser_B.tif") as rb:
    
    # I assume the final raster will have the same characteristic as raster_A
    profile = ra.profile

    # load the data, I assume that you want to compare the 1st band of each file 
    raw_a = ra.read(1)
    raw_b = rb.read(1)

    # create the new dataset 
    data = raw_a * (raw_b >= 2)

    # make sure that data is the same type as the Tiff file
    data = data.astype(profile.dtype)

    with rio.open("new_raster.tif", "w" **profile) as dst:
        dst.write(data, 1)