pdme/tests/subspace_simulation/test_sort_dipoles.py

41 lines
881 B
Python

import numpy
import pdme.subspace_simulation
def test_sort_dipoles_by_freq(snapshot):
orig = numpy.array(
[
[1, 2, 3, 4, 5, 6, 7],
[100, 200, 300, 400, 500, 600, 0.07],
[10, 200, 30, 41, 315, 0.31, 100],
]
)
actual_sorted = pdme.subspace_simulation.sort_array_of_dipoles_by_frequency(orig)
assert actual_sorted.tolist() == snapshot
def test_sort_dipoleses_by_freq(snapshot):
sample_1 = numpy.array(
[
[1, 2, 3, 4, 5, 6, 7],
[100, 200, 300, 400, 500, 600, 0.07],
[10, 200, 30, 41, 315, 0.31, 100],
]
)
sample_2 = numpy.array(
[
[1, 1, 1, 1, 1, 1, 100],
[33, 33.3, 33, 3.3, 0.33, 0.3, 300],
[22, 22.2, 2.2, 222, 22, 2, 200],
]
)
original_samples = numpy.array([sample_1, sample_2])
actual_sorted = pdme.subspace_simulation.sort_array_of_dipoleses_by_frequency(
original_samples
)
assert actual_sorted.tolist() == snapshot