text:
# Face flux

abstract type SpeedFlux end

struct SpeedFluxMax{T} <: SpeedFlux
   Cmax_XF::Array{T,3}
   Cmax_YF::Array{T,3}
   Cmax_ZF::Array{T,3}
end

struct SpeedFluxMinMax{T} <: SpeedFlux
   Cmax_XF::Array{T,3}
   Cmax_YF::Array{T,3}
   Cmax_ZF::Array{T,3}
   Cmin_XF::Array{T,3}
   Cmin_YF::Array{T,3}
   Cmin_ZF::Array{T,3}
end

struct FaceFlux{T}
   Flux_XV::Array{T,4}
   Flux_YV::Array{T,4}
   Flux_ZV::Array{T,4}
end

struct FaceFluxLR{T}
   LFlux_XV::Array{T,4}
   RFlux_XV::Array{T,4}
   LFlux_YV::Array{T,4}
   RFlux_YV::Array{T,4}
   LFlux_ZV::Array{T,4}
   RFlux_ZV::Array{T,4}
end

function init_flux(param::Param)
   (;GridSize, nVar) = param

   Flux_XV = Array{Float64,4}(undef,GridSize+[1,0,0]...,nVar)
   Flux_YV = Array{Float64,4}(undef,GridSize+[0,1,0]...,nVar)
   Flux_ZV = Array{Float64,4}(undef,GridSize+[0,0,1]...,nVar)

   LFlux_XV = Array{Float64,4}(undef,GridSize+[1,0,0]...,nVar)
   RFlux_XV = Array{Float64,4}(undef,GridSize+[1,0,0]...,nVar)
   LFlux_YV = Array{Float64,4}(undef,GridSize+[0,1,0]...,nVar)
   RFlux_YV = Array{Float64,4}(undef,GridSize+[0,1,0]...,nVar)
   LFlux_ZV = Array{Float64,4}(undef,GridSize+[0,0,1]...,nVar)
   RFlux_ZV = Array{Float64,4}(undef,GridSize+[0,0,1]...,nVar)

   Cmax_XF = Array{Float64,3}(undef,GridSize+[1,0,0]...)
   Cmax_YF = Array{Float64,3}(undef,GridSize+[0,1,0]...)
   Cmax_ZF = Array{Float64,3}(undef,GridSize+[0,0,1]...)

   faceFluxLR = FaceFluxLR(LFlux_XV, RFlux_XV, LFlux_YV, RFlux_YV, LFlux_ZV, RFlux_ZV)
   faceFlux = FaceFlux(Flux_XV, Flux_YV, Flux_ZV)
   if param.Scheme == "Rusanov"
      speedFlux = SpeedFluxMax(Cmax_XF, Cmax_YF, Cmax_ZF)
   elseif param.Scheme == "HLLE"
      Cmin_XF = similar(Cmax_XF)
      Cmin_YF = similar(Cmax_YF)
      Cmin_ZF = similar(Cmax_ZF)
      speedFlux = SpeedFluxMinMax(Cmax_XF, Cmax_YF, Cmax_ZF, Cmin_XF, Cmin_YF, Cmin_ZF)
   end

   return faceFluxLR, faceFlux, speedFlux
end

function calc_face_flux!(param::Param, faceValue::FaceState,
   faceFlux::FaceFlux, speedFlux::SpeedFlux, faceFluxLR::FaceFluxLR)

   get_physical_flux!(param, faceValue, faceFlux, speedFlux, faceFluxLR)

   add_numerical_flux!(param, faceValue, faceFlux, speedFlux, faceFluxLR)

   return
end

function get_physical_flux!(param::Param, faceValue::FaceState,
   faceFlux::FaceFlux, speedFlux::SpeedFlux, faceFluxLR::FaceFluxLR)

   (;LState_XV, RState_XV, LState_YV, RState_YV, LState_ZV, RState_ZV) = faceValue
   (;LFlux_XV, RFlux_XV, LFlux_YV, RFlux_YV, LFlux_ZV, RFlux_ZV) = faceFluxLR
   (;Flux_XV, Flux_YV, Flux_ZV) = faceFlux
   (;nI, nJ, nK, nVar) = param

   # Density flux
   @inbounds for k = 1:nK, j = 1:nJ, i = 1:nI+1
      LFlux_XV[i,j,k,Rho_] = LState_XV[i,j,k,Mx_]
      RFlux_XV[i,j,k,Rho_] = RState_XV[i,j,k,Mx_]
   end

   @inbounds for k = 1:nK, j = 1:nJ+1, i = 1:nI
      LFlux_YV[i,j,k,Rho_] = LState_YV[i,j,k,My_]
      RFlux_YV[i,j,k,Rho_] = RState_YV[i,j,k,My_]
   end

   @inbounds for k = 1:nK+1, j = 1:nJ, i = 1:nI
      LFlux_ZV[i,j,k,Rho_] = LState_ZV[i,j,k,Mz_]
      RFlux_ZV[i,j,k,Rho_] = RState_ZV[i,j,k,Mz_]
   end

   # Momentum flux
   @inbounds for k = 1:nK, j = 1:nJ, i = 1:nI+1
      bL = □(LState_XV[i,j,k,Bx_], LState_XV[i,j,k,By_], LState_XV[i,j,k,Bz_])
      bR = □(RState_XV[i,j,k,Bx_], RState_XV[i,j,k,By_], RState_XV[i,j,k,Bz_])

      LFlux_XV[i,j,k,Mx_] = LState_XV[i,j,k,Mx_]^2 / LState_XV[i,j,k,Rho_] +
         LState_XV[i,j,k,P_] + 0.5*bL - LState_XV[i,j,k,Bx_]^2
      RFlux_XV[i,j,k,Mx_] = RState_XV[i,j,k,Mx_]^2 / RState_XV[i,j,k,Rho_] +
         RState_XV[i,j,k,P_] + 0.5*bR - RState_XV[i,j,k,Bx_]^2
      LFlux_XV[i,j,k,My_] = LState_XV[i,j,k,Mx_] * LState_XV[i,j,k,My_] /
         LState_XV[i,j,k,Rho_] - LState_XV[i,j,k,Bx_]*LState_XV[i,j,k,By_]
      RFlux_XV[i,j,k,My_] = RState_XV[i,j,k,Mx_] * RState_XV[i,j,k,My_] /
         RState_XV[i,j,k,Rho_] - RState_XV[i,j,k,Bx_]*RState_XV[i,j,k,By_]
      LFlux_XV[i,j,k,Mz_] = LState_XV[i,j,k,Mx_] * LState_XV[i,j,k,Mz_] /
         LState_XV[i,j,k,Rho_] - LState_XV[i,j,k,Bx_]*LState_XV[i,j,k,Bz_]
      RFlux_XV[i,j,k,Mz_] = RState_XV[i,j,k,Mx_] * RState_XV[i,j,k,Mz_] /
         RState_XV[i,j,k,Rho_] - RState_XV[i,j,k,Bx_]*RState_XV[i,j,k,Bz_]
   end

   @inbounds for k = 1:nK, j = 1:nJ+1, i = 1:nI
      bL = □(LState_YV[i,j,k,Bx_], LState_YV[i,j,k,By_], LState_YV[i,j,k,Bz_])
      bR = □(RState_YV[i,j,k,Bx_], RState_YV[i,j,k,By_], RState_YV[i,j,k,Bz_])

      LFlux_YV[i,j,k,Mx_] = LState_YV[i,j,k,My_] * LState_YV[i,j,k,Mx_] /
         LState_YV[i,j,k,Rho_] - LState_YV[i,j,k,By_]*LState_YV[i,j,k,Bx_]
      RFlux_YV[i,j,k,Mx_] = RState_YV[i,j,k,My_] * RState_YV[i,j,k,Mx_] /
         RState_YV[i,j,k,Rho_] - RState_YV[i,j,k,By_]*RState_YV[i,j,k,Bx_]
      LFlux_YV[i,j,k,My_] = LState_YV[i,j,k,My_]^2 / LState_YV[i,j,k,Rho_] +
         LState_YV[i,j,k,P_] + 0.5*bL - LState_YV[i,j,k,By_]^2
      RFlux_YV[i,j,k,My_] = RState_YV[i,j,k,My_]^2 / RState_YV[i,j,k,Rho_] +
         RState_YV[i,j,k,P_] + 0.5*bR - RState_YV[i,j,k,By_]^2
      LFlux_YV[i,j,k,Mz_] = LState_YV[i,j,k,My_] * LState_YV[i,j,k,Mz_] /
         LState_YV[i,j,k,Rho_] - LState_YV[i,j,k,Bx_]*LState_YV[i,j,k,Bz_]
      RFlux_YV[i,j,k,Mz_] = RState_YV[i,j,k,My_] * RState_YV[i,j,k,Mz_] /
         RState_YV[i,j,k,Rho_] - RState_YV[i,j,k,Bx_]*RState_YV[i,j,k,Bz_]
   end

   @inbounds for k = 1:nK+1, j = 1:nJ, i = 1:nI
      bL = □(LState_ZV[i,j,k,Bx_], LState_ZV[i,j,k,By_], LState_ZV[i,j,k,Bz_])
      bR = □(RState_ZV[i,j,k,Bx_], RState_ZV[i,j,k,By_], RState_ZV[i,j,k,Bz_])

      LFlux_ZV[i,j,k,Mx_] = LState_ZV[i,j,k,Mz_] * LState_ZV[i,j,k,Mx_] /
         LState_ZV[i,j,k,Rho_] - LState_ZV[i,j,k,Bz_]*LState_ZV[i,j,k,Bx_]
      RFlux_ZV[i,j,k,Mx_] = RState_ZV[i,j,k,Mz_] * RState_ZV[i,j,k,Mx_] /
         RState_ZV[i,j,k,Rho_] - RState_ZV[i,j,k,Bz_]*RState_ZV[i,j,k,Bx_]
      LFlux_ZV[i,j,k,My_] = LState_ZV[i,j,k,Mz_] * LState_ZV[i,j,k,My_] /
         LState_ZV[i,j,k,Rho_] - LState_ZV[i,j,k,Bz_]*LState_ZV[i,j,k,By_]
      RFlux_ZV[i,j,k,My_] = RState_ZV[i,j,k,Mz_] * RState_ZV[i,j,k,My_] /
         RState_ZV[i,j,k,Rho_] - RState_ZV[i,j,k,Bz_]*RState_ZV[i,j,k,By_]
      LFlux_ZV[i,j,k,Mz_] = LState_ZV[i,j,k,Mz_]^2 / LState_ZV[i,j,k,Rho_] +
         LState_ZV[i,j,k,P_] + 0.5*bL - LState_ZV[i,j,k,Bz_]^2
      RFlux_ZV[i,j,k,Mz_] = RState_ZV[i,j,k,Mz_]^2 / RState_ZV[i,j,k,Rho_] +
         RState_ZV[i,j,k,P_] + 0.5*bR - RState_ZV[i,j,k,Bz_]^2
   end

   # Magnetic flux
   @inbounds for k = 1:nK, j = 1:nJ, i = 1:nI+1
      LFlux_XV[i,j,k,Bx_] = 0.0
      RFlux_XV[i,j,k,Bx_] = 0.0
      LFlux_XV[i,j,k,By_] = LState_XV[i,j,k,Mx_]*LState_XV[i,j,k,By_] -
         LState_XV[i,j,k,Bx_]*LState_XV[i,j,k,My_]
      RFlux_XV[i,j,k,By_] = RState_XV[i,j,k,Mx_]*RState_XV[i,j,k,By_] -
         RState_XV[i,j,k,Bx_]*RState_XV[i,j,k,My_]
      LFlux_XV[i,j,k,Bz_] = LState_XV[i,j,k,Mx_]*LState_XV[i,j,k,Bz_] -
         LState_XV[i,j,k,Bx_]*LState_XV[i,j,k,Mz_]
      RFlux_XV[i,j,k,Bz_] = RState_XV[i,j,k,Mx_]*RState_XV[i,j,k,Bz_] -
         RState_XV[i,j,k,Bx_]*RState_XV[i,j,k,Mz_]
   end

   @inbounds for k = 1:nK, j = 1:nJ+1, i = 1:nI
      LFlux_YV[i,j,k,Bx_] = LState_YV[i,j,k,My_]*LState_YV[i,j,k,Bx_] -
         LState_YV[i,j,k,By_]*LState_YV[i,j,k,Mx_]
      RFlux_YV[i,j,k,Bx_] = RState_YV[i,j,k,My_]*RState_YV[i,j,k,Bx_] -
         RState_YV[i,j,k,By_]*RState_YV[i,j,k,Mx_]
      LFlux_YV[i,j,k,By_] = 0.0
      RFlux_YV[i,j,k,By_] = 0.0
      LFlux_YV[i,j,k,Bz_] = LState_YV[i,j,k,My_]*LState_YV[i,j,k,Bz_] -
         LState_YV[i,j,k,By_]*LState_YV[i,j,k,Mz_]
      RFlux_YV[i,j,k,Bz_] = RState_YV[i,j,k,My_]*RState_YV[i,j,k,Bz_] -
         RState_YV[i,j,k,By_]*RState_YV[i,j,k,Mz_]
   end

   @inbounds for k = 1:nK+1, j = 1:nJ, i = 1:nI
      LFlux_ZV[i,j,k,Bx_] = LState_ZV[i,j,k,Mz_]*LState_ZV[i,j,k,Bx_] -
         LState_ZV[i,j,k,Bz_]*LState_ZV[i,j,k,Mx_]
      RFlux_ZV[i,j,k,Bx_] = RState_ZV[i,j,k,Mz_]*RState_ZV[i,j,k,Bx_] -
         RState_ZV[i,j,k,Bz_]*RState_ZV[i,j,k,Mx_]
      LFlux_ZV[i,j,k,By_] = LState_ZV[i,j,k,Mz_]*LState_ZV[i,j,k,By_] -
         LState_ZV[i,j,k,Bz_]*LState_ZV[i,j,k,My_]
      RFlux_ZV[i,j,k,By_] = RState_ZV[i,j,k,Mz_]*RState_ZV[i,j,k,By_] -
         RState_ZV[i,j,k,Bz_]*RState_ZV[i,j,k,My_]
      LFlux_ZV[i,j,k,Bz_] = 0.0
      RFlux_ZV[i,j,k,Bz_] = 0.0
   end

   # Pressure flux / energy flux
   if !param.UseConservative
      @inbounds for k = 1:nK, j = 1:nJ, i = 1:nI+1
         LFlux_XV[i,j,k,P_] = LState_XV[i,j,k,P_]*LState_XV[i,j,k,Mx_]/
            LState_XV[i,j,k,Rho_]
         RFlux_XV[i,j,k,P_] = RState_XV[i,j,k,P_]*RState_XV[i,j,k,Mx_]/
            RState_XV[i,j,k,Rho_]
      end

      @inbounds for k = 1:nK, j = 1:nJ+1, i = 1:nI
         LFlux_YV[i,j,k,P_] = LState_YV[i,j,k,P_]*LState_YV[i,j,k,My_]/
            LState_YV[i,j,k,Rho_]
         RFlux_YV[i,j,k,P_] = RState_YV[i,j,k,P_]*RState_YV[i,j,k,My_]/
            RState_YV[i,j,k,Rho_]
      end

      @inbounds for k = 1:nK+1, j = 1:nJ, i = 1:nI
         LFlux_ZV[i,j,k,P_] = LState_ZV[i,j,k,P_]*LState_ZV[i,j,k,Mz_]/
            LState_ZV[i,j,k,Rho_]
         RFlux_ZV[i,j,k,P_] = RState_ZV[i,j,k,P_]*RState_ZV[i,j,k,Mz_]/
            RState_ZV[i,j,k,Rho_]
      end
   else
      # Currently use the same index for pressure/energy
      @inbounds for k = 1:nK, j = 1:nJ, i = 1:nI+1
         u = □(LState_XV[i,j,k,Mx_], LState_XV[i,j,k,My_], LState_XV[i,j,k,Mz_])
         b = □(LState_XV[i,j,k,Bx_], LState_XV[i,j,k,By_], LState_XV[i,j,k,Bz_])
         ub= LState_XV[i,j,k,Mx_]*LState_XV[i,j,k,Bx_] +
            LState_XV[i,j,k,My_]*LState_XV[i,j,k,By_] +
            LState_XV[i,j,k,Mz_]*LState_XV[i,j,k,Bz_]
         LFlux_XV[i,j,k,E_] = LState_XV[i,j,k,Mx_]/LState_XV[i,j,k,Rho_]*
            ((LState_XV[i,j,k,P_]/(γ-1.0) + 0.5/LState_XV[i,j,k,Rho_]*u +
            0.5*b + LState_XV[i,j,k,P_] + 0.5*b)) - ub *
            LState_XV[i,j,k,Bx_]
      end

      @inbounds for k = 1:nK, j = 1:nJ, i = 1:nI+1
         u = □(RState_XV[i,j,k,Mx_], RState_XV[i,j,k,My_], RState_XV[i,j,k,Mz_])
         b = □(RState_XV[i,j,k,Bx_], RState_XV[i,j,k,By_], RState_XV[i,j,k,Bz_])
         ub= RState_XV[i,j,k,Mx_]*RState_XV[i,j,k,Bx_] +
            RState_XV[i,j,k,My_]*RState_XV[i,j,k,By_] +
            RState_XV[i,j,k,Mz_]*RState_XV[i,j,k,Bz_]
         RFlux_XV[i,j,k,E_] = RState_XV[i,j,k,Mx_]/RState_XV[i,j,k,Rho_]*
            ((RState_XV[i,j,k,P_]/(γ-1.0) + 0.5/RState_XV[i,j,k,Rho_]*u +
            0.5*b + RState_XV[i,j,k,P_] + 0.5*b)) - ub *
            RState_XV[i,j,k,Bx_]
      end

      @inbounds for k = 1:nK, j = 1:nJ+1, i = 1:nI
         u = □(LState_YV[i,j,k,Mx_], LState_YV[i,j,k,My_], LState_YV[i,j,k,Mz_])
         b = □(LState_YV[i,j,k,Bx_], LState_YV[i,j,k,By_], LState_YV[i,j,k,Bz_])
         ub= LState_YV[i,j,k,Mx_]*LState_YV[i,j,k,Bx_] +
            LState_YV[i,j,k,My_]*LState_YV[i,j,k,By_] +
            LState_YV[i,j,k,Mz_]*LState_YV[i,j,k,Bz_]
         LFlux_YV[i,j,k,E_] = LState_YV[i,j,k,My_]/LState_YV[i,j,k,Rho_]*
            ((LState_YV[i,j,k,P_]/(γ-1.0) + 0.5/LState_YV[i,j,k,Rho_]*u +
            0.5*b + LState_YV[i,j,k,P_] + 0.5*b)) - ub *
            LState_YV[i,j,k,Bx_]
      end

      @inbounds for k = 1:nK, j = 1:nJ+1, i = 1:nI
         u = □(RState_YV[i,j,k,Mx_], RState_YV[i,j,k,My_], RState_YV[i,j,k,Mz_])
         b = □(RState_YV[i,j,k,Bx_], RState_YV[i,j,k,By_], RState_YV[i,j,k,Bz_])
         ub= RState_YV[i,j,k,Mx_]*RState_YV[i,j,k,Bx_] +
            RState_YV[i,j,k,My_]*RState_YV[i,j,k,By_] +
            RState_YV[i,j,k,Mz_]*RState_YV[i,j,k,Bz_]
         RFlux_YV[i,j,k,E_] = RState_YV[i,j,k,My_]/RState_YV[i,j,k,Rho_]*
            ((RState_YV[i,j,k,P_]/(γ-1.0) + 0.5/RState_YV[i,j,k,Rho_]*u +
            0.5*b + RState_YV[i,j,k,P_] + 0.5*b)) - ub *
            RState_YV[i,j,k,Bx_]
      end

      @inbounds for k = 1:nK+1, j = 1:nJ, i = 1:nI
         u = □(LState_ZV[i,j,k,Mx_], LState_ZV[i,j,k,My_], LState_ZV[i,j,k,Mz_])
         b = □(LState_ZV[i,j,k,Bx_], LState_ZV[i,j,k,By_], LState_ZV[i,j,k,Bz_])
         ub= LState_ZV[i,j,k,Mx_]*LState_ZV[i,j,k,Bx_] +
            LState_ZV[i,j,k,My_]*LState_ZV[i,j,k,By_] +
            LState_ZV[i,j,k,Mz_]*LState_ZV[i,j,k,Bz_]
         LFlux_ZV[i,j,k,E_] = LState_ZV[i,j,k,Mz_]/LState_ZV[i,j,k,Rho_]*
            ((LState_ZV[i,j,k,P_]/(γ-1.0) + 0.5/LState_ZV[i,j,k,Rho_]*u +
            0.5*b + LState_ZV[i,j,k,P_] + 0.5*b)) - ub *
            LState_ZV[i,j,k,Bx_]
      end

      @inbounds for k = 1:nK+1, j = 1:nJ, i = 1:nI
         u = □(RState_ZV[i,j,k,Mx_], RState_ZV[i,j,k,My_], RState_ZV[i,j,k,Mz_])
         b = □(RState_ZV[i,j,k,Bx_], RState_ZV[i,j,k,By_], RState_ZV[i,j,k,Bz_])
         ub= RState_ZV[i,j,k,Mx_]*RState_ZV[i,j,k,Bx_] +
            RState_ZV[i,j,k,My_]*RState_ZV[i,j,k,By_] +
            RState_ZV[i,j,k,Mz_]*RState_ZV[i,j,k,Bz_]
         RFlux_ZV[i,j,k,E_] = RState_ZV[i,j,k,Mz_]/RState_ZV[i,j,k,Rho_]*
            ((RState_ZV[i,j,k,P_]/(γ-1.0) + 0.5/RState_ZV[i,j,k,Rho_]*u +
            0.5*b + RState_ZV[i,j,k,P_] + 0.5*b)) - ub *
            RState_ZV[i,j,k,Bx_]
      end
   end

   # Collect all the physical fluxes
   @inbounds for iVar = 1:nVar, k = 1:nK, j = 1:nJ, i = 1:nI+1
      Flux_XV[i,j,k,iVar] = 0.5*(LFlux_XV[i,j,k,iVar] + RFlux_XV[i,j,k,iVar])
   end
   @inbounds for iVar = 1:nVar, k = 1:nK, j = 1:nJ+1, i = 1:nI
      Flux_YV[i,j,k,iVar] = 0.5*(LFlux_YV[i,j,k,iVar] + RFlux_YV[i,j,k,iVar])
   end
   @inbounds for iVar = 1:nVar, k = 1:nK+1, j = 1:nJ, i = 1:nI
      Flux_ZV[i,j,k,iVar] = 0.5*(LFlux_ZV[i,j,k,iVar] + RFlux_ZV[i,j,k,iVar])
   end

