Threading

Vetorização da Iteração Red-Black

O motivo de vetorizarmos um programa que usa objetos do módulo numpy é que as operações vetoriais e matriciais são encaminhadas para rotinas compilada escritas em C, Fortran, ou até mesmo assembly, que, dependendendo da situação, podem ser ordens de magnitude mais rápidas do que o código correspondente em Python nativo.

Pensando na atualização Red-Black, conceitualmente, todos os pontos de cada cor podem ser atulizados simultaneamente. Vamos usar as figuras abaixo para ajudar a definir alguns conjuntos que vão facilitar a notação. Cada célula representa um nó, e as extremidades estão coloridas de amarelo para mostrar que não estão sendo coloridas.

Número par de nós Red-Black-Even

Número ímpar de nós Red-Black-Even

Os conjuntos de índices são, o conjunto dos nós vermelhos

$$ R = 1,3,\ldots, M_R$$ e $$ B = 2, 4, \ldots, M_B, $$

onde $M_R$ e $M_B$ são os valores máximos do índice de cada subconjunto.

Por inspeção direta, podemos ver que $M_R = N-3$ e $M_B = N-2$ se $N$ é par e $M_R = N-2$ e $M_B = N-3$ se $N$ é ímpar, onde $N$ é o número total de nós.

Vamos também definir os conjuntos imediatamente à esquerda e à direita de cada um destes,

$$ R_{{}+1} = 2,4,\ldots, M_R+1, \quad R_{{}-1} = 0,2,\ldots, M_R-1$$ e $$ B_{{}+1} = 3,5,\ldots, M_B+1, \quad B_{{}-1} = 1,3,\ldots, M_B-1.$$

A atualização de cada ponto é dada por $$ u_i = \frac{u_{i-1} + u_{i+1} - f_i \Delta x^2}{2},$$

adotando a notação de slices do Matlab ou Python, podemos escrever esta atualização simultânea para todos os nós do conjunto vermelho como

$$ u(R) = \frac{1}{2}\left(u(R_{-1}) + u(R_{+1}) - f(R) \Delta x^2\right),$$

e do conjunto preto como

$$ u(B) = \frac{1}{2}\left(u(B_{-1}) + u(B_{+1}) - f(B) \Delta x^2\right).$$

Um ciclo com estas duas atualizações atualiza todos os pontos da malha.

É importante notar que, apesar de dizermos que conceitualmente, a atualização de todos os nós de um conjunto é simultânea, em uma implementação prática com a linguagem Python, por exemplo, o a atualização dentro de cada conjunto é sequencial, feita em uma ordem interna determinada pela implementação dos operadores do Numpy.

Vamos ver uma implementação desta atualização para ver se funciona, depois nos preocupamos com o desempenho. O arquivo de apoio foi simplificado um pouco para diminuir a confusão. A principal inclusão é de uma classe que computa os slices dos vetores para implementação vetorial.

"""
gssup04.py

This just organizes a few support functions for the GS iteration.
This version introduces the Red-Black order range.
"""

import logging
import numpy       as     np
from   math        import pi
from   itertools   import chain

__format = "%(asctime)s: %(message)s"
logging.basicConfig(format=__format, level=logging.INFO, datefmt="%H:%M:%S")

#%%
def zero_rhs(x, L):
    """Sets up a null vector for the RHS of the equation"""
    return np.zeros_like(x)

def sin_rhs(x, L, A):
    """Sets up a rhs that produces a single period of a senoidal function with
       amplitude A as the solution."""
    f = 2*pi/L
    return -f*f*A*np.sin(f*x)

def cub_rhs(x, L, A):
    """Sets up rhs to have a cubic solution with zeros on 0, L/2, L/3
       and value A on L.
       Don't forget to set up the correct boundary conditions."""
    return -9*(A*L - 3*A*x)/(L*L*L)

