From 2c0c91acfaf77a65bef67599068a6eb5842d5e7b Mon Sep 17 00:00:00 2001
From: Matthieu Gomez <gomez.matthieu@gmail.com>
Date: Fri, 14 Jun 2024 20:01:06 -0400
Subject: [PATCH] add experimental jointoperator

---
 Project.toml                   |  4 +++-
 src/InfinitesimalGenerators.jl |  6 +++++-
 src/MarkovProcess.jl           | 13 +++++++++++--
 src/derivatives.jl             | 10 +++++-----
 src/jointoperator.jl           | 20 ++++++++++++++++++++
 5 files changed, 44 insertions(+), 9 deletions(-)
 create mode 100644 src/jointoperator.jl

diff --git a/Project.toml b/Project.toml
index c57a6d0..c1db7ab 100644
--- a/Project.toml
+++ b/Project.toml
@@ -1,9 +1,10 @@
 name = "InfinitesimalGenerators"
 uuid = "2fce0c6f-5f0b-5c85-85c9-2ffe1d5ee30d"
-version = "2.0.0"
+version = "2.1.0"
 
 [deps]
 Arpack = "7d9fca2a-8960-54d3-9f78-7d1dccf2cb97"
+BlockBandedMatrices = "ffab5731-97b5-5995-9138-79e8c1846df0"
 Distributions = "31c24e10-a181-5473-b8eb-7969acd0382f"
 FillArrays = "1a297f60-69ca-5386-bcde-b61e274b549b"
 KrylovKit = "0b1a1467-8014-51b9-945f-bf0ae24f4b77"
@@ -12,6 +13,7 @@ Roots = "f2b01f46-fcfa-551c-844a-d8ac1e96c665"
 
 [compat]
 Arpack = "0.3, 0.4, 0.5"
+BlockBandedMatrices = "0.7, 0.8, 0.9, 0.10, 0.11, 0.12"
 Distributions = "0.22, 0.23, 0.24, 0.25"
 FillArrays = "0.8, 0.9, 0.10, 0.11, 0.12, 0.13, 1"
 KrylovKit = "0.4, 0.5, 0.6"
diff --git a/src/InfinitesimalGenerators.jl b/src/InfinitesimalGenerators.jl
index d0230c3..d73e527 100644
--- a/src/InfinitesimalGenerators.jl
+++ b/src/InfinitesimalGenerators.jl
@@ -6,6 +6,7 @@ using FillArrays: Ones, Zeros
 using KrylovKit: KrylovKit
 using LinearAlgebra: Diagonal, Tridiagonal, I, factorize, ldiv!, diag
 using Roots: fzero
+using BlockBandedMatrices: BandedBlockBandedMatrix, Block
 
 
 
@@ -17,6 +18,8 @@ include("AdditiveFunctional.jl")
 include("feynman_kac.jl")
 include("principal_eigenvalue.jl")
 include("derivatives.jl")
+include("jointoperator.jl")
+
 
 
 
@@ -34,5 +37,6 @@ cgf,
 tail_index,
 AdditiveFunctionalDiffusion,
 FirstDerivative,
-SecondDerivative
+SecondDerivative,
+jointoperator
 end
