From 36454d5044d93b6b178e016b84dd59a5ebaf15e2 Mon Sep 17 00:00:00 2001
From: Deepak Mallubhotla <dmallubhotla+github@gmail.com>
Date: Sun, 9 Apr 2023 16:26:08 -0500
Subject: [PATCH] feat: adds fast calc that allows for variable temp

---
 pdme/util/fast_v_calc.py       | 41 ++++++++++++++++++++++++++++++++++
 tests/util/test_fast_v_calc.py | 26 +++++++++++++++++++++
 2 files changed, 67 insertions(+)

diff --git a/pdme/util/fast_v_calc.py b/pdme/util/fast_v_calc.py
index c42b520..833db57 100644
--- a/pdme/util/fast_v_calc.py
+++ b/pdme/util/fast_v_calc.py
@@ -69,6 +69,47 @@ def fast_vs_for_dipoleses(
 	return numpy.einsum("...j->...", ases * bses)
 
 
+def fast_vs_for_asymmetric_dipoleses(
+	dot_inputs: numpy.ndarray, dipoleses: numpy.ndarray, temp: numpy.ndarray
+) -> numpy.ndarray:
+	"""
+	No error correction here baby.
+	Expects dot_inputs to be numpy array of [rx, ry, rz, f] entries, so a n by 4 where n is number of measurement points.
+
+	Dipoleses are expected to be array of arrays of arrays:
+	list of sets of dipoles which are part of a single arrangement to be added together.
+	Within each dipole, the expected format is [px, py, pz, sx, sy, sz, e1, e2, w]
+	The passed in w is expected to be half the actual. This is bad, but here for historical reasons to be changed later.
+	"""
+	raw_ps = dipoleses[:, :, 0:3]
+	ss = dipoleses[:, :, 3:6]
+	e1s = dipoleses[:, :, 6]
+	e2s = dipoleses[:, :, 7]
+	raw_ws = dipoleses[:, :, 8]
+
+	rs = dot_inputs[:, 0:3]
+	fs = dot_inputs[:, 3]
+
+	diffses = rs[:, None] - ss[:, None, :]
+
+	w1s = numpy.exp(-e1s / temp) * raw_ws
+	w2s = numpy.exp(-e2s / temp) * raw_ws
+
+	mag_prefactor = 4 * ((w1s * w2s) / ((w1s + w2s) ** 2))
+	ws = w1s + w2s
+
+	# some annoying broadcast thing here?
+	ps = (raw_ps.T * mag_prefactor.T).T
+
+	norms = numpy.linalg.norm(diffses, axis=3) ** 3
+
+	ases = (numpy.einsum("abcd,acd->abc", diffses, ps) / norms) ** 2
+
+	bses = (1 / numpy.pi) * (ws[:, None, :] / (fs[:, None] ** 2 + ws[:, None, :] ** 2))
+
+	return numpy.einsum("...j->...", ases * bses)
+
+
 def between(a: numpy.ndarray, low: numpy.ndarray, high: numpy.ndarray) -> numpy.ndarray:
 	"""
 	Intended specifically for the case where a is a list of arrays, and each array must be between the single array low and high, but without error checking.
diff --git a/tests/util/test_fast_v_calc.py b/tests/util/test_fast_v_calc.py
index ab86ded..ecbf764 100644
--- a/tests/util/test_fast_v_calc.py
+++ b/tests/util/test_fast_v_calc.py
@@ -107,3 +107,29 @@ def test_between():
 	expected = numpy.array([False, False, True])
 
 	numpy.testing.assert_array_equal(actual, expected, err_msg="Between calc wrong")
+
+
+def test_fast_v_calc_asymmetric_multidipoles_but_symmetric():
+	# expected format is [px, py, pz, sx, sy, sz, e1, e2, w]
+	d1 = [1, 2, 3, 4, 5, 6, 1, 1, 7 / 2]
+	d2 = [2, 5, 3, 4, -5, -6, 2, 2, 2 / 2]
+
+	dipoles = numpy.array([[d1, d2]])
+
+	dot_inputs = numpy.array([[-1, -1, -1, 11], [2, 3, 1, 5.5]])
+	# expected_ij is for dot i, dipole j
+	expected_11 = 0.00001421963287022476
+	expected_12 = 0.00001107180225755457
+	expected_21 = 0.000345021108583681380388722
+	expected_22 = 0.0000377061050587914705139781
+
+	expected = numpy.array([[expected_11 + expected_12, expected_21 + expected_22]])
+
+	numpy.testing.assert_allclose(
+		pdme.util.fast_v_calc.fast_vs_for_asymmetric_dipoleses(
+			dot_inputs, dipoles, 1e10
+		),
+		expected,
+		err_msg="Voltages at dot aren't as expected for multidipole calc.",
+	)
+