style: fmt updates

pdme/inputs/__init__.py

@@ -1,4 +1,7 @@
-from pdme.inputs.dot_inputs import inputs_with_frequency_range, input_pairs_with_frequency_range
+from pdme.inputs.dot_inputs import (
+	inputs_with_frequency_range,
+	input_pairs_with_frequency_range,
+)
 
 
-__all__ = ['inputs_with_frequency_range', 'input_pairs_with_frequency_range']
+__all__ = ["inputs_with_frequency_range", "input_pairs_with_frequency_range"]

pdme/inputs/dot_inputs.py

@@ -4,10 +4,18 @@ import itertools
 from typing import Sequence, Tuple
 
 
-def inputs_with_frequency_range(dots: Sequence[numpy.typing.ArrayLike], frequency_range: Sequence[float]) -> Sequence[Tuple[numpy.typing.ArrayLike, float]]:
+def inputs_with_frequency_range(
+	dots: Sequence[numpy.typing.ArrayLike], frequency_range: Sequence[float]
+) -> Sequence[Tuple[numpy.typing.ArrayLike, float]]:
 	return list(itertools.chain(*[[(dot, f) for f in frequency_range] for dot in dots]))
 
 
-def input_pairs_with_frequency_range(dots: Sequence[numpy.typing.ArrayLike], frequency_range: Sequence[float]) -> Sequence[Tuple[numpy.typing.ArrayLike, numpy.typing.ArrayLike, float]]:
+def input_pairs_with_frequency_range(
+	dots: Sequence[numpy.typing.ArrayLike], frequency_range: Sequence[float]
+) -> Sequence[Tuple[numpy.typing.ArrayLike, numpy.typing.ArrayLike, float]]:
 	all_pairs = itertools.combinations(dots, 2)
-	return list(itertools.chain(*[[(dot1, dot2, f) for f in frequency_range] for (dot1, dot2) in all_pairs]))
+	return list(
+		itertools.chain(
+			*[[(dot1, dot2, f) for f in frequency_range] for (dot1, dot2) in all_pairs]
+		)
+	)

pdme/measurement/__init__.py

@@ -1,7 +1,22 @@
 from pdme.measurement.dot_measure import DotMeasurement, DotRangeMeasurement
-from pdme.measurement.dot_pair_measure import DotPairMeasurement, DotPairRangeMeasurement
-from pdme.measurement.oscillating_dipole import OscillatingDipole, OscillatingDipoleArrangement
+from pdme.measurement.dot_pair_measure import (
+	DotPairMeasurement,
+	DotPairRangeMeasurement,
+)
+from pdme.measurement.oscillating_dipole import (
+	OscillatingDipole,
+	OscillatingDipoleArrangement,
+)
 from pdme.measurement.input_types import DotInput, DotPairInput
 
 
-__all__ = ['DotMeasurement', 'DotRangeMeasurement', 'DotPairMeasurement', 'DotPairRangeMeasurement', 'OscillatingDipole', 'OscillatingDipoleArrangement', 'DotInput', 'DotPairInput']
+__all__ = [
+	"DotMeasurement",
+	"DotRangeMeasurement",
+	"DotPairMeasurement",
+	"DotPairRangeMeasurement",
+	"OscillatingDipole",
+	"OscillatingDipoleArrangement",
+	"DotInput",
+	"DotPairInput",
+]

pdme/measurement/dot_measure.py

@@ -4,19 +4,20 @@ import numpy.typing
 
 
 @dataclass
-class DotMeasurement():
-	'''
+class DotMeasurement:
+	"""
 	Representation of a dot measuring oscillating dipoles.
 
 	Parameters
 	----------
 	v : float
-		The voltage measured at the dot.
+			The voltage measured at the dot.
 	r : numpy.ndarray
-		The position of the dot.
+			The position of the dot.
 	f : float
-		The measurement frequency.
-	'''
+			The measurement frequency.
+	"""
+
 	v: float
 	r: numpy.ndarray
 	f: float
@@ -26,21 +27,22 @@ class DotMeasurement():
 
 
 @dataclass
-class DotRangeMeasurement():
-	'''
+class DotRangeMeasurement:
+	"""
 	Representation of a dot measuring oscillating dipoles.
 
 	Parameters
 	----------
 	v_low : float
-		The lower range of voltage measured at the dot.
+			The lower range of voltage measured at the dot.
 	v_high : float
-		The upper range of voltage measured at the dot.
+			The upper range of voltage measured at the dot.
 	r : numpy.ndarray
-		The position of the dot.
+			The position of the dot.
 	f : float
-		The measurement frequency.
-	'''
+			The measurement frequency.
+	"""
+
 	v_low: float
 	v_high: float
 	r: numpy.ndarray

pdme/measurement/dot_pair_measure.py

@@ -4,21 +4,22 @@ import numpy.typing
 
 
 @dataclass
-class DotPairMeasurement():
-	'''
+class DotPairMeasurement:
+	"""
 	Representation of a dot measuring oscillating dipoles.
 
 	Parameters
 	----------
 	v : float
-		The voltage measured at the dot.
+			The voltage measured at the dot.
 	r1 : numpy.ndarray
-		The position of the first dot.
+			The position of the first dot.
 	r2 : numpy.ndarray
-		The position of the second dot.
+			The position of the second dot.
 	f : float
-		The measurement frequency.
-	'''
+			The measurement frequency.
+	"""
+
 	v: float
 	r1: numpy.ndarray
 	r2: numpy.ndarray
@@ -30,23 +31,24 @@ class DotPairMeasurement():
 
 
 @dataclass
