Support general p-Norms in GradientNormController

nifty5/minimization/iteration_controllers.py

@@ -89,34 +89,42 @@ class GradientNormController(IterationController):
     name : str, optional
         if supplied, this string and some diagnostic information will be
         printed after every iteration
+    p : float
+        Order of norm, default is the 2-Norm (p=2)
     """
 
     def __init__(self, tol_abs_gradnorm=None, tol_rel_gradnorm=None,
-                 convergence_level=1, iteration_limit=None, name=None):
+                 convergence_level=1, iteration_limit=None, name=None, p=2):
         self._tol_abs_gradnorm = tol_abs_gradnorm
         self._tol_rel_gradnorm = tol_rel_gradnorm
         self._convergence_level = convergence_level
         self._iteration_limit = iteration_limit
         self._name = name
+        self._p = p
 
     def start(self, energy):
         self.energyhistory = []
         self._itcount = -1
         self._ccount = 0
         if self._tol_rel_gradnorm is not None:
-            self._tol_rel_gradnorm_now = self._tol_rel_gradnorm \
-                                       * energy.gradient_norm
+            self._tol_rel_gradnorm_now = self._tol_rel_gradnorm * self._norm(energy)
         return self.check(energy)
 
+    def _norm(self, energy):
+        # FIXME Only p=2 norm is cached in energy class
+        if self._p == 2:
+            return energy.gradient_norm
+        return energy.gradient.norm(self._p)
+
     def check(self, energy):
         self._itcount += 1
 
         inclvl = False
         if self._tol_abs_gradnorm is not None:
-            if energy.gradient_norm <= self._tol_abs_gradnorm:
+            if self._norm(energy) <= self._tol_abs_gradnorm:
                 inclvl = True
         if self._tol_rel_gradnorm is not None:
-            if energy.gradient_norm <= self._tol_rel_gradnorm_now:
+            if self._norm(energy) <= self._tol_rel_gradnorm_now:
                 inclvl = True
         if inclvl:
             self._ccount += 1
@@ -128,7 +136,7 @@ class GradientNormController(IterationController):
             logger.info(
                 "{}: Iteration #{} energy={:.6E} gradnorm={:.2E} clvl={}"
                 .format(self._name, self._itcount, energy.value,
-                        energy.gradient_norm, self._ccount))
+                        self._norm(energy), self._ccount))
         self.energyhistory.append(energy.value)
 
         # Are we done?