import os
import time
from pyBadlands.remote import RemoteModel
from pyBadlands.model import Model
start_time = time.time()
model = Model()
model.load_xml('bench/input.xml')
model.run_to_time(4000)
print 'singlethread finished in %s seconds' % (time.time() - start_time)
for ncpus in range(1, 5):
start_time = time.time()
remote = RemoteModel(maxcpus=ncpus)
# match the child's cwd to ours
remote._view.execute('import os; os.chdir("%s")' % os.getcwd())
remote.load_xml('bench/input.xml')
remote.run_to_time(4000)
print 'mpi (%d cores) finished in %s seconds' % (ncpus, time.time() - start_time)
class SPM(object):
def __init__(self,
mesh,
velocityField,
swarm,
materialField,
airIndex,
sedimentIndex,
XML,
resolution,
checkpoint_interval,
surfElevation=0.,
verbose=True,
restartFolder=None,
restartStep=None):
self.SECONDS_PER_YEAR = 31556925.9747 # Tropical year in seconds
self.verbose = verbose
self.restartStep = restartStep
self.restartFolder = restartFolder
# AutoScaling
self.scaleDIM = 1.0 / scaling_coefficients["[length]"].magnitude
self.scaleTIME = 1.0 / scaling_coefficients["[time]"].magnitude
self.mesh = mesh
self.velocityField = velocityField
self.swarm = swarm
self.material_index = materialField
self.airIndex = airIndex
self.sedimentIndex = sedimentIndex
self.resolution = resolution
self.surfElevation = surfElevation
self.checkpoint_interval = checkpoint_interval / self.scaleTIME / self.SECONDS_PER_YEAR
self.XML = XML
if rank == 0:
self.badlands_model = BadlandsModel()
self.badlands_model.load_xml(self.XML)
if self.restartStep:
self.badlands_model.input.restart = True
self.badlands_model.input.rstep = self.restartStep
self.badlands_model.input.rfolder = self.restartFolder
self.badlands_model.input.outDir = self.restartFolder
self.badlands_model.outputStep = self.restartStep
self.minCoord = self.mesh.minCoord
self.maxCoord = self.mesh.maxCoord
self.time_years = 0.
if self.restartStep:
# Parse xmf for the last timestep time
import xml.etree.ElementTree as etree
xmf = self.restartFolder + "/xmf/tin.time" + str(
self.restartStep) + ".xmf"
tree = etree.parse(xmf)
root = tree.getroot()
self.time_years = float(root[0][0][0].attrib["Value"])
self._tmp = _tempdir
self._demfile = self._tmp + "/dem.csv"
# Create Initial Flat DEM
if rank == 0:
self.dem = self._generate_flat_dem(self.minCoord, self.maxCoord,
self.resolution,
self.surfElevation,
self.scaleDIM)
np.savetxt(self._demfile, self.dem)
# Build Mesh
self.badlands_model.build_mesh(self._demfile, verbose=False)
self.badlands_model.input.disp3d = True # enable 3D displacements
self.badlands_model.input.region = 0 # TODO: check what this does
self.badlands_model.input.tStart = self.time_years
self.badlands_model.tNow = self.time_years
# Override the checkpoint/display interval in the Badlands model to
# ensure BL and UW are synced
self.badlands_model.input.tDisplay = self.checkpoint_interval
# Set Badlands minimal distance between nodes before regridding
self.badlands_model.force.merge3d = self.badlands_model.input.Afactor * self.badlands_model.recGrid.resEdges * 0.5
# Bodge Badlands to perform an initial checkpoint
# FIXME: we need to run the model for at least one iteration before this is generated. It would be nice if this wasn't the case.
self.badlands_model.force.next_display = 0
comm.Barrier()
self._disp_inserted = False
# Transfer the initial DEM state to Underworld
self._update_material_types()
comm.Barrier()
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()
dt_years = Dimensionalize(dt, UnitRegistry.years).magnitude
if rank == 0:
rg = self.badlands_model.recGrid
if self.mesh.dim == 2:
zVals = rg.regZ.mean(axis=1)
np_surface = np.column_stack((rg.regX, zVals)) * self.scaleDIM
if self.mesh.dim == 3:
np_surface = np.column_stack(
(rg.rectX, rg.rectY, rg.rectZ)) * self.scaleDIM
else:
np_surface = None
np_surface = comm.bcast(np_surface, root=0)
comm.Barrier()
# Get Velocity Field at the surface
tracer_velocity_mps = get_UW_velocities(
np_surface, self.velocityField) * self.scaleTIME / self.scaleDIM
if rank == 0:
# Use the tracer vertical velocities to deform the Badlands TIN
# convert from meters per second to meters displacement over the whole iteration
tracer_disp = tracer_velocity_mps * self.SECONDS_PER_YEAR * dt_years
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 _determine_particle_state_2D(self):
if rank == 0:
known_xy = self.badlands_model.recGrid.tinMesh[
'vertices'] * self.scaleDIM # points that we have known elevation for
known_z = self.badlands_model.elevation * self.scaleDIM # elevation for those points
xs = self.badlands_model.recGrid.regX * self.scaleDIM
ys = self.badlands_model.recGrid.regY * self.scaleDIM
else:
known_xy = None
known_z = None
xs = None
ys = None
known_xy = comm.bcast(known_xy, root=0)
known_z = comm.bcast(known_z, root=0)
xs = comm.bcast(xs, root=0)
ys = comm.bcast(ys, root=0)
comm.Barrier()
grid_x, grid_y = np.meshgrid(xs, ys)
interpolate_z = griddata(known_xy,
known_z, (grid_x, grid_y),
method='nearest').T
interpolate_z = interpolate_z.mean(axis=1)
f = interp1d(xs, interpolate_z)
uw_surface = self.swarm.particleCoordinates.data
bdl_surface = f(uw_surface[:, 0])
flags = uw_surface[:, 1] < bdl_surface
return flags
def _determine_particle_state(self):
# Given Badlands' mesh, determine if each particle in 'volume' is above
# (False) or below (True) it.
# To do this, for each X/Y pair in 'volume', we interpolate its Z value
# relative to the mesh in blModel. Then, if the interpolated Z is
# greater than the supplied Z (i.e. Badlands mesh is above particle
# elevation) it's sediment (True). Else, it's air (False).
# TODO: we only support air/sediment layers right now; erodibility
# layers are not implemented
if rank == 0:
known_xy = self.badlands_model.recGrid.tinMesh[
'vertices'] * self.scaleDIM # points that we have known elevation for
known_z = self.badlands_model.elevation * self.scaleDIM # elevation for those points
else:
known_xy = None
known_z = None
known_xy = comm.bcast(known_xy, root=0)
known_z = comm.bcast(known_z, root=0)
comm.Barrier()
volume = self.swarm.particleCoordinates.data
interpolate_xy = volume[:, [0, 1]]
# NOTE: we're using nearest neighbour interpolation. This should be
# sufficient as Badlands will normally run at a much higher resolution
# than Underworld. 'linear' interpolation is much, much slower.
interpolate_z = griddata(points=known_xy,
values=known_z,
xi=interpolate_xy,
method='nearest')
# True for sediment, False for air
flags = volume[:, 2] < interpolate_z
return flags
def _update_material_types(self):
# What do the materials (in air/sediment terms) look like now?
if self.mesh.dim == 3:
material_flags = self._determine_particle_state()
if self.mesh.dim == 2:
material_flags = self._determine_particle_state_2D()
# If any materials changed state, update the Underworld material types
mi = self.material_index.data
# convert air to sediment
for air_material in self.airIndex:
sedimented_mask = np.logical_and(np.in1d(mi, air_material),
material_flags)
mi[sedimented_mask] = self.sedimentIndex
# convert sediment to air
for air_material in self.airIndex:
eroded_mask = np.logical_and(~np.in1d(mi, air_material),
~material_flags)
mi[eroded_mask] = self.airIndex[0]
def _inject_badlands_displacement(self, time, dt, disp, sigma):
"""
Takes a plane of tracer points and their DISPLACEMENTS in 3D over time
period dt applies a gaussian filter on it. Injects it into Badlands as 3D
tectonic movement.
"""
# The Badlands 3D interpolation map is the displacement of each DEM
# node at the end of the time period relative to its starting position.
# If you start a new displacement file, it is treated as starting at
# the DEM starting points (and interpolated onto the TIN as it was at
# that tNow).
# kludge; don't keep adding new entries
if self._disp_inserted:
self.badlands_model.force.T_disp[0, 0] = time
self.badlands_model.force.T_disp[0, 1] = (time + dt)
else:
self.badlands_model.force.T_disp = np.vstack(
([time, time + dt], self.badlands_model.force.T_disp))
self._disp_inserted = True
# Extent the velocity field in the third dimension
if self.mesh.dim == 2:
dispX = np.tile(disp[:, 0], self.badlands_model.recGrid.rny)
dispY = np.zeros((self.badlands_model.recGrid.rnx *
self.badlands_model.recGrid.rny, ))
dispZ = np.tile(disp[:, 1], self.badlands_model.recGrid.rny)
disp = np.zeros((self.badlands_model.recGrid.rnx *
self.badlands_model.recGrid.rny, 3))
disp[:, 0] = dispX
disp[:, 1] = dispY
disp[:, 2] = dispZ
# Gaussian smoothing
if sigma > 0:
dispX = np.copy(disp[:,
0]).reshape(self.badlands_model.recGrid.rnx,
self.badlands_model.recGrid.rny)
dispY = np.copy(disp[:,
1]).reshape(self.badlands_model.recGrid.rnx,
self.badlands_model.recGrid.rny)
dispZ = np.copy(disp[:,
2]).reshape(self.badlands_model.recGrid.rnx,
self.badlands_model.recGrid.rny)
smoothX = gaussian_filter(dispX, sigma)
smoothY = gaussian_filter(dispY, sigma)
smoothZ = gaussian_filter(dispZ, sigma)
disp[:, 0] = smoothX.flatten()
disp[:, 1] = smoothY.flatten()
disp[:, 2] = smoothZ.flatten()
self.badlands_model.force.injected_disps = disp
def _generate_flat_dem(self,
minCoord,
maxCoord,
resolution,
elevation,
scale=1.):
"""
Generate a flat DEM. This can be used as the initial Badlands state.
minCoord: tuple of (X, Y, Z) coordinates defining the minimum bounds of
the DEM. Only the X and Y coordinates are used.
maxCoord: tuple of (X, Y, Z) coordinates defining the maximum bounds of
the DEM. Only the X and Y coordinates are used.
resolution: resolution of the model. Badlands assumes dx = dy.
elevation: the Z parameter that each point is created at
scale: the scaling factor between the 2 codes
For your convenience, minCoord and maxCoord are designed to have the
same formatting as the Underworld FeMesh_Cartesian minCoord and
maxCoord parameters.
IMPORTANT: minCoord and maxCoord are defined in terms of the Underworld
coordinate system, but the returned DEM uses the Badlands coordinate
system.
Note that the initial elevation of the Badlands surface should coincide
with the material transition in Underworld.
"""
# Calculate number of nodes from required resolution.
nx = np.int((maxCoord[0] - minCoord[0]) / resolution) + 1
ny = np.int((maxCoord[1] - minCoord[1]) / resolution) + 1
if self.mesh.dim == 2:
minCoord = (minCoord[0], minCoord[0])
maxCoord = (maxCoord[0], maxCoord[0])
ny = nx
items = []
# FIXME: there should be a fast numpy way to do this
for y in np.linspace(minCoord[1] / scale, maxCoord[1] / scale, ny):
for x in np.linspace(minCoord[0] / scale, maxCoord[0] / scale, nx):
items.append([x, y, elevation / scale])
# NOTE: Badlands uses the difference in X coord of the first two points to determine the resolution.
# This is something we should fix.
# This is why we loop in y/x order instead of x/y order.
return np.array(items)
def __init__(self,
mesh,
velocityField,
swarm,
materialField,
airIndex,
sedimentIndex,
XML,
resolution,
checkpoint_interval,
surfElevation=0.,
verbose=True,
restartFolder=None,
restartStep=None):
self.SECONDS_PER_YEAR = 31556925.9747 # Tropical year in seconds
self.verbose = verbose
self.restartStep = restartStep
self.restartFolder = restartFolder
# AutoScaling
self.scaleDIM = 1.0 / scaling_coefficients["[length]"].magnitude
self.scaleTIME = 1.0 / scaling_coefficients["[time]"].magnitude
self.mesh = mesh
self.velocityField = velocityField
self.swarm = swarm
self.material_index = materialField
self.airIndex = airIndex
self.sedimentIndex = sedimentIndex
self.resolution = resolution
self.surfElevation = surfElevation
self.checkpoint_interval = checkpoint_interval / self.scaleTIME / self.SECONDS_PER_YEAR
self.XML = XML
if rank == 0:
self.badlands_model = BadlandsModel()
self.badlands_model.load_xml(self.XML)
if self.restartStep:
self.badlands_model.input.restart = True
self.badlands_model.input.rstep = self.restartStep
self.badlands_model.input.rfolder = self.restartFolder
self.badlands_model.input.outDir = self.restartFolder
self.badlands_model.outputStep = self.restartStep
self.minCoord = self.mesh.minCoord
self.maxCoord = self.mesh.maxCoord
self.time_years = 0.
if self.restartStep:
# Parse xmf for the last timestep time
import xml.etree.ElementTree as etree
xmf = self.restartFolder + "/xmf/tin.time" + str(
self.restartStep) + ".xmf"
tree = etree.parse(xmf)
root = tree.getroot()
self.time_years = float(root[0][0][0].attrib["Value"])
self._tmp = _tempdir
self._demfile = self._tmp + "/dem.csv"
# Create Initial Flat DEM
if rank == 0:
self.dem = self._generate_flat_dem(self.minCoord, self.maxCoord,
self.resolution,
self.surfElevation,
self.scaleDIM)
np.savetxt(self._demfile, self.dem)
# Build Mesh
self.badlands_model.build_mesh(self._demfile, verbose=False)
self.badlands_model.input.disp3d = True # enable 3D displacements
self.badlands_model.input.region = 0 # TODO: check what this does
self.badlands_model.input.tStart = self.time_years
self.badlands_model.tNow = self.time_years
# Override the checkpoint/display interval in the Badlands model to
# ensure BL and UW are synced
self.badlands_model.input.tDisplay = self.checkpoint_interval
# Set Badlands minimal distance between nodes before regridding
self.badlands_model.force.merge3d = self.badlands_model.input.Afactor * self.badlands_model.recGrid.resEdges * 0.5
# Bodge Badlands to perform an initial checkpoint
# FIXME: we need to run the model for at least one iteration before this is generated. It would be nice if this wasn't the case.
self.badlands_model.force.next_display = 0
comm.Barrier()
self._disp_inserted = False
# Transfer the initial DEM state to Underworld
self._update_material_types()
comm.Barrier()
# parameters for our model composition
MIN_COORD = (0., 0., -80e3)
MAX_COORD = (100e3, 70e3, 20e3
) # this crashes at DEM load? do we have X/Y swapped?
# MAX_COORD = (100e3, 100e3, 20e3)
INITIAL_AIR_ELEVATION = 0.0 # the height at which we transition from sediment to air
# Put a rectangular prism underground
cubeSize = (10e3, 20e3, 20e3)
centre = (50e3, 25e3, -30e3)
linkage = LinkageModel()
### SET UP THE BADLANDS MODEL
badlands_model = BadlandsModel()
badlands_model.load_xml('no_erosion_clock.xml')
linkage.badlands_model = badlands_model
# IMPORTANT: disable BL->UW transfers
# Because erosion is disabled in the Badlands config, this completely uncouples
# the two models. They independently advect the surface. In this way, we can
# ensure that they advect identically.
linkage.disable_material_changes = True
# Load a flat DEM
bl_dem = linkage.generate_flat_dem(minCoord=MIN_COORD,
maxCoord=MAX_COORD,
resolution=(180, 180),
elevation=INITIAL_AIR_ELEVATION)
linkage.load_badlands_dem_array(bl_dem)
# dome.buildGrid(elevation=dome.Z, nameCSV='dem')
import numpy
import os
import pandas
from pyBadlands.model import Model as BadlandsModel
import underworld as uw
from underworld import function as fn
from linkagemodel.linkage import LinkageModel
linkage = LinkageModel()
### SET UP THE BADLANDS MODEL
badlands_model = BadlandsModel()
badlands_model.load_xml('badlands.xml')
linkage.badlands_model = badlands_model
# We're going to load from a DEM; let's just load it to determine the Underworld mesh size
dem = pandas.read_csv('dem.csv', sep=' ', header=None, na_filter=False, dtype=numpy.float, low_memory=False)
# set up min/max coord for UW
min_coord = [dem[0].min(), dem[1].min(), -1e3]
max_coord = [dem[0].max(), dem[1].max(), +1e3]
print 'The model has bounds %s to %s' % (min_coord, max_coord)
### SET UP THE UNDERWORLD MODEL
def _init_model(self):
if self.minCoord:
self.minCoord = tuple([nd(val) for val in self.minCoord])
else:
self.minCoord = self.Model.mesh.minCoord
if self.maxCoord:
self.maxCoord = tuple([nd(val) for val in self.maxCoord])
else:
self.maxCoord = self.Model.mesh.maxCoord
if self.Model.mesh.dim == 2:
self.minCoord = (self.minCoord[0],
self.aspectRatio2d * self.minCoord[0])
self.maxCoord = (self.maxCoord[0],
self.aspectRatio2d * self.maxCoord[0])
if rank == 0:
from pyBadlands.model import Model as BadlandsModel
self.badlands_model = BadlandsModel()
self.badlands_model.load_xml(self.XML)
if self.restartStep:
self.badlands_model.input.restart = True
self.badlands_model.input.rstep = self.restartStep
self.badlands_model.input.rfolder = self.restartFolder
self.badlands_model.input.outDir = self.restartFolder
self.badlands_model.outputStep = self.restartStep
# Parse xmf for the last timestep time
import xml.etree.ElementTree as etree
xmf = (self.restartFolder + "/xmf/tin.time" +
str(self.restartStep) + ".xmf")
tree = etree.parse(xmf)
root = tree.getroot()
self.time_years = float(root[0][0][0].attrib["Value"])
# Create Initial DEM
self._demfile = _tempdir + "/dem.csv"
self.dem = self._generate_dem()
np.savetxt(self._demfile, self.dem)
# Build Mesh
self.badlands_model.build_mesh(self._demfile, verbose=False)
self.badlands_model.input.outDir = self.outputDir
self.badlands_model.input.disp3d = True # enable 3D displacements
self.badlands_model.input.region = 0 # TODO: check what this does
self.badlands_model.input.tStart = self.time_years
self.badlands_model.tNow = self.time_years
# Override the checkpoint/display interval in the Badlands model to
# ensure BL and UW are synced
self.badlands_model.input.tDisplay = (dimensionalise(
self.checkpoint_interval, u.years).magnitude)
# Set Badlands minimal distance between nodes before regridding
self.badlands_model.force.merge3d = (
self.badlands_model.input.Afactor *
self.badlands_model.recGrid.resEdges * 0.5)
# Bodge Badlands to perform an initial checkpoint
# FIXME: we need to run the model for at least one
# iteration before this is generated.
# It would be nice if this wasn't the case.
self.badlands_model.force.next_display = 0
comm.Barrier()
self._disp_inserted = False
# Transfer the initial DEM state to Underworld
self._update_material_types()
comm.Barrier()
#
# To run this on an MPI cluster, use something like:
#
# mpiexec -n 4 python launch_mpi.py 1000
import os
import re
import sys
import time
from pyBadlands.model import Model
base_path = '.'
xml_name = 'hrtest.xml'
run_years = int(sys.argv[1])
# Modify base_path to run from another directory
# Change into HRtest directory if required...
os.chdir(base_path)
start_time = time.time()
model = Model()
print 'loading %s' % xml_name
print model.load_xml(xml_name)
print model.run_to_time(run_years)
print 'run to %s years finished in %s seconds' % (run_years,
time.time() - start_time)
from linkagemodel.linkage import LinkageModel
# We're going to load from a DEM; let's just load it to determine the Underworld mesh size
dem = pandas.read_csv('nodes.csv', sep=' ', header=None, na_filter=False, dtype=numpy.float, low_memory=False)
# set up min/max coord for UW
min_coord = [dem[0].min(), dem[1].min(), -20e3]
max_coord = [dem[0].max(), dem[1].max(), +20e3]
print 'The model has bounds %s to %s' % (min_coord, max_coord)
linkage = LinkageModel()
### SET UP THE BADLANDS MODEL
badlands_model = BadlandsModel()
badlands_model.load_xml('dem.xml') # The XML file specifies the DEM to load
linkage.badlands_model = badlands_model
### SET UP THE UNDERWORLD MODEL
# All output will go to the 'uwout' directory, which we will create
uw_output_path = 'uwout'
try:
os.mkdir(uw_output_path)
except OSError:
# probably already exists
pass
UNDERWORLD_RESOLUTION = 20
class Badlands(SurfaceProcesses):
""" A wrapper class for Badlands """
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_model(self):
if self.minCoord:
self.minCoord = tuple([nd(val) for val in self.minCoord])
else:
self.minCoord = self.Model.mesh.minCoord
if self.maxCoord:
self.maxCoord = tuple([nd(val) for val in self.maxCoord])
else:
self.maxCoord = self.Model.mesh.maxCoord
if self.Model.mesh.dim == 2:
self.minCoord = (self.minCoord[0],
self.aspectRatio2d * self.minCoord[0])
self.maxCoord = (self.maxCoord[0],
self.aspectRatio2d * self.maxCoord[0])
if rank == 0:
from pyBadlands.model import Model as BadlandsModel
self.badlands_model = BadlandsModel()
self.badlands_model.load_xml(self.XML)
if self.restartStep:
self.badlands_model.input.restart = True
self.badlands_model.input.rstep = self.restartStep
self.badlands_model.input.rfolder = self.restartFolder
self.badlands_model.input.outDir = self.restartFolder
self.badlands_model.outputStep = self.restartStep
# Parse xmf for the last timestep time
import xml.etree.ElementTree as etree
xmf = (self.restartFolder + "/xmf/tin.time" +
str(self.restartStep) + ".xmf")
tree = etree.parse(xmf)
root = tree.getroot()
self.time_years = float(root[0][0][0].attrib["Value"])
# Create Initial DEM
self._demfile = _tempdir + "/dem.csv"
self.dem = self._generate_dem()
np.savetxt(self._demfile, self.dem)
# Build Mesh
self.badlands_model.build_mesh(self._demfile, verbose=False)
self.badlands_model.input.outDir = self.outputDir
self.badlands_model.input.disp3d = True # enable 3D displacements
self.badlands_model.input.region = 0 # TODO: check what this does
self.badlands_model.input.tStart = self.time_years
self.badlands_model.tNow = self.time_years
# Override the checkpoint/display interval in the Badlands model to
# ensure BL and UW are synced
self.badlands_model.input.tDisplay = (dimensionalise(
self.checkpoint_interval, u.years).magnitude)
# Set Badlands minimal distance between nodes before regridding
self.badlands_model.force.merge3d = (
self.badlands_model.input.Afactor *
self.badlands_model.recGrid.resEdges * 0.5)
# Bodge Badlands to perform an initial checkpoint
# FIXME: we need to run the model for at least one
# iteration before this is generated.
# It would be nice if this wasn't the case.
self.badlands_model.force.next_display = 0
comm.Barrier()
self._disp_inserted = False
# Transfer the initial DEM state to Underworld
self._update_material_types()
comm.Barrier()
def _generate_dem(self):
"""
Generate a badlands DEM. This can be used as the initial Badlands state.
"""
# Calculate number of nodes from required resolution.
nx = np.int((self.maxCoord[0] - self.minCoord[0]) / self.resolution)
ny = np.int((self.maxCoord[1] - self.minCoord[1]) / self.resolution)
nx += 1
ny += 1
x = np.linspace(self.minCoord[0], self.maxCoord[0], nx)
y = np.linspace(self.minCoord[1], self.maxCoord[1], ny)
coordsX, coordsY = np.meshgrid(x, y)
dem = np.zeros((nx * ny, 3))
dem[:, 0] = coordsX.flatten()
dem[:, 1] = coordsY.flatten()
coordsZ = self.surfElevation.evaluate(dem[:, :2])
dem[:, 2] = coordsZ.flatten()
return dimensionalise(dem, u.meter).magnitude
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 _determine_particle_state_2D(self):
known_xy = None
known_z = None
xs = None
ys = None
fact = dimensionalise(1.0, u.meter).magnitude
if rank == 0:
# points that we have known elevation for
known_xy = self.badlands_model.recGrid.tinMesh['vertices'] / fact
# elevation for those points
known_z = self.badlands_model.elevation / fact
xs = self.badlands_model.recGrid.regX / fact
ys = self.badlands_model.recGrid.regY / fact
known_xy = comm.bcast(known_xy, root=0)
known_z = comm.bcast(known_z, root=0)
xs = comm.bcast(xs, root=0)
ys = comm.bcast(ys, root=0)
comm.Barrier()
grid_x, grid_y = np.meshgrid(xs, ys)
interpolate_z = griddata(known_xy,
known_z, (grid_x, grid_y),
method='nearest').T
interpolate_z = interpolate_z.mean(axis=1)
f = interp1d(xs, interpolate_z)
uw_surface = self.Model.swarm.particleCoordinates.data
bdl_surface = f(uw_surface[:, 0])
flags = uw_surface[:, 1] < bdl_surface
return flags
def _determine_particle_state(self):
# Given Badlands' mesh, determine if each particle in 'volume' is above
# (False) or below (True) it.
# To do this, for each X/Y pair in 'volume', we interpolate its Z value
# relative to the mesh in blModel. Then, if the interpolated Z is
# greater than the supplied Z (i.e. Badlands mesh is above particle
# elevation) it's sediment (True). Else, it's air (False).
# TODO: we only support air/sediment layers right now; erodibility
# layers are not implemented
known_xy = None
known_z = None
fact = dimensionalise(1.0, u.meter).magnitude
if rank == 0:
# points that we have known elevation for
known_xy = self.badlands_model.recGrid.tinMesh['vertices'] / fact
known_z = self.badlands_model.elevation / fact
known_xy = comm.bcast(known_xy, root=0)
known_z = comm.bcast(known_z, root=0)
comm.Barrier()
volume = self.Model.swarm.particleCoordinates.data
interpolate_xy = volume[:, [0, 1]]
# NOTE: we're using nearest neighbour interpolation. This should be
# sufficient as Badlands will normally run at a much higher resolution
# than Underworld. 'linear' interpolation is much, much slower.
interpolate_z = griddata(points=known_xy,
values=known_z,
xi=interpolate_xy,
method='nearest')
# True for sediment, False for air
flags = volume[:, 2] < interpolate_z
return flags
def _update_material_types(self):
# What do the materials (in air/sediment terms) look like now?
if self.Model.mesh.dim == 3:
material_flags = self._determine_particle_state()
if self.Model.mesh.dim == 2:
material_flags = self._determine_particle_state_2D()
# If any materials changed state, update the Underworld material types
mi = self.Model.materialField.data
# convert air to sediment
for air_material in self.airIndex:
sedimented_mask = np.logical_and(np.in1d(mi, air_material),
material_flags)
mi[sedimented_mask] = self.sedimentIndex
# convert sediment to air
for air_material in self.airIndex:
eroded_mask = np.logical_and(~np.in1d(mi, air_material),
~material_flags)
mi[eroded_mask] = self.airIndex[0]
def _inject_badlands_displacement(self, time, dt, disp, sigma):
"""
Takes a plane of tracer points and their DISPLACEMENTS in 3D over time
period dt applies a gaussian filter on it. Injects it into Badlands as 3D
tectonic movement.
"""
# The Badlands 3D interpolation map is the displacement of each DEM
# node at the end of the time period relative to its starting position.
# If you start a new displacement file, it is treated as starting at
# the DEM starting points (and interpolated onto the TIN as it was at
# that tNow).
# kludge; don't keep adding new entries
if self._disp_inserted:
self.badlands_model.force.T_disp[0, 0] = time
self.badlands_model.force.T_disp[0, 1] = (time + dt)
else:
self.badlands_model.force.T_disp = np.vstack(
([time, time + dt], self.badlands_model.force.T_disp))
self._disp_inserted = True
# Extent the velocity field in the third dimension
if self.Model.mesh.dim == 2:
dispX = np.tile(disp[:, 0], self.badlands_model.recGrid.rny)
dispY = np.zeros((self.badlands_model.recGrid.rnx *
self.badlands_model.recGrid.rny, ))
dispZ = np.tile(disp[:, 1], self.badlands_model.recGrid.rny)
disp = np.zeros((self.badlands_model.recGrid.rnx *
self.badlands_model.recGrid.rny, 3))
disp[:, 0] = dispX
disp[:, 1] = dispY
disp[:, 2] = dispZ
# Gaussian smoothing
if sigma > 0:
dispX = np.copy(disp[:,
0]).reshape(self.badlands_model.recGrid.rnx,
self.badlands_model.recGrid.rny)
dispY = np.copy(disp[:,
1]).reshape(self.badlands_model.recGrid.rnx,
self.badlands_model.recGrid.rny)
dispZ = np.copy(disp[:,
2]).reshape(self.badlands_model.recGrid.rnx,
self.badlands_model.recGrid.rny)
smoothX = gaussian_filter(dispX, sigma)
smoothY = gaussian_filter(dispY, sigma)
smoothZ = gaussian_filter(dispZ, sigma)
disp[:, 0] = smoothX.flatten()
disp[:, 1] = smoothY.flatten()
disp[:, 2] = smoothZ.flatten()
self.badlands_model.force.injected_disps = disp
# # To run this on an MPI cluster, use something like: # # mpiexec -n 4 python mpi_example.py 10000 import os import re import sys import time from pyBadlands.model import Model base_path = 'crater' xml_name = 'crater.xml' run_years = int(sys.argv[1]) # change into crater data directory os.chdir(base_path) start_time = time.time() model = Model() print 'loading %s' % xml_name print model.load_xml(xml_name) print model.run_to_time(run_years) print 'run to %s years finished in %s seconds' % (run_years, time.time() - start_time)
conditions=[yBCs])
nodeSolve = uw.systems.Solver(nodeDiffuse)
nodeSolve.solve()
with mesh.deform_mesh():
mesh.data[:,1] = yField.data[:,0]
# If the above is used then it must be repeated for the
# 'mesh0' created further on in this script
########################################################
'''
### SET UP THE BADLANDS MODEL
badlands_model = BadlandsModel()
badlands_model.load_xml('airlayer.xml')
linkage.badlands_model = badlands_model
linkage.material_map = [
[erodedIndex, airIndex],
[sedimentIndex, heavyIndex, lightIndex]
]
dem = linkage.generate_flat_dem(minCoord=MIN_COORD, maxCoord=MAX_COORD, resolution=BADLANDS_RESOLUTION, elevation=AIR_ELEVATION)
print 'dem shape %s' % (dem.shape,)
linkage.load_badlands_dem_array(dem)
# Create the swarm, material index variable and swarm advector
swarm = uw.swarm.Swarm(mesh=mesh)
import os
import time
from pyBadlands.remote import RemoteModel
from pyBadlands.model import Model
start_time = time.time()
model = Model()
model.load_xml('bench/input.xml')
model.run_to_time(4000)
print 'singlethread finished in %s seconds' % (time.time() - start_time)
for ncpus in range(1, 5):
start_time = time.time()
remote = RemoteModel(maxcpus=ncpus)
# match the child's cwd to ours
remote._view.execute('import os; os.chdir("%s")' % os.getcwd())
remote.load_xml('bench/input.xml')
remote.run_to_time(4000)
print 'mpi (%d cores) finished in %s seconds' % (ncpus,
time.time() - start_time)
def __init__(self,
badlands_config,
underworld_resolution,
materials,
min_coords=None,
max_coords=None,
badlands_resolution=None,
surface_elevation=None,
elev_range=None):
self.linkage = LinkageModel()
# Set up the Badlands model
badlands_model = BadlandsModel()
badlands_model.load_xml(badlands_config)
self.linkage.badlands_model = badlands_model
# Load the Badlands config to find out if we're using a DEM file
tree = ET.parse(badlands_config)
root = tree.getroot()
# Is a demfile loaded already through the badlands config?
try:
# This will throw an exception if not defined
demfile = root.find('grid').find('demfile').text
# The DEM sets the bounds of the model
dem = pandas.read_csv(demfile,
sep=' ',
header=None,
na_filter=False,
dtype=numpy.float,
low_memory=False)
min_coords = [dem[0].min(), dem[1].min(), elev_range[0]]
max_coords = [dem[0].max(), dem[1].max(), elev_range[1]]
except (ValueError, AttributeError):
assert min_coords is not None and max_coords is not None and badlands_resolution is not None and surface_elevation is not None, "If no DEM file is loaded in the Badlands XML config, you must supply the min_coords, max_coords, badlands_resolution and surface_elevation parameters"
dem = self.linkage.generate_flat_dem(
minCoord=min_coords,
maxCoord=max_coords,
resolution=badlands_resolution,
elevation=surface_elevation)
self.linkage.load_badlands_dem_array(dem)
# Configure the linkage material map
# TODO this needs to be updated when we have multiple erodibility layers in Badlands
mm = [[], []] # we assume air and sediment layers
for item in materials:
mm[item['bl_layer']].append(item['uw_index'])
self.linkage.material_map = mm
### SET UP THE UNDERWORLD MODEL
mesh = uw.mesh.FeMesh_Cartesian(elementType=("Q1/dQ0"),
elementRes=underworld_resolution,
minCoord=min_coords,
maxCoord=max_coords)
self.linkage.mesh = mesh
# We want to track velocity and pressure.
velocityField = uw.mesh.MeshVariable(mesh=mesh, nodeDofCount=mesh.dim)
self.linkage.velocity_field = velocityField
pressureField = uw.mesh.MeshVariable(mesh=mesh.subMesh, nodeDofCount=1)
self.pressure_field = pressureField
# Set initial states
velocityField.data[:] = [0., 0., 0.]
pressureField.data[:] = 0.
# Create the swarm, material index variable and swarm advector
swarm = uw.swarm.Swarm(mesh=mesh)
self.linkage.swarm = swarm
materialIndex = swarm.add_variable(dataType="int", count=1)
self.linkage.material_index = materialIndex
swarmLayout = uw.swarm.layouts.GlobalSpaceFillerLayout(
swarm=swarm, particlesPerCell=20)
swarm.populate_using_layout(layout=swarmLayout)
self.advector = uw.systems.SwarmAdvector(swarm=swarm,
velocityField=velocityField,
order=2)
# Set viscosities and densities of the model.
viscosityMapFn = 1e19
density_map = {}
for item in materials:
density_map[item['uw_index']] = item['density']
densityFn = fn.branching.map(fn_key=materialIndex, mapping=density_map)
# And the final buoyancy force function.
buoyancyFn = densityFn * 9.8 * [0.0, 0.0, -1.0]
# wall velocity boundary conditions - free slip on all walls
iWalls = mesh.specialSets["MinI_VertexSet"] + mesh.specialSets[
"MaxI_VertexSet"]
jWalls = mesh.specialSets["MinJ_VertexSet"] + mesh.specialSets[
"MaxJ_VertexSet"]
kWalls = mesh.specialSets["MinK_VertexSet"] + mesh.specialSets[
"MaxK_VertexSet"]
velocityBC = uw.conditions.DirichletCondition(variable=velocityField,
indexSetsPerDof=(iWalls,
jWalls,
kWalls))
# combine all the above into Stokes system and get solver
stokesPIC = uw.systems.Stokes(velocityField=velocityField,
pressureField=pressureField,
voronoi_swarm=swarm,
conditions=[
velocityBC,
],
fn_viscosity=viscosityMapFn,
fn_bodyforce=buoyancyFn)
self.solver = uw.systems.Solver(stokesPIC)
# FINISH SETTING UP LINKAGE
self.linkage.update_function = self.update_function
self.linkage.checkpoint_function = self.checkpoint_function
# stuff that gets exported to user
self.swarm = swarm
self.particle_coordinates = swarm.particleCoordinates
self.material_index = materialIndex
### SET UP THE OUTPUT DIRECTORY
self.out_dir = badlands_model.input.outDir
# Store the SimpleMaterialModel instance for later
self.linkage.smm = self
SPHERE_RADIUS = 12e3
# Make sure the sphere starts underground. Badlands sets the initial elevation, so if the sphere pokes through the surface, it will be clipped during model initialisation.
SPHERE_CENTRE = (50e3, 50e3, -18e3)
# Material types
airIndex = 0
heavyIndex = 1
lightIndex = 2
sedimentIndex = 3
erodedIndex = 4
linkage = LinkageModel()
### SET UP THE BADLANDS MODEL
badlands_model = BadlandsModel()
badlands_model.load_xml('rising_ball.xml')
linkage.badlands_model = badlands_model
# Override the DEM with a flat one. This defines how the Badlands surface looks
# at the start.
dem = linkage.generate_flat_dem(minCoord=MIN_COORD, maxCoord=MAX_COORD, resolution=(180, 180), elevation=INITIAL_AIR_ELEVATION)
# The linkage needs to know which Underworld material types correspond to air or sediment in Badlands' worldview.
linkage.material_map = [
[erodedIndex, airIndex], # air layer; eroded material has its own material type
[sedimentIndex, heavyIndex, lightIndex] # sediment layer; sediment has its own material
]
linkage.load_badlands_dem_array(dem)