-class DotPairRangeMeasurement():
-	'''
+class DotPairRangeMeasurement:
+	"""
 	Representation of a dot measuring oscillating dipoles.
 
 	Parameters
 	----------
 	v_low : float
-		The lower range of voltage measured at the dot.
+			The lower range of voltage measured at the dot.
 	v_high : float
-		The upper range of voltage measured at the dot.
+			The upper range of voltage measured at the dot.
 	r1 : numpy.ndarray
-		The position of the first dot.
+			The position of the first dot.
 	r2 : numpy.ndarray
-		The position of the second dot.
+			The position of the second dot.
 	f : float
-		The measurement frequency.
-	'''
+			The measurement frequency.
+	"""
+
 	v_low: float
 	v_high: float
 	r1: numpy.ndarray

pdme/measurement/input_types.py

@@ -9,14 +9,28 @@ DotPairInput = Tuple[numpy.typing.ArrayLike, numpy.typing.ArrayLike, float]
 
 
 def dot_inputs_to_array(dot_inputs: Sequence[DotInput]) -> numpy.ndarray:
-	return numpy.array([numpy.append(numpy.array(input[0]), input[1]) for input in dot_inputs])
+	return numpy.array(
+		[numpy.append(numpy.array(input[0]), input[1]) for input in dot_inputs]
+	)
 
 
 def dot_pair_inputs_to_array(pair_inputs: Sequence[DotPairInput]) -> numpy.ndarray:
-	return numpy.array([[numpy.append(numpy.array(input[0]), input[2]), numpy.append(numpy.array(input[1]), input[2])] for input in pair_inputs])
+	return numpy.array(
+		[
+			[
+				numpy.append(numpy.array(input[0]), input[2]),
+				numpy.append(numpy.array(input[1]), input[2]),
+			]
+			for input in pair_inputs
+		]
+	)
 
 
-def dot_range_measurements_low_high_arrays(dot_range_measurements: Union[Sequence[DotRangeMeasurement], Sequence[DotPairRangeMeasurement]]) -> Tuple[numpy.ndarray, numpy.ndarray]:
+def dot_range_measurements_low_high_arrays(
+	dot_range_measurements: Union[
+		Sequence[DotRangeMeasurement], Sequence[DotPairRangeMeasurement]
+	]
+) -> Tuple[numpy.ndarray, numpy.ndarray]:
 	lows = [measurement.v_low for measurement in dot_range_measurements]
 	highs = [measurement.v_high for measurement in dot_range_measurements]
 	return (numpy.array(lows), numpy.array(highs))

pdme/measurement/oscillating_dipole.py

@@ -3,51 +3,55 @@ import numpy
 import numpy.typing
 from typing import Sequence, List
 from pdme.measurement.dot_measure import DotMeasurement, DotRangeMeasurement
-from pdme.measurement.dot_pair_measure import DotPairMeasurement, DotPairRangeMeasurement
+from pdme.measurement.dot_pair_measure import (
+	DotPairMeasurement,
+	DotPairRangeMeasurement,
+)
 from pdme.measurement.input_types import DotInput, DotPairInput
 
 
 @dataclass
-class OscillatingDipole():
-	'''
+class OscillatingDipole:
+	"""
 	Representation of an oscillating dipole, either known or guessed.
 
 	Parameters
 	----------
 	p : numpy.ndarray
-		The oscillating dipole moment, with overall sign arbitrary.
+			The oscillating dipole moment, with overall sign arbitrary.
 	s : numpy.ndarray
-		The position of the dipole.
+			The position of the dipole.
 	w : float
-		The oscillation frequency.
-	'''
+			The oscillation frequency.
+	"""
+
 	p: numpy.ndarray
 	s: numpy.ndarray
 	w: float
 
 	def __post_init__(self) -> None:
-		'''
+		"""
 		Coerce the inputs into numpy arrays.
-		'''
+		"""
 		self.p = numpy.array(self.p)
 		self.s = numpy.array(self.s)
 
 	def s_at_position(self, r: numpy.ndarray, f: float) -> float:
-		'''
+		"""
 		Returns the noise potential at a point r, at some frequency f.
 
 		Parameters
 		----------
 		r : numpy.ndarray
-			The position of the dot.
+				The position of the dot.
 		f : float
-			The dot frequency to sample.
-		'''
-		return (self._alpha(r))**2 * self._b(f)
+				The dot frequency to sample.
+		"""
+		return (self._alpha(r)) ** 2 * self._b(f)
 
 	def _alpha(self, r: numpy.ndarray) -> float:
 		diff = r - self.s
-		return self.p.dot(diff) / (numpy.linalg.norm(diff)**3)
+		return self.p.dot(diff) / (numpy.linalg.norm(diff) ** 3)
 
 	def _b(self, f: float) -> float:
 		return (1 / numpy.pi) * (self.w / (f**2 + self.w**2))
@@ -59,65 +63,114 @@ class OscillatingDipole():
 		return numpy.concatenate([self.p, self.s, numpy.array([self.w])])
 
 
-class OscillatingDipoleArrangement():
-	'''
+class OscillatingDipoleArrangement:
+	"""
 	A collection of oscillating dipoles, which we are interested in being able to characterise.
 
 	Parameters
 	--------
 	dipoles : Sequence[OscillatingDipole]
-	'''
+	"""
+
 	def __init__(self, dipoles: Sequence[OscillatingDipole]):
 		self.dipoles = dipoles
 
 	def get_dot_measurement(self, dot_input: DotInput) -> DotMeasurement:
 		r = numpy.array(dot_input[0])
 		f = dot_input[1]
-		return DotMeasurement(sum([dipole.s_at_position(r, f) for dipole in self.dipoles]), r, f)
+		return DotMeasurement(
+			sum([dipole.s_at_position(r, f) for dipole in self.dipoles]), r, f
+		)
 
-	def get_dot_pair_measurement(self, dot_pair_input: DotPairInput) -> DotPairMeasurement:
+	def get_dot_pair_measurement(
+		self, dot_pair_input: DotPairInput
+	) -> DotPairMeasurement:
 		r1 = numpy.array(dot_pair_input[0])
 		r2 = numpy.array(dot_pair_input[1])
 		f = dot_pair_input[2]
-		return DotPairMeasurement(sum([dipole.s_for_dot_pair(r1, r2, f) for dipole in self.dipoles]), r1, r2, f)
+		return DotPairMeasurement(
+			sum([dipole.s_for_dot_pair(r1, r2, f) for dipole in self.dipoles]),
+			r1,
+			r2,
+			f,
+		)
 
-	def get_dot_measurements(self, dot_inputs: Sequence[DotInput]) -> List[DotMeasurement]:
-		'''
+	def get_dot_measurements(
+		self, dot_inputs: Sequence[DotInput]
+	) -> List[DotMeasurement]:
+		"""
 		For a series of points, each with three coordinates and a frequency, return a list of the corresponding DotMeasurements.
-		'''
+		"""
 		return [self.get_dot_measurement(dot_input) for dot_input in dot_inputs]
 
-	def get_dot_pair_measurements(self, dot_pair_inputs: Sequence[DotPairInput]) -> List[DotPairMeasurement]:
-		'''
+	def get_dot_pair_measurements(
+		self, dot_pair_inputs: Sequence[DotPairInput]
+	) -> List[DotPairMeasurement]:
+		"""
 		For a series of pairs of points, each with three coordinates and a frequency, return a list of the corresponding DotPairMeasurements.
-		'''
-		return [self.get_dot_pair_measurement(dot_pair_input) for dot_pair_input in dot_pair_inputs]
+		"""
+		return [
+			self.get_dot_pair_measurement(dot_pair_input)
+			for dot_pair_input in dot_pair_inputs
+		]
 
-	def get_percent_range_dot_measurement(self, dot_input: DotInput, low_percent: float, high_percent: float) -> DotRangeMeasurement:
+	def get_percent_range_dot_measurement(
+		self, dot_input: DotInput, low_percent: float, high_percent: float
+	) -> DotRangeMeasurement:
 		r = numpy.array(dot_input[0])
 		f = dot_input[1]
-		return DotRangeMeasurement(low_percent * sum([dipole.s_at_position(r, f) for dipole in self.dipoles]), high_percent * sum([dipole.s_at_position(r, f) for dipole in self.dipoles]), r, f)
+		return DotRangeMeasurement(
+			low_percent * sum([dipole.s_at_position(r, f) for dipole in self.dipoles]),
+			high_percent * sum([dipole.s_at_position(r, f) for dipole in self.dipoles]),
+			r,
+			f,
+		)
 
-	def get_percent_range_dot_measurements(self, dot_inputs: Sequence[DotInput], low_percent: float, high_percent: float) -> List[DotRangeMeasurement]:
-		'''
+	def get_percent_range_dot_measurements(
+		self, dot_inputs: Sequence[DotInput], low_percent: float, high_percent: float
+	) -> List[DotRangeMeasurement]:
+		"""
 		For a series of pairs of points, each with three coordinates and a frequency, and also a lower error range and upper error range, return a list of the corresponding DotPairRangeMeasurements.