\ No newline at end of file
diff --git a/src/MarkovProcess.jl b/src/MarkovProcess.jl
index f4fc7bc..2900211 100644
--- a/src/MarkovProcess.jl
+++ b/src/MarkovProcess.jl
@@ -59,11 +59,20 @@ state_space(X::DiffusionProcess) = X.x
 
     at the border of state space
 """
-generator(X::DiffusionProcess) = generator(X.x, X.μx, X.σx)
+function generator(X::DiffusionProcess)
+    n = length(X.x)
+    𝕋 = Tridiagonal(zeros(n-1), zeros(n), zeros(n-1))
+    generator!(𝕋, X.x, X.μx, X.σx)
+end
 
 function generator(x::AbstractVector, μx::AbstractVector, σx::AbstractVector)
+    generator!(T, x, μx, σx)
+end
+
+function generator!(𝕋, x::AbstractVector, μx::AbstractVector, σx::AbstractVector)
+    # if you use this form, make sure that 𝕋 only has zero
     n = length(x)
-    𝕋 = Tridiagonal(zeros(n-1), zeros(n), zeros(n-1))
+    fill!(𝕋, 0)
     @inbounds for i in 1:n
         Δxp = x[min(i, n-1)+1] - x[min(i, n-1)]
         Δxm = x[max(i-1, 1) + 1] - x[max(i-1, 1)]
diff --git a/src/derivatives.jl b/src/derivatives.jl
index 849ff4f..520d91d 100644
--- a/src/derivatives.jl
+++ b/src/derivatives.jl
@@ -27,14 +27,14 @@ function Base.getindex(d::FirstDerivative{T}, i::Int) where {T}
 			return convert(T, bc[end])
 		else
 			Δxp = x[min(i, length(x)-1)+1] - x[min(i, length(x)-1)]
-			return (y[i+1] - y[i]) / Δxp
+			return convert(T, (y[i+1] - y[i]) / Δxp)
 		end
 	else
 		if i == 1
 			return convert(T, bc[1])
 		else
 			Δxm = x[max(i-1, 1) + 1] - x[max(i-1, 1)]
-			return (y[i] - y[i-1]) / Δxm
+			return convert(T, (y[i] - y[i-1]) / Δxm)
 		end
 	end
 end
@@ -65,11 +65,11 @@ function Base.getindex(d::SecondDerivative{T}, i::Int) where {T}
     Δxm = x[max(i-1, 1) + 1] - x[max(i-1, 1)]
     Δx = (Δxm + Δxp) / 2
     if i == 1 
-    	return y[2] / (Δxp * Δx) + (y[1] - bc[1] * Δxm) / (Δxm * Δx) - 2 * y[1] / (Δxp * Δxm)
+    	return convert(T, y[2] / (Δxp * Δx) + (y[1] - bc[1] * Δxm) / (Δxm * Δx) - 2 * y[1] / (Δxp * Δxm))
     elseif i ==  length(x)
-    	return (y[end] + bc[end] * Δxp) / (Δxp * Δx) + y[end - 1] / (Δxm * Δx) - 2 * y[end] / (Δxp * Δxm)
+    	return convert(T, (y[end] + bc[end] * Δxp) / (Δxp * Δx) + y[end - 1] / (Δxm * Δx) - 2 * y[end] / (Δxp * Δxm))
     else
-    	return y[i + 1] / (Δxp * Δx) + y[i - 1] / (Δxm * Δx) - 2 * y[i] / (Δxp * Δxm)
+    	return convert(T, y[i + 1] / (Δxp * Δx) + y[i - 1] / (Δxm * Δx) - 2 * y[i] / (Δxp * Δxm))
     end
 end
 
diff --git a/src/jointoperator.jl b/src/jointoperator.jl
new file mode 100644
index 0000000..59e6593
--- /dev/null
+++ b/src/jointoperator.jl
@@ -0,0 +1,20 @@
+function jointoperator(operators, Q::Array)
+    N = length(operators)
+    wn = size(operators[1], 1)
+    @assert all(o isa Tridiagonal for o in operators)
+    # check if all os have same size
+    @assert all(size(o) == (wn, wn) for o in operators) 
+    # check if the size of transition matrix is 
+    # same as the number of operators 
+    @assert size(Q,1) == size(Q,2) == N
+    J = BandedBlockBandedMatrix(Zeros(wn * N, wn * N), fill(wn, N) ,fill(wn, N), (N-1, N-1), (1, 1))
+    for i in 1:N
+        J[Block(i,i)] = operators[i]
+    end
+    for i in 1:N
+        for j in 1:N
+            J[Block(i,j)] += Tridiagonal(zeros(wn -1), fill(Q[i, j], wn), zeros(wn-1))
+        end
+    end
+    return J
+end
\ No newline at end of file