TL: adding local work

Author: tlunet

blockops/blockIterationGenerator.py

@@ -4,13 +4,10 @@
 import re
 
 
-SIMPLE_FORM = False
+SIMPLE_FORM = True
 
 class BlockIterationGenerator():
 
-    def __init__(self):
-        pass
-
     def solve(self, lhs, rhs, u):
         tmp = []
         for i in range(len(u)):
@@ -38,8 +35,8 @@ def createZeros(self, n):
         return sy.zeros(n, 1)
 
     def jacobi(self, A, u, f, w=1):
-        D = A.lower_triangular(0).upper_triangular(0)
-        return self.simplifyElementwise(u + w * D.inv() * (f - A * u))
+        Dinv = sy.diag([d**(-1) for d in A.diagonal()[:]], unpack=True)
+        return self.simplifyElementwise(u + w * Dinv * (f - A * u))
 
     def gausseidel(self, A, u, f, A_c, u_res):
         return self.solve(lhs=A_c * (u_res - u), rhs=f - A * u, u=u_res)
@@ -59,13 +56,13 @@ def checkResults(self, u):
                 tmp = gen.generateBlockRule()
 
                 if SIMPLE_FORM:
-                    dico = self.generateData(6, 4)
+                    dico = self.symbols
                     subDico = {}
                     prol = 1
                     rest = 1
                     for i, (phiOp, chiOp, TCtoF, TFtoC, prop) in enumerate(zip(
                             dico['phi'], dico['chi'], dico['T_c_to_f'], dico['T_f_to_c'],
-                            ['F', 'G', 'H', 'K'])):
+                            ['F', 'G', 'H', 'K', 'L', 'M'])):
                         sym = sy.symbols(prop, commutative=False)
                         subDico[prol*phiOp**(-1)*chiOp*rest] = sym
                         subDico[prol * (phiOp**(-1)*chiOp)**2 * rest] = sym**2
@@ -81,31 +78,34 @@ def checkResults(self, u):
             else:
                 gen.check(expr=u[i], n=i + 1)
 
-    def generateData(self, n, L, pre_s=1, post_s=0):
+    def generateSymbols(self, nBlocks, nLevels, nPreSmooth=1, nPostSmooth=0):
         save_symbols = {}
         u_0 = sy.Symbol(r'u_0', commutative=False)