-		'''
-		return [self.get_percent_range_dot_measurement(dot_input, low_percent, high_percent) for dot_input in dot_inputs]
+		"""
+		return [
+			self.get_percent_range_dot_measurement(dot_input, low_percent, high_percent)
+			for dot_input in dot_inputs
+		]
 
-	def get_percent_range_dot_pair_measurement(self, pair_input: DotPairInput, low_percent: float, high_percent: float) -> DotPairRangeMeasurement:
+	def get_percent_range_dot_pair_measurement(
+		self, pair_input: DotPairInput, low_percent: float, high_percent: float
+	) -> DotPairRangeMeasurement:
 		r1 = numpy.array(pair_input[0])
 		r2 = numpy.array(pair_input[1])
 		f = pair_input[2]
-		return DotPairRangeMeasurement(low_percent * sum([dipole.s_for_dot_pair(r1, r2, f) for dipole in self.dipoles]), high_percent * sum([dipole.s_for_dot_pair(r1, r2, f) for dipole in self.dipoles]), r1, r2, f)
+		return DotPairRangeMeasurement(
+			low_percent
+			* sum([dipole.s_for_dot_pair(r1, r2, f) for dipole in self.dipoles]),
+			high_percent
+			* sum([dipole.s_for_dot_pair(r1, r2, f) for dipole in self.dipoles]),
+			r1,
+			r2,
+			f,
+		)
 
-	def get_percent_range_dot_pair_measurements(self, pair_inputs: Sequence[DotPairInput], low_percent: float, high_percent: float) -> List[DotPairRangeMeasurement]:
-		'''
+	def get_percent_range_dot_pair_measurements(
+		self,
+		pair_inputs: Sequence[DotPairInput],
+		low_percent: float,
+		high_percent: float,
+	) -> List[DotPairRangeMeasurement]:
+		"""
 		For a series of pairs of points, each with three coordinates and a frequency, and also a lower error range and upper error range, return a list of the corresponding DotPairRangeMeasurements.
-		'''
-		return [self.get_percent_range_dot_pair_measurement(pair_input, low_percent, high_percent) for pair_input in pair_inputs]
+		"""
+		return [
+			self.get_percent_range_dot_pair_measurement(
+				pair_input, low_percent, high_percent
+			)
+			for pair_input in pair_inputs
+		]
 
 	def to_numpy_array(self) -> numpy.ndarray:
-		'''
+		"""
 		Returns a numpy array with the canonical representation of each dipole in a nx7 numpy array.
-		'''
+		"""
 		return numpy.array([dipole.to_flat_array() for dipole in self.dipoles])

pdme/meta.py

@@ -1,3 +1,3 @@
 from importlib.metadata import version
 
-__version__ = version('pdme')
+__version__ = version("pdme")

pdme/model/__init__.py

@@ -4,4 +4,11 @@ from pdme.model.unrestricted_model import UnrestrictedModel
 from pdme.model.fixed_dipole_model import FixedDipoleModel
 from pdme.model.fixed_magnitude_model import FixedMagnitudeModel
 
-__all__ = ["Model", "Discretisation", "FixedZPlaneModel", "UnrestrictedModel", "FixedDipoleModel", "FixedMagnitudeModel"]
+__all__ = [
+	"Model",
+	"Discretisation",
+	"FixedZPlaneModel",
+	"UnrestrictedModel",
+	"FixedDipoleModel",
+	"FixedMagnitudeModel",
+]

pdme/model/fixed_dipole_model.py

@@ -4,21 +4,36 @@ from dataclasses import dataclass
 from typing import Sequence, Tuple
 import scipy.optimize
 from pdme.model.model import Model, Discretisation
-from pdme.measurement import DotMeasurement, OscillatingDipoleArrangement, OscillatingDipole
+from pdme.measurement import (
+	DotMeasurement,
+	OscillatingDipoleArrangement,
+	OscillatingDipole,
+)
 
 
 class FixedDipoleModel(Model):
-	'''
+	"""
 	Model of oscillating dipole with a fixed dipole moment.
 
 	Parameters
 	----------
 	p : numpy.ndarray
-		The fixed dipole moment.
+			The fixed dipole moment.
 	n : int
-		The number of dipoles to assume.
-	'''
-	def __init__(self, xmin: float, xmax: float, ymin: float, ymax: float, zmin: float, zmax: float, p: numpy.ndarray, n: int) -> None:
+			The number of dipoles to assume.
+	"""
+
+	def __init__(
+		self,
+		xmin: float,
+		xmax: float,
+		ymin: float,
+		ymax: float,
+		zmin: float,
+		zmax: float,
+		p: numpy.ndarray,
+		n: int,
+	) -> None:
 		self.xmin = xmin
 		self.xmax = xmax
 		self.ymin = ymin
@@ -30,12 +45,20 @@ class FixedDipoleModel(Model):
 		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()})'
+		return f"FixedDipoleModel({self.xmin}, {self.xmax}, {self.ymin}, {self.ymax}, {self.zmin}, {self.zmax}, {self.p}, {self.n()})"
 
 	# TODO: this signature doesn't make sense.
 	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)])
+		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 solution_single_dipole(self, pt: numpy.ndarray) -> OscillatingDipole:
 		# assume length is 4.
@@ -44,10 +67,10 @@ class FixedDipoleModel(Model):
 		return OscillatingDipole(self.p, s, w)
 
 	def point_length(self) -> int:
-		'''
-			Dipole is constrained magnitude, but free orientation.
-			Six degrees of freedom: (sx, sy, sz, w).
-		'''
+		"""
+		Dipole is constrained magnitude, but free orientation.
+		Six degrees of freedom: (sx, sy, sz, w).
+		"""
 		return 4
 
 	def n(self) -> int:
@@ -58,45 +81,56 @@ class FixedDipoleModel(Model):
 		w = pt[3]
 
 		diff = dot.r - s
-		alpha = self.p.dot(diff) / (numpy.linalg.norm(diff)**3)
+		alpha = self.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:
+	def jac_for_point_at_dot(
+		self, dot: DotMeasurement, pt: numpy.ndarray
+	) -> numpy.ndarray:
 		s = pt[0:3]
 		w = pt[3]
 
 		diff = dot.r - s
-		alpha = self.p.dot(diff) / (numpy.linalg.norm(diff)**3)
+		alpha = self.p.dot(diff) / (numpy.linalg.norm(diff) ** 3)
 		b = (1 / numpy.pi) * (w / (w**2 + dot.f**2))
 
-		r_divs = (-self.p / (numpy.linalg.norm(diff)**3) + 3 * self.p.dot(diff) * diff / (numpy.linalg.norm(diff)**5)) * 2 * alpha * b
+		r_divs = (
+			(
+				-self.p / (numpy.linalg.norm(diff) ** 3)
+				+ 3 * self.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))
+		w_div = alpha**2 * (1 / numpy.pi) * ((f2 - w2) / ((f2 + w2) ** 2))
 
 		return numpy.concatenate((r_divs, w_div), axis=None)
 
 
 @dataclass
 class FixedDipoleDiscretisation(Discretisation):
-	'''
+	"""
 	Representation of a discretisation of a FixedDipoleDiscretisation.
 	Also captures a rough maximum value of dipole.
 
 	Parameters
 	----------
 	model : FixedDipoleModel
-		The parent model of the discretisation.
+			The parent model of the discretisation.
 	num_x : int
-		The number of partitions of the x axis.
+			The number of partitions of the x axis.
 	num_y : int
-		The number of partitions of the y axis.
+			The number of partitions of the y axis.
 	num_z : int
-		The number of partitions of the z axis.
-	'''
+			The number of partitions of the z axis.
+	"""
+
 	model: FixedDipoleModel
 	num_x: int
 	num_y: int
