Habitat Suitability

import os
import pathlib
import zipfile
from glob import glob
from math import floor, ceil

import pandas as pd
import numpy as np
import geopandas as gpd
import rioxarray as rxr
from rioxarray.merge import merge_arrays
import xarray as xr
import xrspatial

import matplotlib.pyplot as plt
import hvplot.pandas
import hvplot.xarray

import earthaccess

# Images
images_dir = os.path.join(
    # Home directory
    pathlib.Path.home(),
    'Projects',
    # Project directory
    'habitat-suitability',
    'images'
)

# Datasets
datasets_dir = os.path.join(
    # Home directory
    pathlib.Path.home(),
    'Projects',
    # Project directory
    'habitat-suitability',
    'datasets'
)

# Project data directory 
data_dir = os.path.join(
    # Home directory
    pathlib.Path.home(),
    'Projects',
    # Project directory
    'habitat-suitability',
    'data'
)

# Define directories for data
land_units_dir = os.path.join(data_dir, 'usfs-national-lands')
eldorado_elevation_dir = os.path.join(data_dir, 'srtm', 'eldorado')
los_padres_elevation_dir = os.path.join(data_dir, 'srtm', 'los_padres')

os.makedirs(images_dir, exist_ok=True)
os.makedirs(datasets_dir, exist_ok=True)
os.makedirs(data_dir, exist_ok=True)
os.makedirs(eldorado_elevation_dir, exist_ok=True)
os.makedirs(los_padres_elevation_dir, exist_ok=True)

def build_da(urls, bounds):
    """
    Build a DataArray from a list of urls.
    
    Args:
    urls (list): Input list of URLs.
    bounds (tuple): Site boundaries.

    Returns:
    xarray.DataArray: A merged DataArray.
    """
    
    all_das = []

    # Add buffer to bounds for plotting
    buffer = .025
    xmin, ymin, xmax, ymax = bounds
    bounds_buffer = (xmin-buffer, ymin-buffer, xmax+buffer, ymax+buffer)

    for url in urls:
        # Open data granule, mask missing data, scale data, 
        # and remove dimensions of length 1
        tile_da = rxr.open_rasterio(
                url,
                # For the fill/missing value
                mask_and_scale=True
            ).squeeze()
        # Unpack the bounds and crop tile
        cropped_da = tile_da.rio.clip_box(*bounds_buffer)
        all_das.append(cropped_da)

    merged = merge_arrays(all_das)
    return merged

def convert_longitude(longitude):
    """
    Convert longitude values from a range of 0 to 360 to -180 to 180.
    
    Args:
    longitude (float): Input longitude value.

    Returns:
    float: A value in the specified range.
    """
    
    return (longitude - 360) if longitude > 180 else longitude

def export_raster(da, raster_path):
    """
    Export raster DataArray to a raster file.
    
    Args:
    raster (xarray.DataArray): Input raster layer.
    raster_path (str): Output raster directory.

    Returns: None
    """
    
    output_file = os.path.join(data_dir, os.path.basename(raster_path))
    da.rio.to_raster(output_file)

def harmonize_raster_layers(reference_raster, input_rasters, output_dir):
    """
    Harmonize raster layers to ensure consistent spatial resolution 
    and projection.

    Args:
    reference_raster (xarray.DataArray): Input reference raster.
    input_rasters (list): List of site rasters.
    output_dir (str): Path of raster directory.

    Returns:
    list: A list of harmonized rasters.
    """
    harmonized_files = []

    harmonized_files.append(reference_raster)
    # Load the reference raster
    ref_raster = rxr.open_rasterio(reference_raster, masked=True)

    for raster_path in input_rasters:
        # Load the input raster
        input_raster = rxr.open_rasterio(raster_path, masked=True)

        # Reproject and align the input raster to match the reference raster
        harmonized_raster = input_raster.rio.reproject_match(ref_raster)

        # Save the harmonized raster to the output directory
        output_file = os.path.join(output_dir, os.path.basename(raster_path))
        harmonized_raster.rio.to_raster(output_file)
        harmonized_files.append(output_file)

    print('Harmonized rasters: ', len(harmonized_files))
    return harmonized_files

