uxarray-mcp 0.1.0__py3-none-any.whl

uxarray_mcp/domain/vector_calc.py

@@ -0,0 +1,256 @@
+"""Vector calculus and azimuthal averaging on unstructured meshes."""
+
+from __future__ import annotations
+
+from typing import Any
+
+
+def compute_gradient(uxds: Any, variable_name: str) -> dict:
+    """Compute the gradient of a face-centered scalar field.
+
+    Uses UXarray's Green-Gauss finite-volume gradient, which forms a closed
+    control volume around each cell by connecting centroids of neighbouring
+    cells.
+
+    Parameters
+    ----------
+    uxds : ux.UxDataset
+        Loaded UXarray dataset.
+    variable_name : str
+        Face-centered scalar variable to differentiate.
+
+    Returns
+    -------
+    dict
+        Keys: variable_name, zonal_component_name, meridional_component_name,
+        n_face, stats (min/max/mean for each component).
+    """
+    if variable_name not in uxds.data_vars:
+        raise ValueError(
+            f"Variable '{variable_name}' not found. Available: {list(uxds.data_vars)}"
+        )
+    var = uxds[variable_name]
+    if "n_face" not in var.dims and "nCells" not in var.dims:
+        raise ValueError(
+            f"Variable '{variable_name}' is not face-centered. "
+            "Gradient requires face-centered data."
+        )
+
+    import numpy as np
+
+    grad = var.gradient()
+    # gradient() returns a UxDataset with two variables
+    comp_names = list(grad.data_vars)
+
+    def _stats(arr: Any) -> dict:
+        vals = arr.values
+        finite = vals[np.isfinite(vals)]
+        if finite.size == 0:
+            return {"min": None, "max": None, "mean": None}
+        return {
+            "min": float(finite.min()),
+            "max": float(finite.max()),
+            "mean": float(finite.mean()),
+        }
+
+    components = {name: _stats(grad[name]) for name in comp_names}
+
+    return {
+        "variable_name": variable_name,
+        "components": comp_names,
+        "component_stats": components,
+        "n_face": int(uxds.uxgrid.n_face),
+        "interpretation": "zonal (∂/∂x) and meridional (∂/∂y) components of the gradient",
+    }
+
+
+def compute_curl(uxds: Any, u_variable: str, v_variable: str) -> dict:
+    """Compute the curl (relative vorticity) of a 2-D vector field (u, v).
+
+    The curl is the vertical component of ∇ × (u, v):
+        ζ = ∂v/∂x − ∂u/∂y
+
+    For atmospheric wind fields this is the relative vorticity.
+
+    Parameters
+    ----------
+    uxds : ux.UxDataset
+        Loaded UXarray dataset containing both components.
+    u_variable : str
+        Zonal (east–west) component variable name.
+    v_variable : str
+        Meridional (north–south) component variable name.
+
+    Returns
+    -------
+    dict
+        Keys: u_variable, v_variable, n_face, stats (min/max/mean/std of curl).
+    """
+    for name in (u_variable, v_variable):
+        if name not in uxds.data_vars:
+            raise ValueError(
+                f"Variable '{name}' not found. Available: {list(uxds.data_vars)}"
+            )
+    u = uxds[u_variable]
+    v = uxds[v_variable]
+    for name, var in ((u_variable, u), (v_variable, v)):
+        if "n_face" not in var.dims and "nCells" not in var.dims:
+            raise ValueError(
+                f"Variable '{name}' is not face-centered. "
+                "Curl requires face-centered vector components."
+            )
+
+    import numpy as np
+
+    result = u.curl(v)
+    vals = result.values
+    finite = vals[np.isfinite(vals)]
+
+    stats: dict = {}
+    if finite.size > 0:
+        stats = {
+            "min": float(finite.min()),
+            "max": float(finite.max()),
+            "mean": float(finite.mean()),
+            "std": float(finite.std()),
+        }
+    else:
+        stats = {"min": None, "max": None, "mean": None, "std": None}
+
+    return {
+        "u_variable": u_variable,
+        "v_variable": v_variable,
+        "interpretation": "relative vorticity ζ = ∂v/∂x − ∂u/∂y",
+        "n_face": int(uxds.uxgrid.n_face),
+        "stats": stats,
+    }
+
+
+def compute_divergence(uxds: Any, u_variable: str, v_variable: str) -> dict:
+    """Compute the horizontal divergence of a 2-D vector field (u, v).
+
+    Divergence = ∂u/∂x + ∂v/∂y.
+
+    Positive values indicate divergence (outflow), negative values indicate
+    convergence (inflow). Surface wind convergence drives rising motion and
+    convection.
+
+    Parameters
+    ----------
+    uxds : ux.UxDataset
+        Loaded UXarray dataset.
+    u_variable : str
+        Zonal (east–west) component variable name.
+    v_variable : str
+        Meridional (north–south) component variable name.
+
+    Returns
+    -------
+    dict
+        Keys: u_variable, v_variable, n_face, stats (min/max/mean/std).
+    """
+    for name in (u_variable, v_variable):
+        if name not in uxds.data_vars:
+            raise ValueError(
+                f"Variable '{name}' not found. Available: {list(uxds.data_vars)}"
+            )
+    u = uxds[u_variable]
+    v = uxds[v_variable]
+    for name, var in ((u_variable, u), (v_variable, v)):
+        if "n_face" not in var.dims and "nCells" not in var.dims:
+            raise ValueError(
+                f"Variable '{name}' is not face-centered. "
+                "Divergence requires face-centered vector components."
+            )
+
+    import numpy as np
+
+    result = u.divergence(v)
+    vals = result.values
+    finite = vals[np.isfinite(vals)]
+
+    stats: dict = {}
+    if finite.size > 0:
+        stats = {
+            "min": float(finite.min()),
+            "max": float(finite.max()),
+            "mean": float(finite.mean()),
+            "std": float(finite.std()),
+        }
+    else:
+        stats = {"min": None, "max": None, "mean": None, "std": None}
+
+    return {
+        "u_variable": u_variable,
+        "v_variable": v_variable,
+        "interpretation": "horizontal divergence ∂u/∂x + ∂v/∂y",
+        "n_face": int(uxds.uxgrid.n_face),
+        "stats": stats,
+    }
+
+
+def compute_azimuthal_mean(
+    uxds: Any,
+    variable_name: str,
+    center_lon: float,
+    center_lat: float,
+    outer_radius: float,
+    radius_step: float,
+) -> dict:
+    """Compute the azimuthal (radial) mean around a centre point.
+
+    Averages the variable along circles of constant great-circle distance from
+    the centre, producing a radial profile. Useful for analysing tropical
+    cyclones, polar vortex structure, or any feature with approximate radial
+    symmetry.
+
+    Parameters
+    ----------
+    uxds : ux.UxDataset
+        Loaded UXarray dataset.
+    variable_name : str
+        Face-centered variable to average.
+    center_lon : float
+        Longitude of the centre point (degrees).
+    center_lat : float
+        Latitude of the centre point (degrees).
+    outer_radius : float
+        Maximum radius in great-circle degrees.
+    radius_step : float
+        Radial bin width in great-circle degrees.
+
+    Returns
+    -------
+    dict
+        Keys: variable_name, center, radii, azimuthal_mean_values, n_face.
+    """
+    if variable_name not in uxds.data_vars:
+        raise ValueError(
+            f"Variable '{variable_name}' not found. Available: {list(uxds.data_vars)}"
+        )
+    var = uxds[variable_name]
+    if "n_face" not in var.dims and "nCells" not in var.dims:
+        raise ValueError(
+            f"Variable '{variable_name}' is not face-centered. "
+            "Azimuthal mean requires face-centered data."
+        )
+
+    result = var.azimuthal_mean(
+        center_coord=(center_lon, center_lat),
+        outer_radius=outer_radius,
+        radius_step=radius_step,
+    )
+
+    # result is an xr.DataArray with a radius coordinate
+    radii = result.coords[result.dims[0]].values.tolist()
+    values = result.values.tolist()
+
+    return {
+        "variable_name": variable_name,
+        "center": {"lon": center_lon, "lat": center_lat},
+        "outer_radius_deg": outer_radius,
+        "radius_step_deg": radius_step,
+        "radii_deg": radii,
+        "azimuthal_mean_values": values,
+        "n_face": int(uxds.uxgrid.n_face),
+    }

