def test_spherical_grid_eq_equal_length_grids_return_true(self):
r2 = 1. * r # Use this so that x2 == x but is not x
assert r2 is not r
assert np.array_equal(r, r2)
grid = gr.SphericalGrid((r, t, p))
grid2 = gr.SphericalGrid((r2, t, p))
self.assertEqual(grid, grid2)
def test_spherical_grid_gradient_container_type(self):
xm, ym, zm = np.meshgrid(r, t, p, indexing='ij')
tuple_grid = gr.SphericalGrid((r, t, p))
list_grid = gr.SphericalGrid([r, t, p])
tuple_grad = tuple_grid.gradient(xm)
list_grad = list_grid.gradient(ym)
self.assertIs(type(tuple_grad), tuple)
self.assertIs(type(list_grad), list)
def test_spherical_grid_eq_very_different_equal_length_grids_return_false(
self):
r2 = 2. * r
assert not np.array_equal(r, r2)
grid = gr.SphericalGrid((r, t, p))
grid2 = gr.CartesianGrid((r2, t, p))
self.assertNotEqual(grid, grid2)
def test_spherical_grid_init_equal_length_inputs_cartesian_equality(self):
# Initialize Grids with equal length coordinate arrays
grid = gr.SphericalGrid((r, t, p))
log_grid = gr.SphericalGrid((r_log, t_log, p_log))
# Define Expected Behavior
expected_cart = sp.spherical_to_cartesian(grid.mesh)
expected_cart = np.round(expected_cart, 7)
expected_log_cart = sp.spherical_to_cartesian(log_grid.mesh)
expected_log_cart = np.round(expected_log_cart, 7)
cart_equality = np.array_equal(np.round(grid.cartesian, 7),
expected_cart)
log_cart_equality = np.array_equal(np.round(log_grid.cartesian, 7),
expected_log_cart)
self.assertTrue(cart_equality)
self.assertTrue(log_cart_equality)
def test_spherical_grid_mesh_iter_lists(self):
grid = gr.SphericalGrid([r, t, p])
rm, tm, pm = np.meshgrid(r, t, p, indexing='ij')
r_expected = rm.flatten()
t_expected = tm.flatten()
p_expected = pm.flatten()
for c, re, te, pe in zip(grid.mesh_iter(), r_expected, t_expected,
p_expected):
self.assertEqual(c, [re, te, pe])
def test_spherical_grid_iter_lists(self):
grid = gr.SphericalGrid([r, t, p])
xm, ym, zm = sp.spherical_to_cartesian(
np.meshgrid(r, t, p, indexing='ij'))
x_expected = xm.flatten()
y_expected = ym.flatten()
z_expected = zm.flatten()
for c, xe, ye, ze in zip(grid, x_expected, y_expected, z_expected):
self.assertEqual(c, [xe, ye, ze])
def test_spherical_grid_gradient_unequal_length_arrays_linear(self):
rm, tm, pm = np.meshgrid(r_short, t, p, indexing='ij')
rf = rm
grid = gr.SphericalGrid((r_short, t, p))
grad_r = grid.gradient(rf)
grad_re = (np.ones_like(rf), np.zeros_like(tm), np.zeros_like(pm))
grad_rc = np.round(grad_r, 6)
self.assertTrue(np.array_equal(grad_rc, grad_re))
def test_spherical_grid_init_list_implicit_container_type(self):
grid = gr.SphericalGrid([r_short, t_short, p_short])
container_type = grid._container_type
coordinate_type = type(grid.coordinates)
mesh_type = type(grid.mesh)
cartesian_type = type(grid.cartesian)
self.assertIs(container_type, list)
self.assertIs(coordinate_type, list)
self.assertIs(mesh_type, list)
self.assertIs(cartesian_type, list)
def test_spherical_grid_init_unequal_length_inputs_mesh_equality(self):
# Initialize Grids with equal length coordinate arrays
mixed_grid = gr.SphericalGrid((r, t_short, p_short))
# Define Expected Behavior
expected_mixed_mesh = np.meshgrid(r, t_short, p_short, indexing='ij')
# Check Equality
mixed_mesh_eq = np.array_equal(mixed_grid.mesh, expected_mixed_mesh)
self.assertTrue(mixed_mesh_eq)
def test_spherical_grid_init_unequal_length_inputs_cartesian_equality(
self):
# Initialize Grids with equal length coordinate arrays
mixed_grid = gr.SphericalGrid((r, t_short, p_short))
# Define Expected Behavior
expected_mixed_cart = sp.spherical_to_cartesian(mixed_grid.mesh)
expected_mixed_cart = np.round(expected_mixed_cart, 7)
# Check Equality
mixed_cart_eq = np.array_equal(np.round(mixed_grid.cartesian, 7),
expected_mixed_cart)
self.assertTrue(mixed_cart_eq)
def test_spherical_grid_init_equal_length_inputs_mesh_equality(self):
# Initialize Grids with equal length coordinate arrays
grid = gr.SphericalGrid((r, t, p))
log_grid = gr.CartesianGrid((r_log, t_log, p_log))
# Define Expected Behavior
expected_mesh = np.meshgrid(r, t, p, indexing='ij')
expected_log_mesh = np.meshgrid(r_log, t_log, p_log, indexing='ij')
mesh_equality = np.array_equal(grid.mesh, expected_mesh)
log_mesh_equality = np.array_equal(log_grid.mesh, expected_log_mesh)
self.assertTrue(mesh_equality)
self.assertTrue(log_mesh_equality)
def test_spherical_grid_gradient_unequal_length_arrays_quadratic(self):
rm, tm, pm = np.meshgrid(r, t, p, indexing='ij')
rf = rm**2
grid = gr.SphericalGrid((r, t, p))
grad_r = grid.gradient(rf)
grad_re = np.round((2 * rm, np.zeros_like(tm), np.zeros_like(pm)), 6)
grad_rc = np.round(grad_r, 6)
test = grad_re == grad_rc
t0 = np.sum(test[0])
t1 = np.sum(test[1])
t2 = np.sum(test[2])
self.assertTrue(np.array_equal(grad_rc, grad_re))
def test_cartesian_grid_divergence_unequal_length_arrays_azimuthal_vectors(
self):
rl = r_short + 1E-12 * np.ones_like(r_short)
tl = t + 1E-12 * np.ones_like(t)
pl = r + 1E-12 * np.ones_like(p)
rm, tm, pm = np.meshgrid(rl, tl, pl, indexing='ij')
vector_field = (np.zeros_like(rm), np.zeros_like(tm), rm * np.sin(tm))
divergence_expected = np.zeros_like(rm)
grid = gr.SphericalGrid((rl, tl, pl))
divergence_calculated = grid.divergence(vector_field)
# Test Float Equality
divergence_calculated = np.round(divergence_calculated, 7)
self.assertTrue(
np.array_equal(divergence_calculated, divergence_expected))
def test_spherical_grid_init(self):
r_array = (1., 0.5, 0.0)
theta_array = (0.0, pi / 2)
phi_array = (0.0, pi / 2)
correct_grid = sp.spherical_to_cartesian_mesh(
(r_array, theta_array, phi_array))
r = correct_grid[0]
theta = correct_grid[1]
diffs = (np.ones_like(r), np.ones_like(r), np.ones_like(r))
correct_volume = sp.calculate_spherical_volume_element(r, theta, diffs)
grid_object = gr.SphericalGrid((r_array, theta_array, phi_array))
calculated_grid = grid_object.grid
calculated_volume = grid_object.volume
for correct, calculated in zip(correct_grid, calculated_grid):
self.assertTrue(np.array_equal(correct, calculated))
self.assertTrue(
np.array_equal(np.round(calculated_volume, 7),
np.round(correct_volume, 7)))
def test_spherical_grid_volume_unequal_length_arrays(self):
grid = gr.SphericalGrid((r_short, t, p))
dr, dt, dp = 0.1 * np.ones_like(r_short), 0.05 * pi * np.ones_like(
t), 0.1 * pi * np.ones_like(p)
dr[0] = 0.5 * dr[0]
dr[-1] = 0.5 * dr[-1]
dt[0] = 0.5 * dt[0]
dt[-1] = 0.5 * dt[-1]
dp[0] = 0.5 * dp[0]
dp[-1] = 0.5 * dp[-1]
drg, dtg, dpg = np.meshgrid(dr, dt, dp, indexing='ij')
rg, tg, pg = np.meshgrid(r_short, t, p, indexing='ij')
v_e = sp.calculate_spherical_volume_element(rg, tg, (drg, dtg, dpg))
# Round For float equality check
v_e = np.round(v_e, 7)
v_c = np.round(grid.volume, 7)
self.assertTrue(np.array_equal(v_e, v_c))
def test_spherical_grid_getitem_unequal_length_arrays(self):
grid = gr.SphericalGrid((r_short, t, p))
expected_000 = sp.spherical_to_cartesian((r_short[0], t[0], p[0]))
expected_210 = sp.spherical_to_cartesian((r_short[2], t[1], p[0]))
self.assertEqual(grid[0, 0, 0], expected_000)
self.assertEqual(grid[2, 1, 0], expected_210)
def test_spherical_grid_eq_nearly_equal_length_grids_return_false(self):
r2 = r + 1E-7 * np.ones_like(r)
assert not np.array_equal(r, r2)
grid = gr.SphericalGrid((r, t, p))
grid2 = gr.SphericalGrid((r2, t, p))
self.assertNotEqual(grid, grid2)
# Initialize Cartesian Coordinate Arrays
x = np.linspace(-1., 1., 21)
y = np.linspace(-1., 1., 21)
z = np.linspace(-1., 1., 21)
c_grid = gr.CartesianGrid((x, y, z))
x_short = np.linspace(-1., 1., 11)
y_short = np.linspace(-1., 1., 11)
z_short = np.linspace(-1., 1., 11)
short_c_grid = gr.CartesianGrid((x_short, y_short, z_short))
# Initialize Spherical Coordinate Arrays
r = np.linspace(0., 1., 21)
t = pi * np.linspace(0., 1, 21)
p = 2 * pi * np.linspace(0., 1, 21)
s_grid = gr.SphericalGrid((r, t, p))
r_short = np.linspace(0., 1., 11)
t_short = pi * np.linspace(0., 1, 11)
p_short = 2 * pi * np.linspace(0., 1, 11)
short_s_grid = gr.SphericalGrid((r_short, t_short, p_short))
r_log = np.logspace(-3, 0, 11)
t_log = pi * np.logspace(-3, 0, 11)
p_log = 2 * pi * np.logspace(-5, 0, 11)
# Test Grid Classes
class TestCartesianGrid(TestCase):
def test_cartesian_grid_init_equal_length_inputs_mesh_equality(self):
# Initialize Grids with equal length coordinate arrays
short_grid = gr.CartesianGrid((x_short, y_short, z_short))
long_grid = gr.CartesianGrid((x, y, z))
def test_spherical_grid_mesh_iter_types(self):
tuple_grid = gr.SphericalGrid((r, t, p))
list_grid = gr.SphericalGrid([r, t, p])
for tup, l in zip(tuple_grid.mesh_iter(), list_grid.mesh_iter()):
self.assertIs(type(tup), tuple)
self.assertIs(type(l), list)
def test_spherical_grid_mesh_item_unequal_length_arrays(self):
grid = gr.SphericalGrid((r_short, t, p))
expected_000 = (r_short[0], t[0], p[0])
expected_210 = (r_short[2], t[1], p[0])
self.assertEqual(grid.mesh_item(0, 0, 0), expected_000)
self.assertEqual(grid.mesh_item(2, 1, 0), expected_210)
def test_spherical_grid_mesh_item_type_list(self):
grid = gr.SphericalGrid([r, t, p])
getitem_type = type(grid.mesh_item(0, 0, 0))
self.assertIs(getitem_type, list)
self.assertIs(getitem_type, grid._container_type)
res_z = 5
x = np.linspace(-0.5, 0.5, res_x)
y = np.linspace(-0.5, 0.5, res_y)
z = np.linspace(-0.5, 0.5, res_z)
cart_grid = gr.CartesianGrid((x, y, z))
# Make Spherical Grid
res_r = 5
res_theta = 5
res_phi = 5
r = np.linspace(0.05, 0.5, res_r)
theta = np.linspace(0., np.pi, res_theta)
phi = np.linspace(0., 2 * np.pi, res_phi)
spher_grid = gr.SphericalGrid((r, theta, phi))
# Initialize Operators
cart_ops = gr.VectorGrid(cart_grid.cartesian,
cart_grid,
operator_type=la.CartesianSpinOperator)
spher_ops = gr.VectorGrid(spher_grid.mesh,
spher_grid,
operator_type=la.SphericalSpinOperator)
# Calculate Unitary Derivatives for Both Grids
cart_rho_dot = grid_unitary_derivative(cart_ops, ham)
spher_rho_dot = grid_bloch_derivative(spher_ops, ham, g_diss, g_diss, r_pump)
# Mask Cartesian Plots on the Bloch Sphere
x_g, y_g, z_g = cart_grid.mesh
def test_spherical_grid_getitem_type_list(self):
grid = gr.SphericalGrid([x, y, z])
getitem_type = type(grid[0, 0, 0])
self.assertIs(getitem_type, list)
self.assertIs(getitem_type, grid._container_type)
