diff --git a/pdme/model/log_spaced_random_choice_fixed_orientation_model.py b/pdme/model/log_spaced_random_choice_fixed_orientation_model.py index 600be1b..907d26f 100644 --- a/pdme/model/log_spaced_random_choice_fixed_orientation_model.py +++ b/pdme/model/log_spaced_random_choice_fixed_orientation_model.py @@ -5,6 +5,11 @@ from pdme.measurement import ( OscillatingDipole, OscillatingDipoleArrangement, ) +import logging +from typing import Optional +import pdme.subspace_simulation + +_logger = logging.getLogger(__name__) class LogSpacedRandomCountMultipleDipoleFixedMagnitudeFixedOrientationModel( @@ -138,3 +143,51 @@ class LogSpacedRandomCountMultipleDipoleFixedMagnitudeFixedOrientationModel( w = 10 ** rng.uniform(self.wexp_min, self.wexp_max, shape) return numpy.stack([px, py, pz, sx, sy, sz, w], axis=-1) + + def markov_chain_monte_carlo_proposal( + self, + dipole: numpy.ndarray, + stdev: pdme.subspace_simulation.DipoleStandardDeviation, + rng_arg: Optional[numpy.random.Generator] = None, + ) -> numpy.ndarray: + if rng_arg is None: + rng_to_use = self.rng + else: + rng_to_use = rng_arg + + px = dipole[0] + py = dipole[1] + pz = dipole[2] + # won't change p for this model of fixed dipole moment. + + rx = dipole[3] + ry = dipole[4] + rz = dipole[5] + + tentative_rx = rx + stdev.rx_step * rng_to_use.uniform(-1, 1) + if tentative_rx < self.xmin or tentative_rx > self.xmax: + tentative_rx = rx + + tentative_ry = ry + stdev.ry_step * rng_to_use.uniform(-1, 1) + if tentative_ry < self.ymin or tentative_ry > self.ymax: + tentative_ry = ry + tentative_rz = rz + stdev.rz_step * rng_to_use.uniform(-1, 1) + if tentative_rz < self.zmin or tentative_rz > self.zmax: + tentative_rz = rz + + w = dipole[6] + tentative_w = numpy.exp( + numpy.log(w) + (stdev.w_log_step * rng_to_use.uniform(-1, 1)) + ) + tentative_dip = numpy.array( + [ + px, + py, + pz, + tentative_rx, + tentative_ry, + tentative_rz, + tentative_w, + ] + ) + return tentative_dip diff --git a/pdme/model/log_spaced_random_choice_model.py b/pdme/model/log_spaced_random_choice_model.py index f933e0f..a149858 100644 --- a/pdme/model/log_spaced_random_choice_model.py +++ b/pdme/model/log_spaced_random_choice_model.py @@ -5,6 +5,8 @@ from pdme.measurement import ( OscillatingDipole, OscillatingDipoleArrangement, ) +import pdme.subspace_simulation +from typing import Optional class LogSpacedRandomCountMultipleDipoleFixedMagnitudeModel(DipoleModel): @@ -125,3 +127,73 @@ class LogSpacedRandomCountMultipleDipoleFixedMagnitudeModel(DipoleModel): w = 10 ** rng.uniform(self.wexp_min, self.wexp_max, shape) return numpy.stack([px, py, pz, sx, sy, sz, w], axis=-1) + + def markov_chain_monte_carlo_proposal( + self, + dipole: numpy.ndarray, + stdev: pdme.subspace_simulation.DipoleStandardDeviation, + rng_arg: Optional[numpy.random.Generator] = None, + ) -> numpy.ndarray: + if rng_arg is None: + rng_to_use = self.rng + else: + rng_to_use = rng_arg + + px = dipole[0] + py = dipole[1] + pz = dipole[2] + theta = numpy.arccos(pz / self.pfixed) + phi = numpy.arctan2(py, px) + + # need to step phi, theta, rx, ry, rz, w + # then p^\ast is 1/(2 phi_step) and Delta = phi_step(2 * {0, 1} - 1) + delta_phi = stdev.p_phi_step * rng_to_use.uniform(-1, 1) + tentative_phi = phi + delta_phi + + # theta + delta_theta = stdev.p_theta_step * rng_to_use.uniform(-1, 1) + r = (numpy.sin(theta + delta_theta)) / (numpy.sin(theta)) + if r > rng_to_use.uniform(0, 1): + tentative_theta = theta + delta_theta + else: + tentative_theta = theta + + tentative_px = ( + self.pfixed * numpy.sin(tentative_theta) * numpy.cos(tentative_phi) + ) + tentative_py = ( + self.pfixed * numpy.sin(tentative_theta) * numpy.sin(tentative_phi) + ) + tentative_pz = self.pfixed * numpy.cos(tentative_theta) + + rx = dipole[3] + ry = dipole[4] + rz = dipole[5] + + tentative_rx = rx + stdev.rx_step * rng_to_use.uniform(-1, 1) + if tentative_rx < self.xmin or tentative_rx > self.xmax: + tentative_rx = rx + + tentative_ry = ry + stdev.ry_step * rng_to_use.uniform(-1, 1) + if tentative_ry < self.ymin or tentative_ry > self.ymax: + tentative_ry = ry + tentative_rz = rz + stdev.rz_step * rng_to_use.uniform(-1, 1) + if tentative_rz < self.zmin or tentative_rz > self.zmax: + tentative_rz = rz + + w = dipole[6] + tentative_w = numpy.exp( + numpy.log(w) + (stdev.w_log_step * rng_to_use.uniform(-1, 1)) + ) + tentative_dip = numpy.array( + [ + tentative_px, + tentative_py, + tentative_pz, + tentative_rx, + tentative_ry, + tentative_rz, + tentative_w, + ] + ) + return tentative_dip diff --git a/pdme/model/log_spaced_random_choice_xy_model.py b/pdme/model/log_spaced_random_choice_xy_model.py index 4010235..610dec9 100644 --- a/pdme/model/log_spaced_random_choice_xy_model.py +++ b/pdme/model/log_spaced_random_choice_xy_model.py @@ -5,6 +5,8 @@ from pdme.measurement import ( OscillatingDipole, OscillatingDipoleArrangement, ) +import pdme.subspace_simulation +from typing import Optional class LogSpacedRandomCountMultipleDipoleFixedMagnitudeXYModel(DipoleModel): @@ -123,3 +125,59 @@ class LogSpacedRandomCountMultipleDipoleFixedMagnitudeXYModel(DipoleModel): w = 10 ** rng.uniform(self.wexp_min, self.wexp_max, shape) return numpy.stack([px, py, pz, sx, sy, sz, w], axis=-1) + + def markov_chain_monte_carlo_proposal( + self, + dipole: numpy.ndarray, + stdev: pdme.subspace_simulation.DipoleStandardDeviation, + rng_arg: Optional[numpy.random.Generator] = None, + ) -> numpy.ndarray: + if rng_arg is None: + rng_to_use = self.rng + else: + rng_to_use = rng_arg + + px = dipole[0] + py = dipole[1] + pz = dipole[2] + phi = numpy.arctan2(py, px) + + # need to step phi, rx, ry, rz, w + # then p^\ast is 1/(2 phi_step) and Delta = phi_step(2 * {0, 1} - 1) + delta_phi = stdev.p_phi_step * rng_to_use.uniform(-1, 1) + tentative_phi = phi + delta_phi + + tentative_px = self.pfixed * numpy.cos(tentative_phi) + tentative_py = self.pfixed * numpy.sin(tentative_phi) + + rx = dipole[3] + ry = dipole[4] + rz = dipole[5] + + tentative_rx = rx + stdev.rx_step * rng_to_use.uniform(-1, 1) + if tentative_rx < self.xmin or tentative_rx > self.xmax: + tentative_rx = rx + + tentative_ry = ry + stdev.ry_step * rng_to_use.uniform(-1, 1) + if tentative_ry < self.ymin or tentative_ry > self.ymax: + tentative_ry = ry + tentative_rz = rz + stdev.rz_step * rng_to_use.uniform(-1, 1) + if tentative_rz < self.zmin or tentative_rz > self.zmax: + tentative_rz = rz + + w = dipole[6] + tentative_w = numpy.exp( + numpy.log(w) + (stdev.w_log_step * rng_to_use.uniform(-1, 1)) + ) + tentative_dip = numpy.array( + [ + tentative_px, + tentative_py, + pz, + tentative_rx, + tentative_ry, + tentative_rz, + tentative_w, + ] + ) + return tentative_dip diff --git a/pdme/model/model.py b/pdme/model/model.py index d281d80..c38e937 100644 --- a/pdme/model/model.py +++ b/pdme/model/model.py @@ -4,7 +4,8 @@ from pdme.measurement import ( OscillatingDipoleArrangement, ) import logging - +import pdme.subspace_simulation +from typing import List, Tuple, Optional _logger = logging.getLogger(__name__) @@ -32,3 +33,67 @@ class DipoleModel: For a given DipoleModel, gets a set of dipole collections as a monte_carlo_n x dipole_count x 7 numpy array. """ raise NotImplementedError + + def markov_chain_monte_carlo_proposal( + self, + dipole: numpy.ndarray, + stdev: pdme.subspace_simulation.DipoleStandardDeviation, + rng_arg: Optional[numpy.random.Generator] = None, + ) -> numpy.ndarray: + raise NotImplementedError + + def get_mcmc_chain( + self, + seed, + cost_function, + chain_length, + stdevs: pdme.subspace_simulation.MCMCStandardDeviation, + initial_cost: Optional[float] = None, + rng_arg: Optional[numpy.random.Generator] = None, + ) -> List[Tuple[float, numpy.ndarray]]: + """ + performs constrained markov chain monte carlo starting on seed parameter. + The cost function given is used as a constrained to condition the chain; + a new state is only accepted if cost_function(state) < cost_function(previous_state). + The stdevs passed in are the stdevs we're expected to use. + + Because we're using this for subspace simulation where our proposal function is not too important, we're in good shape. + Note that for our adaptive stdevs to work, there's an unwritten contract that we sort each dipole in the state by frequency (increasing). + + The seed is a list of dipoles, and each chain state is a list of dipoles as well. + + initial_cost is a performance guy that lets you pre-populate the initial cost used to define the condition. + Probably premature optimisation. + + Returns a chain of [ (cost: float, state: dipole_ndarray ) ] format. + """ + _logger.debug( + f"Starting Markov Chain Monte Carlo with seed: {seed} for chain length {chain_length} and provided stdevs {stdevs}" + ) + chain: List[Tuple[float, numpy.ndarray]] = [] + current = seed + if initial_cost is None: + cost_to_compare = cost_function(current) + else: + cost_to_compare = initial_cost + current_cost = cost_to_compare + for i in range(chain_length): + dips = [] + for dipole_index, dipole in enumerate(current): + stdev = stdevs[dipole_index] + tentative_dip = self.markov_chain_monte_carlo_proposal( + dipole, stdev, rng_arg + ) + + dips.append(tentative_dip) + dips_array = pdme.subspace_simulation.sort_array_of_dipoles_by_frequency( + dips + ) + tentative_cost = cost_function(dips_array) + if tentative_cost < cost_to_compare: + chain.append((tentative_cost, dips_array)) + current = dips_array + current_cost = tentative_cost + else: + chain.append((current_cost, current)) + return chain diff --git a/pdme/subspace_simulation/__init__.py b/pdme/subspace_simulation/__init__.py index b7467bc..83593e6 100644 --- a/pdme/subspace_simulation/__init__.py +++ b/pdme/subspace_simulation/__init__.py @@ -1,6 +1,10 @@ from dataclasses import dataclass from typing import Sequence import numpy +from pdme.subspace_simulation.mcmc_costs import ( + proportional_cost, + proportional_costs_vs_actual_measurement, +) @dataclass @@ -40,3 +44,12 @@ def sort_array_of_dipoles_by_frequency(configuration) -> numpy.ndarray: Utility function. """ return numpy.array(sorted(configuration, key=lambda l: l[6])) + + +__all__ = [ + "DipoleStandardDeviation", + "MCMCStandardDeviation", + "sort_array_of_dipoles_by_frequency", + "proportional_cost", + "proportional_costs_vs_actual_measurement", +] diff --git a/pdme/subspace_simulation/mcmc_costs.py b/pdme/subspace_simulation/mcmc_costs.py new file mode 100644 index 0000000..4fd79eb --- /dev/null +++ b/pdme/subspace_simulation/mcmc_costs.py @@ -0,0 +1,20 @@ +import numpy +import numpy.typing +import pdme.util.fast_v_calc + + +def proportional_cost(a: numpy.ndarray, b: numpy.ndarray) -> numpy.ndarray: + tops = numpy.max(b / a, axis=-1) + bottoms = numpy.max(a / b, axis=-1) + return numpy.maximum(tops, bottoms) + + +def proportional_costs_vs_actual_measurement( + dot_inputs_array: numpy.ndarray, + actual_measurement_array: numpy.ndarray, + dipoles_to_test: numpy.ndarray, +) -> numpy.ndarray: + vals = pdme.util.fast_v_calc.fast_vs_for_dipoleses( + dot_inputs_array, numpy.array([dipoles_to_test]) + ) + return proportional_cost(actual_measurement_array, vals) diff --git a/tests/model/__snapshots__/test_log_spaced_fixed_orientation_fixedxy_orientation_mcmc.ambr b/tests/model/__snapshots__/test_log_spaced_fixed_orientation_fixedxy_orientation_mcmc.ambr new file mode 100644 index 0000000..a2069a7 --- /dev/null +++ b/tests/model/__snapshots__/test_log_spaced_fixed_orientation_fixedxy_orientation_mcmc.ambr @@ -0,0 +1,55 @@ +# serializer version: 1 +# name: test_log_spaced_fixedxy_orientation_mcmc_basic + list([ + tuple( + array([3984.46179656]), + array([[ 9.55610128, 2.94634152, 0. , 9.21529051, -2.46576127, + 2.42481096, 9.19034554]]), + ), + tuple( + array([8583.99087872]), + array([[ 9.99991539, 0.04113671, 0. , 8.71258954, -2.26599865, + 2.60452102, 6.37042214]]), + ), + tuple( + array([6215.6376616]), + array([[ 9.81950685, -1.89137124, 0. , 8.90637055, -2.48043039, + 2.28444435, 8.84239221]]), + ), + tuple( + array([424.73328466]), + array([[ 1.00028483, 9.94984574, 0. , 8.53064898, -2.59230757, + 2.33774773, 8.6714416 ]]), + ), + tuple( + array([300.92203808]), + array([[ 1.4003442 , 9.90146636, 0. , 8.05557992, -2.6753126 , + 2.65915755, 13.02021385]]), + ), + tuple( + array([2400.01072771]), + array([[ 9.97761813, 0.66868263, 0. , 8.69171028, -2.73145011, + 2.90140456, 19.94999593]]), + ), + tuple( + array([5001.46205113]), + array([[ 9.93976109, -1.09596962, 0. , 8.95245025, -2.59409162, + 2.90140456, 9.75535945]]), + ), + tuple( + array([195.21980745]), + array([[ 0.20690762, 9.99785923, 0. , 9.59636585, -2.83240984, + 2.90140456, 16.14771567]]), + ), + tuple( + array([2698.2588445]), + array([[-9.68130127, -2.50447712, 0. , 8.94823619, -2.92889659, + 2.77065328, 13.63173263]]), + ), + tuple( + array([1193.69854739]), + array([[-6.16597091, -7.87278875, 0. , 9.62210721, -2.75993924, + 2.77065328, 5.64553534]]), + ), + ]) +# --- diff --git a/tests/model/__snapshots__/test_log_spaced_fixed_orientation_free_orientation_mcmc.ambr b/tests/model/__snapshots__/test_log_spaced_fixed_orientation_free_orientation_mcmc.ambr new file mode 100644 index 0000000..3034f5c --- /dev/null +++ b/tests/model/__snapshots__/test_log_spaced_fixed_orientation_free_orientation_mcmc.ambr @@ -0,0 +1,55 @@ +# serializer version: 1 +# name: test_log_spaced_free_orientation_mcmc_basic + list([ + tuple( + array([3167.67112687]), + array([[ 9.60483896, -1.41627817, -2.3960853 , -4.76615152, -1.80902942, + 2.11809123, 16.17452242]]), + ), + tuple( + array([3167.67112687]), + array([[ 9.60483896, -1.41627817, -2.3960853 , -4.76615152, -1.80902942, + 2.11809123, 16.17452242]]), + ), + tuple( + array([3167.67112687]), + array([[ 9.60483896, -1.41627817, -2.3960853 , -4.76615152, -1.80902942, + 2.11809123, 16.17452242]]), + ), + tuple( + array([736.03065271]), + array([[ 4.1660069 , -8.11557337, 4.0965663 , -4.35968351, -1.97945216, + 2.43615641, 12.92143144]]), + ), + tuple( + array([736.03065271]), + array([[ 4.1660069 , -8.11557337, 4.0965663 , -4.35968351, -1.97945216, + 2.43615641, 12.92143144]]), + ), + tuple( + array([736.03065271]), + array([[ 4.1660069 , -8.11557337, 4.0965663 , -4.35968351, -1.97945216, + 2.43615641, 12.92143144]]), + ), + tuple( + array([2248.07799863]), + array([[-1.71755535, -5.59925137, 8.10545419, -4.03306318, -1.81098441, + 2.77407111, 32.28020575]]), + ), + tuple( + array([1663.31067274]), + array([[-5.16785855, 2.7558756 , 8.10545419, -3.34620897, -1.74763642, + 2.42770463, 52.98214008]]), + ), + tuple( + array([1329.27041439]), + array([[ -1.39600464, 9.69718343, -2.00394725, -2.59147366, + -1.91246681, 2.07361175, 123.01833742]]), + ), + tuple( + array([355.76955919]), + array([[ 9.76047401, 0.84696075, -2.00394725, -3.04310053, + -1.99338573, 2.1185589 , 271.35743739]]), + ), + ]) +# --- diff --git a/tests/model/__snapshots__/test_log_spaced_fixed_orientation_mcmc.ambr b/tests/model/__snapshots__/test_log_spaced_fixed_orientation_mcmc.ambr new file mode 100644 index 0000000..20c470c --- /dev/null +++ b/tests/model/__snapshots__/test_log_spaced_fixed_orientation_mcmc.ambr @@ -0,0 +1,55 @@ +# serializer version: 1 +# name: test_log_spaced_fixed_orientation_mcmc_basic + list([ + tuple( + array([50.56831193]), + array([[ 0. , 0. , 10. , -2.3960853 , 4.23246234, + 2.26169242, 39.39900844]]), + ), + tuple( + array([50.56831193]), + array([[ 0. , 0. , 10. , -2.3960853 , 4.23246234, + 2.26169242, 39.39900844]]), + ), + tuple( + array([47.40865455]), + array([[ 0. , 0. , 10. , -2.03666518, 4.14084039, + 2.21309317, 47.82371559]]), + ), + tuple( + array([47.40865455]), + array([[ 0. , 0. , 10. , -2.03666518, 4.14084039, + 2.21309317, 47.82371559]]), + ), + tuple( + array([47.40865455]), + array([[ 0. , 0. , 10. , -2.03666518, 4.14084039, + 2.21309317, 47.82371559]]), + ), + tuple( + array([47.40865455]), + array([[ 0. , 0. , 10. , -2.03666518, 4.14084039, + 2.21309317, 47.82371559]]), + ), + tuple( + array([22.93279028]), + array([[ 0. , 0. , 10. , -1.63019717, 3.97041764, + 2.53115835, 38.2051999 ]]), + ), + tuple( + array([28.81197733]), + array([[ 0. , 0. , 10. , -1.14570315, 4.07709911, + 2.48697441, 49.58615195]]), + ), + tuple( + array([28.81197733]), + array([[ 0. , 0. , 10. , -1.14570315, 4.07709911, + 2.48697441, 49.58615195]]), + ), + tuple( + array([40.97406005]), + array([[ 0. , 0. , 10. , -0.50178755, 3.83878089, + 2.93560796, 82.07827571]]), + ), + ]) +# --- diff --git a/tests/model/test_log_spaced_fixed_orientation_fixedxy_orientation_mcmc.py b/tests/model/test_log_spaced_fixed_orientation_fixedxy_orientation_mcmc.py new file mode 100644 index 0000000..12dab54 --- /dev/null +++ b/tests/model/test_log_spaced_fixed_orientation_fixedxy_orientation_mcmc.py @@ -0,0 +1,90 @@ +from pdme.model import ( + LogSpacedRandomCountMultipleDipoleFixedMagnitudeXYModel, +) +import pdme.inputs +import pdme.measurement.input_types +import pdme.subspace_simulation +import numpy + +SEED_TO_USE = 42 + + +def get_cost_function(): + 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 = LogSpacedRandomCountMultipleDipoleFixedMagnitudeXYModel( + 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(SEED_TO_USE) + + freqs = [0.01, 0.1, 1, 10, 100] + dot_positions = [[-1.5, 0, 0], [-0.5, 0, 0], [0.5, 0, 0], [1.5, 0, 0]] + dot_inputs = pdme.inputs.inputs_with_frequency_range(dot_positions, freqs) + dot_input_array = pdme.measurement.input_types.dot_inputs_to_array(dot_inputs) + + actual_dipoles = model.get_dipoles(0, numpy.random.default_rng(SEED_TO_USE)) + actual_measurements = actual_dipoles.get_dot_measurements(dot_inputs) + actual_measurements_array = numpy.array([m.v for m in actual_measurements]) + + def cost_to_use(sample_dipoles: numpy.ndarray) -> numpy.ndarray: + return pdme.subspace_simulation.proportional_costs_vs_actual_measurement( + dot_input_array, actual_measurements_array, sample_dipoles + ) + + return cost_to_use + + +def test_log_spaced_fixedxy_orientation_mcmc_basic(snapshot): + + 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 = LogSpacedRandomCountMultipleDipoleFixedMagnitudeXYModel( + 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) + + seed = model.get_monte_carlo_dipole_inputs(1, -1)[0] + + cost_function = get_cost_function() + stdev = pdme.subspace_simulation.DipoleStandardDeviation(2, 2, 1, 0.25, 0.5, 1) + stdevs = pdme.subspace_simulation.MCMCStandardDeviation([stdev]) + + chain = model.get_mcmc_chain( + seed, cost_function, 10, stdevs, rng_arg=numpy.random.default_rng(1515) + ) + + assert chain == snapshot diff --git a/tests/model/test_log_spaced_fixed_orientation_free_orientation_mcmc.py b/tests/model/test_log_spaced_fixed_orientation_free_orientation_mcmc.py new file mode 100644 index 0000000..b5f7398 --- /dev/null +++ b/tests/model/test_log_spaced_fixed_orientation_free_orientation_mcmc.py @@ -0,0 +1,90 @@ +from pdme.model import ( + LogSpacedRandomCountMultipleDipoleFixedMagnitudeModel, +) +import pdme.inputs +import pdme.measurement.input_types +import pdme.subspace_simulation +import numpy + +SEED_TO_USE = 42 + + +def get_cost_function(): + 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(SEED_TO_USE) + + freqs = [0.01, 0.1, 1, 10, 100] + dot_positions = [[-1.5, 0, 0], [-0.5, 0, 0], [0.5, 0, 0], [1.5, 0, 0]] + dot_inputs = pdme.inputs.inputs_with_frequency_range(dot_positions, freqs) + dot_input_array = pdme.measurement.input_types.dot_inputs_to_array(dot_inputs) + + actual_dipoles = model.get_dipoles(0, numpy.random.default_rng(SEED_TO_USE)) + actual_measurements = actual_dipoles.get_dot_measurements(dot_inputs) + actual_measurements_array = numpy.array([m.v for m in actual_measurements]) + + def cost_to_use(sample_dipoles: numpy.ndarray) -> numpy.ndarray: + return pdme.subspace_simulation.proportional_costs_vs_actual_measurement( + dot_input_array, actual_measurements_array, sample_dipoles + ) + + return cost_to_use + + +def test_log_spaced_free_orientation_mcmc_basic(snapshot): + + 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) + + seed = model.get_monte_carlo_dipole_inputs(1, -1)[0] + + cost_function = get_cost_function() + stdev = pdme.subspace_simulation.DipoleStandardDeviation(2, 2, 1, 0.25, 0.5, 1) + stdevs = pdme.subspace_simulation.MCMCStandardDeviation([stdev]) + + chain = model.get_mcmc_chain( + seed, cost_function, 10, stdevs, rng_arg=numpy.random.default_rng(1515) + ) + + assert chain == snapshot diff --git a/tests/model/test_log_spaced_fixed_orientation_mcmc.py b/tests/model/test_log_spaced_fixed_orientation_mcmc.py new file mode 100644 index 0000000..46a7f07 --- /dev/null +++ b/tests/model/test_log_spaced_fixed_orientation_mcmc.py @@ -0,0 +1,98 @@ +from pdme.model import ( + LogSpacedRandomCountMultipleDipoleFixedMagnitudeFixedOrientationModel, +) +import pdme.inputs +import pdme.measurement.input_types +import pdme.subspace_simulation +import numpy + +SEED_TO_USE = 42 + + +def get_cost_function(): + 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(SEED_TO_USE) + + freqs = [0.01, 0.1, 1, 10, 100] + dot_positions = [[-1.5, 0, 0], [-0.5, 0, 0], [0.5, 0, 0], [1.5, 0, 0]] + dot_inputs = pdme.inputs.inputs_with_frequency_range(dot_positions, freqs) + dot_input_array = pdme.measurement.input_types.dot_inputs_to_array(dot_inputs) + + actual_dipoles = model.get_dipoles(0, numpy.random.default_rng(SEED_TO_USE)) + actual_measurements = actual_dipoles.get_dot_measurements(dot_inputs) + actual_measurements_array = numpy.array([m.v for m in actual_measurements]) + + def cost_to_use(sample_dipoles: numpy.ndarray) -> numpy.ndarray: + return pdme.subspace_simulation.proportional_costs_vs_actual_measurement( + dot_input_array, actual_measurements_array, sample_dipoles + ) + + return cost_to_use + + +def test_log_spaced_fixed_orientation_mcmc_basic(snapshot): + + 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) + + seed = model.get_monte_carlo_dipole_inputs(1, -1)[0] + + cost_function = get_cost_function() + stdev = pdme.subspace_simulation.DipoleStandardDeviation(2, 2, 1, 0.25, 0.5, 1) + stdevs = pdme.subspace_simulation.MCMCStandardDeviation([stdev]) + + chain = model.get_mcmc_chain( + seed, cost_function, 10, stdevs, rng_arg=numpy.random.default_rng(1515) + ) + + assert chain == snapshot diff --git a/tests/subspace_simulation/__snapshots__/test_costs.ambr b/tests/subspace_simulation/__snapshots__/test_costs.ambr new file mode 100644 index 0000000..d6f0270 --- /dev/null +++ b/tests/subspace_simulation/__snapshots__/test_costs.ambr @@ -0,0 +1,4 @@ +# serializer version: 1 +# name: test_proportional_costs + 7000.0 +# --- diff --git a/tests/subspace_simulation/test_costs.py b/tests/subspace_simulation/test_costs.py new file mode 100644 index 0000000..3625c7b --- /dev/null +++ b/tests/subspace_simulation/test_costs.py @@ -0,0 +1,10 @@ +import pdme.subspace_simulation +import numpy + + +def test_proportional_costs(snapshot): + a = numpy.array([2, 4, 5, 6, 7, 8, 10]) + b = numpy.array([51, 13, 1, 31, 0.001, 3, 1]) + + actual_result = pdme.subspace_simulation.proportional_cost(a, b).tolist() + assert actual_result == snapshot diff --git a/tests/subspace_simulation/test_sort_dipoles.py b/tests/subspace_simulation/test_sort_dipoles.py index 550c358..2d8c28e 100644 --- a/tests/subspace_simulation/test_sort_dipoles.py +++ b/tests/subspace_simulation/test_sort_dipoles.py @@ -11,7 +11,5 @@ def test_sort_dipoles_by_freq(snapshot): ] ) - actual_sorted = ( - pdme.subspace_simulation.sort_array_of_dipoles_by_frequency(orig) - ) + actual_sorted = pdme.subspace_simulation.sort_array_of_dipoles_by_frequency(orig) assert actual_sorted.tolist() == snapshot