pfasst-tikz.tex

% Copyright © 2015 Fabian Köhler <fkoehler1024@googlemail.com>
\usetikzlibrary{arrows.meta}
\usetikzlibrary{decorations.pathreplacing}
\usetikzlibrary{math}

%%%%%%%%%%%%%%%%%
% Configuration %
%%%%%%%%%%%%%%%%%
% config variables
\tikzmath{%
  % number of processors in the diagram (not used in manual block drawing)
  \numProc = 4;
  %
  % width of sweep blocks
  \blockWidth = 3.5;
  %
  % horizontal space between blocks
  \blockHSpace = 1.6;
  %
  % vertical space between blocks
  \blockVSpace = 0.28;
  %
  % height of coarse sweep blocks
  \coarseHeight = 0.6;
  %
  % height of fine sweep blocks
  \fineHeight = 2.4;
  %
  % spacing between y-axis and pfasst drawing
  \axisHSpace = 0.2;
  %
  % spacing between x-axis and pfasst drawing
  \axisVSpace = 0.15;
  %
  % length of ticks
  \ticksLength = 0.4;
  %
  % offset for processor ticks labels
  \processorTicksOffset = 1.2;
  %
  % offset for time ticks labels
  \timeTicksOffset = 0.8;
  %
  % if drawComm == 1 the arrows for communication will be drawn
  \drawComm = 1;
  %
  % ratio at which the triangle for coarse comm arrows will be drawn
  \coarseCommRatio = 0.45;
  %
  % ratio at which the triangle for fine comm arrows will be drawn
  \fineCommRatio = 0.85;
  %
  % radius for the small circles at the end of comm arrows
  \commCircleRadius = 0.03;
  %
  % amplitude of the brace used to highlight the predictor
  \braceAmplitude = 3;
  %
  % offset between brace and label
  \braceLabelOffset = 1;
  %
  % width of blocks in the legend
  \legendBlockWidth  = 1.3;
  %
  % height of blocks in the legend
  \legendBlockHeight = 1;
  %
  % horizontal space between legend blocks and their labels (somewhat), must be adjusted by trial and error for a certain \legendLabelWidth
  \legendBlockHSpace = 1.8;
  %
  % vertical space between legend blocks
  \legendBlockVSpace = 1.8;
  %
  % horizontal space between drawing and legend
  \legendHSpace = 2.7;
  %
  % max width of a legend label
  \legendLabelWidth = 3;
  %
  % variable to highlight that a value is not used
  \unused = 60000;
}



%%%%%%%%%%%%%%%%%
% Define colors %
%%%%%%%%%%%%%%%%%
\definecolor{pfasstCoarseColor}{RGB}{110,159,189}
\definecolor{pfasstFineColor}{RGB}{175,90,80}
\definecolor{pfasstTextColor}{RGB}{70,70,70}
\definecolor{pfasstCoarseCommColor}{RGB}{80,80,80}
\definecolor{pfasstFineCommColor}{RGB}{125,150,110}
\definecolor{pfasstHighlightColor}{RGB}{102,102,102}
\definecolor{pfasstGhostColor}{RGB}{80,80,80}



