Pair-wise equality over large sets of large vectors

Question

I've got an interesting performance tuning/algorithmic problem that I'm running into in an optimization context.

I've got a set of ~16-50 lists of integers (usually in Range[0, 5] but no restricted to that). The data might look like this (although obviously not random)

maxQ = 5;
ndim = 16;
nstates = 100000;
braVecs = RandomInteger[{0, maxQ}, {ndim, nstates}];
ketVecs = braVecs + RandomInteger[{0, 1}, {ndim, nstates}];

now for every element q ∈ Subsets[ndim, 4] I need to determine where every pair of braVecs and ketVecs are the same, except for the indices in q, i.e. I for every possible q I need to calculate this

qComp = Complement[Range[ndim], q];
diffs = braVecs[[qComp]] - ketVecs[[qComp]];
Pick[Range[nstates], Times @@ (1 - Unitize[diffs]), 1]

Just as an example, this is the kind of thing I expect to get out

q = RandomSample[Range[ndim], 4];
qComp = Complement[Range[ndim], q];
diffs = braVecs[[qComp]] - ketVecs[[qComp]];
{q, Pick[Range[nstates], Times @@ (1 - Unitize[diffs]), 1]}

{{2, 9, 6, 4}, {825, 1993, 5577, 5666, 9690, 9856, 11502, 13515, 15680, 18570, 
  19207, 23131, 26986, 27269, 31889, 39396, 39942, 51688, 52520, 54905, 55360, 
  60180, 61682, 66258, 66458, 68742, 71871, 78489, 80906, 90275, 91520, 93184}}

This can obviously be done just by looping, but I sure there is an algorithmically more efficient way to do this, maybe using a Trie or maybe using some Boolean algebra to reuse prior calculations? This is important because my ndim can get up to ~50 and there are a huge number of elements in Subsets[50, 4]... I just don't know what the best way to approach this kind of thing is.

Answer 1

Interesting problem. The best I have been able to come up so far is this:

diffdata = Unitize[Subtract[braVecs, ketVecs]];
diffcounts = Total[diffdata];
F = q \[Function] Random`Private`PositionsOf[
    Plus[
     diffcounts,
     SparseArray[Partition[q, 1] -> -1, ndim, 0].diffdata
     ],
    0
    ];

For ndim = 50, the timings on my machine are as follows:

q = RandomSample[Range[ndim], 4];
t1 = N@First@RepeatedTiming[
    qComp = Complement[Range[ndim], q];
    diffs = braVecs[[qComp]] - ketVecs[[qComp]];
    aa = Pick[Range[nstates], Times @@ (1 - Unitize[diffs]), 1];
    ]
t2 = N@First@RepeatedTiming[bb = F[q];]
aa == bb
t1/t2

0.103461

0.00156106

True

66.2757

Explanation

My first observation in the original code was that each component of Unitize[braVecs[[qComp]] - ketVecs[[qComp]]] was recomputed for each q. Thus I refactored so that diffdata is computed only once.

The next observation was that Times @@ (1 - Unitize[diffs]) requires (ndim -Length[q]-1) nstates multiplications. My first step was to try to replace the multiplications by additions (On recent CPUs, multiplication and addition are equally fast, but my old Haswell has one of the last ones with dedicated addition pipelines.) But then I realized that I can make it only (Length[q]-1) nstates additions by precompting diffcounts and by subtracting Total[diffdata[[q]]]. This becomes more and more advantageous in increasing ndim.

However, diffdata[[q]] intruces a copy operation that I preferred to avoid. This is why I use

SparseArray[Partition[q, 1] -> -1, ndim, 0].diffdata

instead of Total[diffdata[[q]]].

