1"""
2Copyright 2017 Steven Diamond
3
4   Licensed under the Apache License, Version 2.0 (the "License");
5   you may not use this file except in compliance with the License.
6   You may obtain a copy of the License at
7
8       http://www.apache.org/licenses/LICENSE-2.0
9
10   Unless required by applicable law or agreed to in writing, software
11   distributed under the License is distributed on an "AS IS" BASIS,
12   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13   See the License for the specific language governing permissions and
14   limitations under the License.
15"""
16import numbers
17
18import numpy as np
19import scipy.sparse
20
21import cvxpy.cvxcore.python.cvxcore as cvxcore
22import cvxpy.settings as s
23from cvxpy.lin_ops import lin_op as lo
24
25
26def get_parameter_vector(param_size,
27                         param_id_to_col,
28                         param_id_to_size,
29                         param_id_to_value_fn,
30                         zero_offset: bool = False):
31    """Returns a flattened parameter vector
32
33    The flattened vector includes a constant offset (i.e, a 1).
34
35    Parameters
36    ----------
37        param_size: The number of parameters
38        param_id_to_col: A dict from parameter id to column offset
39        param_id_to_size: A dict from parameter id to parameter size
40        param_id_to_value_fn: A callable that returns a value for a parameter id
41        zero_offset: (optional) if True, zero out the constant offset in the
42                     parameter vector
43
44    Returns
45    -------
46        A flattened NumPy array of parameter values, of length param_size + 1
47    """
48    #TODO handle parameters with structure.
49    if param_size == 0:
50        return None
51    param_vec = np.zeros(param_size + 1)
52    for param_id, col in param_id_to_col.items():
53        if param_id == lo.CONSTANT_ID:
54            if not zero_offset:
55                param_vec[col] = 1
56        else:
57            value = param_id_to_value_fn(param_id).flatten(order='F')
58            size = param_id_to_size[param_id]
59            param_vec[col:col + size] = value
60    return param_vec
61
62
63def reduce_problem_data_tensor(A, var_length, quad_form: bool = False):
64    """Reduce a problem data tensor, for efficient construction of the problem data
65
66    If quad_form=False, the problem data tensor A is a matrix of shape (m, p), where p is the
67    length of the parameter vector. The product A@param_vec gives the
68    entries of the problem data matrix for a solver;
69    the solver's problem data matrix has dimensions(n_constr, n_var + 1),
70    and n_constr*(n_var + 1) = m. In other words, each row in A corresponds
71    to an entry in the solver's data matrix.
72
73    If quad_form=True, the problem data tensor A is a matrix of shape (m, p), where p is the
74    length of the parameter vector. The product A@param_vec gives the
75    entries of the quadratic form matrix P for a QP solver;
76    the solver's quadratic matrix P has dimensions(n_var, n_var),
77    and n_var*n_var = m. In other words, each row in A corresponds
78    to an entry in the solver's quadratic form matrix P.
79
80    This function removes the rows in A that are identically zero, since these
81    rows correspond to zeros in the problem data. It also returns the indices
82    and indptr to construct the problem data matrix from the reduced
83    representation of A, and the shape of the problem data matrix.
84
85    Let reduced_A be the sparse matrix returned by this function. Then the
86    problem data can be computed using
87
88       data : = reduced_A @ param_vec
89
90    and the problem data matrix can be constructed with
91
92        problem_data : = scipy.sparse.csc_matrix(
93            (data, indices, indptr), shape = shape)
94
95    Parameters
96    ----------
97        A : A sparse matrix, the problem data tensor; must not have a 0 in its
98            shape
99        var_length: number of variables in the problem
100        quad_form: (optional) if True, consider quadratic form matrix P
101
102    Returns
103    -------
104        reduced_A: A CSR sparse matrix with redundant rows removed
105        indices: CSC indices for the problem data matrix
106        indptr: CSC indptr for the problem data matrix
107        shape: the shape of the problem data matrix
108    """
109    # construct a reduced COO matrix
110    A.eliminate_zeros()
111    A_coo = A.tocoo()
112
113    unique_old_row, reduced_row = np.unique(A_coo.row, return_inverse=True)
114    reduced_A_shape = (unique_old_row.size, A_coo.shape[1])
115
116    # remap the rows
117    reduced_A = scipy.sparse.coo_matrix(
118        (A_coo.data, (reduced_row, A_coo.col)), shape=reduced_A_shape)
119
120    # convert reduced_A to csr
121    reduced_A = reduced_A.tocsr()
122
123    nonzero_rows = unique_old_row
124    n_cols = var_length
125    # add one more column for the offset if not quad_form
126    if not quad_form:
127        n_cols += 1
128    n_constr = A.shape[0] // (n_cols)
129    shape = (n_constr, n_cols)
130    indices = nonzero_rows % (n_constr)
131
132    # cols holds the column corresponding to each row in nonzero_rows
133    cols = nonzero_rows // n_constr
134
135    # construction of the indptr: scan through cols, and find
136    # the structure of the column index pointer
137    indptr = np.zeros(n_cols + 1, dtype=np.int32)
138    positions, counts = np.unique(cols, return_counts=True)
139    indptr[positions+1] = counts
140    indptr = np.cumsum(indptr)
141
142    return reduced_A, indices, indptr, shape
143
144
145def nonzero_csc_matrix(A):
146    # this function returns (rows, cols) corresponding to nonzero entries in
147    # A; an entry that is explicitly set to zero is treated as nonzero
148    assert not np.isnan(A.data).any()
149
150    # scipy drops rows, cols with explicit zeros; use nan as a sentinel
151    # to prevent them from being dropped
152    zero_indices = (A.data == 0)
153    A.data[zero_indices] = np.nan
154
155    # A.nonzero() returns (rows, cols) sorted in C-style order,
156    # but (when A is a csc matrix) A.data is stored in Fortran-order, hence
157    # the sorting below
158    A_rows, A_cols = A.nonzero()
159    ind = np.argsort(A_cols, kind='mergesort')
160    A_rows = A_rows[ind]
161    A_cols = A_cols[ind]
162
163    A.data[zero_indices] = 0
164    return A_rows, A_cols
165
166
167def A_mapping_nonzero_rows(problem_data_tensor, var_length):
168    # get the rows in the map from parameters to problem data that
169    # have any nonzeros
170    problem_data_tensor_csc = problem_data_tensor.tocsc()
171    A_nrows = problem_data_tensor.shape[0] // (var_length + 1)
172    A_ncols = var_length
173    A_mapping = problem_data_tensor_csc[:A_nrows*A_ncols, :-1]
174    # don't call nonzero_csc_matrix, because here we don't want to
175    # count explicit zeros
176    A_mapping_nonzero_rows, _ = A_mapping.nonzero()
177    return np.unique(A_mapping_nonzero_rows)
178
179
180def get_matrix_from_tensor(problem_data_tensor, param_vec,
181                           var_length, nonzero_rows=None,
182                           with_offset=True,
183                           problem_data_index=None):
184    """Applies problem_data_tensor to param_vec to obtain matrix and (optionally)
185    the offset.
186
187    This function applies problem_data_tensor to param_vec to obtain
188    a matrix representation of the corresponding affine map.
189
190    Parameters
191    ----------
192        problem_data_tensor: tensor returned from get_problem_matrix,
193            representing a parameterized affine map
194        param_vec: flattened parameter vector
195        var_length: the number of variables
196        nonzero_rows: (optional) rows in the part of problem_data_tensor
197            corresponding to A that have nonzeros in them (i.e., rows that
198            are affected by parameters); if not None, then the corresponding
199            entries in A will have explicit zeros.
200        with_offset: (optional) return offset. Defaults to True.
201        problem_data_index: (optional) a tuple (indices, indptr, shape) for
202            construction of the CSC matrix holding the problem data and offset
203
204    Returns
205    -------
206        A tuple (A, b), where A is a matrix with `var_length` columns
207        and b is a flattened NumPy array representing the constant offset.
208        If with_offset=False, returned b is None.
209    """
210    if param_vec is None:
211        flat_problem_data = problem_data_tensor
212        if problem_data_index is not None:
213            flat_problem_data = flat_problem_data.toarray().flatten()
214    elif problem_data_index is not None:
215        flat_problem_data = problem_data_tensor @ param_vec
216    else:
217        param_vec = scipy.sparse.csc_matrix(param_vec[:, None])
218        flat_problem_data = problem_data_tensor @ param_vec
219
220
221    if problem_data_index is not None:
222        indices, indptr, shape = problem_data_index
223        M = scipy.sparse.csc_matrix(
224            (flat_problem_data, indices, indptr), shape=shape)
225    else:
226        n_cols = var_length
227        if with_offset:
228            n_cols += 1
229        M = flat_problem_data.reshape((-1, n_cols), order='F').tocsc()
230
231    if with_offset:
232        A = M[:, :-1].tocsc()
233        b = np.squeeze(M[:, -1].toarray().flatten())
234    else:
235        A = M.tocsc()
236        b = None
237
238    if nonzero_rows is not None and nonzero_rows.size > 0:
239        A_nrows, _ = A.shape
240        A_rows, A_cols = nonzero_csc_matrix(A)
241        A_vals = np.append(A.data, np.zeros(nonzero_rows.size))
242        A_rows = np.append(A_rows, nonzero_rows % A_nrows)
243        A_cols = np.append(A_cols, nonzero_rows // A_nrows)
244        A = scipy.sparse.csc_matrix((A_vals, (A_rows, A_cols)),
245                                    shape=A.shape)
246
247    return (A, b)
248
249
250def get_matrix_and_offset_from_unparameterized_tensor(problem_data_tensor,
251                                                      var_length):
252    """Converts unparameterized tensor to matrix offset representation
253
254    problem_data_tensor _must_ have been obtained from calling
255    get_problem_matrix on a problem with 0 parameters.
256
257    Parameters
258    ----------
259        problem_data_tensor: tensor returned from get_problem_matrix,
260            representing an affine map
261        var_length: the number of variables
262
263    Returns
264    -------
265        A tuple (A, b), where A is a matrix with `var_length` columns
266        and b is a flattened NumPy array representing the constant offset.
267    """
268    assert problem_data_tensor.shape[1] == 1
269    return get_matrix_and_offset_from_tensor(
270        problem_data_tensor, None, var_length)
271
272
273def get_problem_matrix(linOps,
274                       var_length,
275                       id_to_col,
276                       param_to_size,
277                       param_to_col,
278                       constr_length):
279    """
280    Builds a sparse representation of the problem data.
281
282    Parameters
283    ----------
284        linOps: A list of python linOp trees representing an affine expression
285        var_length: The total length of the variables.
286        id_to_col: A map from variable id to column offset.
287        param_to_size: A map from parameter id to parameter size.
288        param_to_col: A map from parameter id to column in tensor.
289        constr_length: Summed sizes of constraints input.
290
291    Returns
292    -------
293        A sparse (CSC) matrix with constr_length * (var_length + 1) rows and
294        param_size + 1 columns (where param_size is the length of the
295        parameter vector).
296    """
297    lin_vec = cvxcore.ConstLinOpVector()
298
299    id_to_col_C = cvxcore.IntIntMap()
300    for id, col in id_to_col.items():
301        id_to_col_C[int(id)] = int(col)
302
303    param_to_size_C = cvxcore.IntIntMap()
304    for id, size in param_to_size.items():
305        param_to_size_C[int(id)] = int(size)
306
307    # dict to memoize construction of C++ linOps, and to keep Python references
308    # to them to prevent their deletion
309    linPy_to_linC = {}
310    for lin in linOps:
311        build_lin_op_tree(lin, linPy_to_linC)
312        tree = linPy_to_linC[lin]
313        lin_vec.push_back(tree)
314
315    problemData = cvxcore.build_matrix(lin_vec,
316                                       int(var_length),
317                                       id_to_col_C,
318                                       param_to_size_C,
319                                       s.get_num_threads())
320
321    # Populate tensors with info from problemData.
322    tensor_V = {}
323    tensor_I = {}
324    tensor_J = {}
325    for param_id, size in param_to_size.items():
326        tensor_V[param_id] = []
327        tensor_I[param_id] = []
328        tensor_J[param_id] = []
329        problemData.param_id = param_id
330        for i in range(size):
331            problemData.vec_idx = i
332            prob_len = problemData.getLen()
333            tensor_V[param_id].append(problemData.getV(prob_len))
334            tensor_I[param_id].append(problemData.getI(prob_len))
335            tensor_J[param_id].append(problemData.getJ(prob_len))
336
337    # Reduce tensors to a single sparse CSR matrix.
338    V = []
339    I = []
340    J = []
341    # one of the 'parameters' in param_to_col is a constant scalar offset,
342    # hence 'plus_one'
343    param_size_plus_one = 0
344    for param_id, col in param_to_col.items():
345        size = param_to_size[param_id]
346        param_size_plus_one += size
347        for i in range(size):
348            V.append(tensor_V[param_id][i])
349            I.append(tensor_I[param_id][i] +
350                     tensor_J[param_id][i]*constr_length)
351            J.append(tensor_J[param_id][i]*0 + (i + col))
352    V = np.concatenate(V)
353    I = np.concatenate(I)
354    J = np.concatenate(J)
355    A = scipy.sparse.csc_matrix(
356        (V, (I, J)),
357        shape=(np.int64(constr_length)*np.int64(var_length+1),
358               param_size_plus_one))
359    return A
360
361
362def format_matrix(matrix, shape=None, format='dense'):
363    """Returns the matrix in the appropriate form for SWIG wrapper"""
364    if (format == 'dense'):
365        # Ensure is 2D.
366        if len(shape) == 0:
367            shape = (1, 1)
368        elif len(shape) == 1:
369            shape = shape + (1,)
370        return np.reshape(matrix, shape, order='F')
371    elif(format == 'sparse'):
372        return scipy.sparse.coo_matrix(matrix)
373    elif(format == 'scalar'):
374        return np.asfortranarray([[matrix]])
375    else:
376        raise NotImplementedError()
377
378
379TYPE_MAP = {
380    "VARIABLE": cvxcore.VARIABLE,
381    "PARAM": cvxcore.PARAM,
382    "PROMOTE": cvxcore.PROMOTE,
383    "MUL": cvxcore.MUL,
384    "RMUL": cvxcore.RMUL,
385    "MUL_ELEM": cvxcore.MUL_ELEM,
386    "DIV": cvxcore.DIV,
387    "SUM": cvxcore.SUM,
388    "NEG": cvxcore.NEG,
389    "INDEX": cvxcore.INDEX,
390    "TRANSPOSE": cvxcore.TRANSPOSE,
391    "SUM_ENTRIES": cvxcore.SUM_ENTRIES,
392    "TRACE": cvxcore.TRACE,
393    "RESHAPE": cvxcore.RESHAPE,
394    "DIAG_VEC": cvxcore.DIAG_VEC,
395    "DIAG_MAT": cvxcore.DIAG_MAT,
396    "UPPER_TRI": cvxcore.UPPER_TRI,
397    "CONV": cvxcore.CONV,
398    "HSTACK": cvxcore.HSTACK,
399    "VSTACK": cvxcore.VSTACK,
400    "SCALAR_CONST": cvxcore.SCALAR_CONST,
401    "DENSE_CONST": cvxcore.DENSE_CONST,
402    "SPARSE_CONST": cvxcore.SPARSE_CONST,
403    "NO_OP": cvxcore.NO_OP,
404    "KRON": cvxcore.KRON
405}
406
407
408def get_type(linPy):
409    ty = linPy.type.upper()
410    if ty in TYPE_MAP:
411        return TYPE_MAP[ty]
412    else:
413        raise NotImplementedError("Type %s is not supported." % ty)
414
415
416def set_matrix_data(linC, linPy) -> None:
417    """Calls the appropriate cvxcore function to set the matrix data field of
418       our C++ linOp.
419    """
420    if get_type(linPy) == cvxcore.SPARSE_CONST:
421        coo = format_matrix(linPy.data, format='sparse')
422        linC.set_sparse_data(coo.data, coo.row.astype(float),
423                             coo.col.astype(float), coo.shape[0],
424                             coo.shape[1])
425    else:
426        linC.set_dense_data(format_matrix(linPy.data, shape=linPy.shape))
427        linC.set_data_ndim(len(linPy.data.shape))
428
429
430def set_slice_data(linC, linPy) -> None:
431    """
432    Loads the slice data, start, stop, and step into our C++ linOp.
433    The semantics of the slice operator is treated exactly the same as in
434    Python.  Note that the 'None' cases had to be handled at the wrapper level,
435    since we must load integers into our vector.
436    """
437    for i, sl in enumerate(linPy.data):
438        slice_vec = cvxcore.IntVector()
439        for var in [sl.start, sl.stop, sl.step]:
440            slice_vec.push_back(int(var))
441        linC.push_back_slice_vec(slice_vec)
442
443
444def set_linC_data(linC, linPy) -> None:
445    """Sets numerical data fields in linC."""
446    assert linPy.data is not None
447    if isinstance(linPy.data, tuple) and isinstance(linPy.data[0], slice):
448        slice_data = set_slice_data(linC, linPy)
449    elif isinstance(linPy.data, float) or isinstance(linPy.data,
450                                                   numbers.Integral):
451        linC.set_dense_data(format_matrix(linPy.data, format='scalar'))
452        linC.set_data_ndim(0)
453    else:
454        set_matrix_data(linC, linPy)
455
456
457def make_linC_from_linPy(linPy, linPy_to_linC) -> None:
458    """Construct a C++ LinOp corresponding to LinPy.
459
460    Children of linPy are retrieved from linPy_to_linC.
461    """
462    if linPy in linPy_to_linC:
463        return
464    typ = get_type(linPy)
465    shape = cvxcore.IntVector()
466    lin_args_vec = cvxcore.ConstLinOpVector()
467    for dim in linPy.shape:
468        shape.push_back(int(dim))
469    for argPy in linPy.args:
470        lin_args_vec.push_back(linPy_to_linC[argPy])
471    linC = cvxcore.LinOp(typ, shape, lin_args_vec)
472    linPy_to_linC[linPy] = linC
473
474    if linPy.data is not None:
475        if isinstance(linPy.data, lo.LinOp):
476            linC_data = linPy_to_linC[linPy.data]
477            linC.set_linOp_data(linC_data)
478            linC.set_data_ndim(len(linPy.data.shape))
479        else:
480            set_linC_data(linC, linPy)
481
482
483def build_lin_op_tree(root_linPy, linPy_to_linC) -> None:
484    """Construct C++ LinOp tree from Python LinOp tree.
485
486    Constructed C++ linOps are stored in the linPy_to_linC dict,
487    which maps Python linOps to their corresponding C++ linOps.
488
489    Parameters
490    ----------
491        linPy_to_linC: a dict for memoizing construction and storing
492            the C++ LinOps
493    """
494    bfs_stack = [root_linPy]
495    post_order_stack = []
496    while bfs_stack:
497        linPy = bfs_stack.pop()
498        if linPy not in linPy_to_linC:
499            post_order_stack.append(linPy)
500            for arg in linPy.args:
501                bfs_stack.append(arg)
502            if isinstance(linPy.data, lo.LinOp):
503                bfs_stack.append(linPy.data)
504    while post_order_stack:
505        linPy = post_order_stack.pop()
506        make_linC_from_linPy(linPy, linPy_to_linC)
507