def plot_site(site_da, site_gdf, site_fig_name, plot_title, 
              bar_label, plot_cmap, boundary_clr, tif_file=False):
    """
    Create custom site plot.
    
    Args:
    site_da (xarray.DataArray): Input site raster.
    site_gdf (geopandas.GeoDataFrame): Input site GeoDataFrame.
    site_fig_name (str): Site figure name.
    plot_title (str): Plot title. 
    bar_label (str): Plot bar variable name.
    plot_cmap (str): Plot colormap name.
    boundary_clr (str): Plot site boundary color.
    tif_file (boolean): Indicates a site file.

    Returns:
    matplotlib.pyplot.plot: A plot of site values.
    """
    
    fig = plt.figure(figsize=(8, 6)) 
    ax = plt.axes()

    if tif_file:
        site_da = rxr.open_rasterio(site_da, masked=True)

    # Plot DataArray values
    site_plot = site_da.plot(
                            cmap=plot_cmap, 
                            cbar_kwargs={'label': bar_label}
                        )

    # Plot site boundary
    site_gdf.boundary.plot(ax=plt.gca(), color=boundary_clr)

    plt.title(f'{plot_title}')
    ax.set_xlabel('Longitude')
    ax.set_ylabel('Latitude')

    fig.savefig(f"{images_dir}/{site_fig_name}.png") 

    return site_plot

# Only extract once
usfs_pattern = os.path.join(land_units_dir, '*.shp')

if not glob(usfs_pattern):
    usfs_zip = f'{datasets_dir}/S_USA.NFSLandUnit.zip'

    # Unzip data
    with zipfile.ZipFile(usfs_zip, 'r') as zip:
        zip.extractall(path=land_units_dir)

# Find the extracted .shp file path
usfs_land_path = glob(usfs_pattern)[0]

# Load USFS land units from shapefile
usfs_land_units_gdf = (
    gpd.read_file(usfs_land_path)
)

# Obtain units with location name
valid_units = usfs_land_units_gdf.dropna(subset=['HQ_LOCATIO'])

# Select only CA units
all_ca_units = valid_units[valid_units['HQ_LOCATIO'].str.contains('CA')]

earthaccess.login(strategy="interactive", persist=True)

# Search for Digital Elevation Models

ea_dem = earthaccess.search_datasets(keyword='SRTM DEM', count=15)
for dataset in ea_dem:
    print(dataset['umm']['ShortName'], dataset['umm']['EntryTitle'])

# Los Padres National Forest
los_padres_gdf = all_ca_units.loc[
    all_ca_units['NFSLANDU_2'] == 'Los Padres National Forest'
]

# plot with EsriImagery basemap  
los_padres_gdf.hvplot(
    geo=True, tiles='EsriImagery',
    title='Los Padres National Forest',
    fill_color=None, line_color='lightblue', line_width=3,
    frame_width=600, frame_height=600
)

# Eldorado National Forest
eldorado_gdf = all_ca_units.loc[
    all_ca_units['NFSLANDU_2'] == 'Eldorado National Forest'
]

# plot with EsriImagery basemap  
eldorado_gdf.hvplot(
    geo=True, tiles='EsriImagery',
    title='Eldorado National Forest',
    fill_color=None, line_color='lightblue', line_width=3,
    frame_width=600, frame_height=600
)

