Two dipoles, with normalisation

pathfinder/model/oscillating/model.py

@@ -2,6 +2,7 @@ from typing import Callable, Sequence
 import numpy
 import scipy.optimize
 from pathfinder.model.oscillating.dot import DotMeasurement
+import pathfinder.model.oscillating.util
 
 
 class DotOscillatingDipoleModel():
@@ -39,4 +40,6 @@ class DotOscillatingDipoleModel():
 	def sol(self, initial_dipole=(0.1, 0.1, 0.1), initial_position=(.1, .1, .1), initial_frequency=1, use_root=True):
 		initial = numpy.tile(numpy.concatenate((initial_dipole, initial_position, initial_frequency), axis=None), self.n)
 
-		return scipy.optimize.least_squares(self.costs(), initial, jac=self.jac(), ftol=1e-15, gtol=3e-16)
+		result = scipy.optimize.least_squares(self.costs(), initial, jac=self.jac(), ftol=1e-15, gtol=3e-16)
+		result.pathfinder_x = pathfinder.model.oscillating.util.normalize_oscillating_dipole_list(result.x)
+		return result

pathfinder/model/oscillating/oscillating_dipole.py

@@ -3,7 +3,7 @@ import numpy
 import numpy.typing
 from typing import Sequence, List, Tuple
 from pathfinder.model.oscillating.dot import DotMeasurement
-
+# import pathfinder.model.oscillating.util
 
 DotInput = Tuple[numpy.typing.ArrayLike, float]
 

pathfinder/model/oscillating/util.py

@@ -0,0 +1,19 @@
+import numpy
+import operator
+
+
+# flips px, py, pz
+SIGN_ARRAY = numpy.array((-1, -1, -1, 1, 1, 1, 1))
+
+
+def flip_chunk_to_positive_px(pt: numpy.ndarray) -> numpy.ndarray:
+	if pt[0] > 0:
+		return pt
+	else:
+		return SIGN_ARRAY * pt
+
+
+def normalize_oscillating_dipole_list(pts: numpy.ndarray) -> numpy.ndarray:
+	pt_length = 7
+	chunked_pts = [flip_chunk_to_positive_px(pts[i: i + pt_length]) for i in range(0, len(pts), pt_length)]
+	return numpy.concatenate(sorted(chunked_pts, key=operator.itemgetter(0, 1, 2, 3, 4, 5, 6)), axis=None)

tests/model/oscillating/test_osc_model_sol.py

@@ -3,15 +3,15 @@ import pathfinder.model.oscillating
 
 
 def chunk_n_sort(pts):
-	pt_length = 6
+	pt_length = 7
 	chunked_pts = [pts[i: i + pt_length] for i in range(0, len(pts), pt_length)]
 
-	return sorted(chunked_pts, key=lambda x: x[0])
+	return chunked_pts
 
 
 def print_result(msg, result):
 	print(msg)
-	print(f"\tResult: {chunk_n_sort(result.x)}")
+	print(f"\tResult: {result.pathfinder_x}")
 	print(f"\tSuccess: {result.success}. {result.message}")
 	try:
 		print(f"\tFunc evals: {result.nfev}")
@@ -26,8 +26,8 @@ def print_result(msg, result):
 def test_one_dipole_six_dot_two_frequencies():
 	# setup
 	dot_inputs = [
-		([0, 0, .01], 5), ([-1, 0, -.01], 5), ([-2, 0, 0], 5), ([-1, -1, 0], 5), ([-1, -1, 1], 5), ([0, -2, 0], 5),
-		([0, 0, .01], 1), ([-1, 0, -.01], 1), ([-2, 0, 0], 1), ([-1, -1, 0], 1), ([-1, -1, 1], 1), ([0, -2, 0], 1)
+		([0, 0, .01], 5), ([-1, 0, -.01], 5), ([-2, 0, -.01], 5), ([0, -1, .01], 5), ([-1, -1, 0], 5), ([-2, -1, 0], 5),
+		([0, 0, .01], 1), ([-1, 0, -.01], 1), ([-2, 0, -.01], 1), ([0, -1, .01], 1), ([-1, -1, 0], 1), ([-2, -1, 0], 1),
 	]
 	dipole = pathfinder.model.oscillating.OscillatingDipole([1, 2, 3], [0, 4, -1], 7)
 	expected_result = numpy.array([1, 2, 3, 0, 4, -1, 7])