@@ -115,22 +149,32 @@ class FixedDipoleDiscretisation(Discretisation):
 		# We want to keep w unbounded, restrict sx, sy, sz, px and py based on step.
 		return (
 			[
-				xi * self.x_step + self.model.xmin, yi * self.y_step + self.model.ymin, zi * self.z_step + self.model.zmin,
-				-numpy.inf
+				xi * self.x_step + self.model.xmin,
+				yi * self.y_step + self.model.ymin,
+				zi * self.z_step + self.model.zmin,
+				-numpy.inf,
 			],
 			[
-				(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
-			]
+				(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, ...]) -> scipy.optimize.OptimizeResult:
+	def solve_for_index(
+		self, dots: Sequence[DotMeasurement], index: Tuple[float, ...]
+	) -> scipy.optimize.OptimizeResult:
 		bounds = self.bounds(index)
 		sx_mean = (bounds[0][0] + bounds[1][0]) / 2
 		sy_mean = (bounds[0][1] + bounds[1][1]) / 2
 		sz_mean = (bounds[0][2] + bounds[1][2]) / 2
-		return self.model.solve(dots, initial_pt=numpy.array([sx_mean, sy_mean, sz_mean, .1]), bounds=bounds)
+		return self.model.solve(
+			dots,
+			initial_pt=numpy.array([sx_mean, sy_mean, sz_mean, 0.1]),
+			bounds=bounds,
+		)

pdme/model/fixed_magnitude_model.py

@@ -4,21 +4,36 @@ from dataclasses import dataclass
 from typing import Sequence, Tuple
 import scipy.optimize
 from pdme.model.model import Model, Discretisation
-from pdme.measurement import DotMeasurement, OscillatingDipole, OscillatingDipoleArrangement
+from pdme.measurement import (
+	DotMeasurement,
+	OscillatingDipole,
+	OscillatingDipoleArrangement,
+)
 
 
 class FixedMagnitudeModel(Model):
-	'''
+	"""
 	Model of oscillating dipole with a fixed magnitude, but free rotation.
 
 	Parameters
 	----------
 	pfixed : float
-		The fixed dipole magnitude.
+			The fixed dipole magnitude.
 	n : int
-		The number of dipoles to assume.
-	'''
-	def __init__(self, xmin: float, xmax: float, ymin: float, ymax: float, zmin: float, zmax: float, pfixed: float, n: int) -> None:
+			The number of dipoles to assume.
+	"""
+
+	def __init__(
+		self,
+		xmin: float,
+		xmax: float,
+		ymin: float,
+		ymax: float,
+		zmin: float,
+		zmax: float,
+		pfixed: float,
+		n: int,
+	) -> None:
 		self.xmin = xmin
 		self.xmax = xmax
 		self.ymin = ymin
@@ -30,7 +45,7 @@ class FixedMagnitudeModel(Model):
 		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.pfixed}, {self.n()})'
+		return f"FixedMagnitudeModel({self.xmin}, {self.xmax}, {self.ymin}, {self.ymax}, {self.zmin}, {self.zmax}, {self.pfixed}, {self.n()})"
 
 	def solution_single_dipole(self, pt: numpy.ndarray) -> OscillatingDipole:
 		# assume length is 6, who needs error checking.
@@ -39,18 +54,20 @@ class FixedMagnitudeModel(Model):
 		s = pt[2:5]
 		w = pt[5]
 
-		p = numpy.array([
-			self.pfixed * numpy.sin(p_theta) * numpy.cos(p_phi),
-			self.pfixed * numpy.sin(p_theta) * numpy.sin(p_phi),
-			self.pfixed * numpy.cos(p_theta)
-		])
+		p = numpy.array(
+			[
+				self.pfixed * numpy.sin(p_theta) * numpy.cos(p_phi),
+				self.pfixed * numpy.sin(p_theta) * numpy.sin(p_phi),
+				self.pfixed * numpy.cos(p_theta),
+			]
+		)
 		return OscillatingDipole(p, s, w)
 
 	def point_length(self) -> int:
-		'''
-			Dipole is constrained magnitude, but free orientation.
-			Six degrees of freedom: (p_theta, p_phi, sx, sy, sz, w).
-		'''
+		"""
+		Dipole is constrained magnitude, but free orientation.
+		Six degrees of freedom: (p_theta, p_phi, sx, sy, sz, w).
+		"""
 		return 6
 
 	def get_dipoles(self, frequency: float) -> OscillatingDipoleArrangement:
@@ -59,8 +76,16 @@ class FixedMagnitudeModel(Model):
 		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)])
+		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 get_n_single_dipoles(self, n: int, max_frequency: float) -> numpy.ndarray:
 		# psw
@@ -88,29 +113,35 @@ class FixedMagnitudeModel(Model):
 		s = pt[2:5]
 		w = pt[5]
 
-		p = numpy.array([
-			self.pfixed * numpy.sin(p_theta) * numpy.cos(p_phi),
-			self.pfixed * numpy.sin(p_theta) * numpy.sin(p_phi),
-			self.pfixed * numpy.cos(p_theta)
-		])
+		p = numpy.array(
+			[
+				self.pfixed * numpy.sin(p_theta) * numpy.cos(p_phi),
+				self.pfixed * numpy.sin(p_theta) * numpy.sin(p_phi),
+				self.pfixed * numpy.cos(p_theta),
+			]
+		)
 		diff = dot.r - s
-		alpha = p.dot(diff) / (numpy.linalg.norm(diff)**3)
+		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:
+	def jac_for_point_at_dot(
+		self, dot: DotMeasurement, pt: numpy.ndarray
+	) -> numpy.ndarray:
 		p_theta = pt[0]
 		p_phi = pt[1]
 		s = pt[2:5]
 		w = pt[5]
 
-		p = numpy.array([
-			self.pfixed * numpy.sin(p_theta) * numpy.cos(p_phi),
-			self.pfixed * numpy.sin(p_theta) * numpy.sin(p_phi),
-			self.pfixed * numpy.cos(p_theta)
-		])
+		p = numpy.array(
+			[
+				self.pfixed * numpy.sin(p_theta) * numpy.cos(p_phi),
+				self.pfixed * numpy.sin(p_theta) * numpy.sin(p_phi),
+				self.pfixed * numpy.cos(p_theta),
+			]
+		)
 		diff = dot.r - s
-		alpha = p.dot(diff) / (numpy.linalg.norm(diff)**3)
+		alpha = p.dot(diff) / (numpy.linalg.norm(diff) ** 3)
 		b = (1 / numpy.pi) * (w / (w**2 + dot.f**2))
 
 		theta_div_middle = self.pfixed * (
@@ -118,45 +149,54 @@ class FixedMagnitudeModel(Model):
 			+ diff[1] * numpy.sin(p_phi) * numpy.cos(p_theta)
 			- diff[2] * numpy.sin(p_theta)
 		)
-		theta_div = 2 * alpha * (theta_div_middle) / (numpy.linalg.norm(diff)**3) * b
+		theta_div = 2 * alpha * (theta_div_middle) / (numpy.linalg.norm(diff) ** 3) * b
 
 		phi_div_middle = self.pfixed * (
 			diff[1] * numpy.sin(p_theta) * numpy.cos(p_phi)
 			- diff[0] * numpy.sin(p_theta) * numpy.sin(p_phi)
 		)
-		phi_div = 2 * alpha * (phi_div_middle) / (numpy.linalg.norm(diff)**3) * b
+		phi_div = 2 * alpha * (phi_div_middle) / (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
+		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))
+		w_div = alpha**2 * (1 / numpy.pi) * ((f2 - w2) / ((f2 + w2) ** 2))
 
 		return numpy.concatenate((theta_div, phi_div, r_divs, w_div), axis=None)
 
 
 @dataclass
 class FixedMagnitudeDiscretisation(Discretisation):
-	'''
+	"""
 	Representation of a discretisation of a FixedMagnitudeDiscretisation.
 	Also captures a rough maximum value of dipole.
 
 	Parameters
 	----------
 	model : FixedMagnitudeModel
-		The parent model of the discretisation.
+			The parent model of the discretisation.
 	num_ptheta: int
-		The number of partitions of ptheta.
+			The number of partitions of ptheta.
 	num_pphi: int
-		The number of partitions of pphi.
+			The number of partitions of pphi.
 	num_x : int
-		The number of partitions of the x axis.
+			The number of partitions of the x axis.
 	num_y : int
-		The number of partitions of the y axis.
+			The number of partitions of the y axis.
 	num_z : int
-		The number of partitions of the z axis.
-	'''
+			The number of partitions of the z axis.
+	"""
+
 	model: FixedMagnitudeModel
 	num_ptheta: int
 	num_pphi: int
@@ -179,15 +219,21 @@ class FixedMagnitudeDiscretisation(Discretisation):
 		# We want to keep w unbounded, restrict sx, sy, sz, px and py based on step.
 		return (
 			[
-				numpy.arccos(1 - pthetai * self.h_step), pphii * self.phi_step,
-				xi * self.x_step + self.model.xmin, yi * self.y_step + self.model.ymin, zi * self.z_step + self.model.zmin,
-				-numpy.inf
+				numpy.arccos(1 - pthetai * self.h_step),
+				pphii * self.phi_step,
+				xi * self.x_step + self.model.xmin,
+				yi * self.y_step + self.model.ymin,
+				zi * self.z_step + self.model.zmin,
+				-numpy.inf,
 			],
 			[
-				numpy.arccos(1 - (pthetai + 1) * self.h_step), (pphii + 1) * self.phi_step,
-				(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
-			]
+				numpy.arccos(1 - (pthetai + 1) * self.h_step),
+				(pphii + 1) * self.phi_step,
+				(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 get_model(self) -> Model:
@@ -195,13 +241,23 @@ class FixedMagnitudeDiscretisation(Discretisation):
 
 	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_ptheta, self.num_pphi, self.num_x, self.num_y, self.num_z))  # type:ignore
+		return numpy.ndindex(
+			(self.num_ptheta, self.num_pphi, self.num_x, self.num_y, self.num_z)
+		)  # type:ignore
 
-	def solve_for_index(self, dots: Sequence[DotMeasurement], index: Tuple[float, ...]) -> scipy.optimize.OptimizeResult:
+	def solve_for_index(
+		self, dots: Sequence[DotMeasurement], index: Tuple[float, ...]
+	) -> scipy.optimize.OptimizeResult:
 		bounds = self.bounds(index)
 		ptheta_mean = (bounds[0][0] + bounds[1][0]) / 2
 		pphi_mean = (bounds[0][1] + bounds[1][1]) / 2
 		sx_mean = (bounds[0][2] + bounds[1][2]) / 2
 		sy_mean = (bounds[0][3] + bounds[1][3]) / 2
 		sz_mean = (bounds[0][4] + bounds[1][4]) / 2
-		return self.model.solve(dots, initial_pt=numpy.array([ptheta_mean, pphi_mean, sx_mean, sy_mean, sz_mean, .1]), bounds=bounds)
+		return self.model.solve(
+			dots,
+			initial_pt=numpy.array(
+				[ptheta_mean, pphi_mean, sx_mean, sy_mean, sz_mean, 0.1]
+			),
+			bounds=bounds,
+		)

pdme/model/fixed_z_plane_model.py

@@ -8,26 +8,29 @@ from pdme.measurement import DotMeasurement
 
 
 class FixedZPlaneModel(Model):
-	'''
+	"""
 	Model of oscillating dipoles constrained to lie within a plane.
 	Additionally, each dipole is assumed to be orientated in the plus or minus z direction.
 
 	Parameters
 	----------
 	z : float
-		The z position of the plane where dipoles are constrained to lie.
+			The z position of the plane where dipoles are constrained to lie.
 	xmin : float
-		The minimum x value for dipoles.
+			The minimum x value for dipoles.
 	xmax : float
-		The maximum x value for dipoles.
+			The maximum x value for dipoles.
 	ymin : float
-		The minimum y value for dipoles.
+			The minimum y value for dipoles.
 	ymax : float
-		The maximum y value for dipoles.
+			The maximum y value for dipoles.
 	n : int
-		The number of dipoles to assume.
-	'''
-	def __init__(self, z: float, xmin: float, xmax: float, ymin: float, ymax: float, n: int) -> None:
+			The number of dipoles to assume.
+	"""
+
+	def __init__(
+		self, z: float, xmin: float, xmax: float, ymin: float, ymax: float, n: int
+	) -> None:
 		self.z = z
 		self.xmin = xmin
 		self.xmax = xmax
@@ -36,13 +39,13 @@ class FixedZPlaneModel(Model):
 		self._n = n
 
 	def __repr__(self) -> str:
-		return f'FixedZPlaneModel({self.z}, {self.xmin}, {self.xmax}, {self.ymin}, {self.ymax}, {self.n()})'
+		return f"FixedZPlaneModel({self.z}, {self.xmin}, {self.xmax}, {self.ymin}, {self.ymax}, {self.n()})"
 
 	def point_length(self) -> int:
-		'''
-			Dipole is constrained in this model to have (px, py, pz) = (0, 0, pz) and (sx, sy, sz) = (sx, sy, self.z).
-			With some frequency w, there are four degrees of freedom: (pz, sx, sy, w).
-		'''
+		"""
+		Dipole is constrained in this model to have (px, py, pz) = (0, 0, pz) and (sx, sy, sz) = (sx, sy, self.z).
+		With some frequency w, there are four degrees of freedom: (pz, sx, sy, w).
+		"""
 		return 4
 
 	def n(self) -> int:
@@ -54,48 +57,61 @@ class FixedZPlaneModel(Model):
 		w = pt[3]
 
 		diff = dot.r - s
-		alpha = p.dot(diff) / (numpy.linalg.norm(diff)**3)
+		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:
+	def jac_for_point_at_dot(
+		self, dot: DotMeasurement, pt: numpy.ndarray
+	) -> numpy.ndarray:
 		p = numpy.array([0, 0, pt[0]])
 		s = numpy.array([pt[1], pt[2], self.z])
 		w = pt[3]
 
 		diff = dot.r - s
-		alpha = p.dot(diff) / (numpy.linalg.norm(diff)**3)
+		alpha = p.dot(diff) / (numpy.linalg.norm(diff) ** 3)
 		b = (1 / numpy.pi) * (w / (w**2 + dot.f**2))
 
-		p_divs = 2 * alpha * diff[2] / (numpy.linalg.norm(diff)**3) * b  # only need the z component.
+		p_divs = (
+			2 * alpha * diff[2] / (numpy.linalg.norm(diff) ** 3) * b
+		)  # only need the z component.
 
-		r_divs = (-p[0:2] / (numpy.linalg.norm(diff)**3) + 3 * p.dot(diff) * diff[0:2] / (numpy.linalg.norm(diff)**5)) * 2 * alpha * b
+		r_divs = (
+			(
+				-p[0:2] / (numpy.linalg.norm(diff) ** 3)
+				+ 3 * p.dot(diff) * diff[0:2] / (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))
+		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 FixedZPlaneDiscretisation():
-	'''
+class FixedZPlaneDiscretisation:
+	"""
 	Representation of a discretisation of a FixedZPlaneModel.
 	Also captures a rough maximum value of dipole.
 
 	Parameters
 	----------
 	model : FixedZPlaneModel
-		The parent model of the discretisation.
+			The parent model of the discretisation.
 	num_pz: int
-		The number of partitions of pz.
+			The number of partitions of pz.
 	num_x : int
-		The number of partitions of the x axis.
+			The number of partitions of the x axis.
 	num_y : int
-		The number of partitions of the y axis.
-	'''
+			The number of partitions of the y axis.
+	"""
+
 	model: FixedZPlaneModel
 	num_pz: int
 	num_x: int
@@ -108,24 +124,42 @@ class FixedZPlaneDiscretisation():
 		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, float]) -> Tuple[numpy.ndarray, numpy.ndarray]:
+	def bounds(
+		self, index: Tuple[float, float, float]
+	) -> Tuple[numpy.ndarray, numpy.ndarray]:
 		pzi, xi, yi = index
 
 		# For this model, a point is (pz, sx, sy, w).
 		# We want to keep w bounded, and restrict pz, sx and sy based on step.
 		return (
-			numpy.array((pzi * self.pz_step - self.max_pz, xi * self.x_step + self.model.xmin, yi * self.y_step + self.model.ymin, -numpy.inf)),
-			numpy.array(((pzi + 1) * self.pz_step - self.max_pz, (xi + 1) * self.x_step + self.model.xmin, (yi + 1) * self.y_step + self.model.ymin, numpy.inf))
+			numpy.array(
+				(
+					pzi * self.pz_step - self.max_pz,
+					xi * self.x_step + self.model.xmin,
+					yi * self.y_step + self.model.ymin,
+					-numpy.inf,
+				)
+			),
+			numpy.array(
+				(
+					(pzi + 1) * self.pz_step - 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_pz, self.num_x, self.num_y))  # type:ignore
 
-	def solve_for_index(self, dots: Sequence[DotMeasurement], index: Tuple[float, float, float]) -> scipy.optimize.OptimizeResult:
+	def solve_for_index(
+		self, dots: Sequence[DotMeasurement], index: Tuple[float, float, float]
+	) -> scipy.optimize.OptimizeResult:
 		bounds = self.bounds(index)
 		pz_mean = (bounds[0][0] + bounds[1][0]) / 2
 		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((pz_mean, sx_mean, sy_mean, .1)), bounds=bounds)  # type: ignore
+		return self.model.solve(dots, initial_pt=numpy.array((pz_mean, sx_mean, sy_mean, 0.1)), bounds=bounds)  # type: ignore

pdme/model/model.py

@@ -1,7 +1,11 @@
 import numpy
 import scipy.optimize
 from typing import Callable, Sequence, Tuple, List
-from pdme.measurement import DotMeasurement, OscillatingDipoleArrangement, OscillatingDipole
+from pdme.measurement import (
+	DotMeasurement,
+	OscillatingDipoleArrangement,
+	OscillatingDipole,
+)
 import pdme.util
 import logging
 
@@ -9,7 +13,7 @@ import logging
 _logger = logging.getLogger(__name__)
 
 
-class Model():
+class Model:
 	"""
 	Interface for models.
 	"""
@@ -34,18 +38,20 @@ class Model():
 
 	def solution_as_dipoles(self, pts: numpy.ndarray) -> List[OscillatingDipole]:
 		pt_length = self.point_length()
-		chunked_pts = [pts[i: i + pt_length] for i in range(0, len(pts), pt_length)]
+		chunked_pts = [pts[i : i + pt_length] for i in range(0, len(pts), pt_length)]
 		return [self.solution_single_dipole(pt) for pt in chunked_pts]
 
 	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.
 		pt_length = self.point_length()
-		chunked_pts = [pts[i: i + pt_length] for i in range(0, len(pts), pt_length)]
+		chunked_pts = [pts[i : i + pt_length] for i in range(0, len(pts), pt_length)]
 		return sum(self.v_for_point_at_dot(dot, pt) for pt in chunked_pts) - dot.v
 
-	def costs(self, dots: Sequence[DotMeasurement]) -> Callable[[numpy.ndarray], numpy.ndarray]:
-		'''
+	def costs(
+		self, dots: Sequence[DotMeasurement]
+	) -> Callable[[numpy.ndarray], numpy.ndarray]:
+		"""
 		Returns a function that returns the cost for the given list of DotMeasurements for a particular model-dependent phase space point.
 		Default implementation assumes a single dot cost from which to build the list.
 
@@ -56,7 +62,7 @@ class Model():
 		Returns
 		----------
 		Returns the model's cost function.
-		'''
+		"""
 		_logger.debug(f"Constructing costs for dots: {dots}")
 
 		def costs_to_return(pts: numpy.ndarray) -> numpy.ndarray:
@@ -64,18 +70,24 @@ class Model():
 
 		return costs_to_return
 
-	def jac_for_point_at_dot(self, dot: DotMeasurement, pt: numpy.ndarray) -> numpy.ndarray:
+	def jac_for_point_at_dot(
+		self, dot: DotMeasurement, pt: numpy.ndarray
+	) -> numpy.ndarray:
 		raise NotImplementedError
 
 	def jac_for_dot(self, dot: DotMeasurement, pts: numpy.ndarray) -> numpy.ndarray:
 		# creates numpy.ndarrays in groups of self.point_length().
 		# Will throw problems for irregular points, but that's okay for now.
 		pt_length = self.point_length()
-		chunked_pts = [pts[i: i + pt_length] for i in range(0, len(pts), pt_length)]
-		return numpy.append([], [self.jac_for_point_at_dot(dot, pt) for pt in chunked_pts])
+		chunked_pts = [pts[i : i + pt_length] for i in range(0, len(pts), pt_length)]
+		return numpy.append(
+			[], [self.jac_for_point_at_dot(dot, pt) for pt in chunked_pts]
+		)
 
-	def jac(self, dots: Sequence[DotMeasurement]) -> Callable[[numpy.ndarray], numpy.ndarray]:
-		'''
+	def jac(
+		self, dots: Sequence[DotMeasurement]
+	) -> Callable[[numpy.ndarray], numpy.ndarray]:
+		"""
 		Returns a function that returns the cost function's Jacobian for the given list of DotMeasurements for a particular model-dependent phase space point.
 		Default implementation assumes a single dot jacobian from which to build the list.
 
@@ -86,33 +98,53 @@ class Model():
 		Returns
 		----------
 		Returns the model's cost function's Jacobian.
-		'''
+		"""
+
 		def jac_to_return(pts: numpy.ndarray) -> numpy.ndarray:
 			return numpy.array([self.jac_for_dot(dot, pts) for dot in dots])
 
 		return jac_to_return
 
-	def solve(self, dots: Sequence[DotMeasurement], initial_pt: numpy.ndarray = None, bounds=(-numpy.inf, numpy.inf)) -> scipy.optimize.OptimizeResult:
+	def solve(
+		self,
+		dots: Sequence[DotMeasurement],
+		initial_pt: numpy.ndarray = None,
+		bounds=(-numpy.inf, numpy.inf),
+	) -> scipy.optimize.OptimizeResult:
 		if initial_pt is None:
-			initial = numpy.tile(.1, self.n() * self.point_length())
+			initial = numpy.tile(0.1, self.n() * self.point_length())
 		else:
 			if len(initial_pt) != self.point_length():
-				raise ValueError(f"The initial point {initial_pt} does not have the model's expected length: {self.point_length()}")
+				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, xtol=None, bounds=bounds)
-		result.normalised_x = pdme.util.normalise_point_list(result.x, self.point_length())
+		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
 
 
-class Discretisation():
+class Discretisation:
 	def bounds(self, index: Tuple[float, ...]) -> Tuple:
 		raise NotImplementedError
 
 	def all_indices(self) -> numpy.ndindex:
 		raise NotImplementedError
 
-	def solve_for_index(self, dots: Sequence[DotMeasurement], index: Tuple) -> scipy.optimize.OptimizeResult:
+	def solve_for_index(
+		self, dots: Sequence[DotMeasurement], index: Tuple
+	) -> scipy.optimize.OptimizeResult:
 		raise NotImplementedError
 
 	def get_model(self) -> Model:

pdme/model/unrestricted_model.py

@@ -3,20 +3,35 @@ from dataclasses import dataclass
 from typing import Sequence, Tuple
 import scipy.optimize
 from pdme.model.model import Model, Discretisation
-from pdme.measurement import DotMeasurement, OscillatingDipoleArrangement, OscillatingDipole
+from pdme.measurement import (
+	DotMeasurement,
+	OscillatingDipoleArrangement,
+	OscillatingDipole,
+)
 
 
 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, xmin: float, xmax: float, ymin: float, ymax: float, zmin: float, zmax: float, max_p: float, n: int) -> None:
