Ancestor Reconstruction.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "c79e74a3",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\u001b[36m\u001b[1m[ \u001b[22m\u001b[39m\u001b[36m\u001b[1mInfo: \u001b[22m\u001b[39mFile exists: 1A3N\n",
      "\u001b[36m\u001b[1m[ \u001b[22m\u001b[39m\u001b[36m\u001b[1mInfo: \u001b[22m\u001b[39mFile exists: 1A3N\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "142×146 Matrix{Float64}:\n",
       " -25.0358  -24.7162  -25.0068  -33.7361   …  -28.0232  -66.9887   -13.434\n",
       " -36.1451  -15.4942  -16.2317  -25.3725      -19.0279  -32.9741   -13.6597\n",
       " -33.6691  -21.2741  -14.5955  -25.4042      -20.5477  -44.5028   -11.7251\n",
       " -39.7601  -34.1013  -33.2484   -5.02986     -35.2034  -30.3247    -3.88709\n",
       " -40.4012  -32.9982  -32.0806   -7.61196     -34.7408  -27.4761    -3.17999\n",
       " -39.8164  -35.0303  -33.3116   -9.24855  …  -34.5567  -28.6569    -3.61262\n",
       " -39.9721  -33.9442  -32.4771   -8.79749     -35.2304  -27.4783    -3.5067\n",
       " -40.3626  -33.7294  -29.2066   -8.62386     -34.5341  -29.1127    -3.54312\n",
       " -39.4897  -34.8798  -32.3605   -9.88389     -33.899   -27.2323    -3.96756\n",
       " -42.5864  -31.472   -31.3898   -9.33567     -33.4383  -29.8048    -3.36862\n",
       " -39.4256  -33.5821  -32.9801   -8.75491  …  -35.4168  -27.1953    -3.52709\n",
       " -41.3326  -33.0187  -31.1069   -7.76237     -33.8885  -28.4409    -3.16163\n",
       " -41.5296  -33.063   -31.1095   -7.71654     -33.9511  -28.0034    -3.15441\n",
       "   ⋮                                      ⋱                         ⋮\n",
       " -40.9774  -33.633   -31.1651   -7.74883  …  -33.7832  -26.5294    -3.37743\n",
       " -41.7462  -31.8062  -31.7774   -9.38849     -33.8059  -30.4752    -3.39513\n",
       " -39.0229  -33.8772  -31.5923   -8.10748     -33.8955  -29.7599    -3.47554\n",
       " -41.0386  -33.4216  -28.7636   -8.64531     -34.0483  -29.0397    -3.52754\n",
       " -41.1785  -32.1647  -32.3694   -9.24471     -34.5342  -29.6882    -3.38969\n",
       " -37.5481  -26.5371  -35.3924   -8.84838  …  -32.1875  -27.3223    -3.38098\n",
       " -43.4024  -27.7105  -26.5937  -10.1882      -29.1568  -19.1809    -4.9569\n",
       " -37.2676  -30.633   -33.4871   -7.85298     -31.7102  -26.8061    -3.61648\n",
       " -41.8893  -26.5085  -27.5764  -11.1675      -28.0125  -18.6169    -5.37511\n",
       " -37.8765  -32.0577  -35.9568  -10.1224      -29.0982  -25.5581    -4.17535\n",
       " -60.3151  -35.7187  -37.3722  -17.2967   …  -37.0523  -13.4074    -8.95197\n",
       " -18.4819  -14.2604  -10.9914   -3.84752     -14.249    -9.21093    0.0"
      ]
     },
     "execution_count": 4,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "using TorusEvol\n",
    "using Distributions\n",
    "using TimerOutputs\n",
    "using Plots, StatsPlots\n",
    "using LogExpFunctions\n",
    "\n",
    "chainY = from_pdb(\"1A3N\", \"A\")\n",
    "chainZ = from_pdb(\"1A3N\", \"B\")\n",
    "#Y = slice(data(chainY), 1:5, :)\n",
    "#Z = slice(data(chainZ), 2:7, :)\n",
    "Y = data(chainY)\n",
    "Z = data(chainZ)\n",
    "\n",
    "data(Y)[2][isnan.(data(Y)[2])] .= 0.0\n",
    "data(Z)[2][isnan.(data(Z)[2])] .= 0.0\n",
    "\n",
    "my_diff = jumping(WrappedDiffusion(-1.2, -0.8, 0.7, 0.9, 1.0, 1.0, 0.2), 1.0)\n",
    "my_diff2 = jumping(WrappedDiffusion(-1.0, -1.0, 1.8, 1.8, 1.0, 1.0, 0.1), 1.0)\n",
    "\n",
    "w = WAG_SubstitutionProcess()\n",
    "processes = [w w; my_diff my_diff2]\n",
    "weights = [1.0, 0.0]\n",
    "ξ = MixtureProductProcess(weights, processes)\n",
    "\n",
    "t = 0.4\n",
    "λ = 0.015; μ = 0.0156; r = 0.6\n",
    "align_model = TKF92([t], λ, μ, r; known_ancestor=true)\n",
    "pair_hmm = PairDataHMM(align_model, num_sites(Y), num_sites(Z));\n",
    "pair_chain_dist = ChainJointDistribution(ξ, align_model)\n",
    "emission_lps = fulljointlogpdf(ξ, t, Y, Z)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "id": "25ad056a",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "logpdfsregime (generic function with 1 method)"
      ]
     },
     "execution_count": 20,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# compute emission lp for each site in alignment, given pairwise emission matrix\n",
    "function logpdfs(alignment, emission_lps)\n",
    "    res = Real[]\n",
    "    N, M = size(emission_lps) .- 1\n",
    "    i = 0\n",
    "    j = 0\n",
    "    for v ∈ alignment \n",
    "        if v == [1, 1]\n",
    "            i += 1\n",
    "            j += 1\n",
    "            push!(res, emission_lps[i, j])\n",
    "        elseif v == [1, 0]\n",
    "            i += 1\n",
    "            push!(res, emission_lps[i, M+1])\n",
    "        else\n",
    "            j += 1\n",
    "            push!(res, emission_lps[N+1, j])\n",
    "        end     \n",
    "    end\n",
    "    return res\n",
    "end\n",
    "\n",
    "# compute emission lps for specific regime choices per site\n",
    "function logpdfsregime(alignment, regimes, Y, Z, ξ, t)\n",
    "    M = length(alignment)\n",
    "    C = num_coords(ξ)\n",
    "    res = fill(0, m)\n",
    "    for m ∈ 1:M \n",
    "        e = regimes[m]\n",
    "        v = alignment[m]\n",
    "        for c ∈ 1:C \n",
    "            p = processes(ξ)[c, e]\n",
    "            if v == [1, 1]\n",
    "                res[m] = logpdf(statdist(p), Y[m]) + logpdf(transdist(p, t, Y[m]), Z[m])\n",
    "            elseif v == [1, 0]\n",
    "                res[m] = logpdf(statdist(p), Y[m])\n",
    "            else\n",
    "                res[m] = logpdf(statdist(p), Z[m])\n",
    "            end   \n",
    "        end\n",
    "    end\n",
    "end"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "id": "dd1cc15c",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "2×153 Matrix{Integer}:\n",
       " 1  1  1  1  1  1  1  1  1  1  1  1  1  …  1  1  1  1  1  1  1  1  1  1  1  1\n",
       " 1  1  1  1  1  1  1  1  1  1  1  1  1     1  1  1  1  1  1  1  1  1  1  1  1"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/plain": [
       "3×153 Matrix{Integer}:\n",
       " 1  1  1  1  1  1  1  1  1  1  1  1  1  …  1  1  1  1  1  1  1  1  1  1  1  1\n",
       " 1  1  1  1  1  1  1  1  1  1  1  1  1     1  1  1  1  1  1  1  1  1  1  1  1\n",
       " 1  1  1  1  1  1  1  1  1  1  1  1  1     1  1  1  1  1  1  1  1  1  1  1  1"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "M_YZ = backward_sampling(pair_hmm.α, align_model)\n",
    "function sample_anc_alignment(M_YZ, Y, Z, t, ξ)\n",
    "    emission_lps = fulljointlogpdf(ξ, t, Y, Z)\n",
    "    align_emission_lps = logpdfs(M_YZ, emission_lps)\n",
    "    anc_align_model = TKF92([t/2, t/2], λ, μ, r; known_ancestor=false)\n",
    "    α = forward_anc(M_YZ, anc_align_model, align_emission_lps)\n",
    "    M_XYZ = backward_sampling_anc(α, anc_align_model)\n",
    "    return M_XYZ\n",
    "end\n",
    "display(M_YZ.data)\n",
    "M_XYZ = sample_anc_alignment(M_YZ, Y, Z, t, ξ)\n",
    "display(M_XYZ.data)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 165,
   "id": "bb1cbcc6",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\u001b[36m\u001b[1m[ \u001b[22m\u001b[39m\u001b[36m\u001b[1mInfo: \u001b[22m\u001b[39m[Turing]: progress logging is disabled globally\n",
      "\u001b[36m\u001b[1m[ \u001b[22m\u001b[39m\u001b[36m\u001b[1mInfo: \u001b[22m\u001b[39m[AdvancedVI]: global PROGRESS is set as false\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "sample_anc_coords (generic function with 1 method)"
      ]
     },
     "execution_count": 165,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "using Turing \n",
    "using ReverseDiff\n",
    "using TorusEvol\n",
    "using LinearAlgebra\n",
    "using StatsPlots\n",
    "Turing.setadbackend(:reversediff)\n",
    "Turing.setrdcache(true)\n",
    "Turing.setprogress!(false)\n",
    "\n",
    "@model function bisection_sampler_vec_mv(p, descs, t)\n",
    "    n = size(descs[1], 2)\n",
    "    D = length(descs)\n",
    "    descs[1] ~ filldist(statdist(p), n)\n",
    "    \n",
    "    x = Matrix{Real}(undef, length(p), n)\n",
    "    for i ∈ 1:n\n",
    "        x[:, i] ~ transdist(p, t, vec(descs[1][:, i]))\n",
    "        for d ∈ 2:D\n",
    "            descs[d][:, i] ~ transdist(p, t, vec(x[:, i]))\n",
    "        end\n",
    "    end\n",
    "end\n",
    "\n",
    "@model function bisection_sampler_sub(p, descs, t)\n",
    "    n = size(descs[1], 2)\n",
    "    D = length(descs)\n",
    "    descs[1] ~ filldist(statdist(p), n)\n",
    "    \n",
    "    x = Matrix{Real}(undef, length(p), n)\n",
    "    for i ∈ 1:n\n",
    "        x[:, i] ~ transdist(p, t, vec(descs[1][:, i]))\n",
    "        for d ∈ 2:D\n",
    "            descs[d][:, i] ~ transdist(p, t, vec(x[:, i]))\n",
    "        end\n",
    "    end\n",
    "end\n",
    "\n",
    "@model function bisection_sampler_wn(p, descs, t)\n",
    "    D = length(descs)\n",
    "    descs[1] ~ statdist(p)\n",
    "    \n",
    "    x = Vector{Real}(undef, length(p))\n",
    "    x ~ transdist(p, t, descs[1])\n",
    "    for d ∈ 2:D\n",
    "        descs[d] ~ transdist(p, t, x)\n",
    "    end\n",
    "end\n",
    "\n",
    "torus_proposal(vv) = [MixtureModel([WrappedNormal(v, I), WrappedNormal(v, 20*I)], [0.8, 0.2]) for v ∈ vv]\n",
    "\n",
    "\n",
    "function sample_anc_coords_wn(p, descs, t; burn_in=300)\n",
    "    mhs = MH(:x => v -> MixtureModel([WrappedNormal(v, I), WrappedNormal(v, 20*I)], [0.8, 0.2]))\n",
    "    n = size(descs[1], 2)\n",
    "    D = length(descs)\n",
    "    X = Matrix{Real}(undef, length(p), n)\n",
    "    for i ∈ 1:n\n",
    "        descs_i = [descs[d][:, i] for d ∈ 1:D]\n",
    "        model = bisection_sampler_wn(p, descs_i, t)\n",
    "        chn = sample(model, mhs, burn_in+1)\n",
    "        X[1, i] = get(chn, :x).x[1][burn_in+1]\n",
    "        X[2, i] = get(chn, :x).x[2][burn_in+1]\n",
    "    end\n",
    "    return x\n",
    "end\n",
    "\n",
    "function sample_anc_coords(p, descs, t; burn_in=300)\n",
    "    mhs = MH(:x => vv -> torus_proposal(vv))\n",
    "    #par = PG(100, :x)\n",
    "    #sampler = eltype(p) <: Integer ? par : mhs\n",
    "    model = bisection_sampler_vec(p, descs, t)\n",
    "    chn = sample(model, mhs, burn_in+1)\n",
    "    return chn\n",
    "end\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 167,
   "id": "cc7cae20",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "2×141 Matrix{Real}:\n",
       " -1.36342    -1.10835   -1.18526  …  -2.21012   -0.745595  -2.66263\n",
       "  0.0718734  -0.945002   2.79838     -0.324282  -1.98636    0.0589801"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/plain": [
       "2×141 Matrix{Real}:\n",
       " 0.0     -1.56373  -1.32366  -0.987045  …  -1.86729   -1.18805   -2.78712\n",
       " 2.7491   2.17697   2.91804  -0.760957      0.280312  -0.481882   0.0"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/plain": [
       "2×145 Matrix{Real}:\n",
       " 0.0      -1.24806  -1.47966  -1.29803   …  -1.67459   -1.28189  -2.73098\n",
       " 1.93653   2.47863   3.11704  -0.591874      0.240513   2.62932   0.0"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "X_angles = sample_anc_coords_wn(my_diff, [data(Y)[2], data(Z)[2]], t)\n",
    "display(X_angles)\n",
    "display(data(Y)[2])\n",
    "display(data(Z)[2])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b4e41137",
   "metadata": {},
   "outputs": [],
   "source": [
    "function ancestor_sampling(M_YZ::Alignment, Y::ObservedChain, Z:;ObservedChain, \n",
    "                           t::Real, ξ::MixtureProductProcess)\n",
    "    # step 1 - sample XYZ alignment\n",
    "    M_XYZ = sample_anc_alignment(M_YZ, Y, Z, t, ξ)\n",
    "    \n",
    "    # step 2 - sample coordinates of X\n",
    "    alignment = M_XYZ\n",
    "    X_mask = mask(alignment, [[1], [0,1], [0,1]])\n",
    "    alignmentX = slice(alignment, X_mask)\n",
    "    Y_mask = mask(alignment, [[0,1], [1], [0,1]])\n",
    "    alignmentY = slice(alignment, Y_mask)\n",
    "    Z_mask = mask(alignment, [[0,1], [0,1], [1]])\n",
    "    alignmentZ = slice(alignment, Z_mask)\n",
    "    \n",
    "    M = length(alignment)\n",
    "    regimes = ones(M)\n",
    "    \n",
    "    X_maskX = mask(alignmentX, [[1], [0], [0]])\n",
    "    \n",
    "    XY_maskX = mask(alignmentX, [[1], [1], [0]])\n",
    "    XY_maskY = mask(alignmentY, [[1], [1], [0]])\n",
    "    \n",
    "    XZ_maskX = mask(alignmentX, [[1], [0], [1]])\n",
    "    XZ_maskZ = mask(alignmentZ, [[1], [0], [1]])\n",
    "    \n",
    "    XYZ_maskX = mask(alignmentX, [[1], [1], [1]])\n",
    "    XYZ_maskY = mask(alignmentY, [[1], [1], [1]])\n",
    "    XYZ_maskZ = mask(alignmentZ, [[1], [1], [1]])\n",
    "    \n",
    "    dataY = data(Y)\n",
    "    dataZ = data(Z)\n",
    "    \n",
    "    # Initialise internal coordinates of X\n",
    "    N = sequence_lengths(M_XYZ)[1]\n",
    "    dataX = [similar(dataY[c], size(dataY[c], 1), N) for c ∈ 1:C]\n",
    "    \n",
    "    for c ∈ 1:C, e ∈ 1:E\n",
    "        p = processes(ξ)[c, e]\n",
    "        \n",
    "        # [1, 0, 0] - sample from stationary distribution\n",
    "        dataX100 = @view dataX[c][:, X_maskX .& regimesX .== e]\n",
    "        n100 = size(dataX100, 2)\n",
    "        \n",
    "        dataX100 .= rand(statdist(p), n100)\n",
    "        \n",
    "        \n",
    "        # [1, 1, 0] - observe Y, then sample X from Y\n",
    "        dataY110 = @view dataY[c][:, XY_maskY .& regimesY .== e]\n",
    "        dataX110 = @view dataX[c][:, XY_maskX .& regimesX .== e]\n",
    "        \n",
    "        dataX110 .= sample_anc_coords(p, [dataY110], t)\n",
    "    \n",
    "        \n",
    "        # [1, 0, 1] - observe Z, then sample X from Z\n",
    "        dataZ101 = @view dataZ[c][:, XZ_maskZ .& regimesZ .== e]\n",
    "        dataX101 = @view dataX[c][:, XZ_maskZ .& regimesX .== e]\n",
    "        \n",
    "        dataX101 .= sample_anc_coords(p, [dataZ101], t)\n",
    "        \n",
    "        \n",
    "        # [1, 1, 1] - observe Y, sample X from Y, then observe Z from X\n",
    "        dataY111 = @view dataY[c][:, XYZ_maskY .& regimesY .== e]\n",
    "        dataZ111 = @view dataZ[c][:, XYZ_maskZ .& regimesZ .== e]\n",
    "        dataX111 = @view dataX[c][:, XYZ_maskX .& regimesX .== e]\n",
    "        \n",
    "        dataX111 .= sample_anc_coords(p, [dataY111, dataZ111], t)\n",
    "    end\n",
    "    \n",
    "    X = ObservedChain(dataX)\n",
    "end\n",
    "\n",
    "function trajectory_reconstruction(Y, Z, align_model, ξ, t;levels=1)\n",
    "    # sample alignment\n",
    "    \n",
    "    # for each level\n",
    "    ancestor_sampling(M_YZ, regimes, Y, Z, t, ξ)\n",
    "    trajectory_reconstruction(Y, X, align_model, ξ, t/2; levels=levels-1)\n",
    "    trajectory_reconstruction(X, Z, align_model, ξ, t/2; levels=levels-1)\n",
    "end"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "a3c7b31e",
   "metadata": {},
   "outputs": [
    {
     "ename": "InterruptException",
     "evalue": "InterruptException:",
     "output_type": "error",
     "traceback": [
      "InterruptException:",
      "",
      "Stacktrace:",
      " [1] display_dict(x::Function)",
      "   @ IJulia ~/.julia/packages/IJulia/6TIq1/src/display.jl:102",
      " [2] #invokelatest#2",
      "   @ ./essentials.jl:816 [inlined]",
      " [3] invokelatest",
      "   @ ./essentials.jl:813 [inlined]",
      " [4] execute_request(socket::ZMQ.Socket, msg::IJulia.Msg)",
      "   @ IJulia ~/.julia/packages/IJulia/6TIq1/src/execute_request.jl:112",
      " [5] #invokelatest#2",
      "   @ ./essentials.jl:816 [inlined]",
      " [6] invokelatest",
      "   @ ./essentials.jl:813 [inlined]",
      " [7] eventloop(socket::ZMQ.Socket)",
      "   @ IJulia ~/.julia/packages/IJulia/6TIq1/src/eventloop.jl:8",
      " [8] (::IJulia.var\"#15#18\")()",
      "   @ IJulia ./task.jl:514"
     ]
    }
   ],
   "source": [
    "using Turing \n",
    "using ReverseDiff\n",
    "using TorusEvol\n",
    "#Turing.setadbackend(:reversediff)\n",
    "Turing.setrdcache(true)\n",
    "Turing.setadbackend(:forwarddiff)\n",
    "Turing.setprogress!(true)\n",
    "\n",
    "@model function ancestor_sampling(alignment, regimes, dataY, dataZ, t, ξ)\n",
    "    C = num_coords(ξ)\n",
    "    E = num_regimes(ξ)\n",
    "    \n",
    "    \n",
    "    #dataX = \n",
    "    \n",
    "\n",
    "    X_mask = mask(alignment, [[1], [0,1], [0,1]])\n",
    "    alignmentX = alignment[X_mask]\n",
    "    Y_mask = mask(alignment, [[0,1], [1], [0,1]])\n",
    "    alignmentY = alignment[Y_mask]\n",
    "    Z_mask = mask(alignment, [[0,1], [0,1], [1]])\n",
    "    alignmentZ = alignment[Z_mask]\n",
    "    \n",
    "    M = length(alignment)\n",
    "    regimes = tzeros(Int, M)\n",
    "    regimes ~ filldist(Categorical(weights(ξ)), M)\n",
    "    regimesX = regimes[X_mask]\n",
    "    regimesY = regimes[Y_mask]\n",
    "    regimesZ = regimes[Z_mask]\n",
    "    \n",
    "    X_maskX = mask(alignmentX, [[1], [0], [0]])\n",
    "    \n",
    "    XY_maskX = mask(alignmentX, [[1], [1], [0]])\n",
    "    XY_maskY = mask(alignmentY, [[1], [1], [0]])\n",
    "    \n",
    "    XZ_maskX = mask(alignmentX, [[1], [0], [1]])\n",
    "    XZ_maskZ = mask(alignmentZ, [[1], [0], [1]])\n",
    "    \n",
    "    XYZ_maskX = mask(alignmentX, [[1], [1], [1]])\n",
    "    XYZ_maskY = mask(alignmentY, [[1], [1], [1]])\n",
    "    XYZ_maskZ = mask(alignmentZ, [[1], [1], [1]])\n",
    "    \n",
    "    for c ∈ 1:C, e ∈ 1:E\n",
    "        p = processes(ξ)[c, e]\n",
    "        # [1, 0, 0] - sample from stationary distribution\n",
    "        data[c][:, X_maskX .& regimesX .== e] .~ statdist(p)\n",
    "        \n",
    "        # [1, 1, 0] - observe Y, then sample X from Y\n",
    "        dataY[c][:, XY_maskY .& regimesY .== e] .~ statdist(p)\n",
    "        dataX[c][:, XY_maskX .& regimesX .== e] ~ arraydist(transdist.(Ref(p), Ref(t), eachcol(dataY[c][:, XY_maskY .& regimesY .== e])))\n",
    "        \n",
    "        # [1, 0, 1] - observe Z, then sample X from Z\n",
    "        dataY[c][:, XZ_maskZ .& regimesZ .== e] .~ statdist(p)\n",
    "        dataX[c][:, XZ_maskX .& regimesX .== e] ~ arraydist(transdist.(Ref(p), Ref(t), eachcol(dataZ[c][:, XY_maskZ .& regimesZ .== e])))\n",
    "        \n",
    "        # [1, 1, 1] - observe Y, sample X from Y, then observe Z from X\n",
    "        dataY[c][:, XYZ_maskY .& regimesY .== e] .~ statdist(p)\n",
    "        dataX[c][:, XYZ_maskX .& regimesX .== e] ~ arraydist(transdist.(Ref(p), Ref(t), eachcol(dataY[c][:, XYZ_maskY .& regimesY .== e])))\n",
    "        dataZ[c][:, XYZ_maskZ .& regimesZ .== e] ~ arraydist(transdist.(Ref(p), Ref(t), eachcol(dataX[c][:, XYZ_maskX .& regimesX .== e])))\n",
    "    end\n",
    "end\n",
    "    "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5fe15fb6",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Julia 1.9.0-rc2",
   "language": "julia",
   "name": "julia-1.9"
  },
  "language_info": {
   "file_extension": ".jl",
   "mimetype": "application/julia",
   "name": "julia",
   "version": "1.9.0"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}