def func_jacobian(
self,
variable_dict: Dict[Variable, np.ndarray],
variables: Optional[Tuple[Variable, ...]] = None,
axis: Axis = False,
**kwargs,
) -> Tuple[FactorValue, JacobianValue]:
"""
Call this factor with a set of arguments
Parameters
----------
args
Positional arguments for the underlying factor
kwargs
Keyword arguments for the underlying factor
Returns
-------
An object encapsulating the value for the factor
"""
if variables is None:
variables = self.variables
variable_names = tuple(self._variable_name_kw[v.name]
for v in variables)
kwargs = self.resolve_variable_dict(variable_dict)
vals, *jacs = self._call_factor(kwargs, variables=variable_names)
shift, shape = self._function_shape(**kwargs)
plate_dim = dict(zip(self.plates, shape[shift:]))
det_shapes = {
v: shape[:shift] + tuple(plate_dim[p] for p in v.plates)
for v in self.deterministic_variables
}
var_shapes = {self._kwargs[v]: np.shape(x) for v, x in kwargs.items()}
if not (isinstance(vals, tuple) or self.n_deterministic > 1):
vals = vals,
log_val = (0. if (shape == () or axis is None) else aggregate(
np.zeros(tuple(1 for _ in shape)), axis))
det_vals = {
k: np.reshape(val, det_shapes[k]) if det_shapes[k] else val
for k, val in zip(self._deterministic_variables, vals)
}
fval = FactorValue(log_val, det_vals)
vjacs = {}
for k, _jacs in zip(self._deterministic_variables, jacs):
for v, jac in zip(variables, _jacs):
vjacs.setdefault(v, {})[k] = np.reshape(
jac, det_shapes[k] + var_shapes[v][shift:])
fjac = {
v: FactorValue(np.zeros(np.shape(log_val) + var_shapes[v]),
vjacs[v])
for v in variables
}
return fval, fjac
def func_jacobian(self, variable_dict: Dict[Variable, np.ndarray],
**kwargs) -> FactorValue:
# generate set of factors to call, these are indexed by the
# missing deterministic variables that need to be calculated
log_value = 0.0
det_values = {}
factor_jacs = {}
variables = {**self.fixed_values, **variable_dict}
missing = set(v.name
for v in self.variables).difference(v.name
for v in variables)
if missing:
n_miss = len(missing)
missing_str = ", ".join(missing)
raise ValueError(
f"{self} missing {n_miss} arguments: {missing_str}"
f"factor graph call signature: {self.call_signature}")
for calls in self._call_sequence:
# TODO parallelise this part?
for factor in calls:
ret, factor_jacs[factor] = factor.func_jacobian(
variables, **kwargs)
log_value += ret
det_values.update(ret.deterministic_values)
variables.update(ret.deterministic_values)
jac_out = (FactorValue, ) + tuple(det_values)
jac_variables = tuple(variables)
def graph_vjp(args):
args = args if isinstance(args, tuple) else (args, )
grads = {}
grads.update(zip(jac_out, args))
for calls in self._call_sequence[::-1]:
for factor in calls:
factor_grad = factor_jacs[factor](grads)
for v, v_grad in factor_grad.items():
grads[v] = grads.get(v, 0) + v_grad
return tuple(grads[v] for v in jac_variables)
fval = FactorValue(log_value, det_values)
graph_vjp = VectorJacobianProduct(
jac_out,
graph_vjp,
*jac_variables,
out_shapes=fval.to_dict().shapes,
)
return fval, graph_vjp
def __call__(
self,
variable_dict: Dict[Variable, np.ndarray],
axis: Axis = False,
) -> FactorValue:
values = self.resolve_variable_dict(variable_dict)
val = self._call_factor(values, variables=None)
val = aggregate(val, axis)
return FactorValue(val, {})
def _factor_value(self, raw_fval) -> FactorValue:
"""Converts the raw output of the factor into a `FactorValue`
where the values of the deterministic values are stored in a dict
attribute `FactorValue.deterministic_values`
"""
det_values = VariableData(
nested_filter(is_variable, self.factor_out, raw_fval))
fval = det_values.pop(FactorValue, 0.0)
return FactorValue(fval, det_values)
def func_jacobian(self,
variable_dict: Dict[Variable, np.ndarray],
variables: Optional[Tuple[Variable, ...]] = None,
axis: Axis = False,
**kwargs) -> Tuple[FactorValue, JacobianValue]:
"""
Call the underlying factor
Parameters
----------
variable_dict :
the values to call the function with
variables : tuple of Variables
the variables to calculate gradients and Jacobians for
if variables = None then differentiates wrt all variables
axis : None or False or int or tuple of ints, optional
the axes to reduce the result over.
if axis = None the sum over all dimensions
if axis = False then does not reduce result
Returns
-------
FactorValue, JacobianValue
encapsulating the result of the function call
"""
if variables is None:
variables = self.variables
variable_names = tuple(self._variable_name_kw[v.name]
for v in variables)
kwargs = self.resolve_variable_dict(variable_dict)
val, jacs = self._call_factor(kwargs, variables=variable_names)
grad_axis = tuple(range(np.ndim(val))) if axis is None else axis
fval = FactorValue(aggregate(self._reshape_factor(val, kwargs), axis))
fjac = {
v: FactorValue(aggregate(jac, grad_axis))
for v, jac in zip(variables, jacs)
}
return fval, fjac
def __call__(
self,
variable_dict: Dict[Variable, np.ndarray],
axis: Axis = False,
) -> FactorValue:
"""
Call each function in the graph in the correct order, adding the logarithmic results.
Deterministic values computed in initial factor calls are added to a dictionary and
passed to subsequent factor calls.
Parameters
----------
variable_dict
Positional arguments
axis
Keyword arguments
Returns
-------
Object comprising the log value of the computation and a dictionary containing
the values of deterministic variables.
"""
# generate set of factors to call, these are indexed by the
# missing deterministic variables that need to be calculated
log_value = 0.
det_values = {}
variables = variable_dict.copy()
missing = set(v.name
for v in self.variables).difference(v.name
for v in variables)
if missing:
n_miss = len(missing)
missing_str = ", ".join(missing)
raise ValueError(
f"{self} missing {n_miss} arguments: {missing_str}"
f"factor graph call signature: {self.call_signature}")
for calls in self._call_sequence:
# TODO parallelise this part?
for factor in calls:
ret = factor(variables)
ret_value = self.broadcast_plates(factor.plates, ret.log_value)
log_value = add_arrays(log_value, aggregate(ret_value, axis))
det_values.update(ret.deterministic_values)
variables.update(ret.deterministic_values)
return FactorValue(log_value, det_values)
def __call__(
self,
variable_dict: Dict[Variable, np.ndarray],
axis: Axis = False,
# **kwargs: np.ndarray
) -> FactorValue:
"""
Call this factor with a set of arguments
Parameters
----------
args
Positional arguments for the underlying factor
kwargs
Keyword arguments for the underlying factor
Returns
-------
An object encapsulating the value for the factor
"""
kwargs = self.resolve_variable_dict(variable_dict)
res = self._call_factor(**kwargs)
shift, shape = self._function_shape(**kwargs)
plate_dim = dict(zip(self.plates, shape[shift:]))
det_shapes = {
v: shape[:shift] + tuple(
plate_dim[p] for p in v.plates)
for v in self.deterministic_variables
}
if not (isinstance(res, tuple) or self.n_deterministic > 1):
res = res,
log_val = (
0. if (shape == () or axis is None) else
aggregate(np.zeros(tuple(1 for _ in shape)), axis))
det_vals = {
k: np.reshape(val, det_shapes[k])
if det_shapes[k]
else val
for k, val
in zip(self._deterministic_variables, res)
}
return FactorValue(log_val, det_vals)
def __call__(
self,
variable_dict: Dict[Variable, np.ndarray],
axis: Axis = False,
# **kwargs: np.ndarray
) -> FactorValue:
"""
Call the underlying factor
Parameters
----------
args
Positional arguments for the factor
kwargs
Keyword arguments for the factor
Returns
-------
Object encapsulating the result of the function call
"""
kwargs = self.resolve_variable_dict(variable_dict)
val = self._call_factor(**kwargs)
val = aggregate(self._reshape_factor(val, kwargs), axis)
return FactorValue(val, {})
def numerical_func_jacobian(
factor: "AbstractNode",
values: Dict[Variable, np.array],
variables: Optional[Tuple[Variable, ...]] = None,
axis: Axis = False,
_eps: float = 1e-6,
_calc_deterministic: bool = True,
) -> Tuple[FactorValue, JacobianValue]:
"""Calculates the numerical Jacobian of the passed factor
the arguments passed correspond to the variable that we want
to take the derivatives for
the values must be passed as keywords
_eps specifies the numerical accuracy of the numerical derivative
returns a jac = JacobianValue namedtuple
jac.log_value stores the Jacobian of the factor output as a dictionary
where the keys are the names of the variables
jac.determinisic_values store the Jacobian of the deterministic variables
where the keys are a Tuple[str, str] pair of the variable and deterministic
variables
Example
-------
>>> import numpy as np
>>> x_ = Variable('x')
>>> y_ = Variable('y')
>>> A = np.arange(4).reshape(2, 2)
>>> dot = factor(lambda x: A.dot(x))(x_) == y_
>>> dot.jacobian([x_], {x_: [1, 2]})
JacobianValue(
log_value={'x': array([[0.], [0.]])},
deterministic_values={('x', 'y'): array([[0., 2.], [1., 3.]])})
"""
if variables is None:
variables = factor.variables
# copy the input array
p0 = {v: np.array(x, dtype=float) for v, x in values.items()}
f0 = factor(p0, axis=axis)
log_f0 = f0.log_value
det_vars0 = f0.deterministic_values
fjac = {
v: FactorValue(np.empty(np.shape(log_f0) + np.shape(values[v])))
for v in variables
}
if _calc_deterministic:
for v, grad in fjac.items():
grad.deterministic_values = {
det: np.empty(np.shape(val) + np.shape(values[v]))
for det, val in det_vars0.items()
}
det_slices = {
v: (slice(None), ) * np.ndim(a)
for v, a in det_vars0.items()
}
for v, grad in fjac.items():
x0 = p0[v]
v_jac = grad.deterministic_values
if x0.shape:
inds = tuple(a.ravel() for a in np.indices(x0.shape))
i0 = tuple(slice(None) for _ in range(np.ndim(log_f0)))
for ind in zip(*inds):
x0[ind] += _eps
p0[v] = x0
f = factor(p0, axis=axis)
x0[ind] -= _eps
# print(ind)
grad[i0 + ind] = (f - f0) / _eps
if _calc_deterministic:
det_vars = f.deterministic_values
for det, val in det_vars.items():
v_jac[det][det_slices[det] + ind] = \
(val - det_vars0[det]) / _eps
else:
p0[v] += _eps
f = factor(p0, axis=axis)
p0[v] -= _eps
grad.itemset((f - f0) / _eps)
if _calc_deterministic:
det_vars = f.deterministic_values
for det, val in det_vars.items():
v_jac[det] = (val - det_vars0[det]) / _eps
return f0, fjac