lanl / scico

Scientific Computational Imaging COde
BSD 3-Clause "New" or "Revised" License
90 stars 17 forks source link

Improve `TVNorm` implementation #526

Closed bwohlberg closed 1 month ago

bwohlberg commented 1 month ago

Improve the implementation of scico.functional.TVNorm, recucing memory requirements and computation time.

A number of other minor changes are also included in this PR.

bwohlberg commented 1 month ago

Some timing comparisons. On main:

import scico.numpy as snp
from scico.functional import IsotropicTVNorm
from scico.random import randn

N = 1024
x, key = randn((N, N), seed=123)

TV = IsotropicTVNorm(circular=False, ndims=2, input_shape=x.shape)

y = TV(x)
%timeit TV(x).block_until_ready()
>> 10 ms ± 81.5 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)

y = TV.prox(x)
%timeit TV.prox(x).block_until_ready()
>> 88.1 ms ± 1.53 ms per loop (mean ± std. dev. of 7 runs, 10 loops each)

On the branch for this PR:

import scico.numpy as snp
from scico.functional import IsotropicTVNorm
from scico.random import randn

N = 1024
x, key = randn((N, N), seed=123)

TV = IsotropicTVNorm(circular=False, axes=None, input_shape=x.shape)

y = TV(x)
%timeit TV(x).block_until_ready()
>> 9.99 ms ± 55 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)

y = TV.prox(x)
%timeit TV.prox(x).block_until_ready()
>> 25.8 ms ± 273 µs per loop (mean ± std. dev. of 7 runs, 10 loops each)