%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
% draw functions for elements (blocks, lines, ...) %
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
\tikzmath{%
  % specify how coarse sweep blocks are drawn
  function drawCoarseSweep(\xa, \ya, \xb, \yb) {
      { \fill[fill=pfasstCoarseColor] (\xa em,\ya em) rectangle (\xb em,\yb em); };
  };
  %
  % specify how fine sweep blocks are drawn
  function drawFineSweep(\xa, \ya, \xb, \yb) {
      { \fill[fill=pfasstFineColor] (\xa em,\ya em) rectangle (\xb em,\yb em); };
  };
  %
  % specify how ghost blocks are drawn
  function drawGhost(\xa, \ya, \xb, \yb) {
      { \draw[thick,dotted,pfasstGhostColor] (\xa em,\ya em) rectangle (\xb em,\yb em); };
  };
  %
  % specify how axes are drawn
  function drawAxis(\xa, \ya, \xb, \yb) {
      { \draw[->,>=latex,thick,pfasstTextColor] (\xa em,\ya em) -- (\xb em,\yb em); };
  };
  %
  % draw a label tick for a processor
  function drawProcTick(\x, \proc) {
      % Enable these lines to have ticks for the processors
      % \tmp = -\axisVSpace-\ticksLength;
      % { \draw[pfasstTextColor] (\x em, -\axisVSpace em) -- (\x em,\tmp em); };
      \tmp = -\ticksLength-\processorTicksOffset;
      { \node[text=pfasstTextColor,font=\normalsize] at (\x em,\tmp em) {$P_{\pgfmathprintnumber{\proc}}$}; };
  };
  %
  % draw a labeled tick for a time step
  function drawTimeTick(\x, \i) {
      \tmp = -\axisVSpace-\ticksLength;
      { \draw[pfasstHighlightColor] (\x em,-\axisVSpace em) -- (\x em,\tmp em); };
      \tmp = -\ticksLength-\timeTicksOffset;
      { \node[text=pfasstHighlightColor] at (\x em,\tmp em) {$t_{\pgfmathprintnumber{\i}}$}; };
  };
  %
  % draw coarse communication
  function drawCoarseComm(\xa, \ya, \xb, \yb) {
      \xh = (\coarseCommRatio*\xb+(1-\coarseCommRatio)*\xa);
      \yh = (\coarseCommRatio*\yb+(1-\coarseCommRatio)*\ya);
      { \draw[thick,-{Triangle[open]},pfasstCoarseCommColor] (\xa em,\ya em) -- (\xh em,\yh em); };
      { \draw[thick, pfasstCoarseCommColor] (\xh em,\yh em) -- (\xb em,\yb em); };
      { \fill[draw=pfasstCoarseCommColor,fill=white] (\xa em,\ya em) circle[radius=\commCircleRadius]; };
      { \fill[draw=pfasstCoarseCommColor,fill=white] (\xb em,\yb em) circle[radius=\commCircleRadius]; };
  };
  %
  % draw fine communication
  function drawFineComm(\xa, \ya, \xb, \yb) {
      \xh = (\fineCommRatio*\xb+(1-\fineCommRatio)*\xa);
      \yh = (\fineCommRatio*\yb+(1-\fineCommRatio)*\ya);
      { \draw[thick, -{Triangle[open]},>=latex,pfasstFineCommColor] (\xa em,\ya em) -- (\xh em,\yh em); };
      { \draw[thick, pfasstFineCommColor] (\xh em,\yh em) -- (\xb em,\yb em); };
      { \fill[draw=pfasstFineCommColor,fill=white] (\xa em,\ya em) circle[radius=\commCircleRadius]; };
      { \fill[draw=pfasstFineCommColor,fill=white] (\xb em,\yb em) circle[radius=\commCircleRadius]; };
  };
  %
  % label the y axis
  function labelYAxis(\pfasstSteps) {
      \y = getPfasstHeight(-\timeTicksOffset, \pfasstSteps)/2;
      \tmp = getTimeTickX(0)-\axisHSpace-\timeTicksOffset;
      { \node[text=pfasstTextColor, rotate=90, font=\normalsize] at (\tmp em, \y em) { computation time }; };
  };
  %
  %
  function drawHightlightLine(\xa, \ya, \xb, \yb) {
      { \draw[thick, pfasstHighlightColor, dashed] (\xa em,\ya em) -- (\xb em,\yb em); };
  };
  %
  %
  function drawHighlightBrace(\xa, \ya, \xb, \yb) {
      { \draw[decorate,decoration={brace,mirror,amplitude=\braceAmplitude},thick,pfasstHighlightColor] (\xa em,\ya em) -- (\xb em,\yb em); };
      \y = (\ya+\yb)/2;
      \tmp = \xa+\braceLabelOffset;
      { \node[text=pfasstHighlightColor, font=\normalsize, rotate=90] at (\tmp em,\y em){\small predictor}; };
  };
}



