Optimización del cálculo de la distancia de Python al tiempo que se tienen en cuenta las condiciones de contorno periódicas

He escrito una secuencia de comandos de Python para calcular la distancia entre dos puntos en el espacio 3D y tener en cuenta las condiciones de contorno periódicas. El problema es que necesito hacer este cálculo para muchos, muchos puntos y el cálculo es bastante lento. Aquí está mi función.

def PBCdist(coord1,coord2,UC): dx = coord1[0] - coord2[0] if (abs(dx) > UC[0]*0.5): dx = UC[0] - dx dy = coord1[1] - coord2[1] if (abs(dy) > UC[1]*0.5): dy = UC[1] - dy dz = coord1[2] - coord2[2] if (abs(dz) > UC[2]*0.5): dz = UC[2] - dz dist = np.sqrt(dx**2 + dy**2 + dz**2) return dist 

Entonces llamo a la función como tal

 for i, coord2 in enumerate(coordlist): if (PBCdist(coord1,coord2,UC) < radius): do something with i 

Recientemente leí que puedo boost mucho el rendimiento utilizando la comprensión de lista. Lo siguiente funciona para el caso no PBC, pero no para el caso PBC

 coord_indices = [i for i, y in enumerate([np.sqrt(np.sum((coord2-coord1)**2)) for coord2 in coordlist]) if y < radius] for i in coord_indices: do something 

¿Hay alguna manera de hacer el equivalente de esto para el caso de PBC? ¿Hay alguna alternativa que funcione mejor?

Debería escribir su función distance() de manera que pueda vectorizar el bucle sobre los 5711 puntos. La siguiente implementación acepta una matriz de puntos como el parámetro x0 o x1 :

 def distance(x0, x1, dimensions): delta = numpy.abs(x0 - x1) delta = numpy.where(delta > 0.5 * dimensions, delta - dimensions, delta) return numpy.sqrt((delta ** 2).sum(axis=-1)) 

Ejemplo:

 >>> dimensions = numpy.array([3.0, 4.0, 5.0]) >>> points = numpy.array([[2.7, 1.5, 4.3], [1.2, 0.3, 4.2]]) >>> distance(points, [1.5, 2.0, 2.5], dimensions) array([ 2.22036033, 2.42280829]) 

El resultado es la matriz de distancias entre los puntos pasados ​​como segundo parámetro a la distance() y cada punto en points .

 import numpy as np bounds = np.array([10, 10, 10]) a = np.array([[0, 3, 9], [1, 1, 1]]) b = np.array([[2, 9, 1], [5, 6, 7]]) min_dists = np.min(np.dstack(((a - b) % bounds, (b - a) % bounds)), axis = 2) dists = np.sqrt(np.sum(min_dists ** 2, axis = 1)) 

Aquí a y b son listas de vectores entre los que desea calcular la distancia y los bounds son los límites del espacio (de modo que aquí las tres dimensiones van de 0 a 10 y luego se ajustan). Calcula las distancias entre a[0] y b[0] , a[1] y b[1] , y así sucesivamente.

Estoy seguro de que los expertos en numpy podrían hacerlo mejor, pero esto probablemente será un orden de magnitud más rápido de lo que estás haciendo, ya que la mayor parte del trabajo ahora se realiza en C.

Eche un vistazo al tutorial de alto rendimiento de python de Ian Ozsvalds. Contiene un montón de sugerencias sobre dónde puede ir a continuación.

Incluso:

  • vectorización
  • cython
  • pypy
  • numexpr
  • pycuda
  • multiprocesamiento
  • python paralelo

He encontrado que meshgrid es muy útil para generar distancias. Por ejemplo:

 import numpy as np row_diff, col_diff = np.meshgrid(range(7), range(8)) radius_squared = (row_diff - x_coord)**2 + (col_diff - y_coord)**2 

Ahora tengo una matriz ( radius_squared ) donde cada entrada especifica el cuadrado de la distancia desde la posición de la matriz [x_coord, y_coord] .

Para circularizar la matriz, puedo hacer lo siguiente:

 row_diff, col_diff = np.meshgrid(range(7), range(8)) row_diff = np.abs(row_diff - x_coord) row_circ_idx = np.where(row_diff > row_diff.shape[1] / 2) row_diff[row_circ_idx] = (row_diff[row_circ_idx] - 2 * (row_circ_idx + x_coord) + row_diff.shape[1]) row_diff = np.abs(row_diff) col_diff = np.abs(col_diff - y_coord) col_circ_idx = np.where(col_diff > col_diff.shape[0] / 2) col_diff[row_circ_idx] = (row_diff[col_circ_idx] - 2 * (col_circ_idx + y_coord) + col_diff.shape[0]) col_diff = np.abs(row_diff) circular_radius_squared = (row_diff - x_coord)**2 + (col_diff - y_coord)**2 

Ahora tengo todas las distancias de matrices circularizadas con vector math.