def test_get_all_parameters_with_get_coordinates_of(self):
k = np.array([1, 2, 3]).reshape((-1, 1))
l = np.array([2, 3, 4]).reshape((-1, 1))
array = np.append(k, l, axis=1)
particles = Particles(array, {"x": 0, "y": 1})
received_x_and_y = particles.get_default_coordinates_of("x", "y")
assert np.allclose(received_x_and_y, array)
def test_get_matrix(self):
k = np.array([1, 2, 1]).reshape((-1, 1))
l = np.array([2, 3, 4]).reshape((-1, 1))
array = np.append(k, l, axis=1)
particles = Particles(array, {"x": 0, "y": 1})
array1 = particles.get_matrix()
array1[0][0] = -1
assert np.allclose(array1, particles.particles)
def test_filter_equals(self):
k = np.array([1, 2, 1]).reshape((-1, 1))
l = np.array([2, 3, 4]).reshape((-1, 1))
array = np.append(k, l, axis=1)
particles = Particles(array, {"x": 0, "y": 1})
particles_with_x_eq_1 = particles.filter_equals("x", 1)
assert np.allclose(particles_with_x_eq_1.particles,
np.array([[1, 2], [1, 4]]))
def test_get_values_of(self):
k = np.array([1, 2, 3]).reshape((-1, 1))
l = np.array([2, 3, 4]).reshape((-1, 1))
m = np.array([3, 4, 5]).reshape((-1, 1))
n = np.array([4, 5, 6]).reshape((-1, 1))
array = np.append(k, l, axis=1)
particles = Particles(array, {"x": 0, "y": 1})
self.assertTrue(np.allclose(particles.get_column("x"), k))
def test_get_values_of_should_return_copy_of_vector(self):
k = np.array([1, 2, 3]).reshape((-1, 1))
l = np.array([2, 3, 4]).reshape((-1, 1))
array = np.append(k, l, axis=1)
particles = Particles(array, {"x": 0, "y": 1})
cp = particles.get_column("x")
cp[0][0] = 100
self.assertFalse(np.allclose(particles.get_column("x"), cp))
def test_add_column(self):
k = np.array([1, 2, 3]).reshape((-1, 1))
l = np.array([2, 3, 4]).reshape((-1, 1))
m = np.array([3, 4, 5]).reshape((-1, 1))
array = np.append(k, l, axis=1)
particles = Particles(array, {"x": 0, "y": 1})
particles_with_m = particles.add_column("m", m)
self.assertTrue(np.allclose(particles_with_m.get_column("m"), m))
def test_shift_parameter(self):
k = np.array([1, 2, 1]).reshape((-1, 1))
l = np.array([2, 3, 4]).reshape((-1, 1))
l_shifted = l + 1
array = np.append(k, l, axis=1)
particles = Particles(array, {"x": 0, "y": 1})
particles_with_shifted_y = particles.shift_parameter("y", 1)
assert np.allclose(particles_with_shifted_y.get_column("y"), l_shifted)
def test_get_mapping(self):
k = np.array([1, 2, 1]).reshape((-1, 1))
l = np.array([2, 3, 4]).reshape((-1, 1))
array = np.append(k, l, axis=1)
particles = Particles(array, {"x": 0, "y": 1})
mapping1 = particles.get_mapping()
mapping1["t"] = 1
assert mapping1 == particles.mapping
def test_get_all_parameters_in_different_order_with_get_coordinates_of(
self):
k = np.array([1, 2, 3]).reshape((-1, 1))
l = np.array([2, 3, 4]).reshape((-1, 1))
array = np.append(k, l, axis=1)
particles = Particles(array, {"x": 0, "y": 1})
received_y_and_x = particles.get_default_coordinates_of("y", "x")
l_and_k = np.append(l, k, axis=1)
assert np.allclose(received_y_and_x, l_and_k)
def test_constructor_should_copy_particles(self):
k = np.array([1, 2, 3]).reshape((-1, 1))
l = np.array([2, 3, 4]).reshape((-1, 1))
array = np.append(k, l, axis=1)
particles = Particles(array, {"x": 0, "y": 1})
array[0][0] = 100
self.assertFalse(
np.allclose(particles.get_column("x"), array.T[0].reshape(
(-1, 1))))
def test_add_zeros_column(self):
k = np.array([1, 2, 3]).reshape((-1, 1))
l = np.array([2, 3, 4]).reshape((-1, 1))
m = np.array([0, 0, 0]).reshape((-1, 1))
array = np.append(k, l, axis=1)
particles = Particles(array, {"x": 0, "y": 1})
particles_with_zeros_column = particles.add_zeros_column("z")
array_with_zeros_column = np.append(array, m, axis=1)
assert np.allclose(particles_with_zeros_column.particles,
array_with_zeros_column)
def transport(particles, wanted_xa, configuration_madx_xa1, xa1,
configuration_madx_xa2, xa2):
segments_xa_1 = ptg.transport(configuration_madx_xa1, particles)
segments_xa_2 = ptg.transport(configuration_madx_xa2, particles)
new_places = {}
for place in segments_xa_1.keys():
# (X1 - X2) / (XA1 - XA2) * XA + X2
# todo refactor, it break isolation of Particles object
particles_xa_1 = segments_xa_1[place].get_matrix()
particles_xa_2 = segments_xa_2[place].get_matrix()
ax = (particles_xa_1 - particles_xa_2) / (xa1 - xa2) * (wanted_xa -
xa2)
b = particles_xa_2
particles_with_xa = ax + b
new_places[place] = Particles(particles_with_xa,
segments_xa_1[place].get_mapping())
return new_places
def test_get_number_of_particles(self):
k = np.array([1, 2, 3]).reshape((-1, 1))
l = np.array([2, 3, 4]).reshape((-1, 1))
array = np.append(k, l, axis=1)
particles = Particles(array, {"x": 0, "y": 1})
self.assertTrue(particles.get_number_of_particles() == 3)
def test_constructor(self):
k = np.array([1, 2, 3]).reshape((-1, 1))
l = np.array([2, 3, 4]).reshape((-1, 1))
array = np.append(k, l, axis=1)
particles = Particles(array, {"x": 0, "y": 1})
Particles(particles.particles, particles.mapping)