integrate_query

Functions related to querying posterior realizations in the INTEGRATE module.

Query posterior realizations based on geophysical constraints.

This module provides tools to compute probabilities that posterior realizations from Bayesian inversion satisfy user-defined constraints (e.g., thickness of lithology classes, resistivity thresholds).

integrate.integrate_query.cells_to_polygon(cells, mask)

Union of the masked Voronoi cells -> a single shapely Polygon (largest part).

integrate.integrate_query.find_coherent_area(X, Y, P, p_min, X_center=None, Y_center=None, max_area_m2=None, seed_index=None, xy=None, hull_ratio=0.1, edge_buffer=None, cell_area_k=6.0, elong_max=4.0)

Grow one coherent high-probability area and return its indices + polygon.

Full geometry pipeline in one call. From the sounding coordinates a Voronoi tessellation defines the adjacency graph and the per-sounding cell areas; the cells are clipped to a concave hull of the survey so their areas are meaningful. Edge-affected soundings (huge, badly constrained cells on the sparse survey rim) are dropped – they can neither seed a region nor be reached by growth. A connected region is then grown with grow_connected_region() and its boundary returned as a polygon.

Parameters:

  • X (ndarray (N,) or (N, 1)) – Sounding coordinates [m]; column vectors (raw HDF5) are flattened.

  • Y (ndarray (N,) or (N, 1)) – Sounding coordinates [m]; column vectors (raw HDF5) are flattened.

  • P (ndarray (N,) or (N, 1)) – Per-sounding probability/score, aligned with X/Y. Soundings failing the edge filter are made unreachable internally.

  • p_min (float) – Inclusion probability cutoff (the region “width”/area knob).

  • X_center (float, optional) – Seed the region near this (x, y) – the nearest kept sounding is used as the center. If both are None (and no seed_index/xy), the sounding with the maximum P is used. The resolved center is returned in the output as X_center / Y_center.

  • Y_center (float, optional) – Seed the region near this (x, y) – the nearest kept sounding is used as the center. If both are None (and no seed_index/xy), the sounding with the maximum P is used. The resolved center is returned in the output as X_center / Y_center.

  • max_area_m2 (float, optional) – Hard cap on the region area [m^2]; None = no cap.

  • seed_index (int, optional) – Seed at this exact sounding index (bypasses X_center/Y_center).

  • xy ((float, float), optional) – Deprecated alias for (X_center, Y_center). At most one seed selector (X_center/Y_center/xy, seed_index) may be set.

  • hull_ratio (float, optional) – Concave-hull tightness for the survey outline (0 = tight, 1 = convex).

  • edge_buffer (optional) – Edge-filter parameters; see flag_edge_cells().

  • cell_area_k (optional) – Edge-filter parameters; see flag_edge_cells().

  • elong_max (optional) – Edge-filter parameters; see flag_edge_cells().

Returns:

region – idx (int (M,) region sounding indices), polygon (shapely Polygon outline), seed (int), area (float [m^2]), order (list[int] growth path), X_center / Y_center (float – the sounding coordinates actually used as the region center/seed: the requested X_center/Y_center snapped to its nearest sounding, or the argmax-P sounding if no center was given), plus the reusable Voronoi scaffold: vor (scipy Voronoi), neighbors (list[list[int]] adjacency), cells (per-sounding clipped cell polygons), cell_area ((N,) float), boundary (concave survey outline Polygon), good ((N,) bool edge-filter keep mask) and mask ((N,) bool) – mask[i] = i in idx. For several regions of interest share this scaffold across calls, or pass vor/neighbors/cells/cell_area/boundary/good into grow_connected_region() to grow extra regions without recomputing the geometry.

Return type:

dict

Examples

>>> region = find_coherent_area(X, Y, P_raw, p_min=0.5)
>>> region['idx'], region['polygon'].area, region['area']
>>> # extra region on the same scaffold, no recompute:
>>> idx2, area2, _ = grow_connected_region(
...     np.where(region['good'], P_raw, -np.inf),
...     region['neighbors'], region['cell_area'], p_min=0.5,
...     seed=region['seed'])
integrate.integrate_query.flag_edge_cells(X, Y, vor, cells, boundary, edge_buffer=None, k=6.0, elong_max=None)

