1from pyomo.contrib.pynumero.interfaces.nlp import NLP 2import numpy as np 3import scipy.sparse as sp 4 5class _BaseNLPDelegator(NLP): 6 def __init__(self, original_nlp): 7 """ 8 This is a base class to make it easier to implement NLP 9 classes that wrap other NLP instances. This class 10 simply reproduces the NLP interface by passing the call 11 onto the original_nlp passed in the constructor. This allows 12 new wrapper classes to only implement the methods that change, 13 allowing the others to pass through. 14 15 Parameters 16 ---------- 17 original_nlp : NLP-like 18 The original NLP object that we want to wrap 19 """ 20 super(NLP, self).__init__() 21 self._original_nlp = original_nlp 22 23 def n_primals(self): 24 return self._original_nlp.n_primals() 25 26 def primals_names(self): 27 return self._original_nlp.primals_names() 28 29 def n_constraints(self): 30 return self._original_nlp.n_constraints() 31 32 def constraint_names(self): 33 return self._original_nlp.constraint_names() 34 35 def nnz_jacobian(self): 36 return self._original_nlp.nnz_jacobian() 37 38 def nnz_hessian_lag(self): 39 return self._original_nlp.nnz_hessian_lag() 40 41 def primals_lb(self): 42 return self._original_nlp.primals_lb() 43 44 def primals_ub(self): 45 return self._original_nlp.primals_ub() 46 47 def constraints_lb(self): 48 return self._original_nlp.constraints_lb() 49 50 def constraints_ub(self): 51 return self._original_nlp.constraints_ub() 52 53 def init_primals(self): 54 return self._original_nlp.init_primals() 55 56 def init_duals(self): 57 return self._original_nlp.init_duals() 58 59 def create_new_vector(self, vector_type): 60 return self._original_nlp.create_new_vector(vector_type) 61 62 def set_primals(self, primals): 63 self._original_nlp.set_primals(primals) 64 65 def get_primals(self): 66 return self._original_nlp.get_primals() 67 68 def set_duals(self, duals): 69 self._original_nlp.set_duals(duals) 70 71 def get_duals(self): 72 return self._original_nlp.get_duals() 73 74 def set_obj_factor(self, obj_factor): 75 self._original_nlp.set_obj_factor(obj_factor) 76 77 def get_obj_factor(self): 78 return self._original_nlp.get_obj_factor() 79 80 def get_obj_scaling(self): 81 return self._original_nlp.get_obj_scaling() 82 83 def get_primals_scaling(self): 84 return self._original_nlp.get_primals_scaling() 85 86 def get_constraints_scaling(self): 87 return self._original_nlp.get_constraints_scaling() 88 89 def evaluate_objective(self): 90 return self._original_nlp.evaluate_objective() 91 92 def evaluate_grad_objective(self, out=None): 93 return self._original_nlp.evaluate_grad_objective(out) 94 95 def evaluate_constraints(self, out=None): 96 return self._original_nlp.evaluate_constraints(out) 97 98 def evaluate_jacobian(self, out=None): 99 return self._original_nlp.evaluate_jacobian(out) 100 101 def evaluate_hessian_lag(self, out=None): 102 return self._original_nlp.evaluate_hessian_lag(out) 103 104 def report_solver_status(self, status_code, status_message): 105 self._original_nlp.report_solver_status(status_code, status_message) 106 107 108class RenamedNLP(_BaseNLPDelegator): 109 def __init__(self, original_nlp, primals_name_map): 110 """ 111 This class takes an NLP that and allows one to rename the primals. 112 It is a thin wrapper around the original NLP. 113 114 Parameters 115 ---------- 116 original_nlp : NLP-like 117 The original NLP object that implements the NLP interface 118 119 primals_name_map : dict of str --> str 120 This is a dictionary that maps from the names 121 in the original NLP class to the desired names 122 for this instance. 123 """ 124 super(RenamedNLP, self).__init__(original_nlp) 125 self._primals_name_map = primals_name_map 126 self._new_primals_names = None 127 # Todo: maybe do this on first call instead of __init__? 128 self._generate_new_names() 129 130 def _generate_new_names(self): 131 if self._new_primals_names is None: 132 assert self._original_nlp.n_primals() == len(self._primals_name_map) 133 self._new_primals_names = \ 134 [self._primals_name_map[nm] for nm in self._original_nlp.primals_names()] 135 136 def primals_names(self): 137 return self._new_primals_names 138 139 140class ProjectedNLP(_BaseNLPDelegator): 141 def __init__(self, original_nlp, primals_ordering): 142 """ 143 This class takes an NLP that depends on a set of primals (original 144 space) and converts it to an NLP that depends on a reordered set of 145 primals (projected space). 146 147 This will impact all the returned items associated with primal 148 variables. E.g., the gradient will be in the new primals ordering 149 instead of the original primals ordering. 150 151 Note also that this can include additional primal variables not 152 in the original NLP, or can exclude primal variables that were 153 in the original NLP. 154 155 Parameters 156 ---------- 157 original_nlp : NLP-like 158 The original NLP object that implements the NLP interface 159 160 primals_ordering: list 161 List of strings indicating the desired primal variable 162 ordering for this NLP. The list can contain new variables 163 that are not in the original NLP, thereby expanding the 164 space of the primal variables. 165 """ 166 super(ProjectedNLP, self).__init__(original_nlp) 167 self._primals_ordering = list(primals_ordering) 168 self._original_idxs = None 169 self._projected_idxs = None 170 self._generate_maps() 171 self._projected_primals = self.init_primals() 172 self._jacobian_nz_mask = None 173 self._hessian_nz_mask = None 174 self._nnz_jacobian = None 175 self._nnz_hessian_lag = None 176 177 def _generate_maps(self): 178 if self._original_idxs is None or self._projected_idxs is None: 179 primals_ordering_dict = {k:i for i,k in enumerate(self._primals_ordering)} 180 original_names = self._original_nlp.primals_names() 181 original_idxs = list() 182 projected_idxs = list() 183 for i,nm in enumerate(original_names): 184 if nm in primals_ordering_dict: 185 # we need the reordering for this element 186 original_idxs.append(i) 187 projected_idxs.append(primals_ordering_dict[nm]) 188 self._original_idxs = np.asarray(original_idxs) 189 self._projected_idxs = np.asarray(projected_idxs) 190 self._original_to_projected = np.nan*np.zeros(self._original_nlp.n_primals()) 191 self._original_to_projected[self._original_idxs] = self._projected_idxs 192 193 def n_primals(self): 194 return len(self._primals_ordering) 195 196 def primals_names(self): 197 return list(self._primals_ordering) 198 199 def nnz_jacobian(self): 200 if self._nnz_jacobian is None: 201 J = self.evaluate_jacobian() 202 self._nnz_jacobian = len(J.data) 203 return self._nnz_jacobian 204 205 def nnz_hessian_lag(self): 206 if self._nnz_hessian_lag is None: 207 H = self.evaluate_hessian_lag() 208 self._nnz_hessian_lag = len(H.data) 209 return self._nnz_hessian_lag 210 211 def _project_primals(self, default, original_primals): 212 projected_x = default*np.ones(self.n_primals(), dtype=np.float64) 213 projected_x[self._projected_idxs] = original_primals[self._original_idxs] 214 return projected_x 215 216 def primals_lb(self): 217 return self._project_primals(-np.inf, self._original_nlp.primals_lb()) 218 219 def primals_ub(self): 220 return self._project_primals(np.inf, self._original_nlp.primals_ub()) 221 222 def init_primals(self): 223 # Todo: think about what to do here if an entry is not defined 224 # for now, we default to NaN, but there may be a better way 225 # (e.g., taking a new initial value in the constructor?) 226 return self._project_primals(np.nan, self._original_nlp.init_primals()) 227 228 def create_new_vector(self, vector_type): 229 if vector_type == 'primals': 230 return np.zeros(self.n_primals(), dtype=np.float64) 231 return self._original_nlp.create_new_vector(vector_type) 232 233 def set_primals(self, primals): 234 # here, we keep a local copy of the projected primals so 235 # we can give back these values in get_primals. This might 236 # not be the best idea since we can't really support this 237 # same strategy for other methods (e.g., init_primals) where 238 # we now use NaNs to fill in any "missing" entries. 239 np.copyto(self._projected_primals, primals) 240 original_primals = self._original_nlp.get_primals() 241 original_primals[self._original_idxs] = primals[self._projected_idxs] 242 self._original_nlp.set_primals(original_primals) 243 244 def get_primals(self): 245 original_primals = self._original_nlp.get_primals() 246 self._projected_primals[self._projected_idxs] = original_primals[self._original_idxs] 247 return self._projected_primals 248 249 def get_primals_scaling(self): 250 return self._project_primals(np.nan, self._original_nlp.get_primals_scaling()) 251 252 def evaluate_grad_objective(self, out=None): 253 original_grad_objective = self._original_nlp.evaluate_grad_objective() 254 projected_objective = self._project_primals(0.0, original_grad_objective) 255 if out is None: 256 return projected_objective 257 np.copyto(out, projected_objective) 258 return out 259 260 def evaluate_jacobian(self, out=None): 261 original_jacobian = self._original_nlp.evaluate_jacobian() 262 if out is not None: 263 np.copyto(out.data, original_jacobian.data[self._jacobian_nz_mask]) 264 return out 265 266 row = original_jacobian.row 267 col = original_jacobian.col 268 data = original_jacobian.data 269 270 if self._jacobian_nz_mask is None: 271 # need to remap the irow, jcol to the new space and change the size 272 self._jacobian_nz_mask = np.isin(col, self._original_idxs) 273 274 new_col = col[self._jacobian_nz_mask] 275 new_col = self._original_to_projected[new_col] 276 new_row = row[self._jacobian_nz_mask] 277 new_data = data[self._jacobian_nz_mask] 278 279 return sp.coo_matrix((new_data, (new_row,new_col)), shape=(self.n_constraints(), self.n_primals())) 280 281 def evaluate_hessian_lag(self, out=None): 282 original_hessian = self._original_nlp.evaluate_hessian_lag() 283 if out is not None: 284 np.copyto(out.data, original_hessian.data[self._hessian_nz_mask]) 285 return out 286 287 row = original_hessian.row 288 col = original_hessian.col 289 data = original_hessian.data 290 291 if self._hessian_nz_mask is None: 292 # need to remap the irow, jcol to the new space and change the size 293 self._hessian_nz_mask = np.isin(col, self._original_idxs) & np.isin(row, self._original_idxs) 294 295 new_col = col[self._hessian_nz_mask] 296 new_col = self._original_to_projected[new_col] 297 new_row = row[self._hessian_nz_mask] 298 new_row = self._original_to_projected[new_row] 299 new_data = data[self._hessian_nz_mask] 300 301 return sp.coo_matrix((new_data, (new_row,new_col)), shape=(self.n_primals(), self.n_primals())) 302 303 def report_solver_status(self, status_code, status_message): 304 raise NotImplementedError('Need to think about this...') 305