Fixes to COMPASS to support conda MPI

testing_and_setup/compass/README_ocean.md

@@ -6,11 +6,11 @@ To set up and run ocean test cases from COMPASS, you will need a conda
 environment.  First, install Miniconda3 (if miniconda is not already
 installed), then create a new conda environment as follows:
 ``` bash
-conda create -n compass_0.1.2 -c conda-forge -c e3sm python=3.7 compass=0.1.2
+conda create -n compass_0.1.3 -c conda-forge -c e3sm python=3.7 compass=0.1.3
 ```
 Each time you want to work with COMPASS, you will need to run:
 ```
-conda activate compass_0.1.2
+conda activate compass_0.1.3
 ```
 
 An appropriate conda environment is already available on Los Alamos National

testing_and_setup/compass/load_compass_env.bash → testing_and_setup/compass/load_compass_env.sh

@@ -1,4 +1,4 @@
-version=0.1.2
+version=0.1.3
 
 # The rest of the script should not need to be modified
 if [[ $HOSTNAME = "cori"* ]] || [[ $HOSTNAME = "dtn"* ]]; then

testing_and_setup/compass/ocean/baroclinic_channel/10km/rpe_test/config_mpas_mesh.xml

@@ -1,7 +1,7 @@
 <?xml version="1.0"?>
 <config case="mpas_mesh">
 	<run_script name="run.py">
-		<step executable="planar_hex">
+		<step executable="planar_hex" conda_mpi="true">
 			<argument flag="--nx">16</argument>
 			<argument flag="--ny">50</argument>
 			<argument flag="--dc">10000.0</argument>

testing_and_setup/compass/ocean/baroclinic_channel/1km/rpe_test/config_mpas_mesh.xml

@@ -1,7 +1,7 @@
 <?xml version="1.0"?>
 <config case="mpas_mesh">
 	<run_script name="run.py">
-		<step executable="planar_hex">
+		<step executable="planar_hex" conda_mpi="true">
 			<argument flag="--nx">160</argument>
 			<argument flag="--ny">500</argument>
 			<argument flag="--dc">1000.0</argument>

testing_and_setup/compass/ocean/baroclinic_channel/4km/rpe_test/config_mpas_mesh.xml

@@ -1,7 +1,7 @@
 <?xml version="1.0"?>
 <config case="mpas_mesh">
 	<run_script name="run.py">
-		<step executable="planar_hex">
+		<step executable="planar_hex" conda_mpi="true">
 			<argument flag="--nx">40</argument>
 			<argument flag="--ny">126</argument>
 			<argument flag="--dc">4000.0</argument>

testing_and_setup/compass/ocean/coastal/Maine/init/define_vertical_grid.py

@@ -0,0 +1 @@
+../../../scripts/vertical_grid/define_vertical_grid_64_layers.py

testing_and_setup/compass/ocean/dam_break/default/004m/config_init1.xml

@@ -42,7 +42,7 @@
 		</stream>
 	</streams>
 	<run_script name="run.py">
-		<step executable="planar_hex">
+		<step executable="planar_hex" conda_mpi="true">
 			<argument flag="--nx">325</argument>
 			<argument flag="--ny">700</argument>
 			<argument flag="--dc">0.04</argument>

testing_and_setup/compass/ocean/dam_break/default/012m/config_init1.xml

@@ -41,7 +41,7 @@
 		</stream>
 	</streams>
 	<run_script name="run.py">
-		<step executable="planar_hex">
+		<step executable="planar_hex" conda_mpi="true">
 			<argument flag="--nx">108</argument>
 			<argument flag="--ny">232</argument>
 			<argument flag="--dc">0.12</argument>

testing_and_setup/compass/ocean/drying_slope/meshes/1km/config_init.xml

@@ -1,7 +1,7 @@
 <?xml version="1.0"?>
 <config case="build_1km_mesh">
 	<run_script name="run.py">
-		<step executable="planar_hex">
+		<step executable="planar_hex" conda_mpi="true">
 			<argument flag="--nx">6</argument>
 			<argument flag="--ny">28</argument>
 			<argument flag="--dc">1000.0</argument>

testing_and_setup/compass/ocean/drying_slope/meshes/250m/config_init.xml

@@ -1,7 +1,7 @@
 <?xml version="1.0"?>
 <config case="build_250m_mesh">
 	<run_script name="run.py">
-		<step executable="planar_hex">
+		<step executable="planar_hex" conda_mpi="true">
 			<argument flag="--nx">6</argument>
 			<argument flag="--ny">114</argument>
 			<argument flag="--dc">250.0</argument>

testing_and_setup/compass/ocean/global_ocean/QU60/init/define_vertical_grid.py

