Add hooks for PyTorch velocity emulator

components/mpas-albany-landice/src/Registry.xml

@@ -65,8 +65,16 @@
 
 	<nml_record name="velocity_solver" in_defaults="true">
 		<nml_option name="config_velocity_solver" type="character" default_value="sia" units="unitless"
-		            description="Selection of the method for solving ice velocity. 'L1L2', 'FO', and 'Stokes' require compiling with external dycores. 'none' skips the calculation of velocity so the velocity field will be 0 or set to a field read from an input file.  'simple' gives a simple prescribed velocity field computed at initialization."
-		            possible_values="'sia', 'L1L2', 'FO', 'Stokes', 'simple', 'none'"
+		            description="Selection of the method for solving ice velocity. 'L1L2', 'FO', and 'Stokes' require compiling with external dycores. 'gnn' uses a GNN emulator via FTorch. 'none' skips the calculation of velocity so the velocity field will be 0 or set to a field read from an input file.  'simple' gives a simple prescribed velocity field computed at initialization."
+		            possible_values="'sia', 'L1L2', 'FO', 'Stokes', 'gnn', 'simple', 'none'"
+		/>
+		<nml_option name="config_gnn_model_path" type="character" default_value="gnn_velocity_model.pt" units="unitless"
+		            description="Path to the TorchScript GNN model file for the 'gnn' velocity solver."
+		            possible_values="Any valid file path."
+		/>
+		<nml_option name="config_gnn_device_type" type="character" default_value="cpu" units="unitless"
+		            description="Device type for GNN inference. 'cpu' or 'cuda'."
+		            possible_values="'cpu', 'cuda'"
 		/>
 		<nml_option name="config_sia_tangent_slope_calculation" type="character" default_value="from_vertex_barycentric" units="unitless"
 		            description="Selection of the method for calculating the tangent component of surface slope at edges needed by the SIA velocity solver.

components/mpas-albany-landice/src/build_options.mk

@@ -50,6 +50,27 @@ ifeq "$(SLM)" "true"
     override CPPFLAGS += -DUSE_SEALEVELMODEL
 endif
 
+# Optional FTorch for GNN velocity emulator
+ifeq "$(FTORCH)" "true"
+ifndef FTORCH_ROOT
+$(error FTORCH_ROOT is not set.  Please set FTORCH_ROOT to the FTorch install directory when FTORCH=true)
+endif
+ifndef LIBTORCH_ROOT
+$(error LIBTORCH_ROOT is not set.  Please set LIBTORCH_ROOT to the LibTorch install directory when FTORCH=true)
+endif
+	override CPPFLAGS += -DUSE_FTORCH
+	FCINCLUDES += -I$(FTORCH_ROOT)/include/ftorch
+	override LDFLAGS += -Wl,-rpath,$(FTORCH_ROOT)/lib64 -Wl,-rpath,$(LIBTORCH_ROOT)/lib
+ifeq "$(FTORCH_CUDA)" "true"
+ifndef CUDA_HOME
+$(error CUDA_HOME is not set.  Please set CUDA_HOME to the CUDA toolkit directory when FTORCH_CUDA=true)
+endif
+	LIBS += -Wl,--as-needed,--allow-shlib-undefined -L$(FTORCH_ROOT)/lib64 -L$(LIBTORCH_ROOT)/lib -Wl,-rpath-link,$(CUDA_HOME)/lib64 -Wl,-rpath-link,$(LIBTORCH_ROOT)/lib -lftorch -ltorch -ltorch_cpu -ltorch_cuda -lc10 -lc10_cuda -lgomp -lstdc++ -Wl,--no-as-needed
+else
+	LIBS += -Wl,--as-needed -L$(FTORCH_ROOT)/lib64 -L$(LIBTORCH_ROOT)/lib -lftorch -ltorch -ltorch_cpu -lc10 -lgomp -lstdc++ -Wl,--no-as-needed
+endif
+endif
+
 # ===================================
 report_builds:
 	@echo "CORE=landice"

components/mpas-albany-landice/src/mode_forward/Makefile

@@ -17,6 +17,7 @@ OBJS =  mpas_li_core.o \
     mpas_li_sia.o \
     mpas_li_velocity_simple.o \
     mpas_li_velocity_external.o \
+    mpas_li_velocity_gnn.o \
     mpas_li_subglacial_hydro.o \
     mpas_li_bedtopo.o \
     mpas_li_ocean_extrap.o
@@ -68,6 +69,7 @@ mpas_li_diagnostic_vars.o: mpas_li_thermal.o
 mpas_li_velocity.o: mpas_li_sia.o \
                     mpas_li_velocity_simple.o \
                     mpas_li_velocity_external.o \
+                    mpas_li_velocity_gnn.o \
 		    mpas_li_advection.o \
 		    mpas_li_thermal.o
 
@@ -77,6 +79,8 @@ mpas_li_subglacial_hydro.o: mpas_li_diagnostic_vars.o
 
 mpas_li_velocity_simple.o:
 
+mpas_li_velocity_gnn.o:
+
 mpas_li_statistics.o: mpas_li_diagnostic_vars.o
 
 mpas_li_velocity_external.o: Interface_velocity_solver.o

components/mpas-albany-landice/src/mode_forward/mpas_li_velocity.F

@@ -30,6 +30,7 @@ module li_velocity
 
    use li_velocity_external
    use li_velocity_simple
+   use li_velocity_gnn
    use li_sia
    use li_setup
 
@@ -120,6 +121,8 @@ subroutine li_velocity_init(domain, err)
           call li_sia_init(domain, err)
       case ('simple')
           call li_velocity_simple_init(domain, err)
+      case ('gnn')
+          call li_velocity_gnn_init(domain, err)
       case ('L1L2', 'FO', 'Stokes')
           call li_velocity_external_init(domain, err)
       case default
@@ -193,6 +196,8 @@ subroutine li_velocity_block_init(block, err)
           call li_sia_block_init(block, err)
       case ('simple')
           call li_velocity_simple_block_init(block, err)
+      case ('gnn')
+          ! No block-level init needed for GNN solver
       case ('L1L2', 'FO', 'Stokes')
           call li_velocity_external_block_init(block, err)
       case default
@@ -518,6 +523,10 @@ subroutine li_velocity_solve(domain, solveVelo, updateMask, err)
                endif
             enddo
 
+        case('gnn')
+
+            call li_velocity_gnn_solve(domain, err_tmp)
+
         case default
 
              call mpas_log_write(trim(config_velocity_solver) // ' is not a valid land ice velocity solver option.', MPAS_LOG_ERR)
@@ -750,6 +759,8 @@ subroutine li_velocity_finalize(domain, err)
           call li_sia_finalize(domain, err)
       case ('simple')
           call li_velocity_simple_finalize(err)
+      case ('gnn')
+          call li_velocity_gnn_finalize(err)
       case ('L1L2', 'FO', 'Stokes')
           call li_velocity_external_finalize(err)
       case default