feat: adds tarucha phase calculation, using spin qubit precession rate noise

deepdog/real_spectrum_run.py

@@ -66,7 +66,7 @@ def get_a_result_fast_filter_pairs(input) -> int:
 	return len(current_sample)
 
 
-def get_a_result_fast_filter_pair_phase_only(input) -> int:
+def get_a_result_fast_filter_potential_pair_phase_only(input) -> int:
 	(
 		model,
 		pair_inputs,
@@ -102,6 +102,102 @@ def get_a_result_fast_filter_pair_phase_only(input) -> int:
 	return len(current_sample)
 
 
+def get_a_result_fast_filter_tarucha_spin_qubit_pair_phase_only(input) -> int:
+	(
+		model,
+		pair_inputs,
+		pair_phase_lows,
+		pair_phase_highs,
+		monte_carlo_count,
+		seed,
+	) = input
+
+	def fast_s_spin_qubit_tarucha_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]
+
+		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]
+
+		diffses1 = r1s[:, None] - ss[:, None, :]
+		diffses2 = r2s[:, None] - ss[:, None, :]
+
+		norms1 = numpy.linalg.norm(diffses1, axis=3)
+		norms2 = numpy.linalg.norm(diffses2, axis=3)
+
+		alphses1 = (
+			(
+				3
+				* numpy.transpose(
+					numpy.transpose(
+						numpy.einsum("abcd,acd->abc", diffses1, ps) / (norms1**2)
+					)
+					* numpy.transpose(diffses1)
+				)[:, :, :, 0]
+			)
+			- ps[:, :, 0, numpy.newaxis]
+		) / (norms1**3)
+		alphses2 = (
+			(
+				3
+				* numpy.transpose(
+					numpy.transpose(
+						numpy.einsum("abcd,acd->abc", diffses2, ps) / (norms2**2)
+					)
+					* numpy.transpose(diffses2)
+				)[:, :, :, 0]
+			)
+			- ps[:, :, 0, numpy.newaxis]
+		) / (norms2**3)
+
+		bses = (1 / numpy.pi) * (
+			ws[:, None, :] / (f1s[:, None] ** 2 + ws[:, None, :] ** 2)
+		)
+
+		return numpy.einsum("...j->...", alphses1 * alphses2 * bses)
+
+	rng = numpy.random.default_rng(seed)
+	# TODO: A long term refactor is to pull the frequency stuff out from here. The None stands for max_frequency, which is unneeded in the actually useful models.
+	sample_dipoles = model.get_monte_carlo_dipole_inputs(
+		monte_carlo_count, None, rng_to_use=rng
+	)
+
+	current_sample = sample_dipoles
+
+	for pi, plow, phigh in zip(pair_inputs, pair_phase_lows, pair_phase_highs):
+		if len(current_sample) < 1:
+			break
+
+		###
+		# This should be  abstracted out, but we're going to dump it here for time pressure's sake
+		###
+		# vals = pdme.util.fast_nonlocal_spectrum.signarg(
+		# 	pdme.util.fast_nonlocal_spectrum.fast_s_nonlocal_dipoleses(
+		# 		numpy.array([pi]), current_sample
+		# 	)
+		#
+		vals = pdme.util.fast_nonlocal_spectrum.signarg(
+			fast_s_spin_qubit_tarucha_nonlocal_dipoleses(
+				numpy.array([pi]), current_sample
+			)
+		)
+		current_sample = current_sample[
+			numpy.all(
+				((vals > plow) & (vals < phigh)) | ((vals < plow) & (vals > phigh)),
+				axis=1,
+			)
+		]
+	return len(current_sample)
+
+
 def get_a_result_fast_filter(input) -> int:
 	model, dot_inputs, lows, highs, monte_carlo_count, seed = input
 
@@ -299,6 +395,7 @@ class RealSpectrumRun:
 					seeds = seed_sequence.spawn(self.monte_carlo_cycles)
 
 					if self.use_pair_measurements:
+						_logger.debug("using pair measurements")
 						current_success = sum(
 							pool.imap_unordered(
 								get_a_result_fast_filter_pairs,
@@ -320,9 +417,11 @@ class RealSpectrumRun:
 							)
 						)
 					elif self.use_pair_phase_measurements:
+						_logger.debug("using pair phase measurements")
+						_logger.debug("specifically using tarucha")
 						current_success = sum(
 							pool.imap_unordered(
-								get_a_result_fast_filter_pairs,
+								get_a_result_fast_filter_tarucha_spin_qubit_pair_phase_only,
 								[
 									(
 										model,