Concurrent downloads

CHANGELOG.md

@@ -6,6 +6,10 @@ 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.1.2]
+### Added
+* Concurrent download functionality for burst extractor calls.
+
 ## [0.1.1]
 ### Fixed
 * Typo in the `release-checklist-comment` workflow.

src/hyp3_isce2/burst.py

@@ -1,6 +1,8 @@
 import copy
 import re
+import shutil
 import time
+from concurrent.futures import ThreadPoolExecutor
 from dataclasses import dataclass
 from pathlib import Path
 from typing import Iterator, List, Tuple, Union
@@ -196,7 +198,7 @@ def download_metadata(
     return str(out_file)
 
 
-def download_burst(asf_session: requests.Session, burst_params: BurstParams, out_file: Union[Path, str]) -> str:
+def download_burst(asf_session: requests.Session, burst_params: BurstParams, out_file: Union[Path, str] = None) -> Path:
     """Download a burst geotiff.
 
     Args:
@@ -209,13 +211,18 @@ def download_burst(asf_session: requests.Session, burst_params: BurstParams, out
     """
     content = download_from_extractor(asf_session, burst_params, 'geotiff')
 
+    if not out_file:
+        out_file = (
+            f'{burst_params.granule}_{burst_params.swath}_{burst_params.polarization}_{burst_params.burst_number}.tiff'
+        ).lower()
+
     with open(out_file, 'wb') as f:
         f.write(content)
 
-    return str(out_file)
+    return Path(out_file)
 
 
-def spoof_safe(asf_session: requests.Session, burst: BurstMetadata, base_path: Path = Path('.')) -> Path:
+def spoof_safe(burst: BurstMetadata, burst_tiff_path: Path, base_path: Path = Path('.')) -> Path:
     """Spoof a Sentinel-1 SAFE file for a burst.
 
     The created SAFE file will be saved to the base_path directory. The SAFE will have the following structure:
@@ -230,8 +237,8 @@ def spoof_safe(asf_session: requests.Session, burst: BurstMetadata, base_path: P
             └── noise.xml
 
     Args:
-        asf_session: A requests session with an ASF URS cookie.
         burst: The burst metadata.
+        burst_tiff_path: The path to the burst geotiff.
         base_path: The path to save the SAFE file to.
 
     Returns:
@@ -252,8 +259,7 @@ def spoof_safe(asf_session: requests.Session, burst: BurstMetadata, base_path: P
     etree.ElementTree(burst.noise).write(calibration_path / burst.noise_name, **et_args)
     etree.ElementTree(burst.manifest).write(safe_path / 'manifest.safe', **et_args)
 
-    burst_params = BurstParams(burst.safe_name, burst.swath, burst.polarization, burst.burst_number)
-    download_burst(asf_session, burst_params, measurement_path / burst.measurement_name)
+    shutil.move(str(burst_tiff_path), str(measurement_path / burst.measurement_name))
 
     return safe_path
 
@@ -303,30 +309,34 @@ def get_asf_session() -> requests.Session:
 
 
 def download_bursts(param_list: Iterator[BurstParams]) -> List[BurstMetadata]:
-    """Download bursts and create SAFE files.
+    """Download bursts in parallel and creates SAFE files.
 
     For each burst:
         1. Download metadata
-        2. Create BurstMetadata object
-        3. Create directory structure
-        4. Write metadata
-        5. Download and write geotiff
+        2. Download geotiff
+        3. Create BurstMetadata object
+        4. Create directory structure
+        5. Write metadata
+        6. Move geotiff to correct directory
 
     Args:
         param_list: An iterator of burst search parameters.
 
     Returns:
         A list of BurstMetadata objects.
     """
-    bursts = []
     with get_asf_session() as asf_session:
-        for i, params in enumerate(param_list):
-            print(f'Creating SAFE {i+1}...')
-            metadata_xml = download_metadata(asf_session, params)
-            burst = BurstMetadata(metadata_xml, params)
-            spoof_safe(asf_session, burst)
-            bursts.append(burst)
+        with ThreadPoolExecutor(max_workers=10) as executor:
+            xml_futures = [executor.submit(download_metadata, asf_session, params) for params in param_list]
+            tiff_futures = [executor.submit(download_burst, asf_session, params) for params in param_list]
+            metadata_xmls = [future.result() for future in xml_futures]
+            burst_paths = [future.result() for future in tiff_futures]
 
+    bursts = []
+    for params, metadata_xml, burst_path in zip(param_list, metadata_xmls, burst_paths):
+        burst = BurstMetadata(metadata_xml, params)
+        spoof_safe(burst, burst_path)
+        bursts.append(burst)
     print('SAFEs created!')
 
     return bursts

src/hyp3_isce2/etc/download_example.py

@@ -5,20 +5,23 @@
 
 Example 1: Downloading metadata for a pair of ascending/descending bursts
             (these are the metadata files used in the test suite.)
+
+Example 2: Downloading a pair of ascending/descending bursts and spoofing a SAFE
 """
 
-from hyp3_isce2.burst import (
-    BurstParams,
-    download_metadata,
-    get_asf_session
-)
+from hyp3_isce2.burst import BurstParams, download_bursts, download_metadata, get_asf_session
 
-# Example 1
 ref_desc = BurstParams('S1A_IW_SLC__1SDV_20200604T022251_20200604T022318_032861_03CE65_7C85', 'IW2', 'VV', 3)
 sec_desc = BurstParams('S1A_IW_SLC__1SDV_20200616T022252_20200616T022319_033036_03D3A3_5D11', 'IW2', 'VV', 3)
 ref_asc = BurstParams('S1A_IW_SLC__1SDV_20211229T231926_20211229T231953_041230_04E66A_3DBE', 'IW1', 'VV', 4)
 sec_asc = BurstParams('S1A_IW_SLC__1SDV_20220110T231926_20220110T231953_041405_04EC57_103E', 'IW1', 'VV', 4)
 
-session = get_asf_session()
-download_metadata(session, ref_asc, 'data/reference_ascending.xml')
-download_metadata(session, sec_desc, 'data/secondary_ascending.xml')
+# Example 1
+# Download metadata files
+with get_asf_session() as session:
+    download_metadata(session, ref_asc, 'data/reference_ascending.xml')
+    download_metadata(session, sec_desc, 'data/secondary_ascending.xml')
+
+# Example 2
+# Download with SAFE spoofing
+download_bursts([ref_desc, sec_desc])

tests/test_burst.py

@@ -2,7 +2,6 @@
 
 import numpy as np
 import pytest
-import requests
 from lxml import etree
 from shapely import geometry
 
@@ -54,12 +53,15 @@ def test_create_geometry():
         '*SAFE/annotation/*xml',
         '*SAFE/annotation/calibration/calibration*xml',
         '*SAFE/annotation/calibration/noise*xml',
+        '*SAFE/measurement/*tiff',
     ),
 )
 def test_spoof_safe(tmp_path, mocker, pattern):
+    mock_tiff = tmp_path / 'test.tiff'
+    mock_tiff.touch()
+
     ref_burst = burst.BurstMetadata(load_metadata('reference_descending.xml'), REF_DESC)
-    mocker.patch('hyp3_isce2.burst.download_burst', return_value='')
-    burst.spoof_safe(requests.Session(), ref_burst, tmp_path)
+    burst.spoof_safe(ref_burst, mock_tiff, tmp_path)
     assert len(list(tmp_path.glob(pattern))) == 1