Adding bash script for running test suite

tests/run_tests.sh

@@ -0,0 +1,76 @@
+#!/bin/bash
+
+# Set default parameters
+default_grid_size_lat=1
+default_grid_size_lon=1
+default_run_distributed=false
+
+# Parse command line arguments
+while [[ $# -gt 0 ]]; do
+    case $1 in
+        -h|--help)
+
+            bold=$(tput bold)
+            normal=$(tput sgr0)
+
+            echo "Runs the torch-harmonics test suite."
+            echo "${bold}Arguments:${normal}"
+            echo "  ${bold}-h   | --help:${normal} Prints this text."
+            echo "  ${bold}-d   | --run_distributed:${normal} Run the distributed test suite."
+            echo "  ${bold}-lat | --grid_size_lat:${normal} Number of ranks in latitudinal direction for distributed case."
+            echo "  ${bold}-lon | --grid_size_lon:${normal} Number of ranks in longitudinal direction for distributed case."
+
+            shift
+            exit 0
+            ;;
+        -lat|--grid_size_lat)
+            grid_size_lat="$2"
+            shift 2
+            ;;
+        -lon|--grid_size_lon)
+            grid_size_lon="$2"
+            shift 2
+            ;;
+        -d|--run_distributed)
+            run_distributed=true
+            shift
+            ;;
+        *)
+            echo "Unknown argument: $1"
+            exit 1
+            ;;
+    esac
+done
+
+# Use default values if arguments were not provided
+grid_size_lat=${grid_size_lat:-$default_grid_size_lat}
+grid_size_lon=${grid_size_lon:-$default_grid_size_lon}
+run_distributed=${run_distributed:-$default_run_distributed}
+
+echo "Running sequential tests:"
+python3 -m pytest tests/test_convolution.py tests/test_sht.py
+
+# Run distributed tests if requested
+if [ "$run_distributed" = "true" ]; then
+
+    echo "Running distributed tests with the following parameters:"
+    echo "Grid size latitude: $grid_size_lat"
+    echo "Grid size longitude: $grid_size_lon"
+
+    ngpu=$(( ${grid_size_lat} * ${grid_size_lon} ))
+
+    mpirun --allow-run-as-root -np ${ngpu} bash -c "
+        export CUDA_LAUNCH_BLOCKING=1;
+        export WORLD_RANK=\${OMPI_COMM_WORLD_RANK};
+        export WORLD_SIZE=\${OMPI_COMM_WORLD_SIZE};
+        export RANK=\${OMPI_COMM_WORLD_RANK};
+        export MASTER_ADDR=localhost;
+        export MASTER_PORT=29501;
+        export GRID_H=${grid_size_lat};
+        export GRID_W=${grid_size_lon};
+        python3 -m pytest tests/test_distributed_sht.py
+        python3 -m pytest tests/test_distributed_convolution.py
+        "
+else
+    echo "Skipping distributed tests."
+fi