spatrial_sample_group.py

import random
import numpy as np
import pandas as pd
import rasterio as rio
import geopandas as gpd
from sklearn.cluster import KMeans


def sample_coordinates(bool_array, n):
    true_coords = np.argwhere(bool_array)

    if len(true_coords) < n:
        return true_coords
        # raise ValueError("Not enough True values to sample from")

    kmeans = KMeans(n_clusters=n)
    kmeans.fit(true_coords)

    sampled_coords = list()
    for cluster_idx in range(n):
        cluster_points = true_coords[kmeans.labels_ == cluster_idx]
        sampled_point = cluster_points[random.randint(0, len(cluster_points) - 1)]
        sampled_coords.append(tuple(sampled_point))

    return np.array(sampled_coords)

def sample_raster(src_path):
    starts = [0.0001, 0.0002, 0.0003, 0.0004, 0.0005, 0.0006, 0.0007, 0.0009, 0.001, 0.002]
    stops = starts[1:] + [np.inf]
    groups = [chr(65+i) for i in range(len(starts))]
    frames = list()
    with rio.open(src_path, 'r') as src:
        forward = np.array(src.transform).reshape((3, 3))
        img = src.read(1, masked=True)
        for a, b, g in zip(starts, stops, groups):
            mask = np.logical_and((a <= img), (img < b))
            try:
                coords = sample_coordinates(mask, 15)
            except ValueError:
                print(a, b)
            vals = img[coords[:, 0], coords[:, 1]]
            coords = np.stack([coords[:, 1], coords[:, 0]], axis=1)
            coords = np.concatenate([(coords + 0.5), np.ones((coords.shape[0], 1))], axis=1)
            coords = np.einsum(
                "ij,kj -> ki", forward, coords
            )
            coords[:, -1] = vals
            df = pd.DataFrame(coords, columns=["x", "y", "Value"])
            df["Group"] = g
            gdf = gpd.GeoDataFrame(df, geometry=gpd.points_from_xy(df.x, df.y))
            gdf = gdf.drop(columns=['x', 'y'])
            gdf.set_crs(crs=src.crs, inplace=True)
            frames.append(gdf)
    main_df = pd.concat(frames, ignore_index=True, axis=0)
    return main_df

if __name__ == "__main__":
    geo_df = sample_raster("gs_10m_meanm.tif")
    geo_df.to_file(filename="Sampled_Points.geojson", driver="GeoJSON")