Mapping Invasive Species Using AVIRIS-NG Data

# Modules imported separately - not available at the time of the workshop managed environment 
# !pip3 install --user xvec
# !pip3 install --user shap
# !pip3 install --user xgboost

from os import path
import datetime as dt 
import geopandas as gpd
import s3fs
import pandas as pd
import xarray as xr
from shapely.geometry import box, mapping, Polygon, MultiPolygon
import rioxarray as riox
import numpy as np
import netCDF4 as nc
import hvplot.xarray
import holoviews as hv
import xvec
import shap
import xgboost
import matplotlib.pyplot as plt
from dask.diagnostics import ProgressBar
import warnings
import earthaccess # >0.17
import dask
#our functions
from utils import get_first_xr

warnings.filterwarnings('ignore')
hvplot.extension('bokeh')

# search collections
collections = earthaccess.search_datasets(keyword='BioSCape')
# Print Quantity of Results
print(f'Collections found: {len(collections)}')

collections_info = [
    {
        'entry_title': c.get_umm('EntryTitle'),
        'DOI': c.get_umm('DOI')['DOI']
    }
    for c in collections
]
pd.set_option('display.max_colwidth', 75)
collections_info = pd.DataFrame(collections_info)
collections_info

# let's first create a DataFrame and assign labels to each class
label_df = pd.DataFrame({'LandType': ['Bare ground/Rock','Mature Fynbos', 
              'Recently Burnt Fynbos', 'Wetland', 
              'Forest', 'Pine', 'Eucalyptus' , 'Wattle', 'Water'],
               'class': ['0','1','2','3','4','5','6','7','8']})

label_df

# Now let's open the dataset, project to the South African UTM projection also used by the AVIRIS-NG airborne data, and merge it with the label data frame created above.
class_data = gpd.read_file('data/ct_invasive.gpkg')
# class_data.crs
class_data_utm = (class_data
                 .to_crs("EPSG:32734")
                 .merge(label_df, on='class', how='left')
                 )
class_data_utm

# examine the number of plots of each land type
class_data_utm.groupby(['LandType']).size()

# The group class will be used to determine training and test data
class_data_utm.groupby(['group']).size()

# Let's visualize the plot data in an interactive map, with color by class, using a Google satellite basemap
class_data_utm[['LandType', 'geometry']].explore('LandType', tiles='https://mt1.google.com/vt/lyrs=s&x={x}&y={y}&z={z}', attr='Google')

# Earthdata Authentication
auth = earthaccess.login()

# search granules
granules = earthaccess.search_data(
    doi='10.3334/ORNLDAAC/2427', # BioSCape: AVIRIS-NG L3
    granule_name = "AVIRIS-NG*", # exclude geojson tile
)
print(f"Total Datasets (results) found: {len(granules)}")

# print one granule at index 100
granules[100]

vars(granules[100])

AVNG_Coverage = gpd.GeoDataFrame(granules, geometry=gpd.GeoSeries(granules, crs=4326))
AVNG_Coverage = AVNG_Coverage.join(pd.json_normalize(AVNG_Coverage.pop('umm')))
AVNG_Coverage[['GranuleUR', 'geometry']].explore(tiles='https://mt1.google.com/vt/lyrs=s&x={x}&y={y}&z={z}', attr='Google')

single_granule = earthaccess.search_data(
    doi='10.3334/ORNLDAAC/2427', # BioSCape: AVIRIS-NG L3
    granule_name = "AVIRIS-NG_BIOSCAPE_V02_L3_36_11*", # select only one granule
)
single_granule[0].data_links(access="direct")

# earthaccess open
fh = earthaccess.open(single_granule)
fh

def separate_granule_type(fo):
    """separating granules by type"""
    return ([f for f in fo if f.path.endswith("_RFL.nc")],
            [f for f in fo if f.path.endswith("_QL.tif")],
            [f for f in fo if f.path.endswith("_UNC.nc")])
rfl_f, ql_f, unc_f = separate_granule_type(fh)

# plot a single file netcdf
rfl_netcdf = xr.open_datatree(rfl_f[0],
                              engine='h5netcdf', chunks={})
rfl_netcdf = rfl_netcdf.reflectance.to_dataset()
rfl_netcdf = rfl_netcdf.reflectance.where(rfl_netcdf.reflectance>0)
rfl_netcdf