@@ -0,0 +1 @@
+../../../scripts/vertical_grid/define_vertical_grid_64_layers.py

testing_and_setup/compass/ocean/global_ocean/SOQU60to15/init/define_vertical_grid.py

@@ -0,0 +1 @@
+../../../scripts/vertical_grid/define_vertical_grid_64_layers.py

testing_and_setup/compass/ocean/global_ocean/scripts/create_E3SM_coupling_files.py

@@ -21,9 +21,10 @@
 from mpas_tools.ocean.moc import make_moc_basins_and_transects
 from mpas_tools.io import write_netcdf
 import mpas_tools.conversion
+import pyproj
 from pyremap import MpasMeshDescriptor, ProjectionGridDescriptor, Remapper, \
     get_lat_lon_descriptor, get_polar_descriptor
-import pyproj
+from mpas_tools.scrip.from_mpas import scrip_from_mpas
 # }}}
 
 
@@ -248,25 +249,19 @@ def scrip(config):  # {{{
         nomaskStr = ''
 
     # create links
-    make_link('../mesh.nc', '{}.nc'.format(mesh_name))
+    mesh_file = '{}.nc'.format(mesh_name)
+    make_link('../mesh.nc', mesh_file)
 
     # command line execution
     scrip_file = 'ocean.{}{}.scrip.{}.nc'.format(mesh_name, nomaskStr,
                                                  date_string)
 
-    args = ['create_SCRIP_file_from_MPAS_mesh.py',
-            '-m', '{}.nc'.format(mesh_name),
-            '-s', scrip_file]
-    run_command(args)
+    scrip_from_mpas(mesh_file, scrip_file)
 
     if ice_shelf_cavities:
         scrip_file_mask = 'ocean.{}.mask.scrip.{}.nc'.format(mesh_name,
                                                              date_string)