%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
% Functions to calculate properties %
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
\tikzmath{
  % get the x coordinate where a process starts
  function blockX1(\proc) {
      if \proc == 0 then {
        return \proc*\blockWidth;
      } else {
        return \proc*\blockWidth+\proc*\blockHSpace;
      };
  };
  %
  % get the x coordinate where a process ends
  function blockX2(\proc) {
      if \proc == 0 then {
          return (\proc+1)*\blockWidth;
      } else {
          return (\proc+1)*\blockWidth+\proc*\blockHSpace;
      };
  };
  %
  % get the y coordinate where to start a block
  % the number of coarse and fine blocks in the stack has to be specified
  function blockY(\nCoarse, \nFine) {
      return \nCoarse*\coarseHeight+\nFine*\fineHeight+(\nCoarse+\nFine)*\blockVSpace;
  };
  %
  % number of coarse sweeps in prediction on a processor
  function numPredictBlocks(\proc) {
      return \proc+1;
  };
  %
  % number of coarse sweeps in prediction up to a certain step on a processor
  function numPredictBlocksStep(\proc, \step) {
      if \step >= \proc then {
        return \proc+1;
      } else {
        return \step+1;
      };
  };
  %
  % calculate the width of the pfasst drawing according to the number of processors
  function getPfasstWidth(\unused) {
    return (\numProc-1)*(\blockWidth+\blockHSpace)+\blockWidth;
  };
  %
  % calculate the height of the pfasst drawing according to the number of pfasst steps
  function getPfasstHeight(\pfasstSteps) {
      \predictHeight = (\numProc-1)*(\coarseHeight+\blockVSpace)+\coarseHeight;
      if \pfasstSteps == 0 then {
          return \predictHeight;
      };
      \h = (\pfasstSteps-1)*(\coarseHeight+\fineHeight+2*\blockVSpace)+\coarseHeight+\fineHeight+\blockVSpace;
      return \h+\predictHeight;
  };
  %
  % get the x position of the tick for a certain processor
  function getProcTickX(\proc) {
      if \proc == 0 then {
          return \blockWidth/2;
      } else {
          return \proc*(\blockWidth+\blockHSpace)+\blockWidth/2;
      };
  };
  %
  % get the x position of the tick for a certain time step
  function getTimeTickX(\i) {
      if \i == 0 then {
          return -\blockHSpace/2;
      } else {
          return (\i-1)*(\blockWidth+\blockHSpace)+\blockWidth+\blockHSpace/2;
      };
  };
}


