Mathematica Functions

Table of Contents

Setting the working directory

This sets the working directory to the one in which the notebook is saved. I don't know why Mathematica doesn't do this by default.

SetDirectory[
    ToFileName[
    ("FileName"/.NotebookInformation[EvaluationNotebook[]])[[1]]]]

Fourier Transforms

Virtually every time I use an fft, I want to know the squared amplitude as a function of frequency. This function does just that – if you supply it with a list of x-y pairs: {{x1,y1},{x2,y2},...}, it will give you the Power spectrum {{k1,P(k1)},{k2,P(k2)},...}, with the output in real units.

There's also an inverse function, which takes a power spectrum and gives you real data. Note that the argument takes complex amplitudes, rather than squared amplitudes – this is so that you can specify phases if you want to.

(* myPowerSpectrum: takes an array of {x,y} pairs and returns the
     power spectrum in real units.  This is virtually always what I
     want to do when I take a Fourier transform. *)
myPowerSpectrum[a : {{_, _} ...}] :=
 Module[{time, y, dt, ft, len, front, back, freq},
  (* Extract data and take the Fourier transform *)
  time = a[[All, 1]];
  y = a[[All, 2]];
  dt = time[[2]] - time[[1]];
  ft = Abs[Fourier[y]]^2;

  (* Identify positive and negative frequencies.
     Zero frequency mode appears at position 1 *)
  len = If[EvenQ[Length[ft]], Length[ft] + 1, Length[ft]];
  front = Take[ft, Ceiling[len/2]];
  back = Reverse[Take[ft, -Floor[len/2]]];
  ft = front + PadLeft[back, Length[front]];

  (* Make frequency units and normalize the amplitudes *)
  freq = 2 Pi*Table[i/dt/Length[a], {i, 0, Length[ft]-1}];
  Transpose[{freq, ft/Length[ft]}]]

(* myInversePowerSpectrum: produces an array with a power spectrum
     similar to its argument.  Argument isn't actually a power
     spectrum; they're amplitudes, but the frequency convention is the
     same as with the power spectrum. *)
myInversePowerSpectrum[a : {{_, _} ...}] :=
 Module[{freq, amp, dt, len, time},
  (* Extract data *)
  freq = a[[All, 1]];
  amp = a[[All, 2]];

  (* Len is a bit tricky... the mapping from Length[data]
     to Length[myPowerSpectrum[data]] isn't one-to-one,
     so you can't always get this right. *)
  len = 2*Length[amp] - 1;
  dt = 2 Pi / (len * freq[[2]]);
  time = Table[i dt, {i, 0, len - 1}];

  amp = PadRight[amp, Length[time]]*Sqrt[Length[time]];
  Transpose[{time, InverseFourier[amp]}]]

Here's a basic test of each function, which also gives a sense of how to use them:

(* test power spectrum: should see peaks of height 1 
   at frequencies 1, 2 and 4 *)
tmax = 50;
dt = 0.1;
data = Table[{t, Sin[t] + Sin[2 t] + Sin[4 t]}, {t, 0, tmax, dt}];
ListLogLogPlot[myPowerSpectrum[data],
             Joined -> True,
             PlotStyle -> Thick,
             PlotRange -> All]

(* test inverse power spectrum: plots should overlap -- this test
   won't be exact unless the length of data is even, all the phases
   are zero and you have precisely an integer # of modes in the data.
   These problems shouldn't be important for larger datasets. *)
tmax = (2 - 0.01) Pi;
dt = 0.01 Pi;
data = Table[
   {t, Cos[2 t] + Cos[4 t]},
   {t, 0, tmax, dt}];

ps = myPowerSpectrum[data];
ps = Map[{#[[1]], Sqrt[#[[2]]]} &, ps];
is = Re[myInversePowerSpectrum[ps]];
ListPlot[{is, data}, Joined -> True, PlotRange -> All]

Resampling Data for ListPlot

One of the most annoying things about Mathematica is the fact that ListPlot actually plots every element in a list, and displays it using postscript. This crashes my computer whenever I try to plot a list with more than 103 elements.

As far as I can tell, Mathematica doesn't provide a built-in function to downsample an array, so I wrote these.

resample[a:{{_,_}...},dx_Real:0.01]:=Module[
(* Resample a list for plotting.  Second argument is the
   logarithmic bin size *)
  {},
  Map[Mean,GatherBy[a,Round[#[[1]],dlogx]&]]]

logresample[a:{{_,_}...},dlogx_Real:0.01]:=Module[
(* Resample a list for plotting.  Second argument is the
   logarithmic bin size *)
  {},
  Map[Mean,GatherBy[a,Round[Log[10,#[[1]]],dlogx]&]]]

Reading VTK files

This is a simple function to import athena VTK files into mathematica. Note that this is 42 lines of code, while the equivalent c code is 562!

(* Simple function for extracting numbers from a string.  Very fragile. *)
processLine[line_String] := Map[Read[StringToStream[#], Number] &,
                                Drop[StringSplit[line, " "], 1]]

(* VTK reader: reads any one variable from a vtk file.
     example usage: readVTK["file", "velocty", "vector"]
     returns: an interpolating function of that variable.  Also leaves
       the raw data in the variable =data= *)
readVTK[file_String, label_String, type_String] :=
    Module[ {str, n},
          str = OpenRead[file, BinaryFormat -> True];

          (* Read the header *)
          dim = Map[If[# > 1, # - 1, #] &,
                    processLine[Find[str, "DIMENSION"]]];
          n = Apply[Times, dim];
          If[type=="vector", n = n*3];
          origin  = processLine[Find[str, "ORIGIN"]];
          spacing = processLine[Find[str, "SPACING"]];
          range = Reverse[Transpose[{origin, origin+spacing*dim}]];

          (* Find the data; note that find leaves the '\n' in the stream *)
          Find[str, label]; BinaryRead[str, "Character8"];
          If[type == "scalar", (* LOOKUP_TABLE *)
             Block[{}, Read[str,Record];  BinaryRead[str,"Character8"]]];

          (* Read the data and close the file; vtk files are big-endian *)
          data = BinaryReadList[str, "Real32", n, ByteOrdering -> +1];
          Close[str];

          (* store in a 3D array *)
          If[type == "vector",
             data = Partition[data, 3]];

          (* There must be a better way to do this *)
          data = Partition[data, dim[[1]]];
          data = Partition[data, dim[[2]]];

          (* Interpolate and return the data *)
          If[type=="scalar",
             ListInterpolation[data,range],
             Table[ListInterpolation[data[[All,All,All,i]],range],{i,1,3}]]]

Date: 2013-10-07T11:14-0700

Author: Michael McCourt

Org version 7.9.3f with Emacs version 24

Valid XHTML 1.0 Strict Valid CSS