bayesdiagram.sty

%% This is file `bayesdiagram.sty',
%% Copyright 2012 Yori Zwols
%
% This work may be distributed and/or modified under the
% conditions of the LaTeX Project Public License, either version 1.3
% of this license or (at your option) any later version.
% The latest version of this license is in
%   http://www.latex-project.org/lppl.txt
% and version 1.3 or later is part of all distributions of LaTeX
% version 2005/12/01 or later.
%
% This work has the LPPL maintenance status `maintained'.
% 
% The Current Maintainer of this work is Y. Zwols.
%

\NeedsTeXFormat{LaTeX2e}[1995/12/01]

\ProvidesPackage{bayesdiagram}[2012/05/08 bayesdiagram]

\RequirePackage{tikz}
\RequirePackage{pgfplots}
\RequirePackage{etoolbox}
\RequirePackage{ifthen}

\usetikzlibrary{calc}
\usetikzlibrary{matrix}

\makeatletter

% -----------------------------------------------------------------
%
% INITIALIZATION OF PGF KEYS
%
% -----------------------------------------------------------------
\pgfkeys{/tikz/bayes/label/.initial={}}
\pgfkeys{/tikz/bayes/name/.initial={}}
\pgfkeys{/tikz/bayes/pdf/function/.initial={0}}
\pgfkeys{/tikz/bayes/pdf/name/.initial={}}
\pgfkeys{/tikz/bayes/pdf/xmin/.initial={0.1}}
\pgfkeys{/tikz/bayes/pdf/xmax/.initial={1}}
\pgfkeys{/tikz/bayes/pdf/name pos/.initial={0.5}}
\pgfkeys{/tikz/bayes/pdf/parameter pos/.initial={0.66}}
\pgfkeys{/tikz/bayes/pdf/parameters/.initial={}}
\pgfkeys{/tikz/bayes/pdf/main parameter count/.initial={1}}

\tikzset{/bayes/pdf/.style={}}
\tikzstyle{bayes/node}=[]
\tikzstyle{bayes/parameter node}=[anchor=south west, font=\footnotesize, inner sep=2pt, yshift=2pt]
\tikzstyle{bayes/main parameter node}=[anchor=south west, inner sep=2pt]
\tikzstyle{bayes/distribution name node}=[font=\sffamily\tiny, yshift=1pt]