-        phi = [sy.Symbol(f'\phi_{i}', commutative=False) for i in range(L)]
-        chi = [sy.Symbol(f'\chi_{i}', commutative=False) for i in range(L)]
-        T_c_to_f = [sy.Symbol(f'T_{i + 1}^{i}', commutative=False) for i in range(L)]
-        T_f_to_c = [sy.Symbol(f'T_{i}^{i + 1}', commutative=False) for i in range(L)]
-        A = [sy.Matrix(np.eye(n, dtype=int) * phi[i]) + sy.Matrix(np.eye(n, k=-1, dtype=int) * -chi[i]) for i in
-             range(L)]
+        phi = [sy.Symbol(f'\phi_{i}', commutative=False) for i in range(nLevels)]
+        chi = [sy.Symbol(f'\chi_{i}', commutative=False) for i in range(nLevels)]
+        T_c_to_f = [sy.Symbol(f'T_{i + 1}^{i}', commutative=False) for i in range(nLevels)]
+        T_f_to_c = [sy.Symbol(f'T_{i}^{i + 1}', commutative=False) for i in range(nLevels)]
+        A = [sy.Matrix(np.eye(nBlocks, dtype=int) * phi[i]) + sy.Matrix(np.eye(nBlocks, k=-1, dtype=int) * -chi[i]) for i in
+             range(nLevels)]
         u_k = [
-            [self.createVec('u^0', n=n, ss=save_symbols), self.createVec('u^0', n=n, ss=save_symbols)] if i == 0 else [
-                self.createVec(f'u^0_{i}', n=n, ss=save_symbols),
-                self.createZeros(n)] for i in range(L)]
+            [self.createVec('u^0', n=nBlocks, ss=save_symbols),
+             self.createVec('u^0', n=nBlocks, ss=save_symbols)] if i == 0 else [
+                self.createVec(f'u^0_{i}', n=nBlocks, ss=save_symbols),
+                self.createZeros(nBlocks)] for i in range(nLevels)]
         u_k_1 = [
-            [self.createVec('u^1', n=n, ss=save_symbols), self.createVec('u^1', n=n, ss=save_symbols)] if i == 0 else [
-                self.createVec(f'u^1_{i}', n=n, ss=save_symbols),
-                self.createVec(f'u^1_{i}', n=n, ss=save_symbols)] for i in
-            range(L)]
-        f = [sy.Matrix([[chi[0] * u_0 if i == 0 else 0 for i in range(n)]]).transpose() if i == 0 else None for i in
-             range(L)]
-        pre_smoothing = [pre_s for _ in range(L)]
-        post_smoothing = [post_s for _ in range(L)]
-        return {
-            'L': L,
-            'n': n,
+            [self.createVec('u^1', n=nBlocks, ss=save_symbols),
+             self.createVec('u^1', n=nBlocks, ss=save_symbols)] if i == 0 else [
+                self.createVec(f'u^1_{i}', n=nBlocks, ss=save_symbols),
+                self.createVec(f'u^1_{i}', n=nBlocks, ss=save_symbols)] for i in
+            range(nLevels)]
+        f = [sy.Matrix([[chi[0] * u_0 if i == 0 else 0 for i in range(nBlocks)]]).transpose()
+             if i == 0 else None for i in range(nLevels)]
+        pre_smoothing = [nPreSmooth for _ in range(nLevels)]
+        post_smoothing = [nPostSmooth for _ in range(nLevels)]
+
+        self.symbols = {
+            'L': nLevels,
+            'n': nBlocks,
             'phi': phi,
             'chi': chi,
             'T_c_to_f': T_c_to_f,
@@ -117,21 +117,23 @@ def generateData(self, n, L, pre_s=1, post_s=0):
             'pre_smoothing': pre_smoothing,
             'post_smoothing': post_smoothing,
         }
+        return self.symbols
 
 
 class PararealGenerator(BlockIterationGenerator):
 
-    def __init__(self, n):
-        super().__init__()
-        res = self.parareal(settings=self.generateData(n=n, L=2))
-        self.checkResults(res)
-
-    def parareal(self, settings, overlapping=False):
-        u = settings['u_k']
-        chi = settings['chi']
-        A = settings['A']
-        f = settings['f']
-        u_k_1 = settings['u_k_1']
+    def __init__(self, nBlocks):
+        self.generateSymbols(nBlocks=nBlocks, nLevels=2)
+        self.res = self.parareal()
+        self.checkResults(self.res)
+
+    def parareal(self, overlapping=False):
+        symbols = self.symbols
+        u = symbols['u_k']
+        chi = symbols['chi']
+        A = symbols['A']
+        f = symbols['f']
+        u_k_1 = symbols['u_k_1']
         jac = self.jacobi(u=u[0][1], A=A[0], f=f[0])
         if overlapping:
             jac = self.jacobi(u=jac, A=A[0], f=f[0])
@@ -140,23 +142,25 @@ def parareal(self, settings, overlapping=False):
 
 
 class MultilevelGenerator(BlockIterationGenerator):
