From 439229b351c65e2535c2612cd64848c9879e781b Mon Sep 17 00:00:00 2001 From: healeyq3 Date: Sat, 29 Jun 2024 21:30:58 -0400 Subject: [PATCH 1/3] refactored cone_program for solve_only --- diffcp/cone_program.py | 128 +++++++++++++++++++++++++++++++++++------ 1 file changed, 110 insertions(+), 18 deletions(-) diff --git a/diffcp/cone_program.py b/diffcp/cone_program.py index acdd5a6..a865b4e 100644 --- a/diffcp/cone_program.py +++ b/diffcp/cone_program.py @@ -146,6 +146,62 @@ def DTi(i): return xs, ys, ss, D_batch, DT_batch +def solve_only_wrapper(A, b, c, cone_dict, warm_start, mode, kwargs): + """A wrapper around solve_only for the batch function""" + return solve_only( + A, b, c, cone_dict, warm_start=warm_start, mode=mode, **kwargs) + + +def solve_only_batch(As, bs, cs, cone_dicts, n_jobs_forward=-1, mode="lsqr", + warm_starts=None, **kwargs): + """ + Solves a batch of cone programs. + Uses a ThreadPool to perform operations across + the batch in parallel. + + For more information on the arguments and return values, + see the docstring for `solve_and_derivative_batch` function. + + This function simply contains the first half of + the functionality contained in `solve_and_derivative_batch`. + For differentiating through a cone program, this function is of no use. + This function exists because cvxpylayers utilizes `solve_and_derivative_batch` + to solve an optimization problem and populate the backward function (in PyTorch dialect) + during a forward pass through a cvxpylayer. However, because at inference time + gradient information is no longer desired, the limited functionality provided + by `solve_only_batch` was wanted for computational enhancements. + """ + batch_size = len(As) + if warm_starts is None: + warm_starts = [None] * batch_size + if n_jobs_forward == -1: + n_jobs_forward = mp.cpu_count() + n_jobs_forward = min(batch_size, n_jobs_forward) + + if n_jobs_forward == 1: + #serial + xs, ys, ss = [], [], [] + for i in range(batch_size): + x, y, s = solve_only(As[i], bs[i], cs[i], cone_dicts[i], + warm_starts[i], mode=mode, **kwargs) + xs += [x] + ys += [y] + ss += [s] + else: + # thread pool + pool = ThreadPool(processes=n_jobs_forward) + args = [(A, b, c, cone_dict, warm_start, mode, kwargs) for A, b, c, cone_dict, warm_start in + zip(As, bs, cs, cone_dicts, warm_starts)] + with threadpool_limits(limits=1): + results = pool.starmap(solve_only_wrapper, args) + pool.close() + xs = [r[0] for r in results] + ys = [r[1] for r in results] + ss = [r[2] for r in results] + + return xs, ys, ss + + class SolverError(Exception): pass @@ -225,7 +281,27 @@ def solve_and_derivative(A, b, c, cone_dict, warm_start=None, mode='lsqr', return x, y, s, D, DT -def solve_and_derivative_internal(A, b, c, cone_dict, solve_method=None, +def solve_only(A, b, c, cone_dict, warm_start=None, mode='lsqr', + solve_method='SCS', **kwargs): + """ + Solves a cone program and returns its solution. + + For more information on the arguments and return values, + see the docstring for `solve_and_derivative` function. However, note + that only x, y, and s are being returned from this function. + + This is another function which was created for the benefit of cvxpylayers. + """ + result = solve_internal( + A, b, c, cone_dict, warm_start=warm_start, mode=mode, + solve_method=solve_method, **kwargs) + x = result["x"] + y = result["y"] + s = result["s"] + return x, y, s + + +def solve_internal(A, b, c, cone_dict, solve_method=None, warm_start=None, mode='lsqr', raise_on_error=True, **kwargs): if mode not in ["dense", "lsqr", "lsmr"]: raise ValueError("Unsupported mode {}; the supported modes are " @@ -233,17 +309,16 @@ def solve_and_derivative_internal(A, b, c, cone_dict, solve_method=None, if np.isnan(A.data).any(): raise RuntimeError("Found a NaN in A.") - - # set explicit 0s in A to np.nan - A.data[A.data == 0] = np.nan - - # compute rows and cols of nonzeros in A - rows, cols = A.nonzero() - - # reset np.nan entries in A to 0.0 - A.data[np.isnan(A.data)] = 0.0 - - # eliminate explicit zeros in A, we no longer need them + + ''' + TODO(quill): in solve_and_derivative_internal (sdi) there are more operations + on A to compute "rows" and "columns" variables. Furthermore, when + sdi is called, A.eliminate_zeros() is performed 2x. + An alternative design would be to performs op1, op2, op3 (labeled in sdi) + before calling solve_internal, and then return the A computed here. + Perhaps this is all a non-factor, but I wanted to call attention to it + in case this was worth changing. + ''' A.eliminate_zeros() if solve_method is None: @@ -304,10 +379,6 @@ def solve_and_derivative_internal(A, b, c, cone_dict, solve_method=None, result["DT"] = None return result - x = result["x"] - y = result["y"] - s = result["s"] - elif solve_method == "ECOS": if warm_start is not None: raise ValueError('ECOS does not support warmstart.') @@ -430,8 +501,6 @@ def solve_and_derivative_internal(A, b, c, cone_dict, solve_method=None, result["y"] = np.array(solution.z) result["s"] = np.array(solution.s) - x, y, s = result["x"], result["y"], result["s"] - CLARABEL2SCS_STATUS_MAP = { "Solved": "Solved", "PrimalInfeasible": "Infeasible", @@ -450,7 +519,30 @@ def solve_and_derivative_internal(A, b, c, cone_dict, solve_method=None, } else: raise ValueError("Solver %s not supported." % solve_method) + + return result + +def solve_and_derivative_internal(A, b, c, cone_dict, solve_method=None, + warm_start=None, mode='lsqr', raise_on_error=True, **kwargs): + + result = solve_internal(A, b, c, cone_dict, solve_method=solve_method, + warm_start=warm_start, mode=mode, raise_on_error=raise_on_error, **kwargs) + x = result["x"] + y = result["y"] + s = result["s"] + # set explicit 0s in A to np.nan (op1) + A.data[A.data == 0] = np.nan + + # compute rows and cols of nonzeros in A (op2) + rows, cols = A.nonzero() + + # reset np.nan entries in A to 0.0 (op3) + A.data[np.isnan(A.data)] = 0.0 + + # eliminate explicit zeros in A, we no longer need them + A.eliminate_zeros() + # pre-compute quantities for the derivative m, n = A.shape N = m + n + 1 From d9a7a07507dcbbbc2cfef38c9a72b5a6952cad8f Mon Sep 17 00:00:00 2001 From: healeyq3 Date: Tue, 2 Jul 2024 09:40:28 -0400 Subject: [PATCH 2/3] addressed review suggestions --- diffcp/cone_program.py | 58 ++++++++++++++++++------------------------ 1 file changed, 25 insertions(+), 33 deletions(-) diff --git a/diffcp/cone_program.py b/diffcp/cone_program.py index a865b4e..9919108 100644 --- a/diffcp/cone_program.py +++ b/diffcp/cone_program.py @@ -146,13 +146,13 @@ def DTi(i): return xs, ys, ss, D_batch, DT_batch -def solve_only_wrapper(A, b, c, cone_dict, warm_start, mode, kwargs): +def solve_only_wrapper(A, b, c, cone_dict, warm_start, kwargs): """A wrapper around solve_only for the batch function""" return solve_only( - A, b, c, cone_dict, warm_start=warm_start, mode=mode, **kwargs) + A, b, c, cone_dict, warm_start=warm_start, **kwargs) -def solve_only_batch(As, bs, cs, cone_dicts, n_jobs_forward=-1, mode="lsqr", +def solve_only_batch(As, bs, cs, cone_dicts, n_jobs_forward=-1, warm_starts=None, **kwargs): """ Solves a batch of cone programs. @@ -169,7 +169,7 @@ def solve_only_batch(As, bs, cs, cone_dicts, n_jobs_forward=-1, mode="lsqr", to solve an optimization problem and populate the backward function (in PyTorch dialect) during a forward pass through a cvxpylayer. However, because at inference time gradient information is no longer desired, the limited functionality provided - by `solve_only_batch` was wanted for computational enhancements. + by `solve_only_batch` was wanted for computational efficiency. """ batch_size = len(As) if warm_starts is None: @@ -183,14 +183,14 @@ def solve_only_batch(As, bs, cs, cone_dicts, n_jobs_forward=-1, mode="lsqr", xs, ys, ss = [], [], [] for i in range(batch_size): x, y, s = solve_only(As[i], bs[i], cs[i], cone_dicts[i], - warm_starts[i], mode=mode, **kwargs) + warm_starts[i], **kwargs) xs += [x] ys += [y] ss += [s] else: # thread pool pool = ThreadPool(processes=n_jobs_forward) - args = [(A, b, c, cone_dict, warm_start, mode, kwargs) for A, b, c, cone_dict, warm_start in + args = [(A, b, c, cone_dict, warm_start, kwargs) for A, b, c, cone_dict, warm_start in zip(As, bs, cs, cone_dicts, warm_starts)] with threadpool_limits(limits=1): results = pool.starmap(solve_only_wrapper, args) @@ -281,8 +281,8 @@ def solve_and_derivative(A, b, c, cone_dict, warm_start=None, mode='lsqr', return x, y, s, D, DT -def solve_only(A, b, c, cone_dict, warm_start=None, mode='lsqr', - solve_method='SCS', **kwargs): +def solve_only(A, b, c, cone_dict, warm_start=None, + solve_method='SCS', **kwargs): """ Solves a cone program and returns its solution. @@ -292,8 +292,12 @@ def solve_only(A, b, c, cone_dict, warm_start=None, mode='lsqr', This is another function which was created for the benefit of cvxpylayers. """ + if np.isnan(A.data).any(): + raise RuntimeError("Found a NaN in A.") + A.eliminate_zeros() + result = solve_internal( - A, b, c, cone_dict, warm_start=warm_start, mode=mode, + A, b, c, cone_dict, warm_start=warm_start, solve_method=solve_method, **kwargs) x = result["x"] y = result["y"] @@ -302,24 +306,7 @@ def solve_only(A, b, c, cone_dict, warm_start=None, mode='lsqr', def solve_internal(A, b, c, cone_dict, solve_method=None, - warm_start=None, mode='lsqr', raise_on_error=True, **kwargs): - if mode not in ["dense", "lsqr", "lsmr"]: - raise ValueError("Unsupported mode {}; the supported modes are " - "'dense', 'lsqr' and 'lsmr'".format(mode)) - - if np.isnan(A.data).any(): - raise RuntimeError("Found a NaN in A.") - - ''' - TODO(quill): in solve_and_derivative_internal (sdi) there are more operations - on A to compute "rows" and "columns" variables. Furthermore, when - sdi is called, A.eliminate_zeros() is performed 2x. - An alternative design would be to performs op1, op2, op3 (labeled in sdi) - before calling solve_internal, and then return the A computed here. - Perhaps this is all a non-factor, but I wanted to call attention to it - in case this was worth changing. - ''' - A.eliminate_zeros() + warm_start=None, raise_on_error=True, **kwargs): if solve_method is None: psd_cone = ('s' in cone_dict) and (cone_dict['s'] != []) @@ -524,12 +511,11 @@ def solve_internal(A, b, c, cone_dict, solve_method=None, def solve_and_derivative_internal(A, b, c, cone_dict, solve_method=None, warm_start=None, mode='lsqr', raise_on_error=True, **kwargs): - - result = solve_internal(A, b, c, cone_dict, solve_method=solve_method, - warm_start=warm_start, mode=mode, raise_on_error=raise_on_error, **kwargs) - x = result["x"] - y = result["y"] - s = result["s"] + if mode not in ["dense", "lsqr", "lsmr"]: + raise ValueError("Unsupported mode {}; the supported modes are " + "'dense', 'lsqr' and 'lsmr'".format(mode)) + if np.isnan(A.data).any(): + raise RuntimeError("Found a NaN in A.") # set explicit 0s in A to np.nan (op1) A.data[A.data == 0] = np.nan @@ -542,6 +528,12 @@ def solve_and_derivative_internal(A, b, c, cone_dict, solve_method=None, # eliminate explicit zeros in A, we no longer need them A.eliminate_zeros() + + result = solve_internal(A, b, c, cone_dict, solve_method=solve_method, + warm_start=warm_start, raise_on_error=raise_on_error, **kwargs) + x = result["x"] + y = result["y"] + s = result["s"] # pre-compute quantities for the derivative m, n = A.shape From 095d9105c032ff88d363b5400ee1c11307b94922 Mon Sep 17 00:00:00 2001 From: Parth Nobel Date: Tue, 2 Jul 2024 17:18:21 -0700 Subject: [PATCH 3/3] Update build.yml --- .github/workflows/build.yml | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/.github/workflows/build.yml b/.github/workflows/build.yml index 2d2eb14..12fa1b5 100644 --- a/.github/workflows/build.yml +++ b/.github/workflows/build.yml @@ -24,7 +24,7 @@ jobs: strategy: fail-fast: false matrix: - os: [ ubuntu-20.04, macos-11, windows-2022 ] + os: [ ubuntu-20.04, macos-12, windows-2022 ] python-version: [ 3.8, 3.9, "3.10", "3.11" ] env: @@ -64,13 +64,13 @@ jobs: strategy: fail-fast: false matrix: - os: [ ubuntu-20.04, macos-11, windows-2022 ] - python-version: [ 3.7, 3.9, "3.10", "3.11" ] + os: [ ubuntu-20.04, macos-12, windows-2022 ] + python-version: [ 3.8, 3.9, "3.10", "3.11" ] include: - os: ubuntu-20.04 python-version: 3.8 single_action_config: "True" - - os: macos-11 + - os: macos-12 python-version: 3.8 - os: windows-2019 python-version: 3.8