Source code for pyomo.contrib.pynumero.interfaces.nlp_projections

from pyomo.contrib.pynumero.interfaces.nlp import NLP
import numpy as np
import scipy.sparse as sp

class _BaseNLPDelegator(NLP):
    def __init__(self, original_nlp):
        """
        This is a base class to make it easier to implement NLP
        classes that wrap other NLP instances. This class
        simply reproduces the NLP interface by passing the call
        onto the original_nlp passed in the constructor. This allows
        new wrapper classes to only implement the methods that change,
        allowing the others to pass through.

        Parameters
        ----------
        original_nlp : NLP-like
            The original NLP object that we want to wrap
        """
        super(NLP, self).__init__()
        self._original_nlp = original_nlp

    def n_primals(self):
        return self._original_nlp.n_primals()

    def primals_names(self):
        return self._original_nlp.primals_names()

    def n_constraints(self):
        return self._original_nlp.n_constraints()

    def constraint_names(self):
        return self._original_nlp.constraint_names()

    def nnz_jacobian(self):
        return self._original_nlp.nnz_jacobian()

    def nnz_hessian_lag(self):
        return self._original_nlp.nnz_hessian_lag()

    def primals_lb(self):
        return self._original_nlp.primals_lb()

    def primals_ub(self):
        return self._original_nlp.primals_ub()

    def constraints_lb(self):
        return self._original_nlp.constraints_lb()

    def constraints_ub(self):
        return self._original_nlp.constraints_ub()

    def init_primals(self):
        return self._original_nlp.init_primals()

    def init_duals(self):
        return self._original_nlp.init_duals()

    def create_new_vector(self, vector_type):
        return self._original_nlp.create_new_vector(vector_type)

    def set_primals(self, primals):
        self._original_nlp.set_primals(primals)

    def get_primals(self):
        return self._original_nlp.get_primals()

    def set_duals(self, duals):
        self._original_nlp.set_duals(duals)

    def get_duals(self):
        return self._original_nlp.get_duals()

    def set_obj_factor(self, obj_factor):
        self._original_nlp.set_obj_factor(obj_factor)

    def get_obj_factor(self):
        return self._original_nlp.get_obj_factor()

    def get_obj_scaling(self):
        return self._original_nlp.get_obj_scaling()

    def get_primals_scaling(self):
        return self._original_nlp.get_primals_scaling()

    def get_constraints_scaling(self):
        return self._original_nlp.get_constraints_scaling()

    def evaluate_objective(self):
        return self._original_nlp.evaluate_objective()

    def evaluate_grad_objective(self, out=None):
        return self._original_nlp.evaluate_grad_objective(out)

    def evaluate_constraints(self, out=None):
        return self._original_nlp.evaluate_constraints(out)

    def evaluate_jacobian(self, out=None):
        return self._original_nlp.evaluate_jacobian(out)

    def evaluate_hessian_lag(self, out=None):
        return self._original_nlp.evaluate_hessian_lag(out)

    def report_solver_status(self, status_code, status_message):
        self._original_nlp.report_solver_status(status_code, status_message)


class RenamedNLP(_BaseNLPDelegator):
    def __init__(self, original_nlp, primals_name_map):
        """
        This class takes an NLP that and allows one to rename the primals.
        It is a thin wrapper around the original NLP.

        Parameters
        ----------
        original_nlp : NLP-like
            The original NLP object that implements the NLP interface

        primals_name_map : dict of str --> str
            This is a dictionary that maps from the names
            in the original NLP class to the desired names
            for this instance.
        """
        super(RenamedNLP, self).__init__(original_nlp)
        self._primals_name_map = primals_name_map
        self._new_primals_names = None
        # Todo: maybe do this on first call instead of __init__?
        self._generate_new_names()

    def _generate_new_names(self):
        if self._new_primals_names is None:
            assert self._original_nlp.n_primals() == len(self._primals_name_map)
            self._new_primals_names = \
                [self._primals_name_map[nm] for nm in self._original_nlp.primals_names()]
            
    def primals_names(self):
        return self._new_primals_names


