Adds model generators

pdme/model/fixed_dipole_model.py

@@ -1,9 +1,10 @@
 import numpy
+import numpy.random
 from dataclasses import dataclass
 from typing import Sequence, Tuple
 import scipy.optimize
 from pdme.model.model import Model
-from pdme.measurement import DotMeasurement
+from pdme.measurement import DotMeasurement, OscillatingDipoleArrangement, OscillatingDipole
 
 
 class FixedDipoleModel(Model):
@@ -26,10 +27,15 @@ class FixedDipoleModel(Model):
 		self.zmax = zmax
 		self.p = p
 		self._n = n
+		self.rng = numpy.random.default_rng()
 
 	def __repr__(self) -> str:
 		return f'FixedDipoleModel({self.xmin}, {self.xmax}, {self.ymin}, {self.ymax}, {self.zmin}, {self.zmax}, {self.p}, {self.n()})'
 
+	def get_dipoles(self, frequency: float) -> OscillatingDipoleArrangement:
+		s_pts = numpy.array((self.rng.uniform(self.xmin, self.xmax), self.rng.uniform(self.ymin, self.ymax), self.rng.uniform(self.zmin, self.zmax)))
+		return OscillatingDipoleArrangement([OscillatingDipole(self.p, s_pts, frequency)])
+
 	def point_length(self) -> int:
 		'''
 			Dipole is constrained magnitude, but free orientation.

pdme/model/fixed_magnitude_model.py

@@ -1,9 +1,10 @@
 import numpy
+import numpy.random
 from dataclasses import dataclass
 from typing import Sequence, Tuple
 import scipy.optimize
 from pdme.model.model import Model
-from pdme.measurement import DotMeasurement
+from pdme.measurement import DotMeasurement, OscillatingDipole, OscillatingDipoleArrangement
 
 
 class FixedMagnitudeModel(Model):
@@ -26,6 +27,7 @@ class FixedMagnitudeModel(Model):
 		self.zmax = zmax
 		self.pfixed = pfixed
 		self._n = n
+		self.rng = numpy.random.default_rng()
 
 	def __repr__(self) -> str:
 		return f'FixedMagnitudeModel({self.xmin}, {self.xmax}, {self.ymin}, {self.ymax}, {self.zmin}, {self.zmax}, {self.n()})'
@@ -37,6 +39,15 @@ class FixedMagnitudeModel(Model):
 		'''
 		return 6
 
+	def get_dipoles(self, frequency: float) -> OscillatingDipoleArrangement:
+		theta = numpy.arccos(self.rng.uniform(-1, 1))
+		phi = self.rng.uniform(0, 2 * numpy.pi)
+		px = self.pfixed * numpy.sin(theta) * numpy.cos(phi)
+		py = self.pfixed * numpy.sin(theta) * numpy.sin(phi)
+		pz = self.pfixed * numpy.cos(theta)
+		s_pts = numpy.array((self.rng.uniform(self.xmin, self.xmax), self.rng.uniform(self.ymin, self.ymax), self.rng.uniform(self.zmin, self.zmax)))
+		return OscillatingDipoleArrangement([OscillatingDipole(numpy.array([px, py, pz]), s_pts, frequency)])
+
 	def n(self) -> int:
 		return self._n
 

pdme/model/model.py

@@ -1,7 +1,7 @@
 import numpy
 import scipy.optimize
 from typing import Callable, Sequence
-from pdme.measurement import DotMeasurement
+from pdme.measurement import DotMeasurement, OscillatingDipoleArrangement
 import pdme.util
 import logging
 
@@ -20,6 +20,9 @@ class Model():
 	def v_for_point_at_dot(self, dot: DotMeasurement, pt: numpy.ndarray) -> float:
 		raise NotImplementedError
 
+	def get_dipoles(self, frequency: float) -> OscillatingDipoleArrangement:
+		raise NotImplementedError
+
 	def cost_for_dot(self, dot: DotMeasurement, pts: numpy.ndarray) -> float:
 		# creates numpy.ndarrays in groups of self.point_length().
 		# Will throw problems for irregular points, but that's okay for now.

pdme/model/unrestricted_model.py

@@ -16,17 +16,18 @@ class UnrestrictedModel(Model):
 	n : int
 		The number of dipoles to assume.
 	'''
-	def __init__(self, xmin: float, xmax: float, ymin: float, ymax: float, zmin: float, zmax: float, n: int) -> None:
+	def __init__(self, xmin: float, xmax: float, ymin: float, ymax: float, zmin: float, zmax: float, max_p: float, n: int) -> None:
 		self.xmin = xmin
 		self.xmax = xmax
 		self.ymin = ymin
 		self.ymax = ymax
 		self.zmin = zmin
 		self.zmax = zmax
+		self.max_p = max_p
 		self._n = n
 
 	def __repr__(self) -> str:
-		return f'UnrestrictedModel({self.xmin}, {self.xmax}, {self.ymin}, {self.ymax}, {self.zmin}, {self.zmax}, {self.n()})'
+		return f'UnrestrictedModel({self.xmin}, {self.xmax}, {self.ymin}, {self.ymax}, {self.zmin}, {self.zmax}, {self.max_p}, {self.n()})'
 
 	def point_length(self) -> int:
 		'''
@@ -101,9 +102,9 @@ class UnrestrictedDiscretisation():
 	num_x: int
 	num_y: int
 	num_z: int
-	max_p: int
 
 	def __post_init__(self):
+		self.max_p = self.model.max_p
 		self.cell_count = self.num_x * self.num_y * self.num_z
 		self.x_step = (self.model.xmax - self.model.xmin) / self.num_x
 		self.y_step = (self.model.ymax - self.model.ymin) / self.num_y