Adds some tests and dot methods and jacobian

pathfinder/model/oscillating/__init__.py

@@ -1,3 +1,4 @@
 from pathfinder.model.oscillating.oscillating_dipole import OscillatingDipole, OscillatingDipoleArrangement
+from pathfinder.model.oscillating.dot import DotMeasurement
 
-__all__ = ['OscillatingDipole', 'OscillatingDipoleArrangement']
+__all__ = ['DotMeasurement', 'OscillatingDipole', 'OscillatingDipoleArrangement']

pathfinder/model/oscillating/dot.py

@@ -24,43 +24,58 @@ class DotMeasurement():
 	def __post_init__(self) -> None:
 		self.r = numpy.array(self.r)
 
-	# def v_for_point(self, pt: numpy.ndarray) -> float:
-	# 	'''
-	# 	Returns the voltage for a static dipole represented as a phase space point.
-	#
-	# 	Parameters
-	# 	----------
-	# 	pt: A length 6 numpy array representing a dipole, as (px, py, pz, sx, sy, sz).
-	#
-	# 	Returns
-	# 	----------
-	# 	Returns the voltage from that dipole at the point.
-	# 	'''
-	# 	p = pt[0:3]  # hardcoded here because chances
-	# 	s = pt[3:6]  # are we'll only ever work in 3d.
-	#
-	# 	diff = self.r - s
-	# 	return p.dot(diff) / (numpy.linalg.norm(diff)**3)
-	#
-	# def cost(self, pts: numpy.ndarray) -> float:
-	# 	# 6 because dipole in 3d has 6 degrees of freedom.
-	# 	pt_length = 6
-	# 	# creates numpy.ndarrays in groups of pt_length.
-	# 	# Will throw problems for irregular points, but that's okay for now.
-	# 	chunked_pts = [pts[i: i + pt_length] for i in range(0, len(pts), pt_length)]
-	# 	return sum(self.v_for_point(pt) for pt in chunked_pts) - self.v
-	#
-	# def jac_pt(self, pt: numpy.ndarray) -> numpy.ndarray:
-	# 	p = pt[0:3]  # hardcoded here because chances
-	# 	s = pt[3:6]  # are we'll only ever work in 3d.
-	#
-	# 	diff = self.r - s
-	#
-	# 	p_divs = diff / (numpy.linalg.norm(diff)**3)
-	#
-	# 	r_divs = -p / (numpy.linalg.norm(diff)**3) + 3 * p.dot(diff) * diff / (numpy.linalg.norm(diff)**5)
-	#
-	# 	return numpy.append(p_divs, r_divs)
+	def v_for_point(self, pt: numpy.ndarray) -> float:
+		'''
+		Returns the voltage for an oscillating dipole represented as a phase space point.
+
+		Parameters
+		----------
+		pt: A length 7 numpy array representing a dipole, as (px, py, pz, sx, sy, sz, w).
+
+		Returns
+		----------
+		Returns the voltage from that dipole at the point.
+		'''
+		p = pt[0:3]  # hardcoded here because chances
+		s = pt[3:6]  # are we'll only ever work in 3d.
+		w = pt[6]
+
+		return (self._alpha(p, s))**2 * self._b(w)
+
+	def _alpha(self, p: numpy.ndarray, s: numpy.ndarray) -> float:
+		diff = self.r - s
+		return p.dot(diff) / (numpy.linalg.norm(diff)**3)
+
+	def _b(self, w: float) -> float:
+		return (1 / numpy.pi) * (w / (self.f**2 + w**2))
+
+	def cost(self, pts: numpy.ndarray) -> float:
+		# 7 because dipole in 3d has 7 degrees of freedom.
+		pt_length = 7
+		# creates numpy.ndarrays in groups of pt_length.
+		# Will throw problems for irregular points, but that's okay for now.
+		chunked_pts = [pts[i: i + pt_length] for i in range(0, len(pts), pt_length)]
+		return sum(self.v_for_point(pt) for pt in chunked_pts) - self.v
+
+	def jac_pt(self, pt: numpy.ndarray) -> numpy.ndarray:
+		p = pt[0:3]  # hardcoded here because chances
+		s = pt[3:6]  # are we'll only ever work in 3d.
+		w = pt[6]
+
+		diff = self.r - s
+		alpha = self._alpha(p, s)
+		b = self._b(w)
+
+		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 = self.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)
 	#
 	# def jac(self, pts: numpy.ndarray) -> numpy.ndarray:
 	# 	# 6 because dipole in 3d has 6 degrees of freedom.

tests/model/oscillating/test_oscdot.py

@@ -0,0 +1,34 @@
+import numpy
+import numpy.testing
+
+import pathfinder.model.oscillating as model
+
+
+def test_dot_v_from_dipole():
+
+	dot_position1 = (-1, -1, -1)
+	dot_frequency1 = 11
+	expected_v1 = 0.00001421963287022476
+	dot = model.DotMeasurement(expected_v1, dot_position1, dot_frequency1)
+
+	# and dipole located at (4, 5, 6) with p=(1, 2, 3) and w = 7
+	pt = numpy.array((1, 2, 3, 4, 5, 6, 7))
+
+	numpy.testing.assert_allclose(dot.v_for_point(pt), dot.v, err_msg="v from dipole at a dot was incorrect!")
+	numpy.testing.assert_allclose(dot.cost(pt), 0, atol=1e-12, err_msg="cost should be zero")
+
+
+def test_jac():
+	expected_jac = [
+		3.742008650059147e-6, 4.4904103800709765e-6, 5.238812110082806e-6,
+		-3.12967996186765e-6, -3.1568945702317167e-6, -3.184109178595783e-6,
+		8.603475350051953e-7
+	]
+
+	dot = model.DotMeasurement(50, (-1, -1, -1), 11)
+
+	# dipole located at (4, 5, 6) with p=(1, 2, 3) and w = 7
+	pt = numpy.array((1, 2, 3, 4, 5, 6, 7))
+	assert len(dot.jac_pt(pt)) == 7
+	numpy.testing.assert_allclose(dot.jac_pt(pt), expected_jac, err_msg="Jac pt doesn't match Mathematica result.")
+	# numpy.testing.assert_allclose(dot.jac(pts), jac_row_target, err_msg="whole row should match")