1 files changed, 45 insertions, 0 deletions
diff --git a/Wrappers/Python/ccpi/plugins/regularisers.py b/Wrappers/Python/ccpi/plugins/regularisers.py
index 6ed9fb2..ef79231 100644
--- a/Wrappers/Python/ccpi/plugins/regularisers.py
+++ b/Wrappers/Python/ccpi/plugins/regularisers.py
@@ -91,6 +91,51 @@ class FGP_TV(Function):
             out = x.copy()
             out.fill(res)
         return out
+
+class FGP_dTV(Function):
+    def __init__(self, refdata, regularisation_parameter, iterations,
+                 tolerance, eta_const, methodTV, nonneg, device='cpu'):
+        # set parameters
+        self.lambdaReg = regularisation_parameter
+        self.iterationsTV = iterations
+        self.tolerance = tolerance
+        self.methodTV = methodTV
+        self.nonnegativity = nonneg
+        self.device = device # string for 'cpu' or 'gpu'
+        self.refData = np.asarray(refdata.as_array(), dtype=np.float32)
+        self.eta = eta_const
+        
+    def __call__(self,x):
+        # evaluate objective function of TV gradient
+        EnergyValTV = TV_ENERGY(np.asarray(x.as_array(), dtype=np.float32), np.asarray(x.as_array(), dtype=np.float32), self.lambdaReg, 2)
+        return 0.5*EnergyValTV[0]
+    def proximal(self,x,tau, out=None):
+        pars = {'algorithm' : FGP_dTV, \
+               'input' : np.asarray(x.as_array(), dtype=np.float32),\
+                'regularization_parameter':self.lambdaReg*tau, \
+                'number_of_iterations' :self.iterationsTV ,\
+                'tolerance_constant':self.tolerance,\
+                'methodTV': self.methodTV ,\
+                'nonneg': self.nonnegativity ,\
+                'eta_const' : self.eta,\
+                'refdata':self.refData}
+       #inputData, refdata, regularisation_parameter, iterations,
+       #              tolerance_param, eta_const, methodTV, nonneg, device='cpu' 
+        res , info = regularisers.FGP_dTV(pars['input'],
+              pars['refdata'], 
+              pars['regularization_parameter'],
+              pars['number_of_iterations'],
+              pars['tolerance_constant'],
+              pars['eta_const'], 
+              pars['methodTV'],
+              pars['nonneg'],
+              self.device)
+        if out is not None:
+            out.fill(res)
+        else:
+            out = x.copy()
+            out.fill(res)
+        return out
         
 class SB_TV(Function):
     def __init__(self,lambdaReg,iterationsTV,tolerance,methodTV,printing,device):