def showChannels(vis, options={}, freq=False, subtitle='', useImspect=False):
"""
Plot visibility data with a matplotlib window.
The matplotlib window has better mouse controls and helps
with selecting channel numbers.
Data is from Miriad's uvspec with stokes=i, axes=chan,amp, interval=9999
and options=avall,nobase
set `freq` to true to have the x-axis be frequency instead of channels
"""
defaults = {
'stokes': 'i',
'options': 'avall,nobase',
'axis': 'freq,amp' if freq else 'chan,amp',
'interval': 9999
}
options = {**defaults, **options}
chans, amps = dumpSpec(vis, options)
fig = plawt.plot({
0: {
'x': chans,
'y': amps,
'draw': 'steps-mid',
'line': 'k'
},
'xlabel': 'Frequency' if freq else 'Channel',
'ylabel': 'Amplitude',
'title': '{}: {}'.format(vis, subtitle),
'keepOpen': True
})
print('Right-click to print channel number:')
def onclick(e):
if e.button == 3: print(round(e.xdata))
cid = fig.canvas.mpl_connect('button_press_event', onclick)
plt.show()
fig.canvas.mpl_disconnect(cid)
'line': 'k-',
'label': 'Lyman Series $\log{N_1 / N}$'
},
1: {
'x': T,
'y': excitationEq(2, T),
'line': 'bo-',
'lw': 1,
'label': 'Balmer Series $\log{N_2 / N}$',
'ms': 1,
'mfc': 'k'
},
'xlabel': 'Temperature (K)',
'ylabel': 'Relative Population $N_n / N$',
'title':
'log of Relative Populations of Hydrogen atoms in $n=1$ and $n=2$',
'set_yscale': 'log',
# 'set_xscale': 'log',
'grid': True,
'legend': {
'loc': 4
},
'ylim': (10e-18, 10e2),
'xlim': (2000, 13e3),
'filename': 'test.png',
# 'figsize': (10,5),
'show': False
}
plawt.plot(myplot)
def compareSpectra(visUncorrected,
visCorrected,
combineI=1,
combineV=10,
plotOptions={},
stokesVylim=None,
stokesIylim=None,
legendloc=1,
filterMaxPeak=False):
"""
Compare spectra utility.
Three panel plot with Stokes I, stokes V uncorrected and Stokes V corrected
combine: the number of velocity channels to average together
"""
miriadDefaults = {
'options': 'avall,nobase',
'axis': 'freq,amp',
'interval': 9999,
}
optionsI = {
**miriadDefaults,
**{
'stokes': 'i',
'line': miriad.averageVelocityLine(visUncorrected, factor=combineI)
}
}
optionsV = {
**miriadDefaults,
**{
'stokes': 'v',
'line': miriad.averageVelocityLine(visCorrected, factor=combineV)
}
}
freq1, amps1 = miriad.dumpSpec(visCorrected, optionsI)
freq2, amps2 = miriad.dumpSpec(visUncorrected, optionsV)
freq3, amps3 = miriad.dumpSpec(visCorrected, optionsV)
# hack for a bad channel
if filterMaxPeak:
amps1[amps1.index(max(amps1))] = 0
amps2[amps2.index(max(amps2))] = 0
amps3[amps3.index(max(amps3))] = 0
fontsize = 8
subtitledict = {'verticalalignment': 'center'}
pad = 0.05 # 5%
if stokesVylim is None:
upperlim = max(max(amps2), max(amps3))
lowerlim = min(min(amps2), min(amps3))
upperlim = upperlim + pad * upperlim
stokesVylim = [(lowerlim, upperlim), (lowerlim, upperlim)]
elif stokesVylim is 'separate':
upperlim = max(max(amps2), max(amps3))
lowerlim = min(min(amps2), min(amps3))
upperlim = upperlim + pad * upperlim
stokesVylim = [(lowerlim, max(amps2) + pad * max(amps2)),
(lowerlim, max(amps3) + pad * max(amps3))]
# xlim = (345.65,347.65)
xlim = (min(freq2), max(freq2))
defaults = {
0: {
'x': freq1,
'y': amps1,
'draw': 'steps-mid',
'line': 'k-',
'label': 'Stokes I'
},
'legend': {
'loc': legendloc,
'fontsize': fontsize
},
'title': '',
'xlabel': 'Frequency (GHz)',
'ylabel': 'Visibility Amplitude',
'sharex': True,
'sharey': 'row',
'hspace': 0.0,
'ylim': stokesIylim,
'xlim': xlim,
'minorticks': True,
}
fig = plawt.plot({
**defaults,
**plotOptions
}, {
0: {
'x': freq2,
'y': amps2,
'draw': 'steps-mid',
'line': 'k-',
'label': 'Stokes V before squint correction'
},
'legend': {
'loc': legendloc,
'fontsize': fontsize
},
'ylim': stokesVylim[0],
}, {
0: {
'x': freq3,
'y': amps3,
'draw': 'steps-mid',
'line': 'k-',
'label': 'Stokes V after squint correction'
},
'legend': {
'loc': legendloc,
'fontsize': fontsize
},
'ylim': stokesVylim[1],
})
def compareMapSpectra(uncorrectedMap,
correctedMap,
line,
stokes,
source,
peakStokes='v',
regionWidth=1,
legendloc=1,
useFull=True,
plotOptions={},
imspectOptions={},
maxfitOptions={},
data=None):
"""
maxfit options=abs on the corrected/uncorrected maps for each line
reinvert the visibilities without the 'mfs' to preserve the velocity axis
use velsw to switch from velocity to frequency:
`velsw in=MAPSCorrect/NGC7538S-s4.co3-2.v.cm axis=FREQ`
imspect with region='abspixel,box(x1,y1,x1,y1)' where x1,y1 is the peak found by maxfit
I can also not use mfs at all with no line selection and get the peak with maxfit on that map.
I then use imspec at the new peak and this gives me a better spectra.
The map however has Stokes I and Stokes V peaks that don't align. The peaks are bigger however.
---
uncorrectedMap: string of the path
correctedMap: string of the path
line: string of line (ex 'co3-2') to find the peak of
stokes: string, one of 'i', 'q', 'u', 'v', 'rr', 'rl', 'lr', 'll'
source: string
regionWidth: the width of the box to put around the peak when creating spectra
peakStokes: string of one of the polarizations to take the spectra through the peak of.
ex. 'i' means take the spectra through the peak of Stokes I
"""
# find the peak of Stokes I or V in the corrected line map
# if source == 'orkl_080106' and line != 'sio8-7':
# peakStokes = 'v'
if data is None:
frequencies, amplitudes = getMapData(uncorrectedMap, correctedMap,
line, stokes, peakStokes,
regionWidth, useFull,
imspectOptions, maxfitOptions)
else:
frequencies, amplitudes = data[2], data[3]
if source == 'NGC7538S-s4':
print("Zeroing bad window in NGC7538")
channels = list(
filter(lambda x: x > 346.7 and x < 346.75, frequencies[2]))
for channel in channels:
idx = frequencies[2].index(channel)
amplitudes[2][idx] = 0
amplitudes[3][idx] = 0
legendFontSize = 20
tickParams = [{
'axis': 'both',
'which': 'both',
'direction': 'in',
'right': True,
'top': True,
'labelsize': 16
}, {
'which': 'major',
'length': 7
}]
axisLabelSize = 20
plotDefaults = {
0: {
'x': frequencies[2],
'y': amplitudes[2],
'draw': 'steps-mid',
'line': 'r-',
'label': 'Corrected Stokes I'
},
'nrows': 2,
'ncols': 1,
'legend': {
'loc': legendloc,
'fontsize': legendFontSize
},
'xlabel': 'Frequency (GHz)',
'ylabel': 'Average Intensity (Jy/beam)',
'xlabelsize': axisLabelSize,
'ylabelsize': axisLabelSize,
'sharex': True,
'sharey': 'row',
'minorticks': True,
'tight_layout': {
'pad': 3,
'h_pad': 0.1,
'w_pad': 0
},
'tick_params': tickParams,
'titlesize': 20,
'hspace': 0.05,
'top': 0.93,
}
plotOptions['title'] = plotOptions[
'title'] + ': Peak Stokes {} through {} line'.format(
peakStokes.upper(), line.upper())
plotOptions['subloc'] = 'left'
fig = plawt.plot({
**plotDefaults,
**plotOptions
}, {
0: {
'x': frequencies[3],
'y': amplitudes[3],
'draw': 'steps-mid',
'line': 'b-',
'label': 'Corrected Stokes V'
},
'legend': {
'loc': legendloc,
'fontsize': legendFontSize
},
'tick_params': tickParams
})
def plotAllSources(uncorrectedMaps,
correctedMaps,
sources,
peakStokes='v',
regionWidths=[],
plotOptions={},
imspectOptions=[]):
"""
Plot all spectra for all sources in one figure
locked to 'co3-2' line and peak stokes v
"""
# init defaults
if len(regionWidths) != len(sources):
regionWidths = [1 for s in sources]
if len(imspectOptions) != len(sources):
imspectOptions = [{} for s in sources]
sourceTitles = ['Orion KL', 'NGC7538', 'IRC+10216', 'IRAS2a']
# gather data for all sources
tickParams = {
'axis': 'both',
'which': 'both',
'direction': 'in',
'right': True,
'top': True,
'labelsize': 10
}
plotDefaults = {
'nrows': 4,
'ncols': 2,
'sharex': 'none',
'sharey': 'row',
'xlabel': 'Frequency (GHz)',
'ylabel': 'Intensity (Jy/Beam)',
'minorticks': True,
'tight_layout': {
'pad': 3,
'h_pad': 2.0,
'w_pad': 0
},
'left': 0.11,
'top': 0.92,
'wspace': 0.05,
'xlabelsize': 12,
'ylabelsize': 12,
}
sourcePlots = []
legendlocs = [1, 1, 2, 3]
for i, source in enumerate(sources):
uncMap = uncorrectedMaps[i]
cMap = correctedMaps[i]
useFull = not (
source == 'NGC7538S-s4'
) # if NGC7538, use the integrated map. The peak in the full map looks like squint and not real.
print('useFull:', useFull)
print(source)
freqs, amps = getMapData(uncMap,
cMap,
'co3-2', ['v'],
peakStokes,
regionWidths[i],
useFull=useFull,
imspectOptions=imspectOptions[i])
newPanels = []
for j, freq in enumerate(freqs):
newPanel = plotDefaults if i == 0 and j == 0 else {}
if i == 0:
if j == 0:
newPanel['subtitle'] = 'Uncorrected Spectra'
elif j == 1:
newPanel['subtitle'] = 'Corrected Spectra'
linestyle = 'r-' if j is 0 or j is 2 else 'b-'
linestyle = 'b-'
if source == 'NGC7538S-s4':
print("Zeroing bad window in NGC7538")
channels = list(
filter(lambda x: x > 346.7 and x < 346.75, freqs[j]))
for channel in channels:
idx = freqs[j].index(channel)
amps[j][idx] = 0
idx = freqs[j].index(channels[0]) - 500
amps[j] = amps[j][:idx]
freqs[j] = freqs[j][:idx]
newPanel[0] = {
'x': freqs[j],
'y': amps[j],
'label': sourceTitles[i],
'draw': 'steps-mid',
'line': linestyle
}
newPanel = {
**newPanel,
**{
'legend': {
'loc': legendlocs[i],
'fontsize': 10
},
'tick_params': tickParams,
}
}
newPanels.append(newPanel)
sourcePlots.append(newPanels[0])
sourcePlots.append(newPanels[1])
sourcePlots = [{**sourcePlots[0], **plotOptions}] + sourcePlots[1:]
fig = plawt.plot(*sourcePlots)
#!/usr/bin/python3
import plawt
chans = [3.0, 4.0, 5.0, 6.0, 7.0, 8.0, 9.0, 10.0, 11.0, 12.0, 13.0, 14.0, 15.0, 16.0, 17.0, 18.0, 19.0, 20.0, 21.0, 22.0, 23.0, 24.0, 25.0, 26.0, 27.0, 28.0, 29.0, 30.0, 31.0, 32.0, 33.0, 34.0, 35.0, 36.0, 37.0, 38.0, 39.0, 40.0, 41.0, 42.0, 43.0, 44.0, 45.0, 46.0, 47.0, 48.0, 49.0, 50.0, 51.0, 52.0, 53.0, 54.0, 55.0, 56.0, 57.0, 58.0, 59.0, 60.0, 61.0, 62.0, 63.0, 64.0, 65.0, 66.0, 67.0, 68.0, 69.0, 70.0, 71.0, 72.0, 73.0, 74.0, 75.0, 76.0, 77.0, 78.0, 79.0, 80.0, 81.0, 82.0, 83.0, 84.0, 85.0, 86.0, 87.0, 88.0, 89.0, 90.0, 91.0, 92.0, 93.0, 94.0, 95.0, 96.0, 97.0, 98.0, 99.0, 100.0, 101.0, 102.0, 103.0, 104.0, 105.0, 106.0, 107.0, 108.0, 109.0, 110.0, 111.0, 112.0, 113.0, 114.0, 115.0, 116.0, 117.0, 118.0, 119.0, 120.0, 121.0, 122.0, 123.0, 124.0, 125.0, 126.0, 127.0, 128.0, 129.0, 130.0, 131.0, 132.0, 133.0, 134.0, 135.0, 136.0, 137.0, 138.0, 139.0, 140.0, 141.0, 142.0, 143.0, 144.0, 145.0, 146.0, 147.0, 148.0, 149.0, 150.0, 151.0, 152.0, 153.0, 154.0, 155.0, 156.0, 157.0, 158.0, 159.0, 160.0, 161.0, 162.0, 163.0, 164.0, 165.0, 166.0, 167.0, 168.0, 169.0, 170.0, 171.0, 172.0, 173.0, 174.0, 175.0, 176.0, 177.0, 178.0, 179.0, 180.0, 181.0, 182.0, 183.0, 184.0, 185.0, 186.0, 187.0, 188.0, 189.0, 190.0, 191.0, 192.0, 193.0, 194.0, 195.0, 196.0, 197.0, 198.0, 199.0, 200.0, 201.0, 202.0, 203.0, 204.0, 205.0, 206.0, 207.0, 208.0, 209.0, 210.0, 211.0, 212.0, 213.0, 214.0, 215.0, 216.0, 217.0, 218.0, 219.0, 220.0, 221.0, 222.0, 223.0, 224.0, 225.0, 226.0, 227.0, 228.0, 229.0, 230.0, 231.0, 232.0, 233.0, 234.0, 235.0, 236.0, 237.0, 238.0, 239.0, 240.0, 241.0, 242.0, 243.0, 244.0, 245.0, 246.0, 247.0, 248.0, 249.0, 250.0, 251.0, 252.0, 253.0, 254.0, 255.0, 256.0, 257.0, 258.0, 259.0, 260.0, 261.0, 262.0, 263.0, 264.0, 265.0, 266.0, 267.0, 268.0, 269.0, 270.0, 271.0, 272.0, 273.0, 274.0, 275.0, 276.0, 277.0, 278.0, 279.0, 280.0, 281.0, 282.0, 283.0, 284.0, 285.0, 286.0, 287.0, 288.0, 289.0, 290.0, 291.0, 292.0, 293.0, 294.0, 295.0, 296.0, 297.0, 298.0, 299.0, 300.0, 301.0, 302.0, 303.0, 304.0, 305.0, 306.0, 307.0, 308.0, 309.0, 310.0, 311.0, 312.0, 313.0, 314.0, 315.0, 316.0, 317.0, 318.0, 319.0, 320.0, 321.0, 322.0, 323.0, 324.0, 325.0, 326.0, 327.0, 328.0, 329.0, 330.0, 331.0, 332.0, 333.0, 334.0, 335.0, 336.0, 337.0, 338.0, 339.0, 340.0, 341.0, 342.0, 343.0, 344.0]
amps = [0.6030806, 0.6837087, 0.6371896, 0.5690394, 0.6679561, 0.6748063, 0.6571434, 0.7004986, 0.5838464, 0.6492448, 0.6545905, 0.7203998, 0.6872855, 0.7237729, 0.732918, 0.6665491, 0.6614937, 0.6439695, 0.6943766, 0.6929428, 0.7050078, 0.7587503, 0.6372474, 0.6755744, 0.597022, 0.6071005, 0.6374301, 0.6681313, 0.6941413, 0.6307471, 0.7715253, 0.6422789, 0.6534501, 0.6834064, 0.6233898, 0.6765852, 0.5996928, 0.6059702, 0.532366, 1.530873, 1.995607, 1.151967, 0.7835325, 0.7550735, 0.7925035, 0.7849458, 0.7888911, 0.7411907, 0.7353308, 0.6551155, 0.6743459, 0.6655676, 0.7094763, 0.6414862, 0.6859187, 0.6862016, 0.6435194, 0.6825865, 0.662622, 0.6936107, 0.681871, 0.6526321, 0.6545926, 0.7351808, 0.7302195, 0.6981022, 0.6364391, 0.6797721, 0.7573355, 0.6753334, 0.7105518, 0.6716144, 0.6950731, 0.6657793, 0.6729059, 0.7100132, 0.6249868, 0.6844749, 0.6908333, 0.7762515, 0.662553, 0.6727911, 0.7305974, 0.7004743, 0.7272439, 0.6263704, 0.6362215, 0.6694011, 0.6318359, 0.6848253, 0.6318443, 0.6213173, 0.6706879, 0.6732186, 0.6374716, 0.6313396, 0.9778146, 2.682548, 1.359655, 0.7630816, 0.662514, 0.6784354, 0.703732, 0.7576462, 0.7486892, 0.7133566, 0.7465675, 0.7390884, 0.7255204, 0.7090992, 0.6545323, 0.6366355, 0.7514731, 0.6789933, 0.7204047, 0.6966297, 0.6550111, 0.6339555, 0.7074148, 0.6423266, 0.6131723, 0.6607807, 0.6422971, 0.6608815, 0.6662244, 0.6128898, 0.6245586, 0.6465213, 0.6261068, 0.6248351, 0.6371271, 0.5853242, 0.5785147, 0.6686863, 0.6457974, 0.781877, 0.6621993, 0.6905352, 0.6711752, 0.641484, 0.6282443, 0.5987218, 0.6499262, 0.6514604, 0.6543415, 0.7384826, 0.7560073, 0.7332252, 0.6890889, 0.6860042, 0.5730995, 0.6801083, 0.7284067, 0.6903766, 0.6536254, 0.8050244, 0.8658758, 2.593296, 1.442149, 0.748728, 0.6702675, 0.6008819, 0.6644086, 0.7028915, 0.660691, 0.6759114, 0.7413937, 0.7441263, 0.6403871, 0.6987057, 0.6755445, 0.5951479, 0.6800404, 0.6918594, 0.6380497, 0.5243984, 0.5040802, 0.7274017, 5.604488, 8.731837, 3.425878, 1.002521, 0.730807, 0.6965266, 0.6574699, 0.6511405, 0.7235293, 0.7651855, 0.7430843, 0.6180075, 0.6781835, 0.7199444, 0.6896753, 0.6528503, 0.6537533, 0.7358531, 0.6939768, 0.6576437, 0.6582898, 0.5817757, 0.7327138, 0.6167199, 0.6683732, 0.671705, 0.6849148, 0.6395864, 0.6740314, 0.639993, 0.5927391, 0.670041, 0.5839626, 0.6350645, 0.6338769, 0.6762487, 0.6510934, 0.6616208, 0.6471664, 0.6204606, 0.6427088, 0.6055776, 0.5554015, 0.6364431, 0.6573479, 0.6997938, 0.6386005, 0.6445189, 0.8192143, 0.6585949, 0.630053, 0.6424926, 0.6388685, 0.7741821, 1.046616, 0.7158167, 0.9138164, 1.041885, 0.8185232, 0.6814007, 0.6426553, 0.6461026, 0.6409208, 0.6984751, 0.7041847, 0.671133, 0.6397859, 0.6290236, 0.6276253, 0.668026, 0.6459958, 0.681568, 0.6703814, 0.6646535, 0.6578941, 0.6216274, 0.6333016, 0.6127044, 0.6613026, 0.6662278, 0.6304022, 0.6700389, 0.6811751, 0.6911133, 0.6100067, 0.6399619, 0.6294864, 0.660863, 0.7047668, 0.7167548, 0.7246642, 0.6770457, 0.6603574, 0.7061953, 0.7244675, 0.6993666, 0.666106, 0.74291, 0.6401825, 0.6687987, 0.6503596, 0.6375297, 0.6408038, 0.7029941, 0.7140093, 0.8503788, 0.8528588, 0.7600691, 0.8529701, 0.9322207, 0.6891239, 0.5604228, 0.6766065, 0.7689702, 0.655958, 0.6705654, 0.651525, 0.6022623, 0.5689662, 0.5360836, 0.433879, 0.2912524, 0.3277543, 0.5663058, 1.633974, 4.41852, 7.212595, 4.049658, 0.7425946, 9.426757, 2.546259, 0.7758642, 0.9863672, 0.8462248, 0.9170761, 0.9928068, 1.02269, 0.9133889, 0.8045723, 0.7689427, 0.6578668, 0.6855389, 0.6688508, 0.6766053, 0.6102203, 0.6144144, 0.5309623, 0.7105601, 0.6795527, 0.6335855, 0.6228023, 0.6700085, 0.7074459, 0.7548816, 0.6546933, 0.6945201, 0.5961808, 0.6433087, 0.6827869, 0.8051797, 1.084111, 0.759516, 0.7164305, 0.6180438]
plawt.plot({
0: {'x': chans[:75], 'y': amps[:75], 'draw': 'steps-mid', 'line': 'k-'},
'show': True,
'filename': 'bartest.png'
})
plawt.plot(
{
0: {
'x': x,
'y': y
},
'title': 'Sharing both axes',
'subtitle': 'panel a',
'subtitledict': {
'verticalalignment': 'center'
},
'fontsize': 12,
'subloc': 'left',
'minorticks': True,
'xlabel': 'Velocity',
'ylabel': 'Amplitude',
'sharex': True,
'sharey': True,
'hspace': 0.3,
# 'aspect': 16/9,
'filename': 'subplottest.png'
},
{
0: {
'x': x,
'y': y,
'line': 'bo'
},
'subtitle': 'panel b',
'subtitledict': {
'verticalalignment': 'center'
},
'fontsize': 12,
'subloc': 'left',
},
{
0: {
'x': x,
'y': 2 * y**2 - 1,
'line': 'ro',
'label': 'panel c'
},
'subtitle': 'panel c',
'subtitledict': {
'verticalalignment': 'center'
},
'fontsize': 12,
'subloc': 'left',
'legend': {
'loc': 2,
'fontsize': 10
},
})
def threebodyplots():
simdir = 'simulations'
figdir = 'figures'
for folder in [simdir, figdir]:
if not os.path.exists(folder):
os.mkdir(folder)
for filename in glob.glob(os.path.join(simdir, '*.in')):
title = os.path.splitext(os.path.basename(filename))[0]
simulation = parseNBodyData(os.path.join(simdir, title + '.out'))
numPlanets = simulation['n']
time = simulation['time']
planets = simulation[
'planets'] # # The indexing is: planets[planetid, posOrVel, time, component]
tend = -1 # -1 takes the whole slice
tstep = 20
plawt.plot({
'title': title,
# Paths
0: {
'x': planets[0, 0, :tend, 0],
'y': planets[0, 0, :tend, 1],
'lw': 0.8,
'alpha': 0.5,
'line': 'k-'
},
1: {
'x': planets[1, 0, :tend, 0],
'y': planets[1, 0, :tend, 1],
'lw': 0.8,
'alpha': 0.5,
'line': 'b-'
},
2: {
'x': planets[2, 0, :tend, 0],
'y': planets[2, 0, :tend, 1],
'lw': 0.8,
'alpha': 0.5,
'line': 'r-'
},
# Starting Positions
3: {
'x': planets[0, 0, 0, 0],
'y': planets[0, 0, 0, 1],
'line': 'ko',
'ms': 10
},
4: {
'x': planets[1, 0, 0, 0],
'y': planets[1, 0, 0, 1],
'line': 'bo',
'ms': 10
},
5: {
'x': planets[2, 0, 0, 0],
'y': planets[2, 0, 0, 1],
'line': 'ro',
'ms': 10
},
# Final Positions
6: {
'x': planets[0, 0, -1, 0],
'y': planets[0, 0, -1, 1],
'line': 'k*',
'ms': 14,
'mfc': 'none'
},
7: {
'x': planets[1, 0, -1, 0],
'y': planets[1, 0, -1, 1],
'line': 'b*',
'ms': 14,
'mfc': 'none'
},
8: {
'x': planets[2, 0, -1, 0],
'y': planets[2, 0, -1, 1],
'line': 'r*',
'ms': 14,
'mfc': 'none'
},
# Time Sequence
9: {
'x': planets[0, 0, :tend, 0][::tstep],
'y': planets[0, 0, :tend, 1][::tstep],
'lw': 0,
'line': 'ko',
'ms': 2
},
10: {
'x': planets[1, 0, :tend, 0][::tstep],
'y': planets[1, 0, :tend, 1][::tstep],
'lw': 0,
'line': 'bo',
'ms': 2
},
11: {
'x': planets[2, 0, :tend, 0][::tstep],
'y': planets[2, 0, :tend, 1][::tstep],
'lw': 0,
'line': 'ro',
'ms': 2
},
'xlabel': 'x',
'ylabel': 'y',
'filename': os.path.join(figdir, title + '.png'),
'grid': True,
'show': False
})
def specialplot():
specialdir = 'special'
simulation = parseNBodyData(os.path.join(specialdir, 'star.out'))
numPlanets = simulation['n']
time = simulation['time']
planets = simulation['planets']
tend = -1 # -1 takes the whole slice
tstep = 100
plawt.plot({
'title': 'star',
# Paths
0: {
'x': planets[0, 0, :tend, 0],
'y': planets[0, 0, :tend, 1],
'lw': 0.8,
'alpha': 0.5,
'line': 'k-'
},
1: {
'x': planets[1, 0, :tend, 0],
'y': planets[1, 0, :tend, 1],
'lw': 0.8,
'alpha': 0.5,
'line': 'b-'
},
2: {
'x': planets[2, 0, :tend, 0],
'y': planets[2, 0, :tend, 1],
'lw': 0.8,
'alpha': 0.5,
'line': 'r-'
},
3: {
'x': planets[3, 0, :tend, 0],
'y': planets[3, 0, :tend, 1],
'lw': 0.8,
'alpha': 0.5,
'line': 'g-'
},
4: {
'x': planets[4, 0, :tend, 0],
'y': planets[4, 0, :tend, 1],
'lw': 0.8,
'alpha': 0.5,
'line': 'c-'
},
# Start
5: {
'x': planets[0, 0, 0, 0],
'y': planets[0, 0, 0, 1],
'line': 'ko',
'ms': 10
},
6: {
'x': planets[1, 0, 0, 0],
'y': planets[1, 0, 0, 1],
'line': 'bo',
'ms': 10
},
7: {
'x': planets[2, 0, 0, 0],
'y': planets[2, 0, 0, 1],
'line': 'ro',
'ms': 10
},
8: {
'x': planets[3, 0, 0, 0],
'y': planets[3, 0, 0, 1],
'line': 'go',
'ms': 10
},
9: {
'x': planets[4, 0, 0, 0],
'y': planets[4, 0, 0, 1],
'line': 'co',
'ms': 10
},
# Time
10: {
'x': planets[0, 0, :tend, 0][::tstep],
'y': planets[0, 0, :tend, 1][::tstep],
'lw': 0,
'line': 'ko',
'ms': 2
},
11: {
'x': planets[1, 0, :tend, 0][::tstep],
'y': planets[1, 0, :tend, 1][::tstep],
'lw': 0,
'line': 'bo',
'ms': 2
},
12: {
'x': planets[2, 0, :tend, 0][::tstep],
'y': planets[2, 0, :tend, 1][::tstep],
'lw': 0,
'line': 'ro',
'ms': 2
},
13: {
'x': planets[3, 0, :tend, 0][::tstep],
'y': planets[3, 0, :tend, 1][::tstep],
'lw': 0,
'line': 'go',
'ms': 2
},
14: {
'x': planets[4, 0, :tend, 0][::tstep],
'y': planets[4, 0, :tend, 1][::tstep],
'lw': 0,
'line': 'co',
'ms': 2
},
# Final
15: {
'x': planets[0, 0, -1, 0],
'y': planets[0, 0, -1, 1],
'line': 'k*',
'ms': 14,
'mfc': 'none'
},
16: {
'x': planets[1, 0, -1, 0],
'y': planets[1, 0, -1, 1],
'line': 'b*',
'ms': 14,
'mfc': 'none'
},
17: {
'x': planets[2, 0, -1, 0],
'y': planets[2, 0, -1, 1],
'line': 'r*',
'ms': 14,
'mfc': 'none'
},
18: {
'x': planets[3, 0, -1, 0],
'y': planets[3, 0, -1, 1],
'line': 'g*',
'ms': 14,
'mfc': 'none'
},
19: {
'x': planets[4, 0, -1, 0],
'y': planets[4, 0, -1, 1],
'line': 'c*',
'ms': 14,
'mfc': 'none'
},
'grid': True,
'xlabel': 'x',
'ylabel': 'y',
'filename': os.path.join(specialdir, 'star.png'),
'show': False
})
with imageio.get_writer(os.path.join(specialdir, 'special.mp4'),
mode='I',
fps=24) as writer:
for i, t in enumerate(time):
print("frame #" + str(i))
currplot = {
'show': False,
'keepOpen': True,
'ylim': (-1, 1),
'xlim': (-2, 2),
}
for n, _ in enumerate(planets):
currplot[n] = {
'x': planets[n, 0, i, 0],
'y': planets[n, 0, i, 1],
'line': 'ko'
}
plt = plawt.plot(currplot)
fig = plt.gcf()
fig.canvas.draw()
data = fig.canvas.tostring_rgb()
row, col = fig.canvas.get_width_height()
image = np.fromstring(data, dtype=np.uint8).reshape(col, row, 3)
writer.append_data(image)
plt.close()
axarr[0, 1].text(0, -1.0, 'Text here too')
axarr[1, 0].plot(x, y ** 2)
axarr[1, 0].set_title('Axis [1,0]')
axarr[1, 1].plot(x, y ** 2)
axarr[1, 1].set_title('Axis [1,1]')
f.subplots_adjust(hspace=0.3)
plt.savefig('matplotlibsubplotgrid.png')
##########
plawt.plot({
0: {'x': x, 'y': y},
'nrows':2, 'ncols': 2,
'title': 'Plawt: Grid of subplots',
'subtitle': 'Axis [0,0]',
'hspace': 0.3,
'text': {'x':0, 'y':0.8, 's': 'Text here', 'fontsize': 8},
'filename': 'plawtsubplotgrid.png'
}, {
0: {'x': x, 'y': y},
'subtitle': 'Axis [0,1]',
'text': [{'x':0, 'y':0.8, 's': 'Text here'},
{'x':0, 'y':-1.0, 's':'Text here too'}],
}, {
0: {'x': x, 'y': y**2},
'subtitle': 'Axis [1,0]'
}, {
0: {'x': x, 'y': y**2},
'subtitle': 'Axis [1,1]'
})