[docs]class ProjectedNLP(_BaseNLPDelegator): def __init__(self, original_nlp, primals_ordering): """ This class takes an NLP that depends on a set of primals (original space) and converts it to an NLP that depends on a reordered set of primals (projected space). This will impact all the returned items associated with primal variables. E.g., the gradient will be in the new primals ordering instead of the original primals ordering. Note also that this can include additional primal variables not in the original NLP, or can exclude primal variables that were in the original NLP. Parameters ---------- original_nlp : NLP-like The original NLP object that implements the NLP interface primals_ordering: list List of strings indicating the desired primal variable ordering for this NLP. The list can contain new variables that are not in the original NLP, thereby expanding the space of the primal variables. """ super(ProjectedNLP, self).__init__(original_nlp) self._primals_ordering = list(primals_ordering) self._original_idxs = None self._projected_idxs = None self._generate_maps() self._projected_primals = self.init_primals() self._jacobian_nz_mask = None self._hessian_nz_mask = None self._nnz_jacobian = None self._nnz_hessian_lag = None def _generate_maps(self): if self._original_idxs is None or self._projected_idxs is None: primals_ordering_dict = {k:i for i,k in enumerate(self._primals_ordering)} original_names = self._original_nlp.primals_names() original_idxs = list() projected_idxs = list() for i,nm in enumerate(original_names): if nm in primals_ordering_dict: # we need the reordering for this element original_idxs.append(i) projected_idxs.append(primals_ordering_dict[nm]) self._original_idxs = np.asarray(original_idxs) self._projected_idxs = np.asarray(projected_idxs) self._original_to_projected = np.nan*np.zeros(self._original_nlp.n_primals()) self._original_to_projected[self._original_idxs] = self._projected_idxs
[docs] def n_primals(self): return len(self._primals_ordering)
[docs] def primals_names(self): return list(self._primals_ordering)
[docs] def nnz_jacobian(self): if self._nnz_jacobian is None: J = self.evaluate_jacobian() self._nnz_jacobian = len(J.data) return self._nnz_jacobian
[docs] def nnz_hessian_lag(self): if self._nnz_hessian_lag is None: H = self.evaluate_hessian_lag() self._nnz_hessian_lag = len(H.data) return self._nnz_hessian_lag
def _project_primals(self, default, original_primals): projected_x = default*np.ones(self.n_primals(), dtype=np.float64) projected_x[self._projected_idxs] = original_primals[self._original_idxs] return projected_x
[docs] def primals_lb(self): return self._project_primals(-np.inf, self._original_nlp.primals_lb())
[docs] def primals_ub(self): return self._project_primals(np.inf, self._original_nlp.primals_ub())
[docs] def init_primals(self): # Todo: think about what to do here if an entry is not defined # for now, we default to NaN, but there may be a better way # (e.g., taking a new initial value in the constructor?) return self._project_primals(np.nan, self._original_nlp.init_primals())
[docs] def create_new_vector(self, vector_type): if vector_type == 'primals': return np.zeros(self.n_primals(), dtype=np.float64) return self._original_nlp.create_new_vector(vector_type)
[docs] def set_primals(self, primals): # here, we keep a local copy of the projected primals so # we can give back these values in get_primals. This might # not be the best idea since we can't really support this # same strategy for other methods (e.g., init_primals) where # we now use NaNs to fill in any "missing" entries. np.copyto(self._projected_primals, primals) original_primals = self._original_nlp.get_primals() original_primals[self._original_idxs] = primals[self._projected_idxs] self._original_nlp.set_primals(original_primals)
[docs] def get_primals(self): original_primals = self._original_nlp.get_primals() self._projected_primals[self._projected_idxs] = original_primals[self._original_idxs] return self._projected_primals
[docs] def get_primals_scaling(self): return self._project_primals(np.nan, self._original_nlp.get_primals_scaling())
[docs] def evaluate_grad_objective(self, out=None): original_grad_objective = self._original_nlp.evaluate_grad_objective() projected_objective = self._project_primals(0.0, original_grad_objective) if out is None: return projected_objective np.copyto(out, projected_objective) return out
[docs] def evaluate_jacobian(self, out=None): original_jacobian = self._original_nlp.evaluate_jacobian() if out is not None: np.copyto(out.data, original_jacobian.data[self._jacobian_nz_mask]) return out row = original_jacobian.row col = original_jacobian.col data = original_jacobian.data if self._jacobian_nz_mask is None: # need to remap the irow, jcol to the new space and change the size self._jacobian_nz_mask = np.isin(col, self._original_idxs) new_col = col[self._jacobian_nz_mask] new_col = self._original_to_projected[new_col] new_row = row[self._jacobian_nz_mask] new_data = data[self._jacobian_nz_mask] return sp.coo_matrix((new_data, (new_row,new_col)), shape=(self.n_constraints(), self.n_primals()))
[docs] def evaluate_hessian_lag(self, out=None): original_hessian = self._original_nlp.evaluate_hessian_lag() if out is not None: np.copyto(out.data, original_hessian.data[self._hessian_nz_mask]) return out row = original_hessian.row col = original_hessian.col data = original_hessian.data if self._hessian_nz_mask is None: # need to remap the irow, jcol to the new space and change the size self._hessian_nz_mask = np.isin(col, self._original_idxs) & np.isin(row, self._original_idxs) new_col = col[self._hessian_nz_mask] new_col = self._original_to_projected[new_col] new_row = row[self._hessian_nz_mask] new_row = self._original_to_projected[new_row] new_data = data[self._hessian_nz_mask] return sp.coo_matrix((new_data, (new_row,new_col)), shape=(self.n_primals(), self.n_primals()))
[docs] def report_solver_status(self, status_code, status_message): raise NotImplementedError('Need to think about this...')