def test_hermval2d(self):
x1, x2, x3 = self.x
y1, y2, y3 = self.y
#test exceptions
assert_raises(ValueError, herm.hermval2d, x1, x2[:2], self.c2d)
#test values
tgt = y1 * y2
res = herm.hermval2d(x1, x2, self.c2d)
assert_almost_equal(res, tgt)
#test shape
z = np.ones((2, 3))
res = herm.hermval2d(z, z, self.c2d)
assert_(res.shape == (2, 3))
def test_hermval2d(self):
x1, x2, x3 = self.x
y1, y2, y3 = self.y
#test exceptions
assert_raises(ValueError, herm.hermval2d, x1, x2[:2], self.c2d)
#test values
tgt = y1*y2
res = herm.hermval2d(x1, x2, self.c2d)
assert_almost_equal(res, tgt)
#test shape
z = np.ones((2,3))
res = herm.hermval2d(z, z, self.c2d)
assert_(res.shape == (2,3))
def two_mode_wavefunction(ket, l=4.5, N=100):
"""Calculate the two-mode wavefunction of a state vector.
Args:
ket (array): the two mode state vector of size cutoff^2
l (float): the maximum domain size to calculate the wavefunction on.
N (int): number of discrete points in the domain.
Returns:
tuple (X, Y, Z): the x, y and z 2D-grids of the two-mode wavefunction.
"""
c = int(np.sqrt(ket.shape[0]))
output_state = ket.reshape(c, c)
n = np.arange(c)[:, None]
m = np.arange(c)[None, :]
coefficients = np.real(output_state) / (
(np.sqrt(factorial(n) * (2 ** n))) * (np.sqrt(factorial(m) * (2 ** m))))
x = np.linspace(-l, l, N)
y = np.linspace(-l, l, N)
X, Y = np.meshgrid(x, y)
Z = (np.exp(-X ** 2 / 2))*(np.exp(-Y ** 2 / 2))*hermval2d(X, Y, coefficients)
return X, Y, Z
def test_hermvander2d(self):
# also tests hermval2d for non-square coefficient array
x1, x2, x3 = self.x
c = np.random.random((2, 3))
van = herm.hermvander2d(x1, x2, [1, 2])
tgt = herm.hermval2d(x1, x2, c)
res = np.dot(van, c.flat)
assert_almost_equal(res, tgt)
# get_inds shape
van = herm.hermvander2d([x1], [x2], [1, 2])
assert_(van.shape == (1, 5, 6))
def test_hermvander2d(self) :
# also tests hermval2d for non-square coefficient array
x1, x2, x3 = self.x
c = np.random.random((2, 3))
van = herm.hermvander2d(x1, x2, [1, 2])
tgt = herm.hermval2d(x1, x2, c)
res = np.dot(van, c.flat)
assert_almost_equal(res, tgt)
# check shape
van = herm.hermvander2d([x1], [x2], [1, 2])
assert_(van.shape == (1, 5, 6))
