def __init__(self,
filepath: str = ".",
savepath: str = ".",
nodes: Optional[Union[optplan.ProblemGraphNode, List[
optplan.ProblemGraphNode]]] = None) -> None:
"""Initializes the workspace.
Args:
filepath: Root folder for the workspace. This is used as the root
folder for loading any auxiliary files.
savepath: Path to save any logs.
nodes: Optional list of nodes that exist in the graph.
"""
self._filepath = filepath
# Mapping from node name to the actual node. This is for node
# definition lookup as well as to quickly see which nodes are part of
# the graph.
self._nodes = {}
# List of all the objects that have been instantiated.
self._objects = {}
# TODO(logansu): Accept argument-less variable.
self._objects[VARIABLE_NODE] = problem.Variable(1)
if isinstance(nodes, collections.Iterable):
for node in nodes:
self._add_node(node)
elif nodes:
self._add_node(nodes)
# Setup logging if applicable.
self._logger = None
if savepath:
self._logger = Logger(savepath, self)
def test_eval():
param = parametrization.DirectParam((1, 2, 3, 4, 5))
x = problem.Variable(5)
obj = (x[0] * x[1] - x[2]**2) * (x[3] + 2 * x[4])
assert graph_executor.eval_fun(obj, param) == -98
np.testing.assert_array_equal(graph_executor.eval_grad(obj, param),
[28, 14, -84, -7, -14])
def test_eval_old_in_new():
"""Tests evals when a new-style function depends on old-style function."""
param = parametrization.DirectParam((3, 2))
x = problem.Variable(2)
obj = OldPower(x[0], 2) + x[0]
assert graph_executor.eval_fun(obj, param) == 12
np.testing.assert_array_equal(graph_executor.eval_grad(obj, param), [7, 0])
def test_eval_multiple():
param = parametrization.DirectParam((1, 2, 3, 4, 5))
x = problem.Variable(5)
obj_part1 = x[3] + 2 * x[4]
obj_part2 = x[0] * x[1] - x[2]**2
obj = obj_part1 * obj_part2
assert graph_executor.eval_fun([obj, obj_part1, obj_part2],
param) == [-98, 14, -7]
def test_eval_heavy_compute_end_on_heavy():
param = parametrization.DirectParam((-3, 2))
x = problem.Variable(2)
obj = HeavyIdentity(x[0] * x[1])
np.testing.assert_array_almost_equal(graph_executor.eval_fun(obj, param),
-6)
np.testing.assert_array_equal(graph_executor.eval_grad(obj, param),
[2, -3])
def test_eval_fun_no_in_node_works():
"""Tests when no input node is found.
This can happen if the input node is fed only to the old-style function,
which does not list its dependencies. Consequently, no input node can be
found.
"""
param = parametrization.DirectParam((2, ))
x = problem.Variable(1)
obj = OldPower(x[0], 2) + OldPower(x[0], 3)
assert graph_executor.eval_fun(obj, param) == 12
np.testing.assert_array_equal(graph_executor.eval_grad(obj, param), [16])
def test_eval_heavy_compute():
param = parametrization.DirectParam((-1, 2, 1.2, 3, 5.1))
x = problem.Variable(5)
obj = HeavyIdentity(x[1] * x[0] * x[3])
obj2 = HeavyIdentity(x[3] + x[1]**3)
obj3 = obj * (x[0] + 2) + obj2 - x[4]
np.testing.assert_array_almost_equal(graph_executor.eval_fun(obj3, param),
-0.1)
np.testing.assert_array_equal(graph_executor.eval_grad(obj3, param),
[0, 9, 0, -1, -1])
def test_epsilon_eval():
space = make_simspace()
space_inst = space(1550)
param_shape = np.prod(space.design_dims)
vec = np.ones(param_shape)
vec[10:20] = 0.6
vec[30:40] = 0.2
eps = creator_em.Epsilon(problem.Variable(1), 1550, space)
eps_actual = eps.eval([vec])
eps_expected = (fdfd_tools.vec(space_inst.eps_bg.grids) +
space_inst.selection_matrix @ vec)
np.testing.assert_array_equal(eps_actual, eps_expected)
def test_fdfd_simulation_grad():
# Create a 3x3 2D grid to brute force check adjoint gradients.
shape = [3, 3, 1]
# Setup epsilon (pure vacuum).
epsilon = [np.ones(shape) for i in range(3)]
# Setup dxes. Assume dx = 40.
dxes = [[np.ones(shape[i]) * 40 for i in range(3)] for j in range(2)]
# Setup a point source in the center.
J = [np.zeros(shape).astype(complex) for i in range(3)]
J[2][1, 0, 0] = 1.2j
J[2][1, 1, 0] = 1
# Setup target fields.
target_fields = [np.zeros(shape).astype(np.complex128) for i in range(3)]
target_fields[2][:, :, 0] = 20j + 1
overlap_vec = fdfd_tools.vec(target_fields)
# TODO(logansu): Deal with this.
class SimspaceMock:
@property
def dxes(self):
return dxes
@property
def pml_layers(self):
return [0] * 6
eps_param = parametrization.DirectParam(fdfd_tools.vec(epsilon),
bounds=[0, 100])
eps_fun = problem.Variable(len(fdfd_tools.vec(epsilon)))
sim_fun = creator_em.FdfdSimulation(
eps=eps_fun,
solver=local_matrix_solvers.DirectSolver(),
wlen=1500,
source=J,
simspace=SimspaceMock(),
)
obj_fun = problem.AbsoluteValue(
objective=creator_em.OverlapFunction(sim_fun, overlap_vec))**2
grad_actual = obj_fun.calculate_gradient(eps_param)
def eval_fun(vec: np.ndarray):
eps_param.from_vector(vec)
return obj_fun.calculate_objective_function(eps_param)
grad_brute = eval_grad_brute(fdfd_tools.vec(epsilon), eval_fun)
np.testing.assert_array_almost_equal(grad_actual, grad_brute, decimal=0)
def test_epsilon_grad():
# Compute derivative of `sum(epsilon)`.
space = make_simspace()
space_inst = space(1550)
param_shape = np.prod(space.design_dims)
vec = np.ones(param_shape) * 0.4
vec[10:20] = 0.6
vec[20:40] = 0.2
eps = creator_em.Epsilon(problem.Variable(1), 1550, space)
# Brute force gradient.
fun = lambda vec: np.sum(eps.eval([vec]))
grad_brute = eval_grad_brute(vec, fun)
grad_actual = eps.grad([vec], np.ones(3 * np.prod(space.dims)))
np.testing.assert_array_almost_equal(grad_actual[0], grad_brute, decimal=4)
def test_plane_power_grad():
space = Simspace(
TESTDATA,
optplan.SimulationSpace(
pml_thickness=[0, 0, 0, 0, 0, 0],
mesh=optplan.UniformMesh(dx=40),
sim_region=optplan.Box3d(
center=[0, 0, 0],
extents=[80, 80, 80],
),
eps_bg=optplan.GdsEps(
gds="straight_waveguide.gds",
mat_stack=optplan.GdsMaterialStack(
background=optplan.Material(mat_name="air"),
stack=[
optplan.GdsMaterialStackLayer(
gds_layer=[100, 0],
extents=[-80, 80],
foreground=optplan.Material(mat_name="Si"),
background=optplan.Material(mat_name="air"),
),
],
),
),
))
wlen = 1550
power_fun = poynting.PowerTransmissionFunction(
field=problem.Variable(1),
simspace=space,
wlen=wlen,
plane_slice=grid_utils.create_region_slices(space.edge_coords,
[0, 0, 0], [40, 80, 80]),
axis=gridlock.axisvec2axis([1, 0, 0]),
polarity=gridlock.axisvec2polarity([1, 0, 0]))
field = np.arange(np.prod(space.dims) * 3).astype(np.complex128) * 1j
grad_actual = power_fun.grad([field], 1)
fun = lambda vec: power_fun.eval([vec])
grad_brute = eval_grad_brute_wirt(field, fun)
np.testing.assert_array_almost_equal(grad_actual[0], grad_brute, decimal=4)
def test_stored_energy_grad():
space = Simspace(
TESTDATA,
optplan.SimulationSpace(
pml_thickness=[0, 0, 0, 0, 0, 0],
mesh=optplan.UniformMesh(dx=40),
sim_region=optplan.Box3d(
center=[0, 0, 0],
extents=[80, 80, 80],
),
eps_bg=optplan.GdsEps(
gds="straight_waveguide.gds",
mat_stack=optplan.GdsMaterialStack(
background=optplan.Material(mat_name="air"),
stack=[
optplan.GdsMaterialStackLayer(
gds_layer=[100, 0],
extents=[-80, 80],
foreground=optplan.Material(mat_name="Si"),
background=optplan.Material(mat_name="air"),
),
],
),
),
))
wlen = 1550
energy_fun = stored_energy.StoredEnergyFunction(
input_function=problem.Variable(1),
simspace=space,
center=[0,0,0],
extents=[0,0,0],
epsilon=space._eps_bg)
plane_slice=grid_utils.create_region_slices(space.edge_coords,
[0, 0, 0], [40, 80, 80]),
axis=gridlock.axisvec2axis([1, 0, 0]),
polarity=gridlock.axisvec2polarity([1, 0, 0]))
def test_eval_fun_multiple_variables_raises_value_error():
with pytest.raises(ValueError, match=r"Multiple Variable"):
param = parametrization.DirectParam((1, 2))
x = problem.Variable(2)
y = problem.Variable(2)
graph_executor.eval_fun(x + y, param)
def test_epsilon_string():
space = make_simspace()
space_inst = space(1550)
eps = creator_em.Epsilon(problem.Variable(1), 1550, space)
assert str(eps) == "Epsilon(1550)"
