Some questions about mplot3D

Pau <vim.unix@...83...> writes:

I am trying to generate a 3d-plot
I have two functions that depend on two free parameters,
T_g = (5./512.) * Light_c**5 * a**4 / (Grav_G**3 * m**3)
T_d = 3.e4 * sqrt(a**3/ (Grav_G * m**2.))
These are given in units of time, so that I would like axis y to be
"time", running between 1.0 and 1.e7

The free parameters are "a" and "m". I would like to have the axis

x = a (semi-major axis) , distributed as a = np.arange(1.e10, 1.e15, 5.e11)
z = m (mass) , between 1.e3 and 1.e7

Light_c and Grav_G are two constants

I am trying to follow the examples in
http://matplotlib.sourceforge.net/mpl_toolkits/mplot3d/tutorial.html
but I cannot understand the syntax

Below is a first attempt. What may be tricky is constructing 2-D arrays from
the vectors of a- and m-values. Matlab has a convenient function for this:
http://www.mathworks.com/access/helpdesk/help/techdoc/ref/ndgrid.html

The Numpy way is for instance:

n = 11
a = np.linspace(1e10, 1e15, n)
m = np.linspace(1e3, 1e7, n)
# Now a.shape == m.shape == (11,). Make them (11, 1) and (1, 11):
a, m = np.ix_(a, m)
# But mplot3d expects full arrays of shape (11, 11).
a, m = np.broadcast_arrays(a, m)

I had a Python implementation of ndgrid() lying around and attach it below. It
allows you to do:
a, m = ndgrid(np.linspace(1e10, 1e15, n), np.linspace(1e3, 1e7, n))

It would be nice if plot() could just broadcast arrays as necessary. There are
other rough edges, such as the lack of a zlabel() function (use ax.set_zlabel
()), and no "zscale" property of Axes3D so that setp(ax, "zscale", "log") does
not work.

Hope this helps,
Jon Olav

== Example using plot_wireframe ==

from __future__ import division # http://www.python.org/dev/peps/pep-0238/
import numpy as np
import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import Axes3D

# Oh, if my old science teacher could see me now
Light_c = 3e8 # m / s
Grav_G = 9.81 # m / s**2

n = 11
a = np.linspace(1e10, 1e15, n) # np.arange(1e10, 1e15, 5e11) gives 2000 steps
m = np.linspace(1e3, 1e7, n)
a, m = np.ix_(a, m)
a, m = np.broadcast_arrays(a, m)

T_g = (5/512) * Light_c**5 * a**4 / (Grav_G**3 * m**3)
T_d = 3e4 * np.sqrt(a**3/ (Grav_G * m**2))

fig = plt.figure()
ax = Axes3D(fig)
ax.plot_wireframe(a, m, np.log(T_g))
xlabel("Semi-major axis")
ylabel("Mass")
# There is no zlabel(), but:
ax.set_zlabel("Time, T_g (log-transformed)")
plt.show()

== ndgrid.py ==

"""
n-dimensional gridding like Matlab's NDGRID

Typical usage:

x, y, z = [0, 1], [2, 3, 4], [5, 6, 7, 8]
X, Y, Z = ndgrid(x, y, z)

See ?ndgrid for details.
"""
import numpy as np

def ndgrid(*args, **kwargs):
    """
    n-dimensional gridding like Matlab's NDGRID
    
    The input *args are an arbitrary number of numerical sequences,
    e.g. lists, arrays, or tuples.
    The i-th dimension of the i-th output argument
    has copies of the i-th input argument.
    
    Optional keyword argument:
    same_dtype : If False (default), the result is an ndarray.
                 If True, the result is a lists of ndarrays, possibly with
                 different dtype. This can save space if some *args
                 have a smaller dtype than others.

    Typical usage:
    >>> x, y, z = [0, 1], [2, 3, 4], [5, 6, 7, 8]
    >>> X, Y, Z = ndgrid(x, y, z) # unpacking the returned ndarray into X, Y, Z

    Each of X, Y, Z has shape [len(v) for v in x, y, z].
    >>> X.shape == Y.shape == Z.shape == (2, 3, 4)
    True
    >>> X
    array([[[0, 0, 0, 0],
            [0, 0, 0, 0],
            [0, 0, 0, 0]],
           [[1, 1, 1, 1],
            [1, 1, 1, 1],
            [1, 1, 1, 1]]])
    >>> Y
    array([[[2, 2, 2, 2],
            [3, 3, 3, 3],
            [4, 4, 4, 4]],
           [[2, 2, 2, 2],
            [3, 3, 3, 3],
            [4, 4, 4, 4]]])
    >>> Z
    array([[[5, 6, 7, 8],
            [5, 6, 7, 8],
            [5, 6, 7, 8]],
           [[5, 6, 7, 8],
            [5, 6, 7, 8],
            [5, 6, 7, 8]]])
    
    With an unpacked argument list:
    >>> V = [[0, 1], [2, 3, 4]]
    >>> ndgrid(*V) # an array of two arrays with shape (2, 3)
    array([[[0, 0, 0],
            [1, 1, 1]],
           [[2, 3, 4],
            [2, 3, 4]]])
    
    For input vectors of different data types, same_dtype=False makes ndgrid()
    return a list of arrays with the respective dtype.
    >>> ndgrid([0, 1], [1.0, 1.1, 1.2], same_dtype=False)
    [array([[0, 0, 0], [1, 1, 1]]),
     array([[ 1. , 1.1, 1.2], [ 1. , 1.1, 1.2]])]
    
    Default is to return a single array.
    >>> ndgrid([0, 1], [1.0, 1.1, 1.2])
    array([[[ 0. , 0. , 0. ], [ 1. , 1. , 1. ]],
           [[ 1. , 1.1, 1.2], [ 1. , 1.1, 1.2]]])
    """
    same_dtype = kwargs.get("same_dtype", True)
    V = [np.array(v) for v in args] # ensure all input vectors are arrays
    shape = [len(v) for v in args] # common shape of the outputs
    result = []
    for i, v in enumerate(V):
        # reshape v so it can broadcast to the common shape
        # http://docs.scipy.org/doc/numpy/user/basics.broadcasting.html
        zero = np.zeros(shape, dtype=v.dtype)
        thisshape = np.ones_like(shape)
        thisshape[i] = shape[i]
        result.append(zero + v.reshape(thisshape))
    if same_dtype:
        return np.array(result) # converts to a common dtype
    else:
        return result # keeps separate dtype for each output

if __name__ == "__main__":
    import doctest
    doctest.testmod(optionflags=doctest.NORMALIZE_WHITESPACE)