File size: 6,056 Bytes
9860743 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 |
import numpy as np
import diffcp
# import diffcp_lpgd
import time
from dataclasses import dataclass
from typing import Any
import importlib.util
@dataclass
class ForwardContext:
gp: bool
solve_and_derivative: bool
batch: bool
batch_size: int
batch_sizes: list
compiler: callable
param_ids: list
param_order: list
old_params_to_new_params: dict
cone_dims: list
solver_args: dict
variables: list
var_dict: dict
lpgd: bool
@dataclass
class BackwardContext:
info: dict
gp: bool
batch: bool
batch_size: int
batch_sizes: list
variables: list
compiler: Any
param_ids: list
param_order: list
params: list
old_params_to_new_params: dict
sol: list
lpgd: bool
def forward_numpy(params_numpy, context):
"""Forward pass in numpy."""
info = {}
if context.gp:
param_map = {}
# construct a list of params for the DCP problem
for param, value in zip(context.param_order, params_numpy):
if param in context.old_params_to_new_params:
new_id = context.old_params_to_new_params[param].id
param_map[new_id] = np.log(value)
else:
new_id = param.id
param_map[new_id] = value
params_numpy = [param_map[pid] for pid in context.param_ids]
# canonicalize problem
start = time.time()
As, bs, cs, cone_dicts, shapes = [], [], [], [], []
for i in range(context.batch_size):
params_numpy_i = [
p if sz == 0 else p[i]
for p, sz in zip(params_numpy, context.batch_sizes)]
c, _, neg_A, b = context.compiler.apply_parameters(
dict(zip(context.param_ids, params_numpy_i)),
keep_zeros=True)
A = -neg_A # cvxpy canonicalizes -A
As.append(A)
bs.append(b)
cs.append(c)
cone_dicts.append(context.cone_dims)
shapes.append(A.shape)
info['canon_time'] = time.time() - start
info['shapes'] = shapes
# compute solution and derivative function
start = time.time()
try:
if context.solve_and_derivative:
if context.lpgd:
xs, _, _, _, DT_batch = diffcp.solve_and_derivative_batch(
As, bs, cs, cone_dicts, mode='lpgd', derivative_kwargs=dict(tau=1e-4, rho=0.1), **context.solver_args)
else:
xs, _, _, _, DT_batch = diffcp.solve_and_derivative_batch(
As, bs, cs, cone_dicts, **context.solver_args)
info['DT_batch'] = DT_batch
else:
# xs, _, _ = diffcp.solve_only_batch(
# As, bs, cs, cone_dicts, **context.solver_args)
xs, _, _ = diffcp.solve_only_batch(
As, bs, cs, cone_dicts, **context.solver_args)
except diffcp.SolverError as e:
print(
"Please consider re-formulating your problem so that "
"it is always solvable or increasing the number of "
"solver iterations.")
raise e
info['solve_time'] = time.time() - start
# extract solutions and append along batch dimension
start = time.time()
sol = [[] for i in range(len(context.variables))]
for i in range(context.batch_size):
sltn_dict = context.compiler.split_solution(
xs[i], active_vars=context.var_dict)
for j, v in enumerate(context.variables):
sol[j].append(np.expand_dims(sltn_dict[v.id], axis=0))
sol = [np.concatenate(s, axis=0) for s in sol]
if not context.batch:
sol = [np.squeeze(s, axis=0) for s in sol]
if context.gp:
sol = [np.exp(s) for s in sol]
info['sol'] = sol
return sol, info
def backward_numpy(dvars_numpy, context):
"""Backward pass in numpy."""
info = {}
if context.gp:
# derivative of exponential recovery transformation
dvars_numpy = [dvar*s for dvar, s in zip(dvars_numpy, context.sol)]
if not context.batch:
dvars_numpy = [np.expand_dims(dvar, 0) for dvar in dvars_numpy]
# differentiate from cvxpy variables to cone problem data
dxs, dys, dss = [], [], []
for i in range(context.batch_size):
del_vars = {}
for v, dv in zip(context.variables, [dv[i] for dv in dvars_numpy]):
del_vars[v.id] = dv
dxs.append(context.compiler.split_adjoint(del_vars))
dys.append(np.zeros(context.info['shapes'][i][0]))
dss.append(np.zeros(context.info['shapes'][i][0]))
start = time.time()
dAs, dbs, dcs = context.info['DT_batch'](dxs, dys, dss)
info['dDT_time'] = time.time() - start
# differentiate from cone problem data to cvxpy parameters
start = time.time()
grad = [[] for _ in range(len(context.param_ids))]
for i in range(context.batch_size):
del_param_dict = context.compiler.apply_param_jac(
dcs[i], -dAs[i], dbs[i])
for j, pid in enumerate(context.param_ids):
grad[j].append(np.expand_dims(del_param_dict[pid], 0))
grad = [np.concatenate(g, axis=0) for g in grad]
if context.gp:
# differentiate through the log transformation of params
dcp_grad = grad
grad = []
dparams = {pid: g for pid, g in zip(context.param_ids, dcp_grad)}
for param, value in zip(context.param_order, context.params):
g = 0.0 if param.id not in dparams else dparams[param.id]
if param in context.old_params_to_new_params:
dcp_param_id = context.old_params_to_new_params[param].id
# new_param.value == log(param), apply chain rule
g += (1.0 / value) * dparams[dcp_param_id]
grad.append(g)
info['dcanon_time'] = time.time() - start
if not context.batch:
grad = [g.squeeze(0) for g in grad]
else:
for i, sz in enumerate(context.batch_sizes):
if sz == 0:
grad[i] = grad[i].sum(axis=0)
return grad, info |