meanas/meanas/test/utils.py

import numpy
from numpy.typing import NDArray
from numpy.typing import ArrayLike


def make_prng(seed: int = 12345) -> numpy.random.RandomState:
    return numpy.random.RandomState(seed)


def assert_fields_close(
        x: NDArray,
        y: NDArray,
        *args,
        **kwargs,
        ) -> None:
    x_disp = numpy.moveaxis(x, -1, 0)
    y_disp = numpy.moveaxis(y, -1, 0)
    numpy.testing.assert_allclose(
        x,           # type: ignore
        y,           # type: ignore
        *args,
        verbose=False,
        err_msg=f'Fields did not match:\n{x_disp}\n{y_disp}',
        **kwargs,
        )

def assert_close(
        x: ArrayLike,
        y: ArrayLike,
        *args,
        **kwargs,
        ) -> None:
    numpy.testing.assert_allclose(numpy.asarray(x), numpy.asarray(y), *args, **kwargs)