def __init__(self, center, r1, r2):
"""Create an Annulus shape
Parameters
----------
center : center of the annulus
r1 : Internal radius
r2 : External radius
Returns
-------
An UWGeodynamics Shape object
"""
self.center = center
self.r1 = r1
self.r2 = r2
self.bottom = center[1] - self.r2
self.top = center[1] + self.r2
center = tuple(nd(x) for x in list(self.center))
r1 = nd(self.r1)
r2 = nd(self.r2)
coord = fn.input() - center
self._fn = (fn.math.dot(coord, coord) < r2**2) & (fn.math.dot(
coord, coord) > r1**2)
super(Annulus, self).__init__(argument_fns=None)
self._fncself = self._fn._fncself
def __init__(self, center, radius):
"""Create a Disk shape
Parameters
----------
center : center of the disk
radius : radius of the disk
Returns
-------
An UWGeodynamics Shape
"""
self.center = center
self.radius = radius
self.top = center[1] + self.radius
self.bottom = center[1] - self.radius
center = tuple(nd(x) for x in list(self.center))
radius = nd(self.radius)
coord = fn.input() - center
self._fn = fn.math.dot(coord, coord) < radius**2
super(Disk, self).__init__(argument_fns=None)
self._fncself = self._fn._fncself
def fn_Tukey_window(r, centre, width, top, bottom):
""" Define a tuckey window
A Tukey window is a rectangular window with the first and last r/2
percent of the width equal to parts of a cosine.
see tappered cosine function
"""
centre = nd(centre)
width = nd(width)
top = nd(top)
bottom = nd(bottom)
x = fn.input()[0]
y = fn.input()[1]
start = centre - 0.5 * width
xx = (x - start) / width
x_conditions = [
((0. <= xx) & (xx < r / 2.0),
0.5 * (1.0 + fn.math.cos(2. * np.pi / r * (xx - r / 2.0)))),
((r / 2.0 <= xx) & (xx < 1.0 - r / 2.0), 1.0),
((1.0 - r / 2.0 <= xx) & (xx <= 1.0),
0.5 * (1. + fn.math.cos(2. * np.pi / r * (xx + r / 2.0)))),
(True, 0.0)
]
x_conditions = fn.branching.conditional(x_conditions)
y_conditions = fn.branching.conditional([((y >= bottom) & (y <= top), 1.0),
(True, 0.0)])
return x_conditions * y_conditions
def __init__(self, Model, elementRes, minCoord, maxCoord, velocityField):
self.minCoord = minCoord
self.maxCoord = maxCoord
self.elementRes = elementRes
self.velocityField = velocityField
minCoord = tuple([nd(val) for val in self.minCoord])
maxCoord = tuple([nd(val) for val in self.maxCoord])
self.mesh = uw.mesh.FeMesh_Cartesian(elementType="Q1/dQ0",
elementRes=self.elementRes,
minCoord=minCoord,
maxCoord=maxCoord)
boundaryNodes = (Model.left_wall + Model.right_wall + Model.top_wall +
Model.bottom_wall)
self.Model = Model
# Build a KDTree to handle boundaries
self.boundaries = boundaryNodes.data
x = Model.mesh.data[self.boundaries, 0]
y = Model.mesh.data[self.boundaries, 1]
coords = np.zeros((x.size, 2))
coords[:, 0] = x.ravel()
coords[:, 1] = y.ravel()
self.tree = spatial.cKDTree(coords)
def extract_profile(field, line, nsamples=1000):
""" Extract values along a line
Parameters:
field: The field to extract the data from
line: list of (x,y, [z]) coordinates defining the sampling
line.
nsamples: number of sampling points
"""
if size > 1:
raise NotImplementedError("""The extract_profile function will not work
in parallel""")
coords = np.array([(nd(x), nd(y)) for (x, y) in line])
x = np.linspace(coords[0, 0], coords[-1, 0], nsamples)
if coords[0, 0] == coords[-1, 0]:
y = np.linspace(coords[0, 1], coords[-1, 1], nsamples)
else:
from scipy import interpolate
f = interpolate.interp1d(coords[:, 0], coords[:, 1])
y = f(x)
dx = np.diff(x)
dy = np.diff(y)
distances = np.zeros(x.shape)
distances[1:] = np.sqrt(dx**2 + dy**2)
distances = np.cumsum(distances)
pts = np.array([x, y]).T
values = field.evaluate(pts)
return distances, values
def _cohesionFn(self):
if self.plasticStrain:
cohesion = self.cohesionWeakeningFn(
self.plasticStrain,
Cohesion=nd(self.cohesion),
CohesionSw=nd(self.cohesionAfterSoftening))
else:
cohesion = fn.misc.constant(self.cohesion)
return cohesion
def circle_points_tracers(radius, centre=tuple([0., 0.]), npoints=72):
angles = np.linspace(0, 360, npoints)
radius = nd(radius)
x = radius * np.cos(np.radians(angles)) + nd(centre[0])
y = radius * np.sin(np.radians(angles)) + nd(centre[1])
coords = np.ndarray((x.size, 2))
coords[:, 0] = x
coords[:, 1] = y
return coords
def _effectiveViscosity(self):
if not self.viscosity:
raise ValueError("Can not find viscosity field")
# Maxwell relaxation time
alpha = self.viscosity / nd(self.shear_modulus)
# observation time
dt_e = nd(self.observation_time)
# Calculate effective viscosity
mu_eff = (self.viscosity * dt_e) / (alpha + dt_e)
return mu_eff
def _elastic_stress(self):
# Check that the viscosity field has been properly
# linked
if not self.viscosity:
raise ValueError("Can not find viscosity field")
if not self.previousStress:
raise ValueError("Can not find previous stress field")
elasticStressFn = self.viscosity / (nd(self.shear_modulus) *
nd(self.observation_time))
elasticStressFn *= self.previousStress
return elasticStressFn
def solve(self, dt):
if not self.Model:
raise ValueError("Model is not defined")
self.Model.materialField.data[:] = self._fn.evaluate(self.Model.swarm)
if self.surfaceTracers:
if self.surfaceTracers.swarm.particleCoordinates.data.size > 0:
coords = self.surfaceTracers.swarm.particleCoordinates
coords.data[coords.data[:, -1] > nd(self.threshold),
-1] = nd(self.threshold)
return
def sphere_points_tracers(radius, centre=tuple([0., 0., 0.]), npoints=30):
theta = np.linspace(0, 180, npoints)
phi = np.linspace(0, 360, npoints)
radius = nd(radius)
theta, phi = np.meshgrid(theta, phi)
x = radius * np.sin(np.radians(theta.ravel())) * np.cos(np.radians(phi.ravel()))
y = radius * np.sin(np.radians(theta.ravel())) * np.sin(np.radians(phi.ravel()))
z = radius * np.cos(np.radians(theta.ravel()))
x += nd(centre[0])
y += nd(centre[1])
z += nd(centre[2])
return x, y, z
def __init__(self, velocity):
self._Model = None
self._wall = None
self._time = None
self.velocity = velocity
self.material = None
if isinstance(velocity, (list, tuple)):
self.velocityFn = fn.branching.conditional(velocity)
else:
self.velocityFn = fn.misc.constant(nd(velocity))
self.wall_operators = {
"left": op.le,
"right": op.ge,
"top": op.ge,
"bottom": op.le,
"front": op.ge,
"back": op.le
}
self.wall_direction_axis = {
"left": 0,
"right": 0,
"front": 1,
"back": 1,
"top": -1,
"bottom": -1
}
def __init__(self, normal, origin=None, reverse=False):
""" HalfSpace
Parameters:
-----------
normal: A vector defining the normal to the plan.
origin: Origin
reverse: by default, particles tested against this class are
assigned "True" if they lay on or below the plan.
You can reverse than behavior by setting reverse=True.
Returns:
--------
A UWGeodynamics Shape object.
"""
if isinstance(normal, (tuple, list)):
self.normal = fn.misc.constant([float(nd(val)) for val in normal])
else:
raise ValueError("{0} must be a list or tuple".format(normal))
if isinstance(origin, (tuple, list)):
self.origin = fn.misc.constant([float(nd(val)) for val in origin])
else:
self.origin = fn.misc.constant([0.] * len(normal))
self.reverse = reverse
coords = fn.input()
new_coords = coords - self.origin
func = fn.math.dot(self.normal, new_coords)
# True if below, False if above
if not self.reverse:
conditions = [(func <= 0., True), (func > 0., False)]
else:
conditions = [(func >= 0., True), (func < 0., False)]
self._fn = fn.branching.conditional(conditions)
super(HalfSpace, self).__init__(argument_fns=None)
self._fncself = self._fn._fncself
def __init__(self, value, min_value=None, max_value=None):
self.value = fn.Function.convert(value)
self.min_value = fn.Function.convert(nd(min_value))
self.max_value = fn.Function.convert(nd(max_value))
if self.max_value and self.min_value:
self._fn = fn.misc.min(self.value, self.max_value)
self._fn = fn.misc.max(self._fn, self.min_value)
elif self.max_value:
self._fn = fn.misc.min(self.value, self.max_value)
elif self.min_value:
self._fn = fn.misc.max(self.value, self.min_value)
else:
self._fn = self.value
super(Limiter, self).__init__(
argument_fns=[self.value, self.min_value, self.max_value])
self._fncself = self._fn._fncself
def __init__(self, top, bottom, minX=0., maxX=0., minY=None, maxY=None):
"""Create a Box Shape
Parameters
----------
top : Top of the Box
bottom : Bottom of the Box
minX : Minimum extent of the Box along the x-axis
maxX : Maximum extent of the Box along the x-axis
Only in 3D:
minY : Minimum extent of the Box along the y-axis
maxY : Maximum extent of the Box along the y-axis
Returns
-------
"""
self.top = top
self.bottom = bottom
self.minX = minX
self.maxX = maxX
self.minY = minY
self.maxY = maxY
coord = fn.input()
if (self.minY is not None) and (self.maxY is not None):
func = ((coord[1] <= nd(self.maxY)) &
(coord[1] >= nd(self.minY)) &
(coord[0] <= nd(self.maxX)) &
(coord[0] >= nd(self.minX)) &
(coord[2] <= nd(self.top)) &
(coord[2] >= nd(self.bottom)))
else:
func = ((coord[1] <= nd(self.top)) &
(coord[1] >= nd(self.bottom)) &
(coord[0] <= nd(self.maxX)) &
(coord[0] >= nd(self.minX)))
self._fn = func
super(Box, self).__init__(argument_fns=None)
self._fncself = self._fn._fncself
def __init__(self, vertices):
"""Create a polygon shape
Parameters
----------
vertices : vertices of the polygon as (x,y) pairs.
Returns
-------
Polygon Shape Class
"""
self.vertices = vertices
self.top = min([y for x, y in self.vertices])
self.bottom = max([y for x, y in self.vertices])
vertices = [(nd(x), nd(y)) for x, y in self.vertices]
self._fn = uw.function.shape.Polygon(np.array(vertices))
super(Polygon, self).__init__(argument_fns=None)
self._fncself = self._fn._fncself
def _create_function(self):
# Create wall function
operator = self.wall_operators[self._wall]
axis = self.wall_direction_axis[self._wall]
pos = self.wall_init_pos[self._wall]
condition = [(operator(fn.input()[axis],
(self._time * self.velocityFn + nd(pos))),
True), (True, False)]
return fn.branching.conditional(condition)
def _frictionFn(self):
if self.plasticStrain:
friction = self.frictionWeakeningFn(
self.plasticStrain,
FrictionCoef=self.frictionCoefficient,
FrictionCoefSw=self.frictionAfterSoftening,
epsilon1=self.epsilon1,
epsilon2=self.epsilon2)
else:
friction = fn.math.atan(nd(self.frictionCoefficient))
return friction
def __init__(self,
airIndex,
sedimentIndex,
XML,
resolution,
checkpoint_interval,
surfElevation=0.,
verbose=True,
Model=None,
outputDir="outbdls",
restartFolder=None,
restartStep=None,
timeField=None,
minCoord=None,
maxCoord=None,
aspectRatio2d=1.):
try:
import pyBadlands
except ImportError:
raise ImportError("""pyBadlands import as failed. Please check your
installation, PYTHONPATH and PATH environment
variables""")
self.verbose = verbose
self.outputDir = outputDir
self.restartStep = restartStep
self.restartFolder = restartFolder
self.airIndex = airIndex
self.sedimentIndex = sedimentIndex
self.resolution = nd(resolution)
self.surfElevation = fn.Function.convert(nd(surfElevation))
self.checkpoint_interval = nd(checkpoint_interval)
self.timeField = timeField
self.XML = XML
self.time_years = 0.
self.minCoord = minCoord
self.maxCoord = maxCoord
self.aspectRatio2d = aspectRatio2d
self.Model = Model
def __init__(self, top, bottom):
"""Create a 2D Layer object
Parameters
----------
top : top of the layer
bottom : bottom of the layer
Returns
-------
AN UWGeodynamics Shape object
"""
self.top = top
self.bottom = bottom
coord = fn.input()
self._fn = ((coord[1] <= nd(self.top)) & (coord[1] >= nd(self.bottom)))
super(Layer, self).__init__(argument_fns=None)
self._fncself = self._fn._fncself
def sphere_points_tracers(radius, centre=tuple([0., 0., 0.]), npoints=30):
theta = np.linspace(0, 180, npoints)
phi = np.linspace(0, 360, npoints)
radius = nd(radius)
theta, phi = np.meshgrid(theta, phi)
x = radius * np.sin(np.radians(theta.ravel())) * np.cos(np.radians(phi.ravel()))
y = radius * np.sin(np.radians(theta.ravel())) * np.sin(np.radians(phi.ravel()))
z = radius * np.cos(np.radians(theta.ravel()))
x += nd(centre[0])
y += nd(centre[1])
z += nd(centre[2])
# Finally, returns a 2D array
coords = np.ndarray((x.size, 2))
coords[:, 0] = x
coords[:, 1] = y
coords[:, 2] = z
return coords
def _advect_surface(self, dt):
if self.top:
# Extract top surface
x = self.model.mesh.data[self.top.data][:, 0]
y = self.model.mesh.data[self.top.data][:, 1]
# Extract velocities from top
vx = self.model.velocityField.data[self.top.data][:, 0]
vy = self.model.velocityField.data[self.top.data][:, 1]
# Advect top surface
x2 = x + vx * nd(dt)
y2 = y + vy * nd(dt)
# Spline top surface
f = interp1d(x2, y2, kind='cubic', fill_value='extrapolate')
self.TField.data[self.top.data, 0] = f(x)
uw.barrier()
self.TField.syncronise()
def __init__(self, model):
"""Create a Freesurface processor
Parameters
----------
model : UWGeodynamics Model
"""
self.model = model
minCoord = tuple([nd(val) for val in self.model.minCoord])
maxCoord = tuple([nd(val) for val in self.model.maxCoord])
# Initialize model mesh
self._init_mesh = uw.mesh.FeMesh_Cartesian(
elementType=self.model.elementType,
elementRes=self.model.elementRes,
minCoord=minCoord,
maxCoord=maxCoord,
periodic=self.model.periodic)
# Create the tools
self.TField = self._init_mesh.add_variable(nodeDofCount=1)
self.TField.data[:, 0] = self._init_mesh.data[:, 1].copy()
self.top = self.model.top_wall
self.bottom = self.model.bottom_wall
# Create boundary condition
self._conditions = uw.conditions.DirichletCondition(
variable=self.TField, indexSetsPerDof=(self.top + self.bottom, ))
# Create Eq System
self._system = uw.systems.SteadyStateHeat(temperatureField=self.TField,
fn_diffusivity=1.0,
conditions=self._conditions)
self._solver = uw.systems.Solver(self._system)
def __init__(self,
reference_density,
thermalExpansivity=3e-5 / u.kelvin,
reference_temperature=273.15 * u.degK,
beta=0. / u.pascal,
reference_pressure=0. * u.pascal):
""" The LinearDensity function calculates:
density = rho0 * (1 + (beta * deltaP) - (alpha * deltaT))
where deltaP is the difference between P and the reference P,
and deltaT is the difference between T and the reference T
Parameters
----------
reference_density : reference density
thermalExpansivity : thermal expansivity of the material at the
temperature of reference.
reference_temperature : reference temperature
beta : coefficient of compressibility
reference_pressure : reference pressure
Returns
-------
An UWGeodynamics Linear Density object.
"""
super(LinearDensity, self).__init__()
self.name = "Linear (ref: {0})".format(str(reference_density))
self.reference_density = reference_density
self.reference_temperature = reference_temperature
self.thermalExpansivity = thermalExpansivity
self.reference_pressure = reference_pressure
self._alpha = nd(thermalExpansivity)
self._beta = nd(beta)
self._Tref = nd(reference_temperature)
self._Pref = nd(reference_pressure)
def solve(self, dt, sigma=0):
if rank == 0 and self.verbose:
purple = "\033[0;35m"
endcol = "\033[00m"
print(purple + "Processing surface with Badlands" + endcol)
sys.stdout.flush()
np_surface = None
if rank == 0:
rg = self.badlands_model.recGrid
if self.Model.mesh.dim == 2:
zVals = rg.regZ.mean(axis=1)
np_surface = np.column_stack((rg.regX, zVals))
if self.Model.mesh.dim == 3:
np_surface = np.column_stack((rg.rectX, rg.rectY, rg.rectZ))
np_surface = comm.bcast(np_surface, root=0)
comm.Barrier()
# Get Velocity Field at the surface
nd_coords = nd(np_surface * u.meter)
tracer_velocity = self.Model.velocityField.evaluate_global(nd_coords)
dt_years = dimensionalise(dt, u.years).magnitude
if rank == 0:
tracer_disp = dimensionalise(tracer_velocity * dt,
u.meter).magnitude
self._inject_badlands_displacement(self.time_years, dt_years,
tracer_disp, sigma)
# Run the Badlands model to the same time point
self.badlands_model.run_to_time(self.time_years + dt_years)
self.time_years += dt_years
# TODO: Improve the performance of this function
self._update_material_types()
comm.Barrier()
if rank == 0 and self.verbose:
purple = "\033[0;35m"
endcol = "\033[00m"
print(purple + "Processing surface with Badlands...Done" + endcol)
sys.stdout.flush()
return
def __init__(self, reference_density):
"""Constant density function
Parameters
----------
reference_density : density
Returns
-------
An UWGeodynamics Constant Density object
"""
self.reference_density = reference_density
self._density = nd(reference_density)
self.name = "Constant ({0})".format(str(reference_density))
def add_particles_with_coordinates(self, vertices, **kwargs):
for dim, _ in enumerate(vertices):
vertices[dim] = nd(vertices[dim])
sizes = np.array([np.array(x).size for x in vertices])
points = np.zeros((sizes.max(), len(vertices)))
for dim, _ in enumerate(vertices):
points[:, dim] = vertices[dim]
vals = super(PassiveTracers,
self).add_particles_with_coordinates(points, **kwargs)
self.advector = uw.systems.SwarmAdvector(
swarm=self, velocityField=self.velocityField, order=2)
self._global_indices()
return vals
def _init_model(self):
materialField = self.Model.materialField
materialMap = {}
for material in self.air:
materialMap[material.index] = 1.0
isAirMaterial = fn.branching.map(fn_key=materialField,
mapping=materialMap,
fn_default=0.0)
belowthreshold = [
(((isAirMaterial < 0.5) & (fn.input()[1] > nd(self.threshold))),
self.air[0].index), (True, materialField)
]
self._fn = fn.branching.conditional(belowthreshold)
def effective_density(self):
"""calculate effective_density based
on PT conditions"""
density = nd(self.reference_density)
# Temperature dependency
if not self.temperatureField:
raise RuntimeError("No temperatureField found!")
t_term = self._alpha * (self.temperatureField - self._Tref)
# Pressure dependency
if not self.pressureField:
raise RuntimeError("No pressureField found!")
p_term = self._beta * (self.pressureField - self._Pref)
return density * (1.0 + p_term - t_term)
def _apply_conditions_nodes(self, condition, nodes):
""" Apply condition to a set of nodes
Parameters:
-----------
condition:
condition
nodes: np.array, list, IndexSet, Function
set of nodes
"""
nodes = self._convert_nodes_to_indexSets(nodes)
# Expect a list or tuple of dimension equal to the dof of the
# field on which the condition is used
if isinstance(condition, (float, int, u.Quantity)):
condition = [condition]
# Check that the domain actually contains some boundary nodes
if isinstance(condition, (list, tuple)) and nodes.data.size > 0:
for dim in range(self.field.data.shape[1]):
# Scalar condition
if isinstance(condition[dim], (u.Quantity, float, int)):
self.field.data[nodes.data, dim] = nd(condition[dim])
self._add_to_indices(dim, nodes)
# If it's an Underworld function
elif isinstance(condition[dim], fn.Function):
func = condition[dim]
values = func.evaluate(nodes)
axis = dim if values.shape[1] > 1 else 0
self.field.data[nodes.data, dim] = values[:, axis]
self._add_to_indices(dim, nodes)
elif condition[dim] is not None:
raise ValueError("""Wrong condition""")
return