Skip to content

Commit 882a902

Browse files
pavelkomarovclaude
andcommitted
tvrdiff: add axis parameter for multidimensional support
Use the standard ndindex direct-indexer pattern (same as polydiff, rtsdiff, etc.) instead of Maria's two-function recursive dispatch. The 1D cvxpy solve is per-vector with no batching opportunity, so the loop is the right structure. Also adds tvrdiff to the multidimensionality test. Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
1 parent b9d73b1 commit 882a902

2 files changed

Lines changed: 56 additions & 85 deletions

File tree

pynumdiff/tests/test_diff_methods.py

Lines changed: 5 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -6,7 +6,7 @@
66
from ..finite_difference import finitediff, first_order, second_order, fourth_order
77
from ..polynomial_fit import polydiff, savgoldiff, splinediff
88
from ..basis_fit import spectraldiff, rbfdiff
9-
from ..total_variation_regularization import velocity, acceleration, jerk, iterative_velocity, smooth_acceleration
9+
from ..total_variation_regularization import velocity, acceleration, jerk, iterative_velocity, smooth_acceleration, tvrdiff
1010
from ..kalman_smooth import rtsdiff, constant_velocity, constant_acceleration, constant_jerk, robustdiff
1111
from ..linear_model import lineardiff
1212
# Function aliases for testing cases where parameters change the behavior in a big way, so error limits can be indexed in dict
@@ -331,7 +331,8 @@ def test_diff_method(diff_method_and_params, test_func_and_deriv, request): # re
331331
(spectraldiff, {'high_freq_cutoff': 0.25, 'pad_to_zero_dxdt': False}),
332332
(rbfdiff, {'sigma': 0.5, 'lmbd': 1e-6}),
333333
(splinediff, {'degree': 9, 's': 1e-6}),
334-
(robustdiff, {'order':2, 'log_q':7, 'log_r':2})
334+
(robustdiff, {'order':2, 'log_q':7, 'log_r':2}),
335+
(tvrdiff, {'order': 1, 'gamma': 1e-2})
335336
]
336337

337338
# Similar to the error_bounds table, index by method first. But then we test against only one 2D function,
@@ -348,7 +349,8 @@ def test_diff_method(diff_method_and_params, test_func_and_deriv, request): # re
348349
spectraldiff: [(2, 1), (3, 2)], # lot of Gibbs ringing in 2nd order derivatives along t1 with t_1^2 sin(3 pi t_2 / 2)
349350
rbfdiff: [(0, -1), (1, 0)],
350351
splinediff: [(0, -1), (1, 0)],
351-
robustdiff: [(-2, -3), (0, -1)]
352+
robustdiff: [(-2, -3), (0, -1)],
353+
tvrdiff: [(2, 1), (3, 2)]
352354
}
353355

354356
@mark.parametrize("multidim_method_and_params", multidim_methods_and_params)

pynumdiff/total_variation_regularization.py