Boolean good mask – False for edge-affected soundings on the sparse survey rim.

A sounding is dropped if its Voronoi cell is unbounded (a convex-hull vertex), or if it lies within edge_buffer of the survey outline AND its cell is either much larger than the interior-median cell (> k x) or very elongated (perimeter^2 / (4 pi area) > elong_max). Large cells that are NOT near the rim (interior data gaps) are kept.

Returns (good, edge_affected, info).

integrate.integrate_query.get_prior_model_info(f_prior_h5, im)

Return metadata for prior model im.

Parameters:

  • f_prior_h5 (str) – Path to the prior HDF5 file.

  • im (int) – Model index.

Returns:

info – Keys: ‘name’, ‘is_discrete’, ‘z’, ‘class_id’, ‘class_name’.

Return type:

dict

integrate.integrate_query.grow_connected_region(P, neighbors, cell_area, p_min, max_area_m2=None, seed=None)

Grow one connected high-probability region outward from a seed sounding.

Starting at seed, repeatedly add the adjacent sounding with the highest P, keeping only those with P >= p_min. Growth stops on its own when no untried adjacent sounding is likely enough, or when the accumulated area reaches max_area_m2 (if given).

This is the pure graph step – it needs a pre-built adjacency graph and per-sounding areas but no Voronoi/shapely. Use find_coherent_area() for the full geometry pipeline (Voronoi + edge filter + boundary polygon).

Parameters:

  • P (ndarray (N,)) – Per-sounding score/probability; NaN/-inf soundings are never added (use this to make edge-dropped soundings unreachable).

  • neighbors (list of iterable of int) – Adjacency graph; neighbors[i] = indices of soundings adjacent to i.

  • cell_area (ndarray (N,)) – Per-sounding representative area [m^2]; summed into area.

  • p_min (float) – Inclusion cutoff – the region grows while the best neighbour has P >= p_min.

  • max_area_m2 (float, optional) – Hard cap on the accumulated area [m^2]; None = no cap.

  • seed (int, optional) – Starting sounding index; None -> argmax of P.

Returns:

  • indices (ndarray (M,) int) – Indices of the soundings in the grown region (includes the seed).

  • area (float) – Accumulated area [m^2], sum of cell_area over indices.

  • order (list of int) – indices in the order they were added (growth path; order[0] = seed).

integrate.integrate_query.load_query(path)

Load a query dict from a JSON file.

Parameters:

path (str) – Input JSON file path.

Returns:

query – Query definition dictionary.

Return type:

dict

integrate.integrate_query.prior_describe(f_prior_h5)

Print a human-readable summary of all models in a prior HDF5 file.

Parameters:

f_prior_h5 (str) – Path to the prior HDF5 file.

Examples

>>> ig.prior_describe('prior.h5')
Prior file: prior.h5
N realizations: 1000000
  im=1  Resistivity   CONTINUOUS   depth 0–89 m  (89 layers)
  im=2  Lithology     DISCRETE     depth 0–89 m  (89 layers)
          class 1 = Sand
          class 2 = Grus
  im=3  Waterlevel    SCALAR
integrate.integrate_query.query(f_post_h5, query_dict)

Dispatcher: route to query_probability() or query_percentile() based on query_dict.

If query_dict contains a "metric" key, calls query_percentile(). Otherwise calls query_probability() (backward compatible with all existing "constraints"-based dicts).

Parameters:
Returns:

See the delegated function for details.

Return type:

result, meta

integrate.integrate_query.query_from_text(text, f_prior_h5, model='anthropic/claude-sonnet-4-6', api_key=None, max_tokens=4096, verbose=False, system_prompt=None)

Translate a natural-language query into a query dict using an LLM.

Uses LiteLLM to interpret the user’s text query in the context of the available prior models and the integrate query schema, returning a query dict and a plain-English interpretation of what the LLM understood.

