From 2510d8d0f27b24f6982f148214e30beb6b618f59 Mon Sep 17 00:00:00 2001
From: Jean-Philippe ARGAUD <jean-philippe.argaud@edf.fr>
Date: Sat, 2 Jan 2021 20:50:34 +0100
Subject: [PATCH] Improvement and extension of EnKF algorithm (VAR)

---
 .../daAlgorithms/EnsembleKalmanFilter.py      |  3 +++
 src/daComposant/daCore/NumericObjects.py      | 26 ++++++++++++++++++-
 2 files changed, 28 insertions(+), 1 deletion(-)

diff --git a/src/daComposant/daAlgorithms/EnsembleKalmanFilter.py b/src/daComposant/daAlgorithms/EnsembleKalmanFilter.py
index d242b6b..b1019ee 100644
--- a/src/daComposant/daAlgorithms/EnsembleKalmanFilter.py
+++ b/src/daComposant/daAlgorithms/EnsembleKalmanFilter.py
@@ -37,6 +37,7 @@ class ElementaryAlgorithm(BasicObjects.Algorithm):
                 "StochasticEnKF",
                 "ETKF",
                 "ETKF-KFF",
+                "ETKF-VAR",
                 ],
             )
         self.defineRequiredParameter(
@@ -158,6 +159,8 @@ class ElementaryAlgorithm(BasicObjects.Algorithm):
             NumericObjects.senkf(self, Xb, Y, U, HO, EM, CM, R, B, Q)
         elif self._parameters["Minimizer"] in ["ETKF", "ETKF-KFF"]:
             NumericObjects.etkf(self, Xb, Y, U, HO, EM, CM, R, B, Q, KorV="KalmanFilterFormula")
+        elif self._parameters["Minimizer"] == "ETKF-VAR":
+            NumericObjects.etkf(self, Xb, Y, U, HO, EM, CM, R, B, Q, KorV="Variational")
         else:
             raise ValueError("Error in Minimizer name: %s"%self._parameters["Minimizer"])
         #
diff --git a/src/daComposant/daCore/NumericObjects.py b/src/daComposant/daCore/NumericObjects.py
index 15300a9..8e847b3 100644
--- a/src/daComposant/daCore/NumericObjects.py
+++ b/src/daComposant/daCore/NumericObjects.py
@@ -26,7 +26,7 @@ __doc__ = """
 __author__ = "Jean-Philippe ARGAUD"
 
 import os, time, copy, types, sys, logging
-import math, numpy, scipy
+import math, numpy, scipy, scipy.optimize
 from daCore.BasicObjects import Operator
 from daCore.PlatformInfo import PlatformInfo
 mpr = PlatformInfo().MachinePrecision()
@@ -914,6 +914,30 @@ def etkf(selfA, Xb, Y, U, HO, EM, CM, R, B, Q, KorV="KalmanFilterFormula"):
             #
             Xn = Xfm.reshape((__n,-1)) + EaX @ ( vw.reshape((__m,-1)) + numpy.sqrt(__m-1) * Tdemi @ mU )
         #--------------------------
+        elif KorV == "Variational":
+            RI     = R.getI()
+            HXfm = H((Xfm, Un)) # Eventuellement Hfm
+            def CostFunction(w):
+                _A  = Ynpu.reshape((__p,-1)) - HXfm.reshape((__p,-1)) - (EaHX @ w).reshape((__p,-1))
+                _J  = 0.5 * (__m-1) * w.T @ w + 0.5 * _A.T * RI * _A
+                return float(_J)
+            def GradientOfCostFunction(w):
+                _A  = Ynpu.reshape((__p,-1)) - HXfm.reshape((__p,-1)) - (EaHX @ w).reshape((__p,-1))
+                _GradJ = (__m-1) * w.reshape((__m,1)) - EaHX.T * RI * _A
+                return numpy.ravel(_GradJ)
+            vw = scipy.optimize.fmin_cg(
+                f           = CostFunction,
+                x0          = numpy.zeros(__m),
+                fprime      = GradientOfCostFunction,
+                args        = (),
+                disp        = False,
+                )
+            #
+            Pta = numpy.linalg.inv( (__m-1)*numpy.eye(__m) + EaHX.T * RI * EaHX )
+            EWa = numpy.real(scipy.linalg.sqrtm((__m-1)*Pta)) # Partie imaginaire ~= 10^-18
+            #
+            Xn = Xfm.reshape((__n,-1)) + EaX @ (vw.reshape((__m,-1)) + EWa)
+        #--------------------------
         else:
             raise ValueError("KorV has to be chosen as either \"KalmanFilterFormula\" or \"Variational\".")
         #
-- 
2.39.2