Obtaining joint distributions and conditional distributions using Mathematica

Question

I have two multivariate Gaussian distributions $p(x)$ and $p(z)$ with mean vectors $m_x$ and $m_z$, and covariance matrices $\Sigma_x$ and $\Sigma_z$. my model is a simple linear model $x = W z+n$ where $n$ is a noise vector with mean $0$ and diagonal covariance matrix of the form $\sigma^2 I$, where $I$ is the identity matrix.

I observe the variable $x$. now how can I calculate the joint distribution $p(x,z)$ and the conditional distributions $p(x|z)$ and $p(z|x)$?

Answer 1

To get the joint density of two distributions, you need to use ProductDistribution. For example, consider the two distributions $p(x)$ and $p(y)$:

px = NormalDistribution[2, 2];
py = NormalDistribution[-2, 3];

(* visualize *)
{Plot[PDF[px, x], {x, -3, 7}, PlotRange -> All],  
 Plot[PDF[py, y], {y, -10, 8}, PlotRange -> All]} // GraphicsRow

Now obtain the joint distribution $p(x,y)$ and visualize:

pxy = ProductDistribution[px, py];
Quiet@Plot3D[PDF[pxy, {x, y}], {x, -3, 7}, {y, -10, 8}, PlotRange -> All, AxesLabel -> {"x", "y"}]

The conditional distribution $p(x|y)$ is then simply the ratio of the joint distribution to the marginal of $y$.

pcxy = PDF[pxy, {x, y}]/PDF[MarginalDistribution[pxy, 2], x]
Quiet@Plot3D[pcxy, {x, -3, 12}, {y, -10, 8}, PlotRange -> All, AxesLabel -> {"x", "y"}]

Answer 2

As I use this a lot in my own research, let me answer your question by generalizing it to possibly larger dimensions and with a possibly correlated joint probability.

Let me define

 ConditionalMultinormalDistribution::usage ="ConditionalMultinormalDistribution[pdf,val,idx] 
returns the conditional MultiNormal PDF from the joint PDF pdf while setting the variables 
of index idx to values vals"

so that for example:

m = Table[i, {i, 3}];
S = Table[i + j, {i, 3}, {j, 3}]/20 + DiagonalMatrix[Table[1, {3}]];
pdf = MultinormalDistribution[m, S];
cpdf = ConditionalMultinormalDistribution[pdf, {1, 5}, {1, 3}]

(* NormalDistribution[327/139, 317/278] *)

or slightly less trivially,

m = Table[i, {i, 5}];
S = Table[i + j, {i, 5}, {j, 5}]/20 + DiagonalMatrix[Table[1, {5}]];
pdf = MultinormalDistribution[m, S];
cpdf = ConditionalMultinormalDistribution[pdf, {1, 1, 1}, {1, 3, 5}]

(* MultinormalDistribution[{1, 63/23}, {{35/32, 5/32}, {5/32, 885/736}}] *)

ContourPlot[PDF[cpdf, {x, y}], {x, -2, 4}, {y, 0, 6}, PlotRange -> All,
PlotPoints -> 50, Contours -> 15]

The actual code:

ConditionalMultinormalDistribution[pdf_, val_, idx_] := Module[
    {S = pdf[[2]], m = pdf[[1]], odx, Σa, Σb, Σc, μ2, S2, idx2, val2},
    odx  = Flatten[{Complement[Range[Length[S]], Flatten[{idx}]]}];
    Σa   = (S[[odx]] // Transpose)[[odx]];
    idx2 = Flatten[{idx}];
    val2 = Flatten[{val}];
    Σc   = (S[[odx]] // Transpose)[[idx2]] // Transpose;
    Σb   = (S[[idx2]] // Transpose)[[idx2]];
    μ2   = m[[odx]] + Σc.Inverse[Σb].(val2 - m[[idx]]);
    S2   = Σa - Σc.Inverse[Σb]. Transpose[  Σc];
    S2   = 1/2 (S2 + Transpose[S2]);
    If[
        Length[μ2] == 1, 
        NormalDistribution[μ2 // First, Sqrt[S2 // First // First]], 
        MultinormalDistribution[μ2, S2]
    ]
] /; Head[pdf] == MultinormalDistribution