%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
% Functions to draw individual sweep blocks %
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
\tikzmath{
  % draw a coarse sweep block on stack of a certain processor that contains
  % a given number of coarse and fine blocks
  function coarseSweep(\proc, \nCoarse, \nFine) {
      \xa = blockX1(\proc);
      \xb = blockX2(\proc);
      \ya = blockY(\nCoarse, \nFine);
      \yb = \ya+\coarseHeight;
      drawCoarseSweep(\xa, \ya, \xb, \yb);
      if \drawComm == 1 then {
          if \proc > 0 then {
              if \nCoarse > 0 then {
                  \xa = blockX2(\proc-1);
                  \ya = blockY(\nCoarse-1, \nFine)+\coarseHeight;
                  \xb = blockX1(\proc);
                  \yb = blockY(\nCoarse, \nFine);
                  drawCoarseComm(\xa,\ya,\xb,\yb);
              };
          };
      };
  };
  %
  % draw a fine sweep block on stack of a certain processor that contains
  % a given number of coarse and fine blocks
  function fineSweep(\proc, \nCoarse, \nFine) {
      \xa = blockX1(\proc);
      \xb = blockX2(\proc);
      \ya = blockY(\nCoarse, \nFine);
      \yb = \ya+\fineHeight;
      drawFineSweep(\xa, \ya, \xb, \yb);
      if \drawComm == 1 then {
          if \proc > 0 then {
              if \nFine > 0 then {
                  \xa = blockX2(\proc-1);
                  \ya = blockY(\nCoarse-2, \nFine-1)+\fineHeight;
                  \xb = blockX1(\proc);
                  \yb = blockY(\nCoarse, \nFine);
                  drawFineComm(\xa, \ya, \xb, \yb);
              };
          };
      };
  };
  %
  % draw only a part of a fine sweep (defined by height)
  function partialFineSweep(\proc, \nCoarse, \nFine, \height) {
      \xa = blockX1(\proc);
      \xb = blockX2(\proc);
      \ya = blockY(\nCoarse, \nFine);
      \yb = \ya+\height;
      drawFineSweep(\xa, \ya, \xb, \yb);
      if \drawComm == 1 then {
          if \proc > 0 then {
              if \nFine > 0 then {
                  \xa = blockX2(\proc-1);
                  \ya = blockY(\nCoarse-2, \nFine-1)+\fineHeight;
                  \xb = blockX1(\proc);
                  \yb = blockY(\nCoarse, \nFine);
                  drawFineComm(\xa, \ya, \xb, \yb);
              };
          };
      };
  };
  %
  % draw a ghost of a coarse sweep on the stack of a certain processor that
  % contains a given number of coarse and fine blocks
  function ghostCoarse(\proc, \nCoarse, \nFine) {
    \xa = blockX1(\proc);
    \xb = blockX2(\proc);
    \ya = blockY(\nCoarse, \nFine);
    \yb = \ya+\coarseHeight;
    drawGhost(\xa, \ya, \xb, \yb);
  };
  %
  % draw a ghost of a fine sweep on the stack of a certain processor that
  % contains a given number of coarse and fine blocks
  function ghostFine(\proc, \nCoarse, \nFine) {
      \xa = blockX1(\proc);
      \xb = blockX2(\proc);
      \ya = blockY(\nCoarse, \nFine);
      \yb = \ya+\fineHeight;
      drawGhost(\xa, \ya, \xb, \yb);
  };
}



%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
% Functions to draw bigger parts of pfasst at once %
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
\tikzmath{
  % draw a single step of the prediction on all processors
  function predictStep(\step) {
      for \proc in {0, ..., \numProc-1}{
          \n = numPredictBlocksStep(\proc, \step);
          if \step+1 <= \n then {
              coarseSweep(\proc, \n-1, 0);
          };
      };
  };
  %
  % draw fine sweeps of pfasst step
  function pfasstStepFine(\step) {
      for \proc in {0, ..., \numProc-1}{
          \nCoarse = \step+numPredictBlocks(\proc);
          \nFine   = \step;
          fineSweep(\proc, \nCoarse, \nFine);
      };
  };
  %
  % draw coarse sweeps of pfasst step
  function pfasstStepCoarse(\step) {
      for \proc in {0, ..., \numProc-1}{
          \nCoarse = \step+numPredictBlocks(\proc);
          \nFine   = \step+1;
          coarseSweep(\proc, \nCoarse, \nFine);
      };
  };
  %
  % draw the full prediction on all processors
  % the parameter is not used, but it seems that you cannot use functions
  % without parameters
  function predict(\unused) {
      for \step in {0, ..., \numProc-1}{
          predictStep(\step);
      };
  };
  %
  % draw a full pfasst step
  function pfasstStep(\step) {
      pfasstStepFine(\step);
      pfasstStepCoarse(\step);
  };
  %
  % draw the full algorithm with a specified number of pfasst steps
  function pfasst(\steps) {
      predict(\unused);
      for \step in {0, ..., \steps-1}{
          pfasstStep(\step);
      };
  };
}