end

function add_numerical_flux!(param::Param, faceValue::FaceState,
   faceFlux::FaceFlux, speedFlux::SpeedFlux, faceFluxLR::FaceFluxLR)

   (;nI, nJ, nK, nVar) = param
   (;LState_XV, RState_XV, LState_YV, RState_YV, LState_ZV, RState_ZV) = faceValue
   (;Flux_XV, Flux_YV, Flux_ZV) = faceFlux
   (;Cmax_XF, Cmax_YF, Cmax_ZF) = speedFlux

   if param.Scheme == "Rusanov"
      get_speed_max!(param, faceValue, speedFlux)

      if !param.UseConservative
         @inbounds for iVar = 1:nVar, k = 1:nK, j = 1:nJ, i = 1:nI+1
            Flux_XV[i,j,k,iVar] -= 0.5*Cmax_XF[i,j,k]*(RState_XV[i,j,k,iVar] -
               LState_XV[i,j,k,iVar])
         end

         @inbounds for iVar = 1:nVar, k = 1:nK, j = 1:nJ+1, i = 1:nI
            Flux_YV[i,j,k,iVar] -= 0.5*Cmax_YF[i,j,k]*(RState_YV[i,j,k,iVar] -
               LState_YV[i,j,k,iVar])
         end

         @inbounds for iVar = 1:nVar, k = 1:nK+1, j = 1:nJ, i = 1:nI
            Flux_ZV[i,j,k,iVar] -= 0.5*Cmax_ZF[i,j,k]*(RState_ZV[i,j,k,iVar] -
               LState_ZV[i,j,k,iVar])
         end
      else
         # If I solve energy equation instead of pressure, there's
         # duplicate calculation above, even though the expression
         # looks compact. That's why I use an if-else statement
         @inbounds for iVar = Rho_:Bz_, k = 1:nK, j = 1:nJ, i = 1:nI+1
            Flux_XV[i,j,k,iVar] -= 0.5*Cmax_XF[i,j,k]*(RState_XV[i,j,k,iVar] -
               LState_XV[i,j,k,iVar])
         end

         @inbounds for k = 1:nK, j = 1:nJ, i = 1:nI+1
            uL = □(LState_XV[i,j,k,Mx_],LState_XV[i,j,k,My_],LState_XV[i,j,k,Mz_])
            uR = □(RState_XV[i,j,k,Mx_],RState_XV[i,j,k,My_],RState_XV[i,j,k,Mz_])
            bL = □(LState_XV[i,j,k,Bx_],LState_XV[i,j,k,By_],LState_XV[i,j,k,Bz_])
            bR = □(RState_XV[i,j,k,Bx_],RState_XV[i,j,k,By_],RState_XV[i,j,k,Bz_])

            Flux_XV[i,j,k,E_] -= 0.5*Cmax_XF[i,j,k]*(
               (RState_XV[i,j,k,P_]/(γ-1) + 0.5/RState_XV[i,j,k,Rho_]*uR + 0.5*bR) -
               (LState_XV[i,j,k,P_]/(γ-1) + 0.5/LState_XV[i,j,k,Rho_]*uL + 0.5*bL))
         end

         @inbounds for iVar = Rho_:Bz_, k = 1:nK, j = 1:nJ+1, i = 1:nI
            Flux_YV[i,j,k,iVar] -= 0.5*Cmax_YF[i,j,k]*(RState_YV[i,j,k,iVar] -
               LState_YV[i,j,k,iVar])
         end

         @inbounds for k = 1:nK, j = 1:nJ+1, i = 1:nI
            uL = □(LState_YV[i,j,k,Mx_],LState_YV[i,j,k,My_],LState_YV[i,j,k,Mz_])
            uR = □(RState_YV[i,j,k,Mx_],RState_YV[i,j,k,My_],RState_YV[i,j,k,Mz_])
            bL = □(LState_YV[i,j,k,Bx_],LState_YV[i,j,k,By_],LState_YV[i,j,k,Bz_])
            bR = □(RState_YV[i,j,k,Bx_],RState_YV[i,j,k,By_],RState_YV[i,j,k,Bz_])

            Flux_YV[i,j,k,E_] -= 0.5*Cmax_YF[i,j,k]*(
               (RState_YV[i,j,k,P_]/(γ-1) + 0.5/RState_YV[i,j,k,Rho_]*uR + 0.5*bR) -
               (LState_YV[i,j,k,P_]/(γ-1) + 0.5/LState_YV[i,j,k,Rho_]*uL + 0.5*bL))
         end

         @inbounds for iVar = Rho_:Bz_, k = 1:nK+1, j = 1:nJ, i = 1:nI
            Flux_ZV[i,j,k,iVar] -= 0.5*Cmax_ZF[i,j,k]*(RState_ZV[i,j,k,iVar] -
               LState_ZV[i,j,k,iVar])
         end

         @inbounds for k = 1:nK+1, j = 1:nJ, i = 1:nI
            uL = □(LState_ZV[i,j,k,Mx_],LState_ZV[i,j,k,My_],LState_ZV[i,j,k,Mz_])
            uR = □(RState_ZV[i,j,k,Mx_],RState_ZV[i,j,k,My_],RState_ZV[i,j,k,Mz_])
            bL = □(LState_ZV[i,j,k,Bx_],LState_ZV[i,j,k,By_],LState_ZV[i,j,k,Bz_])
            bR = □(RState_ZV[i,j,k,Bx_],RState_ZV[i,j,k,By_],RState_ZV[i,j,k,Bz_])

            Flux_ZV[i,j,k,E_] -= 0.5*Cmax_ZF[i,j,k]*(
               (RState_ZV[i,j,k,P_]/(γ-1) + 0.5/RState_ZV[i,j,k,Rho_]*uR + 0.5*bR) -
               (LState_ZV[i,j,k,P_]/(γ-1) + 0.5/LState_ZV[i,j,k,Rho_]*uL + 0.5*bL))
         end

      end
   elseif param.Scheme == "HLLE"
      Cmin_XF = speedFlux.Cmin_XF
      Cmin_YF = speedFlux.Cmin_YF
      Cmin_ZF = speedFlux.Cmin_ZF

      LFlux_XV = faceFluxLR.LFlux_XV
      RFlux_XV = faceFluxLR.RFlux_XV
      LFlux_YV = faceFluxLR.LFlux_YV
      RFlux_YV = faceFluxLR.RFlux_YV
      LFlux_ZV = faceFluxLR.LFlux_ZV
      RFlux_ZV = faceFluxLR.RFlux_ZV

      get_speed_maxmin!(param, faceValue, speedFlux)

      if ~param.UseConservative
         @inbounds for iVar = 1:nVar, k = 1:nK, j = 1:nJ, i = 1:nI+1
            Flux_XV[i,j,k,iVar] -= 0.5*(Cmax_XF[i,j,k] + Cmin_XF[i,j,k])/
               (Cmax_XF[i,j,k] - Cmin_XF[i,j,k])*
               (RFlux_XV[i,j,k,iVar] - LFlux_XV[i,j,k,iVar]) -
               Cmax_XF[i,j,k]*Cmin_XF[i,j,k]/(Cmax_XF[i,j,k] - Cmin_XF[i,j,k])*
               (RState_XV[i,j,k,iVar] - LState_XV[i,j,k,iVar])
         end

         @inbounds for iVar = 1:nVar, k = 1:nK, j = 1:nJ+1, i = 1:nI
            Flux_YV[i,j,k,iVar] -= 0.5*(Cmax_YF[i,j,k] + Cmin_YF[i,j,k])/
               (Cmax_YF[i,j,k] - Cmin_YF[i,j,k])*
               (RFlux_YV[i,j,k,iVar] - LFlux_YV[i,j,k,iVar]) -
               Cmax_YF[i,j,k]*Cmin_YF[i,j,k]/(Cmax_YF[i,j,k] - Cmin_YF[i,j,k])*
               (RState_YV[i,j,k,iVar] - LState_YV[i,j,k,iVar])
         end

         @inbounds for iVar = 1:nVar, k = 1:nK+1, j = 1:nJ, i = 1:nI
            Flux_ZV[i,j,k,iVar] -= 0.5*(Cmax_ZF[i,j,k] + Cmin_ZF[i,j,k])/
               (Cmax_ZF[i,j,k] - Cmin_ZF[i,j,k])*
               (RFlux_ZV[i,j,k,iVar] - LFlux_ZV[i,j,k,iVar]) -
               Cmax_ZF[i,j,k]*Cmin_ZF[i,j,k]/(Cmax_ZF[i,j,k] - Cmin_ZF[i,j,k])*
               (RState_ZV[i,j,k,iVar] - LState_ZV[i,j,k,iVar])
         end
      else
         # If I solve energy equation instead of pressure, there's
         # duplicate calculation above, even though the expression
         # looks compact. That's why I use an if-else statement.

         @inbounds for iVar = Rho_:Bz_, k = 1:nK, j = 1:nJ, i = 1:nI+1
            Flux_XV[i,j,k,iVar] -= 0.5*(Cmax_XF[i,j,k] + Cmin_XF[i,j,k])/
               (Cmax_XF[i,j,k] - Cmin_XF[i,j,k])*
               (RFlux_XV[i,j,k,iVar] - LFlux_XV[i,j,k,iVar]) -
               Cmax_XF[i,j,k]*Cmin_XF[i,j,k]/(Cmax_XF[i,j,k] - Cmin_XF[i,j,k])*
               (RState_XV[i,j,k,iVar] - LState_XV[i,j,k,iVar])
         end

         @inbounds for k = 1:nK, j = 1:nJ, i = 1:nI+1
            uL = □(LState_XV[i,j,k,Mx_],LState_XV[i,j,k,My_],LState_XV[i,j,k,Mz_])
            uR = □(RState_XV[i,j,k,Mx_],RState_XV[i,j,k,My_],RState_XV[i,j,k,Mz_])
            bL = □(LState_XV[i,j,k,Bx_],LState_XV[i,j,k,By_],LState_XV[i,j,k,Bz_])
            bR = □(RState_XV[i,j,k,Bx_],RState_XV[i,j,k,By_],RState_XV[i,j,k,Bz_])

            Flux_XV[i,j,k,E_] -= 0.5*(Cmax_XF[i,j,k] + Cmin_XF[i,j,k])/
               (Cmax_XF[i,j,k] - Cmin_XF[i,j,k])*(RFlux_XV[i,j,k,E_] - LFlux_XV[i,j,k,E_]) -
               Cmax_XF[i,j,k]*Cmin_XF[i,j,k] / (Cmax_XF[i,j,k] - Cmin_XF[i,j,k])*(
               (RState_XV[i,j,k,P_] / (γ-1) + 0.5/RState_XV[i,j,k,Rho_]*uR + 0.5*bR) -
               (LState_XV[i,j,k,P_] / (γ-1) + 0.5/LState_XV[i,j,k,Rho_]*uL + 0.5*bL) )
         end

         @inbounds for iVar = Rho_:Bz_, k = 1:nK, j = 1:nJ+1, i = 1:nI
            Flux_YV[i,j,k,iVar] -= 0.5*(Cmax_YF[i,j,k] + Cmin_YF[i,j,k])/
               (Cmax_YF[i,j,k] - Cmin_YF[i,j,k])*
               (RFlux_YV[i,j,k,iVar] - LFlux_YV[i,j,k,iVar]) -
               Cmax_YF[i,j,k]*Cmin_YF[i,j,k] / (Cmax_YF[i,j,k] - Cmin_YF[i,j,k])*
               (RState_YV[i,j,k,iVar] - LState_YV[i,j,k,iVar])
         end

         @inbounds for k = 1:nK, j = 1:nJ+1, i = 1:nI
            uL = □(LState_YV[i,j,k,Mx_],LState_YV[i,j,k,My_],LState_YV[i,j,k,Mz_])
            uR = □(RState_YV[i,j,k,Mx_],RState_YV[i,j,k,My_],RState_YV[i,j,k,Mz_])
            bL = □(LState_YV[i,j,k,Bx_],LState_YV[i,j,k,By_],LState_YV[i,j,k,Bz_])
            bR = □(RState_YV[i,j,k,Bx_],RState_YV[i,j,k,By_],RState_YV[i,j,k,Bz_])

            Flux_YV[i,j,k,E_] -= 0.5*(Cmax_YF[i,j,k] + Cmin_YF[i,j,k])/
               (Cmax_YF[i,j,k] - Cmin_YF[i,j,k])*
               (RFlux_YV[i,j,k,E_] - LFlux_YV[i,j,k,E_]) -
               Cmax_YF[i,j,k]*Cmin_YF[i,j,k]/(Cmax_YF[i,j,k] - Cmin_YF[i,j,k])*(
               (RState_YV[i,j,k,P_] / (γ-1) + 0.5/RState_YV[i,j,k,Rho_]*uR + 0.5*bR) -
               (LState_YV[i,j,k,P_] / (γ-1) + 0.5/LState_YV[i,j,k,Rho_]*uL + 0.5*bL) )
         end

         @inbounds for iVar = Rho_:Bz_, k = 1:nK+1, j = 1:nJ, i = 1:nI
            Flux_ZV[i,j,k,iVar] -= 0.5*(Cmax_ZF[i,j,k] + Cmin_ZF[i,j,k])/
               (Cmax_ZF[i,j,k] - Cmin_ZF[i,j,k])*
               (RFlux_ZV[i,j,k,iVar] - LFlux_ZV[i,j,k,iVar]) -
               Cmax_ZF[i,j,k]*Cmin_ZF[i,j,k]/(Cmax_ZF[i,j,k] - Cmin_ZF[i,j,k])*
               (RState_ZV[i,j,k,iVar] - LState_ZV[i,j,k,iVar])
         end

         @inbounds for k = 1:nK+1, j = 1:nJ, i = 1:nI
            uL = □(LState_ZV[i,j,k,Mx_],LState_ZV[i,j,k,My_],LState_ZV[i,j,k,Mz_])
            uR = □(RState_ZV[i,j,k,Mx_],RState_ZV[i,j,k,My_],RState_ZV[i,j,k,Mz_])
            bL = □(LState_ZV[i,j,k,Bx_],LState_ZV[i,j,k,By_],LState_ZV[i,j,k,Bz_])
            bR = □(RState_ZV[i,j,k,Bx_],RState_ZV[i,j,k,By_],RState_ZV[i,j,k,Bz_])

            Flux_ZV[i,j,k,E_] -= 0.5*(Cmax_ZF[i,j,k] + Cmin_ZF[i,j,k])/
               (Cmax_ZF[i,j,k] - Cmin_ZF[i,j,k])*
               (RFlux_ZV[i,j,k,E_] - LFlux_ZV[i,j,k,E_]) -
               Cmax_ZF[i,j,k]*Cmin_ZF[i,j,k]/(Cmax_ZF[i,j,k] - Cmin_ZF[i,j,k])*(
               (RState_ZV[i,j,k,P_] / (γ-1) + 0.5/RState_ZV[i,j,k,Rho_]*uR + 0.5*bR) -
               (LState_ZV[i,j,k,P_] / (γ-1) + 0.5/LState_ZV[i,j,k,Rho_]*uL + 0.5*bL))
         end
      end

   end
   return
