function augGeneratorMat = formaugmentedgeneratormatrix3d( generatorMat, nPages, stateDim, augStateDim )
    % Given a generator matrix of 'regular' form (i.e. not necessarily
    % sparse or using the augmented state space), return the 'augmented'
    % version.
    %
    % Works for both generator matrices and transition matrices.
    %
    %%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
    % Note generatorMat is NOT in the log domain; neither is
    % augGeneratorMat.
    %
    % stateDim is the number of (non-augmented) states.
    %
    %%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
    % CHANGELOG
    % 01.10.2014 - copied from old version.
    %
    %%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%

    if nargin < 2,
        error('Not enough arguments!')
    end
    if nargin < 3,
        stateDim = size(generatorMat, 1);
    elseif (size(generatorMat, 1) ~= stateDim) || (size(generatorMat, 2) ~= stateDim),
        error('First two dimensions of generatorMat must match stateDim!')
    end
%     if any(any(any(generatorMat > 1))),
%         error('generatorMat contains probabilities greater than 1!')
%     end
%     if any(any(sum(generatorMat, 2) ~= 1)),
%         error('Rows of generatorMat do not sum 1!')
%     end
    
    if nargin < 4,
        augStateDim = stateDim * (stateDim + 1) / 2;
    end

    % Form the triangle numbers plus one.
    % (The (K + 1)th term is included because of the transition of m from K
    % - 1 to K: the augmented state will be the (K + 1)th triangle number.)
    triangle = cumsum(0:stateDim) + 1;
    
    augGeneratorMat = zeros(augStateDim, augStateDim, nPages);

    % Loop over state.
    for i = 1:stateDim,
        augStates = triangle(i):(triangle(i + 1) - 1);
        augGeneratorMat(augStates, augStates, :) = generatorMat(augStates - triangle(i) + 1, augStates - triangle(i) + 1, :);
        % This is the only other (non-zero probability) transition possible
        % from this source state: that which increases m by one and
        % transitions z to the new m.
        if i < stateDim,
            augGeneratorMat(augStates, triangle(i + 2) - 1, :) = sum(generatorMat(augStates - triangle(i) + 1, (i + 1):stateDim, :), 2);
        end
    end
end