%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
% Functions to highlight parts of the algorithm %
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
\tikzmath{
  % highlight the predictor
  function highlightPredictor(\unused) {
      for \proc in {0, ..., \numProc-1}{
          % horizontal line
          \xa = blockX1(\proc);
          \xb = blockX2(\proc);
          \ya = blockY(numPredictBlocks(\proc-1), 0)+\coarseHeight+\blockVSpace/2;
          \yb = \ya;
          drawHightlightLine(\xa, \ya, \xb, \yb);
          if \proc != \numProc-1 then {
              % diagonal line
              \xa = blockX2(\proc);
              \xb = blockX1(\proc+1);
              \yb = \yb+\coarseHeight+\blockVSpace;
              drawHightlightLine(\xa, \ya, \xb, \yb);
          };
      };
      \xa = blockX2(\numProc-1)+\blockHSpace/4;
      \xb = \xa;
      \ya = 0;
      \yb = blockY(numPredictBlocks(\numProc-1)-1, 0)+\coarseHeight+\blockVSpace/2;
      drawHighlightBrace(\xa, \ya, \xb, \yb);
  };
  %
  % draw a legend
  function legend(\pfasstSteps) {
      \baseX = getPfasstWidth(\unused)+\legendHSpace;
      % \baseY = getPfasstHeight(\pfasstSteps)-\legendBlockHeight;
      \baseY = getPfasstHeight(\pfasstSteps);
      \xa = \baseX;
      \ya = \baseY;
      \xb = \baseX+\legendBlockWidth;
      \yb = \baseY-\legendBlockHeight;
      { \fill[pfasstCoarseColor] (\xa em, \ya em) rectangle (\xb em, \yb em); };
      \xl = \xb+\legendBlockHSpace;
      \yl = (\ya+\yb)/2;
      { \node[text=pfasstTextColor, align=left,text width=\legendLabelWidth em] at (\xl em,\yl em) {\small coarse sweep}; };
      \ya = \ya-\legendBlockHeight-\legendBlockVSpace;
      \yb = \yb-\legendBlockHeight-\legendBlockVSpace;
      { \fill[pfasstFineColor] (\xa em, \ya em) rectangle (\xb em, \yb em); };
      \xl = \xb+\legendBlockHSpace;
      \yl = (\ya+\yb)/2;
      { \node[text=pfasstTextColor, align=left,text width=\legendLabelWidth em] at (\xl em,\yl em) {\small fine sweep}; };
      \ya = \ya-\legendBlockHeight-\legendBlockVSpace;
      \yb = \yb-\legendBlockHeight-\legendBlockVSpace;
      drawCoarseComm(\xa, (\ya+\yb)/2, \xb, (\ya+\yb)/2);
      \xl = \xb+\legendBlockHSpace;
      \yl = (\ya+\yb)/2;
      { \node[text=pfasstTextColor, align=left,text width=\legendLabelWidth em] at (\xl em,\yl em) {\small coarse comm.}; };
      \ya = \ya-\legendBlockHeight-\legendBlockVSpace;
      \yb = \yb-\legendBlockHeight-\legendBlockVSpace;
      drawFineComm(\xa, (\ya+\yb)/2, \xb, (\ya+\yb)/2);
      \xl = \xb+\legendBlockHSpace;
      \yl = (\ya+\yb)/2;
      { \node[text=pfasstTextColor, align=left,text width=\legendLabelWidth em] at (\xl em,\yl em) {\small fine comm.}; };
  };
}



%%%%%%%%%%%%%%%%%%%%%%%%%%
% Functions to draw axes %
%%%%%%%%%%%%%%%%%%%%%%%%%%
\tikzmath{
  % draw the x-axis
  function drawXAxis(\unused) {
      for \proc in {0, ..., \numProc-1}{
          \x = getProcTickX(\proc);
          drawProcTick(\x, \proc);
          \x = getTimeTickX(\proc);
          drawTimeTick(\x, \proc);
      };
      \x = getTimeTickX(\numProc);
      drawTimeTick(\x, \numProc);
      drawAxis(getTimeTickX(0)-\axisHSpace, -\axisVSpace, getTimeTickX(\numProc)+10*\axisHSpace, -\axisVSpace);
  };
  % draw the y axis according to the number of pfasst steps
  function drawYAxis(\pfasstSteps) {
      \h = getPfasstHeight(\pfasstSteps)+\axisVSpace;
      drawAxis(getTimeTickX(0)-\axisHSpace, -\axisVSpace, getTimeTickX(0)-\axisHSpace, \h);
      labelYAxis(\pfasstSteps);
  };
  % draw axes according to the number of pfasst steps and number of processors
  function drawAxes(\pfasstSteps) {
      drawXAxis(\unused);
      drawYAxis(\pfasstSteps);
  };
}