Parameters:

  • text (str) – Natural language description of the query, e.g. “What is the probability that cumulative clay thickness exceeds 10 m?”.

  • f_prior_h5 (str) – Path to the prior HDF5 file. Model metadata (class names, depth ranges, discrete/continuous type) is read automatically and included in the LLM prompt so the model knows what constraints are valid.

  • model (str, optional) – LiteLLM model string (default: ‘anthropic/claude-sonnet-4-6’). Any LiteLLM-supported model works, e.g. ‘openai/gpt-4o’.

  • api_key (str, optional) – Provider API key. If None, the relevant environment variable (e.g. ANTHROPIC_API_KEY) is used.

  • verbose (bool, optional) – If True, print the system prompt and LLM response for inspection.

  • system_prompt (str, optional) – Custom system prompt overriding the default built from the prior file. Use to hand-tune the translation instructions (model list and schema context). The effective prompt is still returned for inspection.

Returns:

  • query_dict (dict) – Query dict ready to pass to ig.query(f_post_h5, query_dict).

  • interpretation (str) – Plain English confirmation of what the LLM understood the query to mean. Check this before running ig.query() to catch misunderstandings cheaply.

  • system_prompt (str) – The full system prompt sent to the LLM. Useful for inspection and debugging.

Raises:

  • ImportError – If the litellm package is not installed.

  • ValueError – If the LLM reports the query is unsupported, or if the response cannot be parsed as valid JSON.

Notes

Requires either the api_key parameter or the relevant provider environment variable to be set. Install the dependency with: pip install litellm

Examples

>>> import integrate as ig
>>> query_dict, interpretation, system_prompt = ig.query_from_text(
...     "Probability that cumulative clay thickness > 10 m within 0-30 m",
...     f_prior_h5='prior.h5',
...     api_key='sk-ant-...',
... )
>>> print(interpretation)
>>> P, meta = ig.query('posterior.h5', query_dict)
>>> ig.query_plot(P, meta)
integrate.integrate_query.query_percentile(f_post_h5, query_dict)

Compute per-data-point percentiles of a metric over posterior realizations.

Rather than asking “what fraction of realizations satisfy condition X?”, this asks “what is the p5/p50/p95 of metric X across realizations?”. The metric is defined by the same fields as a probability constraint, minus the comparison fields (thickness_comparison, thickness_threshold, negate).

Parameters:

  • f_post_h5 (str) – Path to the posterior HDF5 file.

  • query_dict (str or dict) – Path to a JSON file, or a dict with a "metric" key and an optional "percentiles" key (default [5, 50, 95]).

Returns:

  • percentile_values (ndarray (N_data, n_percentiles)) – Requested percentile values for each data location.

  • meta (dict) – Keys: ‘X’, ‘Y’, ‘N_data’, ‘N_post’, ‘i_use’, ‘percentiles’.

Examples

>>> query_def = {
...     "metric": {
...         "im": 2, "classes": [1, 2],
...         "thickness_mode": "cumulative",
...         "depth_max": 30.0
...     },
...     "percentiles": [5, 50, 95]
... }
>>> pct_values, meta = query_percentile('f_post.h5', query_def)
>>> # pct_values shape: (N_data, 3) — p5, p50, p95 per location
integrate.integrate_query.query_percentile_plot(percentile_values, meta, query_text=None, interpretation=None, text_panel=False, hardcopy=False, **kwargs)

Plot one probability map per requested percentile as side-by-side subplots.

Parameters:

  • percentile_values (ndarray (N_data, n_percentiles)) – Output of query_percentile().

  • meta (dict) – Metadata dict from query_percentile() containing ‘X’, ‘Y’, ‘percentiles’.

  • query_text (str, optional) – Original query string — shown as figure suptitle.

  • interpretation (str, optional) – LLM interpretation string — shown below query_text if provided.

  • text_panel (bool, optional) – If True, add a narrow text column to the right of the maps.

  • hardcopy (bool or str, optional) – Save figure to disk. True → ‘query_percentile_plot.png’; a string is used as the filename (.png appended if no extension).

  • **kwargs – All remaining keyword arguments are forwarded to plot_xy(), giving full control over cmap, clim, uselog, colorbar, colorbar_label, plotPoints, plotPoints_color, plotPoints_marker, s, etc. clim defaults to [percentile_values.min(), percentile_values.max()] so that all subplots share the same colour scale. cmap defaults to 'viridis'.

