source: inundation/ga/storm_surge/pyvolution/shallow_water_ext.c @ 643

// Python - C extension module for shallow_water.py
//
// To compile (Python2.3):
//  gcc -c domain_ext.c -I/usr/include/python2.3 -o domain_ext.o -Wall -O 
//  gcc -shared domain_ext.o  -o domain_ext.so  
//
// or use python compile.py
//
// See the module shallow_water.py
//
//
// Ole Nielsen, GA 2004
    
    
#include "Python.h"
#include "Numeric/arrayobject.h"
#include "math.h"

//Shared code snippets
#include "util_ext.h"

const double pi = 3.14159265358979;

// Computational function for rotation
int _rotate(double *q, double n1, double n2) {
  /*Rotate the momentum component q (q[1], q[2])
    from x,y coordinates to coordinates based on normal vector (n1, n2).
    
    Result is returned in array 3x1 r     
    To rotate in opposite direction, call rotate with (q, n1, -n2)
    
    Contents of q are changed by this function */    


  double q1, q2;
  
  //Shorthands
  q1 = q[1];  //uh momentum
  q2 = q[2];  //vh momentum

  //Rotate
  q[1] =  n1*q1 + n2*q2;
  q[2] = -n2*q1 + n1*q2; 

  return 0;
}  


        
// Computational function for flux computation (using stage w=z+h)
int flux_function(double *q_left, double *q_right, 
                  double z_left, double z_right, 
                  double n1, double n2, 
                  double epsilon, double g, 
                  double *edgeflux, double *max_speed) {
  
  /*Compute fluxes between volumes for the shallow water wave equation
    cast in terms of the 'stage', w = h+z using 
    the 'central scheme' as described in
    
    Kurganov, Noelle, Petrova. 'Semidiscrete Central-Upwind Schemes For
    Hyperbolic Conservation Laws and Hamilton-Jacobi Equations'.
    Siam J. Sci. Comput. Vol. 23, No. 3, pp. 707-740.
    
    The implemented formula is given in equation (3.15) on page 714 
  */
       
  int i;
       
  double w_left, h_left, uh_left, vh_left, u_left;
  double w_right, h_right, uh_right, vh_right, u_right;
  double s_min, s_max, soundspeed_left, soundspeed_right;
  double denom, z;
  double q_left_copy[3], q_right_copy[3];    
  double flux_right[3], flux_left[3];
  
  //Copy conserved quantities to protect from modification
  for (i=0; i<3; i++) {
    q_left_copy[i] = q_left[i];
    q_right_copy[i] = q_right[i];
  } 
    
  //Align x- and y-momentum with x-axis
  _rotate(q_left_copy, n1, n2);
  _rotate(q_right_copy, n1, n2);    

  z = (z_left+z_right)/2; //Take average of field values

  //Compute speeds in x-direction
  w_left = q_left_copy[0];              // h+z
  h_left = w_left-z;
  uh_left = q_left_copy[1];

  if (h_left < epsilon) {
    h_left = 0.0;  //Could have been negative
    u_left = 0.0;
  } else { 
    u_left = uh_left/h_left;
  }
  
  w_right = q_right_copy[0];
  h_right = w_right-z;
  uh_right = q_right_copy[1];

  if (h_right < epsilon) {
    h_right = 0.0; //Could have been negative
    u_right = 0.0;
  } else { 
    u_right = uh_right/h_right;
  }

  //Momentum in y-direction              
  vh_left  = q_left_copy[2];
  vh_right = q_right_copy[2];   
       

  //Maximal and minimal wave speeds
  soundspeed_left  = sqrt(g*h_left); 
  soundspeed_right = sqrt(g*h_right);
    
  s_max = max(u_left+soundspeed_left, u_right+soundspeed_right);
  if (s_max < 0.0) s_max = 0.0; 
  
  s_min = min(u_left-soundspeed_left, u_right-soundspeed_right);
  if (s_min > 0.0) s_min = 0.0;   
  
  //Flux formulas  
  flux_left[0] = u_left*h_left;
  flux_left[1] = u_left*uh_left + 0.5*g*h_left*h_left;
  flux_left[2] = u_left*vh_left;
  
  flux_right[0] = u_right*h_right;
  flux_right[1] = u_right*uh_right + 0.5*g*h_right*h_right;
  flux_right[2] = u_right*vh_right;
    

  //Flux computation    
  denom = s_max-s_min;
  if (denom == 0.0) {
    for (i=0; i<3; i++) edgeflux[i] = 0.0;
    *max_speed = 0.0;
  } else {    
    for (i=0; i<3; i++) {
      edgeflux[i] = s_max*flux_left[i] - s_min*flux_right[i];
      edgeflux[i] += s_max*s_min*(q_right_copy[i]-q_left_copy[i]);
      edgeflux[i] /= denom;
    }  
        
    //Maximal wavespeed
    *max_speed = max(fabs(s_max), fabs(s_min));
    
    //Rotate back       
    _rotate(edgeflux, n1, -n2);
  }
  return 0;
}
        
