diff --git a/AGENTS.md b/AGENTS.md index ccaa82c..78186c4 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -25,9 +25,9 @@ It supports local execution and optional remote execution on HPC clusters via The tool surface is built by `uxarray_mcp.registry.build_registry()`: -- **`core` (default, ~27 tools)** — 11 front-door gateway tools at the +- **`core` (default, ~31 tools)** — 11 front-door gateway tools at the top level, 12 control/status tools under `session/` and `hpc/` - namespaces, `io-list_datasets`, and 3 prompt helpers under `prompt/`. + namespaces, `io-list_datasets`, and 7 prompt helpers under `prompt/`. No deferred tools, no BM25 discovery. This is what clients see by default when running `uxarray-mcp serve`. @@ -54,7 +54,8 @@ HPC control (`hpc/`): `endpoint_status`, `get_execution_mode`, IO (`io/`): `list_datasets`. Prompt helpers (`prompt/`): `first_look`, `vorticity_analysis`, -`hpc_diagnose`. +`cyclone_structure`, `eddy_activity`, `model_evaluation`, +`climatology_anomaly`, `hpc_diagnose`. Low-level implementation functions such as `inspect_mesh`, `calculate_area`, `plot_mesh`, and `calculate_curl` remain importable from `uxarray_mcp.tools` diff --git a/CHANGELOG.md b/CHANGELOG.md index 1d7477a..a77951a 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -5,6 +5,27 @@ uses Semantic Versioning for public releases. ## Unreleased +### Added +- Guided science workflows as `prompt/` tools, each composing existing + operations around a scientific question: `cyclone_structure` (storm radial + structure), `eddy_activity` (departures from the zonal mean), + `model_evaluation` (bias/RMSE/pattern correlation vs a reference), and + `climatology_anomaly` (time-mean state and anomalies). These join the existing + `vorticity_analysis` workflow. +- `run_analysis` operation `zonal_anomaly` — per-face deviation from the zonal + mean of each latitude band (`UxDataArray.zonal_anomaly`). +- `run_analysis` operation `remap_to_rectilinear` — remap an unstructured + variable onto a regular lon/lat grid (`UxDataArray.remap.to_rectilinear`). +- `gradient` and `curl` operations now accept a `scale_by_radius` flag. It + defaults to `False` to preserve unit-sphere results; set it to `True` to + divide by `uxgrid.sphere_radius` for physical units. + +### Changed +- Bumped the `uxarray` floor to `>=2026.6.0` for the new zonal-anomaly, + rectilinear-remap, and radius-scaled gradient/curl APIs. + +## 0.2.0 — 2026-06-19 + ### Changed - **Server engine**: replaced FastMCP with [toolregistry](https://github.com/Oaklight/ToolRegistry) + @@ -32,6 +53,14 @@ uses Semantic Versioning for public releases. - `fastmcp` dependency. - `@mcp.prompt()` decorators (replaced by `prompt/` namespace tools). +## 0.1.1 — 2026-06-11 + +### Changed +- Pinned Python to `>=3.12,<3.13` to match the supported runtime and avoid + Globus Compute pickle version-mismatch failures. +- Aligned the published package metadata with the current PyPI release for the + conda-forge recipe. + ## 0.1.0 — 2026-06-04 Initial public release. diff --git a/README.md b/README.md index d29e020..8ceb29b 100644 --- a/README.md +++ b/README.md @@ -16,6 +16,13 @@ remotely on an HPC system you have access to. > / vorticity / divergence, subset, remap, plot, and run multi-step workflows. > All as natural-language prompts. +> **Local by default; HPC is opt-in.** Everything runs on your machine unless +> you configure a [Globus Compute](https://www.globus.org/compute) endpoint. +> The remote option only becomes available once such an endpoint exists — +> running one requires an account and allocation on that HPC system, though a +> shared/service-account endpoint can let authorized users submit without their +> own login. + > **⚠️ What the AI can access.** Any file you (or your HPC account) can read. > Any compute the configured endpoint can submit. Outputs are written to your > disk. **See [SECURITY.md](SECURITY.md) before connecting any remote endpoint.** diff --git a/docs/api.rst b/docs/api.rst index 9e865e2..cc0bfbb 100644 --- a/docs/api.rst +++ b/docs/api.rst @@ -24,6 +24,10 @@ These modules contain the pure computation logic, separate from MCP and I/O. :members: :undoc-members: +.. automodule:: uxarray_mcp.domain.vector_calc + :members: + :undoc-members: + .. automodule:: uxarray_mcp.domain.plotting :members: :undoc-members: @@ -33,10 +37,18 @@ Tools MCP tool functions that are exposed to AI agents. +.. automodule:: uxarray_mcp.tools.frontdoor + :members: + :undoc-members: + .. automodule:: uxarray_mcp.tools.inspection :members: :undoc-members: +.. automodule:: uxarray_mcp.tools.vector_calc + :members: + :undoc-members: + .. automodule:: uxarray_mcp.tools.capabilities :members: :undoc-members: diff --git a/docs/architecture.html b/docs/architecture.html index f048dab..64b5c11 100644 --- a/docs/architecture.html +++ b/docs/architecture.html @@ -104,7 +104,7 @@

UXarray MCP Server — Architecture Diagram

- feature/tool-discovery  ·  Python 3.13  ·  FastMCP + main  ·  Python 3.12  ·  toolregistry-server
@@ -173,7 +173,7 @@

UXarray MCP Server — Architecture Diagram