-    def __init__(self, n, L, pre_smoothing=1, post_smoothing=1):
-        super().__init__()
-        res = self.multilevel(
-            setting=self.generateData(n=n, L=L, pre_s=pre_smoothing, post_s=post_smoothing))
-        self.checkResults(res)
-
-    def multilevel(self, setting, l=0):
-        L = setting['L'] - 1
-        T_c_f_s = setting['T_c_to_f']
-        T_f_c_s = setting['T_f_to_c']
-        A = setting['A']
-        u_k = setting['u_k']
-        u_k_1 = setting['u_k_1']
-        f = setting['f']
-        pre = setting['pre_smoothing']
-        post = setting['post_smoothing']
-        chi = setting['chi']
+
+    def __init__(self, nBlocks, nLevels, nPreSmooth=1, nPostSmooth=1):
+        self.generateSymbols(nBlocks=nBlocks, nLevels=nLevels, nPreSmooth=nPreSmooth, nPostSmooth=nPostSmooth)
+        self.res = self.multilevel()
+        self.checkResults(self.res)
+
+    def multilevel(self, l=0):
+        symbols = self.symbols
+
+        L = symbols['L'] - 1
+        T_c_f_s = symbols['T_c_to_f']
+        T_f_c_s = symbols['T_f_to_c']
+        A = symbols['A']
+        u_k = symbols['u_k']
+        u_k_1 = symbols['u_k_1']
+        f = symbols['f']
+        pre = symbols['pre_smoothing']
+        post = symbols['post_smoothing']
+        chi = symbols['chi']
 
         state = {}
         state2 = {}
@@ -179,7 +183,7 @@ def multilevel(self, setting, l=0):
                     else:
                         u_k_1[l][1] = self.jacobi(u=u_k_1[l][1], A=A[l], f=f[l])
             f[l + 1] = (T_f_c_s[l] * (f[l] - A[l] * u_k_1[l][1]))
-            self.multilevel(l=l + 1, setting=setting)
+            self.multilevel(l=l+1)
             tmp = self.solve(lhs=u_k_1[l][0],
                              rhs=u_k_1[l][1] + T_c_f_s[l] * u_k_1[l + 1][0],
                              u=u_k_1[l + 1][0]).subs(state2)
@@ -285,5 +289,6 @@ def u(n, k):
         return tmp
 
 
-# PararealGenerator(n=6)
-MultilevelGenerator(n=6, L=3, pre_smoothing=1, post_smoothing=0)
+if __name__ == "__main__":
+    # PararealGenerator(n=4)
+    MultilevelGenerator(nBlocks=7, nLevels=3, nPreSmooth=1, nPostSmooth=0)

blockops/iteration.py

@@ -8,7 +8,6 @@
 import numpy as np
 import sympy as sy
 from typing import Dict
-import time
 
 from blockops.block import BlockOperator, I
 from blockops.run import PintRun
@@ -289,7 +288,7 @@ def plotSchedule(self, N: int, K, nProc: int, schedulerType: str = 'BLOCK-BY-BLO
         pool = TaskPool(run=run)
         schedule = getSchedule(taskPool=pool, nProc=nProc, nPoints=N + 1, schedulerType=schedulerType)
         return schedule.plotPlotly()
-        
+
         # schedule.plot(figName=None if self.name is None else self.name + f' ({schedule.NAME} schedule)',
         #           figSize=figSize, saveFig=saveFig)
 
@@ -318,22 +317,26 @@ def register(cls):
     'explicit': 'F'}
 
 
+def u(n, k):
+    n = "" if n == 0 else f"+{n}" if n > 0 else f"{n}"
+    k = "" if k == 0 else f"+{k}" if k > 0 else f"{k}"
+    return "u_{n"+n+"}^{k"+k+"}"
+
+
 @register
 class Parareal(BlockIteration):
     needApprox = True
 
-    def __init__(self, implicitForm=True, approxPred=True, **blockOps):
-        if implicitForm:
-            B00 = "(phi**(-1)*chi-phiApprox**(-1)*chi) * u_{n}^k"
-            B01 = "phiApprox**(-1)*chi * u_{n}^{k+1}"
-            predictor = "phiApprox**(-1)*chi" if approxPred else None
-        else:
-            B00 = "(F-G) * u_{n}^k"
-            B01 = "G * u_{n}^{k+1}"
-            predictor = "G" if approxPred else None
-        update = f"{B00} + {B01}"
+    def __init__(self, implicitForm=True, approxPred=True, nOverlap=0, **blockOps):
+        F = "(phi**(-1)*chi)" if implicitForm else "F"
+        G = "(phiApprox**(-1)*chi)" if implicitForm else "G"
+        predictor = G if approxPred else None
+
+        overlapTerm = "" if nOverlap == 0 else f"*{F}**{nOverlap}"
+        update = f"({F}-{G}){overlapTerm} {u(0-nOverlap, 0)} + {G} {u(0, 1)}"
         propagator = DEFAULT_PROP['implicit'] if implicitForm \
             else DEFAULT_PROP['explicit']