h = rfl_netcdf.sel({'wavelength': 660},method='nearest').hvplot('easting', 'northing',
                                                      rasterize=True, data_aspect=1,
                                                      cmap='magma',frame_width=400,clim=(0,0.3))
h

# plot a single file geotif
ql_tif = xr.open_dataset(ql_f[0], engine='rasterio')
h = ql_tif.hvplot.rgb('x', 'y', bands='band', rasterize=True, data_aspect=1, frame_width=400)
h

# search granules that spatially intersects with class_data
granules_cp = earthaccess.search_data(
    doi='10.3334/ORNLDAAC/2427', # BioSCape: AVIRIS-NG L3
    bounding_box = tuple(class_data.total_bounds),
    granule_name = "AVIRIS-NG*", # exclude geojson tile
)

AVNG_Coverage = gpd.GeoDataFrame(granules_cp, geometry=gpd.GeoSeries(granules_cp, crs=4326))
AVNG_Coverage = AVNG_Coverage.join(pd.json_normalize(AVNG_Coverage.pop('umm')))
AVNG_Coverage[['GranuleUR', 'geometry']].explore('GranuleUR', tiles='https://mt1.google.com/vt/lyrs=s&x={x}&y={y}&z={z}', attr='Google')

f_cp = earthaccess.open(granules_cp)
rfl_f, ql_f, unc_f = separate_granule_type(f_cp)
AVNG_CP = AVNG_Coverage.to_crs("EPSG:32734")
ds_all =[]                                     # initialize empty list ds_all to store individual data subsets extracted from each granule
for rfl in rfl_f:
    granuleur = path.basename(rfl.path)[:-7]   # extracts a native identifier from the file path. -7 removes the last 7 characters 
    geof = AVNG_CP[AVNG_CP["GranuleUR"]==granuleur].geometry    # filters AVNG_CP to only the geometry with a matching native ID 
    points = class_data_utm.clip(geof)         # restricts the point set to only those points that fall within the geometry’s boundaries
    n_points = points.shape[0]
    if n_points:
        print(f'got {n_points} points from {granuleur}')
        ds = xr.open_datatree(rfl, engine='h5netcdf', chunks={})  # opens the granule as an xarray datatree
        points = points.to_crs(ds.transverse_mercator.crs_wkt)    # reprojects the points to the coordinate reference system used in the dataset
        ds_all.append(ds.reflectance.to_dataset().xvec.extract_points(points['geometry'],   # extracts the corresponding reflectance values 
                                                                x_coords="easting",         # along a vector (xvec) based on the 
                                                                y_coords="northing",        # provided x (easting) and y (northing) coordinate
                                                                index=True))
ds_all = xr.concat(ds_all, dim='file')

ds_all

# create a data set with just the spectra
# get_first_xr is a helper function from the utils.py import
ds = get_first_xr(ds_all)
ds

class_xr =class_data_utm[['class','group']].to_xarray()
ds = ds.merge(class_xr.astype(int),join='left')
ds

##  PLEASE NOTE THE LONGER RUN TIME
with ProgressBar():
    dsp = ds.persist()

## SAVING THIS OUTPUT (dsp.nc) 
# dsp.drop_vars('geometry').to_netcdf('data/dsp.nc')
# dsp = xr.open_dataset('data/dsp.nc')
dsp

# recall the class types
label_df

dsp_plot = dsp.where(dsp['class']==6, drop=True)
h = dsp_plot['reflectance'].hvplot.line(x='wavelength',by='index',
                                    color='green', alpha=0.5,legend=False)
h

wavelengths_to_drop = dsp.wavelength.where(
    (dsp.wavelength < 450) |
    (dsp.wavelength >= 1340) & (dsp.wavelength <= 1480) |
    (dsp.wavelength >= 1800) & (dsp.wavelength <= 1980) |
    (dsp.wavelength > 2400), drop=True
)

# Use drop_sel() to remove those specific wavelength ranges
dsp = dsp.drop_sel(wavelength=wavelengths_to_drop)

mask = (dsp['reflectance'] > -1).all(dim='wavelength')  # Create a mask where all values along 'z' are non-negative
dsp = dsp.sel(index=mask)
dsp

# Calculate the L2 norm along the 'wavelength' dimension
l2_norm = np.sqrt((dsp['reflectance'] ** 2).sum(dim='wavelength'))  # the square root of the sum of the squares of the values across dim wavelength

