From b776ff43019c219a5edb5c86497e63a2c32437e1 Mon Sep 17 00:00:00 2001
From: pogudingleb <pogudin.gleb@gmail.com>
Date: Thu, 16 Nov 2023 20:54:00 +0100
Subject: [PATCH] adding x(0) to the difference case when necessary

---
 src/discrete.jl                        | 32 ++++++++++++++++++-
 test/local_identifiability_discrete.jl | 44 ++++++++++++++++++++++++--
 2 files changed, 73 insertions(+), 3 deletions(-)

diff --git a/src/discrete.jl b/src/discrete.jl
index 3e84cafc9..1076ca9e3 100644
--- a/src/discrete.jl
+++ b/src/discrete.jl
@@ -208,6 +208,7 @@ function _assess_local_identifiability_discrete(
 
     @debug "Computing the observability matrix"
     prec = length(dds.x_vars) + length(dds.parameters)
+    @debug "The truncation order is $prec"
 
     # Computing the bound from the Schwartz-Zippel-DeMilo-Lipton lemma
     deg_x = _degree_with_common_denom(values(dds.x_equations))
@@ -265,6 +266,8 @@ function _assess_local_identifiability_discrete(
         end
     end
 
+    @debug "Hit the road Jac: $Jac"
+
     @debug "Computing the result"
     base_rank = LinearAlgebra.rank(Jac)
     result = OrderedDict{Any, Bool}()
@@ -277,6 +280,8 @@ function _assess_local_identifiability_discrete(
             Jac[end - k + 1, 1] =
                 output_derivatives[(str_to_var("loc_aux_$i", dds_ext.poly_ring), x)][1]
         end
+        @debug "And Jac for $(funcs_to_check[i]) is $Jac"
+        @debug "$(LinearAlgebra.rank(Jac))"
         result[funcs_to_check[i]] = LinearAlgebra.rank(Jac) == base_rank
     end
 
@@ -311,6 +316,26 @@ function assess_local_identifiability(
     funcs_to_check = Array{}[],
     known_ic = Array{}[],
     p::Float64 = 0.99,
+    loglevel = Logging.Info,
+)
+    restart_logging(loglevel = loglevel)
+    with_logger(_si_logger[]) do
+        return _assess_local_identifiability(dds,
+            measured_quantities = measured_quantities,
+            funcs_to_check = funcs_to_check,
+            known_ic = known_ic,
+            p = p,
+        )
+    end
+end
+
+
+function _assess_local_identifiability(
+    dds::ModelingToolkit.DiscreteSystem;
+    measured_quantities = Array{ModelingToolkit.Equation}[],
+    funcs_to_check = Array{}[],
+    known_ic = Array{}[],
+    p::Float64 = 0.99,
 )
     if length(measured_quantities) == 0
         if any(ModelingToolkit.isoutput(eq.lhs) for eq in ModelingToolkit.equations(dds))
@@ -330,9 +355,13 @@ function assess_local_identifiability(
 
     dds_aux, conversion = mtk_to_si(dds, measured_quantities)
     if length(funcs_to_check) == 0
+        params = parameters(dds)
+        params_from_measured_quantities = union(
+            [filter(s -> !istree(s), get_variables(y)) for y in measured_quantities]...,
+        )
         funcs_to_check = vcat(
             [x for x in states(dds) if conversion[x] in dds_aux.x_vars],
-            parameters(dds),
+            union(params, params_from_measured_quantities)
         )
     end
     funcs_to_check_ = [eval_at_nemo(x, conversion) for x in funcs_to_check]
@@ -343,6 +372,7 @@ function assess_local_identifiability(
     out_dict = OrderedDict(nemo2mtk[param] => result[param] for param in funcs_to_check_)
     if length(known_ic) > 0
         @warn "Since known initial conditions were provided, identifiability of states (e.g., `x(t)`) is at t = 0 only !"
+        out_dict = OrderedDict(substitute(k, Dict(t => 0)) => v for (k, v) in out_dict)
     end
     return out_dict
 end
diff --git a/test/local_identifiability_discrete.jl b/test/local_identifiability_discrete.jl
index ae08eca4f..90af24f17 100644
--- a/test/local_identifiability_discrete.jl
+++ b/test/local_identifiability_discrete.jl
@@ -119,8 +119,8 @@
         Dict(
             :dds => lv,
             :res => OrderedDict(
-                x1 => true,
-                x2 => true,
+                substitute(x1, Dict(t => 0)) => true,
+                substitute(x2, Dict(t => 0)) => true,
                 a => true,
                 b => true,
                 c => true,
@@ -196,6 +196,46 @@
         ),
     )
 
+    @parameters a b
+    @variables t x1(t) y(t)
+    D = Difference(t; dt = 1.0)
+
+    eqs = [D(x1) ~ a]
+
+    @named kic = DiscreteSystem(eqs)
+    push!(
+        cases,
+        Dict(
+            :dds => kic,
+            :res => OrderedDict(
+                x1 => false,
+                a => true,
+                b => false,
+            ),
+            :y => [y ~ x1 + b],
+            :y2 => [x1 + b],
+            :known_ic => Array{}[],
+            :to_check => Array{}[],
+        ),
+    )
+    push!(
+        cases,
+        Dict(
+            :dds => kic,
+            :res => OrderedDict(
+                substitute(x1, Dict(t => 0)) => true,
+                a => true,
+                b => true,
+            ),
+            :y => [y ~ x1 + b],
+            :y2 => [x1 + b],
+            :known_ic => [x1],
+            :to_check => Array{}[],
+        ),
+    )
+
+
+ 
     for c in cases
         @test assess_local_identifiability(
             c[:dds];