+        print(update)
         super().__init__(update, propagator, predictor,
                          rules=None, name='Parareal', **blockOps)
 
@@ -344,12 +347,12 @@ class ABJ(BlockIteration):
 
     def __init__(self, implicitForm=True, approxPred=True, **blockOps):
         if implicitForm:
-            B00 = "phiApprox**(-1)*chi * u_{n}^k"
-            B10 = "(I-phiApprox**-1 * phi) * u_{n+1}^{k}"
+            B00 = "phiApprox**(-1)*chi" + u(0, 0)
+            B10 = "(I-phiApprox**-1 * phi)" + u(1, 0)
             predictor = "phiApprox**(-1)*chi" if approxPred else None
         else:
-            B00 = "G * u_{n}^k"
-            B10 = "(I-G*F**(-1)) * u_{n}^{k+1}"
+            B00 = "G" + u(0, 0)
+            B10 = "(I-G*F**(-1))" + u(1, 0)
             predictor = "G" if approxPred else None
         update = f"{B10} + {B00}"
         blockOps['I'] = I
@@ -365,12 +368,12 @@ class ABGS(BlockIteration):
 
     def __init__(self, implicitForm=True, approxPred=True, **blockOps):
         if implicitForm:
-            B01 = "phiApprox**(-1)*chi * u_{n}^{k+1}"
-            B10 = "(I-phiApprox**-1 * phi) * u_{n+1}^{k}"
+            B01 = "phiApprox**(-1)*chi" + u(0, 1)
+            B10 = "(I-phiApprox**-1 * phi)" + u(1, 0)
             predictor = "phiApprox**(-1)*chi" if approxPred else None
         else:
-            B01 = "G * u_{n}^{k+1}"
-            B10 = "(I-G*F**(-1)) * u_{n}^{k+1}"
+            B01 = "G" + u(0, 1)
+            B10 = "(I-G*F**(-1))" + u(1, 0)
             predictor = "G" if approxPred else None
         update = f"{B10} + {B01}"
         blockOps['I'] = I
@@ -384,20 +387,20 @@ def __init__(self, implicitForm=True, approxPred=True, **blockOps):
 class TMG(BlockIteration):
     needCoarse = True
 
-    def __init__(self, coarsePred=True, **blockOps):
-        omega = blockOps.get('omega', 1)
-        blockOps.update({'omega': omega})
+    def __init__(self, coarsePred=True, omega=1, nRelax=1, **blockOps):
         phiC = "TCtoF * phiCoarse**(-1) * TFtoC"
+
         B01 = f"{phiC}*chi"" * u_{n}^{k+1}"
         B00 = f"omega*(phi**(-1)*chi - {phiC}*chi)"" * u_{n}^{k}"
         B10 = f"(1-omega)*(I-{phiC}*phi)"" * u_{n+1}^{k}"
+
         predictor = f"{phiC}*chi" if coarsePred else None
         update = f"{B10} + {B01} + {B00}"
         blockOps['I'] = I
         rules = [("TFtoC * TCtoF", I)]
         propagator = DEFAULT_PROP['implicit']
         super().__init__(update, propagator, predictor,
-                         rules=rules, name='TMG', **blockOps)
+                         rules=rules, name='TMG', omega=1, **blockOps)
         self.omega = omega
 
 

blockops/problem.py

