pdme/model/__init__.py

@@ -11,10 +11,15 @@ from pdme.model.log_spaced_random_choice_model import (
 	LogSpacedRandomCountMultipleDipoleFixedMagnitudeModel,
 )
 
+from pdme.model.log_spaced_random_choice_fixed_orientation_model import (
+	LogSpacedRandomCountMultipleDipoleFixedMagnitudeFixedOrientationModel,
+)
+
 __all__ = [
 	"DipoleModel",
 	"SingleDipoleFixedMagnitudeModel",
 	"MultipleDipoleFixedMagnitudeModel",
 	"RandomCountMultipleDipoleFixedMagnitudeModel",
 	"LogSpacedRandomCountMultipleDipoleFixedMagnitudeModel",
+	"LogSpacedRandomCountMultipleDipoleFixedMagnitudeFixedOrientationModel",
 ]

pdme/model/log_spaced_random_choice_fixed_orientation_model.py

@@ -0,0 +1,140 @@
+import numpy
+import numpy.random
+from pdme.model.model import DipoleModel
+from pdme.measurement import (
+	OscillatingDipole,
+	OscillatingDipoleArrangement,
+)
+
+
+class LogSpacedRandomCountMultipleDipoleFixedMagnitudeFixedOrientationModel(
+	DipoleModel
+):
+	"""
+	Model of multiple oscillating dipoles with a fixed magnitude and fixed rotation. Spaced log uniformly in relaxation time.
+
+	Parameters
+	----------
+
+	wexp_min: log-10 lower bound for dipole frequency
+	wexp_min: log-10 upper bound for dipole frequency
+
+	pfixed : float
+	The fixed dipole magnitude.
+
+	thetafixed: float
+	The fixed theta (polar angle).
+	Should be between 0 and pi.
+
+	phifixed: float
+	The fixed phi (azimuthal angle).
+	Should be between 0 and 2 pi.
+
+	n_max : int
+	The maximum number of dipoles.
+
+	prob_occupancy : float
+	The probability of dipole occupancy
+	"""
+
+	def __init__(
+		self,
+		xmin: float,
+		xmax: float,
+		ymin: float,
+		ymax: float,
+		zmin: float,
+		zmax: float,
+		wexp_min: float,
+		wexp_max: float,
+		pfixed: float,
+		thetafixed: float,
+		phifixed: float,
+		n_max: int,
+		prob_occupancy: float,
+	) -> None:
+		self.xmin = xmin
+		self.xmax = xmax
+		self.ymin = ymin
+		self.ymax = ymax
+		self.zmin = zmin
+		self.zmax = zmax
+		self.wexp_min = wexp_min
+		self.wexp_max = wexp_max
+		self.pfixed = pfixed
+		self.thetafixed = thetafixed
+		self.phifixed = phifixed
+		self.rng = numpy.random.default_rng()
+		self.n_max = n_max
+
+		px = self.pfixed * numpy.sin(self.thetafixed) * numpy.cos(self.phifixed)
+		py = self.pfixed * numpy.sin(self.thetafixed) * numpy.sin(self.phifixed)
+		pz = self.pfixed * numpy.cos(self.thetafixed)
+
+		self.moment_fixed = numpy.array([px, py, pz])
+		if prob_occupancy >= 1 or prob_occupancy <= 0:
+			raise ValueError(
+				f"The probability of a dipole site occupancy must be between 0 and 1, got {prob_occupancy}"
+			)
+		self.prob_occupancy = prob_occupancy
+
+	def __repr__(self) -> str:
+		return f"LogSpacedRandomCountMultipleDipoleFixedMagnitudeFixedOrientationModel({self.xmin}, {self.xmax}, {self.ymin}, {self.ymax}, {self.zmin}, {self.zmax}, {self.wexp_min}, {self.wexp_max}, {self.pfixed}, {self.thetafixed}, {self.phifixed}, {self.n_max}, {self.prob_occupancy})"
+
+	def get_dipoles(
+		self, max_frequency: float, rng_to_use: numpy.random.Generator = None
+	) -> OscillatingDipoleArrangement:
+		rng: numpy.random.Generator
+		if rng_to_use is None:
+			rng = self.rng
+		else:
+			rng = rng_to_use
+
+		dipoles = []
+
+		n = rng.binomial(self.n_max, self.prob_occupancy)
+
+		for i in range(n):
+			s_pts = numpy.array(
+				(
+					rng.uniform(self.xmin, self.xmax),
+					rng.uniform(self.ymin, self.ymax),
+					rng.uniform(self.zmin, self.zmax),
+				)
+			)
+			frequency = 10 ** rng.uniform(self.wexp_min, self.wexp_max)
+
+			dipoles.append(OscillatingDipole(self.moment_fixed, s_pts, frequency))
+		return OscillatingDipoleArrangement(dipoles)
+
+	def get_monte_carlo_dipole_inputs(
+		self,
+		monte_carlo_n: int,
+		_: float,
+		rng_to_use: numpy.random.Generator = None,
+	) -> numpy.ndarray:
+
+		rng: numpy.random.Generator
+		if rng_to_use is None:
+			rng = self.rng
+		else:
+			rng = rng_to_use
+
+		shape = (monte_carlo_n, self.n_max)
+
+		p_mask = rng.binomial(1, self.prob_occupancy, shape)
+
+		dipoles = numpy.einsum("ij,k->ijk", p_mask, self.moment_fixed)
+		# Is there a better way to create the final array? probably! can create a flatter guy then reshape.
+		# this is easier to reason about.
+		px = dipoles[:, :, 0]
+		py = dipoles[:, :, 1]
+		pz = dipoles[:, :, 2]
+
+		sx = rng.uniform(self.xmin, self.xmax, shape)
+		sy = rng.uniform(self.ymin, self.ymax, shape)
+		sz = rng.uniform(self.zmin, self.zmax, shape)
+
+		w = 10 ** rng.uniform(self.wexp_min, self.wexp_max, shape)
+
+		return numpy.stack([px, py, pz, sx, sy, sz, w], axis=-1)

