feat: adds multidipole nonlocal spectrum calculations

pdme/util/fast_nonlocal_spectrum.py

@@ -50,3 +50,51 @@ def fast_s_nonlocal(
 		_logger.debug(f"bses: {bses}")
 
 	return alphses1 * alphses2 * bses
+
+
+def fast_s_nonlocal_dipoleses(
+	dot_pair_inputs: numpy.ndarray, dipoleses: numpy.ndarray
+) -> numpy.ndarray:
+	"""
+	No error correction here baby.
+	"""
+	ps = dipoleses[:, :, 0:3]
+	ss = dipoleses[:, :, 3:6]
+	ws = dipoleses[:, :, 6]
+
+	_logger.debug(f"ps: {ps}")
+	_logger.debug(f"ss: {ss}")
+	_logger.debug(f"ws: {ws}")
+
+	r1s = dot_pair_inputs[:, 0, 0:3]
+	r2s = dot_pair_inputs[:, 1, 0:3]
+	f1s = dot_pair_inputs[:, 0, 3]
+	f2s = dot_pair_inputs[:, 1, 3]
+
+	if (f1s != f2s).all():
+		raise ValueError(f"Dot pair frequencies are inconsistent: {dot_pair_inputs}")
+
+	diffses1 = r1s[:, None] - ss[:, None, :]
+	diffses2 = r2s[:, None] - ss[:, None, :]
+	if _logger.isEnabledFor(logging.DEBUG):
+		_logger.debug(f"diffses1: {diffses1}")
+		_logger.debug(f"diffses2: {diffses2}")
+
+	norms1 = numpy.linalg.norm(diffses1, axis=3) ** 3
+	norms2 = numpy.linalg.norm(diffses2, axis=3) ** 3
+	if _logger.isEnabledFor(logging.DEBUG):
+		_logger.debug(f"norms1: {norms1}")
+		_logger.debug(f"norms2: {norms2}")
+
+	alphses1 = numpy.einsum("abcd,acd->abc", diffses1, ps) / norms1
+	alphses2 = numpy.einsum("abcd,acd->abc", diffses2, ps) / norms2
+	if _logger.isEnabledFor(logging.DEBUG):
+		_logger.debug(f"alphses1: {alphses1}")
+		_logger.debug(f"alphses2: {alphses2}")
+
+	bses = (1 / numpy.pi) * (ws[:, None, :] / (f1s[:, None]**2 + ws[:, None, :] ** 2))
+	if _logger.isEnabledFor(logging.DEBUG):
+		_logger.debug(f"bses: {bses}")
+
+	_logger.debug(f"Raw pair calc: [{alphses1 * alphses2 * bses}]")
+	return numpy.einsum("...j->...", alphses1 * alphses2 * bses)

tests/util/test_fast_nonlocal_spectrum_pairs.py

@@ -0,0 +1,53 @@
+import numpy
+import pdme.util.fast_nonlocal_spectrum
+import logging
+import pytest
+
+
+def test_fast_nonlocal_calc_multidipole():
+	d1 = [1, 2, 3, 4, 5, 6, 7]
+	d2 = [1, 2, 3, 4, 5, 6, 8]
+	d3 = [2, 5, 3, 4, -5, -6, 2]
+	d4 = [-3, 2, 1, 4, 5, 6, 10]
+
+	dipoleses = numpy.array([[d1, d2], [d3, d4]])
+
+	dot_pairs = numpy.array(
+		[[[-1, -2, -3, 11], [-1, 2, 5, 11]], [[-1, -2, -3, 6], [2, 4, 6, 6]]]
+	)
+	# expected_ij is for pair i, dipole j
+	expected_11 = 0.000021124454334947546213
+	expected_12 = 0.000022184755131682365135
+	expected_13 = 0.0000053860643617855849275
+	expected_14 = -0.0000023069501696755220764
+	expected_21 = 0.00022356021100884617796
+	expected_22 = 0.00021717277640859343002
+	expected_23 = 0.000017558321044891869169
+	expected_24 = -0.000034714318479634499683
+
+	expected = numpy.array(
+		[
+			[expected_11 + expected_12, expected_21 + expected_22],
+			[expected_13 + expected_14, expected_23 + expected_24],
+		]
+	)
+
+	# this is a bit silly but just set the logger to debug so that the coverage stats don't get affected by the debug statements.
+	pdme.util.fast_nonlocal_spectrum._logger.setLevel(logging.DEBUG)
+
+	numpy.testing.assert_allclose(
+		pdme.util.fast_nonlocal_spectrum.fast_s_nonlocal_dipoleses(dot_pairs, dipoleses),
+		expected,
+		err_msg="nonlocal voltages at dot aren't as expected for dipoleses.",
+	)
+
+
+def test_fast_nonlocal_frequency_check_multidipole():
+	d1 = [1, 2, 3, 4, 5, 6, 7]
+
+	dipoles = numpy.array([[d1]])
+
+	dot_pairs = numpy.array([[[-1, -2, -3, 11], [-1, 2, 5, 10]]])
+
+	with pytest.raises(ValueError):
+		pdme.util.fast_nonlocal_spectrum.fast_s_nonlocal_dipoleses(dot_pairs, dipoles)