ssr_utilities.py

#!/usr/bin/env python
############################################################################
#
# MODULE:       ssr_utilities.py
# AUTHOR:       Collin Bode, UC Berkeley                March, 2012
#
# PURPOSE:      Utility functions for SSR model (file-based, WhiteboxTools).
#               Replaces GRASS GIS workspace management with file system ops.
#
# COPYRIGHT:    (c) 2012 Collin Bode
#               This program is free software under the GNU General Public
#               License (>=v2). Read the file COPYING that comes with GRASS
#               for details.
#
#############################################################################

import os
import sys
import re
import shutil
import datetime as dt
import numpy as np
import rasterio
from rasterio.transform import Affine
from rasterio.merge import merge as rasterio_merge
import whitebox
from ssr_params import *

wbt = whitebox.WhiteboxTools()
wbt.verbose = False


def printout(str_text, lf):
    timestamp = dt.datetime.strftime(dt.datetime.now(), "%H:%M:%S")
    lf.write(timestamp + ": " + str_text + '\n')
    print(timestamp + ": " + str_text)


def get_path():
    return os.path.dirname(os.path.realpath(__file__)) + os.sep


def ensure_dir(path):
    os.makedirs(path, exist_ok=True)


def setup_workspace():
    """Create all output subdirectories under workspace."""
    for subdir in [workspace, dir_lidar, dir_lpi, dir_sun, dir_ssr, dir_temp, dir_logs]:
        ensure_dir(os.path.join(workspace, subdir) if subdir != workspace else workspace)


def raster_path(name, subdir=''):
    """Return full path to a raster .tif file."""
    if subdir:
        return os.path.join(workspace, subdir, name + '.tif')
    return os.path.join(workspace, name + '.tif')


def raster_exists(name, subdir=''):
    """Check if a raster file exists."""
    return os.path.exists(raster_path(name, subdir))


def read_raster(path):
    """Read raster and return (data array, profile dict)."""
    with rasterio.open(path) as src:
        data = src.read(1).astype(np.float32)
        profile = src.profile.copy()
    return data, profile


def save_raster(data, profile, path):
    """Save raster as GeoTIFF with LZW compression."""
    out_profile = profile.copy()
    out_profile.update(
        driver='GTiff',
        compress='lzw',
        dtype=data.dtype,
        count=1,
    )
    ensure_dir(os.path.dirname(path))
    with rasterio.open(path, 'w', **out_profile) as dst:
        dst.write(data, 1)


def mosaic_rasters(input_paths, output_path, nodata=np.nan):
    """Merge a list of raster files into a single output raster."""
    datasets = [rasterio.open(p) for p in input_paths]
    mosaic, transform = rasterio_merge(datasets)
    profile = datasets[0].profile.copy()
    profile.update(
        driver='GTiff',
        compress='lzw',
        height=mosaic.shape[1],
        width=mosaic.shape[2],
        transform=transform,
        count=1,
    )
    for ds in datasets:
        ds.close()
    ensure_dir(os.path.dirname(output_path))
    with rasterio.open(output_path, 'w', **profile) as dst:
        dst.write(mosaic[0], 1)


def apply_recode(data, rules_file):
    """Apply recode rules from GRASS r.recode format (min:max:value per line)."""
    output = np.full_like(data, np.nan, dtype=np.float32)
    with open(rules_file) as f:
        for line in f:
            line = line.strip()
            if not line or line.lower() == 'end':
                continue
            parts = line.split(':')
            if len(parts) < 3:
                continue
            from_min = float(parts[0])
            from_max = float(parts[1])
            to_val = float(parts[2])
            mask = (data >= from_min) & (data <= from_max)
            output[mask] = to_val
    return output