Add LinearCombination transform

src/Transforms.jl

@@ -2,11 +2,12 @@ module Transforms
 
 using Tables
 
-export Transform, Power
+export LinearCombination, Transform, Power
 export transform, transform!
 
 include("utils.jl")
 include("transformers.jl")
+include("linear_combination.jl")
 include("power.jl")
 
 end

src/linear_combination.jl

@@ -0,0 +1,77 @@
+"""
+    LinearCombination(coefficients) <: Transform
+
+Calculate the linear combination using the column weights passed in.
+"""
+struct LinearCombination <: Transform
+    coefficients::Vector{Real}
+end
+
+function _check_dimensions_match(LC::LinearCombination, num_inds)
+    num_coefficients = length(LC.coefficients)
+    if num_inds != num_coefficients
+        throw(DimensionMismatch(
+            "Size $num_inds doesn't match number of coefficients ($num_coefficients)"
+        ))
+    end
+end
+
+_sum_row(row, coefficients) = sum(map(*, row, coefficients))
+
+"""
+    apply(x::AbstractVector, LC::LinearCombination; inds=Colon())
+
+Applies the [`LinearCombination`](@ref) to each of the specified indices in `x`.
+
+If no `inds` are specified, then the [`LinearCombination`](@ref) is applied to all elements.
+"""
+function apply(x::AbstractVector, LC::LinearCombination; inds=Colon())
+    # Treat each element as it's own column
+    # Error if dimensions don't match
+    num_inds = inds isa Colon ? length(x) : length(inds)
+    _check_dimensions_match(LC, num_inds)
+
+    return [_sum_row(x[inds], LC.coefficients)]
+end
+
+"""
+    apply(x::AbstractArray, LC::LinearCombination; dims=1, inds=Colon())
+
+Applies the [`LinearCombination`](@ref) to each of the specified indices in `x` along the
+dimension specified, which defaults to applying across the columns of x.
+
+If no `inds` are specified, then the [`LinearCombination`](@ref) is applied to all columns.
+"""
+function apply(x::AbstractArray, LC::LinearCombination; dims=1, inds=Colon())
+    # Get the number of vectors in the dimension not specified
+    other_dim = dims ==  1 ? 2 : 1
+    num_inds = inds isa Colon ? size(x, other_dim) : length(inds)
+    # Error if dimensions don't match
+    _check_dimensions_match(LC, num_inds)
+
+    return [_sum_row(row[inds], LC.coefficients) for row in eachslice(x; dims=dims)]
+end
+
+"""
+    apply(x::Table, LC::LinearCombination; inds=nothing
+
+Applies the [`LinearCombination`](@ref) to each of the specified indices (columns) in `x`.
+
+If no `inds` are specified, then the [`LinearCombination`](@ref) is applied to all columns.
+"""
+function apply(x, LC::LinearCombination; inds=nothing)
+    # Error if dimensions don't match
+    num_inds = inds === nothing ? length(Tables.columnnames(x)) : length(inds)
+    _check_dimensions_match(LC, num_inds)
+
+    # Keep the generic form when not specifying column names
+    # because that is much more performant than selecting each col by name
+    if inds === nothing
+        return [_sum_row(row, LC.coefficients) for row in Tables.rows(x)]
+    else
+        return [
+            _sum_row([row[cname] for cname in inds], LC.coefficients)
+            for row in Tables.rows(x)
+        ]
+    end
+end

test/linear_combination.jl

