Skip to content

Commit 84d945b

Browse files
dance858dance858
committed
first draft
1 parent 90385ff commit 84d945b

60 files changed

Lines changed: 4066 additions & 0 deletions

Some content is hidden

Large Commits have some content hidden by default. Use the searchbox below for content that may be hidden.

‎sparsediffpy/__init__.py‎

Lines changed: 36 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -14,4 +14,40 @@
1414
limitations under the License.
1515
"""
1616

17+
# C extension (low-level, for advanced users)
1718
from sparsediffpy import _sparsediffengine # noqa: F401
19+
20+
# Core classes
21+
from sparsediffpy._core._scope import Scope, Variable, Parameter # noqa: F401
22+
from sparsediffpy._core._expression import Expression # noqa: F401
23+
24+
# Compile
25+
from sparsediffpy._core._compile import compile # noqa: F401
26+
27+
# Named functions
28+
from sparsediffpy._core._functions import ( # noqa: F401
29+
sin,
30+
cos,
31+
exp,
32+
log,
33+
tan,
34+
sinh,
35+
tanh,
36+
asinh,
37+
atanh,
38+
logistic,
39+
normal_cdf,
40+
entr,
41+
xexp,
42+
diag_vec,
43+
power,
44+
sum,
45+
prod,
46+
reshape,
47+
trace,
48+
hstack,
49+
vstack,
50+
quad_form,
51+
quad_over_lin,
52+
rel_entr,
53+
)

‎sparsediffpy/_core/__init__.py‎

Whitespace-only changes.

‎sparsediffpy/_core/_compile.py‎

Lines changed: 353 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,353 @@
1+
"""compile() and CompiledExpression.
2+
3+
The recursive tree walker that converts Python expression nodes to C capsules.
4+
Node-type-to-C-call mapping lives in _registry.py.
5+
"""
6+
7+
import numpy as np
8+
import scipy.sparse
9+
10+
from sparsediffpy import _sparsediffengine as _C
11+
from sparsediffpy._core._constants import Constant, SparseConstant
12+
from sparsediffpy._core._nodes_affine import (
13+
LeftMatMul,
14+
ParamScalarMult,
15+
ParamVectorMult,
16+
RightMatMul,
17+
)
18+
from sparsediffpy._core._registry import (
19+
ATOM_CONVERTERS,
20+
_to_dense_row_major,
21+
make_dense_left_matmul,
22+
make_dense_right_matmul,
23+
make_sparse_left_matmul,
24+
make_sparse_right_matmul,
25+
)
26+
from sparsediffpy._core._scope import Parameter, Variable
27+
28+
29+
def compile(expr):
30+
"""Compile an expression tree into a CompiledExpression.
31+
32+
Walks the Python expression tree, discovers all Variables and Parameters,
33+
builds C capsules bottom-up, creates a C problem, and initializes
34+
sparsity patterns for Jacobian and Hessian computation.
35+
"""
36+
# 1. Collect all Variable and Parameter leaves
37+
variables = []
38+
parameters = []
39+
_collect_leaves(expr, variables, parameters, set())
40+
41+
# 2. Determine the scope
42+
scope = None
43+
for v in variables:
44+
if scope is None:
45+
scope = v._scope
46+
elif v._scope is not scope:
47+
raise ValueError("All variables must belong to the same Scope")
48+
49+
if scope is None:
50+
from sparsediffpy._core._scope import Scope
51+
scope = Scope()
52+
53+
n_vars = scope._next_var_offset
54+
if n_vars == 0:
55+
n_vars = 1 # C layer needs at least 1 variable
56+
57+
# 3. Build C capsules bottom-up
58+
capsule_cache = {}
59+
param_capsules_ordered = []
60+
param_objects_ordered = []
61+
root_capsule = _build_capsule(
62+
expr, n_vars, capsule_cache, param_capsules_ordered, param_objects_ordered
63+
)
64+
65+
# 4. Create dummy zero objective (scalar)
66+
dummy_obj = _C.make_parameter(1, 1, -1, n_vars, np.array([0.0]))
67+
68+
# 5. Create C problem with expr as the single constraint
69+
problem = _C.make_problem(dummy_obj, [root_capsule], False)
70+
71+
# 6. Register parameters if any
72+
if param_capsules_ordered:
73+
_C.problem_register_params(problem, param_capsules_ordered)
74+
75+
# 7. Init sparsity patterns
76+
_C.problem_init_jacobian(problem)
77+
_C.problem_init_hessian(problem)
78+
79+
return CompiledExpression(
80+
problem_capsule=problem,
81+
scope=scope,
82+
param_capsules=param_capsules_ordered,
83+
param_objects=param_objects_ordered,
84+
expr_shape=expr.shape,
85+
n_vars=n_vars,
86+
)
87+
88+
89+
# ---------------------------------------------------------------------------
90+
# Tree walking
91+
# ---------------------------------------------------------------------------
92+
93+
def _collect_leaves(node, variables, parameters, visited):
94+
"""Walk the expression tree to collect Variable and Parameter leaves."""
95+
node_id = id(node)
96+
if node_id in visited:
97+
return
98+
visited.add(node_id)
99+
100+
if isinstance(node, Variable):
101+
variables.append(node)
102+
return
103+
if isinstance(node, Parameter):
104+
parameters.append(node)
105+
return
106+
if isinstance(node, (Constant, SparseConstant)):
107+
return
108+
109+
# Walk children
110+
if hasattr(node, "child"):
111+
_collect_leaves(node.child, variables, parameters, visited)
112+
if hasattr(node, "left"):
113+
_collect_leaves(node.left, variables, parameters, visited)
114+
if hasattr(node, "right"):
115+
_collect_leaves(node.right, variables, parameters, visited)
116+
if hasattr(node, "x") and hasattr(node, "z"):
117+
_collect_leaves(node.x, variables, parameters, visited)
118+
_collect_leaves(node.z, variables, parameters, visited)
119+
elif hasattr(node, "x") and hasattr(node, "y"):
120+
_collect_leaves(node.x, variables, parameters, visited)
121+
_collect_leaves(node.y, variables, parameters, visited)
122+
if hasattr(node, "param_expr"):
123+
_collect_leaves(node.param_expr, variables, parameters, visited)
124+
if hasattr(node, "matrix_expr"):
125+
_collect_leaves(node.matrix_expr, variables, parameters, visited)
126+
if hasattr(node, "children"):
127+
for c in node.children:
128+
_collect_leaves(c, variables, parameters, visited)
129+
130+
131+
# ---------------------------------------------------------------------------
132+
# Capsule building
133+
# ---------------------------------------------------------------------------
134+
135+
def _build_capsule(node, n_vars, cache, param_caps, param_objs):
136+
"""Recursively build C capsules for the expression tree."""
137+
node_id = id(node)
138+
if node_id in cache:
139+
return cache[node_id]
140+
141+
cap = _convert_node(node, n_vars, cache, param_caps, param_objs)
142+
143+
# Post-conversion dimension check
144+
d1_c, d2_c = _C.get_expr_dimensions(cap)
145+
d1_py, d2_py = node.shape
146+
if d1_c != d1_py or d2_c != d2_py:
147+
raise ValueError(
148+
f"Dimension mismatch for {type(node).__name__}: "
149+
f"C dimensions ({d1_c}, {d2_c}) vs Python dimensions ({d1_py}, {d2_py})"
150+
)
151+
152+
cache[node_id] = cap
153+
return cap
154+
155+
156+
def _convert_node(node, n_vars, cache, param_caps, param_objs):
157+
"""Convert a single Python expression node to a C capsule."""
158+
159+
# --- Leaves ---
160+
if isinstance(node, Variable):
161+
return _C.make_variable(
162+
node.shape[0], node.shape[1], node._var_id, n_vars
163+
)
164+
165+
if isinstance(node, Parameter):
166+
cap = _C.make_parameter(
167+
node.shape[0], node.shape[1], node._param_id, n_vars,
168+
node._value_flat,
169+
)
170+
param_caps.append(cap)
171+
param_objs.append(node)
172+
return cap
173+
174+
if isinstance(node, Constant):
175+
return _C.make_parameter(
176+
node.shape[0], node.shape[1], -1, n_vars, node._value_flat
177+
)
178+
179+
if isinstance(node, SparseConstant):
180+
return _C.make_parameter(
181+
node.shape[0], node.shape[1], -1, n_vars, node._to_dense_flat()
182+
)
183+
184+
# --- Matmul and multiply with parameter dispatch ---
185+
# These need special handling because they access matrix_expr / param_expr
186+
# directly rather than going through a uniform children list.
187+
if isinstance(node, LeftMatMul):
188+
return _convert_left_matmul(node, n_vars, cache, param_caps, param_objs)
189+
190+
if isinstance(node, RightMatMul):
191+
return _convert_right_matmul(node, n_vars, cache, param_caps, param_objs)
192+
193+
if isinstance(node, ParamScalarMult):
194+
param_cap = _build_capsule(node.param_expr, n_vars, cache, param_caps, param_objs)
195+
child_cap = _build_capsule(node.child, n_vars, cache, param_caps, param_objs)
196+
return _C.make_param_scalar_mult(param_cap, child_cap)
197+
198+
if isinstance(node, ParamVectorMult):
199+
param_cap = _build_capsule(node.param_expr, n_vars, cache, param_caps, param_objs)
200+
child_cap = _build_capsule(node.child, n_vars, cache, param_caps, param_objs)
201+
return _C.make_param_vector_mult(param_cap, child_cap)
202+
203+
# --- Registry lookup ---
204+
node_type = type(node)
205+
if node_type in ATOM_CONVERTERS:
206+
child_caps = _build_children(node, n_vars, cache, param_caps, param_objs)
207+
return ATOM_CONVERTERS[node_type](node, child_caps)
208+
209+
raise TypeError(f"Unknown expression node type: {node_type.__name__}")
210+
211+
212+
def _build_children(node, n_vars, cache, param_caps, param_objs):
213+
"""Build C capsules for all children of a node, returned as a list."""
214+
caps = []
215+
# Unary: .child
216+
if hasattr(node, "child"):
217+
caps.append(_build_capsule(node.child, n_vars, cache, param_caps, param_objs))
218+
# Binary: .left, .right
219+
if hasattr(node, "left"):
220+
caps.append(_build_capsule(node.left, n_vars, cache, param_caps, param_objs))
221+
if hasattr(node, "right"):
222+
caps.append(_build_capsule(node.right, n_vars, cache, param_caps, param_objs))
223+
# QuadOverLin/RelEntr: .x, .y or .x, .z
224+
if hasattr(node, "x") and not caps:
225+
caps.append(_build_capsule(node.x, n_vars, cache, param_caps, param_objs))
226+
if hasattr(node, "z"):
227+
caps.append(_build_capsule(node.z, n_vars, cache, param_caps, param_objs))
228+
elif hasattr(node, "y"):
229+
caps.append(_build_capsule(node.y, n_vars, cache, param_caps, param_objs))
230+
# HStack: .children
231+
if hasattr(node, "children"):
232+
for c in node.children:
233+
caps.append(_build_capsule(c, n_vars, cache, param_caps, param_objs))
234+
return caps
235+
236+
237+
# ---------------------------------------------------------------------------
238+
# Left/right matmul converters
239+
# ---------------------------------------------------------------------------
240+
241+
def _convert_left_matmul(node, n_vars, cache, param_caps, param_objs):
242+
"""Convert A @ f(x)."""
243+
child_cap = _build_capsule(node.child, n_vars, cache, param_caps, param_objs)
244+
matrix = node.matrix_expr
245+
m, n = matrix.shape
246+
247+
if isinstance(matrix, SparseConstant):
248+
return make_sparse_left_matmul(None, child_cap, matrix)
249+
250+
if isinstance(matrix, Parameter):
251+
param_cap = _build_capsule(matrix, n_vars, cache, param_caps, param_objs)
252+
return make_dense_left_matmul(
253+
param_cap, child_cap, _to_dense_row_major(matrix), m, n
254+
)
255+
256+
if isinstance(matrix, Constant):
257+
return make_dense_left_matmul(
258+
None, child_cap, _to_dense_row_major(matrix), m, n
259+
)
260+
261+
raise TypeError(f"LeftMatMul matrix must be Constant, SparseConstant, or Parameter")
262+
263+
264+
def _convert_right_matmul(node, n_vars, cache, param_caps, param_objs):
265+
"""Convert f(x) @ A."""
266+
child_cap = _build_capsule(node.child, n_vars, cache, param_caps, param_objs)
267+
matrix = node.matrix_expr
268+
m, n = matrix.shape
269+
270+
if isinstance(matrix, SparseConstant):
271+
return make_sparse_right_matmul(None, child_cap, matrix)
272+
273+
if isinstance(matrix, Parameter):
274+
param_cap = _build_capsule(matrix, n_vars, cache, param_caps, param_objs)
275+
return make_dense_right_matmul(
276+
param_cap, child_cap, _to_dense_row_major(matrix), m, n
277+
)
278+
279+
if isinstance(matrix, Constant):
280+
return make_dense_right_matmul(
281+
None, child_cap, _to_dense_row_major(matrix), m, n
282+
)
283+
284+
raise TypeError(f"RightMatMul matrix must be Constant, SparseConstant, or Parameter")
285+
286+
287+
# ---------------------------------------------------------------------------
288+
# CompiledExpression
289+
# ---------------------------------------------------------------------------
290+
291+
class CompiledExpression:
292+
"""A compiled expression ready for evaluation.
293+
294+
Reads variable values from the scope's flat buffer.
295+
Reads parameter values from the Parameter objects.
296+
"""
297+
298+
def __init__(self, problem_capsule, scope, param_capsules, param_objects,
299+
expr_shape, n_vars):
300+
self._problem = problem_capsule
301+
self._scope = scope
302+
self._param_capsules = param_capsules
303+
self._param_objects = param_objects
304+
self._expr_shape = expr_shape
305+
self._n_vars = n_vars
306+
307+
def _sync_params(self):
308+
"""Push current parameter values to the C problem."""
309+
if not self._param_objects:
310+
return
311+
theta_parts = [p._value_flat for p in self._param_objects]
312+
theta = np.concatenate(theta_parts)
313+
_C.problem_update_params(self._problem, theta)
314+
315+
def _set_point(self):
316+
"""Push variable values and evaluate forward pass."""
317+
self._sync_params()
318+
_C.problem_objective_forward(self._problem, self._scope._flat_values)
319+
_C.problem_constraint_forward(self._problem, self._scope._flat_values)
320+
321+
def forward(self):
322+
"""Evaluate the expression at the current variable values."""
323+
self._set_point()
324+
result = _C.problem_constraint_forward(self._problem, self._scope._flat_values)
325+
return result
326+
327+
def jacobian(self):
328+
"""Compute the sparse Jacobian at the current variable values.
329+
330+
Returns scipy.sparse.csr_matrix of shape (expr_size, n_vars).
331+
"""
332+
self._set_point()
333+
data, indices, indptr, (m, n) = _C.problem_jacobian(self._problem)
334+
return scipy.sparse.csr_matrix((data, indices, indptr), shape=(m, n))
335+
336+
def hessian(self, weights):
337+
"""Compute the sparse Hessian of the weighted expression.
338+
339+
The Hessian is of the scalar function w^T f(x), where w is the
340+
weights vector and f is the compiled expression.
341+
342+
Args:
343+
weights: array of length expr_size
344+
345+
Returns scipy.sparse.csr_matrix of shape (n_vars, n_vars).
346+
"""
347+
weights = np.asarray(weights, dtype=np.float64).ravel()
348+
self._set_point()
349+
_C.problem_jacobian(self._problem)
350+
data, indices, indptr, (m, n) = _C.problem_hessian(
351+
self._problem, 0.0, weights
352+
)
353+
return scipy.sparse.csr_matrix((data, indices, indptr), shape=(m, n))

0 commit comments

Comments
 (0)