# Normalize the reflectance by dividing by the L2 norm
dsp['reflectance'] = dsp['reflectance'] / l2_norm  # removes brightness differences due to varying magnitudes and ensures that 
                                                   # the reflectance values across wavelengths are unit-normalized.

dsp_norm_plot = dsp.where(dsp['class']==6, drop=True)
h = dsp_norm_plot['reflectance'].hvplot.line(x='wavelength',by='index',
                                         color='green',ylim=(-0.01,0.15),alpha=0.5,legend=False)
h

import xgboost as xgb
from sklearn.model_selection import GridSearchCV
from sklearn.metrics import accuracy_score, f1_score, precision_score, recall_score, confusion_matrix, ConfusionMatrixDisplay

# recall groups
class_data_utm.groupby(['group']).size()

dtrain = dsp.where(dsp['group']==1,drop=True)  # the training set
dtest = dsp.where(dsp['group']==2,drop=True)   # the evaluation (testing) set

#create separte datasets for labels and features
y_train = dtrain['class'].values.astype(int)  # the variable 'class' is extracted from both the training (dtrain) and testing (dtest) datasets
y_test = dtest['class'].values.astype(int)
X_train = dtrain['reflectance'].values        # The variable 'reflectance' is extracted from both the training (dtrain) and testing (dtest) datasets
X_test = dtest['reflectance'].values          # .values: Converts the 'reflectance' variable into a NumPy array for the machine learning models

# Define the hyperparameter grid
param_grid = {                   # param_grid dictionary specifies the hyperparameters to be tuned and the values to be tested for each hyperparameter
    'max_depth': [5],            # maximum depth of a tree, controlling the complexity of the individual tree
    'learning_rate': [0.1],      # rate at which the model updates weights during training
    'subsample': [0.75],         # fraction of the training data to randomly sample for each tree
    'n_estimators' : [50,100]    # number of boosting rounds (trees) in the ensemble
}

# Create the XGBoost model object
xgb_model = xgb.XGBClassifier(tree_method='hist')

# Create the GridSearchCV object
grid_search = GridSearchCV(xgb_model, param_grid, cv=5, scoring='accuracy')

# Fit the GridSearchCV object to the training data
grid_search.fit(X_train, y_train)

# Print the best set of hyperparameters and the corresponding score
print("Best set of hyperparameters: ", grid_search.best_params_)  # the combination of hyperparameters that achieved the highest accuracy during cross-validation
print("Best score: ", grid_search.best_score_)    # the cross-validation accuracy score corresponding to the best hyperparameters
best_model = grid_search.best_estimator_    # Retrieves the trained model (XGBoost classifier) corresponding to the best combination of hyperparameters 
                                            # found during the grid search

y_pred = best_model.predict(X_test)

# Step 2: Calculate acc and F1 score for the entire dataset
acc = accuracy_score(y_test, y_pred)
print(f"Accuracy: {acc}")

f1 = f1_score(y_test, y_pred, average='weighted')  # 'weighted' accounts for class imbalance
print(f"F1 Score (weighted): {f1}")

# Step 3: Calculate precision and recall for class 6 (Eucalyptus)
precision_class_6 = precision_score(y_test, y_pred, labels=[6], average='macro', zero_division=0)
recall_class_6 = recall_score(y_test, y_pred, labels=[6], average='macro', zero_division=0)

print(f"Precision for Class 6: {precision_class_6}")
print(f"Recall for Class 6: {recall_class_6}")

# Step 4: Plot the confusion matrix
conf_matrix = confusion_matrix(y_test, y_pred)

ConfusionMatrixDisplay(confusion_matrix=conf_matrix).plot()
plt.show()

# list the LandTypes and class
label_df

import shap
shap.initjs()

# Initialize the SHAP Tree Explainer for XGBoost
explainer = shap.TreeExplainer(best_model,feature_names=dsp.wavelength.values.astype(int).astype(str))
shap_values = explainer(X_test)

sel_in = 55

print(f"predicted class: {label_df.loc[y_pred[sel_in],'LandType']}")
print(f"true class: {label_df.loc[y_test[sel_in], 'LandType']}")

#lables for wavelength in plot
feats =  dsp.wavelength.values.astype(int).astype(str)
shap.force_plot(explainer.expected_value[y_test[sel_in]], shap_values.values[sel_in,:,y_test[sel_in]], pd.DataFrame(X_test).iloc[sel_in, :],link='logit',feature_names=list(feats))

sel_in = 59