\newcommand\bayesdiagramset[1]{\pgfkeys{/tikz/bayes/.cd, #1}}

\definecolor{pdfblue}{rgb}{0.45,0.76,0.88}
\bayesdiagramset{pdf/.style={pdfblue, semithick, fill=pdfblue!10!white}}

\newlength\bd@x
\newlength\bd@y
\newlength\bd@tmp
\newcounter{bd@eqn@params}

% -----------------------------------------------------------------
%
% UTILITY COMMANDS
%
% -----------------------------------------------------------------
\newcommand\getwidthofnode[2]{%
    \pgfextractx{#1}{\pgfpointanchor{#2}{east}}%
    \pgfextractx{\pgf@xa}{\pgfpointanchor{#2}{west}}% 
    \addtolength{#1}{-\pgf@xa}%
}


% -----------------------------------------------------------------
%
% CODE TO DRAW RANDOM VARIABLE NODE CONTENT
%
% Arguments for \bd@draw@distribution:
%    1) label prefix for sublabels ("@temp" if unspecified)
%    2) text inside node (usually the name of the random variable)
%
% -----------------------------------------------------------------
\long\def\bd@draw@distribution#1#2{
  \pgfkeysgetvalue{/tikz/bayes/pdf/function}{\bd@pgffunction}%
  \pgfkeysgetvalue{/tikz/bayes/pdf/name}{\bd@name}%
  \pgfkeysgetvalue{/tikz/bayes/pdf/xmin}{\bd@xmin}%
  \pgfkeysgetvalue{/tikz/bayes/pdf/xmax}{\bd@xmax}%
  \pgfkeysgetvalue{/tikz/bayes/pdf/name pos}{\bd@namepos}%
  \pgfkeysgetvalue{/tikz/bayes/pdf/parameter pos}{\bd@parpos}%
  \pgfkeysgetvalue{/tikz/bayes/pdf/main parameter count}{\bd@mainparcount}%
     \begin{axis}[width=4cm, height=3cm,   
                 axis x line=none, axis y line=none, 
                 ymin=0, xmin=\bd@xmin, xmax=\bd@xmax, 
                 clip=false]
        %                    
        % --- draw black line under the graph of the pdf 
        \fill[/tikz/bayes/pdf, draw=none]
                (rel axis cs: 0, 0) 
             -- (rel axis cs: 1, 0)
             -- ($(rel axis cs: 1, 0)+(0, -2pt)$) 
             -- ($(rel axis cs: 0, 0)+(0, -2pt)$);  
        \draw[yshift=-2pt] (rel axis cs: 0, 0) -- (rel axis cs: 1, 0);
        %       
        % --- draw pdf itself
        \addplot[/tikz/bayes/pdf, mark=none, area style, draw=none,  
                 domain=\bd@xmin:\bd@xmax, samples=50, smooth] {\bd@pgffunction}
                 \closedcycle;
        \addplot[/tikz/bayes/pdf, mark=none, area style,  fill=none, 
                 domain=\bd@xmin:\bd@xmax, samples=50, smooth] {\bd@pgffunction};
        %                               
        % --- add distribution name
        \node[anchor=base, bayes/distribution name node] 
              at (rel axis cs: \bd@namepos, 0) {\bd@name};
        %                               
        \pgfmathdeclarefunction{x}{0}%
             {\pgfmathparse{\bd@xmin + \bd@parpos * (\bd@xmax - \bd@xmin)}}
        \pgfmathsetmacro\bd@parx{x}
        \pgfmathsetmacro\bd@pary{\bd@pgffunction}
        %
        % --- add anchors for parameters
        \node[anchor=center] 
            (temp main parameter) at (rel axis cs: \bd@namepos, 0.3) {};
        \edef\tmp{\noexpand\draw
           let \noexpand\p1 = (axis cs: \bd@parx, \bd@pary) in (\noexpand\p1) 
               node (temp aux parameter) {};}
        \tmp
     \end{axis} 
     %
     % --- draw parameters
     \bd@x=0pt
     \foreach [count=\ci] \@param in \bd@parameters {
        \ifthenelse{\ci > \bd@mainparcount}{}{%
          \edef\tmp{\noexpand\node [bayes/main parameter node, 
                    xshift=\bd@x, anchor=base] (#1 parameter \ci) 
                at (temp main parameter.center) {$\@param$};              
                \noexpand\getwidthofnode{\noexpand\bd@tmp}{#1 parameter \ci}
                }
          \tmp
          \pgfmathsetlength{\global\bd@x}{\bd@x+\bd@tmp-4pt}
        }
     }
     \bd@x=2pt
     \foreach [count=\ci] \@param in \bd@parameters {
        \ifthenelse{\ci > \bd@mainparcount}{
          \pgfmathtruncatemacro\@tmp{1+\bd@mainparcount}%
          \ifthenelse{\ci > \@tmp}{%
          \edef\tmp{\noexpand\node [bayes/parameter node, 
                    xshift=\bd@x, anchor=base west] (#1 parameter comma \ci) 
                at (temp aux parameter.center) {$,$};              
                \noexpand\getwidthofnode{\noexpand\bd@tmp}{#1 parameter comma \ci}
                }
          \tmp
          }{}
          %
          \pgfmathsetlength{\global\bd@x}{\bd@x+\bd@tmp-4pt}
          \edef\tmp{\noexpand\node [bayes/parameter node, 
                    xshift=\bd@x, anchor=base west] (#1 parameter \ci) 
                at (temp aux parameter.center) {$\@param$};              
                \noexpand\getwidthofnode{\noexpand\bd@tmp}{#1 parameter \ci}
                }
          \tmp
          \pgfmathsetlength{\global\bd@x}{\bd@x+\bd@tmp-4pt}
        }{}
     }
}

% -----------------------------------------------------------------
%
% MAIN IMPLEMENTATION OF THE \distribution COMMAND
%
% Arguments:
%    1) options
%    2) label of node (\empty if unspecified)
%    3) position (\empty if unspecified)
%    4) parameters of random variable
%
% -----------------------------------------------------------------

\long\def\distribution@impl#1#2#3#4{%
  \begingroup
  \def\bd@parameters{#4}%
   %
  \def\@tmp@nodelabel{#2}%
  \def\@tmp@nodeposition{#3}%
  \ifx\@tmp@nodeposition\@empty%
     \ifx\@tmp@nodelabel\@empty%
        \node[bayes/node, matrix, #1] { \bd@draw@distribution{@temp}{} \\ };%
     \else
        \node[bayes/node, matrix, #1] (#2) { \bd@draw@distribution{#2}{} \\ };%
     \fi
  \else
     \ifx\@tmp@nodelabel\@empty%
        \node[bayes/node, matrix, #1] at (#3) { \bd@draw@distribution{@temp}{} \\ };%
     \else
        \node[bayes/node, matrix, #1] at (#3) (#2) { \bd@draw@distribution{#2}{} \\ };%
     \fi
  \fi
  \endgroup
}


% -----------------------------------------------------------------
%
% INTERFACE FOR \distribution
%
% -----------------------------------------------------------------
\long\def\distribution@named@withcoord#1#2(#3){
  \distribution@impl{#1}{#2}{#3}%
}
\long\def\distribution@named@at#1#2at{%
  \pgfutil@ifnextchar ({%
    \distribution@named@withcoord{#1}{#2}%
  }{}%
}

\long\def\distribution@named#1(#2){%
  \pgfutil@ifnextchar a{%
     \distribution@named@at{#1}{#2}%
  }{%
     \distribution@named@withcoord{#1}{#2}()%
  }%
}

\long\def\distribution@[#1]{%
  \pgfutil@ifnextchar ({%
     \distribution@named{#1}%
  }{%
     \distribution@named{#1}()%
  }%
}

\def\distribution{
   \pgfutil@ifnextchar [{%
      \distribution@%
   }{%
      \distribution@[]%
   }%
}


% -----------------------------------------------------------------
%
% MAIN IMPLEMENTATION OF THE \bayesequation COMMAND
%
% Arguments:
%    1) options
%    2) label of node (\empty if unspecified)
%    3) position (\empty if unspecified)
%    4) parameters of random variable
%
% -----------------------------------------------------------------

\def\param@#1#2{ $ \& %
  \node (#1 parameter \arabic{bd@eqn@params}) {$#2$};  \stepcounter{bd@eqn@params} %
  \& $ }
  
\newcommand\param[2][1]{%
  \def\@space{ }%
  \xappto\bd@param@draw@arrows{\noexpand\draw[->] (H1) -- (\bd@cur@paramlabel\@space parameter \arabic{bd@eqn@params})}%
  \expandafter\param@{\bd@cur@paramlabel}{#2}%
}%;

\long\def\bayesequation@impl#1#2#3#4{
   \def\bd@cur@paramlabel{#2}%
   \def\bd@param@draw@arrows{}%
   \matrix[matrix of nodes, ampersand replacement=\&, 
        column sep=0pt, nodes={inner sep=0pt}, #1] (#2) { 
   \setcounter{bd@eqn@params}{1}%
   $#4$ \\ };
   \bd@param@draw@arrows
}

% -----------------------------------------------------------------
%
% INTERFACE FOR \bayesequation
%
% -----------------------------------------------------------------
\long\def\bayesequation@named@withcoord#1#2(#3){
  \bayesequation@impl{#1}{#2}{#3}%
}
\long\def\bayesequation@named@at#1#2at{%
  \pgfutil@ifnextchar ({%
    \bayesequation@named@withcoord{#1}{#2}%
  }{}%
}

\long\def\bayesequation@named#1(#2){%
  \pgfutil@ifnextchar a{%
     \bayesequation@named@at{#1}{#2}%
  }{%
     \bayesequation@named@withcoord{#1}{#2}()%
  }%
}

\long\def\bayesequation@[#1]{%
  \pgfutil@ifnextchar ({%
     \bayesequation@named{#1}%
  }{%
     \bayesequation@named{#1}()%
  }%
}

\def\bayesequation{
   \pgfutil@ifnextchar [{%
      \bayesequation@%
   }{%
      \bayesequation@[]%
   }%
}


% -----------------------------------------------------------------
%
% DECLARATION OF DISTRIBUTIONS
%
% Note that distributions are declared up to a constant 
% because pgfplots scales the plots automatically
%
% -----------------------------------------------------------------

\pgfmathdeclarefunction{gausspdf}{2}{%
  \pgfmathparse{exp(-((x-#1)^2)/(2*#2^2))}%
}

\pgfmathdeclarefunction{exppdf}{1}{%
  \pgfmathparse{exp(-(#1) * x)}%
}

\pgfmathdeclarefunction{gammapdf}{2}{%
  \pgfmathparse{(x^(#1)) * exp(-(#2)*x)}%
}

\pgfmathdeclarefunction{betapdf}{2}{%
  \pgfmathparse{(x^((#1)-1) * (1-x)^((#2)-1}%
}

\newcommand\bayesdeclaredistribution[2][]{%
    \tikzstyle{#2}=[/tikz/bayes/pdf/.cd, #1, /tikz/.cd]%
}

\bayesdeclaredistribution[function={gausspdf(0, 1)}, 
   name=gaussian, xmin=-3.5, xmax=3.5, 
   main parameter count=1]{gaussian}

\bayesdeclaredistribution[function={gausspdf(0, 1)}, 
   name=normal, xmin=-3.5, xmax=3.5, 
   main parameter count=1]{normal}

\bayesdeclaredistribution[function={gammapdf(2, 3)}, 
   name=gamma, xmin=0, xmax=6, 
   name pos=0.15,
   parameter pos=0.6, 
   main parameter count=1]{gamma}

\bayesdeclaredistribution[function={betapdf(2, 2)},
   name=beta, xmin=0, xmax=1,
   name pos=0.5,
   main parameter count=2]{beta}
   
\bayesdeclaredistribution[function={exppdf(5)},
   name=exponential, name pos=0.2, main parameter count=0]{exponential}