Adds unrestricted model and adds util for normalising

pdme/model/__init__.py

@@ -1,4 +1,5 @@
 from pdme.model.model import Model
 from pdme.model.fixed_z_plane_model import FixedZPlaneModel
+from pdme.model.unrestricted_model import UnrestrictedModel
 
-__all__ = ["Model", "FixedZPlaneModel"]
+__all__ = ["Model", "FixedZPlaneModel", "UnrestrictedModel"]

pdme/model/model.py

@@ -2,6 +2,7 @@ import numpy
 import scipy.optimize
 from typing import Callable, Sequence
 from pdme.measurement import DotMeasurement
+import pdme.util
 import logging
 
 
@@ -83,4 +84,5 @@ class Model():
 			initial = numpy.tile(initial_pt, self.n())
 
 		result = scipy.optimize.least_squares(self.costs(dots), initial, jac=self.jac(dots), ftol=1e-15, gtol=3e-16, bounds=bounds)
+		result.normalised_x = pdme.util.normalise_point_list(result.x, self.point_length())
 		return result

pdme/model/unrestricted_model.py

@@ -0,0 +1,60 @@
+import numpy
+from pdme.model.model import Model
+from pdme.measurement import DotMeasurement
+
+
+class UnrestrictedModel(Model):
+	'''
+	Model of oscillating dipoles with no restrictions.
+	Additionally, each dipole is assumed to be orientated in the plus or minus z direction.
+
+	Parameters
+	----------
+	n : int
+		The number of dipoles to assume.
+	'''
+	def __init__(self, n: int) -> None:
+		self._n = n
+
+	def __repr__(self) -> str:
+		return f'UnrestrictedModel({self.n()})'
+
+	def point_length(self) -> int:
+		'''
+			Dipole is unconstrained in this model.
+			All seven degrees of freedom: (px, py, pz, sx, sy, sz, w).
+		'''
+		return 7
+
+	def n(self) -> int:
+		return self._n
+
+	def v_for_point_at_dot(self, dot: DotMeasurement, pt: numpy.ndarray) -> float:
+		p = pt[0:3]
+		s = pt[3:6]
+		w = pt[6]
+
+		diff = dot.r - s
+		alpha = p.dot(diff) / (numpy.linalg.norm(diff)**3)
+		b = (1 / numpy.pi) * (w / (w**2 + dot.f**2))
+		return alpha**2 * b
+
+	def jac_for_point_at_dot(self, dot: DotMeasurement, pt: numpy.ndarray) -> numpy.ndarray:
+		p = pt[0:3]
+		s = pt[3:6]
+		w = pt[6]
+
+		diff = dot.r - s
+		alpha = p.dot(diff) / (numpy.linalg.norm(diff)**3)
+		b = (1 / numpy.pi) * (w / (w**2 + dot.f**2))
+
+		p_divs = 2 * alpha * diff / (numpy.linalg.norm(diff)**3) * b
+
+		r_divs = (-p / (numpy.linalg.norm(diff)**3) + 3 * p.dot(diff) * diff / (numpy.linalg.norm(diff)**5)) * 2 * alpha * b
+
+		f2 = dot.f**2
+		w2 = w**2
+
+		w_div = alpha**2 * (1 / numpy.pi) * ((f2 - w2) / ((f2 + w2)**2))
+
+		return numpy.concatenate((p_divs, r_divs, w_div), axis=None)

pdme/util/__init__.py

@@ -0,0 +1,3 @@
+from pdme.util.normal_form import normalise_point_list
+
+__all__ = ["normalise_point_list"]

pdme/util/normal_form.py

@@ -0,0 +1,20 @@
+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 normalise_point_list(pts: numpy.ndarray, pt_length) -> numpy.ndarray:
+	chunked_pts = [flip_chunk_to_positive_px(pts[i: i + pt_length]) for i in range(0, len(pts), pt_length)]
+	range_to_length = list(range(pt_length))
+	rotated_range = range_to_length[pt_length - 1:] + range_to_length[0:pt_length - 1]
+	return numpy.concatenate(sorted(chunked_pts, key=lambda x: tuple(round(val, 3) for val in operator.itemgetter(*rotated_range)(x))), axis=None)