Aarondard5/constrained_multiOT

1 2 3#######################################4# PACKAGES5#######################################6 7import pickle, time, warnings8import numpy as np9import networkx as nx10import random11import copy12import math13import scipy as sp14import itertools15 16from scipy.sparse import csr_matrix, lil_matrix, issparse, csc_matrix17from scipy.sparse import diags18from scipy.sparse import identity19from scipy.sparse.linalg import spsolve20from scipy import sparse21 22import quadprog23import scipy.optimize24import scipy.stats25 26 27import matplotlib.pyplot as plt28from matplotlib.ticker import MaxNLocator29 30#######################################31INF = 1e2032 33warnings.filterwarnings("ignore", message="Matrix is exactly singular")34 35 36def initialize_mu(flux, pflux, length, ord_norm=2):37 38    flux_norm = np.linalg.norm(flux, axis=1, ord=ord_norm)39    num = flux_norm ** (2 / (3 - pflux))40    return num41 42 43def remaining_e_ids(tdens_init, removed_e_id):44    return np.setdiff1d(np.arange(len(tdens_init)), removed_e_id)45 46 47def tdensinit(nedge, seed=10, tdens_init=None):48    """initialization of the conductivities: mu_e ~ U(0,1)"""49    prng = np.random.RandomState(seed=seed)50 51    if tdens_init is None:52        print("-" * 15)53        print(f"Using U(0,1) as tdens")54        print("-" * 15)55        tdens_0 = 0.1 + prng.uniform(0, 1, size=nedge)56    else:57        print("=" * 15)58        print(f"Using existing tdens")59        print("-" * 15)60        tdens_0 = (61            0.01 * prng.uniform(0, 1, size=nedge) + tdens_init62        )  # [remaining_e_ids(tdens_init, removed_e_id)]63        # tdens_0 = 0.1 + tdens_init  # [remaining_e_ids(tdens_init, removed_e_id)]64 65    weight = np.ones(nedge) + 0.01 * prng.uniform(0, 1, size=nedge)66 67    return tdens_0, weight68 69 70def dyn(71    g,72    nodes,73    pflux,74    nedge,75    length,76    forcing0,77    tol_var_tdens,78    removed_e_id,79    seed=10,80    verbose=False,81    N_real=1,82    plot_cost=False,83    coupling="l2",84    tdens0=None,85    constraint=True,86    clipping=False,87    capacity=INF,88    time_step=1,89    alpha=1,90    tot_time=300,91):92    """dynamics method"""93 94    capacity = np.full(nedge, capacity)  # set capacity on edges95    print(96        f"capacity shape:{capacity.shape}; min capacity:{np.min(capacity)}; max capacity:{np.max(capacity)}"97    )98 99    print("\ndynamics...")100 101    relax_linsys = 1.0e-5  # relaxation for stiffness matrix102    tot_time = tot_time  # upper bound on number of time steps103    threshold_cost = 1.0e-6  # threshold for stopping criteria using cost104    prng = np.random.RandomState(105        seed=seed106    )  # only needed if spsolve has problems (inside update)107 108    nnode = len(nodes)109    ncomm = forcing0.shape[0]110    forcing = forcing0.transpose()111 112    minCost = 1e14113    minCost_list = []114 115    inc_mat = csr_matrix(nx.incidence_matrix(g, nodelist=nodes, oriented=True))  # B116    117    118    inc_transpose = csr_matrix(inc_mat.transpose())  # B^T119    inv_len_mat = diags(1 / length, 0)  # diag[1/l_e]120 121    for r in range(N_real):122 123        tdens_0, weight = tdensinit(124            nedge, seed=seed + r, tdens_init=tdens0125        )  # conductivities initialization126 127        # forcing = csc_matrix(forcing0.transpose())128        # --------------------------------------------------------------------------------129        tdens = tdens_0.copy()130        td_mat = diags(tdens.astype(float), 0)  # matrix M131        # print(f"tdens:{tdens}")132        stiff = (133            inc_mat * td_mat * inv_len_mat * inc_transpose134        )  # B diag[mu] diag[1/l_e] B^T135        stiff_relax = stiff + relax_linsys * identity(nnode)  # avoid zero kernel136        pot = spsolve(stiff_relax, forcing).reshape((nnode, ncomm))  # pressure137        print(f"stiff_relax:{stiff_relax.shape} --  forcing:{forcing.shape} -- pot:{pot.shape}")138        # --------------------------------------------------------------------------------139        # Run dynamics140        convergence_achieved = False141        cost_update = 0142        cost_update_inter = 0143        cost_list = []144        g_list = []145        time_iteration = 0146 147        fmax = forcing0.max()148 149        not_in_C_id = np.where(compute_capacity_contraint(tdens, capacity) < 0)[0]150        print(f"initial # of active constraints: {len(not_in_C_id)}")151        while not convergence_achieved and time_iteration < tot_time:152 153            time_iteration += 1154 155            # update tdens-pot system156            tdens_old = tdens.copy()157            pot_old = pot.copy()158 159            # print(f'mu:{tdens}')160            # equations update161            tdens, pot, grad, info = update(162                g,163                tdens,164                pot,165                weight,166                inc_mat,167                inc_transpose,168                inv_len_mat,169                forcing,170                time_step,171                pflux,172                relax_linsys,173                nnode,174                coupling=coupling,175                constraint=constraint,176                clipping=clipping,177                capacity=capacity,178                alpha=alpha,179            )180 181            g_ = compute_capacity_contraint(tdens, capacity)182            g_list.append(g_)183 184            # print(f'g_list:{g_list}')185            # singular stiffness matrix186            if info != 0:187                print(f"info = {info}")188                tdens = (189                    tdens_old + prng.rand(*tdens.shape) * np.mean(tdens_old) / 1000.0190                )191                pot = pot_old + prng.rand(*pot.shape) * np.mean(pot_old) / 1000.0192 193            # 1) convergence with conductivities194            # var_tdens = max(np.abs(tdens - tdens_old))/time_step195            # print(time_iteration, var_tdens)196            if verbose > 1:197                print("==========")198 199            # 2) an alternative convergence criteria: using total cost and maximum variation of conductivities200            var_tdens = max(np.abs(tdens - tdens_old)) / time_step201 202            # var_tdens_inter = ([ max(np.abs(tdens_inter[i] - tdens_old_inter[i]))/time_step for i in range(nnode) ] )203            (204                convergence_achieved,205                cost_update,206                abs_diff_cost,207                flux_norm,208                flux_mat,209            ) = cost_convergence(210                threshold_cost,211                cost_update,212                tdens,213                pot,214                inc_mat,215                inv_len_mat,216                length,217                weight,218                pflux,219                convergence_achieved,220                var_tdens,221                coupling=coupling,222            )  # , var_tdens_inter)223 224            if verbose > 1:225                # print(time_iteration, var_tdens/forcing.max(), abs_diff_cost)226 227                # print('\r','It=',it,'err=', abs_diff,'J-J_old=',abs_diff_cost,sep=' ', end='', flush=True)228                print(229                    "\r",230                    "it=%3d, err/max_f=%5.2f, J_diff=%8.2e"231                    % (time_iteration, var_tdens / fmax, abs_diff_cost),232                    sep=" ",233                    end=" ",234                    flush=True,235                )236                time.sleep(0.05)237 238            cost_list.append(cost_update)239 240            if var_tdens < tol_var_tdens:  # or (var_tdens_inter < tol_var_tdens):241                convergence_achieved = True242 243            elif time_iteration >= tot_time:244                convergence_achieved = True245                tdens = tdens_old.copy()246                if verbose > 0:247                    print(248                        "ERROR: convergence dynamics not achieved, iteration time > maxit"249                    )250 251            if convergence_achieved == True and time_iteration < tot_time:252                if constraint == True:  # check that all constraints are satisfied253                    not_in_C_id = np.where(254                        compute_capacity_contraint(tdens, capacity) <= 0255                    )[0]256                    if len(not_in_C_id) > 0:257                        convergence_achieved = False258 259 260        if convergence_achieved:261            if verbose > 0:262                print("cost:", cost_update, " - N_real:", r, "- Best cost", minCost)263            if cost_update < minCost:264                minCost = cost_update265                minCost_list = cost_list266                tdens_best = tdens.copy()267                pot_best = pot.copy()268                flux_best_norm = flux_norm.copy()269                flux_best = flux_mat.copy()270                min_g_list = g_list.copy()271 272        else:273            print("ERROR: convergence dynamics not achieved")274 275    if plot_cost:276        plot_J(minCost_list, int_ticks=True)277        plot_gmu(min_g_list, not_in_C_id, tdens_best, capacity, int_ticks=True)278 279        # plt.plot(f_dot, rhs_ode)280    return (281        tdens_best,282        pot_best,283        flux_best_norm,284        minCost,285        minCost_list,286        flux_best,287        min_g_list,288    )289 290 291def update(292    g,293    tdens,294    pot,295    weight,296    inc_mat,297    inc_transpose,298    inv_len_mat,299    forcing,300    time_step,301    pflux,302    relax_linsys,303    nnode,304    coupling="l1",305    constraint=True,306    clipping=False,307    capacity=None,308    alpha=1,309):310 311    # updating dynamic...312 313    nedge = tdens.shape[0]314    nnode, ncomm = pot.shape315 316    weight = 1.0317    # weight = length318    grad = inv_len_mat * inc_transpose * pot  # discrete gradient319 320    if coupling == "l2":321        rhs_ode = ((tdens**pflux) * ((grad**2).sum(axis=1)) / (weight**2)) - tdens322    if coupling == "l1":323        rhs_ode = ((tdens**pflux) * ((np.abs(grad).sum(axis=1)) ** 2) / (weight**2)) - tdens324    # print(f"rhs_ode:{rhs_ode}")325    if clipping == True:  # use trivial method of clipping326        constraint = False327 328    # --------------------------329    # adding force to enforce constraints330    # --------------------------331 332    if constraint == True:333        """334        this section runs analytic result335        """336        small_err = 1e-05337        g_capacity = compute_capacity_contraint(tdens + small_err, capacity)338        not_in_C_id = np.where(g_capacity <= 0)[0]339        alpha_pos = 1 / time_step340 341        if len(not_in_C_id) > 0:342            for i in np.arange(nedge):343                if rhs_ode[i] >= alpha * g_capacity[i]:344                    rhs_ode[i] = alpha * g_capacity[i]345                else:346                    rhs_ode[i] = rhs_ode[i]  # -alpha_pos * tdens[i]  #347 348    # update conductivity349    if rhs_ode.ndim > 1:350        if alpha == 0:351            tdens = tdens + time_step * np.ravel(rhs_ode[:, 0])352        else:353            tdens = tdens + (alpha * time_step * np.ravel(rhs_ode[:, 0]))354    else:355        if alpha == 0:356            tdens = tdens + time_step * rhs_ode357        else:358            tdens = tdens + (alpha * time_step * rhs_ode)359 360    not_in_C_id = np.where(compute_capacity_contraint(tdens, capacity) < 0)[0]361    td_mat = diags(tdens.astype(float), 0)362 363    # update stiffness matrix364    stiff = inc_mat * td_mat * inv_len_mat * inc_transpose365 366    # spsolve367    stiff_relax = stiff + relax_linsys * identity(nnode)  # avoid zero kernel368    # update potential369    pot = spsolve(stiff_relax, forcing).reshape((nnode, ncomm))  # pressure370    if np.any(np.isnan(pot)):  # or np.any(np.isnan(pot_ctr)): # or np.any(pot != pot)371        info = -1372        pass373    else:374        info = 0375 376    return tdens, pot, grad, info377 378 379def calculate_cost(flux_mat, length, pflux=1.9, coupling="l2", tdens=None):380    if coupling == "l2":381        flux_norm = np.linalg.norm(flux_mat, axis=1)382    if coupling == "l1":383        flux_norm = np.linalg.norm(flux_mat, axis=1, ord=1)384 385    if tdens is None:386        return np.sum(length * (flux_norm ** (2 * (2 - pflux) / (3 - pflux))))387    else:388        normalization = 1.0 / np.sum(tdens)389        return normalization * np.sum(390            length * (flux_norm ** (2 * (2 - pflux) / (3 - pflux)))391        )392 393 394def cost_convergence(395    threshold_cost,396    cost,397    tdens,398    pot,399    inc_mat,400    inv_len_mat,401    length,402    weight,403    pflux,404    convergence_achieved,405    var_tdens,406    coupling="l1",407):  # , var_tdens_inter):408    """computing convergence using total cost: setting a high value for maximum conductivity variability"""409 410    L = len(pot)411    nnode = pot[0].shape[0]412    td_mat = np.diag(tdens.astype(float))413    flux_mat = np.matmul(td_mat * inv_len_mat * np.transpose(inc_mat), pot)414 415    if coupling == "l2":416        flux_norm = np.linalg.norm(flux_mat, axis=1)417    if coupling == "l1":418        flux_norm = np.linalg.norm(flux_mat, axis=1, ord=1)419 420    # cost_update = calculate_cost( flux_mat, length, pflux=pflux, coupling=coupling, tdens=tdens )421    normalization = 1.0 / np.sum(tdens)422    cost_update = normalization * np.sum(423        length * (flux_norm ** (2 * (2 - pflux) / (3 - pflux)))424    )425 426 427    abs_diff_cost = abs(cost_update - cost)428 429    convergence_achieved = bool(convergence_achieved)430    if min(pflux) > 0.0:431        if abs_diff_cost < threshold_cost:432            convergence_achieved = True433    else:434        if abs_diff_cost < threshold_cost and var_tdens < 1:435            convergence_achieved = True436 437    return convergence_achieved, cost_update, abs_diff_cost, flux_norm, flux_mat438 439 440 441 442def plot_J(443    values, indices=None, k_i=5, figsize=(7, 3), int_ticks=False, xlab="Iterations"444):445 446    fig, ax = plt.subplots(1, 1, figsize=figsize)447    # print('\n\nL: \n\n',values[k_i:])448 449    if indices is None:450        ax.plot(values[k_i:])451    else:452        ax.plot(indices[k_i:], values[k_i:])453    ax.set_xlabel(xlab)454    ax.set_ylabel("Best cost")455    if int_ticks:456        ax.xaxis.set_major_locator(MaxNLocator(integer=True))457    ax.grid()458 459    plt.tight_layout()460    plt.show()461 462 463def plot_gmu(464    g_func,465    not_in_C_id,466    tdens,467    capacity,468    indices=None,469    k_i=5,470    figsize=(14, 5),471    int_ticks=False,472    xlab="Edges",473    outfig=None,474):475    """476    plot function g(mu) = c_e - mu_e477    """478    g = compute_capacity_contraint(tdens, capacity)479 480    fig, (ax1, ax2) = plt.subplots(1, 2, figsize=figsize)481    fs = 23482    indices = None483    # for item in g_func:484    not_in_C_id = np.where(g < 0)[0]485 486    if indices is None:487        ax1.plot(g[k_i:])488    else:489        ax1.plot(indices[k_i:], g[k_i:])490    ax1.set_xlabel(xlab)491    ax1.set_ylabel(r"$g = c_e - \mu_e$", fontsize=fs)492    if int_ticks:493        ax1.xaxis.set_major_locator(MaxNLocator(integer=True))494    ax1.grid()495 496    # axis 2497    if indices is None:498        ax2.plot(g[not_in_C_id][k_i:])499    else:500        ax2.plot(indices[k_i:], g[not_in_C_id][k_i:])501    ax2.set_xlabel(xlab)502    ax2.set_ylabel(r"$g(\mu_e) = c_e - \mu_e < 0$", fontsize=fs)503    if int_ticks:504        ax2.xaxis.set_major_locator(MaxNLocator(integer=True))505    ax2.grid()506    507    plt.tight_layout()508    if outfig is not None:509        plt.savefig(outfig, dpi=300)510    return plt.show()511 512 513def capacity_constraint_QP(514    MultiOT_rhs_ode, infeasible_points_ids, capacity, g=0, alpha=1515):516 517    L = infeasible_points_ids.shape[0]518    v = np.zeros(L)519 520    alpha_g = alpha * g[infeasible_points_ids]521 522    mask = MultiOT_rhs_ode[infeasible_points_ids] >= alpha_g523    v[mask] = alpha_g[mask]524    v[np.logical_not(mask)] = -np.abs(525        MultiOT_rhs_ode[infeasible_points_ids][np.logical_not(mask)]526    )527    return v528 529 530def compute_capacity_contraint(tdens, capacities):531    return capacities - tdens532 533 534 535 536