def eval_grad(self, node: ProblemGraphNode,
wrt_nodes: List[ProblemGraphNode]):
self.add_node(node)
from spins.goos import graph_executor
from spins.goos import flows
override_map = {}
for var_name, var_value in self._var_value.items():
# Determine the gradient.
if self._var_frozen[var_name]:
grad_value = flows.NumericFlow(np.zeros_like(var_value))
else:
grad_value = flows.NumericFlow(np.ones_like(var_value))
# Setup the context.
const_flags = flows.NumericFlow.ConstFlags()
frozen_flags = flows.NumericFlow.ConstFlags(False)
frozen_flags.set_all(self._var_frozen[var_name])
context = NodeFlags(const_flags=const_flags,
frozen_flags=frozen_flags)
override_map[self._node_map[var_name]] = (
flows.NumericFlow(var_value), grad_value, context)
return graph_executor.eval_grad(node, wrt_nodes, override_map)
def eval_nodes(self, nodes: List[ProblemGraphNode]) -> List[flows.Flow]:
"""Evaluates nodes.
If the node is not already in the optplan, it is added.
Args:
nodes: List of nodes to evaluate.
Returns:
List of flows, one for each node.
"""
for node in nodes:
self.add_node(node)
from spins.goos import graph_executor
from spins.goos import flows
override_map = {}
for var_name, var_value in self._var_value.items():
# Setup the context.
const_flags = flows.NumericFlow.ConstFlags()
frozen_flags = flows.NumericFlow.ConstFlags(False)
frozen_flags.set_all(self._var_frozen[var_name])
context = NodeFlags(const_flags=const_flags,
frozen_flags=frozen_flags)
override_map[self._node_map[var_name]] = (
flows.NumericFlow(var_value), context)
return graph_executor.eval_fun(nodes, override_map)
def eval(self, inputs: List[goos.ShapeFlow]) -> flows.NumericFlow:
self._grid.clear()
extra_grids = self._render(self._grid, inputs[0])
self._grid.render()
grids = self._grid.grids
for grid in extra_grids:
for i in range(3):
grids[i] += grid.grids[i]
return flows.NumericFlow(grids)
def eval(self, inputs: List[goos.ShapeFlow]) -> flows.NumericFlow:
self._grid.clear()
# Save geometry for backprop.
self._geom = _create_geometry(inputs[0])
extra_grids = self._geom.eval(self._grid, self._render_params)
self._grid.render()
if extra_grids is None:
extra_grids = []
elif not isinstance(extra_grids, list):
extra_grids = [extra_grids]
grids = self._grid.grids
for grid in extra_grids:
for i in range(3):
grids[i] += grid.grids[i]
return flows.NumericFlow(grids)
def run(self, plan: goos.OptimizationPlan, start_iter: int = 0):
variables = plan.get_thawed_vars()
var_shapes = []
initial_val = []
bounds = []
for var in variables:
value = plan.get_var_value(var)
if value.shape:
var_shapes.append(value.shape)
else:
var_shapes.append([1])
initial_val.append(value.flatten())
bound = plan.get_var_bounds(var)
for lower, upper in zip(bound[0].flatten(), bound[1].flatten()):
if lower == -np.inf:
lower = None
if upper == np.inf:
upper = None
bounds.append((lower, upper))
override_map = {
plan._node_map[var_name]: flows.NumericFlow(value)
for var_name, value in plan._var_value.items()
}
# TODO(logansu): Currently we call optimize with every single variable
# in the plan, but we can really reduce the number of elements by
# focusing only the variables that are required to compute the objective
# function.
def unpack(x):
cur_ind = 0
values = []
for shape in var_shapes:
values.append(
np.reshape(x[cur_ind:cur_ind + np.prod(shape)], shape))
cur_ind += np.prod(shape)
return values
def unpack_and_set(x):
values = unpack(x)
for var, value in zip(variables, values):
plan.set_var_value(var, value)
def func(x):
unpack_and_set(x)
val = plan.eval_node(self._obj).array
plan.logger.debug("Function evaluated: %f", val)
return val
def grad(x):
unpack_and_set(x)
grad_flows = plan.eval_grad(self._obj, variables)
val = np.hstack([flow.array_grad.flatten() for flow in grad_flows])
plan.logger.debug("Gradient evaluated, norm: %f",
np.linalg.norm(val))
return val
# To avoid scipy warning, only pass Jacobian to methods that need it.
methods_that_need_jacobian = {
"CG", "BFGS", "L-BFGS-B", "TNC", "SLSQP", "dogleg", "trust-ncg"
}
jac = None
if self._method in methods_that_need_jacobian:
jac = grad
# Handle constraints.
methods_that_accept_constraints = {'SLSQP', 'COBYLA'}
constraints = None
if self._method in methods_that_accept_constraints:
constraints = []
for eq in self._cons_eq:
def cons_fun(x):
unpack_and_set(x)
val = plan.eval_node(eq).array
plan.logger.debug("Eq. cons. function evaluated: %f", val)
return val
def cons_jac(x):
unpack_and_set(x)
grad_flows = plan.eval_grad(eq, variables)
val = []
for flow, var_shape in zip(grad_flows, var_shapes):
# Flatten only the dimension corresponding to the
# variable.
arr = flow.array_grad
new_shape = arr.shape[:-len(var_shape)] + (
np.prod(var_shape), )
val.append(np.reshape(arr, new_shape))
val = np.hstack(val)
plan.logger.debug("Eq. cons. gradient evaluated, norm: %f",
np.linalg.norm(val))
return val
constraints.append({
"type": "eq",
"fun": cons_fun,
"jac": cons_jac
})
for ineq in self._cons_ineq:
# Note the negative sign because of opposite convention
# for inequalities (f >= 0 vs f <= 0).
def cons_fun(x):
unpack_and_set(x)
val = plan.eval_node(ineq).array
plan.logger.debug("Ineq. cons. function evaluated: %f",
val)
return -val
def cons_jac(x):
unpack_and_set(x)
grad_flows = plan.eval_grad(ineq, variables)
val = []
for flow, var_shape in zip(grad_flows, var_shapes):
# Flatten only the dimension corresponding to the
# variable.
arr = flow.array_grad
new_shape = arr.shape[:-len(var_shape)] + (
np.prod(var_shape), )
val.append(np.reshape(arr, new_shape))
val = np.hstack(val)
plan.logger.debug(
"Ineq. cons. gradient evaluated, norm: %f",
np.linalg.norm(val))
return -val
constraints.append({
"type": "ineq",
"fun": cons_fun,
"jac": cons_jac
})
elif (len(self._cons_ineq) > 0) or (len(self._cons_eq) > 0):
plan.logger.warning(
"Using optimizer that cannot handle constraints. Constraints "
"ignored: %d" % (len(self._cons_ineq) + len(self._cons_eq)))
# Keep track of iteration number.
iter_num = start_iter
def callback(x):
# Update the variable values before evaluating monitors.
values = unpack(x)
for var, value in zip(variables, values):
plan.set_var_value(var, value)
# Update iteration number.
nonlocal iter_num
iter_num += 1
if self._iter:
plan.set_var_value(self._iter, iter_num)
plan.write_event({
"state": "optimizing",
"iteration": iter_num
}, self._monitor_list)
# Adjust total number of iterations if we are resuming.
options = copy.deepcopy(self._options)
if "maxiter" in options:
options["maxiter"] -= start_iter
initial_val = np.hstack(initial_val)
self._results = scipy.optimize.minimize(func,
initial_val,
method=self._method,
jac=jac,
callback=callback,
bounds=bounds,
constraints=constraints,
**options)
unpack_and_set(self._results["x"])
def eval(self, inputs: List[flows.NumericFlow],
context: goos.EvalContext) -> flows.NumericFlow:
return flows.NumericFlow(self._value)
def eval(self, inputs: List) -> flows.NumericFlow:
return flows.NumericFlow(goos.get_default_plan().get_var_value(self))
def grad(self, inputs: List[flows.NumericFlow],
grad_val: flows.NumericFlow) -> List[flows.NumericFlow]:
return [flows.NumericFlow(np.zeros_like(self._value))]
def eval(self, inputs: List[flows.NumericFlow]) -> flows.NumericFlow:
return flows.NumericFlow(self._value)
