Skip to content

Commit ec89260

Browse files
committed
simplify things
1 parent d292a9f commit ec89260

6 files changed

Lines changed: 56 additions & 79 deletions

File tree

sparsediffpy/_core/_fn_affine.py

Lines changed: 12 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -30,8 +30,14 @@ def sum(x, axis=None):
3030
axis=0: sum along rows (collapse d1) -> (1, d2)
3131
axis=1: sum along columns (collapse d2) -> (1, d1)
3232
"""
33-
c_axis = -1 if axis is None else axis
34-
return Sum(x, c_axis)
33+
if axis is None:
34+
return Sum(x, -1)
35+
elif axis == 0:
36+
return Sum(x, 0)
37+
elif axis == 1:
38+
return Sum(x, 1)
39+
else:
40+
raise ValueError(f"Invalid axis {axis}, must be None, 0, or 1")
3541

3642

3743
def prod(x, axis=None):
@@ -58,15 +64,12 @@ def hstack(expressions):
5864
"""
5965
exprs = [_wrap_constant(e) for e in expressions]
6066
if not exprs:
61-
raise ValueError("hstack requires at least one expression")
67+
raise ValueError("hstack: empty argument")
6268

6369
d1 = exprs[0].shape[0]
6470
for e in exprs[1:]:
6571
if e.shape[0] != d1:
66-
raise ValueError(
67-
f"hstack: all expressions must have the same number of rows, "
68-
f"got {d1} and {e.shape[0]}"
69-
)
72+
raise ValueError(f"hstack: row mismatch, {d1} vs {e.shape[0]}")
7073

7174
total_d2 = _builtin_sum(e.shape[1] for e in exprs)
7275
return HStack(exprs, (d1, total_d2))
@@ -79,15 +82,12 @@ def vstack(expressions):
7982
"""
8083
exprs = [_wrap_constant(e) for e in expressions]
8184
if not exprs:
82-
raise ValueError("vstack requires at least one expression")
85+
raise ValueError("vstack: empty argument")
8386

8487
d2 = exprs[0].shape[1]
8588
for e in exprs[1:]:
8689
if e.shape[1] != d2:
87-
raise ValueError(
88-
f"vstack: all expressions must have the same number of columns, "
89-
f"got {d2} and {e.shape[1]}"
90-
)
90+
raise ValueError(f"vstack: column mismatch, {d2} vs {e.shape[1]}")
9191

9292
transposed = [Transpose(e) for e in exprs]
9393
total_d1 = _builtin_sum(e.shape[0] for e in exprs)

sparsediffpy/_core/_fn_bivariate.py

Lines changed: 7 additions & 26 deletions
Original file line numberDiff line numberDiff line change
@@ -1,41 +1,22 @@
11
"""Bivariate named functions: sp.quad_form, sp.quad_over_lin, sp.rel_entr."""
22

3-
import numpy as np
43
import scipy.sparse
54

6-
from sparsediffpy._core._expression import _wrap_constant
75
from sparsediffpy._core._nodes_bivariate import QuadOverLin, RelEntr
86
from sparsediffpy._core._nodes_other import QuadForm
97

108

119
def quad_form(x, Q):
12-
"""Quadratic form x' Q x.
13-
14-
x must be a column vector (n, 1).
15-
Q must be a scipy.sparse matrix or np.ndarray of shape (n, n).
16-
"""
17-
x = _wrap_constant(x)
18-
if x.shape[1] != 1:
19-
raise ValueError(f"quad_form: x must be a column vector, got shape {x.shape}")
20-
21-
if not scipy.sparse.issparse(Q):
10+
"""Quadratic form xT Q x with x (n, 1) and Q (n, n)"""
11+
if not isinstance(Q, scipy.sparse.csr_matrix):
2212
Q = scipy.sparse.csr_matrix(Q)
23-
else:
24-
Q = Q.tocsr()
2513

2614
n = x.shape[0]
27-
if Q.shape != (n, n):
28-
raise ValueError(
29-
f"quad_form: Q shape {Q.shape} doesn't match x shape {x.shape}"
30-
)
31-
32-
return QuadForm(
33-
x,
34-
Q_csr_data=np.asarray(Q.data, dtype=np.float64),
35-
Q_csr_indices=np.asarray(Q.indices, dtype=np.int32),
36-
Q_csr_indptr=np.asarray(Q.indptr, dtype=np.int32),
37-
Q_shape=Q.shape,
38-
)
15+
if x.shape[1] != 1 or Q.shape != (n, n):
16+
raise ValueError(f"quad_form: need x (n, 1) and Q (n, n) "
17+
", got x {x.shape} and Q {Q.shape}")
18+
19+
return QuadForm(x, Q)
3920

4021

4122
def quad_over_lin(x, z):

sparsediffpy/_core/_nodes_bivariate.py

Lines changed: 10 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -36,24 +36,25 @@ def __init__(self, left, right, result_shape):
3636

3737

3838
class QuadOverLin(Expression):
39-
"""sum(x^2) / z where z is a scalar variable.
39+
"""sum(x^2) / z where both x and z are plain Variables.
4040
41-
z must be a plain Variable and must not appear in x.
41+
z must be scalar and must not appear in x.
4242
"""
4343
def __init__(self, x, z):
44-
if not is_scalar(z.shape):
45-
raise ValueError(f"quad_over_lin: z must be scalar, got shape {z.shape}")
4644
from sparsediffpy._core._scope import Variable
45+
if not isinstance(x, Variable):
46+
raise ValueError(
47+
"quad_over_lin: x (first argument) must be a plain Variable."
48+
)
4749
if not isinstance(z, Variable):
4850
raise ValueError(
49-
"quad_over_lin: z (second argument) must be a plain Variable. "
50-
"The C engine requires this — compositions like quad_over_lin(x, f(y)) "
51-
"are not supported."
51+
"quad_over_lin: z (second argument) must be a plain scalar Variable."
5252
)
53+
if not is_scalar(z.shape):
54+
raise ValueError(f"quad_over_lin: z must be scalar, got {z.shape}")
5355
if _expr_contains_variable(x, z):
5456
raise ValueError(
55-
"quad_over_lin: the denominator variable z must not appear in "
56-
"the numerator x. Use separate variables for numerator and denominator."
57+
"quad_over_lin: z must not appear in x."
5758
)
5859
self.x = x
5960
self.z = z

