# -*- coding: utf-8 -*- """ Extract selected BedMachine Antarctica NetCDF subdatasets from ./input and write COG outputs to ./output in EPSG:3031. Extracts: 1) bed elevation 2) ice surface elevation (mask excludes ocean) 3) ice thickness (mask excludes ocean) 4) bed elevation uncertainty 5) firn (mask excludes ocean) 6) dataid (integer, not masked) Additionally extracts bed elevation masked to each mask class: 0 = ocean 1 = ice_free_land 2 = grounded_ice 3 = floating_ice 4 = lake_vostok Assumptions: - Exactly one .nc file exists in ./input - BedMachine variable names are: bed, surface, thickness, errbed, firn, dataid, mask """ from pathlib import Path import sys import numpy as np from osgeo import gdal, osr gdal.UseExceptions() # ============================================================================= # CONFIG # ============================================================================= INPUT_DIR = Path.cwd() / "input" OUTPUT_DIR = Path.cwd() / "output" MASKED_OUTPUT_DIR = OUTPUT_DIR / "masked" # Output filenames (all will be written to OUTPUT_DIR as COGs) OUTPUT_NAMES = { "bed": "BedMachine_bed_elevation.tif", "surface": "BedMachine_ice_surface_elevation.tif", "thickness": "BedMachine_ice_thickness.tif", "errbed": "BedMachine_uncertainty.tif", "firn": "BedMachine_firn.tif", "dataid": "BedMachine_dataID.tif", "bed_ocean": "BedMachine_bed_elevation_ocean.tif", "bed_ice_free_land": "BedMachine_bed_elev_ice_free_land.tif", "bed_grounded_ice": "BedMachine_bed_elev_grounded_ice.tif", "bed_floating_ice": "BedMachine_bed_elev_floating_ice.tif", "bed_lake_vostok": "BedMachine_bed_elev_lake_vostok.tif", } # Nodata to use in outputs NODATA_FLOAT = -9999.0 NODATA_INT = -9999 # COG creation options # OVERVIEWS=NONE ensures no overviews are generated COG_CREATION_OPTIONS = [ "COMPRESS=DEFLATE", "PREDICTOR=2", "BIGTIFF=IF_SAFER", "OVERVIEWS=NONE", "RESAMPLING=NEAREST", ] # Set True if you want quick console messages VERBOSE = True # ============================================================================= # HELPERS # ============================================================================= MASK_CLASSES = { 0: "ocean", 1: "ice_free_land", 2: "grounded_ice", 3: "floating_ice", 4: "lake_vostok", } VARIABLES = { "bed": "bed", "surface": "surface", "thickness": "thickness", "errbed": "errbed", "firn": "firn", "dataid": "dataid", "mask": "mask", } def log(msg: str) -> None: if VERBOSE: print(msg) def find_single_nc(input_dir: Path) -> Path: nc_files = sorted(input_dir.glob("*.nc")) if not nc_files: raise FileNotFoundError(f"No .nc files found in: {input_dir}") if len(nc_files) > 1: raise RuntimeError( f"Expected exactly one .nc file in {input_dir}, found {len(nc_files)}:\n" + "\n".join(str(p.name) for p in nc_files) ) return nc_files[0] def nc_subdataset_path(nc_path: Path, var_name: str) -> str: return f'NETCDF:"{nc_path}":{var_name}' def open_ds(path: str): ds = gdal.Open(path) if ds is None: raise RuntimeError(f"Could not open dataset: {path}") return ds def get_proj_wkt_epsg3031() -> str: srs = osr.SpatialReference() srs.ImportFromEPSG(3031) return srs.ExportToWkt() def read_array_and_meta(ds): band = ds.GetRasterBand(1) arr = band.ReadAsArray() nodata = band.GetNoDataValue() gt = ds.GetGeoTransform() proj = ds.GetProjection() xsize = ds.RasterXSize ysize = ds.RasterYSize dtype = band.DataType return arr, nodata, gt, proj, xsize, ysize, dtype def ensure_output_dir(path: Path) -> None: path.mkdir(parents=True, exist_ok=True) def write_cog_from_array( array: np.ndarray, out_path: Path, geotransform, projection_wkt: str, nodata, gdal_dtype, ): """ Write a single-band array to COG using an in-memory dataset. """ mem_driver = gdal.GetDriverByName("MEM") mem_ds = mem_driver.Create("", array.shape[1], array.shape[0], 1, gdal_dtype) mem_ds.SetGeoTransform(geotransform) mem_ds.SetProjection(projection_wkt) band = mem_ds.GetRasterBand(1) band.WriteArray(array) if nodata is not None: band.SetNoDataValue(nodata) band.FlushCache() cog_driver = gdal.GetDriverByName("COG") if cog_driver is None: raise RuntimeError("GDAL COG driver not available in this Python environment.") cog_driver.CreateCopy(str(out_path), mem_ds, options=COG_CREATION_OPTIONS) mem_ds = None def choose_output_dtype(var_key: str): """ Use Float32 for most outputs, Int16 for errbed and dataid. """ if var_key in ["errbed", "dataid"]: return gdal.GDT_Int16, NODATA_INT return gdal.GDT_Float32, NODATA_FLOAT def sanitise_to_output_dtype(array: np.ndarray, var_key: str): if var_key in ["errbed", "dataid"]: return array.astype(np.int16, copy=False) return array.astype(np.float32, copy=False) def build_valid_mask(data_array: np.ndarray, src_nodata): """ True where data is valid. """ if src_nodata is None: if np.issubdtype(data_array.dtype, np.floating): return np.isfinite(data_array) return np.ones(data_array.shape, dtype=bool) if np.issubdtype(data_array.dtype, np.floating): return np.isfinite(data_array) & (data_array != src_nodata) return data_array != src_nodata # ============================================================================= # MAIN # ============================================================================= def main(): ensure_output_dir(OUTPUT_DIR) ensure_output_dir(MASKED_OUTPUT_DIR) nc_path = find_single_nc(INPUT_DIR) log(f"Using NetCDF: {nc_path.name}") # Open core datasets ds_bed = open_ds(nc_subdataset_path(nc_path, VARIABLES["bed"])) ds_surface = open_ds(nc_subdataset_path(nc_path, VARIABLES["surface"])) ds_thickness = open_ds(nc_subdataset_path(nc_path, VARIABLES["thickness"])) ds_errbed = open_ds(nc_subdataset_path(nc_path, VARIABLES["errbed"])) ds_firn = open_ds(nc_subdataset_path(nc_path, VARIABLES["firn"])) ds_dataid = open_ds(nc_subdataset_path(nc_path, VARIABLES["dataid"])) ds_mask = open_ds(nc_subdataset_path(nc_path, VARIABLES["mask"])) # Read arrays + metadata bed, bed_nd, gt, proj, xsize, ysize, _ = read_array_and_meta(ds_bed) surface, surface_nd, gt2, proj2, _, _, _ = read_array_and_meta(ds_surface) thickness, thickness_nd, gt3, proj3, _, _, _ = read_array_and_meta(ds_thickness) errbed, errbed_nd, gt4, proj4, _, _, _ = read_array_and_meta(ds_errbed) firn, firn_nd, gt5, proj5, _, _, _ = read_array_and_meta(ds_firn) dataid, dataid_nd, gt6, proj6, _, _, _ = read_array_and_meta(ds_dataid) mask, mask_nd, gt7, proj7, _, _, _ = read_array_and_meta(ds_mask) # Basic consistency checks if not ( bed.shape == surface.shape == thickness.shape == errbed.shape == firn.shape == dataid.shape == mask.shape ): raise RuntimeError("Input subdatasets do not all have the same raster dimensions.") if not (gt == gt2 == gt3 == gt4 == gt5 == gt6 == gt7): raise RuntimeError("Input subdatasets do not all have the same geotransform.") # Force output projection to EPSG:3031 out_proj = get_proj_wkt_epsg3031() # ------------------------------------------------------------------------- # Write the primary rasters # ------------------------------------------------------------------------- primary_outputs = { "bed": (bed, bed_nd), "surface": (surface, surface_nd), "thickness": (thickness, thickness_nd), "errbed": (errbed, errbed_nd), "firn": (firn, firn_nd), "dataid": (dataid, dataid_nd), } for key, (arr, src_nd) in primary_outputs.items(): out_name = OUTPUT_NAMES[key] out_path = OUTPUT_DIR / out_name gdal_dtype, out_nodata = choose_output_dtype(key) valid = build_valid_mask(arr, src_nd) # ADDITIONAL MASK: exclude ocean for selected variables only # dataid is NOT masked if key in ["surface", "thickness", "firn"]: if mask_nd is None: mask_valid = mask != 0 else: mask_valid = (mask != 0) & (mask != mask_nd) valid = valid & mask_valid out_dtype = np.int16 if key in ["errbed", "dataid"] else np.float32 out_arr = np.full(arr.shape, out_nodata, dtype=out_dtype) out_arr[valid] = sanitise_to_output_dtype(arr[valid], key) log(f"Writing {key} -> {out_name}") write_cog_from_array( array=out_arr, out_path=out_path, geotransform=gt, projection_wkt=out_proj, nodata=out_nodata, gdal_dtype=gdal_dtype, ) # ------------------------------------------------------------------------- # Write bed elevation masked by mask class # ------------------------------------------------------------------------- bed_valid = build_valid_mask(bed, bed_nd) if mask_nd is None: mask_valid = np.ones(mask.shape, dtype=bool) else: mask_valid = mask != mask_nd for class_value, class_name in MASK_CLASSES.items(): out_key = f"bed_{class_name}" out_name = OUTPUT_NAMES[out_key] out_path = MASKED_OUTPUT_DIR / out_name class_match = (mask == class_value) & mask_valid & bed_valid out_arr = np.full(bed.shape, NODATA_FLOAT, dtype=np.float32) out_arr[class_match] = bed[class_match].astype(np.float32, copy=False) log(f"Writing masked bed for {class_name} ({class_value}) -> {out_name}") write_cog_from_array( array=out_arr, out_path=out_path, geotransform=gt, projection_wkt=out_proj, nodata=NODATA_FLOAT, gdal_dtype=gdal.GDT_Float32, ) log("Done.") if __name__ == "__main__": try: main() except Exception as e: print(f"\nERROR: {e}", file=sys.stderr) sys.exit(1)