void _manning_friction(double g, double eps, int N, 
                       double* w, double* z, 
                       double* uh, double* vh, 
                       double* eta, double* xmom, double* ymom) {     

  int k;
  double S, h;
  
  for (k=0; k<N; k++) {
    if (eta[k] > eps) {
      h = w[k]-z[k];
      if (h >= eps) {
        S = -g * eta[k]*eta[k] * sqrt((uh[k]*uh[k] + vh[k]*vh[k]));
        S /= pow(h, 7.0/3);      //Expensive (on Ole's home computer)
        //S /= h*h*(1 + h/3.0 - h*h/9.0); //FIXME: Could use a Taylor expansion


        //Update momentum
        xmom[k] += S*uh[k];
        ymom[k] += S*vh[k];
      }
    }
  }
}
            


int _balance_deep_and_shallow(int N,
                              double* wc,
                              double* zc, 
                              double* hc,                             
                              double* wv, 
                              double* zv, 
                              double* hv,
                              double* xmomc, 
                              double* ymomc, 
                              double* xmomv, 
                              double* ymomv) { 
  
  int k, k3, i;
  double dz, hmin, alpha;
  
  //Compute linear combination between constant levels and and
  //levels parallel to the bed elevation.     
  
  for (k=0; k<N; k++) {
    // Compute maximal variation in bed elevation
    // This quantitiy is
    //     dz = max_i abs(z_i - z_c)
    // and it is independent of dimension
    // In the 1d case zc = (z0+z1)/2
    // In the 2d case zc = (z0+z1+z2)/3

    k3 = 3*k;
    
    //FIXME: Try with this one precomputed
    dz = 0.0;
    hmin = hv[k3];
    for (i=0; i<3; i++) {
      dz = max(dz, fabs(zv[k3+i]-zc[k]));
      hmin = min(hmin, hv[k3+i]);
    }

    
    //Create alpha in [0,1], where alpha==0 means using shallow 
    //first order scheme and alpha==1 means using the stage w as
    //computed by the gradient limiter (1st or 2nd order)
    //
    //If hmin > dz/2 then alpha = 1 and the bed will have no effect
    //If hmin < 0 then alpha = 0 reverting to constant height above bed.
    
    if (dz > 0.0) 
      alpha = max( min( 2*hmin/dz, 1.0), 0.0 );
    else
      alpha = 1.0;  //Flat bed

      
    //Weighted balance between stage parallel to bed elevation 
    //(wvi = zvi + hc) and stage as computed by 1st or 2nd 
    //order gradient limiter
    //(wvi = zvi + hvi) where i=0,1,2 denotes the vertex ids
    //
    //It follows that the updated wvi is
    //  wvi := (1-alpha)*(zvi+hc) + alpha*(zvi+hvi) = 
    //  zvi + hc + alpha*(hvi - hc)
    //
    //Note that hvi = zc+hc-zvi in the first order case (constant).

    if (alpha < 1) {         
      for (i=0; i<3; i++) {
        wv[k3+i] = zv[k3+i] + hc[k] + alpha*(hv[k3+i]-hc[k]);
            
      
        //Update momentum as a linear combination of 
        //xmomc and ymomc (shallow) and momentum
        //from extrapolator xmomv and ymomv (deep).
        xmomv[k3+i] = (1-alpha)*xmomc[k] + alpha*xmomv[k3+i];            
        ymomv[k3+i] = (1-alpha)*ymomc[k] + alpha*ymomv[k3+i];
      } 
    }
  }         
  return 0;
}