#%%
class Problem:
    def __init__(self, L, nnode, U0=None, UL=None, calc_rhs=zero_rhs, **kwargs):
        """
        L: float      -- Length of the domain
        nnode: int    -- Number of nodes
        U0, UL: float -- Dirichlet boundary conditions
        calc_rhs function(x, L, *args) -- computes the rhs of the equation
        """
        self.x = np.linspace(0, L, nnode)
        self.dx = self.x[1] - self.x[0]
        self.dx2 = self.dx*self.dx
        self.u = np.zeros_like(self.x)
        self.L = L
        self.nn = nnode
        if U0 is not None: self.u[0] = U0 
        if UL is not None: self.u[-1] = UL 
        self.f = calc_rhs(self.x, self.L, **kwargs)
        self.hist = None

#%%
class RBSlice:
    """
    Returns array slices for Red-Black update.
    """
    def __init__(self, size):
        Mb, Mr = (size-3, size-2) if size%2 else (size-2, size-3)
        self.R = slice(1,Mr+1,2)
        self.Rm = slice(0,Mr,2)
        self.Rp = slice(2,Mr+2,2)
        self.B = slice(2,Mb+1,2)
        self.Bm = slice(1,Mb,2)
        self.Bp = slice(3,Mb+2,2)

#%%
class RBRange:
    """
    Returns a range **with step 1** in the red-black ordering,
    skipping over 0 and n-1, pretty much just useful for GS
    iterations.
    """
    def __init__(self, size):
        self._s = size
        self.reset()

    def __iter__(self):
        return self

    def reset(self):
        r = range(self._s)
        self.c = chain(r[1:-1:2], r[2:-1:2])

    def __next__(self):
        try:
           return next(self.c)
        except StopIteration:
            self.reset()
            raise StopIteration

def info(*args, **kwargs):
    """Directly logs info calls, 99% of what we're doing."""
    logging.info(*args, **kwargs)

A implementação vetorial pode ser vista no programa abaixo. Note que deixamos apenas a implementação Red-Black normal e a vetorial, para simplificar o programa.

"""
Comparison of convergence rates between forward, backward and
symmetric GS iterations
"""

import numpy  as np
import gssup04 as gs
import matplotlib.pyplot as plt

gs.info("Setting up problems.")
A = 4.0
L = 3.0

nn = 22
maxiter = 1000
tol = 5e-4

probs_rbv = { 'linear': gs.Problem(L, nn, UL=A, calc_rhs=gs.zero_rhs),
          'sinusoidal': gs.Problem(L, nn, calc_rhs=gs.sin_rhs, A=A),
               'cubic': gs.Problem(L, nn, UL=A, calc_rhs=gs.cub_rhs, A=A)}
probs_rb = {  'linear': gs.Problem(L, nn, UL=A, calc_rhs=gs.zero_rhs),
          'sinusoidal': gs.Problem(L, nn, calc_rhs=gs.sin_rhs, A=A),
               'cubic': gs.Problem(L, nn, UL=A, calc_rhs=gs.cub_rhs, A=A)}

rbrange = gs.RBRange(nn)
orders = {'rb': (probs_rb, rbrange)}

#%% Conventional RB iteration
for order, (prob, prob_range) in orders.items():
    gs.info("Running {} ordering.".format(order))
    for name, p in prob.items():
        gs.info("Solving {} problem.".format(name))
        err = np.zeros(maxiter, dtype=np.float64)
        converged = False
        for it in range(maxiter):
            uold = p.u.copy()
            for i in prob_range:
                p.u[i] = 0.5*(p.u[i-1] + p.u[i+1] - p.dx2*p.f[i])  
            err[it] = np.linalg.norm(uold - p.u)/np.linalg.norm(p.u)
            if err[it] <= tol:
                converged = True
                gs.info('Problem {} converged after {} iterations'.\
                                                              format(name, it))
                break
        if not converged:
            gs.info('Problem {} failed to converge after {} iterations'.\
                                                              format(name, it))
        err = np.resize(err, it+1)
        p.hist = err

