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_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_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_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_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_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_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_max_two_array_2d(): with goos.OptimizationPlan() as plan: x = goos.Variable([[1, 2], [3, 4]]) y = goos.Variable([[3, 0.5], [4.5, 6]]) z = goos.dot([[1, 2], [3, 4]], goos.max(x, y)) np.testing.assert_array_equal(z.get().array, 44.5) np.testing.assert_array_equal( z.get_grad([x])[0].array_grad, [[0, 2], [0, 0]]) np.testing.assert_array_equal( z.get_grad([y])[0].array_grad, [[1, 0], [3, 4]])
def test_max_three_scalars(): with goos.OptimizationPlan() as plan: x = goos.Variable(2) y = goos.Variable(3) w = goos.Variable(1) z = goos.max(x, y, w) assert z.get() == 3 assert z.get_grad([x])[0] == 0 assert z.get_grad([y])[0] == 1 assert z.get_grad([w])[0] == 0
def test_optimize_freeze_check(): with goos.OptimizationPlan() as plan: x = goos.Variable([1]) y = goos.Variable([1]) y.freeze() obj = (x + y)**2 + 3 goos.opt.scipy_minimize(obj, method="L-BFGS-B") plan.run() assert x.get().array == -1 assert y.get().array == 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_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 compat_param(param: optplan.Parametrization, initializer: Callable, extents: List[float], pixel_size: List[float], var_name: Optional[str] = None, **kwargs) -> Tuple[goos.Variable, OldStyleParametrization]: """Creates a parametrization based on `optplan.Parametrization`. Note that some arguments of `optplan.Parametrization` may be ignored. Some arguments may need to be omitted (e.g. `optplan.SimulationSpace`). Args: param: Old style parametrization. initializer: A callable that accepts a shape and returns an array with that shape. This is used to initialize the parametrization variable. extents: Extents of the shape. pixel_size: Size of each pixel in the shape. var_name: Name to give variable. **kwargs: Additional arguments to pass to `OldStyleParametrization`. Returns: A tuple `(var, eps)` where `var` is the variable controlling the parametrization and `eps` is the shape node. """ _, var_shape, lower_bound, upper_bound = create_old_param( param, extents, pixel_size) var = goos.Variable(initializer(var_shape), name=var_name, lower_bounds=lower_bound, upper_bounds=upper_bound) return var, OldStyleParametrization(var=var, param=param, extents=extents, pixel_size=pixel_size, **kwargs)
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_parameter_frozen_but_settable(): """Checks the parameters are always frozen but are always settable.""" with goos.OptimizationPlan() as plan: x = goos.Variable(2, lower_bounds=0) param = goos.Variable(3, parameter=True) # No error should be thrown here. param.set(4) obj = (x + param)**2 goos.opt.scipy_minimize(obj, max_iters=3, method="L-BFGS-B") plan.run() # Check that `x` is optimized by `param` is constant. assert param.get() == 4 assert x.get() == 0
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_pixelated_multietch_grating_fit_edge_locs(): # Grating defined in xz-plane. with goos.OptimizationPlan() as plan: height_index = goos.Variable([2, 0, 1, 2], parameter=True) edge_locs_target = goos.Variable([0, 1, 1.5, 2], parameter=True) grating_target = goos.grating.PixelatedGrating( edge_locs_target, 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, use_edge_locs=True) edge_locs = goos.Variable([0, 0.8, 1.6, 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, use_edge_locs=True) from spins.goos_sim import maxwell region = goos.Box3d(center=[0, 0, 0], extents=[3, 11, 0.5]) mesh = maxwell.UniformMesh(dx=0.1) eps_diff = goos.Norm( maxwell.RenderShape( grating, region=region, mesh=mesh, wavelength=1.55) - maxwell.RenderShape( grating_target, region=region, mesh=mesh, wavelength=1.55)) goos.opt.scipy_minimize(eps_diff, "L-BFGS-B", max_iters=20) plan.run() np.testing.assert_almost_equal(grating.get().array, grating_target.get().array)
def test_pixelated_multietch_grating_xz_gradient(use_edge_locs, params, pol): # Brute force calculate the gradient. # Grating defined in xz-plane. with goos.OptimizationPlan() as plan: edge_locs = goos.Variable(params) height_index = goos.Variable([0, 1, 0, 1]) grating = goos.grating.PixelatedGrating( edge_locs, height_index=height_index, height_fracs=[0.4, 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=4, etch_polarity=pol, use_edge_locs=use_edge_locs) from spins.goos_sim import maxwell shape = maxwell.RenderShape(grating, region=goos.Box3d(center=[0, 0, 0], extents=[4, 10, 0.220]), mesh=maxwell.UniformMesh(dx=0.25), wavelength=1550) np.random.seed(247) obj = goos.dot(shape, np.random.random(shape.get().array.shape)) val_orig = edge_locs.get().array step = 1e-5 grad_num = np.zeros_like(val_orig) for i in range(len(val_orig)): delta = np.zeros_like(val_orig) delta[i] = 1 edge_locs.set(val_orig + step * delta) f_plus = obj.get(run=True).array edge_locs.set(val_orig - step * delta) f_minus = obj.get(run=True).array grad_num[i] = (f_plus - f_minus) / (2 * step) grad_backprop = obj.get_grad([edge_locs])[0].array_grad np.testing.assert_allclose(grad_backprop, grad_num)
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_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_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_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_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_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_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_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_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