Returns:

fig

Return type:

matplotlib Figure

integrate.integrate_query.query_plot(P, meta, ip=None, query_dict=None, f_prior_h5=None, f_post_h5=None, title=None, query_text=None, interpretation=None, text_panel=False, hardcopy=False, **kwargs)

Plot query results and optionally detailed model visualization for a data point.

If ip is None, displays the XY probability map showing P(x, y). If ip is provided (together with query_dict and f_prior_h5/f_post_h5), skips the probability map and shows only the detailed single-point visualization of all posterior realizations and the query-matching subset.

Parameters:

  • P (ndarray (N_data,)) – Probability array from query().

  • meta (dict) – Metadata dict from query() containing ‘X’, ‘Y’, ‘i_use’, ‘i_use_query’.

  • ip (int, optional) – Data point index to visualize in detail. If None, only shows probability map.

  • query_dict (dict, optional) – Query dict used in query(). Required for detailed visualization.

  • f_prior_h5 (str, optional) – Path to prior HDF5 file. If not provided, will be extracted from f_post_h5.

  • f_post_h5 (str, optional) – Path to posterior HDF5 file. Used to automatically extract prior file path if f_prior_h5 is not provided.

  • title (str, optional) – Custom title for the probability map. If None, a title is built from query_text and interpretation (if provided), or ‘Query Probability Map’.

  • query_text (str, optional) – The original natural-language query string. Shown in the figure title, or in the text panel if text_panel=True.

  • interpretation (str, optional) – The LLM interpretation string returned by query_from_text(). Shown as a second line in the figure title, or in the text panel if text_panel=True.

  • text_panel (bool, optional) – If True and query_text or interpretation is provided, adds a narrow text column to the right of the probability map. The query text appears at the top and the interpretation below. Default False.

  • hardcopy (bool or str, optional) – Save the probability map figure. If True, saves as ‘query_plot.png’. If a string, uses that as the filename (a ‘.png’ extension is appended if the string has no extension). Default False.

  • **kwargs – All remaining keyword arguments are forwarded to plot_xy(), giving full control over cmap, clim, uselog, colorbar, colorbar_label, plotPoints, plotPoints_color, plotPoints_marker, s, etc. cmap defaults to 'hot_r' and clim defaults to [0, 1].

Examples

>>> P, meta = query(f_post_h5, query_def)
>>> query_plot(P, meta)  # Just probability map
>>> query_plot(P, meta, title='Custom Query Title')  # Custom title
>>> query_plot(P, meta, ip=1000, query_dict=query_def, f_post_h5='posterior.h5')
>>> query_plot(P, meta, ip=1000, query_dict=query_def, f_prior_h5='prior.h5')
>>> # With LLM query text and interpretation:
>>> query_dict, interp = ig.query_from_text(text, f_prior_h5)
>>> P, meta = ig.query(f_post_h5, query_dict)
>>> ig.query_plot(P, meta, query_text=text, interpretation=interp)
integrate.integrate_query.query_probability(f_post_h5, query_dict)

Compute per-data-point probability that posterior realizations satisfy a query.

Parameters:

  • f_post_h5 (str) – Path to the posterior HDF5 file.

  • query_dict (str or dict) – Path to a JSON file, or a dict with a "constraints" key.

Returns:

  • P (ndarray (N_data,)) – Probability [0, 1] for each data location.

  • meta (dict) – Keys: ‘X’, ‘Y’, ‘N_data’, ‘N_post’, ‘i_use’, ‘i_use_query’.

Examples

