diff --git a/tools/utils/ats_xdmf.py b/tools/utils/ats_xdmf.py
index d43979022..4560f4062 100644
--- a/tools/utils/ats_xdmf.py
+++ b/tools/utils/ats_xdmf.py
@@ -68,18 +68,99 @@ def find_domains(directory='.'):
f"No ats_vis_*data.h5 files found in '{directory}'.")
return sorted(found, key=natural_sort_key)
+def resolve_vis_input(domain_or_path, directory='.', prefix='ats_vis'):
+ """Resolve (directory, domain, filename) from a domain name or a filepath.
+
+ Accepts three calling conventions:
+
+ 1. Domain name — ``resolve_vis_input('surface')``
+ Returns ``(directory, 'surface', 'ats_vis_surface_data.h5')``.
+
+ 2. Domain name with explicit directory / prefix overrides
+ — ``resolve_vis_input('surface', directory='/run', prefix='myprefix')``
+ Returns ``('/run', 'surface', 'myprefix_surface_data.h5')``.
+
+ 3. File path — ``resolve_vis_input('/run/ats_vis_surface_data.h5')``
+ The argument is treated as a filepath when it contains a path separator
+ or ends in ``.h5`` or ``.xmf``. ``directory`` and ``prefix`` args are
+ ignored. ``domain`` is returned as ``None`` because the filename stem
+ is an opaque combination of prefix and domain (they cannot be separated
+ without additional information). For ``.xmf`` files the companion
+ ``_data.h5`` file is located in the same directory.
+
+ Parameters
+ ----------
+ domain_or_path : str
+ ATS domain name OR path to any vis-related file.
+ directory : str, optional
+ Input directory, used only in convention 1 / 2. Default ``'.'``.
+ prefix : str, optional
+ Filename prefix, used only in convention 1 / 2. Default ``'ats_vis'``.
+
+ Returns
+ -------
+ directory : str
+ domain : str or None
+ filename : str
+ Basename of the data ``.h5`` file (not a full path).
+ """
+ arg = domain_or_path
+
+ # Detect filepath: contains a separator, or ends with a known extension.
+ is_path = (os.sep in arg or '/' in arg or
+ arg.endswith('.h5') or arg.endswith('.xmf'))
+
+ if is_path:
+ directory = os.path.dirname(os.path.abspath(arg)) if os.path.dirname(arg) else '.'
+ basename = os.path.basename(arg)
+
+ # Normalise .xmf files to their companion _data.h5
+ if basename.endswith('.xmf'):
+ # Per-cycle: foo_data.h5.42.xmf -> foo_data.h5
+ # VisIt master: foo_data.VisIt.xmf -> foo_data.h5
+ # General rule: strip everything from the first '.xmf'-adjacent
+ # suffix by finding the _data.h5 sibling.
+ import glob as _glob
+ candidates = _glob.glob(os.path.join(directory, '*_data.h5'))
+ # Keep only candidates whose basename is a prefix of this xmf name
+ matches = [os.path.basename(c) for c in candidates
+ if basename.startswith(os.path.basename(c))]
+ if len(matches) == 1:
+ basename = matches[0]
+ elif not matches:
+ raise RuntimeError(
+ f"Cannot find companion _data.h5 for {arg!r} "
+ f"in {directory!r}.")
+ else:
+ raise RuntimeError(
+ f"Ambiguous companion _data.h5 for {arg!r}: {matches}")
+
+ # domain is None — stem is prefix+domain combined, indistinguishable
+ return directory, None, basename
+
+ else:
+ # Convention 1 / 2: construct filename from prefix and domain
+ domain = arg
+ if domain == 'domain':
+ fname = f'{prefix}_data.h5'
+ else:
+ fname = f'{prefix}_{domain}_data.h5'
+ return directory, domain, fname
+
+
def time_unit_conversion(value, input_unit, output_unit):
time_in_seconds = {
- 'yr': 365.25 * 24 * 3600,
+ 'y': 365.25 * 24 * 3600,
'noleap': 365 * 24 *3600,
'd': 24 * 3600,
- 'hr': 3600,
+ 'h': 3600,
+ 'min': 60,
's': 1
}
if input_unit not in time_in_seconds:
- raise ValueError("Invalid input time unit : must be one of 'yr', 'noleap', 'd', 'hr', or 's'")
+ raise ValueError("Invalid input time unit : must be one of 'y', 'noleap', 'd', 'h', 'min', or 's'")
if output_unit not in time_in_seconds:
- raise ValueError("Invalid output time unit : must be one of 'yr', 'noleap', 'd', 'hr', or 's'")
+ raise ValueError("Invalid output time unit : must be one of 'y', 'noleap', 'd', 'h', 'min', or 's'")
value2sec = value * time_in_seconds[input_unit]
output_value = value2sec / time_in_seconds[output_unit]
@@ -143,8 +224,8 @@ def __init__(self, directory='.', domain=None, filename=None, mesh_filename=None
else:
warnings.warn(
f"HDF5 file {self.fname!r} has no 'time unit' attribute; "
- "assuming 'yr'. This file may be from an old version of ATS.")
- self.input_time_unit = 'yr'
+ "assuming 'y'. This file may be from an old version of ATS.")
+ self.input_time_unit = 'y'
self.output_time_unit = output_time_unit if output_time_unit is not None else self.input_time_unit
diff --git a/tools/utils/combine_vis.py b/tools/utils/combine_vis.py
deleted file mode 100644
index ace6c5010..000000000
--- a/tools/utils/combine_vis.py
+++ /dev/null
@@ -1,566 +0,0 @@
-#!/usr/bin/env python
-"""Combine ATS visualization output from restarted/continuation runs into a single
-self-contained XDMF dataset.
-
-Takes N input directories (in chronological order) and produces one output
-directory containing:
- - A combined HDF5 data file with renumbered timestep keys (0, 1, 2, ...)
- - One per-step XMF per selected cycle
- - A master VisIt.xmf
- - The mesh H5 and XMF files (copied from the first run that has them)
-
-Overlapping cycles at restart boundaries are deduplicated: for each run except
-the last, any cycle whose time >= the start time of the next run is dropped.
-
-Examples
---------
-Combine three restart directories for the "domain" domain::
-
- combine_vis.py domain run0/ run1/ run2/ --output combined/
-
-Same for the surface domain, with a 1-hour overlap tolerance::
-
- combine_vis.py surface run0/ run1/ run2/ --output combined/ --eps 3600 --time-unit s
-"""
-
-import sys
-import os
-import shutil
-import argparse
-import warnings
-import xml.etree.ElementTree as ET
-
-import numpy as np
-import h5py
-
-try:
- sys.path.insert(0, os.path.join(os.environ['ATS_SRC_DIR'], 'tools', 'utils'))
-except KeyError:
- pass
-
-import ats_xdmf
-
-
-# ---------------------------------------------------------------------------
-# Helpers
-# ---------------------------------------------------------------------------
-
-def _validate_time_ordering(vis_files, directories, time_unit):
- """Warn loudly if runs are not in chronological order.
-
- The only condition we can guarantee is that start[i+1] > start[i].
- """
- for i in range(len(vis_files) - 1):
- t0 = vis_files[i].times[0]
- t1 = vis_files[i + 1].times[0]
- if t1 <= t0:
- warnings.warn(
- f"\n *** TIME ORDERING WARNING ***\n"
- f" Run '{directories[i+1]}' starts at t={t1:.6g} {time_unit}\n"
- f" but run '{directories[i]}' starts at t={t0:.6g} {time_unit}.\n"
- f" Directories may be out of order — combined output may be wrong.",
- stacklevel=2)
-
-
-def _select_cycles(vis_files, directories, eps):
- """For each run, select cycles to include after deduplication.
-
- Parameters
- ----------
- vis_files : list of ats_xdmf.VisFile
- directories : list of str
- eps : float
- Overlap tolerance in the VisFile's output_time_unit. Cycles from run i
- with time >= (start_time_of_run_i+1 - eps) are dropped.
-
- Returns
- -------
- list of (vf, directory, selected_cycle_strs, selected_times)
- """
- run_data = []
- for i, (vf, directory) in enumerate(zip(vis_files, directories)):
- cycles = list(vf.cycles)
- times = np.array(vf.times)
-
- if i < len(vis_files) - 1:
- next_start = vis_files[i + 1].times[0]
- cutoff = next_start - eps
- mask = times < cutoff
- sel_cycles = [c for c, m in zip(cycles, mask) if m]
- sel_times = times[mask]
- else:
- sel_cycles = cycles
- sel_times = times
-
- if len(sel_cycles) == 0:
- warnings.warn(
- f"Run '{directory}': all cycles filtered out after deduplication. "
- "Check that directories are in chronological order and that --eps is appropriate.")
- else:
- run_data.append((vf, directory, sel_cycles, sel_times))
-
- return run_data
-
-
-def _compute_selected_vars(vis_files, include_vars, exclude_vars):
- """Return the intersection of variable sets across runs, then apply filters.
-
- Parameters
- ----------
- vis_files : list of ats_xdmf.VisFile
- include_vars : list of str or None
- exclude_vars : list of str or None
-
- Returns
- -------
- list of str
- """
- var_sets = [set(vf.variables()) for vf in vis_files]
- common_vars = set.intersection(*var_sets)
-
- # Warn about any per-run differences
- all_vars = set.union(*var_sets)
- if common_vars != all_vars:
- for i, (vf, vs) in enumerate(zip(vis_files, var_sets)):
- missing = common_vars - vs
- if missing:
- warnings.warn(
- f"Run {i}: missing variables present in other runs: {sorted(missing)}. "
- "These will be excluded from the combined output.")
-
- # Apply user filters using the first VisFile's matching logic, then restrict
- # to common_vars.
- filtered = vis_files[0].variables(names=include_vars, exclude=exclude_vars)
- selected = [v for v in filtered if v in common_vars]
-
- if not selected:
- raise RuntimeError(
- "No variables selected after filtering. Check --include / --exclude.")
-
- return selected
-
-
-def _write_combined_h5(run_data, out_h5_path, selected_vars):
- """Write combined HDF5 with sequentially renumbered cycle keys.
-
- Parameters
- ----------
- run_data : list of (vf, directory, sel_cycles, sel_times)
- out_h5_path : str
- selected_vars : list of str
-
- Returns
- -------
- steps : list of (run_idx, old_key_str, new_key_int, h5_native_time)
- One entry per written cycle step.
- """
- steps = []
- new_key = 0
-
- with h5py.File(out_h5_path, 'w') as dst:
- # Propagate time unit from first run
- first_vf = run_data[0][0]
- if 'time unit' in first_vf.d.attrs:
- dst.attrs['time unit'] = first_vf.d.attrs['time unit']
-
- # Create variable groups up front
- for var in selected_vars:
- dst.create_group(var)
-
- for run_idx, (vf, directory, sel_cycles, sel_times) in enumerate(run_data):
- src = vf.d
- for old_key_str in sel_cycles:
- new_key_str = str(new_key)
- h5_native_time = None
-
- for var in selected_vars:
- if var not in src:
- warnings.warn(f"Variable '{var}' missing in '{directory}'; skipping.")
- continue
- if old_key_str not in src[var]:
- warnings.warn(
- f"Cycle key '{old_key_str}' missing for variable '{var}' "
- f"in '{directory}'; skipping.")
- continue
-
- ds = dst[var].create_dataset(new_key_str, data=src[var][old_key_str][:])
- if 'Time' in src[var][old_key_str].attrs:
- t_val = src[var][old_key_str].attrs['Time']
- ds.attrs['Time'] = t_val
- if h5_native_time is None:
- h5_native_time = float(t_val)
-
- steps.append((run_idx, old_key_str, new_key, h5_native_time))
- new_key += 1
-
- return steps
-
-
-def _write_step_xmf(in_xmf_path, out_xmf_path, old_key_str, new_key_int,
- h5_name, h5_native_time, selected_vars):
- """Write a per-step XMF with updated key references.
-
- Rewrites DataItem text from ``h5name:VARNAME/{old_key}`` to
- ``h5name:VARNAME/{new_key}``. Removes Attribute elements for variables
- not in selected_vars. Updates the ```` element.
-
- Parameters
- ----------
- in_xmf_path : str
- Source per-step XMF from the input run directory.
- out_xmf_path : str
- Destination path in the output directory.
- old_key_str : str
- Original HDF5 dataset key (cycle number string, e.g. '20').
- new_key_int : int
- New sequential key.
- h5_name : str
- Data filename, e.g. 'ats_vis_data.h5'.
- h5_native_time : float or None
- Time value (in H5 native units) to write into ````.
- selected_vars : list of str or None
- If None, keep all Attribute elements.
- """
- if not os.path.isfile(in_xmf_path):
- warnings.warn(
- f"Per-step XMF not found: {in_xmf_path!r}. "
- "Writing a minimal placeholder XMF.")
- _write_minimal_step_xmf(out_xmf_path, h5_name, new_key_int,
- h5_native_time, selected_vars)
- return
-
- # Register the XInclude namespace so it round-trips properly
- ET.register_namespace('xi', 'http://www.w3.org/2001/XInclude')
-
- tree = ET.parse(in_xmf_path)
- root = tree.getroot()
-
- old_suffix = '/' + old_key_str
- new_suffix = '/' + str(new_key_int)
-
- for grid in root.iter('Grid'):
- # Update
- time_elem = grid.find('Time')
- if time_elem is not None and h5_native_time is not None:
- time_elem.set('Value', f'{h5_native_time:.17e}')
-
- # Update Attribute elements
- to_remove = []
- for attr in list(grid):
- if attr.tag != 'Attribute':
- continue
- var_name = attr.get('Name')
- if selected_vars is not None and var_name not in selected_vars:
- to_remove.append(attr)
- continue
- di = attr.find('DataItem')
- if di is not None and di.text:
- text = di.text.strip()
- # text is like: ats_vis_data.h5:VARNAME/OLD_KEY
- # Update h5 filename to the combined file (same base name, same domain)
- parts = text.split(':', 1)
- if len(parts) == 2:
- new_text = h5_name + ':' + parts[1]
- else:
- new_text = text
- # Replace the trailing /OLD_KEY with /NEW_KEY
- if new_text.endswith(old_suffix):
- new_text = new_text[:-len(old_suffix)] + new_suffix
- # Preserve leading/trailing whitespace pattern
- leading = len(di.text) - len(di.text.lstrip())
- trailing = len(di.text) - len(di.text.rstrip())
- di.text = di.text[:leading] + new_text + di.text[len(di.text) - trailing:]
- for elem in to_remove:
- grid.remove(elem)
-
- tree.write(out_xmf_path, xml_declaration=True, encoding='ASCII')
-
-
-def _write_minimal_step_xmf(out_xmf_path, h5_name, new_key_int,
- h5_native_time, selected_vars):
- """Write a skeleton per-step XMF when the source XMF is missing."""
- xdmf = ET.Element('Xdmf')
- xdmf.set('Version', '2.0')
- domain_elem = ET.SubElement(xdmf, 'Domain')
- grid = ET.SubElement(domain_elem, 'Grid')
- grid.set('GridType', 'Uniform')
- if h5_native_time is not None:
- time_elem = ET.SubElement(grid, 'Time')
- time_elem.set('Value', f'{h5_native_time:.17e}')
- tree = ET.ElementTree(xdmf)
- tree.write(out_xmf_path, xml_declaration=True, encoding='ASCII')
-
-
-def _write_visit_xmf(out_directory, h5_name, n_total_steps):
- """Write the master VisIt.xmf with xi:include for each step XMF.
-
- Parameters
- ----------
- out_directory : str
- h5_name : str
- Data filename, e.g. 'ats_vis_data.h5'.
- n_total_steps : int
- Number of steps (0 … n_total_steps-1).
- """
- stem = h5_name[:-3] # strip '.h5'
- out_path = os.path.join(out_directory, f'{stem}.VisIt.xmf')
-
- xi_ns = 'http://www.w3.org/2001/XInclude'
- ET.register_namespace('xi', xi_ns)
-
- xdmf = ET.Element('Xdmf')
- xdmf.set('Version', '2.0')
- xdmf.set('xmlns:xi', xi_ns)
- domain_elem = ET.SubElement(xdmf, 'Domain')
- coll = ET.SubElement(domain_elem, 'Grid')
- coll.set('GridType', 'Collection')
- coll.set('CollectionType', 'Temporal')
-
- for new_key in range(n_total_steps):
- href = f'{h5_name}.{new_key}.xmf'
- xi_elem = ET.SubElement(coll, f'{{{xi_ns}}}include')
- xi_elem.set('href', href)
- xi_elem.set('xpointer', 'xpointer(//Xdmf/Domain/Grid)')
-
- tree = ET.ElementTree(xdmf)
- tree.write(out_path, xml_declaration=True, encoding='ASCII')
- print(f"Wrote VisIt XMF: {out_path}")
-
-
-def _copy_mesh(directories, out_directory, domain):
- """Copy mesh H5 and XMF files from the first run that has them.
-
- Parameters
- ----------
- directories : list of str
- out_directory : str
- domain : str
- """
- mesh_name = ats_xdmf.valid_mesh_filename(domain)
- stem = mesh_name[:-3] # strip '.h5'
-
- # Find the first directory that has the mesh H5
- mesh_src_dir = None
- for d in directories:
- if os.path.isfile(os.path.join(d, mesh_name)):
- mesh_src_dir = d
- break
-
- if mesh_src_dir is None:
- warnings.warn(
- f"Mesh HDF5 '{mesh_name}' not found in any input directory. "
- "Output will not render correctly in VisIt.")
- return
-
- for fname in [mesh_name,
- f'{mesh_name}.0.xmf',
- f'{stem}.VisIt.xmf']:
- src_path = os.path.join(mesh_src_dir, fname)
- dst_path = os.path.join(out_directory, fname)
- if os.path.isfile(src_path):
- shutil.copyfile(src_path, dst_path)
- print(f"Copied mesh file: {dst_path}")
- else:
- warnings.warn(f"Mesh file not found, skipping: {src_path}")
-
-
-# ---------------------------------------------------------------------------
-# Public API
-# ---------------------------------------------------------------------------
-
-def combineVisFiles(directories, domain, out_directory,
- eps=1.0, time_unit='d',
- include_vars=None, exclude_vars=None,
- dry_run=False):
- """Combine ATS visualization output from multiple restarted runs.
-
- Parameters
- ----------
- directories : list of str
- Input run directories in chronological order.
- domain : str
- ATS domain name (e.g. 'domain', 'surface').
- out_directory : str
- Output directory (created if needed).
- eps : float
- Overlap tolerance in *time_unit*. Cycles from run i with time >=
- (first_time_of_run_i+1 - eps) are dropped. Default 1.0.
- time_unit : str
- Time unit for *eps* and for printing. One of 's', 'hr', 'd', 'yr',
- 'noleap'. Default 'd'.
- include_vars : list of str or None
- If given, only these variables are written.
- exclude_vars : list of str or None
- If given, these variables are excluded.
- dry_run : bool
- If True, print what would be written and return without writing files.
- """
- if not directories:
- raise ValueError("Must provide at least one directory.")
-
- # Validate inputs
- for d in directories:
- if not os.path.isdir(d):
- raise RuntimeError(f"Input directory not found: {d!r}")
-
- if os.path.abspath(out_directory) in [os.path.abspath(d) for d in directories]:
- raise RuntimeError(
- f"Output directory '{out_directory}' is the same as one of the input "
- "directories. Choose a different output path.")
-
- h5_name = ats_xdmf.valid_data_filename(domain)
-
- # Open VisFiles
- vis_files = []
- try:
- for d in directories:
- vis_files.append(
- ats_xdmf.VisFile(d, domain=domain, output_time_unit=time_unit))
-
- # Validate time ordering
- _validate_time_ordering(vis_files, directories, time_unit)
-
- # Compute selected variables (intersection + user filter)
- selected_vars = _compute_selected_vars(vis_files, include_vars, exclude_vars)
-
- # Deduplicate cycles across runs
- run_data = _select_cycles(vis_files, directories, eps)
- if not run_data:
- raise RuntimeError("No cycles selected across all runs.")
-
- total_cycles = sum(len(sel) for _, _, sel, _ in run_data)
-
- if dry_run:
- print(f"Combined output would contain {total_cycles} cycles "
- f"from {len(run_data)} run(s):")
- new_key = 0
- for (vf, directory, sel_cycles, sel_times) in run_data:
- print(f"\n Run: {directory}")
- for old_key, t in zip(sel_cycles, sel_times):
- print(f" new key {new_key:6d} (old {old_key:>8}) "
- f"t = {t:.6g} {time_unit}")
- new_key += 1
- print(f"\n{len(selected_vars)} variables selected:")
- for v in selected_vars:
- print(f" {v}")
- return
-
- # Write output
- os.makedirs(out_directory, exist_ok=True)
-
- # Write combined HDF5
- out_h5 = os.path.join(out_directory, h5_name)
- print(f"Writing combined HDF5: {out_h5}")
- steps = _write_combined_h5(run_data, out_h5, selected_vars)
-
- # Write per-step XMFs
- for (run_idx, old_key_str, new_key_int, h5_native_time) in steps:
- directory = run_data[run_idx][1]
- in_xmf = os.path.join(directory, f'{h5_name}.{old_key_str}.xmf')
- out_xmf = os.path.join(out_directory, f'{h5_name}.{new_key_int}.xmf')
- _write_step_xmf(in_xmf, out_xmf, old_key_str, new_key_int,
- h5_name, h5_native_time, selected_vars)
-
- # Write master VisIt.xmf
- _write_visit_xmf(out_directory, h5_name, len(steps))
-
- # Copy mesh files
- _copy_mesh(directories, out_directory, domain)
-
- print(f"\nDone. Combined {len(steps)} cycles from {len(run_data)} run(s).")
- print(f"Output directory: {out_directory}")
- stem = h5_name[:-3]
- print(f"Open in VisIt: {os.path.join(out_directory, stem + '.VisIt.xmf')}")
-
- finally:
- for vf in vis_files:
- try:
- vf.close()
- except Exception:
- pass
-
-
-# ---------------------------------------------------------------------------
-# CLI
-# ---------------------------------------------------------------------------
-
-def main():
- parser = argparse.ArgumentParser(
- description=__doc__,
- formatter_class=argparse.RawDescriptionHelpFormatter)
-
- parser.add_argument('domain', metavar='DOMAIN',
- help='ATS domain name (e.g. "domain", "surface"), or '
- '"*" to process all domains found in the first '
- 'input directory.')
- parser.add_argument('directories', metavar='DIRECTORY', nargs='+',
- help='Input run directories in chronological order. '
- 'Glob patterns (e.g. \'run*\') are accepted if quoted '
- 'and are automatically natural-sorted. Explicit lists '
- 'are used as-is unless --sort is given.')
- parser.add_argument('--output', '-o', dest='output', required=True,
- help='Output directory (created if needed).')
- parser.add_argument('--time-unit', dest='time_unit', default='d',
- choices=['s', 'hr', 'd', 'yr', 'noleap'],
- help='Time unit for --eps and printed times (default: d).')
- parser.add_argument('--eps', dest='eps', type=float, default=1.0,
- help='Overlap tolerance in --time-unit units. Cycles '
- 'from run i with time >= (start_of_run_i+1 - eps) '
- 'are dropped. Default: 1.0.')
-
- var_group = parser.add_mutually_exclusive_group()
- var_group.add_argument('--include', dest='include_vars',
- action='append', metavar='VAR',
- help='Include this variable in output (repeat for multiple).')
- var_group.add_argument('--exclude', dest='exclude_vars',
- action='append', metavar='VAR',
- help='Exclude this variable from output (repeat for multiple).')
-
- parser.add_argument('--sort', dest='sort', action='store_true', default=False,
- help='Natural-sort the directory list before combining. '
- 'Applied automatically when a glob pattern is given.')
- parser.add_argument('--dry-run', dest='dry_run', action='store_true',
- default=False,
- help='Print what would be written; do not create files.')
-
- args = parser.parse_args()
-
- # Expand any glob patterns and natural-sort; trust explicit lists unless --sort
- import glob as _glob
- directories = []
- is_glob = False
- for d in args.directories:
- if any(c in d for c in ('*', '?', '[', ']')):
- expanded = sorted(_glob.glob(d), key=ats_xdmf.natural_sort_key)
- if not expanded:
- raise RuntimeError(f"Glob pattern '{d}' matched no directories.")
- directories.extend(expanded)
- is_glob = True
- else:
- directories.append(d)
-
- if is_glob or args.sort:
- directories = sorted(directories, key=ats_xdmf.natural_sort_key)
- print(f"Directories (natural-sorted): {directories}")
-
- if args.domain == '*':
- domains = ats_xdmf.find_domains(directories[0])
- print(f"Found domains: {domains}")
- else:
- domains = [args.domain]
-
- for domain in domains:
- combineVisFiles(
- directories=directories,
- domain=domain,
- out_directory=args.output,
- eps=args.eps,
- time_unit=args.time_unit,
- include_vars=args.include_vars,
- exclude_vars=args.exclude_vars,
- dry_run=args.dry_run,
- )
-
-
-if __name__ == '__main__':
- main()
diff --git a/tools/utils/query_vis.py b/tools/utils/query_vis.py
index 1bc5a8bc8..689e2c70c 100644
--- a/tools/utils/query_vis.py
+++ b/tools/utils/query_vis.py
@@ -24,11 +24,21 @@ def _display_name(varname, domain):
return name
-def queryVisFiles(directory, domain, time_unit=None):
- """Print a summary of ATS visualization output for one domain."""
- h5_name = ats_xdmf.valid_data_filename(domain)
-
- vf = ats_xdmf.VisFile(directory, domain=domain, output_time_unit=time_unit)
+def queryVisFiles(directory, domain, filename, time_unit=None):
+ """Print a summary of ATS visualization output for one file.
+
+ Parameters
+ ----------
+ directory : str
+ domain : str or None
+ ATS domain name, or None when opening by raw filename (filepath input).
+ filename : str
+ Basename of the data HDF5 file.
+ time_unit : str or None
+ Display unit. None uses the native unit from the file.
+ """
+ vf = ats_xdmf.VisFile(directory, domain=domain, filename=filename,
+ output_time_unit=time_unit)
n_cycles = len(vf.cycles)
first_cycle = int(vf.cycles[0]) if n_cycles > 0 else None
@@ -39,7 +49,6 @@ def queryVisFiles(directory, domain, time_unit=None):
diffs = np.diff(vf.times) if n_cycles > 1 else np.array([])
- # Cell count from first variable at first cycle
n_cells = None
all_vars = vf.variables()
if all_vars and n_cycles > 0:
@@ -48,7 +57,11 @@ def queryVisFiles(directory, domain, time_unit=None):
except (KeyError, IndexError):
pass
- print(f"Domain: {domain} ({h5_name})")
+ label = domain if domain is not None else filename
+ print(f"Domain: {label} ({filename})")
+ for key, val in vf.d.attrs.items():
+ val_str = val.decode() if isinstance(val, bytes) else str(val)
+ print(f" {key}: {val_str}")
if n_cycles > 0:
print(f" Cycles: {n_cycles} ({first_cycle} \u2026 {last_cycle})")
print(f" Time: {t_first:.3f} \u2026 {t_last:.3f} {vf.output_time_unit}")
@@ -81,27 +94,40 @@ def main():
description=__doc__,
formatter_class=argparse.RawDescriptionHelpFormatter)
- parser.add_argument('domain', metavar='DOMAIN', nargs='?', default='*',
- help='ATS domain name (e.g. "surface", "domain"), or '
- '"*" to summarize all domains found (default: *)')
+ parser.add_argument('domain', metavar='DOMAIN_OR_FILE', nargs='?', default='*',
+ help='ATS domain name, path to a vis file (.h5 or .xmf), '
+ 'or "*" to summarize all domains found (default: *). '
+ 'A filepath implies its own directory; -d and -p are ignored.')
parser.add_argument('-d', '--directory', dest='directory', default='.',
help='Directory containing visualization files '
- '(default: current directory)')
+ '(default: current directory). Ignored when a filepath is given.')
+ parser.add_argument('-p', '--prefix', dest='prefix', default='ats_vis',
+ help='Filename prefix (default: ats_vis). Ignored when a filepath is given.')
parser.add_argument('--time-unit', dest='time_unit', default=None,
choices=['s', 'hr', 'd', 'yr', 'noleap'],
help='Time unit for display (default: native unit from file)')
args = parser.parse_args()
- if args.domain == '*':
+ arg = args.domain
+ is_path = (os.sep in arg or '/' in arg or
+ arg.endswith('.h5') or arg.endswith('.xmf'))
+
+ if is_path:
+ directory, domain, filename = ats_xdmf.resolve_vis_input(arg)
+ queryVisFiles(directory, domain, filename, time_unit=args.time_unit)
+ elif arg == '*':
domains = ats_xdmf.find_domains(args.directory)
+ for i, domain in enumerate(domains):
+ if i > 0:
+ print()
+ _, _, filename = ats_xdmf.resolve_vis_input(
+ domain, directory=args.directory, prefix=args.prefix)
+ queryVisFiles(args.directory, domain, filename, time_unit=args.time_unit)
else:
- domains = [args.domain]
-
- for i, domain in enumerate(domains):
- if i > 0:
- print()
- queryVisFiles(args.directory, domain, time_unit=args.time_unit)
+ _, domain, filename = ats_xdmf.resolve_vis_input(
+ arg, directory=args.directory, prefix=args.prefix)
+ queryVisFiles(args.directory, domain, filename, time_unit=args.time_unit)
if __name__ == '__main__':
diff --git a/tools/utils/resample_vis.py b/tools/utils/resample_vis.py
new file mode 100644
index 000000000..faca8081d
--- /dev/null
+++ b/tools/utils/resample_vis.py
@@ -0,0 +1,695 @@
+#!/usr/bin/env python
+"""Subset and/or combine ATS visualization output by time/cycle/index and variable.
+
+With a single input directory, subsets that run's visualization output.
+With multiple input directories (restart/continuation runs), combines them into
+a single self-contained XDMF dataset after deduplicating restart seams, and then
+subsets the combined dataset.
+
+In all cases the time/cycle/index filter is applied to the merged dataset.
+
+Produces a self-contained output directory suitable for downloading from HPC
+and opening locally in VisIt.
+
+Examples
+--------
+Subset a single run by time range::
+
+ subset_vis.py surface run0/ -o out/ --times 0:2.5 --time-unit y
+
+Combine three directories into a single combined result::
+
+ subset_vis.py domain run0/ run1/ run2/ -o combined/
+
+Combine with glob expansion::
+
+ subset_vis.py surface 'run*/' -o combined/ --sort
+
+Variable subset with custom output prefix::
+
+ subset_vis.py surface run0/ -o out/ \\
+ --out-prefix ats_vis_ponded_depth --include ponded_depth
+
+Downsample to every 10th visualized step::
+
+ subset_vis.py surface run0/ run1/ -o out/ --indices 0:10 --dry-run
+"""
+
+import sys
+import os
+import shutil
+import argparse
+import warnings
+import xml.etree.ElementTree as ET
+
+import numpy as np
+import h5py
+
+try:
+ sys.path.insert(0, os.path.join(os.environ['ATS_SRC_DIR'], 'tools', 'utils'))
+except KeyError:
+ pass
+
+import ats_xdmf
+
+
+# ---------------------------------------------------------------------------
+# Argument parsing helpers
+# ---------------------------------------------------------------------------
+
+def _parse_slice_or_list(s, cast):
+ """Parse a slice or comma-separated list string, returning a slice or list.
+
+ Slice syntax: 'START:STOP' or 'START:STOP:STEP' (any field may be empty)
+ List syntax: 'v1,v2,v3'
+
+ cast is applied to each non-empty value (e.g. float, int).
+ """
+ if ':' in s:
+ parts = s.split(':')
+ if len(parts) not in (2, 3):
+ raise argparse.ArgumentTypeError(
+ f"Cannot parse slice argument: {s!r}")
+ vals = [cast(p) if p else None for p in parts]
+ return slice(*vals)
+ else:
+ try:
+ return [cast(v) for v in s.split(',')]
+ except ValueError:
+ raise argparse.ArgumentTypeError(
+ f"Cannot parse list argument: {s!r}")
+
+
+# ---------------------------------------------------------------------------
+# Internal helpers
+# ---------------------------------------------------------------------------
+
+_TIME_IN_SECONDS = {
+ 'y': 365.25 * 86400,
+ 'noleap': 365 * 86400,
+ 'd': 86400,
+ 'h': 3600,
+ 'min': 60,
+ 's': 1,
+}
+
+
+def _default_time_tolerance(effective_unit):
+ """Return 0.1 s expressed in effective_unit."""
+ return 0.1 / _TIME_IN_SECONDS[effective_unit]
+
+
+def _validate_time_ordering(vis_files, directories, time_unit):
+ """Warn loudly if runs are not in chronological order."""
+ for i in range(len(vis_files) - 1):
+ t0 = vis_files[i].times[0]
+ t1 = vis_files[i + 1].times[0]
+ if t1 <= t0:
+ warnings.warn(
+ f"\n *** TIME ORDERING WARNING ***\n"
+ f" Run '{directories[i+1]}' starts at t={t1:.6g} {time_unit}\n"
+ f" but run '{directories[i]}' starts at t={t0:.6g} {time_unit}.\n"
+ f" Directories may be out of order — combined output may be wrong.",
+ stacklevel=3)
+
+
+def _select_cycles(vis_files, directories, eps):
+ """For each run, select cycles to include after deduplication.
+
+ Parameters
+ ----------
+ vis_files : list of ats_xdmf.VisFile
+ directories : list of str
+ eps : float
+ Overlap tolerance in the VisFile's output_time_unit. Cycles from run i
+ with time >= (start_time_of_run_i+1 - eps) are dropped.
+
+ Returns
+ -------
+ list of (vf, directory, selected_cycle_strs, selected_times_array)
+ """
+ run_data = []
+ for i, (vf, directory) in enumerate(zip(vis_files, directories)):
+ cycles = list(vf.cycles)
+ times = np.array(vf.times)
+
+ if i < len(vis_files) - 1:
+ next_start = vis_files[i + 1].times[0]
+ cutoff = next_start - eps
+ mask = times < cutoff
+ sel_cycles = [c for c, m in zip(cycles, mask) if m]
+ sel_times = times[mask]
+ else:
+ sel_cycles = cycles
+ sel_times = times
+
+ if len(sel_cycles) == 0:
+ warnings.warn(
+ f"Run '{directory}': all cycles filtered out after deduplication. "
+ "Check that directories are in chronological order and that "
+ "--time-tolerance is appropriate.")
+ else:
+ run_data.append((vf, directory, sel_cycles, sel_times))
+
+ return run_data
+
+
+def _compute_selected_vars(vis_files, include_vars, exclude_vars):
+ """Return the intersection of variable sets across runs, then apply filters."""
+ var_sets = [set(vf.variables()) for vf in vis_files]
+ common_vars = set.intersection(*var_sets)
+ all_vars = set.union(*var_sets)
+
+ if common_vars != all_vars:
+ for vf, vs in zip(vis_files, var_sets):
+ missing = common_vars - vs
+ if missing:
+ warnings.warn(
+ f"Directory '{vf.directory}': missing variables present in "
+ f"other runs: {sorted(missing)}. These will be excluded.")
+
+ filtered = vis_files[0].variables(names=include_vars, exclude=exclude_vars)
+ selected = [v for v in filtered if v in common_vars]
+
+ if not selected:
+ raise RuntimeError(
+ "No variables selected after filtering. Check --include / --exclude.")
+
+ return selected
+
+
+def _apply_filter(combined, times_spec, cycles_spec, indices_spec,
+ time_tolerance, effective_unit):
+ """Apply --times / --cycles / --indices filter to the flat combined list.
+
+ Parameters
+ ----------
+ combined : list of (vf, directory, cycle_str, time_float)
+ times_spec, cycles_spec, indices_spec : slice/list/None
+ time_tolerance : float
+ effective_unit : str
+
+ Returns
+ -------
+ filtered : list of (vf, directory, cycle_str, time_float)
+ """
+ if indices_spec is not None:
+ if isinstance(indices_spec, slice):
+ return combined[indices_spec]
+ else:
+ n = len(combined)
+ result = []
+ for idx in indices_spec:
+ if -n <= idx < n:
+ result.append(combined[idx])
+ else:
+ warnings.warn(f"Index {idx} out of range (0..{n-1}); skipping.")
+ return result
+
+ if cycles_spec is not None:
+ if isinstance(cycles_spec, slice):
+ start = cycles_spec.start
+ stop = cycles_spec.stop
+ step = cycles_spec.step
+ result = []
+ for entry in combined:
+ c = int(entry[2])
+ lo = (start is None) or (c >= start)
+ hi = (stop is None) or (c <= stop)
+ result.append(entry) if (lo and hi) else None
+ if step is not None and step != 1:
+ # Resample by step within the already-filtered list
+ result = result[::step]
+ return result
+ else:
+ cycle_set = set(cycles_spec)
+ return [e for e in combined if int(e[2]) in cycle_set]
+
+ if times_spec is not None:
+ all_times = np.array([e[3] for e in combined])
+ if isinstance(times_spec, slice):
+ lo = times_spec.start
+ hi = times_spec.stop
+ interval = times_spec.stop - times_spec.start if (
+ times_spec.start is not None and times_spec.stop is not None
+ and times_spec.step is not None) else None
+
+ if times_spec.step is None:
+ # Simple range
+ mask = np.ones(len(combined), dtype=bool)
+ if lo is not None:
+ mask &= (all_times >= lo - time_tolerance)
+ if hi is not None:
+ mask &= (all_times <= hi + time_tolerance)
+ return [e for e, m in zip(combined, mask) if m]
+ else:
+ # Evenly-spaced targets: START:STOP:INTERVAL
+ targets = np.arange(lo, hi + times_spec.step * 0.5,
+ times_spec.step)
+ selected_indices = set()
+ for target in targets:
+ diffs = np.abs(all_times - target)
+ best = int(np.argmin(diffs))
+ if diffs[best] <= time_tolerance:
+ selected_indices.add(best)
+ return [e for i, e in enumerate(combined) if i in selected_indices]
+ else:
+ # Explicit list of times
+ selected_indices = set()
+ for target in times_spec:
+ diffs = np.abs(all_times - target)
+ best = int(np.argmin(diffs))
+ if diffs[best] <= time_tolerance:
+ selected_indices.add(best)
+ else:
+ warnings.warn(
+ f"Requested time {target} {effective_unit} not found within "
+ f"tolerance {time_tolerance} {effective_unit}; skipping.")
+ return [e for i, e in enumerate(combined) if i in selected_indices]
+
+ # No filter — return all
+ return combined
+
+
+# ---------------------------------------------------------------------------
+# Output writers
+# ---------------------------------------------------------------------------
+
+def _write_output_h5(combined, out_h5_path, selected_vars):
+ """Write HDF5 with original cycle keys (no renumbering)."""
+ with h5py.File(out_h5_path, 'w') as dst:
+ # Merge file-level attributes from all source VisFiles (superset)
+ seen_vfs = {}
+ for vf, _, _, _ in combined:
+ if id(vf) not in seen_vfs:
+ seen_vfs[id(vf)] = vf
+ for key, val in vf.d.attrs.items():
+ dst.attrs[key] = val
+
+ for var in selected_vars:
+ dst.create_group(var)
+
+ for vf, directory, cycle_str, _ in combined:
+ src = vf.d
+ for var in selected_vars:
+ if var not in src:
+ warnings.warn(f"Variable '{var}' missing in '{directory}'; skipping.")
+ continue
+ if cycle_str not in src[var]:
+ warnings.warn(
+ f"Cycle key '{cycle_str}' missing for variable '{var}' "
+ f"in '{directory}'; skipping.")
+ continue
+ ds = dst[var].create_dataset(cycle_str, data=src[var][cycle_str][:])
+ if 'Time' in src[var][cycle_str].attrs:
+ ds.attrs['Time'] = src[var][cycle_str].attrs['Time']
+
+ print(f"Wrote data HDF5: {out_h5_path}")
+
+
+def _write_step_xmf(in_xmf_path, out_xmf_path, cycle_str,
+ in_h5_name, out_h5_name, selected_vars):
+ """Write a per-step XMF, updating h5 filename and removing unselected vars."""
+ if not os.path.isfile(in_xmf_path):
+ warnings.warn(f"Per-step XMF not found: {in_xmf_path!r}; skipping.")
+ return
+
+ ET.register_namespace('xi', 'http://www.w3.org/2001/XInclude')
+ tree = ET.parse(in_xmf_path)
+ root = tree.getroot()
+
+ for grid in root.iter('Grid'):
+ to_remove = [
+ e for e in list(grid)
+ if e.tag == 'Attribute' and e.get('Name') not in selected_vars
+ ]
+ for e in to_remove:
+ grid.remove(e)
+
+ for item in root.iter('DataItem'):
+ text = item.text or ''
+ if in_h5_name in text:
+ item.text = text.replace(in_h5_name, out_h5_name)
+
+ tree.write(out_xmf_path, xml_declaration=True, encoding='ASCII')
+
+
+def _write_visit_xmf(out_directory, out_h5_name, cycle_strs):
+ """Write the master VisIt.xmf with xi:include for each step XMF."""
+ stem = out_h5_name[:-3] # strip '.h5'
+ out_path = os.path.join(out_directory, f'{stem}.VisIt.xmf')
+
+ xi_ns = 'http://www.w3.org/2001/XInclude'
+ ET.register_namespace('xi', xi_ns)
+
+ xdmf = ET.Element('Xdmf')
+ xdmf.set('Version', '2.0')
+ xdmf.set('xmlns:xi', xi_ns)
+ domain_elem = ET.SubElement(xdmf, 'Domain')
+ coll = ET.SubElement(domain_elem, 'Grid')
+ coll.set('GridType', 'Collection')
+ coll.set('CollectionType', 'Temporal')
+ coll.set('Name', 'Mesh')
+
+ for cycle_str in cycle_strs:
+ href = f'{out_h5_name}.{int(cycle_str)}.xmf'
+ xi_elem = ET.SubElement(coll, f'{{{xi_ns}}}include')
+ xi_elem.set('href', href)
+ xi_elem.set('xpointer', 'xpointer(//Xdmf/Domain/Grid)')
+
+ tree = ET.ElementTree(xdmf)
+ tree.write(out_path, xml_declaration=True, encoding='ASCII')
+ print(f"Wrote VisIt XMF: {out_path}")
+
+
+def _copy_mesh(directories, out_directory, domain):
+ """Copy mesh H5 and XMF files from the first directory that has them."""
+ mesh_name = ats_xdmf.valid_mesh_filename(domain)
+ stem = mesh_name[:-3] # strip '.h5'
+
+ mesh_src_dir = None
+ for d in directories:
+ if os.path.isfile(os.path.join(d, mesh_name)):
+ mesh_src_dir = d
+ break
+
+ if mesh_src_dir is None:
+ warnings.warn(
+ f"Mesh HDF5 '{mesh_name}' not found in any input directory. "
+ "Output will not render correctly in VisIt.")
+ return
+
+ for fname in [mesh_name, f'{mesh_name}.0.xmf', f'{stem}.VisIt.xmf']:
+ src_path = os.path.join(mesh_src_dir, fname)
+ dst_path = os.path.join(out_directory, fname)
+ if os.path.isfile(src_path):
+ shutil.copyfile(src_path, dst_path)
+ print(f"Copied mesh file: {dst_path}")
+ else:
+ warnings.warn(f"Mesh file not found, skipping: {src_path}")
+
+
+# ---------------------------------------------------------------------------
+# Core public function
+# ---------------------------------------------------------------------------
+
+def subsetVisFiles(directories, domain, out_directory, out_stem,
+ in_prefix='ats_vis',
+ times=None, time_unit=None, time_tolerance=None,
+ cycles=None, indices=None,
+ include_vars=None, exclude_vars=None,
+ dry_run=False):
+ """Subset and/or combine ATS visualization output files.
+
+ Parameters
+ ----------
+ directories : list of str
+ Input directory or directories in chronological order.
+ A single-element list is the subset case.
+ domain : str
+ ATS domain name (e.g. 'surface', 'domain').
+ out_directory : str
+ Directory for output files (created if needed).
+ out_stem : str
+ Output filename stem (everything before ``_data.h5``).
+ in_prefix : str
+ Input filename prefix. Default ``'ats_vis'``.
+ times : slice or list or None
+ Time filter (result of _parse_slice_or_list with float cast), or None.
+ time_unit : str or None
+ Time unit for display and filtering. None uses the native unit from file.
+ time_tolerance : float or None
+ Tolerance for time matching and restart deduplication, in time_unit.
+ Default: 0.1 s expressed in the effective unit.
+ cycles : slice or list or None
+ Cycle filter, or None.
+ indices : slice or list or None
+ Index filter (0-based into the combined timeline), or None.
+ include_vars : list of str or None
+ If given, only these variables are copied.
+ exclude_vars : list of str or None
+ If given, these variables are excluded.
+ dry_run : bool
+ If True, print what would be done and return without writing files.
+ """
+ if not directories:
+ raise ValueError("Must provide at least one directory.")
+
+ for d in directories:
+ if not os.path.isdir(d):
+ raise RuntimeError(f"Input directory not found: {d!r}")
+
+ if not dry_run:
+ if os.path.abspath(out_directory) in [os.path.abspath(d) for d in directories]:
+ raise RuntimeError(
+ f"Output directory '{out_directory}' is the same as one of the "
+ "input directories. Choose a different output path.")
+
+ _, _, in_filename = ats_xdmf.resolve_vis_input(
+ domain, directory=directories[0], prefix=in_prefix)
+ out_h5_name = f'{out_stem}_data.h5'
+
+ vis_files = []
+ try:
+ # Open first file to resolve effective time unit
+ vf0 = ats_xdmf.VisFile(directories[0], domain=domain,
+ filename=in_filename,
+ output_time_unit=time_unit)
+ effective_unit = vf0.output_time_unit
+ vis_files = [vf0] + [
+ ats_xdmf.VisFile(d, domain=domain, filename=in_filename,
+ output_time_unit=effective_unit)
+ for d in directories[1:]
+ ]
+
+ if time_tolerance is None:
+ time_tolerance = _default_time_tolerance(effective_unit)
+
+ if len(vis_files) > 1:
+ _validate_time_ordering(vis_files, directories, effective_unit)
+
+ selected_vars = _compute_selected_vars(vis_files, include_vars, exclude_vars)
+
+ # Deduplicate restart seams
+ run_data = _select_cycles(vis_files, directories, eps=time_tolerance)
+ if not run_data:
+ raise RuntimeError("No cycles selected after deduplication.")
+
+ # Build flat combined timeline
+ combined = []
+ for vf, directory, sel_cycles, sel_times in run_data:
+ for c, t in zip(sel_cycles, sel_times):
+ combined.append((vf, directory, c, t))
+
+ # Apply user filter
+ combined = _apply_filter(combined, times, cycles, indices,
+ time_tolerance, effective_unit)
+
+ if not combined:
+ raise RuntimeError("No cycles selected — check your filter arguments.")
+
+ # Dry-run output
+ if dry_run:
+ n_dirs = len({e[1] for e in combined})
+ by_dir = {}
+ for vf, directory, cycle_str, t in combined:
+ by_dir.setdefault(directory, []).append((cycle_str, t))
+
+ print(f"{n_dirs} director{'y' if n_dirs == 1 else 'ies'}, "
+ f"{len(combined)} cycles selected "
+ f"(after deduplication and filtering):")
+ for directory, entries in by_dir.items():
+ cycles_in_dir = [e[0] for e in entries]
+ times_in_dir = [e[1] for e in entries]
+ print(f" {directory} {len(entries)} cycles "
+ f"cycle {int(cycles_in_dir[0])} … {int(cycles_in_dir[-1])} "
+ f"t = {times_in_dir[0]:.3f} … {times_in_dir[-1]:.3f} "
+ f"{effective_unit}")
+ print(f"\n{len(selected_vars)} variables selected:")
+ for v in selected_vars:
+ print(f" {v}")
+ return
+
+ os.makedirs(out_directory, exist_ok=True)
+
+ # Write per-step XMFs
+ for vf, directory, cycle_str, _ in combined:
+ N = int(cycle_str)
+ in_xmf_path = os.path.join(directory, f'{in_filename}.{N}.xmf')
+ out_xmf_path = os.path.join(out_directory, f'{out_h5_name}.{N}.xmf')
+ _write_step_xmf(in_xmf_path, out_xmf_path, cycle_str,
+ in_filename, out_h5_name, selected_vars)
+
+ # Write master VisIt.xmf
+ cycle_strs = [e[2] for e in combined]
+ _write_visit_xmf(out_directory, out_h5_name, cycle_strs)
+
+ # Write HDF5 data file
+ out_h5 = os.path.join(out_directory, out_h5_name)
+ _write_output_h5(combined, out_h5, selected_vars)
+
+ # Copy mesh files
+ _copy_mesh(directories, out_directory, domain)
+
+ n_dirs_used = len({e[1] for e in combined})
+ print(f"\nDone. {len(combined)} cycles from "
+ f"{n_dirs_used} director{'y' if n_dirs_used == 1 else 'ies'}.")
+ print(f"Output directory: {out_directory}")
+ print(f"Open in VisIt: "
+ f"{os.path.join(out_directory, out_stem + '.VisIt.xmf')}")
+
+ finally:
+ for vf in vis_files:
+ try:
+ vf.close()
+ except Exception:
+ pass
+
+
+# ---------------------------------------------------------------------------
+# CLI
+# ---------------------------------------------------------------------------
+
+def main():
+ parser = argparse.ArgumentParser(
+ description=__doc__,
+ formatter_class=argparse.RawDescriptionHelpFormatter)
+
+ parser.add_argument('domain', metavar='DOMAIN',
+ help='ATS domain name (e.g. "surface", "domain"), or '
+ '"*" to process all domains found in the first '
+ 'input directory.')
+ parser.add_argument('directories', metavar='DIRECTORY', nargs='+',
+ help='One or more input directories in chronological order. '
+ 'A single directory subsets that run; multiple directories '
+ 'are combined after restart deduplication. '
+ 'Glob patterns (e.g. \'run*/\') are accepted if quoted '
+ 'and are automatically natural-sorted.')
+
+ parser.add_argument('-p', '--prefix', dest='prefix', default='ats_vis',
+ help='Input filename prefix (default: ats_vis)')
+ parser.add_argument('-o', '--output-directory', dest='output', default=None,
+ required=False,
+ help='Output directory (default: DIRECTORY/subset for a '
+ 'single input, required for multiple inputs).')
+ parser.add_argument('--out-prefix', dest='out_prefix', default=None,
+ help='Output filename prefix (default: same as -p). '
+ 'Output stem is {out_prefix}_{domain}.')
+
+ # Time/cycle/index selection (mutually exclusive)
+ filter_group = parser.add_mutually_exclusive_group()
+ filter_group.add_argument(
+ '--times', dest='times', default=None,
+ metavar='SLICE_OR_LIST',
+ help='Time filter. Formats: '
+ '"START:STOP" (range), '
+ '"START:STOP:INTERVAL" (evenly-spaced targets), '
+ '"t1,t2,t3" (specific times). '
+ 'Units set by --time-unit.')
+ filter_group.add_argument(
+ '--cycles', dest='cycles', default=None,
+ metavar='SLICE_OR_LIST',
+ help='Cycle filter. Formats: '
+ '"START:STOP" or "START:STOP:STEP" (inclusive both ends), '
+ '"c1,c2,c3" (specific cycles).')
+ filter_group.add_argument(
+ '--indices', dest='indices', default=None,
+ metavar='SLICE_OR_LIST',
+ help='Index filter (0-based into the combined timeline). '
+ 'Formats: "START:STOP:STEP" (numpy-style, exclusive end), '
+ '"i1,i2,i3". Supports negative indices.')
+
+ parser.add_argument('--time-unit', dest='time_unit', default=None,
+ choices=['s', 'min', 'h', 'd', 'y', 'noleap'],
+ help='Time unit for display and --times values '
+ '(default: native unit from file).')
+ parser.add_argument('--time-tolerance', dest='time_tolerance',
+ type=float, default=None,
+ help='Tolerance for time matching and restart boundary '
+ 'deduplication, in --time-unit units. '
+ 'Default: 0.1 s expressed in the effective unit '
+ '(~1e-9 y, ~3e-8 d, 2.8e-5 h, 0.1 s).')
+
+ # Variable selection (mutually exclusive)
+ var_group = parser.add_mutually_exclusive_group()
+ var_group.add_argument('--include', dest='include_vars',
+ action='append', metavar='VAR',
+ help='Include this variable (repeat for multiple).')
+ var_group.add_argument('--exclude', dest='exclude_vars',
+ action='append', metavar='VAR',
+ help='Exclude this variable (repeat for multiple).')
+
+ parser.add_argument('--sort', dest='sort', action='store_true', default=False,
+ help='Natural-sort the directory list. '
+ 'Applied automatically when glob patterns are given.')
+ parser.add_argument('--dry-run', dest='dry_run', action='store_true',
+ default=False,
+ help='Print selected cycles and variables; do not write files.')
+
+ args = parser.parse_args()
+
+ import glob as _glob
+ directories = []
+ is_glob = False
+ for d in args.directories:
+ if any(c in d for c in ('*', '?', '[', ']')):
+ expanded = sorted(_glob.glob(d), key=ats_xdmf.natural_sort_key)
+ if not expanded:
+ raise RuntimeError(f"Glob pattern '{d}' matched no directories.")
+ directories.extend(expanded)
+ is_glob = True
+ else:
+ directories.append(d)
+
+ if is_glob or args.sort:
+ directories = sorted(directories, key=ats_xdmf.natural_sort_key)
+ print(f"Directories (natural-sorted): {directories}")
+
+ # Default output directory: DIRECTORY/subset for single input
+ if args.output is None:
+ if len(directories) > 1:
+ parser.error(
+ "argument -o/--output-directory is required when multiple "
+ "input directories are given.")
+ out_directory = os.path.join(directories[0], 'subset')
+ else:
+ out_directory = args.output
+
+ out_prefix = args.out_prefix or args.prefix
+
+ # Parse filter arguments
+ times_spec = cycles_spec = indices_spec = None
+ if args.times is not None:
+ times_spec = _parse_slice_or_list(args.times, float)
+ elif args.cycles is not None:
+ cycles_spec = _parse_slice_or_list(args.cycles, int)
+ elif args.indices is not None:
+ indices_spec = _parse_slice_or_list(args.indices, int)
+
+ if args.domain == '*':
+ domains = ats_xdmf.find_domains(directories[0])
+ print(f"Found domains: {domains}")
+ else:
+ domains = [args.domain]
+
+ for domain in domains:
+ out_stem = out_prefix if domain == 'domain' else f'{out_prefix}_{domain}'
+ subsetVisFiles(
+ directories=directories,
+ domain=domain,
+ out_directory=out_directory,
+ out_stem=out_stem,
+ in_prefix=args.prefix,
+ times=times_spec,
+ time_unit=args.time_unit,
+ time_tolerance=args.time_tolerance,
+ cycles=cycles_spec,
+ indices=indices_spec,
+ include_vars=args.include_vars,
+ exclude_vars=args.exclude_vars,
+ dry_run=args.dry_run,
+ )
+
+
+if __name__ == '__main__':
+ main()
diff --git a/tools/utils/subset_vis.py b/tools/utils/subset_vis.py
deleted file mode 100644
index d1eeeebdc..000000000
--- a/tools/utils/subset_vis.py
+++ /dev/null
@@ -1,339 +0,0 @@
-#!/usr/bin/env python
-"""Subset ATS visualization output by time/cycle/index range and variable list.
-
-Produces a small, self-contained directory suitable for downloading from HPC
-and opening locally in VisIt.
-"""
-
-import sys
-import os
-import shutil
-import argparse
-import warnings
-import xml.etree.ElementTree as ET
-
-import numpy as np
-import h5py
-
-try:
- sys.path.insert(0, os.path.join(os.environ['ATS_SRC_DIR'], 'tools', 'utils'))
-except KeyError:
- pass
-
-import ats_xdmf
-
-
-# ---------------------------------------------------------------------------
-# Argument parsing helpers
-# ---------------------------------------------------------------------------
-
-def _parse_slice_or_list(s, cast):
- """Parse a slice or comma-separated list string, returning a slice or list.
-
- Slice syntax: 'START:STOP' or 'START:STOP:STEP' (any field may be empty)
- List syntax: 'v1,v2,v3'
-
- cast is applied to each non-empty value (e.g. float, int).
- """
- if ':' in s:
- parts = s.split(':')
- if len(parts) not in (2, 3):
- raise argparse.ArgumentTypeError(
- f"Cannot parse slice argument: {s!r}")
- vals = [cast(p) if p else None for p in parts]
- return slice(*vals)
- else:
- try:
- return [cast(v) for v in s.split(',')]
- except ValueError:
- raise argparse.ArgumentTypeError(
- f"Cannot parse list argument: {s!r}")
-
-
-# ---------------------------------------------------------------------------
-# Core subset function
-# ---------------------------------------------------------------------------
-
-def subsetVisFiles(in_directory, domain, out_directory,
- times=None, time_unit='s', time_tolerance=1.0,
- cycles=None, indices=None,
- include_vars=None, exclude_vars=None,
- dry_run=False):
- """Subset ATS visualization output files by time/cycle/index and variable.
-
- Parameters
- ----------
- in_directory : str
- Directory containing input visualization files.
- domain : str
- ATS domain name (e.g. 'surface', 'domain'). Used to construct
- filenames via ats_xdmf.valid_data_filename().
- out_directory : str
- Directory for output files (created if needed). The same domain
- name is used for output filenames.
- times : slice or list or None
- Result of _parse_slice_or_list() for --times, or None.
- time_unit : str
- One of 's', 'hr', 'd', 'yr'. Units for time filtering values.
- time_tolerance : float
- Tolerance for time matching, in time_unit.
- cycles : slice or list or None
- Result of _parse_slice_or_list() for --cycles, or None.
- indices : slice or list or None
- Result of _parse_slice_or_list() for --indices, or None.
- include_vars : list of str or None
- If given, only these variables are copied.
- exclude_vars : list of str or None
- If given, these variables are excluded.
- dry_run : bool
- If True, print what would be done and return without writing files.
- """
- h5_name = ats_xdmf.valid_data_filename(domain)
- mesh_name = ats_xdmf.valid_mesh_filename(domain)
-
- in_h5 = os.path.join(in_directory, h5_name)
- if not os.path.isfile(in_h5):
- raise RuntimeError(f"No HDF5 data file found at {in_h5!r}.")
-
- vf = ats_xdmf.VisFile(in_directory, domain=domain,
- output_time_unit=time_unit)
-
- # Apply cycle/time/index filter
- if times is not None:
- vf.filterTimes(times, time_unit=time_unit, tolerance=time_tolerance)
- elif cycles is not None:
- vf.filterCycles(cycles)
- elif indices is not None:
- vf.filterIndices(indices)
- # else: all cycles
-
- if not vf.cycles:
- raise RuntimeError("No cycles selected — check your filter arguments.")
-
- # Select variables from the h5 file
- selected_vars = vf.variables(names=include_vars, exclude=exclude_vars)
-
- # Dry run
- if dry_run:
- print(f"Selected {len(vf.cycles)} cycles:")
- for cycle, t in zip(vf.cycles, vf.times):
- print(f" cycle {int(cycle):8d} t = {t:.6g} {time_unit}")
- print(f"\nSelected {len(selected_vars)} variables:")
- for v in selected_vars:
- print(f" {v}")
- vf.close()
- return
-
- # Create output directory
- os.makedirs(out_directory, exist_ok=True)
-
- # Write per-step XMFs for selected cycles
- for cycle in vf.cycles:
- N = int(cycle)
- in_xmf_path = os.path.join(in_directory, f'{h5_name}.{N}.xmf')
- out_xmf_path = os.path.join(out_directory, f'{h5_name}.{N}.xmf')
-
- ET.register_namespace('', 'http://www.w3.org/2001/XInclude')
- tree = ET.parse(in_xmf_path)
- root = tree.getroot()
-
- # Remove unselected Attribute elements
- for grid in root.iter('Grid'):
- to_remove = [
- e for e in list(grid)
- if e.tag == 'Attribute' and e.get('Name') not in selected_vars
- ]
- for e in to_remove:
- grid.remove(e)
-
- tree.write(out_xmf_path, xml_declaration=True, encoding='ASCII')
-
- # Write master VisIt.xmf
- _write_visit_xmf(out_directory, h5_name, vf.cycles)
-
- # Write subset H5 data file
- out_h5 = os.path.join(out_directory, h5_name)
- _write_subset_h5(in_h5, out_h5, vf.cycles, selected_vars)
-
- # Copy mesh H5
- in_mesh_h5 = os.path.join(in_directory, mesh_name)
- out_mesh_h5 = os.path.join(out_directory, mesh_name)
- if os.path.isfile(in_mesh_h5):
- shutil.copyfile(in_mesh_h5, out_mesh_h5)
- print(f"Copied mesh: {out_mesh_h5}")
- else:
- warnings.warn(f"Mesh HDF5 not found: {in_mesh_h5}")
-
- # Copy mesh XMFs
- stem = mesh_name[:-3] # strip '.h5'
- for suffix in [f'{mesh_name}.0.xmf', f'{stem}.VisIt.xmf']:
- in_mesh_xmf = os.path.join(in_directory, suffix)
- if os.path.isfile(in_mesh_xmf):
- out_mesh_xmf = os.path.join(out_directory, suffix)
- shutil.copyfile(in_mesh_xmf, out_mesh_xmf)
- print(f"Copied mesh XMF: {out_mesh_xmf}")
-
- vf.close()
- print(f"Done. Output in: {out_directory}")
- print(f" {len(vf.cycles)} cycles, {len(selected_vars)} variables")
- print(f" Open in VisIt: {os.path.join(out_directory, h5_name[:-3] + '.VisIt.xmf')}")
-
-
-def _write_visit_xmf(out_directory, h5_name, cycles):
- """Write the master VisIt.xmf file with xi:include for each step XMF."""
- stem = h5_name[:-3] # strip '.h5'
- out_path = os.path.join(out_directory, f'{stem}.VisIt.xmf')
-
- xi_ns = 'http://www.w3.org/2001/XInclude'
- xdmf = ET.Element('Xdmf')
- xdmf.set('xmlns:xi', xi_ns)
- xdmf.set('Version', '2.0')
- domain_elem = ET.SubElement(xdmf, 'Domain')
- coll = ET.SubElement(domain_elem, 'Grid')
- coll.set('GridType', 'Collection')
- coll.set('CollectionType', 'Temporal')
- coll.set('Name', 'Mesh')
-
- for cycle in cycles:
- href = f'{h5_name}.{int(cycle)}.xmf'
- xi_elem = ET.SubElement(coll, f'{{{xi_ns}}}include')
- xi_elem.set('href', href)
-
- ET.register_namespace('xi', xi_ns)
- tree = ET.ElementTree(xdmf)
- tree.write(out_path, xml_declaration=True, encoding='ASCII')
- print(f"Wrote VisIt XMF: {out_path}")
-
-
-def _write_subset_h5(in_h5, out_h5, cycles, selected_vars):
- """Copy selected variables and cycles from in_h5 to out_h5."""
- with h5py.File(in_h5, 'r') as src, h5py.File(out_h5, 'w') as dst:
- # Copy the time unit attribute
- if 'time unit' in src.attrs:
- dst.attrs['time unit'] = src.attrs['time unit']
- for var in selected_vars:
- if var not in src:
- warnings.warn(f"Variable {var!r} not found in {in_h5}")
- continue
- grp = dst.create_group(var)
- for cycle in cycles:
- key = str(cycle)
- if key not in src[var]:
- warnings.warn(
- f"Cycle {key} not found for variable {var!r} in {in_h5}")
- continue
- ds = grp.create_dataset(key, data=src[var][key][:])
- if 'Time' in src[var][key].attrs:
- ds.attrs['Time'] = src[var][key].attrs['Time']
- print(f"Wrote data HDF5: {out_h5}")
-
-
-# ---------------------------------------------------------------------------
-# CLI
-# ---------------------------------------------------------------------------
-
-def main():
- parser = argparse.ArgumentParser(
- description=__doc__,
- formatter_class=argparse.RawDescriptionHelpFormatter)
-
- parser.add_argument('domain', metavar='DOMAIN',
- help='ATS domain name (e.g. "surface", "domain"), or '
- '"*" to process all domains found in the input '
- 'directory.')
- parser.add_argument('-d', '--directory', dest='directory', default='.',
- help='Directory containing input visualization files '
- '(default: current directory)')
- parser.add_argument('--output', '-o', dest='output', default=None,
- help='Output directory '
- '(default: DIRECTORY/subset)')
-
- # Time/cycle/index selection (mutually exclusive)
- filter_group = parser.add_mutually_exclusive_group()
- filter_group.add_argument(
- '--times', dest='times', default=None,
- metavar='SLICE_OR_LIST',
- help='Time filter. Formats: '
- '"START:STOP" (range, inclusive), '
- '"START:STOP:INTERVAL" (evenly-spaced targets), '
- '"t1,t2,t3" (specific times). '
- 'Units set by --time-unit.')
- filter_group.add_argument(
- '--cycles', dest='cycles', default=None,
- metavar='SLICE_OR_LIST',
- help='Cycle filter (ATS cycle numbers). Formats: '
- '"START:STOP" or "START:STOP:STEP" (inclusive both ends), '
- '"c1,c2,c3" (specific cycles).')
- filter_group.add_argument(
- '--indices', dest='indices', default=None,
- metavar='SLICE_OR_LIST',
- help='Index filter (0-based position into sorted cycle list). '
- 'Formats: "START:STOP:STEP" (numpy-style, exclusive end), '
- '"i1,i2,i3". Supports negative indices.')
-
- parser.add_argument('--time-unit', dest='time_unit', default='s',
- choices=['s', 'hr', 'd', 'yr', 'noleap'],
- help='Time unit for --times values (default: s)')
- parser.add_argument('--time-tolerance', dest='time_tolerance',
- type=float, default=1.0,
- help='Tolerance for nearest-time matching, in '
- '--time-unit units (default: 1.0 s)')
-
- # Variable selection (mutually exclusive)
- var_group = parser.add_mutually_exclusive_group()
- var_group.add_argument('--include', dest='include_vars',
- action='append', metavar='VAR',
- help='Include this variable in output (repeat for multiple).')
- var_group.add_argument('--exclude', dest='exclude_vars',
- action='append', metavar='VAR',
- help='Exclude this variable from output (repeat for multiple).')
-
- parser.add_argument('--dry-run', dest='dry_run', action='store_true',
- default=False,
- help='Print selected cycles and variables; do not '
- 'write any files.')
-
- args = parser.parse_args()
-
- in_directory = args.directory
-
- out_directory = args.output
- if out_directory is None:
- out_directory = os.path.join(in_directory, 'subset')
-
- # Parse filter arguments into slice objects or lists
- times_spec = None
- cycles_spec = None
- indices_spec = None
-
- if args.times is not None:
- times_spec = _parse_slice_or_list(args.times, float)
- elif args.cycles is not None:
- cycles_spec = _parse_slice_or_list(args.cycles, int)
- elif args.indices is not None:
- indices_spec = _parse_slice_or_list(args.indices, int)
-
- if args.domain == '*':
- domains = ats_xdmf.find_domains(in_directory)
- print(f"Found domains: {domains}")
- else:
- domains = [args.domain]
-
- for domain in domains:
- subsetVisFiles(
- in_directory=in_directory,
- domain=domain,
- out_directory=out_directory,
- times=times_spec,
- time_unit=args.time_unit,
- time_tolerance=args.time_tolerance,
- cycles=cycles_spec,
- indices=indices_spec,
- include_vars=args.include_vars,
- exclude_vars=args.exclude_vars,
- dry_run=args.dry_run,
- )
-
-
-if __name__ == '__main__':
- main()