Merge pull request #50 from forrestfwilliams/develop

Release v0.4.0

Author: williamh890

CHANGELOG.md

@@ -6,6 +6,14 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/),
 and this project adheres to [PEP 440](https://www.python.org/dev/peps/pep-0440/)
 and uses [Semantic Versioning](https://semver.org/spec/v2.0.0.html).
 
+## [0.4.0]
+
+### Changed
+* For the S1RTC platform `chipdata` now processes and download all nessary RTC prior to getting individual chip data.
+
+### Fixed
+* Fix fmask path when chipping HLS data
+
 ## [0.3.0]
 
 ### Added

requirements-static.txt

@@ -1,2 +1,2 @@
-ruff==0.12.12
-mypy==1.17.1
+ruff==0.13.2
+mypy==1.18.2

src/satchip/chip_data.py

@@ -10,14 +10,11 @@
 import satchip
 from satchip import utils
 from satchip.chip_hls import get_hls_data
-from satchip.chip_sentinel1rtc import get_s1rtc_data
+from satchip.chip_sentinel1rtc import get_rtc_paths_for_chips, get_s1rtc_chip_data
 from satchip.chip_sentinel2 import get_s2l2a_data
 from satchip.terra_mind_grid import TerraMindGrid
 
 
-GET_DATA_FNS = {'S2L2A': get_s2l2a_data, 'S1RTC': get_s1rtc_data, 'HLS': get_hls_data}
-
-
 def fill_missing_times(data_chip: xr.DataArray, times: np.ndarray) -> xr.DataArray:
     missing_times = np.setdiff1d(times, data_chip.time.data)
     missing_shape = (len(missing_times), len(data_chip.band), data_chip.y.size, data_chip.x.size)
@@ -42,28 +39,35 @@ def chip_data(
     strategy: str,
     max_cloud_pct: int,
     output_dir: Path,
-    scratch_dir: Path | None = None,
+    scratch_dir: Path,
 ) -> xr.Dataset:
-    get_data_fn = GET_DATA_FNS[platform]
     labels = utils.load_chip(label_path)
     date = labels.time.data[0].astype('M8[ms]').astype(datetime)
     bounds = labels.attrs['bounds']
+
     grid = TerraMindGrid([bounds[1] - 1, bounds[3] + 1], [bounds[0] - 1, bounds[2] + 1])  # type: ignore
     terra_mind_chips = [c for c in grid.terra_mind_chips if c.name in list(labels.sample.data)]
 
-    opts = {'strategy': strategy, 'date_start': date_start, 'date_end': date_end}
+    opts: utils.ChipDataOpts = {'strategy': strategy, 'date_start': date_start, 'date_end': date_end}
     if platform in ['S2L2A', 'HLS']:
         opts['max_cloud_pct'] = max_cloud_pct
 
+    if platform == 'S1RTC':
+        rtc_paths_for_chips = get_rtc_paths_for_chips(terra_mind_chips, bounds, scratch_dir, opts)
+
     data_chips = []
-    if scratch_dir is not None:
-        for chip in tqdm(terra_mind_chips):
-            data_chips.append(get_data_fn(chip, scratch_dir, opts=opts))
-    else:
-        with TemporaryDirectory() as tmp_dir:
-            scratch_dir = Path(tmp_dir)
-            for chip in tqdm(terra_mind_chips):
-                data_chips.append(get_data_fn(chip, scratch_dir, opts=opts))
+    for chip in tqdm(terra_mind_chips):
+        if platform == 'S1RTC':
+            rtc_paths = rtc_paths_for_chips[chip.name]
+            chip_data = get_s1rtc_chip_data(chip, rtc_paths, scratch_dir, opts=opts)
+        elif platform == 'S2L2A':
+            chip_data = get_s2l2a_data(chip, scratch_dir, opts=opts)
+        elif platform == 'HLS':
+            chip_data = get_hls_data(chip, scratch_dir, opts=opts)
+        else:
+            raise Exception(f'Unknown platform {platform}')
+
+        data_chips.append(chip_data)
 
     times = np.unique(np.concatenate([dc.time.data for dc in data_chips]))
     for i, data_chip in enumerate(data_chips):
@@ -99,17 +103,24 @@ def main() -> None:
     assert 0 <= args.maxcloudpct <= 100, 'maxcloudpct must be between 0 and 100'
     date_start, date_end = [datetime.strptime(d, '%Y%m%d') for d in args.daterange.split('-')]
     assert date_start < date_end, 'start date must be before end date'
-    chip_data(
+
+    params = (
         args.labelpath,
         args.platform,
         date_start,
         date_end,
         args.strategy,
         args.maxcloudpct,
         args.outdir,
-        args.scratchdir,
     )
 
+    if args.scratchdir is not None:
+        chip_data(*params, args.scratchdir)
+    else:
+        with TemporaryDirectory() as tmp_dir:
+            scratch_dir = Path(tmp_dir)
+            chip_data(*params, scratch_dir)
+
 
 if __name__ == '__main__':
     main()

src/satchip/chip_hls.py

@@ -8,6 +8,7 @@
 import xarray as xr
 from earthaccess.results import DataGranule
 
+from satchip import utils
 from satchip.chip_xr_base import create_template_da
 from satchip.terra_mind_grid import TerraMindChip
 
@@ -76,7 +77,7 @@ def get_scenes(
         n_products = len(list(scratch_dir.glob(f'{product_id}*')))
         if n_products < 15:
             earthaccess.download([item], scratch_dir, pqdm_kwargs={'disable': True})
-        fmask_path = scratch_dir / f'{product_id}.v2.0.FMask.tif'
+        fmask_path = scratch_dir / f'{product_id}.v2.0.Fmask.tif'
         assert fmask_path.exists(), f'File not found: {fmask_path}'
         qual_da = rioxarray.open_rasterio(fmask_path).rio.clip_box(*roi.bounds, crs='EPSG:4326')  # type: ignore
         bit_masks = np.unpackbits(qual_da.data[0][..., np.newaxis], axis=-1)
@@ -93,7 +94,7 @@ def get_scenes(
     return valid_scenes
 
 
-def get_hls_data(chip: TerraMindChip, scratch_dir: Path, opts: dict) -> xr.DataArray:
+def get_hls_data(chip: TerraMindChip, scratch_dir: Path, opts: utils.ChipDataOpts) -> xr.DataArray:
     """Returns XArray DataArray of a Harmonized Landsat Sentinel-2 image for the given bounds and
     closest collection after date.
     """

src/satchip/chip_sentinel1rtc.py

@@ -1,79 +1,162 @@
 from datetime import datetime, timedelta
 from pathlib import Path
 
-import asf_search as search
+import asf_search as asf
+import hyp3_sdk
 import numpy as np
 import rioxarray
 import shapely
 import xarray as xr
-from asf_search import ASFSearchResults, S1Product, constants
-from hyp3_sdk import Batch, HyP3, Job
-from hyp3_sdk.util import extract_zipped_product
 
+from satchip import utils
 from satchip.chip_xr_base import create_template_da
 from satchip.terra_mind_grid import TerraMindChip
 
 
-def get_pct_intersect(product: S1Product, roi: shapely.geometry.Polygon) -> int:
+def get_rtc_paths_for_chips(
+    terra_mind_chips: list[TerraMindChip], bounds: list[float], scratch_dir: Path, opts: utils.ChipDataOpts
+) -> dict[str, list[Path]]:
+    _check_bounds_size(bounds)
+    granules = _get_granules(bounds, opts['date_start'], opts['date_end'])
+    slcs_for_chips = _get_slcs_for_each_chip(terra_mind_chips, granules, opts['strategy'])
+    assert len(slcs_for_chips) == len(terra_mind_chips)
+
+    rtc_paths_for_chips = _get_rtcs_for(slcs_for_chips, scratch_dir)
+    return rtc_paths_for_chips
+
+
+def _check_bounds_size(bounds: list[float]) -> None:
+    min_lon, min_lat, max_lon, max_lat = bounds
+    MAX_BOUND_AREA_DEGREES = 3
+    bounds_area_degrees = (max_lon - min_lon) * (max_lat - min_lat)
+
+    err_message = f'Bounds area is to large ({bounds_area_degrees}). Must be less than {MAX_BOUND_AREA_DEGREES} degrees'
+    assert bounds_area_degrees < MAX_BOUND_AREA_DEGREES, err_message
+
+
+def _get_granules(bounds: list[float], date_start: datetime, date_end: datetime) -> list[asf.S1Product]:
+    date_start = date_start
+    date_end = date_end + timedelta(days=1)  # inclusive end
+    roi = shapely.box(*bounds)
+    search_results = asf.geo_search(
+        intersectsWith=roi.wkt,
+        start=date_start,
+        end=date_end,
+        beamMode=asf.constants.BEAMMODE.IW,
+        polarization=asf.constants.POLARIZATION.VV_VH,
+        platform=asf.constants.PLATFORM.SENTINEL1,
+        processingLevel=asf.constants.PRODUCT_TYPE.SLC,
+    )
+
+    return list(search_results)
+
+
+def _get_slcs_for_each_chip(
+    chips: list[TerraMindChip], granules: list[asf.S1Product], strategy: str, intersection_pct: int = 95
+) -> dict[str, list[asf.S1Product]]:
+    slcs_for_chips: dict[str, list[asf.S1Product]] = {}
+
+    for chip in chips:
+        chip_roi = shapely.box(*chip.bounds)
+        intersecting = [granule for granule in granules if _get_pct_intersect(granule, chip_roi) > intersection_pct]
+        intersecting = sorted(intersecting, key=lambda g: (-_get_pct_intersect(g, chip_roi), g.properties['startTime']))
+
+        if len(intersecting) < 1:
+            raise ValueError(f'No products found for chip {chip.name} in given date range')
+
+        if strategy == 'BEST':
+            slcs_for_chips[chip.name] = intersecting[:1]
+        else:
+            slcs_for_chips[chip.name] = intersecting
+
+    return slcs_for_chips
+
+
+def _get_pct_intersect(product: asf.S1Product, roi: shapely.geometry.Polygon) -> int:
     footprint = shapely.geometry.shape(product.geometry)
     intersection = int(np.round(100 * roi.intersection(footprint).area / roi.area))
     return intersection
 
 
-def download_hyp3_rtc(job: Job, scratch_dir: Path) -> tuple[Path, Path]:
+def _get_rtcs_for(slcs_for_chips: dict[str, list[asf.S1Product]], scratch_dir: Path) -> dict[str, list[Path]]:
+    flat_slcs = sum(slcs_for_chips.values(), [])
+    slc_names = set(granule.properties['sceneName'] for granule in flat_slcs)
+
+    finished_rtc_jobs = _process_rtcs(slc_names)
+
+    paths_for_slc_name: dict[str, Path] = {}
+    for job in finished_rtc_jobs:
+        rtc_path = _download_hyp3_rtc(job, scratch_dir)
+        slc_name = job.job_parameters['granules'][0]
+
+        paths_for_slc_name[slc_name] = rtc_path
+
+    rtc_paths_for_chips: dict[str, list[Path]] = {}
+    for chip_name, chip_slcs in slcs_for_chips.items():
+        rtc_paths = [paths_for_slc_name[name.properties['sceneName']] for name in chip_slcs]
+        rtc_paths_for_chips[chip_name] = rtc_paths
+
+    return rtc_paths_for_chips
+
+
+def _process_rtcs(slc_names: set[str]) -> hyp3_sdk.Batch:
+    hyp3 = hyp3_sdk.HyP3()
+    jobs_by_scene_name = _get_rtc_jobs_by_scene_name(hyp3)
+
+    hyp3_jobs = []
+    for slc_name in slc_names:
+        if slc_name in jobs_by_scene_name:
+            job = jobs_by_scene_name[slc_name]
+            hyp3_jobs.append(job)
+        else:
+            new_batch = hyp3.submit_rtc_job(slc_name, radiometry='gamma0', resolution=20)
+            hyp3_jobs.append(list(new_batch)[0])
+
+    batch = hyp3_sdk.Batch(hyp3_jobs)
+    batch = hyp3.watch(batch)
+    assert all([j.succeeded() for j in batch]), 'One or more HyP3 jobs failed'
+
+    return batch
+
+
+def _get_rtc_jobs_by_scene_name(hyp3: hyp3_sdk.HyP3) -> dict[str, hyp3_sdk.Job]:
+    jobs_by_scene_name = {}
+
+    for job in hyp3.find_jobs(job_type='RTC_GAMMA'):
+        if not _is_valid_rtc_job(job):
+            continue
+
+        name = job.job_parameters['granules'][0]
+        jobs_by_scene_name[name] = job
+
+    return jobs_by_scene_name
+
+
+def _is_valid_rtc_job(job: hyp3_sdk.Job) -> bool:
+    return (
+        not job.failed()
+        and not job.expired()
+        and job.job_parameters['radiometry'] == 'gamma0'
+        and job.job_parameters['resolution'] == 20
+    )
+
+
+def _download_hyp3_rtc(job: hyp3_sdk.Job, scratch_dir: Path) -> tuple[Path, Path]:
     output_path = scratch_dir / job.to_dict()['files'][0]['filename']
     output_dir = output_path.with_suffix('')
     output_zip = output_path.with_suffix('.zip')
     if not output_dir.exists():
         job.download_files(location=scratch_dir)
-        extract_zipped_product(output_zip)
+        hyp3_sdk.util.extract_zipped_product(output_zip)
     vv_path = list(output_dir.glob('*_VV.tif'))[0]
     vh_path = list(output_dir.glob('*_VH.tif'))[0]
     return vv_path, vh_path
 
 
-def get_hyp3_rtcs(items: ASFSearchResults, roi: shapely.geometry.Polygon, strategy: str, scratch_dir: Path) -> list:
-    valid_items = [item for item in items if get_pct_intersect(item, roi) > 95]
-    valid_items = sorted(items, key=lambda x: (-get_pct_intersect(x, roi), x.properties['startTime']))
-    if strategy == 'BEST':
-        valid_items = valid_items[:1]
-    hyp3 = HyP3()
-    old_jobs = [j for j in hyp3.find_jobs(job_type='RTC_GAMMA') if not j.failed() and not j.expired()]
-    old_jobs = [j for j in old_jobs if j.job_parameters['radiometry'] == 'gamma0']  # type: ignore
-    old_jobs = [j for j in old_jobs if j.job_parameters['resolution'] == 20]  # type: ignore
-    jobs = []
-    for item in valid_items:
-        scene_name = item.properties['sceneName']
-        matching_jobs = [j for j in old_jobs if j.job_parameters['granules'] == [scene_name]]  # type: ignore
-        if len(matching_jobs) == 0:
-            new_batch = hyp3.submit_rtc_job(scene_name, radiometry='gamma0', resolution=20)
-            jobs.append(list(new_batch)[0])
-        else:
-            jobs.append(matching_jobs[0])
-    jobs = Batch(jobs)
-    hyp3.watch(jobs)
-    assert all([j.succeeded() for j in jobs]), 'One or more HyP3 jobs failed'
-    paths = [download_hyp3_rtc(job, scratch_dir) for job in jobs]
-    return paths
-
-
-def get_s1rtc_data(chip: TerraMindChip, scratch_dir: Path, opts: dict) -> xr.DataArray:
-    date_start = opts['date_start']
-    date_end = opts['date_end'] + timedelta(days=1)  # inclusive end
+def get_s1rtc_chip_data(
+    chip: TerraMindChip, image_sets: list[Path], scratch_dir: Path, opts: utils.ChipDataOpts
+) -> xr.DataArray:
     roi = shapely.box(*chip.bounds)
-    search_results = search.geo_search(
-        intersectsWith=roi.wkt,
-        start=date_start,
-        end=date_end,
-        beamMode=constants.BEAMMODE.IW,
-        polarization=constants.POLARIZATION.VV_VH,
-        platform=constants.PLATFORM.SENTINEL1,
-        processingLevel=constants.PRODUCT_TYPE.SLC,
-    )
-    if len(search_results) == 0:
-        raise ValueError(f'No products found for chip {chip.name} in date range {date_start} to {date_end}')
-    strategy = opts.get('strategy', 'BEST').upper()
-    image_sets = get_hyp3_rtcs(search_results, roi, strategy, scratch_dir)
     das = []
     template = create_template_da(chip)
     for image_set in image_sets:

src/satchip/chip_sentinel2.py

@@ -11,6 +11,7 @@
 from pystac.item import Item
 from pystac_client import Client
 
+from satchip import utils
 from satchip.chip_xr_base import create_template_da
 from satchip.terra_mind_grid import TerraMindChip
 
@@ -127,7 +128,7 @@ def get_scenes(
     return valid_scenes
 
 
-def get_s2l2a_data(chip: TerraMindChip, scratch_dir: Path, opts: dict) -> xr.DataArray:
+def get_s2l2a_data(chip: TerraMindChip, scratch_dir: Path, opts: utils.ChipDataOpts) -> xr.DataArray:
     """Get XArray DataArray of Sentinel-2 L2A image for the given bounds and best collection parameters.
 
     Args:

src/satchip/utils.py

@@ -1,10 +1,22 @@
+import datetime
 from pathlib import Path
+from typing import TypedDict
 
 import xarray as xr
 import zarr
 from pyproj import CRS, Transformer
 
 
+class ChipDataRequiredOpts(TypedDict):
+    strategy: str
+    date_start: datetime.datetime
+    date_end: datetime.datetime
+
+
+class ChipDataOpts(ChipDataRequiredOpts, total=False):
+    max_cloud_pct: int
+
+
 def get_epsg4326_point(x: float, y: float, in_epsg: int) -> tuple[float, float]:
     if in_epsg == 4326:
         return x, y