diff --git a/pdme/model/__init__.py b/pdme/model/__init__.py index a7d7f51..e3c98ad 100644 --- a/pdme/model/__init__.py +++ b/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", ] diff --git a/pdme/model/log_spaced_random_choice_fixed_orientation_model.py b/pdme/model/log_spaced_random_choice_fixed_orientation_model.py new file mode 100644 index 0000000..600be1b --- /dev/null +++ b/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) diff --git a/tests/model/test_log_spaced_random_count_multiple_dipole_fixed_orientation_fixed_magnitude_model.py b/tests/model/test_log_spaced_random_count_multiple_dipole_fixed_orientation_fixed_magnitude_model.py new file mode 100644 index 0000000..4d99970 --- /dev/null +++ b/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) + )