int _protect(int N, 
             double minimum_allowed_height,        
             double* wc,
             double* zc, 
             double* xmomc, 
             double* ymomc) {
  
  int k; 
  double hc;
  
  //Protect against initesimal and negative heights
  
  for (k=0; k<N; k++) {
    hc = wc[k] - zc[k];
    
    if (hc < minimum_allowed_height) {    
      wc[k] = zc[k];       
      xmomc[k] = 0.0;
      ymomc[k] = 0.0;      
    }
    
  }
  return 0;
}



int _assign_wind_field_values(int N,
                              double* xmom_update,
                              double* ymom_update, 
                              double* s_vec,
                              double* phi_vec,
                              double cw) {

  //Assign windfield values to momentum updates

  int k;
  double S, s, phi, u, v;
  
  for (k=0; k<N; k++) {
    
    s = s_vec[k];
    phi = phi_vec[k];

    //Convert to radians
    phi = phi*pi/180;

    //Compute velocity vector (u, v)
    u = s*cos(phi);
    v = s*sin(phi);
        
    //Compute wind stress
    S = cw * sqrt(u*u + v*v);
    xmom_update[k] += S*u;
    ymom_update[k] += S*v;        
  }
  return 0;  
}            



///////////////////////////////////////////////////////////////////
// Gateways to Python

PyObject *gravity(PyObject *self, PyObject *args) {
  //
  //  gravity(g, h, v, x, xmom, ymom)
  //
  
  
  PyArrayObject *h, *v, *x, *xmom, *ymom;
  int k, i, N, k3, k6;
  double g, avg_h, zx, zy;
  double x0, y0, x1, y1, x2, y2, z0, z1, z2;
    
  if (!PyArg_ParseTuple(args, "dOOOOO",
                        &g, &h, &v, &x, 
                        &xmom, &ymom))  
    return NULL;  

  N = h -> dimensions[0];
  for (k=0; k<N; k++) {
    k3 = 3*k;  // base index
    k6 = 6*k;  // base index    
    
    avg_h = 0.0;
    for (i=0; i<3; i++) {
      avg_h += ((double *) h -> data)[k3+i];
    }   
    avg_h /= 3;
        
    
    //Compute bed slope
    x0 = ((double*) x -> data)[k6 + 0];
    y0 = ((double*) x -> data)[k6 + 1];    
    x1 = ((double*) x -> data)[k6 + 2];
    y1 = ((double*) x -> data)[k6 + 3];        
    x2 = ((double*) x -> data)[k6 + 4];
    y2 = ((double*) x -> data)[k6 + 5];           


    z0 = ((double*) v -> data)[k3 + 0];
    z1 = ((double*) v -> data)[k3 + 1];
    z2 = ((double*) v -> data)[k3 + 2];        

    _gradient(x0, y0, x1, y1, x2, y2, z0, z1, z2, &zx, &zy);

    //Update momentum
    ((double*) xmom -> data)[k] += -g*zx*avg_h;
    ((double*) ymom -> data)[k] += -g*zy*avg_h;        
  }
    
  return Py_BuildValue("");
}


PyObject *manning_friction(PyObject *self, PyObject *args) {
  //
  // manning_friction(g, eps, h, uh, vh, eta, xmom_update, ymom_update)
  //
  
  
  PyArrayObject *w, *z, *uh, *vh, *eta, *xmom, *ymom;
  int N;
  double g, eps;
    
  if (!PyArg_ParseTuple(args, "ddOOOOOOO",
                        &g, &eps, &w, &z, &uh, &vh, &eta, 
                        &xmom, &ymom))  
    return NULL;  

  N = w -> dimensions[0];    
  _manning_friction(g, eps, N,
                    (double*) w -> data,
                    (double*) z -> data,                    
                    (double*) uh -> data, 
                    (double*) vh -> data, 
                    (double*) eta -> data,
                    (double*) xmom -> data, 
                    (double*) ymom -> data);

  return Py_BuildValue("");
}                   

