from datetime import datetime
import pyspark.sql.functions as f
# Date range for imagery to be used by the model
start_date = datetime(2021, 1, 1)
end_date = datetime(2022, 1, 1)
# The Meta CHM v1 recipe runs on 60cm NAIP imagery. The index holds one row per NAIP scene, with its
# footprint, resolution (res) and acquisition time.
MODEL_RES = 0.6
naip_index = (
sedona.read.format("geoparquet")
.load("s3://wherobots-examples/rasterflow/indexes/naip_index.parquet")
.select("geometry", "res", "year", "time")
)
# `time` is stored with nanosecond precision, which Spark reads as a BIGINT of nanoseconds since the epoch
if dict(naip_index.dtypes)["time"] == "bigint":
naip_index = naip_index.withColumn("time", f.timestamp_micros((f.col("time") / 1000).cast("long")))
# The AOI as a one-row DataFrame, so the checks below can join against it
aoi_df = sedona.createDataFrame([(aoi,)], ["wkt"]).select(ST_GeomFromWKT("wkt").alias("aoi"))
# Scenes at the model's resolution touching the AOI (the index and wkls are both EPSG:4326)
nearby = (
naip_index.filter(f.col("res") == MODEL_RES)
.join(aoi_df, ST_Intersects("geometry", "aoi"))
.select("geometry", "year", "time")
.cache()
)
# Narrowed to the requested date range
covering = nearby.filter(f.col("time").between(start_date, end_date))
# The recipe needs the AOI fully inside the union of the matching scene footprints
footprint = covering.agg(ST_Union_Aggr("geometry").alias("footprint"))
coverage = aoi_df.crossJoin(footprint).select(
ST_Contains("footprint", "aoi").alias("covered"),
# Share of the AOI under the footprint, as geodesic area on the WGS84 spheroid
(1.0 - ST_AreaSpheroid(ST_Difference("aoi", "footprint")) / ST_AreaSpheroid("aoi")).alias("fraction"),
).first()
if not coverage["covered"]:
fraction = max(0.0, coverage["fraction"] or 0.0)
# Years that would work, applying the same full-coverage test as the gate above
by_year = nearby.groupBy("year").agg(ST_Union_Aggr("geometry").alias("footprint"))
available = [
row.year
for row in by_year.crossJoin(aoi_df).filter(ST_Contains("footprint", "aoi")).orderBy("year").collect()
]
raise ValueError(
f"{MODEL_RES:g}m NAIP covers {fraction:.1%} of this AOI between "
f"{start_date:%Y-%m-%d} and {end_date:%Y-%m-%d}; the recipe needs the AOI fully covered. "
+ (
f"Years with complete {MODEL_RES:g}m coverage over this AOI: {available}."
if available
else f"No year has complete {MODEL_RES:g}m NAIP coverage over this AOI."
)
)
years = sorted(row.year for row in covering.select("year").distinct().collect())
print(f"NAIP coverage is available. Found {covering.count()} NAIP scenes at {MODEL_RES:g}m covering the AOI")
print(f"Acquisition years: {years}")