-        args = ['create_SCRIP_file_from_MPAS_mesh.py',
-                '-m', mesh_name + '.nc',
-                '-s', scrip_file_mask,
-                '--landice']
-        run_command(args)
+        scrip_from_mpas(mesh_file, scrip_file, useLandIceMask=True)
 
     # make links to output directories
     os.chdir('../assembled_files_for_upload/inputdata/ocn/mpas-o/{}'.format(
@@ -391,30 +386,38 @@ def mapping(config, atm_scrip_tag):  # {{{
     atm_scrip_file = '{}.nc'.format(atm_scrip_tag)
     make_link('{}/{}'.format(atm_scrip_path, atm_scrip_file), atm_scrip_file)
 
+    if 'CONDA_PREFIX' not in os.environ:
+        raise ValueError('A COMPASS conda environment needs to be loaded.')
+    mpirun_path = '{}/bin/mpirun'.format(os.environ['CONDA_PREFIX'])
+    if os.path.exists(mpirun_path):
+        prefix = [mpirun_path, '-n', nprocs]
+    else:
+        prefix = list()
+
     for method, short in [['conserve', 'aave'], ['bilinear', 'blin'],
                           ['patch', 'patc']]:
 
         # Ocean to atmosphere
         mapping_file = 'map_{}{}_TO_{}_{}.{}.nc'.format(
             mesh_name, nomaskStr, atm_scrip_tag, short, date_string)
-        args = ['mpirun', '-n', nprocs, 'ESMF_RegridWeightGen',
+        args = ['ESMF_RegridWeightGen',
                 '--method', method,
                 '--source', ocn_scrip_file,
                 '--destination', atm_scrip_file,
                 '--weight', mapping_file,
                 '--ignore_unmapped']
-        run_command(args)
+        run_command(prefix + args)
 
         # Atmosphere to ocean
         mapping_file = 'map_{}_TO_{}{}_{}.{}.nc'.format(
             atm_scrip_tag, mesh_name, nomaskStr, short, date_string)
-        args = ['mpirun', '-n', nprocs, 'ESMF_RegridWeightGen',
+        args = ['ESMF_RegridWeightGen',
                 '--method', method,
                 '--source', atm_scrip_file,
                 '--destination', ocn_scrip_file,
                 '--weight', mapping_file,
                 '--ignore_unmapped']
-        run_command(args)
+        run_command(prefix + args)
 
     if ice_shelf_cavities:
         print("\n Mapping files with masks for ice shelf cavities")
@@ -429,24 +432,25 @@ def mapping(config, atm_scrip_tag):  # {{{
             # Ocean to atmosphere
             mapping_file = 'map_{}.mask_TO_{}_{}.{}.nc'.format(
                 mesh_name, atm_scrip_tag, short, date_string)
-            args = ['mpirun', '-n', nprocs, 'ESMF_RegridWeightGen',
+            args = ['ESMF_RegridWeightGen',
                     '--method', method,
                     '--source', ocn_scrip_file,
                     '--destination', atm_scrip_file,
                     '--weight', mapping_file,
                     '--ignore_unmapped']
-            run_command(args)
+            run_command(prefix + args)
 
             # Atmosphere to ocean
             mapping_file = 'map_{}_TO_{}.mask_{}.{}.nc'.format(
                 atm_scrip_tag, mesh_name, short, date_string)
-            args = ['mpirun', '-n', nprocs, 'ESMF_RegridWeightGen',
+            args = ['ESMF_RegridWeightGen',
                     '--method', method,
                     '--source', atm_scrip_file,
                     '--destination', ocn_scrip_file,
                     '--weight', mapping_file,
                     '--ignore_unmapped']
-            run_command(args)
+            run_command(prefix + args)
+
 # }}}
 
 
@@ -603,6 +607,14 @@ def salinity_restoring(config):  # {{{
                                                      date_string)
     make_link('../scrip/{}'.format(ocn_scrip_file), ocn_scrip_file)
 
+    if 'CONDA_PREFIX' not in os.environ:
+        raise ValueError('A COMPASS conda environment needs to be loaded.')
+    mpirun_path = '{}/bin/mpirun'.format(os.environ['CONDA_PREFIX'])
+    if os.path.exists(mpirun_path):
+        prefix = [mpirun_path, '-n', nprocs]
+    else:
+        prefix = list()
+
     # execute commands
     salinity_restoring_output_file = \
         'sss.PHC2_monthlyClimatology.{}.{}.nc'.format(mesh_name, date_string)
@@ -612,13 +624,13 @@ def salinity_restoring(config):  # {{{
         # mapping file, 1x1 to ocean mesh
         map_Levitus_file = 'map_Levitus_1x1_TO_{}_{}.{}.nc'.format(
             mesh_name, short, date_string)
-        args = ['mpirun', '-n', nprocs, 'ESMF_RegridWeightGen',
+        args = ['ESMF_RegridWeightGen',
                 '--method', method,
                 '--source', 'grid_Levitus_1x1_scrip_file.nc',
                 '--destination', ocn_scrip_file,
                 '--weight', map_Levitus_file,
                 '--ignore_unmapped']
-        run_command(args)
+        run_command(prefix + args)
 
     # remap from 1x1 to model grid
     args = ['ncremap',
@@ -1001,7 +1013,8 @@ def make_analysis_lat_lon_map(config, mesh_name):
     remapper = Remapper(inDescriptor, outDescriptor, mappingFileName)
 
     mpiTasks = config.getint('main', 'nprocs')
-    remapper.build_mapping_file(method='bilinear', mpiTasks=mpiTasks)
+    remapper.build_mapping_file(method='bilinear', mpiTasks=mpiTasks,
+                                tempdir='.')
     # }}}
 
 
@@ -1033,7 +1046,8 @@ def make_analysis_polar_map(config, mesh_name, projection):
     remapper = Remapper(inDescriptor, outDescriptor, mappingFileName)
 
     mpiTasks = config.getint('main', 'nprocs')
-    remapper.build_mapping_file(method='bilinear', mpiTasks=mpiTasks)
+    remapper.build_mapping_file(method='bilinear', mpiTasks=mpiTasks,
+                                tempdir='.')
     # }}}
 
 

testing_and_setup/compass/ocean/global_ocean/scripts/cull_mesh.py

@@ -25,7 +25,6 @@
     unicode_literals
 
 import os
-import subprocess
 from optparse import OptionParser
 import xarray
 
@@ -34,6 +33,7 @@
 from mpas_tools.io import write_netcdf
 from mpas_tools.ocean.coastline_alteration import widen_transect_edge_masks, \
     add_critical_land_blockages, add_land_locked_cells_to_mask
+from mpas_tools.viz.paraview_extractor import extract_vtk
 
 
 parser = OptionParser()
@@ -158,20 +158,13 @@
     write_netcdf(dsLandIceCulledMesh, 'no_ISC_culled_mesh.nc',
                  format=netcdfFormat)
 
-args = ['paraview_vtk_field_extractor.py',
-        '--ignore_time',
-        '-d', 'maxEdges=',
-        '-v', 'allOnCells',
-        '-f', 'culled_mesh.nc',
-        '-o', 'culled_mesh_vtk']
-print("running", ' '.join(args))
-subprocess.check_call(args, env=os.environ.copy())
-
-args = ['paraview_vtk_field_extractor.py',
-        '--ignore_time',
-        '-d', 'maxEdges=',
-        '-v', 'allOnCells',
-        '-f', 'no_ISC_culled_mesh.nc',
-        '-o', 'no_ISC_culled_mesh_vtk']
-print("running", ' '.join(args))
-subprocess.check_call(args, env=os.environ.copy())
+extract_vtk(ignore_time=True, dimension_list=['maxEdges='],
+            variable_list=['allOnCells'], filename_pattern='culled_mesh.nc',
+            out_dir='culled_mesh_vtk')
+
+if options.with_cavities:
+    extract_vtk(ignore_time=True, dimension_list=['maxEdges='],
+                variable_list=['allOnCells'],
+                filename_pattern='no_ISC_culled_mesh.nc',
+                out_dir='no_ISC_culled_mesh_vtk')
+

testing_and_setup/compass/ocean/internal_waves/5km/rpe_test/config_mpas_mesh.xml

@@ -1,7 +1,7 @@
 <?xml version="1.0"?>
 <config case="mpas_mesh">
 	<run_script name="run.py">
-		<step executable="planar_hex">
+		<step executable="planar_hex" conda_mpi="true">
 			<argument flag="--nx">4</argument>
 			<argument flag="--ny">50</argument>
 			<argument flag="--dc">5000.0</argument>

testing_and_setup/compass/ocean/jigsaw_to_MPAS/build_mesh.py

@@ -12,14 +12,14 @@
 from __future__ import absolute_import, division, print_function, \
     unicode_literals
 
-import subprocess
 import os
 import xarray
 import argparse
 import matplotlib.pyplot as plt
 
 from mpas_tools.conversion import convert
 from mpas_tools.io import write_netcdf
+from mpas_tools.viz.paraview_extractor import extract_vtk
 
 from jigsaw_to_MPAS.jigsaw_driver import jigsaw_driver
 from jigsaw_to_MPAS.triangle_jigsaw_to_netcdf import jigsaw_to_netcdf
@@ -121,15 +121,9 @@ def build_mesh(
                                    floodplain_elevation=floodplain_elevation)
 
     print('Step 8. Create vtk file for visualization')
-    args = ['paraview_vtk_field_extractor.py',
-            '--ignore_time',
-            '-l',
-            '-d', 'maxEdges=0',
-            '-v', 'allOnCells',
-            '-f', 'base_mesh.nc',
-            '-o', 'base_mesh_vtk']
-    print("running", ' '.join(args))
-    subprocess.check_call(args, env=os.environ.copy())
+    extract_vtk(ignore_time=True, lonlat=True, dimension_list=['maxEdges='],
+                variable_list=['allOnCells'], filename_pattern='base_mesh.nc',
+                out_dir='base_mesh_vtk')
 
     print("***********************************************")
     print("**    The global mesh file is base_mesh.nc   **")

testing_and_setup/compass/ocean/lock_exchange/0.5km/default/config_init1.xml

@@ -47,7 +47,7 @@ v		<option name="config_vert_levels">1</option>
 	</streams>
 
 	<run_script name="run.py">
-		<step executable="planar_hex">
+		<step executable="planar_hex" conda_mpi="true">
 			<argument flag="--nx">6</argument>
 			<argument flag="--ny">30</argument>
 			<argument flag="--dc">1000.0</argument>

testing_and_setup/compass/ocean/lock_exchange/0.5km/rpe_test/config_mpas_mesh.xml

@@ -1,7 +1,7 @@
 <?xml version="1.0"?>
 <config case="mpas_mesh">
 	<run_script name="run.py">
-		<step executable="planar_hex">
+		<step executable="planar_hex" conda_mpi="true">
 			<argument flag="--nx">4</argument>
 			<argument flag="--ny">128</argument>
 			<argument flag="--dc">500.0</argument>

testing_and_setup/compass/ocean/overflow/1km/rpe_test/config_mpas_mesh.xml

@@ -1,7 +1,7 @@
 <?xml version="1.0"?>
 <config case="mpas_mesh">
 	<run_script name="run.py">
-		<step executable="planar_hex">
+		<step executable="planar_hex" conda_mpi="true">
 			<argument flag="--nx">4</argument>
 			<argument flag="--ny">200</argument>
 			<argument flag="--dc">1000.0</argument>

testing_and_setup/compass/ocean/surface_waves/direct/1km/config_init1.xml

@@ -54,7 +54,7 @@
 	</streams>
 
 	<run_script name="run.py">
-		<step executable="planar_hex">
+		<step executable="planar_hex" conda_mpi="true">
 			<argument flag="--nx">6</argument>
 			<argument flag="--ny">30</argument>
 			<argument flag="--dc">1000.0</argument>