def get_phase(self, param: Parametrization):
struct = param.get_structure()
efields = self.sim.simulate(struct)
Cx = np.array(self.objective_C.H.dot(efields))
Cx = np.squeeze(Cx)
return np.angle(Cx)
def get_overlap_norm2(self, param: Parametrization):
""" Get the overlap norm squared given a parametrization. """
# Get the fields and structure.
struct = param.get_structure()
efields = self.sim.simulate(struct)
# Calculate overlap.
overlap2 = np.abs(np.array(self.objective_C.H.dot(efields)))**2
return overlap2
def eval_fun(fun_list: Union[List["OptimizationFunction"],
"OptimizationFunction"],
param: Parametrization) -> Union[List[np.ndarray], np.ndarray]:
"""Evaluates `fun` with parametrization `param`.
This function evaluates `fun_list` by building up a computational graph and
serially processes nodes that are quick to compute and parallely processes
nodes that are computationally demanding.
This function assumes that `param.get_structure()` will replace up to one
`Variable` in the computational graph. If no `Variable` nodes exist, the
value will not be used. If multiple `Variable` nodes exist, a `ValueError`
is raised.
Args:
fun_list: List of functions to evaluate or a single function.
param: Parametrization to evaluate with.
Returns:
List of function values for each function in `fun_list` if `fun_list`
is an iterable. Otherwise, returns the value of the single function.
Raises:
ValueError: If more than one `Variable` node was found.
"""
if isinstance(fun_list, problem.OptimizationFunction):
fun_list = [fun_list]
single_output = True
else:
single_output = False
fun_map, graph, in_nodes, out_nodes, heavy_nodes = _create_computational_graph(
fun_list)
# Run computation based on topological sort.
top_sorted_nodes = _top_sort_affinity(graph, heavy_nodes)
# Initialize any known function values. For now this is the structure,
# but in the future this could include other variables and cached values.
fun_vals = {}
if len(in_nodes) == 1:
fun_vals[list(in_nodes)[0]] = param.get_structure()
elif len(in_nodes) > 1:
raise ValueError("Multiple Variable nodes detected.")
_eval_fun_vals(fun_vals, fun_map, graph, top_sorted_nodes, param)
ret_vals = [fun_vals[out_node] for out_node in out_nodes]
if single_output:
return ret_vals[0]
return ret_vals
def calculate_gradient(self, param: Parametrization) -> np.ndarray:
struct = param.get_structure()
efields = self.sim.simulate(struct)
total_df_dz = self.calculate_df_dz(efields, struct)
# TODO(logansu): Cache gradient calculation.
dz_dp = aslinearoperator(param.calculate_gradient())
df_dp = np.conj(
dz_dp.adjoint()
@ np.conj(total_df_dz +
self.calculate_partial_df_dz(efields, struct)))
return df_dp
def eval_grad(fun: "OptimizationFunction",
param: Parametrization) -> np.ndarray:
"""Evalutes the gradient of `fun` with respect to `param`.
This function evaluates the gradient by reverse-mode autodiff using similar
technique described in `eval_fun`. Reverse-mode autodiff occurs in two
phases. In the first "forward" phase, all the values of the functions
at each node are computed. In the "backward" phase, the gradients are
computed starting from the output function moving towards the inputs.
This function assumes that `param.get_structure()` will replace up to one
`Variable` in the computational graph. If no `Variable` nodes exist, the
value will not be used. If multiple `Variable` nodes exist, a `ValueError`
is raised.
Args:
fun: The function to evaluate.
param: Parametrization to evaluate with.
Returns:
The function value.
Raises:
ValueError: If more than one `Variable` node was found.
"""
fun_map, graph, in_nodes, out_nodes, heavy_nodes = _create_computational_graph(
[fun])
out_node = out_nodes[0]
# Forward pass (phase 1): Compute function values. This is a near identical
# behavior to `eval_fun`.
# Dictionary mapping function ids to values.
fun_vals = {}
# Check that we have at most one input node.
# It is actually possible to have no input node because the old-style
# function might be masking it (i.e. not listing it as a dependency).
if len(in_nodes) == 1:
in_node = list(in_nodes)[0]
fun_vals[in_node] = param.get_structure()
elif len(in_nodes) > 1:
raise ValueError("Multiple Variable nodes detected.")
top_sorted_nodes = _top_sort_affinity(graph, heavy_nodes)
_eval_fun_vals(fun_vals, fun_map, graph, top_sorted_nodes, param)
# Backward pass (phase 2): The gradients are computed moving from
# the last node to the first.
# Dictionary mapping function ids to gradients.
grad_vals = {node: 0.0 for node in fun_map.keys()}
grad_vals[out_node] = 1
def _eval_grad(node, input_vals, grad_val):
if _is_old_fun(fun_map[node]):
return fun_map[node].calculate_gradient(param)
return fun_map[node].grad(input_vals, grad_val)
# We need to maintain a list of gradients computed using the old-style
# functions as the old-style functions don't list dependencies. These
# gradients are then summed with the new-style gradients at the end.
old_node_grads = []
try:
node_iter = iter(top_sorted_nodes)
node = next(node_iter)
while True:
while True:
# Postpone any heavy computation.
if _is_heavy_fun(fun_map[node]):
break
if _is_old_fun(fun_map[node]):
old_grad = fun_map[node].calculate_gradient(param)
# Use `np.dot` since either operand could be a scalar
# and matrix multiplication(@) disallows scalars.
old_node_grads.append(np.dot(grad_vals[node], old_grad))
else:
# Get the gradients of the output function with respect to
# the inputs of function `node`.
input_vals = [
fun_vals[next_node] for next_node in graph[node]
]
node_grad_vals = fun_map[node].grad(
input_vals, grad_vals[node])
# Update the gradient values.
for next_node, next_fun_grad in zip(
graph[node], node_grad_vals):
grad_vals[next_node] += next_fun_grad
node = next(node_iter)
# At this point, we have processed all the possible nodes.
# We now parallelize the computation of the `heavy_compute_nodes`.
heavy_node_block = []
last_node_reached = False
try:
while _is_heavy_fun(fun_map[node]):
heavy_node_block.append(node)
node = next(node_iter)
except StopIteration:
last_node_reached = True
# Parallelize heavy compute.
arg_list = []
for node_ in heavy_node_block:
arg_list.append((
node_,
[fun_vals[in_node] for in_node in graph[node_]],
grad_vals[node_],
))
with concurrent.futures.ThreadPoolExecutor() as executor:
heavy_node_vals = executor.map(lambda args: _eval_grad(*args),
arg_list)
for node_, node_grad_vals in zip(heavy_node_block,
heavy_node_vals):
for next_node, next_fun_grad in zip(graph[node_],
node_grad_vals):
grad_vals[next_node] += next_fun_grad
if last_node_reached:
break
except StopIteration:
pass
# Compute the final gradients. This consists of any gradients computed
# using old-style functions as well as the new-style function.
total_grads = old_node_grads
if len(in_nodes) >= 1:
total_grads += [
np.array(grad_vals[in_node]) @ param.calculate_gradient()
]
return np.sum(total_grads, axis=0)
def get_electric_fields(self, param: Parametrization):
return fdfd_tools.unvec(self.sim.simulate(param.get_structure()),
self.sim.get_dims())
def calculate_objective_function(self, param: Parametrization) -> np.array:
vec = param.get_structure()
return self.radius**2 - (np.ones_like(vec) @ (vec - self.r0)**2)
def calculate_gradient(self, param: Parametrization) -> np.array:
vec = param.get_structure()
return -2 * (vec - self.r0)
def calculate_objective_function(self, param: Parametrization) -> np.array:
vec = param.get_structure()
return -np.ones_like(vec) @ vec
def calculate_gradient(self, param: Parametrization) -> np.array:
vec = param.get_structure()
return -np.ones_like(vec)
def get_electric_fields(self, param: Parametrization):
# TODO(logansu): Refactor please
return fdfd_tools.unvec(self.sim.simulate(param.get_structure()),
self.sim.get_dims())