diff --git a/backend/connectors/nmose/source.py b/backend/connectors/nmose/source.py index fa785e0a..881e0427 100644 --- a/backend/connectors/nmose/source.py +++ b/backend/connectors/nmose/source.py @@ -39,13 +39,16 @@ def get_records(self, *args, **kw) -> List[Dict]: # if config.end_date: # params["endDt"] = config.end_dt.date().isoformat() + # The OSE POD FeatureServer was renamed from "OSE_PODs" to + # "OSE_Points_of_Diversion" (the old name now 400s "Invalid URL"). url: str = ( - "https://services2.arcgis.com/qXZbWTdPDbTjl7Dy/arcgis/rest/services/OSE_PODs/FeatureServer/0/query" + "https://services2.arcgis.com/qXZbWTdPDbTjl7Dy/arcgis/rest/services/" + "OSE_Points_of_Diversion/FeatureServer/0/query" ) params["where"] = "pod_status = 'ACT' AND pod_basin NOT IN ('SP', 'SD', 'LWD')" params["outFields"] = ( - "OBJECTID,pod_basin,pod_status,easting,northing,datum,utm_accura,status,county" + "OBJECTID,pod_basin,pod_status,easting,northing,datum,utm_accura,status,county," "pod_name,pod_nbr,pod_suffix,pod_file,depth_well,aquifer,elevation" ) diff --git a/backend/persisters/ogc_features.py b/backend/persisters/ogc_features.py index 63403ad9..12d51cb7 100644 --- a/backend/persisters/ogc_features.py +++ b/backend/persisters/ogc_features.py @@ -51,7 +51,11 @@ def _point_geometry(lat, lon, elev=None) -> dict: def _dump_collection( - path: str, collection_id: str, features: list, meta: dict, extra: Optional[dict] = None + path: str, + collection_id: str, + features: list, + meta: dict, + extra: Optional[dict] = None, ) -> dict: """Build the OGC FeatureCollection envelope around *features*, write it to *path*, and return it. *extra* injects collection-level keys (e.g. @@ -69,7 +73,11 @@ def _dump_collection( "numberMatched": len(features), "numberReturned": len(features), "links": [ - {"href": meta.get("href", ""), "rel": "self", "type": "application/geo+json"} + { + "href": meta.get("href", ""), + "rel": "self", + "type": "application/geo+json", + } ], "features": features, } @@ -99,8 +107,11 @@ def _tds_class(value: Optional[float]) -> str: def _make_feature(record, collection_id: str) -> dict: """Build one OGC-compliant Feature from a SummaryRecord or SiteRecord.""" - props = {k: getattr(record, k) for k in record.keys - if k not in ("latitude", "longitude", "elevation")} + props = { + k: getattr(record, k) + for k in record.keys + if k not in ("latitude", "longitude", "elevation") + } if getattr(record, "parameter_name", None) == TDS: props["tds_class"] = _tds_class(_num(getattr(record, "latest_value", None))) @@ -123,7 +134,7 @@ def _num(value) -> Optional[float]: return None try: return float(value) - except (TypeError, ValueError): + except TypeError, ValueError: return None @@ -272,12 +283,16 @@ def dump_major_chemistry_collection(path: str, records: list, meta: dict) -> dic props[f"{analyte}_units"] = vals["units"] props[f"{analyte}_date"] = vals["date"] - features.append({ - "type": "Feature", - "id": _feature_id(source, rid), - "geometry": _point_geometry(well["latitude"], well["longitude"], well["elevation"]), - "properties": props, - }) + features.append( + { + "type": "Feature", + "id": _feature_id(source, rid), + "geometry": _point_geometry( + well["latitude"], well["longitude"], well["elevation"] + ), + "properties": props, + } + ) return _dump_collection(path, collection_id, features, meta) @@ -368,16 +383,18 @@ def dump_trend_collection( if source_datastream_link: props["source_datastream_link"] = source_datastream_link - features.append({ - "type": "Feature", - "id": _feature_id(props["source"], props["id"]), - "geometry": _point_geometry( - site.get("latitude"), - site.get("longitude"), - site.get("elevation"), - ), - "properties": props, - }) + features.append( + { + "type": "Feature", + "id": _feature_id(props["source"], props["id"]), + "geometry": _point_geometry( + site.get("latitude"), + site.get("longitude"), + site.get("elevation"), + ), + "properties": props, + } + ) return _dump_collection( path, collection_id, features, meta, extra={"trend_method": method} @@ -412,7 +429,9 @@ def dump_mcl_exceedance_collection( well = wells.get(key) if well is None: well = { - "source": source, "id": rid, "name": getattr(r, "name", None), + "source": source, + "id": rid, + "name": getattr(r, "name", None), "latitude": getattr(r, "latitude", None), "longitude": getattr(r, "longitude", None), "elevation": getattr(r, "elevation", None), @@ -431,7 +450,9 @@ def dump_mcl_exceedance_collection( features = [] for (source, rid), well in wells.items(): props = { - "source": source, "id": rid, "name": well["name"], + "source": source, + "id": rid, + "name": well["name"], "well_depth": well["well_depth"], "well_depth_units": well["well_depth_units"], } @@ -462,14 +483,16 @@ def dump_mcl_exceedance_collection( props["exceedance_count"] = len(exceeded) props["exceeded_analytes"] = sorted(exceeded) - features.append({ - "type": "Feature", - "id": _feature_id(source, rid), - "geometry": _point_geometry( - well["latitude"], well["longitude"], well["elevation"] - ), - "properties": props, - }) + features.append( + { + "type": "Feature", + "id": _feature_id(source, rid), + "geometry": _point_geometry( + well["latitude"], well["longitude"], well["elevation"] + ), + "properties": props, + } + ) return _dump_collection( path, collection_id, features, meta, extra={"mcl_thresholds": thresholds} @@ -501,10 +524,12 @@ def dump_monitoring_recency_collection( features = [] for site, obs_list in zip(site_records, timeseries_records): epochs = [ - e for e in ( + e + for e in ( _parse_epoch_seconds(o.get("date_measured"), o.get("time_measured")) for o in obs_list - ) if e is not None + ) + if e is not None ] record_count = len(epochs) if record_count: @@ -533,17 +558,22 @@ def dump_monitoring_recency_collection( "days_since_last": days_since_last, "status": status, } - features.append({ - "type": "Feature", - "id": _feature_id(props["source"], props["id"]), - "geometry": _point_geometry( - site.get("latitude"), site.get("longitude"), site.get("elevation") - ), - "properties": props, - }) + features.append( + { + "type": "Feature", + "id": _feature_id(props["source"], props["id"]), + "geometry": _point_geometry( + site.get("latitude"), site.get("longitude"), site.get("elevation") + ), + "properties": props, + } + ) return _dump_collection( - path, collection_id, features, meta, + path, + collection_id, + features, + meta, extra={"stale_threshold_days": stale_days}, ) @@ -604,12 +634,14 @@ def dump_timeseries_collection( props = {k: getattr(obs, k) for k in obs.keys} props["datetime"] = dt - features.append({ - "type": "Feature", - "id": feature_id, - "geometry": geometry, - "properties": props, - }) + features.append( + { + "type": "Feature", + "id": feature_id, + "geometry": geometry, + "properties": props, + } + ) return _dump_collection(path, collection_id, features, meta) @@ -662,9 +694,7 @@ def dump_hardness_collection(path: str, records: list, meta: dict) -> dict: if ca is None or mg is None: hardness = None else: - hardness = round( - _HARDNESS_CA_FACTOR * ca + _HARDNESS_MG_FACTOR * mg, 1 - ) + hardness = round(_HARDNESS_CA_FACTOR * ca + _HARDNESS_MG_FACTOR * mg, 1) props = { "source": source, @@ -683,7 +713,10 @@ def dump_hardness_collection(path: str, records: list, meta: dict) -> dict: features.append(_well_feature(source, rid, well, props)) return _dump_collection( - path, collection_id, features, meta, + path, + collection_id, + features, + meta, extra={"hardness_method": HARDNESS_METHOD_DESCRIPTION}, ) @@ -765,14 +798,20 @@ def dump_water_type_collection(path: str, records: list, meta: dict) -> dict: "anion_meq_total": round(anion_total, 3), } if cation_total <= 0 or anion_total <= 0: - props.update({ - "water_type": "insufficient", - "dominant_cation": None, - "dominant_anion": None, - "ca_pct": None, "mg_pct": None, "na_k_pct": None, - "hco3_pct": None, "cl_pct": None, "so4_pct": None, - "charge_balance_pct": None, - }) + props.update( + { + "water_type": "insufficient", + "dominant_cation": None, + "dominant_anion": None, + "ca_pct": None, + "mg_pct": None, + "na_k_pct": None, + "hco3_pct": None, + "cl_pct": None, + "so4_pct": None, + "charge_balance_pct": None, + } + ) else: ca_pct = 100 * ca / cation_total mg_pct = 100 * mg / cation_total @@ -780,30 +819,34 @@ def dump_water_type_collection(path: str, records: list, meta: dict) -> dict: hco3_pct = 100 * hco3_co3 / anion_total cl_pct = 100 * cl / anion_total so4_pct = 100 * so4 / anion_total - dom_cation = _dominant( - {"Ca": ca_pct, "Mg": mg_pct, "Na+K": na_k_pct} - ) + dom_cation = _dominant({"Ca": ca_pct, "Mg": mg_pct, "Na+K": na_k_pct}) dom_anion = _dominant({"HCO3": hco3_pct, "Cl": cl_pct, "SO4": so4_pct}) - props.update({ - "water_type": f"{dom_cation}-{dom_anion}", - "dominant_cation": dom_cation, - "dominant_anion": dom_anion, - "ca_pct": round(ca_pct, 1), - "mg_pct": round(mg_pct, 1), - "na_k_pct": round(na_k_pct, 1), - "hco3_pct": round(hco3_pct, 1), - "cl_pct": round(cl_pct, 1), - "so4_pct": round(so4_pct, 1), - "charge_balance_pct": round( - 100 * (cation_total - anion_total) - / (cation_total + anion_total), - 1, - ), - }) + props.update( + { + "water_type": f"{dom_cation}-{dom_anion}", + "dominant_cation": dom_cation, + "dominant_anion": dom_anion, + "ca_pct": round(ca_pct, 1), + "mg_pct": round(mg_pct, 1), + "na_k_pct": round(na_k_pct, 1), + "hco3_pct": round(hco3_pct, 1), + "cl_pct": round(cl_pct, 1), + "so4_pct": round(so4_pct, 1), + "charge_balance_pct": round( + 100 + * (cation_total - anion_total) + / (cation_total + anion_total), + 1, + ), + } + ) features.append(_well_feature(source, rid, well, props)) return _dump_collection( - path, collection_id, features, meta, + path, + collection_id, + features, + meta, extra={"water_type_method": WATER_TYPE_METHOD_DESCRIPTION}, ) @@ -877,7 +920,10 @@ def dump_sar_collection(path: str, records: list, meta: dict) -> dict: features.append(_well_feature(source, rid, well, props)) return _dump_collection( - path, collection_id, features, meta, + path, + collection_id, + features, + meta, extra={"sar_method": SAR_METHOD_DESCRIPTION}, ) @@ -957,7 +1003,10 @@ def dump_ion_balance_collection(path: str, records: list, meta: dict) -> dict: features.append(_well_feature(source, rid, well, props)) return _dump_collection( - path, collection_id, features, meta, + path, + collection_id, + features, + meta, extra={"ion_balance_method": ION_BALANCE_METHOD_DESCRIPTION}, ) @@ -1004,7 +1053,8 @@ def dump_data_density_collection( mean_interval_days = ( round((span_years * 365.25) / (record_count - 1), 1) - if record_count > 1 else None + if record_count > 1 + else None ) observations_per_year = ( round(observation_count / span_years, 2) if span_years > 0 else None @@ -1028,7 +1078,10 @@ def dump_data_density_collection( features.append(_site_feature(site, props)) return _dump_collection( - path, collection_id, features, meta, + path, + collection_id, + features, + meta, extra={"data_density_method": DATA_DENSITY_METHOD_DESCRIPTION}, ) @@ -1085,18 +1138,15 @@ def dump_waterlevel_change_collection( dtw_end = round(dtw_end, 2) target = end_e - target_span # Closest daily point to the window-start target, excluding the end. - cand_epoch, cand_val = min( - pairs[:-1], key=lambda p: abs(p[0] - target) - ) + cand_epoch, cand_val = min(pairs[:-1], key=lambda p: abs(p[0] - target)) if abs(cand_epoch - target) <= tolerance: start_e, dtw_start = cand_epoch, round(cand_val, 2) change_ft = round(dtw_end - dtw_start, 2) - actual_window_years = round( - (end_e - start_e) / _SECONDS_PER_YEAR, 3 - ) + actual_window_years = round((end_e - start_e) / _SECONDS_PER_YEAR, 3) n_in_window = sum(1 for p in pairs if start_e <= p[0] <= end_e) direction = ( - "declining" if change_ft > 0 + "declining" + if change_ft > 0 else "rising" if change_ft < 0 else "stable" ) status = "ok" @@ -1124,7 +1174,10 @@ def dump_waterlevel_change_collection( features.append(_site_feature(site, props)) return _dump_collection( - path, collection_id, features, meta, + path, + collection_id, + features, + meta, extra={ "change_method": WATERLEVEL_CHANGE_METHOD_TEMPLATE.format( window=window_years @@ -1198,15 +1251,12 @@ def dump_waterlevel_status_collection( min_v, max_v = values[0], values[-1] mid = record_count // 2 median_v = ( - values[mid] if record_count % 2 - else (values[mid - 1] + values[mid]) / 2 + values[mid] if record_count % 2 else (values[mid - 1] + values[mid]) / 2 ) if record_count >= _STATUS_MIN_RECORDS: less = sum(1 for v in values if v < latest_v) equal = sum(1 for v in values if v == latest_v) - dtw_percentile = round( - 100 * (less + 0.5 * equal) / record_count, 1 - ) + dtw_percentile = round(100 * (less + 0.5 * equal) / record_count, 1) status = _waterlevel_status(dtw_percentile) props = { @@ -1218,7 +1268,9 @@ def dump_waterlevel_status_collection( "well_depth_units": site.get("well_depth_units"), "record_count": record_count, "observation_count": observation_count, - "first_observation_datetime": _iso_utc(pairs[0][0]) if record_count else None, + "first_observation_datetime": ( + _iso_utc(pairs[0][0]) if record_count else None + ), "last_observation_datetime": _iso_utc(latest_e), "span_years": round(span_years, 3), "latest_dtw": None if latest_v is None else round(latest_v, 2), @@ -1233,7 +1285,10 @@ def dump_waterlevel_status_collection( features.append(_site_feature(site, props)) return _dump_collection( - path, collection_id, features, meta, + path, + collection_id, + features, + meta, extra={"status_method": WATERLEVEL_STATUS_METHOD_DESCRIPTION}, ) @@ -1321,7 +1376,10 @@ def dump_seasonal_amplitude_collection( features.append(_site_feature(site, props)) return _dump_collection( - path, collection_id, features, meta, + path, + collection_id, + features, + meta, extra={ "seasonal_amplitude_method": SEASONAL_AMPLITUDE_METHOD_TEMPLATE.format( min_days=min_days_per_year @@ -1372,7 +1430,8 @@ def dump_depletion_projection_collection( record_count = len(pairs) span_years = ( (pairs[-1][0] - pairs[0][0]) / _SECONDS_PER_YEAR - if record_count >= 2 else 0.0 + if record_count >= 2 + else 0.0 ) latest_e = latest_v = None @@ -1400,9 +1459,7 @@ def dump_depletion_projection_collection( status = "dtw exceeds well depth" else: years_to_depletion = round(remaining / slope, 1) - latest_year = datetime.fromtimestamp( - latest_e, tz=timezone.utc - ).year + latest_year = datetime.fromtimestamp(latest_e, tz=timezone.utc).year projected_year = int(latest_year + years_to_depletion) status = "projected" @@ -1428,7 +1485,10 @@ def dump_depletion_projection_collection( features.append(_site_feature(site, props)) return _dump_collection( - path, collection_id, features, meta, + path, + collection_id, + features, + meta, extra={"depletion_method": DEPLETION_PROJECTION_METHOD_DESCRIPTION}, ) @@ -1460,9 +1520,7 @@ def _wqi_class(wqi: Optional[float]) -> str: return "poor" -def dump_wqi_collection( - path: str, records: list, meta: dict, thresholds: dict -) -> dict: +def dump_wqi_collection(path: str, records: list, meta: dict, thresholds: dict) -> dict: """ Write an OGC FeatureCollection of per-well CCME water quality index, one Feature per well. *records* is a flat list of SummaryRecord (one per @@ -1512,16 +1570,129 @@ def dump_wqi_collection( wqi = 100 - math.sqrt(f1 * f1 + f2 * f2 + f3 * f3) / 1.732 wqi = round(min(100.0, max(0.0, wqi)), 1) - props.update({ - "wqi": wqi, - "wqi_class": _wqi_class(wqi), - "n_analytes_tested": len(tests), - "n_exceeding": len(exceeded), - "exceeded_analytes": exceeded, - }) + props.update( + { + "wqi": wqi, + "wqi_class": _wqi_class(wqi), + "n_analytes_tested": len(tests), + "n_exceeding": len(exceeded), + "exceeded_analytes": exceeded, + } + ) features.append(_well_feature(source, rid, well, props)) return _dump_collection( - path, collection_id, features, meta, + path, + collection_id, + features, + meta, extra={"wqi_method": WQI_METHOD_DESCRIPTION, "mcl_thresholds": thresholds}, ) + + +def dump_well_correlation_collection( + path: str, + sites: list, + meta: dict, + *, + max_link_distance_m: float = None, + depth_tolerance_ft: float = None, + elevation_tolerance_ft: float = None, + pod_link_distance_m: float = None, +) -> dict: + """ + Write an OGC FeatureCollection of the cross-agency well correlation layer, + one Feature per input well. + + *sites* is a flat list of site payload dicts gathered from **every** source + (``source``, ``id``, ``latitude``, ``longitude``, ``well_depth``, + ``usgs_site_id``, ``alternate_site_id``, ...). They are correlated by + :func:`backend.well_correlation.correlate_wells`, which links wells across + agencies (explicit id references + spatial/depth agreement) and associates + each with any OSE POD. + + Per well the feature carries: ``cluster_id``, ``cluster_size``, + ``n_agencies``, ``linked_site_ids`` (comma-joined "SOURCE:id" of the other + wells in the cluster), ``linked_by_agency`` (JSON string), ``n_linked``, + ``ose_pod_ids`` (comma-joined), ``ose_pod_link_method``, ``match_method``, + ``match_confidence``, and ``is_ose_pod``. Nested collections are serialized + to strings so the layer publishes cleanly to GeoPackage/GeoServer. + + The collection carries ``correlation_method``. + + §V: MUST include top-level id, type, numberReturned, timeStamp. + §V: Each Feature MUST have top-level id. + """ + # Imported here (not at module top) to keep the correlation engine — pure + # stdlib — decoupled from the serialization layer's heavier imports. + from backend.well_correlation import ( + CORRELATION_METHOD_DESCRIPTION, + DEFAULT_DEPTH_TOLERANCE_FT, + DEFAULT_ELEVATION_TOLERANCE_FT, + DEFAULT_MAX_LINK_DISTANCE_M, + DEFAULT_POD_LINK_DISTANCE_M, + correlate_wells, + ) + + collection_id = meta.get("id", "collection") + correlations = correlate_wells( + sites, + max_link_distance_m=( + DEFAULT_MAX_LINK_DISTANCE_M + if max_link_distance_m is None + else max_link_distance_m + ), + depth_tolerance_ft=( + DEFAULT_DEPTH_TOLERANCE_FT + if depth_tolerance_ft is None + else depth_tolerance_ft + ), + elevation_tolerance_ft=( + DEFAULT_ELEVATION_TOLERANCE_FT + if elevation_tolerance_ft is None + else elevation_tolerance_ft + ), + pod_link_distance_m=( + DEFAULT_POD_LINK_DISTANCE_M + if pod_link_distance_m is None + else pod_link_distance_m + ), + ) + + features = [] + for c in correlations: + props = { + "source": c["source"], + "id": c["id"], + "name": c["name"], + "well_depth": c["well_depth"], + "cluster_id": c["cluster_id"], + "cluster_size": c["cluster_size"], + "n_agencies": c["n_agencies"], + "n_linked": len(c["linked_site_ids"]), + "linked_site_ids": ",".join(c["linked_site_ids"]), + "linked_by_agency": json.dumps(c["linked_by_agency"], default=str), + "ose_pod_ids": ",".join(c["ose_pod_ids"]), + "ose_pod_link_method": c["ose_pod_link_method"], + "match_method": c["match_method"], + "match_confidence": c["match_confidence"], + "is_ose_pod": c["is_ose_pod"], + } + features.append( + { + "type": "Feature", + "id": _feature_id(c["source"] or "", c["id"] or ""), + "geometry": _point_geometry( + c["latitude"], c["longitude"], c["elevation"] + ), + "properties": props, + } + ) + + return _dump_collection( + path, + collection_id, + features, + meta, + extra={"correlation_method": CORRELATION_METHOD_DESCRIPTION}, + ) diff --git a/backend/unifier.py b/backend/unifier.py index 8c3c4cf0..10906f59 100644 --- a/backend/unifier.py +++ b/backend/unifier.py @@ -146,7 +146,7 @@ def _site_wrapper(site_source, parameter_source, persister, config, raise_errors try: sites = site_source.read() - except (USGSRateLimitError, PartialOrNoDataError): + except USGSRateLimitError, PartialOrNoDataError: config.warn(incomplete_sites_record_msg) sites = [] @@ -176,10 +176,12 @@ def _site_wrapper(site_source, parameter_source, persister, config, raise_errors summary_records = parameter_source.read( site_records, use_summarize, start_ind, end_ind ) - except (USGSRateLimitError, PartialOrNoDataError): + except USGSRateLimitError, PartialOrNoDataError: # remove partial records to prevent incomplete data from being saved persister.sites = persister.sites[:initial_sites_len] - persister.timeseries = persister.timeseries[:initial_timeseries_len] + persister.timeseries = persister.timeseries[ + :initial_timeseries_len + ] persister.records = persister.records[:initial_records_len] config.warn(incomplete_parameter_record_msg) break @@ -193,10 +195,12 @@ def _site_wrapper(site_source, parameter_source, persister, config, raise_errors results = parameter_source.read( site_records, use_summarize, start_ind, end_ind ) - except (USGSRateLimitError, PartialOrNoDataError): + except USGSRateLimitError, PartialOrNoDataError: # remove partial records to prevent incomplete data from being saved persister.sites = persister.sites[:initial_sites_len] - persister.timeseries = persister.timeseries[:initial_timeseries_len] + persister.timeseries = persister.timeseries[ + :initial_timeseries_len + ] persister.records = persister.records[:initial_records_len] config.warn(incomplete_parameter_record_msg) break @@ -354,6 +358,29 @@ def unify_source_both(config, source_key): return summary_persister, timeseries_persister +def collect_sites(config): + """Gather site records from **every** enabled source (sites only, no + parameter data) and return them as a flat list of ``_payload`` dicts. + + Used by the cross-agency well-correlation product, which needs the location + of every well across all agencies — including OSE PODs, which carry no + parameter time series. Runs in ``sites_only`` mode over + ``all_site_sources()`` so a source that provides no parameter (e.g. the OSE + POD source) still contributes its sites. Errors on individual sources are + swallowed by ``_site_wrapper`` so one dead source does not abort the sweep. + """ + config.validate() + + config.sites_only = True + persister = make_persister(config) + config._persister = persister + + for site_source, _ in config.all_site_sources(): + _site_wrapper(site_source, None, persister, config) + + return [s._payload for s in persister.sites] + + def get_county_bounds(county): config = Config() config.county = county diff --git a/backend/well_correlation.py b/backend/well_correlation.py new file mode 100644 index 00000000..c688d9ef --- /dev/null +++ b/backend/well_correlation.py @@ -0,0 +1,691 @@ +# =============================================================================== +# Copyright 2025 Jake Ross +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# =============================================================================== +"""Cross-agency well correlation. + +Every agency (NMBGMR, USGS-NWIS, OSE ISC Seven Rivers, PVACD, ...) assigns its +own identifier to a physical well, and the same well is often monitored by more +than one agency. A few wells carry explicit cross-references +(``alternate_site_id`` / ``usgs_site_id``), but most do not. This module links +wells that are the *same physical well* across agencies so every well can be +traced back — ideally to an OSE Point of Diversion (POD). + +Two kinds of evidence link wells: + +1. **Explicit id references** (high confidence). A site's ``usgs_site_id`` or the + tokens in its ``alternate_site_id`` are matched against the ids (and USGS ids) + of every other site. A match is a hard link — the agencies themselves assert + the wells are the same. + +2. **Spatial + attribute agreement** (lower confidence). OSE PODs and many older + records are not accurately surveyed, so latitude/longitude cannot be trusted + as the exact well position. Two wells from *different* agencies within + ``max_link_distance_m`` of each other are candidate matches; the confidence is + raised when their reported well depths agree (``depth_tolerance_ft``) and + lowered when depth is unavailable to corroborate. + +Links are unioned into connected components — each component is one inferred +physical well. Every component is then associated with any OSE POD(s) it +contains or that sit within ``pod_link_distance_m`` of it. + +:func:`correlate_wells` returns one crosswalk record per *input* well (every +well is emitted, matched or not), carrying the ids of the associated wells from +other agencies and the linked OSE POD(s), plus the method and a confidence. + +The algorithm is pure and deterministic: given the same sites it always returns +the same crosswalk, independent of input order. It depends only on the standard +library so it is cheap to unit-test in isolation from the connectors. +""" + +from __future__ import annotations + +import hashlib +import math +from collections import defaultdict +from typing import Any, Iterable, Optional + +# Default correlation parameters. Latitude/longitude are not trustworthy to the +# meter for these sources, so the spatial thresholds are deliberately generous; +# tighten them via the keyword arguments to correlate_wells. +DEFAULT_MAX_LINK_DISTANCE_M = 150.0 +DEFAULT_DEPTH_TOLERANCE_FT = 50.0 +DEFAULT_ELEVATION_TOLERANCE_FT = 20.0 +DEFAULT_POD_LINK_DISTANCE_M = 150.0 +DEFAULT_POD_SOURCE = "NMOSEPOD" + +# Confidence score for an explicit id reference (the agencies assert the match). +_CONF_EXPLICIT = 0.95 + +# Spatial confidence is built up from corroborating attributes rather than being +# a fixed value: proximity alone is weak, and each independent attribute that +# agrees (depth, elevation, name/id tokens) raises confidence toward — but never +# to — the explicit level. A single disagreeing attribute (depth or elevation +# out of tolerance) rejects the pair outright: same location + clearly different +# construction means stacked/adjacent wells, not one well. +_CONF_SPATIAL_BASE = 0.35 # proximity within threshold, nothing else known +_CONF_BONUS_CLOSE = 0.10 # within 1/3 of the distance threshold +_CONF_BONUS_DEPTH = 0.25 # reported well depths agree +_CONF_BONUS_ELEVATION = 0.12 # reported surface elevations agree +_CONF_BONUS_NAME = 0.10 # name / id share a token +_CONF_SPATIAL_CAP = 0.9 # spatial evidence never reaches explicit (0.95) + +# Prefixes stripped when normalizing an identifier for cross-reference matching +# (e.g. NWIS ids are "USGS-08313000" but NMBGMR records the bare "08313000"). +_ID_PREFIXES = ("USGS-", "USGS:", "NWIS-", "NWIS:") + +_EARTH_RADIUS_M = 6371008.8 + +CORRELATION_METHOD_DESCRIPTION = ( + "Cross-agency well correlation: wells are linked by explicit id references " + "(usgs_site_id / alternate_site_id, confidence 0.95) and by spatial " + "proximity between agencies (default <=150 m). Spatial confidence starts low " + "(0.35) and is raised by each corroborating attribute that agrees — well " + "depth (<=50 ft), surface elevation (<=20 ft), a shared name/id token, and " + "very close proximity — up to a 0.9 cap; a depth or elevation that disagrees " + "beyond tolerance rejects the pair. Linked wells are unioned into components " + "(one inferred physical well) and associated with any OSE POD they contain " + "or that lies within the spatial threshold. Coordinates are approximate." +) + + +# --------------------------------------------------------------------------- +# Small helpers +# --------------------------------------------------------------------------- +def _num(value: Any) -> Optional[float]: + """Coerce to float, or None when missing/unparseable.""" + if value is None: + return None + try: + return float(value) + except TypeError, ValueError: + return None + + +def _normalize_id(value: Any) -> str: + """Uppercase, strip surrounding whitespace, and collapse internal spaces. + Returns "" for missing values.""" + if value is None: + return "" + s = str(value).strip().upper() + return " ".join(s.split()) + + +def _match_keys(value: Any) -> set[str]: + """The set of normalized forms an identifier can be matched by: the + normalized value itself and, if it carries a known agency prefix, the value + with that prefix removed. Empty/degenerate ids yield no keys.""" + norm = _normalize_id(value) + if not norm: + return set() + keys = {norm} + for prefix in _ID_PREFIXES: + p = prefix.upper() + if norm.startswith(p): + stripped = norm[len(p) :].strip() + if stripped: + keys.add(stripped) + return keys + + +def _parse_references(value: Any) -> list[str]: + """Split a free-form alternate-id field into individual reference tokens. + Agencies pack multiple ids into one field with assorted delimiters.""" + if value is None: + return [] + text = str(value) + for sep in (",", ";", "|", "/", "\n", "\t"): + text = text.replace(sep, " ") + return [tok for tok in (t.strip() for t in text.split(" ")) if tok] + + +def _name_tokens(*values: Any) -> set[str]: + """Alphanumeric tokens (len >= 3) drawn from a well's name and id, for + cheap fuzzy corroboration. Short tokens are dropped because they collide too + easily to be evidence.""" + tokens: set[str] = set() + for value in values: + if value is None: + continue + cur = "" + for ch in str(value).upper(): + if ch.isalnum(): + cur += ch + else: + if len(cur) >= 3: + tokens.add(cur) + cur = "" + if len(cur) >= 3: + tokens.add(cur) + return tokens + + +def haversine_m(lat1: float, lon1: float, lat2: float, lon2: float) -> float: + """Great-circle distance in meters between two WGS84 points.""" + p1 = math.radians(lat1) + p2 = math.radians(lat2) + dphi = math.radians(lat2 - lat1) + dlam = math.radians(lon2 - lon1) + a = math.sin(dphi / 2) ** 2 + math.cos(p1) * math.cos(p2) * math.sin(dlam / 2) ** 2 + return 2 * _EARTH_RADIUS_M * math.asin(math.sqrt(a)) + + +def site_key(site: dict) -> str: + """Stable "SOURCE:id" key for a site (matches the OGC feature id format).""" + source = _normalize_id(site.get("source")) + rid = str(site.get("id") if site.get("id") is not None else "").strip() + return f"{source}:{rid}" + + +# --------------------------------------------------------------------------- +# Union-Find +# --------------------------------------------------------------------------- +class _UnionFind: + def __init__(self) -> None: + self._parent: dict[str, str] = {} + + def add(self, x: str) -> None: + self._parent.setdefault(x, x) + + def find(self, x: str) -> str: + self.add(x) + root = x + while self._parent[root] != root: + root = self._parent[root] + # Path compression. + while self._parent[x] != root: + self._parent[x], x = root, self._parent[x] + return root + + def union(self, a: str, b: str) -> None: + ra, rb = self.find(a), self.find(b) + if ra == rb: + return + # Deterministic root choice (smaller key wins) so component ids are + # independent of input order. + if rb < ra: + ra, rb = rb, ra + self._parent[rb] = ra + + def components(self) -> dict[str, list[str]]: + groups: dict[str, list[str]] = defaultdict(list) + for x in self._parent: + groups[self.find(x)].append(x) + return groups + + +# --------------------------------------------------------------------------- +# Spatial index +# --------------------------------------------------------------------------- +class _GridIndex: + """Uniform lat/lon grid for near-neighbor candidate search. Cell size is + derived from the link distance so only the 3x3 block around a point needs to + be scanned. Avoids the O(n^2) all-pairs comparison on large site sets.""" + + def __init__(self, distance_m: float) -> None: + # Latitude degrees per meter is ~constant; use it (with a safety factor) + # as the cell size so a cell spans at least the link distance in both + # axes at NM latitudes. + deg_per_m = 1.0 / 111_320.0 + self._cell = max(distance_m * deg_per_m, 1e-6) + self._cells: dict[tuple[int, int], list[int]] = defaultdict(list) + + def _cell_of(self, lat: float, lon: float) -> tuple[int, int]: + return (int(math.floor(lat / self._cell)), int(math.floor(lon / self._cell))) + + def add(self, idx: int, lat: float, lon: float) -> None: + self._cells[self._cell_of(lat, lon)].append(idx) + + def neighbors(self, lat: float, lon: float) -> Iterable[int]: + ci, cj = self._cell_of(lat, lon) + for di in (-1, 0, 1): + for dj in (-1, 0, 1): + yield from self._cells.get((ci + di, cj + dj), ()) + + +# --------------------------------------------------------------------------- +# Core +# --------------------------------------------------------------------------- +def _explicit_links(sites: list[dict]) -> list[tuple[int, int, str, float]]: + """Edges asserted by explicit id references. For each site, its usgs_site_id + and every alternate_site_id token is matched against an index of the ids + (and USGS ids) of all sites; a hit links the two sites. Self-links and + within-source links contribute nothing new but are harmless (dropped). + + Returns ``(i, j, "explicit", confidence)`` tuples.""" + # Map every normalized match key -> the site indices reachable by it. + index: dict[str, set[int]] = defaultdict(set) + for i, s in enumerate(sites): + for key in _match_keys(s.get("id")): + index[key].add(i) + for key in _match_keys(s.get("usgs_site_id")): + index[key].add(i) + + edges: list[tuple[int, int, str, float]] = [] + for i, s in enumerate(sites): + refs: list[str] = [] + if s.get("usgs_site_id"): + refs.append(str(s["usgs_site_id"])) + refs.extend(_parse_references(s.get("alternate_site_id"))) + + seen: set[str] = set() + for ref in refs: + for key in _match_keys(ref): + if key in seen: + continue + seen.add(key) + for j in index.get(key, ()): # noqa: E501 + if j != i: + edges.append((i, j, "explicit", _CONF_EXPLICIT)) + return edges + + +def _spatial_confidence( + distance_m: float, + max_link_distance_m: float, + si: dict, + sj: dict, + depth_tolerance_ft: float, + elevation_tolerance_ft: float, +) -> Optional[float]: + """Confidence (0..1) that two nearby cross-agency wells are the same physical + well, or ``None`` if an attribute disagrees beyond tolerance (reject). + + Proximity alone is weak evidence because coordinates are approximate; each + independent attribute that agrees adds confidence. Depth or elevation present + on *both* sites and out of tolerance is disqualifying.""" + conf = _CONF_SPATIAL_BASE + if distance_m <= max_link_distance_m / 3.0: + conf += _CONF_BONUS_CLOSE + + depth1, depth2 = _num(si.get("well_depth")), _num(sj.get("well_depth")) + if depth1 is not None and depth2 is not None: + if abs(depth1 - depth2) > depth_tolerance_ft: + return None # same spot, clearly different construction -> reject + conf += _CONF_BONUS_DEPTH + + elev1, elev2 = _num(si.get("elevation")), _num(sj.get("elevation")) + if elev1 is not None and elev2 is not None: + if abs(elev1 - elev2) > elevation_tolerance_ft: + return None + conf += _CONF_BONUS_ELEVATION + + tokens_i = _name_tokens(si.get("name"), si.get("id")) + tokens_j = _name_tokens(sj.get("name"), sj.get("id")) + if tokens_i & tokens_j: + conf += _CONF_BONUS_NAME + + return min(conf, _CONF_SPATIAL_CAP) + + +def _spatial_links( + sites: list[dict], + keys: list[str], + max_link_distance_m: float, + depth_tolerance_ft: float, + elevation_tolerance_ft: float, +) -> list[tuple[int, int, str, float]]: + """Candidate edges from spatial proximity between *different* sources, scored + by attribute agreement (see :func:`_spatial_confidence`). Same-source pairs + are never linked spatially (an agency's own two wells are distinct sites, not + duplicates). Returns ``(i, j, "spatial", confidence)`` tuples. + + Spatial links are restricted to **mutual nearest neighbors within each agency + pair**: well *i* links to well *j* of agency *B* only if *j* is *i*'s best + (highest-confidence, then nearest) *B* candidate **and** *i* is *j*'s best + candidate in *i*'s agency. Proximity is not transitive, so without this a + dense monitoring area collapses into one giant cluster as wells chain + neighbor-to-neighbor; mutual best-match caps each well at one link per other + agency and prevents a well from absorbing a whole neighborhood.""" + located: list[int] = [] + grid = _GridIndex(max_link_distance_m) + src_of: dict[int, str] = {} + for i, s in enumerate(sites): + lat, lon = _num(s.get("latitude")), _num(s.get("longitude")) + if lat is None or lon is None: + continue + grid.add(i, lat, lon) + located.append(i) + src_of[i] = _normalize_id(s.get("source")) + + # best[i][other_source] = (j, confidence, distance) — i's top candidate in + # each other agency. + best: dict[int, dict[str, tuple[int, float, float]]] = defaultdict(dict) + for i in located: + si = sites[i] + lat1, lon1 = _num(si["latitude"]), _num(si["longitude"]) + src1 = src_of[i] + for j in grid.neighbors(lat1, lon1): + if j == i or src_of.get(j) == src1: + continue + sj = sites[j] + lat2, lon2 = _num(sj["latitude"]), _num(sj["longitude"]) + dist = haversine_m(lat1, lon1, lat2, lon2) + if dist > max_link_distance_m: + continue + conf = _spatial_confidence( + dist, + max_link_distance_m, + si, + sj, + depth_tolerance_ft, + elevation_tolerance_ft, + ) + if conf is None: + continue + srcj = src_of[j] + cur = best[i].get(srcj) + # Higher confidence wins; ties broken by distance, then site key for + # determinism (independent of neighbor iteration order). + cand = (j, conf, dist) + if cur is None or _better_candidate(cand, cur, keys): + best[i][srcj] = cand + + edges: list[tuple[int, int, str, float]] = [] + seen: set[tuple[int, int]] = set() + for i in located: + src1 = src_of[i] + for srcj, (j, conf, _dist) in best[i].items(): + # Mutual: j's best candidate in i's agency must be i. + back = best.get(j, {}).get(src1) + if back is None or back[0] != i: + continue + pair = (i, j) if i < j else (j, i) + if pair in seen: + continue + seen.add(pair) + edges.append((pair[0], pair[1], "spatial", conf)) + return edges + + +def _better_candidate(cand: tuple, cur: tuple, keys: list[str]) -> bool: + """True if spatial candidate *cand* (j, conf, dist) should replace *cur*. + Higher confidence, then shorter distance, then smaller site key — fully + deterministic regardless of iteration order.""" + j_c, conf_c, dist_c = cand + j_p, conf_p, dist_p = cur + if conf_c != conf_p: + return conf_c > conf_p + if dist_c != dist_p: + return dist_c < dist_p + return keys[j_c] < keys[j_p] + + +def _component_id(member_keys: list[str]) -> str: + """Stable short id for a component from its sorted member keys.""" + joined = "|".join(sorted(member_keys)) + return "wc_" + hashlib.sha1(joined.encode("utf-8")).hexdigest()[:12] + + +def correlate_wells( + sites: Iterable[dict], + *, + max_link_distance_m: float = DEFAULT_MAX_LINK_DISTANCE_M, + depth_tolerance_ft: float = DEFAULT_DEPTH_TOLERANCE_FT, + elevation_tolerance_ft: float = DEFAULT_ELEVATION_TOLERANCE_FT, + pod_link_distance_m: float = DEFAULT_POD_LINK_DISTANCE_M, + pod_source: str = DEFAULT_POD_SOURCE, +) -> list[dict]: + """Correlate wells across agencies and link each to OSE POD(s). + + Parameters + ---------- + sites: + Site dicts with at least ``source`` and ``id``; optionally ``name``, + ``latitude``, ``longitude``, ``elevation``, ``well_depth``, + ``usgs_site_id`` and ``alternate_site_id``. + max_link_distance_m: + Max great-circle distance for a cross-agency spatial link. + depth_tolerance_ft: + Max well-depth difference for depth to corroborate a spatial link (and + beyond which a spatial candidate is rejected). + elevation_tolerance_ft: + Max surface-elevation difference for elevation to corroborate a spatial + link (and beyond which a spatial candidate is rejected). + pod_link_distance_m: + Max distance for associating an OSE POD with a component spatially. + pod_source: + The ``source`` value identifying OSE POD sites (default ``"NMOSEPOD"``). + + Returns + ------- + list[dict] + One crosswalk record per input well. Duplicate ``(source, id)`` inputs + are collapsed to the first occurrence. Records are sorted by well key + for stable output. Each record has: + + ``source``, ``id``, ``name``, ``latitude``, ``longitude``, + ``cluster_id``, ``cluster_size``, ``n_agencies``, + ``linked_site_ids`` (list of "SOURCE:id" for other wells in the + cluster), ``linked_by_agency`` (dict agency -> [ids]), + ``ose_pod_ids`` (list), ``ose_pod_link_method`` + (``explicit`` / ``spatial`` / ``none``), ``match_method`` + (``explicit`` / ``spatial`` / ``mixed`` / ``unmatched``), + ``match_confidence`` (0..1), ``is_ose_pod`` (bool). + """ + # De-duplicate on (source, id), keeping the first occurrence, so a site that + # appears twice does not inflate a cluster or the output. + unique: list[dict] = [] + seen_keys: set[str] = set() + for s in sites: + k = site_key(s) + if k in seen_keys: + continue + seen_keys.add(k) + unique.append(s) + + sites = unique + keys = [site_key(s) for s in sites] + pod_norm = _normalize_id(pod_source) + + uf = _UnionFind() + for k in keys: + uf.add(k) + + # Track, per unordered site-index pair, the strongest evidence linking them + # as (method, confidence). explicit outranks spatial; among same method the + # higher confidence wins. + edge_evidence: dict[tuple[int, int], tuple[str, float]] = {} + rank = {"explicit": 2, "spatial": 1} + + def _record_edge(i: int, j: int, method: str, conf: float) -> None: + pair = (i, j) if i < j else (j, i) + prev = edge_evidence.get(pair) + if ( + prev is None + or rank[method] > rank[prev[0]] + or (rank[method] == rank[prev[0]] and conf > prev[1]) + ): + edge_evidence[pair] = (method, conf) + + for i, j, method, conf in _explicit_links(sites): + _record_edge(i, j, method, conf) + for i, j, method, conf in _spatial_links( + sites, + keys, + max_link_distance_m, + depth_tolerance_ft, + elevation_tolerance_ft, + ): + _record_edge(i, j, method, conf) + + for i, j in edge_evidence: + uf.union(keys[i], keys[j]) + + # Component membership as site indices. + key_to_index = {k: i for i, k in enumerate(keys)} + comp_members: dict[str, list[int]] = defaultdict(list) + for k in keys: + comp_members[uf.find(k)].append(key_to_index[k]) + + # Per-component: which methods appear among its internal edges and the best + # spatial-edge confidence (spatial confidence is graduated by attribute + # agreement, so the cluster reports its strongest spatial link). + comp_methods: dict[str, set[str]] = defaultdict(set) + comp_best_spatial: dict[str, float] = defaultdict(float) + for (i, j), (method, conf) in edge_evidence.items(): + root = uf.find(keys[i]) + comp_methods[root].add(method) + if method == "spatial": + comp_best_spatial[root] = max(comp_best_spatial[root], conf) + + results: list[dict] = [] + for root, member_idxs in comp_members.items(): + member_keys = [keys[i] for i in member_idxs] + cluster_id = _component_id(member_keys) + methods = comp_methods.get(root, set()) + + # OSE PODs in this component (explicit membership). + pod_ids_explicit = [ + str(sites[i].get("id")) + for i in member_idxs + if _normalize_id(sites[i].get("source")) == pod_norm + ] + + agencies = {_normalize_id(sites[i].get("source")) for i in member_idxs} + + if "explicit" in methods and "spatial" in methods: + cluster_method = "mixed" + confidence = _CONF_EXPLICIT + elif "explicit" in methods: + cluster_method = "explicit" + confidence = _CONF_EXPLICIT + elif "spatial" in methods: + cluster_method = "spatial" + confidence = round(comp_best_spatial[root], 3) + else: + cluster_method = "unmatched" + confidence = 0.0 + + for i in member_idxs: + s = sites[i] + k = keys[i] + others = [ok for ok in member_keys if ok != k] + by_agency: dict[str, list[str]] = defaultdict(list) + for oi in member_idxs: + if oi == i: + continue + by_agency[_normalize_id(sites[oi].get("source"))].append( + str(sites[oi].get("id")) + ) + + is_pod = _normalize_id(s.get("source")) == pod_norm + # POD linkage from this well's perspective. + pod_ids = list(pod_ids_explicit) + pod_method = ( + "explicit" if pod_ids and not (is_pod and len(pod_ids) == 1) else "none" + ) + if is_pod: + # A POD well is trivially "its own" POD; report others in cluster. + pod_ids = [p for p in pod_ids_explicit if p != str(s.get("id"))] + pod_method = "explicit" if pod_ids else "none" + + results.append( + { + "source": s.get("source"), + "id": s.get("id"), + "name": s.get("name"), + "latitude": _num(s.get("latitude")), + "longitude": _num(s.get("longitude")), + "elevation": _num(s.get("elevation")), + "well_depth": _num(s.get("well_depth")), + "cluster_id": cluster_id, + "cluster_size": len(member_idxs), + "n_agencies": len(agencies), + "linked_site_ids": sorted(others), + "linked_by_agency": { + a: sorted(v) for a, v in sorted(by_agency.items()) + }, + "ose_pod_ids": sorted(set(pod_ids)), + "ose_pod_link_method": pod_method, + "match_method": ( + cluster_method if len(member_idxs) > 1 else "unmatched" + ), + "match_confidence": confidence if len(member_idxs) > 1 else 0.0, + "is_ose_pod": is_pod, + } + ) + + # Second pass: spatially associate PODs with components that contain no POD, + # so wells with no explicit POD reference still get a candidate POD. This is + # done after clustering so a POD links to the whole component, not one member. + _link_pods_spatially(sites, keys, uf, results, pod_norm, pod_link_distance_m) + + results.sort(key=lambda r: f"{_normalize_id(r['source'])}:{r['id']}") + return results + + +def _link_pods_spatially( + sites: list[dict], + keys: list[str], + uf: _UnionFind, + results: list[dict], + pod_norm: str, + pod_link_distance_m: float, +) -> None: + """For components that contain no explicit OSE POD, attach the nearest POD(s) + within ``pod_link_distance_m`` of any member as *candidate* POD links. POD + coordinates are approximate, so this is advisory (method ``spatial``).""" + pod_idxs = [ + i for i, s in enumerate(sites) if _normalize_id(s.get("source")) == pod_norm + ] + if not pod_idxs: + return + + grid = _GridIndex(pod_link_distance_m) + for i in pod_idxs: + lat, lon = _num(sites[i].get("latitude")), _num(sites[i].get("longitude")) + if lat is None or lon is None: + continue + grid.add(i, lat, lon) + + # Which components already have an explicit POD? + comp_has_pod: dict[str, bool] = defaultdict(bool) + for i in pod_idxs: + comp_has_pod[uf.find(keys[i])] = True + + # Nearest POD per component (that lacks one), scanning each member's cell. + comp_best: dict[str, tuple[float, str]] = {} + result_by_key = {f"{_normalize_id(r['source'])}:{r['id']}": r for r in results} + for idx, k in enumerate(keys): + root = uf.find(k) + if comp_has_pod.get(root): + continue + lat, lon = _num(sites[idx].get("latitude")), _num(sites[idx].get("longitude")) + if lat is None or lon is None: + continue + for pj in grid.neighbors(lat, lon): + plat = _num(sites[pj].get("latitude")) + plon = _num(sites[pj].get("longitude")) + d = haversine_m(lat, lon, plat, plon) + if d > pod_link_distance_m: + continue + pod_id = str(sites[pj].get("id")) + best = comp_best.get(root) + if best is None or d < best[0] or (d == best[0] and pod_id < best[1]): + comp_best[root] = (d, pod_id) + + for idx, k in enumerate(keys): + root = uf.find(k) + best = comp_best.get(root) + if best is None: + continue + r = result_by_key.get(k) + if r is None: + continue + r["ose_pod_ids"] = [best[1]] + r["ose_pod_link_method"] = "spatial" + + +# ============= EOF ============================================= diff --git a/orchestration/assets/products.py b/orchestration/assets/products.py index cddb3451..c65ec7ac 100644 --- a/orchestration/assets/products.py +++ b/orchestration/assets/products.py @@ -62,6 +62,7 @@ the bulk of the remaining redundant API pulls for analyte products; it touches DIE core, not this asset graph, so it is intentionally out of scope here. """ + import tempfile import traceback from collections import namedtuple @@ -90,10 +91,11 @@ dump_water_type_collection, dump_waterlevel_change_collection, dump_waterlevel_status_collection, + dump_well_correlation_collection, dump_wqi_collection, ) from backend.record import ParameterRecord, SiteRecord, SummaryRecord -from backend.unifier import unify_source_both +from backend.unifier import collect_sites, unify_source_both from orchestration.logging_bridge import forward_die_logs from orchestration.resources.die_config import DIEConfigResource from orchestration.resources.gcs import GCSResource @@ -188,10 +190,28 @@ def _param_source_keys(product: dict, parameter: str) -> list[str]: return agencies +# Products that are not built on the per-parameter shared-source graph. The well +# correlation product needs the *sites* of every source (parameter-independent, +# including the OSE POD source which has no parameter data), so its combine +# gathers sites itself rather than reading per-parameter shared source assets. +_STANDALONE_OUTPUT_TYPES = {"ogc_well_correlation"} + + +def is_standalone(product: dict) -> bool: + """True for products whose combine gathers its own data (no shared source + assets, no parameter cohort).""" + return product.get("output_type") in _STANDALONE_OUTPUT_TYPES + + def product_source_specs(product: dict) -> list[SourceSpec]: """Every shared source asset this product depends on, one per (parameter, source) pair. ``scope`` is constant for a product; the parameter - and source vary. Returned in a stable order.""" + and source vary. Returned in a stable order. + + Standalone products (see :func:`is_standalone`) depend on no shared source + asset and return an empty list.""" + if is_standalone(product): + return [] scope = _spatial_scope(product) specs: list[SourceSpec] = [] for param in _product_params(product): @@ -400,45 +420,71 @@ def _combine_asset( dump_summary_collection(str(out), records, meta) elif output_type == "ogc_data_density": dump_data_density_collection( - str(out), all_sites, all_timeseries, meta, + str(out), + all_sites, + all_timeseries, + meta, parameter_name=product.get("parameter"), ) elif output_type == "ogc_waterlevel_change": dump_waterlevel_change_collection( - str(out), all_sites, all_timeseries, meta, + str(out), + all_sites, + all_timeseries, + meta, window_years=float(product.get("window_years", 5)), ) elif output_type == "ogc_waterlevel_trend": # all_sites/all_timeseries are index-aligned payload dicts (see # source asset); consumed as dicts to keep memory bounded. dump_trend_collection( - str(out), all_sites, all_timeseries, meta, - slope_units="ft/year", reducer="min", + str(out), + all_sites, + all_timeseries, + meta, + slope_units="ft/year", + reducer="min", ) elif output_type == "ogc_analyte_trend": dump_trend_collection( - str(out), all_sites, all_timeseries, meta, - slope_units="mg/L/year", reducer="mean", + str(out), + all_sites, + all_timeseries, + meta, + slope_units="mg/L/year", + reducer="mean", method=ANALYTE_TREND_METHOD_DESCRIPTION, parameter_name=product.get("parameter"), ) elif output_type == "ogc_waterlevel_status": dump_waterlevel_status_collection( - str(out), all_sites, all_timeseries, meta, + str(out), + all_sites, + all_timeseries, + meta, ) elif output_type == "ogc_seasonal_amplitude": dump_seasonal_amplitude_collection( - str(out), all_sites, all_timeseries, meta, + str(out), + all_sites, + all_timeseries, + meta, min_days_per_year=int(product.get("min_days_per_year", 4)), ) elif output_type == "ogc_depletion_projection": dump_depletion_projection_collection( - str(out), all_sites, all_timeseries, meta, + str(out), + all_sites, + all_timeseries, + meta, ) elif output_type == "ogc_monitoring_recency": run_date = datetime.now(timezone.utc).strftime("%Y-%m-%d") dump_monitoring_recency_collection( - str(out), all_sites, all_timeseries, meta, + str(out), + all_sites, + all_timeseries, + meta, run_date=run_date, stale_days=int(product.get("stale_days", 365)), ) @@ -476,6 +522,103 @@ def _combine_asset( return _combine_asset +def _build_correlation_combine_asset(product: dict, group: str) -> dg.AssetsDefinition: + """Build the combine asset for the well-correlation product (keyed + ``[product_id]``). + + Unlike the parameter products, this asset depends on **no** shared source + assets: correlation needs the location of every well from every source + (including OSE PODs, which carry no parameter data). It gathers all sites + itself via ``collect_sites`` (sites-only over ``all_site_sources``), runs the + cross-agency correlation, writes the OGC GeoJSON, and uploads it to GCS.""" + pid = product["id"] + + @dg.asset( + key=dg.AssetKey(pid), + group_name=group, + description=( + f"**{product.get('title', pid)}** — standalone combine asset " + f"(`ogc_well_correlation`). {product.get('description', '').rstrip('.')}. " + f"Gathers sites from every source (sites-only, incl. OSE PODs), " + f"correlates wells across agencies, links each to an OSE POD, and " + f"uploads the OGC GeoJSON to GCS. Depends on no shared source asset." + ), + check_specs=[dg.AssetCheckSpec(name=_CHECK_NAME, asset=dg.AssetKey(pid))], + ) + def _combine_asset( + context: dg.AssetExecutionContext, + die_config: DIEConfigResource, + gcs: GCSResource, + ) -> Iterator[dg.Output | dg.AssetCheckResult]: + error = "" + feature_count = 0 + info: dict = {} + try: + with forward_die_logs(context): + config = die_config.get_config(product) + sites = collect_sites(config) + + meta = { + "id": pid, + "title": product.get("title", pid), + "description": product.get("description", ""), + } + with tempfile.TemporaryDirectory() as tmpdir: + out = Path(tmpdir) / "collection.geojson" + coll = dump_well_correlation_collection( + str(out), + sites, + meta, + max_link_distance_m=_num_opt(product.get("max_link_distance_m")), + depth_tolerance_ft=_num_opt(product.get("depth_tolerance_ft")), + elevation_tolerance_ft=_num_opt( + product.get("elevation_tolerance_ft") + ), + pod_link_distance_m=_num_opt(product.get("pod_link_distance_m")), + ) + feature_count = len(coll.get("features", [])) + info = gcs.upload_product(str(out), pid) + except Exception: + error = traceback.format_exc() + context.log.error(f"Well correlation combine failed for {pid}:\n{error}") + + metadata: dict = {"error": error} + if info: + metadata.update( + { + "feature_count": dg.MetadataValue.int( + info.get("feature_count", feature_count) + ), + "latest_uri": dg.MetadataValue.url(info.get("latest_uri", "")), + "skipped_unchanged": dg.MetadataValue.bool( + bool(info.get("skipped")) + ), + } + ) + if info.get("dated_uri"): + metadata["dated_uri"] = dg.MetadataValue.url(info["dated_uri"]) + + yield dg.Output(None, metadata=metadata) + yield dg.AssetCheckResult( + asset_key=dg.AssetKey(pid), + check_name=_CHECK_NAME, + passed=error == "" and feature_count > 0, + severity=dg.AssetCheckSeverity.WARN, + metadata={ + "feature_count": feature_count, + "error": error or ("no features" if feature_count == 0 else ""), + }, + ) + + return _combine_asset + + +def _num_opt(value): + """None passes through; otherwise coerce to float (products.yaml overrides + for the correlation thresholds are optional).""" + return None if value is None else float(value) + + def _geojson_to_geopackage(geojson_path: Path, layer_name: str, out_dir: Path): """Convert a GeoJSON file to a GeoPackage whose layer (table) is named *layer_name* (so the published GeoServer layer is named *layer_name*). @@ -577,8 +720,13 @@ def build_product_pipeline_assets( publish asset. The shared source assets it consumes (``specs``) are built once by :func:`build_shared_source_asset` in ``definitions.py``, not here, so products sharing a source share one asset. The combine's group follows its - parameter family (waterlevels vs analytes).""" - group = "waterlevels" if product.get("parameter") == WATERLEVELS else "analytes" - combine = _build_combine_asset(product, specs, group) + parameter family (waterlevels vs analytes); standalone products (well + correlation) get their own ``sites`` group and a self-contained combine.""" + if is_standalone(product): + group = "sites" + combine = _build_correlation_combine_asset(product, group) + else: + group = "waterlevels" if product.get("parameter") == WATERLEVELS else "analytes" + combine = _build_combine_asset(product, specs, group) geoserver = _build_geoserver_asset(product, group) return [combine, geoserver] diff --git a/orchestration/config/products.yaml b/orchestration/config/products.yaml index 5f80b555..e4a6b7d9 100644 --- a/orchestration/config/products.yaml +++ b/orchestration/config/products.yaml @@ -267,3 +267,24 @@ products: state: NM sources: exclude: [] + + # One feature per well linking it to the same physical well at other agencies + # and, ultimately, to an OSE POD. Standalone product (no `parameter`): its + # combine gathers sites from EVERY source (sites-only, incl. the OSE POD + # source) and correlates them via backend/well_correlation.py — explicit id + # references (usgs_site_id / alternate_site_id) plus spatial+depth agreement, + # since coordinates are only approximate. Generated quarterly (1st of Jan/Apr/ + # Jul/Oct, 06:00 America/Denver). Correlation thresholds are optional overrides + # (defaults live in backend/well_correlation.py). + - id: nm_well_correlation + output_type: ogc_well_correlation + title: "NM Cross-Agency Well Correlation" + description: "Per-well links to associated wells at other agencies and to OSE PODs, all NM sources" + schedule: "0 6 1 */3 *" + spatial_filter: + state: NM + sources: + exclude: [] + # max_link_distance_m: 150 + # depth_tolerance_ft: 50 + # pod_link_distance_m: 150 diff --git a/orchestration/definitions.py b/orchestration/definitions.py index 04347525..f37cfaad 100644 --- a/orchestration/definitions.py +++ b/orchestration/definitions.py @@ -13,10 +13,12 @@ from orchestration.assets.products import ( build_product_pipeline_assets, build_shared_source_asset, + is_standalone, product_source_specs, shared_source_key, ) + class _TolerantGCSPickleIOManager(GCSPickleIOManager): """GCS pickle IO manager that tolerates a missing source input. @@ -39,7 +41,7 @@ class _TolerantGCSPickleIOManager(GCSPickleIOManager): def load_input(self, context: dg.InputContext): try: return super().load_input(context) - except (NotFound, FileNotFoundError): + except NotFound, FileNotFoundError: context.log.warning( f"Source input {context.asset_key.to_user_string()!r} not found " "in GCS; treating as empty. Run sources_job before the product " @@ -75,6 +77,7 @@ def load_input(self, context: dg.InputContext): "ogc_waterlevel_status", "ogc_seasonal_amplitude", "ogc_depletion_projection", + "ogc_well_correlation", } @@ -141,7 +144,7 @@ def _cron_sort_key(cron: str) -> tuple[int, int]: parts = cron.split() try: return (int(parts[1]), int(parts[0])) - except (IndexError, ValueError): + except IndexError, ValueError: return (6, 0) @@ -159,7 +162,9 @@ def _build_cohorts(products_config: dict, specs_by_pid: dict) -> dict: cohort = cohorts.setdefault(name, {"members": [], "cron": None}) cohort["members"].append(pid) cron = product.get("schedule", "0 6 * * *") - if cohort["cron"] is None or _cron_sort_key(cron) < _cron_sort_key(cohort["cron"]): + if cohort["cron"] is None or _cron_sort_key(cron) < _cron_sort_key( + cohort["cron"] + ): cohort["cron"] = cron return cohorts @@ -213,12 +218,56 @@ def _build_schedules( ] +def _build_standalone_jobs_and_schedules( + products_config: dict, +) -> tuple[dict[str, "UnresolvedAssetJobDefinition"], list[dg.ScheduleDefinition]]: + """Jobs and schedules for standalone products (see ``is_standalone``). + + These products depend on no shared source asset, so they are not part of any + cohort. Each gets its own job (selecting just its combine + geoserver assets) + and its own schedule at the product's cron. The well correlation product runs + quarterly (see products.yaml).""" + jobs: dict[str, "UnresolvedAssetJobDefinition"] = {} + schedules: list[dg.ScheduleDefinition] = [] + for product in _products(products_config): + if not is_standalone(product): + continue + pid = product["id"] + job_name = f"{pid}_job" + jobs[job_name] = dg.define_asset_job( + name=job_name, + selection=dg.AssetSelection.keys( + dg.AssetKey(pid), dg.AssetKey([pid, "geoserver"]) + ), + description=( + f"Materialize the standalone {pid} product: gather all sites, " + f"correlate wells across agencies, publish to GeoServer." + ), + ) + schedules.append( + dg.ScheduleDefinition( + name=f"schedule_{pid}", + job=jobs[job_name], + cron_schedule=product.get("schedule", "0 6 1 */3 *"), + execution_timezone="America/Denver", + ) + ) + return jobs, schedules + + _products_config = _load_products() -_source_assets, _pipeline_assets, _specs_by_pid, _all_specs = _build_graph(_products_config) +_source_assets, _pipeline_assets, _specs_by_pid, _all_specs = _build_graph( + _products_config +) _assets = _source_assets + _pipeline_assets _cohorts = _build_cohorts(_products_config, _specs_by_pid) _cohort_jobs = _build_cohort_jobs(_cohorts, _specs_by_pid) _schedules = _build_schedules(_cohorts, _cohort_jobs) +_standalone_jobs, _standalone_schedules = _build_standalone_jobs_and_schedules( + _products_config +) +_cohort_jobs.update(_standalone_jobs) +_schedules += _standalone_schedules defs = dg.Definitions( assets=_assets, diff --git a/orchestration/resources/die_config.py b/orchestration/resources/die_config.py index 6d43d012..47fff1b6 100644 --- a/orchestration/resources/die_config.py +++ b/orchestration/resources/die_config.py @@ -72,6 +72,8 @@ def get_config(self, product: dict, parameter: Optional[str] = None) -> Config: payload[f"use_source_{s}"] = False config = Config(payload=payload) - config.parameter = parameter or product["parameter"] + # An empty parameter is valid for sites-only flows (e.g. the well + # correlation product), so fall back to "" when the product has none. + config.parameter = parameter or product.get("parameter", "") config.finalize() return config diff --git a/tests/test_persisters/test_ogc_features.py b/tests/test_persisters/test_ogc_features.py index 9de2ac8a..24e5cc52 100644 --- a/tests/test_persisters/test_ogc_features.py +++ b/tests/test_persisters/test_ogc_features.py @@ -11,74 +11,85 @@ def _make_summary_record( - source="nmbgmr_amp", rid="RA-1234", lat=35.0, lon=-106.5, - parameter_name="waterlevels", parameter_units="ft", latest_value=220.0, + source="nmbgmr_amp", + rid="RA-1234", + lat=35.0, + lon=-106.5, + parameter_name="waterlevels", + parameter_units="ft", + latest_value=220.0, latest_units="ft", ): - return SummaryRecord({ - "source": source, - "id": rid, - "name": "Test Well", - "usgs_site_id": "", - "alternate_site_id": "", - "latitude": lat, - "longitude": lon, - "horizontal_datum": "WGS84", - "elevation": 1650.0, - "elevation_units": "ft", - "well_depth": None, - "well_depth_units": "ft", - "parameter_name": parameter_name, - "parameter_units": parameter_units, - "nrecords": 10, - "min": 200.0, - "max": 250.0, - "mean": 225.0, - "earliest_date": "1990-01-01", - "earliest_time": "00:00:00", - "earliest_value": 200.0, - "earliest_units": "ft", - "latest_date": "2024-01-01", - "latest_time": "00:00:00", - "latest_value": latest_value, - "latest_units": latest_units, - }) + return SummaryRecord( + { + "source": source, + "id": rid, + "name": "Test Well", + "usgs_site_id": "", + "alternate_site_id": "", + "latitude": lat, + "longitude": lon, + "horizontal_datum": "WGS84", + "elevation": 1650.0, + "elevation_units": "ft", + "well_depth": None, + "well_depth_units": "ft", + "parameter_name": parameter_name, + "parameter_units": parameter_units, + "nrecords": 10, + "min": 200.0, + "max": 250.0, + "mean": 225.0, + "earliest_date": "1990-01-01", + "earliest_time": "00:00:00", + "earliest_value": 200.0, + "earliest_units": "ft", + "latest_date": "2024-01-01", + "latest_time": "00:00:00", + "latest_value": latest_value, + "latest_units": latest_units, + } + ) def _make_site_record(source="nmbgmr_amp", rid="RA-1234", lat=35.0, lon=-106.5): - return SiteRecord({ - "source": source, - "id": rid, - "name": "Test Well", - "latitude": lat, - "longitude": lon, - "elevation": 1650.0, - "elevation_units": "ft", - "horizontal_datum": "WGS84", - "vertical_datum": "", - "usgs_site_id": "", - "alternate_site_id": "", - "formation": "", - "aquifer": "", - "well_depth": None, - "well_depth_units": "ft", - }) + return SiteRecord( + { + "source": source, + "id": rid, + "name": "Test Well", + "latitude": lat, + "longitude": lon, + "elevation": 1650.0, + "elevation_units": "ft", + "horizontal_datum": "WGS84", + "vertical_datum": "", + "usgs_site_id": "", + "alternate_site_id": "", + "formation": "", + "aquifer": "", + "well_depth": None, + "well_depth_units": "ft", + } + ) def _make_wl_record(source="nmbgmr_amp", rid="RA-1234", date="2024-01-15", value=212.4): - return ParameterRecord({ - "source": source, - "id": rid, - "parameter_name": "waterlevels", - "parameter_value": value, - "parameter_units": "ft", - "date_measured": date, - "time_measured": "00:00:00", - "source_parameter_name": "depth_to_water", - "source_parameter_units": "ft", - "conversion_factor": 1.0, - "record_type": "waterlevels", - }) + return ParameterRecord( + { + "source": source, + "id": rid, + "parameter_name": "waterlevels", + "parameter_value": value, + "parameter_units": "ft", + "date_measured": date, + "time_measured": "00:00:00", + "source_parameter_name": "depth_to_water", + "source_parameter_units": "ft", + "conversion_factor": 1.0, + "record_type": "waterlevels", + } + ) class TestDumpSummaryCollection: @@ -150,8 +161,11 @@ def test_tds_class_boundaries(self, tmp_path): ] records = [ _make_summary_record( - rid=f"RA-{i}", parameter_name="tds", parameter_units="mg/L", - latest_value=value, latest_units="mg/L", + rid=f"RA-{i}", + parameter_name="tds", + parameter_units="mg/L", + latest_value=value, + latest_units="mg/L", ) for i, (value, _) in enumerate(cases) ] @@ -215,21 +229,25 @@ def test_ogc_required_fields(self, tmp_path): assert "numberReturned" in result -def _make_chem_record(source, rid, analyte, value, units="mg/L", date="2024-05-01", well_depth=None): - return SummaryRecord({ - "source": source, - "id": rid, - "name": f"Well {rid}", - "latitude": 34.0, - "longitude": -106.0, - "elevation": None, - "well_depth": well_depth, - "well_depth_units": "ft", - "parameter_name": analyte, - "latest_value": value, - "latest_units": units, - "latest_date": date, - }) +def _make_chem_record( + source, rid, analyte, value, units="mg/L", date="2024-05-01", well_depth=None +): + return SummaryRecord( + { + "source": source, + "id": rid, + "name": f"Well {rid}", + "latitude": 34.0, + "longitude": -106.0, + "elevation": None, + "well_depth": well_depth, + "well_depth_units": "ft", + "parameter_name": analyte, + "latest_value": value, + "latest_units": units, + "latest_date": date, + } + ) class TestMajorChemistryCollection: @@ -240,7 +258,9 @@ def test_pivots_analytes_into_one_feature_per_well(self, tmp_path): _make_chem_record("WQP", "W2", "calcium", 55.0), ] out = tmp_path / "mc.geojson" - result = dump_major_chemistry_collection(str(out), records, {"id": "nm_major_chemistry"}) + result = dump_major_chemistry_collection( + str(out), records, {"id": "nm_major_chemistry"} + ) assert result["numberReturned"] == 2 # two distinct wells by_id = {f["id"]: f for f in result["features"]} @@ -259,7 +279,9 @@ def test_pivots_analytes_into_one_feature_per_well(self, tmp_path): def test_geometry_and_required_fields(self, tmp_path): out = tmp_path / "mc.geojson" result = dump_major_chemistry_collection( - str(out), [_make_chem_record("NMBGMR", "W1", "sodium", 30.0)], {"id": "nm_major_chemistry"} + str(out), + [_make_chem_record("NMBGMR", "W1", "sodium", 30.0)], + {"id": "nm_major_chemistry"}, ) assert result["type"] == "FeatureCollection" assert "timeStamp" in result @@ -273,9 +295,14 @@ def test_geometry_and_required_fields(self, tmp_path): # The trend dumper consumes payload dicts directly (no record rebuild). def _trend_site(source="NMBGMR", rid="W1", well_depth=100.0): return { - "source": source, "id": rid, "name": f"Well {rid}", - "latitude": 34.0, "longitude": -106.0, "elevation": None, - "well_depth": well_depth, "well_depth_units": "ft", + "source": source, + "id": rid, + "name": f"Well {rid}", + "latitude": 34.0, + "longitude": -106.0, + "elevation": None, + "well_depth": well_depth, + "well_depth_units": "ft", } @@ -285,16 +312,34 @@ def _trend_obs(date, value): class TestWaterLevelTrendCollection: def test_classifies_trends_and_carries_method(self, tmp_path): - increasing = [_trend_obs(f"{2010 + i}-01-01", 50.0 + 0.5 * i) for i in range(12)] + increasing = [ + _trend_obs(f"{2010 + i}-01-01", 50.0 + 0.5 * i) for i in range(12) + ] stable = [_trend_obs(f"{2010 + i}-01-01", 50.0) for i in range(12)] - sparse = [_trend_obs("2010-01-01", 50.0), _trend_obs("2011-01-01", 51.0), _trend_obs("2012-01-01", 52.0)] + sparse = [ + _trend_obs("2010-01-01", 50.0), + _trend_obs("2011-01-01", 51.0), + _trend_obs("2012-01-01", 52.0), + ] decreasing = [_trend_obs(f"{2010 + i}-01-01", 60.0 - 1.0 * i) for i in range(5)] - sites = [_trend_site(rid="A"), _trend_site(rid="B"), _trend_site("NWIS", "C"), _trend_site("PVACD", "D")] + sites = [ + _trend_site(rid="A"), + _trend_site(rid="B"), + _trend_site("NWIS", "C"), + _trend_site("PVACD", "D"), + ] series = [increasing, stable, sparse, decreasing] out = tmp_path / "tr.geojson" - result = dump_trend_collection(str(out), sites, series, {"id": "nm_waterlevel_trends"}, slope_units="ft/year", reducer="min") + result = dump_trend_collection( + str(out), + sites, + series, + {"id": "nm_waterlevel_trends"}, + slope_units="ft/year", + reducer="min", + ) assert result["numberReturned"] == 4 assert "trend_method" in result and result["trend_method"] @@ -304,18 +349,29 @@ def test_classifies_trends_and_carries_method(self, tmp_path): assert round(by_id["NMBGMR:A"]["slope_per_year"], 2) == 0.5 assert by_id["NMBGMR:B"]["trend_category"] == "stable" assert by_id["NWIS:C"]["trend_category"] == "not enough data" # only 3 records - assert by_id["PVACD:D"]["trend_category"] == "decreasing" # 5 records / 4 yr span + assert ( + by_id["PVACD:D"]["trend_category"] == "decreasing" + ) # 5 records / 4 yr span def test_required_fields_and_geometry(self, tmp_path): sites = [_trend_site(rid="W1")] series = [[_trend_obs("2010-01-01", 50.0)]] out = tmp_path / "tr.geojson" - result = dump_trend_collection(str(out), sites, series, {"id": "nm_waterlevel_trends"}, slope_units="ft/year", reducer="min") + result = dump_trend_collection( + str(out), + sites, + series, + {"id": "nm_waterlevel_trends"}, + slope_units="ft/year", + reducer="min", + ) assert result["type"] == "FeatureCollection" assert "timeStamp" in result feat = result["features"][0] assert feat["geometry"]["coordinates"] == [-106.0, 34.0] - assert feat["properties"]["trend_category"] == "not enough data" # single record + assert ( + feat["properties"]["trend_category"] == "not enough data" + ) # single record class TestWaterLevelTrendDailyMin: @@ -329,7 +385,12 @@ def test_downsamples_to_daily_min(self, tmp_path): ] out = tmp_path / "tr.geojson" result = dump_trend_collection( - str(out), [_trend_site(rid="W1")], [obs], {"id": "nm_waterlevel_trends"}, slope_units="ft/year", reducer="min" + str(out), + [_trend_site(rid="W1")], + [obs], + {"id": "nm_waterlevel_trends"}, + slope_units="ft/year", + reducer="min", ) props = result["features"][0]["properties"] assert props["observation_count"] == 3 @@ -341,13 +402,20 @@ class TestSourceDatastreamLink: def test_trend_feature_includes_source_datastream_link(self, tmp_path): site = _trend_site(source="PVACD", rid="W1") obs = [ - {**_trend_obs(f"{2010 + i}-01-01", 50.0 + 0.5 * i), - "source_datastream_link": "https://st2/FROST-Server/v1.1/Datastreams(42)"} + { + **_trend_obs(f"{2010 + i}-01-01", 50.0 + 0.5 * i), + "source_datastream_link": "https://st2/FROST-Server/v1.1/Datastreams(42)", + } for i in range(12) ] out = tmp_path / "tr.geojson" result = dump_trend_collection( - str(out), [site], [obs], {"id": "nm_waterlevel_trends"}, slope_units="ft/year", reducer="min" + str(out), + [site], + [obs], + {"id": "nm_waterlevel_trends"}, + slope_units="ft/year", + reducer="min", ) assert ( result["features"][0]["properties"]["source_datastream_link"] @@ -359,7 +427,12 @@ def test_trend_feature_omits_link_when_absent(self, tmp_path): obs = [_trend_obs(f"{2010 + i}-01-01", 50.0) for i in range(12)] out = tmp_path / "tr.geojson" result = dump_trend_collection( - str(out), [site], [obs], {"id": "nm_waterlevel_trends"}, slope_units="ft/year", reducer="min" + str(out), + [site], + [obs], + {"id": "nm_waterlevel_trends"}, + slope_units="ft/year", + reducer="min", ) assert "source_datastream_link" not in result["features"][0]["properties"] @@ -367,7 +440,9 @@ def test_summary_feature_includes_source_datastream_link(self, tmp_path): rec = _make_summary_record(source="PVACD", rid="W1") rec.update(source_datastream_link="https://st2/Datastreams(9)") out = tmp_path / "s.geojson" - result = dump_summary_collection(str(out), [rec], {"id": "nm_waterlevels_summary"}) + result = dump_summary_collection( + str(out), [rec], {"id": "nm_waterlevels_summary"} + ) assert ( result["features"][0]["properties"]["source_datastream_link"] == "https://st2/Datastreams(9)" @@ -381,25 +456,39 @@ def test_summary_feature_includes_source_datastream_link(self, tmp_path): def _mcl_record(source, rid, analyte, value, date="2024-05-01"): - return SummaryRecord({ - "source": source, "id": rid, "name": f"Well {rid}", - "latitude": 34.0, "longitude": -106.0, "elevation": None, - "well_depth": None, "well_depth_units": "ft", - "parameter_name": analyte, "latest_value": value, - "latest_date": date, - }) + return SummaryRecord( + { + "source": source, + "id": rid, + "name": f"Well {rid}", + "latitude": 34.0, + "longitude": -106.0, + "elevation": None, + "well_depth": None, + "well_depth_units": "ft", + "parameter_name": analyte, + "latest_value": value, + "latest_date": date, + } + ) class TestMCLExceedanceCollection: def test_flags_exceedances(self, tmp_path): recs = [ - _mcl_record("WQP", "W1", "arsenic", 0.02, date="2024-03-10"), # > 0.01 -> exceeds - _mcl_record("WQP", "W1", "nitrate", 5.0, date="2024-06-20"), # < 10 -> ok + _mcl_record( + "WQP", "W1", "arsenic", 0.02, date="2024-03-10" + ), # > 0.01 -> exceeds + _mcl_record("WQP", "W1", "nitrate", 5.0, date="2024-06-20"), # < 10 -> ok ] - thresholds = {"arsenic": {"mcl": 0.01, "type": "primary"}, - "nitrate": {"mcl": 10.0, "type": "primary"}} + thresholds = { + "arsenic": {"mcl": 0.01, "type": "primary"}, + "nitrate": {"mcl": 10.0, "type": "primary"}, + } out = tmp_path / "mcl.geojson" - result = dump_mcl_exceedance_collection(str(out), recs, {"id": "nm_mcl"}, thresholds) + result = dump_mcl_exceedance_collection( + str(out), recs, {"id": "nm_mcl"}, thresholds + ) props = result["features"][0]["properties"] assert props["arsenic_exceeds"] is True assert props["arsenic_date"] == "2024-03-10" @@ -433,7 +522,12 @@ def test_status_active_and_stale(self, tmp_path): ] out = tmp_path / "rec.geojson" result = dump_monitoring_recency_collection( - str(out), sites, series, {"id": "nm_rec"}, run_date="2024-06-01", stale_days=365 + str(out), + sites, + series, + {"id": "nm_rec"}, + run_date="2024-06-01", + stale_days=365, ) by_id = {f["id"]: f["properties"] for f in result["features"]} assert by_id["PVACD:A"]["status"] == "active" @@ -452,8 +546,13 @@ def test_daily_mean_and_units(self, tmp_path): obs.append(_trend_obs(f"{2010 + i}-01-01", 0.005 + 0.001 * i)) out = tmp_path / "at.geojson" result = dump_trend_collection( - str(out), [site], [obs], {"id": "nm_arsenic_trend"}, - slope_units="mg/L/year", reducer="mean", parameter_name="arsenic", + str(out), + [site], + [obs], + {"id": "nm_arsenic_trend"}, + slope_units="mg/L/year", + reducer="mean", + parameter_name="arsenic", ) props = result["features"][0]["properties"] assert props["parameter_name"] == "arsenic" @@ -517,12 +616,12 @@ class TestWaterTypeCollection: def test_classifies_ca_hco3(self, tmp_path): # Ca-dominant cation, HCO3-dominant anion. recs = [ - _make_chem_record("WQP", "W1", "calcium", 100.0), # 4.99 meq - _make_chem_record("WQP", "W1", "magnesium", 1.0), # 0.08 meq - _make_chem_record("WQP", "W1", "sodium", 1.0), # 0.04 meq + _make_chem_record("WQP", "W1", "calcium", 100.0), # 4.99 meq + _make_chem_record("WQP", "W1", "magnesium", 1.0), # 0.08 meq + _make_chem_record("WQP", "W1", "sodium", 1.0), # 0.04 meq _make_chem_record("WQP", "W1", "bicarbonate", 300.0), # 4.92 meq - _make_chem_record("WQP", "W1", "chloride", 1.0), # 0.03 meq - _make_chem_record("WQP", "W1", "sulfate", 1.0), # 0.02 meq + _make_chem_record("WQP", "W1", "chloride", 1.0), # 0.03 meq + _make_chem_record("WQP", "W1", "sulfate", 1.0), # 0.02 meq ] out = tmp_path / "wt.geojson" result = dump_water_type_collection(str(out), recs, {"id": "nm_water_type"}) @@ -536,9 +635,9 @@ def test_classifies_ca_hco3(self, tmp_path): def test_mixed_when_no_majority(self, tmp_path): # Cations split roughly evenly across Ca / Mg / Na+K -> mixed cation. recs = [ - _make_chem_record("WQP", "W1", "calcium", 20.04), # 1.0 meq + _make_chem_record("WQP", "W1", "calcium", 20.04), # 1.0 meq _make_chem_record("WQP", "W1", "magnesium", 12.15), # 1.0 meq - _make_chem_record("WQP", "W1", "sodium", 22.99), # 1.0 meq + _make_chem_record("WQP", "W1", "sodium", 22.99), # 1.0 meq _make_chem_record("WQP", "W1", "bicarbonate", 305.1), # ~5 meq -> HCO3 ] out = tmp_path / "wt.geojson" @@ -558,8 +657,8 @@ def test_insufficient_when_no_anions(self, tmp_path): def test_charge_balance_reported(self, tmp_path): recs = [ - _make_chem_record("WQP", "W1", "calcium", 20.04), # 1.0 meq cation - _make_chem_record("WQP", "W1", "chloride", 35.45), # 1.0 meq anion + _make_chem_record("WQP", "W1", "calcium", 20.04), # 1.0 meq cation + _make_chem_record("WQP", "W1", "chloride", 35.45), # 1.0 meq anion ] out = tmp_path / "wt.geojson" result = dump_water_type_collection(str(out), recs, {"id": "nm_water_type"}) @@ -611,9 +710,9 @@ def test_change_over_window(self, tmp_path): ) props = result["features"][0]["properties"] assert props["status"] == "ok" - assert props["dtw_end"] == 60.0 # 2020 value - assert props["dtw_start"] == 55.0 # 2015 value - assert props["change_ft"] == 5.0 # deeper -> declining + assert props["dtw_end"] == 60.0 # 2020 value + assert props["dtw_start"] == 55.0 # 2015 value + assert props["change_ft"] == 5.0 # deeper -> declining assert props["direction"] == "declining" assert round(props["actual_window_years"]) == 5 assert "change_method" in result @@ -646,8 +745,11 @@ def test_single_reading_insufficient(self, tmp_path): site = _trend_site(rid="W1") out = tmp_path / "ch.geojson" result = dump_waterlevel_change_collection( - str(out), [site], [[_trend_obs("2020-01-01", 50.0)]], - {"id": "nm_change"}, window_years=5, + str(out), + [site], + [[_trend_obs("2020-01-01", 50.0)]], + {"id": "nm_change"}, + window_years=5, ) props = result["features"][0]["properties"] assert props["status"] == "insufficient" @@ -709,7 +811,7 @@ def test_missing_sodium_insufficient(self, tmp_path): def test_one_divalent_ion_suffices(self, tmp_path): # Missing Mg treated as 0 when Ca present. recs = [ - _make_chem_record("WQP", "W1", "sodium", 229.9), # 10 meq + _make_chem_record("WQP", "W1", "sodium", 229.9), # 10 meq _make_chem_record("WQP", "W1", "calcium", 40.08), # 2 meq ] out = tmp_path / "sar.geojson" @@ -733,7 +835,7 @@ def test_zero_denominator_insufficient(self, tmp_path): class TestIonBalanceCollection: def test_balanced(self, tmp_path): recs = [ - _make_chem_record("WQP", "W1", "calcium", 20.04), # 1 meq cation + _make_chem_record("WQP", "W1", "calcium", 20.04), # 1 meq cation _make_chem_record("WQP", "W1", "chloride", 35.45), # 1 meq anion ] out = tmp_path / "ib.geojson" @@ -747,7 +849,7 @@ def test_balanced(self, tmp_path): def test_suspect_when_large_imbalance(self, tmp_path): recs = [ - _make_chem_record("WQP", "W1", "calcium", 40.08), # 2 meq cations + _make_chem_record("WQP", "W1", "calcium", 40.08), # 2 meq cations _make_chem_record("WQP", "W1", "chloride", 35.45), # 1 meq anions ] out = tmp_path / "ib.geojson" @@ -849,7 +951,9 @@ def test_much_below_normal_when_deepest(self, tmp_path): obs = [_trend_obs(f"{2010 + i}-01-01", 50.0 + i) for i in range(10)] obs.append(_trend_obs("2021-01-01", 100.0)) out = tmp_path / "st.geojson" - result = dump_waterlevel_status_collection(str(out), [site], [obs], {"id": "nm_st"}) + result = dump_waterlevel_status_collection( + str(out), [site], [obs], {"id": "nm_st"} + ) props = result["features"][0]["properties"] assert props["latest_dtw"] == 100.0 assert props["dtw_percentile"] > 90 @@ -863,7 +967,9 @@ def test_much_above_normal_when_shallowest(self, tmp_path): obs = [_trend_obs(f"{2010 + i}-01-01", 50.0 + i) for i in range(10)] obs.append(_trend_obs("2021-01-01", 10.0)) # shallowest ever out = tmp_path / "st.geojson" - result = dump_waterlevel_status_collection(str(out), [site], [obs], {"id": "nm_st"}) + result = dump_waterlevel_status_collection( + str(out), [site], [obs], {"id": "nm_st"} + ) props = result["features"][0]["properties"] assert props["dtw_percentile"] < 10 assert props["status"] == "much above normal" @@ -874,7 +980,9 @@ def test_normal_mid_record(self, tmp_path): obs = [_trend_obs(f"{2010 + i}-01-01", 50.0 + i) for i in range(11)] obs.append(_trend_obs("2021-06-01", 55.0)) out = tmp_path / "st.geojson" - result = dump_waterlevel_status_collection(str(out), [site], [obs], {"id": "nm_st"}) + result = dump_waterlevel_status_collection( + str(out), [site], [obs], {"id": "nm_st"} + ) props = result["features"][0]["properties"] assert props["status"] == "normal" @@ -882,7 +990,9 @@ def test_insufficient_below_min_records(self, tmp_path): site = _trend_site(rid="W1") obs = [_trend_obs(f"{2010 + i}-01-01", 50.0) for i in range(5)] out = tmp_path / "st.geojson" - result = dump_waterlevel_status_collection(str(out), [site], [obs], {"id": "nm_st"}) + result = dump_waterlevel_status_collection( + str(out), [site], [obs], {"id": "nm_st"} + ) props = result["features"][0]["properties"] assert props["status"] == "insufficient" assert props["dtw_percentile"] is None @@ -891,7 +1001,9 @@ def test_insufficient_below_min_records(self, tmp_path): def test_empty_well(self, tmp_path): site = _trend_site(rid="W1") out = tmp_path / "st.geojson" - result = dump_waterlevel_status_collection(str(out), [site], [[]], {"id": "nm_st"}) + result = dump_waterlevel_status_collection( + str(out), [site], [[]], {"id": "nm_st"} + ) props = result["features"][0]["properties"] assert props["status"] == "insufficient" assert props["latest_dtw"] is None @@ -946,7 +1058,9 @@ def test_projects_declining_well(self, tmp_path): site = _trend_site(rid="W1", well_depth=100.0) obs = [_trend_obs(f"{2010 + i}-01-01", 50.0 + i) for i in range(12)] out = tmp_path / "dp.geojson" - result = dump_depletion_projection_collection(str(out), [site], [obs], {"id": "nm_dp"}) + result = dump_depletion_projection_collection( + str(out), [site], [obs], {"id": "nm_dp"} + ) props = result["features"][0]["properties"] assert props["status"] == "projected" assert props["trend_category"] == "increasing" @@ -961,7 +1075,9 @@ def test_not_declining(self, tmp_path): site = _trend_site(rid="W1", well_depth=100.0) obs = [_trend_obs(f"{2010 + i}-01-01", 60.0 - i) for i in range(12)] out = tmp_path / "dp.geojson" - result = dump_depletion_projection_collection(str(out), [site], [obs], {"id": "nm_dp"}) + result = dump_depletion_projection_collection( + str(out), [site], [obs], {"id": "nm_dp"} + ) props = result["features"][0]["properties"] assert props["status"] == "not declining" assert props["years_to_depletion"] is None @@ -970,16 +1086,22 @@ def test_no_well_depth(self, tmp_path): site = _trend_site(rid="W1", well_depth=None) obs = [_trend_obs(f"{2010 + i}-01-01", 50.0 + i) for i in range(12)] out = tmp_path / "dp.geojson" - result = dump_depletion_projection_collection(str(out), [site], [obs], {"id": "nm_dp"}) + result = dump_depletion_projection_collection( + str(out), [site], [obs], {"id": "nm_dp"} + ) props = result["features"][0]["properties"] assert props["status"] == "no well depth" assert props["slope_ft_per_year"] == 1.0 # trend still reported def test_dtw_exceeds_well_depth(self, tmp_path): site = _trend_site(rid="W1", well_depth=55.0) - obs = [_trend_obs(f"{2010 + i}-01-01", 50.0 + i) for i in range(12)] # latest 61 + obs = [ + _trend_obs(f"{2010 + i}-01-01", 50.0 + i) for i in range(12) + ] # latest 61 out = tmp_path / "dp.geojson" - result = dump_depletion_projection_collection(str(out), [site], [obs], {"id": "nm_dp"}) + result = dump_depletion_projection_collection( + str(out), [site], [obs], {"id": "nm_dp"} + ) props = result["features"][0]["properties"] assert props["status"] == "dtw exceeds well depth" assert props["years_to_depletion"] is None @@ -988,7 +1110,97 @@ def test_not_enough_data(self, tmp_path): site = _trend_site(rid="W1", well_depth=100.0) obs = [_trend_obs("2020-01-01", 50.0)] out = tmp_path / "dp.geojson" - result = dump_depletion_projection_collection(str(out), [site], [obs], {"id": "nm_dp"}) + result = dump_depletion_projection_collection( + str(out), [site], [obs], {"id": "nm_dp"} + ) props = result["features"][0]["properties"] assert props["status"] == "not enough data" assert props["trend_category"] is None + + +class TestWellCorrelationCollection: + def _sites(self): + return [ + { + "source": "NMBGMR", + "id": "W1", + "name": "Well 1", + "latitude": 34.0, + "longitude": -106.0, + "elevation": 1600.0, + "well_depth": 100, + "usgs_site_id": "08313000", + "alternate_site_id": "P1", + }, + { + "source": "USGS-NWIS", + "id": "USGS-08313000", + "name": "NWIS", + "latitude": 34.0001, + "longitude": -106.0, + "well_depth": 100, + }, + { + "source": "NMOSEPOD", + "id": "P1", + "name": None, + "latitude": 34.0, + "longitude": -106.0, + }, + ] + + def test_ogc_required_fields_and_one_feature_per_well(self, tmp_path): + from backend.persisters.ogc_features import dump_well_correlation_collection + + out = tmp_path / "wc.geojson" + coll = dump_well_correlation_collection( + str(out), self._sites(), {"id": "nm_well_correlation"} + ) + assert coll["type"] == "FeatureCollection" + assert coll["id"] == "nm_well_correlation" + assert coll["numberReturned"] == 3 + assert "timeStamp" in coll + assert "correlation_method" in coll + assert len(coll["features"]) == 3 + for f in coll["features"]: + assert "id" in f + assert f["geometry"]["type"] == "Point" + + def test_cluster_and_pod_links(self, tmp_path): + from backend.persisters.ogc_features import dump_well_correlation_collection + + out = tmp_path / "wc.geojson" + coll = dump_well_correlation_collection( + str(out), self._sites(), {"id": "nm_well_correlation"} + ) + props = {f["properties"]["id"]: f["properties"] for f in coll["features"]} + w1 = props["W1"] + # all three collapse to one cluster + assert w1["cluster_size"] == 3 + assert w1["n_agencies"] == 3 + assert "P1" in w1["ose_pod_ids"] + assert "USGS-NWIS:USGS-08313000" in w1["linked_site_ids"] + # nested map serialized to a JSON string + assert isinstance(w1["linked_by_agency"], str) + assert json.loads(w1["linked_by_agency"])["NMOSEPOD"] == ["P1"] + + def test_writes_valid_geojson_file(self, tmp_path): + from backend.persisters.ogc_features import dump_well_correlation_collection + + out = tmp_path / "wc.geojson" + dump_well_correlation_collection( + str(out), self._sites(), {"id": "nm_well_correlation"} + ) + with open(out) as f: + data = json.load(f) + assert data["type"] == "FeatureCollection" + + def test_empty_sites(self, tmp_path): + from backend.persisters.ogc_features import dump_well_correlation_collection + + out = tmp_path / "wc.geojson" + coll = dump_well_correlation_collection( + str(out), [], {"id": "nm_well_correlation"} + ) + assert coll["numberReturned"] == 0 + assert coll["features"] == [] diff --git a/tests/test_well_correlation.py b/tests/test_well_correlation.py new file mode 100644 index 00000000..1a099409 --- /dev/null +++ b/tests/test_well_correlation.py @@ -0,0 +1,262 @@ +from backend.well_correlation import ( + correlate_wells, + haversine_m, + site_key, + _match_keys, + _parse_references, + _normalize_id, +) + + +def _site(source, rid, lat=None, lon=None, **kw): + d = {"source": source, "id": rid, "latitude": lat, "longitude": lon} + d.update(kw) + return d + + +def _by_key(results): + return {f"{_normalize_id(r['source'])}:{r['id']}": r for r in results} + + +class TestHelpers: + def test_match_keys_strips_usgs_prefix(self): + assert _match_keys("USGS-08313000") == {"USGS-08313000", "08313000"} + assert _match_keys(" abc123 ") == {"ABC123"} + assert _match_keys(None) == set() + assert _match_keys("") == set() + + def test_parse_references_splits_delimiters(self): + assert _parse_references("a, b; c|d/e") == ["a", "b", "c", "d", "e"] + assert _parse_references(None) == [] + assert _parse_references(" ") == [] + + def test_site_key(self): + assert site_key({"source": "nmbgmr", "id": "X1"}) == "NMBGMR:X1" + + def test_haversine_known_distance(self): + # ~111 km per degree of latitude. + d = haversine_m(34.0, -106.0, 35.0, -106.0) + assert abs(d - 111_195) < 500 + + +class TestExplicitLinking: + def test_usgs_site_id_links_across_prefix(self): + sites = [ + _site("NMBGMR", "WELL1", usgs_site_id="08313000"), + _site("USGS-NWIS", "USGS-08313000"), + ] + res = _by_key(correlate_wells(sites)) + a = res["NMBGMR:WELL1"] + b = res["USGS-NWIS:USGS-08313000"] + assert a["cluster_id"] == b["cluster_id"] + assert a["cluster_size"] == 2 + assert a["match_method"] == "explicit" + assert a["match_confidence"] == 0.95 + assert b["linked_site_ids"] == ["NMBGMR:WELL1"] + assert a["linked_by_agency"]["USGS-NWIS"] == ["USGS-08313000"] + + def test_alternate_site_id_multi_token(self): + sites = [ + _site("NMBGMR", "W", alternate_site_id="POD-9; OTHER-1"), + _site("NMOSEPOD", "POD-9"), + _site("ISC", "OTHER-1"), + ] + res = _by_key(correlate_wells(sites)) + assert res["NMBGMR:W"]["cluster_size"] == 3 + assert res["NMBGMR:W"]["n_agencies"] == 3 + + +class TestSpatialLinking: + def test_close_across_agencies_with_depth(self): + # ~30 m apart (within threshold/3 -> close bonus), depths agree. + sites = [ + _site("NMBGMR", "A", 34.00000, -106.00000, well_depth=100), + _site("PVACD", "B", 34.00027, -106.00000, well_depth=110), + ] + res = _by_key(correlate_wells(sites)) + assert res["NMBGMR:A"]["cluster_size"] == 2 + assert res["NMBGMR:A"]["match_method"] == "spatial" + # base 0.35 + close 0.10 + depth 0.25 + assert res["NMBGMR:A"]["match_confidence"] == 0.7 + + def test_spatial_only_without_depth_lower_confidence(self): + sites = [ + _site("NMBGMR", "A", 34.00000, -106.00000), + _site("PVACD", "B", 34.00027, -106.00000), + ] + res = _by_key(correlate_wells(sites)) + # base 0.35 + close 0.10, no corroborating attributes + assert res["NMBGMR:A"]["match_confidence"] == 0.45 + + def test_elevation_agreement_adds_confidence(self): + base = [ + _site("NMBGMR", "A", 34.0, -106.0, elevation=1600.0), + _site("PVACD", "B", 34.00027, -106.0, elevation=1605.0), + ] + c_elev = _by_key(correlate_wells(base))["NMBGMR:A"]["match_confidence"] + no_elev = _by_key( + correlate_wells( + [ + _site("NMBGMR", "A", 34.0, -106.0), + _site("PVACD", "B", 34.00027, -106.0), + ] + ) + )["NMBGMR:A"]["match_confidence"] + assert c_elev > no_elev + # base 0.35 + close 0.10 + elevation 0.12 + assert c_elev == 0.57 + + def test_elevation_disagreement_rejects(self): + sites = [ + _site("NMBGMR", "A", 34.0, -106.0, elevation=1600.0), + _site("PVACD", "B", 34.00027, -106.0, elevation=1900.0), + ] + res = _by_key(correlate_wells(sites)) + assert res["NMBGMR:A"]["cluster_size"] == 1 + + def test_name_token_overlap_adds_confidence(self): + sites = [ + _site("NMBGMR", "SMITH RANCH 1", 34.0, -106.0, name="Smith Ranch"), + _site("PVACD", "SMITH-RANCH", 34.00027, -106.0, name="Smith Ranch Well"), + ] + res = _by_key(correlate_wells(sites)) + a = res["NMBGMR:SMITH RANCH 1"] + assert a["cluster_size"] == 2 + # base 0.35 + close 0.10 + name 0.10 + assert a["match_confidence"] == 0.55 + + def test_confidence_capped_below_explicit(self): + # Every corroborating signal present -> capped at 0.9 (< explicit 0.95). + sites = [ + _site( + "NMBGMR", + "RANCHWELL", + 34.0, + -106.0, + well_depth=100, + elevation=1600.0, + name="Ranch Well", + ), + _site( + "PVACD", + "RANCHWELL", + 34.00009, + -106.0, + well_depth=105, + elevation=1602.0, + name="Ranch Well", + ), + ] + res = _by_key(correlate_wells(sites)) + assert res["NMBGMR:RANCHWELL"]["match_confidence"] == 0.9 + + def test_disagreeing_depth_blocks_link(self): + sites = [ + _site("NMBGMR", "A", 34.0, -106.0, well_depth=100), + _site("PVACD", "B", 34.00027, -106.0, well_depth=900), + ] + res = _by_key(correlate_wells(sites)) + assert res["NMBGMR:A"]["cluster_size"] == 1 + assert res["NMBGMR:A"]["match_method"] == "unmatched" + + def test_far_apart_not_linked(self): + sites = [ + _site("NMBGMR", "A", 34.0, -106.0), + _site("PVACD", "B", 34.1, -106.0), # ~11 km + ] + res = _by_key(correlate_wells(sites)) + assert res["NMBGMR:A"]["cluster_size"] == 1 + + def test_same_source_never_spatially_linked(self): + sites = [ + _site("NMBGMR", "A", 34.0, -106.0), + _site("NMBGMR", "B", 34.00009, -106.0), # ~10 m, same agency + ] + res = _by_key(correlate_wells(sites)) + assert res["NMBGMR:A"]["cluster_size"] == 1 + + +class TestPODLinking: + def test_pod_in_cluster_explicit(self): + sites = [ + _site("NMBGMR", "W", usgs_site_id=None, alternate_site_id="P1"), + _site("NMOSEPOD", "P1"), + ] + res = _by_key(correlate_wells(sites)) + w = res["NMBGMR:W"] + assert w["ose_pod_ids"] == ["P1"] + assert w["ose_pod_link_method"] == "explicit" + pod = res["NMOSEPOD:P1"] + assert pod["is_ose_pod"] is True + + def test_pod_linked_spatially_when_no_explicit(self): + sites = [ + _site("NMBGMR", "W", 34.0, -106.0), + _site("NMOSEPOD", "P1", 34.00027, -106.0), # ~30 m, no explicit ref + ] + res = _by_key(correlate_wells(sites)) + w = res["NMBGMR:W"] + # NMBGMR and NMOSEPOD are different agencies within distance -> they also + # cluster spatially, so POD is a member (explicit-membership) link. + assert "P1" in w["ose_pod_ids"] + + def test_pod_spatial_association_without_clustering(self): + # POD far enough that it does not cluster (>150 m) but within POD link + # distance widened here, so it is only an advisory spatial POD link. + sites = [ + _site("NMBGMR", "W", 34.0, -106.0), + _site("NMOSEPOD", "P1", 34.0027, -106.0), # ~300 m + ] + res = _by_key( + correlate_wells(sites, max_link_distance_m=150, pod_link_distance_m=500) + ) + w = res["NMBGMR:W"] + assert w["cluster_size"] == 1 # not clustered + assert w["ose_pod_ids"] == ["P1"] + assert w["ose_pod_link_method"] == "spatial" + + +class TestGeneral: + def test_every_input_well_emitted(self): + sites = [ + _site("A", "1", 34.0, -106.0), + _site("B", "2", 35.0, -107.0), + _site("C", "3"), + ] + res = correlate_wells(sites) + assert len(res) == 3 + + def test_duplicate_collapsed(self): + sites = [ + _site("A", "1", 34.0, -106.0), + _site("A", "1", 34.0, -106.0), + ] + res = correlate_wells(sites) + assert len(res) == 1 + + def test_order_independent(self): + sites = [ + _site("NMBGMR", "W", usgs_site_id="08313000"), + _site("USGS-NWIS", "USGS-08313000"), + _site("PVACD", "B", 34.0, -106.0), + ] + r1 = correlate_wells(sites) + r2 = correlate_wells(list(reversed(sites))) + assert r1 == r2 + + def test_transitive_clustering(self): + # A-B explicit, B-C spatial => all one cluster (mixed method). + sites = [ + _site("NMBGMR", "A", 34.0, -106.0, alternate_site_id="BID"), + _site("PVACD", "BID", 34.0, -106.0, well_depth=100), + _site("ISC", "C", 34.00027, -106.0, well_depth=100), + ] + res = _by_key(correlate_wells(sites)) + cids = { + res["NMBGMR:A"]["cluster_id"], + res["PVACD:BID"]["cluster_id"], + res["ISC:C"]["cluster_id"], + } + assert len(cids) == 1 + assert res["NMBGMR:A"]["match_method"] == "mixed" + assert res["NMBGMR:A"]["match_confidence"] == 0.95