def create_polaris_urls(soil_prop, stat, soil_depth, gdf_bounds):
    """
    Create POLARIS dataset URLs using site bounds.

    Args:
    soil_prop (str): Soil property.
    stat (str): Summary statistic. 
    soil_depth (str): Soil depth (cm).
    gdf_bounds (numpy.ndarray): Array of site boundaries.

    Returns:
    list: A list of POLARIS datset URLs. 
    """

    # Get latitude and longitude bounds from site
    min_lon, min_lat, max_lon, max_lat = gdf_bounds

    site_min_lon = floor(min_lon) 
    site_min_lat = floor(min_lat)  
    site_max_lon = ceil(max_lon)  
    site_max_lat = ceil(max_lat)

    all_soil_urls = []

    for lon in range(site_min_lon, site_max_lon): 
        for lat in range(site_min_lat, site_max_lat):
            current_max_lon = lon + 1
            current_max_lat = lat + 1

            soil_template = (
                "http://hydrology.cee.duke.edu/POLARIS/PROPERTIES/v1.0/"
                "{soil_prop}/"
                "{stat}/"
                "{soil_depth}/"
                "lat{min_lat}{max_lat}_lon{min_lon}{max_lon}.tif"
            )

            soil_url = soil_template.format(
                soil_prop=soil_prop, stat=stat, soil_depth=soil_depth,
                min_lat=lat, max_lat=current_max_lat, 
                min_lon=lon, max_lon=current_max_lon
            )

            all_soil_urls.append(soil_url)

    return all_soil_urls

# Set site parameters
soil_prop = 'ph'
soil_prop_stat = 'mean'
# cm (the minimum depth for blue oaks is 1-2 ft)
soil_depth = '30_60'

# Build soil DataArrays 

# Eldorado site
ed_polaris_urls = create_polaris_urls(
                    soil_prop, soil_prop_stat, soil_depth, 
                    eldorado_gdf.total_bounds
                  )
eldorado_soil_da = build_da(
                    ed_polaris_urls, 
                    tuple(eldorado_gdf.total_bounds)
                  )
export_raster(eldorado_soil_da, "eldorado_soil_ph.tif")

# Los Padres site
lp_polaris_urls = create_polaris_urls(
                    soil_prop, soil_prop_stat, soil_depth, 
                    los_padres_gdf.total_bounds
                  )
los_padres_soil_da = build_da(
                      lp_polaris_urls, 
                      tuple(los_padres_gdf.total_bounds)
                    )
export_raster(los_padres_soil_da, "los_padres_soil_ph.tif")

# Plot Los Padres soils

los_padres_soil_plt = plot_site(
    los_padres_soil_da, los_padres_gdf, 
    'Los-Padres-National-Forest-Soil', 
    'Los Padres National Forest Soil', 
    'pH', 'viridis', 'lightblue'
)

los_padres_soil_plt

# Plot Eldorado soils

eldorado_soil_plt = plot_site(
    eldorado_soil_da, eldorado_gdf, 
    'Eldorado-National-Forest-Soil', 
    'Eldorado National Forest Soil', 
    'pH', 'viridis', 'lightblue'
)

eldorado_soil_plt

def select_dem(bounds, site_gdf, download_dir):
    """
    Create elevation DataArray from NASA Shuttle Radar Topography Mission data.

    Args:
    bounds (tuple): Input site boundaries.
    site_gdf (geopandas.GeoDataFrame): Land unit GeoDataFrame.
    download_dir (str): Path of download directory.

    Returns:
    xarray.DataArray: A site elevation raster.
    """

    # Returns data granules for given bounds
    strm_granules = earthaccess.search_data(
        # SRTMGL1: NASA Shuttle Radar Topography Mission 
        # Global 1 arc second V003
        short_name="SRTMGL1",
        bounding_box=bounds
    )

    # Download data granules
    earthaccess.download(strm_granules, download_dir)

    # Set SRTM data dir. hgt = height 
    strm_pattern = os.path.join(download_dir, '*.hgt.zip')

    # Build merged elevation DataArray
    strm_da = build_da(glob(strm_pattern), tuple(site_gdf.total_bounds))

    return strm_da  

eldorado_elev_da = select_dem(
                        tuple(eldorado_gdf.total_bounds), eldorado_gdf, 
                        eldorado_elevation_dir
                    )
export_raster(eldorado_elev_da, "eldorado_elevation.tif")

