Adds a lot of stuff for discretisation

pdme/model/fixed_z_plane_model.py

@@ -1,6 +1,7 @@
 import numpy
 from dataclasses import dataclass
-from typing import Tuple
+from typing import Tuple, Sequence
+import scipy.optimize
 
 from pdme.model.model import Model
 from pdme.measurement import DotMeasurement
@@ -81,7 +82,8 @@ class FixedZPlaneModel(Model):
 @dataclass
 class FixedZPlaneDiscretisation():
 	'''
-	Representation of a discretisation of a FixedZPlaneModel
+	Representation of a discretisation of a FixedZPlaneModel.
+	Also captures a rough maximum value of dipole.
 
 	Parameters
 	----------
@@ -95,19 +97,27 @@ class FixedZPlaneDiscretisation():
 	model: FixedZPlaneModel
 	num_x: int
 	num_y: int
+	max_pz: int
 
 	def __post_init__(self):
 		self.cell_count = self.num_x * self.num_y
 		self.x_step = (self.model.xmax - self.model.xmin) / self.num_x
 		self.y_step = (self.model.ymax - self.model.ymin) / self.num_y
 
-	def bounds(self, index: Tuple[float, float]) -> Tuple:
+	def bounds(self, index: Tuple[float, float]) -> Tuple[numpy.ndarray, numpy.ndarray]:
 		xi, yi = index
 
 		# For this model, a point is (pz, sx, sy, w).
-		# We want to keep pz and w bounded, and restrict sx and sy.
-		return ([-numpy.inf, xi * self.x_step + self.model.xmin, yi * self.y_step + self.model.ymin, -numpy.inf], [numpy.inf, (xi + 1) * self.x_step + self.model.xmin, (yi + 1) * self.y_step + self.model.ymin, numpy.inf])
+		# We want to keep w bounded, and restrict sx and sy based on step and pz generally.
+		return (numpy.array((-self.max_pz, xi * self.x_step + self.model.xmin, yi * self.y_step + self.model.ymin, -numpy.inf)), numpy.array((self.max_pz, (xi + 1) * self.x_step + self.model.xmin, (yi + 1) * self.y_step + self.model.ymin, numpy.inf)))
 
 	def all_indices(self) -> numpy.ndindex:
 		# see https://github.com/numpy/numpy/issues/20706 for why this is a mypy problem.
 		return numpy.ndindex((self.num_x, self.num_y))  # type:ignore
+
+	def solve_for_index(self, dots: Sequence[DotMeasurement], index: Tuple[float, float]) -> scipy.optimize.OptimizeResult:
+		bounds = self.bounds(index)
+		sx_mean = (bounds[0][1] + bounds[1][1]) / 2
+		sy_mean = (bounds[0][2] + bounds[1][2]) / 2
+		# I don't care about the typing here at the moment.
+		return self.model.solve(dots, initial_pt=numpy.array(.1, sx_mean, sy_mean, .1), bounds=bounds)  # type: ignore

pdme/model/model.py

@@ -83,6 +83,6 @@ class Model():
 				raise ValueError(f"The initial point {initial_pt} does not have the model's expected length: {self.point_length()}")
 			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 = scipy.optimize.least_squares(self.costs(dots), initial, jac=self.jac(dots), ftol=1e-15, gtol=3e-16, xtol=None, bounds=bounds)
 		result.normalised_x = pdme.util.normalise_point_list(result.x, self.point_length())
 		return result

pdme/model/unrestricted_model.py

@@ -1,4 +1,7 @@
 import numpy
+from dataclasses import dataclass
+from typing import Sequence, Tuple
+import scipy.optimize
 from pdme.model.model import Model
 from pdme.measurement import DotMeasurement
 
@@ -13,11 +16,17 @@ class UnrestrictedModel(Model):
 	n : int
 		The number of dipoles to assume.
 	'''
-	def __init__(self, n: int) -> None:
+	def __init__(self, xmin: float, xmax: float, ymin: float, ymax: float, zmin: float, zmax: float, n: int) -> None:
+		self.xmin = xmin
+		self.xmax = xmax
+		self.ymin = ymin
+		self.ymax = ymax
+		self.zmin = zmin
+		self.zmax = zmax
 		self._n = n
 
 	def __repr__(self) -> str:
-		return f'UnrestrictedModel({self.n()})'
+		return f'UnrestrictedModel({self.xmin}, {self.xmax}, {self.ymin}, {self.ymax}, {self.zmin}, {self.zmax}, {self.n()})'
 
 	def point_length(self) -> int:
 		'''
@@ -58,3 +67,64 @@ class UnrestrictedModel(Model):
 		w_div = alpha**2 * (1 / numpy.pi) * ((f2 - w2) / ((f2 + w2)**2))
 
 		return numpy.concatenate((p_divs, r_divs, w_div), axis=None)
+
+
+@dataclass
+class UnrestrictedDiscretisation():
+	'''
+	Representation of a discretisation of a UnrestrictedModel.
+	Also captures a rough maximum value of dipole.
+
+	Parameters
+	----------
+	model : UnrestrictedModel
+		The parent model of the discretisation.
+	num_x : int
+		The number of partitions of the x axis.
+	num_y : int
+		The number of partitions of the y axis.
+	num_z : int
+		The number of partitions of the z axis.
+	max_p : int
+		The maximum p coordinate in any direction.
+	'''
+	model: UnrestrictedModel
+	num_x: int
+	num_y: int
+	num_z: int
+	max_p: int
+
+	def __post_init__(self):
+		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
+		self.z_step = (self.model.zmax - self.model.zmin) / self.num_z
+
+	def bounds(self, index: Tuple[float, float, float]) -> Tuple:
+		xi, yi, zi = index
+
+		# For this model, a point is (px, py, pz, sx, sx, sy, w).
+		# We want to keep w unbounded, restrict sx, sy, sz based on step and all of p generally.
+		return (
+			[
+				-self.max_p, -self.max_p, -self.max_p,
+				xi * self.x_step + self.model.xmin, yi * self.y_step + self.model.ymin, zi * self.z_step + self.model.zmin,
+				-numpy.inf
+			],
+			[
+				self.max_p, self.max_p, self.max_p,
+				(xi + 1) * self.x_step + self.model.xmin, (yi + 1) * self.y_step + self.model.ymin, (zi + 1) * self.z_step + self.model.zmin,
+				numpy.inf
+			]
+		)
+
+	def all_indices(self) -> numpy.ndindex:
+		# see https://github.com/numpy/numpy/issues/20706 for why this is a mypy problem.
+		return numpy.ndindex((self.num_x, self.num_y, self.num_z))  # type:ignore
+
+	def solve_for_index(self, dots: Sequence[DotMeasurement], index: Tuple[float, float, float]) -> scipy.optimize.OptimizeResult:
+		bounds = self.bounds(index)
+		sx_mean = (bounds[0][3] + bounds[1][3]) / 2
+		sy_mean = (bounds[0][4] + bounds[1][4]) / 2
+		sz_mean = (bounds[0][5] + bounds[1][5]) / 2
+		return self.model.solve(dots, initial_pt=numpy.array([.1, .1, .1, sx_mean, sy_mean, sz_mean, .1]), bounds=bounds)