uxarray_mcp/domain/zonal.py

@@ -0,0 +1,66 @@
+"""Shared zonal mean computation logic."""
+
+from __future__ import annotations
+
+from typing import Any, Optional
+
+
+def compute_zonal_mean_stats(
+    uxds: Any,
+    variable_name: str,
+    lat_spec: Optional[tuple | float | list] = None,
+    conservative: bool = False,
+) -> dict:
+    """Compute zonal mean statistics from a loaded UXarray dataset.
+
+    Parameters
+    ----------
+    uxds : ux.UxDataset
+        Loaded UXarray dataset.
+    variable_name : str
+        Name of face-centered variable to average.
+    lat_spec : tuple | float | list | None
+        Latitude specification for zonal bands.
+    conservative : bool
+        If True, use area-weighted conservative averaging.
+
+    Returns
+    -------
+    dict
+        Keys: variable_name, latitudes, zonal_mean_values, conservative, grid_info
+    """
+    if variable_name not in uxds.data_vars:
+        available = list(uxds.data_vars.keys())
+        raise ValueError(
+            f"Variable '{variable_name}' not found. Available variables: {available}"
+        )
+
+    var = uxds[variable_name]
+
+    if "n_face" not in var.dims and "nCells" not in var.dims:
+        raise ValueError(
+            f"Variable '{variable_name}' is not face-centered. "
+            "Zonal mean only supports face-centered data."
+        )
+
+    if lat_spec is not None:
+        zonal_result = var.zonal_mean(lat=lat_spec, conservative=conservative)
+    else:
+        zonal_result = var.zonal_mean(conservative=conservative)
+
+    latitudes = zonal_result.coords["latitudes"].values.tolist()
+    zonal_mean_values = zonal_result.values.tolist()
+
+    grid_info = {
+        "n_face": int(uxds.uxgrid.n_face),
+        "n_node": int(uxds.uxgrid.n_node),
+        "n_edge": int(uxds.uxgrid.n_edge),
+    }
+
+    return {
+        "variable_name": variable_name,
+        "latitudes": latitudes,
+        "zonal_mean_values": zonal_mean_values,
+        "conservative": conservative,
+        "grid_info": grid_info,
+    }

