[Tutor] Question in regards to loops and matplotlib

[Oscar Benjamin]

On 1 August 2013 06:21, Zachary Rizer <protonlenny at gmail.com> wrote:
> So I just started coding, and I'm so glad I chose Python to start me off! I
> really enjoy the clean layout, easy syntax, and power of the language. I'm
> doing astronomy research and I just started teaching myself matplotlib along
> with my general python work. I feel like I'm catching on quick, but I also
> feel that this particular plot script is a little rough. The plot looks
> correct, but the code seems really long for what I'm trying to do. So any
> tricks to make this more efficient would be greatly appreciated!

Hi Zachary,

You can definitely make this shorter. I can't run your code since you
didn't provide any data. It's good to use dummy data in examples
posted to mailing lists so that others can run the code. Instead I'll
show an example with my own dummy data:


#!/usr/bin/env python
#
# Plot heights and weights of subjects from different groups.

from matplotlib import pyplot as plt

# Dummy data
subjects = [
    # Group, height (m), weight (kg), diagnosis
    {'type': 'patient', 'height': 1.8, 'weight': 90 , 'diag': 'acute'  },
    {'type': 'patient', 'height': 1.6, 'weight': 85 , 'diag': 'acute'  },
    {'type': 'patient', 'height': 1.9, 'weight': 120, 'diag': 'chronic'},
    {'type': 'patient', 'height': 2.0, 'weight': 110, 'diag': 'chronic'},
    {'type': 'control', 'height': 1.7, 'weight': 60 , 'diag': 'N/A'    },
    {'type': 'control', 'height': 2.1, 'weight': 100, 'diag': 'N/A'    },
]

# Have just one place where we define the properties of each group
groups = {
    'control': {
        'indicator': lambda s: s['type'] == 'control',
        'color': 'blue',
        'marker': '*',
        'label':'Controls'
     },
    'chronic': {
        'indicator': lambda s: s['diag'] == 'chronic',
        'color': 'magenta',
        'marker': 's',
        'label':'Chronic patients'
    },
    'acute'  : {
        'indicator': lambda s: s['diag'] == 'acute',
        'color': 'red',
        'marker': 'o',
        'label':'Acute patients'
    },
}

# Now we can reuse the same code for every subject
for groupdata in groups.values():
    groupdata['subjects'] = []

# Distribute the subjects to each group
for subject in subjects:
    for groupdata in groups.values():
        isingroup = groupdata['indicator']
        if isingroup(subject):
            groupdata['subjects'].append(subject)
            break

# Now create/format the figure
fig = plt.figure(figsize=(6, 5))
ax = fig.add_axes([0.15, 0.15, 0.70, 0.70])
ax.set_xlabel(r'Weight ($kg$)')
ax.set_ylabel(r'Height ($m$)')
ax.set_title('Height vs. weight by subject group')
ax.set_xlim([50, 150])
ax.set_ylim([1.5, 2.2])

# Plot each group separately with its own format settings
for group, groupdata in groups.items():
    xdata = [s['weight'] for s in groupdata['subjects']]
    ydata = [s['height'] for s in groupdata['subjects']]
    ax.plot(xdata, ydata, linestyle='none', color=groupdata['color'],
            marker=groupdata['marker'], label=groupdata['label'])

ax.legend(loc='upper left', numpoints=1)

# Show the figure in GUI by default.
# Save to e.g. 'plot.pdf' if filename is given.
import sys
if len(sys.argv) > 1:
    imgname = sys.argv[1]
    fig.savefig(imgname)
else:
    plt.show()


Oscar