From 66cd0aa2d53b22b3efae1fa5f264c3c137a4120f Mon Sep 17 00:00:00 2001
From: PharmCat <v.s.arnautov@yandex.ru>
Date: Fri, 19 Jan 2024 22:36:48 +0300
Subject: [PATCH] types

---
 src/reml.jl     | 9 ++++++---
 src/vartypes.jl | 8 ++++----
 test/test.jl    | 4 +++-
 3 files changed, 13 insertions(+), 8 deletions(-)

diff --git a/src/reml.jl b/src/reml.jl
index 02dbaece..5b8500c9 100644
--- a/src/reml.jl
+++ b/src/reml.jl
@@ -124,7 +124,8 @@ function reml_sweep_β(lmm, data, θ::Vector{T}; maxthreads::Int = 16) where T #
             logdetθ₂ = logdet(rθs₂)
             =#
             β       .= NaN
-            return   Inf, β, Inf, Inf, false
+            θ₂      .= NaN
+            return   Inf, β, θs₂, Inf, false
         end
         # θ₃
         #@inbounds @simd for i = 1:n
@@ -184,7 +185,8 @@ function reml_sweep_β_nlopt(lmm, data, θ::Vector{T}; maxthreads::Int = 16) whe
         noerror = all(erroracc)
         if !noerror
             β = fill(NaN, lmm.rankx)
-            return   Inf, β, Inf, Inf, false
+            θ₂ .= NaN
+            return   Inf, β, θ₂, Inf, false
         end
         θ₁   = sum(accθ₁)
         θ₂tc = sum(accθ₂)
@@ -194,7 +196,8 @@ function reml_sweep_β_nlopt(lmm, data, θ::Vector{T}; maxthreads::Int = 16) whe
         ldθ₂, info = LinearAlgebra.LAPACK.potrf!('U', θ₂tc)
         if info != 0
             β = fill(NaN, lmm.rankx)
-            return   Inf, β, Inf, Inf, false
+            θ₂ .= NaN
+            return   Inf, β, θ₂, Inf, false
         end
         LinearAlgebra.LAPACK.potrs!('U', θ₂tc, βtc)
         β = βtc
diff --git a/src/vartypes.jl b/src/vartypes.jl
index ba30346a..1b4707b4 100644
--- a/src/vartypes.jl
+++ b/src/vartypes.jl
@@ -37,11 +37,11 @@ struct ZERO <: AbstractCovarianceType end
 Make covariance type with AbstractCovmatMethod.
 
 """
-struct CovarianceType
-    s::AbstractCovarianceType
+struct CovarianceType{T <: AbstractCovarianceType}
+    s::T
     z::Bool
-    function CovarianceType(s::AbstractCovarianceType, z::Bool)
-        new(s, z)
+    function CovarianceType(s::T, z::Bool) where T <: AbstractCovarianceType
+        new{T}(s, z)
     end
     function CovarianceType(s::AbstractCovarianceType)
         CovarianceType(s,  true)
diff --git a/test/test.jl b/test/test.jl
index 1489afa0..54bc17cb 100644
--- a/test/test.jl
+++ b/test/test.jl
@@ -14,7 +14,9 @@ include("testdata.jl")
     random = Metida.VarEffect(Metida.@covstr(formulation|nosubj), Metida.DIAG),
     )
     Metida.fit!(lmm)
-    @test Metida.m2logreml(lmm) ≈ 25.129480634331067 atol=1E-6
+    @test Metida.m2logreml(lmm, ) ≈ 25.129480634331067 atol=1E-6
+    # Test -2 reml for provided theta
+    @test Metida.m2logreml(lmm, Metida.theta(lmm)) ≈ 25.129480634331067 atol=1E-6
 
     # Casuistic case - random
     lmm = Metida.LMM(@formula(var~sequence+period+formulation), df0;