Merge pull request #116 from sandialabs/plumbing/variable_properties

Plumbing/variable properties

Author: ralberd

optimism/FunctionSpace.py

@@ -193,7 +193,14 @@ def interpolate_to_points(functionSpace, nodalField):
     return jax.vmap(interpolate_to_element_points, (None, 0, 0))(nodalField, functionSpace.shapes, functionSpace.mesh.conns)
 
 
-def integrate_over_block(functionSpace, U, stateVars, dt, func, block,
+def vmapPropValue(propArray):
+    numAxes = len(propArray.shape)
+    if numAxes > 1:
+        return 0
+    else:
+        return None
+
+def integrate_over_block(functionSpace, U, stateVars, props, dt, func, block,
                          *params, modify_element_gradient=default_modify_element_gradient):
     """Integrates a density function over a block of the mesh.
 
@@ -219,11 +226,11 @@ def integrate_over_block(functionSpace, U, stateVars, dt, func, block,
       block of elements.
     """
 
-    vals = evaluate_on_block(functionSpace, U, stateVars, dt, func, block, *params, modify_element_gradient=modify_element_gradient)
+    vals = evaluate_on_block(functionSpace, U, stateVars, props, dt, func, block, *params, modify_element_gradient=modify_element_gradient)
     return np.dot(vals.ravel(), functionSpace.vols[block].ravel())
 
 
-def evaluate_on_block(functionSpace, U, stateVars, dt, func, block,
+def evaluate_on_block(functionSpace, U, stateVars, props, dt, func, block,
                       *params, modify_element_gradient=default_modify_element_gradient):
     """Evaluates a density function at every quadrature point in a block of the mesh.
 
@@ -249,23 +256,60 @@ def evaluate_on_block(functionSpace, U, stateVars, dt, func, block,
       density functional ``func`` at every quadrature point in the block.
     """
     fs = functionSpace
-    compute_elem_values = jax.vmap(evaluate_on_element, (None, None, 0, None, 0, 0, 0, 0, None, None, *tuple(0 for p in params)))
 
-    blockValues = compute_elem_values(U, fs.mesh.coords, stateVars[block], dt, fs.shapes[block],
+    compute_elem_values = jax.vmap(evaluate_on_element, (None, None, 0, vmapPropValue(props), None, 0, 0, 0, 0, None, None, *tuple(0 for p in params)))
+    
+    blockValues = compute_elem_values(U, fs.mesh.coords, stateVars[block], props, dt, fs.shapes[block],
+                                      fs.shapeGrads[block], fs.vols[block],
+                                      fs.mesh.conns[block], func, modify_element_gradient, *params)
+    return blockValues
+
+
+def evaluate_on_block_heterogeneous_props(
+    functionSpace, U, stateVars, props, dt, func, block,
+    *params, modify_element_gradient=default_modify_element_gradient
+):
+    """Evaluates a density function at every quadrature point in a block of the mesh.
+
+    Args:
+      functionSpace: Function space object to do the evaluation with.
+      U: The vector of dofs for the primal field in the functional.
+      stateVars: Internal state variable array.
+      dt: Current time increment
+      func: Lagrangian density function to evaluate, Must have the signature
+        ``func(u, dudx, q, x, *params) -> scalar``, where ``u`` is the primal field, ``q`` is the
+        value of the internal variables, ``x`` is the current point coordinates, and ``*params`` is
+        a variadic set of additional parameters, which correspond to the ``*params`` argument.
+      block: Group of elements to evaluate over. This is an array of element indices. For
+        performance, the elements within the block should be numbered consecutively.
+      *params: Optional parameter fields to pass into Lagrangian density function. These are
+        represented as a single value per element.
+      modify_element_gradient: Optional function that modifies the gradient at the element level.
+        This can be to set the particular 2D mode, and additionally to enforce volume averaging
+        on the gradient operator. This is a keyword-only argument.
+
+    Returns:
+      An array of shape (numElements, numQuadPtsPerElement) that contains the scalar values of the
+      density functional ``func`` at every quadrature point in the block.
+    """
+    fs = functionSpace
+    compute_elem_values = jax.vmap(evaluate_on_element, (None, None, 0, 0, None, 0, 0, 0, 0, None, None, *tuple(0 for p in params)))
+    
+    blockValues = compute_elem_values(U, fs.mesh.coords, stateVars[block], props[block], dt, fs.shapes[block],
                                       fs.shapeGrads[block], fs.vols[block],
                                       fs.mesh.conns[block], func, modify_element_gradient, *params)
     return blockValues
 
 
-def integrate_element_from_local_field(elemNodalField, elemNodalCoords, elemStates, dt, elemShapes, elemShapeGrads, elemVols, func, modify_element_gradient=default_modify_element_gradient):
+def integrate_element_from_local_field(elemNodalField, elemNodalCoords, elemStates, elemProps, dt, elemShapes, elemShapeGrads, elemVols, func, modify_element_gradient=default_modify_element_gradient):
     """Integrate over element with element nodal field as input.
     This allows element residuals and element stiffness matrices to computed.
     """
     elemVals = jax.vmap(interpolate_to_point, (None,0))(elemNodalField, elemShapes)
     elemGrads = jax.vmap(compute_quadrature_point_field_gradient, (None,0))(elemNodalField, elemShapeGrads)
     elemGrads = modify_element_gradient(elemGrads, elemShapes, elemVols, elemNodalField, elemNodalCoords)
     elemPoints = jax.vmap(interpolate_to_point, (None,0))(elemNodalCoords, elemShapes)
-    fVals = jax.vmap(func, (0, 0, 0, 0, None))(elemVals, elemGrads, elemStates, elemPoints, dt)
+    fVals = jax.vmap(func, (0, 0, 0, None, 0, None))(elemVals, elemGrads, elemStates, elemProps, elemPoints, dt)
     return np.dot(fVals, elemVols)
 
 
@@ -299,12 +343,13 @@ def integrate_element(U, coords, elemStates, elemShapes, elemShapeGrads, elemVol
     return np.dot(fVals, elemVols)
 
 
-def evaluate_on_element(U, coords, elemStates, dt, elemShapes, elemShapeGrads, elemVols, elemConn, kernelFunc, modify_element_gradient, *params):
+def evaluate_on_element(U, coords, elemStates, props, dt, elemShapes, elemShapeGrads, elemVols, elemConn, kernelFunc, modify_element_gradient, *params):
     elemVals = interpolate_to_element_points(U, elemShapes, elemConn)
     elemGrads = compute_element_field_gradient(U, coords, elemShapes, elemShapeGrads, elemVols, elemConn, modify_element_gradient)
     elemXs = interpolate_to_element_points(coords, elemShapes, elemConn)
-    vmapArgs = 0, 0, 0, 0, None, *tuple(None for p in params)
-    fVals = jax.vmap(kernelFunc, vmapArgs)(elemVals, elemGrads, elemStates, elemXs, dt, *params)
+    # vmapArgs = 0, 0, 0, vmapPropValue(props), 0, None, *tuple(None for p in params)
+    vmapArgs = 0, 0, 0, None, 0, None, *tuple(None for p in params)
+    fVals = jax.vmap(kernelFunc, vmapArgs)(elemVals, elemGrads, elemStates, props, elemXs, dt, *params)
     return fVals