@@ -71,7 +71,7 @@ class BlockProblem(ParamClass):
     **schemeArgs :
         Additional keyword arguments used for the time-discretization scheme.
     """
-    def __init__(self, lam, tEnd, nBlocks, scheme, u0=1, **schemeArgs):
+    def __init__(self, lam, tEnd, nBlocks, scheme, u0=1.0, **schemeArgs):
 
         # Initialize parameters
         self.initialize(locals())
@@ -114,6 +114,10 @@ def __init__(self, lam, tEnd, nBlocks, scheme, u0=1, **schemeArgs):
         if np.size(lam) == 1:
             self.u0 = self.u0.squeeze(axis=0)
 
+    @property
+    def lamDt(self):
+        return self.lam*self.dt
+
     @property
     def points(self):
         return self.scheme.points
@@ -166,7 +170,7 @@ def setApprox(self, scheme, **schemeArgs):
             raise ValueError(f'{scheme} scheme is not implemented')
         self.schemeApprox = SCHEMES[scheme](**schemeArgs)
         self.phiApprox, _ = self.schemeApprox.getBlockOperators(
-            self.lam*self.dt, r'\tilde{\phi}', r'\tilde{\chi}')
+            self.lamDt, r'\tilde{\phi}', r'\tilde{\chi}')
         self.propApprox = self.phiApprox**(-1) * self.chi
 
         # Eventually set the coarse approximate block operators
@@ -186,7 +190,7 @@ def approxIsSet(self):
     # -------------------------------------------------------------------------
     # Method for coarse level operators
     # -------------------------------------------------------------------------
-    def setCoarseLevel(self, nPoints, **schemeArgs):
+    def setCoarseLevel(self, nPoints, tType="GEOM", **schemeArgs):
         # Retrieve parameters and BlockScheme class from fine level
         params = self.scheme.getParamsValue()
         params.update(schemeArgs)
@@ -196,11 +200,11 @@ def setCoarseLevel(self, nPoints, **schemeArgs):
         # Build coarse block operators
         self.schemeCoarse = BlockScheme(**params)
         self.phiCoarse, self.chiCoarse = self.schemeCoarse.getBlockOperators(
-                self.lam*self.dt, r'\phi_C', r'\chi_C')
+                self.lamDt, r'\phi_C', r'\chi_C')
 
         # Build transfer operators
         TFtoC, TCtoF = self.scheme.getTransferMatrices(
-            self.pointsCoarse, vectorized=self.nLam > 1)
+            self.pointsCoarse, self.lamDt, tType=tType)
         self.TFtoC = BlockOperator('T_F^C', matrix=TFtoC, cost=0)
         self.TCtoF = BlockOperator('T_C^F', matrix=TCtoF, cost=0)
 
@@ -355,7 +359,7 @@ def getError(self, uNum='fine', uRef='exact') -> np.ndarray:
     # -------------------------------------------------------------------------
     # Method for block iterations
     # -------------------------------------------------------------------------
-    def getBlockIteration(self, algo: str) -> BlockIteration:
+    def getBlockIteration(self, algo:str, **algoArgs) -> BlockIteration:
         """
         Generate a block iteration object associated to the block problem.
 
@@ -370,24 +374,27 @@ def getBlockIteration(self, algo: str) -> BlockIteration:
             The block iteration object.
         """
         try:
+
             BlockIter = ALGORITHMS[algo]
             if BlockIter.needApprox and not self.approxIsSet:
                 raise ValueError(f'{algo} need an approximate block operator')
             if BlockIter.needCoarse and not self.coarseIsSet:
                 raise ValueError(f'{algo} need a coarse block operator')
-            blockIter = BlockIter(
+
+            blockIter = BlockIter(**algoArgs,
                 phi=self.phi, phiApprox=self.phiApprox, chi=self.chi,
                 phiCoarse=self.phiCoarse, chiCoarse=self.chiCoarse,
                 TFtoC=self.TFtoC, TCtoF=self.TCtoF,
                 phiCoarseApprox=self.phiCoarseApprox)
             blockIter.prob = self
+
             return blockIter
+
         except KeyError:
             raise NotImplementedError(
                 f'block iteration for {algo} not implemented')
 
 
-
 if __name__ == '__main__':
     # Quick module testing