end

"Calculate the maximum speed in each direction."
function get_speed_max!(param::Param,faceValue::FaceState,speedFlux::SpeedFlux)
   (;nI, nJ, nK) = param
   (;LS_XV, RS_XV, LS_YV, RS_YV, LS_ZV, RS_ZV) = faceValue
   (;Cmax_XF, Cmax_YF, Cmax_ZF) = speedFlux

   @inbounds for k = 1:nK, j = 1:nJ, i = 1:nI+1
      Cs2_XF = γ*(LS_XV[i,j,k,P_] + RS_XV[i,j,k,P_]) /
         (LS_XV[i,j,k,Rho_] + RS_XV[i,j,k,Rho_])
      Ca2_XF = □(LS_XV[i,j,k,Bx_] + RS_XV[i,j,k,Bx_],
         LS_XV[i,j,k,By_] + RS_XV[i,j,k,By_],
         LS_XV[i,j,k,Bz_] + RS_XV[i,j,k,Bz_]) /
         (2.0*(LS_XV[i,j,k,Rho_] + RS_XV[i,j,k,Rho_]))
      Can2_XF = (LS_XV[i,j,k,Bx_] + RS_XV[i,j,k,Bx_])^2 /
         (2.0*(LS_XV[i,j,k,Rho_] + RS_XV[i,j,k,Rho_]))

      Cmax_XF[i,j,k] = 0.5 * abs(LS_XV[i,j,k,Mx_]/LS_XV[i,j,k,Rho_] +
         RS_XV[i,j,k,Mx_]/RS_XV[i,j,k,Rho_]) + √( 0.5*(Cs2_XF + Ca2_XF +
         √((Cs2_XF + Ca2_XF)^2 - 4.0*Cs2_XF*Can2_XF)) )
   end

   @inbounds for k = 1:nK, j = 1:nJ+1, i = 1:nI
      Cs2_YF = γ*(LS_YV[i,j,k,P_] + RS_YV[i,j,k,P_]) /
         (LS_YV[i,j,k,Rho_] + RS_YV[i,j,k,Rho_])
      Ca2_YF = □(LS_YV[i,j,k,Bx_] + RS_YV[i,j,k,Bx_],
         LS_YV[i,j,k,By_] + RS_YV[i,j,k,By_],
         LS_YV[i,j,k,Bz_] + RS_YV[i,j,k,Bz_]) /
         (2.0*(LS_YV[i,j,k,Rho_] + RS_YV[i,j,k,Rho_]))
      Can2_YF = (LS_YV[i,j,k,By_] + RS_YV[i,j,k,By_])^2 /
         (2.0*(LS_YV[i,j,k,Rho_] + RS_YV[i,j,k,Rho_]))
      Cmax_YF[i,j,k] = 0.5 * abs(LS_YV[i,j,k,My_]/LS_YV[i,j,k,Rho_] +
         RS_YV[i,j,k,My_]/RS_YV[i,j,k,Rho_]) + √( 0.5*(Cs2_YF + Ca2_YF +
         √((Cs2_YF + Ca2_YF)^2 - 4.0*Cs2_YF*Can2_YF)) )
   end

   @inbounds for k = 1:nK+1, j = 1:nJ, i = 1:nI
      Cs2_ZF = γ*(LS_ZV[i,j,k,P_] + RS_ZV[i,j,k,P_]) /
         (LS_ZV[i,j,k,Rho_] + RS_ZV[i,j,k,Rho_])
      Ca2_ZF = □(LS_ZV[i,j,k,Bx_] + RS_ZV[i,j,k,Bx_],
         LS_ZV[i,j,k,By_] + RS_ZV[i,j,k,By_],
         LS_ZV[i,j,k,Bz_] + RS_ZV[i,j,k,Bz_]) /
         (2.0*(LS_ZV[i,j,k,Rho_] + RS_ZV[i,j,k,Rho_]))
      Can2_ZF = (LS_ZV[i,j,k,Bz_] + RS_ZV[i,j,k,Bz_])^2 /
         (2.0*(LS_ZV[i,j,k,Rho_] + RS_ZV[i,j,k,Rho_]))
      Cmax_ZF[i,j,k] = 0.5 * abs(LS_ZV[i,j,k,Mz_]/LS_ZV[i,j,k,Rho_] +
         RS_ZV[i,j,k,Mz_]/RS_ZV[i,j,k,Rho_]) + √( 0.5*(Cs2_ZF + Ca2_ZF +
         √((Cs2_ZF + Ca2_ZF)^2 - 4.0*Cs2_ZF*Can2_ZF)) )
   end
   return
end

function get_speed_maxmin!(param::Param, faceValue::FaceState, speedFlux::SpeedFluxMinMax)
   (;nI, nJ, nK) = param
   # Aliases
   LS_XV, RS_XV = faceValue.LState_XV, faceValue.RState_XV
   LS_YV, RS_YV = faceValue.LState_YV, faceValue.RState_YV
   LS_ZV, RS_ZV = faceValue.LState_ZV, faceValue.RState_ZV
   (;Cmax_XF, Cmax_YF, Cmax_ZF, Cmin_XF, Cmin_YF, Cmin_ZF) = speedFlux

   # There must be better ways to do this: LS_XV and RS_XV is just a shift of state_GV,
   # so some repetitive calculation can be avoided!

   @inbounds for k = 1:nK, j = 1:nJ, i = 1:nI+1
      Cs2_LXF = γ*LS_XV[i,j,k,P_]/LS_XV[i,j,k,Rho_]
      Ca2_LXF = □(LS_XV[i,j,k,Bx_], LS_XV[i,j,k,By_], LS_XV[i,j,k,Bz_]) / LS_XV[i,j,k,Rho_]

      Can2_LXF = LS_XV[i,j,k,Bx_]^2 / LS_XV[i,j,k,Rho_]
      u_LXF = LS_XV[i,j,k,Mx_]/LS_XV[i,j,k,Rho_]
      c_LXF = √( 0.5*(Cs2_LXF + Ca2_LXF + √((Cs2_LXF + Ca2_LXF)^2 - 4.0*Cs2_LXF*Can2_LXF)) )

      Cs2_RXF = γ*RS_XV[i,j,k,P_]/RS_XV[i,j,k,Rho_]
      Ca2_RXF = □(RS_XV[i,j,k,Bx_], RS_XV[i,j,k,By_], RS_XV[i,j,k,Bz_]) / RS_XV[i,j,k,Rho_]
      Can2_RXF = RS_XV[i,j,k,Bx_]^2 / RS_XV[i,j,k,Rho_]
      u_RXF = RS_XV[i,j,k,Mx_]/RS_XV[i,j,k,Rho_]
      c_RXF = √( 0.5*(Cs2_RXF + Ca2_RXF + √((Cs2_RXF + Ca2_RXF)^2 - 4.0*Cs2_RXF*Can2_RXF)) )

      Cmax_XF[i,j,k] = max(0, u_LXF+c_LXF, u_RXF+c_RXF)
      Cmin_XF[i,j,k] = min(0, u_LXF-c_LXF, u_RXF-c_RXF)
   end

   @inbounds for k = 1:nK, j = 1:nJ+1, i = 1:nI
      Cs2_LYF = γ*LS_YV[i,j,k,P_]/LS_YV[i,j,k,Rho_]
      Ca2_LYF = □(LS_YV[i,j,k,Bx_], LS_YV[i,j,k,By_], LS_YV[i,j,k,Bz_]) / LS_YV[i,j,k,Rho_]
      Can2_LYF = LS_YV[i,j,k,Bx_]^2 / LS_YV[i,j,k,Rho_]
      u_LYF = LS_YV[i,j,k,My_]/LS_YV[i,j,k,Rho_]
      c_LYF = √( 0.5*(Cs2_LYF + Ca2_LYF + √((Cs2_LYF + Ca2_LYF)^2 - 4.0*Cs2_LYF*Can2_LYF)) )

      Cs2_RYF = γ*RS_YV[i,j,k,P_]/RS_YV[i,j,k,Rho_]
      Ca2_RYF = □(RS_YV[i,j,k,Bx_], RS_YV[i,j,k,By_], RS_YV[i,j,k,Bz_]) / RS_YV[i,j,k,Rho_]
      Can2_RYF = RS_YV[i,j,k,Bx_]^2 / RS_YV[i,j,k,Rho_]
      u_RYF = RS_YV[i,j,k,Mx_]/RS_YV[i,j,k,Rho_]
      c_RYF = √( 0.5*(Cs2_RYF + Ca2_RYF + √((Cs2_RYF + Ca2_RYF)^2 - 4.0*Cs2_RYF*Can2_RYF)) )

      Cmax_YF[i,j,k] = max(0, u_LYF+c_LYF, u_RYF+c_RYF)
      Cmin_YF[i,j,k] = min(0, u_LYF-c_LYF, u_RYF-c_RYF)
   end

   @inbounds for k = 1:nK+1, j = 1:nJ, i = 1:nI
      Cs2_LZF = γ*LS_ZV[i,j,k,P_]/LS_ZV[i,j,k,Rho_]
      Ca2_LZF = □(LS_ZV[i,j,k,Bx_], LS_ZV[i,j,k,By_], LS_ZV[i,j,k,Bz_]) / LS_ZV[i,j,k,Rho_]
      Can2_LZF = LS_ZV[i,j,k,Bx_]^2 / LS_ZV[i,j,k,Rho_]
      u_LZF = LS_ZV[i,j,k,Mx_]/LS_ZV[i,j,k,Rho_]
      c_LZF = √( 0.5*(Cs2_LZF + Ca2_LZF + √((Cs2_LZF + Ca2_LZF)^2 - 4.0*Cs2_LZF*Can2_LZF)) )

      Cs2_RZF = γ*RS_ZV[i,j,k,P_]/RS_ZV[i,j,k,Rho_]
      Ca2_RZF = □(RS_ZV[i,j,k,Bx_], RS_ZV[i,j,k,By_], RS_ZV[i,j,k,Bz_]) / RS_ZV[i,j,k,Rho_]
      Can2_RZF = RS_ZV[i,j,k,Bx_]^2 / RS_ZV[i,j,k,Rho_]
      u_RZF = RS_ZV[i,j,k,Mz_]/RS_ZV[i,j,k,Rho_]
      c_RZF = √( 0.5*(Cs2_RZF + Ca2_RZF + √((Cs2_RZF + Ca2_RZF)^2 - 4.0*Cs2_RZF*Can2_RZF)) )

      Cmax_ZF[i,j,k] = max(0, u_LZF+c_LZF, u_RZF+c_RZF)
      Cmin_ZF[i,j,k] = min(0, u_LZF-c_LZF, u_RZF-c_RZF)
   end
   return
