Cómo crear una matriz tridiagonal en Python con NumPy y SciPy

Introducción

En este rápido apunte vamos a ver cómo construir una matriz tridiagonal en Python utilizando NumPy y SciPy. Una matriz tridiagonal es una matriz cuadrada que solamente tiene elementos distintos de cero en su diagonal principal y en las dos diagonales adyacentes a esta (la superdiagonal y la subdiagonal). Las matrices tridiagonales aparecen mucho en cálculo numérico, por ejemplo en la discretización de ecuaciones diferenciales, y tienen la característica de ser matrices dispersas (en lugar de densas) al ser la mayoría de sus elementos cero.
Sin que sirva de precedente, hoy vamos a escribir código que sea compatible tanto con Python 2 como con Python 3. Es un cambio nimio, pero merece la pena ir acostumbrándose a pensar que tarde o temprano habrá que abandonar Python 2 🙂
En esta entrada se ha usado python 2.7.3, numpy 1.6.1 y scipy 0.10.1 y es compatible con python 3.2.3.

Problema matemático

Como ejemplo, vamos a construir la matriz del sistema de ecuaciones diferenciales ordinarias que resulta de discretizar la ecuación del calor unidimensional
$latex {displaystyle frac{partial u}{partial t} = mathcal{L}(u) = frac{partial^2 u}{partial x^2}}$
en el dominio $latex x in [0, 1]$. Aproximando el operador $latex mathcal{L}$ mediante diferencias centradas de orden 2 en los puntos de colocación $latex x_0,, x_1,, dots,, x_N$, se obtiene para los puntos interiores
$latex {displaystyle frac{d u_j}{d t} = frac{1}{Delta x^2} (u_{j + 1} – 2 u_j + u_{j – 1})}$
que, expresado matricialmente, queda
$latex {displaystyle frac{d U}{d t} = A U + b}$
con
$latex U = begin{pmatrix} u_1 \ u_2 \ vdots \ u_{N – 1} end{pmatrix}$
y
$latex
A = begin{pmatrix}
-2 & 1 & 0 & 0 & dots & 0 & 0 & 0 \
1 & -2 & 1 & 0 & dots & 0 & 0 & 0 \
0 & 1 & -2 & 1 & dots & 0 & 0 & 0 \
vdots & vdots & vdots & vdots & ddots & vdots & vdots & vdots \
0 & 0 & 0 & 0 & dots & 0 & 1 & -2
end{pmatrix}
$
Esta matriz es la que vamos a construir.

Matriz tridiagonal

Para crear la matriz tridiagonal, vamos a utilizar la función spdiags del módulo sparse de SciPy, que contiene diversas funciones para trabajar con matrices dispersas. Los argumentos de esta función son una lista de las diagonales de la matriz y una lista de posiciones donde colocar cada una de esas diagonales. En este caso nos interesan las posiciones k = 0 (diagonal principal), k = 1 (superdiagonal) y k = -1 (subdiagonal).
Nota: En SciPy 0.11 se introdujo scipy.sparse.diags, que es el nuevo método recomendado para construir matrices disperas a partir de las diagonales.
Para escribir código compatible con Python 3, al principio del programa escribiremos la línea
[sourcecode language=”python”]
from __future__ import print_function
[/sourcecode]
que sustituye la sentencia print de Python 2 por la nueva función print de Python 3. El módulo __future__ de Python 2 contiene varias sentencias que importan características de versiones de Python posteriores, de tal manera que podemos aprovechar nuevas posibilidades o simplemente hacer nuestro código más portable. Para lo que vamos a hacer hoy no necesitamos más.
En primer lugar construimos las tres diagonales, haciendo uso de la función ones de NumPy que nos da un array lleno de unos, y utilizando la función vstack de Numpy las apilamos (de ahí el nombre) por filas:
[sourcecode language=”python”]
# Número de puntos de colocación
N = 100
dl = np.ones(N – 1)
du = np.ones(N – 1)
d0 = -2 * np.ones(N – 1)
d = np.vstack((dl, d0, du))
[/sourcecode]
Ahora simplemente tenemos que llamar a la función spdiags para construir la matriz:
[sourcecode language=”python”]
A = sparse.spdiags(d, (-1, 0, 1), N – 1, N – 1)
[/sourcecode]
El código quedará finalmente de esta manera:
[sourcecode language=”python”]
# coding: utf-8
#
# Crea una matriz tridiagonal
# Juan Luis Cano Rodríguez <juanlu001@gmail.com>
from __future__ import print_function
import numpy as np
from scipy import sparse
# Número de puntos de colocación
N = 100
dl = du = np.ones(N – 1)
d0 = -2 * np.ones(N – 1)
d = np.vstack((dl, d0, du))
A = sparse.spdiags(d, (-1, 0, 1), N – 1, N – 1)
print(A.todense())
[/sourcecode]
Si lo ejecutamos desde IPython:
[sourcecode language=”python”]
In [1]: %run tridiag.py
[[-2. 1. 0. …, 0. 0. 0.]
[ 1. -2. 1. …, 0. 0. 0.]
[ 0. 1. -2. …, 0. 0. 0.]
…,
[ 0. 0. 0. …, -2. 1. 0.]
[ 0. 0. 0. …, 1. -2. 1.]
[ 0. 0. 0. …, 0. 1. -2.]]
In [2]: A
Out[2]:
<99×99 sparse matrix of type ‘<class ‘numpy.float64′>’
with 295 stored elements (3 diagonals) in DIAgonal format>
In [3]: type(A)
Out[3]: scipy.sparse.dia.dia_matrix
In [4]: sparse.dia?
Type: module
Base Class: <class ‘module’>
String Form:<module ‘scipy.sparse.dia’ from ‘/usr/lib/python3.2/site-packages/scipy/sparse/dia.py’>
Namespace: Interactive
File: /usr/lib/python3.2/site-packages/scipy/sparse/dia.py
Docstring: Sparse DIAgonal format
[/sourcecode]
¡Y ya está! 🙂 Espero que te haya resultado útil, no dudes mandarnos tus dudas y sugerencias así como de difundir el artículo.
¡Un saludo!