# Copyright 2021 CR-Suite Development Team
#
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#
# Unless required by applicable law or agreed to in writing, software
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
[docs]def smooth_quad_matrix(P=None, q=None, r=None): r"""Quadratic function and its gradient :math:f(x) = \frac{1}{2} x^T P x + \langle q, x \rangle + r """ if P is not None: P = jnp.asarray(P) P = cnb.promote_arg_dtypes(P) if q is not None: q = jnp.asarray(q) q = cnb.promote_arg_dtypes(q) @jit def func(x): x = jnp.asarray(x) x = cnb.promote_arg_dtypes(x) if P is None: v = 0.5 * cnb.arr_rdot(x, x) else: v = 0.5 * cnb.arr_rdot(P @ x, x) if q is not None: v = v + cnb.arr_rdot(q, x) if r is not None: v = v + r return v @jit def gradient(x): x = jnp.asarray(x) x = cnb.promote_arg_dtypes(x) if P is None: g = x else: g = P @ x if q is not None: g = g + q return g return build2(func, gradient)
[docs]def smooth_quad_error(A, b): r"""Quadratic error function and its gradient :math:f(x) = \frac{1}{2} \| A x - b \|_2^2 """ @jit def func(x): x = jnp.asarray(x) x = cnb.promote_arg_dtypes(x) r = A @ x - b return 0.5 * jnp.dot(r, r) @jit def gradient(x): x = jnp.asarray(x) x = cnb.promote_arg_dtypes(x) r = A @ x - b g = r.T @ A return g @jit def grad_val(x): x = jnp.asarray(x) x = cnb.promote_arg_dtypes(x) r = A @ x - b v = 0.5 * jnp.dot(r, r) g = r.T @ A return g, v return build3(func, gradient, grad_val)