public
Created

List comprehension for Mathematica

  • Download Gist
tableGen.m
Objective-C
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60
(* http://stackoverflow.com/questions/6367932/generate-a-list-in-mathematica-with-a-conditional-tested-for-each-element/6368770#6368770 *)
 
(* The code *)
 
TableIf::usage = "TableIf[expr,{i,\!\(\*SubscriptBox[\(i\), \(max\)]\)},addif] will \
generate a list of values expr when i runs from 1 to \
\!\(\*SubscriptBox[\(i\), \(max\)]\), only including elements if \
addif[expr] returns true. Note that addif can have dependence on the \
iterator variables.
The default of addif is True&.
TableIf[expr,{\!\(\*SubscriptBox[\(i\), \(max\)]\)},addif] iterates \
\!\(\*SubscriptBox[\(i\), \(max\)]\) times with no named index.
TableIf[expr,{i,\!\(\*SubscriptBox[\(i\), \
\(min\)]\),\!\(\*SubscriptBox[\(i\), \(max\)]\)},addif] starts with \
i=\!\(\*SubscriptBox[\(i\), \(min\)]\).
TableIf[expr,{i,\!\(\*SubscriptBox[\(i\), \
\(min\)]\),\!\(\*SubscriptBox[\(i\), \(max\)]\),di},addif] uses steps di.
TableIf[expr,{i,{\!\(\*SubscriptBox[\(i\), \
\(1\)]\),\!\(\*SubscriptBox[\(i\), \(2\)]\),...}},addif] uses the \
successive values \!\(\*SubscriptBox[\(i\), \(1\)]\), \
\!\(\*SubscriptBox[\(i\), \(2\)]\), ....
TableIf[expr,{i,\!\(\*SubscriptBox[\(i\), \
\(max\)]\)},{j,\!\(\*SubscriptBox[\(j\), \(max\)]\)},addif], will \
give a multidimensional list.";
 
Unprotect[TableIf]
ClearAll[TableIf];
SetAttributes[TableIf, HoldAll];
SyntaxInformation[TableIf] = {"ArgumentsPattern" -> {_, {_, _., _., _.} .., _.},
"LocalVariables" -> {"Table", {2, -2}}};
 
(* If not given a pure function for addif, then turn it into a function with trivial dependence on expr. *)
(* This allows tests depending only on the iterators. *)
(* Iterators must be lists. Can have any number of iterators. *)
(* The test addif must match the pattern Except[_List] is so that it will never match an iterator. *)
 
TableIf[expr_, iter : {__} .., addif : Except[_List] : (True &)] :=
Module[{indices, indexedRes, sowTag, Q},
If[Head[addif] === Function || Head[addif] === CompiledFunction,
Q = addif, Q = Function[{}, addif]];
SetDelayed@@Prepend[Thread[Map[Take[#, 1] &, List @@ Hold @@@ Hold[iter]], Hold], indices];
indexedRes = Last@Reap[Do[If[Q[#], Sow[{#, indices}, sowTag]] &[expr], iter], sowTag];
indexedRes = If[# === {}, #, First@#]&@indexedRes;
Map[First, SplitBy[indexedRes, Table[With[{i = i},
Function[Slot[1][[2, i]]]], {i, Length[Hold[iter]] - 1}]], {-3}]]
 
(* Catch bad iterators, use General::itform *)
 
TableIf[expr_, iter__, addif : Except[_List] : (True &)] :=
Module[{pos = First@Flatten@Position[{iter}, Except[{Repeated[_, 4]}], {1}, 1, Heads -> False]},
Message[TableIf::itform, Extract[{iter}, pos], pos + 1];
HoldForm[TableIf[expr, iter, addif]]]
 
SetAttributes[TableIf, ReadProtected];
Protect[TableIf];
 
(* Some test code *)
i = x;
TableIf[i^j, {i, 8}, {j, 8}, Mod[i, 4] == Mod[j, 4] && 50 < # < 1000 &]
TableIf[f[i, j], {i, 5}, {j, 5}, (i == 4 || i == 1) && j < 3]

Please sign in to comment on this gist.

Something went wrong with that request. Please try again.