Skip to content

Instantly share code, notes, and snippets.

@simonjtyler
Created June 16, 2011 15:05
Show Gist options
  • Star 1 You must be signed in to star a gist
  • Fork 0 You must be signed in to fork a gist
  • Save simonjtyler/1029438 to your computer and use it in GitHub Desktop.
Save simonjtyler/1029438 to your computer and use it in GitHub Desktop.
List comprehension for Mathematica
(* 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]
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment