B @`,@sddlZddlZddlZddlZddlmZddlmZmZm Z m Z m Z ddl m Z ddlmZmZmZmZmZmZmZmZddlmZmZddlmZmZmZddlmZdd l m!Z!dd l"m#Z#m$Z$m%Z%dd l&m'Z'ej()e*Z+e+,d d ,dd dZ-dZ.Gddde Z/Gddde/e Z0Gddde/e Z1ddZ2ie2_3e0e2j3d<e1e2j3d<ddZ4ddZ5dd Z6d!d"Z7d#d$Z8d%d&Z9d'd(Z:d)d*Z;d+d,ZGd1d2d2e Z?Gd3d4d4e#e Z@d=d5d6ZAd7d8ZBd>d9d:ZCe%Dd;e@eEdg@g?) r#r7r$ _z_previous_primal_residual_norm_dual_residual_norm_mu_tol_rho_increase_factor_rho_decrease_factor)r%r&r')r(r r!r$VszAdaptiveRhoStrategy.__init__cCs`xZ|jjddD]F}x@|jD]6}|j|j}||j}x|jD]}||||<q@WqWqWdS)N) scenario_treestagesnodesr9name_standard_variable_ids)r%r*r-stage tree_node z_prev_nodez_nodeid_r r r! snapshot_z`s     zAdaptiveRhoStrategy.snapshot_zc Csx|jjddD]}xz|jD]p}|j|j}||j}xT|jD]J}d} x2|jD](} | || j|j|||d7} qPWt| ||<q@WqWqWdS)Nr@g) rArBrCr:rDrEZ _scenariosmathsqrt) r%r*r+r-rFrG prnorm_noderIrJ prnorm_varscenarior r r!compute_primal_residual_normhs     z0AdaptiveRhoStrategy.compute_primal_residual_normc Csx|jjddD]}xz|jD]p}|j|j}|j|j}||j}||j} x>|jD]4} tt |j | | t || || || <qVWqWqWdS)Nr@) rArBrCr;rDr9rErMrNlen scenariosabs) r%r*r-r.rFrG drnorm_noderHrIrho_noderJr r r!compute_dual_residual_normus      z.AdaptiveRhoStrategy.compute_dual_residual_normcCspi|_i|_i|_xL|jjddD]8}x2|jD](}i|j|j<i|j|j<i|j|j<q0Wq$W|||dS)Nr@)r9r:r;rArBrCrDrK)r%r*r+r,r-r.rFrGr r r!r/s   zAdaptiveRhoStrategy.initializec Cshdd}d}|||||||||d}x"|jjddD] } x| jD]} |j| j} |j| j} || j} x| j D]}| j |\}}| |}| |}d}||j |kr||j kr| ||j 9<d}n2||j |kr ||j kr | ||j9<d}|r|dk r|r*d}t|td ||t|||| |fqWqVWqFW|||dS) Nz-Updating Rho Values: %21s %25s %16s %16s %16s)ActionVariablezPrimal Residualz Dual ResidualzNew RhoTr8r@Z IncreasingZ DecreasingFz%21s %25s %16g %16g %16g)rRrX get_optionrArBrCr:rDr;rE _variable_idsr<r=r>r?printrrK)r%r*r+r,r-r. first_lineZ first_printr8rFrGrOrVrWrJrDindexrPZ drnorm_varactionr r r!r0sH        zAdaptiveRhoStrategy.update_rho)N) r1r2r3r4r"r$rKrRrXr/r0r5r r )r(r!r7Os     r7cOs*|tjkrtj|||Std|dS)Nz(No rho strategy registered with name: %s)RhoStrategyFactoryregistered_types ValueError)rDr&r'r r r!ras raZfixedadaptivecCsV|drtd|j|j}t|jds.t|jd|jdk rRd|jj|j<dS)Nr8z-Cleaning up admm modifications to scenario %sz.admmT) r[r]rD _instancehasattrAssertionErrorZ del_component preprocessorobjective_updated)r*rQinstancer r r!EXTERNAL_cleanup_for_admms    rkc Cs|drtd|jtdd}t|jdr2t|jd||j}t |j |d|_ t dd|_ t dd|_ t |j d|_t |j d|_t |j d|_d}|jtkrd }|j|j||j||j_x|jdd D]}|rttdd}||j|tdt|jd |_t|jddd |_t|jddd |_t|jddd |_x~|jD]t}|j|\}} |j || } |j j|j|| |j|7_|j j|j|d | |j|d 7_qVWqW|j!dk rd|j!j"|j<dS)Nr8z+Initializing scenario %s for admm algorithmT)Zconcretez.admm)r/gr@)Zorderedr/)r/Zmutableg@rL)#r[r]rDrrfrergZ add_component_instance_cost_expressionrZ probabilitycost_expressionlagrangian_expressionpenalty_expression cost_termlagrangian_term penalty_termZobjective_senserZ_instance_objectiveexpr node_list is_leaf_nodersortedrEZnode_index_setrr-r,r.r\find_componentrhri) r*rQ admm_blockZuser_cost_expressionZobjective_directionrG node_blockrJvarnamer_varr r r!EXTERNAL_initialize_for_admms`               . r}cCsZ|drtd|jt|jds(tt|jd}|j|j_ |j dk rVd|j j |j<dS)Nr8z.Activating admm lagrangian term on scenario %sz.admmT) r[r]rDrfrerggetattrrorrrtrhri)r*rQryr r r!!EXTERNAL_activate_lagrangian_terms     rcCsX|drtd|jt|jds(tt|jd}d|j_|j dk rTd|j j |j<dS)Nr8z0Deactivating admm lagrangian term on scenario %sz.admmgT) r[r]rDrfrergr~rrrtrhri)r*rQryr r r!#EXTERNAL_deactivate_lagrangian_term&s    rcCsZ|drtd|jt|jds(tt|jd}|j|j_ |j dk rVd|j j |j<dS)Nr8z+Activating admm penalty term on scenario %sz.admmT) r[r]rDrfrergr~rprsrtrhri)r*rQryr r r!EXTERNAL_activate_penalty_term2s     rcCsX|drtd|jt|jds(tt|jd}d|j_|j dk rTd|j j |j<dS)Nr8z-Deactivating admm penalty term on scenario %sz.admmgT) r[r]rDrfrergr~rsrtrhri)r*rQryr r r! EXTERNAL_deactivate_penalty_term?s    rcCsZ|drtd|jt|jds(tt|jd}|j|j_ |j dk rVd|j j |j<dS)Nr8z=Activating admm probability-weighted cost term on scenario %sz.admmT) r[r]rDrfrergr~rnrqrtrhri)r*rQryr r r!0EXTERNAL_activate_probability_weighted_cost_termKs     rcCsZ|drtd|jt|jds(tt|jd}|j|j_ |j dk rVd|j j |j<dS)Nr8z?Deactivating admm probability-weighted cost term on scenario %sz.admmT) r[r]rDrfrergr~rmrqrtrhri)r*rQryr r r!2EXTERNAL_deactivate_probability_weighted_cost_termXs     rcCs|drtd|jt|jds(tt|jd}xL|jddD]:}|rTt| |j}|dk slt|j ||jqDW|j dk rd|j j |j<dS)Nr8z+Updating admm rho parameter for scenario %sz.admmr@T)r[r]rDrfrergr~rurvrxr. store_valuesrhri)r*rQr.ryrGrzr r r!EXTERNAL_update_rhoes       rcCs|drtd|jt|jds(tt|jd}xL|jddD]:}|rTt| |j}|dk slt|j ||jqDW|j dk rd|j j |j<dS)Nr8z)Updating admm y parameter for scenario %sz.admmr@T)r[r]rDrfrergr~rurvrxr,rrhri)r*rQr,ryrGrzr r r!EXTERNAL_update_yvs       rcCs|drtd|jt|jds(tt|jd}xL|jddD]:}|rTt| |j}|dk slt|j ||jqDW|j dk rd|j j |j<dS)Nr8z)Updating admm z parameter for scenario %sz.admmr@T)r[r]rDrfrergr~rurvrxr-rrhri)r*rQr-ryrGrzr r r!EXTERNAL_update_zs       rcseZdZed(ddZddZddZdd Zfd d Zd d Z ddZ ddZ ddZ ddZ ddZddZddZddZddZd d!Zd)d$d%Zd&d'ZZS)* ADMMAlgorithmNcCs0|dkrt}t|dtdtj|ddd|S)Nr8)ap_group subproblem_F)options_prefixZsetup_argparse)r r _admm_group_labelrregister_options)rrr r r!r"szADMMAlgorithm._declare_optionscCs|S)Nr )r%r r r! __enter__szADMMAlgorithm.__enter__cGs |dS)N)close)r%r&r r r!__exit__szADMMAlgorithm.__exit__cCs6|jdk r||jdk r&|jd|_d|_dS)N)_managercleanup_subproblems_manager_solverr)r%r r r!rs    zADMMAlgorithm.closecshtt|j||t|ts,td|jj|jsBt d|jj||_ t |j |j dd|_ |dS)NzY%s requires an instance of the ScenarioTreeManagerSolver interface as the second argumentzC%s requires a scenario tree manager that has been fully initializedr)r)r#rr$ isinstancer TypeErrorr(r1Z initializedrcrr_optionsrinitialize_subproblems)r%r*r&r')r(r r!r$s    zADMMAlgorithm.__init__cCs,|drtd|jjdttjdddS)Nr8z!Initializing subproblems for admmr}T)invocation_type oneway_call)r[r]rinvoke_functionthisfiler PerScenario)r%r r r!rs z$ADMMAlgorithm.initialize_subproblemscCs,|drtd|jjdttjdddS)Nr8z Cleaning up subproblems for admmrkT)rr)r[r]rrrrr)r%r r r!rs z!ADMMAlgorithm.cleanup_subproblemscCs|jjdttjdddS)NrT)rr)rrrrr)r%r r r!activate_lagrangian_terms z&ADMMAlgorithm.activate_lagrangian_termcCs|jjdttjdddS)NrT)rr)rrrrr)r%r r r!deactivate_lagrangian_terms z(ADMMAlgorithm.deactivate_lagrangian_termcCs|jjdttjdddS)NrT)rr)rrrrr)r%r r r!activate_penalty_terms z#ADMMAlgorithm.activate_penalty_termcCs|jjdttjdddS)NrT)rr)rrrrr)r%r r r!deactivate_penalty_terms z%ADMMAlgorithm.deactivate_penalty_termcCs|jjdttjdddS)NrT)rr)rrrrr)r%r r r!'activate_probability_weighted_cost_terms z5ADMMAlgorithm.activate_probability_weighted_cost_termcCs|jjdttjdddS)NrT)rr)rrrrr)r%r r r!)deactivate_probability_weighted_cost_terms z7ADMMAlgorithm.deactivate_probability_weighted_cost_termc Cs|jjdttj|fddx8|jjjD]*}|jjdtt|j||jfddq&W|jjdttj|fdd|j d}x|jjjD]r}||j}|j }||j 7}xR|j ddD]@} | rt|| j} || j} x| jD]} | | | | <qWqWqW|S)NrT)rZ function_argsrrrgr@)rrrrrrArTZ OnScenariorDrZsolve_scenarios_xZ _objectiverurvrgrE) r%r+r,r-r.rQ objective x_scenarioZx_scenario_solutionrGx_nodeZx_node_solutionrJr r r! run_x_updates@          zADMMAlgorithm.run_x_updatecCshd}d}d}d}x0|jjjddD]} x| jD]} || j} || j} x| jD]} | | }| | }d}x"| jD]}||j| j| 7}qxW|tt | j}| | s| | rt t |}|t | j|d||d7}x:| jD]0}|j| j| }||d7}|||d7}qW|| | <||d7}qXWq6Wq&Wt|t|t|t|fS)Ngr@rL)rrArBrCrDrErTrfloatrSis_variable_binaryis_variable_integerintroundrMrN)r%r+r,r-r.primal_residual dual_residualx_scalez_scalerFrGrIrWrJZ z_var_prevZrho_varZz_varrQZx_varr r r! run_z_update*s>         zADMMAlgorithm.run_z_updatec Csd}x|jjjD]}||j}|j}x|jddD]|} | rDt|| j} || j} || j} || j} xB| jD]8}| || || || |7<|| |d7}qtWq4WqWt |S)Ngr@rL) rrArTrDrrurvrgrErMrN)r%r+r,r-r.y_scalerQ y_scenariorrGrWrIy_noderrJr r r! run_y_updateKs       $zADMMAlgorithm.run_y_update?c Cs~i}xH|jjjddD]2}x,|jD]"}tdd|jD}||j<q$WqWi}i} x|jjjD]} i} || j<i} | | j<x| jddD]}| rt i} | |j<i}| |j<||j}xZ|jD]P}d| |<t |tkr|| j|j|||<n|||<||||7<qWqWqbWx|jjjddD]t}xl|jD]b}||j}xP|jD]F}t ||dkrb|j |\}}td|t|t||fqbWqLWq@Wi}i}x|jjjddD]}x|jD]}i}||j<i}||j<xd|jD]Z}t |tkr.||j|||<n|||<t |tkrX||j|||<n|||<qWqWqW||| |fS)Nr@css|]}|dfVqdS)gNr ).0rJr r r! esz:ADMMAlgorithm.initialize_algorithm_data..gư>z\Initial lagrange multipler estimates for non-anticipative variable %s do not sum to zero: %s)rrArBrCdictrErDrTrurvrgtyperUr\rcrrepr)r%rho_inity_initz_initZy_sumrFrGZ y_sum_noder+r,rQrrrrrJrDr_r.r-rIrWr r r!initialize_algorithm_data[sV &       .  z'ADMMAlgorithm.initialize_algorithm_datacCsdd}xT|jjjddD]>}x8|jD].}||j}x|jD]}|||d7}q:Wq$WqWt|S)Ngr@rL)rrArBrCrDrErMrN)r%vZvnormrFrGZv_noderJr r r!compute_nodevector_norms   z%ADMMAlgorithm.compute_nodevector_norm)N)rrr)r1r2r3r4r"rrrr$rrrrrrrrrrrrrr5r r )r(r!rs*    %! 9rcsTeZdZedddZfddZddZedd Zed d Z dddZ Z S) ADMMSolverNc Cs|dkrt}t|dtdt|dtdtddddtdt|dtdtd dddtdt|xtj D]}||qnWd tj kst t|d td t d t ttj dddtd|S) Nmax_iterations)r"primal_residual_relative_toleranceg-C6?z4Relative primal-residual tolerance. Default is 1e-4.r)domain descriptiondoc visibility dual_residual_relative_tolerancez2Relative dual-residual tolerance. Default is 1e-4.rd rho_strategyztt|||t|_t|_t|_d|_ dS)N) r#rr$rset_options_to_defaultrobjective_historyprimal_residual_historydual_residual_history iterations)r%)r(r r!r$s zADMMSolver.__init__cCs||_dS)N)rr)r%r r r!rsz!ADMMSolver.set_options_to_defaultcCs|jS)N)r)r%r r r!rszADMMSolver.optionscCsdS)Nadmmr )r%r r r!rDszADMMSolver.name?Fc) CsHt|jjdkrtdt}|j}t|j}t|j} d} d} d} d} d}x~|jjddD]j}xd|jD]Z}| d7} | t|j7} x>|jD]4}||r| d7} q| |r|d7}q| d7} qWqlWq`W| d}| d}| d}t |_ t |_ t |_d|_|rtd d d d d dd}t||j}|j|||d\}}}}t| d|j}||||||x@t|D] }|||||}|||||\}}}} |||||}!|jd7_td|| }"td|!}#|t||"}$|t| |#}%||j |<|$|j |<|%|j|<|rd|ddkrBt|td|||$|%t||f|$|kr|%|kr|rtd|jPn| |||||qW|rtd|WdQRX|r2td tdddtdd|td|$|td|%|||}&|$|%}'td|'|&t!}(||(_"||(_#|(S)NrLz7ADMM solver does not yet handle more than 2 time-stagesrr@rlrrrrz${0:^4} {1:>16} {2:>8} {3:>8} {4:>12}iterrZpr_resZdu_resz lg(||rho||))rrrrg? z%4d %16.7e %8.2e %8.2e %12.2ez Number of Iterations....: %sz) Maximum number of iterations reached: %sz' {0:^24} {1:^24}z(scaled)z (unscaled)z+Objective..........: {0:^24} {1:^24.16e}-z/Primal residual....: {0:^24.16e} {1:^24.16e}z/Dual residual......: {0:^24.16e} {1:^24.16e}z/Overall error......: {0:^24.16e} {1:^24.16e})$rSrArBrctimerTrCrErrr[rrrrrr]formatrrrrar/rrrrmaxrMrNlogrr0rrxhat))r%spr.rroutput_solver_log start_timerAZ num_scenariosZ num_stagesZ num_na_nodesZnum_na_variablesZnum_na_continuous_variablesZnum_na_binary_variablesZnum_na_integer_variablesrFrGrJZrel_tol_primalZ rel_tol_dualrZ label_colsrr+r,r-rirZunscaled_primal_residualZunscaled_dual_residualrrrZprimal_rel_scaleZdual_rel_scalerrZ unscaled_errerrresultsr r r! _solve_impls                          zADMMSolver._solve_impl)N)rrrF) r1r2r3r4r"r$rpropertyrrDrr5r r )r(r!rs 1   rc Cs|dkrt}t|t|t|dt|dt|dt|dt|dt|dt|dtdtd dd d d td |S)Nr8 disable_gcprofile tracebackroutput_scenario_tree_solution default_rhog?zIThe default rho value for all non-anticipative variables. Default is 1.0.r)rrrr)z-rz --default-rho)Zap_argsr) r rrrr r r rr)rr r r!runadmm_register_optionss8  rc Cst}t|}|tdtdt}|j|dd}|j|||j|jd}|j }|` tdt||j rtd|j |j |j WdQRXtdtdt|d S) zQ Construct a senario tree manager and solve it with the ADMM solver. rz7Running ADMM solver for stochastic programming problemsT)sparse)rr.rz&Final solution (scenario tree format):Nz!Total execution time=%.2f secondsr)rrr/r]rZextract_user_options_to_dictZsolverrrrrAZsnapshotSolutionFromScenariosZpprintSolutionZ pprintCosts)rrrrZ admm_optionsrrr r r!runadmms6    rc Cs\ddl}yt|tddd}Wn$tk r@}z|jSd}~XYnXtt|d|j|j|j dS)NrrzbOptimize a stochastic program using the Alternating Direction Method of Multipliers (ADMM) solver.)progrz runadmm: )Z error_labelrZ profile_countr) Z pyomo.environrr SystemExitcoderrrrr)r&ZpyomorZ_excr r r!mains rr__main__)N)N)GossysrrMZpyomo.common.collectionsrZ pyomo.corerrrrrZ!pyomo.pysp.util.configured_objectrZpyomo.pysp.util.configr r r r r rrrZpyomo.pysp.util.miscrrZpyomo.pysp.scenariotree.managerrrrZ&pyomo.pysp.scenariotree.manager_solverrZpyomo.pysp.phutilsrZpyomo.pysp.solvers.spsolverrrrZ six.movesrpathabspath__file__rreplacerZ_rho_group_labelrr6r7rarbrkr}rrrrrrrrrrrrrrZregister_solverr1exitr r r r! s\  (      j   A       f ") #