IMP logo
IMP Reference Guide  2.20.1
The Integrative Modeling Platform
core/model_numpy.py

Demonstration using NumPy to manipulate IMP data.

1 ## \example core/model_numpy.py
2 # Demonstration using NumPy to manipulate IMP data.
3 #
4 
5 import IMP.container
6 import sys
7 
8 IMP.setup_from_argv(sys.argv, "model numpy")
9 
10 if not IMP.IMP_KERNEL_HAS_NUMPY:
11  print("IMP was not built with NumPy support.")
12  sys.exit(0)
13 
14 # Make 100 xyz particles randomly distributed in a cubic box
15 m = IMP.Model()
18  IMP.algebra.Vector3D(10,10,10))
19 for i in range(100):
20  p = m.add_particle("p")
21  sc.add(p)
24  d.set_coordinates_are_optimized(True)
25 
26 # Get all pairs of particles that are close to each other
28 m.update()
29 
30 # contents is a NumPy 2xN array of the indices of each close particle pair
31 contents = nbl.get_contents()
32 
33 # Get direct access to the coordinates and radii of all particles in the Model.
34 # These are NumPy views, not copies; they share memory with the Model. This
35 # allows us very efficient read-write access to the data, with a caveat: if
36 # the Model memory is reallocated (typically when particles or attributes are
37 # added) any views become invalid, so we should not keep them around long-term.
38 # These arrays are indexed by particle indices.
39 xyz, r = m.get_spheres_numpy()
40 
41 import numpy
42 
43 # Once we have coordinates and particle indices, we can use any NumPy
44 # function on them. Here we calculate the Cartesian distance between each
45 # pair of particles in the close pair container.
46 dists = numpy.linalg.norm(numpy.diff(xyz[contents], axis=1), axis=2).flatten()
47 
48 print("Distribution of center-center particle distances:")
49 print("Min %.2f Max %.2f Mean %.2f" % (dists.min(), dists.max(), dists.mean()))