pdme/inputs/__init__.py

@@ -0,0 +1,4 @@
+from pdme.inputs.dot_inputs import inputs_with_frequency_range, input_pairs_with_frequency_range
+
+
+__all__ = ['inputs_with_frequency_range', 'input_pairs_with_frequency_range']

pdme/inputs/dot_inputs.py

@@ -0,0 +1,13 @@
+import numpy
+import numpy.typing
+import itertools
+from typing import Sequence, Tuple
+
+
+def inputs_with_frequency_range(dots: Sequence[numpy.typing.ArrayLike], frequency_range: Sequence[float]) -> Sequence[Tuple[numpy.typing.ArrayLike, float]]:
+	return list(itertools.chain(*[[(dot, f) for f in frequency_range] for dot in dots]))
+
+
+def input_pairs_with_frequency_range(dots: Sequence[numpy.typing.ArrayLike], frequency_range: Sequence[float]) -> Sequence[Tuple[numpy.typing.ArrayLike, numpy.typing.ArrayLike, float]]:
+	all_pairs = itertools.combinations(dots, 2)
+	return list(itertools.chain(*[[(dot1, dot2, f) for f in frequency_range] for (dot1, dot2) in all_pairs]))

tests/inputs/test_dot_inputs.py

@@ -0,0 +1,31 @@
+import pdme.inputs
+
+
+def test_inputs_with_frequency_range():
+	dot1 = [1, 2, 3]
+	dot2 = [2, 4, 6]
+
+	frequencies = [5, 7, 9]
+
+	expected = [
+		([1, 2, 3], 5), ([1, 2, 3], 7), ([1, 2, 3], 9), ([2, 4, 6], 7), ([2, 4, 6], 9), ([2, 4, 6], 5)
+	]
+
+	actual = pdme.inputs.inputs_with_frequency_range([dot1, dot2], frequencies)
+
+	assert sorted(actual) == sorted(expected), "Did not actually match dot inputs!"
+
+
+def test_input_pairs_with_frequency_range():
+	dot1 = [1, 2, 3]
+	dot2 = [2, 4, 6]
+
+	frequencies = [5, 7, 9]
+
+	expected = [
+		([1, 2, 3], [2, 4, 6], 5), ([1, 2, 3], [2, 4, 6], 7), ([1, 2, 3], [2, 4, 6], 9)
+	]
+
+	actual = pdme.inputs.input_pairs_with_frequency_range([dot1, dot2], frequencies)
+
+	assert sorted(actual) == sorted(expected), "Did not actually match dot inputs!"