@@ -0,0 +1,202 @@
+@testset "linear combination" begin
+
+    lc = LinearCombination([1, -1])
+    @test lc isa Transform
+
+    @testset "Vector" begin
+        x = [1, 2]
+        expected = [-1]
+
+        @testset "all inds" begin
+            @test Transforms.apply(x, lc) == expected
+            @test lc(x) == expected
+        end
+
+        @testset "dims not supported" begin
+            @test_throws  MethodError Transforms.apply(x, lc; dims=1)
+        end
+
+        @testset "dimension mismatch" begin
+            x = [1, 2, 3]
+            @test_throws DimensionMismatch Transforms.apply(x, lc)
+        end
+
+        @testset "specified inds" begin
+            x = [1, 2, 3]
+            inds = [2, 3]
+            expected = [-1]
+
+            @test Transforms.apply(x, lc; inds=inds) == expected
+            @test lc(x; inds=inds) == expected
+        end
+    end
+
+    @testset "Matrix" begin
+        M = [1 1; 2 2; 3 5]
+        expected = [0, 0, -2]
+
+        @testset "all inds" begin
+            @test Transforms.apply(M, lc) == expected
+            @test lc(M) == expected
+        end
+
+        # TODO: Colon() is not supported by eachslice(A, dims=:), do we care?
+        @testset "dims" begin
+            @testset "dims = 1" begin
+                @test Transforms.apply(M, lc; dims=1) == expected
+                @test lc(M; dims=1) == expected
+            end
+
+            @testset "dims = 2" begin
+                # There are 3 rows so trying to apply along dim 2 without specifying inds
+                # won't work
+                @test_throws DimensionMismatch Transforms.apply(M, lc; dims=2)
+
+                @test Transforms.apply(M, lc; dims=2, inds=[2, 3]) == [-1, -3]
+                @test lc(M; dims=2, inds=[1, 3]) == [-2, -4]
+            end
+        end
+
+        @testset "dimension mismatch" begin
+            M = [1 1 1; 2 2 2]
+            @test_throws DimensionMismatch Transforms.apply(M, lc)
+        end
+
+        @testset "specified inds" begin
+            M = [1 1 5; 2 2 4]
+            inds = [2, 3]
+            expected = [-4, -2]
+
+            @test Transforms.apply(M, lc; inds=inds) == expected
+            @test lc(M; inds=inds) == expected
+        end
+    end
+
+    @testset "AxisArray" begin
+        A = AxisArray([1 2; 4 5], foo=["a", "b"], bar=["x", "y"])
+        expected = [-1, -1]
+
+        @testset "all inds" begin
+            @test Transforms.apply(A, lc) == expected
+            @test lc(A) == expected
+        end
+
+        @testset "dims" begin
+            @testset "dims = 1" begin
+                @test Transforms.apply(A, lc; dims=1) == expected
+                @test lc(A; dims=1) == expected
+            end
+
+            @testset "dims = 2" begin
+                @test Transforms.apply(A, lc; dims=2) == [-3, -3]
+                @test lc(A; dims=2) == [-3, -3]
+            end
+        end
+
+        @testset "dimension mismatch" begin
+            A = AxisArray([1 2 3; 4 5 5], foo=["a", "b"], bar=["x", "y", "z"])
+            @test_throws DimensionMismatch Transforms.apply(A, lc)
+        end
+
+        @testset "specified inds" begin
+            A = AxisArray([1 2 3; 4 5 5], foo=["a", "b"], bar=["x", "y", "z"])
+            inds = [1, 2]
+            expected = [-1, -1]
+
+            @test Transforms.apply(A, lc; inds=inds) == expected
+            @test lc(A; inds=inds) == expected
+        end
+    end
+
+    @testset "AxisKey" begin
+        A = KeyedArray([1 2; 4 5], foo=["a", "b"], bar=["x", "y"])
+        expected = [-1, -1]
+
+        @testset "all inds" begin
+            @test Transforms.apply(A, lc) == expected
+            @test lc(A) == expected
+        end
+
+        @testset "dims" begin
+            @testset "dims = 1" begin
+                @test Transforms.apply(A, lc; dims=1) == expected
+                @test lc(A; dims=1) == expected
+            end
+
+            @testset "dims = 2" begin
+                @test Transforms.apply(A, lc; dims=2) == [-3, -3]
+                @test lc(A; dims=2) == [-3, -3]
+            end
+        end
+
+        @testset "dimension mismatch" begin
+            A = KeyedArray([1 2 3; 4 5 6], foo=["a", "b"], bar=["x", "y", "z"])
+            @test_throws DimensionMismatch Transforms.apply(A, lc)
+        end
+
+        @testset "specified inds" begin
+            A = KeyedArray([1 2 3; 4 5 5], foo=["a", "b"], bar=["x", "y", "z"])
+            inds = [1, 2]
+            expected = [-1, -1]
+
+            @test Transforms.apply(A, lc; inds=inds) == expected
+            @test lc(A; inds=inds) == expected
+        end
+    end
+
+    @testset "NamedTuple" begin
+        nt = (a = [1, 2, 3], b = [4, 5, 6])
+        expected = [-3, -3, -3]
+
+        @testset "all cols" begin
+            @test Transforms.apply(nt, lc) == expected
+            @test lc(nt) == expected
+        end
+
+        @testset "dims not supported" begin
+            @test_throws MethodError Transforms.apply(nt, lc; dims=1)
+        end
+
+        @testset "dimension mismatch" begin
+            nt = (a = [1, 2, 3], b = [4, 5, 6], c = [1, 1, 1])
+            @test_throws DimensionMismatch Transforms.apply(nt, lc)
+        end
+
+        @testset "specified cols" begin
+            nt = (a = [1, 2, 3], b = [4, 5, 6], c = [1, 1, 1])
+            inds = [:a, :b]
+            expected = [-3, -3, -3]
+
+            @test Transforms.apply(nt, lc; inds=inds) == expected
+            @test lc(nt; inds=inds) == expected
+        end
+    end
+
+    @testset "DataFrame" begin
+        df = DataFrame(:a => [1, 2, 3], :b => [4, 5, 6])
+        expected = [-3, -3, -3]
+
+        @testset "all cols" begin
+            @test Transforms.apply(df, lc) == expected
+            @test lc(df) == expected
+        end
+
+        @testset "dims not supported" begin
+            @test_throws MethodError Transforms.apply(df, lc; dims=1)
+        end
+
+        @testset "dimension mismatch" begin
+            df = DataFrame(:a => [1, 2, 3], :b => [4, 5, 6], :c => [1, 1, 1])
+            @test_throws DimensionMismatch Transforms.apply(df, lc)
+        end
+
+        @testset "specified cols" begin
+            df = DataFrame(:a => [1, 2, 3], :b => [4, 5, 6], :c => [1, 1, 1])
+            inds = [:b, :c]
+            expected = [3, 4, 5]
+
+            @test Transforms.apply(df, lc; inds=inds) == expected
+            @test lc(df; inds=inds) == expected
+        end
+    end
+end

test/runtests.jl

@@ -6,5 +6,6 @@ using Transforms: _try_copy
 using Test
 
 @testset "Transforms.jl" begin
+    include("linear_combination.jl")
     include("power.jl")
 end