tantri/tests/dipoles/test_dipoles.py

84 lines
1.4 KiB
Python
Executable File

from tantri.dipoles import (
DipoleTO,
DipoleTimeSeries,
DotPosition,
DipoleMeasurementType,
)
import numpy
def test_dipoles_1(snapshot):
rng = numpy.random.default_rng(1234)
dots = [DotPosition(numpy.array([0, 0, 0]), "dot1")]
d1 = DipoleTO(
numpy.array([0, 0, 10]),
numpy.array([5, 3, 2]),
15,
)
d2 = DipoleTO(
numpy.array([0, 0, 10]),
numpy.array([-2, 3, 2]),
0.1,
)
d3 = DipoleTO(
numpy.array([0, 0, 10]),
numpy.array([2, 1, 2]),
6,
)
d4 = DipoleTO(
numpy.array([0, 0, 10]),
numpy.array([-5, -5, 2]),
50,
)
ts_gen = DipoleTimeSeries(
[d1, d2, d3, d4],
dots,
DipoleMeasurementType.ELECTRIC_POTENTIAL,
0.01,
rng_to_use=rng,
)
time_series = [ts_gen.transition() for i in range(50)]
assert time_series == snapshot
def test_dipoles_1_field(snapshot):
rng = numpy.random.default_rng(1234)
dots = [DotPosition(numpy.array([0, 0, 0]), "dot1")]
d1 = DipoleTO(
numpy.array([0, 0, 10]),
numpy.array([5, 3, 2]),
15,
)
d2 = DipoleTO(
numpy.array([0, 0, 10]),
numpy.array([-2, 3, 2]),
0.1,
)
d3 = DipoleTO(
numpy.array([0, 0, 10]),
numpy.array([2, 1, 2]),
6,
)
d4 = DipoleTO(
numpy.array([0, 0, 10]),
numpy.array([-5, -5, 2]),
50,
)
ts_gen = DipoleTimeSeries(
[d1, d2, d3, d4],
dots,
DipoleMeasurementType.X_ELECTRIC_FIELD,
0.01,
rng_to_use=rng,
)
time_series = [ts_gen.transition() for i in range(50)]
assert time_series == snapshot