From 869ee526975f7d7fc781608ae7967142f1a5096e Mon Sep 17 00:00:00 2001 From: Sahil Jhawar Date: Thu, 9 Jul 2026 11:53:22 +0200 Subject: [PATCH] feat: use new ace url --- swvo/io/solar_wind/ace.py | 163 ++++++++++++++++++-------------- tests/io/solar_wind/test_ace.py | 125 +++++++++++++++++------- 2 files changed, 182 insertions(+), 106 deletions(-) diff --git a/swvo/io/solar_wind/ace.py b/swvo/io/solar_wind/ace.py index 94bc2f4..6df6207 100644 --- a/swvo/io/solar_wind/ace.py +++ b/swvo/io/solar_wind/ace.py @@ -48,28 +48,30 @@ class SWACE(BaseIO): ENV_VAR_NAME = "RT_SW_ACE_STREAM_DIR" - URL = "https://services.swpc.noaa.gov/text/" - NAME_MAG = "ace-magnetometer.txt" - NAME_SWEPAM = "ace-swepam.txt" + URL = "https://sohoftp.nascom.nasa.gov/sdb/goes/ace/daily/" + NAME_MAG = "{date:%Y%m%d}_ace_mag_1m.txt" + NAME_SWEPAM = "{date:%Y%m%d}_ace_swepam_1m.txt" SWEPAM_FIELDS = ["speed", "proton_density", "temperature"] MAG_FIELDS = ["bx_gsm", "by_gsm", "bz_gsm", "bavg"] LABEL = "ace" - def download_and_process(self, request_time: datetime) -> None: + def download_and_process(self, start_time: datetime, end_time: datetime) -> None: """ Download and process ACE data, splitting data across midnight into appropriate day files. Parameters ---------- - request_time : datetime - The time for which the data is requested. Must be in the past and within the last two hours. + start_time : datetime + Start time of the data to download. Must be timezone-aware. + end_time : datetime + End time of the data to download. Must be timezone-aware. Raises ------ AssertionError - If the request_time is in the future. + If the requested interval is invalid or extends into a future UTC date. FileNotFoundError If the downloaded files are empty. @@ -78,56 +80,28 @@ def download_and_process(self, request_time: datetime) -> None: None """ - current_time = datetime.now(timezone.utc) - - assert request_time < current_time, ( - f"Request time: {request_time} cannot be in the future. Current time: {current_time}!" - ) - # assert request_time < (datetime.now(timezone.utc).replace(second=0, microsecond=0) - timedelta(minutes = 121)), "Request time cannot be in the past!" + start_time = enforce_utc_timezone(start_time) + end_time = enforce_utc_timezone(end_time) - if current_time - request_time > timedelta(hours=2): - logger.debug("We can only download and process ACE RT data for the last two hours!") - return + assert start_time < end_time, "Start time must be before end time!" temporary_dir = Path("./temp_sw_ace_wget") temporary_dir.mkdir(exist_ok=True, parents=True) - self._download(temporary_dir, self.NAME_MAG) - self._download(temporary_dir, self.NAME_SWEPAM) - logger.debug("Processing file ...") - processed_df = self._process_single_file(temporary_dir) - - unique_dates = np.unique(processed_df.index.date) # ty: ignore[unresolved-attribute] - - for date in unique_dates: - file_path = self.data_dir / date.strftime("%Y/%m") / f"ACE_SW_NOWCAST_{date.strftime('%Y%m%d')}.csv" - tmp_path = file_path.with_suffix(file_path.suffix + ".tmp") - - try: - day_start = enforce_utc_timezone(datetime.combine(date, datetime.min.time())) - day_end = enforce_utc_timezone(datetime.combine(date, datetime.max.time())) - - day_data = processed_df[(processed_df.index >= day_start) & (processed_df.index <= day_end)] - - if file_path.exists(): - logger.debug(f"Found previous file for {date}. Loading and combining ...") - previous_df = self._read_single_file(file_path) + try: + for date in pd.date_range(start=start_time.date(), end=end_time.date(), freq="D"): + request_date = date.to_pydatetime().replace(tzinfo=timezone.utc) + mag_file_name = self._dated_file_name(self.NAME_MAG, request_date) + swepam_file_name = self._dated_file_name(self.NAME_SWEPAM, request_date) - previous_df.drop("file_name", axis=1, inplace=True) - day_data = day_data.combine_first(previous_df) + self._download(temporary_dir, mag_file_name) + self._download(temporary_dir, swepam_file_name) - logger.debug(f"Saving processed file for {date}") - file_path.parent.mkdir(parents=True, exist_ok=True) - day_data.to_csv(tmp_path, index=True, header=True) - tmp_path.replace(file_path) - - except Exception as e: - logger.error(f"Failed to process file for {date}: {e}") - if tmp_path.exists(): - tmp_path.unlink() - continue - - rmtree(temporary_dir, ignore_errors=True) + logger.debug(f"Processing ACE source files for {request_date.date()} ...") + processed_df = self._process_single_file(temporary_dir, mag_file_name, swepam_file_name) + self._save_processed_data(processed_df) + finally: + rmtree(temporary_dir, ignore_errors=True) def _download(self, temporary_dir: Path, file_name: str) -> None: """Download a file from ACE server. @@ -156,6 +130,10 @@ def _download(self, temporary_dir: Path, file_name: str) -> None: if (temporary_dir / file_name).stat().st_size == 0: raise FileNotFoundError(f"Error while downloading file: {self.URL + file_name}!") + def _dated_file_name(self, template: str, request_time: datetime) -> str: + request_time = enforce_utc_timezone(request_time) + return template.format(date=request_time) + def read( self, start_time: datetime, @@ -184,8 +162,8 @@ def read( Raises ------ - AssertionError - Raises `AssertionError` if the start time is before the end time. + ValueError + Raises `ValueError` if the start time is after the end time. """ if start_time > end_time: @@ -197,9 +175,8 @@ def read( end_time = enforce_utc_timezone(end_time) if propagation: - logger.info("Shiting start day by -1 day to account for propagation") + logger.info("Shifting start day by -1 day to account for propagation") start_time = start_time - timedelta(days=1) - file_paths, _ = self._get_processed_file_list(start_time, end_time) t = pd.date_range( @@ -222,16 +199,13 @@ def read( }, ) - for file_path in file_paths: - if not file_path.exists() and download: - file_date = enforce_utc_timezone(datetime.strptime(file_path.stem.split("_")[-1], "%Y%m%d")) - hour_now = datetime.now(timezone.utc).hour - file_date = file_date.replace(hour=hour_now, minute=0, second=0, microsecond=0) - try: - self.download_and_process(file_date) - except AssertionError as e: - logger.error(f"`download_and_process` failed because: {e}") + if download and any(not file_path.exists() for file_path in file_paths): + try: + self.download_and_process(start_time, end_time) + except AssertionError as e: + logger.error(f"`download_and_process` failed because: {e}") + for file_path in file_paths: if not file_path.exists(): warnings.warn(f"File {file_path} not found") continue @@ -327,7 +301,43 @@ def _read_single_file(self, file_path) -> pd.DataFrame: return df - def _process_single_file(self, temporary_dir: Path) -> pd.DataFrame: + def _save_processed_data(self, processed_df: pd.DataFrame) -> None: + unique_dates = np.unique(processed_df.index.date) # ty: ignore[unresolved-attribute] + + for date in unique_dates: + file_path = self.data_dir / date.strftime("%Y/%m") / f"ACE_SW_NOWCAST_{date.strftime('%Y%m%d')}.csv" + tmp_path = file_path.with_suffix(file_path.suffix + ".tmp") + + try: + day_start = enforce_utc_timezone(datetime.combine(date, datetime.min.time())) + day_end = enforce_utc_timezone(datetime.combine(date, datetime.max.time())) + + day_data = processed_df[(processed_df.index >= day_start) & (processed_df.index <= day_end)] + + if file_path.exists(): + logger.debug(f"Found previous file for {date}. Loading and combining ...") + previous_df = self._read_single_file(file_path) + + previous_df.drop("file_name", axis=1, inplace=True) + day_data = day_data.combine_first(previous_df) + + logger.debug(f"Saving processed file for {date}") + file_path.parent.mkdir(parents=True, exist_ok=True) + day_data.to_csv(tmp_path, index=True, header=True) + tmp_path.replace(file_path) + + except Exception as e: + logger.error(f"Failed to process file for {date}: {e}") + if tmp_path.exists(): + tmp_path.unlink() + continue + + def _process_single_file( + self, + temporary_dir: Path, + mag_file_name: str | None = None, + swepam_file_name: str | None = None, + ) -> pd.DataFrame: """Process mag and swepam ACE file to a DataFrame. Returns @@ -335,14 +345,14 @@ def _process_single_file(self, temporary_dir: Path) -> pd.DataFrame: pd.DataFrame ACE data. """ - data_mag = self._process_mag_file(temporary_dir) - data_swepam = self._process_swepam_file(temporary_dir) + data_mag = self._process_mag_file(temporary_dir, mag_file_name) + data_swepam = self._process_swepam_file(temporary_dir, swepam_file_name) data = pd.concat([data_swepam, data_mag], axis=1) return data - def _process_mag_file(self, temporary_dir: Path) -> pd.DataFrame: + def _process_mag_file(self, temporary_dir: Path, file_name: str | None = None) -> pd.DataFrame: """ Reads magnetic instrument last available real time ACE data. @@ -369,8 +379,9 @@ def _process_mag_file(self, temporary_dir: Path) -> pd.DataFrame: "lon", ] + data_file = self._resolve_data_file(temporary_dir, file_name, "*_ace_mag_1m.txt") data_mag = pd.read_csv( - temporary_dir / self.NAME_MAG, + data_file, comment="#", skiprows=2, sep=r"\s+", @@ -402,7 +413,7 @@ def _process_mag_file(self, temporary_dir: Path) -> pd.DataFrame: return data_mag - def _process_swepam_file(self, temporary_dir: Path) -> pd.DataFrame: + def _process_swepam_file(self, temporary_dir: Path, file_name: str | None = None) -> pd.DataFrame: """ This method reads faraday cup SWEPAM instrument daily file from ACE original data. @@ -427,8 +438,9 @@ def _process_swepam_file(self, temporary_dir: Path) -> pd.DataFrame: "temperature", ] + data_file = self._resolve_data_file(temporary_dir, file_name, "*_ace_swepam_1m.txt") data_sw = pd.read_csv( - temporary_dir / self.NAME_SWEPAM, + data_file, comment="#", skiprows=2, sep=r"\s+", @@ -454,6 +466,17 @@ def _process_swepam_file(self, temporary_dir: Path) -> pd.DataFrame: return data_sw + def _resolve_data_file(self, temporary_dir: Path, file_name: str | None, pattern: str) -> Path: + if file_name is not None: + return temporary_dir / file_name + + matches = sorted(temporary_dir.glob(pattern)) + if len(matches) == 1: + return matches[0] + if len(matches) > 1: + raise ValueError(f"Found multiple ACE source files matching {pattern} in {temporary_dir}") + raise FileNotFoundError(f"No ACE source file matching {pattern} found in {temporary_dir}") + def _to_date(self, x) -> datetime: """ Converts into a proper datetime format. diff --git a/tests/io/solar_wind/test_ace.py b/tests/io/solar_wind/test_ace.py index aa8a5a3..9be7702 100644 --- a/tests/io/solar_wind/test_ace.py +++ b/tests/io/solar_wind/test_ace.py @@ -38,35 +38,20 @@ def swace_instance(self): @pytest.fixture def sample_mag_data(self): """Sample magnetometer data""" - return """ + return """:Data_list: 20240315_ace_mag_1m.txt +:Created: 2024 Mar 15 0102 UT - 2024 03 15 0100 -1.0000000 3 2.31 1.42 -0.89 2.84 31.42 31.63 - 2024 03 15 0101 -1.0000000 3 2.28 1.39 -0.87 2.80 31.38 31.59""" + 2024 03 15 0100 60384 3600 0 1.42 -0.89 2.84 3.31 31.42 31.63 + 2024 03 15 0101 60384 3660 0 1.39 -0.87 2.80 3.26 31.38 31.59""" @pytest.fixture def sample_swepam_data(self): """Sample SWEPAM data""" - return """ + return """:Data_list: 20240315_ace_swepam_1m.txt +:Created: 2024 Mar 15 0102 UT - 2024 03 15 0100 -1.0000000 3 4.2 380 145000 - 2024 03 15 0101 -1.0000000 3 4.1 375 144000""" - - @pytest.fixture - def mock_download_response(self, sample_mag_data, sample_swepam_data): - def mock_download(url, output): - output_path = Path(output) - if SWACE.NAME_MAG in url: - file_path = output_path / SWACE.NAME_MAG - content = sample_mag_data - else: - file_path = output_path / SWACE.NAME_SWEPAM - content = sample_swepam_data - - with open(file_path, "w") as f: - f.write(content) - return file_path - - return mock_download + 2024 03 15 0100 60384 3600 0 4.2 380 145000 + 2024 03 15 0101 60384 3660 0 4.1 375 144000""" def test_initialization_with_env_var(self): with patch.dict("os.environ", {SWACE.ENV_VAR_NAME: str(DATA_DIR)}): @@ -96,21 +81,84 @@ def test_get_processed_file_list(self, swace_instance): assert len(time_intervals) == 366 assert all(isinstance(interval, tuple) for interval in time_intervals) - def test_download_and_process(self, swace_instance): - current_time = datetime.now(timezone.utc) - swace_instance.download_and_process(current_time) + def test_dated_source_file_names(self, swace_instance): + request_time = datetime(2024, 3, 15, 1, 0, tzinfo=timezone.utc) - expected_file = ( - DATA_DIR / current_time.strftime("%Y/%m") / f"ACE_SW_NOWCAST_{current_time.strftime('%Y%m%d')}.csv" - ) + assert swace_instance._dated_file_name(swace_instance.NAME_MAG, request_time) == "20240315_ace_mag_1m.txt" + assert swace_instance._dated_file_name(swace_instance.NAME_SWEPAM, request_time) == "20240315_ace_swepam_1m.txt" + + def test_download_and_process(self, swace_instance, sample_mag_data, sample_swepam_data): + start_time = datetime(2024, 3, 15, 1, 0, tzinfo=timezone.utc) + end_time = datetime(2024, 3, 15, 1, 1, tzinfo=timezone.utc) + mag_file_name = swace_instance._dated_file_name(swace_instance.NAME_MAG, start_time) + swepam_file_name = swace_instance._dated_file_name(swace_instance.NAME_SWEPAM, start_time) + + def mock_get(url, **kwargs): + mock_response = Mock() + if url.endswith(mag_file_name): + mock_response.content = sample_mag_data.encode() + else: + mock_response.content = sample_swepam_data.encode() + mock_response.raise_for_status = Mock() + return mock_response + + with patch("requests.get", side_effect=mock_get) as mock_requests_get: + swace_instance.download_and_process(start_time, end_time) + + expected_file = DATA_DIR / "2024/03" / "ACE_SW_NOWCAST_20240315.csv" assert expected_file.exists() data = pd.read_csv(expected_file) - assert len(data) > 0 + assert len(data) == 2 assert all(field in data.columns for field in SWACE.MAG_FIELDS + SWACE.SWEPAM_FIELDS) + assert data["speed"].iloc[0] == 380.0 + assert data["pdyn"].iloc[0] == pytest.approx(2e-6 * 4.2 * 380.0**2) + assert mock_requests_get.call_args_list[0].args[0] == SWACE.URL + mag_file_name + assert mock_requests_get.call_args_list[1].args[0] == SWACE.URL + swepam_file_name + + def test_download_and_process_multiple_days(self, swace_instance): + start_time = datetime(2024, 3, 15, 23, 0, tzinfo=timezone.utc) + end_time = datetime(2024, 3, 16, 1, 0, tzinfo=timezone.utc) + + def mag_data(date: datetime): + return f""":Data_list: {date:%Y%m%d}_ace_mag_1m.txt +:Created: {date:%Y} Mar {date:%d} 0102 UT + + {date:%Y %m %d} 0000 60384 0 0 1.42 -0.89 2.84 3.31 31.42 31.63 + {date:%Y %m %d} 0001 60384 60 0 1.39 -0.87 2.80 3.26 31.38 31.59""" + + def swepam_data(date: datetime): + return f""":Data_list: {date:%Y%m%d}_ace_swepam_1m.txt +:Created: {date:%Y} Mar {date:%d} 0102 UT + + {date:%Y %m %d} 0000 60384 0 0 4.2 380 145000 + {date:%Y %m %d} 0001 60384 60 0 4.1 375 144000""" + + def mock_get(url, **kwargs): + file_name = Path(url).name + file_date = datetime.strptime(file_name.split("_", maxsplit=1)[0], "%Y%m%d").replace(tzinfo=timezone.utc) + mock_response = Mock() + mock_response.content = ( + mag_data(file_date) if "_ace_mag_" in file_name else swepam_data(file_date) + ).encode() + mock_response.raise_for_status = Mock() + return mock_response + + with patch("requests.get", side_effect=mock_get) as mock_requests_get: + swace_instance.download_and_process(start_time, end_time) + + assert (DATA_DIR / "2024/03" / "ACE_SW_NOWCAST_20240315.csv").exists() + assert (DATA_DIR / "2024/03" / "ACE_SW_NOWCAST_20240316.csv").exists() + assert [call.args[0] for call in mock_requests_get.call_args_list] == [ + SWACE.URL + "20240315_ace_mag_1m.txt", + SWACE.URL + "20240315_ace_swepam_1m.txt", + SWACE.URL + "20240316_ace_mag_1m.txt", + SWACE.URL + "20240316_ace_swepam_1m.txt", + ] def test_process_mag_file(self, swace_instance, sample_mag_data): - test_file = DATA_DIR / SWACE.NAME_MAG + request_time = datetime(2024, 3, 15, 1, 0, tzinfo=timezone.utc) + test_file = DATA_DIR / swace_instance._dated_file_name(swace_instance.NAME_MAG, request_time) test_file.parent.mkdir(exist_ok=True) with open(test_file, "w") as f: @@ -124,7 +172,8 @@ def test_process_mag_file(self, swace_instance, sample_mag_data): assert data["bx_gsm"].iloc[0] == 1.42 def test_process_swepam_file(self, swace_instance, sample_swepam_data): - test_file = DATA_DIR / SWACE.NAME_SWEPAM + request_time = datetime(2024, 3, 15, 1, 0, tzinfo=timezone.utc) + test_file = DATA_DIR / swace_instance._dated_file_name(swace_instance.NAME_SWEPAM, request_time) test_file.parent.mkdir(exist_ok=True) with open(test_file, "w") as f: @@ -135,7 +184,8 @@ def test_process_swepam_file(self, swace_instance, sample_swepam_data): assert isinstance(data, pd.DataFrame) assert all(field in data.columns for field in SWACE.SWEPAM_FIELDS) assert len(data) == 2 - assert data["speed"].iloc[0] == 145000.0 + assert data["speed"].iloc[0] == 380.0 + assert data["temperature"].iloc[0] == 145000.0 def test_read_with_no_data(self, swace_instance): start_time = datetime(2020, 1, 1, tzinfo=timezone.utc) @@ -184,9 +234,13 @@ def test_read_with_existing_data(self, swace_instance): assert all(col in data.columns for col in SWACE.MAG_FIELDS + SWACE.SWEPAM_FIELDS) def test_cleanup_after_download(self, swace_instance, sample_mag_data, sample_swepam_data): + start_time = datetime(2024, 3, 15, 1, 0, tzinfo=timezone.utc) + end_time = datetime(2024, 3, 15, 1, 1, tzinfo=timezone.utc) + mag_file_name = swace_instance._dated_file_name(swace_instance.NAME_MAG, start_time) + def mock_get(url, **kwargs): mock_response = Mock() - if SWACE.NAME_MAG in url: + if url.endswith(mag_file_name): mock_response.content = sample_mag_data.encode() else: mock_response.content = sample_swepam_data.encode() @@ -194,8 +248,7 @@ def mock_get(url, **kwargs): return mock_response with patch("requests.get", side_effect=mock_get): - current_time = datetime.now(timezone.utc) - swace_instance.download_and_process(current_time) + swace_instance.download_and_process(start_time, end_time) temp_dir = Path("./temp_sw_ace_wget") assert not temp_dir.exists(), "Temporary directory should be cleaned up"