Skip to content

Commit 87236d7

Browse files
dance858dance858
committed
problem class
1 parent 68c932c commit 87236d7

4 files changed

Lines changed: 473 additions & 0 deletions

File tree

‎sparsediffpy/__init__.py‎

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -24,6 +24,9 @@
2424
# Compile
2525
from sparsediffpy._core._compile import compile # noqa: F401
2626

27+
# Problem
28+
from sparsediffpy._core._problem import Problem # noqa: F401
29+
2730
# Elementwise functions
2831
from sparsediffpy._core._fn_elementwise import ( # noqa: F401
2932
sin, cos, exp, log, tan, sinh, tanh, asinh, atanh,

‎sparsediffpy/_core/_problem.py‎

Lines changed: 194 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,194 @@
1+
"""Problem class: wraps a C problem capsule (objective + list of constraints).
2+
3+
Takes SparseDiffPy expressions (not CVXPY, not pre-built capsules). CVXPY-facing
4+
adapters live in downstream libraries (e.g. DNLP).
5+
"""
6+
7+
import numpy as np
8+
9+
from sparsediffpy import _sparsediffengine as _C
10+
from sparsediffpy._core._compile import _build_capsule, _collect_leaves
11+
12+
13+
class Problem:
14+
"""A compiled NLP-style problem: one scalar objective plus a list of constraints.
15+
16+
Method names mirror DNLP's `C_problem` so a CVXPY adapter can return a
17+
Problem and existing solver callsites keep working.
18+
"""
19+
20+
def __init__(self, objective, constraints=None, verbose=False):
21+
constraints = list(constraints) if constraints else []
22+
23+
if objective.shape != (1, 1):
24+
raise ValueError(
25+
f"Objective must be scalar (shape (1, 1)), got {objective.shape}"
26+
)
27+
28+
variables, parameters = [], []
29+
visited = set()
30+
_collect_leaves(objective, variables, parameters, visited)
31+
for c in constraints:
32+
_collect_leaves(c, variables, parameters, visited)
33+
34+
if not variables:
35+
raise ValueError("Problem must contain at least one Variable")
36+
37+
scope = variables[0]._scope
38+
for v in variables[1:]:
39+
if v._scope is not scope:
40+
raise ValueError("All variables must belong to the same Scope")
41+
42+
n_vars = scope._next_var_offset
43+
44+
# One shared cache across objective + all constraints: CSE in both
45+
# directions (within an expression, and across the obj/constraint
46+
# boundary) is safe, and each Parameter capsule is appended to
47+
# param_caps exactly once.
48+
cache = {}
49+
param_caps, param_objs = [], []
50+
obj_cap = _build_capsule(objective, n_vars, cache, param_caps, param_objs)
51+
constraint_caps = [
52+
_build_capsule(c, n_vars, cache, param_caps, param_objs)
53+
for c in constraints
54+
]
55+
56+
self._capsule = _C.make_problem(obj_cap, constraint_caps, verbose)
57+
if param_caps:
58+
_C.problem_register_params(self._capsule, param_caps)
59+
60+
self._scope = scope
61+
self._param_capsules = param_caps
62+
self._param_objects = param_objs
63+
self._n_vars = n_vars
64+
self._total_constraint_size = sum(c.size for c in constraints)
65+
self._jacobian_coo_initialized = False
66+
self._hessian_coo_initialized = False
67+
68+
if param_caps:
69+
self._sync_params()
70+
71+
# ------------------------------------------------------------------
72+
# Internal
73+
# ------------------------------------------------------------------
74+
75+
def _sync_params(self):
76+
"""Push current Parameter values to the C problem.
77+
78+
Called once at construction. After construction, callers invoke
79+
update_params(theta) explicitly (matching DNLP's solver-loop contract).
80+
"""
81+
for p in self._param_objects:
82+
if p._value_flat is None:
83+
raise ValueError(
84+
f"Parameter with shape {p.shape} has no value set. "
85+
f"Assign a value via parameter.value = ... before constructing Problem."
86+
)
87+
theta = np.concatenate([p._value_flat for p in self._param_objects])
88+
_C.problem_update_params(self._capsule, theta)
89+
self._scope._params_dirty = False
90+
91+
# ------------------------------------------------------------------
92+
# Parameter updates
93+
# ------------------------------------------------------------------
94+
95+
def update_params(self, theta):
96+
"""Update parameter values in the C DAG from a flat theta vector.
97+
98+
Sparsity structures (Jacobian/Hessian) remain valid after this call.
99+
"""
100+
theta = np.asarray(theta, dtype=np.float64)
101+
_C.problem_update_params(self._capsule, theta)
102+
self._scope._params_dirty = False
103+
104+
# ------------------------------------------------------------------
105+
# Sparsity initialization (COO)
106+
# ------------------------------------------------------------------
107+
108+
def init_jacobian_coo(self):
109+
"""Fill sparsity for the constraint Jacobian in COO format.
110+
111+
Must be called once before get_jacobian_sparsity_coo() or eval_jacobian_vals().
112+
"""
113+
_C.problem_init_jacobian_coo(self._capsule)
114+
self._jacobian_coo_initialized = True
115+
116+
def init_hessian_coo_lower_tri(self):
117+
"""Fill sparsity for the Lagrangian Hessian (lower triangle, COO).
118+
119+
Must be called once before get_problem_hessian_sparsity_coo() or
120+
eval_hessian_vals_coo_lower_tri().
121+
"""
122+
_C.problem_init_hessian_coo_lower_triangular(self._capsule)
123+
self._hessian_coo_initialized = True
124+
125+
# ------------------------------------------------------------------
126+
# Forward evaluation
127+
# ------------------------------------------------------------------
128+
129+
def objective_forward(self, u):
130+
"""Evaluate the objective at variable values `u`. Returns a float."""
131+
u = np.asarray(u, dtype=np.float64)
132+
return _C.problem_objective_forward(self._capsule, u)
133+
134+
def constraint_forward(self, u):
135+
"""Evaluate constraints at variable values `u`. Returns an np.ndarray."""
136+
u = np.asarray(u, dtype=np.float64)
137+
return _C.problem_constraint_forward(self._capsule, u)
138+
139+
def gradient(self):
140+
"""Compute gradient of the objective. Call objective_forward first."""
141+
return _C.problem_gradient(self._capsule)
142+
143+
# ------------------------------------------------------------------
144+
# Jacobian (COO path)
145+
# ------------------------------------------------------------------
146+
147+
def get_jacobian_sparsity_coo(self):
148+
"""Return the sparsity pattern (rows, cols) of the constraint Jacobian.
149+
150+
Call init_jacobian_coo() first.
151+
"""
152+
rows, cols, _shape = _C.get_jacobian_sparsity_coo(self._capsule)
153+
return rows, cols
154+
155+
def eval_jacobian_vals(self):
156+
"""Evaluate the constraint Jacobian and return its nonzero values.
157+
158+
Values correspond to the sparsity pattern from get_jacobian_sparsity_coo().
159+
Call constraint_forward() first to set the evaluation point.
160+
"""
161+
return _C.problem_eval_jacobian_vals(self._capsule)
162+
163+
# ------------------------------------------------------------------
164+
# Lagrangian Hessian (COO lower-triangular path)
165+
# ------------------------------------------------------------------
166+
167+
def get_problem_hessian_sparsity_coo(self):
168+
"""Return the sparsity pattern (rows, cols) of the lower-triangular
169+
Lagrangian Hessian.
170+
171+
Call init_hessian_coo_lower_tri() first.
172+
"""
173+
rows, cols, _shape = _C.get_problem_hessian_sparsity_coo(self._capsule)
174+
return rows, cols
175+
176+
def eval_hessian_vals_coo_lower_tri(self, obj_factor, lagrange):
177+
"""Evaluate the lower-triangular Lagrangian Hessian values.
178+
179+
Computes obj_factor * H_f + sum_i lagrange[i] * H_gi, where f is the
180+
objective and g_i are the constraints. Values correspond to the sparsity
181+
pattern from get_problem_hessian_sparsity_coo().
182+
183+
Call objective_forward() and constraint_forward() first to set the
184+
evaluation point.
185+
"""
186+
lagrange = np.asarray(lagrange, dtype=np.float64)
187+
if lagrange.size != self._total_constraint_size:
188+
raise ValueError(
189+
f"lagrange length {lagrange.size} != total_constraint_size "
190+
f"{self._total_constraint_size}"
191+
)
192+
return _C.problem_eval_hessian_vals_coo(
193+
self._capsule, float(obj_factor), lagrange
194+
)

‎tests/problem/__init__.py‎

Whitespace-only changes.

0 commit comments

Comments
 (0)