python: Parametric 2D Surface Class

parasurf.py

import numpy as np
from scipy.interpolate import RectBivariateSpline, bisplev


class ParaSurf(object):
    def __init__(self, u, v, xyz, bbox=[-0.25, 1.25, -0.5, 1.5], ku=3, kv=3):
        """Parametric (u,v) surface approximation over a rectangular mesh.

        Parameters
        ----------
        u, v : array_like
            1-D arrays of coordinates in strictly ascending order.
        xyz : array_like
            3-D array of (x, y, z) data with shape (3, u.size, v.size).
        bbox : array_like, optional
            Sequence of length 4 specifying the boundary of the rectangular
            approximation domain. See scipy.interpolate.RectBivariateSpline
            for more info
        ku, kv : ints, optional
            Degrees of the bivariate spline. Default is 3.
        """
        self._create_srf(u, v, xyz, bbox, ku, kv)

        self.bbox = bbox
        self.u = u
        self.v = v
        self.ku = ku
        self.kv = kv

    def _create_srf(self, u, v, xyz, bbox, ku, kv):
        # Create surface definitions
        xsrf = RectBivariateSpline(u, v, xyz[0], bbox=bbox, kx=ku, ky=kv, s=0)
        ysrf = RectBivariateSpline(u, v, xyz[1], bbox=bbox, kx=ku, ky=kv, s=0)
        zsrf = RectBivariateSpline(u, v, xyz[2], bbox=bbox, kx=ku, ky=kv, s=0)

        # Create the 5-element list suitable for feeding
        # to bisplev() later (we need bisplev() to get partial
        # derivatives, for surface normals and surface tangents
        # along u and v)
        self._xsrfdef = xsrf.tck + xsrf.degrees
        self._ysrfdef = ysrf.tck + ysrf.degrees
        self._zsrfdef = zsrf.tck + zsrf.degrees

        self._xsrf = xsrf
        self._ysrf = ysrf
        self._zsrf = zsrf
        self._origxyzs = xyz

    def ev(self, u, v, mesh=True):
        """Get point(s) on surface at (u, v).

        Parameters
        ----------
        u, v: array-like or scalar
            u and v may be scalar or vector

        mesh : boolean
            If True, will expand the u and v values into a mesh.
            For example, with u = [0, 1] and v = [0, 1]: if 'mesh'
            is True, the surface will be evaluated at [0, 0], [0, 1],
            [1, 0] and [1, 1], while if it is False, the evalation
            will only be made at [0, 0] and [1, 1]

        Returns
        -------
        If scalar values are passed for *both* u and v, returns
        a 1-D 3-element array [x,y,z]. Otherwise, returns an array
        of shape 3 x len(u) x len(v), suitable for feeding to Mayavi's
        mlab.mesh() plotting function (as mlab.mesh(*arr)).
        """
        u = np.array([u]).reshape(-1,)
        v = np.array([v]).reshape(-1,)
        if mesh:
            # I'm still not sure why we're required to flip u and v
            # below, but trust me, it doesn't work otherwise.
            V, U = np.meshgrid(v, u)
            U = U.ravel()
            V = V.ravel()
        else:
            if len(u) != len(v): # *Need* to mesh this, like above!
                V, U = np.meshgrid(v, u)
                U = U.ravel()
                V = V.ravel()
            else:
                U, V = u, v
        x = self._xsrf.ev(U, V)
        y = self._ysrf.ev(U, V)
        z = self._zsrf.ev(U, V)

        if u.shape == (1,) and v.shape == (1,):
            # Scalar u and v values; return 1-D 3-element array.
            return np.array([x, y, z]).ravel()
        else:
            # u and/or v passed as lists; return 3 x m x n array,
            # where m is len(u) and n is len(v). This format
            # is compatible with enthought.mayavi's mlab.mesh()
            # function.
            return np.array([x, y, z]).reshape(3, len(u), -1)

    def utan(self, u, v):
        dxdu = bisplev(u, v, self._xsrfdef, dx=1, dy=0)
        dydu = bisplev(u, v, self._ysrfdef, dx=1, dy=0)
        dzdu = bisplev(u, v, self._zsrfdef, dx=1, dy=0)
        du = np.array([dxdu, dydu, dzdu])
        du /= np.sqrt((du**2).sum(axis=0))
        return du

    def vtan(self, u, v):
        dxdv = bisplev(u, v, self._xsrfdef, dx=0, dy=1)
        dydv = bisplev(u, v, self._ysrfdef, dx=0, dy=1)
        dzdv = bisplev(u, v, self._zsrfdef, dx=0, dy=1)
        dv = np.array([dxdv, dydv, dzdv])
        dv /= np.sqrt((dv**2).sum(axis=0))
        return dv

    def normal(self, u, v):
        """Get normal(s) at (u, v).
        """
        u = np.array([u]).reshape(-1,)
        v = np.array([v]).reshape(-1,)

        dxdus = bisplev(u, v, self._xsrfdef, dx=1)
        dydus = bisplev(u, v, self._ysrfdef, dx=1)
        dzdus = bisplev(u, v, self._zsrfdef, dx=1)
        dxdvs = bisplev(u, v, self._xsrfdef, dy=1)
        dydvs = bisplev(u, v, self._ysrfdef, dy=1)
        dzdvs = bisplev(u, v, self._zsrfdef, dy=1)

        if u.shape == (1,) and v.shape == (1,):
            # Scalar u and v values; return 1-D 3-element array.
            normal = np.cross([dxdus, dydus, dzdus], [dxdvs, dydvs, dzdvs])
            normal /= np.sqrt(np.dot(normal, normal))
            return normal
        else:
            # u and/or v passed as lists; return 3 x m x n array,
            # where m is len(u) and n is len(v). This format
            # is compatible with enthought.mayavi's mlab.mesh()
            # function.
            normals = np.cross([dxdus, dydus, dzdus], [dxdvs, dydvs, dzdvs],
                                axisa=0, axisb=0)
            normals /= np.sqrt((normals**2).sum(axis=2))[:, :, np.newaxis]
            return normals.T

    @property
    def xyzs(self):
        return self._origxyzs

    def offset(self, offsetamt):
        """Offset the original surface by the given amount
        (but what direction?)
        """
        normals = self.normal(self.u, self.v)
        offpts = self._origxyzs + offsetamt * normals
        return self.__class__(self.u, self.v, offpts,
                              self.bbox, self.ku, self.kv)

    def plane_intersect_pts(self, planedef):
        """Given a plane definition (see below), returns the set of
        intersection points, one for each u-isocurve.

        Parameters
        ----------
        planedef : array_like
            Normalized [A, B, C, D] values from the plane equation
            `Ax + By + Cz - D = 0`.

            As a reminder, A, B and C are the (normalized) values
            of the plane normal vector, and D is the dot product
            of *any* origin-to-plane vector and the normalized
            plane normal vector.

        Returns
        -------
        pts: Array of intersection points
        """

        # Strategy: rotate the bispline coefficients and the
        # original xyz values to put them into 'cut plane coords'.
        # By doing this, we can restrict our evals and optimizations
        # to the z surface only, which speeds things up by a factor
        # of ~3.
        import msk.xy_rots_from_vector as mxyr
        zht = planedef[3]
        tr = mxyr.trsf_from_zvec(planedef[0:3])
        rmat = tr.mat[0:3, 0:3]
        coeffs = np.array([
                    self._xsrf.get_coeffs(),
                    self._ysrf.get_coeffs(),
                    self._zsrf.get_coeffs()])
        rotcoeffs = np.dot(rmat.T, coeffs)
        xyzsshape = self.xyzs.shape
        rotxyzs = np.dot(rmat.T, self.xyzs.reshape(3,-1)).reshape(xyzsshape)
        # Redefine (temporarily) the z-surface coefficients; we'll
        # reset them below when we're done.
        self._zsrf.tck = (self._zsrf.tck[0], self._zsrf.tck[1], rotcoeffs[2])

        import scipy.optimize as spo
        def _get_intersection_v(u, vi=0.0, vf=1.0):
            def intersectfunc(u):
                def thefunc(v):
                    z = self._zsrf.ev(u, v)
                    d = zht - z
                    return d
                return thefunc
            v = spo.brentq(intersectfunc(u), vi, vf, xtol=1e-6)
            return v

        def _get_search_zones():
            """Find the nearest v values above and below the plane
            for each u value (for brentq limits).
            """
            # Check to see if we need to extend the set of
            # surface points before looking for sign changes.
            # (i.e. is zht above the tip or below the base?)
            # Note that we *could* optimize this further by
            # adjusting for each case separately, but I don't think
            # it's worth the extra code.
            if zht > rotxyzs[2,:,-1].min() or zht < rotxyzs[2,:,0].max():
                extvs = np.r_[self.bbox[2], self.bbox[3]]
                # Note (v, u) order in meshgrid(). This is to get
                # the expected array shape on evaluation later.
                V, U = np.meshgrid(extvs, self.u)
                U = U.ravel()
                V = V.ravel()
                extzs = self._zsrf.ev(U, V).reshape(len(self.u), len(extvs))
                vs = np.r_[extvs[0], self.v, extvs[1]]
                zs = np.c_[extzs[:,0], rotxyzs[2], extzs[:,1]]
            else:
                vs = self.v
                zs = rotxyzs[2]

            # Subtract the zht from the surface z values. This
            # will tell us which points are above/below the plane.
            signs = np.sign(zs - zht)
            # Fix the case where we hit the number exactly
            signs[np.where(signs[:] == 0)] = 1
            us, vcrossings = np.where(np.diff(signs) != 0)
            vlimits = [[vs[vcndx], vs[vcndx+1]] for vcndx in vcrossings]
            us = self.u[us]
            return us, vlimits

        us, vlimits = _get_search_zones()
        V = [_get_intersection_v(u, *vs) for u, vs in zip(us, vlimits)]
        # Reset the z-surface coefficients to the original, unrotated
        # values.
        self._zsrf.tck = (self._zsrf.tck[0], self._zsrf.tck[1], coeffs[2])
        pts = self.ev(us, V, mesh=False)

        return pts

    def mplot(self, ures=1, vres=1, **kwargs):
        """Plot the surface using Mayavi's `mesh()` function

        Parameters
        ----------
        ures, vres: int
            Specifies the oversampling of the original
            surface in u and v directions. For example:
            if `ures` = 2, and `self.u` = [0, 1, 2, 3],
            then the surface will be resampled at
            [0, 0.5, 1, 1.5, 2, 2.5, 3] prior to
            plotting.

        kwargs: dict
            See Mayavi docs for `mesh()`

        Returns
        -------
            None
        """
        from enthought.mayavi import mlab

        # Set some Mayavi defaults
        _def_color = (0, 0, 1) # Blue
        if not kwargs.has_key('color'):
            kwargs['color'] = _def_color

        # Make new u and v values of (possibly) higher resolution
        # the original ones.
        u, v = self.u, self.v
        lu, lv = len(u), len(v)
        nus = np.array(list(enumerate(u))).T
        nvs = np.array(list(enumerate(v))).T
        newundxs = np.linspace(0, lu - 1, ures * lu - (ures - 1))
        newvndxs = np.linspace(0, lv - 1, vres * lv - (vres - 1))
        hru = np.interp(newundxs, *nus)
        hrv = np.interp(newvndxs, *nvs)
        # Sample the surface at the new u, v values and plot
        meshpts = self.ev(hru, hrv, mesh=True)
        mlab.mesh(*meshpts, **kwargs)