end


text:
using Revise
using Test
using ModelExploration
using Catlab.CategoricalAlgebra, Catlab.Present, Catlab.Theories
#include("../src/Petri.jl")
include("../src/ModelSelection.jl")
using AlgebraicPetri: LabelledPetriNetUntyped

infectious_type_ = LabelledPetriNet([:Pop],
  :interaction=>((:Pop, :Pop)=>(:Pop, :Pop)),
  :t_infection=>(:Pop=>:Pop),
  :t_strata=>(:Pop=>:Pop))
#Graph(infectious_type)
infectious_type = map(infectious_type_, Name=name->nothing)

IO_help(i::Int) = let d = Dict([j=>j%2+1 for j in 1:2*i]); (I=d,O=d) end
make_slice(p::LabelledPetriNet, n::NamedTuple) =
    Slice{ACSetTransformation}(
        homomorphism(p, infectious_type; initial=n,
                       type_components=(Name=x->nothing,)))

SIR = LabelledPetriNet([:S, :I, :R],
  :inf => ((:S, :I)=>(:I, :I)),
  :rec => (:I=>:R),
  :id => (:S => :S),
  :id => (:I => :I),
  :id => (:R => :R)
)
SIR_type = make_slice(SIR, merge((T=[1,2,3,3,3],), IO_help(1)))
#
SIR_travel_ban = LabelledPetriNet([:S, :I, :R],
  :inf => ((:S, :I)=>(:I, :I)),
  :rec => (:I=>:R),
  :id => (:S => :S),
  :id => (:R => :R)
)
SIR_tb_type = make_slice(SIR_travel_ban, merge((T=[1,2,3,3],), IO_help(1)))
#
SIS = LabelledPetriNet([:S, :I],
  :inf => ((:S, :I)=>(:I, :I)),
  :rec => (:I=>:S),
  :id => (:S => :S),
  :id => (:I => :I),
)
SIS_type = make_slice(SIS, merge((T=[1,2,3,3],),IO_help(1)))
#
SVIIvR = LabelledPetriNet([:S, :I, :R, :Iv, :V],
  :inf => ((:S, :I)=>(:I, :I)),
  :inf => ((:V,:I)=>(:I,:I)),
  :inf => (:S, :Iv) => (:Iv,:Iv),
  :inf => ((:V, :Iv) => (:Iv,:Iv)),
  :rec => (:I=>:R),
  :rec => (:Iv=>:R),
  :vax => (:S=>:V),
  :id => (:S => :S),
  :id => (:I => :I),
  :id => (:V => :V),
  :id => (:Iv => :Iv),
  :id => (:R => :R)
)

SVIIvR_type = make_slice(SVIIvR, merge((T=[1,1,1,1,2,2,2,3,3,3,3,3],),IO_help(4)))
#
quarantine = LabelledPetriNet([:Q, :not_Q],
    :interaction => ((:not_Q, :not_Q) => (:not_Q, :not_Q)),
    :enter_quarantine => (:not_Q => :Q),
    :exit_quarantine => (:Q => :not_Q),
    :id => (:Q => :Q),
    :id => (:not_Q => :not_Q),
)
quarantine_type = make_slice(quarantine, (T=[1, 3, 3, 2, 2], I=Dict(1=>1,2=>2), O=Dict(1=>1,2=>2)))
#
age_stratification = LabelledPetriNet([:Child, :Adult],
    :interaction => ((:Child, :Child) => (:Child, :Child)),
    :interaction => ((:Adult, :Adult) => (:Adult, :Adult)),
    :interaction => ((:Child, :Adult) => (:Child, :Adult)),
    :id => (:Child => :Child),
    :id => (:Adult => :Adult),
)
age_s_type = make_slice(age_stratification, merge((T=[1,1,1,2,2],),IO_help(3)))

#
flux_metapopulation = LabelledPetriNet([:Patch1, :Patch2],
    :interaction => ((:Patch1, :Patch1) => (:Patch1, :Patch1)),
    :interaction => ((:Patch2, :Patch2) => (:Patch2, :Patch2)),
	:travel => (:Patch1 => :Patch2),
    :travel => (:Patch2 => :Patch1),
    :id => (:Patch1 => :Patch1),
    :id => (:Patch2 => :Patch2),
)
flux_m_type = make_slice(flux_metapopulation, merge((T=[1,1,3,3,2,2],), IO_help(2)))
simple_trip =  LabelledPetriNet([:P11, :P21, :P12, :P22],
    :interaction11_11 => ((:P11, :P11) => (:P11, :P11)),
    :interaction12_12 => ((:P12, :P12) => (:P12, :P12)),
    :interaction21_21 => ((:P21, :P21) => (:P21, :P21)),
    :interaction22_22 => ((:P22, :P22) => (:P22, :P22)),
    :interaction11_21 => ((:P11, :P21) => (:P11, :P21)),
    :interaction22_12 => ((:P22, :P12) => (:P22, :P12)),
    :interaction12_22 => ((:P12, :P22) => (:P12, :P22)),
    :interaction21_11 => ((:P21, :P11) => (:P21, :P11)),
	:travel11_12 => (:P11 => :P12),
	:travel12_11 => (:P12 => :P11),
	:travel21_22 => (:P21 => :P22),
	:travel22_21 => (:P22 => :P21),
    :id => (:P11 => :P11),
    :id => (:P21 => :P21),
    :id => (:P12 => :P12),
    :id => (:P22 => :P22),
)
simple_t_type = make_slice(simple_trip,
    merge((T=[1,1,1,1,1,1,1,1,3,3,3,3,2,2,2,2],), IO_help(8)))

@present ThFour(FreeCategory) begin
    (X1,X2,X3,X4)::Ob
end
@present ThOne(FreeCategory) begin
  X::Ob
end

const ACSetCat{S} = TypeCat{S, ACSetTransformation}
const ACSetCatSlice{S} = TypeCat{S, ACSetTransformation}
const Petri = ACSetCat{LabelledPetriNet}
const PetriHom = SliceCat{ACSetTransformation}

One, Four = FinCat.([ThOne,ThFour])

to_diag(x) = Diagram(FinDomFunctor(x, nothing, Four, PetriHom()))
to_slicehom(x) = SliceDiagHom(Literal(to_diag(x)))

diag_disease = to_slicehom(Dict(:X1=>SIR_type,:X2=>SIS_type,
                               :X3=>SVIIvR_type,:X4=>SIR_tb_type));
diag_strata = to_slicehom(Dict(:X1=>quarantine_type, :X2=>age_s_type,
                               :X3=>flux_m_type, :X4=>simple_t_type));

pb = PullbackSpace(to_model_hom(diag_disease), to_model_hom(diag_strata));
upb = unfold(pb);
aupb = apex(upb);
@test length(ob_generators(dom(diagram(aupb)))) == 16

true_model = ob_map(aupb, Symbol("(X3, X1)"))
true_rxn = MakeReactionSystem(true_model)

p_real = vcat(repeat([1e-4], 14), repeat([0.01], 6))
u0 = [0.0,0.0,0.0,0.0,0.0,999.0,1.0,0.0,0.0,0.0] #|> gpu
tspan = (0.0,250.0)

#sample_data, sample_times, prob_real, sol_real = generate_data(true_model, p_real, u0, tspan, 50)=#

#=true_model = SIR
true_rxn = MakeReactionSystem(SIR)
p_real = [.1/1000, .01]
u0 = [999.0, 1.0, 0.0]
tspan = (0.0,250.0)=#

sample_data, sample_times, prob_real, sol_real = generate_data(true_model, p_real, u0, tspan, 50)

plt = plot(sol_real, lw=2, label=reshape(map(string, true_model[:, :sname]), 1, ns(true_model)))
plot!(sample_times, sample_data, seriestype=:scatter, label="")
display(plt)

#small training example to precompile everything
#full_train(SIR, [999.0,1.0,0.0], tspan, sample_data, sample_times);

#assign inital concentrations to models TODO: be more smart about this
models = [
  ob_map(aupb, Symbol("(X1, X1)")) => [0.0, 0.0, 0.0, 999.0, 1.0, 0.0],
  ob_map(aupb, Symbol("(X1, X2)")) => [500.0, 0.0, 0.0, 499.0, 1.0, 0.0],
  # ob_map(aupb, Symbol("(X1, X3)")) => [500.0, 1.0, 0.0, 499.0, 0.0, 0.0],
  # ob_map(aupb, Symbol("(X1, X4)")) => [250.0, 1.0, 0.0, 250.0, 0.0, 0.0, 250.0, 0.0, 0.0, 249.0, 0.0, 0.0],
  # ob_map(aupb, Symbol("(X2, X1)")) => [0.0, 0.0, 999.0, 1.0],
  # ob_map(aupb, Symbol("(X2, X2)")) => [500.0, 0.0, 499.0, 1.0],
  # ob_map(aupb, Symbol("(X2, X3)")) => [500.0, 1.0, 499.0, 0.0],
  # ob_map(aupb, Symbol("(X2, X4)")) => [250.0, 1.0, 250.0, 0.0, 250.0, 0.0, 249.0, 0.0],
  ob_map(aupb, Symbol("(X3, X1)")) => [0.0, 0.0, 0.0, 0.0, 0.0, 999.0, 1.0, 0.0, 0.0, 0.0],
  ob_map(aupb, Symbol("(X3, X2)")) => [500.0, 0.0, 0.0, 0.0, 0.0, 499.0, 1.0, 0.0, 0.0, 0.0],
  # ob_map(aupb, Symbol("(X3, X3)")) => [500.0, 1.0, 0.0, 0.0, 0.0, 499.0, 0.0, 0.0, 0.0, 0.0],
  #ob_map(aupb, Symbol("(X3, X4)")) => [250.0, 1.0, 0.0, 0.0, 0.0,
  #                                     250.0, 0.0, 0.0, 0.0, 0.0,
  #                                     250.0, 0.0, 0.0, 0.0, 0.0,
  #                                     249.0, 0.0, 0.0, 0.0, 0.0],
  # ob_map(aupb, Symbol("(X4, X1)")) => [0.0, 0.0, 0.0, 999.0, 1.0, 0.0],
  # ob_map(aupb, Symbol("(X4, X2)")) => [500.0, 0.0, 0.0, 499.0, 1.0, 0.0],
  # ob_map(aupb, Symbol("(X4, X3)")) => [500.0, 1.0, 0.0, 499.0, 0.0, 0.0],
  # ob_map(aupb, Symbol("(X4, X4)")) => [250.0, 1.0, 0.0, 250.0, 0.0, 0.0, 250.0, 0.0, 0.0, 249.0, 0.0, 0.0]
]

function explore(models, tspan, sample_data, sample_times)
  losses = zeros(length(models))
  Threads.@threads for i in 1:length(models)
    model, u0 = models[i]
    sol, loss = full_train(model, u0, tspan, sample_data, sample_times)
    losses[i] = loss
    println(loss)
    #=plt = plot(sol, lw=2, label=reshape(map(string, model[:, :sname]), 1, ns(model)))
    plot!(sample_times, sample_data, seriestype=:scatter, label="")
    display(plt)=#
  end
  return losses
end

#=println(min_loss)
plt = plot(best_sol, lw=2, label=reshape(map(string, best_model[:, :sname]), 1, ns(best_model)))
plot!(sample_times, sample_data, seriestype=:scatter, label="")
display(plt)=#

# Pushout example
Infect = LabelledPetriNet([:I],)
Death = LabelledPetriNet([:I,:D],     :death => (:I => :D))
ISIR = only(homomorphisms(Infect, SIR))
IID = only(homomorphisms(Infect, Death))

I = Diagram(FinDomFunctor(Dict(:X=>Infect), nothing, One, Petri()))
D = FinDomFunctor(Dict(:X1=>SIR,:X2=>SIS,:X3=>SVIIvR,:X4=>SIR_travel_ban),
                  nothing, Four,Petri())
ID = LitModelHom(DiagramHom(FinFunctor(Dict(:X=>:X1), nothing, One, Four), Dict(:X=>ISIR),
                  I, Diagram(D)))
Q = Diagram(FinDomFunctor(Dict(:X=>Death), nothing, One, Petri()))
IQ = LitModelHom(DiagramHom(id(One),Dict(:X=>IID),I,Q))

#po = PushoutSpace(ID,IQ) # fails b/c no chase of ACSets only C-sets
                          # and no FinCats with Attr

# a pushout over a product space: (A+B*C) = A+B * A+C

ext: jl


text:
# This file is a part of SimilaritySearch.jl
# License is Apache 2.0: https://www.apache.org/licenses/LICENSE-2.0.txt

using SimilaritySearch
using LinearAlgebra
using Test

function test_vectors(create_index, dist::Function, ksearch, nick)
    @testset "indexing vectors with $nick and $dist" begin
        n = 1000 # number of items in the dataset
        m = 100  # number of queries
        dim = 3  # vector's dimension

        db = [rand(Float32, dim) |> normalize! for i in 1:n]
        queries = [rand(Float32, dim) |> normalize! for i in 1:m]

        index = create_index(db)
        optimize!(index, dist, recall=0.9, k=10)
        perf = Performance(dist, index.db, queries, expected_k=10)
        p = probe(perf, index, dist)
        @show dist, p
        @test p.recall > 0.8

        @info "adding more items"
        for item in queries
            push!(index, dist, item)
        end
        perf = Performance(dist, index.db, queries, expected_k=1)
        p = probe(perf, index, dist)
        @show dist, p
        @test p.recall > 0.999
        return p
    end
end

function test_sequences(create_index, dist::Function, ksearch, nick)
    @testset "indexing sequences with $nick and $dist" begin
        n = 1000 # number of items in the dataset
        m = 100  # number of queries
        dim = 5  # the length of sequences
        V = collect(1:10)  # vocabulary of the sequences

        function create_item()
            s = rand(V, dim)
            if dist == jaccard_distance || dist == dice_distance || dist == intersection_distance
                sort!(s)
                s = unique(s)
            end

            return s
        end

        db = [create_item() for i in 1:n]
        queries = [create_item() for i in 1:m]

        @info "inserting items into the index"
        index = create_index(db)
        # optimize!(index, recall=0.9, k=10)
        perf = Performance(dist, index.db, queries, expected_k=10)
        p = probe(perf, index, dist)
        @show dist, p
        @test p.recall > 0.1  ## Performance object tests object identifiers, but sequence distances have a lot of distance collisions

        # for item in queries
        #     push!(index, dist, item)
        # end
        # perf = Performance(dist, index.db, queries, expected_k=1)
        # p = probe(perf, index, dist)
        # @show dist, p
        # @test p.recall > 0.999
        # return p
    end
end

@testset "indexing vectors" begin
    # NOTE: The following algorithms are complex enough to say we are testing it doesn't have syntax errors, a more grained test functions are required
    ksearch = 10
    σ = 127
    κ = 3

    for dist in [
        l2_distance, # 1.0 -> metric, < 1.0 if dist is not a metric
        l1_distance,
        linf_distance,
        lp_distance(3),
        lp_distance(0.5),
        angle_distance
    ]
        p = test_vectors((db) -> fit(Knr, dist, db, numrefs=σ, k=κ), dist, ksearch, "KNR")
    end
end

