SetDirectory["~/projects/MarkovJumpToolbox"];
<<MaTeX`
<<MarkovJumpToolbox`

xmax = 2.;
xmin=-2.;
xbins = 10.;
ymax = 2.;
ymin = -2;
ybins = 10.;
Diffusion = {1., 1.};
A = 0.;
B = 4.;
grid = GenerateGrid[xmax, xmin, xbins,ymax, ymin, ybins];
{edges, W} = GenerateW4well[grid, Diffusion, A,B];
wellcenters = {{1, 1}, {1, -1}, {-1, -1}, {-1, 1}};
σ = 1;
U[x_, y_]:=-B Sum[Exp[-((x-xv[[1]])^2+(y-xv[[2]])^2)/σ^2], {xv, wellcenters}];

Plot3D[
U[x,y], {x, xmin, xmax}, {y, ymin, ymax}
]

discretizedEnergies = Table[
{x,y,U[x,y]}, {x, grid[[4]]}, {y, grid[[5]]}
];

energyplot = ListContourPlot[Flatten[discretizedEnergies,1], InterpolationOrder->0, PlotLegends->True]

(* Discrete time transition matrix *)
M = MatrixExp[W*0.01];

(* Take NumberOfSteps steps of dynamics to generate a trajectory *)
newsite = 12; (* Start the system in site number 12 -- this is arbitrary *)
NumberOfSteps = 20000;
traj = Table[oldsite = newsite; newsite = CompiledBinarySearch[Accumulate[Normal[Transpose[M][[oldsite]]]], RandomReal[]],
{t, 1, NumberOfSteps}
];

(* Movie of where the trajectory visited *)
Dynamic[
Show[
ListContourPlot[
Flatten[discretizedEnergies,1], InterpolationOrder->0, PlotLegends->True
],
ListPlot[
{grid[[3,traj[[Clock[{1, NumberOfSteps,1}, 50]]]]]},
PlotMarkers->{Black, ■, 36}
]
]
]

xpositions = Table[grid[[3,traj[[i]]]][[1]], {i, 1, NumberOfSteps}];
Mean[xpositions]

(* Histogram of where the trajectory visited *)

binpositions = Table[grid[[3,i]], {i, 1, xbins * ybins}];
binheights = HistogramList[traj,{Table[i, {i, 1, xbins * ybins+1}]}, "PDF"][[2]];

ListContourPlot[
Flatten[#]&/@Transpose[{binpositions, binheights}],
PlotLegends->True
]

p = Table[0, {i, 1, xbins * ybins}];
p[[12]] = 1;
p

(* Function to plot probability distribution p *)
DistributionPlotter[p_]:=ListContourPlot[
Flatten[#]&/@Transpose[{binpositions,p}],
PlotLegends->True,
PlotRange->All,
InterpolationOrder->0
]

(* Initial configuration *)
DistributionPlotter[p]

(* After 1 step of dynamics *)
DistributionPlotter[M.p]

(* After 10 steps of dynamics *)
DistributionPlotter[MatrixPower[M,10].p]

(* After 100 steps of dynamics *)
DistributionPlotter[MatrixPower[M,100].p]

(* After 1000 steps of dynamics *)
DistributionPlotter[MatrixPower[M,1000].p]

(* After 1000000 steps of dynamics *)
DistributionPlotter[MatrixPower[M,1000000].p]

{eigenvalues, eigenvectors} = Eigensystem[M];

eigenvalues

UnnormalizedSteadyState = Chop[eigenvectors[[1]]];
NormalizedSteadyState = UnnormalizedSteadyState / Total[UnnormalizedSteadyState]

(* Here I compute the spatial averages <x> and <y> in a somewhat slick way. grid[[3]] is a list of the (x,y) corrdinates for each cell (in our standard ordering convention. NormalizedSteadyState is a list of the normalized probability for those cells in the steady state.*)
{avgx, avgy} =  Total[grid[[3]] * NormalizedSteadyState]
{avgx2, avgy2} = Total[grid[[3]]^2 * NormalizedSteadyState]

(* The steady state distribution from the largest eigenvalue *)
densityplot = DistributionPlotter[NormalizedSteadyState]

For reasons we could talk more about, the stationary state for this dynamics is the Boltzmann distribution...

discretizedProbabilities = Table[
{x,y,Exp[-U[x,y]]}, {x, grid[[4]]}, {y, grid[[5]]}
];

NormalizedDiscretizedProbabilities = Table[
{x,y,Exp[-U[x,y]]/Total[Total[discretizedProbabilities]][[3]]}, {x, grid[[4]]}, {y, grid[[5]]}
];

Probplot = ListContourPlot[Flatten[NormalizedDiscretizedProbabilities,1], InterpolationOrder->0, PlotLegends->True]