diff --git a/pdme/model/__init__.py b/pdme/model/__init__.py index 6655463..77d4b0c 100644 --- a/pdme/model/__init__.py +++ b/pdme/model/__init__.py @@ -3,9 +3,13 @@ from pdme.model.fixed_magnitude_model import SingleDipoleFixedMagnitudeModel from pdme.model.multidipole_fixed_magnitude_model import ( MultipleDipoleFixedMagnitudeModel, ) +from pdme.model.random_count_multidipole_fixed_magnitude_model import ( + RandomCountMultipleDipoleFixedMagnitudeModel, +) __all__ = [ "DipoleModel", "SingleDipoleFixedMagnitudeModel", "MultipleDipoleFixedMagnitudeModel", + "RandomCountMultipleDipoleFixedMagnitudeModel", ] diff --git a/pdme/model/random_count_multidipole_fixed_magnitude_model.py b/pdme/model/random_count_multidipole_fixed_magnitude_model.py new file mode 100644 index 0000000..0c11d8f --- /dev/null +++ b/pdme/model/random_count_multidipole_fixed_magnitude_model.py @@ -0,0 +1,119 @@ +import numpy +import numpy.random +from pdme.model.model import DipoleModel +from pdme.measurement import ( + OscillatingDipole, + OscillatingDipoleArrangement, +) + + +class RandomCountMultipleDipoleFixedMagnitudeModel(DipoleModel): + """ + Model of multiple oscillating dipoles with a fixed magnitude, but free rotation. + + Parameters + ---------- + pfixed : float + The fixed dipole magnitude. + + 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, + pfixed: 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.pfixed = pfixed + self.rng = numpy.random.default_rng() + self.n_max = n_max + 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"RandomCountMultipleDipoleFixedMagnitudeModel({self.xmin}, {self.xmax}, {self.ymin}, {self.ymax}, {self.zmin}, {self.zmax}, {self.pfixed}, {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): + theta = numpy.arccos(rng.uniform(-1, 1)) + phi = 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( + ( + rng.uniform(self.xmin, self.xmax), + rng.uniform(self.ymin, self.ymax), + rng.uniform(self.zmin, self.zmax), + ) + ) + frequency = rng.uniform(0, max_frequency) + + dipoles.append( + OscillatingDipole(numpy.array([px, py, pz]), s_pts, frequency) + ) + return OscillatingDipoleArrangement(dipoles) + + def get_monte_carlo_dipole_inputs( + self, + monte_carlo_n: int, + max_frequency: 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) + theta = 2 * numpy.pi * rng.random(shape) + phi = numpy.arccos(2 * rng.random(shape) - 1) + + p_mask = rng.binomial(1, 0.5, shape) + p_magnitude = self.pfixed * p_mask + + px = p_magnitude * numpy.cos(theta) * numpy.sin(phi) + py = p_magnitude * numpy.sin(theta) * numpy.sin(phi) + pz = p_magnitude * numpy.cos(phi) + + 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 = rng.uniform(1, max_frequency, shape) + + return numpy.stack([px, py, pz, sx, sy, sz, w], axis=-1) diff --git a/tests/model/test_random_count_multiple_dipole_fixed_magnitude_model.py b/tests/model/test_random_count_multiple_dipole_fixed_magnitude_model.py new file mode 100644 index 0000000..cb2bf77 --- /dev/null +++ b/tests/model/test_random_count_multiple_dipole_fixed_magnitude_model.py @@ -0,0 +1,217 @@ +from pdme.model import RandomCountMultipleDipoleFixedMagnitudeModel +import numpy +import logging +import pytest + + +_logger = logging.getLogger(__name__) + + +def test_random_count_multiple_dipole_wrong_probability(): + with pytest.raises(ValueError): + RandomCountMultipleDipoleFixedMagnitudeModel(-10, 10, -5, 5, 2, 3, 10, 5, 2) + + +def test_repr_random_count_multiple_dipole_fixed_mag(): + model = RandomCountMultipleDipoleFixedMagnitudeModel( + -10, 10, -5, 5, 2, 3, 10, 5, 0.5 + ) + assert ( + repr(model) + == "RandomCountMultipleDipoleFixedMagnitudeModel(-10, 10, -5, 5, 2, 3, 10, 5, 0.5)" + ), "Repr should be what I want." + + +def test_random_count_multiple_dipole_fixed_mag_model_get_dipoles(): + + p_fixed = 10 + + model = RandomCountMultipleDipoleFixedMagnitudeModel( + -10, 10, -5, 5, 2, 3, p_fixed, 1, 0.5 + ) + + dipole_arrangement = model.get_dipoles(5, numpy.random.default_rng(1234)) + dipoles = dipole_arrangement.dipoles + + assert len(dipoles) == 1, "Should have only had one dipole generated." + expected_p = numpy.array([8.60141814, -4.50270821, -2.3960853]) + expected_s = numpy.array([-4.76615152, -1.80902942, 2.11809123]) + expected_w = 1.2088314662639255 + + numpy.testing.assert_allclose( + dipoles[0].p, + expected_p, + err_msg="Random multiple dipole p wasn't as expected", + ) + numpy.testing.assert_allclose( + dipoles[0].s, + expected_s, + err_msg="Random multiple dipole s wasn't as expected", + ) + numpy.testing.assert_allclose( + dipoles[0].w, expected_w, err_msg="Random multiple dipole w wasn't as expected" + ) + numpy.testing.assert_allclose( + numpy.linalg.norm(dipoles[0].p), + p_fixed, + err_msg="Should have had the expected dipole moment magnitude.", + ) + + +def test_random_count_multiple_dipole_fixed_mag_model_get_dipoles_multiple(): + + p_fixed = 10 + dipole_count = 5 + + model = RandomCountMultipleDipoleFixedMagnitudeModel( + -10, 10, -5, 5, 2, 3, p_fixed, dipole_count, 0.5 + ) + + dipole_arrangement = model.get_dipoles(20, numpy.random.default_rng(1234)) + dipoles = dipole_arrangement.dipoles + + assert ( + len(dipoles) == dipole_count + ), "Should have had multiple dipole based on count generated." + + +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 + max_frequency = 5 + + model = RandomCountMultipleDipoleFixedMagnitudeModel( + x_min, x_max, y_min, y_max, z_min, z_max, p_fixed, 1, 0.5 + ) + model.rng = numpy.random.default_rng(1234) + + dipole_arrangement = model.get_dipoles(5) + dipoles = dipole_arrangement.dipoles + + assert len(dipoles) == 1, "Should have only had one dipole generated." + expected_p = numpy.array([8.60141814, -4.50270821, -2.3960853]) + expected_s = numpy.array([-4.76615152, -1.80902942, 2.11809123]) + expected_w = 1.2088314662639255 + + numpy.testing.assert_allclose( + dipoles[0].p, + expected_p, + err_msg="Random multiple dipole p wasn't as expected", + rtol=1e-06, + ) + numpy.testing.assert_allclose( + dipoles[0].s, expected_s, err_msg="Random multiple dipole s wasn't as expected" + ) + numpy.testing.assert_allclose( + dipoles[0].w, expected_w, err_msg="Random multiple dipole w wasn't as expected" + ) + 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 < max_frequency and dipoles[0].w > 0 + ), "Dipole frequency should have been between 0 and max." + + numpy.testing.assert_allclose( + numpy.linalg.norm(dipoles[0].p), + p_fixed, + err_msg="Should have had the expected dipole moment magnitude.", + ) + + +def test_random_count_multiple_dipole_fixed_mag_model_get_n_dipoles(): + + x_min = -10 + x_max = 10 + y_min = -5 + y_max = 5 + z_min = 2 + z_max = 3 + p_fixed = 10 + max_frequency = 5 + + model = RandomCountMultipleDipoleFixedMagnitudeModel( + x_min, x_max, y_min, y_max, z_min, z_max, p_fixed, 1, 0.5 + ) + model.rng = numpy.random.default_rng(1234) + + dipole_array = model.get_monte_carlo_dipole_inputs(1, max_frequency) + expected_dipole_array = numpy.array( + [ + [ + [ + 9.60483896, + -1.41627817, + -2.3960853, + -4.76615152, + -1.80902942, + 2.11809123, + 1.96706517, + ] + ] + ] + ) + + numpy.testing.assert_allclose( + dipole_array, + expected_dipole_array, + err_msg="Should have had the expected output dipole array.", + rtol=1e-6, + ) + numpy.testing.assert_allclose( + model.get_monte_carlo_dipole_inputs( + 1, max_frequency, numpy.random.default_rng(1234) + ), + expected_dipole_array, + err_msg="Should have had the expected output dipole array, even with explicitly passed rng.", + ) + + +def test_random_count_multiple_dipole_shape(): + + x_min = -10 + x_max = 10 + y_min = -5 + y_max = 5 + z_min = 2 + z_max = 3 + p_fixed = 10 + max_frequency = 5 + num_dipoles = 13 + monte_carlo_n = 11 + + model = RandomCountMultipleDipoleFixedMagnitudeModel( + x_min, x_max, y_min, y_max, z_min, z_max, p_fixed, num_dipoles, 0.5 + ) + model.rng = numpy.random.default_rng(1234) + + actual_shape = model.get_monte_carlo_dipole_inputs( + monte_carlo_n, max_frequency + ).shape + + numpy.testing.assert_equal( + actual_shape, + (monte_carlo_n, num_dipoles, 7), + err_msg="shape was wrong for monte carlo outputs", + )