control

import%20marimo%0A%0A__generated_with%20%3D%20%220.25.2%22%0Aapp%20%3D%20marimo.App(width%3D%22medium%22)%0A%0A%0A%40app.cell(hide_code%3DTrue)%0Adef%20_()%3A%0A%20%20%20%20import%20statistics%0A%20%20%20%20import%20sys%0A%20%20%20%20import%20time%0A%20%20%20%20from%20pathlib%20import%20Path%0A%0A%20%20%20%20import%20marimo%20as%20mo%0A%20%20%20%20import%20numpy%20as%20np%0A%0A%20%20%20%20start%20%3D%20Path(__file__).resolve().parent%20if%20%22__file__%22%20in%20globals()%20else%20Path.cwd().resolve()%0A%20%20%20%20for%20candidate%20in%20(start%2C%20*start.parents)%3A%0A%20%20%20%20%20%20%20%20if%20(candidate%20%2F%20%22fomtools%22).is_dir()%3A%0A%20%20%20%20%20%20%20%20%20%20%20%20repo_root%20%3D%20candidate%0A%20%20%20%20%20%20%20%20%20%20%20%20break%0A%20%20%20%20else%3A%0A%20%20%20%20%20%20%20%20repo_root%20%3D%20start%0A%0A%20%20%20%20if%20str(repo_root)%20not%20in%20sys.path%3A%0A%20%20%20%20%20%20%20%20sys.path.insert(0%2C%20str(repo_root))%0A%0A%20%20%20%20examples_dir%20%3D%20repo_root%20%2F%20%22examples%22%0A%0A%20%20%20%20from%20docs.tutorials.helper_functions.control%20import%20admm_step%2C%20solve_admm_numpy%0A%20%20%20%20from%20docs.tutorials.helper_functions.codegen_benchmarks%20import%20run_codegen_benchmark%0A%20%20%20%20from%20fomtools.codegen%20import%20CodegenConfig%2C%20TerminationSpec%0A%20%20%20%20import%20fomtools%20as%20ft%0A%0A%20%20%20%20timing_file%20%3D%20start%20%2F%20%22_recorded_timings.json%22%0A%0A%20%20%20%20np.set_printoptions(precision%3D5)%0A%20%20%20%20np.set_printoptions(suppress%3DTrue)%0A%20%20%20%20return%20(%0A%20%20%20%20%20%20%20%20CodegenConfig%2C%0A%20%20%20%20%20%20%20%20TerminationSpec%2C%0A%20%20%20%20%20%20%20%20admm_step%2C%0A%20%20%20%20%20%20%20%20examples_dir%2C%0A%20%20%20%20%20%20%20%20ft%2C%0A%20%20%20%20%20%20%20%20mo%2C%0A%20%20%20%20%20%20%20%20np%2C%0A%20%20%20%20%20%20%20%20run_codegen_benchmark%2C%0A%20%20%20%20%20%20%20%20solve_admm_numpy%2C%0A%20%20%20%20%20%20%20%20statistics%2C%0A%20%20%20%20%20%20%20%20time%2C%0A%20%20%20%20%20%20%20%20timing_file%2C%0A%20%20%20%20)%0A%0A%0A%40app.cell(hide_code%3DTrue)%0Adef%20_(mo)%3A%0A%20%20%20%20mo.md(r%22%22%22%0A%20%20%20%20%23%20Optimal%20Control%3A%20Python%20vs%20Codegen%0A%0A%20%20%20%20%5BBack%20to%20the%20code%20generation%20tutorials%5D(.%2F)%0A%0A%20%20%20%20In%20this%20tutorial%20we%20benchmark%20a%20small%20finite-horizon%20control%20problem%20for%20a%20linear%20dynamical%20system.%0A%20%20%20%20The%20solver%20uses%20the%20alternating%20direction%20method%20of%20multipliers%3A%20each%20iteration%20solves%20a%20small%20linear%20system%2C%20projects%20onto%20box%20constraints%2C%20and%20updates%20the%20auxiliary%20variables.%0A%0A%20%20%20%20We%20compare%20a%20direct%20NumPy%20implementation%20with%20the%20verified%20generated%20C%20solver.%0A%20%20%20%20The%20ADMM%20step%20and%20NumPy%20reference%20solver%20are%20shared%20with%20the%20%5Bcontrol%20verification%20tutorial%5D(..%2Fverification%2Fcontrol.html).%0A%20%20%20%20The%20generated%20C%20is%20timed%20after%20code%20generation%2C%20then%20timed%20again%20after%20the%20C-editing%20optimization%20pass.%0A%20%20%20%20%22%22%22)%0A%20%20%20%20return%0A%0A%0A%40app.cell(hide_code%3DTrue)%0Adef%20_(mo)%3A%0A%20%20%20%20mo.md(r%22%22%22%0A%20%20%20%20We%20fix%20the%20initial%20state%2C%20input%20bounds%2C%20prediction%20horizon%2C%20and%20quadratic%20costs.%0A%20%20%20%20Both%20implementations%20run%20the%20same%20update%20rule%20for%20the%20same%20number%20of%20iterations.%0A%20%20%20%20%22%22%22)%0A%20%20%20%20return%0A%0A%0A%40app.cell%0Adef%20_(examples_dir%2C%20np)%3A%0A%20%20%20%20A_dyn%20%3D%20np.array(%0A%20%20%20%20%20%20%20%20%5B%0A%20%20%20%20%20%20%20%20%20%20%20%20%5B1.0%2C%200.0%2C%200.0%2C%200.0%2C%200.0%2C%200.0%2C%200.1%2C%200.0%2C%200.0%2C%200.0%2C%200.0%2C%200.0%5D%2C%0A%20%20%20%20%20%20%20%20%20%20%20%20%5B0.0%2C%201.0%2C%200.0%2C%200.0%2C%200.0%2C%200.0%2C%200.0%2C%200.1%2C%200.0%2C%200.0%2C%200.0%2C%200.0%5D%2C%0A%20%20%20%20%20%20%20%20%20%20%20%20%5B0.0%2C%200.0%2C%201.0%2C%200.0%2C%200.0%2C%200.0%2C%200.0%2C%200.0%2C%200.1%2C%200.0%2C%200.0%2C%200.0%5D%2C%0A%20%20%20%20%20%20%20%20%20%20%20%20%5B0.0488%2C%200.0%2C%200.0%2C%201.0%2C%200.0%2C%200.0%2C%200.0016%2C%200.0%2C%200.0%2C%200.0992%2C%200.0%2C%200.0%5D%2C%0A%20%20%20%20%20%20%20%20%20%20%20%20%5B0.0%2C%20-0.0488%2C%200.0%2C%200.0%2C%201.0%2C%200.0%2C%200.0%2C%20-0.0016%2C%200.0%2C%200.0%2C%200.0992%2C%200.0%5D%2C%0A%20%20%20%20%20%20%20%20%20%20%20%20%5B0.0%2C%200.0%2C%200.0%2C%200.0%2C%200.0%2C%201.0%2C%200.0%2C%200.0%2C%200.0%2C%200.0%2C%200.0%2C%200.0992%5D%2C%0A%20%20%20%20%20%20%20%20%20%20%20%20%5B0.0%2C%200.0%2C%200.0%2C%200.0%2C%200.0%2C%200.0%2C%201.0%2C%200.0%2C%200.0%2C%200.0%2C%200.0%2C%200.0%5D%2C%0A%20%20%20%20%20%20%20%20%20%20%20%20%5B0.0%2C%200.0%2C%200.0%2C%200.0%2C%200.0%2C%200.0%2C%200.0%2C%201.0%2C%200.0%2C%200.0%2C%200.0%2C%200.0%5D%2C%0A%20%20%20%20%20%20%20%20%20%20%20%20%5B0.0%2C%200.0%2C%200.0%2C%200.0%2C%200.0%2C%200.0%2C%200.0%2C%200.0%2C%201.0%2C%200.0%2C%200.0%2C%200.0%5D%2C%0A%20%20%20%20%20%20%20%20%20%20%20%20%5B0.9734%2C%200.0%2C%200.0%2C%200.0%2C%200.0%2C%200.0%2C%200.0488%2C%200.0%2C%200.0%2C%200.9846%2C%200.0%2C%200.0%5D%2C%0A%20%20%20%20%20%20%20%20%20%20%20%20%5B0.0%2C%20-0.9734%2C%200.0%2C%200.0%2C%200.0%2C%200.0%2C%200.0%2C%20-0.0488%2C%200.0%2C%200.0%2C%200.9846%2C%200.0%5D%2C%0A%20%20%20%20%20%20%20%20%20%20%20%20%5B0.0%2C%200.0%2C%200.0%2C%200.0%2C%200.0%2C%200.0%2C%200.0%2C%200.0%2C%200.0%2C%200.0%2C%200.0%2C%200.9846%5D%2C%0A%20%20%20%20%20%20%20%20%5D%0A%20%20%20%20)%0A%20%20%20%20B_dyn%20%3D%20np.array(%0A%20%20%20%20%20%20%20%20%5B%0A%20%20%20%20%20%20%20%20%20%20%20%20%5B0.0%2C%20-0.0726%2C%200.0%2C%200.0726%5D%2C%0A%20%20%20%20%20%20%20%20%20%20%20%20%5B-0.0726%2C%200.0%2C%200.0726%2C%200.0%5D%2C%0A%20%20%20%20%20%20%20%20%20%20%20%20%5B-0.0152%2C%200.0152%2C%20-0.0152%2C%200.0152%5D%2C%0A%20%20%20%20%20%20%20%20%20%20%20%20%5B-0.0%2C%20-0.0006%2C%20-0.0%2C%200.0006%5D%2C%0A%20%20%20%20%20%20%20%20%20%20%20%20%5B0.0006%2C%200.0%2C%20-0.0006%2C%200.0%5D%2C%0A%20%20%20%20%20%20%20%20%20%20%20%20%5B0.0106%2C%200.0106%2C%200.0106%2C%200.0106%5D%2C%0A%20%20%20%20%20%20%20%20%20%20%20%20%5B0.0%2C%20-1.4512%2C%200.0%2C%201.4512%5D%2C%0A%20%20%20%20%20%20%20%20%20%20%20%20%5B-1.4512%2C%200.0%2C%201.4512%2C%200.0%5D%2C%0A%20%20%20%20%20%20%20%20%20%20%20%20%5B-0.3049%2C%200.3049%2C%20-0.3049%2C%200.3049%5D%2C%0A%20%20%20%20%20%20%20%20%20%20%20%20%5B-0.0%2C%20-0.0236%2C%200.0%2C%200.0236%5D%2C%0A%20%20%20%20%20%20%20%20%20%20%20%20%5B0.0236%2C%200.0%2C%20-0.0236%2C%200.0%5D%2C%0A%20%20%20%20%20%20%20%20%20%20%20%20%5B0.2107%2C%200.2107%2C%200.2107%2C%200.2107%5D%2C%0A%20%20%20%20%20%20%20%20%5D%0A%20%20%20%20)%0A%0A%20%20%20%20n_state%2C%20n_input%20%3D%20B_dyn.shape%0A%20%20%20%20horizon%20%3D%202%0A%20%20%20%20K%20%3D%20200%0A%0A%20%20%20%20Q%20%3D%20np.diag(%5B0.0%2C%200.0%2C%2010.0%2C%2010.0%2C%2010.0%2C%2010.0%2C%200.0%2C%200.0%2C%200.0%2C%205.0%2C%205.0%2C%205.0%5D)%0A%20%20%20%20R%20%3D%200.1%20*%20np.eye(n_input)%0A%20%20%20%20x_initial%20%3D%202.0%20*%20np.array(%5B1.0%2C%201.0%2C%201.0%2C%200.0%2C%200.0%2C%200.0%2C%200.0%2C%200.0%2C%200.0%2C%200.0%2C%200.0%2C%200.0%5D)%0A%20%20%20%20u_min%20%3D%20-0.99%20*%20np.ones(n_input)%0A%20%20%20%20u_max%20%3D%202.41%20*%20np.ones(n_input)%0A%0A%20%20%20%20n_control%20%3D%20(horizon%20-%201)%20*%20n_input%0A%20%20%20%20n_total%20%3D%20n_state%20%2B%20n_control%0A%20%20%20%20Sx%20%3D%20np.vstack(%5Bnp.linalg.matrix_power(A_dyn%2C%20i)%20for%20i%20in%20range(horizon)%5D)%0A%20%20%20%20Sz_mid%20%3D%20np.zeros(((horizon%20-%201)%20*%20n_state%2C%20n_control))%0A%20%20%20%20for%20i%20in%20range(horizon%20-%201)%3A%0A%20%20%20%20%20%20%20%20Sz_mid%5Bi%20*%20n_state%20%3A%20(i%20%2B%201)%20*%20n_state%2C%20%3A%5D%20%3D%20np.linalg.matrix_power(A_dyn%2C%20i)%20%40%20B_dyn%0A%20%20%20%20Sz%20%3D%20np.vstack(%5Bnp.zeros((n_state%2C%20n_control))%2C%20Sz_mid%5D)%0A%0A%20%20%20%20Q_block%20%3D%20np.kron(np.eye(horizon)%2C%20Q)%0A%20%20%20%20R_block%20%3D%20np.kron(np.eye(horizon%20-%201)%2C%20R)%0A%20%20%20%20P%20%3D%20np.block(%0A%20%20%20%20%20%20%20%20%5B%0A%20%20%20%20%20%20%20%20%20%20%20%20%5BSx.T%20%40%20Q_block%20%40%20Sx%2C%20Sx.T%20%40%20Q_block%20%40%20Sz%5D%2C%0A%20%20%20%20%20%20%20%20%20%20%20%20%5BSz.T%20%40%20Q_block%20%40%20Sx%2C%20Sz.T%20%40%20Q_block%20%40%20Sz%20%2B%20R_block%5D%2C%0A%20%20%20%20%20%20%20%20%5D%0A%20%20%20%20)%0A%0A%20%20%20%20M%20%3D%20np.eye(n_total)%0A%20%20%20%20lower%20%3D%20np.hstack(%5Bx_initial%2C%20np.kron(np.ones(horizon%20-%201)%2C%20u_min)%5D)%0A%20%20%20%20upper%20%3D%20np.hstack(%5Bx_initial%2C%20np.kron(np.ones(horizon%20-%201)%2C%20u_max)%5D)%0A%20%20%20%20sigma%20%3D%201e-4%0A%20%20%20%20rho%20%3D%201.0%0A%20%20%20%20lhs%20%3D%20P%20%2B%20sigma%20*%20np.eye(n_total)%20%2B%20rho%20*%20(M.T%20%40%20M)%0A%20%20%20%20q%20%3D%20np.zeros(n_total)%0A%0A%20%20%20%20z0%20%3D%20np.zeros(n_total)%0A%20%20%20%20v0%20%3D%200.5%20*%20np.ones(n_total)%0A%20%20%20%20out_dir%20%3D%20examples_dir%20%2F%20%22_codegen_speedups%22%20%2F%20%22control_codegen_speedup%22%0A%20%20%20%20return%20K%2C%20M%2C%20lhs%2C%20lower%2C%20n_total%2C%20out_dir%2C%20q%2C%20rho%2C%20sigma%2C%20upper%2C%20v0%2C%20z0%0A%0A%0A%40app.cell(hide_code%3DTrue)%0Adef%20_(%0A%20%20%20%20CodegenConfig%2C%0A%20%20%20%20K%2C%0A%20%20%20%20M%2C%0A%20%20%20%20TerminationSpec%2C%0A%20%20%20%20admm_step%2C%0A%20%20%20%20ft%2C%0A%20%20%20%20lhs%2C%0A%20%20%20%20lower%2C%0A%20%20%20%20n_total%2C%0A%20%20%20%20q%2C%0A%20%20%20%20rho%2C%0A%20%20%20%20sigma%2C%0A%20%20%20%20upper%2C%0A%20%20%20%20v0%2C%0A%20%20%20%20z0%2C%0A)%3A%0A%20%20%20%20verifier%20%3D%20ft.Verifier(%0A%20%20%20%20%20%20%20%20method%3D%22mip%22%2C%0A%20%20%20%20%20%20%20%20solver%3D%22highs%22%2C%0A%20%20%20%20%20%20%20%20obbt_solver%3DNone%2C%0A%20%20%20%20)%0A%0A%20%20%20%20z%20%3D%20verifier.iterate(n_total%2C%20lb%3Dz0%2C%20ub%3Dz0)%0A%20%20%20%20v%20%3D%20verifier.iterate(n_total%2C%20lb%3Dv0%2C%20ub%3Dv0)%0A%20%20%20%20def%20project(vp%2C%20v_prev)%3A%0A%20%20%20%20%20%20%20%20return%20vp.satlin(v_prev%2C%20l%3Dlower%2C%20u%3Dupper)%0A%0A%20%20%20%20z_next%2C%20v_next%2C%20_%20%3D%20admm_step(%0A%20%20%20%20%20%20%20%20verifier%2C%20M%2C%20q%2C%20rho%2C%20sigma%2C%20z%2C%20v%2C%20project%2C%20lhs%0A%20%20%20%20)%0A%0A%20%20%20%20config%20%3D%20CodegenConfig(%0A%20%20%20%20%20%20%20%20termination%3DTerminationSpec(max_iters%3DK)%2C%0A%20%20%20%20%20%20%20%20num_verify_samples%3D4%2C%0A%20%20%20%20)%0A%20%20%20%20next_state%20%3D%20%5Bz_next%2C%20v_next%5D%0A%20%20%20%20s_state%20%3D%20%5Bz%2C%20v%5D%0A%20%20%20%20return%20config%2C%20next_state%2C%20s_state%2C%20verifier%0A%0A%0A%40app.cell%0Adef%20_(%0A%20%20%20%20K%2C%0A%20%20%20%20M%2C%0A%20%20%20%20lhs%2C%0A%20%20%20%20lower%2C%0A%20%20%20%20q%2C%0A%20%20%20%20rho%2C%0A%20%20%20%20sigma%2C%0A%20%20%20%20solve_admm_numpy%2C%0A%20%20%20%20statistics%2C%0A%20%20%20%20time%2C%0A%20%20%20%20upper%2C%0A%20%20%20%20v0%2C%0A%20%20%20%20z0%2C%0A)%3A%0A%20%20%20%20def%20control_solve()%3A%0A%20%20%20%20%20%20%20%20return%20solve_admm_numpy(%0A%20%20%20%20%20%20%20%20%20%20%20%20M%2C%20q%2C%20rho%2C%20sigma%2C%20lhs%2C%20lower%2C%20upper%2C%20z0%2C%20v0%2C%20max_iters%3DK%0A%20%20%20%20%20%20%20%20)%0A%0A%20%20%20%20for%20_%20in%20range(3)%3A%0A%20%20%20%20%20%20%20%20control_solve()%0A%0A%20%20%20%20python_times%20%3D%20%5B%5D%0A%20%20%20%20for%20_%20in%20range(30)%3A%0A%20%20%20%20%20%20%20%20t0%20%3D%20time.perf_counter()%0A%20%20%20%20%20%20%20%20control_solve()%0A%20%20%20%20%20%20%20%20python_times.append((time.perf_counter()%20-%20t0)%20*%201.0e6)%0A%0A%20%20%20%20python_us%20%3D%20statistics.median(python_times)%0A%20%20%20%20return%20(python_us%2C)%0A%0A%0A%40app.cell(hide_code%3DTrue)%0Adef%20_(%0A%20%20%20%20config%2C%0A%20%20%20%20next_state%2C%0A%20%20%20%20out_dir%2C%0A%20%20%20%20python_us%2C%0A%20%20%20%20run_codegen_benchmark%2C%0A%20%20%20%20s_state%2C%0A%20%20%20%20timing_file%2C%0A%20%20%20%20verifier%2C%0A)%3A%0A%20%20%20%20timings%20%3D%20run_codegen_benchmark(%0A%20%20%20%20%20%20%20%20config%3Dconfig%2C%0A%20%20%20%20%20%20%20%20example%3D%22control%22%2C%0A%20%20%20%20%20%20%20%20name%3D%22control_codegen_speedup%22%2C%0A%20%20%20%20%20%20%20%20next_state%3Dnext_state%2C%0A%20%20%20%20%20%20%20%20out_dir%3Dout_dir%2C%0A%20%20%20%20%20%20%20%20params%3DNone%2C%0A%20%20%20%20%20%20%20%20python_us%3Dpython_us%2C%0A%20%20%20%20%20%20%20%20state%3Ds_state%2C%0A%20%20%20%20%20%20%20%20timing_file%3Dtiming_file%2C%0A%20%20%20%20%20%20%20%20verifier%3Dverifier%2C%0A%20%20%20%20)%0A%20%20%20%20return%20(timings%2C)%0A%0A%0A%40app.cell(hide_code%3DTrue)%0Adef%20_(K%2C%20mo%2C%20out_dir%2C%20timings)%3A%0A%20%20%20%20if%20timings%5B%22live%22%5D%3A%0A%20%20%20%20%20%20%20%20_footer%20%3D%20f%22%22%22%0A%20%20%20%20Generated%20solver%20directory%3A%20%60%7Bout_dir%7D%60%0A%0A%20%20%20%20Generated%20solver%20source%3A%20%60%7Btimings%5B%22solver_path%22%5D%7D%60%0A%20%20%20%20%22%22%22%0A%20%20%20%20else%3A%0A%20%20%20%20%20%20%20%20_footer%20%3D%20%22%22%0A%20%20%20%20summary%20%3D%20mo.md(f%22%22%22%0A%20%20%20%20%23%23%20Timing%20summary%0A%0A%20%20%20%20One%20complete%20solve%20runs%20%60%7BK%7D%60%20iterations.%0A%0A%20%20%20%20%7C%20implementation%20%7C%20median%20%2F%20mean%20per%20solve%20%7C%0A%20%20%20%20%7C---%7C---%3A%7C%0A%20%20%20%20%7C%20direct%20NumPy%20Python%20%7C%20%60%7Btimings%5B%22python_us%22%5D%3A.3f%7D%20%C2%B5s%60%20%7C%0A%20%20%20%20%7C%20generated%20C%20after%20code%20generation%20%7C%20%60%7Btimings%5B%22codegen_us%22%5D%3A.3f%7D%20%C2%B5s%60%20%7C%0A%20%20%20%20%7C%20generated%20C%20after%20optimization%20iteration%20%7C%20%60%7Btimings%5B%22iteration_us%22%5D%3A.3f%7D%20%C2%B5s%60%20%7C%0A%0A%20%20%20%20The%20speedup%20from%20code%20generation%20is%20%60%7Btimings%5B%22codegen_speedup%22%5D%3A.1f%7Dx%60%2C%20while%20the%20speedup%20from%20the%20optimization%20iteration%20is%20%60%7Btimings%5B%22iteration_speedup%22%5D%3A.2f%7Dx%60%2C%20totaling%20a%20speedup%20of%20%60%7Btimings%5B%22total_speedup%22%5D%3A.1f%7Dx%60.%0A%20%20%20%20%7B_footer%7D%22%22%22%0A%20%20%20%20)%0A%20%20%20%20summary%0A%20%20%20%20return%0A%0A%0Aif%20__name__%20%3D%3D%20%22__main__%22%3A%0A%20%20%20%20app.run()%0A