>>> query_def = {
...     "constraints": [{
...         "im": 2, "classes": [2],
...         "thickness_mode": "cumulative",
...         "thickness_comparison": ">",
...         "thickness_threshold": 10.0,
...         "depth_min": 0.0, "depth_max": 30.0
...     }]
... }
>>> P, meta = query_probability('f_post.h5', query_def)
integrate.integrate_query.query_test_llm(model='anthropic/claude-sonnet-4-6', api_key=None, verbose=1)

Test whether a given LLM model and API key are working correctly.

Sends a minimal JSON-generation prompt and checks that the response is valid JSON. Prints a summary and returns a status dict.

Parameters:

  • model (str, optional) – LiteLLM model string (default: ‘anthropic/claude-sonnet-4-6’).

  • api_key (str, optional) – Provider API key. If None, the relevant environment variable is used.

  • verbose (int, optional) – 0 = silent, 1 = summary only (default), 2 = full response included.

Returns:

result – Keys: ‘ok’ (bool), ‘model’, ‘response’ (str or None), ‘error’ (str or None).

Return type:

dict

integrate.integrate_query.region_volumes(pct, area)

Volume at each percentile, summed over an area’s soundings.

Parameters:

  • pct (ndarray (N_sounding, N_pct)) – Per-sounding posterior thickness percentiles [m] (e.g. the array returned by ig.query(f_post_h5, {"metric": ..., "percentiles": ...})).

  • area (dict) – Must provide mask ((N,) bool – which soundings to sum over) and cell_area ((N,) float [m^2] – the representative area of each sounding). A dict returned by find_coherent_area() works directly; so does any {'mask': ..., 'cell_area': ...} (e.g. one built from the soundings inside a hand-drawn polygon).

Returns:

sum(pct[mask] * cell_area[mask]) – one volume [m^3] per percentile.

Return type:

ndarray (N_pct,)

integrate.integrate_query.save_query(query, path)

Save a query dict to a JSON file.

Parameters:

  • query (dict) – Query definition dictionary.

  • path (str) – Output JSON file path.

integrate.integrate_query.title_from_json(file_json, f_prior_h5=None, model='anthropic/claude-sonnet-4-6', api_key=None, showInfo=1)

Return a plain-language description of what a query JSON dict will do.

Uses an LLM to produce a short human-readable summary suitable for a figure title or log message. If the LLM is unavailable (missing package, no API key, network error), returns an empty string.

Parameters:

  • file_json (str or dict) – Path to a query JSON file, or a query dict directly (e.g. from ig.load_query()).

  • f_prior_h5 (str, optional) – Path to the prior HDF5 file. When provided, real model names, depth ranges, and class labels are included in the prompt so the description uses geological names instead of numeric model/class IDs.

  • model (str, optional) – LiteLLM model string (default: ‘anthropic/claude-sonnet-4-6’).

  • api_key (str, optional) – Provider API key. If None, the relevant environment variable is used.

  • showInfo (int, optional) – 0 = silent; 1 = print a message when the LLM cannot be reached (default); 2 = also print the exception detail.

Returns:

description – One-sentence plain-English summary of the query, or an empty string if the LLM could not be reached.

Return type:

str

Examples

>>> description = ig.title_from_json('my_query.json')
>>> description = ig.title_from_json('my_query.json', f_prior_h5='prior.h5')
>>> query = ig.load_query('query_ex1.json')
>>> title = ig.title_from_json(query, f_prior_h5='prior.h5')
>>> title = ig.title_from_json(query, showInfo=0)  # silent on failure
integrate.integrate_query.voronoi_cells_ordered(X, Y, boundary)

Per-sounding Voronoi cell polygons, in input order, clipped to boundary.

Always returns exactly len(X) entries with cells[i] belonging to sounding i (each cell is matched to the sounding point it contains, or the nearest sounding as a fallback), so cells / cell_area stay aligned with X / Y / P.

integrate.integrate_query.voronoi_graph(X, Y)

(vor, neighbors): scipy Voronoi object + neighbour-index lists (cells sharing an edge).

Fails on exactly-duplicated coordinates – deduplicate X, Y first if that happens.