PyObject *rotate(PyObject *self, PyObject *args, PyObject *kwargs) {
  //
  // r = rotate(q, normal, direction=1)
  //
  // Where q is assumed to be a Float numeric array of length 3 and
  // normal a Float numeric array of length 2.

  
  PyObject *Q, *Normal;
  PyArrayObject *q, *r, *normal;
  
  static char *argnames[] = {"q", "normal", "direction", NULL};
  int dimensions[1], i, direction=1;
  double n1, n2;

  // Convert Python arguments to C  
  if (!PyArg_ParseTupleAndKeywords(args, kwargs, "OO|i", argnames, 
                                   &Q, &Normal, &direction))  
    return NULL;  

  //Input checks (convert sequences into numeric arrays)
  q = (PyArrayObject *) 
    PyArray_ContiguousFromObject(Q, PyArray_DOUBLE, 0, 0);
  normal = (PyArrayObject *) 
    PyArray_ContiguousFromObject(Normal, PyArray_DOUBLE, 0, 0);
  
  //Allocate space for return vector r (don't DECREF) 
  dimensions[0] = 3;
  r = (PyArrayObject *) PyArray_FromDims(1, dimensions, PyArray_DOUBLE);

  //Copy
  for (i=0; i<3; i++) {
    ((double *) (r -> data))[i] = ((double *) (q -> data))[i];  
  }
  
  //Get normal and direction
  n1 = ((double *) normal -> data)[0]; 
  n2 = ((double *) normal -> data)[1];   
  if (direction == -1) n2 = -n2;

  //Rotate
  _rotate((double *) r -> data, n1, n2);

  //Release numeric arrays
  Py_DECREF(q);    
  Py_DECREF(normal);
        
  //return result using PyArray to avoid memory leak
  return PyArray_Return(r);
}    