+			The number of dipoles to assume.
+	"""
+
+	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
@@ -28,13 +43,13 @@ class UnrestrictedModel(Model):
 		self.rng = numpy.random.default_rng()
 
 	def __repr__(self) -> str:
-		return f'UnrestrictedModel({self.xmin}, {self.xmax}, {self.ymin}, {self.ymax}, {self.zmin}, {self.zmax}, {self.max_p}, {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:
-		'''
-			Dipole is unconstrained in this model.
-			All seven degrees of freedom: (px, py, pz, sx, sy, sz, w).
-		'''
+		"""
+		Dipole is unconstrained in this model.
+		All seven degrees of freedom: (px, py, pz, sx, sy, sz, w).
+		"""
 		return 7
 
 	def n(self) -> int:
@@ -46,27 +61,37 @@ class UnrestrictedModel(Model):
 		w = pt[6]
 
 		diff = dot.r - s
-		alpha = p.dot(diff) / (numpy.linalg.norm(diff)**3)
+		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:
+	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)
+		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
+		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
+		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))
+		w_div = alpha**2 * (1 / numpy.pi) * ((f2 - w2) / ((f2 + w2) ** 2))
 
 		return numpy.concatenate((p_divs, r_divs, w_div), axis=None)
 
@@ -77,35 +102,44 @@ class UnrestrictedModel(Model):
 		px = p * numpy.sin(theta) * numpy.cos(phi)
 		py = p * numpy.sin(theta) * numpy.sin(phi)
 		pz = p * 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)])
+		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)]
+		)
 
 
 @dataclass
 class UnrestrictedDiscretisation(Discretisation):
-	'''
+	"""
 	Representation of a discretisation of a UnrestrictedModel.
 	Also captures a rough maximum value of dipole.
 
 	Parameters
 	----------
 	model : UnrestrictedModel
