Simplify/reduce arguments needed to load IBTrACS

Author: leosaffin

huracanpy/_data/_load.py

@@ -25,9 +25,7 @@ def load(
     variable_names=None,
     rename=dict(),
     baselon=None,
-    ibtracs_online=False,
     ibtracs_subset="wmo",
-    ibtracs_clean=True,
     tempest_extremes_unstructured=False,
     tempest_extremes_header_str="start",
     track_calendar=None,
@@ -37,13 +35,15 @@ def load(
 
     The optional parameters for different sources of tracks (currently **IBTrACS**,
     **TRACK** and **TempestExtremes**) are named {source}_{parameter} (in lower case),
-    e.g. "ibtracs_online".
+    e.g. "ibtracs_subset".
 
     Parameters
     ----------
     filename : str, optional
         The file to be loaded. If `source="ibtracs"`, this is not needed as the data is
-        either included in huracanpy or downloaded when called
+        either included in huracanpy or downloaded when called. If the filename is
+        provided for an online IBTrACS subset, then the raw downloaded data will be
+        saved there.
     source : str, optional
         If the file is not a CSV or NetCDF (identified by the file extension) then the
         source needs to be specified to decide how to load the data
@@ -76,15 +76,12 @@ def load(
     baselon : scalar, optional
         Force the loaded longitudes into the range (baselon, baselon + 360). e.g.
         (0, 360) or (-180, 180)
-    ibtracs_online : bool, default=False
-        * **False**: Use a small subset of the IBTrACS data included in this package
-        * **True**: Download the IBTrACS data
-    ibtracs_subset : str, default="ALL"
+    ibtracs_subset : str, default="wmo"
         IBTrACS subset. When loading offline data it is one of
 
-        * **WMO**: Data with the wmo_* variables. The data as reported by the WMO agency
+        * **wmo**: Data with the wmo_* variables. The data as reported by the WMO agency
           responsible for each basin, so methods are not consistent across basins
-        * **USA** or **JTWC**: Data with the usa_* variables. The data as recorded by
+        * **usa** or **JTWC**: Data with the usa_* variables. The data as recorded by
           the USA/Joint Typhoon Warning Centre. Methods are consistent across basins,
           but may not be complete.
 
@@ -98,10 +95,6 @@ def load(
         * **since1980**: Entire IBTrACS database since 1980 (advent of satellite era,
           considered reliable from then on)
 
-    ibtracs_clean : bool, default=True
-        If downloading IBTrACS data, this parameter says whether to delete the
-        downloaded file after loading it into memory.
-
     tempest_extremes_unstructured : bool, default=False,
         By default the first two columns in TempestExtremes files are the i, j indices
         of the closest gridpoint, but for unstructured grids it is a single lookup index
@@ -169,42 +162,7 @@ def load(
                 tempest_extremes_header_str,
             )
         elif source.lower() == "ibtracs":
-            if ibtracs_online:
-                if filename is None:
-                    filename = "ibtracs.csv"
-
-                with ibtracs.online(ibtracs_subset, filename, ibtracs_clean) as f:
-                    # Put IBTrACS specific arguments to read_csv second, so it
-                    # overwrites any arguments passed
-                    kwargs = {
-                        **kwargs,
-                        **dict(
-                            header=0,
-                            skiprows=[1],
-                            converters={
-                                "SID": str,
-                                "SEASON": int,
-                                "BASIN": str,
-                                "SUBBASIN": str,
-                                "LON": float,
-                                "LAT": float,
-                            },
-                        ),
-                    }
-                    return load(
-                        filename=f,
-                        source="csv",
-                        rename=rename,
-                        baselon=baselon,
-                        **kwargs,
-                    )
-            else:
-                return load(
-                    filename=ibtracs.offline(ibtracs_subset),
-                    rename=rename,
-                    baselon=baselon,
-                    **kwargs,
-                )
+            data = ibtracs.load(ibtracs_subset, filename, **kwargs)
         else:
             raise ValueError(f"Source {source} unsupported or misspelled")
 

huracanpy/_data/ibtracs.py

@@ -1,22 +1,44 @@
 """Module for functions related to the ibtracs database"""
 
 from urllib.request import urlretrieve
-import os
 import warnings
 import pathlib
-from contextlib import contextmanager
+
+from . import _csv
 
 here = pathlib.Path(__file__).parent
 ibdata_dir = here / "_ibtracs_files/"
 
 wmo_file = str(ibdata_dir / "wmo.csv")
 usa_file = str(ibdata_dir / "usa.csv")
 
-
-@contextmanager
-def online(subset, filename="ibtracs.csv", clean=True):
+online_default_kwargs = (
+    dict(
+        header=0,
+        skiprows=[1],
+        converters={
+            "SID": str,
+            "SEASON": int,
+            "BASIN": str,
+            "SUBBASIN": str,
+            "LON": float,
+            "LAT": float,
+        },
+    ),
+)
+
+
+def load(subset, filename, **kwargs):
+    if subset.lower() in ["wmo", "usa"]:
+        return offline(subset)
+    else:
+        return online(subset, filename=filename, **kwargs)
+
+
+def online(subset, filename=None, **kwargs):
     """
-    Downloads an load into the current workspace the specified ibtracs subset from the IBTrACS archive online.
+    Downloads and load into the current workspace the specified ibtracs subset from the
+    IBTrACS archive online.
 
     Parameters
     ----------
@@ -26,35 +48,34 @@ def online(subset, filename="ibtracs.csv", clean=True):
         * ALL: Entire IBTrACS database
         * Specific basins: EP, NA, NI, SA, SI, SP, WP
         * last3years: self-explanatory
-        * since1980: Entire IBTrACS database since 1980 (advent of satellite era, considered reliable from then on)
-
-    filename : str
-        (temporary) file to which to save the data
-        The default is "tmp/ibtracs_ACTIVE.csv".
+        * since1980: Entire IBTrACS database since 1980 (advent of satellite era,
+                     considered reliable from then on)
 
-    clean : bool
-        If True (default), remove the temporary file after loading the data.
+    filename : str, optional
+        file to which to save the raw data. None to use a temporary file. Default is
+        None
 
     Returns
     -------
-    ib : the IBTrACS subset requested
-
+    xarray.DataArray
+        the IBTrACS subset requested
     """
-    # TODO: Make it so that the user does not need to specify the filename
+    # Put IBTrACS specific arguments to read_csv second, so it
+    # overwrites any arguments passed
+    kwargs = {**kwargs, **online_default_kwargs}
+
     url = (
-        "https://www.ncei.noaa.gov/data/international-best-track-archive-for-climate-stewardship-ibtracs/v04r01/access/csv/ibtracs."
-        + subset
-        + ".list.v04r01.csv"
+        "https://www.ncei.noaa.gov/data/"
+        "international-best-track-archive-for-climate-stewardship-ibtracs/"
+        f"v04r01/access/csv/ibtracs.{subset}.list.v04r01.csv"
     )
 
+    # filename=None downloads the data to a temporary file
     # Ruff (Flake8 bandit) complains that this url isn't checked, but it explicitly has
     # "https:/" at the start anyway
-    urlretrieve(url, filename)  # noqa: S310
-
-    yield filename
+    filename, _ = urlretrieve(url, filename)  # noqa: S310
 
-    if clean:
-        os.remove(filename)  # Somehow, this is slower than the rest ofthe function (??)
+    return _csv.load(filename, **kwargs)
 
 
 def offline(subset="wmo"):
@@ -103,9 +124,9 @@ def offline(subset="wmo"):
                       which means in particular that wind speeds are in knots and averaged over different time periods.\n\
                     For more information, see the IBTrACS column documentation at https://www.ncei.noaa.gov/sites/default/files/2021-07/IBTrACS_v04_column_documentation.pdf"
         )
-        return wmo_file
+        return _csv.load(wmo_file)
     if subset.lower() in ["usa", "jtwc"]:
-        return usa_file
+        return _csv.load(usa_file)
 
 
 # TODOS:

scripts/prepare_ibtracs_offline.py

@@ -8,7 +8,7 @@
 
 
 def prepare_offline(wmo=True, usa=True):
-    ib = ibtracs.online("since1980", "tmp/ibtracs.csv")
+    ib = huracanpy.load(source="ibtracs", ibtracs_subset="since1980")
 
     # Remove season with tracks that are still provisional
     first_season_provi = ib.where(