def test_rename(temp_context): class NewConstant(goos.Function): node_type = "goos.new_constant" def __init__(self, value: np.ndarray): super().__init__() self._value = np.array(value) def eval(self, inputs: List[flows.NumericFlow]) -> flows.NumericFlow: return flows.NumericFlow(self._value) def grad(self, inputs: List[flows.NumericFlow], grad_val: flows.NumericFlow) -> List[flows.NumericFlow]: return [flows.NumericFlow(np.zeros_like(self._value))] with goos.OptimizationPlan() as plan: orig_node = NewConstant(3, name="old_node") new_node = goos.rename(orig_node, "new_node") # Make sure that the name is correct. assert new_node._goos_name == "new_node" # Ensure that the node type is correct. assert isinstance(new_node, NewConstant) # Ensure that basic operations hold. assert (new_node + 4).get() == 7
def test_product_double_array(): with goos.OptimizationPlan() as plan: x = goos.Variable([3, 1]) res = goos.Norm(x * goos.Constant([2, 4]))**2 np.testing.assert_allclose(res.get().array, 52) np.testing.assert_allclose(res.get_grad([x])[0].array_grad, [24, 32])
def __init__(self, eps: goos.Shape, source: goos.Function, solver: str, wavelengths: List[float], simulation_space: simspace.SimulationSpace, background: goos.material.Material, outputs: List[SimOutput]) -> None: # Determine the output flow types. output_flow_types = [out.Attributes.output_type for out in outputs] super().__init__([eps, source], flow_types=output_flow_types) # Internally setup a plan to handle the simulation. self._plan = goos.OptimizationPlan() self._sims = [] self._eps = eps with self._plan: for wlen in wavelengths: # TODO(logansu): Use a placeholder dummy for the shape to # reduce runtime and save memory. eps_rendered = render.RenderShape( self._eps, region=simulation_space.sim_region, mesh=simulation_space.mesh, background=background, wavelength=wlen) sim_result = SimulateNode( eps=eps_rendered, source=source, wavelength=wlen, simulation_space=simulation_space, solver=solver, outputs=outputs, ) self._sims.append(sim_result)
def test_power(): with goos.OptimizationPlan() as plan: x = goos.Variable(3.0) y = x**2 assert y.get() == 9.0 assert y.get_grad([x]) == [6.0]
def test_render_prism(): with goos.OptimizationPlan() as plan: rect = goos.Cuboid(extents=goos.Constant([80, 80, 100]), pos=goos.Constant([0, 0, 0]), material=goos.material.Material(index=3.45)) render = maxwell.RenderShape( rect, region=goos.Box3d(center=[0, 0, 0], extents=[160, 200, 40]), mesh=maxwell.UniformMesh(dx=40), background=goos.material.Material(index=1.0), wavelength=1550) np.testing.assert_array_almost_equal( render.get().array, [[[[1.], [1.], [1.], [1.], [1.]], [[1.], [1.], [11.9025], [11.9025], [1.]], [[1.], [1.], [11.9025], [11.9025], [1.]], [[1.], [1.], [1.], [1.], [1.]]], [[[1.], [1.], [1.], [1.], [1.]], [[1.], [3.725625], [6.45125], [3.725625], [1.]], [[1.], [6.45125], [11.9025], [6.45125], [1.]], [[1.], [3.725625], [6.45125], [3.725625], [1.]]], [[[1.], [1.], [1.], [1.], [1.]], [[1.], [1.], [6.45125], [6.45125], [1.]], [[1.], [1.], [11.9025], [11.9025], [1.]], [[1.], [1.], [6.45125], [6.45125], [1.]]]])
def test_product_double(): with goos.OptimizationPlan() as plan: x = goos.Variable(3) res = x * 2 assert res.get() == 6 assert res.get_grad([x])[0] == 2
def test_sum_single(): with goos.OptimizationPlan() as plan: x = goos.Variable(2.0) res = goos.Sum([x]) assert res.get() == 2 assert res.get_grad([x]) == [1]
def test_pixelated_multietch_grating_xz_polarity(): use_edge_locs = True params = [0, 1, 1.5, 2] # Grating defined in xz-plane. with goos.OptimizationPlan() as plan: edge_locs = goos.Variable(params) height_index = goos.Variable([2, 0, 1, 2]) grating = goos.grating.PixelatedGrating( edge_locs, height_index=height_index, height_fracs=[0, 0.5, 1], pos=[0, 0, 0], extents=[2, 10, 0.220], material=goos.material.Material(index=1), material2=goos.material.Material(index=3.5), grating_dir=0, grating_dir_spacing=0.5, etch_dir_divs=2, etch_polarity=-1, use_edge_locs=use_edge_locs).get() np.testing.assert_equal(grating.pixel_size, [0.5, 10, 0.110]) np.testing.assert_equal(grating.extents, [2, 10, 0.220]) np.testing.assert_almost_equal(grating.array, [[[1, 1]], [[1, 1]], [[0, 0]], [[1, 0]]], decimal=4)
def main(save_folder: str): goos.util.setup_logging(save_folder) with goos.OptimizationPlan(save_path=save_folder) as plan: x = goos.Variable(3, name="x") y = goos.Variable(1, name="y") obj = goos.rename((x + y)**2 + (y - 2)**2, name="obj") # First optimize only `x`. y.freeze() goos.opt.scipy_minimize(obj, "L-BFGS-B", max_iters=10) # Now do co-optimization. y.thaw() goos.opt.scipy_minimize(obj, "L-BFGS-B", max_iters=10) plan.save() plan.run() # More efficient to call `eval_nodes` when evaluating multiple nodes # at the same time. x_val, y_val, obj_val = plan.eval_nodes([x, y, obj]) print("x: {}, y: {}, obj: {}".format(x_val.array, y_val.array, obj_val.array))
def test_render_pixelated_cont_shape(): with goos.OptimizationPlan() as plan: def initializer(size): return [[1, 0], [0.5, 0.75], [0, 0.8], [1, 0.8]] var, design = goos.pixelated_cont_shape( initializer=initializer, pos=goos.Constant([200, 0, 0]), extents=[1000, 1000, 220], material=goos.material.Material(index=2), material2=goos.material.Material(index=4), pixel_size=[250, 500, 220]) rect = goos.Cuboid(extents=goos.Constant([1000, 1000, 220]), pos=goos.Constant([200, 0, 0]), material=goos.material.Material(index=2)) render = maxwell.RenderShape( design, region=goos.Box3d(center=[0, 0, 0], extents=[1500, 1500, 200]), mesh=maxwell.UniformMesh(dx=40), background=goos.material.Material(index=1.0), wavelength=1550) np.testing.assert_almost_equal( np.real(render.get().array[2][10:, 8, 3]), [ 1., 8.5, 16., 16., 16., 16., 16., 14.5, 10., 10., 10., 10., 10., 10., 4., 4., 4., 4., 4., 4., 13., 16., 16., 16., 16., 16., 8.5 ])
def test_slice(): with goos.OptimizationPlan() as plan: x = goos.Variable([[0, 1, 2, 4, 5], [6, 7, 8, 9, 10], [11, 12, 13, 14, 15], [16, 17, 18, 19, 20], [21, 22, 23, 24, 25]]) t = goos.Slice(x, ['c', 'c']) np.testing.assert_allclose(t.get().array, 13) g = np.array([[0, 0, 0, 0, 0], [0, 0, 0, 0, 0], [0, 0, 1, 0, 0], [0, 0, 0, 0, 0], [0, 0, 0, 0, 0]]) np.testing.assert_allclose(t.get_grad([x])[0].array_grad, g) t = goos.Slice(x, [[1, 4], 'c']) np.testing.assert_allclose(t.get().array, [[8], [13], [18]]) g = np.array([[0, 0, 0, 0, 0], [0, 0, 1, 0, 0], [0, 0, 1, 0, 0], [0, 0, 1, 0, 0], [0, 0, 0, 0, 0]]) np.testing.assert_allclose(t.get_grad([x])[0].array_grad, g) t = goos.Slice(x, [3, [1, 3]]) np.testing.assert_allclose(t.get().array, [[17, 18]]) g = np.array([[0, 0, 0, 0, 0], [0, 0, 0, 0, 0], [0, 0, 0, 0, 0], [0, 1, 1, 0, 0], [0, 0, 0, 0, 0]]) np.testing.assert_allclose(t.get_grad([x])[0].array_grad, g) t = goos.Slice(x, [3, None]) np.testing.assert_allclose(t.get().array, [[16, 17, 18, 19, 20]]) g = np.array([[0, 0, 0, 0, 0], [0, 0, 0, 0, 0], [0, 0, 0, 0, 0], [1, 1, 1, 1, 1], [0, 0, 0, 0, 0]]) np.testing.assert_allclose(t.get_grad([x])[0].array_grad, g)
def test_render_cylinder(): with goos.OptimizationPlan() as plan: cyl = goos.Cylinder(pos=goos.Constant([0, 0, 0]), radius=goos.Constant(60), height=goos.Constant(60), material=goos.material.Material(index=3.45)) render = maxwell.RenderShape( cyl, region=goos.Box3d(center=[0, 0, 0], extents=[200, 200, 40]), mesh=maxwell.UniformMesh(dx=40), background=goos.material.Material(index=1.0), wavelength=1550) np.testing.assert_array_almost_equal( render.get().array, [[[[1.], [1.], [1.], [1.], [1.]], [[1.], [2.20393657], [8.15133225], [8.14708434], [2.20801505]], [[1.], [4.81696873], [9.176875], [9.176875], [4.81254501]], [[1.], [2.20393657], [8.15133225], [8.14708434], [2.20801505]], [[1.], [1.], [1.], [1.], [1.]]], [[[1.], [1.], [1.], [1.], [1.]], [[1.], [2.20452237], [4.81426614], [2.20733312], [1.]], [[1.], [8.1503488], [9.176875], [8.14865466], [1.]], [[1.], [8.1503488], [9.176875], [8.14865466], [1.]], [[1.], [2.20452237], [4.81426614], [2.20733312], [1.]]], [[[1.], [1.], [1.], [1.], [1.]], [[1.], [1.06734618], [5.08385966], [5.07819579], [1.07209387]], [[1.], [5.0825484], [11.9025], [11.9025], [5.08028954]], [[1.], [5.0825484], [11.9025], [11.9025], [5.08028954]], [[1.], [1.06734618], [5.08385966], [5.07819579], [1.07209387]]]])
def main(save_folder: str, min_feature: float = 100, sim_3d: bool = False, visualize: bool = False) -> None: goos.util.setup_logging(save_folder) with goos.OptimizationPlan(save_path=save_folder) as plan: wg_in = goos.Cuboid(pos=goos.Constant([-2000, 0, 0]), extents=goos.Constant([3000, 400, 220]), material=goos.material.Material(index=3.45)) wg_out = goos.Cuboid(pos=goos.Constant([0, 2000, 0]), extents=goos.Constant([400, 3000, 220]), material=goos.material.Material(index=3.45)) def initializer(size): return np.random.random(size) * 0.2 + 0.5 # Continuous optimization. var, design = goos.cubic_param_shape( initializer=initializer, pos=goos.Constant([0, 0, 0]), extents=[2000, 2000, 220], pixel_spacing=40, control_point_spacing=1.5 * min_feature, material=goos.material.Material(index=1), material2=goos.material.Material(index=3.45), var_name="var_cont") sigmoid_factor = goos.Variable(4, parameter=True, name="discr_factor") design = goos.cast(goos.Sigmoid(sigmoid_factor * (2 * design - 1)), goos.Shape) eps = goos.GroupShape([wg_in, wg_out, design]) # This node is purely for debugging purposes. eps_rendered = maxwell.RenderShape( design, region=goos.Box3d(center=[0, 0, 0], extents=[3000, 3000, 0]), mesh=maxwell.UniformMesh(dx=40), wavelength=1550, ) if visualize: goos.util.visualize_eps(eps_rendered.get().array[2]) obj, sim = make_objective_fdtd(eps, "cont", sim_3d=sim_3d) for factor in [4, 8]: sigmoid_factor.set(factor) goos.opt.scipy_minimize( obj, "L-BFGS-B", monitor_list=[sim["eps"], sim["field"], sim["overlap"], obj], max_iters=6, name="opt_cont{}".format(factor)) plan.save() plan.run() if visualize: goos.util.visualize_eps(eps_rendered.get().array[2])
def test_sigmoid_with_scalar_ops(): with goos.OptimizationPlan() as plan: x = goos.Variable([0, 1]) y = goos.dot([1, 2], goos.Sigmoid(2 * x - 1)) np.testing.assert_allclose(y.get().array, 1.7310585786300) np.testing.assert_allclose( y.get_grad([x])[0].array_grad, [0.39322386648296, 0.7864477329659])
def test_norm(vec, order, expected_val, expected_grad): with goos.OptimizationPlan() as plan: var = goos.Variable(vec) norm = goos.Norm(var, order=order) assert norm.get() == expected_val np.testing.assert_almost_equal( norm.get_grad([var])[0].array_grad, expected_grad)
def test_sum_triple(): with goos.OptimizationPlan() as plan: x = goos.Variable(2.0) y = goos.Variable(3.0) res = x + y + 5 assert res.get() == 10 assert res.get_grad([x, y]) == [1, 1]
def test_sigmoid_array(): with goos.OptimizationPlan() as plan: x = goos.Variable([0, 1]) y = goos.dot([1, 2], goos.Sigmoid(x)) np.testing.assert_allclose(y.get().array, 1.962117) np.testing.assert_allclose( y.get_grad([x])[0].array_grad, [0.25, 0.39322386648])
def test_plan(tmp_path): plan_dir = tmp_path / "test_plan" plan_dir.mkdir() with goos.OptimizationPlan() as plan: x = goos.Variable(3.0, name="x") y = goos.Variable(2.0, name="y") z = x + y assert z.get() == 5 assert z.get_grad([x, y]) == [1, 1] assert (x + x + y + 2).get_grad([x, y]) == [2, 1] assert (x**2).get_grad([x]) == [6] x.set(4) assert z.get() == 5 assert z.get(run=True) == 6 assert z.get() == 6 y.set(x) assert z.get(run=True) == 8 with goos.OptimizationPlan(): assert z.get() == 5 goos.opt.scipy_minimize((x + y**2 + 1)**2 + (y + 1)**2, "CG") plan.run() plan.save(plan_dir) np.testing.assert_almost_equal(x.get().array, -2, decimal=4) np.testing.assert_almost_equal(y.get().array, -1, decimal=4) with goos.OptimizationPlan() as plan: plan.load(plan_dir) x = plan.get_node("x") y = plan.get_node("y") assert x.get() == 3 assert y.get() == 2 plan.run() np.testing.assert_almost_equal(x.get().array, -2, decimal=4) np.testing.assert_almost_equal(y.get().array, -1, decimal=4)
def test_product_triple(): with goos.OptimizationPlan() as plan: x = goos.Variable(3) y = goos.Variable(4) res = x * 2 * y assert res.get() == 24 assert res.get_grad([x, y])[0] == 8 assert res.get_grad([x, y])[1] == 6
def test_optimize_simple(): with goos.OptimizationPlan() as plan: x = goos.Variable([1]) obj = (x + 1)**2 + 3 goos.opt.scipy_minimize(obj, method="L-BFGS-B") plan.run() assert x.get().array == -1
def test_optimize_2d(): with goos.OptimizationPlan() as plan: x = goos.Variable([[1, 2], [3, 4]]) obj = goos.Norm(x - goos.Constant([[3, 2], [-4, 2]]))**2 + 3 goos.opt.scipy_minimize(obj, method="L-BFGS-B") plan.run() np.testing.assert_almost_equal(x.get().array, [[3, 2], [-4, 2]])
def test_sum_double_array(): with goos.OptimizationPlan() as plan: x = goos.Variable([2.0, 1.0]) y = goos.Variable([3.0, -1.0]) res = goos.dot([1, 1], x + y) np.testing.assert_allclose(res.get().array, 5) np.testing.assert_allclose(res.get_grad([x])[0].array_grad, [1, 1]) np.testing.assert_allclose(res.get_grad([y])[0].array_grad, [1, 1])
def test_optimize_ineq_constraints(): with goos.OptimizationPlan() as plan: x = goos.Variable(1) obj = (x + 1)**2 + 3 goos.opt.scipy_maximize(obj, constraints_ineq=[x - 5], method="SLSQP") plan.run() assert x.get().array == 5
def test_sigmoid_scalar(): with goos.OptimizationPlan() as plan: x = goos.Variable(0.5) y = goos.Sigmoid(x) np.testing.assert_allclose(y.get().array, 1 / (1 + np.exp(-0.5))) np.testing.assert_allclose( y.get_grad([x])[0].array_grad, np.exp(-0.5) / (1 + np.exp(-0.5))**2)
def test_optimize_upper_bounds(): with goos.OptimizationPlan() as plan: x = goos.Variable([1], upper_bounds=5) obj = (x + 1)**2 + 3 goos.opt.scipy_maximize(obj, method="L-BFGS-B") plan.run() assert x.get().array == 5
def test_max_two_array(): with goos.OptimizationPlan() as plan: x = goos.Variable([2, 1]) y = goos.Variable([3, 0.5]) z = goos.dot(goos.max(x, y), [1, -1]) np.testing.assert_array_equal(z.get().array, 2) np.testing.assert_array_equal(z.get_grad([x])[0].array_grad, [0, -1]) np.testing.assert_array_equal(z.get_grad([y])[0].array_grad, [1, 0])
def test_max_two_scalar(): with goos.OptimizationPlan() as plan: x = goos.Variable(2) y = goos.Variable(3) z = goos.max(x, y) assert z.get() == 3 assert z.get_grad([x])[0] == 0 assert z.get_grad([y])[0] == 1
def test_optimizer_eq_constraints(): with goos.OptimizationPlan() as plan: x = goos.Variable(1) y = goos.Variable(2) obj = (x * y - 12)**2 goos.opt.scipy_minimize(obj, constraints_eq=[y - 3], method="SLSQP") plan.run() np.testing.assert_allclose(x.get().array, 4)
def test_dot(): with goos.OptimizationPlan() as plan: x = goos.Variable([3, 2, 1]) y = goos.Variable([-1, 1, 4]) res = goos.dot(x, y) assert res.get() == 3 grad = res.get_grad([x, y]) np.testing.assert_array_equal(grad[0].array_grad, [-1, 1, 4]) np.testing.assert_array_equal(grad[1].array_grad, [3, 2, 1])
def test_plan_heavy(tmp_path): # This is the same test as `test_plan` excep that we replace some nodes # with heavy node implementations by directly manipulating the flags. plan_dir = tmp_path / "test_plan" plan_dir.mkdir() with goos.OptimizationPlan() as plan: x = goos.Variable(3.0, name="x") y = goos.Variable(2.0, name="y") z = x + y z.parallelize() assert z.get() == 5 assert z.get_grad([x, y]) == [1, 1] assert (x + x + y + 2).get_grad([x, y]) == [2, 1] assert (x**2).get_grad([x]) == [6] x.set(4) assert z.get() == 5 assert z.get(run=True) == 6 assert z.get() == 6 y.set(x) assert z.get(run=True) == 8 with goos.OptimizationPlan(): assert z.get() == 5 first_part = x + y**2 + 1 first_part.parallelize() second_part = y + 1 second_part.parallelize() obj = first_part**2 + second_part**2 goos.opt.scipy_minimize(obj, "CG") plan.run() plan.save(plan_dir) np.testing.assert_almost_equal(x.get().array, -2, decimal=4) np.testing.assert_almost_equal(y.get().array, -1, decimal=4)