Merge pull request #709 from rhayes777/features/nested

Features/nested

Author: matthewghgriffiths

autofit/graphical/factor_graphs/abstract.py

@@ -18,6 +18,10 @@
     FlattenArrays,
     nested_filter,
     nested_update,
+    nested_zip,
+    nested_set, 
+    nested_map, 
+    nested_items, 
     is_variable,
     Status, 
 )
@@ -92,7 +96,7 @@ def resolve_variable_dict(
     def resolve_args(
             self, values: Dict[Variable, np.ndarray]
     ) -> Tuple[np.ndarray, ...]:
-        return (values[k] for k in self.args)
+        return nested_update(self.args, values)
 
     @cached_property
     def fixed_values(self) -> VariableData:
@@ -103,7 +107,7 @@ def variables(self) -> Set[Variable]:
         """
         Dictionary mapping the names of variables to those variables
         """
-        return frozenset(self._kwargs.values())
+        return frozenset(self.flat_args)
 
     @property
     def free_variables(self) -> Set[Variable]:
@@ -121,12 +125,16 @@ def kwargs(self, kwargs):
         self._kwargs = kwargs
 
     @property
-    def args(self) -> Tuple[Variable, ...]:
+    def args(self) -> Tuple[Any, ...]:
         return tuple(self.kwargs.values())
 
     @property
     def arg_names(self) -> Tuple[str, ...]:
         return tuple(self.kwargs)
+    
+    @property 
+    def flat_args(self) -> Tuple[Variable, ...]:
+        return tuple(x for x, in nested_zip(self._kwargs))
 
     @property
     def factor_out(self):
@@ -245,7 +253,7 @@ def _unique_representation(self):
         return (
             self._factor,
             self.arg_names,
-            self.args,
+            self.flat_args,
             self.deterministic_variables,
         )
 
@@ -272,19 +280,19 @@ def _numerical_factor_jacobian(
         factor._factor(*args), jax.jacobian(factor._factor, range(len(args)))(*args)
         """
         eps = eps or self.eps
-        args = tuple(np.array(value, dtype=np.float64) for value in args)
+
+        args = nested_map(lambda _, val: np.array(val, dtype=np.float64), self.args, args)
 
         raw_fval0 = self._factor_args(*args)
         fval0 = self._factor_value(raw_fval0).to_dict()
 
         jac = {
-            v0: tuple(
-                np.empty_like(val, shape=np.shape(val) + np.shape(value))
-                for value in args
-            )
-            for v0, val in fval0.items()
+            v0: nested_map(
+                lambda _, v: np.empty_like(val, shape=np.shape(val) + np.shape(v)),
+                self.args, args
+            ) for v0, val in fval0.items()
         }
-        for i, val in enumerate(args):
+        for ks, _, val in nested_items(self.args, args):
             with np.nditer(val, op_flags=["readwrite"], flags=["multi_index"]) as it:
                 for x_i in it:
                     val[it.multi_index] += eps
@@ -293,7 +301,9 @@ def _numerical_factor_jacobian(
                     x_i -= eps
                     indexes = (Ellipsis,) + it.multi_index
                     for v0, jac_v0v_i in jac_v1_i.items():
-                        jac[v0][i][indexes] = jac_v0v_i
+                        key_path = (v0, *ks, indexes)
+                        nested_set(jac, key_path, jac_v0v_i)
+                        # jac[v0][i][indexes] = jac_v0v_i
 
         # This replicates the output of normal
         # jax.jacobian(self.factor, len(self.args))(*args)
@@ -304,7 +314,7 @@ def _numerical_factor_jacobian(
     def numerical_func_jacobian(
             self, values: VariableData, **kwargs
     ) -> tuple:
-        args = (values[k] for k in self.args)
+        args = self.resolve_args(values)
         raw_fval, raw_jac = self._numerical_factor_jacobian(*args, **kwargs)
         fval = self._factor_value(raw_fval)
         jvp = self._jac_out_to_jvp(raw_jac, values=fval.to_dict().merge(values))

autofit/graphical/factor_graphs/factor.py

@@ -13,6 +13,8 @@
 
 from autofit.graphical.utils import (
     nested_filter,
+    to_variabledata, 
+    nested_zip, 
     is_variable,
     try_getitem,
 )
@@ -191,7 +193,6 @@ def __init__(
             **kwargs,
         )
 
-        # self.factor_out = factor_out
         self.eps = eps
         self._set_factor(factor)
         self._set_jacobians(
@@ -319,16 +320,16 @@ def _factor_value(self, raw_fval) -> 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))
+        det_values = to_variabledata(self.factor_out, raw_fval)
         fval = det_values.pop(FactorValue, 0.0)
         return FactorValue(fval, det_values)
 
     def __call__(self, values: VariableData) -> FactorValue:
         """Calls the factor with the values specified by the dictionary of
         values passed, returns a FactorValue with the value returned by the
         factor, and any deterministic factors"""
-        args = [values[v] for v in self.args]
-        key = self._key("__call__", *args)
+        args = self.resolve_args(values)
+        key = self._key("__call__", *(val for _, val in nested_zip(self.args, args)))
 
         if key not in self._cache:
             raw_fval = self._factor_args(*args)
@@ -351,7 +352,7 @@ def _vjp_func_jacobian(
         from autofit.graphical.factor_graphs.jacobians import (
             VectorJacobianProduct,
         )
-        raw_fval, fvjp = self._factor_vjp(*(values[v] for v in self.args))
+        raw_fval, fvjp = self._factor_vjp(*self.resolve_args(values))
         fval = self._factor_value(raw_fval)
 
         fvjp_op = VectorJacobianProduct(
@@ -380,7 +381,7 @@ def _key(*args):
     def _jvp_func_jacobian(
             self, values: VariableData, **kwargs
     ) -> Tuple[FactorValue, "JacobianVectorProduct"]:
-        args = list(values[k] for k in self.args)
+        args = self.resolve_args(values)
         key = self._key("_jvp_func_jacobian", *args)
 
         if key not in self._cache:
@@ -402,7 +403,7 @@ def _unpack_jacobian_out(self, raw_jac: Any) -> Dict[Variable, VariableData]:
         jac = {}
         for v0, vjac in nested_filter(is_variable, self.factor_out, raw_jac):
             jac[v0] = VariableData()
-            for v1, j in zip(self.args, vjac):
+            for v1, j in nested_zip(self.args, vjac):
                 jac[v0][v1] = j
 
         return jac

autofit/graphical/factor_graphs/jacobians.py

@@ -20,6 +20,7 @@
     nested_filter,
     nested_update,
     is_variable,
+    to_variabledata,
 )
 from autofit.mapper.variable import (
     Variable,
@@ -114,9 +115,11 @@ def grad(self, values=None):
         if values:
             grad.update(values)
 
-        for v, g in self(grad).items():
+        jac = self(grad)
+        for v, g in jac.items():
             grad[v] = grad.get(v, 0) + g
 
+        grad.pop(FactorValue)
         return grad
 
 
@@ -138,13 +141,18 @@ def factor_out(self):
 
 class VectorJacobianProduct(AbstractJacobian):
     def __init__(
-            self, factor_out, vjp: Callable, *variables: Variable, out_shapes=None
+            self, factor_out, vjp: Callable, *args: Variable, out_shapes=None
     ):
         self.factor_out = factor_out
         self.vjp = vjp
-        self._variables = variables
+        self._args = args
+        self._variables = tuple(v for v, in nested_filter(is_variable, args))
         self.out_shapes = out_shapes
 
+    @property 
+    def args(self):
+        return self._args 
+
     @property
     def variables(self):
         return self._variables
@@ -172,7 +180,7 @@ def _get_cotangent(self, values):
     def __call__(self, values: Union[VariableData, FactorValue]) -> VariableData:
         v = self._get_cotangent(values)
         grads = self.vjp(v)
-        return VariableData(zip(self.variables, grads))
+        return to_variabledata(self.args, grads)
 
     __rmul__ = __call__
 

autofit/graphical/laplace/newton.py

@@ -70,7 +70,7 @@ def diag_sr1_update(
     d = dzk.dot(dzk)
     if d > tol * dk.norm() ** 2 * zk.norm() ** 2:
         alpha = -zk.dot(dk) / d
-        Bk = Bk.diagonalupdate(alpha * (zk ** 2))
+        Bk = Bk.diagonalupdate((zk ** 2) * alpha)
 
     state1.hessian = Bk
     return state1
@@ -93,7 +93,7 @@ def diag_sr1_update_(
         else:
             alpha[v] = 0.0
 
-    Bk = Bk.diagonalupdate(alpha * (zk ** 2))
+    Bk = Bk.diagonalupdate((zk ** 2) * alpha)
 
     state1.hessian = Bk
     return state1
@@ -184,7 +184,7 @@ def diag_quasi_deterministic_update(
     zk2 = zk ** 2
     zk4 = (zk2 ** 2).sum()
     alpha = (dk.dot(Bxk.dot(dk)) - zk.dot(Bzk.dot(zk))) / zk4
-    state1.det_hessian = Bzk.diagonalupdate(float(alpha) * zk2)
+    state1.det_hessian = Bzk.diagonalupdate(zk2 * alpha)
 
     return state1