import inspect
import itertools
import xarray as xr
import numpy as np
import matplotlib.pyplot as plt
import sys
[docs]
def printCones(*args, max_elements=10):
"""
Prints a message including the name of the function from which it was called.
"""
# Get the current frame in the execution stack.
# The '1' indicates that we want the frame of the caller,
# not the frame of printCones itself.
caller_frame = inspect.currentframe().f_back
# Get the name of the function from the caller's frame
caller_name = caller_frame.f_code.co_name
if args:
args_str = " ".join(str(arg) for arg in args)
message = f"{caller_name} : \n\t {args_str}"
else:
message = f"{caller_name} : \n\t"
# Print the message with the calling function's name
print(message)
sys.stdout.flush()
[docs]
def global2Local(globalID, globalCellList):
"""
Find the rank and local cell index corresponding to a global cell id
:param globalID: The globall cell id to locate
:type globalID: int
:param globalCellList: list of global cell ids per rank
:type globalCellList: list
:returns: (rank, localID) if found, otherwise None
:rtype: tuple:
"""
for irank, cellList in enumerate(globalCellList):
if (any(globalID == e for e in cellList)):
return irank, cellList.index(globalID)
[docs]
def is_cell_set(filepath):
"""
Check the OpenFOAM 'class' header of a polyMesh/sets/ file to confirm
it is a cellSet, filtering out faceSet/pointSet artifacts (e.g.
'wrongFaces' written by checkMesh) that are not DA clippings.
:param filepath: path to the set file
:type filepath: str
:returns: True if the file's FoamFile class is cellSet
:rtype: bool
"""
with open(filepath) as f:
for line in f:
stripped = line.strip()
if stripped.startswith("class"):
return stripped.replace(";", "").split()[-1] == "cellSet"
return False
[docs]
def local2Global(localID, globalCellList, rank):
"""
Convert a local cell id on a given rank to its global cell id
:param localID: the local cell id
:type localID: int
:param globalCellList: List of the global cell ids per rank
:type globalCellList: list
:param rank: The rank the local cell id belongs to
:type rank: int
:returns: The global cell id
:rtype: int
"""
return globalCellList[rank][localID]
[docs]
def mergeClips(clipList, globalCellList):
"""
Given a list of conesClip, finds clippings that shares cells and merge them
:param clipList: the list of clippings
:type clipList: list
:param globalCellList: the list of global cell ids
:type globalCellList: list
:returns: the list of global cells id of merged clips
:rtype: list
"""
# First we flat the global cell lists
globalCells = []
for clip in clipList:
flatList = list(itertools.chain(*clip.globalCellList))
flatList.sort()
globalCells.append(flatList)
printCones("Flatten lists", globalCells)
# Then we merge lists if they share elements
element_to_list_index = {}
for i, sublist in enumerate(globalCells):
for element in sublist:
if element not in element_to_list_index:
element_to_list_index[element] = [i]
else:
element_to_list_index[element].append(i)
num_list = len(globalCells)
parent = list(range(num_list))
def find(i):
if parent[i] == i:
return i
parent[i] = find(parent[i])
return parent[i]
def union(i, j):
root_i = find(i)
root_j = find(j)
if root_i != root_j:
parent[root_i] = root_j
return True
return False
shared_elements_log = {}
for element, indices in element_to_list_index.items():
unique_indices = list(set(indices))
if len(unique_indices) > 1:
# if a cell appears in multiple clips, merge those clips
first_index = unique_indices[0]
for i in range(1, len(unique_indices)):
union(first_index, unique_indices[i])
clip_tuple = tuple(sorted(unique_indices))
if clip_tuple not in shared_elements_log:
shared_elements_log[clip_tuple] = []
shared_elements_log[clip_tuple].append(element)
# group lists by their representative parent
merged_groups = {}
merged_clip_indices = {}
for i in range(num_list):
root = find(i)
if root not in merged_groups:
merged_groups[root] = set()
merged_clip_indices[root] = []
merged_groups[root].update(globalCells[i])
merged_clip_indices[root].append(i)
result = [list(group) for group in merged_groups.values()]
for res in result:
res.sort()
printCones("Original clips", len(clipList))
printCones("Merged clips", len(result))
for clip_tuple, elements in shared_elements_log.items():
printCones(f"Clips {clip_tuple} share {len(elements)} elements. {elements}")
for root, indices in merged_clip_indices.items():
if len(indices) > 1:
printCones("Clips merged together:", indices)
return result
def _cones_decode_obs(ds):
"""Decode an observation dataset, keeping the time axis in seconds
"""
import numpy as np
import xarray as xr
decoded = xr.decode_cf(ds, decode_timedelta=False)
time = decoded["time"]
if np.issubdtype(time.dtype, np.timedelta64):
decoded["time"] = time /np.timedelta64(1,"s")
elif np.issubdtype(time.dtype, np.datetime64):
origin = time.values[0]
decoded["time"] = (time - origin) /np.timedelta64(1, "s")
return decoded
[docs]
def cones_open_netcdf_dataset(file, **kwargs):
"""
Open an observation netCDF file, working around files whose variables
carry a 'dtype' (or other xarray-reserved) key directly in their attrs
instead of in their encoding. Some xarray versions refuse to CF-decode
such variables and raise "failed to prevent overwriting existing key
... in attrs" instead of silently moving the key to encoding.
"""
ds = xr.open_dataset(file, engine="netcdf4", decode_cf=False, **kwargs)
for var in ds.variables:
if 'dtype' in ds[var].attrs:
ds[var].encoding['dtype'] = ds[var].attrs.pop('dtype')
return _cones_decode_obs(ds)
[docs]
def cones_count_netcdf_obs(file):
"""
Count the number of observations stored in an observation netCDF file
:param file: Path to the observation netCDF file
:type file: str
:returns: The number of observations
:rtype: int
"""
ds = cones_open_netcdf_dataset(file)
num_obs = ds.dims['obs']
printCones(num_obs)
return int(num_obs)
[docs]
def cones_get_netcdf_obs_coord(file, index):
"""
Read the (x, y, z) coordinates of one observation from a netCDF file
:param file: Path to the observation netCDF file
:type file: str
:param index: Index of the observation
:type index: int
:returns: The [x, y, z] coordinates, or an error message if index is out of bounds
:rtype: list
"""
ds = cones_open_netcdf_dataset(file)
num_obs = ds.dims['obs']
if index >= num_obs:
return f"Error: Observation Index {index} is out of bounds. Max index is {num_obs}"
x = ds.x.values[index]
y = ds.y.values[index]
z = ds.z.values[index]
return [x, y, z]
[docs]
def cones_get_netcdf_var(file, index, var):
"""
Read the value of a given observation variable at a given index
:param file: Path to the observation netCDF file
:type file: str
:param index: Index of the observation
:type index: int
:param var: Name of the variable to read
:type var: str
:returns: The variable value, or an error message on failure
:raises ValueError: if 'var' is not a variable of the dataset
"""
ds = cones_open_netcdf_dataset(file)
if var not in ds.variables:
raise ValueError(f"variable '{var}' not found in {file}")
try:
val = ds[var].isel(obs=index).values
return val
except Exception as e:
return f"Error accessing indices: {e}"
[docs]
def cones_get_netcdf_obs_var(file):
"""
List the observed variable names stored in a netCDF file, excluding
coordinates (x, y, z) and standard deviation companions (*_std)
:param file: path to the observation netCDF file
:type file: str
:returns: The list of observed variable names
:rtype: list
"""
ds = cones_open_netcdf_dataset(file)
actual_vars = [v for v in ds.data_vars if v not in ['x', 'y', 'z'] and not v.endswith('_std')]
return actual_vars
[docs]
def cones_live_plot(i, log_file, truth, test_name):
"""
Matplotlib FuncAnimation callback: redraws the live convergence plot of a
parameter's ensemble mean and std read from 'log_file'
:param i: Frame index passed by the FuncAnimation (unused)
:param log_file: Path to the parameter log file (columns: iter, time, mean, std)
:type log_file: str
:param truth: Ground-truth value plotted as a reference line
:type truth: float
:param test_name: Title of the plot
:type test_name: str
:returns: None
"""
try:
data = np.loadtxt(log_file, ndmin=2)
iters, _, mean, std = data[:, 0], data[:, 1], data[:, 2], data[:, 3]
plt.cla()
plt.title(rf"\textbf{{{test_name}}}", fontsize=14, pad=15)
plt.plot(iters, mean, color='blue', label=r'$\bar{\theta}$', lw=1.5)
plt.fill_between(iters, mean - std, mean + std, color='blue', alpha=0.2, label=r'$\pm \sigma$')
plt.axhline(y=truth, color='red', linestyle='--', label=r'Ground Truth')
plt.xlabel(r'Assimilation Iteration')
plt.ylabel(r'Parameter value')
plt.legend(loc='best', frameon=True)
plt.grid(True, linestyle='--', alpha=0.5)
plt.draw()
except Exception:
pass
[docs]
def matrix_to_txt(matrix, name_of_file):
"""
Dump a matrix to a comma-separated text file for debugging purposes
:param matrix: The matrix to dump
:type matrix: numpy ndarray
:param name_of_file: Output file name, without extension (".txt" is appended)
:type name_of_file: str
:returns: None
"""
np.savetxt(name_of_file+".txt", matrix, fmt="%2e", delimiter=",")
return
def _cones_allocated_cores():
"""Cores in the whole job allocation, not just on this node.
The local figure from sched_getaffinity describes only the node the driver
happens to run on. On a multi-node job it understates the allocation.
:returns: (cores, provenance string), or (None, None) when no scheduler is
visible
"""
import os
spec = os.environ.get("SLURM_JOB_CPUS_PER_NODE")
if spec:
total = 0
for group in spec.split(","):
group = group.strip()
if not group:
continue
count, _, repeat = group.partition("(x")
try:
n = int(count)
total += n * int(repeat.rstrip(")")) if repeat else n
except ValueError:
total = 0
break
if total:
return total, "SLURM_JOB_CPUS_PER_NODE=%s" % spec
for var in ("SLURM_NPROCS", "SLURM_NTASKS", "PBS_NP", "OAR_NP"):
value = os.environ.get(var)
if value:
try:
return int(value), "%s=%s" % (var, value)
except ValueError:
pass
for var in ("PBS_NODEFILE", "OAR_NODEFILE"):
path = os.environ.get(var)
if path:
try:
with open(path) as fh:
n = sum(1 for line in fh if line.strip())
except OSError:
continue
if n:
return n, "%s listing %d slot(s)" % (var, n)
return None, None
def _cones_repo_root(anchor_file):
"""Repository root, given a file inside src/ ."""
import os
return os.path.dirname(os.path.dirname(os.path.abspath(anchor_file)))
def _cones_file_stamp(path):
"""sha256 (truncated) and mtime of a file, or why it is unavailable."""
import hashlib
import os
import time
try:
with open(path, "rb") as fh:
digest = hashlib.sha256(fh.read()).hexdigest()[:16]
stamp = time.strftime("%Y-%m-%d %H:%M",
time.localtime(os.path.getmtime(path)))
return "sha256:%s mtime:%s" % (digest, stamp)
except Exception as exc:
return "unavailable (%s)" % exc
def _cones_source_digest(anchor_file):
"""Fallback identity when git is unavailable: digest of the python tree."""
import hashlib
import os
root = os.path.join(_cones_repo_root(anchor_file), "src")
acc = hashlib.sha256()
try:
for base, dirs, files in os.walk(root):
dirs[:] = sorted(d for d in dirs
if d not in ("__pycache__", "archive", "bk"))
for name in sorted(f for f in files if f.endswith(".py")):
path = os.path.join(base, name)
acc.update(os.path.relpath(path, root).encode())
with open(path, "rb") as fh:
acc.update(fh.read())
return "sha256:%s (python sources)" % acc.hexdigest()[:16]
except Exception as exc:
return "unavailable (%s)" % exc
def _cones_git_state(anchor_file):
"""Git identity of the source tree, as a list of header lines."""
import subprocess
repo = _cones_repo_root(anchor_file)
def _run(args):
return subprocess.run(args, cwd=repo, capture_output=True, text=True,
timeout=15).stdout.strip()
try:
full = _run(["git", "rev-parse", "HEAD"])
if not full:
return ["commit : not a git checkout",
"source digest : %s" % _cones_source_digest(anchor_file)]
lines = ["commit : %s" % full,
"describe : %s" % _run(["git", "describe", "--always",
"--dirty", "--tags"]),
"branch : %s" % _run(["git", "rev-parse",
"--abbrev-ref", "HEAD"]),
"commit date : %s" % _run(["git", "show", "-s",
"--format=%ci", "HEAD"])]
dirty = _run(["git", "status", "--porcelain", "--untracked-files=no"])
if dirty:
entries = dirty.splitlines()
lines.append("UNCOMMITTED : %d tracked file(s) modified -- the"
" commit above does NOT describe this run"
% len(entries))
for entry in entries[:12]:
lines.append(" %s" % entry)
if len(entries) > 12:
lines.append(" ... and %d more"
% (len(entries) - 12))
else:
lines.append("uncommitted : none")
return lines
except Exception as exc:
return ["commit : unavailable (%s)" % exc]
def _cones_find_repo_root(start_path):
"""
Walk upward from start_path until a .git is found
Unlike _cones_repo_root, this works from any file at any nesting
depth.
Returns None if no .git rutnrs up within a few levels
"""
import os
path = os.path.dirname(os.path.abspath(start_path))
for _ in range(6):
if os.path.isdir(os.path.join(path, ".git")):
return path
parent = os.path.dirname(path)
if parent == path:
break
path = parent
return None
def _cones_version():
"""CONES version, read from the repository's git tags.
`git describe --tags --always` reports the nearest tag reachable from
HEAD -- e.g. "v1.4.1" exactly on a tagged commit, or
"v1.4.0-69-gb86973ae" when HEAD is 69 commits past the last tag (also
naming the commit, unlike a bare tag). Computed once at import time;
never raises, so a source tree without git (e.g. a tarball release)
still imports cleanly.
"""
import subprocess
repo = _cones_find_repo_root(__file__)
if repo is None:
return "unknown (not a git checkout)"
try:
out = subprocess.run(["git", "describe", "--tags", "--always"],
cwd=repo, capture_output=True, text=True,
timeout=15).stdout.strip()
return out if out else "unknown (no commits)"
except Exception:
return "unknown (git unavailable)"
CONES_VERSION=_cones_version()
def _cones_build_state(anchor_file):
"""Identity of the compiled C++ artefacts plus a staleness check.
The git hash describes the sources; the C++ side is built separately with
wmake into $(FOAM_USER_LIBBIN)/libconestoolbox and
$(FOAM_USER_APPBIN)/conesIcoFoam, so a run started after editing EnKF.C
without rebuilding would otherwise be recorded under a commit that does
not describe what executed.
"""
import glob
import hashlib
import os
import time
targets = []
for env_var, pattern in (("FOAM_USER_LIBBIN", "libconestoolbox*"),
("FOAM_USER_APPBIN", "conesIcoFoam"),
("FOAM_APPBIN", "conesFoam")):
root = os.environ.get(env_var)
if root:
targets.extend(sorted(glob.glob(os.path.join(root, pattern))))
if not targets:
return ["compiled C++ : no artefact found via FOAM_USER_LIBBIN /"
" FOAM_USER_APPBIN (variables not exported to this process?)"]
newest, newest_path = 0.0, None
src_root = os.path.join(_cones_repo_root(anchor_file), "solvers")
try:
for base, dirs, files in os.walk(src_root):
dirs[:] = [d for d in dirs if d not in ("bk", "archive", "Make")]
for name in files:
if name.endswith((".C", ".H")):
path = os.path.join(base, name)
mtime = os.path.getmtime(path)
if mtime > newest:
newest, newest_path = mtime, path
except Exception:
pass
lines = []
for path in targets:
try:
with open(path, "rb") as fh:
digest = hashlib.sha256(fh.read()).hexdigest()[:16]
built = os.path.getmtime(path)
lines.append("compiled C++ : %s" % path)
lines.append(" sha256:%s built:%s"
% (digest, time.strftime("%Y-%m-%d %H:%M",
time.localtime(built))))
if newest_path and newest > built:
lines.append(" STALE BUILD: %s is newer (%s)"
" -- rebuild with wmake before trusting this run"
% (os.path.relpath(newest_path, src_root),
time.strftime("%Y-%m-%d %H:%M",
time.localtime(newest))))
except Exception as exc:
lines.append("compiled C++ : %s -- unavailable (%s)"
% (path, exc))
return lines
def _cones_cpu_count():
"""Cores usable by this process, and the machine total. They differ under
cgroup limits, taskset, or a scheduler allocation."""
import os
total = os.cpu_count()
try:
usable = len(os.sched_getaffinity(0))
except AttributeError:
usable = total
return usable, total
def _cones_environment_state():
"""User, machine and scheduler identity, as header lines. Only an explicit
allowlist of scheduler variables is read: the environment is never dumped
wholesale, since it can carry credentials."""
import getpass
import os
import platform
import socket
try:
user = getpass.getuser()
except Exception:
user = os.environ.get("USER") or os.environ.get("LOGNAME") or "unknown"
try:
fqdn = socket.getfqdn()
except Exception:
fqdn = platform.node()
usable, total = _cones_cpu_count()
lines = ["user : %s" % user,
"host : %s" % platform.node(),
"fqdn : %s" % fqdn,
"platform : %s" % platform.platform(),
"working dir : %s" % os.getcwd(),
"cpu cores : %s usable of %s on this node"
% (usable, total)]
schedulers = (
("SLURM", [("job id", "SLURM_JOB_ID"),
("job name", "SLURM_JOB_NAME"),
("partition", "SLURM_JOB_PARTITION"),
("account", "SLURM_JOB_ACCOUNT"),
("nodes", "SLURM_JOB_NODELIST"),
("ntasks", "SLURM_NTASKS"),
("cpus per task", "SLURM_CPUS_PER_TASK"),
("submit dir", "SLURM_SUBMIT_DIR")]),
("PBS", [("job id", "PBS_JOBID"),
("job name", "PBS_JOBNAME"),
("queue", "PBS_QUEUE"),
("nodefile", "PBS_NODEFILE")]),
("OAR", [("job id", "OAR_JOB_ID"),
("nodefile", "OAR_NODEFILE")]),
)
for name, keys in schedulers:
present = [(label, os.environ[var]) for label, var in keys
if os.environ.get(var)]
if present:
lines.append("scheduler : %s" % name)
for label, value in present:
lines.append(" %-14s: %s" % (label, value))
break
else:
lines.append("scheduler : none detected (interactive run)")
return lines
def _cones_obs_summary(path, obs_window, dt):
"""Dimensions and time span of the observation database, and how many
analyses that record can cover at the configured window."""
try:
import netCDF4
with netCDF4.Dataset(path) as ds:
n_time = len(ds.dimensions["time"])
n_obs = len(ds.dimensions["obs"])
t = ds.variables["time"][:]
t0, t1 = float(t[0]), float(t[-1])
lines = ["n_obs=%d n_time=%d t=[%.4f, %.4f] s"
% (n_obs, n_time, t0, t1)]
if obs_window and dt:
covered = int((t1 - t0) / (float(obs_window) * float(dt)))
lines.append("analyses covered by this record at the configured "
"window: %d" % covered)
lines.append("past that point the observation is held at the last "
"available sample")
return lines
except Exception as exc:
return ["unavailable (%s)" % exc]
[docs]
def cones_run_id(conesSet, driver_file):
"""
Unique, self-describing identifier for this run
Three fields joined by underscores:
<UTC timestamp> when the run started
cfg<8 hex> digest of conesDict, identifies the configuration
g<7 jex> short git commit of the sources, or "nogit"
:param conesSet: settings object in use
:param driver_file: pass __file__ from the calling driver
:returns: the run identifier
:rtype: str
"""
import hashlib
import subprocess
import time
ts = time.strftime("%Y%m%dT%H%M%SZ", time.gmtime())
try:
with open(conesSet.conesDict, "rb") as fh:
cfg = hashlib.sha256(fh.read()).hexdigest()[:8]
except Exception:
cfg = "nocfg"
try:
commit = subprocess.run(["git", "rev-parse", "--short=7", "HEAD"],
cwd=_cones_repo_root(driver_file),
capture_output=True, text=True,
timeout=15).stdout.strip()
code = "g" + commit if commit else "nogit"
except Exception:
code = "nogit"
return "%s_cfg%s_%s" % (ts, cfg, code)
def _cones_refresh_latest_symlink(target, link="par.log"):
"""Point `link` at `target` so that scripts expecting a fixed file name
keep working. Atomic, and silent where symlinks are unavailable."""
import os
try:
tmp = link + ".tmp"
if os.path.lexists(tmp):
os.remove(tmp)
os.symlink(os.path.basename(target), tmp)
os.replace(tmp, link)
return True
except Exception:
return False
def _cones_rank_layout(n_proc_app):
"""
One-line summary of the ranks-per-model layout.
"""
try:
counts = {}
for value in n_proc_app:
key = int(value)
counts[key] = counts.get(key, 0) + 1
except TypeError: # a scalar: cannot wrap, pass through
text = "%s" % n_proc_app
else:
if len(counts) == 1:
ranks, models = next(iter(counts.items()))
text = "%d for each of %d model instance(s)" % (ranks, models)
else:
groups = sorted(counts.items(), reverse=True)
parts = ["%d model(s) x %d rank(s)" % (m, r) for r, m in groups[:8]]
if len(groups) > 8:
parts.append("and %d further group(s)" % (len(groups) - 8))
text = "irregular -- " + ", ".join(parts)
return " ".join(text.split())[:200] # never multi-line, never unbounded
[docs]
def cones_prior_digest(params, decimals=6):
"""Short reproducible digest of a realised parameter ensemble.
:param params: parameter ensemble, either (n_ensemble, n_parameters)
:param decimals: rounding applied before hashing
:returns: first 16 hexadecimal characters of the sha256 digest
:rtype: str
"""
import hashlib
import numpy as np
a = np.atleast_2d(np.asarray(params, dtype=float))
if a.shape[0] < a.shape[1]:
a = a.T
a = np.ascontiguousarray(np.round(a, decimals))
return hashlib.sha256(a.tobytes()).hexdigest()[:16]
[docs]
def cones_write_prior_block(path, params, label="forecast", decimals=6):
"""Append a digest and summary of a parameter ensemble to the run log.
:param path: run log to append to; nothing is written if it is None
:param params: parameter ensemble as gathered, either orientation
:param label: name of the ensemble, for drivers carrying more than one
(fine and coarse for MGEnKF, principal, control and ancillary for
MFEnKF)
:returns: the digest, or None if nothing was written
:rtype: str or None
"""
import numpy as np
if path is None or params is None:
return None
a = np.atleast_2d(np.asarray(params, dtype=float))
if a.shape[0] < a.shape[1]:
a = a.T
n_e, n_p = a.shape
digest = cones_prior_digest(a, decimals)
lines = ["#",
"# Realised %s ensemble at the first analysis" % label,
"# prior digest : %s (sha256 of the ensemble rounded to %d"
" decimals, shape (n_e, n_par) = (%d, %d))"
% (digest, decimals, n_e, n_p),
"# parameter mean sd min max"]
for j in range(n_p):
col = a[:, j]
lines.append("# %-9d %12.5f %12.5f %12.5f %12.5f"
% (j, col.mean(),
col.std(ddof=1) if n_e > 1 else 0.0,
col.min(), col.max()))
lines.append("#")
with open(path, "a") as handle:
handle.write("\n".join(lines) + "\n")
return digest