# Plot Eldorado elevation

eldorado_elev_plt = plot_site(
    eldorado_elev_da, eldorado_gdf, 
    'Eldorado-National-Forest-Elevation', 
    'Eldorado National Forest Elevation', 
    'Meters', 'terrain', 'black'
)

eldorado_elev_plt

los_padres_elev_da = select_dem(
                        tuple(los_padres_gdf.total_bounds), los_padres_gdf, 
                        los_padres_elevation_dir
                    )
export_raster(los_padres_elev_da, "los_padres_elevation.tif")

# Plot Los Padres elevation

los_padres_elev_plt = plot_site(
    los_padres_elev_da, los_padres_gdf, 
    'Los-Padres-National-Forest-Elevation', 
    'Los Padres National Forest Elevation', 
    'Meters', 'terrain', 'black'
)

los_padres_elev_plt

def calculate_aspect(elev_da):
    """
    Create aspect DataArray from site elevation.
    
    Args:
    elev_da (xarray.DataArray): Input raster layer.

    Returns:
    xarray.DataArray: A raster of site aspect. 
    """

    # Calculate aspect (degrees)
    aspect_da = xrspatial.aspect(elev_da)
    aspect_da = aspect_da.where(aspect_da >= 0)

    return aspect_da

# Calculate aspect

eldorado_aspect_da = calculate_aspect(eldorado_elev_da)
export_raster(eldorado_aspect_da, "eldorado_aspect.tif")

los_padres_aspect_da = calculate_aspect(los_padres_elev_da)
export_raster(los_padres_aspect_da, "los_padres_aspect.tif")

# Plot Eldorado aspect

eldorado_aspect_plt = plot_site(
    eldorado_aspect_da, eldorado_gdf, 
    'Eldorado-National-Forest-Aspect-Degrees', 
    'Eldorado National Forest Aspect', 
    'Degrees', 'terrain', 'black'
)

eldorado_aspect_plt

# Plot Los Padres aspect

los_padres_aspect_plt = plot_site(
    los_padres_aspect_da, los_padres_gdf,
    'Los-Padres-National-Forest-Aspect-Degrees', 
    'Los Padres National Forest Aspect', 
    'Degrees', 'terrain', 'black'
)

los_padres_aspect_plt

def get_projected_climate(site_gdf_dict, emissions_scenarios, 
                          start_year, end_year):
    """
    Create DataFrame of projected site climate.
    
    Args:
    site_gdf_dict (dict): Input dictionary of site scenario temperatures.
    emissions_scenarios (list): List of climate scenarios.
    start_year (str): Projected start year.
    end_year (str): Projected end year. 

    Returns:
    pandas.DataFrame: Projected site average max temperatures.
    """

    maca_da_list = []

    for site_name, site_gdf in site_gdf_dict.items():
        for scenario in emissions_scenarios:
            
            MACA_URL = (
                'http://thredds.northwestknowledge.net:8080/thredds/dodsC' 
                '/MACAV2/CanESM2/macav2metdata_tasmax_CanESM2_r1i1p1_'
                f'{scenario}_{start_year}_{end_year}_CONUS_monthly.nc'
            )

            maca_da = xr.open_dataset(MACA_URL).squeeze()

            bounds = site_gdf.to_crs(maca_da.rio.crs).total_bounds

            # update coordinate range
            maca_da = maca_da.assign_coords(
                lon=("lon", [convert_longitude(l) for l in maca_da.lon.values])
            )

            maca_da = maca_da.rio.set_spatial_dims(x_dim='lon', y_dim='lat')

            maca_da = maca_da.rio.clip_box(*bounds)

            maca_da_list.append(dict(
                site_name=site_name,
                climate_scenario=scenario,
                da=maca_da)
            )

    maca_df = pd.DataFrame(maca_da_list)
    return maca_df 

