refactor: removes redundant calculation and uses pdme

deepdog/direct_monte_carlo/dmc_filters.py

@@ -54,109 +54,13 @@ class SingleDotSpinQubitFrequencyFilter(DirectMonteCarloFilter):
 			self.measurements
 		)
 
-	# oh no not this again
-	def fast_s_spin_qubit_tarucha_apsd_dipoleses(
-		self, dot_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]
-
-		# 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]
-		rs = dot_inputs[:, 0:3]
-		# axes are [measurement_idx]
-		fs = dot_inputs[:, 3]
-
-		# 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
-		diffses = rs[:, None] - ss[:, None, :]
-
-		# norms takes out axis 3, the last one, giving [A, measurement_idx, j]
-		norms = numpy.linalg.norm(diffses, axis=3)
-
-		# _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)
-		alphses = (
-			(
-				3
-				* numpy.transpose(
-					numpy.transpose(
-						numpy.einsum("abcd,acd->abc", diffses, ps) / (norms**2)
-					)
-					* numpy.transpose(diffses)
-				)[:, :, :, 0]
-			)
-			- ps[:, numpy.newaxis, :, 0]
-		) / (norms**3)
-
-		bses = (
-			2
-			* numpy.pi
-			* ws[:, None, :]
-			/ ((2 * numpy.pi * fs[:, None]) ** 2 + 4 * ws[:, None, :] ** 2)
-		)
-
-		return numpy.einsum("...j->...", alphses * alphses * bses)
-
 	def filter_samples(self, samples: ndarray) -> ndarray:
 		current_sample = samples
 		for di, low, high in zip(self.dot_inputs_array, self.lows, self.highs):
 
 			if len(current_sample) < 1:
 				break
-			vals = self.fast_s_spin_qubit_tarucha_apsd_dipoleses(
+			vals = pdme.util.fast_v_calc.fast_efieldxs_for_dipoleses(
 				numpy.array([di]), current_sample
 			)
 			# _logger.info(vals)

tests/direct_monte_carlo/test_eletric_field_x_dmc_filter.py

@@ -0,0 +1,137 @@
+import pdme.measurement
+import pdme.measurement.input_types
+from pdme.model import (
+	LogSpacedRandomCountMultipleDipoleFixedMagnitudeModel,
+	LogSpacedRandomCountMultipleDipoleFixedMagnitudeXYModel,
+	LogSpacedRandomCountMultipleDipoleFixedMagnitudeFixedOrientationModel,
+)
+import deepdog.direct_monte_carlo.dmc_filters
+import numpy.random
+import numpy.testing
+import logging
+
+_logger = logging.getLogger(__name__)
+
+
+def fixed_z_model_func(
+	xmin,
+	xmax,
+	ymin,
+	ymax,
+	zmin,
+	zmax,
+	wexp_min,
+	wexp_max,
+	pfixed,
+	n_max,
+	prob_occupancy,
+):
+	return LogSpacedRandomCountMultipleDipoleFixedMagnitudeFixedOrientationModel(
+		xmin,
+		xmax,
+		ymin,
+		ymax,
+		zmin,
+		zmax,
+		wexp_min,
+		wexp_max,
+		pfixed,
+		0,
+		0,
+		n_max,
+		prob_occupancy,
+	)
+
+
+def get_model(orientation):
+	model_funcs = {
+		"fixedz": fixed_z_model_func,
+		"free": LogSpacedRandomCountMultipleDipoleFixedMagnitudeModel,
+		"fixedxy": LogSpacedRandomCountMultipleDipoleFixedMagnitudeXYModel,
+	}
+	model = model_funcs[orientation](
+		-10,
+		10,
+		-17.5,
+		17.5,
+		5,
+		7.5,
+		-5,
+		6.5,
+		10**3,
+		2,
+		0.99999999,
+	)
+	model.n = 2
+	model.rng = numpy.random.default_rng(1234)
+
+	return (
+		f"connors_geom-5height-orientation_{orientation}-pfixexp_{3}-dipole_count_{2}",
+		model,
+	)
+
+
+def test_electric_field_x_dmc_filter():
+
+	dipoles_raw = [
+		[(1, 2, 3), (4, 5, 6), 1],
+		[(-1, 5, 2), (6, 5, 4), 10],
+	]
+	dipoles = [
+		pdme.measurement.OscillatingDipole(numpy.array(d[0]), numpy.array(d[1]), d[2])
+		for d in dipoles_raw
+	]
+
+	_logger.debug(f"dipoles: {dipoles}")
+	dot_inputs_raw = [
+		([-1, -1, 0], 1),
+		([-1, -1, 0], 2),
+		([-1, -1, 0], 3),
+		([-1, -1, 0], 4),
+	]
+	dot_inputs_array = pdme.measurement.input_types.dot_inputs_to_array(dot_inputs_raw)
+	_logger.debug(f"dot_inputs_array: {dot_inputs_array}")
+
+	arrangement = pdme.measurement.OscillatingDipoleArrangement(dipoles)
+	measurements = []
+	for input in dot_inputs_raw:
+		ex = sum(
+			[
+				dipole.s_electric_fieldx_at_position(*input)
+				for dipole in arrangement.dipoles
+			]
+		)
+		ex_low = ex * 0.5
+		ex_high = ex * 1.5
+		meas = pdme.measurement.DotRangeMeasurement(ex_low, ex_high, input[0], input[1])
+		measurements.append(meas)
+
+	filter = deepdog.direct_monte_carlo.dmc_filters.SingleDotSpinQubitFrequencyFilter(
+		measurements
+	)
+
+	samples = numpy.array(
+		[
+			[
+				[1, 2, 3, 4, 5, 6, 1],
+				[-1, 5, 2, 6, 5, 4, 10],
+			],
+			[
+				[10, 20, 30, 40, 50, 60, 1],
+				[-1, 5, 2, 6, 5, 4, 1],
+			],
+			[
+				[1, 1, 1, 1, 1, 1, 1],
+				[2, 2, 2, 2, 2, 2, 1],
+			],
+		]
+	)
+
+	expected = samples[
+		0:1
+	]  # only expect to see the first guy, because that's what generated our thing
+	filtered = filter.filter_samples(samples)
+	assert len(filtered) != len(samples), "Should have filtered some out!"
+	numpy.testing.assert_array_equal(
+		filtered, expected, "The filter should have only returned the first one"
+	)