Comparing values in two arrays from rasterio and performing operations

Question

I have a GeoTIFF that has two bands in it, one a "precip reflectivity" (intensity) and "precip type" (snow/rain/etc). I'm wanting to adjust precip values that are snow so I can color them differently on my final map. Here's what I'm currently doing and what I'm trying to do. The tif that I am using is available here.

I'm loading in the two bands using rasterio like so:

import rasterio


src = rasterio.open("stack.tif")
ref=src.read(1) # This is the precip reflectivity
ptype=src.read(2) # This is the precip type

Currently, ref and ptype are two arrays that look like the following (-999 is a nodata value):

[[-999. -999. -999. ... -999. -999. -999.]
 [-999. -999. -999. ... -999. -999. -999.]
 [-999. -999. -999. ... -999. -999. -999.]
 ...
 [-999. -999. -999. ... -999. -999. -999.]
 [-999. -999. -999. ... -999. -999. -999.]
 [-999. -999. -999. ... -999. -999. -999.]]

I'm want to eventually correctly color with a colormap. For values that are "snow" (or type 3 for ptype I want to add 200 to the value to denote for snow and for all other values I want to keep the values as is.

That leads me to this, is there a best way to compare these values or any examples?

Answer 1

As stated in the docs, reading a dataset with Rasterio returns an numpy.ndarray.

Therefore, you can take advantage of boolean comparison to do something like:

import numpy as np
import rasterio

fill_val = -999.0
snow_val = 3

# open original dataset
with rasterio.open("stack.tif") as src:
    ref = src.read(1) # This is the precip reflectivity
    ptype = src.read(2) # This is the precip type
    kwargs = src.meta.copy()

# modify valid `ref` data based on values in the `ptype` layer
ref[np.logical_and(ptype == snow_val, ref != fill_val)] += 200

# write the result to a new raster dataset with same structure
with rasterio.open("stack_modified.tif", "w", **kwargs) as dst:
    dst.write(ref, 1)
    dst.write(ptype, 2)

Note that it's important to not modify the nodata (fill) values, as this would falsely cause them to be recognized as valid data.