function [P,varargout] = laplace(xs,A,b,AEq,bEq,varargin)
% LAPLACE Smooth compact 1,2 or 3 dimensional atomic probability distribution subject to linear and moment constraints 
% 
% LAPLACE('demo') or LAPLACE() runs univariate, bivariate and trivariate examples
%
% P = LAPLACE({xs1,xs2},[],[],[],[],[mu1 mu2],[mu11,mu12,mu22]) returns an
% matrix of size length(xs1) x length(xs2). The value P(i,j) is interpreted 
% as a probability assigned to lattice point (xs1(i), xs2(j)). This 2d atomic
% probability minimizes the square of the discrete laplacian assuming zero
% boundary condition, subject to moment constraints E[x1]=mu1, E[x2]=mu2,
% E[x1.*x2] = mu11, E[x1.*x2]=mu12, E[x2.*x2]=mu22 where E denotes the expectation
% operator E(f) = sum(sum(f.*P)) and x1,x2,x3 are arrays of n-d coordinates 
% function generated by [x1,x2,x3] = ndgrid(xs{:}). 
% 
% P = LAPLACE({xs1,xs2},A,b,AEq,bEq) returns a matrix of size length(xs1) x length(xs2). 
% The value P(i,j) is interpreted as a probability assigned to lattice point (xs1(i), xs2(j)).
% The 2-d atomic probability minimizes the square of the discrete laplacian subject to inequality
% constraint A*P(:) < b and/or equality constraint AEq*P(:) = bEq. Supplied vectors
% b and bEq should be of length length(xs1) x length(xs2). A,b, AEq and bEq
% may be empty. In particular ...
%
% P = LAPLACE(xs1,[],[],[f1;f2;...],[Ef1;Ef2;...]) where f1,f2,... are row vectors
% of size 1 x length(xs1) representing functions on grid points xs1, and Ef1, Ef2
% are scalars representing their means returns a smooth compact probability
% distribution with E[f1] = Ef1, E[f2] = Ef2 and so forth. 
%
% P = LAPLACE(xs1,[],[],[],[],mu1,mu11,mu111,mu1111) returns a
% vector of size length(xs1) x 1. The element P(i) is interpreted 
% as a probability assigned to lattice point xs1(i). The 1-d atomic
% probability minimizes the square of the discrete laplacian assuming zero
% boundary condition, subject to moment constraints E[x1]=mu1,
% E[x1.*x1]=mu11, E[x1.*x1.*x1] = mu111 and E[x1.*x1.*x1.*x1]=mu1111.
% Addtional moment constraints may be supplied. 
% 
% P = LAPLACE({xs1,xs2,xs3},[],[],[],[],varargin) returns an
% matrix of size length(xs1) x length(xs2) x length(xs3). The value P(i,j,k) is interpreted 
% as a probability assigned to lattice point (xs1(i), xs2(j), xs3(j)). The 3d atomic
% probability minimizes the sums squares of projections of the discrete laplacian assuming zero
% boundary condition, subject to moment constraints contained in varargin.
% Each constraint should be a vector of length NCHOOSEK(K+2,2) and contain
% moment values (as defined above) for all polynomials in x1, x2 and x3 of
% degree k sorted by reverse lexicographic ordering of their coefficients.  
% For example, P = LAPLACE({xs1,xs2,xs3},[],[],[],[],[mu11,mu12,mu13,mu22,mu23,mu33])
% minimizes E[x1.^2] = mu11, E[x1.x2]= mu12 and so forth. The degree of the
% implied polynomials is inferred from the length of the moment vector. 
% 
% P = LAPLACE(...,A,b,AEq,bEq,...) applies additional user supplied
% inequality and equality constraints on P(:), as per the convention in
% most optimization functions including QUADPROG
%
%    
% ARGUMENTS:
%   xs      n x 1 cell array  of vectors  - Lattice points at which the probability distribution is supported (compatible with ndgrid(xs{:}))
%   A,b,AEq,bEq                           - Inequality and equality constraints (can be empty), as per QUADPROG 
%   varargin                              - Listing of desired values for moments (see explanation below)
%   P       m1 x m2 x ... x mn            - Lattice point probabilities (n-dimensional)
%   varargout                             - Evaluation of moments inferred by P in the same order as they are supplied
%
% MOMENT CONSTRAINTS: 
%   If varargin{i} has length nchoosek(k+n-1,n-1) is it intepreted as a set of moment
%   constraints on polynomials of degreee k in the n coordinate functions, listed by 
%   reverse lexicographic ordering. [A shortcoming of this convention is
%   that *all* degree k moments must be supplied]. 
%
%           if n=2 then [0.1 0.3] means E[x1]=0.1 and E[x2]=0.3
%           if n=2 then [0.5 0.2 0.1] means E[x1.^2] = 0.5, E[x1.*x2] = 0.2 and E[x2.*x2] = 0.1
%           if n=3 then [0.5 0.2 0.1] means E[x1] = 0.5, E[x2] = 0.2 and E[x3] = 0.1
%
% LIMITATIONS:
%   The discrete laplacian and implicit boundary condition may not be the most sensible
%   objective function, especially in the 3-d case as currently implemented (I'm hoping someone
%   improves this). One might also replace the atomic measure with mixture of gaussians
%   centred at the lattice points. 
%
% RELATED:
%   ksdensity.m - stats toolbox
%
% Peter Cotton