The undocumented function Random`Private`PositionsOf does the same as Pick here, but is a bit faster (because it does not have to create Range[nstates], first.

Final Remark

Afterall, your biggest problem is to store all the results. Even with only ndim = 16, my laptop had to give up (16 GB of RAM + compression + swapping).

Answer 2

I realized that building a partial prefix trie could actually have significant benefits if done right. There are twos key realizations:

1. by the time you've computed the set of indices that agree over 4 of the dimensions, the number of remaining indices to compare drops by ~1/16th
2. since Complement[Range[ndim], q] is sorted, every prefix string will start with either 1, 2, 3, or 4, except for a single term starting with 5. This allows us to compute a significantly smaller number of prefixes and still get the benefit of the pre-computation.

The downside, of course, is that the code to build the prefix trie is more involved. First we define a function to refine our set of indices based on including the next set of pair-wise equalities. We precompute diffdata just as Henrik did.

diffdata = Unitize[Subtract[braVecs, ketVecs]];
pullSubtrieInds[trie_, i_] :=
 Pick[trie["idx"], diffdata[[i, trie["idx"]]], 0]

Next we have a function to recursively build out most of the trie

constructSubtrie[trie_, i_, depth_, maxDepth_] :=
  Module[{newTrie = trie},
   newTrie[i] = <|"idx" -> pullSubtrieInds[trie, i]|>;
   If[depth < maxDepth,
    Do[
     newTrie[i] = constructSubtrie[newTrie[i], j, depth + 1, maxDepth];,
     {j, i + 1, ndim}
     ]
    ];
   newTrie
   ];

and now we add a small modification to add the remaining prefix starting with 5, which simply doesn't loop as much

addFinalSubtrieElem[trie_, i_, depth_, maxDepth_] :=
  Module[{newTrie = trie},
   newTrie[i] = <|"idx" -> pullSubtrieInds[trie, i]|>;
   If[depth < maxDepth,
    newTrie[i] = constructSubtrie[newTrie[i], i + 1, depth + 1, maxDepth];
    ];
   newTrie
   ];

and then we can build the total trie like

constructSubtrie[trie_, maxDepth_] :=
  (* 
  after doing the complement operation, 
  we get lists of indices that are at minimum ndim-4 elements long, 
  so after sorting
  they all must contain 1 through 5 at some point and the one with the 5 only \
occuring once, so we just compute that one directly... *)
  
  Module[{newTrie = trie, subtrie},
   Do[
    newTrie = constructSubtrie[newTrie, i, 1, maxDepth], 
    {i, 4}
    ];
   (* add on the single set of indices that can start with 5 *)
   
   addFinalSubtrieElem[newTrie, 5, 1, maxDepth]
   ] ;

The trie itself is a bit slow to build at ndim=25

AbsoluteTiming[
  $preindexTrieDepth = 4;
  $preindexTrie = 
   constructSubtrie[<|"idx" -> Range[nstates]|>, $preindexTrieDepth];
  ][[1]]

2.93388

but could be made faster or slower based on using a larger/smaller $preindexTrieDepth. We'll still see a payoff in the end.

Next we have a function that actually uses this trie to pull indices

getSubsetDiffs[q_, prefTrie_, prefTrieDepth_] :=
 Module[{qComp = Complement[Range[ndim], q], initInds, qRest},
  initInds = prefTrie[Sequence @@ qComp[[;; prefTrieDepth]]]["idx"];
  qRest = qComp[[1 + prefTrieDepth ;;]];
  Pick[initInds, Plus @@ diffdata[[qRest, initInds]], 0]
  ]

and now some timing comparisons

getSubsetDiffsOG[q_] :=
 Random`Private`PositionsOf[
  Plus @@ diffdata[[qComp]], 
  0
  ]
AbsoluteTiming[
  Do[
   getSubsetDiffsOG[q],
   {q, subsets}
   ]
  ][[1]]

40.5721

diffcounts = Total[diffdata];
F = q \[Function]
   Random`Private`PositionsOf[
    Plus[diffcounts, SparseArray[Partition[q, 1] -> -1, ndim, 0].diffdata], 0];
AbsoluteTiming[
  Do[
   F[q],
   {q, subsets}
   ]
  ][[1]]

19.7401

AbsoluteTiming[
  Do[
   getSubsetDiffs[q, $preindexTrie, $preindexTrieDepth],
   {q, subsets}
   ]
  ][[1]]

5.81803

which shows a factor of ~7 speedup over the totally naive method and a factor of ~3 speedup over Henrik's approach, after accounting for the time required to build the trie.

40.572074`/(5.818029`+2.93388)

6.97351

19.7401`/(5.818029`+2.93388)

3.39292

Finally to confirm that things are coming out right

Catch[
 Do[
  If[Total[F[q] - getSubsetDiffs[q, $preindexTrie, $preindexTrieDepth]] > 0, 
   Throw[False]
   ],
  {q, Subsets[Range[ndim], {4}]}
  ];
 True
 ]

True

This could also likely be made faster by using Henrik's indexing strategy. I'm currently just doing naive subsampling of diffdata and using a Pick call.

---

Full Code Block

diffdata = Unitize[Subtract[braVecs, ketVecs]];
pullSubtrieInds[trie_, i_] :=
 Pick[trie["idx"], diffdata[[i, trie["idx"]]], 0]
ClearAll[constructSubtrie]
constructSubtrie[trie_, i_, depth_, maxDepth_] :=
  Module[{newTrie = trie},
   newTrie[i] = <|"idx" -> pullSubtrieInds[trie, i]|>;
   If[depth < maxDepth,
    Do[
     newTrie[i] = constructSubtrie[newTrie[i], j, depth + 1, maxDepth];,
     {j, i + 1, ndim}
     ]
    ];
   newTrie
   ];
addFinalSubtrieElem[trie_, i_, depth_, maxDepth_] :=
  Module[{newTrie = trie},
   newTrie[i] = <|"idx" -> pullSubtrieInds[trie, i]|>;
   If[depth < maxDepth,
    newTrie[i] = constructSubtrie[newTrie[i], i + 1, depth + 1, maxDepth];
    ];
   newTrie
   ];
constructSubtrie[trie_, maxDepth_] :=
  (* 
  after doing the complement operation, 
  we get lists of indices that are at minimum ndim-4 elements long, 
  so after sorting
  they all must contain 1 through 5 at some point and the one with the 5 only \
occuring once, so we just compute that one directly... *)
  
  Module[{newTrie = trie, subtrie},
   Do[
    newTrie = constructSubtrie[newTrie, i, 1, maxDepth], 
    {i, 4}
    ];
   (* add on the single set of indices that can start with 5 *)
   
   addFinalSubtrieElem[newTrie, 5, 1, maxDepth]
   ] ;
getSubsetDiffs[q_, prefTrie_, prefTrieDepth_] :=
 (* we assume sorting *)
 
 Module[{qComp = Complement[Range[ndim], q], initInds, qRest},
  initInds = prefTrie[Sequence @@ qComp[[;; prefTrieDepth]]]["idx"];
  qRest = qComp[[1 + prefTrieDepth ;;]];
  Pick[initInds, Plus @@ diffdata[[qRest, initInds]], 0]
  ]