print(f"predicted class: {label_df.loc[y_pred[sel_in],'LandType']}")
print(f"true class: {label_df.loc[y_test[sel_in], 'LandType']}")

#lables for wavelength in plot
feats =  dsp.wavelength.values.astype(int).astype(str)
shap.force_plot(explainer.expected_value[y_test[sel_in]], shap_values.values[sel_in,:,y_test[sel_in]], pd.DataFrame(X_test).iloc[sel_in, :],link='logit',feature_names=list(feats))

sel_in = 55
shap.plots.beeswarm(shap_values[:,:,y_test[sel_in]])

#plot spectra and importance
importance = np.abs(shap_values[sel_in,:,y_test[sel_in]].values)
# Create the base plot
fig, ax = plt.subplots(figsize=(12, 6))

# Get the wavelength and importance data
wavelength = dsp['wavelength'].values
importance = importance  # Make sure this aligns with your wavelength data

# Create a colormap
cmap = plt.get_cmap('hot').reversed()  # You can choose a different colormap if you prefer

# Normalize importance values to [0, 1] for colormap
norm = plt.Normalize(importance.min(), importance.max())

# Add shading
for i in range(len(wavelength) - 1):
    ax.fill_between([wavelength[i], wavelength[i+1]], 0, 1, 
                    color=cmap(norm(importance[i])), alpha=0.3)

# Add a colorbar to show the importance scale
sm = plt.cm.ScalarMappable(cmap=cmap, norm=norm)
sm.set_array([])
cbar = plt.colorbar(sm, ax=ax, label='Importance')

# Add white blocks to obscure specified regions
ax.fill_between([0,450], 0, 1, color='white')
ax.fill_between([1340,1480], 0, 1, color='white')
ax.fill_between([1800,1980], 0, 1, color='white')
ax.fill_between([2400,2500], 0, 1, color='white')
ax.set_xlim(420,2400)

plot_xr = xr.DataArray(X_test[sel_in], coords=[wavelength], dims=["wavelength"])
plot_xr.plot.line(x='wavelength', color='green', ylim=(0, 0.2), ax=ax,zorder=0)

plt.title('Reflectance with Importance Shading')
plt.xlabel('Wavelength')
plt.ylabel('Reflectance (normalized)')

plt.tight_layout()
plt.show()

#south africa protected areas
SAPAD = (gpd.read_file('data/SAPAD_2024.gpkg')
         .query("SITE_TYPE!='Marine Protected Area'")
        )
SAPAD.crs

SAPAD.keys()

m = SAPAD[['SITE_TYPE', 'geometry']].explore('SITE_TYPE')
m

#south africa protected areas
SAPAD = (gpd.read_file('data/SAPAD_2024.gpkg')
         .query("SITE_TYPE!='Marine Protected Area'")
        )
SAPAD = SAPAD.to_crs("EPSG:32734")

# Get the bounding box of the training data
bbox = class_data_utm.total_bounds  # (minx, miny, maxx, maxy)
gdf_bbox = gpd.GeoDataFrame({'geometry': [box(*bbox)]}, crs=class_data_utm.crs)  # Specify the CRS
gdf_bbox['geometry'] = gdf_bbox.buffer(500)

# protected areas that intersect with the training data
SAPAD_CT = SAPAD.overlay(gdf_bbox,how='intersection')

#keep only AVIRIS scenes that intersects with CT protected areas
AVNG_sapad = AVNG_CP[AVNG_CP.intersects(SAPAD_CT.unary_union)]

# plot
m = AVNG_sapad[['GranuleUR','geometry']].explore('GranuleUR')
m

def predict_on_chunk(chunk, model):
    probabilities = model.predict_proba(chunk)
    return probabilities