@testset "indexing sequences" begin
    # NOTE: The following algorithms are complex enough to say we are testing it doesn't have syntax errors, a more grained test functions are required
    ksearch = 10
    σ = 127
    κ = 3

    # metric distances should achieve recall=1 (perhaps lesser because of numerical inestability)
    for dist in [
        jaccard_distance,
        dice_distance,
        intersection_distance,
        common_prefix_distance,
        levenshtein_distance,
        lcs_distance,
        hamming_distance,
    ]
        p = test_sequences((db) -> fit(Knr, dist, db, numrefs=σ, k=κ), dist, ksearch, "KNR")
    end
end


text:
function _minerva_parse_species_bound_type_sentence(sentence::Array{MinervaToken,1};
    logger::Union{Nothing, SimpleLogger} = nothing)::VLResult

    # flat sentence in case of error -
    flat_sentence_buffer = ""
    number_of_tokens = length(sentence)
    for (index,token) in enumerate(sentence)

        # lexeme -
        lexeme = token.lexeme
        flat_sentence_buffer *= "$(lexeme)"
        if (index<number_of_tokens)
            flat_sentence_buffer *=" "
        end
    end

    # grab the a token -
    first_token = pop!(sentence)
    if (isa(first_token.type, BIOLOGICAL_SYMBOL) == true)
        return _start_bound_assignment_recursive_descent(sentence)
    else

        # formulate and throw error excpetion -
        error_message = "Error in sentence: $(flat_sentence_buffer). Expected biological symbol but found $(first_token.lexeme)"
        error_object = ErrorException(error_message)
        return VLResult(error_object)
    end
end

# --- PRIVATE RECURSIVE DESCENT METHODS BELOW ----------------------------------------------------------------------------------- #
function _start_bound_assignment_recursive_descent(sentence::Array{MinervaToken,1})::VLResult

    # flat sentence in case of error -
    flat_sentence_buffer = ""
    number_of_tokens = length(sentence)
    for (index,token) in enumerate(sentence)

        # lexeme -
        lexeme = token.lexeme
        flat_sentence_buffer *= "$(lexeme)"
        if (index<number_of_tokens)
            flat_sentence_buffer *=" "
        end
    end

    # grab the a token -
    next_token = pop!(sentence)
    if (isa(next_token.type,IS) == true)
        return _is_token_species_bound(sentence)
    elseif (isa(next_token.type, BOUND) == true || isa(next_token.type, UNBOUND) == true || isa(next_token.type, SOURCE) == true || isa(next_token.type, SINK) == true)
        return _bound_type_token_species_bound(sentence)
    else

        # formulate and throw error excpetion -
        error_message = "Error in sentence: $(flat_sentence_buffer). Expected a species bound type or 'is' but found $(next_token.lexeme) instead"
        error_object = ErrorException(error_message)
        return VLResult(error_object)
    end
end

function _is_token_species_bound(sentence::Array{MinervaToken,1})::VLResult

    # flat sentence in case of error -
    flat_sentence_buffer = ""
    number_of_tokens = length(sentence)
    for (index,token) in enumerate(sentence)

        # lexeme -
        lexeme = token.lexeme
        flat_sentence_buffer *= "$(lexeme)"
        if (index<number_of_tokens)
            flat_sentence_buffer *=" "
        end
    end

    # grab the next token -
    next_token = pop!(sentence)
    if (isa(next_token.type, A) == true)
        return _a_token_species_bound(sentence)
    elseif (isa(next_token.type, BOUND) == true || isa(next_token.type, UNBOUND) == true ||
        isa(next_token.type, SOURCE) == true || isa(next_token.type, SINK) == true)
        return _bound_type_token_species_bound(sentence)
    else
        # formulate and throw error excpetion -
        error_message = "Error in sentence: $(flat_sentence_buffer). Expected a species bound type or 'a' but found $(next_token.lexeme) instead"
        error_object = ErrorException(error_message)
        return VLResult(error_object)
    end
end

function _bound_type_token_species_bound(sentence::Array{MinervaToken,1})::VLResult

    # flat sentence in case of error -
    flat_sentence_buffer = ""
    number_of_tokens = length(sentence)
    for (index,token) in enumerate(sentence)

        # lexeme -
        lexeme = token.lexeme
        flat_sentence_buffer *= "$(lexeme)"
        if (index<number_of_tokens)
            flat_sentence_buffer *=" "
        end
    end

    # grab the a token -
    next_token = pop!(sentence)
    if (isa(next_token.type, SEMICOLON) == true)
        return VLResult(nothing)
    else
        # formulate and throw error excpetion -
        error_message = "Error in sentence: $(flat_sentence_buffer). Expected a ';' but found '$(next_token.lexeme)'"
        error_object = ErrorException(error_message)
        return VLResult(error_object)
    end
end

function _a_token_species_bound(sentence::Array{MinervaToken,1})::VLResult

    # flat sentence in case of error -
    flat_sentence_buffer = ""
    number_of_tokens = length(sentence)
    for (index,token) in enumerate(sentence)

        # lexeme -
        lexeme = token.lexeme
        flat_sentence_buffer *= "$(lexeme)"
        if (index<number_of_tokens)
            flat_sentence_buffer *=" "
        end
    end

    # grab the next token -
    next_token = pop!(sentence)
    if (isa(next_token.type, BOUND) == true || isa(next_token.type, UNBOUND) == true ||
        isa(next_token.type, SOURCE) == true || isa(next_token.type, SINK) == true)
        return _bound_type_token_species_bound(sentence)
    else

        # formulate and throw error excpetion -
        error_message = "Error in sentence: $(flat_sentence_buffer). Expected a bount type but found '$(next_token.lexeme)'"
        error_object = ErrorException(error_message)
        return VLResult(error_object)
    end
end
# --- PRIVATE RECURSIVE DESCENT METHODS ABOVE ----------------------------------------------------------------------------------- #

max_stars_repo_name: varnerlab/CellFreeModelGenerationKit

max_issues_repo_name: varnerlab/CellFreeModelGenerationKit

max_forks_repo_name: varnerlab/CellFreeModelGenerationKit


text:
"""
Dark matter halo models and relations.

Significant credit to Benedikt Diemer for developing [Colossus](https://bitbucket.org/bdiemer/colossus).

While much of this is inspired by Colossus; any implementation faults are my own.

The default halo mass definition is "200c", i.e. the virial radius of the halo is defined
to be the radius that encloses a density that is 200 times the critical density of the
universe.

The default cosmology (from `Slomo.CosmologyTools`) is currently set to the values
from Planck 2018 (h = 0.6766, Neff = 3.046, OmegaM = 0.3111).

# Common HaloModel methods
* [`scale_radius`](@ref): scale radius of the halo
* [`virial_radius`](@ref): virial radius of the halo
* [`virial_mass`](@ref): virial mass of the halo
* [`concentration`](@ref): halo concentration (defined as rvir / rs)

# HaloModel subtypes
* [`NFWModel`](@ref): Navarro, Frenk, and White (1997) profile
* [`CoreNFWModel`](@ref): Core-NFW (Read et al. 2016) profile
* [`GNFWModel`](@ref): generalized NFW model (i.e., αβγ where α = 1, β = 3)
* [`ABGModel`](@ref): α-β-γ profile (double power law)
* [`EinastoModel`](@ref): Einasto (1965) profile
* [`SolNFWModel`](@ref): Soliton + NFW model (Marsh & Pop 2015)
* [`SolABGModel`](@ref): Soliton + αβγ model (Wasserman et al. 2019)

# Halo relations
* [`hmcr`](@ref): halo mass-concentration relation from Dutton & Maccio (2014)
* [`shmr`](@ref): stellar-halo mass relation from Rodriguez-Puebla et al. (2017)
* [`abg_from_logshm`](@ref): αβγ scaling with log(Mstar / Mhalo) from DiCintio et al. (2014)
"""
module Halos

using Roots
using ForwardDiff

using Slomo.Utils: log_gauss
using Slomo.CosmologyTools: Ωm, ρm, ρcrit, default_cosmo
import Slomo.Models: DensityModel, mass, density, NotImplemented
using Slomo.Utils: hyp2f1, gamma, gamma_inc

export virial_radius, virial_mass, scale_radius, concentration
export NFWModel, CoreNFWModel, ABGModel, GNFWModel, EinastoModel, SolNFWModel, SolABGModel
export NFW_from_virial, CoreNFW_from_virial, ABG_from_virial, GNFW_from_virial, Einasto_from_virial
export SolNFW_from_virial, SolABG_from_virial
export hmcr, hmcr_prior, shmr, shmr_prior, abg_from_logshm

const default_mdef = "200c"

abstract type HaloModel <: DensityModel end

"""
    scale_radius(halo::HaloModel)

Return the scale radius of the halo.  The convention adopted here is that the scale radius
occurs where the logarithmic slope of the halo density profile is equal to -2.  For the
standard NFW model this is equal to `rs`, but other halo models may compute some other
value to stay consistent with this convention (e.g., the `ABGModel`).
"""
function scale_radius(halo::HaloModel)::Float64
    throw(NotImplemented("needs to be implemented for subtypes of HaloModel"))
end

"""
    Δρ_from_mdef(mdef; <keyword arguments>)

Parse the halo mass definition to get the spherical overdensity in Msun / kpc3.
If mdef is "vir", then use the spherical overdensity relation of Bryan & Norman 1998.

`mdef` should be either "vir" or of the form "[0-9]*[c|m]" where "c" refers to a multiple
of the critical density and "m" refers to a multiple of the mass density, e.g.,

    mdef in ["vir", "200c", "500c", "200m", "500m", ...]

# Arguments
- `cosmo::AbstractCosmology = default_cosmo`: cosmology under which to evaluate the overdensity
- `z::Real = 0.0`: redshift at which to evaluate the overdensity
"""
function Δρ_from_mdef(mdef; cosmo = default_cosmo, z = 0.0)
    mdef = lowercase(strip(mdef))
    if mdef == "vir"
        x = Ωm(z; cosmo = cosmo) .- 1.0
        Δ = @. 18π^2 + 82.0 * x - 39.0 * x^2
        ρ = ρcrit(z; cosmo = cosmo)
    else
        mdef_type = mdef[end]
        Δ = parse(Float64, mdef[1:end - 1])
        if mdef_type == 'm'
            ρ = ρm(z; cosmo = cosmo)
        elseif mdef_type == 'c'
            ρ = ρcrit(z; cosmo = cosmo)
        else
            throw("$mdef not recognized as a valid halo mass definition")
        end
    end
    return @. Δ * ρ
end

"""
    Rvir_from_Mvir(Mvir; <keyword arguments>)

Compute the virial radius from the virial mass, using the definition

```math
M_\\mathrm{vir} = 4\\pi \\Delta\\rho / (3 R_\\mathrm{vir}^3)
```

# Arguments
- `mdef::AbstractString = default_mdef`: halo mass definition (e.g., "200c", "vir")
- `cosmo::AbstractCosmology = default_cosmo`: cosmology under which to evaluate the overdensity
- `z::Real = 0.0`: redshift at which to evaluate the overdensity
"""
function Rvir_from_Mvir(Mvir; mdef = default_mdef, cosmo = default_cosmo, z = 0.0)
    Δρ = Δρ_from_mdef(mdef; cosmo = cosmo, z = z)
    return @. (Mvir * 3.0 / 4π / Δρ) ^ (1 / 3)
end

"""
    Mvir_from_Rvir(Rvir; <keyword arguments>)

Compute the virial mass from the virial radius, using the definition

```math
M_\\mathrm{vir} = 4\\pi \\Delta\\rho / (3 R_\\mathrm{vir}^3)
```

# Arguments
- `mdef::AbstractString = default_mdef`: halo mass definition (e.g., "200c", "vir")
- `cosmo::AbstractCosmology = default_cosmo`: cosmology under which to evaluate the overdensity
- `z::Real = 0.0`: redshift at which to evaluate the overdensity
"""
function Mvir_from_Rvir(Rvir; mdef = default_mdef, cosmo = default_cosmo, z = 0.0)
    Δρ = Δρ_from_mdef(mdef; cosmo = cosmo, z = z)
    return @. 4π / 3.0 * Rvir ^ 3 * Δρ
end

"""
    virial_radius(halo::HaloModel, <keyword arguments>)

Compute the virial radius of the halo via root finding.

# Arguments
- `mdef::AbstractString = default_mdef`: halo mass definition (e.g., "200c", "vir")
- `cosmo::AbstractCosmology = default_cosmo`: cosmology under which to evaluate the overdensity
- `z::Real = 0.0`: redshift at which to evaluate the overdensity
- `xstart::Real = 10.0`: factor of the scale radius to start the search for the root
- `rtol::Real = 1e-2`: relative error tolerance for the root finding
- `maxevals:Int = 100`: maximum number of function evalutions for the root finding
"""
function virial_radius(halo::HaloModel;
                       mdef = default_mdef,
                       cosmo = default_cosmo,
                       z = 0.0,
                       xstart = 10.0,
                       rtol = 1e-2,
                       maxevals = 100)
    Δρ = Δρ_from_mdef(mdef; cosmo = cosmo, z = z)
    # function to find roots
    f(r) = 3.0 / 4π * mass(halo, r) / r^3 - Δρ
    # derivative
    fp(r) = 3.0 * (r^-1 * density(halo, r) - r^-4 * mass(halo, r))
    rstart = xstart * scale_radius(halo)
    return fzero(f, fp, rstart; rtol = rtol, maxevals = maxevals)
end

"""
    virial_mass(halo::HaloModel; <keyword arguments>)

Compute the virial mass via root finding.

# Arguments
- `mdef::AbstractString = default_mdef`: halo mass definition (e.g., "200c", "vir")
- `cosmo::AbstractCosmology = default_cosmo`: cosmology under which to evaluate the overdensity
- `z::Real = 0.0`: redshift at which to evaluate the overdensity
- `xstart::Real = 10.0`: factor of the scale radius to start the search for the root
- `rtol::Real = 1e-2`: relative error tolerance for the root finding
- `maxevals:Int = 100`: maximum number of function evalutions for the root finding
"""
function virial_mass(halo::HaloModel;
                     mdef = default_mdef,
                     cosmo = default_cosmo,
                     z = 0.0,
                     xstart = 10.0,
                     rtol = 1e-2,
                     maxevals = 100)
    Rvir = virial_radius(halo;
                         mdef = mdef, cosmo = cosmo, z = z,
                         xstart = xstart, rtol = rtol, maxevals = maxevals)
    return Mvir_from_Rvir(Rvir, mdef = mdef, cosmo = cosmo, z = z)
end

"""
    virial_mass(Rvir::Real; <keyword arguments>)

Compute the virial mass from the virial radius.

# Arguments
- `mdef::AbstractString = default_mdef`: halo mass definition (e.g., "200c", "vir")
- `cosmo::AbstractCosmology = default_cosmo`: cosmology under which to evaluate the overdensity
- `z::Real = 0.0`: redshift at which to evaluate the overdensity
"""
function virial_mass(Rvir::Real;
                     mdef = default_mdef,
                     cosmo = default_cosmo,
                     z = 0.0)
    return Mvir_from_Rvir(Rvir, mdef = mdef, cosmo = cosmo, z = z)
end

"""
    concentration(halo::HaloModel; <keyword arguments>)

Compute the halo concentration.

# Arguments
- `mdef::AbstractString = default_mdef`: halo mass definition (e.g., "200c", "vir")
- `cosmo::AbstractCosmology = default_cosmo`: cosmology under which to evaluate the overdensity
- `z::Real = 0.0`: redshift at which to evaluate the overdensity
- `xstart::Real = 10.0`: factor of the scale radius to start the search for the root
- `rtol::Real = 1e-2`: relative error tolerance for the root finding
- `maxevals:Int = 100`: maximum number of function evalutions for the root finding
"""
function concentration(halo::HaloModel;
                       mdef = default_mdef,
                       cosmo = default_cosmo,
                       z = 0.0,
                       xstart = 10.0,
                       rtol = 1e-2,
                       maxevals = 100)
    Rvir = virial_radius(halo;
                         mdef = mdef, cosmo = cosmo, z = z,
                         xstart = xstart, rtol = rtol, maxevals = maxevals)
    rs = scale_radius(halo)
    return Rvir / rs
