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

\({displaystyle frac{partial u}{partial t} = mathcal{L}(u) = frac{partial^2 u}{partial x^2}}\)

en el dominio \(x in [0, 1]\). Aproximando el operador \(mathcal{L}\) mediante diferencias centradas de orden 2 en los puntos de colocación \(x_0,, x_1,, dots,, x_N\), se obtiene para los puntos interiores

\({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

\({displaystyle frac{d U}{d t} = A U + b}\)

con

\(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

from __future__ import print_function

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:

# 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))

Ahora simplemente tenemos que llamar a la función spdiags para construir la matriz:

A = sparse.spdiags(d, (-1, 0, 1), N - 1, N - 1)

El código quedará finalmente de esta manera:

# 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())

Si lo ejecutamos desde IPython:

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]:
<99x99 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

¡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!

Juan Luis Cano

Estudiante de ingeniería aeronáutica y con pasión por la programación y el software libre. Obsesionado con mejorar los pequeños detalles y con ganas de cambiar el mundo. Divulgando Python en español a través de Pybonacci y la asociación Python España.

More Posts - Website

Follow Me:
TwitterLinkedIn

One thought on “Cómo crear una matriz tridiagonal en Python con NumPy y SciPy

  1. Pingback: Python | TagHall

Leave a Reply

Your email address will not be published. Required fields are marked *