report nrmse etc.

fhchl · Aug 10, 2023 · 359906c · 359906c
1 parent e9693fc
commit 359906c
Show file tree

Hide file tree

Showing 2 changed files with 67 additions and 15 deletions.
diff --git a/dynax/estimation.py b/dynax/estimation.py
@@ -21,7 +21,7 @@
 
 from .evolution import AbstractEvolution
 from .system import DynamicalSystem
-from .util import value_and_jacfwd
+from .util import mse, nmse, nrmse, value_and_jacfwd
 
 
 def _get_bounds(module: eqx.Module) -> tuple[list, list]:
@@ -209,9 +209,16 @@ def residuals(params):
     )
 
     # Unscale mse to Least-Squares cost.
-    res.jac = res.jac * np.sqrt(res_size / 2)  # TODO: check this
+    res.fun = res.fun[: y.size].reshape(y.shape) / np.sqrt(2 / res_size) / weight
+    res.jac = res.jac * np.sqrt(res_size / 2)
     res.cost = res.cost * res_size / 2
 
+    # Compute normalized root-mean-squared error
+    y_pred = y - res.fun
+    res.mse = np.atleast_1d(mse(y, y_pred))
+    res.nmse = np.atleast_1d(nmse(y, y_pred))
+    res.nrmse = np.atleast_1d(nrmse(y, y_pred))
+
     # Compute covariance matrix.
     # pcov = H^{-1} ~= inv(J^T J). Use regularized inverse.
     _, s, VT = svd(res.jac, full_matrices=False)