# RCP 4.5 (Medium Emissions Scenario) and 8.5 (High Emissions Scenario)
emissions_scenarios = ['rcp45', 'rcp85']
site_gdf_dict = {'eldorado': eldorado_gdf, 'los padres': los_padres_gdf}
start_year = '2096'
end_year = '2099'

site_proj_temp = get_projected_climate(
                    site_gdf_dict, emissions_scenarios, 
                    start_year, end_year
                )

# Convert temperature from Kelvin to Fahrenheit

converted_das = []

rcp_rasters = [
    'eldorado_rcp45.tif','eldorado_rcp85.tif',
    'los_padres_rcp45.tif','los_padres_rcp85.tif'
]

for da, raster_name in zip(site_proj_temp.da.values, rcp_rasters): 
    da = da.assign(fahrenheit=da["air_temperature"] * 1.8 - 459.67)
    da = da.rio.write_crs(eldorado_soil_da.rio.crs)

    da = da.rename({'lon': 'x','lat': 'y'})

    da = da.fahrenheit
    converted_das.append(da)
    export_raster(da, raster_name)

[eldorado_rcp45_da, eldorado_rcp85_da, 
 los_padres_rcp45_da, los_padres_rcp85_da] = converted_das

# Select monthly average of daily maximum near-surface air temperature

air_eld_45 = eldorado_rcp45_da.sel(
                x=-120.9,y=38.4, method='nearest'
            )
air_eld_85 = eldorado_rcp85_da.sel(
                x=-120.9,y=38.4, method='nearest'
            )

air_lp_45 = los_padres_rcp45_da.sel(
                x=-118.7,y=36.4, method='nearest'
            )
air_lp_85 = los_padres_rcp85_da.sel(
                x=-118.7,y=36.4, method='nearest'
            )

# Plot Eldorado monthly mean max temperature data (2096-2099)

(
    air_eld_45.hvplot.line(
        title='Eldorado RCP 4.5 Average Max Temperature (2096-2099)',
        frame_width=550, frame_height=320
    )
    +
    air_eld_85.hvplot.line(
        title='Eldorado RCP 8.5 Average Max Temperature (2096-2099)',
        frame_width=550, frame_height=320, line_color='orange'
    )
)

# Plot Los Padres monthly mean max temperature data (2096-2099)

(
    air_lp_45.hvplot.line(
        title='Los Padres RCP 4.5 Average Max Temperature (2096-2099)',
        frame_width=550, frame_height=320
    )
    +
    air_lp_85.hvplot.line(
        title='Los Padres RCP 8.5 Average Max Temperature (2096-2099)',
        frame_width=550, frame_height=320, line_color='orange'
    )
)

def select_annual_temp(year, da, raster_name):
    """
    Obtain annual temperature average and export raster.
    
    Args:
    year (str): Target year for temperature selection. 
    da (xarray.DataArray): DataArray of projected temperature values.
    raster_name (str): The name of model raster.

    Returns: None
    """
    da_annual = (
        da
        # select year 
        .sel(time=year)
        # take annual average
        .mean('time')
        # write projection EPSG:4326
        .rio.write_crs(4326)
    )

    export_raster(da_annual, raster_name)

# Eldorado projected temperature
select_annual_temp('2099', eldorado_rcp45_da, 'air_eld_45_2099.tif')

# Los Padres projected temperature 
select_annual_temp('2099', los_padres_rcp45_da, 'air_lp_45_2099.tif')

def calculate_suitability_score(raster, optimal_value, tolerance_range):
    """ 
    Calculate a fuzzy suitability score (0–1) for each raster cell based on 
    proximity to the optimal value.

    Args:
    raster (xarray.DataArray): Input raster layer. 
    optimal_value (float): Optimal value for the variable.
    tolerance_range (float): Suitable values range. 

    Returns:
    xarray.DataArray: A raster of suitability scores (0-1).
    """

    # Calculate using a fuzzy Gaussian function to assign scores 
    # between 0 and 1
    suitability = np.exp(
                    -((raster - optimal_value) ** 2) 
                    / (2 * tolerance_range ** 2)
                )

    # Suitability scores (0–1) 
    return suitability