end

include("nfw.jl")
include("relations.jl")
include("abg.jl")
include("einasto.jl")
include("soliton.jl")

end


text:
using DataFrames, CSV, FileIO, Dates
using BenchmarkTools
using Agtor

import Agtor: run_model
import OrderedCollections: LittleDict

# Start month/day of growing season
# Represents when water allocations are announced.
const gs_start = (5, 1)

"""
Allocation callback for example.
"""
function allocation_callback!(zone, dt_i)

    # Resetting allocations for example run
    if monthday(dt_i) == gs_start
        # Example annual water allocations
        allocs = (surface_water=150.0, groundwater=40.0)

        update_available_water!(zone, allocs)
    end
end

"""
Run scenarios in sequence, and save results as they complete.

Compared to storing all results and saving in one go, this approach
is slower overall, but does not use as much memory.
"""
function example_scenario_run(data_dir::String="test/data/")::Nothing
    z1 = setup_zone(data_dir)
    z1.manager = BaseManager("test")

    scen_data = joinpath(data_dir, "scenarios", "sampled_params.csv")
    samples = DataFrame!(CSV.File(scen_data))

    @sync for (row_id, r) in enumerate(eachrow(samples))
        tmp_z = deepcopy(z1)
        update_model!(tmp_z, r)

        results = run_model(tmp_z, run_timestep!; post=allocation_callback!)

        # Save results as they complete
        @async save_results!("async_save.jld2", row_id, results)
    end
end

"""
Same example as above, but collect all results before saving.

Should be faster than the first example (for the small number of scenarios run) as all
results are kept in memory until the end.
"""
function example_batch_save(; data_dir::String="test/data/", output_fn::String="test_batch_save.jld2")::Nothing
    z1 = setup_zone(data_dir)
    z1.manager = BaseManager("test")

    scen_data = joinpath(data_dir, "scenarios", "sampled_params.csv")
    samples = DataFrame!(CSV.File(scen_data))

    results = run_scenarios!(samples, z1, run_timestep!; post=allocation_callback!)

    save_results!(output_fn, results)
end

@testset "Ensure irrigation occurs" begin
    zone = setup_zone()

    # Create temporary directory that is cleaned up on process exit
    temp_dir = mktempdir()
    fn = joinpath(temp_dir, "test_batch_save.jld2")

    example_batch_save(output_fn=fn)
    saved_results = load(fn)

    df = collate_results(saved_results, "zone_results", "irrigated_yield_sum")

    # compiled_results = for (idx, res) in results
    #     res_df = collate_results(res, "zone_results", "irrigated_yield_sum")
    #     res_stats = [scenario_stats(res_df, "$(i)/z") for i in 1:length(samples)]
    #     result_set = DataFrame(res_stats)
    #     mean(result_set[:total])
    # end

    scen_stats = scenario_stats(df, "1/zone")

    @test !isnan(scen_stats.mean)
    @test !isnan(scen_stats.median)
    @test !isnan(scen_stats.total)

    @test scen_stats.mean != 0.0
    @test scen_stats.median != 0.0
    @test scen_stats.total != 0.0
end

max_stars_repo_name: ConnectedSystems/agtor-jl

max_issues_repo_name: ConnectedSystems/agtor-jl

max_forks_repo_name: ConnectedSystems/agtor-jl


text:
using Statistics
using Printf

using OffsetArrays
using DifferentialEquations
using Turing
using CairoMakie
using ElectronDisplay

# Numerical solution of the 2D advection-diffusion equation

@inline advective_flux_x(i, j, u, ϕ) = u * (ϕ[i-1, j] + ϕ[i, j]) / 2
@inline advective_flux_y(i, j, v, ϕ) = v * (ϕ[i, j-1] + ϕ[i, j]) / 2

@inline ∂x_advective_flux(i, j, Δx, u, ϕ) = (advective_flux_x(i+1, j, u, ϕ) - advective_flux_x(i, j, u, ϕ)) / Δx
@inline ∂y_advective_flux(i, j, Δy, v, ϕ) = (advective_flux_y(i, j+1, v, ϕ) - advective_flux_y(i, j, v, ϕ)) / Δy

@inline div_advective_flux(i, j, Δx, Δy, u, v, ϕ) = ∂x_advective_flux(i, j, Δx, u, ϕ) + ∂y_advective_flux(i, j, Δy, v, ϕ)

@inline ∂xᶠᵃ(i, j, Δx, ϕ) = (ϕ[i, j] - ϕ[i-1, j]) / Δx
@inline ∂yᵃᶠ(i, j, Δy, ϕ) = (ϕ[i, j] - ϕ[i, j-1]) / Δy

@inline diffusive_flux_x(i, j, Δx, κˣ, ϕ) = κˣ * ∂xᶠᵃ(i, j, Δx, ϕ)
@inline diffusive_flux_y(i, j, Δy, κʸ, ϕ) = κʸ * ∂yᵃᶠ(i, j, Δy, ϕ)

@inline ∂x_diffusive_flux(i, j, Δx, κˣ, ϕ) = (diffusive_flux_x(i+1, j, Δx, κˣ, ϕ) - diffusive_flux_x(i, j, Δx, κˣ, ϕ)) / Δx
@inline ∂y_diffusive_flux(i, j, Δy, κʸ, ϕ) = (diffusive_flux_y(i, j+1, Δy, κʸ, ϕ) - diffusive_flux_y(i, j, Δy, κʸ, ϕ)) / Δy

@inline div_diffusive_flux(i, j, Δx, Δy, κˣ, κʸ, ϕ) = ∂x_diffusive_flux(i, j, Δx, κˣ, ϕ) + ∂y_diffusive_flux(i, j, Δy, κʸ, ϕ)

function advection_diffusion_2d(∂ϕ∂t, ϕ, p, t)
    Nx, Ny, Hx, Hy, Δx, Δy, u, v, κˣ, κʸ = p
    Nx, Ny, Hx, Hy = Int.((Nx, Ny, Hx, Hy))

    # filling periodic halos
    ϕ[-Hx+1:0, :] .= ϕ[Nx:Nx, :]
    ϕ[:, -Hy+1:0] .= ϕ[:, Ny:Ny]

    ϕ[Nx+1:Nx+Hx, :] .= ϕ[1:1, :]
    ϕ[:, Ny+1:Ny+Hy] .= ϕ[:, 1:1]

    for i in 1:Nx, j in 1:Ny
        ∂ϕ∂t[i, j] = - div_advective_flux(i, j, Δx, Δy, u, v, ϕ) + div_diffusive_flux(i, j, Δx, Δy, κˣ, κʸ, ϕ)
    end

    return
end

# Initial condition

ϕ_initial_condition(x, y) = exp(-100 * (x - 0.5)^2) * exp(-100 * (y - 0.5)^2)

# Generating a high-resolution solution and animating it

Nx = Ny = 256
Hx = Hy = 1

Lx = Ly = 1
Δx = Lx / Nx
Δy = Ly / Ny

u = 0.25
v = 0.15

κˣ = 1e-3
κʸ = 1e-3

p = [Nx, Ny, Hx, Hy, Δx, Δy, u, v, κˣ, κʸ]

xC = range(Δx/2, Lx - Δx/2, length=Nx)
xF = range(0, Lx, length=Nx+1)

yC = range(Δy/2, Ly - Δy/2, length=Ny)
yF = range(0, Ly, length=Ny+1)

