From e9f327b430d1983fa120b778e14090fe1ac911ae Mon Sep 17 00:00:00 2001
From: Henry Addison <henryaddison@users.noreply.github.com>
Date: Mon, 5 Aug 2024 20:43:00 +0100
Subject: [PATCH] helper scripts for running sampling on JASMIN/LOTUS

---
 bin/jasmin/lotus-sampling-wrapper | 21 +++++++++++++++++++++
 bin/jasmin/queue-sampling         | 17 +++++++++++++++++
 2 files changed, 38 insertions(+)
 create mode 100755 bin/jasmin/lotus-sampling-wrapper
 create mode 100755 bin/jasmin/queue-sampling

diff --git a/bin/jasmin/lotus-sampling-wrapper b/bin/jasmin/lotus-sampling-wrapper
new file mode 100755
index 00000000..fbf71bbd
--- /dev/null
+++ b/bin/jasmin/lotus-sampling-wrapper
@@ -0,0 +1,21 @@
+#!/bin/bash
+# Wrapper script around commands for training a model to queue on LOTUS on JASMIN
+
+source ~/.bashrc
+mamba activate mv-mlde
+
+set -euo pipefail
+
+cd /home/users/vf20964/code/mlde
+
+export DERIVED_DATA=/gws/nopw/j04/bris_climdyn/henrya/bp-backups/
+export KK_SLACK_WH_URL=https://hooks.slack.com
+export WANDB_EXPERIMENT_NAME="ml-downscaling-emulator"
+
+checkpoint=$1
+em=$2
+dataset=$3
+split=$4
+workdir=$5
+
+python bin/predict.py --checkpoint ${checkpoint} --num-samples 1 --dataset ${dataset} --split ${split} --ensemble-member ${em} ${workdir}
diff --git a/bin/jasmin/queue-sampling b/bin/jasmin/queue-sampling
new file mode 100755
index 00000000..f51e97e6
--- /dev/null
+++ b/bin/jasmin/queue-sampling
@@ -0,0 +1,17 @@
+#!/bin/bash
+# Script for queueing a training job on LOTUS on JASMIN via lotus-wrapper script
+
+set -euo pipefail
+
+SCRIPT_DIR=$( cd -- "$( dirname -- "${BASH_SOURCE[0]}" )" &> /dev/null && pwd )
+
+# 12em NCSN++
+# /gws/nopw/j04/bris_climdyn/henrya/workdirs/score-sde/deterministic/ukcp_local_pr_12em_cncsnpp/bham_pSTV
+
+# 12em Plain U-Net
+# /gws/nopw/j04/bris_climdyn/henrya/workdirs/score-sde/deterministic/ukcp_local_pr_12em_plain_unet/bham_pSTV
+
+# 1em Plain U-Net
+# /gws/nopw/j04/bris_climdyn/henrya/workdirs/score-sde/deterministic/ukcp_local_pr_1em_plain_unet/bham_pSTV
+
+for em in 01; do sbatch --gres=gpu:1 --partition=orchid --account=orchid --time=01:00:00 --mem=16G -- ${SCRIPT_DIR}/lotus-sampling-wrapper epoch_100 ${em} bham64_ccpm-4x_1em_vort850_pr val /gws/nopw/j04/bris_climdyn/henrya/workdirs/score-sde/deterministic/ukcp_local_pr_1em_plain_unet/bham_pSTV; done