tests/model/test_log_spaced_random_count_multiple_dipole_fixed_orientation_fixed_magnitude_model.py

@@ -0,0 +1,151 @@
+from pdme.model import (
+	LogSpacedRandomCountMultipleDipoleFixedMagnitudeFixedOrientationModel,
+)
+
+import numpy
+import logging
+import pytest
+
+
+_logger = logging.getLogger(__name__)
+
+
+def test_random_count_fixedorientation_multiple_dipole_wrong_probability():
+	with pytest.raises(ValueError):
+		LogSpacedRandomCountMultipleDipoleFixedMagnitudeFixedOrientationModel(
+			-10, 10, -5, 5, 2, 3, 1, 2, 10, 0, 0, 5, 2
+		)
+
+
+def test_repr_random_count_multiple_dipole_fixed_orientation_mag():
+	model = LogSpacedRandomCountMultipleDipoleFixedMagnitudeFixedOrientationModel(
+		-10, 10, -5, 5, 2, 3, 1, 2, 10, 0, 1, 5, 0.5
+	)
+	assert (
+		repr(model)
+		== "LogSpacedRandomCountMultipleDipoleFixedMagnitudeFixedOrientationModel(-10, 10, -5, 5, 2, 3, 1, 2, 10, 0, 1, 5, 0.5)"
+	), "Repr should be same as instantiation."
+
+
+def test_random_count_multiple_dipole_fixed_mag_model_get_dipoles_invariant():
+
+	x_min = -10
+	x_max = 10
+	y_min = -5
+	y_max = 5
+	z_min = 2
+	z_max = 3
+	p_fixed = 10
+	theta = 0
+	phi = 0
+	max_frequency = 5
+
+	model = LogSpacedRandomCountMultipleDipoleFixedMagnitudeFixedOrientationModel(
+		x_min,
+		x_max,
+		y_min,
+		y_max,
+		z_min,
+		z_max,
+		0,
+		max_frequency,
+		p_fixed,
+		theta,
+		phi,
+		1,
+		0.5,
+	)
+	model.rng = numpy.random.default_rng(1234)
+
+	for i in range(10):
+		dipole_arrangement = model.get_dipoles(max_frequency)
+		dipoles = dipole_arrangement.dipoles
+
+		assert len(dipoles) in (
+			0,
+			1,
+		), "Should have either zero or one dipole generated."
+
+		if len(dipoles) > 0:
+			min_s = numpy.array([x_min, y_min, z_min])
+			max_s = numpy.array([x_max, y_max, z_max])
+
+			numpy.testing.assert_equal(
+				numpy.logical_and(min_s < dipoles[0].s, max_s > dipoles[0].s),
+				True,
+				f"Dipole location [{dipoles[0].s}] should have been between min [{min_s}] and max [{max_s}] bounds.",
+			)
+			assert (
+				dipoles[0].w < 10 ** max_frequency and dipoles[0].w > 10**0
+			), "Dipole frequency should have been between 0 and max."
+
+			numpy.testing.assert_allclose(
+				dipoles[0].p,
+				numpy.array([0, 0, p_fixed]),
+				err_msg="Should have had the expected dipole moment.",
+			)
+
+	custom_rng = numpy.random.default_rng(1234)
+	for i in range(10):
+		dipole_arrangement = model.get_dipoles(max_frequency, custom_rng)
+		dipoles = dipole_arrangement.dipoles
+
+		assert len(dipoles) in (
+			0,
+			1,
+		), "Should have either zero or one dipole generated."
+
+		if len(dipoles) > 0:
+			min_s = numpy.array([x_min, y_min, z_min])
+			max_s = numpy.array([x_max, y_max, z_max])
+
+			numpy.testing.assert_equal(
+				numpy.logical_and(min_s < dipoles[0].s, max_s > dipoles[0].s),
+				True,
+				f"Dipole location [{dipoles[0].s}] should have been between min [{min_s}] and max [{max_s}] bounds.",
+			)
+			assert (
+				dipoles[0].w < 10 ** max_frequency and dipoles[0].w > 10**0
+			), "Dipole frequency should have been between 0 and max."
+
+			numpy.testing.assert_allclose(
+				dipoles[0].p,
+				numpy.array([0, 0, p_fixed]),
+				err_msg="Should have had the expected dipole moment.",
+			)
+
+
+def test_random_count_multiple_dipole_fixed_or_fixed_mag_model_get_n_dipoles():
+	# TODO: this test is a bit garbage just calls things without testing.
+	x_min = -10
+	x_max = 10
+	y_min = -5
+	y_max = 5
+	z_min = 2
+	z_max = 3
+	p_fixed = 10
+	theta = numpy.pi / 2
+	phi = 0
+	max_frequency = 5
+
+	model = LogSpacedRandomCountMultipleDipoleFixedMagnitudeFixedOrientationModel(
+		x_min,
+		x_max,
+		y_min,
+		y_max,
+		z_min,
+		z_max,
+		0,
+		max_frequency,
+		p_fixed,
+		theta,
+		phi,
+		1,
+		0.5,
+	)
+	model.rng = numpy.random.default_rng(1234)
+
+	model.get_monte_carlo_dipole_inputs(1, max_frequency)
+	model.get_monte_carlo_dipole_inputs(
+		1, max_frequency, numpy.random.default_rng(1234)
+	)