CLAUDE DESKTOP MCP Client MCP SERVER - FastMCP · Provenance · Dynamic Registration + toolregistry-server · Provenance · Policy Tags SCIENTIFIC AGENT Analyze → Plan → Execute → Verify LOCAL UXARRAY diff --git a/docs/conf.py b/docs/conf.py index 3241bec..99914c1 100644 --- a/docs/conf.py +++ b/docs/conf.py @@ -46,7 +46,12 @@ } templates_path = ["_templates"] -exclude_patterns = ["_build", "Thumbs.db", ".DS_Store"] +exclude_patterns = [ + "_build", + "Thumbs.db", + ".DS_Store", + "onepager-uxarray-mcp.md", +] # -- Options for HTML output ------------------------------------------------- diff --git a/docs/index.rst b/docs/index.rst index acf515f..8544c04 100644 --- a/docs/index.rst +++ b/docs/index.rst @@ -3,24 +3,33 @@ UXarray MCP Server An MCP server that connects AI assistants like Claude to `UXarray `_ for analyzing unstructured -climate and weather meshes. Supports local execution and remote HPC -via Globus Compute. +climate and weather meshes. It runs locally by default; HPC execution via +Globus Compute is optional and only active once you configure an endpoint. + +The same tools are served over MCP (for AI clients) or OpenAPI/REST (for +HTTP clients) from a single install. .. toctree:: :maxdepth: 2 - :caption: User Guide + :caption: Getting started getting-started + tools + serving + workflows + scientific-agent + +.. toctree:: + :maxdepth: 2 + :caption: Running on HPC (optional) + remote-hpc - operating-an-endpoint globus-compute - tools + operating-an-endpoint hpc improv ucar chrysalis - workflows - scientific-agent .. toctree:: :maxdepth: 2 diff --git a/docs/serving.md b/docs/serving.md new file mode 100644 index 0000000..02dfbcf --- /dev/null +++ b/docs/serving.md @@ -0,0 +1,102 @@ +# Serving and Transports + +The server is built on +[toolregistry](https://github.com/Oaklight/ToolRegistry) + +[toolregistry-server](https://github.com/Oaklight/toolregistry-server). The same +tool implementations are exposed over multiple protocols from a single process, +so the server is not tied to any one client. Claude Desktop is the default +target, but the tools are equally usable over REST or from plain Python. + +For the Claude Desktop quickstart, see {doc}`getting-started`. This page covers +the additional surfaces the engine provides. + +## The `serve` command + +```bash +uxarray-mcp serve [--profile {core,deferred-full}] + [--transport {stdio,sse,http}] + [--host HOST] [--port PORT] +``` + +With no flags, `uxarray-mcp serve` starts MCP over stdio with the `core` +profile — the same behavior earlier releases had. Existing Claude Desktop +configurations work unchanged. + +The server is assembled by `UXarrayApp`, an `App` subclass from +toolregistry-server that carries a single `ServerIdentity` (name, version, +description) shared by every transport. The same object also drives the +`openapi` command below. + +## Profiles + +The visible tool surface is selected by `--profile`: + +| Profile | Visible tools | Use when | +|---|---|---| +| `core` (default) | ~31 | Most clients. Front-door gateway tools, session/HPC control, `list_datasets`, and prompts. Predictable, conservative. | +| `deferred-full` | ~31 visible (+32 deferred) | You want the full low-level surface. The raw implementation tools load with `defer=True`, so they do not appear in the initial list; agents find them with `discover_tools`. | + +```bash +uxarray-mcp serve --profile deferred-full +``` + +In `core`, low-level operations are still reachable through the front-door +dispatchers (for example `run_analysis(operation="curl", ...)`). + +## Transports + +`--transport` selects the MCP transport: + +| Transport | Description | +|---|---| +| `stdio` (default) | Subprocess transport used by Claude Desktop and Claude Code. | +| `sse` | Server-Sent Events over HTTP. Bind with `--host`/`--port`. | +| `http` | Streamable HTTP. Bind with `--host`/`--port`. | + +```bash +uxarray-mcp serve --transport sse --host 127.0.0.1 --port 8001 +uxarray-mcp serve --transport http --port 8000 +``` + +## OpenAPI / REST + +The same tools can be exposed as an OpenAPI/REST service for clients that speak +HTTP rather than MCP — curl, cloud assistants, chat UIs, or scripts. Install the +optional extra and start the `openapi` command: + +```bash +pip install "uxarray-mcp[openapi]" + +uxarray-mcp openapi [--profile {core,deferred-full}] [--host HOST] [--port PORT] +``` + +The tool implementations, provenance, and HPC dispatch are shared across MCP and +REST; only the protocol adapter differs. MCP and OpenAPI can run as separate +processes from the same install — for example `uxarray-mcp serve` for AI clients +and `uxarray-mcp openapi` for HTTP clients, behind a reverse proxy if needed. + +## Tool discovery (`deferred-full`) + +In the `deferred-full` profile, deferred tools are searchable via +`discover_tools`, which ranks tools by a BM25 match over names, docstrings, and +domain search hints. For example, a query like `"compute vorticity wind curl"` +surfaces `calculate_curl`. Operators can also promote deferred tools to the +visible set from the admin panel. + +## Using the tools without an MCP client + +Because the tools are ordinary Python functions, you can drive them directly — +no AI client, no transport: + +```python +from uxarray_mcp.tools import inspect_mesh +result = inspect_mesh(file_path="grid.nc") # dict with a _provenance block + +# Or call by name through the same registry the server uses: +from uxarray_mcp.app import make_registry +reg = make_registry(profile="core") +run_analysis = reg.get_callable("run_analysis") +stats = run_analysis(operation="calculate_area", grid_path="grid.nc") +``` + +This is the basis for the REST surface and for pipeline/post-processing use. diff --git a/docs/tools.md b/docs/tools.md index c637000..09aa828 100644 --- a/docs/tools.md +++ b/docs/tools.md @@ -7,6 +7,11 @@ clients should use the intent-shaped tools below. Most tools return structured dictionaries with a `_provenance` block. Plotting returns MCP content blocks: an inline PNG plus JSON metadata. +The visible tool set depends on the profile (`core` by default, or +`deferred-full`), and the server can expose these tools over MCP stdio/SSE/HTTP +or OpenAPI/REST. See {doc}`serving` for profiles, transports, and tool +discovery. + ## Front-Door Tools ### `get_capabilities` @@ -43,10 +48,12 @@ Supported operations: | `validate_dataset` | NaN/Inf/fill-value checks | | `calculate_area` | Face area statistics | | `calculate_zonal_mean` | Latitude-band mean for a face-centered variable | +| `zonal_anomaly` | Per-face deviation from its latitude-band zonal mean | | `gradient`, `curl`, `divergence`, `azimuthal_mean` | Vector/radial diagnostics | | `subset_bbox`, `subset_polygon`, `cross_section` | Spatial selections | | `compare_fields`, `bias`, `rmse`, `pattern_correlation` | Same-grid comparisons | | `remap_variable`, `regrid_dataset` | UXarray-backed remapping | +| `remap_to_rectilinear` | Remap a variable onto a regular lon/lat grid | | `temporal_mean`, `anomaly` | Time-dimension summaries | | `ensemble_mean`, `ensemble_spread` | Multi-file ensemble summaries | | `export` | Write a persisted result or dataset to NetCDF/CSV | @@ -57,11 +64,40 @@ Common parameters include `grid_path`, `data_path`, `variable_name`, `endpoint`. Each operation validates the parameters it requires and returns a clear error if one is missing. +`gradient` and `curl` accept `scale_by_radius` (default `False`). When `False`, +results stay on the unit sphere (the historical behavior). Set it to `True` to +divide by `uxgrid.sphere_radius` for physical units; the grid must define +`sphere_radius`. + +`zonal_anomaly` and `remap_to_rectilinear` are backed by +`UxDataArray.zonal_anomaly` and `UxDataArray.remap.to_rectilinear`, available in +the pinned UXarray (`>=2026.6.0`). + +> **Local-only:** `zonal_anomaly` and `remap_to_rectilinear` currently run on the +> local machine and do not have a remote (Globus Compute) execution path yet, so +> they cannot operate on files that live only on an HPC filesystem. The other +> compute operations accept `use_remote=True`. Uniform HPC routing across all +> operations is tracked as a design item. + Examples: ```python run_analysis(operation="inspect_mesh", grid_path="healpix:4") run_analysis(operation="calculate_area", grid_path="/path/grid.nc") +run_analysis( + operation="zonal_anomaly", + grid_path="/path/grid.nc", + data_path="/path/data.nc", + variable_name="temperature", +) +run_analysis( + operation="remap_to_rectilinear", + grid_path="/path/grid.nc", + data_path="/path/data.nc", + variable_name="temperature", + target_lon=[0, 30, 60, 90, 120, 150, 180, 210, 240, 270, 300, 330], + target_lat=[-60, -30, 0, 30, 60], +) run_analysis( operation="calculate_zonal_mean", grid_path="/path/grid.nc", @@ -134,13 +170,29 @@ dispatcher either falls back locally or reports a structured readiness error. ## MCP Prompts -Prompts are user-invokable slash commands. In Claude Code or Claude Desktop -they appear as `/first_look`, `/vorticity_analysis`, and `/hpc_diagnose`. +Prompts are user-invokable slash commands that return a guided, multi-step +analysis plan (instruction text, not results) — the assistant then runs the +chained operations and interprets them. In Claude Code or Claude Desktop they +appear as `/first_look`, `/vorticity_analysis`, etc. + +General: - `/first_look path` calls `get_capabilities` and `analyze_dataset`. -- `/vorticity_analysis grid_path data_path u_var v_var` calls - `run_analysis(operation="curl")` and - `run_analysis(operation="divergence")`. -- `/hpc_diagnose [endpoint]` calls - `diagnose_endpoint(action="status")` and +- `/hpc_diagnose [endpoint]` calls `diagnose_endpoint(action="status")` and `diagnose_endpoint(action="validate")`. + +Science workflows (each composes existing `run_analysis` operations around a +scientific question): + +- `/vorticity_analysis grid_path data_path u_var v_var` — rotation and + divergence of a wind field (`curl` + `divergence`). +- `/cyclone_structure grid_path data_path variable_name center_lon center_lat [u_var v_var outer_radius]` + — radial structure of a storm/vortex (`azimuthal_mean` + `subset_bbox`, + optionally `curl`). +- `/eddy_activity grid_path data_path variable_name` — departures from the + latitudinal background state (`calculate_zonal_mean` + `zonal_anomaly` + + `gradient`). +- `/model_evaluation grid_path data_path_a data_path_b variable_name` — verify a + field against a reference (`bias` + `rmse` + `pattern_correlation`). +- `/climatology_anomaly data_path variable_name [grid_path]` — time-mean state + and departures (`temporal_mean` + `anomaly`, optionally `calculate_zonal_mean`). diff --git a/pyproject.toml b/pyproject.toml index 156c240..5873c66 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -29,7 +29,7 @@ dependencies = [ "holoviews>=1.19.0", "matplotlib>=3.9.0", "pyyaml>=6.0", - "uxarray>=2025.12.0", + "uxarray>=2026.6.0", ] [project.optional-dependencies] diff --git a/src/uxarray_mcp/domain/__init__.py b/src/uxarray_mcp/domain/__init__.py index d3c3b29..c529deb 100644 --- a/src/uxarray_mcp/domain/__init__.py +++ b/src/uxarray_mcp/domain/__init__.py @@ -13,7 +13,7 @@ compute_divergence, compute_gradient, ) -from .zonal import compute_zonal_mean_stats +from .zonal import compute_zonal_anomaly_stats, compute_zonal_mean_stats __all__ = [ "load_grid", @@ -21,6 +21,7 @@ "compute_area_stats", "compute_variable_info", "compute_zonal_mean_stats", + "compute_zonal_anomaly_stats", "compute_gradient", "compute_curl", "compute_divergence", diff --git a/src/uxarray_mcp/domain/vector_calc.py b/src/uxarray_mcp/domain/vector_calc.py index 8e74301..a9471e7 100644 --- a/src/uxarray_mcp/domain/vector_calc.py +++ b/src/uxarray_mcp/domain/vector_calc.py @@ -5,7 +5,9 @@ from typing import Any -def compute_gradient(uxds: Any, variable_name: str) -> dict: +def compute_gradient( + uxds: Any, variable_name: str, scale_by_radius: bool = False +) -> dict: """Compute the gradient of a face-centered scalar field. Uses UXarray's Green-Gauss finite-volume gradient, which forms a closed @@ -18,6 +20,10 @@ def compute_gradient(uxds: Any, variable_name: str) -> dict: Loaded UXarray dataset. variable_name : str Face-centered scalar variable to differentiate. + scale_by_radius : bool, default False + When ``True``, divide the unit-sphere derivatives by + ``uxgrid.sphere_radius`` to return physical units (requires a grid with + ``sphere_radius``). The default ``False`` keeps the unit-sphere result. Returns ------- @@ -38,8 +44,8 @@ def compute_gradient(uxds: Any, variable_name: str) -> dict: import numpy as np - grad = var.gradient() - # gradient() returns a UxDataset with two variables + # gradient() returns a UxDataset with zonal and meridional components. + grad = var.gradient(scale_by_radius=scale_by_radius) comp_names = list(grad.data_vars) def _stats(arr: Any) -> dict: @@ -60,11 +66,14 @@ def _stats(arr: Any) -> dict: "components": comp_names, "component_stats": components, "n_face": int(uxds.uxgrid.n_face), + "scale_by_radius": bool(scale_by_radius), "interpretation": "zonal (∂/∂x) and meridional (∂/∂y) components of the gradient", } -def compute_curl(uxds: Any, u_variable: str, v_variable: str) -> dict: +def compute_curl( + uxds: Any, u_variable: str, v_variable: str, scale_by_radius: bool = False +) -> dict: """Compute the curl (relative vorticity) of a 2-D vector field (u, v). The curl is the vertical component of ∇ × (u, v): @@ -80,6 +89,10 @@ def compute_curl(uxds: Any, u_variable: str, v_variable: str) -> dict: Zonal (east–west) component variable name. v_variable : str Meridional (north–south) component variable name. + scale_by_radius : bool, default False + When ``True``, divide the unit-sphere result by ``uxgrid.sphere_radius`` + to return physical units (``1/s`` for wind in ``m/s``; requires a grid + with ``sphere_radius``). The default ``False`` keeps the unit sphere. Returns ------- @@ -102,7 +115,7 @@ def compute_curl(uxds: Any, u_variable: str, v_variable: str) -> dict: import numpy as np - result = u.curl(v) + result = u.curl(v, scale_by_radius=scale_by_radius) vals = result.values finite = vals[np.isfinite(vals)] @@ -122,6 +135,7 @@ def compute_curl(uxds: Any, u_variable: str, v_variable: str) -> dict: "v_variable": v_variable, "interpretation": "relative vorticity ζ = ∂v/∂x − ∂u/∂y", "n_face": int(uxds.uxgrid.n_face), + "scale_by_radius": bool(scale_by_radius), "stats": stats, } diff --git a/src/uxarray_mcp/domain/zonal.py b/src/uxarray_mcp/domain/zonal.py index ee7c546..ee46f8d 100644 --- a/src/uxarray_mcp/domain/zonal.py +++ b/src/uxarray_mcp/domain/zonal.py @@ -64,3 +64,92 @@ def compute_zonal_mean_stats( "conservative": conservative, "grid_info": grid_info, } + + +def compute_zonal_anomaly_stats( + uxds: Any, + variable_name: str, + lat_spec: Optional[tuple | float | list] = None, + conservative: bool = False, +) -> dict: + """Compute zonal-anomaly statistics from a loaded UXarray dataset. + + The zonal anomaly is each face value minus the zonal mean of its latitude + band, producing a per-face field with the same shape as the input variable. + + Parameters + ---------- + uxds : ux.UxDataset + Loaded UXarray dataset. + variable_name : str + Name of the face-centered variable. + lat_spec : tuple | float | list | None + Latitude band specification passed through to ``zonal_anomaly``. A + ``(start, end, step)`` tuple or explicit band edges. ``None`` uses the + UXarray default ``(-90, 90, 10)``. + conservative : bool + If True, use area-weighted band means. + + Returns + ------- + dict + Keys: variable_name, conservative, n_face, stats (min/max/mean/std of + the anomaly field), grid_info. + + Raises + ------ + NotImplementedError + If the installed UXarray does not provide ``UxDataArray.zonal_anomaly``. + """ + import numpy as np + + 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 anomaly only supports face-centered data." + ) + + if not hasattr(var, "zonal_anomaly"): + raise NotImplementedError( + "zonal_anomaly requires a UXarray release that provides " + "UxDataArray.zonal_anomaly. Upgrade uxarray to use this operation." + ) + + if lat_spec is not None: + result = var.zonal_anomaly(lat=lat_spec, conservative=conservative) + else: + result = var.zonal_anomaly(conservative=conservative) + + vals = result.values + finite = vals[np.isfinite(vals)] + stats: dict[str, float | None] + 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 { + "variable_name": variable_name, + "conservative": conservative, + "n_face": int(uxds.uxgrid.n_face), + "stats": stats, + "interpretation": "per-face deviation from the zonal mean of its latitude band", + "grid_info": { + "n_face": int(uxds.uxgrid.n_face), + "n_node": int(uxds.uxgrid.n_node), + "n_edge": int(uxds.uxgrid.n_edge), + }, + } diff --git a/src/uxarray_mcp/registry.py b/src/uxarray_mcp/registry.py index 97dbd0b..a63c1c2 100644 --- a/src/uxarray_mcp/registry.py +++ b/src/uxarray_mcp/registry.py @@ -102,6 +102,7 @@ "calculate_ensemble_spread", "calculate_area", "calculate_zonal_mean", + "calculate_zonal_anomaly", ), "shape": ( "subset_bbox", @@ -109,6 +110,7 @@ "extract_cross_section", "remap_variable", "regrid_dataset", + "remap_to_rectilinear", ), "inspect": ( "inspect_mesh", @@ -189,6 +191,196 @@ def vorticity_analysis(grid_path: str, data_path: str, u_var: str, v_var: str) - ) +def cyclone_structure( + grid_path: str, + data_path: str, + variable_name: str, + center_lon: float, + center_lat: float, + u_var: str = "", + v_var: str = "", + outer_radius: float = 10.0, +) -> str: + """Generate a guided plan to characterise a cyclone / vortex structure. + + Builds a radial picture of a storm or vortex around a centre point using the + azimuthal (radial) mean, optionally adding the rotational field. Returns + instructional text — the LLM runs the operations and interprets them. + + Args: + grid_path: Path to the mesh grid file. + data_path: Path to the data file. + variable_name: Field to profile radially (e.g. wind speed, pressure). + center_lon: Longitude of the storm centre (degrees). + center_lat: Latitude of the storm centre (degrees). + u_var: Optional zonal wind component for a vorticity check. + v_var: Optional meridional wind component for a vorticity check. + outer_radius: Maximum radius in great-circle degrees. + + Returns: + Multi-step cyclone-structure analysis plan as a string. + """ + steps = [ + "1. Call `run_analysis` with " + f'operation="azimuthal_mean", grid_path="{grid_path}", ' + f'data_path="{data_path}", variable_name="{variable_name}", ' + f"center_lon={center_lon}, center_lat={center_lat}, " + f"outer_radius={outer_radius}, radius_step=0.5 to build the radial " + "profile.", + "2. Call `run_analysis` with " + f'operation="subset_bbox", grid_path="{grid_path}", ' + f'data_path="{data_path}", variable_name="{variable_name}", ' + f"lon_bounds=[{center_lon - outer_radius}, {center_lon + outer_radius}], " + f"lat_bounds=[{center_lat - outer_radius}, {center_lat + outer_radius}] " + "to isolate the storm region.", + ] + if u_var and v_var: + steps.append( + "3. Call `run_analysis` with " + f'operation="curl", grid_path="{grid_path}", data_path="{data_path}", ' + f'u_variable="{u_var}", v_variable="{v_var}" to confirm the ' + "rotational signature (relative vorticity)." + ) + steps.append( + f"{len(steps) + 1}. Interpret the radial profile: locate the radius of " + "maximum wind / minimum pressure, the storm's radial extent, and any " + 'asymmetry. Plot the subset with `plot_dataset(plot_type="variable")`.' + ) + return ( + f"Characterise the cyclone/vortex near ({center_lon}, {center_lat}) in " + f"`{data_path}`.\n\n" + "\n".join(steps) + ) + + +def eddy_activity( + grid_path: str, + data_path: str, + variable_name: str, +) -> str: + """Generate a guided plan to assess eddy / wave activity. + + Quantifies departures from the latitudinal background state — the signature + of eddies, stationary waves, and storm tracks — using the zonal anomaly and + its gradient. Returns instructional text. + + Args: + grid_path: Path to the mesh grid file. + data_path: Path to the data file. + variable_name: Face-centered field to analyse (e.g. geopotential height, + temperature). + + Returns: + Multi-step eddy-activity analysis plan as a string. + """ + return ( + f"Assess eddy/wave activity for `{variable_name}` in `{data_path}`.\n\n" + "1. Call `run_analysis` with " + f'operation="calculate_zonal_mean", grid_path="{grid_path}", ' + f'data_path="{data_path}", variable_name="{variable_name}" to establish ' + "the latitudinal background state.\n" + "2. Call `run_analysis` with " + f'operation="zonal_anomaly", grid_path="{grid_path}", ' + f'data_path="{data_path}", variable_name="{variable_name}" to isolate ' + "departures from each latitude band (the eddy field).\n" + "3. Call `run_analysis` with " + f'operation="gradient", grid_path="{grid_path}", data_path="{data_path}", ' + f'variable_name="{variable_name}" to highlight sharp gradients and ' + "fronts associated with the waves.\n" + "4. Interpret the anomaly amplitude (std/max) as eddy strength, note " + "where activity concentrates, and plot the anomaly field with " + '`plot_dataset(plot_type="variable")`.' + ) + + +def model_evaluation( + grid_path: str, + data_path_a: str, + data_path_b: str, + variable_name: str, +) -> str: + """Generate a guided plan to evaluate a model field against a reference. + + Computes the standard verification triple — bias, RMSE, and pattern + correlation — between two same-grid fields and guides interpretation. + Returns instructional text. + + Args: + grid_path: Path to the shared mesh grid file. + data_path_a: Model / candidate dataset. + data_path_b: Reference / observation dataset. + variable_name: Field to compare. + + Returns: + Multi-step model-evaluation plan as a string. + """ + return ( + f"Evaluate `{variable_name}` in `{data_path_a}` against reference " + f"`{data_path_b}`.\n\n" + "1. Call `run_analysis` with " + f'operation="bias", grid_path="{grid_path}", ' + f'data_path_a="{data_path_a}", data_path_b="{data_path_b}", ' + f'variable_name="{variable_name}" for the mean signed error.\n' + "2. Call `run_analysis` with " + f'operation="rmse", grid_path="{grid_path}", ' + f'data_path_a="{data_path_a}", data_path_b="{data_path_b}", ' + f'variable_name="{variable_name}" for the magnitude of the error.\n' + "3. Call `run_analysis` with " + f'operation="pattern_correlation", grid_path="{grid_path}", ' + f'data_path_a="{data_path_a}", data_path_b="{data_path_b}", ' + f'variable_name="{variable_name}" for spatial-pattern skill.\n' + "4. Interpret together: bias = systematic offset, RMSE = typical error " + "size, pattern correlation = structural agreement (1.0 = perfect). Call " + "out whether errors are a uniform offset or a structural mismatch." + ) + + +def climatology_anomaly( + data_path: str, + variable_name: str, + grid_path: str = "", +) -> str: + """Generate a guided plan for a climatology and anomaly analysis. + + Establishes the time-mean state and the departures from it over a time + series, then summarises the anomaly latitudinally. Returns instructional + text. + + Args: + data_path: Path to a time-series data file. + variable_name: Field to analyse. + grid_path: Optional mesh grid file (needed for the zonal summary). + + Returns: + Multi-step climatology/anomaly plan as a string. + """ + grid_arg = f'grid_path="{grid_path}", ' if grid_path else "" + last = ( + "4. Call `run_analysis` with " + f'operation="calculate_zonal_mean", grid_path="{grid_path}", ' + f'data_path="{data_path}", variable_name="{variable_name}" to summarise ' + "the anomaly by latitude.\n" + if grid_path + else "" + ) + return ( + f"Build a climatology and anomalies for `{variable_name}` in " + f"`{data_path}`.\n\n" + "1. Call `run_analysis` with " + f'operation="temporal_mean", data_path="{data_path}", ' + f'variable_name="{variable_name}" to compute the time-mean climatology.\n' + "2. Call `run_analysis` with " + f'operation="anomaly", {grid_arg}data_path="{data_path}", ' + f'variable_name="{variable_name}" to compute departures from the mean ' + "state.\n" + f'3. Plot the anomaly with `plot_dataset(plot_type="variable")`' + + (".\n" + last if last else " and interpret the spatial structure.\n") + + ( + f"{5 if grid_arg and last else 4}. Interpret where and when the " + "field departs most from its climatology." + ) + ) + + def hpc_diagnose(endpoint: str = "") -> str: """Generate a step-by-step prompt for HPC endpoint diagnosis. @@ -214,7 +406,15 @@ def hpc_diagnose(endpoint: str = "") -> str: _PROMPT_TOOLS: dict[str, tuple[str, ...]] = { - "prompt": ("first_look", "vorticity_analysis", "hpc_diagnose"), + "prompt": ( + "first_look", + "vorticity_analysis", + "cyclone_structure", + "eddy_activity", + "model_evaluation", + "climatology_anomaly", + "hpc_diagnose", + ), } # Map prompt tool names to their implementing functions (defined above @@ -222,6 +422,10 @@ def hpc_diagnose(endpoint: str = "") -> str: _PROMPT_FUNCS: dict[str, object] = { "first_look": first_look, "vorticity_analysis": vorticity_analysis, + "cyclone_structure": cyclone_structure, + "eddy_activity": eddy_activity, + "model_evaluation": model_evaluation, + "climatology_anomaly": climatology_anomaly, "hpc_diagnose": hpc_diagnose, } @@ -261,6 +465,10 @@ def hpc_diagnose(endpoint: str = "") -> str: # Prompt tools are always read-only (they just return text) "first_look": ({ToolTag.READ_ONLY}, set()), "vorticity_analysis": ({ToolTag.READ_ONLY}, set()), + "cyclone_structure": ({ToolTag.READ_ONLY}, set()), + "eddy_activity": ({ToolTag.READ_ONLY}, set()), + "model_evaluation": ({ToolTag.READ_ONLY}, set()), + "climatology_anomaly": ({ToolTag.READ_ONLY}, set()), "hpc_diagnose": ({ToolTag.READ_ONLY}, set()), } @@ -271,6 +479,7 @@ def hpc_diagnose(endpoint: str = "") -> str: "calculate_gradient", "calculate_azimuthal_mean", "calculate_zonal_mean", + "calculate_zonal_anomaly", "calculate_temporal_mean", "calculate_anomaly", "calculate_ensemble_mean", @@ -281,6 +490,7 @@ def hpc_diagnose(endpoint: str = "") -> str: "calculate_pattern_correlation", "remap_variable", "regrid_dataset", + "remap_to_rectilinear", "subset_polygon", "extract_cross_section", "plot_mesh", @@ -339,6 +549,7 @@ def _apply_tags( "calculate_gradient": "spatial derivative slope field gradient", "calculate_azimuthal_mean": "radial profile cyclone storm azimuthal", "calculate_zonal_mean": "latitudinal average belt zonal", + "calculate_zonal_anomaly": "zonal anomaly deviation latitude band eddy wave departure", "calculate_temporal_mean": "time average climatology", "calculate_anomaly": "deviation departure climatology", "calculate_ensemble_mean": "model average multi-member", @@ -353,6 +564,7 @@ def _apply_tags( "extract_cross_section": "transect slice latitude longitude", "remap_variable": "interpolation target grid", "regrid_dataset": "interpolation target grid all variables", + "remap_to_rectilinear": "remap rectilinear regular lon lat structured grid interpolation", "inspect_mesh": "topology nodes faces edges grid summary", "inspect_variable": "data variable metadata stats", "validate_dataset": "data quality NaN Inf fill check", diff --git a/src/uxarray_mcp/remote/agent.py b/src/uxarray_mcp/remote/agent.py index b3b548b..ee84b26 100644 --- a/src/uxarray_mcp/remote/agent.py +++ b/src/uxarray_mcp/remote/agent.py @@ -214,20 +214,38 @@ async def calculate_zonal_mean_remote( @action async def calculate_gradient_remote( - self, grid_path: str, data_path: str, variable_name: str + self, + grid_path: str, + data_path: str, + variable_name: str, + scale_by_radius: bool = False, ) -> Dict[str, Any]: """Compute spatial gradient on HPC.""" return await self._run_on_hpc( - remote_calculate_gradient, grid_path, data_path, variable_name + remote_calculate_gradient, + grid_path, + data_path, + variable_name, + scale_by_radius, ) @action async def calculate_curl_remote( - self, grid_path: str, data_path: str, u_variable: str, v_variable: str + self, + grid_path: str, + data_path: str, + u_variable: str, + v_variable: str, + scale_by_radius: bool = False, ) -> Dict[str, Any]: """Compute relative vorticity (curl) on HPC.""" return await self._run_on_hpc( - remote_calculate_curl, grid_path, data_path, u_variable, v_variable + remote_calculate_curl, + grid_path, + data_path, + u_variable, + v_variable, + scale_by_radius, ) @action diff --git a/src/uxarray_mcp/remote/compute_functions.py b/src/uxarray_mcp/remote/compute_functions.py index 154d961..a68dbaa 100644 --- a/src/uxarray_mcp/remote/compute_functions.py +++ b/src/uxarray_mcp/remote/compute_functions.py @@ -1035,9 +1035,10 @@ def _surface(mod): def remote_calculate_gradient( - grid_path: str, data_path: str, variable_name: str + grid_path: str, data_path: str, variable_name: str, scale_by_radius: bool = False ) -> Dict[str, Any]: """Compute the spatial gradient of a face-centered scalar field on HPC.""" + import inspect as _inspect import os import numpy as np @@ -1062,7 +1063,14 @@ def remote_calculate_gradient( "Gradient requires face-centered data." ) - grad = var.gradient() + # Honor scale_by_radius when the worker's UXarray supports it; otherwise + # fall back to the unit-sphere call so older workers keep working. + applied_scale = False + if "scale_by_radius" in _inspect.signature(var.gradient).parameters: + grad = var.gradient(scale_by_radius=scale_by_radius) + applied_scale = bool(scale_by_radius) + else: + grad = var.gradient() comp_names = list(grad.data_vars) def _stats(arr: Any) -> Dict[str, Any]: @@ -1081,17 +1089,23 @@ def _stats(arr: Any) -> Dict[str, Any]: "components": comp_names, "component_stats": {name: _stats(grad[name]) for name in comp_names}, "n_face": int(uxds.uxgrid.n_face), + "scale_by_radius": applied_scale, "interpretation": "zonal (d/dx) and meridional (d/dy) components of the gradient", } def remote_calculate_curl( - grid_path: str, data_path: str, u_variable: str, v_variable: str + grid_path: str, + data_path: str, + u_variable: str, + v_variable: str, + scale_by_radius: bool = False, ) -> Dict[str, Any]: """Compute relative vorticity (curl) of a 2-D wind field on HPC. zeta = dv/dx - du/dy """ + import inspect as _inspect import os import numpy as np @@ -1118,7 +1132,12 @@ def remote_calculate_curl( "Curl requires face-centered vector components." ) - result = u.curl(v) + applied_scale = False + if "scale_by_radius" in _inspect.signature(u.curl).parameters: + result = u.curl(v, scale_by_radius=scale_by_radius) + applied_scale = bool(scale_by_radius) + else: + result = u.curl(v) vals = result.values finite = vals[np.isfinite(vals)] stats: Dict[str, Any] = ( @@ -1136,6 +1155,7 @@ def remote_calculate_curl( "v_variable": v_variable, "interpretation": "relative vorticity zeta = dv/dx - du/dy", "n_face": int(uxds.uxgrid.n_face), + "scale_by_radius": applied_scale, "stats": stats, } diff --git a/src/uxarray_mcp/tools/__init__.py b/src/uxarray_mcp/tools/__init__.py index 8e083ae..bd750da 100644 --- a/src/uxarray_mcp/tools/__init__.py +++ b/src/uxarray_mcp/tools/__init__.py @@ -11,6 +11,7 @@ export_to_netcdf, extract_cross_section, regrid_dataset, + remap_to_rectilinear, remap_variable, subset_bbox, subset_polygon, @@ -33,7 +34,7 @@ plot_dataset, run_analysis, ) -from .inspection import validate_dataset +from .inspection import calculate_zonal_anomaly, validate_dataset from .orchestration import analyze_dataset # Public tool surface for inspection and plotting. Each function is a @@ -90,6 +91,7 @@ "inspect_variable", "calculate_area", "calculate_zonal_mean", + "calculate_zonal_anomaly", "validate_dataset", "plot_mesh", "plot_mesh_geo", @@ -104,6 +106,7 @@ "calculate_pattern_correlation", "remap_variable", "regrid_dataset", + "remap_to_rectilinear", "calculate_temporal_mean", "calculate_anomaly", "calculate_ensemble_mean", diff --git a/src/uxarray_mcp/tools/advanced.py b/src/uxarray_mcp/tools/advanced.py index dd0e02d..90a3c1b 100644 --- a/src/uxarray_mcp/tools/advanced.py +++ b/src/uxarray_mcp/tools/advanced.py @@ -760,6 +760,127 @@ def regrid_dataset( return result +def remap_to_rectilinear( + variable_name: str, + target_lon: Sequence[float], + target_lat: Sequence[float], + grid_path: str | None = None, + data_path: str | None = None, + backend: str = "uxarray", + session_id: str | None = None, + dataset_handle: str | None = None, + result_name: str | None = None, +) -> dict[str, Any]: + """Remap a face-centered variable onto a regular lon/lat (rectilinear) grid. + + Uses ``UxDataArray.remap.to_rectilinear`` to interpolate unstructured data + onto 1-D longitude/latitude coordinate arrays, producing a structured grid + suitable for downstream lon/lat workflows. + + Parameters + ---------- + variable_name : str + Face-centered variable to remap. + target_lon, target_lat : sequence of float + 1-D longitude and latitude cell-center coordinates in degrees. + grid_path, data_path : str, optional + Source grid and data files (or resolve from session/dataset handle). + backend : str + Remapping backend: ``"uxarray"`` (default) or ``"yac"``. + session_id, dataset_handle, result_name : optional + Session/result-handle plumbing. + + Returns + ------- + dict + Keys: ``variable_name``, ``backend``, ``target_shape`` (n_lat, n_lon), + ``stats`` (min/max/mean of the remapped field), ``result_handle``, and + ``_provenance``. + + Raises + ------ + NotImplementedError + If the installed UXarray lacks ``remap.to_rectilinear``. + """ + tracker = OperationTracker("remap_to_rectilinear", session_id=session_id) + resolved_grid, resolved_data = _resolve_paths( + session_id=session_id, + dataset_handle=dataset_handle, + grid_path=grid_path, + data_path=data_path, + ) + if resolved_data is None: + raise ValueError("data_path is required for remapping.") + + _, uxda, selected = _load_dataarray(resolved_grid, resolved_data, variable_name) + + if not hasattr(uxda.remap, "to_rectilinear"): + raise NotImplementedError( + "remap_to_rectilinear requires a UXarray release that provides " + "remap.to_rectilinear. Upgrade uxarray to use this operation." + ) + + lon = list(target_lon) + lat = list(target_lat) + tracker.stage("remapping", f"Remapping {selected} to {len(lat)}x{len(lon)} grid.") + remapped = uxda.remap.to_rectilinear(lon, lat, backend=backend) + + # remapped is a plain xarray.DataArray with lat/lon axes. + remapped_ds = remapped.to_dataset(name=selected) + result_handle = _persist_dataset_result( + dataset=remapped_ds, + session_id=session_id, + name=result_name or f"rectilinear:{selected}", + kind="rectilinear_variable", + summary=summarize_dataset(remapped_ds), + metadata={ + "source_grid": resolved_grid, + "variable_name": selected, + "backend": backend, + "n_lon": len(lon), + "n_lat": len(lat), + }, + ) + + vals = np.asarray(remapped.values, dtype=float) + finite = vals[np.isfinite(vals)] + stats = ( + { + "min": float(finite.min()), + "max": float(finite.max()), + "mean": float(finite.mean()), + } + if finite.size > 0 + else {"min": None, "max": None, "mean": None} + ) + + tracker.succeed("Rectilinear remap complete.") + result: dict[str, Any] = { + "variable_name": selected, + "backend": backend, + "target_shape": [len(lat), len(lon)], + "stats": stats, + "result_handle": result_handle, + } + result = attach_provenance( + result, + tool="remap_to_rectilinear", + inputs={ + "variable_name": variable_name, + "target_lon": lon, + "target_lat": lat, + "grid_path": grid_path, + "data_path": data_path, + "backend": backend, + "session_id": session_id, + "dataset_handle": dataset_handle, + }, + selected_variable=selected, + ) + result["_provenance"]["operation_id"] = tracker.operation_id + return result + + def calculate_temporal_mean( data_path: str, variable_name: str, diff --git a/src/uxarray_mcp/tools/capabilities.py b/src/uxarray_mcp/tools/capabilities.py index d942177..0d397e5 100644 --- a/src/uxarray_mcp/tools/capabilities.py +++ b/src/uxarray_mcp/tools/capabilities.py @@ -10,6 +10,30 @@ from uxarray_mcp.remote.config import load_config +def _uxarray_supports(attr_path: str) -> bool: + """Return True if ``UxDataArray`` provides the given (possibly dotted) attr. + + Used to advertise capability-gated operations only when the installed + UXarray actually ships the underlying method (e.g. ``zonal_anomaly`` or + ``remap.to_rectilinear``). + """ + obj: Any = ux.UxDataArray + for part in attr_path.split("."): + if part == "remap": + # The remap accessor class holds the remap methods. + try: + from uxarray.remap.accessor import RemapAccessor + + obj = RemapAccessor + except Exception: + return False + continue + if not hasattr(obj, part): + return False + obj = getattr(obj, part) + return True + + def get_capabilities( grid_path: str, data_path: Optional[str] = None, @@ -154,6 +178,14 @@ def get_capabilities( "run_analysis:ensemble_spread", "run_analysis:export", ] + # Capability-gated operations: only advertise when the installed + # UXarray actually provides the underlying method. + if _uxarray_supports("zonal_anomaly"): + applicable_mcp.append("run_analysis:zonal_anomaly") + applicable_uxarray.append("var.zonal_anomaly()") + if _uxarray_supports("remap.to_rectilinear"): + applicable_mcp.append("run_analysis:remap_to_rectilinear") + applicable_uxarray.append("var.remap.to_rectilinear(lon, lat)") applicable_uxarray += [ "var.zonal_mean()", "var.integrate()", diff --git a/src/uxarray_mcp/tools/frontdoor.py b/src/uxarray_mcp/tools/frontdoor.py index d771f2a..9ce02df 100644 --- a/src/uxarray_mcp/tools/frontdoor.py +++ b/src/uxarray_mcp/tools/frontdoor.py @@ -46,6 +46,11 @@ def run_analysis( session_id: str | None = None, dataset_handle: str | None = None, result_name: str | None = None, + scale_by_radius: bool = False, + lat_spec: tuple | float | list[Any] | None = None, + conservative: bool = False, + target_lon: list[float] | None = None, + target_lat: list[float] | None = None, use_remote: bool = False, endpoint: str | None = None, ) -> dict[str, Any]: @@ -53,12 +58,17 @@ def run_analysis( Supported operations: ``inspect_mesh``, ``inspect_variable``, ``validate_dataset``, - ``calculate_area``, ``calculate_zonal_mean``, ``gradient``, ``curl``, - ``divergence``, ``azimuthal_mean``, ``subset_bbox``, ``subset_polygon``, - ``cross_section``, ``compare_fields``, ``bias``, ``rmse``, - ``pattern_correlation``, ``remap_variable``, ``regrid_dataset``, - ``temporal_mean``, ``anomaly``, ``ensemble_mean``, ``ensemble_spread``, - and ``export``. + ``calculate_area``, ``calculate_zonal_mean``, ``zonal_anomaly``, + ``gradient``, ``curl``, ``divergence``, ``azimuthal_mean``, + ``subset_bbox``, ``subset_polygon``, ``cross_section``, ``compare_fields``, + ``bias``, ``rmse``, ``pattern_correlation``, ``remap_variable``, + ``regrid_dataset``, ``remap_to_rectilinear``, ``temporal_mean``, + ``anomaly``, ``ensemble_mean``, ``ensemble_spread``, and ``export``. + + ``gradient`` and ``curl`` accept ``scale_by_radius`` (default False keeps the + historical unit-sphere result). ``zonal_anomaly`` accepts ``lat_spec`` and + ``conservative``. ``remap_to_rectilinear`` accepts ``target_lon`` and + ``target_lat`` (1-D coordinate arrays). """ from uxarray_mcp.tools.advanced import ( calculate_anomaly, @@ -71,12 +81,13 @@ def run_analysis( compare_fields, extract_cross_section, regrid_dataset, + remap_to_rectilinear, remap_variable, subset_bbox, subset_polygon, write_result, ) - from uxarray_mcp.tools.inspection import validate_dataset + from uxarray_mcp.tools.inspection import calculate_zonal_anomaly, validate_dataset from uxarray_mcp.tools.remote_tools import ( calculate_area, calculate_zonal_mean, @@ -129,11 +140,20 @@ def run_analysis( endpoint=endpoint, session_id=session_id, ) + if op == "zonal_anomaly": + return calculate_zonal_anomaly( + _require(grid_path, "grid_path", op), + _require(data_path, "data_path", op), + _require(variable_name, "variable_name", op), + lat_spec=lat_spec, + conservative=conservative, + ) if op == "gradient": return calculate_gradient( _require(grid_path, "grid_path", op), _require(data_path, "data_path", op), _require(variable_name, "variable_name", op), + scale_by_radius=scale_by_radius, use_remote=use_remote, endpoint=endpoint, session_id=session_id, @@ -144,6 +164,7 @@ def run_analysis( _require(data_path, "data_path", op), _require(u_variable, "u_variable", op), _require(v_variable, "v_variable", op), + scale_by_radius=scale_by_radius, use_remote=use_remote, endpoint=endpoint, session_id=session_id, @@ -257,6 +278,17 @@ def run_analysis( dataset_handle=dataset_handle, result_name=result_name, ) + if op == "remap_to_rectilinear": + return remap_to_rectilinear( + variable_name=_require(variable_name, "variable_name", op), + target_lon=_require(target_lon, "target_lon", op), + target_lat=_require(target_lat, "target_lat", op), + grid_path=grid_path, + data_path=data_path, + session_id=session_id, + dataset_handle=dataset_handle, + result_name=result_name, + ) if op == "temporal_mean": return calculate_temporal_mean( data_path=_require(data_path, "data_path", op), diff --git a/src/uxarray_mcp/tools/inspection.py b/src/uxarray_mcp/tools/inspection.py index a4d2fe1..8686900 100644 --- a/src/uxarray_mcp/tools/inspection.py +++ b/src/uxarray_mcp/tools/inspection.py @@ -10,6 +10,7 @@ from uxarray_mcp.domain import ( compute_area_stats, compute_variable_info, + compute_zonal_anomaly_stats, compute_zonal_mean_stats, load_dataset, load_grid, @@ -321,6 +322,62 @@ def _calculate_zonal_mean_local( ) +def calculate_zonal_anomaly( + grid_path: str, + data_path: str, + variable_name: str, + lat_spec: Optional[tuple | float | list] = None, + conservative: bool = False, +) -> Dict[str, Any]: + """Compute the zonal anomaly of a face-centered variable. + + The zonal anomaly is each face value minus the zonal mean of its latitude + band, highlighting departures from the latitudinal background state + (e.g. waves, eddies, stationary anomalies). + + Args: + grid_path: Path to the mesh grid file. + data_path: Path to the data file with variables. + variable_name: Name of the face-centered variable. + lat_spec: Latitude band specification. ``None`` uses the UXarray default + ``(-90, 90, 10)``; a ``(start, end, step)`` tuple or explicit band + edges are also accepted. + conservative: If True, use area-weighted band means. + + Returns: + Dictionary with ``variable_name``, ``conservative``, ``n_face``, + ``stats`` (min/max/mean/std of the anomaly field), ``grid_info``, and + ``_provenance``. + + Raises: + NotImplementedError: If the installed UXarray lacks + ``UxDataArray.zonal_anomaly``. + + Example: + >>> calculate_zonal_anomaly("grid.nc", "data.nc", "temperature") + {"stats": {"min": -12.4, "max": 9.8, ...}, ...} + """ + if not Path(grid_path).exists(): + raise FileNotFoundError(f"Grid file not found: {grid_path}") + if not Path(data_path).exists(): + raise FileNotFoundError(f"Data file not found: {data_path}") + + uxds = load_dataset(grid_path, data_path) + result = compute_zonal_anomaly_stats(uxds, variable_name, lat_spec, conservative) + + return attach_provenance( + result, + tool="calculate_zonal_anomaly", + inputs={ + "grid_path": grid_path, + "data_path": data_path, + "variable_name": variable_name, + "lat_spec": lat_spec, + "conservative": conservative, + }, + ) + + def validate_dataset(grid_path: str, data_path: str) -> Dict[str, Any]: """ Validate a mesh dataset for common data quality issues. diff --git a/src/uxarray_mcp/tools/vector_calc.py b/src/uxarray_mcp/tools/vector_calc.py index 2053144..2193965 100644 --- a/src/uxarray_mcp/tools/vector_calc.py +++ b/src/uxarray_mcp/tools/vector_calc.py @@ -85,6 +85,7 @@ def calculate_gradient( grid_path: str, data_path: str, variable_name: str, + scale_by_radius: bool = False, use_remote: bool = False, endpoint: Optional[str] = None, session_id: Optional[str] = None, @@ -102,6 +103,14 @@ def calculate_gradient( Path to the data file containing the variable. variable_name : str Name of the face-centered scalar variable. + scale_by_radius : bool + If True, divide unit-sphere derivatives by ``uxgrid.sphere_radius`` for + physical units (requires a grid with ``sphere_radius``). Default False + preserves the unit-sphere result. Local execution passes this to the + pinned UXarray directly. The remote worker, which may run an older + UXarray, applies it capability-safely and falls back to the unit sphere + if unsupported; the result reports the ``scale_by_radius`` actually + applied. use_remote : bool If True and an HPC endpoint is configured, execute remotely. endpoint : str, optional @@ -129,12 +138,13 @@ def calculate_gradient( "grid_path": grid_path, "data_path": data_path, "variable_name": variable_name, + "scale_by_radius": scale_by_radius, } def _local(): uxds = load_dataset(grid_path, data_path) return attach_provenance( - compute_gradient(uxds, variable_name), + compute_gradient(uxds, variable_name, scale_by_radius=scale_by_radius), tool="calculate_gradient", inputs=inputs, ) @@ -147,7 +157,9 @@ def _local(): session_id=session_id, local_call=_local, remote_call=lambda agent: _run_sync( - lambda: agent.calculate_gradient_remote(grid_path, data_path, variable_name) + lambda: agent.calculate_gradient_remote( + grid_path, data_path, variable_name, scale_by_radius + ) ), ) @@ -157,6 +169,7 @@ def calculate_curl( data_path: str, u_variable: str, v_variable: str, + scale_by_radius: bool = False, use_remote: bool = False, endpoint: Optional[str] = None, session_id: Optional[str] = None, @@ -181,6 +194,14 @@ def calculate_curl( Zonal (east-west) component, e.g. ``"uReconstructZonal"``. v_variable : str Meridional (north-south) component, e.g. ``"uReconstructMeridional"``. + scale_by_radius : bool + If True, divide the unit-sphere result by ``uxgrid.sphere_radius`` for + physical units (requires a grid with ``sphere_radius``). Default False + preserves the unit-sphere result. Local execution passes this to the + pinned UXarray directly. The remote worker, which may run an older + UXarray, applies it capability-safely and falls back to the unit sphere + if unsupported; the result reports the ``scale_by_radius`` actually + applied. use_remote : bool If True and an HPC endpoint is configured, execute remotely. endpoint : str, optional @@ -212,12 +233,13 @@ def calculate_curl( "data_path": data_path, "u_variable": u_variable, "v_variable": v_variable, + "scale_by_radius": scale_by_radius, } def _local(): uxds = load_dataset(grid_path, data_path) return attach_provenance( - compute_curl(uxds, u_variable, v_variable), + compute_curl(uxds, u_variable, v_variable, scale_by_radius=scale_by_radius), tool="calculate_curl", inputs=inputs, ) @@ -231,7 +253,7 @@ def _local(): local_call=_local, remote_call=lambda agent: _run_sync( lambda: agent.calculate_curl_remote( - grid_path, data_path, u_variable, v_variable + grid_path, data_path, u_variable, v_variable, scale_by_radius ) ), ) diff --git a/tests/conftest.py b/tests/conftest.py index d9f6bf0..bf4733f 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -1,6 +1,7 @@ import sys from unittest.mock import MagicMock +import numpy as np import pytest # Mock uxarray if it's not installed, so we can run logic tests without heavy dependencies @@ -244,3 +245,47 @@ def ensemble_data_files(tmp_path): coords={"sample": [0, 1]}, ).to_netcdf(second) return [str(first), str(second)] + + +@pytest.fixture +def healpix_field_dataset(): + """HEALPix UxDataset with face-centered u, v, and a temperature field.""" + grid = ux.Grid.from_healpix(zoom=2) + n = grid.n_face + rng = np.random.default_rng(7) + return ux.UxDataset( + { + "u": ux.UxDataArray( + xr.DataArray(rng.standard_normal(n), dims=["n_face"]), uxgrid=grid + ), + "v": ux.UxDataArray( + xr.DataArray(rng.standard_normal(n), dims=["n_face"]), uxgrid=grid + ), + "temperature": ux.UxDataArray( + xr.DataArray(250 + 30 * rng.standard_normal(n), dims=["n_face"]), + uxgrid=grid, + ), + }, + uxgrid=grid, + ) + + +@pytest.fixture +def structured_mesh_files(tmp_path): + """Coarse global UGRID grid + face-centered data written to disk. + + ``Grid.from_structured`` produces node coordinates that survive a NetCDF + round-trip, making it a reliable file-based fixture for remapping. + """ + lon = np.arange(0, 360, 20.0) + lat = np.arange(-80, 81, 20.0) + grid = ux.Grid.from_structured(lon=lon, lat=lat) + grid_file = tmp_path / "grid.nc" + data_file = tmp_path / "data.nc" + grid.to_xarray().to_netcdf(grid_file) + + rng = np.random.default_rng(11) + xr.Dataset( + {"temperature": (["n_face"], 250 + 30 * rng.random(grid.n_face))} + ).to_netcdf(data_file) + return str(grid_file), str(data_file) diff --git a/tests/test_advanced_tools.py b/tests/test_advanced_tools.py index cde076d..752c7b7 100644 --- a/tests/test_advanced_tools.py +++ b/tests/test_advanced_tools.py @@ -1,7 +1,11 @@ """Tests for subsetting, comparison, remapping, temporal, and export tools.""" +import warnings from pathlib import Path +import numpy as np +import pytest + from uxarray_mcp.tools import ( calculate_anomaly, calculate_bias, @@ -23,6 +27,7 @@ subset_polygon, write_result, ) +from uxarray_mcp.tools.frontdoor import run_analysis def test_spatial_subset_tools_persist_results(state_dir, synthetic_mesh_with_data): @@ -223,3 +228,35 @@ def test_export_tools_support_result_handles_and_dataset_handles( assert Path(exported_netcdf["output_path"]).exists() assert Path(exported_csv["output_path"]).exists() assert Path(written["output_path"]).exists() + + +def test_remap_to_rectilinear_dispatch(state_dir, structured_mesh_files): + """remap_to_rectilinear returns a regular grid with a result handle.""" + grid_file, data_file = structured_mesh_files + with warnings.catch_warnings(): + warnings.simplefilter("ignore") + result = run_analysis( + operation="remap_to_rectilinear", + grid_path=grid_file, + data_path=data_file, + variable_name="temperature", + target_lon=list(np.arange(0, 360, 30.0)), + target_lat=list(np.arange(-60, 61, 30.0)), + ) + assert result["target_shape"] == [5, 12] + assert set(result["stats"]) == {"min", "max", "mean"} + assert result["result_handle"] + assert "_provenance" in result + + +def test_remap_to_rectilinear_missing_target_coords_raises( + state_dir, structured_mesh_files +): + grid_file, data_file = structured_mesh_files + with pytest.raises(ValueError): + run_analysis( + operation="remap_to_rectilinear", + grid_path=grid_file, + data_path=data_file, + variable_name="temperature", + ) diff --git a/tests/test_calculate_zonal_mean.py b/tests/test_calculate_zonal_mean.py index badae8e..e78b47f 100644 --- a/tests/test_calculate_zonal_mean.py +++ b/tests/test_calculate_zonal_mean.py @@ -1,11 +1,13 @@ -"""Tests for the calculate_zonal_mean tool.""" +"""Tests for the calculate_zonal_mean and zonal_anomaly tools.""" from unittest.mock import MagicMock, patch import numpy as np import pytest +from uxarray_mcp.domain.zonal import compute_zonal_anomaly_stats from uxarray_mcp.tools import calculate_zonal_mean +from uxarray_mcp.tools.frontdoor import run_analysis class TestCalculateZonalMeanUnit: @@ -369,3 +371,44 @@ def test_zonal_mean_with_synthetic_data(self, synthetic_mesh_with_data): assert len(result["latitudes"]) == len(result["zonal_mean_values"]) assert result["conservative"] is False assert result["grid_info"]["n_face"] == 1 + + +class TestZonalAnomaly: + """Tests for the zonal_anomaly operation (per-band departure).""" + + def test_domain_returns_stats(self, healpix_field_dataset): + result = compute_zonal_anomaly_stats(healpix_field_dataset, "temperature") + assert result["variable_name"] == "temperature" + assert set(result["stats"]) == {"min", "max", "mean", "std"} + assert result["n_face"] == int(healpix_field_dataset.uxgrid.n_face) + + def test_anomaly_mean_near_zero(self, healpix_field_dataset): + result = compute_zonal_anomaly_stats(healpix_field_dataset, "temperature") + assert abs(result["stats"]["mean"]) < 1.0 + + def test_missing_variable_raises(self, healpix_field_dataset): + with pytest.raises(ValueError): + compute_zonal_anomaly_stats(healpix_field_dataset, "nope") + + def test_run_analysis_dispatch(self, structured_mesh_files): + grid_file, data_file = structured_mesh_files + result = run_analysis( + operation="zonal_anomaly", + grid_path=grid_file, + data_path=data_file, + variable_name="temperature", + ) + assert "stats" in result + assert "_provenance" in result + + def test_capability_guard_when_unsupported( + self, healpix_field_dataset, monkeypatch + ): + """The domain layer raises a clear error if zonal_anomaly is absent.""" + monkeypatch.delattr( + type(healpix_field_dataset["temperature"]), "zonal_anomaly", raising=False + ) + if hasattr(healpix_field_dataset["temperature"], "zonal_anomaly"): + pytest.skip("could not remove zonal_anomaly for negative test") + with pytest.raises(NotImplementedError): + compute_zonal_anomaly_stats(healpix_field_dataset, "temperature") diff --git a/tests/test_server.py b/tests/test_server.py index 246a4ae..bdafd65 100644 --- a/tests/test_server.py +++ b/tests/test_server.py @@ -26,8 +26,9 @@ EXPECTED_FRONTDOOR = 11 EXPECTED_CONTROL = 12 # 8 session + 4 hpc EXPECTED_CORE_EXTRA = 1 # list_datasets -EXPECTED_PROMPTS = 3 # first_look, vorticity_analysis, hpc_diagnose -EXPECTED_DEFERRED = 30 +EXPECTED_PROMPTS = 7 # first_look, vorticity_analysis, cyclone_structure, +# eddy_activity, model_evaluation, climatology_anomaly, hpc_diagnose +EXPECTED_DEFERRED = 32 # +zonal_anomaly, +remap_to_rectilinear # --------------------------------------------------------------------------- @@ -180,7 +181,15 @@ def test_prompt_tools_registered_in_core(): registry = make_registry(profile="core") tools = registry.list_tools() sep = registry._name_sep - for name in ("first_look", "vorticity_analysis", "hpc_diagnose"): + for name in ( + "first_look", + "vorticity_analysis", + "cyclone_structure", + "eddy_activity", + "model_evaluation", + "climatology_anomaly", + "hpc_diagnose", + ): assert f"prompt{sep}{name}" in tools, f"prompt tool {name} missing" @@ -204,6 +213,39 @@ def test_prompt_tool_returns_text(): assert "/tmp/test.nc" in text +def test_science_workflow_prompts_reference_real_operations(): + """Each guided science prompt should chain real run_analysis operations.""" + from uxarray_mcp.registry import ( + climatology_anomaly, + cyclone_structure, + eddy_activity, + model_evaluation, + ) + + cyclone = cyclone_structure( + "g.nc", "d.nc", "wind", 280.0, 25.0, u_var="u", v_var="v" + ) + assert 'operation="azimuthal_mean"' in cyclone + assert 'operation="curl"' in cyclone # wind components present + + # Without wind, the vorticity step is omitted. + cyclone_no_wind = cyclone_structure("g.nc", "d.nc", "pressure", 280.0, 25.0) + assert 'operation="curl"' not in cyclone_no_wind + + eddy = eddy_activity("g.nc", "d.nc", "z500") + assert 'operation="zonal_anomaly"' in eddy + assert 'operation="gradient"' in eddy + + evaluation = model_evaluation("g.nc", "model.nc", "ref.nc", "t2m") + for op in ("bias", "rmse", "pattern_correlation"): + assert f'operation="{op}"' in evaluation + + clim = climatology_anomaly("d.nc", "t2m", grid_path="g.nc") + assert 'operation="temporal_mean"' in clim + assert 'operation="anomaly"' in clim + assert 'operation="calculate_zonal_mean"' in clim # grid_path provided + + # --------------------------------------------------------------------------- # Policy tags # --------------------------------------------------------------------------- diff --git a/tests/test_vector_calc.py b/tests/test_vector_calc.py index 81d0d0e..a6f541f 100644 --- a/tests/test_vector_calc.py +++ b/tests/test_vector_calc.py @@ -2,7 +2,8 @@ from __future__ import annotations -from unittest.mock import patch +import warnings +from unittest.mock import MagicMock, patch import numpy as np import pytest @@ -323,6 +324,64 @@ def test_accepts_use_remote_endpoint_session_params(self): assert "session_id" in sig.parameters +# --------------------------------------------------------------------------- +# scale_by_radius opt-in +# --------------------------------------------------------------------------- + + +class TestScaleByRadius: + def test_gradient_default_keeps_unit_sphere(self, healpix_wind_dataset): + result = compute_gradient(healpix_wind_dataset, "temperature") + assert result["scale_by_radius"] is False + + def test_curl_default_keeps_unit_sphere(self, healpix_wind_dataset): + result = compute_curl(healpix_wind_dataset, "u", "v") + assert result["scale_by_radius"] is False + + def test_gradient_records_scale_by_radius_flag(self, healpix_wind_dataset): + with warnings.catch_warnings(): + warnings.simplefilter("ignore") # grid has no sphere_radius + result = compute_gradient( + healpix_wind_dataset, "temperature", scale_by_radius=True + ) + assert result["scale_by_radius"] is True + + def test_remote_gradient_threads_scale_by_radius(self): + """The remote dispatch must forward scale_by_radius to the agent.""" + from uxarray_mcp.tools import vector_calc + + agent = MagicMock() + agent.config.endpoint_id = "fake-endpoint" + agent.config.endpoint_name = "fake" + agent.config.timeout_seconds = 60 + agent.calculate_gradient_remote.return_value = { + "components": [], + "component_stats": {}, + "n_face": 1, + "scale_by_radius": True, + "_provenance": {"warnings": []}, + } + + with ( + patch("uxarray_mcp.remote.agent.get_agent", return_value=agent), + patch.object( + vector_calc, "_endpoint_manager_is_up", return_value=(True, "ok") + ), + patch.object(vector_calc, "_run_sync", side_effect=lambda f: f()), + ): + vector_calc.calculate_gradient( + "/hpc/grid.nc", + "/hpc/data.nc", + "t", + scale_by_radius=True, + use_remote=True, + endpoint="improv", + ) + + args, kwargs = agent.calculate_gradient_remote.call_args + assert (True in args) or (kwargs.get("scale_by_radius") is True) + + # --------------------------------------------------------------------------- # Server registration # --------------------------------------------------------------------------- diff --git a/uv.lock b/uv.lock index 43e961f..7cbb38e 100644 --- a/uv.lock +++ b/uv.lock @@ -2248,14 +2248,14 @@ wheels = [ [[package]] name = "toolregistry-server" -version = "0.3.3" +version = "0.4.0" source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "toolregistry" }, ] -sdist = { url = "https://files.pythonhosted.org/packages/69/43/f452d737c86ca18f234bd6ccac7a5051a7068da3c8a8e3c2715edfa5528c/toolregistry_server-0.3.3.tar.gz", hash = "sha256:1db4ad838ea8dd2fd7334433249f9cc226295aeb351e83ea63bee8c774d109f2", size = 63105, upload-time = "2026-05-31T09:16:04.505Z" } +sdist = { url = "https://files.pythonhosted.org/packages/2d/bb/3678a93918a49faf225439511b1ff49fa65844cbc5ca4d7dc59ff494c793/toolregistry_server-0.4.0.tar.gz", hash = "sha256:c79f355dfba3a4ce1710cd2e24edbd217c2b41b4717d82782534563fe30122a1", size = 71492, upload-time = "2026-06-22T10:28:49.163Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/ca/5b/bf7eaa663be13ff6396b88d52f85fa7ebb627995c0764016b36815109565/toolregistry_server-0.3.3-py3-none-any.whl", hash = "sha256:09198b80838307e17195ecad589509402ac417fdbfe66d13e3c6051ad6057ff3", size = 50051, upload-time = "2026-05-31T09:16:03.579Z" }, + { url = "https://files.pythonhosted.org/packages/fa/35/b67b516511af5431cbdd5bd5c91775bcf3e1ed107d39b9bcee1e8acb80ce/toolregistry_server-0.4.0-py3-none-any.whl", hash = "sha256:7933c29050bd93e04b56ffec999d71489adda0670f7c1ee9b9589944d4f9fc89", size = 58369, upload-time = "2026-06-22T10:28:48.024Z" }, ] [package.optional-dependencies] @@ -2402,7 +2402,7 @@ wheels = [ [[package]] name = "uxarray" -version = "2026.4.1" +version = "2026.6.0" source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "antimeridian" }, @@ -2430,9 +2430,9 @@ dependencies = [ { name = "spatialpandas" }, { name = "xarray" }, ] -sdist = { url = "https://files.pythonhosted.org/packages/1c/f6/61ff9b8d6b9d4223d3798ee14129c27cc9fad472529b9ea6891e8ef9f75c/uxarray-2026.4.1.tar.gz", hash = "sha256:7934f85430b791186a684022d0e736e26275ac38f597c38b7000c8bea889ab0e", size = 190414, upload-time = "2026-04-23T01:45:32.207Z" } +sdist = { url = "https://files.pythonhosted.org/packages/8f/13/cbf510afda42e223f7ae50046e4f657c437b09bcf796d96c930a5e523bc0/uxarray-2026.6.0.tar.gz", hash = "sha256:f83d34fd8dae54fd8d87f59ade921f70396bce538a8efe73d20d39b63f9b7f5e", size = 202826, upload-time = "2026-06-18T22:22:05.925Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/a9/f7/f4d5689dc20b0c112650f225e68c3ec3e2d57bf2a30a6696d74e7d1a5f1c/uxarray-2026.4.1-py3-none-any.whl", hash = "sha256:ed50b3a7871b7359823095c5f7220ab72f885a312cc774e39c43836333a83e16", size = 213343, upload-time = "2026-04-23T01:45:30.545Z" }, + { url = "https://files.pythonhosted.org/packages/38/41/c8a20e78aa886cdfc87979dc6851ea7c67faf989e42a405678f14dc53328/uxarray-2026.6.0-py3-none-any.whl", hash = "sha256:4a20777953dbd0eba4f8f7add28091d6c3f28ec1789b6638c60bf5865782b8c3", size = 227792, upload-time = "2026-06-18T22:22:04.541Z" }, ] [[package]] @@ -2480,9 +2480,9 @@ requires-dist = [ { name = "pyyaml", specifier = ">=6.0" }, { name = "sphinx", marker = "extra == 'docs'", specifier = ">=7.0" }, { name = "sphinx-book-theme", marker = "extra == 'docs'", specifier = ">=1.1.0" }, - { name = "toolregistry-server", extras = ["mcp"], specifier = ">=0.3.3" }, - { name = "toolregistry-server", extras = ["openapi"], marker = "extra == 'openapi'", specifier = ">=0.3.3" }, - { name = "uxarray", specifier = ">=2025.12.0" }, + { name = "toolregistry-server", extras = ["mcp"], specifier = ">=0.4.0" }, + { name = "toolregistry-server", extras = ["openapi"], marker = "extra == 'openapi'", specifier = ">=0.4.0" }, + { name = "uxarray", specifier = ">=2026.6.0" }, ] provides-extras = ["openapi", "hpc", "docs"]