Export Pits/Depressions/Flats Extent

[cmu002]

Is it possible to export the pits/depressions/flats raster as a geotiff? Or are they locked in memory? I haven't seen any mention of this process in the documentation or tutorials. I've provided some example for reference:

# Define DEM Working Directory  
os.chdir(r'C:\Path\to\workspace')

#Print working directory 
print("Current working directory: {0}".format(os.getcwd()))

#Import DEM
grid = Grid.from_raster('DEM.tif')
dem = grid.read_raster('DEM.tif')

#Set NoData value to NaN
dem[dem==int(dem.nodata)]=np.nan

# Plot the figure
fig, ax = plt.subplots(figsize=(25,10))
fig.patch.set_alpha(0)

plt.imshow(dem, extent=grid.extent, cmap='terrain', zorder=1)
plt.colorbar(label='Elevation (m)')
plt.grid(zorder=0)
plt.title('Digital elevation map', size=14)
plt.xlabel('Longitude')
plt.ylabel('Latitude')
plt.tight_layout()

#Detect the pits in the DEM
pits = grid.detect_pits(dem)

#Detect Depressions 
depressions = grid.detect_depressions(dem)

# Fill Depressions
flooded_dem = grid.fill_depressions(pit_filled_dem)
depressions = grid.detect_depressions(flooded_dem)
assert not depressions.any()

# Detect flats
flats = grid.detect_flats(flooded_dem)

#Export rasters
grid.to_raster(data = depressions, 
               file_name = r'C:\Path\to\workspace\Depressions.tif', 
               apply_output_mask = True)

grid.to_raster(data = flats, 
               file_name = r'C:\Path\to\workspace\Flats.tif', 
               apply_output_mask = True)

grid.to_raster(data = pits, 
               file_name = r'C:\Path\to\workspace\Pits.tif', 
               apply_output_mask = True)

Any help or guidance would be much appreciated!

[zachariahBinx]

You should be able to do something like this:

    # Load inputs
    grid = Grid.from_raster(dem_input)
    dem = grid.read_raster(dem_input)
    rast = rasterio.open(dem_input)

    # Set variables
    prof = rast.profile
    transform = rast.transform

    #Detect the pits in the DEM
    pits = grid.detect_pits(dem)

    test = "C:/Users/TEMP/flats.tif"
    with rasterio.open(test,
                       'w',
                       driver='GTiff',
                       height=prof['height'],
                       width=prof['width'],
                       count=1,
                       dtype=prof['dtype'],
                       crs=rast.crs,
                       transform=transform) as output_raster:

        output_raster.write(pits, 1)