Investigate binning: should it scale better and run faster?

[ojeda-e]

On small fixtures like the ones used in #200, or in tests, curvature analysis is already fast (~0.6 ms/frame according to the benchmark results). However, I have been working on large systems (5 microsec simulation time) using MC and I noticed the binning method takes longer than I expected. This problem shows up only when scale increases. This makes sense, because cost scales ~ n_frames × n_atoms. A 5 microsec trajectory with thousands of frames on a large bilayer runs a bit slow, at least from what I expect from binning.

I am wondering if, with a fixed grid size, 2x the number of reference atoms roughly runs in 2x runtime is an indication that binning accumulation is contributing to the runtime. If runtime scales linearly with reference atom count at fixed grid size, per-atom binning accumulation is likely a major cost, then profiling get_z_surface vs the rest of the binning pipeline will confirm this before implementing a solution.

If the per-atom calculation is a significant contributor to binning cost, a likely candidate is the accumulation loop in get_z_surface:

membrane-curvature/membrane_curvature/binning_surface.py

Lines 145 to 150 in 8b59e9e

	for index_l, index_m, index_z in zip(cell_x_floor, cell_y_floor, z_coords):
	try:
	# negative coordinates
	if index_l < 0 or index_m < 0:
	negative_coord_warning.warn(index_l, index_m)
	continue

Even though surface_method='fourier' is the default, and in the benchmarks binning is faster than fourier (expected), an improvement for the binning method is still worth it.

[ojeda-e]

I think I was on the right track here.

In the function get_z_surface I have the code to populate the grid as:

membrane-curvature/membrane_curvature/binning_surface.py

Lines 143 to 157 in 2a5fcf6

	for index_l, index_m, index_z in zip(cell_x_floor, cell_y_floor, z_coords):
	try:
	# negative coordinates
	if index_l < 0 or index_m < 0:
	negative_coord_warning.warn(index_l, index_m)
	continue
	grid_z_coordinates[index_l, index_m] += index_z
	grid_norm_unit[index_l, index_m] += 1
	# too large positive coordinates
	except IndexError:
	positive_coord_warning.warn(index_l, index_m, x_range[1], y_range[1])
	z_surface = normalized_grid(grid_z_coordinates, grid_norm_unit)

which runs once per atom and per frame - in _single_frame in base.py. Meaning, MC runs scale with n_atoms * n_frames. That's why when I calculate with a trajectory of 5 microsecs (2500 frames) and ~900 beads of headgroups and a long trajectory, I experience longer runtimes. Because I am currently doing 2500*900=2.25M operations to fill the grid.

The calculations with binning should run with np arrays instead of looping through atoms one by one.

Additionally, with changes introduced in #238, improving this becomes even more relevant.