#%% Vector RB iteration
s = gs.RBSlice(probs_rbv['linear'].x.size) 
order = 'rbw'
prob = probs_rbv
gs.info("Running {} ordering.".format(order))
for name, p in prob.items():
    gs.info("Solving {} problem.".format(name))
    err = np.zeros(maxiter, dtype=np.float64)
    converged = False
    for it in range(maxiter):
        uold = p.u.copy()
        p.u[s.R] = 0.5*(p.u[s.Rm]+p.u[s.Rp] - p.dx2*p.f[s.R])
        p.u[s.B] = 0.5*(p.u[s.Bm]+p.u[s.Bp] - p.dx2*p.f[s.B]) 
        err[it] = np.linalg.norm(uold - p.u)/np.linalg.norm(p.u)
        if err[it] <= tol:
            converged = True
            gs.info('Problem {} converged after {} iterations'.\
                                                          format(name, it))
            break
    if not converged:
        gs.info('Problem {} failed to converge after {} iterations'.\
                                                          format(name, it))
    err = np.resize(err, it+1)
    p.hist = err
orders['rbv'] = (probs_rbv, None)       

figs = { 'linear': (0,3),
     'sinusoidal': (1,4),
          'cubic': (2,5)}

styles = {'rb': 'r+',
         'rbv': 'b'}
for name in probs_rb.keys():
    plt.figure(figs[name][0])
    plt.title("Solution for problem {}".format(name))
    plt.figure(figs[name][1])
    plt.yscale('log')
    plt.title("{} problem convergence history".format(name))

# Plot graphs correctly
for order, (prob, _) in orders.items():
    for name, p in prob.items():
        plt.figure(figs[name][0])
        plt.plot(p.x, p.u, styles[order])     
        plt.figure(figs[name][1])
        plt.plot(p.hist, styles[order]) 

Vamos agora ver uma comparação dos tempos de execução para um único problema, para um número fixo de iterações, em um contexto bem simplificado para evitar distrações.

"""
Comparison of convergence rates between forward, backward and
symmetric GS iterations
"""

import gssup04 as     gs
from   timeit  import timeit

gs.info("Setting up problems.")
A = 4.0
L = 3.0

nn = 1000
maxiter = 100
tol = 5e-4

prbv =  gs.Problem(L, nn, calc_rhs=gs.sin_rhs, A=A)
prb =  gs.Problem(L, nn, calc_rhs=gs.sin_rhs, A=A)

rbrange = gs.RBRange(nn)
def run_rb():
    for it in range(maxiter):
        for i in rbrange:
            prb.u[i] = 0.5*(prb.u[i-1] + prb.u[i+1] - prb.dx2*prb.f[i])  

#%% Vector RB iteration
s = gs.RBSlice(prbv.x.size) 
def run_rbv():
    for it in range(maxiter):
        prbv.u[s.R] = 0.5*(prbv.u[s.Rm]+prbv.u[s.Rp] - prbv.dx2*prbv.f[s.R])
        prbv.u[s.B] = 0.5*(prbv.u[s.Bm]+prbv.u[s.Bp] - prbv.dx2*prbv.f[s.B]) 

gs.info("Running  Classical RB iteration.")
trb = timeit(run_rb, number=1000)
gs.info("Time for classic Red-Black: {} seconds".format(trb))
gs.info("Running Vector RB iteration.")
trbv = timeit(run_rbv, number=1000)    
gs.info("Time for vector Red-Black: {} seconds".format(trbv))
gs.info("Ratio Classic/Vector: {}".format(trb/trbv))

É interessante rodar o programa acima para vários tamanhos de vetor. No meu laptop pessoal, para problemas de tamanho 1000, obtive os resultados abaixo.

40:24: Time for classic Red-Black: 106.02983521402348 seconds
40:24: Running Vector RB iteration.
40:26: Time for vector Red-Black: 1.5082596670254134 seconds
40:26: Ratio Classic/Vector: 70.29945673952503

Acho que isto deixa bem claro que tentar paralelizar qualquer coisa que não seja vetorizada é uma tolice sem tamanho.