sparsediffpy/_core/_nodes_other.py

Lines changed: 3 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -4,13 +4,10 @@
44

55

66
class QuadForm(Expression):
7-
"""x' Q x where Q is a constant sparse matrix."""
8-
def __init__(self, child, Q_csr_data, Q_csr_indices, Q_csr_indptr, Q_shape):
7+
"""x' Q x where Q is a constant CSR sparse matrix."""
8+
def __init__(self, child, Q):
99
self.child = child
10-
self.Q_csr_data = Q_csr_data
11-
self.Q_csr_indices = Q_csr_indices
12-
self.Q_csr_indptr = Q_csr_indptr
13-
self.Q_shape = Q_shape
10+
self.Q = Q
1411
self.shape = (1, 1)
1512

1613

sparsediffpy/_core/_registry.py

Lines changed: 7 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -10,10 +10,8 @@
1010
import numpy as np
1111

1212
from sparsediffpy import _sparsediffengine as _C
13-
from sparsediffpy._core._constants import Constant, SparseConstant
1413
from sparsediffpy._core._nodes_affine import (
15-
Add, Broadcast, DiagVec, HStack, Index, LeftMatMul, Neg,
16-
ParamScalarMult, ParamVectorMult, Reshape, RightMatMul, Sum, Trace,
14+
Add, Broadcast, DiagVec, HStack, Index, Neg, Reshape, Sum, Trace,
1715
Transpose,
1816
)
1917
from sparsediffpy._core._nodes_bivariate import (
@@ -26,7 +24,6 @@
2624
from sparsediffpy._core._nodes_other import (
2725
Prod, ProdAxisOne, ProdAxisZero, QuadForm,
2826
)
29-
from sparsediffpy._core._scope import Parameter
3027

3128

3229
# ---------------------------------------------------------------------------
@@ -38,8 +35,7 @@ def make_sparse_left_matmul(param_node, child_cap, matrix):
3835
return _C.make_left_matmul(
3936
param_node, child_cap, "sparse",
4037
matrix._csr_data, matrix._csr_indices, matrix._csr_indptr,
41-
matrix.shape[0], matrix.shape[1],
42-
)
38+
matrix.shape[0], matrix.shape[1])
4339

4440

4541
def make_dense_left_matmul(param_node, child_cap, A_flat, m, n):
@@ -112,10 +108,13 @@ def convert_rel_entr(node, child_caps):
112108

113109