def build_habitat_suitability_model(input_rasters, optimal_values, 
                                    tolerance_ranges, output_dir, raster_name):
    """ 
    Build a habitat suitability model by combining fuzzy suitability scores 
    for each variable. 

    Args:
    input_rasters (list): List of input rasters.
    optimal_values (list): List of optimal values for variables.
    tolerance_ranges (list): List of tolerance values for variables.
    output_dir (str): The model output directory.
    raster_name (str): The name of model raster.

    Returns:
    str: The path of the suitability raster.
    """
    
    os.makedirs(output_dir, exist_ok=True)

    # Load and calculate suitability scores for each raster
    suitability_layers = []
    suit_zip = zip(input_rasters, optimal_values, tolerance_ranges)
    for raster_path, optimal_value, tolerance_range in suit_zip:
        raster = rxr.open_rasterio(raster_path, masked=True).squeeze()
        suitability_layer = calculate_suitability_score(
                                raster, optimal_value, tolerance_range
                            )
        suitability_layers.append(suitability_layer)

    # Combine suitability scores by multiplying across all layers
    combined_suitability = suitability_layers[0]
    for layer in suitability_layers[1:]:
        combined_suitability *= layer

    # Save the combined suitability raster
    output_file = os.path.join(output_dir, f"{raster_name}.tif")
    combined_suitability.rio.to_raster(output_file)
    print(f"Combined suitability raster saved to: {raster_name}.tif")

    # Path to the final combined suitability raster
    return output_file

# Optimal values for species for each variable
# Variables: elevation, temperature, aspect, soil pH
optimal_values = [1500.0, 90.0, 0.0, 6.75]

# Tolerance ranges (acceptable deviation) for each variable
tolerance_ranges = [1500.0, 20.0, 65.0, 0.75]

eldorado_45_input_rasters = [
    f"{data_dir}/air_eld_45_2099.tif",
    f"{data_dir}/eldorado_aspect.tif",
    f"{data_dir}/eldorado_soil_ph.tif"
]

# Ensure consistent spatial resolution and projection
eldorado_rcp45_rasters = harmonize_raster_layers(
                            f"{data_dir}/eldorado_elevation.tif", 
                            eldorado_45_input_rasters, data_dir
                        )

eldorado_rcp45_combined_suitability = build_habitat_suitability_model(
    eldorado_rcp45_rasters, optimal_values, tolerance_ranges, 
    data_dir, 'eldorado_rcp45_suitability'
)

eldorado_rcp45_suitability_plt = plot_site(
    f"{data_dir}/eldorado_rcp45_suitability.tif", eldorado_gdf, 
    'Eldorado-National-Forest-Suitability', 
    'Eldorado National Forest Suitability: RCP 4.5 (2099)', 
    'Suitability', 'YlGn', 'black', tif_file=True
)

eldorado_rcp45_suitability_plt

los_padres_45_input_rasters = [
    f"{data_dir}/air_lp_45_2099.tif",
    f"{data_dir}/los_padres_aspect.tif",
    f"{data_dir}/los_padres_soil_ph.tif"
]

los_padres_rcp45_rasters = harmonize_raster_layers(
                            f"{data_dir}/los_padres_elevation.tif", 
                            los_padres_45_input_rasters, data_dir
                        )

los_padres_rcp45_combined_suitability = build_habitat_suitability_model(
    los_padres_rcp45_rasters, optimal_values, tolerance_ranges, 
    data_dir, 'los_padres_rcp45_suitability'
)

los_padres_rcp45_combined_suitability_plt = plot_site(
    f"{data_dir}/los_padres_rcp45_suitability.tif", los_padres_gdf, 
    'Los-Padres-National-Forest-Suitability', 
    'Los Padres National Forest Suitability: RCP 4.5 (2099)', 
    'Suitability', 'YlGn', 'black', tif_file=True
)

los_padres_rcp45_combined_suitability_plt