Vectorize VEGAS over dimensionality about torchquad

[gomezzz]

@FHof I think you also did this, right?

[FHof]

No, I haven't vectorized the VEGASMap over the dimension everywhere, but this vectorization could be possible.
I think the _add_at_indices function currently only works with 1D tensors, but it could be extended for multiple dimensions.
Here are the current for loops over the dimension:

torchquad/torchquad/integration/vegas_map.py

Lines 58 to 60 in 32a22f2

	for i in range(self.dim):
	ID_i = ID[:, i]
	res[:, i] = self.x_edges[i, ID_i] + self.dx_edges[i, ID_i] * offset[:, i]

torchquad/torchquad/integration/vegas_map.py

Lines 74 to 76 in 32a22f2

	for i in range(self.dim):
	ID_i = ID[:, i]
	jac *= self.N_intervals * self.dx_edges[i][ID_i]

torchquad/torchquad/integration/vegas_map.py

Lines 111 to 114 in 32a22f2

	for i in range(self.dim):
	ID_i = ID[:, i]
	_add_at_indices(self.weights[i], ID_i, jf_vec2)
	_add_at_indices(self.counts[i], ID_i, ones)

torchquad/torchquad/integration/vegas_map.py

Lines 212 to 275 in 32a22f2

	for i in range(self.dim): # Update per dim
	delta_d = delta_weights[i]
	# For each inner edge, determine how many delta_d fit into the
	# accumulated smoothed weights.
	# With torch, CUDA and a high number of points the cumsum operation
	# with float32 precision is too inaccurate which leads to wrong
	# indices, so cast to float64 here.
	delta_d_multiples = astype(
	anp.cumsum(astype(smoothed_weights[i, :-1], "float64"), axis=0)
	/ delta_d,
	"int64",
	)
	# For each number of delta_d multiples in {0, 1, …, N_intervals},
	# determine how many intervals belong to it (num_sw_per_dw)
	# and the sum of smoothed weights in these intervals (val_sw_per_dw)
	dtype_int = delta_d_multiples.dtype
	num_sw_per_dw = anp.zeros(
	[self.N_intervals + 1], dtype=dtype_int, like=delta_d
	)
	_add_at_indices(
	num_sw_per_dw,
	delta_d_multiples,
	anp.ones(delta_d_multiples.shape, dtype=dtype_int, like=delta_d),
	is_sorted=True,
	)
	val_sw_per_dw = anp.zeros(
	[self.N_intervals + 1], dtype=self.dtype, like=delta_d
	)
	_add_at_indices(
	val_sw_per_dw, delta_d_multiples, smoothed_weights[i], is_sorted=True
	)
	# The cumulative sum of the number of smoothed weights per delta_d
	# multiple determines the old inner edges indices for the new inner
	# edges calculation
	indices = anp.cumsum(num_sw_per_dw[:-2], axis=0)
	# d_accu_i is used for the interpolation in the new inner edges
	# calculation when adding it to the old inner edges
	d_accu_i = anp.cumsum(delta_d - val_sw_per_dw[:-2], axis=0)
	# EQ 22
	self.x_edges[i][1:-1] = (
	self.x_edges[i][indices]
	+ d_accu_i / smoothed_weights[i][indices] * self.dx_edges[i][indices]
	)
	finite_edges = anp.isfinite(self.x_edges[i])
	if not anp.all(finite_edges):
	# With float64 precision the delta_d_multiples calculation
	# usually doesn't have rounding errors.
	# If it is nonetheless too inaccurate, few values in
	# smoothed_weights[i][indices] can be zero, which leads to
	# invalid edges.
	num_edges = self.x_edges.shape[1]
	logger.warning(
	f"{num_edges - anp.sum(finite_edges)} out of {num_edges} calculated VEGASMap edges were infinite"
	)
	# Replace inf edges with the average of their two neighbours
	middle_edges = 0.5 * (self.x_edges[i][:-2] + self.x_edges[i][2:])
	self.x_edges[i][1:-1] = anp.where(
	finite_edges[1:-1], self.x_edges[i][1:-1], middle_edges
	)
	if not anp.all(anp.isfinite(self.x_edges[i])):
	raise RuntimeError("Could not replace all infinite edges")
	self.dx_edges[i] = self.x_edges[i][1:] - self.x_edges[i][:-1]

From the functions listed in the issue description, I have only vectorized VEGASStratification._get_indices and VEGASMap._smooth_map over the dimension.

I mainly vectorized the code with respect to values depending on the number of evaluations. With a four-dimensional integrand and a high number of points, this made VEGAS about 300 times faster: FHof#13 (comment)

[gomezzz]

Ok, in any event, I think the performance should for now be quite a lot better. I will close it for now and we can reopen if we find it to be too slow at a later point. Thanks!