def test_unit_attribute_error_case():
"""
Verify that the unit attribute raises an exception if the units on all
axes are not the same.
"""
grid = grids.AbstractGrid(
[-1 * u.cm, 0 * u.rad, -2 * u.cm],
[1 * u.cm, 2 * np.pi * u.rad, 2 * u.cm],
num=5,
)
with pytest.raises(ValueError):
grid.unit
def test_AbstractGrid():
grid = grids.AbstractGrid(-1 * u.cm, 1 * u.cm, num=(10, 20, 5))
array = grid.grid
units = grid.units
pts0, pts1, pts2 = grid.pts0, grid.pts1, grid.pts2
# Test wrong number of positional arguments: 1 or more than 3
with pytest.raises(TypeError):
grid = grids.AbstractGrid(1 * u.cm, num=10)
with pytest.raises(TypeError):
grid = grids.AbstractGrid(-1 * u.cm, 1 * u.cm, 1 * u.cm, 1 * u.cm)
# Test unequal lengths of arguments raises error
with pytest.raises(ValueError):
grid = grids.AbstractGrid(-1 * u.m, [2 * u.m, 3 * u.m], num=10)
with pytest.raises(ValueError):
grid = grids.AbstractGrid(
np.random.randn(2, 5, 3) * u.m,
np.random.randn(2, 5, 3) * u.m,
np.random.randn(2, 5, 4) * u.m,
)
# Test incompatible units
with pytest.raises(ValueError):
grid = grids.AbstractGrid(1 * u.cm, 1 * u.eV, num=10)
# Test adding a quantity
q1 = np.random.randn(10, 20, 5) * u.kg
grid.add_quantity("test quantity", q1)
# Test adding a quantity of incompatible size
q2 = np.random.randn(5, 20, 5) * u.kg
with pytest.raises(ValueError):
grid.add_quantity("test quantity2", q2)
def test_AbstractGrid():
# Create grid with single u.Quantity args
grid = grids.AbstractGrid(-1 * u.cm, 1 * u.cm, num=10)
assert grid.is_uniform
# Create grid with lists of u.Quantity args
grid = grids.AbstractGrid(
[-1 * u.cm, -1 * u.cm, -1 * u.cm],
[1 * u.cm, 1 * u.cm, 1 * u.cm],
num=[10, 10, 10],
)
assert grid.is_uniform
# Create grid with arrays of u.quantities
grid = grids.AbstractGrid(np.array([-1] * 3) * u.cm,
np.array([1] * 3) * u.cm,
num=[10, 10, 10])
assert grid.is_uniform
print(grid)
array = grid.grid
units = grid.units
pts0, pts1, pts2 = grid.pts0, grid.pts1, grid.pts2
# Test wrong number of positional arguments: 1 or more than 3
with pytest.raises(TypeError):
grid = grids.AbstractGrid(1 * u.cm, num=10)
with pytest.raises(TypeError):
grid = grids.AbstractGrid(-1 * u.cm, 1 * u.cm, 1 * u.cm, 1 * u.cm)
# Test unequal lengths of arguments raises error
with pytest.raises(ValueError):
grid = grids.AbstractGrid(-1 * u.m, [2 * u.m, 3 * u.m], num=10)
with pytest.raises(ValueError):
grid = grids.AbstractGrid(
np.random.randn(2, 5, 3) * u.m,
np.random.randn(2, 5, 3) * u.m,
np.random.randn(2, 5, 4) * u.m,
)
# Test incompatible units
with pytest.raises(ValueError):
grid = grids.AbstractGrid(1 * u.cm, 1 * u.eV, num=10)
# Test adding a quantity
q = np.random.randn(10, 10, 10) * u.T
grid.add_quantities(B_x=q)
# Test setting a subset of a quantity array
grid["B_x"][0, 0, 0] = 21 * u.T
assert grid["B_x"][0, 0, 0] == 21 * u.T
# Test accessing a quantity using __getitem__ or directly
Bx = grid.ds["B_x"]
Bx = grid["B_x"]
# Assert that the array returned is a u.Quantity
assert isinstance(Bx, u.Quantity)
# Assert that the array returned has the right shape
assert Bx.shape == grid.shape
# Test require_quantities
# Test with a key that is there
req_q = ["B_x"]
grid.require_quantities(req_q, replace_with_zeros=False)
req_q = ["B_x", "B_y"]
# Test with a key that is not there, but can be replaced
# Do not replace
with pytest.raises(KeyError):
grid.require_quantities(req_q, replace_with_zeros=False)
# Do replace
with pytest.warns(RuntimeWarning,
match="This quantity will be assumed to be zero"):
grid.require_quantities(req_q, replace_with_zeros=True)
req_q = ["B_x", "B_y"]
# Test with a key that is not there, but cannot be replaced because
# it's not a recognized key
req_q = ["B_x", "not_a_recognized_key"]
with pytest.raises(KeyError):
with pytest.warns(RuntimeWarning,
match="This quantity will be assumed to be zero"):
grid.require_quantities(req_q, replace_with_zeros=True)
# Test adding a quantity with wrong units
q = np.random.randn(10, 10, 10) * u.kg
with pytest.raises(ValueError):
grid.add_quantities(B_x=q)
# Testing adding a quantity with an unrecognized key name
with pytest.warns(UserWarning):
grid.add_quantities(not_a_recognized_key=q)
# Test adding a quantity of incompatible size
q = np.random.randn(5, 20, 5) * u.T
with pytest.raises(ValueError):
grid.add_quantities(B_x=q)
# Test adding multiple quantites at once
q = np.random.randn(10, 10, 10) * u.T
grid.add_quantities(B_x=q, B_y=q, B_z=q)
print(grid)
def test_AbstractGrid():
# Create grid with single u.Quantity args
grid = grids.AbstractGrid(-1 * u.cm, 1 * u.cm, num=10)
assert grid.is_uniform
# Create grid with lists of u.Quantity args
grid = grids.AbstractGrid(
[-1 * u.cm, -1 * u.cm, -1 * u.cm],
[1 * u.cm, 1 * u.cm, 1 * u.cm],
num=[10, 10, 10],
)
assert grid.is_uniform
# Create grid with arrays of u.quantities
grid = grids.AbstractGrid(
np.array([-1] * 3) * u.cm, np.array([1] * 3) * u.cm, num=[10, 10, 10]
)
assert grid.is_uniform
print(grid)
array = grid.grid
units = grid.units
pts0, pts1, pts2 = grid.pts0, grid.pts1, grid.pts2
# Test wrong number of positional arguments: 1 or more than 3
with pytest.raises(TypeError):
grid = grids.AbstractGrid(1 * u.cm, num=10)
with pytest.raises(TypeError):
grid = grids.AbstractGrid(-1 * u.cm, 1 * u.cm, 1 * u.cm, 1 * u.cm)
# Test unequal lengths of arguments raises error
with pytest.raises(ValueError):
grid = grids.AbstractGrid(-1 * u.m, [2 * u.m, 3 * u.m], num=10)
with pytest.raises(ValueError):
grid = grids.AbstractGrid(
np.random.randn(2, 5, 3) * u.m,
np.random.randn(2, 5, 3) * u.m,
np.random.randn(2, 5, 4) * u.m,
)
# Test incompatible units
with pytest.raises(ValueError):
grid = grids.AbstractGrid(1 * u.cm, 1 * u.eV, num=10)
# Test adding a quantity
q = np.random.randn(10, 10, 10) * u.T
grid.add_quantities(B_x=q)
# Test accessing a quantity using __getitem__ or directly
Bx = grid.ds["B_x"]
Bx = grid["B_x"]
# Test adding a quantity with wrong units
q = np.random.randn(10, 10, 10) * u.kg
with pytest.raises(ValueError):
grid.add_quantities(B_x=q)
# Testing adding a quantity with an unrecognized key name
with pytest.warns(UserWarning):
grid.add_quantities(not_a_recognized_key=q)
# Test adding a quantity of incompatible size
q = np.random.randn(5, 20, 5) * u.T
with pytest.raises(ValueError):
grid.add_quantities(B_x=q)
# Test adding multiple quantites at once
q = np.random.randn(10, 10, 10) * u.T
grid.add_quantities(B_x=q, B_y=q, B_z=q)
print(grid)