Skip to content

Commit

Permalink
Add correlate 1D implementation (#618)
Browse files Browse the repository at this point in the history 
The implementation is based on the existing convolve 1D function. This commit also adds tests and reorders the convolve tests a bit to better match the structure of the new correlate tests.
  • Loading branch information
@AngelEzquerra
AngelEzquerra authored Jan 14, 2024
1 parent 297ce90 commit 88e80a8
Show file tree
Hide file tree
Showing 2 changed files with 139 additions and 9 deletions.
89 changes: 89 additions & 0 deletions src/arraymancer/tensor/math_functions.nim
View file
Original file line number Diff line number Diff line change
Expand Up @@ -186,3 +186,92 @@ proc convolve*[T: SomeNumber | Complex32 | Complex64](
"convolve input tensors must be 1D, but second input tensor is multi-dimensional (shape=" & $t2.shape & ")")

convolveImpl(f, g, mode=mode)

type CorrelateMode* = ConvolveMode

proc correlate*[T: SomeNumber](
t1, t2: Tensor[T],
mode = CorrelateMode.valid): Tensor[T] {.noinit.} =
## Returns the cross-correlation of two one-dimensional real tensors.
##
## The correlation is defined as the integral of the product of the two tensors
## after the second one is shifted n positions, for all values of n in which
## the tensors overlap (since the integral will be zero outside of that window).
##
## Inputs:
## - t1, t2: Input tensors of size N and M respectively.
## - mode: Correlation mode (full, same, valid):
## - `full`: It returns the correlation at each point
## of overlap, with an output shape of (N+M-1,). At the end-points of
## the correlation, the signals do not overlap completely, and boundary
## effects may be seen.
## - `same`: Returns an output of length max(M, N).
## Boundary effects are still visible.
## - `valid`: This is the default mode. Returns output of length max(M, N) - min(M, N) + 1.
## The correlation is only given for points where the signals overlap
## completely. Values outside the signal boundary have no effect.
##
## Output:
## - Correlation tensor of same type as the inputs and size according to the mode.
##
## Notes:
## - The API of this function is the same as the one of numpy.correlate.
## - Note that (as with np.correlate) the default correlation mode is `valid`,
## which is different than the default convolution mode (`full`).

# Ensure that both arrays are 1-dimensional
let f = if t1.rank > 1: t1.squeeze else: t1
let g = if t2.rank > 1: t2.squeeze else: t2
if f.rank > 1:
raise newException(ValueError,
"correlate input tensors must be 1D, but first input tensor is multi-dimensional (shape=" & $t1.shape & ")")
if g.rank > 1:
raise newException(ValueError,
"correlate input tensors must be 1D, but second input tensor is multi-dimensional (shape=" & $t2.shape & ")")
mixin `|-`
mixin `_`
convolveImpl(f, g[_|-1], mode=mode)

proc correlate*[T: Complex32 | Complex64](
t1, t2: Tensor[T],
mode = CorrelateMode.valid): Tensor[T] {.noinit.} =
## Returns the cross-correlation of two one-dimensional complex tensors.
##
## The correlation is defined as the integral of the product of the two tensors
## after the second one is conjugated and shifted n positions, for all values
## of n in which the tensors overlap (since the integral will be zero outside of
## that window).
##
## Inputs:
## - t1, t2: Input tensors of size N and M respectively.
## - mode: Correlation mode (full, same, valid):
## - `full`: It returns the correlation at each point
## of overlap, with an output shape of (N+M-1,). At the end-points of
## the correlation, the signals do not overlap completely, and boundary
## effects may be seen.
## - `same`: Returns an output of length max(M, N).
## Boundary effects are still visible.
## - `valid`: This is the default mode. Returns output of length max(M, N) - min(M, N) + 1.
## The correlation is only given for points where the signals overlap
## completely. Values outside the signal boundary have no effect.
##
## Output:
## - Correlation tensor of same type as the inputs and size according to the mode.
##
## Notes:
## - The API of this function is the same as the one of numpy.correlate.
## - Note that (as with np.correlate) the default correlation mode is `valid`,
## which is different than the default convolution mode (`full`).

# Ensure that both arrays are 1-dimensional
let f = if t1.rank > 1: t1.squeeze else: t1
let g = if t2.rank > 1: t2.squeeze else: t2
if f.rank > 1:
raise newException(ValueError,
"correlate input tensors must be 1D, but first input tensor is multi-dimensional (shape=" & $t1.shape & ")")
if g.rank > 1:
raise newException(ValueError,
"correlate input tensors must be 1D, but second input tensor is multi-dimensional (shape=" & $t2.shape & ")")
mixin `|-`
mixin `_`
convolveImpl(f, g[_|-1].conjugate, mode=mode)
59 changes: 50 additions & 9 deletions tests/tensor/test_math_functions.nim
View file
Original file line number Diff line number Diff line change
Expand Up @@ -108,25 +108,66 @@ proc main() =
block:
let a = arange(4)
let b = 2 * ones[int](7) - arange(7)
let expected = [0, 2, 5, 8, 2, -4, -10, -16, -17, -12].toTensor
let expected_full = [0, 2, 5, 8, 2, -4, -10, -16, -17, -12].toTensor
let expected_same = [2, 5, 8, 2, -4, -10, -16].toTensor
let expected_valid = [8, 2, -4, -10].toTensor

check: convolve(a, b) == expected
# Test that input order doesn't matter
check: convolve(b, a) == expected
# Test the `same` mode with different input sizes
# Test all the convolution modes
check: convolve(a, b, mode=ConvolveMode.full) == expected_full
check: convolve(a, b, mode=ConvolveMode.same) == expected_same
check: convolve(a, b, mode=ConvolveMode.valid) == expected_valid

# Test that the default convolution mode is `full`
check: convolve(a, b) == expected_full

# Test that input order doesn't matter
check: convolve(b, a) == expected_full

# Test the `same` mode with different input sizes
let a2 = arange(5)
let b2 = 2 * ones[int](8) - arange(8)
let expected_same_a2b = [ 5, 8, 10, 0, -10, -20, -25].toTensor
let expected_same_ab2 = [ 2, 5, 8, 2, -4, -10, -16, -22].toTensor
let expected_same_a2b = [5, 8, 10, 0, -10, -20, -25].toTensor
let expected_same_ab2 = [2, 5, 8, 2, -4, -10, -16, -22].toTensor

check: convolve(a2, b, mode=ConvolveMode.same) == expected_same_a2b
check: convolve(a, b2, mode=ConvolveMode.same) == expected_same_ab2

# Test the `valid` mode
check: convolve(b, a, mode=ConvolveMode.valid) == expected_valid
test "1-D correlation":
block:
let a = [2, 8, -8, -6, 4].toTensor
let b = [-7, -7, 6, 0, 6, -7, 5, -6, 2].toTensor
let expected_full = [4, 4, -54, 62, -40, 50, 26, -30, -94, -36, 122, 14, -28].toTensor
let expected_same = [-54, 62, -40, 50, 26, -30, -94, -36, 122].toTensor
let expected_valid = [-40, 50, 26, -30, -94].toTensor

# Test all the correlation modes
check: correlate(a, b, mode=ConvolveMode.full) == expected_full
check: correlate(a, b, mode=ConvolveMode.same) == expected_same
check: correlate(a, b, mode=ConvolveMode.valid) == expected_valid

# Test that the default correlation mode is `valid`
check: correlate(a, b) == expected_valid

block: # Complex Tensor correlation
let ca = complex(
[6.3, 7.1, -6.0, 1.7].toTensor,
[-8.1, -9.2, 0.0, 3.5].toTensor)
let cb = complex(
[-3.9, 1.6, -1.6].toTensor,
[1.0, 5.2, 2.1].toTensor)
let expected_full = complex(
[-27.09, -62.72, -59.55, -41.86, 44.32, -3.13].toTensor,
[-0.27, -45.91, -13.75, 50.81, 2.76, -15.35].toTensor)
let expected_same = expected_full[1..^2]
let expected_valid = expected_full[2..^3]

# Test all the correlation modes
check: expected_full.mean_absolute_error(correlate(ca, cb, mode=ConvolveMode.full)) < 1e-9
check: expected_same.mean_absolute_error(correlate(ca, cb, mode=ConvolveMode.same)) < 1e-9
check: expected_valid.mean_absolute_error(correlate(ca, cb, mode=ConvolveMode.valid)) < 1e-9

# Test that the default correlation mode is `valid`
check: expected_valid.mean_absolute_error(correlate(ca, cb)) < 1e-9

main()
GC_fullCollect()

0 comments on commit 88e80a8

Please sign in to comment.