def predict_xr(file,geometries):

    rfl_f, ql_f, unc_f = separate_granule_type(file)
    native_id = path.basename(rfl_f[0].path)[:-7]
    
    #part 1 - opening file
    #open the file
    print(f'file: {rfl_f[0]}')
    ds = xr.open_datatree(rfl_f[0], engine='h5netcdf', chunks='auto')
    
    #get the geometries of the protected areas for masking
    ds_crs = ds.transverse_mercator.crs_wkt
    geometries = geometries.to_crs(ds_crs).geometry.apply(mapping)
    
    #condition to use for masking no data later
    condition = (ds['reflectance'] > -1).any(dim='wavelength')
    
    #stack the data into a single dimension. This will be important for applying the model later
    ds = ds.reflectance.to_dataset().stack(sample=('easting','northing'))
    
    #part 2 - pre-processing
    #remove bad wavelenghts
    wavelengths_to_drop = ds.wavelength.where(
        (ds.wavelength < 450) |
        (ds.wavelength >= 1340) & (ds.wavelength <= 1480) |
        (ds.wavelength >= 1800) & (ds.wavelength <= 1980) |
        (ds.wavelength > 2400), drop=True
    )
    # Use drop_sel() to remove those specific wavelength ranges
    ds = ds.drop_sel(wavelength=wavelengths_to_drop)
    
    #normalise the AVIRIS data
    l2_norm = np.sqrt((ds['reflectance'] ** 2).sum(dim='wavelength'))
    ds['reflectance'] = ds['reflectance'] / l2_norm
    
     
    #part 3 - apply the model over chunks
    result = xr.apply_ufunc(
        predict_on_chunk,
        ds['reflectance'].chunk(dict(wavelength=-1)),
        input_core_dims=[['wavelength']],    # input dim with features
        output_core_dims=[['class']],        # name for the new output dim
        exclude_dims=set(('wavelength',)),   # dims to drop in result
        output_sizes={'class': 9},           # length of the new dimension
        output_dtypes=[np.float32],
        dask="parallelized",
        kwargs={'model': best_model}
    )
    
    #part 4 - post-processing
    result = result.where((result >= 0) & (result <= 1), np.nan)  # valid values
    result = result.unstack('sample')                             # remove the stack
    result = result.rio.set_spatial_dims(x_dim='easting',y_dim='northing') # set the spatial dims
    result = result.rio.write_crs(ds_crs)                         # set the CRS
    result = result.rio.clip(geometries)                          # clip to the protected areas and no data
    result = result.transpose('class', 'northing', 'easting')     # transpose the data - rio expects it this way
    
    return result   

test  = predict_xr(earthaccess.open(single_granule),SAPAD_CT)
test

# Again, recall the labels and LandType classes
label_df

# Reproject the dataset to WGS84 (EPSG:4326)
test = test.rio.reproject("EPSG:4326", nodata=np.nan)

# Enhanced Plot with Titles and Styling
h = test.isel({'class': 6}).hvplot(
    tiles=hv.element.tiles.EsriImagery(),   # Use Esri Imagery as a tile background
    project=True,                          # Project the data to work seamlessly with the background
    rasterize=True,                        # Rasterize for efficient plotting
    clim=(0.01, 1),                           # Define color limits to focus on the range of interest
    cmap='turbo',                          # Colormap for better visualization
    frame_width=800,                       # Increase the width for better spatial resolution
    data_aspect=1,                         # Maintain a realistic and proportional aspect ratio
    alpha=0.5,                             # Adjust opacity for optimal visualization
    colorbar=True                          # Enable the colorbar (hvplot handles this properly)
)

# Add Descriptive Titles (handled internally by hvplot)
h = h.opts(
    title="Class6: Eucalyptus - Probability Mapping"  # Set title directly using Holoviews options
)
h 

grid_pred = []
# loop over tile names intersecting the protected areas
for t in AVNG_sapad['GranuleUR'].to_list():
    # search for each tile
    g = earthaccess.search_data(
        doi='10.3334/ORNLDAAC/2427', # BioSCape: AVIRIS-NG L3
        granule_name = f"{t}*", # select the tile
    )
    # run prediction and append the output xarray to grid_pred
    grid_pred.append(predict_xr(earthaccess.open(g),SAPAD_CT))

from rioxarray.merge import merge_arrays

# merged = merge_arrays(grid_pred, nodata=np.nan)
# merged = merged.rio.reproject("EPSG:4326")
# merged.rio.to_raster('data/ct_invasive.tif',driver="COG")

import xarray as xr
import rioxarray as riox
import hvplot.xarray
import holoviews as hv
hvplot.extension('bokeh')

merged = xr.open_dataset('../../shared-public/BioSCape_AVIRIS-NG_L3/ct_invasive_2.tif', engine='rasterio', chunks='auto')
merged

h = merged.isel({'band':6}).hvplot(x='x',y='y',
                               tiles='ESRI',geo=True,
                                project=True,rasterize=True,
                                cmap='turbo',clim=(0,1), frame_width=400,data_aspect=1,alpha=0.5)
h = h.opts(
    title="Class6: Eucalyptus - Probability Mapping"  # Set title directly using Holoviews options
)
h