fix: fixes the broken implementation of the tarucha frequency calculation

pdme/util/fast_nonlocal_spectrum.py

@@ -100,6 +100,132 @@ def fast_s_nonlocal_dipoleses(
 	return numpy.einsum("...j->...", alphses1 * alphses2 * bses)
 
 
+def fast_s_spin_qubit_tarucha_nonlocal_dipoleses(
+	dot_pair_inputs: numpy.ndarray, dipoleses: numpy.ndarray
+) -> numpy.ndarray:
+	"""
+	No error correction here baby.
+	"""
+
+	# We're going to annotate the indices on this class.
+	# Let's define some indices:
+	# A -> index of dipoleses configurations
+	# j -> within a particular configuration, indexes dipole j
+	# measurement_index -> if we have 100 frequencies for example, indexes which one of them it is
+	# If we need to use numbers, let's use A -> 2, j -> 10, measurement_index -> 9 for consistency with
+	# my other notes
+
+	# axes are [dipole_config_idx A, dipole_idx j, {px, py, pz}3]
+	ps = dipoleses[:, :, 0:3]
+	# axes are [dipole_config_idx A, dipole_idx j, {sx, sy, sz}3]
+	ss = dipoleses[:, :, 3:6]
+	# axes are [dipole_config_idx A, dipole_idx j, w], last axis is just 1
+	ws = dipoleses[:, :, 6]
+
+	_logger.debug(f"ps: {ps}")
+	_logger.debug(f"ss: {ss}")
+	_logger.debug(f"ws: {ws}")
+
+	# dot_index is either 0 or 1 for dot1 or dot2
+	# hopefully this adhoc grammar is making sense, with the explicit labelling of the values of the last axis in cartesian space
+	# axes are [measurement_idx, {dot_index}, {rx, ry, rz}] where the inner {dot_index} is gone
+	# [measurement_idx, cartesian3]
+	r1s = dot_pair_inputs[:, 0, 0:3]
+	r2s = dot_pair_inputs[:, 1, 0:3]
+	# axes are [measurement_idx]
+	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}")
+
+	# first operation!
+	# r1s has shape [measurement_idx, rxs]
+	# None inserts an extra axis so the r1s[:, None] has shape
+	# [measurement_idx, 1]([rxs]) with the last rxs hidden
+	#
+	# ss has shape [ A, j, {sx, sy, sz}3], so second term has shape [A, 1, j]([sxs])
+	# these broadcast from right to left
+	# [	 measurement_idx, 1, rxs]
+	# [A,	  1,			   j, sxs]
+	# resulting in [A, measurement_idx, j, cart3] sxs rxs are both cart3
+	diffses1 = r1s[:, None] - ss[:, None, :]
+	diffses2 = r2s[:, None] - ss[:, None, :]
+	if _logger.isEnabledFor(logging.DEBUG):
+		_logger.warning(f"diffses1: {diffses1}")
+		_logger.warning(f"diffses2: {diffses2}")
+
+	# norms takes out axis 3, the last one, giving [A, measurement_idx, j]
+	norms1 = numpy.linalg.norm(diffses1, axis=3)
+	norms2 = numpy.linalg.norm(diffses2, axis=3)
+	if _logger.isEnabledFor(logging.DEBUG):
+		_logger.warning(f"norms1: {norms1}")
+		_logger.warning(f"norms2: {norms2}")
+
+	# _logger.info(f"norms1: {norms1}")
+	# _logger.info(f"norms1 shape: {norms1.shape}")
+	#
+	# diffses1 (A, measurement_idx, j, xs)
+	# ps:  (A, j, px)
+	# result is (A, measurement_idx, j)
+	# intermediate_dot_prod = numpy.einsum("abcd,acd->abc", diffses1, ps)
+	# _logger.info(f"dot product shape: {intermediate_dot_prod.shape}")
+
+	# transpose makes it (j, measurement_idx, A)
+	# transp_intermediate_dot_prod = numpy.transpose(numpy.einsum("abcd,acd->abc", diffses1, ps) / (norms1**3))
+
+	# transpose of diffses has shape (xs, j, measurement_idx, A)
+	# numpy.transpose(diffses1)
+	# _logger.info(f"dot product shape: {transp_intermediate_dot_prod.shape}")
+
+	# inner transpose is (j, measurement_idx, A) * (xs, j, measurement_idx, A)
+	# next transpose puts it back to (A, measurement_idx, j, xs)
+	# p_dot_r_times_r_term = 3 * numpy.transpose(numpy.transpose(numpy.einsum("abcd,acd->abc", diffses1, ps) / (norms1**3)) * numpy.transpose(diffses1))
+	# _logger.info(f"p_dot_r_times_r_term: {p_dot_r_times_r_term.shape}")
+
+	# only x axis puts us at (A, measurement_idx, j)
+	# p_dot_r_times_r_term_x_only = p_dot_r_times_r_term[:, :, :, 0]
+	# _logger.info(f"p_dot_r_times_r_term_x_only.shape: {p_dot_r_times_r_term_x_only.shape}")
+
+	# now to complete the numerator we subtract the ps, which are (A, j, px):
+	# slicing off the end gives us (A, j), so we newaxis to get (A, 1, j)
+	# _logger.info(ps[:, numpy.newaxis, :, 0].shape)
+	alphses1 = (
+		(
+			3
+			* numpy.transpose(
+				numpy.transpose(
+					numpy.einsum("abcd,acd->abc", diffses1, ps) / (norms1**2)
+				)
+				* numpy.transpose(diffses1)
+			)[:, :, :, 0]
+		)
+		- ps[:, numpy.newaxis, :, 0]
+	) / (norms1**3)
+	alphses2 = (
+		(
+			3
+			* numpy.transpose(
+				numpy.transpose(
+					numpy.einsum("abcd,acd->abc", diffses2, ps) / (norms2**2)
+				)
+				* numpy.transpose(diffses2)
+			)[:, :, :, 0]
+		)
+		- ps[:, numpy.newaxis, :, 0]
+	) / (norms2**3)
+	if _logger.isEnabledFor(logging.DEBUG):
+		_logger.warning(f"alphses1: {alphses1}")
+		_logger.warning(f"alphses2: {alphses2}")
+
+	bses = (1 / numpy.pi) * (ws[:, None, :] / (f1s[:, None] ** 2 + ws[:, None, :] ** 2))
+	if _logger.isEnabledFor(logging.DEBUG):
+		_logger.warning(f"bses: {bses}")
+
+	_logger.warning(f"Raw pair calc: [{alphses1 * alphses2 * bses}]")
+	return numpy.einsum("...j->...", alphses1 * alphses2 * bses)
+
+
 def signarg(x, **kwargs):
 	"""
 	uses numpy.sign to implement Arg for real numbers only. Should return pi for negative inputs, 0 for positive.

tests/util/__snapshots__/test_fast_nonlocal_spectrum_pairs.ambr

@@ -11,3 +11,15 @@
     ]),
   ])
 # ---
+# name: test_fast_spin_qubit_frequency_tarucha_calc
+  list([
+    list([
+      0.24485375860291592,
+      0.009068657726033923,
+    ]),
+    list([
+      0.0001499877991005419,
+      0.00011239982757520363,
+    ]),
+  ])
+# ---

tests/util/test_fast_nonlocal_spectrum_pairs.py

@@ -3,6 +3,8 @@ import pdme.util.fast_nonlocal_spectrum
 import logging
 import pytest
 
+_logger = logging.getLogger(__name__)
+
 
 def test_fast_nonlocal_calc_multidipole():
 	d1 = [1, 2, 3, 4, 5, 6, 7]
@@ -74,3 +76,23 @@ def test_fast_nonlocal_calc_multidipole_phase_snapshot(snapshot):
 		pdme.util.fast_nonlocal_spectrum.fast_s_nonlocal_dipoleses(dot_pairs, dipoleses)
 	)
 	assert actual_phases.tolist() == snapshot
+
+
+def test_fast_spin_qubit_frequency_tarucha_calc(snapshot):
+	d1 = [1, 2, 3, 0, 0, 0, 5]
+	d2 = [6, 7, 8, 5, 4, 3, 8]
+	dipoleses = numpy.array([[d1], [d2]])
+
+	dot_pairs = numpy.array(
+		[[[1, 0, 0, 1], [1, 0, 0, 1]], [[1, 0, 0, 1], [3, 0, 0, 1]]]
+	)
+
+	actual = (
+		pdme.util.fast_nonlocal_spectrum.fast_s_spin_qubit_tarucha_nonlocal_dipoleses(
+			dot_pairs, dipoleses
+		)
+	)
+
+	pdme.util.fast_nonlocal_spectrum._logger.setLevel(logging.DEBUG)
+	_logger.info(actual)
+	assert actual.tolist() == snapshot