Python scripts I use for drawing meshes from QGIS shapefiles

geodat.py

import numpy as np
from scipy.io import *
import struct
import sys
import math


def read_geodat(filename):
    """
    Read in one of Ian's geodat files.

    Parameters:
    ==========
    filename: the name of the geodat file to read; expects that there is
              also a file called "filename.geodat", which contains info
              on grid size, spacing, etc.

    Returns:
    =======
    x, y: the grid coordinates of the output field
    data: output field
    """

    def _read_geodat(filename):
        # read in the .geodat file, which stores the number of pixels,
        # the pixel size and the location of the lower left corner
        geodatfile = open(filename, "r")
        xgeo = np.zeros((3, 2))
        i = 0
        while True:
            line = geodatfile.readline().split()
            if len(line) == 2:
                try:
                    xgeo[i, 0] = float(line[0])
                    xgeo[i, 1] = float(line[1])
                    i += 1
                except ValueError:
                    i = i
            if len(line) == 0:
                break
        geodatfile.close()
        xgeo[2, :] = xgeo[2, :] * 1000.0
        return xgeo

    # Get the data about the grid from the .geodat file
    xgeo = _read_geodat(filename + ".geodat")
    nx, ny = map(int, xgeo[0,:])
    dx, dy = xgeo[1, :]
    xo, yo = xgeo[2, :]

    x = np.zeros(nx)
    for i in range(nx):
        x[i] = xo + i * dx

    y = np.zeros(ny)
    for i in range(ny):
        y[i] = yo + i * dy

    # Open the x-velocity file and read in all the binary data
    data_file = open(filename, "rb")
    raw_data = data_file.read()
    data_file.close()

    # Unpack that binary data to an array of floats, knowing that it is
    # in big-endian format. Why? God only knows.
    nvals = len(raw_data)/4
    arr = np.zeros(nvals)
    for i in range(nvals):
        arr[i] = struct.unpack('>f', raw_data[4*i: 4*(i+1)])[0]

    # Fairly certain that this is right, but plot it and compare against
    # matlab to be certain
    data = arr.reshape((ny, nx))

    # Find weird points.
    if np.min(data) == -2.0e+9:
        for i in range(1, ny - 1):
            for j in range(1, nx - 1):
                # A point with no data but which has four cardinal neighbors
                # that does can rasonably have data interpolated from them
                if data[i, j] == -2e+9:
                    nbrs = [ [i+1, i-1, i,   i  ],
                             [j,   j,   j+1, j-1] ]
                    k = sum( data[nbrs[0], nbrs[1]] != -2e+9 )
                    if k == 4:
                        data[i,j] = sum( data[nbrs[0],nbrs[1]] )/4.0

                # A point which does have data but for which only one of its
                # neighbors neighbors does should not have data
                else:
                    nbrs = [ [i+1, i+1, i+1, i,   i,   i-1, i-1, i-1],
                             [j+1, j,   j-1, j+1, j-1, j+1, j,   j-1] ]
                    k = sum( data[nbrs[0],nbrs[1]]!=-2e+9 )
                    if k <= 1:
                        data[i,j] = -2e+9
                        data[i,j] = -2e+9

        for i in range(1, ny - 1):
            for j in range(1, nx - 1):
                if data[i,j] != -2e+9:
                    nbrs = [ [i+1, i+1, i+1, i,   i,   i-1, i-1, i-1],
                             [j+1, j,   j-1, j+1, j-1, j+1, j,   j-1] ]
                    k = sum( data[nbrs[0], nbrs[1]] != -2e+9 )
                    if k < 4:
                        data[i,j] = -2e+9
                        data[i,j] = -2e+9

    return x, y, data

streamlines.py

import numpy as np
from shapefile import *

# -------------------------------
def interpolate(x, y, x0, y0, q):
    """
    Interpolate the value of the field `q` defined on the grid `x`, `y`
    to the point `x0`, `y0`
    """
    dx = x[1]-x[0]
    dy = y[1]-y[0]

    i = int( (y0-y[0])/dy )
    j = int( (x0-x[0])/dx )

    alpha_x = (x0 - x[j]) / dx
    alpha_y = (y0 - y[i]) / dy

    p = (q[i,j]
            + alpha_x * (q[i, j + 1] - q[i, j])
            + alpha_y * (q[i + 1, j] - q[i, j])
            + alpha_x * alpha_y * (q[i + 1, j + 1] + q[i, j]
                                     - q[i + 1, j] - q[i, j + 1]))

    return p


