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