Lines changed: 51 additions & 82 deletions
Original file line numberDiff line numberDiff line change
@@ -52,50 +52,7 @@ def iterative_velocity(x, dt, params=None, options=None, num_iterations=None, ga
5252

5353
return x_hat, dxdt_hat
5454

55-
#N-d case:
5655
def tvrdiff(x, dt, order, gamma, huberM=float('inf'), solver=None, axis=0):
57-
"""
58-
Generalized total variation regularized derivatives (cvxpy). Supports multidimensionality by differentiating along
59-
'axis', independently for each vector obtained by fixing all other indices.
60-
61-
:param np.array[float] x: data to differentiate
62-
:param float dt: step size
63-
:param int order: 1, 2, or 3, the derivative to regularize
64-
:param float gamma: regularization parameter
65-
:param float huberM: Huber loss parameter, in units of scaled median absolute deviation of input data.
66-
:math:`M = \\infty` reduces to :math:`\\ell_2` loss squared on first, fidelity cost term, and
67-
:math:`M = 0` reduces to :math:`\\ell_1` loss, which seeks sparse residuals.
68-
:param str solver: Solver to use. Solver options include: 'MOSEK', 'CVXOPT', 'CLARABEL', 'ECOS'.
69-
If not given, fall back to CVXPY's default.
70-
71-
:return: - **x_hat** (np.array) -- estimated (smoothed) x
72-
- **dxdt_hat** (np.array) -- estimated derivative of x
73-
"""
74-
75-
x0 = np.moveaxis(x, axis, 0)
76-
77-
# end quick if it's just 1d case
78-
if x0.ndim == 1:
79-
x_hat0, dxdt0 = tvrdiff(x0, dt, order, gamma, huberM, solver)
80-
return x_hat0, dxdt0
81-
82-
x_hat0 = np.empty_like(x0, dtype=float)
83-
dxdt0 = np.empty_like(x0, dtype=float)
84-
rest = x0.shape[1:]
85-
print(rest)
86-
87-
# had to loop in python:(
88-
for i in np.ndindex(rest):
89-
slice = (slice(None),) + i
90-
x_hat0[slice], dxdt0[slice] = tvrdiff(x0[slice], dt, order, gamma, huberM, solver)
91-
92-
x_hat = np.moveaxis(x_hat0, 0, axis)
93-
dxdt_hat = np.moveaxis(dxdt0, 0, axis)
94-
95-
return x_hat, dxdt_hat
96-
97-
# 1-d case:
98-
def tvrdiff(x, dt, order, gamma, huberM=float('inf'), solver=None):
9956
"""Generalized total variation regularized derivatives. Use convex optimization (cvxpy) to solve for a
10057
total variation regularized derivative. Other convex-solver-based methods in this module call this function.
10158
@@ -108,50 +65,62 @@ def tvrdiff(x, dt, order, gamma, huberM=float('inf'), solver=None):
10865
:math:`M = 0` reduces to :math:`\\ell_1` loss, which seeks sparse residuals.
10966
:param str solver: Solver to use. Solver options include: 'MOSEK', 'CVXOPT', 'CLARABEL', 'ECOS'.
11067
If not given, fall back to CVXPY's default.
68+
:param int axis: data dimension along which differentiation is performed
11169
11270
:return: - **x_hat** (np.array) -- estimated (smoothed) x
11371
- **dxdt_hat** (np.array) -- estimated derivative of x
11472
"""
73+
x = np.asarray(x, dtype=float)
74+
75+
# had to loop in python :(
76+
x_hat = np.empty_like(x)
77+
dxdt_hat = np.empty_like(x)
78+
for vec_idx in np.ndindex(x.shape[:axis] + x.shape[axis+1:]):
79+
s = vec_idx[:axis] + (slice(None),) + vec_idx[axis:]
80+
xv = x[s]
81+
82+
# Normalize for numerical consistency with convex solver
83+
mu = np.mean(xv)
84+
sigma = median_abs_deviation(xv, scale='normal') # robust alternative to std()
85+
if sigma == 0: sigma = 1 # safety guard
86+
y = (xv-mu)/sigma
87+
88+
# Define the variables for the highest order derivative and the integration constants
89+
deriv_values = cvxpy.Variable(len(y)) # values of the order^th derivative, in which we're penalizing variation
90+
integration_constants = cvxpy.Variable(order) # constants of integration that help get us back to x
91+
92+
# Recursively integrate the highest order derivative to get back to the position. This is a first-
93+
# order scheme, but it's very fast and tends to do not markedly worse than 2nd order. See #116
94+
# I also tried a trapezoidal integration rule here, and it works no better. See #116 too.
95+
hx = deriv_values # variables are integrated to produce the signal estimate variables, \hat{x} in the math
96+
for i in range(order):
97+
hx = cvxpy.cumsum(hx) + integration_constants[i] # cumsum is like integration assuming dt = 1
98+
99+
# Compare the recursively integrated position to the noisy position. \ell_2 doesn't get scaled by 1/2 here,
100+
# so cvxpy's doubled Huber is already the right scale, and \ell_1 should be scaled by 2\sqrt{2} to match.
101+
fidelity_cost = cvxpy.sum_squares(y - hx) if huberM == float('inf') \
102+
else np.sqrt(8)*cvxpy.norm(y - hx, 1) if huberM == 0 \
103+
else utility.huber_const(huberM)*cvxpy.sum(cvxpy.huber(y - hx, huberM)) # data is already scaled, so M rather than M*sigma
104+
# Set up and solve the optimization problem
105+
prob = cvxpy.Problem(cvxpy.Minimize(fidelity_cost + gamma*cvxpy.sum(cvxpy.tv(deriv_values)) ))
106+
prob.solve(solver=solver)
107+
108+
# Recursively integrate the final derivative values to get back to the function and derivative values
109+
v = deriv_values.value
110+
for i in range(order-1): # stop one short to get the first derivative
111+
v = np.cumsum(v) + integration_constants.value[i]
112+
dxdt_v = v/dt # v only holds the dx values; to get deriv scale by dt
113+
x_v = np.cumsum(v) + integration_constants.value[order-1] # smoothed data
114+
115+
# Due to the first-order nature of the derivative, it has a slight lag. Average together every two values
116+
# to better center the answer. But this leaves us one-short, so devise a good last value.
117+
dxdt_v = (dxdt_v[:-1] + dxdt_v[1:])/2
118+
dxdt_v = np.hstack((dxdt_v, 2*dxdt_v[-1] - dxdt_v[-2])) # last value = penultimate value [-1] + diff between [-1] and [-2]
119+
120+
x_hat[s] = x_v*sigma+mu
121+
dxdt_hat[s] = dxdt_v*sigma # derivative is linear, so scale derivative by scatter
115122

116-
# Normalize for numerical consistency with convex solver
117-
mu = np.mean(x)
118-
sigma = median_abs_deviation(x, scale='normal') # robust alternative to std()
119-
if sigma == 0: sigma = 1 # safety guard
120-
y = (x-mu)/sigma
121-
122-
# Define the variables for the highest order derivative and the integration constants
123-
deriv_values = cvxpy.Variable(len(y)) # values of the order^th derivative, in which we're penalizing variation
124-
integration_constants = cvxpy.Variable(order) # constants of integration that help get us back to x
125-
126-
# Recursively integrate the highest order derivative to get back to the position. This is a first-
127-
# order scheme, but it's very fast and tends to do not markedly worse than 2nd order. See #116
128-
# I also tried a trapezoidal integration rule here, and it works no better. See #116 too.
129-
hx = deriv_values # variables are integrated to produce the signal estimate variables, \hat{x} in the math
130-
for i in range(order):
131-
hx = cvxpy.cumsum(hx) + integration_constants[i] # cumsum is like integration assuming dt = 1
132-
133-
# Compare the recursively integrated position to the noisy position. \ell_2 doesn't get scaled by 1/2 here,
134-
# so cvxpy's doubled Huber is already the right scale, and \ell_1 should be scaled by 2\sqrt{2} to match.
135-
fidelity_cost = cvxpy.sum_squares(y - hx) if huberM == float('inf') \
136-
else np.sqrt(8)*cvxpy.norm(y - hx, 1) if huberM == 0 \
137-
else utility.huber_const(huberM)*cvxpy.sum(cvxpy.huber(y - hx, huberM)) # data is already scaled, so M rather than M*sigma
138-
# Set up and solve the optimization problem
139-
prob = cvxpy.Problem(cvxpy.Minimize(fidelity_cost + gamma*cvxpy.sum(cvxpy.tv(deriv_values)) ))
140-
prob.solve(solver=solver)
141-
142-
# Recursively integrate the final derivative values to get back to the function and derivative values
143-
v = deriv_values.value
144-
for i in range(order-1): # stop one short to get the first derivative
145-
v = np.cumsum(v) + integration_constants.value[i]
146-
dxdt_hat = v/dt # v only holds the dx values; to get deriv scale by dt
147-
x_hat = np.cumsum(v) + integration_constants.value[order-1] # smoothed data
148-
149-
# Due to the first-order nature of the derivative, it has a slight lag. Average together every two values
150-
# to better center the answer. But this leaves us one-short, so devise a good last value.
151-
dxdt_hat = (dxdt_hat[:-1] + dxdt_hat[1:])/2
152-
dxdt_hat = np.hstack((dxdt_hat, 2*dxdt_hat[-1] - dxdt_hat[-2])) # last value = penultimate value [-1] + diff between [-1] and [-2]
153-
154-
return x_hat*sigma+mu, dxdt_hat*sigma # derivative is linear, so scale derivative by scatter
123+
return x_hat, dxdt_hat
155124

156125

157126
def velocity(x, dt, params=None, options=None, gamma=None, solver=None):

0 commit comments

Comments
 (0)