From 6af5d6f78c81f56b52e8782da73ffb01a82e3f26 Mon Sep 17 00:00:00 2001
From: Deepak <dmallubhotla+github@gmail.com>
Date: Mon, 30 Aug 2021 16:08:33 -0500
Subject: [PATCH] Adds two dipole tests

---
 pathfinder/model/model.py     |  4 ++--
 tests/model/test_model_sol.py | 44 +++++++++++++++++++++++++++++++++++
 2 files changed, 46 insertions(+), 2 deletions(-)

diff --git a/pathfinder/model/model.py b/pathfinder/model/model.py
index 1e81d01..117f18b 100644
--- a/pathfinder/model/model.py
+++ b/pathfinder/model/model.py
@@ -36,8 +36,8 @@ class DotDipoleModel():
 
 		return jac_to_return
 
-	def sol(self):
-		initial = numpy.zeros(6 * self.n) + .1
+	def sol(self, initial_dipole=(0.1, 0.1, 0.1), initial_position=(.1, .1, .1)):
+		initial = numpy.tile(numpy.append(initial_dipole, initial_position), self.n)
 		if self.m == 6 * self.n:
 			# We are perfectly specified
 			return scipy.optimize.root(self.costs(), initial, jac=self.jac(), tol=1e-12)
diff --git a/tests/model/test_model_sol.py b/tests/model/test_model_sol.py
index 4246b5b..d61c7e5 100644
--- a/tests/model/test_model_sol.py
+++ b/tests/model/test_model_sol.py
@@ -2,6 +2,12 @@ import numpy
 import pathfinder.model
 
 
+def chunk_n_sort(pts):
+	pt_length = 6
+	chunked_pts = [pts[i: i + pt_length] for i in range(0, len(pts), pt_length)]
+	return numpy.sort(chunked_pts, axis=0)
+
+
 def print_result(msg, result):
 	print(msg)
 	print(f"\tResult: {result.x}")
@@ -30,3 +36,41 @@ def test_one_dipole_six_dot():
 	print_result("one 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)
+
+
+def test_two_dipole_twelve_dot():
+	dot_positions = [[x, y, 0] for x in (-4, -1, 1, 4) for y in (-1, -5, -9)]
+	dot_positions.append((1, -1, 2))
+	dot_positions.append((-1, -1, -2))
+	dipole = pathfinder.model.StaticDipole([0, 8, 0], [0, 5, 0])
+	dipole2 = pathfinder.model.StaticDipole([8, 1, 1], [2, 2, -2])
+	expected_result = numpy.array([0, 8, 0, 0, 5, 0, 8, 1, 1, 2, 2, -2])
+	dipole_arrangement = pathfinder.model.StaticDipoleArrangement([dipole, dipole2])
+	dot_measurements = dipole_arrangement.get_dot_measurements(dot_positions)
+
+	model = pathfinder.model.DotDipoleModel(dot_measurements, 2)
+	res = model.sol()
+	print(model.costs()(res.x))
+
+	print_result("two dipole twelve dot", res)
+	assert res.success, "The solution for two dipoles and twelve dots should have succeeded."
+	numpy.testing.assert_allclose(chunk_n_sort(res.x), chunk_n_sort(expected_result), err_msg="Dipole wasn't as expected.", rtol=1e-6, atol=1e-4)
+
+
+def test_two_dipole_twelve_dot_2():
+	dot_positions = [[x, y, 0] for x in (-4, -1, 1, 4) for y in (-1, -5, -9)]
+	dot_positions.append((1, -1, 2))
+	dot_positions.append((-1, -1, -2))
+	dipole = pathfinder.model.StaticDipole([0, 8, 0], [0, 5, 0])
+	dipole2 = pathfinder.model.StaticDipole([8, 1, 1], [2, 2, -2])
+	expected_result = numpy.array([0, 8, 0, 0, 5, 0, 8, 1, 1, 2, 2, -2])
+	dipole_arrangement = pathfinder.model.StaticDipoleArrangement([dipole, dipole2])
+	dot_measurements = dipole_arrangement.get_dot_measurements(dot_positions)
+
+	model = pathfinder.model.DotDipoleModel(dot_measurements, 2)
+	res = model.sol()
+	print(model.costs()(res.x))
+
+	print_result("two dipole twelve dot", res)
+	# assert res.success, "The solution for two dipoles and twelve dots should have succeeded."
+	numpy.testing.assert_allclose(chunk_n_sort(res.x), chunk_n_sort(expected_result), err_msg="Dipole wasn't as expected.", rtol=1e-6, atol=1e-4)