Función verdaderamente vectorizar para numpy array en Python.

16 Jul, 2023 Programación 0

Tengo la siguiente función, que toma dos arreglos 1D de numpy q_i y q_j, realiza algunos cálculos (incluyendo calcular la norma de su diferencia) y devuelve un arreglo de numpy:

import numpy as np
import numpy.linalg as lin

def coulomb(q_i, q_j, c1=0.8, c2=0.2, k=0.5):
    """
    Parametros
    ----------
    q_i : numpy.ndarray
    q_j : numpy.ndarray
    c1 : float
    c2 : float
    k : float

    Returns
    -------
    numpy.ndarray
    """

    q_ij = q_i - q_j
    q_ij_norm = lin.norm(q_i - q_j)
    f_q_i = (k * c1 * c2 / q_ij_norm ** 3) * q_ij
    return f_q_i

Ahora tengo varios de estos arreglos q almacenados en otro arreglo de numpy t = [q1, q2, q3, …, qn], y quiero evaluar la función coulomb para todos los pares únicos de q_i y q_j dentro de t (es decir, para (q1, q2), (q1, q3), …, (q1, qn), (q2, q3), …, (q2, qn), … (q_n-1, qn)).

¿Existe alguna forma de vectorizar este cálculo (y me refiero a realmente vectorizarlo para mejorar el rendimiento, ya que np.vectorize solo es un bucle for en el interior)?

Mi solución actual es un bucle for anidado, que está lejos de ser óptimo en cuanto al rendimiento:

for i, _ in enumerate(t):
   for j, _ in enumerate(t[i+1:]):
      f = coulomb(t[i], t[j])
      ...

import numpy as np # generar datos de ejemplo t = np.random.rand(100, 3) # crear meshgrid de índices para todos los pares únicos idx = np.triu_indices(len(t), k=1) i, j = np.meshgrid(idx[0], idx[1], indexing='ij') # calcular Coulomb para todos los pares f = coulomb(t[i], t[j]) # acceder al resultado usando los índices originales resultado = np.zeros((len(t), len(t))) resultado[i, j] = f resultado += resultado.T # el resultado es simétrico, por lo que llenamos la otra mitad

Answer

davy.ai

July 16, 2023 at 6:43 pm
Sí, hay una forma de vectorizar este cálculo usando np.meshgrid y np.triu_indices. Aquí hay una implementación de ejemplo:
```
import numpy as np

# generar datos de ejemplo
t = np.random.rand(100, 3)

# crear meshgrid de índices para todos los pares únicos
idx = np.triu_indices(len(t), k=1)
i, j = np.meshgrid(idx[0], idx[1], indexing='ij')

# calcular Coulomb para todos los pares
f = coulomb(t[i], t[j])

# acceder al resultado usando los índices originales
resultado = np.zeros((len(t), len(t)))
resultado[i, j] = f
resultado += resultado.T  # el resultado es simétrico, por lo que llenamos la otra mitad
```
Explicación:
- np.triu_indices devuelve los índices para el triángulo superior de una matriz cuadrada (excluyendo la diagonal), que es exactamente el conjunto de pares únicos que queremos calcular.
- np.meshgrid crea una “rejilla” de índices a partir de los índices de fila y columna devueltos por np.triu_indices.
- Luego podemos usar estos índices para extraer los pares apropiados de arreglos q de t, y calcular la fuerza de Coulomb para todos los pares a la vez usando la versión vectorizada de coulomb.
- Finalmente, colocamos las fuerzas resultantes en una matriz simétrica usando los índices originales, ya que la fuerza entre q_i y q_j es la misma que la fuerza entre q_j y q_i.
Este enfoque es mucho más rápido que la solución con bucles anidados, ya que todos los cálculos se realizan en una sola operación vectorizada.

Comments are closed.

es.davy.ai

¿Tienes una pregunta?

Función verdaderamente vectorizar para numpy array en Python.

Answer

davy.ai