PyObject *compute_fluxes(PyObject *self, PyObject *args) {
  /*Compute all fluxes and the timestep suitable for all volumes
    in domain.

    Compute total flux for each conserved quantity using "flux_function"

    Fluxes across each edge are scaled by edgelengths and summed up
    Resulting flux is then scaled by area and stored in
    explicit_update for each of the three conserved quantities 
    level, xmomentum and ymomentum

    The maximal allowable speed computed by the flux_function for each volume
    is converted to a timestep that must not be exceeded. The minimum of
    those is computed as the next overall timestep.
    
    Python call:
    domain.timestep = compute_fluxes(timestep, 
                                     domain.epsilon, 
                                     domain.g,
                                     domain.neighbours,
                                     domain.neighbour_edges,
                                     domain.normals,
                                     domain.edgelengths,                       
                                     domain.radii,
                                     domain.areas,
                                     Level.edge_values, 
                                     Xmom.edge_values, 
                                     Ymom.edge_values,  
                                     Bed.edge_values,   
                                     Level.boundary_values,
                                     Xmom.boundary_values,
                                     Ymom.boundary_values,
                                     Level.explicit_update,
                                     Xmom.explicit_update,
                                     Ymom.explicit_update)
        

    Post conditions:
      domain.explicit_update is reset to computed flux values
      domain.timestep is set to the largest step satisfying all volumes. 

            
  */

  
  PyArrayObject *neighbours, *neighbour_edges,
    *normals, *edgelengths, *radii, *areas,
    *level_edge_values, 
    *xmom_edge_values, 
    *ymom_edge_values,  
    *bed_edge_values,   
    *level_boundary_values,
    *xmom_boundary_values,
    *ymom_boundary_values,
    *level_explicit_update,
    *xmom_explicit_update,
    *ymom_explicit_update;

    
  //Local variables  
  double timestep, max_speed, epsilon, g;
  double normal[2], ql[3], qr[3], zl, zr;
  double flux[3], edgeflux[3]; //Work arrays for summing up fluxes

  int number_of_elements, k, i, j, m, n;
  int ki, nm, ki2; //Index shorthands
  
 
  // Convert Python arguments to C  
  if (!PyArg_ParseTuple(args, "dddOOOOOOOOOOOOOOOO",
                        &timestep,
                        &epsilon,
                        &g,
                        &neighbours, 
                        &neighbour_edges,
                        &normals, 
                        &edgelengths, &radii, &areas,
                        &level_edge_values, 
                        &xmom_edge_values, 
                        &ymom_edge_values,  
                        &bed_edge_values,   
                        &level_boundary_values,
                        &xmom_boundary_values,
                        &ymom_boundary_values,
                        &level_explicit_update,
                        &xmom_explicit_update,
                        &ymom_explicit_update)) {
    PyErr_SetString(PyExc_RuntimeError, "Input arguments failed");
    return NULL;
  }

  number_of_elements = level_edge_values -> dimensions[0];
  
    
  for (k=0; k<number_of_elements; k++) {

    //Reset work array
    for (j=0; j<3; j++) flux[j] = 0.0;
                          
    //Loop through neighbours and compute edge flux for each
    for (i=0; i<3; i++) {
      ki = k*3+i;
      ql[0] = ((double *) level_edge_values -> data)[ki];
      ql[1] = ((double *) xmom_edge_values -> data)[ki];
      ql[2] = ((double *) ymom_edge_values -> data)[ki];            
      zl =    ((double *) bed_edge_values -> data)[ki];                  
      
      //Quantities at neighbour on nearest face
      n = ((int *) neighbours -> data)[ki];
      if (n < 0) {
        m = -n-1; //Convert negative flag to index
        qr[0] = ((double *) level_boundary_values -> data)[m];  
        qr[1] = ((double *) xmom_boundary_values -> data)[m];   
        qr[2] = ((double *) ymom_boundary_values -> data)[m];   
        zr = zl; //Extend bed elevation to boundary
      } else {    
        m = ((int *) neighbour_edges -> data)[ki];
        
        nm = n*3+m;     
        qr[0] = ((double *) level_edge_values -> data)[nm];
        qr[1] = ((double *) xmom_edge_values -> data)[nm];
        qr[2] = ((double *) ymom_edge_values -> data)[nm];            
        zr =    ((double *) bed_edge_values -> data)[nm];                  
      }
      
      // Outward pointing normal vector   
      // normal = domain.normals[k, 2*i:2*i+2]
      ki2 = 2*ki; //k*6 + i*2
      normal[0] = ((double *) normals -> data)[ki2];
      normal[1] = ((double *) normals -> data)[ki2+1];      

      //Edge flux computation
      flux_function(ql, qr, zl, zr, 
                    normal[0], normal[1],
                    epsilon, g, 
                    edgeflux, &max_speed);

                    
      //flux -= edgeflux * edgelengths[k,i]
      for (j=0; j<3; j++) { 
        flux[j] -= edgeflux[j]*((double *) edgelengths -> data)[ki];
      }
      
      //Update timestep
      //timestep = min(timestep, domain.radii[k]/max_speed)
      if (max_speed > epsilon) {
        timestep = min(timestep, ((double *) radii -> data)[k]/max_speed);
      }   
    } // end for i
    
    //Normalise by area and store for when all conserved
    //quantities get updated
    // flux /= areas[k]
    for (j=0; j<3; j++) { 
      flux[j] /= ((double *) areas -> data)[k];
    }

    ((double *) level_explicit_update -> data)[k] = flux[0];
    ((double *) xmom_explicit_update -> data)[k] = flux[1];
    ((double *) ymom_explicit_update -> data)[k] = flux[2];        

  } //end for k

  return Py_BuildValue("d", timestep);
}    



PyObject *protect(PyObject *self, PyObject *args) {
  //
  //    protect(minimum_allowed_height, wc, zc, xmomc, ymomc)
  

  PyArrayObject 
  *wc,            //Level at centroids
  *zc,            //Elevation at centroids    
  *xmomc,         //Momentums at centroids
  *ymomc; 

    
  int N; 
  double minimum_allowed_height;
  
  // Convert Python arguments to C  
  if (!PyArg_ParseTuple(args, "dOOOO",  
                        &minimum_allowed_height,
                        &wc, &zc, &xmomc, &ymomc))
    return NULL;

  N = wc -> dimensions[0];
    
  _protect(N,
           minimum_allowed_height,
           (double*) wc -> data,
           (double*) zc -> data, 
           (double*) xmomc -> data, 
           (double*) ymomc -> data);
 
  return Py_BuildValue("");  
}



