def test_translate(self):
""" Make sure all lines in chain can be translated correctly.
"""
l1 = ElementaryLine([0.0, 1.2, 0.6])
l2 = ElementaryLine([0.0, 1.0, 0.8])
chain = EPChain([l1, l2])
chain.translate(1.0, "x").translate(1.0, "y")
self.assertTupleEqual(l1.eigen_points.A, (1.0, 1.0))
self.assertTupleEqual(l2.eigen_points.A, (4.0, 1.6000000000000001))
def test_append(self):
""" Make sure the elementary line can be appended properly.
"""
l1 = ElementaryLine([0.0, 1.2, 0.7])
l2 = ElementaryLine([0.0, 1.0, 0.5])
l = EPChain([l1])
l.append(l2)
# Check the l2 should be translated.
self.assertTupleEqual(l2.eigen_points.A, (3.0, 0.69999999999999996))
# Appending a repeated line will raise an exception.
self.assertRaises(ValueError, l.append, l2)
def test_iterable(self):
""" Make sure the chain is iterable.
"""
l1 = ElementaryLine([0.0, 1.2, 0.7])
l2 = ElementaryLine([0.0, 1.0, 0.5])
chain = EPChain([l1, l2])
for l in chain:
self.assertTrue(isinstance(l, ElementaryLine))
def test_membership(self):
""" Test the `in` operator.
"""
l1 = ElementaryLine([0.0, 1.2, 0.7])
l2 = ElementaryLine([0.0, 1.0, 0.5])
chain = EPChain([l1])
self.assertTrue(l1 in chain)
self.assertFalse(l2 in chain)
def test_construction_and_query(self):
""" Test we can construct EPChain object correctly.
"""
l1 = ElementaryLine([0.0, 1.2, 0.7])
l2 = ElementaryLine([0.0, 1.0, 0.5])
l = EPChain([l1, l2])
# Now the l2 should be translated.
self.assertTupleEqual(l2.eigen_points.A, (3.0, 0.69999999999999996))
def test_contains(self):
canvas = EPCanvas()
l1 = ElementaryLine([0.0, 1.2, 0.6])
l2 = ElementaryLine([0.0, 1.0, 0.8])
chain = EPChain([l1])
canvas.add_chain(chain)
self.assertTrue(l1 in canvas)
self.assertTrue(chain in canvas)
self.assertFalse(l2 in canvas)
plt.close(canvas.figure)
def test_x_y_data(self):
""" Test we can get correct data points.
"""
l1 = ElementaryLine([0.0, 1.2, 0.5], n=2)
l2 = ElementaryLine([0.0, 0.8], n=2)
chain = EPChain([l1, l2])
ref_x = [0.0, 1.0, 1.0, 2.0, 2.0, 3.0, 3.0, 4.0, 4.0, 5.0, 5.0, 6.0]
ret_x = chain.x.tolist()
self.assertListEqual(ref_x, ret_x)
ref_y = [
0.0, 0.0, 1.7213277274276194e-18, 0.5, 0.5, 0.5, 0.5, 0.5, 0.5,
1.3, 1.3, 1.3
]
ret_y = chain.y.tolist()
self.assertListEqual(ref_y, ret_y)
def test_add_chain(self):
""" Test energy profile chain can be added correctly to canvas.
"""
canvas = EPCanvas()
self.assertFalse(canvas.lines)
self.assertFalse(canvas.chains)
l1 = ElementaryLine([0.0, 1.2, 0.6])
l2 = ElementaryLine([0.0, 1.0, 0.8])
chain = EPChain([l1, l2])
canvas.add_chain(chain)
self.assertEqual(len(canvas.lines), 2)
for l in canvas.lines:
self.assertTrue(isinstance(l, ElementaryLine))
self.assertEqual(len(canvas.chains), 1)
self.assertTrue(isinstance(canvas.chains[0], EPChain))
# Exception is expected if add the chain again.
self.assertRaises(ValueError, canvas.add_chain, chain)
plt.close(canvas.figure)
locs["peak_widths"], locs["colors"], locs["initial_xs"],
locs["initial_ys"])):
lines = []
for hline_length, peak_width, energies in zip(hline_lengths, peak_widths,
energy_tuples):
# Create an elementary line.
line = ElementaryLine(energies,
hline_length=hline_length,
peak_width=peak_width,
line_width=line_width,
color=color,
shadow_depth=locs["shadow_depth"],
shadow_color=locs["shadow_color"])
lines.append(line)
chain = EPChain(lines)
chain.translate(initial_x, "x")
chain.translate(initial_y, "y")
chains.append(chain)
# Export data.
filename = "data_{}.csv".format(idx)
chain.export(filename)
print("INFO: write data for chain_{} to {}".format(idx, filename))
# Create energy profile canvas.
figsize = locs.get("figsize", (30, 20))
dpi = locs.get("dpi", 50)
canvas = EPCanvas(figsize=figsize, dpi=dpi)