46 lines
1.3 KiB
Python
46 lines
1.3 KiB
Python
import numpy
|
|
|
|
from ..fdmath import unvec, vec
|
|
from ._test_builders import complex_ramp, real_ramp
|
|
from .utils import assert_close
|
|
|
|
|
|
SHAPE = (2, 3, 2)
|
|
FIELD = real_ramp((3, *SHAPE))
|
|
COMPLEX_FIELD = complex_ramp((3, *SHAPE), imag_scale=0.5)
|
|
|
|
|
|
def test_vec_and_unvec_return_none_for_none_input() -> None:
|
|
assert vec(None) is None
|
|
assert unvec(None, SHAPE) is None
|
|
|
|
|
|
def test_real_field_round_trip_preserves_shape_and_values() -> None:
|
|
vector = vec(FIELD)
|
|
assert vector is not None
|
|
restored = unvec(vector, SHAPE)
|
|
assert restored is not None
|
|
assert restored.shape == (3, *SHAPE)
|
|
assert_close(restored, FIELD)
|
|
|
|
|
|
def test_complex_field_round_trip_preserves_shape_and_values() -> None:
|
|
vector = vec(COMPLEX_FIELD)
|
|
assert vector is not None
|
|
restored = unvec(vector, SHAPE)
|
|
assert restored is not None
|
|
assert restored.shape == (3, *SHAPE)
|
|
assert_close(restored, COMPLEX_FIELD)
|
|
|
|
|
|
def test_unvec_with_two_components_round_trips_vector() -> None:
|
|
vector = numpy.arange(2 * numpy.prod(SHAPE), dtype=float)
|
|
field = unvec(vector, SHAPE, nvdim=2)
|
|
assert field is not None
|
|
assert field.shape == (2, *SHAPE)
|
|
assert_close(vec(field), vector)
|
|
|
|
|
|
def test_vec_accepts_arraylike_input() -> None:
|
|
arraylike = [[[1, 2], [3, 4]], [[5, 6], [7, 8]]]
|
|
assert_close(vec(arraylike), numpy.ravel(arraylike, order='C'))
|