From 1bd3cb85e8e01484e45779ca173f9b2343a7410c Mon Sep 17 00:00:00 2001
From: bmooremaley <ben.mooremaley@gmail.com>
Date: Mon, 18 Nov 2024 10:06:37 -0600
Subject: [PATCH] Preserved ESMF option for generating weights files from a
 latlon source grid in remap_topography

---
 compass/ocean/mesh/remap_topography.cfg | 13 +++--
 compass/ocean/mesh/remap_topography.py  | 65 +++++++++++++++++++++----
 2 files changed, 65 insertions(+), 13 deletions(-)

diff --git a/compass/ocean/mesh/remap_topography.cfg b/compass/ocean/mesh/remap_topography.cfg
index ba7a6ff81..a3711336c 100644
--- a/compass/ocean/mesh/remap_topography.cfg
+++ b/compass/ocean/mesh/remap_topography.cfg
@@ -2,9 +2,15 @@
 [remap_topography]
 
 # the name of the topography file in the bathymetry database
+#topo_filename = BedMachineAntarctica_v3_and_GEBCO_2023_ne3000_20241004.nc
+#src_scrip_filename = ne3000_20241004.scrip.nc
 topo_filename = BedMachineAntarctica_v3_and_GEBCO_2023_ne3000_20241004.nc
 src_scrip_filename = ne3000_20241004.scrip.nc
 
+# weight generator function:
+#    `tempest` for cubed-sphere bathy or `esmf` for latlon bathy
+weight_generator = tempest
+
 # variable names in topo_filename
 lon_var = lon
 lat_var = lat
@@ -20,10 +26,11 @@ description = Bathymetry is from GEBCO 2023, combined with BedMachine
               Antarctica v3 around Antarctica.
 
 # the target and minimum number of MPI tasks to use in remapping
-ntasks = 4096
-min_tasks = 360
+ntasks = 1280
+min_tasks = 256
 
 # remapping method {'bilinear', 'neareststod', 'conserve'}
+# must use 'conserve' for tempestremap
 method = conserve
 
 # threshold of what fraction of an MPAS cell must contain ocean in order to
@@ -34,7 +41,7 @@ renorm_threshold = 0.01
 ice_density = 910.0
 
 # smoothing parameters
-# no smoothing:
+# no smoothing (required for esmf):
 #     expandDist = 0   [m]
 #     expandFactor = 1 [cell fraction]
 expandDist = 0
diff --git a/compass/ocean/mesh/remap_topography.py b/compass/ocean/mesh/remap_topography.py
index 9c290ab91..d061aa618 100644
--- a/compass/ocean/mesh/remap_topography.py
+++ b/compass/ocean/mesh/remap_topography.py
@@ -107,10 +107,19 @@ def run(self):
         """
         Run this step of the test case
         """
+        config = self.config
+        weight_generator = config.get('remap_topography', 'weight_generator')
+
         self._create_target_scrip_file()
-        self._partition_scrip_file('source.scrip.nc')
-        self._partition_scrip_file('target.scrip.nc')
-        self._create_weights()
+        if weight_generator == 'tempest':
+            self._partition_scrip_file('source.scrip.nc')
+            self._partition_scrip_file('target.scrip.nc')
+            self._create_weights_tempest()
+        elif weight_generator == 'esmf':
+            self._create_weights_esmf()
+        else:
+            msg = f'Unsupported weight generator function {weight_generator}'
+            raise ValueError(msg)
         self._remap_to_target()
         self._modify_remapped_bathymetry()
 
@@ -145,7 +154,7 @@ def _partition_scrip_file(self, in_filename):
         logger = self.logger
         logger.info('Partition SCRIP file')
 
-        # Convert to 64-bit NetCDF
+        # Convert to NetCDF3 64-bit
         args = [
             'ncks', '-5',
             in_filename,
@@ -172,21 +181,54 @@ def _partition_scrip_file(self, in_filename):
 
         logger.info('  Done.')
 
-    def _create_weights(self):
+    def _create_weights_tempest(self):
         """
-        Create mapping weights file using mbtempest
+        Create mapping weights file using TempestRemap
         """
         logger = self.logger
         logger.info('Create weights file')
 
+        config = self.config
+        method = config.get('remap_topography', 'method')
+        if method != 'conserve':
+            raise ValueError(f'Unsupported method {method} for TempestRemap')
+
         args = [
             'mbtempest', '--type', '5',
             '--load', f'source.scrip.64bit.p{self.ntasks}.h5m',
             '--load', f'target.scrip.64bit.p{self.ntasks}.h5m',
-            '--file', f'mapfv_source_to_target.nomask_{self.ntasks}_gnom.nc',
-            '--intx', 'moab_intx_source_target.h5m',
-            '--weights', '--verbose', '--gnomonic', '--boxeps', '1e-9',
+            '--file', f'map_source_to_target_{method}.nc',
+            '--weights', '--gnomonic',
+            '--boxeps', '1e-9',
+        ]
+
+        run_command(
+            args, self.cpus_per_task, self.ntasks,
+            self.openmp_threads, self.config, self.logger,
+        )
+
+        logger.info('  Done.')
+
+    def _create_weights_esmf(self):
+        """
+        Create mapping weights file using ESMF_RegridWeightGen
+        """
+        logger = self.logger
+        logger.info('Create weights file')
+
+        config = self.config
+        method = config.get('remap_topography', 'method')
+
+        args = [
+            'ESMF_RegridWeightGen',
+            '--source', 'source.scrip.nc',
+            '--destination', 'target.scrip.nc',
+            '--weight', f'map_source_to_target_{method}.nc',
+            '--method', method,
+            '--netcdf4',
+            '--ignore_unmapped',
         ]
+
         run_command(
             args, self.cpus_per_task, self.ntasks,
             self.openmp_threads, self.config, self.logger,
@@ -201,13 +243,16 @@ def _remap_to_target(self):
         logger = self.logger
         logger.info('Remap to target')
 
+        config = self.config
+        method = config.get('remap_topography', 'method')
+
         # Build command args
         # Unused options:
         # -P mpas, handles some MPAS-specific index ordering, CF, etc...
         # -C climatology, basically bypasses fill values
         args = [
             'ncremap',
-            '-m', f'mapfv_source_to_target.nomask_{self.ntasks}_gnom.nc',
+            '-m', f'map_source_to_target_{method}.nc',
             '--vrb=1',
             'topography.nc', 'topography_ncremap.nc',
         ]