momentConstraintMultiplier = 1;

if nargin==0 || ischar(xs) && strcmpi(xs,'demo'),
    laplaceDemo;
else
    
    %% Limitations and checks
    [m,n,xs] = laplaceArgChecks(xs,A,b,AEq,bEq,varargin{:});
    
    %% Establish grid points
    X = cell(n,1);
    switch n
        case 1,
            if ~iscell(xs),
                error('laplace: Violation of programmer intent - xs should have been set to cell by laplaceArgChecks');
            elseif iscell(xs),
                X = xs;
            else
                error('Not sure how to interpret xs');
            end
        otherwise
            [X{:}] = ndgrid(xs{:});
    end
    
    %% Ensure probs add to one.
    if isempty(AEq),
        AEq = ones(1,prod(m));
        bEq = 1;
    else
        AEq = [AEq;ones(1,prod(m))];
        bEq = [bEq;1];
    end
    
    %% Interpret additional arguments as equality constraints
    nsumk_ = nan(n,1);
    for k=1:10,
        nsumk_(k) = nsumk(n,k);
    end
    nArgs = length(varargin);
    moments = cell(nArgs,1);
    for argNo=1:length(varargin),
        arg_ = varargin{argNo};
        if isvector(arg_),
            switch n
                case 1,
                    degree_ = argNo;
                    expon = degree_;
                    nExpon= 1;
                otherwise
                    degree_ = find(nsumk_==length(arg_),1);
                    if isempty(degree_),
                        error([num2str(k),'th argument appears to have the wrong number of moments']);
                    else
                        [~,expon] = nsumk(n,degree_); % Listing of all non-negs adding to degree_
                        nExpon = size(expon,1);
                        expon = expon(nExpon:-1:1,:); % Reverse lexicographic ordering is more natural for moments
                    end
            end
            moments_ = nan(prod(m),nExpon);
            for i=1:nExpon,
                % Tabulate pointwise value of polynomial in x1, x2 and x3
                xi = X{1}(:);
                moments_(:,i) = xi.^expon(i,1);
                for j=2:n,
                    xj = X{j}(:);
                    moments_(:,i) = moments_(:,i).*(xj.^expon(i,j));
                end
            end
            AEq = [AEq;momentConstraintMultiplier*moments_'];
            bEq = [bEq;momentConstraintMultiplier*arg_(:)];
            moments{argNo} = moments_;
        end
    end
    
    %% Optimize discrete laplacian
    switch n
        case 1,
            f = @(p) del(p,xs);
        otherwise
            unrol = @(P) P(:);
            f = @(p) unrol(del(reshape(p,m'),xs));
    end
    [C,shouldBeZero] = inferAffineMap(f,prod(m));
    if any(abs(shouldBeZero)),
        error('laplace: Violation of programmer intent (bug). Discrete laplacian should not have constant term');
    end
    H = full(C'*C);
    options = optimset('LargeScale','off','MaxIter',500,'TolX',1e-2,'MaxFunEvals',1000);
    lb = zeros(prod(m),1);
    ub = ones(prod(m),1);
    p = quadprog(H,zeros(prod(m),1),A,b,AEq,bEq,lb,ub,zeros(prod(m),1),options);
    if n>1,
        P = reshape(p,m');
    else
        P = p;
    end
    
    %% Return moments implied by the atomic measure
    vargout = cell(length(varargin),1);
    for argNo=1:length(varargin),
        if ~isempty(moments{argNo}),
            vargout{argNo} = (p'*moments{argNo})';
        end
    end
    varargout = cell(length(varargin),1);
    [varargout{:}] = vargout{:};
end

end

function L = del(P,xs)
% function L = del(P,xs)
%
% Augmented discrete laplacian for 1 or 2 dimensions
% Projected, augmented discrete laplacian for 3 dimensions

% Tabulate grid steps
hs = cell(1,length(xs));
for dimNo=1:length(xs),
    hs{dimNo} = xs{dimNo}(2)-xs{dimNo}(1);
end

if isvector(P),
    L = del1augmented(P,hs);
elseif length(size(P))==2,
    L = del2augmented(P,hs);
elseif length(size(P))==3,
    L = del3augmented(P,hs);
else
    error('laplace is not implemented for dimension four'); % TODO: 
end

    function d = del1augmented(P,hs)
        d = del2([0;0;P(:);0;0],hs{:});
    end

    function d = del2augmented(P,hs)
        % Enlarges the area and then computes del2
        P_ = zeros(size(P)+4);
        P_(3:end-2,3:end-2) = P;
        d = del2(P_,hs{:});
    end

    function d = del3augmented(P,hs)
        % Enlarges the area and then computes del2
        P_ = zeros(size(P)+4);
        P_(3:end-2,3:end-2,3:end-2) = P;
        d = del3(P_,hs{:});
    end
    function L = del3(X,varargin)
        % FIXME: For now this returns the laplacians of 2-d marginals
        L = zeros(0,1);
        for k=1:3,
            idx_ = setdiff(1:3,k);
            d_ = del2(squeeze(sum(X,k)),varargin{idx_});
            if k==1,
                L = d_(:);
            else
                L = [L;d_(:)];
            end
        end
    end

end

function [m,n,xs] = laplaceArgChecks(xs,A,b,AEq,bEq,varargin)
% function [n,m,xs] = laplaceArgChecks(xs,A,b,AEq,bEq,varargin)
% 
% Ensure arguments make sense, and infer dimension where necessary

%% Allow passing of single vector for lattice rather than cell array
if ~iscell(xs) && isvector(xs),
    % Try to infer dimension
    if ~isempty(A),
        n_ = log(size(A,2))/log(length(xs));
        n = round(n_);
        if abs(n-n_)>1e-6,
            error('error infering dimension of problem, A looks to be inconsistent');
        end
    elseif ~isempty(AEq),
        n_ = log(size(AEq,2))/log(length(xs));
        n = round(n_);
        if abs(n-n_)>1e-6,
            error('error infering dimension of problem, AEq looks to be inconsistent');
        end
    else
        n = length(varargin{1});
        if n>3,
            error('inferred dimension exceeds 3 - better to be explicit');
        end
    end
    oldLattice = xs;
    xs = cell(n,1);
    for k1=1:n,
        xs{k1} = oldLattice;
    end
end

%% Check dimensions
n = length(xs);
m = nan(n,1);
for k2=1:n,
    latticeValues = xs{k2};
    m(k2) = length(latticeValues);
    if ~isvector(latticeValues) || any(any(diff(latticeValues)<eps)) || any(any(abs(diff(diff(latticeValues)))>1e-6)),
        error(['laplace anticipates an evenly spaced vector of support points in increasing order: e.g. (-4:0.5:3). Problem with lattice{',num2str(k2),'}']);
    end
end
if ~isempty(A) && size(A,2)~=prod(m),
    error('laplace anticipates size(A,2)=prod(m), where m(k) is the number of lattice points in dimension k');
end
if ~isempty(AEq) && size(AEq,2)~=prod(m),
    error('laplace anticipates size(AEq,2)=prod(m), where m(k) is the number of lattice points in dimension k');
end
if ~isempty(b) && size(b,1)~=size(bEq,1),
    error('laplace anticipates size(b,1)=prod(m), where m(k) is the number of lattice points in dimension k');
end
if ~isempty(bEq) && size(bEq,1)~=size(AEq,1),
    error('laplace anticipates size(bEq,1)=prod(m), where m(k) is the number of lattice points in dimension k');
end


end

function [A,b] = inferAffineMap(f,n)
% f  function_handle
% A  m x n
%
% Deduce affine mapping given its function handle

b = feval(f,zeros(n,1));
m = length(b);
A = sparse(m,n);
for k=1:n,
    ek = zeros(n,1);
    ek(k) = 1;
    A(:,k) = feval(f,ek)-b;
end
end

function [m,x] = nsumk(n,k)
% NSUMK Number and listing of non-negative integer n-tuples summing to k
%    M = NSUMK(N,K) where N and K are positive integers returns M=nchoosek(K+N-1,N-1)
%    This is the number of ordered N-tuples of non-negative integers summing to K
%
%    [M,X] = NSUMK(N,K) produces a matrix X with
%    nchoosek(K+N-1,N-1) rows and n columns. Each row comprises
%    non-negative integers summing to k. The ordering of rows follows the
%    same convention as NCHOOSEK, which is undocumented but with some
%    reliability appears to be lexicographic. The reverse of this presumed ordering
%    is a natural way to list coefficients of polynomials in N variables of degree K.
%    As per nchoosek, this syntax is only practical for situations where N is
%    less than about 15.
% 
%  EXAMPLES:   m = nsumk(5,2)   
%              [~,x] = nsumk(5,2) returns a 15 x 5 matrix x in which rows sum to 2


if isscalar(n) && isscalar(k) && nargout<=1
    m = nchoosek(k+n-1,n-1);
elseif isscalar(n) && isscalar(k) && nargout==2,
    m = nchoosek(k+n-1,n-1);
    dividers = [zeros(m,1),nchoosek((1:(k+n-1))',n-1),ones(m,1)*(k+n)];
    x = diff(dividers,1,2)-1;
else
    error('nsumk anticipates scalar k and n');
end
end

function [y,expon] = kmoments(x,k)
%  KMOMENTS - Evaluate sample mean of all k'th degree polynomials in N-variables 
%    Y = KMOMENTS(X,K) where X is an NSAMPLES by N matrix produces a NSUMK(N,K) x 1
%    vector Y comprising sample means for all K-degree polynomials in N variables
%    Samples for the j'th variable comprise the j'th column of X, and the
%    implied listing of k-degree polynomials follows reverse lexicographic ordering
% 
%    [Y,expon] = KMOMENTS(X,K) also returns the reverse-lex ordered polynomial exponents
%    in a matrix of size NCHOOSEK(K+N-1,N-1) x N


%% Generate reverse lex ordering
n = size(x,2);
switch n
    case 1,
        expon = k; m=1;
    otherwise
        [~,expon] = nsumk(n,k);
        m = size(expon,1);
        expon = expon(m:-1:1,:);
end

y = ones(m,1);

%% Compute k-th degree polys for a sample x
for i=1:m,
    ys = ones(size(x,1),1);
    for j=1:n,
        ys = ys.*x(:,j).^expon(i,j);
    end
    y(i) = mean(ys);
end

end


function laplaceDemo
close all;

%% UNIVARIATE EXAMPLE: 
% An example motivated by finance
% We are inferring a risk-neutral density from market prices
% Prices of call options are linear constraints on the atomic density
% We throw in a few moment constraints as well 
% Don't take this too seriously in the tails

matchExactMoments = true;  % <-- Set to false to verify the dangers in fitting the wrong moments 
compareShortAndLongSample = false; % <-- Set to true to use moments and 'prices' from 'true' distribution
xs = (-5:0.05:5)';

% Generate an imagined risk-neutral density sample
shortSample_ = [-0.1+0.2*randn(10,1);0.3*randn(50,1);0.5+0.5*randn(100,1)];
if compareShortAndLongSample,
   longSample_ = [-0.1+0.2*randn(10*10000,1);0.3*randn(50*10000,1);0.5+0.5*randn(100*10000,1)];
else
   longSample_ = shortSample_;
end
pShort = hist(shortSample_,xs);pShort = pShort/sum(pShort);
pLong = hist(longSample_,xs);pLong = pLong/sum(pLong);
   
% Generate imagined market prices for securities based on the risk-neutral density sample
strikes = [-1:0.5:1];
security = @(x) max(0,x*ones(1,length(strikes))-ones(size(x,1),1)*strikes); % (Call option) Takes  column vector x of length nSamples ->  nSamples x nStrikes matrix
securityMarketPrice = mean(security(longSample_))';
securityPayoff = security(xs)';

% Infer the density subject to price and moment constraints
allMoments = cell(5,1);
for mNo=1:length(allMoments),
    if matchExactMoments,
        allMoments{mNo} = kmoments(longSample_,mNo);
    else
        allMoments{mNo} = kmoments(shortSample_,mNo);
    end
end
for nMomentsUsed = 0:5,
    P = laplace(xs,[],[],securityPayoff,securityMarketPrice,allMoments{1:nMomentsUsed});
    subplot(2,3,nMomentsUsed+1);plot(xs,P,xs,pLong,xs,pShort);legend('laplace','long sample','short sample');ylabel(['Matched ',num2str(nMomentsUsed), ' moments']);
    grid; plotxs = axis;axis([-2,2,plotxs(3),plotxs(4)]);
    drawnow; pause(1);
end


%% BIVARIATE EXAMPLE
% Inferring a bivariate distribution from moments
figure;
xs = {(-4:0.5:4)',(-4:0.5:4)};
sample_ = mvnrnd([-0.1 0.3],[1.2,0.5;0.5,0.98],500);
sample_(:,1) = sample_(:,1);
sample_(:,2) = -0.5+sample_(:,2)+0.5./(1+abs(sample_(:,1)));
pEmpirical = hist3(sample_,xs);
pEmpirical = pEmpirical/sum(pEmpirical(:));
moments1 = kmoments(sample_,1);
moments2 = kmoments(sample_,2);
moments3 = kmoments(sample_,3);
moments4 = kmoments(sample_,4);
[P,m1,m2,m3,m4] = laplace(xs,[],[],[],[],moments1,moments2,moments3,moments4);
subplot(1,2,1);surf(pEmpirical);zlabel('Empirical');
subplot(1,2,2);surf(P);zlabel('Laplace minimizing');
bivariateMomentComparison = [[moments1;moments2;moments3;moments4],[m1;m2;m3;m4]],
drawnow; pause(1);


%% TRIVARIATE EXAMPLE
% Inferring a tri-variate distribution from moments
figure;
disp('Trivariate example takes a little while...');
nTrivariateMomentsToUse = 3; 
xs = {(-3:1:3)',(-3:1:3),(-3:1:3)};
sample_ = mvnrnd([-0.1 0.3 0.1],[1.2,0.5,0.2;0.5,0.98,0.35;0.2,0.35,1.2],50);
triMoments = cell(nTrivariateMomentsToUse,1);
for mNo=1:length(triMoments),
   triMoments{mNo}= kmoments(sample_,mNo);
end
P = laplace(xs,[],[],[],[],triMoments{:});
Pxy = squeeze(sum(P,3));
pxyEmpirical = hist3(sample_(:,[1:2]),xs(1:2)); pxyEmpirical = pxyEmpirical/sum(pxyEmpirical(:));
subplot(1,2,1);surf(pxyEmpirical);zlabel('Empirical (xy projection)');
subplot(1,2,2);surf(Pxy);zlabel('Laplace minimizing (xy projection)');

end