uxarray_mcp/provenance.py

@@ -0,0 +1,79 @@
+"""Provenance tracking for UXarray MCP tool outputs.
+
+Appends a _provenance key to every tool result so scientific workflows
+can trace what ran, when, where, and with what software.
+"""
+
+from __future__ import annotations
+
+import sys
+from datetime import datetime, timezone
+from typing import Any
+
+
+def _get_uxarray_version() -> str:
+    """Return the installed UXarray version or ``unknown`` on failure."""
+    try:
+        import uxarray
+
+        return uxarray.__version__
+    except Exception:
+        return "unknown"
+
+
+def attach_provenance(
+    result: dict[str, Any],
+    tool: str,
+    inputs: dict[str, Any],
+    venue: str = "local",
+    warnings: list[str] | None = None,
+    validation_summary: dict[str, Any] | None = None,
+    selected_variable: str | None = None,
+    artifacts: list[dict[str, Any]] | None = None,
+) -> dict[str, Any]:
+    """Attach a _provenance block to a tool result dict.
+
+    Parameters
+    ----------
+    result : dict
+        The tool output to annotate.
+    tool : str
+        Name of the tool that produced the result.
+    inputs : dict
+        The input arguments passed to the tool.
+    venue : str
+        Execution venue, for example ``"local"`` or ``"hpc:<endpoint-name>"``.
+        Public results should not expose raw endpoint UUIDs.
+    warnings : list[str] | None
+        Any warnings generated during execution.
+    validation_summary : dict | None
+        Summary from validate_dataset: passed, n_variables_checked,
+        n_variables_failed. Included when a validation step ran upstream.
+    selected_variable : str | None
+        The variable name that was analysed, when applicable.
+    artifacts : list[dict] | None
+        Computational outputs produced by this run. Each entry is a dict
+        with at minimum a "type" key describing what was computed
+        (e.g. mesh_topology, face_areas, zonal_mean, validation).
+
+    Returns
+    -------
+    dict
+        The result dict with a _provenance key added.
+    """
+    provenance: dict[str, Any] = {
+        "tool": tool,
+        "inputs": inputs,
+        "execution_venue": venue,
+        "timestamp_utc": datetime.now(timezone.utc).isoformat(),
+        "uxarray_version": _get_uxarray_version(),
+        "python_version": f"{sys.version_info.major}.{sys.version_info.minor}.{sys.version_info.micro}",
+        "warnings": warnings if warnings is not None else [],
+        "artifacts": artifacts if artifacts is not None else [],
+    }
+    if selected_variable is not None:
+        provenance["selected_variable"] = selected_variable
+    if validation_summary is not None:
+        provenance["validation_summary"] = validation_summary
+    result["_provenance"] = provenance
+    return result

uxarray_mcp/remote/__init__.py

@@ -0,0 +1,6 @@
+"""Remote execution support for UXarray MCP server via Globus Compute and Academy."""
+
+from .agent import UXarrayComputeAgent
+from .config import load_config
+
+__all__ = ["load_config", "UXarrayComputeAgent"]