From 7a9fa1ba04a12586ef09c89e35e706817012faab Mon Sep 17 00:00:00 2001 From: Deepak Mallubhotla Date: Sun, 22 May 2022 14:53:47 -0500 Subject: [PATCH] feat: Adds log spaced in frequency space model --- pdme/model/__init__.py | 5 + pdme/model/log_spaced_random_choice_model.py | 127 ++++++++++ ...t_multiple_dipole_fixed_magnitude_model.py | 229 ++++++++++++++++++ 3 files changed, 361 insertions(+) create mode 100644 pdme/model/log_spaced_random_choice_model.py create mode 100644 tests/model/test_log_spaced_random_count_multiple_dipole_fixed_magnitude_model.py diff --git a/pdme/model/__init__.py b/pdme/model/__init__.py index 77d4b0c..a7d7f51 100644 --- a/pdme/model/__init__.py +++ b/pdme/model/__init__.py @@ -7,9 +7,14 @@ from pdme.model.random_count_multidipole_fixed_magnitude_model import ( RandomCountMultipleDipoleFixedMagnitudeModel, ) +from pdme.model.log_spaced_random_choice_model import ( + LogSpacedRandomCountMultipleDipoleFixedMagnitudeModel, +) + __all__ = [ "DipoleModel", "SingleDipoleFixedMagnitudeModel", "MultipleDipoleFixedMagnitudeModel", "RandomCountMultipleDipoleFixedMagnitudeModel", + "LogSpacedRandomCountMultipleDipoleFixedMagnitudeModel", ] diff --git a/pdme/model/log_spaced_random_choice_model.py b/pdme/model/log_spaced_random_choice_model.py new file mode 100644 index 0000000..2d885db --- /dev/null +++ b/pdme/model/log_spaced_random_choice_model.py @@ -0,0 +1,127 @@ +import numpy +import numpy.random +from pdme.model.model import DipoleModel +from pdme.measurement import ( + OscillatingDipole, + OscillatingDipoleArrangement, +) + + +class LogSpacedRandomCountMultipleDipoleFixedMagnitudeModel(DipoleModel): + """ + Model of multiple oscillating dipoles with a fixed magnitude, but free rotation. Spaced logarithmically. + + 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. + + 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, + 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.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"LogSpacedRandomCountMultipleDipoleFixedMagnitudeModel({self.xmin}, {self.xmax}, {self.ymin}, {self.ymax}, {self.zmin}, {self.zmax}, {self.wexp_min}, {self.wexp_max}, {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 = 10 ** rng.uniform(self.wexp_min, self.wexp_max) + + 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, + _: 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 = 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_magnitude_model.py b/tests/model/test_log_spaced_random_count_multiple_dipole_fixed_magnitude_model.py new file mode 100644 index 0000000..e1444d6 --- /dev/null +++ b/tests/model/test_log_spaced_random_count_multiple_dipole_fixed_magnitude_model.py @@ -0,0 +1,229 @@ +from pdme.model import LogSpacedRandomCountMultipleDipoleFixedMagnitudeModel +import numpy +import logging +import pytest + + +_logger = logging.getLogger(__name__) + + +def test_random_count_multiple_dipole_wrong_probability(): + with pytest.raises(ValueError): + LogSpacedRandomCountMultipleDipoleFixedMagnitudeModel( + -10, 10, -5, 5, 2, 3, 1, 2, 10, 5, 2 + ) + + +def test_repr_random_count_multiple_dipole_fixed_mag(): + model = LogSpacedRandomCountMultipleDipoleFixedMagnitudeModel( + -10, 10, -5, 5, 2, 3, 1, 2, 10, 5, 0.5 + ) + assert ( + repr(model) + == "LogSpacedRandomCountMultipleDipoleFixedMagnitudeModel(-10, 10, -5, 5, 2, 3, 1, 2, 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 = LogSpacedRandomCountMultipleDipoleFixedMagnitudeModel( + -10, 10, -5, 5, 2, 3, 0, 5, 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 = 10**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 = LogSpacedRandomCountMultipleDipoleFixedMagnitudeModel( + -10, 10, -5, 5, 2, 3, 0, 5, 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 = LogSpacedRandomCountMultipleDipoleFixedMagnitudeModel( + x_min, x_max, y_min, y_max, z_min, z_max, 0, max_frequency, 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 = 10**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 < 10 ** max_frequency and dipoles[0].w > 10**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 = LogSpacedRandomCountMultipleDipoleFixedMagnitudeModel( + x_min, x_max, y_min, y_max, z_min, z_max, 0, max_frequency, 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, + 10**1.2088314662639255, + ] + ] + ] + ) + + 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 = LogSpacedRandomCountMultipleDipoleFixedMagnitudeModel( + x_min, + x_max, + y_min, + y_max, + z_min, + z_max, + 0, + max_frequency, + 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", + )