# ---------------------------------------------
def streamline(x, y, vx, vy, x0, y0, sign = 1):
    """
    Given the x/y velocity fields `vx`, `vy`, defined at the grid points
    `x`, `y`, generate a streamline originating at the point `x0`, `y0`.
    The streamline will go backwards if the optional argument `sign` = -1.

    The algorithm we use is an adaptive forward Euler method. It's not very
    good. Willie hears ye; Willie don't care.

    Parameters:
    ==========
    x, y: coordinates at which the fields are defined
    vx, vy: velocities in the x, y directions
    x0, y0: starting coordinate of the streamline
    sign: optional; =1 if the streamline is forward, -1 if backward

    Returns:
    =======
    X, Y: coordinates of the resultant streamline
    """
    nx = len(x)
    ny = len(y)

    u = interpolate(x, y, x0, y0, vx)
    v = interpolate(x, y, x0, y0, vy)

    speed = np.sqrt(u**2 + v**2)

    X = [ x0 ]
    Y = [ y0 ]

    k = 0

    while (speed > 5.0 and k < 10000):
        k += 1
        dt = sign * 50.0 / speed

        x0 = x0 + dt * u
        y0 = y0 + dt * v

        X.append(x0)
        Y.append(y0)

        u = interpolate(x, y, x0, y0, vx)
        v = interpolate(x, y, x0, y0, vy)
        speed = np.sqrt(u**2 + v**2)

    return X, Y


# ----------------------------------------------------
def streamlines_from_shapefile(x, y, vx, vy, filename, sign = 1):
    """
    Given an ESRI shapefile, read in all the points it contains and
    generate streamlines from them.

    Parameters:
    ==========
    x, y, vx, vy: same as in last function
    filename: name of .shp file from which we get start points

    Returns:
    =======
    lines: array of streamlines
    """

    sf = Reader(filename)
    shape = sf.shapes()[0]
    num_pts = len(shape.points)

    X0 = np.zeros(num_pts, dtype = np.float64)
    Y0 = np.zeros(num_pts, dtype = np.float64)

    for i in range(num_pts):
        X0[i] = shape.points[i][0]
        Y0[i] = shape.points[i][1]

    lines = []

    for i in range(num_pts):
        X, Y = streamline(x, y, vx, vy, X0[i], Y0[i], sign)

        line = []
        for k in range(len(X)):
            line.append([X[k], Y[k]])
        lines.append(line)

    return lines


# --------------------------------------------
def write_streamlines(lines, filename):
    """
    Given an array of streamlines, write them out to a shapefile.

    Parameters:
    ==========
    lines: array of arrays of tuples, coordinates along each streamline
    filename: name for output file
    """

    strm = Writer()
    strm.autoBalance = 1
    strm.field('FIELD', 'C', '1')

    for line in lines:
        strm.line(parts = [line])
        strm.record('')

    strm.save(filename)


# -------------------------------------
def make_streamlines(velocity_filename,
                     initial_shapefile,
                     streamlines_shapefile,
                     inflow = 1):
    """
    Parameters:
    ==========
    velocity_filename:     stem of the geodat filenames for the ice
                           velocities
    initial_shapefile:     shapefile containing the points from which
                           to create the streamlines
    streamlines_shapefile: shapefile to write streamlines to
    inflow:                optional argument; specify whether the start
                           points are at the glacier inflow or outflow
    """

    x, y, vx = read_geodat(velocity_filename + ".vx")
    x, y, vy = read_geodat(velocity_filename + ".vy")

    ny, nx = np.shape(vx)
    for i in range(ny):
        for j in range(nx):
            if vx[i, j] == -2.0e+9:
                vx[i, j] = 0.0
                vy[i, j] = 0.0

    lines = streamlines_from_shapefile(x, y, vx, vy,
                                      initial_shapefile, sign)

    write_streamlines(lines, streamlines_shapefile)