@@ -39,4 +39,26 @@ def test_one_dipole_six_dot_two_frequencies():
 
 	print_result("one oscillating dipole six dots", res)
 	assert res.success, "The solution for a single dipole and six dots should have succeeded."
-	numpy.testing.assert_allclose(res.x, expected_result, err_msg="Dipole wasn't as expected.", rtol=1e-6, atol=1e-6)
+	numpy.testing.assert_allclose(res.pathfinder_x, expected_result, err_msg="Dipole wasn't as expected.", rtol=1e-6, atol=1e-6)
+
+
+def test_two_dipole_six_dot_two_frequencies():
+	# setup
+	dot_inputs = [
+		([0, 0, .01], 5), ([-1, 0, -.01], 5), ([-2, 0, -.01], 5), ([0, -1, .01], 5), ([-1, -1, 0], 5), ([-2, -1, 0], 5),
+		([0, -2, .01], 5), ([-1, -2, -.01], 5), ([-2, -2, -.01], 5), ([0, -3, .01], 5), ([-1, -3, 0], 5), ([-2, -3, 0], 5),
+		([0, 0, .01], 1), ([-1, 0, -.01], 1), ([-2, 0, -.01], 1), ([0, -1, .01], 1), ([-1, -1, 0], 1), ([-2, -1, 0], 1),
+		([0, -2, .01], 1), ([-1, -2, -.01], 1), ([-2, -2, -.01], 1), ([0, -3, .01], 1), ([-1, -3, 0], 1), ([-2, -3, 0], 1),
+	]
+	dipole = pathfinder.model.oscillating.OscillatingDipole([1, 2, 3], [0, 4, -1], 7)
+	dipole2 = pathfinder.model.oscillating.OscillatingDipole([-1, 2, 0], [2, -1, 1], 4)
+	expected_result = numpy.array([1, -2, 0, 2, -1, 1, 4, 1, 2, 3, 0, 4, -1, 7])
+	dipole_arrangement = pathfinder.model.oscillating.OscillatingDipoleArrangement([dipole, dipole2])
+	dot_measurements = dipole_arrangement.get_dot_measurements(dot_inputs)
+
+	model = pathfinder.model.oscillating.DotOscillatingDipoleModel(dot_measurements, 2)
+	res = model.sol()
+
+	print_result("two oscillating dipole six dots", res)
+	assert res.success, "The solution for two dipole and six dots should have succeeded."
+	numpy.testing.assert_allclose(res.pathfinder_x, expected_result, err_msg="Dipole wasn't as expected.", rtol=1e-6, atol=1e-6)

tests/model/oscillating/test_util.py

@@ -0,0 +1,20 @@
+import pathfinder.model.oscillating.util as util
+import numpy
+
+
+def test_util_flip():
+	pt = numpy.array((1, 2, 3, 4, 5, 6, 7))
+	expected = (1, 2, 3, 4, 5, 6, 7)
+	numpy.testing.assert_allclose(util.flip_chunk_to_positive_px(pt), expected, err_msg="Point should remain unchanged.")
+
+
+def test_util_flip_negative():
+	pt = numpy.array((-1, -2, 3, 4, 5, 6, 7))
+	expected = (1, 2, -3, 4, 5, 6, 7)
+	numpy.testing.assert_allclose(util.flip_chunk_to_positive_px(pt), expected, err_msg="Point should remain unchanged.")
+
+
+def test_util_normalise_two():
+	pt = numpy.array((1, 2, 3, 4, 5, 6, 7, -.5, 1, 1.5, 2, 2.5, 3, 3.5))
+	expected = (.5, -1, -1.5, 2, 2.5, 3, 3.5, 1, 2, 3, 4, 5, 6, 7)
+	numpy.testing.assert_allclose(util.normalize_oscillating_dipole_list(pt), expected, err_msg="Sort should have happened.")