if __name__ == '__main__':
    from enthought.mayavi import mlab

    # Set up a test surface (wavy cylinder)
    a = np.linspace(0, 2 * np.pi, 360)
    x, y = np.cos(a), np.sin(a)
    z = np.zeros(len(x)) # Seed value
    xyz = np.array([x, y, z])
    xyz = np.array([xyz + i * np.array([[0, 0, .03]]).T
                    for i in range(100) ]).T.swapaxes(0, 1)
    f = 1.3 + .13 * np.sin(4 * np.linspace(0, 2 * np.pi, 100))
    xyz[0:2,:,:] *= f

    srf = ParaSurf(np.linspace(0, 1, len(x)),
                   np.linspace(0, 1, xyz.shape[2]), xyz,
                   bbox=[0.0, 1.0, -0.15, 1.15])
    mlab.mesh(*srf.xyzs, color=(0, 1, 0), opacity=1.0)

    # Create a funky spiral around the surface and plot.
    u = np.linspace(0, 4, 4 * len(x)) % 1
    v = np.linspace(0, 1, 4 * len(y))
    pts = srf.ev(u, v, mesh=False)
    mlab.plot3d(*pts, tube_radius=0.02, color=(1, 1, 1)) # White line

    # Create a test plane and cut the surface with it
    ppt = np.array([0, 0, 2.1]) # Point on plane
    pn = np.array([0.1, 0.6, 0.8]) # Normal to plane
    pn /= np.sqrt(np.dot(pn, pn)) # Create unit vector

    D = np.dot(ppt, pn)
    A, B, C = pn
    planedef = np.array([A, B, C, D])

    # Plot the plane
    def get_z(x, y): return (D - A * x - B * y) / C
    X, Y = 2 * np.mgrid[-1:1:2j, -1:1:2j]
    Z = get_z(X, Y)
    mlab.mesh(X, Y, Z, color=(0, 0, 1), opacity=0.5) # Blue plane

    # Get the plane-surface intersection and plot
    pipts = srf.plane_intersect_pts(planedef)
    mlab.points3d(*pipts, scale_factor=0.03, color=(1, 0, 0))

    # Plot a u-isospline on the surface, using the full
    # surface extensions
    V = np.linspace(srf.bbox[2], srf.bbox[3], 200)
    pts = srf.ev(0.0, V)
    #mlab.plot3d(*pts, tube_radius=0.02, color=(1, 1, 0)) # Yellow line
    mlab.points3d(*pts, scale_factor=0.03, color=(1, 1, 0)) # Yellow dots

Hi,
I have been working through the example using matplotlib rather than mlab and hit a snag at line 182 as "import msk.xy_rots_from_vector as mxyr" fails... After a bit of a google I failed to find a msk module... any chance you could share a link to the module / library?
Thanks!

Edit: the plane surface intersection feature is what brought me here, so it would be great to get it working.
Thanks again!