From b998dc7b5c15f622405acbe6d40133218514e976 Mon Sep 17 00:00:00 2001
From: zhengp0 <zhengp@uw.edu>
Date: Wed, 14 Feb 2024 14:47:18 -0800
Subject: [PATCH] add degree of freedom computation

---
 setup.cfg              |  2 +-
 src/limetr/__init__.py | 28 ++++++++++++++++++++++++++++
 tests/test_limetr.py   | 16 +++++++++++++++-
 3 files changed, 44 insertions(+), 2 deletions(-)

diff --git a/setup.cfg b/setup.cfg
index f7373f1..14f751e 100644
--- a/setup.cfg
+++ b/setup.cfg
@@ -1,6 +1,6 @@
 [metadata]
 name = limetr
-version = 0.0.7
+version = 0.0.8
 description =
 long_description = file: README.rst
 long_description_content_type = text/x-rst
diff --git a/src/limetr/__init__.py b/src/limetr/__init__.py
index 4440ffe..805eb68 100644
--- a/src/limetr/__init__.py
+++ b/src/limetr/__init__.py
@@ -584,3 +584,31 @@ def get_R2(model: LimeTr):
 
 def get_rmse(model: LimeTr):
     return np.sqrt(get_marginal_rvar(model))
+
+
+def get_degree_of_freedom(model: LimeTr):
+    k = model.k_total
+
+    c_mat = np.empty(shape=(0, k))
+    c_vec = np.empty(shape=(0,))
+    if model.use_uprior:
+        c_mat_sub = np.identity(k)
+        c_mat = np.vstack([c_mat, -c_mat_sub, c_mat_sub])
+        c_vec = np.hstack([c_vec, -model.uprior[0], model.uprior[1]])
+    if model.use_lprior or model.use_constraints:
+        c_mat_sub = model.jacobian(model.soln)
+        c_mat = np.vstack([c_mat, -c_mat_sub, c_mat_sub])
+        c_vec = np.hstack([c_vec, -model.cl, model.cu])
+
+    # determine active constraints
+    index_active = np.isclose(c_mat.dot(model.soln), c_vec)
+    c_mat_active = c_mat[index_active]
+
+    if c_mat_active.size == 0:
+        return k
+
+    lb, ub = c_mat_active.min(axis=0), c_mat_active.max(axis=0)
+    lb[np.isclose(lb, 0.0)] = 0.0
+    ub[np.isclose(ub, 0.0)] = 0.0
+
+    return k - 0.5 * ((lb < 0) + (ub > 0)).sum()
diff --git a/tests/test_limetr.py b/tests/test_limetr.py
index cd8278c..f67ce19 100644
--- a/tests/test_limetr.py
+++ b/tests/test_limetr.py
@@ -1,6 +1,12 @@
 import numpy as np
 import pytest
-from limetr import LimeTr, get_conditional_R2, get_marginal_R2, get_rmse
+from limetr import (
+    LimeTr,
+    get_conditional_R2,
+    get_marginal_R2,
+    get_rmse,
+    get_degree_of_freedom,
+)
 from scipy.linalg import block_diag
 
 
@@ -347,3 +353,11 @@ def test_performance():
     assert isinstance(conditional_r2, float)
     assert isinstance(marginal_r2, float)
     assert isinstance(rmse, float)
+
+
+def test_degree_of_freedom():
+    model = lmtr_test_problem()
+    model.optimize()
+
+    k = get_degree_of_freedom(model)
+    assert k == model.k