Restrict number of input arguments for all measures (#327)

* Structure OCE code better.

It was hard to debug.

* Restrict number of possible input variables for all measures

* Faster tests for OCE

No need to use so many samples

* Don't restrict type of input for PMI

* Revert shortening of type aliases

Author: kahaaga

src/causal_graphs/oce/OCE.jl

@@ -153,60 +153,89 @@ function prepare_embeddings(alg::OCE, x, i)
     return τs, js, 𝒫s
 end
 
-
-function select_parent!(alg::OCE, parents, τs, js, 𝒫s, xᵢ, i::Int; verbose = true)
-    # Have any parents been identified yet?
+function pairwise_test(parents::OCESelectedParents)
+    # Have any parents been identified yet? If not, then we're doing pairwise tests.
     pairwise = isempty(parents.parents)
 
+    return pairwise
+end
+
+function select_parent!(alg::OCE, parents, τs, js, 𝒫s, xᵢ, i::Int; verbose = true)
     # If there are no potential parents to pick from, return immediately.
     isempty(𝒫s) && return false
 
-    # Configure estimation and independence testing function calls, which differ in the
-    # number of arguments depending on whether we're doing the pairwise or conditional case.
-    if !pairwise
-        P = StateSpaceSet(parents.parents...)
-        f = (measure, est, xᵢ, Pⱼ) -> estimate(measure, est, xᵢ, Pⱼ, P)
-        findep = (test, xᵢ, Pix) -> independence(test, xᵢ, Pix, P)
-    else
-        f = (measure, est, xᵢ, Pⱼ) -> estimate(measure, est, xᵢ, Pⱼ)
-        findep = (test, xᵢ, Pix) -> independence(test, xᵢ, Pix)
-    end
+    # Anonymous two-argument functions for computing raw measure and performing
+    # independence tests, taking care of conditioning on parents when necessary.
+    compute_raw_measure, test_independence = rawmeasure_and_independencetest(alg, parents)
 
     # Compute the measure without significance testing first. This avoids unnecessary
     # independence testing, which takes a lot of time.
     Is = zeros(length(𝒫s))
     for (i, Pⱼ) in enumerate(𝒫s)
-        Is[i] = f(alg.utest.measure, alg.utest.est, xᵢ, Pⱼ)
+        Is[i] = compute_raw_measure(xᵢ, Pⱼ)
     end
 
-    # Sort variables according to maximal measure and select the first lagged variable that
-    # gives significant association with the target variable.
-    maximize_sortidxs = sortperm(Is, rev = true)
-    n_checked = 0
-    n_potential_vars = length(𝒫s)
+    # First sort variables according to maximal measure. Then, we select the first lagged
+    # variable that gives significant association with the target variable.
+    idxs_that_maximize_measure = sortperm(Is, rev = true)
+
+    n_checked, n_potential_vars = 0, length(𝒫s)
     while n_checked < n_potential_vars
         n_checked += 1
-        ix = maximize_sortidxs[n_checked]
+        ix = idxs_that_maximize_measure[n_checked]
         if Is[ix] > 0
-            # findep takes into account the conditioning set too if it is non-empty.
-            result = findep(alg.utest, xᵢ, 𝒫s[ix])
+            result = test_independence(xᵢ, 𝒫s[ix])
             if pvalue(result) < alg.α
-                if verbose && !pairwise
-                    println("\tx$i(0) !⫫ x$(js[ix])($(τs[ix])) | $(selected(parents))")
-                elseif verbose && pairwise
-                    println("\tx$i(0) !⫫ x$(js[ix])($(τs[ix])) | ∅")
-                end
-                push!(parents.parents, 𝒫s[ix])
-                push!(parents.parents_js, js[ix])
-                push!(parents.parents_τs, τs[ix])
-                deleteat!(𝒫s, ix)
-                deleteat!(js, ix)
-                deleteat!(τs, ix)
+                print_status(IndependenceStatus(), parents, τs, js, ix, i; verbose)
+                update_parents_and_selected!(parents, 𝒫s, τs, js, ix)
                 return true
             end
         end
     end
-    # If we reach this stage, no variables have been selected. Print an informative message.
+
+    # If we reach this stage, no variables have been selected.
+    print_status(NoVariablesSelected(), parents, τs, js, i; verbose)
+    return false
+end
+
+# For pairwise cases, we don't need to condition on any parents. For conditional
+# cases, we must condition on the parents that have already been selected (`P`).
+# The measures, estimators and independence tests are different for the pairwise
+# and conditional case.
+# This just defines the functions `compute_raw_measure` and
+# `test_independence` so that they only need two input arguments, ensuring
+# that `P` is always conditioned on when relevant. The two functions are returned.
+function rawmeasure_and_independencetest(alg, parents::OCESelectedParents)
+    if pairwise_test(parents)
+        measure, est = alg.utest.measure, alg.utest.est
+        compute_raw_measure = (xᵢ, Pⱼ) -> estimate(measure, est, xᵢ, Pⱼ)
+        test_independence = (xᵢ, Pix) -> independence(alg.utest, xᵢ, Pix)
+    else
+        measure, est = alg.ctest.measure, alg.ctest.est
+        P = StateSpaceSet(parents.parents...)
+        compute_raw_measure = (xᵢ, Pⱼ) -> estimate(measure, est, xᵢ, Pⱼ, P)
+        test_independence = (xᵢ, Pix) -> independence(alg.ctest, xᵢ, Pix, P)
+    end
+    return compute_raw_measure, test_independence
+end
+
+function update_parents_and_selected!(parents::OCESelectedParents, 𝒫s, τs, js, ix::Int)
+    push!(parents.parents, 𝒫s[ix])
+    push!(parents.parents_js, js[ix])
+    push!(parents.parents_τs, τs[ix])
+    deleteat!(𝒫s, ix)
+    deleteat!(js, ix)
+    deleteat!(τs, ix)
+end
+
+###################################################################
+# Pretty printing
+###################################################################
+struct NoVariablesSelected end
+function print_status(::NoVariablesSelected, parents::OCESelectedParents,
+        τs, js, i::Int; verbose = true)
+
+    pairwise = pairwise_test(parents)
     if verbose && !pairwise
         # No more associations were found
         s = ["x$i(1) ⫫ x$j($τ) | $(selected(parents)))" for (τ, j) in zip(τs, js)]
@@ -216,7 +245,18 @@ function select_parent!(alg::OCE, parents, τs, js, 𝒫s, xᵢ, i::Int; verbose
         s = ["x$i(0) ⫫ x$j($τ) | ∅)" for (τ, j) in zip(τs, js)]
         println("\t$(join(s, "\n\t"))")
     end
-    return false
+end
+
+struct IndependenceStatus end
+function print_status(::IndependenceStatus, parents::OCESelectedParents,
+        τs, js, ix::Int, i::Int; verbose)
+    if verbose
+        if pairwise_test(parents)
+            println("\tx$i(0) !⫫ x$(js[ix])($(τs[ix])) | ∅")
+        else
+            println("\tx$i(0) !⫫ x$(js[ix])($(τs[ix])) | $(selected(parents))")
+        end
+    end
 end
 
 """

src/core.jl

@@ -6,6 +6,10 @@ const ArrayOrStateSpaceSet{D, T, N} = Union{AbstractArray{T, N}, AbstractStateSp
 export AssociationMeasure
 export DirectedAssociationMeasure
 
+
+# Any non-bivariate association measures must implement:
+# - [`min_inputs_vars`](@ref).
+# - [`max_inputs_vars`](@ref).
 """
     AssociationMeasure
 
@@ -33,3 +37,37 @@ with logarithms to base `b`. This can be used to convert the "unit" of an entrop
 function _convert_logunit(h::Real, base_from, base_to)
     h / log(base_from, base_to)
 end
+
+# Default to bivariate measures. Other measures override it.
+"""
+    min_inputs_vars(m::AssociationMeasure) → nmin::Int
+
+Return the minimum number of variables is that the measure can be computed for.
+
+For example, [`CMIShannon`](@ref) requires 3 input variables.
+"""
+min_inputs_vars(m::AssociationMeasure) = 2
+
+# Default to bivariate measures. Other measures override it.
+
+"""
+    max_inputs_vars(m::AssociationMeasure) → nmax::Int
+
+Return the maximum number of variables is that the measure can be computed for.
+
+For example, [`MIShannon`](@ref) cannot be computed for more than 2 variables.
+"""
+max_inputs_vars(m::AssociationMeasure) = 2
+
+function verify_number_of_inputs_vars(measure::AssociationMeasure, n::Int)
+    T = typeof(measure)
+    nmin = min_inputs_vars(measure)
+    if n < nmin
+        throw(ArgumentError("$T requires at least $nmin inputs. Got $n inputs."))
+    end
+
+    nmax = max_inputs_vars(measure)
+    if n > nmax
+        throw(ArgumentError("$T accepts a maximum of $nmax inputs. Got $n inputs."))
+    end
+end

src/independence_tests/independence.jl

@@ -27,8 +27,8 @@ Returns a test `summary`, whose type depends on `test`.
 - [`LocalPermutationTest`](@ref).
 - [`JointDistanceDistributionTest`](@ref).
 """
-function independence(test, args...; kwargs...)
-    error("No concrete implementation for $(typeof(test)) test yet")
+function independence(test::IndependenceTest, x...)
+    throw(ArgumentError("No concrete implementation for $(typeof(test)) test yet"))
 end
 
 function pvalue_text_summary(test::IndependenceTestResult)

src/independence_tests/local_permutation/LocalPermutationTest.jl

@@ -164,6 +164,10 @@ end
 # KD-trees and do marginal searches for all marginals all the time.
 function independence(test::LocalPermutationTest, x, y, z)
     measure, est, nshuffles = test.measure, test.est, test.nshuffles
+
+    # Make sure that the measure is compatible with the input data.
+    verify_number_of_inputs_vars(measure, 3)
+
     X, Y, Z = StateSpaceSet(x), StateSpaceSet(y), StateSpaceSet(z)
     @assert length(X) == length(Y) == length(Z)
     N = length(X)
@@ -178,6 +182,7 @@ end
 # computing the test statistic.
 function permuted_Îs(X, Y, Z, measure, est, test)
     rng, kperm, nshuffles, replace, w = test.rng, test.kperm, test.nshuffles, test.replace, test.w
+
     N = length(X)
     test.kperm < N || throw(ArgumentError("kperm must be smaller than input data length"))
 

src/independence_tests/surrogate/SurrogateTest.jl

@@ -129,6 +129,10 @@ end
 # conditional mutual information.
 function independence(test::SurrogateTest, x, y, z)
     (; measure, est, rng, surrogate, nshuffles) = test
+
+    # Make sure that the measure is compatible with the input data.
+    verify_number_of_inputs_vars(measure, 3)
+
     X, Y, Z = StateSpaceSet(x), StateSpaceSet(y), StateSpaceSet(z)
     @assert length(X) == length(Y) == length(Z)
     N = length(x)
@@ -145,6 +149,10 @@ end
 
 function independence(test::SurrogateTest, x, y)
     (; measure, est, rng, surrogate, nshuffles) = test
+
+    # Make sure that the measure is compatible with the input data.
+    verify_number_of_inputs_vars(measure, 2)
+
     X, Y = StateSpaceSet(x), StateSpaceSet(y)
     @assert length(X) == length(Y)
     N = length(x)

src/methods/correlation/distance_correlation.jl

@@ -34,6 +34,8 @@ is computed.
 """
 struct DistanceCorrelation <: AssociationMeasure end
 
+max_inputs_vars(::DistanceCorrelation) = 3
+
 """
     distance_correlation(x, y) → dcor ∈ [0, 1]
     distance_correlation(x, y, z) → pdcor

src/methods/correlation/partial_correlation.jl

@@ -40,6 +40,9 @@ where ``\\hat{P} = \\hat{\\Sigma}^{-1}`` is the sample precision matrix.
 """
 struct PartialCorrelation <: AssociationMeasure end
 
+min_inputs_vars(::PartialCorrelation) = 3
+max_inputs_vars(::PartialCorrelation) = Inf
+
 """
     partial_correlation(x::VectorOrStateSpaceSet, y::VectorOrStateSpaceSet,
         z::VectorOrStateSpaceSet...)

src/methods/infomeasures/condmutualinfo/CMIRenyiJizba.jl

@@ -36,6 +36,9 @@ struct CMIRenyiJizba{E <: Renyi} <: ConditionalMutualInformation{E}
     end
 end
 
+min_inputs_vars(::CMIRenyiJizba) = 3
+max_inputs_vars(::CMIRenyiJizba) = 3
+
 function estimate(measure::CMIRenyiJizba, est::Contingency, x, y, z)
     c = _contingency_matrix(measure, est, x, y, z)
     pxz = probabilities(c, dims = [1, 3])

src/methods/infomeasures/condmutualinfo/CMIRenyiPoczos.jl

@@ -9,7 +9,7 @@ defined in (Póczos & Schneider, 2012)[^Póczos2012].
 ## Usage
 
 - Use with [`independence`](@ref) to perform a formal hypothesis test for pairwise dependence.
-- Use with [`condmutualinfo`](@ref) to compute the raw conditional mutual information. 
+- Use with [`condmutualinfo`](@ref) to compute the raw conditional mutual information.
 
 ## Definition
 
@@ -37,3 +37,6 @@ struct CMIRenyiPoczos{E <: Renyi} <: ConditionalMutualInformation{E}
         new{E}(e)
     end
 end
+
+min_inputs_vars(::CMIRenyiPoczos) = 3
+max_inputs_vars(::CMIRenyiPoczos) = 3

src/methods/infomeasures/condmutualinfo/CMIRenyiSarbu.jl

@@ -9,7 +9,7 @@ The Rényi conditional mutual information from Sarbu (2014)[^Sarbu2014]).
 ## Usage
 
 - Use with [`independence`](@ref) to perform a formal hypothesis test for pairwise dependence.
-- Use with [`condmutualinfo`](@ref) to compute the raw conditional mutual information. 
+- Use with [`condmutualinfo`](@ref) to compute the raw conditional mutual information.
 
 ## Discrete description
 
@@ -41,12 +41,15 @@ struct CMIRenyiSarbu{E <: Renyi} <: ConditionalMutualInformation{E}
     end
 end
 
-function estimate(measure::CMIRenyiSarbu, est::Contingency{<:ProbabilitiesEstimator}, x...)
-    return estimate(measure, contingency_matrix(est.est, x...))
+min_inputs_vars(::CMIRenyiSarbu) = 3
+max_inputs_vars(::CMIRenyiSarbu) = 3
+
+function estimate(measure::CMIRenyiSarbu, est::Contingency{<:ProbabilitiesEstimator}, x, y, z)
+    return estimate(measure, contingency_matrix(est.est, x, y, z))
 end
 
-function estimate(measure::CMIRenyiSarbu, est::Contingency{<:Nothing}, x...)
-    return estimate(measure, contingency_matrix(x...))
+function estimate(measure::CMIRenyiSarbu, est::Contingency{<:Nothing}, x, y, z)
+    return estimate(measure, contingency_matrix(x, y, z))
 end
 
 function estimate(