1"""
2OSQP solver pure python implementation
3"""
4from builtins import object
5import osqppurepy._osqp as _osqp # Internal low level module
6from warnings import warn
7import numpy as np
8from scipy import sparse
9
10
11class OSQP(object):
12    def __init__(self):
13        self._model = _osqp.OSQP()
14
15    def version(self):
16        return self._model.version()
17
18    def setup(self, P=None, q=None, A=None, l=None, u=None, **settings):
19        """
20        Setup OSQP solver problem of the form
21
22        minimize     1/2 x' * P * x + q' * x
23        subject to   l <= A * x <= u
24
25        solver settings can be specified as additional keyword arguments
26        """
27
28        #
29        # Get problem dimensions
30        #
31
32        if P is None:
33            if q is not None:
34                n = len(q)
35            elif A is not None:
36                n = A.shape[1]
37            else:
38                raise ValueError("The problem does not have any variables")
39        else:
40            n = P.shape[0]
41        if A is None:
42            m = 0
43        else:
44            m = A.shape[0]
45
46        #
47        # Create parameters if they are None
48        #
49
50        if (A is None and (l is not None or u is not None)) or \
51                (A is not None and (l is None and u is None)):
52            raise ValueError("A must be supplied together " +
53                             "with at least one bound l or u")
54
55        # Add infinity bounds in case they are not specified
56        if A is not None and l is None:
57            l = -np.inf * np.ones(A.shape[0])
58        if A is not None and u is None:
59            u = np.inf * np.ones(A.shape[0])
60
61        # Create elements if they are not specified
62        if P is None:
63            P = sparse.csc_matrix((np.zeros((0,), dtype=np.double),
64                                  np.zeros((0,), dtype=np.int),
65                                  np.zeros((n+1,), dtype=np.int)),
66                                  shape=(n, n))
67        if q is None:
68            q = np.zeros(n)
69
70        if A is None:
71            A = sparse.csc_matrix((np.zeros((0,), dtype=np.double),
72                                  np.zeros((0,), dtype=np.int),
73                                  np.zeros((n+1,), dtype=np.int)),
74                                  shape=(m, n))
75            l = np.zeros(A.shape[0])
76            u = np.zeros(A.shape[0])
77
78        #
79        # Check vector dimensions (not checked from C solver)
80        #
81
82        # Check if second dimension of A is correct
83        # if A.shape[1] != n:
84        #     raise ValueError("Dimension n in A and P does not match")
85        if len(q) != n:
86            raise ValueError("Incorrect dimension of q")
87        if len(l) != m:
88            raise ValueError("Incorrect dimension of l")
89        if len(u) != m:
90            raise ValueError("Incorrect dimension of u")
91
92        #
93        # Check or Sparsify Matrices
94        #
95        if not sparse.issparse(P) and isinstance(P, np.ndarray) and \
96                len(P.shape) == 2:
97            raise TypeError("P is required to be a sparse matrix")
98        if not sparse.issparse(A) and isinstance(A, np.ndarray) and \
99                len(A.shape) == 2:
100            raise TypeError("A is required to be a sparse matrix")
101
102        # Convert matrices in CSC form and to individual pointers
103        if not sparse.isspmatrix_csc(P):
104            warn("Converting sparse P to a CSC " +
105                 "(compressed sparse column) matrix. (It may take a while...)")
106            P = P.tocsc()
107        if not sparse.isspmatrix_csc(A):
108            warn("Converting sparse A to a CSC " +
109                 "(compressed sparse column) matrix. (It may take a while...)")
110            A = A.tocsc()
111
112        # Check if P an A have sorted indices
113        if not P.has_sorted_indices:
114            P.sort_indices()
115        if not A.has_sorted_indices:
116            A.sort_indices()
117
118        # Convert infinity values to OSQP Infinity
119        u = np.minimum(u, self._model.constant('OSQP_INFTY'))
120        l = np.maximum(l, -self._model.constant('OSQP_INFTY'))
121
122        self._model.setup((n, m), P.data, P.indices, P.indptr, q,
123                          A.data, A.indices, A.indptr,
124                          l, u, **settings)
125
126    def update(self, q=None, l=None, u=None, P=None, A=None):
127        """
128        Update OSQP problem arguments
129        """
130
131        # Get problem dimensions
132        (n, m) = (self._model.work.data.n, self._model.work.data.m)
133
134        if P is not None:
135            if P.shape != (n, n):
136                raise ValueError("P must have shape (n x n)")
137            if A is None:
138                self._model.update_P(P)
139
140        if A is not None:
141            if A.shape != (m, n):
142                raise ValueError("A must have shape (m x n)")
143            if P is None:
144                self._model.update_A(A)
145
146        if P is not None and A is not None:
147            self._model.update_P_A(P, A)
148
149        if q is not None:
150            if q.shape != (n,):
151                raise ValueError("q must have shape (n,)")
152            self._model.update_lin_cost(q)
153
154        if l is not None:
155            if l.shape != (m,):
156                raise ValueError("l must have shape (m,)")
157
158            # Convert values to OSQP_INFTY
159            l = np.maximum(l, -self._model.constant('OSQP_INFTY'))
160
161            if u is None:
162                self._model.update_lower_bound(l)
163
164        if u is not None:
165            if u.shape != (m,):
166                raise ValueError("u must have shape (m,)")
167
168            # Convert values to OSQP_INFTY
169            u = np.minimum(u, self._model.constant('OSQP_INFTY'))
170
171            if l is None:
172                self._model.update_upper_bound(u)
173
174        if l is not None and u is not None:
175            self._model.update_bounds(l, u)
176
177        if q is None and l is None and u is None and P is None and A is None:
178            raise ValueError("No updatable data has been specified!")
179
180    def update_settings(self, **kwargs):
181        """
182        Update OSQP solver settings
183
184        It is possible to change: 'max_iter', 'eps_abs', 'eps_rel', 'rho, 'alpha',
185                                  'delta', 'polish', 'polish_refine_iter',
186                                  'verbose', 'scaled_termination',
187                                  'check_termination'
188        """
189
190        # get arguments
191        max_iter = kwargs.pop('max_iter', None)
192        eps_abs = kwargs.pop('eps_abs', None)
193        eps_rel = kwargs.pop('eps_rel', None)
194        rho = kwargs.pop('rho', None)
195        alpha = kwargs.pop('alpha', None)
196        delta = kwargs.pop('delta', None)
197        polish = kwargs.pop('polish', None)
198        polish_refine_iter = kwargs.pop('polish_refine_iter', None)
199        verbose = kwargs.pop('verbose', None)
200        scaled_termination = kwargs.pop('scaled_termination', None)
201        check_termination = kwargs.pop('check_termination', None)
202        warm_start = kwargs.pop('warm_start', None)
203
204        # update them
205        if max_iter is not None:
206            self._model.update_max_iter(max_iter)
207
208        if eps_abs is not None:
209            self._model.update_eps_abs(eps_abs)
210
211        if eps_rel is not None:
212            self._model.update_eps_rel(eps_rel)
213
214        if rho is not None:
215            self._model.update_rho(rho)
216
217        if alpha is not None:
218            self._model.update_alpha(alpha)
219
220        if delta is not None:
221            self._model.update_delta(delta)
222
223        if polish is not None:
224            self._model.update_polish(polish)
225
226        if polish_refine_iter is not None:
227            self._model.update_polish_refine_iter(polish_refine_iter)
228
229        if verbose is not None:
230            self._model.update_verbose(verbose)
231
232        if scaled_termination is not None:
233            self._model.update_scaled_termination(scaled_termination)
234
235        if check_termination is not None:
236            self._model.update_check_termination(check_termination)
237
238        if warm_start is not None:
239            self._model.update_warm_start(warm_start)
240
241        if max_iter is None and \
242           eps_abs is None and \
243           eps_rel is None and \
244           rho is None and \
245           alpha is None and \
246           delta is None and \
247           polish is None and \
248           polish_refine_iter is None and \
249           verbose is None and \
250           scaled_termination is None and \
251           check_termination is None and \
252           warm_start is None:
253            raise ValueError("No updatable settings has been specified!")
254
255    def solve(self):
256        """
257        Solve QP Problem
258        """
259        # Solve QP
260        return self._model.solve()
261
262    def constant(self, constant_name):
263        """
264        Return solver constant
265        """
266        return self._model.constant(constant_name)
267
268    def warm_start(self, x=None, y=None):
269        """
270        Warm start primal or dual variables
271        """
272        # get problem dimensions
273        (n, m) = (self._model.work.data.n, self._model.work.data.m)
274
275        if x is not None:
276            if len(x)!=n:
277                raise ValueError("Wrong dimension for variable x")
278
279            if y is None:
280                self._model.warm_start_x(x)
281
282        if y is not None:
283            if len(y)!=m:
284                raise ValueError("Wrong dimension for variable y")
285
286            if x is None:
287                self._model.warm_start_y(y)
288
289        if x is not None and y is not None:
290            self._model.warm_start(x, y)
291