From fc34d09a4d81710576853f20d905cfb3889ce720 Mon Sep 17 00:00:00 2001
From: PharmCat <v.s.arnautov@yandex.ru>
Date: Sat, 16 Nov 2024 06:43:26 +0300
Subject: [PATCH] acov experimental

---
 src/rmat.jl      | 32 ++++++++++++++++++++++++++++++++
 src/varstruct.jl |  3 +++
 src/vartypes.jl  | 26 ++++++++++++++++++++++++++
 3 files changed, 61 insertions(+)

diff --git a/src/rmat.jl b/src/rmat.jl
index 8c4a098..f0e22b9 100644
--- a/src/rmat.jl
+++ b/src/rmat.jl
@@ -342,6 +342,38 @@ function rmat!(mx, θ, rz::AbstractMatrix, ::UN_, ::Int)
     mulαβαtinc!(mx, rz, rcov)
     return mx
 end
+
+# ACOV (AR)
+function rmat!(mx, θ, ::AbstractMatrix, ::ACOV_{AR_}, ::Int)
+    s  = size(mx, 1)
+    ρ  = θ[1]
+    if s > 1
+        for n = 2:s
+            mxnn = mx[n, n]
+            @inbounds @simd for m = 1:n-1
+                mxmm = mx[m, m]
+                mx[m, n] += mxnn * mxmm * ρ ^ (n - m)
+            end
+        end
+    end
+    return mx
+end
+
+# ACOV (CS)
+function rmat!(mx, θ, ::AbstractMatrix, ::ACOV_{CS_}, ::Int)
+    s  = size(mx, 1)
+    ρ  = θ[1]
+    if s > 1
+        for n = 2:s
+            mxnn = mx[n, n]
+            @inbounds @simd for m = 1:n-1
+                mxmm = mx[m, m]
+                mx[m, n] += mxnn * mxmm * ρ 
+            end
+        end
+    end
+    return mx
+end
 ###############################################################################
 ###############################################################################
 ###############################################################################
diff --git a/src/varstruct.jl b/src/varstruct.jl
index 8dd78c7..e09a20d 100644
--- a/src/varstruct.jl
+++ b/src/varstruct.jl
@@ -352,6 +352,9 @@ struct CovStructure{T, T2} <: AbstractCovarianceStructure
                         subjblockdict = sabjcrossdicts(subjblockdict, dicts[i])
                     end
                 end
+                if isa(repeated[1].covtype.s, ACOV_)
+                    @warn "Using ACOV covariance additional effect at first position is meaningless."
+                end
             else
                 subjblockdict = nothing
             end
diff --git a/src/vartypes.jl b/src/vartypes.jl
index 2248d75..83b5fea 100644
--- a/src/vartypes.jl
+++ b/src/vartypes.jl
@@ -32,6 +32,11 @@ struct SPGAUD_ <: AbstractCovarianceType end
 struct UN_ <: AbstractCovarianceType end
 struct ZERO <: AbstractCovarianceType end
 
+struct ACOV_{C <: AbstractCovarianceType} <: AbstractCovarianceType
+    c::C
+    a::Int
+end
+
 ################################################################################
 #                          COVARIANCE TYPE
 ################################################################################
@@ -399,10 +404,20 @@ function Unstructured()
 end
 const UN = Unstructured()
 
+
+"""
+    ACOV(c)
+"""
+function ACOV(c)
+    CovarianceType(ACOV_(c.s, 0))
+end
+
 function RZero()
     CovarianceType(ZERO(), false)
 end
 
+#######################################################################################
+
 function covstrparam(ct::SI_, ::Int)::Tuple{Int, Int}
     return (1, 0)
 end
@@ -446,6 +461,10 @@ function covstrparam(ct::SPPOWD_, ::Int)::Tuple{Int, Int}
     return (2, 1)
 end
 
+function covstrparam(ct::ACOV_{<:Union{CS_, AR_, SPPOW_}}, ::Int)::Tuple{Int, Int}
+    return (0, 1)
+end
+
 function covstrparam(ct::ZERO, ::Int)::Tuple{Int, Int}
     return (0, 0)
 end
@@ -555,6 +574,10 @@ function rcoefnames(s, t, ct::UN_)
     return v
 end
 
+function rcoefnames(s, t, ct::ACOV_{<: Union{CS_, AR_}})
+    return ["ρ "]
+end
+
 function rcoefnames(s, t, ct::AbstractCovarianceType)
     v = Vector{String}(undef, t)
     v .= "Val "
@@ -643,6 +666,9 @@ end
 function Base.show(io::IO, ct::UN_)
     print(io, "UN")
 end
+function Base.show(io::IO, ct::ACOV_)
+    print(io, "ACOV(", ct.c, ")")
+end
 function Base.show(io::IO, ct::ZERO)
     print(io, "No effect")
 end