"""Loaders for the bundled datasets and for live public data.
Bundled data is small, cited and versioned with the code (see ``data/SOURCES.md``).
Live data is downloaded on demand and cached under ``data/raw/``, which is not
committed.
"""
from __future__ import annotations
from dataclasses import dataclass
from pathlib import Path
from typing import TYPE_CHECKING, Any
import numpy as np
from mescla.types import EndMembers, WaterChemistry
if TYPE_CHECKING: # pragma: no cover
import pandas as pd
__all__ = [
"DATA_DIR",
"BesosEndMembers",
"cache_dir",
"fetch_water_quality_portal",
"load_grafton_nh",
"load_grafton_nh_frame",
"load_hubbard_brook",
"load_hubbard_brook_frame",
"load_tubau_besos",
]
DATA_DIR = Path(__file__).parent / "data"
[docs]
def cache_dir() -> Path:
"""Directory for downloaded data. Override with ``$MESCLA_DATA``."""
import os
path = Path(os.environ.get("MESCLA_DATA", Path.cwd() / "data" / "raw"))
path.mkdir(parents=True, exist_ok=True)
return path
[docs]
@dataclass(frozen=True, eq=False)
class BesosEndMembers:
"""The published Besòs River end-members and their standard-deviation scheme.
Attributes
----------
endmembers : EndMembers
W1 (wet period), D1 and D2 (dry periods), 20 species.
sd_endmembers : ndarray
Standard deviations for the end-members, already multiplied out from the
published fractional factors.
sd_factor_river, sd_factor_groundwater : ndarray
The published fractional factors, per species, for observation points.
conservative : tuple of str
The species the authors treated as conservative.
published_ratios : dict
The reported overall groundwater contributions, for checking against.
"""
endmembers: EndMembers
sd_endmembers: np.ndarray
sd_factor_river: np.ndarray
sd_factor_groundwater: np.ndarray
conservative: tuple[str, ...]
published_ratios: dict[str, float]
table: Any
[docs]
def conservative_only(self) -> EndMembers:
"""The four species the paper treats as conservative: Cl, EC, Na, SO4."""
return self.endmembers.select_species(list(self.conservative))
[docs]
def load_tubau_besos() -> BesosEndMembers:
"""The Besòs Delta end-members of Tubau et al. (2014).
Three river end-members retained by an EMMA of 56 river samples: one from wet
periods (W1, dilute) and two from dry periods (D1, D2, concentrated), separated
mainly by ammonium against calcium and magnesium.
The value of this dataset is less the concentrations than the **standard-deviation
scheme**, which is the best-documented published example of the judgement that
decides a MIX result. Note the very large factors on non-conservative species: that
is the published device for carrying a species through so its composition is
predicted without letting it drive the fit.
Returns
-------
BesosEndMembers
Notes
-----
The 56 river and 51 groundwater samples behind the study are not published in the
paper and so are not bundled. See ``data/SOURCES.md``.
References
----------
Tubau, I., Vázquez-Suñé, E., Jurado, A. and Carrera, J. (2014), Sci. Total Environ.
470-471, 1120-1131, doi:10.1016/j.scitotenv.2013.10.121
"""
import pandas as pd
table = pd.read_csv(DATA_DIR / "tubau2014_besos_endmembers.csv")
species = tuple(table["species"])
concentrations = table[["W1", "D1", "D2"]].to_numpy(dtype=float).T
# The published factors are fractions of the mean river concentration.
mean_river = concentrations.mean(axis=0)
sd_endmembers = np.maximum(
table["sd_factor_endmember"].to_numpy(dtype=float) * mean_river, 1e-9
)
sd_endmembers = np.tile(sd_endmembers, (3, 1))
return BesosEndMembers(
endmembers=EndMembers(
concentrations,
species,
("W1", "D1", "D2"),
# Everything is mg/L except conductivity, which is uS/cm. Stating it
# per species keeps the table honest and stops a blanket conversion
# from quietly relabelling that column.
{"EC": "uS/cm"},
sd_endmembers,
),
sd_endmembers=sd_endmembers,
sd_factor_river=table["sd_factor_river"].to_numpy(dtype=float),
sd_factor_groundwater=table["sd_factor_groundwater"].to_numpy(dtype=float),
conservative=tuple(table.loc[table["conservative"] == "yes", "species"]),
published_ratios={"W1": 0.26, "D1": 0.09, "D2": 0.65},
table=table,
)
[docs]
def load_grafton_nh(site_type: str | None = None) -> WaterChemistry:
"""Real major-ion chemistry from Grafton County, New Hampshire.
130 complete analyses (Ca, Mg, Na, K, Cl, SO4, mg/L) from 66 stream and 64
groundwater-well samples collected 1990-2021, retrieved from the Water Quality
Portal. Public domain (U.S. Geological Survey).
This is the county containing the Hubbard Brook Experimental Forest, and the
chemistry spans the range that makes it a useful teaching set: dilute upland
streams at one extreme, road-salt-affected and mineralised groundwater at the
other.
Parameters
----------
site_type : {"Stream", "Well"}, optional
Restrict to one kind of water.
Returns
-------
WaterChemistry
With ``.attrs``-style metadata available through :func:`load_grafton_nh_frame`.
"""
frame = load_grafton_nh_frame(site_type)
species = ["Ca", "Mg", "Na", "K", "Cl", "SO4"]
labels = [f"{row.site_type[0]}{i:03d}" for i, row in enumerate(frame.itertuples(), 1)]
return WaterChemistry(frame[species].to_numpy(dtype=float), tuple(species), tuple(labels))
[docs]
def load_grafton_nh_frame(site_type: str | None = None) -> pd.DataFrame:
"""The Grafton County dataset as a DataFrame, with site, date and site type."""
import pandas as pd
frame = pd.read_csv(DATA_DIR / "grafton_nh_major_ions.csv", parse_dates=["date"])
if site_type is not None:
frame = frame[frame["site_type"] == site_type].reset_index(drop=True)
if frame.empty:
raise ValueError(f"no samples with site_type={site_type!r}")
return frame
[docs]
def fetch_water_quality_portal(
characteristics: tuple[str, ...] = (
"Calcium", "Magnesium", "Sodium", "Potassium", "Chloride", "Sulfate",
),
countycode: str | None = None,
statecode: str | None = None,
siteid: tuple[str, ...] = (),
start_date: str = "01-01-1990",
organisation_prefix: str = "USGS-",
timeout: int = 300,
refresh: bool = False,
) -> pd.DataFrame:
"""Download major-ion chemistry from the Water Quality Portal, wide and complete.
A live, reproducible route to real US water chemistry: public domain, no
registration, and the same query that produced the bundled Grafton County dataset.
Parameters
----------
characteristics : tuple of str
WQP characteristic names. Spelling matters -- it is "Sulfate", not "Sulphate".
countycode, statecode : str, optional
e.g. ``"US:33:009"`` (Grafton County, NH) or ``"US:33"``.
siteid : tuple of str
Specific monitoring locations, e.g. ``("USGS-01075098",)``.
start_date : str, default "01-01-1990"
``MM-DD-YYYY``, as the portal expects.
organisation_prefix : str, default "USGS-"
Keep only sites from this organisation. Pass ``""`` for all.
timeout : int, default 300
refresh : bool, default False
Ignore any cached copy.
Returns
-------
DataFrame
One row per (site, date), one column per characteristic, in mg/L, with only
samples complete in every requested characteristic.
Notes
-----
Please cite: National Water Quality Monitoring Council, Water Quality Portal,
https://www.waterqualitydata.us/, and Read et al. (2017), WRR 53, 1735-1745.
"""
import hashlib
import urllib.parse
import urllib.request
import pandas as pd
if not (countycode or statecode or siteid):
raise ValueError("supply at least one of countycode, statecode or siteid")
query: list[tuple[str, str]] = [("characteristicName", c) for c in characteristics]
if countycode:
query.append(("countycode", countycode))
if statecode:
query.append(("statecode", statecode))
query += [("siteid", s) for s in siteid]
query += [("startDateLo", start_date), ("mimeType", "csv"), ("zip", "no")]
url = "https://www.waterqualitydata.us/data/Result/search?" + urllib.parse.urlencode(query)
key = hashlib.sha256(url.encode()).hexdigest()[:16]
cached = cache_dir() / f"wqp_{key}.csv"
if cached.exists() and not refresh:
raw = pd.read_csv(cached, low_memory=False)
else:
with urllib.request.urlopen(url, timeout=timeout) as response:
cached.write_bytes(response.read())
raw = pd.read_csv(cached, low_memory=False)
raw = raw[raw["ResultMeasure/MeasureUnitCode"].isin(["mg/l", "mg/L"])].copy()
if organisation_prefix:
raw = raw[raw["MonitoringLocationIdentifier"].str.startswith(organisation_prefix)]
raw["value"] = pd.to_numeric(raw["ResultMeasureValue"], errors="coerce")
wide = raw.pivot_table(
index=["MonitoringLocationIdentifier", "ActivityStartDate"],
columns="CharacteristicName",
values="value",
aggfunc="mean",
).dropna()
return wide.reset_index().rename(
columns={"MonitoringLocationIdentifier": "site", "ActivityStartDate": "date"}
)
HUBBARD_BROOK_PORTAL = (
"https://portal.edirepository.org/nis/mapbrowse?packageid=knb-lter-hbr.20.9"
)
HUBBARD_BROOK_DOI = "10.6073/pasta/8d2d88dc718b6c5a2183cd88aae26fb1"
#: Species retained from the bundled record. ``Al`` and ``SiO2`` are columns of the
#: source file but were never measured in precipitation (every value is the missing
#: code), and ``PO4`` is missing in a fifth of months, so none of the three is here.
#: ``H`` is not in the source either -- it is derived from pH on load, because pH is a
#: logarithm and **mixing is linear in concentration**.
HUBBARD_BROOK_SPECIES = ("Ca", "Mg", "K", "Na", "NH4", "H", "SO4", "NO3", "Cl")
#: The Hubbard Brook missing-value code. It is a sentinel, not a measurement.
HUBBARD_BROOK_MISSING = -3.0
#: Hubbard Brook water year: June to May, following the start of the record in
#: June 1963. A water year is labelled by the calendar year it begins in.
HUBBARD_BROOK_WATER_YEAR_START = 6
def _hubbard_brook_raw() -> pd.DataFrame:
"""The bundled CSV with the missing code replaced by ``nan`` and H+ derived."""
import numpy as np
import pandas as pd
frame = pd.read_csv(DATA_DIR / "hubbard_brook_ws6_precipitation.csv")
chemistry = [c for c in frame.columns if c not in ("ws", "year", "mo", "precip")]
frame[chemistry] = frame[chemistry].replace(HUBBARD_BROOK_MISSING, np.nan)
# pH -> H+ in mg/L. Concentrations mix linearly; a pH does not.
frame["H"] = 10.0 ** (-frame["pH"]) * 1.008 * 1000.0
frame["date"] = pd.to_datetime(
{"year": frame["year"], "month": frame["mo"], "day": 1}
)
frame["water_year"] = np.where(
frame["mo"] >= HUBBARD_BROOK_WATER_YEAR_START, frame["year"], frame["year"] - 1
)
return frame.rename(columns={"precip": "precip_mm"})
def _hubbard_brook_annual(monthly: pd.DataFrame) -> pd.DataFrame:
"""Volume-weighted water-year means of ``monthly``, complete years only.
Volume weighting is not a refinement, it is the only average that conserves
mass: a 300 mm month and a 20 mm month deliver very different amounts of solute
at the same concentration, and an unweighted mean of concentrations silently
reweights the record towards dry months.
"""
import numpy as np
import pandas as pd
species = list(HUBBARD_BROOK_SPECIES)
rows = []
for water_year, group in monthly.groupby("water_year"):
complete = group.dropna(subset=species)
if len(complete) != 12: # partial years cannot be volume-weighted honestly
continue
weights = complete["precip_mm"].to_numpy(dtype=float)
row = {s: float(np.average(complete[s], weights=weights)) for s in species}
row["water_year"] = int(water_year)
row["precip_mm"] = float(weights.sum())
rows.append(row)
return pd.DataFrame(rows).set_index("water_year")[["precip_mm", *species]]
[docs]
def load_hubbard_brook_frame(annual: bool = False) -> pd.DataFrame:
"""The bundled Hubbard Brook precipitation record as a DataFrame.
Parameters
----------
annual : bool, default False
``False`` returns the monthly record as published, with ``nan`` for the
missing code and columns ``date``, ``water_year``, ``precip_mm`` alongside
the chemistry. ``True`` returns **volume-weighted water-year means** over
complete years only, indexed by water year.
Returns
-------
DataFrame
"""
monthly = _hubbard_brook_raw()
if annual:
return _hubbard_brook_annual(monthly)
columns = ["date", "water_year", "precip_mm", *HUBBARD_BROOK_SPECIES, "pH"]
return monthly[columns]
[docs]
def load_hubbard_brook(annual: bool = False, complete_only: bool = True) -> WaterChemistry:
"""Bulk precipitation chemistry at Hubbard Brook Watershed 6, 1963-2014.
The canonical long-term record for atmospheric deposition: volume-weighted
monthly concentrations from bulk collectors in a New Hampshire northern-hardwood
catchment, the record on which acid rain was first described in North America.
Bundled under CC-BY; see ``data/SOURCES.md``.
Nine species are retained (Ca, Mg, K, Na, NH4, H, SO4, NO3, Cl), where ``H`` is
derived from the published pH on load. **Never put pH itself into a mixing
calculation** -- it is a logarithm, and mixing is linear in concentration.
Parameters
----------
annual : bool, default False
Return volume-weighted water-year means (June-May, complete years only)
rather than the monthly record. Notebook 09 shows why this matters: the
monthly data fails every Hooper diagnostic and the annual means do not.
complete_only : bool, default True
Drop months with any missing species. Only applies when ``annual`` is False,
since the annual means are built from complete years by construction.
Returns
-------
WaterChemistry
Concentrations in mg/L, indexed ``YYYY-MM`` (monthly) or ``YYYY`` (annual).
Notes
-----
These are **atmospheric sources scavenged into rainwater**, not a mixture of
three groundwaters, and notebook 09 works through what that does to a
volumetric mixing model. The record is also the clearest available
demonstration that end-member compositions need not hold still: non-sea-salt
sulphate falls roughly fivefold across it.
References
----------
Likens, G. (2016). Chemistry of Bulk Precipitation at Hubbard Brook Experimental
Forest, Watershed 6, 1963 - present ver 9. Environmental Data Initiative,
doi:10.6073/pasta/8d2d88dc718b6c5a2183cd88aae26fb1 (CC-BY).
"""
species = list(HUBBARD_BROOK_SPECIES)
if annual:
frame = load_hubbard_brook_frame(annual=True)
labels = tuple(str(y) for y in frame.index)
else:
frame = load_hubbard_brook_frame()
if complete_only:
frame = frame.dropna(subset=species)
labels = tuple(frame["date"].dt.strftime("%Y-%m"))
return WaterChemistry(frame[species].to_numpy(dtype=float), tuple(species), labels)