# geoML - machine learning models for geospatial data
# Copyright (C) 2019 Ítalo Gomes Gonçalves
#
# This program is free software: you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR a PARTICULAR PURPOSE. See the
# GNU General Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with this program. If not, see <https://www.gnu.org/licenses/>.
# ProjectedVGP stays out on purpose: importable for old saves, unadvertised
# until its latent side (geoml.latent.fourier) earns tests.
__all__ = ["GP", "GPEnsemble", "Normalizer", "StructuralField", "GPOptions",
"VGPNetwork", "refine", "cross_validate", "conformalize",
"ConformalCalibration"]
import numpy as np
from collections.abc import Sequence
from typing import Any as _Any, Literal as _Literal, cast as _cast
import geoml._progress as _progress
import geoml._types as _types
import geoml.data as _data
import geoml.kernels as _kr
import geoml.parameter as _gpr
import geoml.latent as _latent
import geoml.likelihood as _lk
import geoml.warping as _warp
import geoml.math.tf as _tftools
import geoml.metrics as _metrics
import geoml.persistence as _persistence
import geoml.stats.random as _srandom
import geoml
import numpy as _np
import pandas as _pd
import scipy.spatial as _spatial
import tensorflow as _tf
import copy as _copy
import itertools as _iter
import os as _os
import shutil as _shutil
import tempfile as _tempfile
import time as _time
import warnings
import tensorflow_probability as _tfp
_tfd = _tfp.distributions
# How much `predict_measurements` will hold in RAM before it refuses. It
# returns the whole answer as one array per variable -- every batch kept in
# a list and concatenated at the end, so the chunks and the result are both
# live for that moment -- and the answer is `n_sim * n_nodes` values a row
# for each column: 5 KB a row at the defaults, measured. **The peak is twice
# that**: a 92 MB answer was measured to cost 182 MB of resident memory, the
# concatenate holding the parts and the whole at once. Ungoverned, a
# cross-validation fold over a few million rows asks for tens of gigabytes
# and the Linux OOM killer ends the session, which is what happened on
# 2026-09-04 (63 GB resident, the notebook's kernel killed outright).
# A fixed ceiling rather than a share of free memory: a limit that moves
# with the machine makes the same script fail in one place and not another,
# and there is no portable way to ask. Generous on purpose -- 400 000 rows
# of a scalar variable at the defaults -- since the point is to stop the
# catastrophe, not to police ordinary use.
MEASUREMENT_LIMIT = 2 * 1024 ** 3
class _ModelOptions:
def __init__(self, verbose=True, prediction_batch_size=20000,
training_batch_size=2000):
self.verbose = verbose
self.training_batch_size = training_batch_size
self.prediction_batch_size = prediction_batch_size
# Drawn, not taken: `geoml.set_seed` is the one knob, so the same call
# that fixes the initial parameters must fix training's Monte Carlo
# draws and the simulation stream, which read this number through
# stateless TensorFlow sampling. A saved model keeps the number it
# drew (`persistence` restores the options `vars` wholesale, skipping
# this constructor), so its simulations replay exactly on reload.
self.seed = int(_srandom.rng().integers(2 ** 31 - 1))
def __repr__(self):
return "%s(%s)" % (self.__class__.__name__, ", ".join(
"%s=%r" % item for item in vars(self).items()))
def batch_index(self, n_data, batch_size=None):
if batch_size is None:
batch_size = self.training_batch_size
return _data.batch_index(n_data, batch_size)
[docs]
class GPOptions(_ModelOptions):
# Class defaults, so that a model saved before these options existed still
# opens: `persistence` rebuilds options with `__new__` plus a `vars()`
# update, never calling `__init__`, so the instance dict has no entry and
# the lookup falls through to here.
jit_predict = False
qmc_simulations = False
# what a model saved before 0.9.0 was trained under; a new one takes the
# `__init__` defaults below
expert_propagation = "consensus"
propagation = "marginal"
training_tolerance = None
[docs]
def __init__(self, verbose: bool = True,
prediction_batch_size: int = 20000, jitter: float = 1e-9,
training_batch_size: int = 2000, training_samples: int = 20,
jit_predict: bool = False, qmc_simulations: bool = False,
expert_propagation: _Literal["consensus", "independent"]
= "independent",
training_tolerance: float | None = None,
propagation: _Literal["joint", "marginal"] = "joint"):
"""
Configuration of Gaussian process models.
This object can be passed on to models based on the Gaussian process in order to
control their behavior.
The `seed` that training and the simulations read is not a parameter:
it is drawn from the package generator when the object is built, so
`geoml.set_seed` before construction governs it the way it already
governs parameter initialization, and a saved model keeps the number
it drew.
Parameters
----------
verbose
Whether to show the training process on screen.
prediction_batch_size : int
Batch size for prediction/inference.
jitter : float
Small value added to covariance matrices for numerical stability.
training_batch_size : int
Number of data points per batch during training.
training_samples : int
Number of Monte Carlo samples to be drawn when training requires it.
jit_predict : bool
Whether to compile the prediction graph with XLA. Worth 3-5x on a
grid of any size, and more on a GPU, but XLA compiles per distinct
batch shape (so a grid that `prediction_batch_size` does not divide
pays that cost twice) and refuses to run anything it cannot
compile, rather than falling back. Prediction only: the same
treatment makes training both slower and unstable.
qmc_simulations : bool
Whether to draw the posterior simulations from a seeded-scramble
Sobol sequence instead of pseudo-random normals, so the same
number of them covers the predictive distribution evenly rather
than by chance. Measured on the Walker Lake model at 16-256
simulations: the ensemble mean lands 7-37x closer to the exact
posterior mean, proportions below a cut-off and the outer
quantiles about a quarter closer -- the accuracy of half again to
twice the simulations -- while the ensemble's own spread and the
correlation between locations gain nothing, and the cost is not
measurable. Deterministic given `seed`, batch-invariant either
way. The sequence covers at most 21201 dimensions (`size` times
the inducing points of a node), beyond which scipy refuses.
expert_propagation : str
How a deep network's experts see each other's inducing sets,
for every `BasicGP` in the network at once. `"independent"`
(the default since 0.9.0) lets each expert speak for its own
set alone -- O(K) in the expert count -- so duplicated points
in overlapping sets may disagree, and the data-side weighting
arbitrates. `"consensus"` (the default before, and what a model
saved before 0.9.0 keeps) predicts every expert's set from
every other and combines by precision weighting -- O(K^2).
Measured (Walker deep model and a 3000-point synthetic,
K = 5-40): training 1.6x to 6.3x faster and prediction up to 8x
as K grows; quality within a few percent of consensus and
sometimes ahead, the consensus coupling appearing to slow
optimization at large K. Only deep (multi-layer) networks are
affected: below a terminal node the propagation never runs. An
operation over GP nodes -- `Add`, `LinearCombination` -- still
makes one layer: the answer and the time are the same under
both rules until a GP node sits on top. The expected kernel
(`propagation="joint"`) takes the independent rule only.
training_tolerance : float, optional
When to stop training before its iteration count runs out, as a
fraction. The bound is smoothed, and training stops once its
gain over the last twenty iterations falls below this share of
everything gained since the call began. `None` (the default)
trains for exactly as long as it is told, which is what every
version before 0.6.5 did. 0.01 is a reasonable setting.
The count passed to `train_full`/`train_svi` remains the cap,
and the criterion only ever ends training sooner. It is
deliberately unable to stop a model that begins the call already
converged: with nothing gained, there is nothing to take a
fraction of, and the run goes to its cap. Training in phases is
the pattern this protects -- a smaller learning rate makes
progress the previous phase could not, and each phase is judged
against its own starting point.
propagation : str
How a GP node reads an input that is another node's uncertain
output. `"joint"` (the default since 0.9.0) averages the
node's kernel over its input's distribution -- the expected
kernel -- each node handing its children the covariance between
locations along with the variances, so two locations whose
inputs move together stay correlated. `"marginal"` (what a
model saved before 0.9.0 keeps) hands on each location's
variance alone and widens the range by half of it, which
understates a deep model's uncertainty by orders of magnitude
and lets training buy a small noise with it. Only a GP node
above another GP node, or above an uncertain input, is
affected. Under `"joint"` the expert propagation must be
`"independent"`, and a GP node reading an uncertain input takes
the Gaussian, exponential, Matérn or rational quadratic kernel.
"""
if expert_propagation not in ("consensus", "independent"):
raise ValueError(
"expert_propagation must be 'consensus' or 'independent', "
"got %r" % (expert_propagation,))
if propagation not in ("joint", "marginal"):
raise ValueError(
"propagation must be 'joint' or 'marginal', got %r"
% (propagation,))
if propagation == "joint" and expert_propagation == "consensus":
raise ValueError(
"the expected kernel (propagation='joint') takes each "
"expert's own chain, which is the independent rule; pass "
"expert_propagation='independent', or propagation='marginal' "
"for the consensus")
super().__init__(verbose, prediction_batch_size,
training_batch_size)
self.jitter = jitter
self.training_samples = training_samples
self.jit_predict = jit_predict
self.qmc_simulations = qmc_simulations
self.expert_propagation = expert_propagation
self.training_tolerance = training_tolerance
self.propagation = propagation
class _Convergence:
"""Has the bound stopped improving? The rule behind `training_tolerance`.
Smooths the bound with an exponential moving average and compares its
gain over the last window against the gain since training started,
stopping when the first is a small enough fraction of the second.
Two things about that ratio are deliberate. Measuring progress from the
**start of the phase** -- the calls one optimizer makes on one set of
trained parameters, which a model carries from call to call
(`VGPNetwork._training_phase`) -- keeps the phased pattern working: a
model trained, given a smaller learning rate and trained again plateaus
and then improves, and each phase is judged on its own terms, while
training split into chunks stops where one call would. And measuring
the recent gain against the total one rather than against the bound's
value is the only scale-free choice available: an ELBO's magnitude means
nothing on its own, since it grows with the number of data points and
with whatever normalization the likelihood carries.
The comparison is over a window rather than between consecutive
iterations, which matters more than it sounds. For a bound approaching
its limit with time constant `tau`, a per-iteration test fires once
`exp(-t/tau) < tolerance * tau` -- so a slowly converging fit stops
almost at once, and how early depends on `tau`, which nobody knows in
advance. Over a window of `w`, it fires at
`t > tau * log((exp(w/tau) - 1) / tolerance)`, which for `w` near `tau`
is a few time constants whatever `tau` is.
"""
# The window, in iterations (full batch) or epochs (SVI), used both as
# the average's span and as how far back the comparison reaches.
_WINDOW = 20
def __init__(self, tolerance, window=None):
self.tolerance = tolerance or 0.0
self.window = self._WINDOW if window is None else int(window)
self.start = None
self.trail = []
def stop(self, value):
"""Whether to stop, having seen `value` as the latest bound."""
if self.tolerance <= 0.0:
return False
if self.start is None:
self.start, smoothed = value, value
else:
weight = 2.0 / (self.window + 1.0)
smoothed = weight * value + (1.0 - weight) * self.trail[-1]
self.trail.append(smoothed)
return self.settled()
def settled(self):
"""Whether the bounds seen so far say stop, without adding one."""
# nothing to look back at yet, which is the burn-in: a flat stretch
# before any real progress has a small numerator and a small
# denominator, and their ratio is not evidence of anything
if self.tolerance <= 0.0 or len(self.trail) <= self.window:
return False
recent = abs(self.trail[-1] - self.trail[-1 - self.window])
total = abs(self.trail[-1] - self.start)
return recent < self.tolerance * total
class _GPModel(_gpr.Parametric):
def __init__(self, options=None):
super().__init__()
# Not a default argument: one `GPOptions()` would then be built at
# import time and shared by every model that did not bring its own,
# so setting an option on one model would set it on all of them.
self.options = GPOptions() if options is None else options
self._pre_computations = {}
self._n_dim = None
self._compiled = {}
def __str__(self):
# a model is summarized, not written as a call: `pretty_print()` is
# still there for the parameter tree
return self.__repr__()
@property
def n_dim(self):
return self._n_dim
def set_learning_rate(self, rate):
"""
Resets the model's optimizer with the provided learning rate. Will erase the optimizer's memory.
Parameters
----------
rate : float
The learning rate to use.
"""
raise NotImplementedError
def save(self, path: _types.PathLike) -> _types.PathLike:
"""Save the model's structure, parameters and training data.
The model is written as the constructor calls that built it, with
their arguments, and the parameters' values, so `open` rebuilds the
same model -- its seed included, so its simulations replay.
Parameters
----------
path
Directory to write to, overwritten if it already exists.
Returns
-------
path
The location written to.
See Also
--------
geoml.persistence.save_model
"""
import geoml.persistence as _persistence
return _persistence.save_model(self, path)
@classmethod
def open(cls, path: _types.PathLike,
data: "_data._SpatialData | None" = None) -> "_GPModel":
"""Load a model saved with :meth:`save`.
The model comes back with its variables and their types, ready to
predict or to be trained further.
Parameters
----------
path
A directory written by :meth:`save`.
data
A data object to build the model around, in place of the one it
was trained on. Used to refit the same structure on a subset,
as cross-validation does.
Returns
-------
_GPModel
The reopened model.
See Also
--------
geoml.persistence.load_model
"""
import geoml.persistence as _persistence
return _persistence.load_model(path, data=data)
[docs]
class GP(_GPModel):
"""
Basic Gaussian process model.
"""
[docs]
def __init__(self, data, variable, covariance, warping=None,
directional_data=None, interpolation=False,
use_trend=False, options=None):
"""
Basic Gaussian process model.
This model is based on the standard Gaussian process for a single output variable. It supports warping
for non-Gaussian variables and directional data as gradients of the modelled field, but not both simultaneously.
Parameters
----------
data
A `PointData` object from the ´data´ module.
variable : str
The name of the variable to be modelled. Must be a continuous variable.
covariance
The covariance function to build the covariance matrices.
warping
An object from the `warping` module. If None, the data is assumed to have zero mean and unit variance.
directional_data
A `DirectionalData` object from the ´data´ module. The corresponding variable will be used as the gradient
of the modelled field.
interpolation : bool
If `True`, will assume that the data is noiseless and try to honor the data points.
use_trend : bool
If `True`, will model a linear trend in the data in addition to the GP.
options : GPOptions
Additional configurations.
"""
super().__init__(options)
self.data = data
self.variable = variable
self.covariance = self._register(covariance)
self.covariance.set_limits(data)
if warping is None:
warping = _warp.ZScore(1)
self.warping = self._register(warping)
keep = ~ _np.isnan(self.data.variables[self.variable].measurements.values)
if _np.sum(keep) > 0:
self.warping.initialize(
self.data.variables[self.variable]
.measurements.values.to_numpy()[keep, None])
self.directional_data = directional_data
self.use_trend = use_trend
self._add_parameter("noise", _gpr.PositiveParameter(0.1, 1e-6, 10))
if interpolation:
self.parameters["noise"].set_value(1e-6)
self.parameters["noise"].fix()
self.training_log = []
self.optimizer = _tf.keras.optimizers.Adam(
_tf.keras.optimizers.schedules.ExponentialDecay(1e-2, 1, 0.999),
amsgrad=True
)
self._pre_computations.update({
"log_likelihood": _tf.Variable(_tf.constant(0.0, _tf.float64)),
})
# Everything `refresh()` computes, cleared here so that a model whose
# parameters changed cannot be predicted from a stale factorization.
self.cov: _Any = None
self.cov_chol: _Any = None
self.cov_inv: _Any = None
self.scale: _Any = None
self.alpha: _Any = None
self.x: _Any = None
self.y: _Any = None
self.x_dir: _Any = None
self.y_dir: _Any = None
self.directions: _Any = None
self.y_warped: _Any = None
self.log_derivative: _Any = None
self.trend: _Any = None
self.mat_a_inv: _Any = None
self.trend_chol = None
self.beta = None
def __repr__(self):
s = "Gaussian process model\n\n"
s += "Variable: " + self.variable + "\n\n"
s += "Kernel:\n"
s += str(self.covariance)
s += "\nWarping:\n"
s += str(self.warping)
return s
[docs]
def set_learning_rate(self, rate):
self.optimizer = _tf.keras.optimizers.Adam(
_tf.keras.optimizers.schedules.ExponentialDecay(rate, 1, 0.999),
amsgrad=True
)
[docs]
def refresh(self, jitter=1e-9):
"""
Updates the model's internal state.
If called within TensorFlow's eager mode, will allow inspection of the internal tensors.
Parameters
----------
jitter : float
Small value added to the covariance matrices for numerical stability.
"""
keep = ~ _np.isnan(
self.data.variables[self.variable].measurements.values)
with _tf.name_scope("GP_refresh"):
self.y = _tf.constant(self.data.variables[self.variable]
.measurements.values.to_numpy()[keep, None],
_tf.float64)
self.x = _tf.constant(self.data.coordinates[keep, :],
_tf.float64)
if self.directional_data is not None:
self.x_dir = _tf.constant(self.directional_data.coordinates,
_tf.float64)
self.directions = _tf.constant(self.directional_data.directions,
_tf.float64)
cov = self.covariance.self_covariance_matrix(self.x)
cov_d1 = self.covariance.covariance_matrix_d1(
self.x, self.x_dir, self.directions)
cov_d2 = self.covariance.self_covariance_matrix_d2(
self.x_dir, self.directions)
self.cov = _tf.concat([
_tf.concat([cov, cov_d1], axis=1),
_tf.concat([_tf.transpose(cov_d1), cov_d2], axis=1)
], axis=0)
self.y_dir = _tf.constant(
self.directional_data.variables[self.variable]
.measurements.values.to_numpy(),
_tf.float64
)
self.y_warped = _tf.concat([
# self.warping.forward(self.y),
self.y,
self.y_dir[:, None]
], axis=0)
self.log_derivative = _tf.constant(0.0, _tf.float64)
eye = _tf.eye(_np.sum(keep) + self.directional_data.n_data,
dtype=_tf.float64)
noise = _tf.concat([
_tf.ones([_np.sum(keep)], _tf.float64),
_tf.zeros([self.directional_data.n_data], _tf.float64)
], axis=0)
else:
self.cov = self.covariance.self_covariance_matrix(self.x)
self.y_warped, self.log_derivative = self.warping.forward(self.y)
eye = _tf.eye(_np.sum(keep), dtype=_tf.float64)
noise = _tf.ones([_np.sum(keep)], _tf.float64)
self.scale = _tf.sqrt(_tf.linalg.diag_part(self.cov))
self.cov = self.cov / self.scale[:, None] / self.scale[None, :]
noise = self.parameters["noise"].get_value() * noise
noise = noise / self.scale**2
self.cov_chol = _tf.linalg.cholesky(
self.cov + _tf.linalg.diag(noise + jitter))
self.cov_inv = _tf.linalg.cholesky_solve(self.cov_chol, eye)
self.alpha = _tf.matmul(
self.cov_inv, self.y_warped / self.scale[:, None])
if self.use_trend:
self.trend = _tf.concat([
_tf.ones([_np.sum(keep), 1], _tf.float64), self.x
], axis=1)
if self.directional_data is not None:
trend_grad = _tf.concat([
_tf.zeros([self.directional_data.n_data, 1], _tf.float64),
self.directions
], axis=1)
self.trend = _tf.concat([self.trend, trend_grad], axis=0)
self.trend = self.trend / self.scale[:, None]
mat_a = _tf.matmul(
self.trend, _tf.matmul(self.cov_inv, self.trend), True)
eye = _tf.eye(self.data.n_dim + 1, dtype=_tf.float64)
mat_a_inv = _tf.linalg.inv(mat_a + eye * jitter)
self.mat_a_inv = mat_a_inv
self.trend_chol = _tf.linalg.cholesky(mat_a_inv)
self.beta = _tf.matmul(
mat_a_inv, _tf.matmul(self.trend, self.alpha, True))
[docs]
@_tf.function
def log_likelihood(self, jitter=1e-9):
"""
Computes the model's log-likelihood with the current parameters.
Parameters
----------
jitter : float
Small value added to the covariance matrices for numerical stability.
"""
self.refresh(jitter)
with _tf.name_scope("GP_log_likelihood"):
fit = -0.5 * _tf.reduce_sum(self.y_warped * self.alpha)
det = - _tf.reduce_sum(_tf.math.log(
_tf.linalg.diag_part(self.cov_chol))) \
- _tf.reduce_sum(_tf.math.log(self.scale))
const = -0.5 * _tf.cast(_tf.shape(self.cov)[0], _tf.float64)\
* _np.log(2 * _np.pi)
log_lik = fit + det + const
# log_derivative = self.warping.log_derivative(self.y)
log_lik = log_lik + _tf.reduce_sum(self.log_derivative)
if self.use_trend:
det_2 = _tf.reduce_sum(_tf.math.log(
_tf.linalg.diag_part(self.trend_chol)))
fit_2 = _tf.reduce_sum(
_tf.matmul(self.trend_chol,
_tf.matmul(self.trend, self.alpha, True),
True)**2)
const_2 = 0.5 * _tf.constant(self.data.n_dim + 1, _tf.float64) \
* _np.log(2 * _np.pi)
log_lik = log_lik + det_2 + fit_2 + const_2
self._pre_computations["log_likelihood"].assign(log_lik)
return log_lik
[docs]
@_tf.function
def predict_raw(self, x_new, jitter=1e-9, n_sim=50):
self.refresh(jitter)
with _tf.name_scope("Prediction"):
noise = self.parameters["noise"].get_value()
# covariance
cov_new = self.covariance.covariance_matrix(x_new, self.x)
if self.directional_data is not None:
cov_new_d1 = self.covariance.covariance_matrix_d1(
x_new, self.x_dir, self.directions)
cov_new = _tf.concat([cov_new, cov_new_d1], axis=1)
cov_new = cov_new / self.scale[None, :]
# prediction
mu = _tf.matmul(cov_new, self.alpha)
point_var = self.covariance.point_variance(x_new)[:, None]
explained_var = _tf.reduce_sum(
_tf.matmul(cov_new, self.cov_inv) * cov_new,
axis=1, keepdims=True)
var = _tf.maximum(point_var - explained_var, 0.0) + noise
# trend
if self.use_trend:
trend_new = _tf.concat([
_tf.ones([_tf.shape(x_new)[0], 1], _tf.float64), x_new
], axis=1)
trend_pred = trend_new - _tf.matmul(
cov_new, _tf.matmul(self.cov_inv, self.trend))
mu = mu + _tf.matmul(trend_pred, self.beta)
trend_var = _tf.reduce_sum(
_tf.matmul(trend_pred, self.mat_a_inv) * trend_pred,
axis=1, keepdims=True
)
var = var + trend_var
# weights
weights = (explained_var / (noise + 1e-6)) ** 2
# warping
# distribution = _tfd.Normal(mu, _tf.sqrt(var))
# sims = distribution.sample(50)
rnd = _tf.random.stateless_normal([1, n_sim], seed=[0, 0], dtype=_tf.float64)
sims = rnd * _tf.sqrt(var) + mu
sims = _tf.transpose(sims)
sims = _tf.map_fn(lambda x: self.warping.backward(x[:, None]), sims)
sims = _tf.transpose(sims, [1, 2, 0])
avg_sim = _tf.reduce_mean(sims, axis=2)
out = {"mean": mu[:, 0],
"variance": var[:, 0],
"weights": _tf.squeeze(weights),
"simulations": sims[:, 0, :],
"average_sim": avg_sim[:, 0],
}
#
# def prob_fn(q):
# p = distribution.cdf(self.warping.forward(q))
# return p
#
# if quantiles is not None:
# prob = _tf.map_fn(prob_fn, quantiles)
# prob = _tf.transpose(prob)
#
# # single point case
# prob = _tf.cond(
# _tf.less(_tf.rank(prob), 2),
# lambda: _tf.expand_dims(prob, 0),
# lambda: prob)
#
# out["probabilities"] = _tf.squeeze(prob)
#
# def quant_fn(p):
# q = self.warping.backward(distribution.quantile(p))
# return q
#
# if probabilities is not None:
# quant = _tf.map_fn(quant_fn, probabilities)
# quant = _tf.transpose(quant)
#
# # single point case
# quant = _tf.cond(
# _tf.less(_tf.rank(quant), 2),
# lambda: _tf.expand_dims(quant, 0),
# lambda: quant)
#
# out["quantiles"] = _tf.squeeze(quant)
return out
[docs]
def predict(self, newdata, n_sim=50):
"""
Makes a prediction on the specified coordinates.
Parameters
----------
newdata :
A reference to a spatial points object of compatible dimension.
The object's variables will be updated.
"""
if self.data.n_dim != newdata.n_dim:
raise ValueError("dimension of newdata is incompatible with model")
if self.variable not in newdata.variables.keys():
self.data.variables[self.variable].copy_to(newdata)
prediction_input = self.data.variables[self.variable].prediction_input()
# prediction in batches
newdata.variables[self.variable].allocate_simulations(n_sim)
batch_id = self.options.batch_index(newdata.n_data,
self.options.prediction_batch_size)
n_batches = len(batch_id)
for i, batch in enumerate(batch_id):
if self.options.verbose:
print("\rProcessing batch %s of %s "
% (str(i + 1), str(n_batches)), end="")
output = self.predict_raw(
_tf.constant(newdata.coordinates[batch], _tf.float64),
jitter=self.options.jitter, n_sim=n_sim, **prediction_input)
newdata.variables[self.variable].update(batch, **output)
if self.options.verbose:
print("\n")
[docs]
def train(self, max_iter=1000):
"""
Model training.
The standard GP does not support batches of data, allways using the full data instead. This is feasible for
up to a few thousand data points.
Parameters
----------
max_iter : int
The number of iterations to train.
"""
model_variables = [pr.variable for pr in self._all_parameters
if not pr.fixed]
def loss():
return - self.log_likelihood(self.options.jitter)
for i in range(max_iter):
# self.optimizer.minimize(loss, model_variables)
_tftools.training_step(self.optimizer, loss, model_variables)
for pr in self._all_parameters:
pr.refresh()
current_log_lik = self._pre_computations["log_likelihood"].numpy()
self.training_log.append(current_log_lik)
if self.options.verbose:
print("\rIteration %s | Log-likelihood: %s" %
(str(i + 1), str(current_log_lik)), end="")
if self.options.verbose:
print("\n")
[docs]
class VGPNetwork(_GPModel):
"""Variational Gaussian process network.
A generalization of the standard Gaussian process: variables of any kind
(continuous, categorical, compositional, directional) are modelled
through a network of latent Gaussian processes, fitted by maximizing the
evidence lower bound on inducing points.
Parameters
----------
data
The training data, a container from the :mod:`geoml.data` module.
variables
The name of a variable in `data` to model, a list of names, or a
mapping from each name to its likelihood -- in which case
`likelihoods` is left out.
likelihoods
A likelihood from :mod:`geoml.likelihood`, or a list of one per
variable, in the same order as `variables`.
latent_network
The tree's leaves, from :mod:`geoml.latent`: a list of nodes, one
per likelihood and in the same order, each sized as its likelihood;
or a single node serving every likelihood, sized as their sizes
summed and split among them in order. Several leaves may share
parents, or sit on separate trees with roots of their own.
directional_data
Structural measurements, whose variable is taken as the gradient of
the modelled field.
options
Training and prediction settings.
Attributes
----------
leaves : list
The tree's output nodes, one per likelihood or one for all.
latent_network
The single leaf, where there is one. A model with several leaves
raises here and points at `leaves`.
training_log : list of float
The evidence lower bound at each iteration of the last training run.
Examples
--------
>>> import geoml
>>> geoml.set_seed(1234)
>>> walker, grid = geoml.datasets.walker()
>>> inducing = geoml.data.inducing.from_kmeans(walker, 100, seed=0)
>>> gp = geoml.latent.BasicGP(
... geoml.latent.BasicInput(inducing), size=1)
>>> model = geoml.models.VGPNetwork(
... walker, "V", geoml.likelihood.Gaussian(), gp)
>>> model.train_full(max_iter=100)
>>> model.predict(grid, n_sim=20)
"""
def __init__(self, data: "_data.PointData",
variables: "str | Sequence[str] | dict",
likelihoods: "_lk._Likelihood | Sequence[_lk._Likelihood] | None" = None,
latent_network: "_latent.network._LatentVariable | Sequence[_latent.network._LatentVariable] | None" = None,
directional_data: "_data.DirectionalData | None" = None,
options: "GPOptions | None" = None):
super().__init__(options=options)
self.data = data
# `variables={"Rock": lik, ...}` names each likelihood beside its
# variable, which is the spelling that cannot get the order wrong
if isinstance(variables, dict):
if likelihoods is not None:
raise ValueError(
"the likelihoods were given twice: once beside each "
"variable and once as `likelihoods`; give one or the "
"other")
names, likelihoods = list(variables.keys()), list(variables.values())
else:
names = [variables] if isinstance(variables, str) \
else list(variables)
if likelihoods is None:
raise ValueError(
"no likelihoods: pass one per variable, or name each "
"beside its variable as `variables={name: likelihood}`")
if isinstance(likelihoods, _lk._Likelihood):
likelihoods = [likelihoods]
likelihoods = list(likelihoods)
self.variables: list[str] = names
self.likelihoods: "list[_lk._Likelihood]" = likelihoods
if len(self.likelihoods) != len(self.variables):
raise ValueError(
"%d variable(s) but %d likelihood(s); each variable takes "
"exactly one" % (len(self.variables), len(self.likelihoods)))
self.lik_sizes = [lik.size for lik in self.likelihoods]
# The tree's leaves, and which likelihoods each one serves. A list
# is one leaf per likelihood; a single node serves them all and is
# split among them by size, which is the shape every model had before
# a list was accepted. Nothing in between: a leaf serving some of
# the likelihoods but not others would need the join this replaces.
if latent_network is None:
raise ValueError("no latent network: pass the tree's leaves")
self.leaves: "list[_latent.network._LatentVariable]"
if isinstance(latent_network, (list, tuple)):
self.leaves = list(latent_network)
if len(self.leaves) != len(self.likelihoods):
raise ValueError(
"%d leaves for %d likelihood(s); a list of leaves takes "
"one per likelihood, in the same order, or pass a single "
"node to serve them all"
% (len(self.leaves), len(self.likelihoods)))
self._leaf_groups = [[i] for i in range(len(self.likelihoods))]
else:
self.leaves = [_cast(_latent.network._LatentVariable,
latent_network)]
self._leaf_groups = [list(range(len(self.likelihoods)))]
for leaf, group in zip(self.leaves, self._leaf_groups):
self._register(leaf)
wanted = sum(self.lik_sizes[i] for i in group)
if leaf.size != wanted:
raise ValueError(
"leaf %s has size %d where its likelihood%s need%s %d"
% (leaf.name, leaf.size,
"s" if len(group) > 1 else "",
"" if len(group) > 1 else "s", wanted))
# The likelihoods are registered AFTER the tree. A save file stores
# the parameters by position in `all_parameters`, which is
# registration order, so this order is part of the format: every
# model saved since the first release put the tree's parameters
# first, and registering the likelihoods before the leaves (as the
# leaves refactor briefly did) made every older save refuse to open.
for likelihood in self.likelihoods:
self._register(likelihood)
# A leaf that is not Gaussian -- a mixture, a product, an
# exponential -- has a mean and variance the training quadrature
# would read as a Gaussian's, blurring what makes it not one; its
# likelihoods train on its realizations instead, where they can.
self._latent_gaussian = [True] * len(self.likelihoods)
for leaf, group in zip(self.leaves, self._leaf_groups):
if getattr(leaf, "gaussian", True):
continue
for i in group:
self._latent_gaussian[i] = False
if not self.likelihoods[i]._MONTE_CARLO:
warnings.warn(
"leaf %s is not Gaussian, but %s can only train on "
"its mean and variance; it will be trained as if "
"it were Gaussian"
% (leaf.name, type(self.likelihoods[i]).__name__),
stacklevel=2)
# the cached refresh trace lives on the model, there being no single
# node to hang it on once there are several leaves
self._refresh_graph = None
if getattr(self.options, "propagation", "marginal") == "joint":
self._check_expected_kernel()
self.var_lengths = [data.variables[v].length for v in self.variables]
self.y, self.has_value = self._stacked_measurements(data)
# a variable rather than a number: the traced bound scales a batch
# by it, and `_set_data` changes it in place so the trace stays valid
self.total_data = _tf.Variable(
float(_np.sum(self.has_value)), dtype=_tf.float64, trainable=False)
# the one traced training step this model reuses -- see
# `_training_step`
self._step = None
# what the next training call goes on with -- see `_training_phase`
self._phase = None
# initializing likelihoods -- declustered where the data carries
# the column `container.decluster()` keeps, so the warpings start
# on the field's distribution rather than the sampling's
stored = data.metadata.get("declustering")
stored = None if stored is None \
else _np.asarray(stored.values, dtype=float).ravel()
for i, v in enumerate(self.variables):
y, has_value = data.variables[v].get_measurements()
has_value = _np.all(has_value == 1.0, axis=1)
y = y[has_value, :]
self.likelihoods[i].initialize(
y, weights=None if stored is None else stored[has_value])
# directions
self.directional_likelihood = _lk.GradientIndicator()
# self.directional_likelihood = _lk.Gaussian()
# self.directional_likelihood.parameters["noise"].set_value(1e-6)
self.directional_data = directional_data
self.total_data_dir = 0
self.y_dir = None
self.has_value_dir = None
self.var_lengths_dir = None
if directional_data is not None:
if self.data.n_dim != directional_data.n_dim:
raise ValueError("the directional data must have the"
"same number of dimensions as the"
"point data")
self.var_lengths_dir = [1] * sum(self.var_lengths)
y_dir, has_value_dir = [], []
for v, s in zip(variables, self.var_lengths):
if v in directional_data.variables.keys():
y_v, h_v = directional_data.variables[v].get_measurements()
y_dir.append(_np.tile(y_v, [1, s]))
has_value_dir.append(_np.tile(h_v, [1, s]))
else:
y_dir.append(_np.zeros([directional_data.n_data, s]))
has_value_dir.append(_np.ones([directional_data.n_data, s]))
y_dir = _np.concatenate(y_dir, axis=1)
has_value_dir = _np.concatenate(has_value_dir, axis=1)
self.y_dir = y_dir.copy()
self.has_value_dir = has_value_dir.copy()
self.total_data_dir = _np.sum(self.has_value_dir)
# optimizer
self.training_log = []
self.optimizer = _tf.keras.optimizers.Adam(
_tf.keras.optimizers.schedules.ExponentialDecay(1e-2, 1, 0.999),
amsgrad=True
)
# intermediate tensors
self.elbo = _tf.Variable(_tf.constant(0.0, _tf.float64))
self.kl_div = _tf.Variable(_tf.constant(0.0, _tf.float64))
def __repr__(self):
s = "Variational Gaussian process model\n\n"
s += "Variables:\n "
for v, lik in zip(self.variables, self.likelihoods):
s += "\t" + v + " (" + lik.__class__.__name__ + ")\n"
s += "\nLatent layer:\n"
if len(self.leaves) == 1:
s += str(self.leaves[0])
else:
# one leaf per likelihood, each written under the variable it
# serves; a parent two leaves share appears under both, which is
# what the tree looks like from either leaf
for v, leaf in zip(self.variables, self.leaves):
s += "[%s]\n%s\n" % (v, str(leaf))
return s
[docs]
def to_dot(self, legend=True, rankdir="BT"):
"""
Writes the model as a Graphviz diagram: coordinates, latent network,
warpings and output variables.
See `geoml.viz.graphviz.to_dot`.
"""
# imported here because that module reads the modules this one needs
import geoml.viz.graphviz as _gv
return _gv.to_dot(self, legend=legend, rankdir=rankdir)
[docs]
def set_learning_rate(self, rate):
self.optimizer = _tf.keras.optimizers.Adam(
_tf.keras.optimizers.schedules.ExponentialDecay(rate, 1, 0.999),
amsgrad=True
)
@property
def latent_network(self):
"""The tree's single leaf, where there is one.
Every model built before a list of leaves was accepted has exactly
one, and everything that read it keeps working. A model with several
leaves has no single output node to hand back, and says so rather
than returning sometimes a node and sometimes a list.
"""
if len(self.leaves) == 1:
return self.leaves[0]
raise AttributeError(
"this model has %d leaves, one per likelihood; read `leaves`"
% len(self.leaves))
def _nodes(self):
"""Every node of the tree, each once, leaves included.
The union over the leaves' ancestors: a parent shared by two leaves
is one node, and must be refreshed, priced in the KL and reset by
cross-validation exactly once. Order is fixed by first sighting so
that anything zipped against it lines up from one call to the next.
"""
seen, nodes = set(), []
for leaf in self.leaves:
for node in [leaf] + leaf.get_unique_parents():
if id(node) not in seen:
seen.add(id(node))
nodes.append(node)
return nodes
def _propagation(self):
"""The context every training and prediction call runs under: how
the experts propagate, and whether uncertainty travels with its
covariance between locations (`GPOptions.propagation`). The network
is checked against the expected kernel's refusals each time, since
the options can be changed on a live model."""
joint = getattr(self.options, "propagation", "marginal") == "joint"
if joint:
self._check_expected_kernel()
return _latent.propagation_rule(self.options.expert_propagation,
joint=joint)
def _propagation_key(self):
"""What the propagation adds to a cached trace's key."""
return (self.options.expert_propagation,
getattr(self.options, "propagation", "marginal"))
def _check_expected_kernel(self):
"""Refuses, under `GPOptions(propagation="joint")`, what the expected
kernel cannot read: the consensus rule; a GP node on an uncertain
input whose kernel is no scale mixture of Gaussians; `Cosine` in any
GP node; and the nodes with no expected-kernel path where a random
input or a GP node above would need one."""
network = _latent.network
if self.options.expert_propagation != "independent":
raise ValueError(
"the expected kernel (propagation='joint') takes each "
"expert's own chain, which is the independent rule; set "
"expert_propagation='independent', or propagation='marginal' "
"for the consensus")
random = (network._GPNode, network.GPWalk,
network.GradientConstrainedInput, network.GaussianInput)
def uncertain(node):
return any(isinstance(p, random)
for p in node.get_unique_parents())
for node in self._nodes():
name = "%s (%s)" % (node.name, type(node).__name__)
if isinstance(node, network._GPNode):
kernel = node.kernel
if isinstance(kernel, _kr.Cosine):
raise ValueError(
"%s: a Cosine kernel is a covariance along one axis "
"only, and is refused in a network under the "
"expected kernel" % name)
if isinstance(node, network.UncertainInputGP):
if any(isinstance(p, random[:3])
for p in node.get_unique_parents()) \
or network._feeds_gp(node):
raise ValueError(
"%s integrates over its own input's marginal; "
"under the expected kernel it reads an input or "
"a GaussianInput and no GP node reads it -- a "
"BasicGP takes the rest" % name)
elif uncertain(node) \
and not network._expected_kernel_supported(kernel):
raise ValueError(
"%s reads an uncertain input, and its %s kernel is "
"no scale mixture of Gaussians, which the expected "
"kernel needs: take Gaussian, Exponential, Matern32, "
"Matern52 or RationalQuadratic"
% (name, type(kernel).__name__))
elif isinstance(node, network.GPWalk) and not \
network._expected_kernel_supported(node.field.kernel):
# the walk reads its field where the walkers have been
# carried, which is uncertain whatever the field's input
raise ValueError(
"%s reads its field at uncertain positions, and the "
"field's %s kernel is no scale mixture of Gaussians, "
"which the expected kernel needs: take Gaussian, "
"Exponential, Matern32, Matern52 or RationalQuadratic"
% (name, type(node.field.kernel).__name__))
elif isinstance(node, network.RadialTrend) and uncertain(node):
raise ValueError(
"%s drops its input's variance, and under the expected "
"kernel takes a certain input only" % name)
elif isinstance(node, network.GradientConstrainedInput) \
and network._feeds_gp(node):
raise ValueError(
"%s is a random input whose covariance between "
"locations the expected kernel does not carry; a GP "
"node may not read it under propagation='joint'" % name)
def _refresh(self, jitter):
for leaf in self.leaves:
leaf.refresh(jitter)
def _jitter_by_likelihood(self, per_leaf, like):
"""`_by_likelihood` for the leaves' latent jitters, `[n, size]` or
None each: zeros shaped as `like` for a leaf with none, and None for
every likelihood where no leaf has any -- the path a model with
certain inputs has always taken."""
if all(j is None for j in per_leaf):
return [None] * len(self.likelihoods)
return self._by_likelihood(
[_tf.zeros_like(l) if j is None else j
for j, l in zip(per_leaf, like)])
def _by_likelihood(self, per_leaf, axis=1):
"""Per-likelihood tensors from per-leaf ones, in likelihood order.
A leaf serving one likelihood hands its tensor straight over; one
serving several is split among them by size along `axis` -- the
split every model used to make of its single node, now made only
where a leaf actually serves more than one.
"""
out: "list[_Any]" = [None] * len(self.likelihoods)
for tensor, group in zip(per_leaf, self._leaf_groups):
if len(group) == 1:
out[group[0]] = tensor
continue
sizes = [self.lik_sizes[i] for i in group]
for i, piece in zip(group, _tf.split(tensor, sizes, axis=axis)):
out[i] = piece
return out
@_tf.function
def _training_elbo(self, x, y, has_value, training_inputs,
x_dir=None, directions=None, y_dir=None,
has_value_directions=None, x_var=None,
samples=20, seed=0, jitter=1e-6):
self._refresh(jitter)
# ELBO
elbo = self._log_lik(x, y, has_value, training_inputs,
x_var=x_var, samples=samples, seed=seed)
# ELBO for directions
if x_dir is not None:
elbo = elbo + self._log_lik_directions(
x_dir, directions, y_dir, has_value_directions)
# KL-divergence, over every node once: a parent two leaves share
# is one distribution and pays one price
kl = _tf.add_n([node.kl_divergence() for node in self._nodes()])
# The MAP term: point-estimated parameters that declare a prior pay
# its log-density here, making the objective a bound on
# `log p(y, theta)`. Zero unless something declares one, so a model
# without priors trains on exactly the objective it always did.
log_prior = self.log_prior()
self.elbo.assign(elbo - kl + log_prior)
self.kl_div.assign(kl)
return elbo - kl + log_prior
@_tf.function
def _log_lik(self, x, y, has_value, training_inputs, x_var=None,
samples=20, seed=0):
with _tf.name_scope("batched_elbo"):
elbo = self._data_log_lik(x, y, has_value, training_inputs,
x_var=x_var, samples=samples, seed=seed)
# batch weight
batch_size = _tf.reduce_sum(has_value)
elbo = elbo * self.total_data / batch_size
return elbo
def _data_log_lik(self, x, y, has_value, training_inputs, x_var=None,
samples=20, seed=0):
"""The expected log-likelihood summed over the rows given, unscaled;
`_log_lik` scales it to the whole data set. Not traced on its own,
so a caller tracing it under an expert subset gets that subset."""
with _tf.name_scope("data_log_lik"):
# prediction, one leaf at a time, then handed to the likelihoods
# each leaf serves
mus, vars_, simss = [], [], []
for leaf in self.leaves:
mu, var, sims, _ = leaf.predict(
x, x_var=x_var, n_sim=samples, seed=[seed, 0])
mus.append(_tf.transpose(mu[:, :, 0]))
vars_.append(_tf.transpose(var))
simss.append(_tf.transpose(sims, [1, 0, 2]))
# likelihood
y_s = _tf.split(y, self.var_lengths, axis=1)
mu = self._by_likelihood(mus)
var = self._by_likelihood(vars_)
hv = _tf.split(has_value, self.var_lengths, axis=1)
sims = self._by_likelihood(simss)
elbo = _tf.constant(0.0, _tf.float64)
for likelihood, mu_i, var_i, y_i, hv_i, sim_i, inp, gaussian \
in zip(self.likelihoods, mu, var, y_s, hv, sims,
training_inputs, self._latent_gaussian):
if not gaussian and likelihood._MONTE_CARLO:
inp = dict(inp, latent_gaussian=False)
elbo = elbo + likelihood.log_lik(
mu_i, var_i, y_i, hv_i, samples=sim_i, **inp)
return elbo
@_tf.function
def _log_lik_directions(self, x_dir, directions, y_dir, has_value):
with _tf.name_scope("batched_elbo_directions"):
# prediction, per leaf; the directional likelihood is one column
# at a time, so the leaves' columns are joined and split by
# column rather than by likelihood
mus, vars_ = [], []
for leaf in self.leaves:
mu, var, _ = leaf.predict_directions(x_dir, directions)
mus.append(_tf.transpose(mu[:, :, 0]))
vars_.append(_tf.transpose(var))
mu = _tf.concat(mus, axis=1) if len(mus) > 1 else mus[0]
var = _tf.concat(vars_, axis=1) if len(vars_) > 1 else vars_[0]
# likelihood
y_s = _tf.split(y_dir, self.var_lengths_dir, axis=1)
mu = _tf.split(mu, self.var_lengths_dir, axis=1)
var = _tf.split(var, self.var_lengths_dir, axis=1)
hv = _tf.split(has_value, self.var_lengths_dir, axis=1)
elbo = _tf.constant(0.0, _tf.float64)
for mu_i, var_i, y_i, hv_i in zip(mu, var, y_s, hv):
elbo = elbo + self.directional_likelihood.log_lik(
mu_i, var_i, y_i, hv_i)
# batch weight
batch_size = _tf.cast(_tf.shape(x_dir)[0], _tf.float64)
elbo = elbo * self.total_data_dir / batch_size
return elbo
def _stacked_measurements(self, data):
"""The measurements of every modelled variable side by side, and
the mask of which cells hold one."""
y, has_value = [], []
for v in self.variables:
y_v, h_v = data.variables[v].get_measurements()
y.append(y_v)
has_value.append(h_v)
return _np.concatenate(y, axis=1), _np.concatenate(has_value, axis=1)
def _set_data(self, data):
"""Point the model at another container of the same variables.
For a workflow that trains one model on several subsets of its data
-- cross-validation swaps the training rows in fold by fold --
without rebuilding it, so the traced training step, the refresh and
the prediction graphs are all reused. Replaces exactly what the
constructor derived from the data: the container, the stacked
measurements and their mask, and the count the minibatch bound
scales by. The likelihoods are not re-initialized: the warpings
keep their state, as a reloaded model keeps it.
"""
if data.n_dim != self.data.n_dim:
raise ValueError(
"the new data has %d dimensions where the model's has %d"
% (data.n_dim, self.data.n_dim))
lengths = [data.variables[v].length for v in self.variables]
if lengths != self.var_lengths:
raise ValueError(
"the new data's variables have lengths %s where the model's "
"have %s" % (lengths, self.var_lengths))
self.data = data
self.y, self.has_value = self._stacked_measurements(data)
self.total_data.assign(float(_np.sum(self.has_value)))
# its active sets were formed on the old rows
self._by_expert = None
def _reset_optimizer(self):
"""Zero the optimizer's memory in place -- its moment estimates and
the step count its learning-rate schedule reads -- leaving the
optimizer object, and so the traced step that captured it, where
they are. `set_learning_rate` replaces the object instead, which a
new rate needs and which costs a retrace. Either way the next
training call starts a phase of its own (`_training_phase`)."""
for variable in self.optimizer.variables:
variable.assign(_tf.zeros(variable.shape, variable.dtype))
self._phase = None
self._by_expert = None
def _training_phase(self, variables, kind):
"""The stopping rule and the minibatch order a training call goes on
with.
A phase is the run of `kind` calls ("full" or "svi") one optimizer
makes on one set of trained parameters. Training in chunks -- how a
caller reports progress or honours a cancel between them -- then
takes the steps one call would and stops where it would: the
optimizer's state and the seed carry over by themselves, and the
rule's trail and the batch order are carried here. A new optimizer
(`set_learning_rate`), `_reset_optimizer`, other trained parameters
or the other trainer start a new phase, judged from its own start.
"""
if self._phase is not None:
cached, optimizer, cached_kind, converged, rng = self._phase
if optimizer is self.optimizer and cached_kind == kind \
and len(cached) == len(variables) \
and all(a is b for a, b in zip(cached, variables)):
# read afresh, so that a lowered tolerance trains on
converged.tolerance = self.options.training_tolerance or 0.0
return converged, rng
converged = _Convergence(self.options.training_tolerance)
# a generator of its own, so the batch order is reproducible from
# options.seed without reaching any draw made outside training
rng = _np.random.default_rng(self.options.seed)
self._phase = (tuple(variables), self.optimizer, kind, converged, rng)
return converged, rng
@property
def converged(self) -> bool:
"""Whether the bound has settled in the current phase of training.
True once `options.training_tolerance` has stopped a call, and until
the optimizer is replaced or reset or other parameters are trained.
A training call made while it holds returns without a step.
"""
if self._phase is None:
return False
cached, optimizer, _, converged, _ = self._phase
variables = self.get_unfixed_variables()
return optimizer is self.optimizer and len(cached) == len(variables) \
and all(a is b for a, b in zip(cached, variables)) \
and converged.settled()
def _training_step(self, variables, training_inputs):
"""
One traced training step: the ELBO, its gradient and the update.
Keeping all three inside a single `tf.function` matters more than it
looks. Run eagerly — as it is when the tape and `apply_gradients` sit
outside a graph — Adam issues its update variable by variable, at a
cost of roughly 1.5 ms each whatever the variable's size. A GP node
holds three parameters *per expert*, so a network of 40 experts spends
more time in the optimizer than in the model itself: 278 ms a step
against 98 ms traced.
Built once per model and kept on it, so `train_full` and `train_svi`
calls -- and the folds of a cross-validation, which swap the data in
under one model -- share a trace. A trace is never returned to the
process: TensorFlow's graph machinery for a step with gradients
stays resident after the function dies (its concrete functions,
their gradient rewrites and the optimizer's branch graphs hold
each other in cycles the garbage collector cannot break), measured
at 200 MB a step trace and 1.5-2.6 GB a model on the Tom v6 model,
so a new one per call is a leak. The step is rebuilt only when the
variables or the optimizer object change (`set_learning_rate`
replaces the optimizer, and a new rate is baked into the trace), or
when a variable hands the bound a payload -- `RockTypeVariable`'s
boundary column rides the trace as a constant, so it cannot serve
another data set. `reduce_retracing` lets the last, shorter batch
of an epoch and the folds' differing row counts share one relaxed
trace instead of a static trace each.
"""
payload = any(len(inp) > 0 for inp in training_inputs)
if self._step is not None and not payload:
cached_variables, optimizer, rule, step = self._step
if optimizer is self.optimizer and rule == self._propagation_key() \
and len(cached_variables) == len(variables) \
and all(a is b for a, b in zip(cached_variables, variables)):
return step
directions = {}
if self.directional_data is not None:
directions = dict(
x_dir=_tf.constant(
self.directional_data.coordinates, _tf.float64),
directions=_tf.constant(
self.directional_data.directions, _tf.float64),
y_dir=_tf.constant(self.y_dir, _tf.float64),
has_value_directions=_tf.constant(
self.has_value_dir, _tf.float64))
@_tf.function(reduce_retracing=True)
def step(x, y, has_value, x_var):
with _tf.GradientTape() as tape:
loss = - self._training_elbo(
x, y, has_value, training_inputs, x_var=x_var,
samples=self.options.training_samples,
jitter=self.options.jitter,
seed=self.options.seed,
**directions)
self.optimizer.apply_gradients(
zip(tape.gradient(loss, variables), variables))
if not payload:
self._step = (tuple(variables), self.optimizer,
self._propagation_key(), step)
return step
[docs]
def train_full(self, max_iter: int = 1000) -> None:
"""Train on the whole data set at every iteration.
Feasible while the data and the latent network fit in memory
together; past that, use :meth:`train_svi`. The evidence lower bound
of each iteration is appended to `training_log`.
Consecutive calls with one optimizer on the same trained parameters
continue each other: `train_full(100)` twice takes the steps of
`train_full(200)` and stops where it would. `set_learning_rate`
starts afresh.
Parameters
----------
max_iter
Number of iterations, and a cap rather than a count when
`options.training_tolerance` asks training to stop once the
bound settles; a call made once it has settled (`converged`)
takes none.
See Also
--------
train_svi : minibatch training, for larger data sets.
"""
training_inputs = [self.data.variables[v].training_input()
for v in self.variables]
model_variables = self.get_unfixed_variables()
step = self._training_step(model_variables, training_inputs)
converged, _ = self._training_phase(model_variables, "full")
if converged.settled():
if self.options.verbose:
print("The bound has settled; nothing to train.")
return
# the whole data set every iteration, so it is converted once
x = _tf.constant(self.data.coordinates, _tf.float64)
y = _tf.constant(self.y, _tf.float64)
has_value = _tf.constant(self.has_value, _tf.float64)
x_var = _tf.constant(self.data.get_batched_variance()[0], _tf.float64)
# the propagation rule is read when the step traces (and re-traces),
# which happens inside the loop
with self._propagation(), \
_progress.reporting("train", max_iter, "iteration") as report:
for i in range(max_iter):
step(x, y, has_value, x_var)
for pr in self._all_parameters:
pr.refresh()
current_elbo = self.elbo.numpy()
self.training_log.append(current_elbo)
if self.options.verbose:
print("\rIteration %s | ELBO: %s" %
(str(i+1), str(current_elbo)), end="")
# after the log and the refresh, so a callback that cancels
# leaves the model at a completed iteration
report(i + 1, bound=float(current_elbo))
if converged.stop(current_elbo):
if self.options.verbose:
print("\nStopped at iteration %s: the bound has "
"settled" % str(i + 1), end="")
break
if self.options.verbose:
print("\n")
[docs]
def train_svi(self, epochs: int = 100) -> None:
"""Train in minibatches, by stochastic variational inference.
Each epoch visits the data once, in batches of
`options.training_batch_size`, drawn in an order reproducible from
the model's seed. The bound of each batch is appended to
`training_log`.
Consecutive calls with one optimizer on the same trained parameters
continue each other, the batch order included: `train_svi(5)` twice
takes the steps of `train_svi(10)` and stops where it would.
`set_learning_rate` starts afresh.
Parameters
----------
epochs
Number of passes over the data, and a cap rather than a count
when `options.training_tolerance` asks training to stop once the
bound settles; a call made once it has settled (`converged`)
takes none. The criterion reads one value an epoch, the mean
over its batches.
See Also
--------
train_full : one gradient step per iteration on all the data.
"""
model_variables = self.get_unfixed_variables()
# The variables' own training inputs are not indexed by batch here, as
# they are in `train_full` -- see the commented-out attempt below.
step = self._training_step(
model_variables, [{} for _ in self.variables])
# judged once an epoch, on the mean over its batches: a single
# batch's bound is an estimate with noise of its own, and smoothing
# that would only measure how the batches were drawn
converged, rng = self._training_phase(model_variables, "svi")
if converged.settled():
if self.options.verbose:
print("The bound has settled; nothing to train.")
return
# an epoch is many gradient steps, and a caller watching a long run
# wants to hear from it oftener than once a pass; the batch count is
# the same every epoch, so the total is known before the first
n_batches = len(self.options.batch_index(self.data.n_data))
done = 0
with self._propagation(), \
_progress.reporting(
"train", epochs * n_batches, "batch") as report:
for i in range(epochs):
current_elbo = []
shuffled = rng.choice(
self.data.n_data, self.data.n_data, replace=False)
batches = self.options.batch_index(self.data.n_data)
for batch in batches:
# training_inputs = [
# self.data.variables[v].training_input(idx)
# for v in self.variables]
idx = shuffled[batch]
step(_tf.constant(self.data.coordinates[idx], _tf.float64),
_tf.constant(self.y[idx], _tf.float64),
_tf.constant(self.has_value[idx], _tf.float64),
_tf.constant(self.data.get_batched_variance(idx)[0],
_tf.float64))
for pr in self._all_parameters:
pr.refresh()
current_elbo.append(self.elbo.numpy())
self.training_log.append(current_elbo[-1])
done += 1
report(done, bound=float(current_elbo[-1]))
total_elbo = _np.mean(current_elbo)
if self.options.verbose:
print("\rEpoch %s | ELBO: %s" %
(str(i + 1), str(total_elbo)), end="")
if converged.stop(total_elbo):
if self.options.verbose:
print("\nStopped at epoch %s: the bound has settled"
% str(i + 1), end="")
break
if self.options.verbose:
print("\n")
[docs]
def predict_raw(self, *args, **kwargs):
"""`_predict_raw` in a graph, XLA-compiled when the options ask for it.
`jit_compile` is settled when a `tf.function` is built, and the
simulation draws are baked into the graph, so honouring either option
means holding one function per combination of settings rather than a
flag on a single one. Each is traced at most once per model, and
`None` (rather than `False`) leaves the uncompiled path exactly as it
was. The `simulation_rule`/`propagation_rule` contexts wrap the call
rather than the trace because a retrace (a new batch shape, a new
`n_sim`) can happen on any call, and has to see the flags the cache
key promised. The propagation rule joins the key for the networks
whose graphs refresh internally (`ProjectedVGP`); on this class the
graph reads snapshotted state, and the extra key is merely unused.
"""
jit = bool(self.options.jit_predict)
qmc = bool(self.options.qmc_simulations)
rule = self._propagation_key()
# an expert subset is read at trace time too; `None` (every expert)
# keeps the key every other caller has always used
subset = _latent.network._subset_key()
slots = _latent.network._slots_key()
key = (jit, qmc, rule) if subset is None and slots is None \
else (jit, qmc, rule, subset, slots)
traced = self._compiled.get(key)
if traced is None:
# `reduce_retracing` relaxes the batch shape, which is the one
# thing that varies without meaning anything: a grid divides
# into equal batches and a short last one, and every container
# of a different size starts the count again. Measured on eight
# predictions of differing size: eight traces and 7.6 s become
# two and 2.2 s, bit-identical, and XLA agrees either way. What
# remains keyed on the value -- `n_sim`, `include_noise` -- is
# baked into the graph and has to retrace.
traced = _tf.function(self._predict_raw, jit_compile=jit or None,
reduce_retracing=True)
self._compiled[key] = traced
with _latent.simulation_rule(qmc), self._propagation():
return traced(*args, **kwargs)
def _predict_raw(self, x_new, variable_inputs, x_var=None,
n_sim=1, seed=0, include_noise=True, n_splits=None):
# The posterior is refreshed once by `predict` and snapshotted into
# Variables; this cached graph reads that state, so it is not recomputed
# per batch.
with _tf.name_scope("Prediction"):
mus, vars_, sims, exp_vars, jitters = [], [], [], [], []
for leaf in self.leaves:
predicted = leaf.predict(
x_new, x_var=x_var, n_sim=n_sim, seed=[seed, 0])
mu, var, sim, exp_var = predicted
mus.append(_tf.transpose(mu[:, :, 0]))
vars_.append(_tf.transpose(var))
sims.append(_tf.transpose(sim, [1, 0, 2]))
exp_vars.append(_tf.transpose(exp_var))
jitters.append(None if predicted.jitter is None
else _tf.transpose(predicted.jitter))
pred_mu = self._by_likelihood(mus)
pred_var = self._by_likelihood(vars_)
pred_sim = self._by_likelihood(sims)
pred_exp_var = self._by_likelihood(exp_vars)
# what the realizations leave out at an uncertain input, which
# a likelihood integrates beside its noise; a model with
# certain inputs passes nothing, as it always has
pred_jitter = self._jitter_by_likelihood(jitters, exp_vars)
output = []
for mu, var, sim, exp_var, jitter, lik, v_inp in zip(
pred_mu, pred_var, pred_sim, pred_exp_var, pred_jitter,
self.likelihoods, variable_inputs):
extra = {} if jitter is None else {"jitter": jitter}
output.append(
lik.predict(
mu, var, sim, exp_var,
include_noise=include_noise,
n_splits=n_splits, **v_inp, **extra
)
)
return output
[docs]
def predict(self, newdata: "_data._SpatialData",
n_sim: "int | None" = None,
include_noise: bool = True,
where: _types.Where = None) -> None:
"""Predict at new locations, writing the answer into the container.
The variables the model was trained on are created on `newdata` if
absent and filled in place: prediction, latent moments, simulations,
and the columns each variable kind adds to those.
Parameters
----------
newdata
The locations to predict, of the same dimension as the training
data. Modified in place.
n_sim
Number of realizations to draw per location. `None` takes the
number `newdata` already holds, or 20 where it holds none.
include_noise
Whether to integrate the likelihood noise out of the answer. The
prediction then reports the value the ground would show once
measurement error and sub-resolution variability are averaged
over -- a deterministic correction, and a large one on a skewed
variable. Turn it off to see the latent field alone.
where
One boolean per location, the indices of the locations to visit,
the name of a boolean metadata column holding the same, or `None`
for all of them. Locations left out keep whatever they hold,
including their simulations, and a location never visited stays
missing.
Raises
------
ValueError
If `where` names some locations of a variable that already holds
simulations, and `n_sim` asks for a different number of them.
Notes
-----
A location's simulated values do not depend on what else is in its
batch, so predicting a subset gives the same answer as predicting
everything and reading that subset back.
Where `newdata` holds measurements -- the training data, a
validation set -- two things are written beside the prediction as
metadata, one column per variable component, for the figures that
compare a model with its data to read without the model:
`pit_<variable>[_<component>]`, where each measurement falls in the
predictive distribution of a measurement (0 to 1; uniform on data
the model did not see, if it is well calibrated), and
`warped_<variable>_<i>`, the measurements through the likelihood's
warping, as the model sees them. A block model, whose locations are
not points, gets neither.
"""
self._predict(newdata, n_sim, include_noise, where)
[docs]
def predict_node(self, node: "_latent.network._LatentVariable",
newdata: "_data._SpatialData",
n_sim: "int | None" = None,
name: "str | None" = None,
labels: "Sequence[str] | None" = None,
where: _types.Where = None) -> "_data.LatentVariable":
"""Predict what a node inside the tree says, into a container.
The same prediction :meth:`predict` makes, stopped at `node` instead
of carried on to the likelihoods: the node's mean and variance at
every location, and its realizations, written as a
:class:`~geoml.data.LatentVariable` with one part per output of the
node. Where a `GPWalk` moved the coordinates, what a shared parent
says before two leaves diverge, what a `Linear` trend adds.
Realization `s` of the node is the one realization `s` of the first
leaf above it was built from, wherever operation nodes (`Add`,
`Linear`, `Scale`, ...) are all that stand between them; a GP node
above reads the node's mean and variance, never its realizations.
Parameters
----------
node
A node of this model's tree, above its input.
newdata
The locations to predict, of the same dimension as the training
data: points, a grid, a section or a mesh. Modified in place.
n_sim
Number of realizations per location. `None` takes the number the
variable already holds, or 20.
name
The variable to write, the node's own name by default. A
`LatentVariable` of that name is written into; any other
variable of that name is refused.
labels
One name per output of the node, `"0"`, `"1"`, ... by default.
where
One boolean per location, the indices of the locations to visit,
the name of a boolean metadata column holding the same, or `None`
for all of them. Locations left out keep what they hold.
Returns
-------
LatentVariable
The variable written, as held by `newdata`.
Raises
------
ValueError
If `node` is an input or not in this model's tree, if `newdata`
is a block model, if `labels` is not one per output, if `name`
is held by another kind of variable or one with other labels,
or if `where` names some locations of a variable holding a
different number of realizations.
See Also
--------
predict : the prediction carried on to the likelihoods.
Notes
-----
Point support only. A block's value is the mean over its sub-blocks
taken by the likelihood, and a node has none; predict on a
`Grid3D` over the same ground instead.
"""
if not any(n is node for n in self._nodes()):
raise ValueError("%s is not a node of this model's tree"
% getattr(node, "name", node))
if isinstance(node, _latent.network._RootLatentVariable):
raise ValueError(
"%s is an input, whose output is the transformed coordinates; "
"apply its transform to the coordinates instead" % node.name)
if isinstance(newdata, (_data.Blocks1D, _data.Blocks2D,
_data.Blocks3D, _data.RotatedBlocks3D,
_data.BlockSet3D)):
raise ValueError(
"a node predicts at points, and %s is a block model: a "
"block's value is the mean over its sub-blocks, which a "
"likelihood takes and a node has not; predict on a Grid3D "
"over the same ground instead" % type(newdata).__name__)
if self.data.n_dim != newdata.n_dim:
raise ValueError("dimension of newdata is incompatible with model")
called = str(node.name) if name is None else name
labels = [str(i) for i in range(node.size)] if labels is None \
else [str(label) for label in labels]
if len(labels) != node.size:
raise ValueError("%s has %d output(s) and %d label(s) were given"
% (node.name, node.size, len(labels)))
variable = newdata.variables.get(called)
if variable is not None and (
not isinstance(variable, _data.LatentVariable)
or variable.labels != labels):
raise ValueError(
"%r is already a %s with labels %s; name the node's variable "
"otherwise" % (called, type(variable).__name__,
getattr(variable, "labels", None)))
where = _where_mask(newdata, where)
n_sim = _simulation_count(newdata, [called], n_sim,
keep=where is not None)
if variable is None:
variable = _data.LatentVariable(called, newdata, labels)
newdata.variables[called] = variable
variable.allocate_simulations(n_sim)
elif where is None or _stored_n_sim(variable) is None:
variable.allocate_simulations(n_sim)
seed = [self.options.seed + self._seed_offset(node), 0]
traced = self._node_function(node)
def call(x, x_var, n_splits):
with _latent.simulation_rule(bool(self.options.qmc_simulations)), \
self._propagation():
return traced(x, x_var, n_sim, seed)
for batch, (mean, variance, sims) in self._over_batches(
newdata, call, where):
variable.update(batch, latent_mean=mean.numpy(),
latent_variance=variance.numpy(),
simulations=sims.numpy())
return variable
def _node_function(self, node):
"""`node.predict` in a graph, one per node and setting, as
`predict_raw` holds one for the leaves."""
jit = bool(self.options.jit_predict)
key = ("node", id(node), jit, bool(self.options.qmc_simulations),
self._propagation_key())
traced = self._compiled.get(key)
if traced is None:
def body(x, x_var, n_sim, seed):
mu, var, sims, _ = node.predict(x, x_var=x_var, n_sim=n_sim,
seed=seed)
return (_tf.transpose(mu[:, :, 0]), _tf.transpose(var),
_tf.transpose(sims, [1, 0, 2]))
traced = _tf.function(body, jit_compile=jit or None,
reduce_retracing=True)
self._compiled[key] = traced
return traced
def _seed_offset(self, node):
"""What the model adds to its seed on the way from the first leaf
that reaches `node` down to it: parent `i` of an operation that
shifts its parents' seeds draws at `seed + i`."""
def search(current, offset):
if current is node:
return offset
parents = getattr(current, "parents", None)
if parents is None:
parent = getattr(current, "parent", None)
parents = [] if parent is None else [parent]
shifts = current._SHIFTS_PARENT_SEEDS
for i, parent in enumerate(parents):
found = search(parent, offset + (i if shifts else 0))
if found is not None:
return found
return None
for leaf in self.leaves:
found = search(leaf, 0)
if found is not None:
return found
raise ValueError("%s is not a node of this model's tree" % node.name)
# ------------------------------------------------------------------ #
# expert by expert
# ------------------------------------------------------------------ #
def _expert_structure(self):
"""The inputs whose experts an expert-by-expert pass visits, in the
network's order, and each one's GP nodes, after refusing what it
cannot handle: directional data, an input other than a
`BasicInput`, and the nodes with no expert-by-expert path --
`AdditiveGP`, `UncertainInputGP`, `GradientConstrainedInput`,
`RadialTrend`, `GaussianMixture`."""
network = _latent.network
if self.directional_data is not None:
raise ValueError("training and predicting by expert does not "
"take directional data")
nodes = self._nodes()
refused = (network.AdditiveGP, network.UncertainInputGP,
network.GradientConstrainedInput, network.RadialTrend,
network.GaussianMixture)
for node in nodes:
if isinstance(node, refused):
raise ValueError("training and predicting by expert does not "
"take a %s (%s)"
% (type(node).__name__, node.name))
roots = [n for n in nodes
if isinstance(n, network._RootLatentVariable)]
for root in roots:
if not isinstance(root, network.BasicInput):
raise ValueError("training and predicting by expert needs a "
"BasicInput at every root, and %s is a %s"
% (root.name, type(root).__name__))
gp = {id(r): [n for n in nodes if isinstance(n, network.BasicGP)
and n.root is r] for r in roots}
return roots, gp
def _expert_blocks(self):
"""Each input with the first column of its experts in the weight
table and their number: the table holds every input's experts side
by side, in the network's order."""
roots, _ = self._expert_structure()
blocks, start = [], 0
for root in roots:
# every root here is a BasicInput, whose count is set when it
# is built
count = _cast(int, root.n_experts)
blocks.append((root, start, count))
start += count
return blocks
def _transformed(self, container, rows=None, root=None):
"""The locations of `container` (every row of a block's fan-out) in
an input's transformed space, the space its kernels measure; the
first input by default."""
if root is None:
root = self._expert_structure()[0][0]
rows = _np.arange(container.n_data) if rows is None else rows
coords, _ = container.get_batched_coordinates(rows)
out = []
for band in self.options.batch_index(
len(coords), batch_size=self.options.prediction_batch_size):
x = _tf.constant(coords[band], _tf.float64)
out.append(_np.asarray(root.propagate(x)[0]))
return _np.concatenate(out, axis=0)
[docs]
def expert_weights(self, container: "_data._SpatialData | None" = None
) -> _types.FloatArray:
"""Each location's weight for each expert, as the model blends them.
One sweep, an expert at a time, so that memory holds one expert's
computation whatever their number: at every location the expert's
explained variance gives its raw weight, the raw weights are
normalized over the experts as the prediction normalizes them, and
averaged over the outputs of every GP node on the expert's input. A
node that reads another reads, in the sweep, what the one expert
alone gives below it -- which is the blend wherever that expert
carries weight.
Parameters
----------
container
The locations, the training data by default.
Returns
-------
ndarray
Of shape `(n, n_experts)`, rows summing to one -- on a network
of several inputs, every input's experts side by side in the
network's order, each input's summing to one.
"""
_, gp = self._expert_structure()
# every container holding locations is point-based
data = _cast("_data.containers._PointBased",
self.data if container is None else container)
n_rows = data.n_data
bands = self.options.batch_index(
n_rows, batch_size=self.options.prediction_batch_size)
blocks = []
for root, _, count in self._expert_blocks():
nodes = gp[id(root)]
n_out = sum(n.size for n in nodes)
raw = _np.zeros([n_rows, count, n_out])
for j in range(count):
with _latent.expert_subset({root: (j,)}):
for node in nodes:
node.refresh(self.options.jitter)
for band in bands:
rows = _np.arange(n_rows)[band]
coords, _ = data.get_batched_coordinates(rows)
variance, _ = data.get_batched_variance(rows)
x = _tf.constant(coords, _tf.float64)
x_var = _tf.constant(variance, _tf.float64)
column = 0
for node in nodes:
_, var = node.interpolate(
*node.parent.propagate(x, x_var))
var = _np.asarray(var).T
raw[band, j, column:column + node.size] = \
(1 - var) / (var + 1e-6) + 1e-6
column += node.size
weights = raw / raw.sum(axis=1, keepdims=True)
blocks.append(weights.mean(axis=2))
return _np.concatenate(blocks, axis=1)
@staticmethod
def _expert_sets(overlap, coverage, blocks):
"""Each expert's active sets, one per input: the experts of that
input holding `coverage` of the weight the expert's ground carries,
`overlap[g, h]` being the weight expert `h` has where expert `g`
has it (the data's weights, multiplied and summed); an expert is in
its own set. Returns the sets, as tuples of local indices per input,
and the largest share any of them leaves out."""
sets, dropped = [], []
for g in range(overlap.shape[0]):
per_input, worst = [], 0.0
for _, start, count in blocks:
row = overlap[g, start:start + count]
order = _np.argsort(-row)
share = _np.cumsum(row[order]) / row.sum()
n = int(_np.searchsorted(share, coverage)) + 1
keep = set(order[:n].tolist())
if start <= g < start + count:
keep.add(g - start)
keep = tuple(sorted(keep))
per_input.append(keep)
worst = max(worst, 1.0 - row[list(keep)].sum() / row.sum())
sets.append(tuple(per_input))
dropped.append(worst)
return sets, _np.asarray(dropped)
def _expert_plan(self, coverage, sets=None):
"""The weight table of the training data, and from it each expert's
batch distribution, active sets and shares of the KL terms. Given
`sets`, the active sets are kept and only the share of weight each
leaves out is measured afresh."""
blocks = self._expert_blocks()
table = self.expert_weights()
measured = _np.any(self.has_value > 0, axis=1)
w = table * measured[:, None]
total = w.sum(axis=0)
overlap = w.T @ w
formed, dropped = self._expert_sets(overlap, coverage, blocks)
if sets is None:
sets = formed
else:
dropped = _np.asarray([
max(1.0 - overlap[g, start:start + count][list(s)].sum()
/ overlap[g, start:start + count].sum()
for s, (_, start, count) in zip(per_input, blocks))
for g, per_input in enumerate(sets)])
# each expert's KL split among the batches it is active on, by the
# weight it carries there, so an epoch counts it once
carried = _np.zeros_like(overlap)
for g, per_input in enumerate(sets):
for s, (_, start, _) in zip(per_input, blocks):
columns = [start + k for k in s]
carried[g, columns] = overlap[g, columns]
shares = carried / carried.sum(axis=0, keepdims=True)
return dict(table=table, measured=measured, total=total,
sets=sets, dropped=dropped, shares=shares,
blocks=blocks)
@staticmethod
def _partition(table, rows, order, rng, coverage, blocks,
quotas="equal"):
"""An epoch's batches under `sampling="partition"`: the experts, in
`order`, each draw their share of the rows still unused this epoch,
by their weight and without replacement, so that the batches split
the rows and every row is read once. A row an expert's quota leaves
behind falls to a later expert's batch. The shares are equal, or
under `quotas="weight"` proportional to the weight each expert
carries over the rows. Each batch's active sets are formed from the
rows it drew: on each input, the experts holding `coverage` of the
weight the batch carries, its own expert always. Returns, per batch,
the expert, the rows, the sets, each row's share of weight its sets
hold (the least over the inputs), the mean weight the rows have for
the expert and the share of the expert's KL the batch carries."""
left = _np.ones(len(rows), dtype=bool)
order = _np.asarray(order)
if quotas == "weight":
mass = table[rows][:, order].sum(axis=0)
exact = mass / mass.sum() * len(rows)
counts = _np.floor(exact).astype(int)
# the rows rounding leaves over go to the largest remainders
short = len(rows) - counts.sum()
counts[_np.argsort(counts - exact)[:short]] += 1
else:
counts = [len(part) for part in
_np.array_split(_np.arange(len(rows)), len(order))]
batches = []
for g, quota in zip(order, counts):
available = _np.flatnonzero(left)
weight = table[rows[available], g]
pick = available[rng.choice(len(available), quota, replace=False,
p=weight / weight.sum())]
left[pick] = False
batches.append(VGPNetwork._batch(
table, rows[pick], g, coverage, blocks,
1.0 / _np.sum(order == g)))
return batches
@staticmethod
def _assignment(table, rows, rng, coverage, blocks, target):
"""An epoch's batches under `sampling="assignment"`: every row draws
one expert from its own weights, so that a row lands in an expert's
batch with the probability its weight for the expert gives, and is
read once. An expert's rows are split into as many batches as
`target` rows each makes, rounded, at least one, so that a crowded
expert steps more often; an expert that drew no rows takes no step,
and its KL is left out of the epoch. The batches come in a random
order, as `_partition` returns them."""
weight = table[rows]
cumulative = _np.cumsum(weight, axis=1)
drawn = rng.random(len(rows)) * cumulative[:, -1]
chosen = _np.minimum((cumulative < drawn[:, None]).sum(axis=1),
weight.shape[1] - 1)
batches = []
for g in _np.unique(chosen):
mine = rows[rng.permutation(_np.flatnonzero(chosen == g))]
k = max(1, int(round(len(mine) / target)))
batches.extend(VGPNetwork._batch(table, part, g, coverage,
blocks, 1.0 / k)
for part in _np.array_split(mine, k))
return [batches[i] for i in rng.permutation(len(batches))]
@staticmethod
def _batch(table, idx, g, coverage, blocks, kl):
"""Expert `g`'s batch of rows `idx` under a partition or an
assignment: on each input, the active experts are those holding
`coverage` of the weight the batch carries, its own expert always.
`kl` is the share of the expert's KL the batch carries."""
sets, cover = [], []
for _, start, count in blocks:
block = table[idx, start:start + count]
carried = block.sum(axis=0)
ranked = _np.argsort(-carried)
share = _np.cumsum(carried[ranked]) / carried.sum()
keep = set(ranked[:int(_np.searchsorted(share, coverage)) + 1]
.tolist())
if start <= g < start + count:
keep.add(int(g - start))
keep = tuple(sorted(keep))
sets.append(keep)
cover.append(block[:, list(keep)].sum(axis=1))
return dict(expert=int(g), rows=idx, sets=tuple(sets),
cover=_np.min(cover, axis=0),
own=float(table[idx, g].mean()), kl=kl)
@staticmethod
def _active_shares(table, batches, blocks):
"""Under `stepping="active"`: each expert's KL shared among the
epoch's batches it is active on, by the weight it carries in each,
so that an epoch counts it once. Returns, per batch, an array per
input over its active experts."""
carried = []
total = _np.zeros(table.shape[1])
for b in batches:
per_input = []
for s, (_, start, _) in zip(b["sets"], blocks):
columns = [start + k for k in s]
weight = table[b["rows"]][:, columns].sum(axis=0)
total[columns] += weight
per_input.append((columns, weight))
carried.append(per_input)
return [[weight / total[columns] for columns, weight in per_input]
for per_input in carried]
@staticmethod
def _expert_adam():
"""The optimizer training by expert gives the shared parameters (and
each expert, on the path that traces a step per set): `train_svi`'s
rate, decaying 0.999 a step, with AMSGrad."""
return _tf.keras.optimizers.Adam(
_tf.keras.optimizers.schedules.ExponentialDecay(1e-2, 1, 0.999),
amsgrad=True)
def _by_expert_state(self, roots, gp, slots, decay="steps"):
"""What training by expert keeps between calls -- the optimizers,
the batch generator, the active sets, the stopping rule -- made
afresh when the trained parameters or the path change, and dropped
by new data or a reset optimizer. The working copy and the traced
steps outlive it, on the model."""
local_ids = set()
local = []
for root in roots:
for k in range(root.n_experts):
variables = []
for node in gp[id(root)]:
for name in ("alpha_white_%d", "delta_%d", "bias_%d"):
parameter = node.parameters[name % k]
local_ids.add(id(parameter.variable))
if not parameter.fixed:
variables.append(parameter.variable)
local.append(variables)
shared = [v for v in self.get_unfixed_variables()
if id(v) not in local_ids]
key = (bool(slots), tuple(id(v) for vs in local for v in vs),
tuple(id(v) for v in shared), decay)
state = self.__dict__.get("_by_expert")
if state is not None and state["key"] == key:
# read afresh, so that a lowered tolerance trains on
state["converged"].tolerance = \
self.options.training_tolerance or 0.0
return state
held = {id(v) for v in shared}
state = dict(
key=key, local=local, shared=shared,
shared_parameters=[pr for pr in self._all_parameters
if id(pr.variable) in held],
shared_optimizer=self._expert_adam(),
# the steps train_svi would have taken by now, which sets both
# rates under decay="epochs"
clock=_tf.Variable(0.0, dtype=_tf.float64, trainable=False),
rng=_np.random.default_rng(self.options.seed),
converged=_Convergence(self.options.training_tolerance),
sets=None)
if decay == "epochs":
# the shared parameters on train_svi's clock too: on their own
# steps, one an epoch, they kept moving at the full rate while
# the experts' had decayed to nothing, and the fit went back
clock = state["clock"]
state["shared_optimizer"] = _tf.keras.optimizers.Adam(
lambda: 1e-2 * _tf.pow(_tf.constant(0.999, _tf.float32),
_tf.cast(clock, _tf.float32)),
amsgrad=True)
if shared:
state["shared_optimizer"].build(shared)
if slots:
store = self._expert_store(roots, gp)
for name, entry in store.items():
if isinstance(name, tuple):
entry.assign(_tf.zeros_like(entry))
continue
for kind in ("alpha", "delta", "bias"):
for prefix in ("m_", "v_", "vh_"):
moment = entry[prefix + kind]
moment.assign(_tf.zeros_like(moment))
state["store"] = store
else:
optimizers = [self._expert_adam() for _ in local]
for optimizer, variables in zip(optimizers, local):
optimizer.build(variables)
state.update(optimizers=optimizers, steps={})
self._by_expert = state
return state
def _expert_store(self, roots, gp):
"""The working copy training by expert steps under slots: each GP
node's `alpha_white`, `delta` and `bias` as raw values stacked over
its input's experts and padded, one padding expert after the last,
beside their bounds, which experts may move them, and AMSGrad's
moments; and each expert's step count, per input (`("t", id)`).
Keyed by the nodes' ids: two trees may number their nodes alike.
Made once per model, so that the traced steps reading it stay valid;
`_sync_expert_store` fills it from the parameters and writes it
back."""
store = self.__dict__.get("_expert_store_vars")
if store is not None:
return store
store = {}
for root in roots:
n_experts = root.n_experts
m = max(root.n_ip)
store[("t", id(root))] = _tf.Variable(
_np.zeros(n_experts + 1), trainable=False, dtype=_tf.float64)
for node in gp[id(root)]:
shapes = dict(alpha=[n_experts + 1, node.size, m, 1],
delta=[n_experts + 1, node.size, m],
bias=[n_experts + 1])
entry = {}
for kind, shape in shapes.items():
entry[kind] = _tf.Variable(_np.zeros(shape),
dtype=_tf.float64)
for prefix in ("m_", "v_", "vh_", "lo_", "hi_"):
entry[prefix + kind] = _tf.Variable(
_np.zeros(shape), trainable=False,
dtype=_tf.float64)
entry["free_" + kind] = _tf.Variable(
_np.zeros(n_experts + 1), trainable=False,
dtype=_tf.float64)
store[id(node)] = entry
self._expert_store_vars = store
return store
@staticmethod
def _sync_expert_store(roots, gp, store, back=False):
"""Fills the working copy from the parameters -- values, bounds and
which experts may move -- or, with `back`, writes it into them."""
names = (("alpha", "alpha_white_%d"), ("delta", "delta_%d"),
("bias", "bias_%d"))
for root in roots:
for node in gp[id(root)]:
entry = store[id(node)]
for kind, name in names:
if back:
values = entry[kind].numpy()
for i, n in enumerate(root.n_ip):
value = values[i] if kind == "bias" \
else values[i, :, :n]
node.parameters[name % i].variable.assign(value)
continue
value = _np.zeros(entry[kind].shape)
lo, hi = _np.zeros_like(value), _np.zeros_like(value)
free = _np.zeros(root.n_experts + 1)
for i, n in enumerate(root.n_ip):
parameter = node.parameters[name % i]
at = i if kind == "bias" \
else (i, slice(None), slice(0, n))
value[at] = parameter.variable.numpy()
lo[at] = parameter.min_transformed.numpy()
hi[at] = parameter.max_transformed.numpy()
free[i] = 0.0 if parameter.fixed else 1.0
entry[kind].assign(value)
entry["lo_" + kind].assign(lo)
entry["hi_" + kind].assign(hi)
entry["free_" + kind].assign(free)
@staticmethod
def _expert_amsgrad(store, names, steps, ids, mask, grads, clock=None):
"""One step on the slots' rows of one input's working copy: Keras'
Adam with `amsgrad=True` at a rate decaying 0.999 a step, the steps
counted per expert (`steps`) as each expert's own optimizer would
count them -- or, given a `clock`, as many as `train_svi` would have
taken by now. A row in an empty slot, or one whose parameter is
fixed, keeps what it holds; a padded entry has no gradient, and
stays where it is."""
t = _tf.gather(steps, ids)
decayed = t if clock is None else clock + 0.0 * t
# the rate and the betas raised to the step in single precision,
# as Keras computes them: in double precision the steps come out
# 6.7e-6 longer, and the two paths drift apart
rate = _tf.cast(_tf.constant(1e-2, _tf.float32) * _tf.pow(
_tf.constant(0.999, _tf.float32),
_tf.cast(decayed, _tf.float32)), _tf.float64)
step = t + 1.0
beta_1 = _tf.cast(_tf.constant(0.9, _tf.float32), _tf.float64)
beta_2 = _tf.cast(_tf.constant(0.999, _tf.float32), _tf.float64)
alpha = rate * _tf.sqrt(1.0 - _tf.pow(beta_2, step)) \
/ (1.0 - _tf.pow(beta_1, step))
index = ids[:, None]
grads = iter(grads)
for name in names:
entry = store[name]
for kind in ("alpha", "delta", "bias"):
grad = next(grads)
variable = entry[kind]
if grad is None:
grad = _tf.zeros_like(variable)
grad = _tf.gather(_tf.convert_to_tensor(grad), ids)
shape = [-1] + [1] * (len(variable.shape) - 1)
moving = _tf.reshape(
mask * _tf.gather(entry["free_" + kind], ids), shape) > 0
x0 = _tf.gather(variable, ids)
m0 = _tf.gather(entry["m_" + kind], ids)
v0 = _tf.gather(entry["v_" + kind], ids)
vh0 = _tf.gather(entry["vh_" + kind], ids)
m1 = m0 + (grad - m0) * (1.0 - 0.9)
v1 = v0 + (_tf.square(grad) - v0) * (1.0 - 0.999)
vh1 = _tf.maximum(vh0, v1)
x1 = x0 - m1 * _tf.reshape(alpha, shape) \
/ (_tf.sqrt(vh1) + 1e-7)
x1 = _tf.clip_by_value(
x1, _tf.gather(entry["lo_" + kind], ids),
_tf.gather(entry["hi_" + kind], ids))
for target, new, old in ((variable, x1, x0),
(entry["m_" + kind], m1, m0),
(entry["v_" + kind], v1, v0),
(entry["vh_" + kind], vh1, vh0)):
target.scatter_nd_update(index,
_tf.where(moving, new, old))
steps.scatter_nd_update(index, t + mask)
[docs]
def train_by_expert(self, epochs: int = 10,
batch_size: "int | None" = None,
coverage: float = 0.99,
global_update: str = "epoch",
weights_every: int = 1,
visits: int = 1,
slots: bool = True,
decay: str = "steps",
sampling: str = "assignment",
quotas: str = "equal",
stepping: str = "active"
) -> "dict[str, _Any]":
"""Train an expert at a time, so memory does not grow with their
number.
Each epoch every data row draws one expert from its own weights, so
that a row lands in an expert's batch with the probability its
weight for the expert gives and is read once an epoch. An expert's
rows are split into batches of about N / (J x `visits`) rows, so a
crowded expert takes more of them, and the batches come in a random
order. The experts active on a batch -- on each input, those holding
`coverage` of the weight the batch carries, its own expert always --
are the only ones computed, and each takes a step on the batch's
bound: its data term as it stands, less each active expert's share
of its KL divergence, shared among the batches the expert is active
on by the weight it carries in each, and a share of the priors. An
epoch's batches add up to the bound. A network on several inputs
assigns each row to one expert among every input's.
Under `sampling="replacement"` an epoch is `visits` rounds instead,
each visiting every expert once in an order of its own: for expert
`j` a batch of `batch_size` rows is drawn with probability
proportional to their weight for `j`, and the batch's data term is
`W_j` times its mean log-likelihood, `W_j` the expert's total weight,
over the number of inputs, so that the experts' terms add up to the
whole data term over an epoch in expectation. Under
`sampling="partition"` the experts, in a random order, each draw
their share of the rows still unused -- by their weight, without
replacement -- so that a row a dense expert's quota leaves behind
falls to a later expert's batch. Under either split of the rows,
`stepping="own"` steps the batch's own expert alone, its KL counted
whole in its own batches.
Under `options.training_tolerance` training stops once the epochs'
bound settles, as `train_svi`'s does, and a call made after that
takes no step.
Parameters
----------
epochs
Number of epochs.
batch_size
Rows per batch under `sampling="replacement"`,
`options.training_batch_size` by default; the other samplings
size their batches from `visits`.
coverage
The share of a batch's expert weight its active experts hold;
the rest are left out of its blend.
global_update
How the parameters every expert shares step: `"batch"` on each
batch's gradient; `"round"` once a round, on the gradients
added up over it; `"epoch"` once an epoch.
weights_every
Epochs between sweeps of the weight table, which sets the batch
distributions, the KL shares and, in slots, the active sets;
without slots the sets are formed from the first sweep and kept.
visits
About how many batches an expert takes an epoch: under an
assignment or a partition the batches hold N / (J x `visits`)
rows; under `"replacement"` the rounds per epoch, which with
`batch_size` divided by as much read as many rows in `visits`
times the steps.
slots
Compute the active experts in as many slots as the largest set
holds, so that one traced step serves every set; `False` traces
a step per set.
decay
How the learning rates decay, 0.999 a step: `"steps"` counts
each expert's own steps and the shared parameters' own, as an
optimizer of their own would; `"epochs"` counts, for both, the
steps `train_svi` would have taken by the same point, batches
of `options.training_batch_size` over the whole data an epoch,
so that the two decay together as they do there. Each expert
steps on every batch whose set it is in and the shared
parameters as `global_update` says, so on their own counts the
first falls faster than `train_svi`'s rate and the second
slower; on `train_svi`'s count both have decayed before
training by expert, which needs more epochs, has converged.
`"epochs"` needs slots.
sampling
`"assignment"` has each row choose its expert, an epoch as many
batches as the experts' rows make, and no rounds for
`global_update="round"`; `"replacement"` draws each expert's
batch from every row; `"partition"` has the experts split the
rows among themselves, as above.
quotas
Under `"partition"`, how many rows each expert draws: `"equal"`
shares, or by `"weight"`, each expert's share of the weight the
rows carry, so that a dense expert leaves fewer rows behind.
stepping
Under `"partition"` or `"assignment"`, which experts step on a
batch: its `"own"` expert alone, its KL counted whole in its
own batches, or every `"active"` one, each expert's KL shared
among the batches it is active on by the weight it carries
there.
Returns
-------
dict
The run's record: the bound per epoch, seconds per epoch, the
active sets (on a network of several inputs, one per input
for each expert), the weight they drop, and the traces the
steps took; under `"partition"` and `"assignment"` also, per
epoch, the batches, the fewest and most an expert stepped on,
the active sets' sizes, the fewest and most rows a batch drew,
the share of each row's weight its batch's sets hold, and the
weight the first and the last quarter of the batches have for
their own expert.
"""
if global_update not in ("batch", "round", "epoch"):
raise ValueError("global_update is 'batch', 'round' or 'epoch', "
"got %r" % (global_update,))
if decay not in ("steps", "epochs"):
raise ValueError("decay is 'steps' or 'epochs', got %r"
% (decay,))
if decay == "epochs" and not slots:
raise ValueError("decay='epochs' needs slots: without, each "
"expert's optimizer keeps its own schedule")
if sampling not in ("replacement", "partition", "assignment"):
raise ValueError("sampling is 'replacement', 'partition' or "
"'assignment', got %r" % (sampling,))
if quotas not in ("equal", "weight"):
raise ValueError("quotas is 'equal' or 'weight', got %r"
% (quotas,))
if stepping not in ("own", "active"):
raise ValueError("stepping is 'own' or 'active', got %r"
% (stepping,))
assignment = sampling == "assignment"
if assignment and global_update == "round":
raise ValueError("sampling='assignment' has no rounds; "
"global_update is 'batch' or 'epoch'")
# the rows split among the batches, the batch's own expert stepping
partition = sampling != "replacement"
network = _latent.network
roots, gp = self._expert_structure()
n_experts = sum(count for _, _, count in self._expert_blocks())
batch_size = batch_size or self.options.training_batch_size
state = self._by_expert_state(roots, gp, slots, decay)
def readable(sets):
# one input: its set alone, as a single-input record reads
return sets if sets is None or len(roots) > 1 \
else [s[0] for s in sets]
record: "dict[str, _Any]" = dict(
bound=[], seconds=[], subsets=readable(state["sets"]),
dropped=None, traces=0, partition=[])
if state["converged"].settled():
if self.options.verbose:
print("The bound has settled; nothing to train.")
return record
local, shared = state["local"], state["shared"]
rng = state["rng"]
nodes = self._nodes()
# the nodes whose KL is a term per expert, shared out among the
# batches each expert is active on: the GP nodes and the walks
per_expert = {
id(r): [n for n in nodes if getattr(n, "root", None) is r
and isinstance(n, (network.BasicGP, network.GPWalk))]
for r in roots}
held = {id(n) for ns in per_expert.values() for n in ns}
others = [n for n in nodes if id(n) not in held]
names = {id(r): [id(node) for node in gp[id(r)]] for r in roots}
# the batches of an epoch, among which the KL of the rest and the
# priors are shared out (`n_batches`: under an assignment the count
# changes from epoch to epoch); every row's weights sum to one per
# input
per_epoch = n_experts * visits
n_inputs = len(roots)
def objective(x, y, has_value, x_var, scale, shares, n_batches):
self._refresh(self.options.jitter)
data = self._data_log_lik(
x, y, has_value, [{} for _ in self.variables],
x_var=x_var, samples=self.options.training_samples,
seed=self.options.seed)
kl = _tf.constant(0.0, _tf.float64)
for root, share in zip(roots, shares):
for node in per_expert[id(root)]:
terms = node.expert_kl_terms()
if isinstance(terms, list):
terms = _tf.stack(terms)
kl = kl + _tf.reduce_sum(terms * share)
for node in others:
kl = kl + node.kl_divergence() / n_batches
return scale * data - kl + self.log_prior() / n_batches
def slot_step(sizes):
# kept on the model across calls: one trace serves every set of
# experts, and a trace with gradients is never given back
steps = self.__dict__.setdefault("_expert_steps", {})
key = (sizes, visits, state["key"][2], decay)
if key in steps:
return steps[key]
store = state["store"]
stacks = {n: (store[n]["alpha"], store[n]["delta"],
store[n]["bias"])
for r in roots for n in names[id(r)]}
copies = [store[n][k] for r in roots for n in names[id(r)]
for k in ("alpha", "delta", "bias")]
@_tf.function(reduce_retracing=True)
def step(ids, mask, moving, x, y, has_value, x_var, scale,
shares, n_batches, clock):
context = {root: network._Slots(size, i, m, stacks)
for root, size, i, m in zip(roots, sizes, ids,
mask)}
with network.expert_slots(context):
with _tf.GradientTape() as tape:
# negated inside the tape, which records only what
# is computed under it
loss = - objective(x, y, has_value, x_var, scale,
shares, n_batches)
grads = tape.gradient(loss, copies + shared)
start = 0
# the slots that step: every active expert, or the batch's
# own alone under a partition
for root, i, m in zip(roots, ids, moving):
count = 3 * len(names[id(root)])
self._expert_amsgrad(
store, names[id(root)], store[("t", id(root))], i,
m, grads[start:start + count],
None if decay == "steps" else clock)
start += count
shared_grads = [
_tf.zeros_like(v) if g is None
else _tf.convert_to_tensor(g)
for g, v in zip(grads[len(copies):], shared)]
return - loss, shared_grads
steps[key] = step
return step
def set_step(sets, own=None):
key = (sets, visits, own)
if key in state["steps"]:
return state["steps"][key]
experts = [start + k for s, (_, start, _) in
zip(sets, plan["blocks"]) for k in s]
if own is not None:
experts = [own]
variables = [v for g in experts for v in local[g]]
counts = [len(local[g]) for g in experts]
optimizers = state["optimizers"]
@_tf.function(reduce_retracing=True)
def step(x, y, has_value, x_var, scale, shares, n_batches):
with _tf.GradientTape() as tape:
loss = - objective(x, y, has_value, x_var, scale, shares,
n_batches)
grads = tape.gradient(loss, variables + shared)
grads = [_tf.zeros_like(v) if g is None else g
for g, v in zip(grads, variables + shared)]
start = 0
for g, count in zip(experts, counts):
optimizers[g].apply_gradients(zip(
grads[start:start + count],
variables[start:start + count]))
start += count
return - loss, grads[len(variables):]
state["steps"][key] = step
return step
converged = state["converged"]
store = state.get("store")
state.setdefault("batches", 0)
svi_steps = self.data.n_data / self.options.training_batch_size
used = {}
plan: "dict[str, _Any]" = {}
rows = _np.zeros(0, int)
sizes: "tuple[int, ...]" = ()
done = 0
try:
if slots:
self._sync_expert_store(roots, gp, store)
with self._propagation(), \
_progress.reporting(
"train", None if assignment else epochs * per_epoch,
"batch") as report:
for epoch in range(epochs):
start_time = _time.perf_counter()
if not plan or epoch % weights_every == 0:
if slots:
# the sweep reads the parameters
self._sync_expert_store(roots, gp, store,
back=True)
# in slots the active sets are formed afresh with
# every weight table, and follow the weights as the
# ranges grow; the slots only ever widen, since a
# new number of them is a new trace. On the path
# tracing a step per set they are formed once and
# kept: re-formed, a 16-expert run traced 76 steps
# and grew 6.6 GB
plan = self._expert_plan(
coverage, None if slots else state["sets"])
state["sets"] = plan["sets"]
rows = _np.flatnonzero(plan["measured"])
if not partition:
widest = [max(len(sets[r])
for sets in plan["sets"])
for r in range(n_inputs)]
previous = state.get("sizes") or [0] * n_inputs
state["sizes"] = sizes = tuple(
max(a, b) for a, b in zip(previous, widest))
record["subsets"] = readable(plan["sets"])
record["dropped"] = plan["dropped"]
total = 0.0
summed = None
if assignment:
batches = self._assignment(
plan["table"], rows, rng, coverage,
plan["blocks"], len(rows) / per_epoch)
order = [b["expert"] for b in batches]
else:
order = _np.concatenate([rng.permutation(n_experts)
for _ in range(visits)])
if sampling == "partition":
batches = self._partition(
plan["table"], rows, order, rng, coverage,
plan["blocks"], quotas)
n_batches = _tf.constant(float(len(order)),
_tf.float64)
if partition:
widest = [max(len(b["sets"][r]) for b in batches)
for r in range(n_inputs)]
previous = state.get("sizes") or [0] * n_inputs
state["sizes"] = sizes = tuple(
max(a, b) for a, b in zip(previous, widest))
quarter = max(1, len(batches) // 4)
cover = _np.concatenate([b["cover"] for b in batches])
if stepping == "active":
active = self._active_shares(
plan["table"], batches, plan["blocks"])
steps = _np.zeros(n_experts, int)
for b in batches:
for s, (_, start, _) in zip(b["sets"],
plan["blocks"]):
steps[[start + k for k in s]] += 1
else:
steps = _np.bincount(order, minlength=n_experts)
record["partition"].append(dict(
batches=len(batches),
steps=[int(steps.min()), int(steps.max())],
set_mean=[float(_np.mean(
[len(b["sets"][r]) for b in batches]))
for r in range(n_inputs)],
set_max=list(widest),
rows=[min(len(b["rows"]) for b in batches),
max(len(b["rows"]) for b in batches)],
cover_mean=float(cover.mean()),
cover_p01=float(_np.quantile(cover, 0.01)),
own_first=float(_np.mean(
[b["own"] for b in batches[:quarter]])),
own_last=float(_np.mean(
[b["own"] for b in batches[-quarter:]]))))
for position, g in enumerate(order):
if partition:
batch = batches[position]
idx, sets = batch["rows"], batch["sets"]
if stepping == "active":
# every active expert steps on the batch's
# sum, its KL shared by the weight it carries
shares = active[position]
moving = [_np.ones(len(s)) for s in sets]
else:
# the batch's own expert alone steps, and its
# KL is counted whole in its own batches
shares = [_np.asarray(
[batch["kl"] if start + k == g else 0.0
for k in s])
for s, (_, start, _)
in zip(sets, plan["blocks"])]
moving = [_np.asarray(
[1.0 if start + k == g else 0.0
for k in s])
for s, (_, start, _)
in zip(sets, plan["blocks"])]
scale = _tf.constant(1.0, _tf.float64)
else:
p = plan["table"][rows, g] \
/ plan["table"][rows, g].sum()
idx = rows[rng.choice(len(rows), batch_size, p=p)]
sets = plan["sets"][g]
shares = [plan["shares"][g, [start + k for k in s]]
/ visits for s, (_, start, _)
in zip(sets, plan["blocks"])]
moving = [_np.ones(len(s)) for s in sets]
scale = _tf.constant(
plan["total"][g]
/ (batch_size * visits * n_inputs),
_tf.float64)
inputs = (
_tf.constant(self.data.coordinates[idx],
_tf.float64),
_tf.constant(self.y[idx], _tf.float64),
_tf.constant(self.has_value[idx], _tf.float64),
_tf.constant(
self.data.get_batched_variance(idx)[0],
_tf.float64))
if slots:
pads = [size - len(s)
for size, s in zip(sizes, sets)]
step = slot_step(sizes)
bound, grads = step(
[_tf.constant(list(s) + [r.n_experts] * pad,
_tf.int32)
for s, r, pad in zip(sets, roots, pads)],
[_tf.constant([1.0] * len(s) + [0.0] * pad,
_tf.float64)
for s, pad in zip(sets, pads)],
[_tf.constant(_np.concatenate(
[m, _np.zeros(pad)]), _tf.float64)
for m, pad in zip(moving, pads)],
*inputs, scale,
[_tf.constant(_np.concatenate(
[share, _np.zeros(pad)]), _tf.float64)
for share, pad in zip(shares, pads)],
n_batches,
# the steps train_svi would have taken
_tf.constant(
state["batches"] * svi_steps / per_epoch,
_tf.float64))
else:
with _latent.expert_subset(
dict(zip(roots, sets))):
step = set_step(
sets, int(g) if partition
and stepping == "own" else None)
bound, grads = step(
*inputs, scale,
[_tf.constant(share, _tf.float64)
for share in shares], n_batches)
for pr in self._all_parameters:
pr.refresh()
used[id(step)] = step
total += float(bound)
summed = grads if summed is None else \
[a + b for a, b in zip(summed, grads)]
if global_update == "batch" or position + 1 == len(
order) or (global_update == "round" and
(position + 1) % n_experts == 0):
if shared:
state["clock"].assign(
state["batches"] * svi_steps / per_epoch)
state["shared_optimizer"].apply_gradients(
zip(summed, shared))
for pr in state["shared_parameters"]:
pr.refresh()
summed = None
done += 1
state["batches"] += 1
report(done, bound=float(bound))
record["bound"].append(total)
record["seconds"].append(_time.perf_counter() - start_time)
self.training_log.append(total)
if self.options.verbose:
print("\rEpoch %d | bound: %s" % (epoch + 1, total),
end="")
if converged.stop(total):
if self.options.verbose:
print("\nStopped at epoch %d: the bound has "
"settled" % (epoch + 1), end="")
break
finally:
# what training got to, cancelled or not
if slots:
self._sync_expert_store(roots, gp, store, back=True)
if self.options.verbose:
print("\n")
record["traces"] = sum(int(f.experimental_get_tracing_count())
for f in used.values())
return record
def _slot_variables(self, root, size):
"""The Variables a prediction's cached traces read an input's slots
from, one pair per input and number of slots, kept so that the
traces stay valid. The ids are int64: TensorFlow keeps an int32
Variable in host memory, which an XLA-compiled prediction on a GPU
cannot read."""
held = self.__dict__.setdefault("_slot_vars", {})
key = (id(root), size)
if key not in held:
held[key] = (
_tf.Variable(_np.zeros(size, _np.int64), trainable=False),
_tf.Variable(_np.zeros(size), trainable=False,
dtype=_tf.float64))
return held[key]
[docs]
def predict_by_expert(self, newdata: "_data._SpatialData",
n_sim: "int | None" = None,
coverage: float = 0.99, neighbours: int = 8,
include_noise: bool = True,
grouping: str = "home",
where: _types.Where = None,
slots: "bool | int" = True,
pack: bool = True) -> "dict[str, _Any]":
"""Predict with only the experts active at each location.
The training data's expert weights are carried to the locations by
an inverse-distance average of the `neighbours` nearest data, in
each input's transformed space. A location's own experts, on each
input, are those holding `coverage` of its weight, a block's the
union over its sub-blocks. Each group of locations is predicted
with only its experts, the rest left out of the blend.
Under `slots` every group is computed in as many slots, on each
input, as the largest active set training forms from the data, so
that one traced refresh and one traced prediction serve them all and
memory stays where training's was. A location takes its own
experts, or its leading ones where it needs more than there are
slots -- the weight that leaves out is in `left_out`. With `pack`,
groups are merged wherever the experts of both fit in the slots,
which costs nothing a padded slot would not; a location then blends
experts beyond its own, so its answer depends, within the weight
`coverage` leaves out, on what is predicted beside it. Without, it
depends on the location alone.
Without slots each group is a trace of its own, and `grouping`
keeps them few: under `"exact"` the locations are grouped by their
own experts; under `"home"` a location takes, on each input, the
active set training forms for the first of its leading experts
whose set holds `coverage` of its weight, then the union of its
leading experts' sets, and its own only where neither does.
Parameters
----------
newdata
The locations to predict. Modified in place, as by `predict`.
n_sim
Number of realizations, as for `predict`.
coverage
The share of a location's expert weight its experts must hold.
neighbours
Data averaged into a location's weights.
include_noise
As for `predict`.
grouping
`"home"` or `"exact"`, without slots.
where
The locations to predict, as for `predict`; the rest keep what
they hold.
slots
Compute the groups in a fixed number of slots -- `True` for as
many as training's largest active set on each input, or a
number for every input, memory growing with it times
`options.prediction_batch_size`; `False` traces one prediction
per group.
pack
Under `slots`, merge groups whose experts fit the slots
together.
Returns
-------
dict
The groups (their active experts -- on a network of several
inputs, one set per input -- and sizes), the number of slots,
and the largest share of weight left out at each location, NaN
where none was predicted.
"""
if grouping not in ("home", "exact"):
raise ValueError("grouping is 'home' or 'exact', got %r"
% (grouping,))
blocks = self._expert_blocks()
roots = [root for root, _, _ in blocks]
mask = _where_mask(newdata, where)
locations = _np.arange(newdata.n_data) if mask is None \
else _np.flatnonzero(mask)
table = self.expert_weights()
rows_per = newdata.rows_per_location
measured = _np.any(self.has_value > 0, axis=1)
w = table * measured[:, None]
home, _ = self._expert_sets(w.T @ w, coverage, blocks)
# per input: the locations' weights, carried from the data in that
# input's own transformed space, and its training sets
per_input = []
for r, (root, start, count) in enumerate(blocks):
tree = _spatial.cKDTree(self._transformed(self.data, root=root))
targets = self._transformed(newdata, locations, root=root)
distance, nearest = tree.query(targets, k=neighbours)
inverse = 1.0 / (distance ** 2 + 1e-12)
weights = _np.einsum(
"rn,rnj->rj", inverse,
table[:, start:start + count][nearest]) \
/ inverse.sum(axis=1, keepdims=True)
weights = weights.reshape([len(locations), rows_per, -1])
home_sets = [sets[r] for sets in home[start:start + count]]
# in slots, as many as training's largest active set, so that
# prediction holds what training held, or as many as asked for
size = max(len(s) for s in home_sets) if slots is True \
else int(slots)
per_input.append((weights, home_sets, size))
def own_set(weights, at):
keep = set()
for row in weights[at]:
order = _np.argsort(-row)
n = int(_np.searchsorted(_np.cumsum(row[order]),
coverage)) + 1
keep |= set(order[:n].tolist())
return tuple(sorted(keep))
def lost(weights, at, subset):
return 1.0 - weights[at][:, list(subset)].sum(axis=1).min()
def choose(weights, home_sets, size, at):
mean = weights[at].mean(axis=0)
subset = None
if slots:
# a location's own experts -- grouping into few sets only
# ever saved traces, and in slots there is one -- its
# leading ones where it needs more than there are slots
subset = own_set(weights, at)
if len(subset) > size:
subset = tuple(sorted(_np.argsort(-mean)[:size].tolist()))
elif grouping == "home":
# the first of its own leading experts whose set holds
# `coverage` of the location's weight
leading = _np.argsort(-mean)[:4]
for k in leading:
if lost(weights, at, home_sets[k]) <= 1.0 - coverage:
subset = home_sets[k]
break
# then the union of the leading experts' sets, a family of
# few members many locations share -- between drillholes a
# block blends experts no one set covers, and its own set
# was a group of its own nearly every time
for n in (2, 3):
if subset is not None:
break
union = tuple(sorted(set().union(
*(home_sets[k] for k in leading[:n]))))
if lost(weights, at, union) <= 1.0 - coverage:
subset = union
if subset is None:
subset = own_set(weights, at)
return subset
groups = {}
left_out = _np.full(newdata.n_data, _np.nan)
for at, loc in enumerate(locations):
key = tuple(choose(weights, home_sets, size, at)
for weights, home_sets, size in per_input)
groups.setdefault(key, []).append(loc)
left_out[loc] = max(lost(weights, at, subset)
for (weights, _, _), subset
in zip(per_input, key))
sizes = [size for _, _, size in per_input]
groups = sorted(groups.items(), key=lambda g: -len(g[1]))
if slots and pack:
packed = []
for key, locs in groups:
for entry in packed:
unions = [a.union(b) for a, b in zip(entry[0], key)]
if all(len(u) <= size for u, size in zip(unions, sizes)):
entry[0] = unions
entry[1] = entry[1] + locs
break
else:
packed.append([[set(s) for s in key], list(locs)])
groups = [(tuple(tuple(sorted(s)) for s in e[0]), e[1])
for e in packed]
n_sim = _simulation_count(newdata, self.variables, n_sim,
keep=mask is not None)
network = _latent.network
with _progress.reporting("predict_by_expert", len(groups),
"group") as report:
for done, (key, locs) in enumerate(groups):
if slots:
context = {}
for root, subset, size in zip(roots, key, sizes):
ids, filled = self._slot_variables(root, size)
pad = size - len(subset)
ids.assign(list(subset) + [root.n_experts] * pad)
filled.assign([1.0] * len(subset) + [0.0] * pad)
context[root] = network._Slots(size, ids, filled)
context = network.expert_slots(context)
else:
context = _latent.expert_subset(dict(zip(roots, key)))
with context:
self._predict(newdata, n_sim, include_noise,
_np.sort(_np.asarray(locs)),
check_measurements=False)
report(done + 1)
single = len(roots) == 1
return dict(subsets=[g[0][0] if single else g[0] for g in groups],
sizes=[len(g[1]) for g in groups],
slots=(sizes[0] if single else tuple(sizes))
if slots else None,
left_out=left_out)
def _predict(self, newdata, n_sim, include_noise, where,
check_measurements=True):
if self.data.n_dim != newdata.n_dim:
raise ValueError("dimension of newdata is incompatible with model")
where = _where_mask(newdata, where)
n_sim = _simulation_count(newdata, self.variables, n_sim,
keep=where is not None)
# managing variables
variable_inputs = []
for v in self.variables:
fresh = v not in newdata.variables.keys()
if fresh:
self.data.variables[v].copy_to(newdata)
else:
# the shares are computed at the training variable's cut-offs
# and filed by position under the target's, so a target that
# declares others takes the model's first
for path, old, new in newdata.variables[v]._adopt_cutoffs(
self.data.variables[v]):
warnings.warn(
"%s declared the cut-offs %s, and the model computes "
"its shares at %s; it declares the model's now, and "
"the shares of the ones it gave up were dropped"
% (path, old, new))
# Allocate when the variable is new, whether or not only some
# locations are being visited: what is not visited stays NaN,
# which is what `unpredicted` reads and what the reporting layer
# already skips. Never reallocate an existing one under `where` --
# that would wipe the simulations the untouched locations hold,
# which is the whole point of naming only some.
if where is None or fresh \
or _stored_n_sim(newdata.variables[v]) is None:
newdata.variables[v].allocate_simulations(n_sim)
variable_inputs.append(self.data.variables[v].prediction_input())
def batch_pred(x, x_var, n_splits):
return self.predict_raw(
x,
variable_inputs,
x_var=x_var,
seed=self.options.seed,
n_sim=n_sim,
include_noise=include_noise,
n_splits=n_splits
)
for batch, output in self._over_batches(newdata, batch_pred, where):
for v, lik, upd in zip(self.variables, self.likelihoods, output):
# A vector variable hands the latent moments to its
# components only when column i is component i's own, which
# an elementwise warping guarantees; under a rotation or a
# projection no latent column belongs to any one component,
# and storing one under a component's name would be wrong
# in a way nobody would catch.
elementwise = lik._SINGLE_WARPING \
and lik.warping.elementwise
newdata.variables[v].update(batch, elementwise=elementwise,
**upd)
if check_measurements:
self._check_measurements(newdata, where, n_sim)
def _check_measurements(self, newdata, where, n_sim):
"""Writes `pit_*` and `warped_*` where `newdata` holds measurements
(see `predict`), at the locations `where` names. A column already
there keeps its values elsewhere, so a prediction finished in parts
fills it in parts."""
if newdata.rows_per_location != 1:
return
visit = _np.ones(newdata.n_data, dtype=bool) if where is None \
else where
truths = {v: _measured_truth(newdata.variables[v])
for v, _ in self._measured_variables()}
rows = visit & _np.any(
[has.any(axis=1) for _, has, _ in truths.values()] or
[_np.zeros(newdata.n_data, dtype=bool)], axis=0)
if not rows.any():
return
def column(name):
if name in newdata.metadata:
return _np.asarray(newdata.get_metadata(name),
dtype=float).copy()
return _np.full(newdata.n_data, _np.nan)
# a pass of its own over the measured rows, silenced: it is part of
# the prediction a caller asked for, not a second one to report
pit = {}
with _progress.progress(None):
for batch, samples in self.measurement_batches(
newdata, n_sim=n_sim, where=rows):
for v, sample in samples.items():
truth, has, components = truths[v]
for c, component in enumerate(components):
measured = has[batch, c]
if not measured.any():
continue
key = _pit_column(v, component)
if key not in pit:
pit[key] = column(key)
pit[key][batch[measured]] = _pit(
sample[measured, c, :], truth[batch, c][measured])
for key, values in pit.items():
newdata.add_metadata(key, values)
for v, lik in self._measured_variables():
values, has_value = newdata.variables[v].get_measurements()
values = _np.asarray(values, dtype=float)
if values.ndim == 1:
values = values[:, None]
has_value = _np.asarray(has_value)
if has_value.ndim == 1:
has_value = has_value[:, None]
if not lik._SINGLE_WARPING:
# a mixture of likelihoods has a warped space per component
# and no one to store
continue
# the whole row or nothing: a warping may mix the columns
full = rows & _np.all(has_value == 1.0, axis=1)
if not full.any():
continue
warped, _ = lik.warping.forward(values[full])
warped = _np.asarray(warped, dtype=float)
for i in range(warped.shape[1]):
key = "warped_%s_%d" % (v, i)
stored = column(key)
stored[full] = warped[:, i]
newdata.add_metadata(key, stored)
def _over_batches(self, newdata, call, where=None, with_rows=False):
"""Runs `call(coordinates, variance, n_splits)` over `newdata` --
`call(coordinates, variance, n_splits, rows)` under `with_rows`, for
a caller that keeps something per location and must hand each
batch its own slice.
A discretized block fans out into several rows before it reaches the
model, so the batch is measured in those rows: otherwise
`prediction_batch_size` would mean `prod(discretization)` times as much
work here as it does on a grid of points.
Yields `(rows, result)`, the rows being indices into `newdata` so that
a caller can write each result back to the location it came from.
"""
batch_size = max(1, self.options.prediction_batch_size
// newdata.rows_per_location)
rows = _np.arange(newdata.n_data)
if where is not None:
where = _np.asarray(where)
rows = _np.flatnonzero(where) if where.dtype == bool else where
batch_id = [rows[batch] for batch in
self.options.batch_index(len(rows), batch_size=batch_size)]
# Refresh the posterior once (the parameters are fixed during
# prediction) and snapshot each node's state into Variables, so the
# cached `predict_raw` graph reads current values without recomputing the
# posterior (Cholesky factorizations, etc.) on every batch. The refresh
# itself is traced -- see `latent.refresh_cached`.
with self._propagation():
if all(hasattr(leaf, "cache_prediction_state")
for leaf in self.leaves):
_latent.refresh_cached(self.leaves, self.options.jitter,
owner=self)
else:
self._refresh(self.options.jitter)
for i, batch in enumerate(batch_id):
if self.options.verbose:
print("\rProcessing batch %s of %s "
% (str(i + 1), str(len(batch_id))), end="")
# `i` batches are done, and what `done` counts is what a cancel
# here would leave written: the consumer writes each batch back
# before asking for the next, so the report for batch `i` is made
# once batch `i - 1` has landed in the container.
_progress.emit("predict", i, len(batch_id), "batch")
data_coords, splits = newdata.get_batched_coordinates(batch)
data_var, _ = newdata.get_batched_variance(batch)
x = _tf.constant(data_coords, _tf.float64)
x_var = _tf.constant(data_var, _tf.float64)
# under the rule the state was refreshed under: a caller that
# reads the leaves itself (`measurement_batches`,
# `responsibilities`) would otherwise run them under the module's
# default against a refresh made under the model's
with self._propagation():
out = call(x, x_var, splits, batch) if with_rows \
else call(x, x_var, splits)
yield batch, out
_progress.emit("predict", len(batch_id), len(batch_id), "batch")
if self.options.verbose:
print("\n")
[docs]
def predict_measurements(self, newdata: "_data._SpatialData",
n_sim: int = 20,
n_nodes: int = 32) -> "dict[str, _np.ndarray]":
"""Draw the predictive distribution of a measurement per location.
:meth:`predict` reports the ground, with the likelihood noise
integrated out, so its simulations describe a quantity no sample
observes. This keeps the noise instead, giving `n_sim * n_nodes`
equally likely readings per location -- the distribution to compare
against measured data, as an accuracy plot or a cross-validation
does.
Nothing is stored: `newdata` is not modified.
Parameters
----------
newdata
Locations to ask about, of the same dimension as the training
data. Intended for the locations that carry measurements, not
for a block model.
n_sim
Latent realizations per location.
n_nodes
Equal-share noise values per realization, so that the two axes
pool into one sample. The nodes are rotated at random per
location and realization from the model's seed, so the sample
is unbiased in every moment, reaches the tails, and is
independent between locations.
Returns
-------
dict of str to ndarray
One `(n_data, size, n_sim * n_nodes)` array per variable whose
likelihood carries a warping. Categorical variables are skipped:
their noise lives in the probabilities, leaving no value for a
measurement to scatter around.
Raises
------
MemoryError
If the answer would not fit. It is held whole, by design, and
costs `n_sim * n_nodes * 8` bytes a row for each column of each
variable -- 5 KB a row at the defaults, and twice that at the
peak, while the batches and the assembled whole are both live.
The ceiling is `models.MEASUREMENT_LIMIT`, and the message names
the ways under it.
See Also
--------
predict : the ground, with the noise integrated out.
"""
wanted = self._measured_variables()
# Refused before any work rather than discovered part way through:
# this door holds the answer whole, and its size is known from the
# shapes alone. Sized in bytes rather than rows because a vector
# variable costs a column each and the two node counts multiply.
wanted_bytes = sum(newdata.n_data * lik.size * n_sim * n_nodes * 8
for _, lik in wanted)
if wanted_bytes > MEASUREMENT_LIMIT:
raise MemoryError(
"measurement samples for %d location(s) at n_sim=%d and "
"n_nodes=%d come to %.1f GB, past the %.1f GB "
"`models.MEASUREMENT_LIMIT` holds. Ask about fewer "
"locations -- more cross-validation folds make each one "
"smaller -- or lower n_sim or n_nodes, which multiply: the "
"sample is their product. `measurement_batches` has no "
"ceiling: it hands the same samples over a batch at a time"
% (newdata.n_data, n_sim, n_nodes,
wanted_bytes / 1024 ** 3,
MEASUREMENT_LIMIT / 1024 ** 3))
chunks = {v: [] for v, _ in wanted}
for _, batch in self.measurement_batches(newdata, n_sim, n_nodes):
for v, values in batch.items():
chunks[v].append(values)
return {v: _np.concatenate(parts, axis=0)
for v, parts in chunks.items()}
def _measured_variables(self):
"""The variables a measurement can be described for, with their
likelihoods. A likelihood with no warping has none -- a categorical
one's noise lives in the probabilities -- so it is passed over
rather than asked."""
return [(v, lik) for v, lik in zip(self.variables, self.likelihoods)
if lik.warped]
[docs]
def measurement_batches(self, newdata: "_data._SpatialData",
n_sim: int = 20, n_nodes: int = 32,
where: _types.Where = None):
"""The measurement samples, a batch of locations at a time.
What :meth:`predict_measurements` returns whole, yielded in the
pieces it assembles it from. The samples are the largest thing this
model produces -- `n_sim * n_nodes` values a row for each column,
5 KB a row at the defaults -- and every statistic taken of them
(coverage, CRPS, the point errors, a PIT) reduces the sample axis
one row at a time, so a caller that accumulates as it goes never
holds more than a batch. That is what :func:`cross_validate` does,
and it is why this door carries no size ceiling where the other one
must.
Parameters
----------
newdata
Locations to ask about, as for :meth:`predict_measurements`.
n_sim
Latent realizations per location.
n_nodes
Equal-share noise values per realization.
where
One boolean per location, or the indices of the locations to
ask about; all of them by default. A location's sample does not
depend on which others are asked about with it.
Yields
------
rows : ndarray
Indices into `newdata` of the locations in this batch.
samples : dict of str to ndarray
One `(len(rows), size, n_sim * n_nodes)` array per variable
whose likelihood carries a warping, in the variable's own
units.
See Also
--------
predict_measurements : the same samples, assembled and returned.
"""
if newdata.rows_per_location != 1:
raise ValueError(
"a measurement is of a point, and %s fans each location out "
"into %d rows; ask this of the data the model was trained "
"from, or of a validation set"
% (type(newdata).__name__, newdata.rows_per_location))
if self.data.n_dim != newdata.n_dim:
raise ValueError("dimension of newdata is incompatible with model")
wanted = self._measured_variables()
# One uniform per location, component and realization, from a
# stream seeded by the model, so a location's sample is the same
# whatever batch computed it -- and drawn a batch at a time rather
# than held whole, this door having promised never to hold more
# than a batch. It rotates the equal-share noise nodes (see
# `_Likelihood._measurement_nodes`): without it every location in a
# column carried the same noise value, and the strata's midpoints
# carried 0.89 of a Laplace's noise variance in warped space, 35-47%
# of the `noise_variance` column through Jura's spline in data
# units, at the default n_nodes -- measured 2026-09-09,
# `docs/cross-validation.md`. The stream restarts from the seed on
# every call, so two containers of one size get the same uniforms
# row for row: nothing to a per-row score, and a caller reading
# samples jointly across calls should know.
def rotation(k, size, rows, stream=1):
"""The rows' slice of the k-th variable's `(n_data, size,
n_sim)` stream, drawn without generating what comes before it:
PCG64 spends exactly one 64-bit output per double, so advancing
by the rows' offset lands where a whole draw would. Stream 1
rotates the noise, stream 2 the latent jitter."""
first, last = int(rows[0]), int(rows[-1])
# the generator `default_rng` would build, named so that its
# `advance` is on the type
bits = _np.random.PCG64(
_np.random.SeedSequence([self.options.seed, stream, k]))
bits.advance(first * size * n_sim)
block = _np.random.Generator(bits).random(
(last - first + 1, size, n_sim))
return block[_np.asarray(rows) - first]
def batch_measure(x, x_var, n_splits, rows):
per_leaf, jitters, like = [], [], []
with _latent.simulation_rule(self.options.qmc_simulations):
for leaf in self.leaves:
predicted = leaf.predict(
x, x_var=x_var, n_sim=n_sim,
seed=[self.options.seed, 0])
per_leaf.append(_tf.transpose(predicted[2], [1, 0, 2]))
like.append(_tf.transpose(predicted[3]))
jitters.append(None if predicted.jitter is None
else _tf.transpose(predicted.jitter))
sims = self._by_likelihood(per_leaf)
jitters = self._jitter_by_likelihood(jitters, like)
measured = [(sim, jitter, lik) for sim, jitter, lik
in zip(sims, jitters, self.likelihoods) if lik.warped]
return [lik.measurement_samples(
sim, n_nodes,
shift=_tf.constant(rotation(k, lik.size, rows),
_tf.float64),
**({} if jitter is None else {
"jitter": jitter,
"jitter_shift": _tf.constant(
rotation(k, 1, rows, stream=2),
_tf.float64)}))
for k, (sim, jitter, lik) in enumerate(measured)]
for rows, output in self._over_batches(newdata, batch_measure,
where=where, with_rows=True):
# in the variable's own units, here rather than at the end: a
# composition's parts reach the model as fractions of the whole,
# and a streaming caller compares them against assays batch by
# batch, so the conversion cannot wait for an assembly that
# never happens
yield rows, {
v: self.data.variables[v].from_model_units(
_np.asarray(values))
for (v, _), values in zip(wanted, output)}
[docs]
def responsibilities(self, newdata: "_data._SpatialData",
store: bool = True) -> "dict[str, _np.ndarray]":
"""Posterior probability that each measurement came from each noise
component of a mixture likelihood.
One answer per location: a :class:`~geoml.likelihood.Mixture` is a
mixture over the row, so a measurement wrong in one component of a
vector variable is a wrong measurement.
Parameters
----------
newdata
Point data carrying the variables' measurements -- the training
data, a validation set, or the out-of-fold container
:func:`cross_validate` returns. Modified in place if `store`.
store
Whether to file the answer on each variable, as
`<variable>/responsibilities/<component>`, as well as return it.
Returns
-------
dict of str to ndarray
One `(n_data, n_components)` array per variable whose likelihood
is a mixture, rows summing to one, and missing at locations
without a measurement.
Raises
------
ValueError
If no variable has a mixture likelihood, if `newdata` lacks one
of those variables, or if it fans each location into several
rows, as a block model does.
Notes
-----
Read out of fold. At a training location the model interpolates its
own measurement, so an outlier is partly absorbed into the fit and
its responsibility understates it.
"""
if newdata.rows_per_location != 1:
raise ValueError(
"a measurement is of a point, and %s fans each location out "
"into %d rows; ask this of the data the model was trained "
"from, or of a validation set"
% (type(newdata).__name__, newdata.rows_per_location))
if self.data.n_dim != newdata.n_dim:
raise ValueError("dimension of newdata is incompatible with model")
mixtures = (_lk.Mixture, _lk.LikelihoodMixture)
wanted = [(v, lik) for v, lik in zip(self.variables, self.likelihoods)
if isinstance(lik, mixtures)]
if not wanted:
raise ValueError(
"no variable in this model has a mixture likelihood, and "
"responsibilities are a mixture's answer: with one noise "
"mechanism every measurement came from it")
absent = [v for v, _ in wanted if v not in newdata.variables.keys()]
if absent:
raise ValueError(
"responsibilities are of measurements, and %s carries no %s"
% (type(newdata).__name__, ", ".join(str(v) for v in absent)))
def batch_moments(x, x_var, n_splits):
mus, vars_ = [], []
for leaf in self.leaves:
mu, var, _, _ = leaf.predict(
x, x_var=x_var, n_sim=1, seed=[self.options.seed, 0])
mus.append(_tf.transpose(mu[:, :, 0]))
vars_.append(_tf.transpose(var))
mu = self._by_likelihood(mus)
var = self._by_likelihood(vars_)
return [(m, v) for m, v, lik
in zip(mu, var, self.likelihoods)
if isinstance(lik, mixtures)]
chunks = {v: ([], []) for v, _ in wanted}
for _, output in self._over_batches(newdata, batch_moments):
for (v, _), (mu, var) in zip(wanted, output):
chunks[v][0].append(_np.asarray(mu))
chunks[v][1].append(_np.asarray(var))
out = {}
for v, lik in wanted:
mu = _np.concatenate(chunks[v][0], axis=0)
var = _np.concatenate(chunks[v][1], axis=0)
y, has_value = newdata.variables[v].get_measurements()
answer = lik.responsibilities(mu, var, y)
answer[~_np.all(has_value == 1.0, axis=1)] = _np.nan
if store:
newdata.variables[v].set_responsibilities(answer)
out[v] = answer
return out
def _where_mask(container, where):
"""`where` as one boolean per location, or None for all of them.
A string names a boolean metadata column -- a stored filter, which is
what a block model has instead of being subsettable -- and indices
become the mask they select.
"""
if where is None:
return None
if isinstance(where, str):
mask = _np.asarray(container.get_metadata(where)).ravel().astype(bool)
else:
mask = _np.asarray(where)
if mask.dtype != bool:
index = _np.zeros(container.n_data, dtype=bool)
index[mask] = True
mask = index
if mask.shape != (container.n_data,):
raise ValueError(
"`where` needs one value per location: got %d for %d"
% (mask.size, container.n_data))
return mask
def _stored_n_sim(variable):
"""How many realizations `variable` holds, or None if it holds none.
A vector or categorical variable keeps them on its components."""
stores = [variable.simulations] \
if getattr(variable, "simulations", None) is not None else []
stores += [c.simulations
for c in getattr(variable, "components", {}).values()
if getattr(c, "simulations", None) is not None]
return stores[0].shape[1] if stores else None
def _simulation_count(container, names, n_sim, keep):
"""The number of realizations a prediction into `container` draws.
`None` takes the number the container's variables already hold, or 20.
Under `keep` -- some locations visited, the rest keeping what they
hold -- an existing store is written into, never reallocated, so a
different number is refused before any batch runs: it would otherwise
fail part way through, or with one realization be copied into every
column without a word.
"""
stored = {name: _stored_n_sim(container.variables[name])
for name in names if name in container.variables}
stored = {name: n for name, n in stored.items() if n is not None}
if n_sim is None:
return next(iter(stored.values()), 20)
if keep:
for name, count in stored.items():
if count != n_sim:
raise ValueError(
"%s holds %d realization(s) and the prediction was asked "
"for %d; predicting some locations writes into the "
"realizations the others keep, so the numbers must agree "
"(pass n_sim=None to take the stored one)"
% (name, count, n_sim))
return n_sim
[docs]
def refine(model, blocks: "_data.BlockSet3D", n_sim: "int | None" = None,
split_on: "str | Sequence[str] | None" = None,
tolerance: "float | None" = None, include_noise: bool = True,
where: _types.Where = None,
meshes: "Sequence[_data.Mesh3D] | None" = None,
verbose: bool = False,
by_expert: bool = False,
expert_options: "dict[str, _Any] | None" = None
) -> "_data.BlockSet3D":
"""Predict on a block model, cutting finer wherever it cannot decide.
Predicts on the coarse blocks, splits the ones still in doubt, predicts
only what the split created, and repeats. Three criteria mark a block for
splitting: the prediction at its sub-blocks falls on both sides of a
cut-off or a category boundary
(:meth:`~geoml.data.BlockSet3D.needs_splitting`); a neighbour is more
than one level finer than it (:meth:`~geoml.data.BlockSet3D.unbalanced`);
or a mesh passes through it (:meth:`~geoml.data.BlockSet3D.crossed_by`).
The loop ends by itself. Each pass takes the blocks it splits one level
finer, and no criterion marks a block already at the lattice's
`max_levels`, so within that many passes every block is either settled or
as fine as the block model was built to go.
Parameters
----------
model
A trained model with a `predict` method, normally a
:class:`VGPNetwork`.
blocks
The coarse block model to start from. Not modified: each pass builds
a new one.
n_sim
Realizations to draw at every pass. `None` takes the number `blocks`
already holds, or 20 where it holds none.
split_on
Which variables have a say in the decision. All of them by default.
tolerance
Deprecated, and without effect: a block is divided or it is not. It
was the share of realizations that had to find a block divided, and
will be removed.
include_noise
Passed to :meth:`VGPNetwork.predict`.
where
One boolean per block, the indices of the blocks worth modelling, or
the name of a boolean metadata column holding the same. The rest are
never predicted and never cut, and keep their missing values. Given
once, against the blocks as they stand: the mask is carried across
each split.
meshes
Surfaces and closed bodies whose crossings force a split. Costs a
side test per sub-block of every splittable block, each pass.
verbose
Print what each pass cut.
by_expert
Predict each pass with :meth:`VGPNetwork.predict_by_expert`, so
that memory does not grow with the number of experts.
expert_options
Keywords for :meth:`VGPNetwork.predict_by_expert` under
`by_expert` -- `coverage`, `neighbours`, `grouping`, `slots`,
`pack`.
Returns
-------
BlockSet3D
The refined model, predicted throughout.
Notes
-----
What decides a split carries no noise: the likelihood noise is
integrated out rather than drawn, so a block never straddles a cut-off on
account of spread that splitting cannot resolve.
Only the blocks a pass creates are predicted. A block that was not split
is the same block on the same support, and its value still stands.
Written out, the loop is three calls -- predict, ask, split -- which is
the way to stop part way and inspect a pass. See
`docs/variable-block-models.md`.
"""
keep = _where_mask(blocks, where)
if tolerance is not None:
warnings.warn(
"`tolerance` has no effect since 0.8.5 -- a block is divided "
"where the prediction's sub-blocks straddle a cut-off, which is "
"yes or no -- and will be removed", FutureWarning, stacklevel=2)
# No total: the passes are not bounded by `max_levels`, since each one
# marks a different set and a block still at level 0 can be marked by
# any later pass -- as the field sharpens around its neighbours,
# `unbalanced` reaches it. The loop stops when nothing is marked, and
# there is no honest count to promise before that.
with _progress.reporting("refine", None, "pass") as report:
return _refine_passes(model, blocks, n_sim, split_on,
include_noise, keep, meshes, verbose, report,
by_expert, expert_options)
def _refine_passes(model, blocks, n_sim, split_on, include_noise,
keep, meshes, verbose, report, by_expert=False,
expert_options=None):
"""The body of :func:`refine`, one pass at a time. Separate only so that
the reporting block can wrap a function that returns from its middle."""
predict = model.predict
if by_expert:
def predict(*args, **kwargs):
return model.predict_by_expert(*args, **kwargs,
**(expert_options or {}))
predict(blocks, n_sim=n_sim, include_noise=include_noise, where=keep)
report(0)
step = 0
while True:
undecided = blocks.needs_splitting(split_on)
uneven = blocks.unbalanced()
crossed = _np.zeros(blocks.n_data, dtype=bool)
for mesh in (meshes or []):
crossed |= blocks.crossed_by(mesh)
mask = undecided | uneven | crossed
if keep is not None:
# a block nobody asked for holds nothing to decide and draws no
# surface, so cutting it would only make more of nothing
mask = mask & keep
if not _np.any(mask):
return blocks
step += 1
# `split` keeps the unsplit blocks first, then each parent's children
# in sub-block order, which is how the mask follows them across
children = blocks.rows_per_location
blocks = blocks.split(mask)
visit = blocks.unpredicted()
if keep is not None:
keep = _np.concatenate(
[keep[~mask], _np.repeat(keep[mask], children)])
visit = visit & keep
predict(blocks, n_sim=n_sim, include_noise=include_noise,
where=visit)
report(step)
if verbose:
print("pass %d: cut %d block(s) (%d undecided, %d crossed by a "
"mesh, %d to level a jump), %d now"
% (step, int(_np.count_nonzero(mask)),
int(_np.count_nonzero(undecided)),
int(_np.count_nonzero(crossed & ~undecided)),
int(_np.count_nonzero(uneven & ~undecided & ~crossed)),
blocks.n_data))
# What encodes the data in a trained VGP, and how each piece starts over.
# Every node of the latent network registers these per inducing set (see
# `BasicGP._set_parameters`); everything else is a hyperparameter.
_VARIATIONAL_STATE = {
"alpha_white_": lambda shape: _srandom.rng().normal(scale=1e-3,
size=shape),
"delta_": _np.ones,
"bias_": _np.zeros,
}
def _terminal_gp_nodes(model):
"""The GP nodes nearest the likelihoods: from each leaf down through
operation nodes, stopping at the first GP node.
A leaf is often an operation node with no variational state of its own
-- `Linear(cat, size=2)` over a rock GP, `LinearCombination(trend,
metal_gp)` -- so the state that conditions on a likelihood's data is
found below it. A node can be terminal for one likelihood and interior
for another; it counts once.
"""
found = []
def visit(node):
if isinstance(node, _latent.network._GPNode):
if not any(node is seen for seen in found):
found.append(node)
elif isinstance(node, _latent.network._Operation):
for parent in node.parents:
visit(parent)
elif isinstance(node, _latent.network._FunctionalLatentVariable):
visit(node.parent)
# a root has no state to free
for leaf in model.leaves:
visit(leaf)
return found
def _fresh_variational_state(model, nodes=None):
"""Freeze what one fold cannot change; forget what it can.
The variational state -- `alpha_white_*`, `delta_*` and `bias_*` on every
node of the latent network -- is where the data lives in a trained VGP, so
a fold model gets it factory-fresh: re-initialized, it is structurally
ignorant of the held-out rows, and no iterations are spent forgetting
them. Everything else (kernels, warpings, likelihood noise) is fixed where
it stands -- the concession kriging cross-validation makes when it keeps
the fitted variogram, made once and said out loud in `cross_validate`'s
docstring. The fresh values are drawn from the package generator, so
`geoml.set_seed` makes the whole procedure reproducible.
`nodes` restricts the forgetting to those nodes (`refit="leaves"` passes
the terminal GP nodes); every other node's state is frozen with the rest.
"""
for parameter in model._all_parameters:
parameter.fix()
for node in model._nodes() if nodes is None else nodes:
for name, parameter in node.parameters.items():
for prefix, init in _VARIATIONAL_STATE.items():
if name.startswith(prefix):
shape = _np.asarray(parameter.get_value()).shape
parameter.set_value(init(shape))
parameter.unfix()
break
[docs]
def cross_validate(model: VGPNetwork, folds: str = "fold",
refit: str = "variational", iterations: int = 200,
method: str = "full", epochs: int = 50,
n_sim: int = 20, n_nodes: int = 32,
path: _types.PathLike | None = None,
expert_options: "dict[str, _Any] | None" = None
) -> "tuple[_data._SpatialData, _pd.DataFrame]":
"""Score a model on folds it never saw, with one short refit per fold.
The trained model is saved once and one fold model is rebuilt from the
file, around the data with the first fold removed; every later fold
swaps its own training rows into that same model, restores the file's
parameters and zeroes the optimizer's memory, all in place, so each
fold starts exactly where a reloaded model would while the graphs
traced for the first serve them all (a model rebuilt per fold leaves
its graph machinery resident for the life of the process -- see
Notes). Under `refit="variational"` the variational state -- the part
of a trained model that encodes the data -- is re-initialized and every
other parameter frozen, so the fold model starts ignorant of the
held-out rows, and only that state is refitted. The fold model then
predicts its held-out rows, and only those, into one shared copy of
the training data. Folds partition the data, so every location ends up
predicted by a model that never saw it.
Scores are of *measurements*: the held-out values are samples, so each
fold model is asked through :meth:`VGPNetwork.predict_measurements`.
Categorical variables get no rows -- subset the returned container by
fold and use the variable's own `compute_metrics`.
Parameters
----------
model
A trained model. It is saved and copied; the original is untouched.
folds
Name of the metadata column holding the fold labels, as
:meth:`~geoml.data.PointData.spatial_k_fold` writes. Any labelling
works: a hole-id column gives leave-one-hole-out.
refit
`"variational"` to refit the variational state alone, re-initialized
on every node; `"leaves"` to re-initialize and refit it on the
terminal GP nodes only -- the GP nearest each likelihood, found from
the leaf down through operation nodes -- keeping the interior (a
`GPWalk`'s field, a shared parent) as all the data taught it --
measured to score 20% past the honest reference on Jura, the
interior remembering the held-out rows, so a diagnostic of that
memory rather than a score (E2 in `docs/cross-validation.md`); or
`"all"` to warm-start every trainable parameter from its trained
value and continue on the reduced data.
iterations
Training iterations per fold, under `method="full"`. Ignored under
`method="svi"`, which counts in `epochs`.
method
`"full"` to refit each fold on all its data at every iteration, or
`"svi"` to refit in minibatches of
`options.training_batch_size`, or `"by_expert"` to refit with
:meth:`VGPNetwork.train_by_expert` and predict the held-out rows
with :meth:`VGPNetwork.predict_by_expert` -- the same choice as
:meth:`VGPNetwork.train_full` against
:meth:`VGPNetwork.train_svi`, and worth making for the same
reason: a fold refit costs the whole reduced data set per
iteration whatever is frozen, so a model too large to train
full-batch is too large to cross-validate that way.
epochs
Passes over each fold's data, under `method="svi"` or
`method="by_expert"`. A separate
argument from `iterations` because the two count different things:
an epoch is one visit to the data in batches, so it is many
gradient steps, and the numbers that make sense for one are wrong
for the other.
n_sim
Latent realizations, for the out-of-fold predictions and the
measurement samples alike.
n_nodes
Noise values per realization in the measurement samples.
path
Where to keep the saved model and its fold copies. A temporary
directory, removed at the end, unless one is given.
expert_options
Keywords for :meth:`VGPNetwork.train_by_expert` under
`method="by_expert"`.
Returns
-------
oof : container
A copy of the training data -- every variable and metadata column,
the fold labels included -- carrying out-of-fold predictions and
simulations, plus one metadata column per scored component
(`pit_<variable>`, or `pit_<variable>_<component>`) holding where
each measurement fell inside its own predictive distribution. It
saves with `to_zarr` and opens with `open` like any container.
scores : pandas.DataFrame
One row per scored component and fold, in the folds' sorted order,
then one per component pooled over every fold, whose `fold` is
`"all"`. The columns: `n`, the held-out measurements scored;
`rmse`, `mae` and `bias` of the mean of the measurement samples
against them, the bias being that mean minus the measurement;
`crps`, from the samples; `goodness`, of the samples' interval
coverage; `variable`; `component`, the variable's own name for a
scalar one; and `fold`, the fold's label.
See Also
--------
conformalize : calibrate interval widths on the PIT columns.
geoml.data.PointData.spatial_k_fold : folds that mimic a prediction task.
Notes
-----
The hyperparameters and the warping were fitted on all the data,
including each fold's -- the concession kriging cross-validation also
makes when it keeps the variogram fixed. Design record and measurements:
`docs/cross-validation.md`.
Under `method="svi"` the convergence rule, if `options.training_tolerance`
is set, reads one value an epoch -- the mean bound over its batches --
rather than one an iteration, so it needs a few epochs before it can
fire at all. On a short refit that is worth knowing: too few epochs and
the rule never speaks; the cap does the stopping.
Memory is flat across folds by construction. TensorFlow keeps the graph
machinery of a differentiated function resident after the function
dies, so a fold model built and dropped per fold cost 2.6 GB a fold on
a 5000-row copy of a real model and took a five-fold run on the full
data past a 62 GB machine; one model with its rows swapped costs one
model, and the same run measured 3.2 GB in total and 3.8x faster, the
rebuild, the retrace and any XLA compilation being paid once.
"""
data = model.data
labels = _np.asarray(data.get_metadata(folds))
if _pd.isna(labels).any():
raise ValueError(
"every location needs a fold; column '%s' has missing entries"
% folds)
fold_names = _np.unique(labels)
if fold_names.size < 2:
raise ValueError(
"cross-validation needs at least 2 folds; column '%s' holds %d"
% (folds, fold_names.size))
if refit not in ("variational", "leaves", "all"):
raise ValueError(
"refit must be 'variational', 'leaves' or 'all', got %r"
% (refit,))
if method not in ("full", "svi", "by_expert"):
raise ValueError(
"method must be 'full', 'svi' or 'by_expert', got %r"
% (method,))
cleanup = path is None
if cleanup:
path = _tempfile.mkdtemp(prefix="geoml_cv_")
saved = _os.path.join(path, "model")
# one shared answer sheet: every fold writes only the rows it held out,
# and the folds partition the data, so nothing stale survives the loop
oof = data[_np.ones(data.n_data, dtype=bool)]
for v in model.variables:
oof.variables[v].allocate_simulations(n_sim)
rows = []
acc = {}
pit = {}
with _progress.reporting("cross_validate", len(fold_names),
"fold") as report:
try:
_persistence.save_model(model, saved)
fold_model = None
for done, fold in enumerate(fold_names):
report(done)
held = labels == fold
if fold_model is None:
# Rebuilt from the file once, around the first fold's rows.
# Every later fold swaps its rows in, restores the file's
# parameters and zeroes the optimizer's memory, all in
# place, so it starts exactly where a reloaded model would
# while the graphs traced for the first fold serve it. A
# model rebuilt per fold left 2.6 GB of graph machinery
# behind each time (TensorFlow never returns it; see
# `_training_step`), which is what took a five-fold run on
# the Tom v6 model past the machine's memory.
fold_model = _cast(
VGPNetwork,
_persistence.load_model(saved, data=data[~held]))
value, shape, position, _, _ = \
fold_model.get_parameter_values(complete=True)
else:
fold_model._set_data(data[~held])
fold_model.update_parameters(value, shape, position)
fold_model._reset_optimizer()
if refit == "variational":
_fresh_variational_state(fold_model)
elif refit == "leaves":
_fresh_variational_state(
fold_model, nodes=_terminal_gp_nodes(fold_model))
# the batch size rides in the model's own options, so the fold
# copy already has whatever the original was trained with
if method == "svi":
fold_model.train_svi(epochs=epochs)
elif method == "by_expert":
fold_model.train_by_expert(epochs=epochs,
**(expert_options or {}))
else:
fold_model.train_full(max_iter=iterations)
if method == "by_expert":
fold_model.predict_by_expert(oof, n_sim, where=held)
else:
fold_model._predict(oof, n_sim, True, held,
check_measurements=False)
held_points = oof[held]
held_rows = _np.flatnonzero(held)
# the truth is small -- one value a row per column -- so
# it is read once for the fold and indexed per batch; only
# the samples are large enough to be worth streaming
truths = {v: _measured_truth(held_points.variables[v])
for v, _ in fold_model._measured_variables()}
# One pass over the fold's samples, a batch at a time, folding
# each into sufficient statistics -- counts and sums, never
# means of means, since batches differ in size. The samples are
# the only large thing here and this way none of them outlives
# its batch.
fold_acc = {}
for batch, samples in fold_model.measurement_batches(
held_points, n_sim=n_sim, n_nodes=n_nodes):
for v, sample in samples.items():
y_true, has_value, components = truths[v]
for c, component in enumerate(components):
measured = has_value[batch, c]
if not measured.any():
continue
truth = y_true[batch, c][measured]
draw = sample[measured, c, :]
_accumulate(fold_acc.setdefault(
(v, component), _fresh_scores()), truth, draw)
# what `conformalize` calibrates on
stored = pit.setdefault(
(v, component),
_np.full(data.n_data, _np.nan))
stored[held_rows[batch][measured]] = _pit(
draw, truth)
for key, a in fold_acc.items():
v, component = key
rows.append(dict(_scores_from(a), variable=v,
component=component, fold=fold))
pooled = acc.setdefault(key, _fresh_scores())
_merge(pooled, a)
report(len(fold_names))
finally:
if cleanup:
_shutil.rmtree(path, ignore_errors=True)
for (v, component), a in acc.items():
rows.append(dict(_scores_from(a), variable=v, component=component,
fold="all"))
for (v, component), column in pit.items():
oof.add_metadata(_pit_column(v, component), column)
return oof, _pd.DataFrame(rows)
def _fresh_scores():
"""An empty accumulator for one component's held-out scores.
Sufficient statistics rather than the scores themselves, so that a fold
read in batches and a pooling read fold by fold are the same arithmetic:
sums and a count, never a mean of means, which would weight a short
batch like a long one.
"""
return {"n": 0, "sse": 0.0, "sae": 0.0, "se": 0.0, "crps": 0.0,
"observed": None, "nominal": None}
def _accumulate(a, truth, draw):
"""Fold one batch of measured rows into an accumulator."""
point = draw.mean(axis=1)
n = int(truth.size)
nominal, observed = _metrics.coverage(truth, draw)
a["n"] += n
a["sse"] += float(((point - truth) ** 2).sum())
a["sae"] += float(_np.abs(point - truth).sum())
a["se"] += float((point - truth).sum())
a["crps"] += float(_metrics.crps(truth, draw)) * n
a["nominal"] = nominal
a["observed"] = observed * n if a["observed"] is None \
else a["observed"] + observed * n
def _merge(into, other):
"""Fold one accumulator into another -- a fold into the pooled row."""
for key in ("n", "sse", "sae", "se", "crps"):
into[key] += other[key]
into["nominal"] = other["nominal"]
into["observed"] = other["observed"] if into["observed"] is None \
else into["observed"] + other["observed"]
def _scores_from(a):
"""The reported scores of one accumulator."""
n = a["n"]
return {
"n": n,
"rmse": float(_np.sqrt(a["sse"] / n)),
"mae": a["sae"] / n,
"bias": a["se"] / n,
"crps": a["crps"] / n,
"goodness": _metrics.goodness(a["nominal"], a["observed"] / n),
}
def _measured_truth(variable):
"""`(values, measured, components)` of a variable's measurements: the
values `(n_data, n_columns)` in the variable's own units, as the
measurement samples come, a boolean per value, and a name per column.
`get_measurements` is the model's door, and hands a composition's parts
over as fractions of the whole; compared with samples in the parts'
own units, a composition declared in percent was scored against numbers
a hundred times too small.
"""
values, has_value = variable.get_measurements()
values = _np.asarray(variable.from_model_units(
_np.asarray(values, dtype=float)), dtype=float)
has_value = _np.asarray(has_value)
if values.ndim == 1:
values = values[:, None]
if has_value.ndim == 1:
has_value = has_value[:, None]
has_value = _np.broadcast_to(has_value == 1, values.shape)
parts = getattr(variable, "labels", None)
return values, has_value, [variable.name] if parts is None \
else list(parts)
def _pit(draws, truth):
"""Where each value of `truth` falls among its row of `draws`, as a
share -- mid-rank, so ties split evenly."""
return (draws < truth[:, None]).mean(axis=1) \
+ 0.5 * (draws == truth[:, None]).mean(axis=1)
def _pit_column(name, component):
"""The metadata column a component's PITs are kept in, by `predict`
and, out of fold, by `cross_validate`."""
if component is None or component == name:
return "pit_%s" % name
return "pit_%s_%s" % (name, component)
[docs]
class StructuralField(_GPModel):
"""Structural field modeling based on gradient data."""
def __init__(self, tangents, covariance, normals=None, mean_vector=None,
options=None):
super().__init__(options=options)
self.tangents = tangents
self.normals = normals
self.covariance = self._register(covariance)
self.covariance.set_limits(self.tangents)
if mean_vector is None:
# initialized as vertical
mean_vector = _np.zeros(self.tangents.n_dim)
mean_vector[-1] = 1
self.training_log = []
self.optimizer = _tf.keras.optimizers.Adam(
_tf.keras.optimizers.schedules.ExponentialDecay(1e-2, 1, 0.999),
amsgrad=True
)
self._add_parameter(
"mean_vector",
_gpr.UnitColumnNormParameter(
_np.array(mean_vector, ndmin=2).T,
- _np.ones([self.tangents.n_dim, 1]),
_np.ones([self.tangents.n_dim, 1]))
)
self._add_parameter("noise",
_gpr.PositiveParameter(1e-4, 1e-6, 10, fixed=True))
if self.normals is None:
# noise not used
self.parameters["noise"].fix()
# pre_computations
self._pre_computations.update({
"log_likelihood": _tf.Variable(_tf.constant(0.0, _tf.float64)),
})
# cleared here and filled by `refresh()`, as in `GP`
self.cov: _Any = None
self.cov_chol: _Any = None
self.cov_inv: _Any = None
self.scale: _Any = None
self.alpha: _Any = None
self.y: _Any = None
self.all_coordinates: _Any = None
self.all_directions: _Any = None
def __repr__(self):
s = "Gaussian process structural field model\n\n"
s += "Kernel:\n"
s += str(self.covariance)
return s
[docs]
def set_learning_rate(self, rate):
self.optimizer = _tf.keras.optimizers.Adam(
_tf.keras.optimizers.schedules.ExponentialDecay(rate, 1, 0.999),
amsgrad=True
)
[docs]
def refresh(self, jitter=1e-9):
with _tf.name_scope("structural_field_refresh"):
mean_vector = self.parameters["mean_vector"].get_value()
all_coordinates = self.tangents.coordinates
all_directions = self.tangents.directions
all_data = self.tangents.n_data
noise = _tf.zeros([self.tangents.n_data], dtype=_tf.float64)
y = _tf.zeros([self.tangents.n_data], dtype=_tf.float64)
if self.normals is not None:
all_coordinates = _np.concatenate([
all_coordinates, self.normals.coordinates
], axis=0)
all_directions = _np.concatenate([
all_directions, self.normals.directions
], axis=0)
all_data += self.normals.n_data
noise = _tf.concat([
noise,
_tf.ones([self.normals.n_data], dtype=_tf.float64)
], axis=0) * self.parameters["noise"].get_value()
y = _tf.concat([
y, _tf.ones([self.normals.n_data], dtype=_tf.float64)
], axis=0)
self.all_coordinates = _tf.constant(all_coordinates, _tf.float64)
self.all_directions = _tf.constant(all_directions, _tf.float64)
self.cov = self.covariance.self_covariance_matrix_d2(
self.all_coordinates, self.all_directions
)
self.scale = _tf.reduce_max(_tf.linalg.diag_part(self.cov))
self.cov = self.cov / self.scale
eye = _tf.eye(all_data, dtype=_tf.float64)
noise = _tf.linalg.diag(noise + jitter)
self.cov_chol = _tf.linalg.cholesky(self.cov + noise)
self.cov_inv = _tf.linalg.cholesky_solve(self.cov_chol, eye)
y = y[:, None] - _tf.matmul(all_directions, mean_vector)
# y = y / _tf.sqrt(self.scale)
self.alpha = _tf.matmul(self.cov_inv, y)
self.y = y
[docs]
@_tf.function
def log_likelihood(self, jitter=1e-9):
self.refresh(jitter)
with _tf.name_scope("structural_field_log_likelihood"):
fit = -0.5 * _tf.reduce_sum(self.y * self.alpha)
det = - _tf.reduce_sum(_tf.math.log(
_tf.linalg.diag_part(self.cov_chol)))
const = -0.5 * _tf.constant(
self.tangents.n_data * self.tangents.n_dim * _np.log(2*_np.pi),
_tf.float64)
log_lik = fit + det + const
self._pre_computations["log_likelihood"].assign(log_lik)
return log_lik
[docs]
def train(self, max_iter=1000):
model_variables = [pr.variable for pr in self._all_parameters
if not pr.fixed]
def loss():
return - self.log_likelihood(self.options.jitter)
for i in range(max_iter):
# self.optimizer.minimize(loss, model_variables)
_tftools.training_step(self.optimizer, loss, model_variables)
for pr in self._all_parameters:
pr.refresh()
current_log_lik = self._pre_computations["log_likelihood"].numpy()
self.training_log.append(current_log_lik)
if self.options.verbose:
print("\rIteration %s | Log-likelihood: %s" %
(str(i + 1), str(current_log_lik)), end="")
if self.options.verbose:
print("\n")
[docs]
@_tf.function
def predict_raw(self, x_new, jitter=1e-9):
self.refresh(jitter)
with _tf.name_scope("Prediction"):
mean_vector = self.parameters["mean_vector"].get_value()
# mean of field
cov_new = self.covariance.covariance_matrix_d1(
x_new, self.all_coordinates, self.all_directions) / self.scale
mu = _tf.matmul(cov_new, self.alpha)
mu = mu + _tf.matmul(x_new, mean_vector)
# variance of gradient along mean direction
cov_new = self.covariance.covariance_matrix_d2(
x_new, self.all_coordinates,
_tf.transpose(mean_vector), self.all_directions
) / self.scale
point_var = self.covariance.point_variance(x_new)[:, None]
explained_var = _tf.reduce_sum(
_tf.matmul(cov_new, self.cov_inv) * cov_new,
axis=1, keepdims=True)
var = _tf.maximum(point_var - explained_var, 0.0)
return mu, var
[docs]
@_tf.function
def predict_raw_directions(self, x_new, x_new_dir, jitter=1e-9):
self.refresh(jitter)
with _tf.name_scope("Prediction"):
mean_vector = self.parameters["mean_vector"].get_value()
cov_new = self.covariance.covariance_matrix_d2(
x_new, self.all_coordinates,
x_new_dir, self.all_directions) / self.scale
mu = _tf.matmul(cov_new, self.alpha)
mu = mu + _tf.matmul(x_new_dir, mean_vector)
point_var = self.covariance.point_variance(x_new)[:, None]
explained_var = _tf.reduce_sum(
_tf.matmul(cov_new, self.cov_inv) * cov_new,
axis=1, keepdims=True)
var = _tf.maximum(point_var - explained_var, 0.0)
return mu, var
[docs]
def predict(self, newdata, variable):
"""
Makes a prediction on the specified coordinates.
Parameters
----------
newdata :
A reference to a spatial points object of compatible dimension.
The object's variables are updated.
variable : str
Name of output variable.
"""
if self.tangents.n_dim != newdata.n_dim:
raise ValueError("dimension of newdata is incompatible with model")
# managing variables
newdata.add_continuous_variable(variable)
# prediction in batches
batch_id = self.options.batch_index(newdata.n_data)
n_batches = len(batch_id)
for i, batch in enumerate(batch_id):
if self.options.verbose:
print("\rProcessing batch %s of %s "
% (str(i + 1), str(n_batches)), end="")
if isinstance(newdata, _data.DirectionalData):
mu, var = self.predict_raw_directions(
_tf.constant(newdata.coordinates[batch], _tf.float64),
_tf.constant(newdata.directions[batch], _tf.float64),
jitter=self.options.jitter
)
else:
mu, var = self.predict_raw(
_tf.constant(newdata.coordinates[batch], _tf.float64),
jitter=self.options.jitter)
# `update` names the columns the way the likelihoods fill them:
# `average_sim` is the prediction and `mean`/`variance` the
# latent pair. This model has no likelihood and no warping to
# separate them -- the potential field is the latent value --
# so the prediction is the mean itself.
output = {"average_sim": _tf.squeeze(mu),
"mean": _tf.squeeze(mu),
"variance": _tf.squeeze(var)}
newdata.variables[variable].update(batch, **output)
if self.options.verbose:
print("\n")
class _EnsembleModel(_GPModel):
def __init__(self, options=None):
super().__init__(options)
self.models = []
self.variable = None
def set_learning_rate(self, rate):
for model in self.models:
model.set_learning_rate(rate)
def predict(self, newdata):
"""
Makes a prediction on the specified coordinates.
Parameters
----------
newdata :
A reference to a spatial points object of compatible dimension.
The object's variables are updated.
"""
if self.n_dim != newdata.n_dim:
raise ValueError("dimension of newdata is incompatible with model")
if self.variable not in newdata.variables.keys():
self.models[0].data.variables[self.variable].copy_to(newdata)
prediction_input = self.models[0].data \
.variables[self.variable].prediction_input()
# prediction in batches
batch_id = self.options.batch_index(newdata.n_data,
self.options.prediction_batch_size)
n_batches = len(batch_id)
for i, batch in enumerate(batch_id):
if self.options.verbose:
print("\rProcessing batch %s of %s "
% (str(i + 1), str(n_batches)), end="")
outputs = [model.predict_raw(
_tf.constant(newdata.coordinates[batch], _tf.float64),
jitter=self.options.jitter, **prediction_input)
for model in self.models]
output = self.combine(outputs)
newdata.variables[self.variable].update(batch, **output)
if self.options.verbose:
print("\n")
@staticmethod
def combine(outputs):
weights = _tf.stack([out["weights"] for out in outputs], axis=1)
# weights = 1 / (1 - weights + 1e-6)
weights = weights / _tf.reduce_sum(weights, axis=1, keepdims=True)
mu = _tf.stack([out["mean"] for out in outputs], axis=1)
mu = _tf.reduce_sum(weights * mu, axis=1)
var = _tf.stack([out["variance"] for out in outputs], axis=1)
var = _tf.reduce_sum(weights ** 2 * var, axis=1)
combined = {"mean": mu, "variance": var}
if "probabilities" in outputs[0].keys():
prob = _tf.stack([out["probabilities"] for out in outputs], axis=2)
prob = _tf.reduce_sum(weights[:, None, :] * prob, axis=2)
combined["probabilities"] = prob
if "quantiles" in outputs[0].keys():
quant = _tf.stack([out["quantiles"] for out in outputs], axis=2)
quant = _tf.reduce_sum(weights[:, None, :] * quant, axis=2)
combined["quantiles"] = quant
return combined
[docs]
class GPEnsemble(_EnsembleModel):
"""An ensemble of Gaussian processes."""
[docs]
def __init__(self, data, variable, covariance, warping=None, directional_data=None,
use_trend=False, options=None):
"""
An ensemble of Gaussian processes.
This model combines independent GPs into a consolidated prediction using the Product of Experts approach.
It is preferable to divide the data spatially instead of randomly, so that each expert can focus on a
specific region of the space.
Parameters
----------
data
A list or tuple of `PointData` objects.
variable : str
The name of the variable to be modelled. Must be present in all data objects.
covariance
The covariance function to build the covariance matrices.
warping
An object from the `warping` module. If None, the data is assumed to have zero mean and unit variance.
directional_data
A `DirectionalData` object from the ´data´ module. The corresponding variable will be used as the gradient
of the modelled field.
use_trend : bool
If `True`, will model a linear trend in the data in addition to the GP.
options : GPOptions
Additional configurations.
"""
super().__init__(options)
if not isinstance(data, (tuple, list)):
raise ValueError("data must be a list or tuple containing"
"data objects")
if directional_data is None:
directional_data = [None for _ in data]
elif not isinstance(directional_data, (tuple, list)):
raise ValueError("directional_data must be a list or tuple containing"
"data objects or None")
dims = set([d.n_dim for d in data])
if len(dims) != 1:
raise Exception("all data objects must have the same dimension")
self._n_dim = list(dims)[0]
self.models = [GP(
data=d,
variable=variable,
covariance=_copy.deepcopy(covariance),
warping=_copy.deepcopy(warping),
directional_data=dd,
use_trend=use_trend,
options=options)
for d, dd in zip(data, directional_data)]
for model in self.models:
self._register(model)
self.variable = variable
def __repr__(self):
s = "Gaussian process ensemble\n\n" \
"Models: %d\n\n" % len(self.models)
s += "Variable: " + self.variable + "\n"
return s
[docs]
def train(self, max_iter=1000):
for i, model in enumerate(self.models):
if self.options.verbose:
print("Training model %d of %d" % (i + 1, len(self.models)))
model.train(max_iter=max_iter)
# class VGPNetworkEnsemble(_EnsembleModel):
# def __init__(self, data, variables, likelihoods, latent_networks,
# directional_data=None, options=GPOptions()):
# super().__init__(options)
# if not isinstance(data, (tuple, list)):
# raise ValueError("data must be a list or tuple containing"
# "data objects")
# if not isinstance(latent_networks, (tuple, list)):
# raise ValueError("latent_trees must be a list or tuple containing"
# "latent variable objects")
# if directional_data is None:
# directional_data = [None for _ in data]
# elif not isinstance(directional_data, (tuple, list)):
# raise ValueError("directional_data must be a list or tuple"
# "containing data objects or None")
#
# dims = set([d.n_dim for d in data])
# if len(dims) != 1:
# raise Exception("all data objects must have the same dimension")
# self._n_dim = list(dims)[0]
#
# self.models = [VGPNetwork(
# data=d,
# variables=variables,
# likelihoods=_copy.deepcopy(likelihoods),
# # likelihoods=lik,
# latent_network=l,
# directional_data=dd,
# options=options)
# # for d, l, dd, lik in zip(
# # data, latent_networks, directional_data, likelihoods)
# for d, l, dd in zip(
# data, latent_networks, directional_data)
# ]
# for model in self.models:
# self._register(model)
#
# if not (isinstance(variables, (list, tuple))):
# variables = [variables]
# self.variables = variables
#
# def __repr__(self):
# s = "Gaussian process ensemble\n\n" \
# "Models: %d\n\n" % len(self.models)
# s += "Variables:\n "
# for v, lik in zip(self.variables, self.models[0].likelihoods):
# s += "\t" + v + " (" + lik.__class__.__name__ + ")\n"
# return s
#
# # def train_full(self, cycles=10, max_iter_per_model=100):
# # for c in range(cycles):
# # for i, model in enumerate(self.models):
# # if self.options.verbose:
# # print("Cycle %d of %d - training model %d of %d" %
# # (c + 1, cycles, i + 1, len(self.models)))
# #
# # model.train_full(max_iter=max_iter_per_model)
# #
# # def train_svi(self, cycles=10, epochs_per_model=10):
# # for c in range(cycles):
# # for i, model in enumerate(self.models):
# # if self.options.verbose:
# # print("Cycle %d of %d - training model %d of %d" %
# # (c + 1, cycles, i + 1, len(self.models)))
# #
# # model.train_svi(epochs=epochs_per_model)
#
# def train_full(self, max_iter=1000):
# for model in self.models:
# model.train_full(max_iter)
#
# def train_svi(self, epochs=100):
# for model in self.models:
# model.train_svi(epochs)
#
# def predict(self, newdata, n_sim=20):
# """
# Makes a prediction on the specified coordinates.
#
# Parameters
# ----------
# newdata :
# A reference to a spatial points object of compatible dimension.
# The object's variables are updated.
# n_sim : int
# Number of predictive samples to draw.
# """
# if self.n_dim != newdata.n_dim:
# raise ValueError("dimension of newdata is incompatible with model")
#
# # managing variables
# variable_inputs = []
# for v in self.variables:
# if v not in newdata.variables.keys():
# self.models[0].data.variables[v].copy_to(newdata)
# newdata.variables[v].allocate_simulations(n_sim)
# variable_inputs.append(
# self.models[0].data.variables[v].prediction_input())
#
# # prediction in batches
# batch_id = self.options.batch_index(
# newdata.n_data, batch_size=self.options.prediction_batch_size)
# n_batches = len(batch_id)
#
# def batch_pred(model, x):
# out = model.predict_raw(
# x,
# variable_inputs,
# seed=self.options.seed,
# n_sim=n_sim,
# jitter=self.options.jitter
# )
# return out
#
# @_tf.function
# def combined_pred(x):
# outputs = [batch_pred(model, x) for model in self.models]
# return self.combine(outputs)
#
# for i, batch in enumerate(batch_id):
# if self.options.verbose:
# print("\rProcessing batch %s of %s "
# % (str(i + 1), str(n_batches)), end="")
#
# # outputs = [batch_pred(
# # model,
# # _tf.constant(newdata.coordinates[batch], _tf.float64))
# # for model in self.models]
# #
# # output = self.combine(outputs)
#
# output = combined_pred(
# _tf.constant(newdata.coordinates[batch], _tf.float64))
#
# for v, upd in zip(self.variables, output):
# newdata.variables[v].update(batch, **upd)
#
# if self.options.verbose:
# print("\n")
#
# @_tf.function
# def combine(self, outputs):
# combined = [{} for _ in self.variables]
# for i, variable in enumerate(self.variables):
# var_keys = outputs[0][i].keys()
#
# weights = _tf.stack([out[i]["weights"] for out in outputs], axis=1)
# weights = weights + 1e-6
# weights = weights / _tf.reduce_sum(weights, axis=1, keepdims=True)
#
# for key in var_keys:
# if key != "weights":
# tensor = _tf.stack([out[i][key] for out in outputs], axis=1)
# if "variance" in key:
# w = weights**2
# else:
# w = weights
#
# w = _tf.cond(_tf.greater_equal(_tf.rank(tensor), 3),
# lambda: _tf.expand_dims(w, axis=-1),
# lambda: w)
# w = _tf.cond(_tf.greater_equal(_tf.rank(tensor), 4),
# lambda: _tf.expand_dims(w, axis=-1),
# lambda: w)
#
# tensor = _tf.reduce_sum(w * tensor, axis=1)
# combined[i][key] = tensor
#
# return combined
[docs]
class Normalizer(_GPModel):
"""Trainable data normalizer."""
[docs]
def __init__(self, warping, options=None):
"""
Trainable data normalizer.
This model will fit a `warping` object to a data vector, allowing its transformation to a Gaussian
distribution with zero mean and unit variance.
Parameters
----------
warping
An object from the `warping` module.
options : GPOptions
Additional configurations.
"""
super().__init__(options)
self.warping = self._register(warping)
self.training_log = []
# self.optimizer = _tf.keras.optimizers.Nadam(
# _tf.keras.optimizers.schedules.ExponentialDecay(1e-1, 1, 0.99),
# )
self.optimizer = _tf.keras.optimizers.Adam(
_tf.keras.optimizers.schedules.ExponentialDecay(1e-1, 1, 0.999),
amsgrad=True#, clipvalue=0.001
)
self.objective = _tf.Variable(_tf.constant(0.0, _tf.float64))
[docs]
def normalize(self, x, max_iter=250):
"""
Model training.
Parameters
----------
x : array-like
The data vector to train on.
max_iter : int
The number of iterations to run.
"""
if len(self.all_parameters) == 0:
warnings.warn("No trainable parameters.")
return None
model_variables = self.get_unfixed_variables()
if len(model_variables) == 0:
warnings.warn("All parameters are fixed. Unfix one or more"
"parameters to continue.")
return None
self.warping.initialize(x)
def loss():
x_warp = self.warping.forward(x)
mean = _tf.reduce_mean(x_warp, axis=0, keepdims=True)
var = _tf.math.reduce_variance(x_warp, axis=0, keepdims=True)
std = _tf.sqrt(var)
log_derivative = self.warping.log_derivative(x)
# kl = _tf.reduce_sum(_tf.math.log(std) + (1 + mean ** 2) / (2 * var) - 0.5)
kl = _tf.reduce_sum(var + mean**2 - 1 - _tf.math.log(var)) * 0.5
density = _tf.reduce_mean(_tf.reduce_sum(-x_warp**2 / 2, axis=1, keepdims=True) + log_derivative)
obj = density #- kl
self.objective.assign(obj)
return -obj
for i in range(max_iter):
# self.optimizer.minimize(loss, model_variables)
_tftools.training_step(self.optimizer, loss, model_variables)
for pr in self._all_parameters:
pr.refresh()
current_elbo = self.objective.numpy()
self.training_log.append(current_elbo)
if self.options.verbose:
print("\rIteration %s | Objective: %s" %
(str(i + 1), str(current_elbo)), end="")
class ProjectedVGP(VGPNetwork):
@_tf.function
def _log_lik(self, x, y, has_value, training_inputs, x_var=None,
samples=20, seed=0):
with _tf.name_scope("batched_elbo"):
# prediction, per leaf
per_leaf = [leaf.predict(x, n_sim=samples, seed=[seed, 0])
for leaf in self.leaves]
sims = self._by_likelihood(per_leaf)
mu = [s[:, :, 0] for s in sims] # dummy
var = [s[:, :, 0] ** 2 for s in sims] # dummy
# likelihood
y_s = _tf.split(y, self.var_lengths, axis=1)
hv = _tf.split(has_value, self.var_lengths, axis=1)
elbo = _tf.constant(0.0, _tf.float64)
for likelihood, mu_i, var_i, y_i, hv_i, sim_i, inp in zip(
self.likelihoods, mu, var, y_s,
hv, sims, training_inputs):
elbo = elbo + likelihood.log_lik(
mu_i, var_i, y_i, hv_i, samples=sim_i, **inp)
# batch weight
batch_size = _tf.reduce_sum(has_value)
elbo = elbo * self.total_data / batch_size
return elbo
def predict_node(self, node, newdata, n_sim=None, name=None,
labels=None, where=None):
# a projected node's `predict` returns realizations alone
raise NotImplementedError(
"predict_node is for inducing-point networks; a ProjectedVGP's "
"nodes report realizations only")
def _predict_raw(self, x_new, variable_inputs, x_var=None,
n_sim=1, seed=0, jitter=1e-6, include_noise=True):
# `predict_raw` is inherited: it compiles this the way the options ask.
self._refresh(jitter)
with _tf.name_scope("Prediction"):
per_leaf = [leaf.predict(x_new, n_sim=n_sim, seed=[seed, 0])
for leaf in self.leaves]
pred_sim = self._by_likelihood(per_leaf)
pred_mu = [s[:, :, 0] for s in pred_sim] # dummy
pred_var = [s[:, :, 0] ** 2 for s in pred_sim] # dummy
pred_exp_var = pred_var # dummy
output = []
for mu, var, sim, exp_var, lik, v_inp in zip(
pred_mu, pred_var, pred_sim, pred_exp_var,
self.likelihoods, variable_inputs):
output.append(lik.predict(mu, var, sim, exp_var, include_noise=include_noise, **v_inp))
return output
def search_throw(model: "VGPNetwork", fault: "_gpr.Parametric",
throws: _types.ArrayLike, iterations: int = 50,
learning_rate: float = 0.1, parameter: str = "throw",
refine: int = 0) -> "tuple[float, list]":
"""
Chooses a fault's throw among candidates, on the bound.
A fault's throw does not always train from zero: the bound is flat
in it away from the right value, with a categorical likelihood or with
several faults it was measured not to move at all. This is the search
that gradient descent cannot do. For each candidate the model is
restored to the state it started in, the throw set, a short burst of
training run with a fresh optimizer, and the bound at its end recorded;
the best candidate is set on the restored model, which is then trained
as usual. A workflow rather than something a model is, like `refine`
and `cross_validate`.
Parameters
----------
model
The model, built and possibly trained; left in its starting state
with the chosen throw set and a fresh optimizer at
`learning_rate`.
fault
The `FaultDisplacement` (any transform holding the parameter).
throws
The candidate values.
iterations
Training iterations per candidate.
learning_rate
The rate of the bursts, and of the optimizer the model is left
with.
parameter
The parameter searched, `"throw"` or `"strike_slip"`.
refine
Further passes, each over five candidates spanning one step of the
previous pass either side of its best.
Returns
-------
best : float
The candidate chosen.
passes : list of (candidates, scores)
One pair of arrays per pass, the score being the bound each
candidate reached, mean of its burst's last fifth.
"""
candidates = _np.asarray(throws, dtype=float).ravel()
value, shape, position, _, _ = model.get_parameter_values(complete=True)
start = len(model.training_log)
passes = []
best = float(candidates[0])
for level in range(refine + 1):
scores = []
for throw in candidates:
model.update_parameters(value, shape, position)
fault.parameters[parameter].set_value(float(throw))
model.set_learning_rate(learning_rate)
model.train_full(max_iter=iterations)
burst = model.training_log[-max(1, iterations // 5):]
scores.append(float(_np.mean(burst)))
scores = _np.asarray(scores)
passes.append((candidates, scores))
best = float(candidates[int(_np.nanargmax(scores))])
if level < refine:
spacing = (candidates.max() - candidates.min()) \
/ max(len(candidates) - 1, 1)
candidates = _np.linspace(best - spacing, best + spacing, 5)
del model.training_log[start:]
model.update_parameters(value, shape, position)
fault.parameters[parameter].set_value(best)
model.set_learning_rate(learning_rate)
return best, passes