def params_to_grid(self, **targ):
"""Convert a set of parameters to grid pixel coordinates.
:param targ:
The target parameter location, as keyword arguments. The elements
of ``labels`` must be present as keywords.
:returns x:
The target parameter location in pixel coordinates.
"""
# Get bin index
inds = np.array([
np.digitize([targ[p]], bins=self.gridpoints[p], right=False) - 1
for p in self.labels
])
inds = np.squeeze(inds)
# Get fractional index.
try:
find = np.asarray([
(targ[p] - self.gridpoints[p][i]) / self.binwidths[p][i]
for i, p in zip(inds, self.labels)
])
except (IndexError):
pstring = "{0}: min={2} max={3} targ={1}\n"
s = [
pstring.format(p, targ[p], *self.gridpoints[p][[0, -1]])
for p in self.labels
]
raise ValueError("At least one parameter outside grid.\n{}".format(
' '.join(s)))
return inds + np.squeeze(find)
def antiderivative(self, xs):
"""
Computes the antiderivative of first order of this spline
"""
# Retrieve parameters
x, y, coefficients = self._x, self._y, self._coefficients
# In case of quadratic, we redefine the knots
if self.k == 2:
knots = (x[1:] + x[:-1]) / 2.0
# We add 2 artificial knots before and after
knots = np.concatenate([
np.array([x[0] - (x[1] - x[0]) / 2.0]),
knots,
np.array([x[-1] + (x[-1] - x[-2]) / 2.0]),
])
else:
knots = x
# Determine the interval that x lies in
ind = np.digitize(xs, knots) - 1
# Include the right endpoint in spline piece C[m-1]
ind = np.clip(ind, 0, len(knots) - 2)
t = xs - knots[ind]
if self.k == 1:
a = y[:-1]
b = coefficients
h = np.diff(knots)
cst = np.concatenate(
[np.zeros(1), np.cumsum(a * h + b * h**2 / 2)])
return cst[ind] + a[ind] * t + b[ind] * t**2 / 2
if self.k == 2:
h = np.diff(knots)
dt = x - knots[:-1]
b = coefficients[:-1]
b1 = coefficients[1:]
a = y - b * dt - (b1 - b) * dt**2 / (2 * h)
c = (b1 - b) / (2 * h)
cst = np.concatenate(
[np.zeros(1),
np.cumsum(a * h + b * h**2 / 2 + c * h**3 / 3)])
return cst[ind] + a[ind] * t + b[ind] * t**2 / 2 + c[ind] * t**3 / 3
if self.k == 3:
h = np.diff(knots)
c = coefficients[:-1]
c1 = coefficients[1:]
a = y[:-1]
a1 = y[1:]
b = (a1 - a) / h - (2 * c + c1) * h / 3.0
d = (c1 - c) / (3 * h)
cst = np.concatenate([
np.zeros(1),
np.cumsum(a * h + b * h**2 / 2 + c * h**3 / 3 + d * h**4 / 4),
])
return (cst[ind] + a[ind] * t + b[ind] * t**2 / 2 +
c[ind] * t**3 / 3 + d[ind] * t**4 / 4)
def digitize(module, x, bins):
"""
Calculate bin indices.
"""
if module in [np, ma]:
return np.digitize(x, bins)
elif module == torch:
return torch.bucketize(x, bins)
elif module == tf:
return tf.bucketize(x, bins)
elif module == jnp:
return jnp.digitize(x, bins)
raise UnknownModuleException(f"Module {module.__name__} not supported.")
def _compute_coeffs(self, xs):
"""Compute the spline coefficients for a given x."""
# Retrieve parameters
x, y, coefficients = self._x, self._y, self._coefficients
# In case of quadratic, we redefine the knots
if self.k == 2:
knots = (x[1:] + x[:-1]) / 2.0
# We add 2 artificial knots before and after
knots = np.concatenate(
[
np.array([x[0] - (x[1] - x[0]) / 2.0]),
knots,
np.array([x[-1] + (x[-1] - x[-2]) / 2.0]),
]
)
else:
knots = x
# Determine the interval that x lies in
ind = np.digitize(xs, knots) - 1
# Include the right endpoint in spline piece C[m-1]
ind = np.clip(ind, 0, len(knots) - 2)
t = xs - knots[ind]
h = np.diff(knots)[ind]
if self.k == 1:
a = y[ind]
result = (t, a, coefficients[ind])
if self.k == 2:
dt = (x - knots[:-1])[ind]
b = coefficients[ind]
b1 = coefficients[ind + 1]
a = y[ind] - b * dt - (b1 - b) * dt ** 2 / (2 * h)
c = (b1 - b) / (2 * h)
result = (t, a, b, c)
if self.k == 3:
c = coefficients[ind]
c1 = coefficients[ind + 1]
a = y[ind]
a1 = y[ind + 1]
b = (a1 - a) / h - (2 * c + c1) * h / 3.0
d = (c1 - c) / (3 * h)
result = (t, a, b, c, d)
return result
def ppf(cf):
x = norm.ppf((cf + 0.5) / (1 << post_prec), mean, stdd)
# Binary search is faster than using the actual gaussian cdf for the
# precisions we typically use, however the cdf is O(1) whereas search
# is O(precision), so for high precision cdf will be faster.
idxs = jnp.digitize(x, std_gaussian_bins(prior_prec)) - 1
# This loop works around an issue which is extremely rare when we use
# float64 everywhere but is common if we work with float32: due to the
# finite precision of floating point arithmetic, norm.[cdf,ppf] are not
# perfectly inverse to each other.
idxs_ = lax.while_loop(
lambda idxs: ~jnp.all((cdf(idxs) <= cf) & (cf < cdf(idxs + 1))),
lambda idxs: jnp.select(
[cf < cdf(idxs), cf >= cdf(idxs + 1)],
[idxs - 1, idxs + 1 ], idxs),
idxs)
return idxs_
def make_graph_mnist(image, patch_size, bins=(0., .3, 1.)):
"""Makes a graph object to hold an MNIST sample.
Args:
image: Should be squeezable to a 2d array
patch_size: size of patches for node features.
bins: Used for binning the pixel values. The highest bin must be greater
than the highest value in image.
Returns:
graph representing the image.
"""
return image_graph.ImageGraph.create(
# The threshold value .3 was selected to keep information
# while not introducing noise
jnp.digitize(image, bins).squeeze(),
get_start_pixel_fn=lambda _: (14, 14), # start in the center
num_colors=len(bins), # number of bins + 'out of bounds' pixel
patch_size=patch_size)
def lib_as_grid(self):
"""Convert the library parameters to pixel indices in each dimension,
and build and store a KDTree for the pixel coordinates.
"""
# Get the unique gridpoints in each param
self.gridpoints = {}
self.binwidths = {}
for p in self.labels:
self.gridpoints[p] = np.unique(self.libparams[p])
self.binwidths[p] = np.diff(self.gridpoints[p])
# Digitize the library parameters
X = np.array([
np.digitize(self.libparams[p], bins=self.gridpoints[p], right=True)
for p in self.labels
])
self.X = X.T
# Build the KDTree
startime = datetime.now()
self._kdt = KDTree(self.X, leafsize=1000) # , metric='euclidean')
print('built KDTree: {}'.format(datetime.now() - startime))
def logacia(Tarr, nus, nucia, tcia, logac):
"""interpolated function of log10(alpha_CIA)
Args:
Tarr: temperature array
nus: wavenumber array
nucia: CIA wavenumber (cm-1)
tcia: CIA temperature (K)
logac: log10 cia coefficient
Returns:
logac(Tarr, nus)
Example:
>>> nucia,tcia,ac=read_cia("../../data/CIA/H2-H2_2011.cia",nus[0]-1.0,nus[-1]+1.0)
>>> logac=jnp.array(np.log10(ac))
>>> logacia(Tarr,nus,nucia,tcia,logac)
"""
def fcia(x, i): return jnp.interp(x, tcia, logac[:, i])
vfcia = vmap(fcia, (None, 0), 0)
mfcia = vmap(vfcia, (0, None), 0)
inus = jnp.digitize(nus, nucia)
return mfcia(Tarr, inus)
def make_graph_pathfinder(
image,
patch_size,
bins,
):
"""Makes a graph holding a pathfinder image.
Args:
image: Should be squeezable to a 2d array
patch_size: size of patches for node features.
bins: Used for binning the pixel values. The highest bin must be greater
than the highest value in image.
Returns:
graph representing the image.
"""
# TODO(gnegiar): Allow multiple start nodes.
# The threshold value .3 was selected to keep information
# while not introducing noise
def _get_start_pixel_fn(image, thresh=.5 * len(bins)):
"""Detects a probable start point in a Pathfinder image example."""
thresh_image = np.where(image > thresh, 1, 0)
distance = ndi.distance_transform_edt(thresh_image)
idx = distance.argmax()
coords = np.unravel_index(idx, thresh_image.shape)
return coords
# TODO(gnegiar): Allow continuous features in models.
return image_graph.ImageGraph.create(
jnp.digitize(image, bins).squeeze(),
# Set thresh to .5 by leveraging the image discretization.
get_start_pixel_fn=functools.partial(_get_start_pixel_fn,
thresh=.5 * len(bins)),
num_colors=len(bins), # number of bins + 'out of bounds' pixel
patch_size=patch_size)
def digitize(x, bins, right=False): if isinstance(x, JaxArray): x = x.value if isinstance(bins, JaxArray): bins = bins.value return JaxArray(jnp.digitize(x, bins=bins, right=right))