def subplot(potential, orbits, codes, fit_orbit, labels): hor, vert = labels X = numpy.linspace(-160, 160, 100) | units.parsec Y = numpy.linspace(-160, 160, 100) | units.parsec X, Y = quantities.meshgrid(X, Y) if labels == 'xy': pot_args = [X, Y, 0 | units.parsec] fit_horizontal = fit_orbit[0] fit_vertical = fit_orbit[1] elif labels == 'xz': pot_args = [X, 0 | units.parsec, Y] fit_horizontal = fit_orbit[0] fit_vertical = fit_orbit[2] elif labels == 'yz': pot_args = [0 | units.parsec, X, Y] fit_horizontal = fit_orbit[1] fit_vertical = fit_orbit[2] phi = potential.get_potential_at_point(None, *pot_args) aplot.imshow_color_plot(X, Y, phi, cmap="Blues") del aplot.UnitlessArgs.arg_units[2] aplot.scatter(0 | units.parsec, 0 | units.parsec, c='black') aplot.plot(fit_horizontal, fit_vertical, c="green", ls="--", label="Kruijssen et al. 2014") colors = cycle(['red', 'black', 'yellow', 'grey', 'green']) for full_orbit, code in zip(orbits, codes): horizontal = full_orbit.x if hor == 'x' else full_orbit.y vertical = full_orbit.y if vert == 'y' else full_orbit.z color = next(colors) aplot.plot(horizontal, vertical, c=color, label=code) aplot.scatter(horizontal[-1], vertical[-1], c=color, edgecolor=color) pyplot.axis('equal') aplot.xlim([-150, 150] | units.parsec) aplot.ylim([-150, 150] | units.parsec) aplot.xlabel(hor) aplot.ylabel(vert)
def test8(self): """ Test the imshow color plot """ if not HAS_MATPLOTLIB: return self.skip() pyplot.clf() x = numpy.linspace(0, 100, 100) | units.m y = numpy.linspace(0, 200, 100) | units.m X, Y = quantities.meshgrid(x, y) Z = X**2 + Y**2 figure, bar = aplot.imshow_color_plot(X, Y, Z, add_colorbar=True) self.assertEqual("[m]", self.xaxis().get_label_text()) self.assertEqual("[m]", self.yaxis().get_label_text())
def test8(self): """ Test the imshow color plot """ if not HAS_MATPLOTLIB: return self.skip() pyplot.clf() x = numpy.linspace(0, 100, 100) | units.m y = numpy.linspace(0, 200, 100) | units.m X, Y = quantities.meshgrid(x, y) Z = X**2 + Y**2 figure, bar = aplot.imshow_color_plot(X, Y, Z, add_colorbar=True) self.assertEquals("[m]", self.xaxis().get_label_text()) self.assertEquals("[m]", self.yaxis().get_label_text()) self.assertEquals("[m**2]", bar._label)