PyObject *balance_deep_and_shallow(PyObject *self, PyObject *args) {
  //
  //    balance_deep_and_shallow(wc, zc, hc, wv, zv, hv,
  //                             xmomc, ymomc, xmomv, ymomv)
  

  PyArrayObject 
    *wc,            //Level at centroids
    *zc,            //Elevation at centroids    
    *hc,            //Height at centroids        
    *wv,            //Level at vertices
    *zv,            //Elevation at vertices
    *hv,            //Heights at vertices    
    *xmomc,         //Momentums at centroids and vertices
    *ymomc, 
    *xmomv, 
    *ymomv;    
    
  int N; //, err;
  
  // Convert Python arguments to C  
  if (!PyArg_ParseTuple(args, "OOOOOOOOOO",  
                        &wc, &zc, &hc, 
                        &wv, &zv, &hv,
                        &xmomc, &ymomc, &xmomv, &ymomv))
    return NULL;

  N = wc -> dimensions[0];
    
  _balance_deep_and_shallow(N,
                            (double*) wc -> data,
                            (double*) zc -> data, 
                            (double*) hc -> data,                           
                            (double*) wv -> data, 
                            (double*) zv -> data, 
                            (double*) hv -> data,
                            (double*) xmomc -> data, 
                            (double*) ymomc -> data, 
                            (double*) xmomv -> data, 
                            (double*) ymomv -> data);  
  
 
  return Py_BuildValue("");  
}







PyObject *assign_windfield_values(PyObject *self, PyObject *args) {
  //
  //      assign_windfield_values(xmom_update, ymom_update,
  //                              s_vec, phi_vec, self.const)

  

  PyArrayObject   //(one element per triangle)
  *s_vec,         //Speeds 
  *phi_vec,       //Bearings 
  *xmom_update,   //Momentum updates
  *ymom_update; 

    
  int N; 
  double cw;
  
  // Convert Python arguments to C  
  if (!PyArg_ParseTuple(args, "OOOOd",  
                        &xmom_update,
                        &ymom_update,                   
                        &s_vec, &phi_vec, 
                        &cw))
    return NULL;

  N = xmom_update -> dimensions[0];
    
  _assign_wind_field_values(N,
           (double*) xmom_update -> data,
           (double*) ymom_update -> data, 
           (double*) s_vec -> data, 
           (double*) phi_vec -> data,
           cw);
 
  return Py_BuildValue("");  
}




//////////////////////////////////////////      
// Method table for python module
static struct PyMethodDef MethodTable[] = {
  /* The cast of the function is necessary since PyCFunction values
   * only take two PyObject* parameters, and rotate() takes
   * three.
   */
  
  {"rotate", (PyCFunction)rotate, METH_VARARGS | METH_KEYWORDS, "Print out"},
  {"compute_fluxes", compute_fluxes, METH_VARARGS, "Print out"},    
  {"gravity", gravity, METH_VARARGS, "Print out"},      
  {"manning_friction", manning_friction, METH_VARARGS, "Print out"},        
  {"balance_deep_and_shallow", balance_deep_and_shallow, 
   METH_VARARGS, "Print out"},          
  {"protect", protect, METH_VARARGS | METH_KEYWORDS, "Print out"},   
  {"assign_windfield_values", assign_windfield_values, 
   METH_VARARGS | METH_KEYWORDS, "Print out"},     
  //{"distribute_to_vertices_and_edges", 
  // distribute_to_vertices_and_edges, METH_VARARGS},    
  //{"update_conserved_quantities", 
  // update_conserved_quantities, METH_VARARGS},       
  //{"set_initialcondition", 
  // set_initialcondition, METH_VARARGS},    
  {NULL, NULL}
};
        
// Module initialisation   
void initshallow_water_ext(void){
  Py_InitModule("shallow_water_ext", MethodTable);
  
  import_array();     //Necessary for handling of NumPY structures  
}