ϕ₀ = OffsetArray(zeros(Nx+2Hx, Ny+2Hy), -Hx+1:Nx+Hx, -Hy+1:Ny+Hy)
@. ϕ₀[1:Nx, 1:Ny] = ϕ_initial_condition(xC, yC')

tspan = (0.0, 10.0)
prob = ODEProblem(advection_diffusion_2d, ϕ₀, tspan, p, saveat=0.1)
sol = solve(prob, Tsit5())

function animate_advection_diffusion_2d(ϕ, x, y, t; filepath, resolution=(1280, 720), colormap=:matter, colorrange=(0, 1), framerate=15)
    Nx, Ny, Nt = size(ϕ)

    n = Observable(1)
    ϕₙ = @lift ϕ[:, :, $n]

    title = @lift @sprintf("2D advection-diffusion: t = %.1f", t[$n])

    fig = Figure(; resolution)
    ax = Axis(fig[1, 1], aspect=1, title=title, xlabel="x", ylabel="y")

    heatmap!(ax, x, y, ϕₙ; colormap, colorrange)
    Colorbar(fig[1, 2], colormap=colormap, limits=colorrange)

    record(fig, filepath, 1:Nt; framerate) do time_index
        @info "Animating 2D advection-diffusion: frame $time_index/$(length(t))"
        n[] = time_index
    end
end

ϕ_correct = Array(sol)[Hx+1:Nx+Hx, Hy+1:Ny+Hy, :]

animate_advection_diffusion_2d(ϕ_correct, xC, yC, sol.t, filepath="advection_diffusion_2d.mp4")

# Coarse graining the high-resolution solution to produce a "true solution" at the coarse resolution.

prob = remake(prob, saveat=1.0)
sol = solve(prob, Tsit5())

ϕ_correct = Array(sol)[Hx+1:Nx+Hx, Hy+1:Ny+Hy, :]

N_thinning = 8
Nx_coarse = Int(Nx / N_thinning)
Ny_coarse = Int(Ny / N_thinning)

Tx_coarse = Nx_coarse + 2Hx
Ty_coarse = Ny_coarse + 2Hy

Δx_coarse = Lx / Nx_coarse
Δy_coarse = Ly / Ny_coarse

Nt = length(sol.t)

ϕ_coarse_correct = zeros(Nx_coarse, Ny_coarse, Nt)

for i in 1:Nx_coarse, j in 1:Ny_coarse
    is = 1+(i-1)*N_thinning:i*N_thinning
    js = 1+(j-1)*N_thinning:j*N_thinning
    ϕ_coarse_correct[i, j, :] .= dropdims(mean(ϕ_correct[is, js, :], dims=(1, 2)), dims=(1, 2))
end

xC_coarse = range(Δx_coarse/2, Lx - Δx_coarse/2, length=Nx_coarse)
yC_coarse = range(Δy_coarse/2, Lx - Δy_coarse/2, length=Ny_coarse)

animate_advection_diffusion_2d(ϕ_coarse_correct, xC_coarse, yC_coarse, sol.t, filepath="advection_diffusion_2d_coarse.mp4", framerate=2)

# Setting up inference

is = -Hx+1:Nx_coarse+Hx
js = -Hy+1:Ny_coarse+Hy
ϕ₀_coarse = OffsetArray(zeros(Tx_coarse, Ty_coarse), is, js)
@. ϕ₀_coarse[1:Nx_coarse, 1:Ny_coarse] = ϕ_initial_condition(xC_coarse, yC_coarse')

p_coarse = [Nx_coarse, Ny_coarse, Hx, Hy, Δx_coarse, Δy_coarse, u, v, κˣ, κʸ]
prob_coarse = ODEProblem(advection_diffusion_2d, ϕ₀_coarse, tspan, p_coarse, saveat=1.0)

function predict_advection_diffusion_2d(prob, u, v, κˣ, κʸ)
    Nx, Ny, Hx, Hy, Δx, Δy, _ = prob.p
    p = [Nx, Ny, Hx, Hy, Δx, Δy, u, v, κˣ, κʸ]

    prob = remake(prob, p=p)
    prediction = solve(prob, Tsit5())

    return prediction
end

@model function probabilistic_advection_diffusion_2d(data, prob)
    σ ~ InverseGamma(3, 0.5)
    u ~ Normal(0, 1)
    v ~ Normal(0, 1)
    κˣ ~ InverseGamma(3, 0.1)
    κʸ ~ InverseGamma(3, 0.1)

    prediction = predict_advection_diffusion_2d(prob, u, v, κˣ, κʸ)

    Nx, Ny, Hx, Hy, _ = prob.p
    Nx, Ny, Hx, Hy = Int.((Nx, Ny, Hx, Hy))

    # for n in 1:length(prediction)
    #     prediction_n = prediction[n][1:Nx, 1:Ny][:]
    #     data[:, n] ~ MvNormal(prediction_n, σ)
    # end

    for n in 1:length(prediction), j in 1:Ny
        data[:, j, n] ~ MvNormal(prediction[n][1:Nx, j], σ)
    end

    # prediction_flattened = Array(prediction)[Hx+1:Nx+Hx, Hy+1:Ny+Hy, :][:]
    # data ~ MvNormal(prediction_flattened, σ)
end

function flatten_solution(ϕ)
    Nx, Ny, Nt = size(ϕ)
    ϕ_flat = zeros(Nx * Ny, Nt)
    for n in 1:Nt
        ϕ_flat[:, n] .= ϕ[1:Nx, 1:Ny, n][:]
    end
    return ϕ_flat
end

function plot_chains_and_densities(chains, vars; filepath)
    n_chains = size(chains, 3)

    fig = Figure(resolution=(800, 200*length(vars)))

    for (n, var) in enumerate(vars)
        ax1 = Axis(fig[n, 1], ylabel="$var", xlabel="iteration")
        for c in 1:n_chains
            lines!(ax1, chains[var].data[:, c])
        end

        ax2 = Axis(fig[n, 2], ylabel="P($var)", xlabel="$var", yticklabelsvisible=false)
        density!(ax2, chains[var][:].data)
    end

    save(filepath, fig, px_per_unit=2)

    return
end

prior_chain = sample(probabilistic_advection_diffusion_2d(ϕ_coarse_correct[:], prob_coarse), Prior(), 100)
plot_chains_and_densities(prior_chain, [:u, :v, :κˣ, :κʸ, :σ], filepath="priors.png")

function callback(rng, model, sampler, sample, state, iter; kwargs...)
    n = Threads.threadid()

    u  = sample.θ.u[1][1]
    v  = sample.θ.v[1][1]
    κˣ = sample.θ.κˣ[1][1]
    κʸ = sample.θ.κʸ[1][1]
    σ  = sample.θ.σ[1][1]

    @info "Chain $n, Iteration $iter: u=$u, v=$v, κˣ=$κˣ, κʸ=$κʸ, σ=$σ"
end

# n_chains = 8
# @assert Threads.nthreads() >= n_chains

stat_model = probabilistic_advection_diffusion_2d(ϕ_coarse_correct, prob_coarse)
# stat_model = probabilistic_advection_diffusion_2d(ϕ_coarse_correct |> flatten_solution, prob_coarse)

# chain = sample(stat_model, MH(), 100, progress=true; callback)
chain = sample(stat_model, NUTS(), 100, progress=true; callback)
# chain = sample(probabilistic_advection_diffusion_2d(ϕ_coarse_correct[:], prob_coarse), MH(), 1000, progress=true; callback)
# chains = sample(probabilistic_advection_diffusion_2d(ϕ_coarse_correct, prob_coarse), sampler, MCMCThreads(), n_samples, n_chains, progress=true, callback=callback_multi_chain)

plot_chains_and_densities(chain, [:u, :v, :κˣ, :κʸ, :σ], filepath="posteriors.png")


text:
using LinearAlgebra
using Test

using SetProg, SetProg.Sets
using Polyhedra
using MultivariatePolynomials

using JuMP
const MOIT = MOI.Test

const quartic_inner_poly = [3.1518541833100864, -0.1617384194869734]
const quartic_inner_obj = 6.447419478140056
const quartic_inner_α = 5.6567546886722795
const quartic_inner_convexity = [12.0, 0.0, quartic_inner_α, 0.0, quartic_inner_poly[1]+2quartic_inner_poly[2],
                                 quartic_inner_α, 8.48516455194103, 0.0, 0.0, 12.0]

const quartic_outer_β = 0.30177574048813055
const quartic_outer_γ = 0.5936049698923986
const quartic_outer_λ = -0.09857757888257276
const quartic_outer_obj = 1.611854896946893
const quartic_outer_α = 0.7928996242545062
const quartic_outer_convexity = [3.621308885857567, 0.0, quartic_outer_α, 0.0, 0.08578956499151169,
                                 quartic_outer_α, 1.5, 0.0, 0.0, 3.6212933687704307]

@testset "Square" begin
    config = MOIT.Config()
    @testset "Ellipsoid" begin
        @testset "John" begin
            # Q = [1 0
            #      0 1]
            # t = √det(Q) = 1
            Q = [1.0, 0.0, 1.0]
            t = 1.0
            @testset "Homogeneous" begin
                Tests.john_homogeneous_square_test(bridged_mock(mock -> MOI.Utilities.mock_optimize!(mock, [Q; t])),
                                                   config)
            end
            @testset "Non-homogeneous" begin
                @testset "Ellipsoid" begin
                    β = -1.0
                    b = [0.0, 0.0]
                    Tests.john_nonhomogeneous_ell_square_test(bridged_mock(mock -> MOI.Utilities.mock_optimize!(mock, [Q; β; b; t])),
                                                              config)
                end
                @testset "PolySet" begin
                    β = 1.0
                    b = [0.0, 0.0]
                    Tests.john_nonhomogeneous_quad_square_test(bridged_mock(mock -> begin
                        # β-1 b[1] b[2]
                        #  .  Q[1] Q[2]
                        #  .   .   Q[3]
                        MOI.Utilities.mock_optimize!(mock, [β-1; b[1]; Q[1]; b[2]; Q[2]; Q[3]; 2Q])
                    end), config)
                end
            end
            @testset "Piecewise" begin
                Q1 = ones(3)
                Q2 = [1.0, -1.0, 1.0]
                Tests.piecewise_semiell_inner_homogeneous_◇_square_test(
                    bridged_mock(mock -> MOI.Utilities.mock_optimize!(
                        mock,
                        [Q1; Q2; Q2; Q1])),
                    config)
                Q = [1.0, 0.0, 1.0]
                Tests.piecewise_semiell_inner_homogeneous_□_square_test(
                    bridged_mock(mock -> MOI.Utilities.mock_optimize!(
                        mock,
                        [Q; Q; Q; Q])),
                    config)
            end
        end
        @testset "Löwner" begin
            # Q = [√2  0
            #       0 √2]
            # t = √det(Q) = 2                                                                  Q11  Q12  Q22  t
            Tests.löwner_homogeneous_square_test(bridged_mock(mock -> MOI.Utilities.mock_optimize!(mock, [0.5, 0.0, 0.5, 0.5])),
                                                 config)
        end
    end
    @testset "Quartic" begin
        @testset "Inner" begin
            # The PSD matrix for the variable  is 3 x 3 so 3 * (3+1) / 2 = 6
            # The PSD matrix for the convexity is 4 x 4 so 4 * (5+1) / 2 = 10
            # entries
            # 1 variable for t
            # hence 17 variables
            sol = [1.0; 0.0; quartic_inner_poly; 0.0; 1.0;
                   quartic_inner_convexity; quartic_inner_obj]
            Tests.quartic_inner_homogeneous_square_test(bridged_mock(mock -> MOI.Utilities.mock_optimize!(mock, sol)),
                                                        config)
        end
        @testset "Outer" begin
            sol = [quartic_outer_β; 0.0; quartic_outer_γ; quartic_outer_λ; 0.0; quartic_outer_β;
                   quartic_outer_convexity; quartic_outer_obj]
            Tests.quartic_outer_homogeneous_square_test(bridged_mock(mock -> MOI.Utilities.mock_optimize!(mock, sol)),
                                                        config)
        end
    end
end


text:
test_private_key_contents = """
-----BEGIN RSA PRIVATE KEY-----
MIIEpAIBAAKCAQEAwaj5gy1tKhNCXYsw5awTpg/SrIgniRWiHQYvwYrjBMTzWDwi
wa7ZqJ/tcLCMxey8d6LkxVramchb6FGlYt6Qao2rPe4xmH1v3wzgfV8HrmS3oFr9
xp4N6y8rkVw1oVHlv7sPNmU7sc/1i+HqSQLJkzUqcBjz5SC1lzvexgqougrxu5Mf
it6ApdauE88K/nHUNKv/93Yj3BgzL+AVSDIA/o9lxEYyYFY034hQEXiMA47g/PSX
JmRpfCYhZPo4JcFpJaHLcxYalYTxdnT/9x0giAD3kseYrm8fL+pCCqLeXiEYZ1pk
Ff6usaraTjKlbdyKgfm/ndW7rZWsmBEEE+BE0QIDAQABAoIBAETi5l5S+lHNJ/4+
VNZdr1+SoU3d4LJSj6fjW3ls/fPkWus3+OUMXUAuIzfClVCSNTgZKWdKOhgZvReW
c04H/Tbpzc/pupqThAfDfNcbxo+yI99Wfo74CdtunADI3ArWJvNchrzCpclKP/3S
DpSJ+Unzz3Swa5FFYVK6aTPo9ru2Kjv8U4boiwwla8RrSjPGHV2RnIj246m+bFGG
jvYxZJCIVZPFjhQYOGyvrtEsBXUsNCAkxMdeMgGb3i3DRZP0Mn06gspGDQ9EL7ZU
0LSUFcNPB+1GrG0QIDAuRx8diAFwr1/T7zcpMh4bYbjyGCpvMB9V7/qDloPEXxUD
DcsGBkECgYEA69HoHOSDSEo53XZvzNd+uSz5NUe7ksentV2fSDh2fouqjID0875L
C1YetJpKTLLj4ezQ0MENOo9EOax5q05c8x8tzDqewKKk1bVAbhhZ6H1ReDUwSJR6
hgqQ8H7dTOKlx2mhFwAq04ZA9RsjOyfb2gbWXTK8UWJx9Rtn705VyM8CgYEA0jt5
SqcaMDlYYRxUA00/4CECzTSy32/IRtOgjdgH0oxnrkIkEYwC8UH85W1l314tyxhM
1eZzMWMnfG+Ru3/I5Ibs4XlhIddGuwCVqBTTnxk3Xl+3AjlEv8wbvfvRyhH207V5
LgojhgyEbaxt7nDLKNqnhWDeixG9mgqa/np5QF8CgYEAg71sMlfmj4gH85SquIDa
ppBdWEqOm102b7exuCLy75cO3BdewJxaSCcAWFypjVeSvzARJ83xmflXX/Det7i7
bIbKiwFklrFF7K2pJNgtwhfgX4s3QhTQrflRJ0l4T+9+NXzpJXROeITT3RogAHgv
9glD0CT712trxkZdNe9JNTkCgYA23sD0hBKovX7csbkXt4TRaaxChNuex50dJ0UD
YHwCLxUsfSTO0508/L0V4GyzVPBczrXXKjOskzWOG8yMhHiZdMHEw5pfa8GRM0Dc
QrtNDhbOcerhELNp38rp1hzWQcQLU5USOCE3bnLt8RAsKJJN0mvguvy6jWHQH65+
dtbIvwKBgQCzuaM2U50rGRgzFVzkCNEf6HCgKbFei31TeXvgb4VrhJ3Tu+RmPXJH
GS2zLFh88/2D1f88F6jr2M5G601vJjAQPypDGw4RHJtsS+KjAw8TjAYr/M2f1dco
HpUHZbGip3HcM84jooRnYlyhQTMiJhXxuno3xRqLGUfinpwqqyB6YA==
-----END RSA PRIVATE KEY-----
"""

test_private_key = joinpath(mktempdir(), "test_private_key.pem")

open(test_private_key, "w") do io
    s = convert(String, strip(test_private_key_contents))::String
    println(io, s)
end


text:
include("area.jl")
include("vectorin.jl")
include("volume.jl")

function Assemblein(IK,JK,VK,con,Node, Ele, Cen, Face_in, EE, NU, FF)
    @inbounds @sync @distributed for i = 1:size(Face_in,1)

        v1 = Volume(Node[Face_in[i,1],1],Node[Face_in[i,1],2],Node[Face_in[i,1],3],
                    Node[Face_in[i,2],1],Node[Face_in[i,2],2],Node[Face_in[i,2],3],
                    Node[Face_in[i,3],1],Node[Face_in[i,3],2],Node[Face_in[i,3],3],
                    Cen[Face_in[i,4],1],Cen[Face_in[i,4],2],Cen[Face_in[i,4],3])
        v2 = Volume(Node[Face_in[i,1],1],Node[Face_in[i,1],2],Node[Face_in[i,1],3],
                    Node[Face_in[i,2],1],Node[Face_in[i,2],2],Node[Face_in[i,2],3],
                    Node[Face_in[i,3],1],Node[Face_in[i,3],2],Node[Face_in[i,3],3],
                    Cen[Face_in[i,5],1],Cen[Face_in[i,5],2],Cen[Face_in[i,5],3])
        v = v1 + v2

        A1 = area(Node[Face_in[i,1],1],Node[Face_in[i,1],2],Node[Face_in[i,1],3],
                  Node[Face_in[i,2],1],Node[Face_in[i,2],2],Node[Face_in[i,2],3],
                  Cen[Face_in[i,4],1],Cen[Face_in[i,4],2],Cen[Face_in[i,4],3])

        A2 = area(Node[Face_in[i,2],1],Node[Face_in[i,2],2],Node[Face_in[i,2],3],
                  Node[Face_in[i,3],1],Node[Face_in[i,3],2],Node[Face_in[i,3],3],
                  Cen[Face_in[i,4],1],Cen[Face_in[i,4],2],Cen[Face_in[i,4],3])

        A3 = area(Node[Face_in[i,1],1],Node[Face_in[i,1],2],Node[Face_in[i,1],3],
                  Node[Face_in[i,3],1],Node[Face_in[i,3],2],Node[Face_in[i,3],3],
                  Cen[Face_in[i,4],1],Cen[Face_in[i,4],2],Cen[Face_in[i,4],3])

        A4 = area(Node[Face_in[i,1],1],Node[Face_in[i,1],2],Node[Face_in[i,1],3],
                  Node[Face_in[i,2],1],Node[Face_in[i,2],2],Node[Face_in[i,2],3],
                  Cen[Face_in[i,5],1],Cen[Face_in[i,5],2],Cen[Face_in[i,5],3])

        A5 = area(Node[Face_in[i,2],1],Node[Face_in[i,2],2],Node[Face_in[i,2],3],
                  Node[Face_in[i,3],1],Node[Face_in[i,3],2],Node[Face_in[i,3],3],
                  Cen[Face_in[i,5],1],Cen[Face_in[i,5],2],Cen[Face_in[i,5],3])

        A6 = area(Node[Face_in[i,1],1],Node[Face_in[i,1],2],Node[Face_in[i,1],3],
                  Node[Face_in[i,3],1],Node[Face_in[i,3],2],Node[Face_in[i,3],3],
                  Cen[Face_in[i,5],1],Cen[Face_in[i,5],2],Cen[Face_in[i,5],3])

        n = Vectorin(Node[Face_in[i,1],1],Node[Face_in[i,1],2],Node[Face_in[i,1],3],
                     Node[Face_in[i,2],1],Node[Face_in[i,2],2],Node[Face_in[i,2],3],
                     Node[Face_in[i,3],1],Node[Face_in[i,3],2],Node[Face_in[i,3],3],
                     Cen[Face_in[i,4],1],Cen[Face_in[i,4],2],Cen[Face_in[i,4],3],
                     Cen[Face_in[i,5],1],Cen[Face_in[i,5],2],Cen[Face_in[i,5],3])

        bCx = (1/v)*(n[1,1]*(5/12)*A1+n[2,1]*(1/12)*A2+n[3,1]*(5/12)*A3+n[4,1]*(5/12)*A4+n[5,1]*(1/12)*A5+n[6,1]*(5/12)*A6)
        bCy = (1/v)*(n[1,2]*(5/12)*A1+n[2,2]*(1/12)*A2+n[3,2]*(5/12)*A3+n[4,2]*(5/12)*A4+n[5,2]*(1/12)*A5+n[6,2]*(5/12)*A6)
        bCz = (1/v)*(n[1,3]*(5/12)*A1+n[2,3]*(1/12)*A2+n[3,3]*(5/12)*A3+n[4,3]*(5/12)*A4+n[5,3]*(1/12)*A5+n[6,3]*(5/12)*A6)
        bC = [bCx 0 0; 0 bCy 0; 0 0 bCz; bCy bCx 0; 0 bCz bCy; bCz 0 bCx]
        bDx = (1/v)*(n[1,1]*(5/12)*A1+n[2,1]*(5/12)*A2+n[3,1]*(1/12)*A3+n[4,1]*(5/12)*A4+n[5,1]*(5/12)*A5+n[6,1]*(1/12)*A6)
        bDy = (1/v)*(n[1,2]*(5/12)*A1+n[2,2]*(5/12)*A2+n[3,2]*(1/12)*A3+n[4,2]*(5/12)*A4+n[5,2]*(5/12)*A5+n[6,2]*(1/12)*A6)
        bDz = (1/v)*(n[1,3]*(5/12)*A1+n[2,3]*(5/12)*A2+n[3,3]*(1/12)*A3+n[4,3]*(5/12)*A4+n[5,3]*(5/12)*A5+n[6,3]*(1/12)*A6)
        bD = [bDx 0 0; 0 bDy 0; 0 0 bDz; bDy bDx 0; 0 bDz bDy; bDz 0 bDx]
        bBx = (1/v)*(n[1,1]*(1/12)*A1+n[2,1]*(5/12)*A2+n[3,1]*(5/12)*A3+n[4,1]*(1/12)*A4+n[5,1]*(5/12)*A5+n[6,1]*(5/12)*A6)
        bBy = (1/v)*(n[1,2]*(1/12)*A1+n[2,2]*(5/12)*A2+n[3,2]*(5/12)*A3+n[4,2]*(1/12)*A4+n[5,2]*(5/12)*A5+n[6,2]*(5/12)*A6)
        bBz = (1/v)*(n[1,3]*(1/12)*A1+n[2,3]*(5/12)*A2+n[3,3]*(5/12)*A3+n[4,3]*(1/12)*A4+n[5,3]*(5/12)*A5+n[6,3]*(5/12)*A6)
        bB = [bBx 0 0; 0 bBy 0; 0 0 bBz; bBy bBx 0; 0 bBz bBy; bBz 0 bBx]
        bAx = (1/v)*(n[1,1]*(1/12)*A1+n[2,1]*(1/12)*A2+n[3,1]*(1/12)*A3)
        bAy = (1/v)*(n[1,2]*(1/12)*A1+n[2,2]*(1/12)*A2+n[3,2]*(1/12)*A3)
        bAz = (1/v)*(n[1,3]*(1/12)*A1+n[2,3]*(1/12)*A2+n[3,3]*(1/12)*A3)
        bA = [bAx 0 0; 0 bAy 0; 0 0 bAz; bAy bAx 0; 0 bAz bAy; bAz 0 bAx]
        bEx = (1/v)*(n[4,1]*(1/12)*A4+n[5,1]*(1/12)*A5+n[6,1]*(1/12)*A6)
        bEy = (1/v)*(n[4,2]*(1/12)*A4+n[5,2]*(1/12)*A5+n[6,2]*(1/12)*A6)
        bEz = (1/v)*(n[4,3]*(1/12)*A4+n[5,3]*(1/12)*A5+n[6,3]*(1/12)*A6)
        bE = [bEx 0 0; 0 bEy 0; 0 0 bEz; bEy bEx 0; 0 bEz bEy; bEz 0 bEx]
        B = hcat(bC,bD,bB,bA,bE)

        D = (EE/((1+NU)*(1-2*NU)))*[1-NU NU NU 0 0 0; NU 1-NU NU 0 0 0; NU NU 1-NU 0 0 0; 0 0 0 (1-2*NU)/2 0 0; 0 0 0 0 (1-2*NU)/2 0; 0 0 0 0 0 (1-2*NU)/2]
        ke = B'*D*B*v
        DOF = [3*Face_in[i, 1]-2,3*Face_in[i, 1]-1,3*Face_in[i, 1],
               3*Face_in[i, 2]-2,3*Face_in[i, 2]-1,3*Face_in[i, 2],
               3*Face_in[i, 3]-2,3*Face_in[i, 3]-1,3*Face_in[i, 3],
               3*Face_in[i, 6]-2,3*Face_in[i, 6]-1,3*Face_in[i, 6],
               3*Face_in[i, 7]-2,3*Face_in[i, 7]-1,3*Face_in[i, 7]]
        xi = con+i*225-224
        for n1 = 1:15
            for n2 = 1:15
                IK[xi] = DOF[n1]
                JK[xi] = DOF[n2]
                VK[xi] = ke[n1,n2]
                xi += 1
            end
        end
    end
    return IK, JK, VK
end

max_stars_repo_path: src/multi-core juSFEM/assemblein.jl

max_stars_repo_name: Elliothuo/juSFEM

max_issues_repo_path: src/multi-core juSFEM/assemblein.jl

max_issues_repo_name: Elliothuo/juSFEM

max_forks_repo_path: src/multi-core juSFEM/assemblein.jl

max_forks_repo_name: Elliothuo/juSFEM


text:
# TODO: Implement scaled Cosine and generalized Cosine
# TODO: This is all wrong. Replace entirely

immutable Cosine <: ContinuousUnivariateDistribution
end

rand(d::Cosine) = asin(2.0 * rand() - 1.0)

function cdf(d::Cosine, x::Real)
    x < zero(x) ? 0.0 : (x > one(x) ? 1.0 : 0.5(1 + sin(x)))
end

entropy(d::Cosine) = log(4.0 * pi) - 1.0

insupport(::Cosine, x::Real) = zero(x) <= x <= one(x)
insupport(::Type{Cosine}, x::Real) = zero(x) <= x <= one(x)

kurtosis(d::Cosine) = -1.5

mean(d::Cosine) = 0.5

median(d::Cosine) = 0.5

# mgf(d::Cosine, t::Real)
# cf(d::Cosine, t::Real)

mode(d::Cosine) = 0.5
modes(d::Cosine) = [0.5]

pdf(d::Cosine, x::Real) = zero(x) <= x <= one(x) ? 0.5 * cos(x) : 0.0

quantile(d::Cosine, p::Real) = asin(2.p - 1.0)

rand(d::Cosine) = sin(rand() * pi / 2.0)^2

skewness(d::Cosine) = 0.0

var(d::Cosine) = (pi^2 - 8.0) / (4.0 * pi^2)


text:
include("general_U4.jl")
include("QFTCircuit.jl")
include("hamiltonians.jl")
include("variational_circuit.jl")
include("supremacy_circuit.jl")
include("google53.jl")
include("openbox.jl")


text:
# This file was formerly a part of Julia. License is MIT: https://julialang.org/license
# TODO: parameterize conv tests
using Test, DSP, OffsetArrays
import DSP: filt, filt!, deconv, conv, xcorr
using DSP: optimalfftfiltlength, unsafe_conv_kern_os!, _conv_kern_fft!, _conv_similar,
    nextfastfft

@testset "filt" begin
    # Filter
    b = [1., 2., 3., 4.]
    x = [1., 1., 0., 1., 1., 0., 0., 0.]
    @test filt(b, 1., x)  == [1., 3., 5., 8., 7., 5., 7., 4.]
    @test filt(b, [1., -0.5], x)  == [1., 3.5, 6.75, 11.375, 12.6875, 11.34375, 12.671875, 10.3359375]
    # With ranges
    @test filt(b, 1., 1.0:10.0) == [1., 4., 10., 20., 30., 40., 50., 60., 70., 80.]
    @test filt(1.:4., 1., 1.0:10.0) == [1., 4., 10., 20., 30., 40., 50., 60., 70., 80.]
    # Across an array is the same as channel-by-channel
    @test filt(b, 1., [x 1.0:8.0]) == [filt(b, 1., x) filt(b, 1., 1.0:8.0)]
    @test filt(b, [1., -0.5], [x 1.0:8.0]) == [filt(b, [1., -0.5], x) filt(b, [1., -0.5], 1.0:8.0)]
    si = zeros(3)
    @test filt(b, 1., [x 1.0:8.0], si) == [filt(b, 1., x, si) filt(b, 1., 1.0:8.0, si)]
    @test si == zeros(3) # Will likely fail if/when arrayviews are implemented
    si = [zeros(3) ones(3)]
    @test filt(b, 1., [x 1.0:8.0], si) == [filt(b, 1., x, zeros(3)) filt(b, 1., 1.0:8.0, ones(3))]
    # With initial conditions: a lowpass 5-pole butterworth filter with W_n = 0.25,
    # and a stable initial filter condition matched to the initial value
    b = [0.003279216306360201,0.016396081531801006,0.03279216306360201,0.03279216306360201,0.016396081531801006,0.003279216306360201]
    a = [1.0,-2.4744161749781606,2.8110063119115782,-1.703772240915465,0.5444326948885326,-0.07231566910295834]
    si = [0.9967207836936347,-1.4940914728163142,1.2841226760316475,-0.4524417279474106,0.07559488540931815]
    @test filt(b, a, ones(10), si) ≈ ones(10) # Shouldn't affect DC offset

    @test_throws ArgumentError filt!([1, 2], [1], [1], [1])
end

@testset "conv" begin

    @test optimalfftfiltlength(1, 3) == 1 # Should be at least the first input

    @testset "conv-1D" begin
        # Convolution
        a = [1, 2, 1, 2]
        b = [1, 2, 3]
        expectation = [1, 4, 8, 10, 7, 6]
        im_expectation = [1, 3, 6, 6, 5, 3]
        a32 = convert(Array{Int32}, a)
        @test conv(a, b) == expectation
        @test conv(a32, b) == expectation
        fa = convert(Array{Float64}, a)
        f32a = convert(Array{Float32}, a)
        fb = convert(Array{Float64}, b)
        fexp = convert(Array{Float64}, expectation)
        im_fexp = convert(Array{Float64}, im_expectation)
        @test conv(fa, fb) ≈ fexp
        @test conv(complex.(fa, 1.), complex.(fb)) ≈ complex.(fexp, im_fexp)

        @test conv(fa, b) ≈ fexp
        @test conv(f32a, b) ≈ fexp
        @test conv(fb, a) ≈ fexp

        offset_arr = OffsetArray{Int}(undef, -1:2)
        offset_arr[:] = a
        @test conv(offset_arr, 1:3) == OffsetVector(expectation, 0:5)
        # Issue #352
        @test conv([1//2, 1//3, 1//4], [1, 2]) ≈ [1//2, 4//3, 11//12, 1//2]
        # Non-numerical arrays should not be convolved
        @test_throws MethodError conv([sin], [cos])
    end

    @testset "conv-2D" begin
        a =[1 2 1;
            2 3 1;
            1 2 1]

        b = [3 2;
             0 1]

        expectation = [3 8 7 2;
                       6 14 11 3;
                       3 10 10 3;
                       0 1 2 1]
        im_expectation = [3 5 5 2;
                          3 6 6 3;
                          3 6 6 3;
                          0 1 1 1]

        # Integers
        # Real Integers
        a32 = convert(Array{Int32}, a)
        @test conv(a, b) == expectation
        @test conv(a32, b) == expectation
        # Floats
        fa = convert(Array{Float64}, a)
        f32a = convert(Array{Float32}, a)
        fb = convert(Array{Float64}, b)
        fexp = convert(Array{Float64}, expectation)
        im_fexp = convert(Array{Float64}, im_expectation)
        # Real
        @test conv(fa, fb) == fexp
        # Complex
        @test conv(complex.(fa, 1), complex.(fb)) == complex.(fexp,
                                                               im_fexp)
        @test conv(fa, b) ≈ fexp
        @test conv(f32a, b) ≈ fexp
        @test conv(fb, a) ≈ fexp

        offset_arr = OffsetArray{Int}(undef, -1:1, -1:1)
        offset_arr[:] = a
        @test conv(offset_arr, b) == OffsetArray(expectation, 0:3, 0:3)
    end

    @testset "seperable conv" begin
        u = [1, 2, 3, 2, 1]
        v = [6, 7, 3, 2]
        A = [1 2 3 4 5 6 7;
             8 9 10 11 12 13 14;
             15 16 17 18 19 20 21;
             22 23 24 25 26 27 28]
        exp = [6 19 35 53 71 89 107 77 33 14;
               60 148 217 285 339 393 447 315 134 56;
               204 478 658 822 930 1038 1146 798 338 140;
               468 1062 1400 1684 1828 1972 2116 1456 614 252;
               636 1426 1848 2188 2332 2476 2620 1792 754 308;
               624 1388 1778 2082 2190 2298 2406 1638 688 280;
               354 785 1001 1167 1221 1275 1329 903 379 154;
               132 292 371 431 449 467 485 329 138 56]
        @test_broken conv(u, v, A) == exp

        fu = convert(Array{Float64}, u)
        fv = convert(Array{Float64}, v)
        fA = convert(Array{Float64}, A)
        fexp = convert(Array{Float64}, exp)
        @test conv(fu, fv, fA) ≈ fexp

    end

    @testset "conv-ND" begin
        # is it safe to assume that if conv works for
        # int/float/complex in 1 and 2 D, it does in ND?
        a = convert(Array, reshape(1:27, (3, 3, 3)))
        b = ones(Int64, 2, 2, 2)
        exp = convert(Array, reshape([1, 3, 5, 3,
                                      5, 12, 16, 9,
                                      11, 24, 28, 15,
                                      7, 15, 17, 9,

                                      11, 24, 28, 15,
                                      28, 60, 68, 36,
                                      40, 84, 92, 48,
                                      23, 48, 52, 27,
                                      29, 60, 64, 33,

                                      64, 132, 140, 72,
                                      76, 156, 164, 84,
                                      41, 84, 88, 45,
                                      19, 39, 41, 21,

                                      41, 84, 88, 45,
                                      47, 96,  100, 51,
                                      25, 51, 53, 27], (4, 4, 4)))
        @test conv(a, b) == exp

        #6D, trivial, just to see if it works
        a = ones(2, 2, 2, 2, 2, 2)
        b = ones(1, 1, 1, 1, 1, 1)
        @test conv(a, b) == a

        # promote dims to largest
        a = cat([fill(n, 3, 3) for n in 1:6]..., dims=3)
        b = ones(Int64, 2, 2)
        expf1 = conv(a[:, :, 1], b)
        exp = cat([expf1 * n for n in 1:6]..., dims=3)
        @test conv(a, b) == exp
        fb = convert(Array{Float64}, b)
        @test conv(a, fb) ≈ convert(Array{Float64}, exp)
    end

    @testset "Overlap-Save" begin
        to_sizetuple = (diml, N) -> ntuple(_ -> diml, N)
        function os_test_data(eltype, diml, N)
            s = to_sizetuple(diml, N)
            arr = rand(eltype, s)
            s, arr
        end
        function test_os(eltype, nu, nv, N, nfft)
            nffts = to_sizetuple(nfft, N)
            su, u = os_test_data(eltype, nu, N)
            sv, v = os_test_data(eltype, nv, N)
            sout = su .+ sv .- 1
            out = _conv_similar(u, sout, axes(u), axes(v))
            unsafe_conv_kern_os!(out, u, v, su, sv, sout, nffts)
            os_out = copy(out)
            fft_nfft = nextfastfft(sout)
            _conv_kern_fft!(out, u, v, su, sv, sout, fft_nfft)
            @test out ≈ os_out
        end
        Ns = [1, 2, 3]
        eltypes = [Float32, Float64, Complex{Float64}]
        nlarge = 128

        regular_nsmall = [12, 128]
        for numdim in Ns
            for elt in eltypes
                for nsmall in regular_nsmall
                    nfft = optimalfftfiltlength(nsmall, nlarge)
                    test_os(elt, nlarge, nsmall, Val{numdim}(), nfft)
                end
            end
        end

        # small = 12, fft = 256 exercises output being smaller than the normal
        # valid region of OS convolution block
        #
        # small = 13, fft = 32 makes sout evenly divisble by nfft - sv + 1
        # small = 12, fft = 32 does not
        adversarial_nsmall_nfft = [(12, 256), (13, 32), (12, 32)]
        for (nsmall, nfft) in adversarial_nsmall_nfft
            test_os(Float64, nlarge, nsmall, Val{1}(), nfft)
        end

        # Should not bug:
        conv(zeros(4, 7, 1), zeros(3, 3, 3))

        # three blocks need to be padded in the following case:
        test_os(Float64, 25, 4, Val{1}(), 16)
    end
end

@testset "xcorr" begin
    a = [1, 2, 3]
    b = [4, 5]
    exp = [5, 14, 23, 12]
    @test xcorr([1, 2], [3, 4]) == [4, 11, 6]
    @test xcorr(a, b) == exp
    @test xcorr(a, b, padmode = :longest) == [0, 5, 14, 23, 12]
    @test xcorr(a, b, padmode = :none) == exp
    @test xcorr([1, 2], [3, 4, 5]) == [5, 14, 11, 6]
    @test xcorr([1, 2], [3, 4, 5], padmode = :longest) == [5, 14, 11, 6, 0]
    @test xcorr([1, 2], [3, 4, 5], padmode = :none) == [5, 14, 11, 6]
    @test xcorr([1.0im], [1.0im]) == [1]
    @test xcorr([1, 2, 3]*1.0im, ComplexF64[4, 5]) ≈ [5, 14, 23, 12]*im
    @test xcorr([1, 2]*1.0im, ComplexF64[3, 4, 5]) ≈ [5, 14, 11, 6]*im
    @test xcorr(ComplexF64[1, 2, 3], [4, 5]*1.0im) ≈ -[5, 14, 23, 12]*im
    @test xcorr(ComplexF64[1, 2], [3, 4, 5]*1.0im) ≈ -[5, 14, 11, 6]*im
    @test xcorr([1, 2, 3]*1.0im, [4, 5]*1.0im) ≈ [5, 14, 23, 12]
    @test xcorr([1, 2]*1.0im, [3, 4, 5]*1.0im) ≈ [5, 14, 11, 6]

    # Base Julia issue #17351
    let
        x = rand(10)
        u = rand(10, 3)
        su = view(u, :, 1)
        @test size(@inferred(xcorr(su, x))) == (19,)
        @test size(@inferred(xcorr(x, su))) == (19,)
    end

    # xcorr only supports 1d inputs at the moment
    @test_throws MethodError xcorr(rand(2, 2), rand(2, 2))

    off_a = OffsetVector(a, -1:1)
    off_b = OffsetVector(b, 0:1)
    @test xcorr(off_a, off_b, padmode = :none) == OffsetVector(exp, -2:1)

    # Issue #288
    @test xcorr(off_a, off_b, padmode = :longest) == OffsetVector(vcat(0, exp), -3:1)
end

@testset "deconv" begin
    # Test for issue #188: deconv mutates inputs
    let b = [4.0, 2.0, 1.0]; a = [2.0, 1.0]
        bb = b[:]
        aa = a[:]
        c = deconv(b,a)
        @test c == [2.0, 0.0]
        @test a == aa
        @test b == bb
    end
end


text:
module DifferenceEquations

using Distributions
using LinearAlgebra
using CommonSolve

import SciMLBase
import SciMLBase: SciMLProblem, solve

abstract type DifferenceProblem <: SciMLProblem end
abstract type AbstractStateSpaceProblem{isinplace} <: DifferenceProblem end

# Yuck hate this so much
promote_noise(x, y) = [x], [y]
promote_noise(x, y::AbstractArray) = [x], y
promote_noise(x::AbstractArray, y) = x, [y]
promote_noise(x::AbstractArray, y::AbstractArray) = x, y

include("noise.jl")
include("alg.jl")
include("linear.jl")
include("nonlinear.jl")
include("solution.jl")

# Exports
export StateSpaceProblem,
    ConditionalGaussian,
    LinearGaussian,
    LinearStateSpaceProblem,
    StandardGaussian,
    DefinedNoise,
    solve

end # module