114110
def convert_quad_form(node, child_caps):
111+
Q = node.Q
115112
return _C.make_quad_form(
116113
child_caps[0],
117-
node.Q_csr_data, node.Q_csr_indices, node.Q_csr_indptr,
118-
node.Q_shape[0], node.Q_shape[1],
114+
np.asarray(Q.data, dtype=np.float64),
115+
np.asarray(Q.indices, dtype=np.int32),
116+
np.asarray(Q.indptr, dtype=np.int32),
117+
Q.shape[0], Q.shape[1],
119118
)
120119

121120

tests/test_validation.py

Lines changed: 17 additions & 18 deletions
Original file line numberDiff line numberDiff line change
@@ -82,13 +82,13 @@ def test_reshape_size_mismatch(self, scope):
8282
def test_quad_form_wrong_Q_size(self, scope):
8383
x = scope.Variable(3, 1)
8484
Q = np.eye(4)
85-
with pytest.raises(ValueError, match="doesn't match"):
85+
with pytest.raises(ValueError, match="need x"):
8686
sp.quad_form(x, Q)
8787

8888
def test_quad_form_non_column(self, scope):
8989
x = scope.Variable(1, 3)
9090
Q = np.eye(3)
91-
with pytest.raises(ValueError, match="column vector"):
91+
with pytest.raises(ValueError, match="need x"):
9292
sp.quad_form(x, Q)
9393

9494
def test_pow_non_numeric_exponent(self, scope):
@@ -97,23 +97,23 @@ def test_pow_non_numeric_exponent(self, scope):
9797
x ** "two"
9898

9999
def test_hstack_empty(self):
100-
with pytest.raises(ValueError, match="at least one"):
100+
with pytest.raises(ValueError, match="empty argument"):
101101
sp.hstack([])
102102

103103
def test_vstack_empty(self):
104-
with pytest.raises(ValueError, match="at least one"):
104+
with pytest.raises(ValueError, match="empty argument"):
105105
sp.vstack([])
106106

107107
def test_hstack_mismatched_rows(self, scope):
108108
x = scope.Variable(3, 1)
109109
y = scope.Variable(2, 1)
110-
with pytest.raises(ValueError, match="same number of rows"):
110+
with pytest.raises(ValueError, match="row mismatch"):
111111
sp.hstack([x, y])
112112

113113
def test_vstack_mismatched_cols(self, scope):
114114
X = scope.Variable(3, 2)
115115
Y = scope.Variable(3, 3)
116-
with pytest.raises(ValueError, match="same number of columns"):
116+
with pytest.raises(ValueError, match="column mismatch"):
117117
sp.vstack([X, Y])
118118

119119
def test_restricted_domain_on_index(self, scope):
@@ -130,25 +130,24 @@ def test_restricted_domain_on_index(self, scope):
130130
with pytest.raises(ValueError, match="cannot be applied directly"):
131131
sp.entr(x[1:3])
132132

133-
def test_quad_over_lin_z_must_be_variable(self, scope):
133+
def test_quad_over_lin_args_must_be_variables(self, scope):
134134
x = scope.Variable(3, 1)
135135
z = scope.Variable(1, 1)
136-
# This should work — z is a plain variable, not in numerator
136+
# This should work — both are plain variables
137137
sp.quad_over_lin(x, z)
138138

139-
# This should fail — z is an expression, not a plain variable
140-
with pytest.raises(ValueError, match="must be a plain Variable"):
139+
# x is a composition — fails
140+
with pytest.raises(ValueError, match="x.*must be a plain Variable"):
141+
sp.quad_over_lin(sp.sin(x), z)
142+
143+
# z is a composition — fails
144+
with pytest.raises(ValueError, match="z.*must be a plain"):
141145
sp.quad_over_lin(x, sp.exp(z))
142146

143-
def test_quad_over_lin_z_not_in_numerator(self, scope):
144-
x = scope.Variable(3, 1)
147+
def test_quad_over_lin_z_not_in_x(self, scope):
145148
z = scope.Variable(1, 1)
146-
# z appears in numerator via broadcast: x + z
147-
with pytest.raises(ValueError, match="denominator variable z must not appear in the numerator"):
148-
sp.quad_over_lin(x + z, z)
149-
150-
# z appears directly as scalar in numerator
151-
with pytest.raises(ValueError, match="denominator variable z must not appear in the numerator"):
149+
# z used as both args
150+
with pytest.raises(ValueError, match="z must not appear in x"):
152151
sp.quad_over_lin(z, z)
153152

154153
def test_prod_must_be_variable(self, scope):

0 commit comments

Comments
 (0)