from netCDF4 import Dataset
import os
import math
import errno
import numpy as np
import subprocess
import argparse
import shutil
import glob
from .regrid import regrid_to_other_mesh
from .mask import extend_seaice_mask
from .regions import make_regions_file
[docs]
def gen_seaice_mesh_partition(meshFilename, regionFilename, nProcsArray,
mpasCullerLocation, outputPrefix, plotting,
metis, cullEquatorialRegion):
"""
Generate graph partition(s) for the given MPAS-Seaice mesh and the given
number(s) of processors and a file defining regions that each processor
should own part of (typically a polar region and an equatorial region)
Parameters
----------
meshFilename : str
The name of a file containing the MPAS-Seaice mesh
regionFilename : str
The name of a file containing a ``region`` field defining different
regions that each processor should own part of
nProcsArray : list or int
The number(s) of processors to create graph partitions for
mpasCullerLocation : str or None
The directory for the ``MpasCellCuller.x`` tool or ``None`` to look in
the user's path
outputPrefix : str
The prefix to prepend to each graph partition file
plotting : bool
Whether to create a NetCDF file ``partition_diag.nc`` to use for
plotting the partitions
metis : str
The exectable to use for partitioning in each region
cullEquatorialRegion : bool
Whether to remove the equatorial region from the paritions
"""
# arguments
meshToolsDir = mpasCullerLocation
if meshToolsDir is None:
culler = shutil.which("MpasCellCuller.x")
if culler is not None:
meshToolsDir = os.path.dirname(culler)
else:
# no directory was provided and none
this_dir = os.path.dirname(os.path.realpath(__file__))
meshToolsDir = os.path.abspath(os.path.join(
this_dir, "..", "..", "..", "mesh_tools",
"mesh_conversion_tools"))
culler = os.path.join(meshToolsDir, "MpasCellCuller.x")
if not os.path.exists(culler):
raise FileNotFoundError(
"MpasCellCuller.x does not exist at the requested location.")
plotFilename = "partition_diag.nc"
# get regions
regionFile = Dataset(regionFilename, "r")
nRegions = regionFile.nRegions
region = regionFile.variables["region"][:]
regionFile.close()
# diagnostics
if plotting:
shutil.copyfile(meshFilename, plotFilename)
# load mesh file
mesh = Dataset(meshFilename, "r")
nCells = len(mesh.dimensions["nCells"])
mesh.close()
cellidsInRegion = []
for iRegion in range(0, nRegions):
# tmp file basename
tmp = "%s_%2.2i_tmp" % (meshFilename, iRegion)
# create precull file
tmpFilenamesPrecull = tmp + "_precull.nc"
shutil.copyfile(meshFilename, tmpFilenamesPrecull)
# make cullCell variable
cullCell = np.ones(nCells)
for iCell in range(0, nCells):
if region[iCell] == iRegion:
cullCell[iCell] = 0
# cull the mesh
tmpFilenamesPostcull = tmp + "_postcull.nc"
_cull_mesh(meshToolsDir, tmpFilenamesPrecull, tmpFilenamesPostcull,
cullCell)
# get the cell IDs for this partition
cellid = _get_cell_ids(tmpFilenamesPostcull, meshFilename)
cellidsInRegion.append(cellid)
# preserve the initial graph file
os.rename("culled_graph.info", f"culled_graph_{iRegion}_tmp.info")
if not isinstance(nProcsArray, (list, tuple, set)):
# presumably, it's a single integer
nProcsArray = [nProcsArray]
for nProcs in nProcsArray:
if cullEquatorialRegion:
nBlocks = nRegions * nProcs
else:
nBlocks = nProcs
combinedGraph = np.zeros(nCells)
for iRegion in range(0, nRegions):
# partition the culled grid
try:
graphFilename = "culled_graph_%i_tmp.info" % iRegion
subprocess.call([metis, graphFilename, str(nProcs)])
except OSError as e:
if e.errno == errno.ENOENT:
raise FileNotFoundError(
"metis program %s not found" % metis)
else:
print("metis error")
raise
cellid = cellidsInRegion[iRegion]
# load this partition
graph = _load_partition(
"culled_graph_%i_tmp.info.part.%i" %
(iRegion, nProcs))
# add this partition to the combined partition
for iCellPartition in range(0, len(graph)):
if cullEquatorialRegion:
combinedGraph[cellid[iCellPartition]] = \
graph[iCellPartition] + nProcs * iRegion
else:
combinedGraph[cellid[iCellPartition]] = \
graph[iCellPartition]
# output the cell partition file
cellPartitionFile = open("%s.part.%i" % (outputPrefix, nBlocks), "w")
for iCell in range(0, nCells):
cellPartitionFile.write("%i\n" % (combinedGraph[iCell]))
cellPartitionFile.close()
# output block partition file
if cullEquatorialRegion:
blockPartitionFile = open(
"%s.part.%i" %
(outputPrefix, nProcs), "w")
for iRegion in range(0, nRegions):
for iProc in range(0, nProcs):
blockPartitionFile.write("%i\n" % iProc)
blockPartitionFile.close()
# diagnostics
if plotting:
plottingFile = Dataset(plotFilename, "a")
partitionVariable = plottingFile.createVariable(
"partition_%i" % nProcs, "i4", ("nCells",))
partitionVariable[:] = combinedGraph
plottingFile.close()
subprocess.run("rm *tmp*", shell=True)
[docs]
def prepare_partitions():
"""
An entry point for performing preparatory work for making seaice partitions
"""
# parsing input
parser = argparse.ArgumentParser(
description="Perform preparatory work for making seaice partitions.")
parser.add_argument("-i", "--inputmesh", dest="meshFilenameSrc",
required=True,
help="MPAS mesh file for source regridding mesh.")
parser.add_argument("-p", "--presence", dest="filenameData",
required=True,
help="Input ice presence file for source mesh.")
parser.add_argument("-m", "--outputmesh", dest="meshFilenameDst",
required=True,
help="MPAS mesh file for destination regridding mesh.")
parser.add_argument("-o", "--outputDir", dest="outputDir",
required=True,
help="Output directory for temporary files and "
"partition files.")
parser.add_argument("-w", "--weightsFilename", dest="weightsFilename",
required=False,
help="A mapping file between the input and output "
"MPAS meshes. One will be generated if it is "
"not supplied.")
args = parser.parse_args()
# Check if output directory exists
if not os.path.isdir(args.outputDir):
raise FileNotFoundError("ERROR: Output directory does not exist.")
# 1) Regrid the ice presence from the input data mesh to the grid of choice
print("Regrid to desired mesh...")
filenameOut = args.outputDir + "/icePresent_regrid.nc"
regrid_to_other_mesh(
meshFilenameSrc=args.meshFilenameSrc,
filenameData=args.filenameData,
meshFilenameDst=args.meshFilenameDst,
filenameOut=filenameOut,
generateWeights=(args.weightsFilename is None),
weightsFilename=args.weightsFilename)
# 2) create icePresence variable
print("fix_regrid_output...")
# check executable exists
if shutil.which("fix_regrid_output.exe") is not None:
# it's in the system path
executable = "fix_regrid_output.exe"
elif os.path.exists("./fix_regrid_output.exe"):
# found in local path
executable = "./fix_regrid_output.exe"
else:
raise FileNotFoundError("fix_regrid_output.exe could not be found.")
inputFile = args.outputDir + "/icePresent_regrid.nc"
outputFile = args.outputDir + "/icePresent_regrid_modify.nc"
subprocess.call([executable, inputFile, args.meshFilenameDst, outputFile])
# 3) create variable icePresenceExtended
print("extend_seaice_mask...")
filenamePresence = args.outputDir + "/icePresent_regrid_modify.nc"
extend_seaice_mask(args.meshFilenameDst, filenamePresence, 0.0, False)
# 4) Make the regions file from the icePresenceExtended variable
print("make_regions_file...")
filenameIcePresent = args.outputDir + "/icePresent_regrid_modify.nc"
filenameOut = args.outputDir + "/regions.nc"
make_regions_file(filenameIcePresent=filenameIcePresent,
filenameMesh=args.meshFilenameDst,
regionType="two_region_eq",
varname="icePresenceExtended",
limit=0.5, filenameOut=filenameOut)
[docs]
def create_partitions():
"""
An entry point for creating sea-ice partitions
"""
# parsing input
parser = argparse.ArgumentParser(description='Create sea-ice partitions.')
parser.add_argument(
'-m', '--outputmesh', dest="meshFilename", required=True,
help='MPAS mesh file for destination regridding mesh.')
parser.add_argument(
'-o', '--outputDir', dest="outputDir", required=True,
help='Output directory for temporary files and partition files.')
parser.add_argument(
'-c', '--cullerDir', dest="mpasCullerLocation", required=False,
help='Location of MPAS MpasCellCuller.x executable.')
parser.add_argument(
'-p', '--prefix', dest="outputPrefix", required=False,
help='prefix for output partition filenames.', default="graph.info")
parser.add_argument(
'-x', '--plotting', dest="plotting", required=False,
help='create diagnostic plotting file of partitions',
action='store_true')
parser.add_argument(
'-g', '--metis', dest="metis", required=False,
help='name of metis utility', default="gpmetis")
parser.add_argument(
'-n', '--nProcs', dest="nProcsArray", nargs='*', required=False,
help='list of the number of processors to create partition for.',
type=int)
parser.add_argument(
'-f', '--nProcsFile', dest="nProcsFile", required=False,
help='number of processors to create partition for.')
args = parser.parse_args()
# number of processors
if args.nProcsArray is None and args.nProcsFile is None:
raise ValueError("Must specify nProcs or nProcsFile")
if args.nProcsArray is not None and args.nProcsFile is not None:
raise ValueError("Can't specify both nProcs or nProcsFile")
if args.nProcsFile is not None:
with open(args.nProcsFile, "r") as fileNProcs:
nProcsLines = fileNProcs.readlines()
nProcsArray = [int(line) for line in nProcsLines
if line.split() != '']
else:
nProcsArray = args.nProcsArray
# create partitions
regionFilename = args.outputDir + "/regions.nc"
outputPrefix = args.outputDir + "/" + args.outputPrefix
gen_seaice_mesh_partition(args.meshFilename, regionFilename, nProcsArray,
args.mpasCullerLocation, outputPrefix,
args.plotting, args.metis,
cullEquatorialRegion=False)
def simple_partitions():
"""
An entry point for creating sea-ice partitions on LCRC (Anvil and
Chrysalis)
"""
data_dir = '/lcrc/group/e3sm/public_html/mpas_standalonedata/' \
'mpas-seaice/partition'
# parsing input
parser = argparse.ArgumentParser(
description='Create sea-ice partitions on LCRC.')
parser.add_argument(
'-m', '--mesh', dest="meshFilename", required=True,
help='MPAS-Seaice mesh file.')
parser.add_argument(
'-p', '--prefix', dest="outputPrefix", required=True,
help='prefix for output partition filenames.')
parser.add_argument(
'-n', '--nprocs', dest="nProcsArray", nargs='*', required=True,
help='list of the number of processors to create partition for.',
type=int)
parser.add_argument(
'-d', '--datadir', dest="dataDir", required=False,
default=data_dir,
help='Directory with seaice_QU60km_polar.nc and '
'icePresent_QU60km_polar.nc.')
args = parser.parse_args()
meshFilenameDst = os.path.abspath(args.meshFilename)
tmpdir = 'tmp_seaice_part_dir'
try:
shutil.rmtree(tmpdir)
except FileNotFoundError:
pass
os.makedirs(tmpdir)
cwd = os.getcwd()
os.chdir(tmpdir)
# make a local link to the mesh file
basename = os.path.basename(meshFilenameDst)
command = ['ln', '-s', meshFilenameDst, basename]
subprocess.run(command, check=True)
meshFilenameDst = basename
# 1) Regrid the ice presence from the input data mesh to the grid of choice
print("Regrid to desired mesh...")
filenameOut = "icePresent_regrid.nc"
meshFilenameSrc = os.path.join(args.dataDir, 'seaice_QU60km_polar.nc')
filenameData = os.path.join(args.dataDir, 'icePresent_QU60km_polar.nc')
regrid_to_other_mesh(
meshFilenameSrc=meshFilenameSrc,
filenameData=filenameData,
meshFilenameDst=meshFilenameDst,
filenameOut=filenameOut,
generateWeights=True,
weightsFilename=None)
# 2) create icePresence variable
print("fix_regrid_output...")
inputFile = "icePresent_regrid.nc"
outputFile = "icePresent_regrid_modify.nc"
subprocess.call(["fix_regrid_output.exe", inputFile, meshFilenameDst,
outputFile])
# 3) create variable icePresenceExtended
print("extend_seaice_mask...")
filenamePresence = "icePresent_regrid_modify.nc"
extend_seaice_mask(meshFilenameDst, filenamePresence, 0.0, False)
# 4) Make the regions file from the icePresenceExtended variable
print("make_regions_file...")
filenameIcePresent = "icePresent_regrid_modify.nc"
filenameOut = "regions.nc"
make_regions_file(filenameIcePresent=filenameIcePresent,
filenameMesh=meshFilenameDst,
regionType="two_region_eq",
varname="icePresenceExtended",
limit=0.5,
filenameOut=filenameOut)
nProcsArray = args.nProcsArray
# create partitions
regionFilename = "regions.nc"
outputPrefix = os.path.join(cwd, args.outputPrefix)
gen_seaice_mesh_partition(meshFilename=meshFilenameDst,
regionFilename=regionFilename,
nProcsArray=nProcsArray,
mpasCullerLocation=None,
outputPrefix=outputPrefix,
plotting=False,
metis="gpmetis",
cullEquatorialRegion=False)
for file in glob.glob(f'{outputPrefix}*'):
command = ['chmod', 'ug+rw', file]
subprocess.run(command, check=True)
command = ['chmod', 'o+r', file]
subprocess.run(command, check=True)
os.chdir(cwd)
shutil.rmtree(tmpdir)
# ---------------------------------------------------------------------
# Private functions
# ---------------------------------------------------------------------
def _degree_to_radian(degree):
return (degree * math.pi) / 180.0
def _add_cell_cull_array(filename, cullCell):
mesh = Dataset(filename, "a")
cullCellVariable = mesh.createVariable("cullCell", "i4", ("nCells",))
cullCellVariable[:] = cullCell
mesh.close()
def _cull_mesh(meshToolsDir, filenameIn, filenameOut, cullCell):
_add_cell_cull_array(filenameIn, cullCell)
executable = "MpasCellCuller.x"
if meshToolsDir is not None:
executable = os.path.join(meshToolsDir, executable)
subprocess.run([executable, filenameIn, filenameOut, "-c"], check=True)
def _load_partition(graphFilename):
graphFile = open(graphFilename, "r")
lines = graphFile.readlines()
graphFile.close()
partition = []
for line in lines:
partition.append(int(line))
return partition
def _get_cell_ids(culledFilename, originalFilename):
culledFile = Dataset(culledFilename, "r")
nCellsCulled = len(culledFile.dimensions["nCells"])
originalFile = Dataset(originalFilename, "r")
nCellsOriginal = len(originalFile.dimensions["nCells"])
cellid = np.zeros(nCellsCulled, dtype=int)
cellMapFile = open("cellMapForward.txt", "r")
cellMapLines = cellMapFile.readlines()
cellMapFile.close()
iCellOriginal = 0
for cellMapLine in cellMapLines:
if iCellOriginal % 1000 == 0:
print(iCellOriginal, " of ", nCellsOriginal)
try:
cellMap = int(cellMapLine)
except ValueError:
# There are blank lines to skip
continue
if cellMap != -1:
cellid[cellMap] = iCellOriginal
iCellOriginal = iCellOriginal + 1
return cellid
def _get_cell_ids_orig(culledFilename, originalFilename):
culledFile = Dataset(culledFilename, "r")
latCellCulled = culledFile.variables["latCell"][:]
lonCellCulled = culledFile.variables["lonCell"][:]
nCellsCulled = len(culledFile.dimensions["nCells"])
originalFile = Dataset(originalFilename, "r")
latCellOriginal = originalFile.variables["latCell"][:]
lonCellOriginal = originalFile.variables["lonCell"][:]
nCellsOriginal = len(originalFile.dimensions["nCells"])
cellid = np.zeros(nCellsCulled, dtype=int)
for iCellCulled in range(0, nCellsCulled):
if iCellCulled % 1000 == 0:
print("iCellCulled: ", iCellCulled, "of ", nCellsCulled)
for iCellOriginal in range(0, nCellsOriginal):
if (latCellCulled[iCellCulled] == latCellOriginal[iCellOriginal] and
lonCellCulled[iCellCulled] == lonCellOriginal[iCellOriginal]):
cellid[iCellCulled] = iCellOriginal
break
return cellid