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)"