3D scatter plot

plot.py

#!/usr/bin/env python

import sys, os
import numpy
import matplotlib
import matplotlib.font_manager as Fm
import matplotlib.backends.backend_pdf as PltPDF
from mpl_toolkits.mplot3d import Axes3D

def loadData(filename):
    Data = []
    with open(filename, 'r') as f:
        for Line in f:
            if not Line.strip().startswith('#'):
                Data.append(tuple(float(x) for x in Line.split()))

    return Data

def forceRange(data, ranges):
    NDims = len(data[0])
    NRanges = len(ranges)
    Result = []

    for Sample in data:
        SampleGood = True
        for DimIdx in range(NRanges):
            Range = ranges[DimIdx]
            if Range is None:
                continue

            if len(Range) != 2:
                raise ValueError("Expect 2-tuple, got {}.".format(repr(Range)))

            if not (Range[0] <= Sample[DimIdx] <= Range[1]):
                SampleGood = False
                break

        if SampleGood is True:
            Result.append(Sample)

    return Result

def finalizeData(data):
    return tuple(zip(*data))

def makeTransparent0Map(colormap):
    """Copy and modify a colormap so that value 0 is transparent."""
    ColorMap = matplotlib.cm.get_cmap("PuRd", 12)
    Colors = ColorMap(numpy.linspace(0, 1, 12))
    Colors[0, 3] = 0.0
    return matplotlib.colors.LinearSegmentedColormap.from_list("Alpha00", Colors)

def plot():
    Data = loadData("plot.esp")
    Fig = matplotlib.figure.Figure()
    matplotlib.backends.backend_agg.FigureCanvasAgg(Fig) # ???
    Ax = Fig.add_subplot(111, projection='3d')

    # Data = forceRange(Data, ((-2, 2), (-2, 2), (-2, 2),))
    Data = finalizeData(Data)

    NewColorMap = makeTransparent0Map("PuRd")

    Plot = Ax.scatter(Data[0], Data[1], Data[2], c=Data[3],
                      marker='o', cmap=NewColorMap)
    Fig.colorbar(Plot)
    Fig.savefig("test.pdf")

if __name__ == "__main__":
    plot()