-		The parent model of the discretisation.
+			The parent model of the discretisation.
 	num_px: int
-		The number of partitions of the px.
+			The number of partitions of the px.
 	num_py: int
-		The number of partitions of the py.
+			The number of partitions of the py.
 	num_pz: int
-		The number of partitions of pz.
+			The number of partitions of pz.
 	num_x : int
-		The number of partitions of the x axis.
+			The number of partitions of the x axis.
 	num_y : int
-		The number of partitions of the y axis.
+			The number of partitions of the y axis.
 	num_z : int
-		The number of partitions of the z axis.
+			The number of partitions of the z axis.
 	max_p : int
-		The maximum p coordinate in any direction.
-	'''
+			The maximum p coordinate in any direction.
+	"""
+
 	model: UnrestrictedModel
 	num_px: int
 	num_py: int
@@ -131,22 +165,34 @@ class UnrestrictedDiscretisation(Discretisation):
 		# We want to keep w unbounded, restrict sx, sy, sz, px and py based on step.
 		return (
 			[
-				pxi * self.px_step - self.max_p, pyi * self.py_step - self.max_p, pzi * self.pz_step - 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
+				pxi * self.px_step - self.max_p,
+				pyi * self.py_step - self.max_p,
+				pzi * self.pz_step - 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,
 			],
 			[
-				(pxi + 1) * self.px_step - self.max_p, (pyi + 1) * self.py_step - self.max_p, (pzi + 1) * self.pz_step - 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
-			]
+				(pxi + 1) * self.px_step - self.max_p,
+				(pyi + 1) * self.py_step - self.max_p,
+				(pzi + 1) * self.pz_step - 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_px, self.num_py, self.num_pz, self.num_x, self.num_y, self.num_z))  # type:ignore
+		return numpy.ndindex(
+			(self.num_px, self.num_py, self.num_pz, self.num_x, self.num_y, self.num_z)
+		)  # type:ignore
 
-	def solve_for_index(self, dots: Sequence[DotMeasurement], index: Tuple[float, ...]) -> scipy.optimize.OptimizeResult:
+	def solve_for_index(
+		self, dots: Sequence[DotMeasurement], index: Tuple[float, ...]
+	) -> scipy.optimize.OptimizeResult:
 		bounds = self.bounds(index)
 		px_mean = (bounds[0][0] + bounds[1][0]) / 2
 		py_mean = (bounds[0][1] + bounds[1][1]) / 2
@@ -154,4 +200,10 @@ class UnrestrictedDiscretisation(Discretisation):
 		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([px_mean, py_mean, pz_mean, sx_mean, sy_mean, sz_mean, .1]), bounds=bounds)
+		return self.model.solve(
+			dots,
+			initial_pt=numpy.array(
+				[px_mean, py_mean, pz_mean, sx_mean, sy_mean, sz_mean, 0.1]
+			),
+			bounds=bounds,
+		)

pdme/util/fast_nonlocal_spectrum.py

@@ -5,10 +5,12 @@ import logging
 _logger = logging.getLogger(__name__)
 
 
-def fast_s_nonlocal(dot_pair_inputs: numpy.ndarray, dipoles: numpy.ndarray) -> numpy.ndarray:
-	'''
-		No error correction here baby.
-	'''
+def fast_s_nonlocal(
+	dot_pair_inputs: numpy.ndarray, dipoles: numpy.ndarray
+) -> numpy.ndarray:
+	"""
+	No error correction here baby.
+	"""
 	ps = dipoles[:, 0:3]
 	ss = dipoles[:, 3:6]
 	ws = dipoles[:, 6]
@@ -31,19 +33,19 @@ def fast_s_nonlocal(dot_pair_inputs: numpy.ndarray, dipoles: numpy.ndarray) -> n
 		_logger.debug(f"diffses1: {diffses1}")
 		_logger.debug(f"diffses2: {diffses2}")
 
-	norms1 = numpy.linalg.norm(diffses1, axis=2)**3
-	norms2 = numpy.linalg.norm(diffses2, axis=2)**3
+	norms1 = numpy.linalg.norm(diffses1, axis=2) ** 3
+	norms2 = numpy.linalg.norm(diffses2, axis=2) ** 3
 	if _logger.isEnabledFor(logging.DEBUG):
 		_logger.debug(f"norms1: {norms1}")
 		_logger.debug(f"norms2: {norms2}")
 
-	alphses1 = numpy.einsum('...ji, ...i', diffses1, ps) / norms1
-	alphses2 = numpy.einsum('...ji, ...i', diffses2, ps) / norms2
+	alphses1 = numpy.einsum("...ji, ...i", diffses1, ps) / norms1
+	alphses2 = numpy.einsum("...ji, ...i", diffses2, ps) / norms2
 	if _logger.isEnabledFor(logging.DEBUG):
 		_logger.debug(f"alphses1: {alphses1}")
 		_logger.debug(f"alphses2: {alphses2}")
 
-	bses = (1 / numpy.pi) * (ws[:, None] / (f1s**2 + ws[:, None]**2))
+	bses = (1 / numpy.pi) * (ws[:, None] / (f1s**2 + ws[:, None] ** 2))
 	if _logger.isEnabledFor(logging.DEBUG):
 		_logger.debug(f"bses: {bses}")
 

pdme/util/fast_v_calc.py

@@ -5,10 +5,12 @@ import logging
 _logger = logging.getLogger(__name__)
 
 
-def fast_vs_for_dipoles(dot_inputs: numpy.ndarray, dipoles: numpy.ndarray) -> numpy.ndarray:
-	'''
-		No error correction here baby.
-	'''
+def fast_vs_for_dipoles(
+	dot_inputs: numpy.ndarray, dipoles: numpy.ndarray
+) -> numpy.ndarray:
+	"""
+	No error correction here baby.
+	"""
 	ps = dipoles[:, 0:3]
 	ss = dipoles[:, 3:6]
 	ws = dipoles[:, 6]
@@ -23,18 +25,18 @@ def fast_vs_for_dipoles(dot_inputs: numpy.ndarray, dipoles: numpy.ndarray) -> nu
 	diffses = rs - ss[:, None]
 
 	_logger.debug(f"diffses: {diffses}")
-	norms = numpy.linalg.norm(diffses, axis=2)**3
+	norms = numpy.linalg.norm(diffses, axis=2) ** 3
 	_logger.debug(f"norms: {norms}")
-	ases = (numpy.einsum('...ji, ...i', diffses, ps) / norms)**2
+	ases = (numpy.einsum("...ji, ...i", diffses, ps) / norms) ** 2
 	_logger.debug(f"ases: {ases}")
 
-	bses = (1 / numpy.pi) * (ws[:, None] / (fs**2 + ws[:, None]**2))
+	bses = (1 / numpy.pi) * (ws[:, None] / (fs**2 + ws[:, None] ** 2))
 	_logger.debug(f"bses: {bses}")
 	return ases * bses
 
 
 def between(a: numpy.ndarray, low: numpy.ndarray, high: numpy.ndarray) -> numpy.ndarray:
-	'''
-		Intended specifically for the case where a is a list of arrays, and each array must be between the single array low and high, but without error checking.
-	'''
+	"""
+	Intended specifically for the case where a is a list of arrays, and each array must be between the single array low and high, but without error checking.
+	"""
 	return numpy.all(numpy.logical_and(low < a, high > a), axis=1)

pdme/util/normal_form.py

@@ -26,7 +26,20 @@ def flip_chunk_to_positive_px(pt: numpy.ndarray) -> numpy.ndarray:
 
 
 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)]
+	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)
+	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,
+	)