"""Omega surfaces"""
import numpy as np
from time import time
from scipy.sparse import csc_matrix
from scipy.sparse.linalg import lsqr
from scipy.sparse.linalg import spsolve
from neutralocean.surface.isopycnal import _isopycnal
from neutralocean.surface._vertsolve import _make_vertsolve
from neutralocean.ppinterp import make_pp
from neutralocean.bfs import bfs_conncomp1, bfs_conncomp1_wet
from neutralocean.grid.graph import edges_to_graph
from neutralocean.ntp import ntp_epsilon_errors, ntp_epsilon_errors_norms
from neutralocean.lib import (
xr_to_np,
_xrs_in,
_xr_out,
_process_pin_cast,
_process_casts,
_process_eos,
aggsum,
)
from neutralocean.mixed_layer import mixed_layer
[docs]def omega_surf(S, T, P, grid, **kw):
"""Calculate an omega surface from structured ocean data.
Given 3D salinity, temperature, and pressure or depth data arranged on a
rectilinear grid, calculate a 2D omega surface [1]_ [2]_, which is a
highly accurate approximately neutral surface.
Parameters
----------
S, T, P : ndarray or xarray.DataArray
See `potential_surf`
p_init : ndarray, Default None
Pressure or depth on the initial approximately neutral surface.
See Examples section.
ref : float, or tuple of float of length 2
If `p_init` is None, the reference value(s) for the initial potential
density surface or in-situ density (specific volume) anomaly surface
that initializes the omega surface algorithm. If `ref` is a scalar, a
potential density urface is used, and if `ref` is None, the reference
`P` is `pin_p` (i.e. taken local to the pinning location). If `ref`
is a tuple of length two, a in-situ density anomaly surface is used,
and if `ref` is (None, None), then the reference `S` and `T` values
are taken local to the pinning location (pressure or depth `pin_p` on
the pinning cast `pin_cast`).
See Examples section.
isoval : float
Isovalue for the initial potential density or in-situ density anomaly
surface when `p_init` is not given. Units are same as returned by
the function specified by `eos`.
See Examples section.
pin_p, pin_cast :
See `potential_surf`
Returns
-------
s, t, p : ndarray or xarray.DataArray
practical / Absolute salinity, potential / Conservative temperature,
and pressure / depth on surface
d : dict
Diagnostics.
The first four listed below give information going into the `i`'th
iteration (e.g. the 0'th element is about the initial surface).
The rest give information about what the `i`'th iteration did (and
hence their 0'th elements are irrelevant).
`"e_MAV"` : array of float
Mean Absolute Value of the ϵ neutrality error on the surface,
area-weighted. Units are those of `eos` return values divided by
those of `dist*` inputs.
`"e_RMS"` : array of float
As `"e_MAV"` but for the Root Mean Square.
`"n_wet"`: array of float
Number of wet casts (surface points).
`"timer"` : array of float
Time spent on each iteration, excluding set-up (approximately) and diagnostics.
`"ϕ_MAV"` : array of float
Mean Absolute Value of the Locally Referenced Potential Density
perturbation, per iteration
`"Δp_MAV"` : array of float
Mean Absolute Value of the pressure or depth change from one
iteration to the next
`"Δp_RMS"` : array of float
Root Mean Square of the pressure or depth change from one
iteration to the next
`"Δp_Linf"` : array of float
Maximum absolute value (infinity norm) of the pressure or depth
change from one iteration to the next
`"n_newly_wet"` : array of int
Number of casts that are newly wet, per iteration
`"timer_bfs"` : array of float
Time spent in Breadth-First Search including wetting, per iteration.
`"timer_mat"` : array of float
Time spent building and solving the matrix problem, per iteration.
`"timer_update"` : array of float
Time spent vertically updating the surface.
Other Parameters
----------------
grid, vert_dim, grav, rho_c, interp, diags, output, TOL_P_SOLVER :
See `potential_surf`
eos : str or tuple of functions, Default 'gsw'
As in `potential_surf`, excluding the option to pass a single function.
The omega surface algorithm relies on knowing the partial derivatives
of the equation of state with respect to salinity and temperature, so
the `eos_s_t` function is also required if given as a tuple of functions.
ITER_MIN : int, Default 1
Minimum number of iterations.
ITER_MAX : int, Default 10
Maximum number of iterations.
ITER_START_WETTING : int, Default 1
Iteration on which wetting begins. Set to `np.inf` (`ITER_MAX` + 1
would also do) to deactivate.
ITER_STOP_WETTING : int, Default 5
The last iteration on which to perform wetting. This can be useful to
avoid pesky water columns that repeatedly wet then dry.
TOL_LRPD_MAV : float, Default 1e-7
Exit iterations when the mean absolute value of the Locally Referenced
Potential Density perturbation that updates the surface from one
iteration to the next is less than this value. Units are [kg m-3],
even if `eos` returns a specific volume. Set to 0 to deactivate.
TOL_P_CHANGE_RMS : float, Default 0.0
Exit iterations when the root mean square of the pressure or depth
change on the surface from one iteration to the next is less than this
value. Set to 0 to deactivate. Units are the same as `P` [dbar or m].
p_ml : ndarray or dict, Default None
If a dict, the pressure or depth at the base of the mixed layer is
computed using `mixed_layer` with p_ml passed as keyword arguments,
enabling control over the parameters in that function.
See `mixed_layer` for details.
If an ndarray (of the same shape as the lateral dimensions of `S`),
the pressure or depth at the base of the mixed layer in each water
column.
When the surface's pressure is shallower than `p_ml` in any water
column, it is set to NaN (a "dry" water column). This is not applied
to the initial surface, but only to the surface after the first
iteration, as the initial surface could be very far from neutral.
If None, the mixed layer is not removed.
OMEGA_FORMULATION : str, Default 'poisson'
Specify how the matrix problem is set up and solved. Options are
- `'poisson'`, to solve the Poisson problem as in [1]_ with Cholesky, or
- `'gradient'`, to solve the overdetermined gradient equations as in [2]_
using LSQR.
Examples
--------
omega surfaces require a pinning cast and initial surface. The surface is
iteratively updated while remaining fixed at the pinning cast. The
initial surface can be provided directly, as the surface with pressure or
depth given by `p_init`, in the following method:
>>> omega_surf(S, T, P, pin_cast, p_init, ...)
Alternatively, a potential density surface or an in-situ density (specific
volume) anomaly surface can be used as the initial surface. To do this,
use one of the following two methods:
>>> omega_surf(S, T, P, ref, isoval, pin_cast, ...)
>>> omega_surf(S, T, P, ref, pin_p, pin_cast, ...)
For more info on these methods, see the Examples section of `potential_surf`.
Note that `pin_cast` is always a required input. Note that `ref` is
needed to distinguish which of the two types of isopycnals will be used as
the initial surface.
Notes
-----
See `potential_surf` Notes.
.. [1] Stanley, G. J., McDougall, T. J., & Barker, P. M. (2021). Algorithmic
Improvements to Finding Approximately Neutral Surfaces. Journal of
Advances in Modeling Earth Systems, 13(5), e2020MS002436.
.. [2] Klocker, A., McDougall, T. J., & Jackett, D. R. (2009). A new method
for forming approximately neutral surfaces. Ocean Science, 5 (2), 155-172.
"""
ref = kw.get("ref")
pin_p = kw.get("pin_p")
pin_cast = kw.get("pin_cast")
p_init = kw.get("p_init")
vert_dim = kw.get("vert_dim", -1)
p_ml = kw.get("p_ml")
diags = kw.get("diags", True)
output = kw.get("output", True)
eos = kw.get("eos", "gsw")
rho_c = kw.get("rho_c")
grav = kw.get("grav")
ITER_MIN = kw.get("ITER_MIN", 1)
ITER_MAX = kw.get("ITER_MAX", 10)
ITER_START_WETTING = kw.get("ITER_START_WETTING", 1)
ITER_STOP_WETTING = kw.get("ITER_STOP_WETTING", 5)
TOL_P_SOLVER = kw.get("TOL_P_SOLVER", 1e-4)
TOL_LRPD_MAV = kw.get("TOL_LRPD_MAV", 1e-7)
TOL_P_CHANGE_RMS = kw.get("TOL_P_CHANGE_RMS", 0.0)
OMEGA_FORMULATION = kw.get("OMEGA_FORMULATION", "poisson")
interp = kw.get("interp", "linear")
# Build function that calculates coefficients of a piecewise polynomial
# interpolant, doing 1 problem at a time, and knowing there will be no nans
# in the input data.
ppc_fn = make_pp(interp, kind="1", out="coeffs", nans=False)
# Build function that evaluates two piecewise polynomial interpolants (with
# the same set of independent data), doing n problems at a time.
interp_two = make_pp(interp, kind="u", out="interp", num_dep_vars=2)
sxr, txr, pxr = _xrs_in(S, T, P, vert_dim) # before _process_casts
pin_cast = _process_pin_cast(pin_cast, S) # call before _process_casts
S, T, P = _process_casts(S, T, P, vert_dim)
eos, eos_s_t = _process_eos(eos, grav, rho_c, need_s_t=True)
# Save shape of horizontal dimensions, then flatten horiz dims to 1D.
surf_shape = S.shape[0:-1]
N = np.prod(surf_shape) # number of nodes (water columns)
S, T, P = (np.reshape(X, (N, -1)) for X in (S, T, P))
# Update pinning cast to a linear index, unless already so
if isinstance(pin_cast, (tuple, list)):
pin_cast = np.ravel_multi_index(pin_cast, surf_shape)
# Prepare grid ratios for matrix problem.
distratio = grid["distperp"] / grid["dist"]
# Pre-calculate grid adjacency needed for Breadth First Search
edges = grid["edges"]
graph = edges_to_graph(edges, N)
if OMEGA_FORMULATION.lower() == "poisson":
global_solver = _omega_matsolve_poisson
elif OMEGA_FORMULATION.lower() == "gradient":
global_solver = _omega_matsolve_gradient
distratio = np.sqrt(distratio)
else:
raise ValueError(
f"Unknown OMEGA_FORMULATION. Given {OMEGA_FORMULATION}"
)
if eos(34.5, 3.0, 1000.0) < 1.0:
# Convert from a density tolerance [kg m^-3] to a specific volume tolerance [m^3 kg^-1]
TOL_LRPD_MAV = TOL_LRPD_MAV / 1000.0**2
# Pre-allocate arrays for diagnostics
if diags:
d = {
"e_MAV": np.zeros(ITER_MAX + 1, dtype=np.float64),
"e_RMS": np.zeros(ITER_MAX + 1, dtype=np.float64),
"timer": np.zeros(ITER_MAX + 1, dtype=np.float64),
"ϕ_MAV": np.zeros(ITER_MAX + 1, dtype=np.float64),
"Δp_MAV": np.zeros(ITER_MAX + 1, dtype=np.float64),
"Δp_RMS": np.zeros(ITER_MAX + 1, dtype=np.float64),
"Δp_Linf": np.zeros(ITER_MAX + 1, dtype=np.float64),
"n_wet": np.zeros(ITER_MAX + 1, dtype=int),
"n_newly_wet": np.zeros(ITER_MAX + 1, dtype=int),
"timer_bfs": np.zeros(ITER_MAX + 1, dtype=np.float64),
"timer_mat": np.zeros(ITER_MAX + 1, dtype=np.float64),
"timer_update": np.zeros(ITER_MAX + 1, dtype=np.float64),
}
else:
d = dict()
timer = time()
if p_init is None:
# Calculate an initial "potential" or "anomaly" surface
if isinstance(ref, (tuple, list)) and len(ref) == 2:
ans_type = "anomaly"
else:
ans_type = "potential"
# Update arguments with pre-processed values
kw["vert_dim"] = -1 # Since S, T, P already reordered
kw["diags"] = False # Will make our own diags next
kw["eos"] = eos
kw["pin_cast"] = pin_cast # update with the 1D value
s, t, p, _ = _isopycnal(ans_type, S, T, P, **kw)
else:
# Handling and error checking on p_init
p_init = xr_to_np(p_init)
if not isinstance(p_init, np.ndarray):
raise TypeError(
'If provided, "p_init" or "p_init.values" must be an ndarray'
)
if p_init.shape != surf_shape:
raise ValueError(
f'"p_init" should contain a 2D array of size {surf_shape};'
f" found size {p_init.shape}"
)
p_init = np.reshape(p_init, -1) # now reshape to 1D
if pin_p is not None and pin_p != p_init[pin_cast]:
raise ValueError("pin_p does not match p_init at pin_cast")
p = p_init.copy()
# Interpolate S and T onto the surface
# TODO: Update this to handle ice shelf cavity friendly interpolation
s, t = interp_two(p, P, S, T, 0)
pin_p = p[pin_cast]
if np.isnan(s[pin_cast]):
raise RuntimeError("The initial surface is NaN at the reference cast.")
if ITER_MAX > 1 and isinstance(p_ml, dict):
# Compute the mixed layer from parameter inputs
p_ml = mixed_layer(S, T, P, eos, **p_ml)
if p_ml is None:
# Prepare array as needed for bfs_conncomp1_wet
p_ml = np.full(p.shape, -np.inf)
# p_ml = np.broadcast_to(-np.inf, (ni, nj)) # DEV: Doesn't work with @numba.njit
else:
# Ensure p_ml is 1D array, in case it was given as a 2D array.
p_ml = np.reshape(p_ml, -1)
# ensure same nan structure between s, t, and p. Just in case user gives
# np.full((ni,nj), 1000) for a 1000dbar isobaric surface, for example
p[np.isnan(s)] = np.nan
if diags:
d["timer"][0] = time() - timer
e_RMS, e_MAV = ntp_epsilon_errors_norms(s, t, p, grid, eos_s_t)
d["e_RMS"][0], d["e_MAV"][0] = e_RMS, e_MAV
n_wet = np.sum(np.isfinite(p))
d["n_wet"][0] = n_wet
if output:
print(
"iter |"
" MAV(ϕ) |"
" RMS(Δp) |"
" # wet casts (# new) |"
" RMS(ϵ) |"
" time (s)"
)
print(
f"{0:4d} |"
f" |"
f" {d['n_wet'][0]:11d} |"
f" {e_RMS:.8e} |"
f" {d['timer'][0]:.3f}"
)
vertsolve = _make_vertsolve(eos, ppc_fn, "omega")
# --- Begin iterations
# Note: the surface exists wherever p is non-nan. The nan structure of s
# and t is made to match that of p when the vertical solve step is done.
Δp_RMS = 0.0 # ensure this is defined; needed if TOL_P_CHANGE_RMS == 0
for iter_ in range(1, ITER_MAX + 1):
timer = time()
# --- Remove the Mixed Layer
if p_ml is not None and iter_ > 1:
p[p < p_ml] = np.nan
# --- Determine the connected component containing the reference cast, via Breadth First Search
timer_loc = time()
if iter_ >= ITER_START_WETTING and iter_ <= ITER_STOP_WETTING:
bfsqu, n_newly_wet = bfs_conncomp1_wet(
graph.indptr,
graph.indices,
pin_cast,
s,
t,
p,
S,
T,
P,
TOL_P_SOLVER,
eos,
ppc_fn,
p_ml,
)
else:
bfsqu = bfs_conncomp1(
graph.indptr, graph.indices, pin_cast, np.isfinite(p)
)
n_newly_wet = 0
timer_bfs = time() - timer_loc
# --- Solve global matrix problem for the exactly determined Poisson equation
timer_loc = time()
# Pre-sort the BFS queue: tests in both MATLAB and Python on OCCA data
# show this gives an overall speedup for both Poisson and Gradient formulations.
bfsqu = np.sort(bfsqu)
ϕ = global_solver(s, t, p, edges, distratio, bfsqu, pin_cast, eos_s_t)
timer_mat = time() - timer_loc
# --- Update the surface (mutating s, t, p by vertsolve)
timer_loc = time()
p_old = p.copy() # Record old surface for pinning and diagnostics
vertsolve(s, t, p, S, T, P, ϕ, TOL_P_SOLVER)
# DEV: time seems indistinguishable from using factory function as above
# _vertsolve_omega(s, t, p, S, T, P, Sppc, Tppc, ϕ, TOL_P_SOLVER, eos)
# Force p to stay constant at the reference column, identically.
# This avoids any intolerance from the vertical solver.
p[pin_cast] = pin_p
timer_update = time() - timer_loc
# --- Closing Remarks
ϕ_MAV = np.nanmean(abs(ϕ))
if diags or TOL_P_CHANGE_RMS > 0:
Δp = p - p_old
Δp_RMS = np.sqrt(np.nanmean(Δp**2))
if diags:
d["timer"][iter_] = time() - timer
Δp_MAV = np.nanmean(abs(Δp))
Δp_Linf = np.nanmax(abs(Δp))
# Diagnostics about what THIS iteration did
d["ϕ_MAV"][iter_] = ϕ_MAV
d["Δp_MAV"][iter_] = Δp_MAV
d["Δp_RMS"][iter_] = Δp_RMS
d["Δp_Linf"][iter_] = Δp_Linf
d["n_newly_wet"][iter_] = n_newly_wet
d["timer_mat"][iter_] = timer_mat
d["timer_update"][iter_] = timer_update
d["timer_bfs"][iter_] = timer_bfs
# Diagnostics about the state AFTER this iteration
e_RMS, e_MAV = ntp_epsilon_errors_norms(s, t, p, grid, eos_s_t)
d["e_RMS"][iter_], d["e_MAV"][iter_] = e_RMS, e_MAV
n_wet = np.sum(np.isfinite(p))
d["n_wet"][iter_] = n_wet
if output:
print(
f"{iter_:4d} |"
f" {ϕ_MAV:.8e} |"
f" {Δp_RMS:.8e} |"
f" {n_wet:11d} ({n_newly_wet:5}) |"
f" {e_RMS:.8e} |"
f" {d['timer'][iter_]:.3f}"
)
# --- Check for convergence
if (
ϕ_MAV < TOL_LRPD_MAV or Δp_RMS < TOL_P_CHANGE_RMS
) and iter_ >= ITER_MIN:
break
if diags:
# Trim diagnostics
for k, v in d.items():
d[k] = v[0 : iter_ + (k in ("e_MAV", "e_RMS"))]
# Reshape (from 1D arrays) and put into DataArrays if appropriate
s, t, p = (
_xr_out(np.reshape(x, surf_shape), xxr)
for (x, xxr) in ((s, sxr), (t, txr), (p, pxr))
)
return s, t, p, d
def _omega_matsolve_gradient(s, t, p, edges, sqrtdistratio, m, mref, eos_s_t):
"""Solve the Gradient formulation of the omega-surface global matrix problem
Parameters
----------
s, t, p, edges, m, mref, eos_s_t :
See `_omega_matsolve_poisson`
sqrtdistratio : array
The square-root of the distance of the interface between adjacent
water columns divided by the square-root of the distance between
adjacent water columns. That is, `np.sqrt(distperp / dist)`,
where `distperp` and `dist` are as in the `geoemtry` input to
`omega_surf`.
Returns
-------
ϕ : ndarray
See `_omega_matsolve_poisson`
"""
N = len(m) # Number of water columns
ϕ = np.full(p.size, np.nan, dtype=p.dtype)
# If there is only one water column, there are no equations to solve,
# then the solution is simply phi = 0 at that water column, and nan elsewhere.
# Note, N >= 1 should be guaranteed by omega_surf(), so N <= 1 should imply
# N == 1. If N >= 1 weren't guaranteed (m empty), this would throw an error.
if N <= 1: # There are definitely no equations to solve
ϕ[m[0]] = 0.0 # Leave this isolated pixel at current pressure
return ϕ.reshape(p.shape)
a, b, e, fac, ref = _omega_matsolve_helper(
s, t, p, edges, sqrtdistratio, m, mref, eos_s_t
)
rhs = np.concatenate((-e, [0.0])) # add 0 for pinning equation
# Build columns for matrix, including extra entry for pinning equation.
# Note m[ref] is the reference cast, so the ref'th entry in the solution
# vector of the matrix column corresponds to ϕ at the reference cast.
c = np.concatenate((a, b, [ref]))
# E = number rows in matrix. Round down ignores pinning equation
E = len(c) // 2
# r = [0, 1, ..., E-1, 0, 1, ..., E-1, E]
r = np.concatenate((np.tile(np.arange(E), 2), [E]))
# When distratio = 1, v is [1, 1, ..., 1, -1, -1, ..., -1, 1e-2]
v = np.concatenate((-fac, fac, [1e-2]))
mat = csc_matrix((v, (r, c)), shape=(E + 1, N))
sol = lsqr(mat, rhs)[0]
# Heave solution to be exactly 0 at pinning cast (to fix any intolerance
# caused by lsqr converging before reaching the exact solution)
sol -= sol[ref]
ϕ[m] = sol
ϕ = ϕ.reshape(p.shape)
return ϕ
def _omega_matsolve_poisson(s, t, p, edges, distratio, m, mref, eos_s_t):
"""Solve the Poisson formulation of the omega-surface global matrix problem
Parameters
----------
s, t, p : ndarray
Salinity, temperature, pressure on the surface
edges : tuple of length 2 of 1D arrays
See grid['edges'] from `omega_surf`
distratio : array
The distance of the interface between adjacent water columns divided by
the distance between adjacent water columns, in the same order as `edges`.
That is, `grid['distperp'] / grid['dist']` where `grid` is as input
to `omega_surf`.
m : array
Linear indices to the nodes in order of the BFS, ie all casts in this
connected component.
That is, `m = qu[0:qt+1]` where `qu` and `qt` are outputs from
`bfs_conncomp1` in bfs.py.
mref : int
Linear index to the reference cast, at which ϕ will be zero
eos_s_t : function
Function returning the partial derivatives of the equation of state
with respect to S and T.
Returns
-------
ϕ : ndarray
Locally referenced potential density (LRPD) perturbation.
Vertically heaving the surface so that its LRPD in water column m
increases by the m'th element of ϕ will yield a more neutral surface.
"""
N = len(m) # number of water columns
ϕ = np.full(p.size, np.nan, dtype=p.dtype) # prealloc space
# If there is only one water column, there are no equations to solve,
# then the solution is simply phi = 0 at that water column, and nan elsewhere.
# Note, N >= 1 should be guaranteed by omega_surf(), so N <= 1 should imply
# N == 1. If N >= 1 weren't guaranteed (m empty), this would throw an error.
if N <= 1: # There are definitely no equations to solve
ϕ[m[0]] = 0.0 # Leave this isolated pixel at current pressure
return ϕ.reshape(p.shape)
# Get list of edges (a,b), ϵ neutrality errors, and geometric factors
# for the current connected component containing the reference cast, and
# map everything onto a set of N
a, b, e, fac, ref = _omega_matsolve_helper(
s, t, p, edges, distratio, m, mref, eos_s_t
)
# Divergence of ϵ,
# D = ∑_{n ∈ N(m)} ϵₘₙ
# = ∑_{j=1}^E δ_{m, aⱼ} ϵ_{m, bⱼ} + ∑_{j=1}^E δ_{m, bⱼ} ϵ_{m, aⱼ}
# = ∑_{j=1}^E δ_{m, aⱼ} ϵ_{m, bⱼ} - ∑_{j=1}^E δ_{m, bⱼ} ϵ_{aⱼ, m}
# achieved by
D = aggsum(e, a, N) - aggsum(e, b, N)
# Prepare diagonal entries of negative Laplacian. For uniform geometry,
# this simply counts the number of edges incident upon each node. For
# rectilinear grids, this value will be 4 for a typical node, but can be
# less near boundaries of the connected component.
diag = aggsum(fac, a, N) + aggsum(fac, b, N)
# Build the rows, columns, and values of the sparse matrix
r = np.concatenate((a, b, np.arange(N)))
c = np.concatenate((b, a, np.arange(N)))
v = np.concatenate((-fac, -fac, diag)) # negative Laplacian
# Build the negative Laplacian sparse matrix with N rows and N columns
L = csc_matrix((v, (r, c)), shape=(N, N))
# Pinning surface at reference cast by ADDING the equation
# 1 * ϕ[ref] = 0 to the ref'th equation. Note, m[ref] is the reference cast.
# If the BFS queue (m) were not sorted, then ref == 0 and m[0] would be the
# reference cast, since the BFS is initialized from the reference cast.
L[ref, ref] += 1
# Alternative pinning strategy: change the mref'th equation to be 1 * ϕ[mref] = 0.
# Then, since ϕ[mref] = 0, values of L in mref'th column are irrelevant, so set
# these to zero to maintain symmetry.
# Resulting output is same as above to many sig figs.
# L[:, ref] = 0
# L[ref, :] = 0
# L[ref, ref] = 1
# D[ref] = 0
# Solve the matrix problem, L ϕ = D
ϕ[m] = spsolve(L, D)
# Fix machine precision errors by applying pinning condition exactly
ϕ[mref] = 0.0
ϕ = ϕ.reshape(p.shape)
return ϕ
def _omega_matsolve_helper(s, t, p, edges, distratio, m, mref, eos_s_t):
"""
Compute ϵ neutrality errors and geometry on the list of edges in the graph
that are between pairs of "wet" casts. Then, prune the list of edges to
just these "wet" edges that are within the connected component containing
the reference cast (whose linear index in the full space is `mref`), and
remap all nodes (water columns) from being labelled (0, 1, ... ni*nj-1)
for the entire 2D grid, to being labelled (0, 1, ..., N-1) where N is the
number of wet casts in this connected component. The connected component
is determined by `m`, which is a list of nodes in the label set of the
full space, (0, 1, ... ni*nj-1).
"""
N = len(m)
# `remap` changes from linear indices (0, 1, ..., ni*nj-1) for the entire
# space (including land), into linear indices (0, 1, ..., N-1) for the
# current connected component
remap = np.full(p.size, -1, dtype=int)
remap[m] = np.arange(N)
# Select a subset of edges, namely those between two "wet" casts
a, b = edges
ge = (remap[a] >= 0) & (remap[b] >= 0) # good edges
a, b = a[ge], b[ge]
e = ntp_epsilon_errors(s, t, p, (a, b), eos_s_t)
if np.isscalar(distratio):
fac = np.full(e.shape, distratio)
if distratio != 1.0:
e *= fac # scale e
else:
fac = distratio[ge]
e *= fac # scale e
# Henceforth we only refer to nodes in the connected component, so remap edges now
a, b = remap[a], remap[b]
# Also remap the index for the reference cast from 2D space to 1D list of wet casts
ref = remap[mref]
return a, b, e, fac, ref