Expose n_epochs on novel_train and simple_mlp_train

scripts/create_datasets/test_resources.sh

@@ -10,50 +10,99 @@ set -e
 
 RAW_DATA=resources_test/common
 OUTPUT_DIR=resources_test/task_predict_modality
+DATASET_DIR=$OUTPUT_DIR/openproblems_neurips2021
 
 mkdir -p $OUTPUT_DIR
 
-export NXF_VER=22.04.5
+export NXF_VER=25.10.7
+
+# skip a step when its output exists and no input is newer; set FORCE=1 to
+# regenerate everything. usage: up_to_date <output> <input>...
+FORCE=${FORCE:-0}
+up_to_date() {
+  local out=$1
+  shift
+
+  if [ "$FORCE" -ne 0 ] || [ ! -e "$out" ]; then
+    return 1
+  fi
+
+  # a directory output only counts once something has been written into it
+  if [ -d "$out" ] && [ -z "$(ls -A "$out" 2>/dev/null)" ]; then
+    return 1
+  fi
+
+  local input
+  for input in "$@"; do
+    if [ "$input" -nt "$out" ]; then
+      return 1
+    fi
+  done
+}
 
 echo "Preprocess datasets"
-nextflow run . \
-  -main-script target/nextflow/workflows/process_datasets/main.nf \
-  -profile docker \
-  -entry auto \
-  -c common/nextflow_helpers/labels_ci.config \
-  --input_states "resources_test/common/openproblems_neurips2021/**/state.yaml" \
-  --rename_keys 'input_mod1:output_mod1;input_mod2:output_mod2' \
-  --settings '{"output_train_mod1": "$id/train_mod1.h5ad", "output_train_mod2": "$id/train_mod2.h5ad", "output_test_mod1": "$id/test_mod1.h5ad", "output_test_mod2": "$id/test_mod2.h5ad"}' \
-  --publish_dir "$OUTPUT_DIR" \
-  --output_state '$id/state.yaml'
+RAW_STATES=($RAW_DATA/openproblems_neurips2021/*/state.yaml)
+if up_to_date $DATASET_DIR/bmmc_cite/normal/state.yaml "${RAW_STATES[@]}" &&
+   up_to_date $DATASET_DIR/bmmc_cite/swap/state.yaml "${RAW_STATES[@]}" &&
+   up_to_date $DATASET_DIR/bmmc_multiome/normal/state.yaml "${RAW_STATES[@]}" &&
+   up_to_date $DATASET_DIR/bmmc_multiome/swap/state.yaml "${RAW_STATES[@]}"; then
+  echo "  already up to date, skipping"
+else
+  nextflow run . \
+    -main-script target/nextflow/workflows/process_datasets/main.nf \
+    -profile docker \
+    -entry auto \
+    -c common/nextflow_helpers/labels_ci.config \
+    --input_states "resources_test/common/openproblems_neurips2021/**/state.yaml" \
+    --rename_keys 'input_mod1:output_mod1;input_mod2:output_mod2' \
+    --settings '{"output_train_mod1": "$id/train_mod1.h5ad", "output_train_mod2": "$id/train_mod2.h5ad", "output_test_mod1": "$id/test_mod1.h5ad", "output_test_mod2": "$id/test_mod2.h5ad"}' \
+    --publish_dir "$OUTPUT_DIR" \
+    --output_state '$id/state.yaml'
+fi
 
 echo "Run one method"
 
 for name in bmmc_cite/normal bmmc_cite/swap bmmc_multiome/normal bmmc_multiome/swap; do
+  STATE=$DATASET_DIR/$name/state.yaml
+
   echo "Run KNN on $name"
-  viash run src/methods/knnr_py/config.vsh.yaml -- \
-    --input_train_mod1 $OUTPUT_DIR/openproblems_neurips2021/$name/train_mod1.h5ad \
-    --input_train_mod2 $OUTPUT_DIR/openproblems_neurips2021/$name/train_mod2.h5ad \
-    --input_test_mod1 $OUTPUT_DIR/openproblems_neurips2021/$name/test_mod1.h5ad \
-    --output $OUTPUT_DIR/openproblems_neurips2021/$name/prediction.h5ad
+  if up_to_date $DATASET_DIR/$name/prediction.h5ad $STATE; then
+    echo "  already up to date, skipping"
+  else
+    viash run src/methods/knnr_py/config.vsh.yaml -- \
+      --input_train_mod1 $DATASET_DIR/$name/train_mod1.h5ad \
+      --input_train_mod2 $DATASET_DIR/$name/train_mod2.h5ad \
+      --input_test_mod1 $DATASET_DIR/$name/test_mod1.h5ad \
+      --output $DATASET_DIR/$name/prediction.h5ad
+  fi
 
   echo "pre-train simple_mlp on $name"
-  [ -d $OUTPUT_DIR/openproblems_neurips2021/$name/models/simple_mlp/ ] && rm -r $OUTPUT_DIR/openproblems_neurips2021/$name/models/simple_mlp/
-  mkdir -p $OUTPUT_DIR/openproblems_neurips2021/$name/models/simple_mlp/
-  viash run src/methods/simple_mlp/simple_mlp_train/config.vsh.yaml -- \
-    --input_train_mod1 $OUTPUT_DIR/openproblems_neurips2021/$name/train_mod1.h5ad \
-    --input_train_mod2 $OUTPUT_DIR/openproblems_neurips2021/$name/train_mod2.h5ad \
-    --input_test_mod1 $OUTPUT_DIR/openproblems_neurips2021/$name/test_mod1.h5ad \
-    --output $OUTPUT_DIR/openproblems_neurips2021/$name/models/simple_mlp/
+  if up_to_date $DATASET_DIR/$name/models/simple_mlp/ $STATE; then
+    echo "  already up to date, skipping"
+  else
+    rm -rf $DATASET_DIR/$name/models/simple_mlp/
+    mkdir -p $DATASET_DIR/$name/models/simple_mlp/
+    viash run src/methods/simple_mlp/simple_mlp_train/config.vsh.yaml -- \
+      --input_train_mod1 $DATASET_DIR/$name/train_mod1.h5ad \
+      --input_train_mod2 $DATASET_DIR/$name/train_mod2.h5ad \
+      --input_test_mod1 $DATASET_DIR/$name/test_mod1.h5ad \
+      --n_epochs 2 \
+      --output $DATASET_DIR/$name/models/simple_mlp/
+  fi
 
   echo "pre-train novel on $name"
-  [ -d $OUTPUT_DIR/openproblems_neurips2021/$name/models/novel/ ] && rm -r $OUTPUT_DIR/openproblems_neurips2021/$name/models/novel/
-  mkdir -p $OUTPUT_DIR/openproblems_neurips2021/$name/models/novel/
-  viash run src/methods/novel/novel_train/config.vsh.yaml -- \
-    --input_train_mod1 $OUTPUT_DIR/openproblems_neurips2021/$name/train_mod1.h5ad \
-    --input_train_mod2 $OUTPUT_DIR/openproblems_neurips2021/$name/train_mod2.h5ad \
-    --input_test_mod1 $OUTPUT_DIR/openproblems_neurips2021/$name/test_mod1.h5ad \
-    --output $OUTPUT_DIR/openproblems_neurips2021/$name/models/novel
+  if up_to_date $DATASET_DIR/$name/models/novel/ $STATE; then
+    echo "  already up to date, skipping"
+  else
+    rm -rf $DATASET_DIR/$name/models/novel/
+    mkdir -p $DATASET_DIR/$name/models/novel/
+    viash run src/methods/novel/novel_train/config.vsh.yaml -- \
+      --input_train_mod1 $DATASET_DIR/$name/train_mod1.h5ad \
+      --input_train_mod2 $DATASET_DIR/$name/train_mod2.h5ad \
+      --input_test_mod1 $DATASET_DIR/$name/test_mod1.h5ad \
+      --n_epochs 2 \
+      --output $DATASET_DIR/$name/models/novel
+  fi
 
 done
 

src/api/file_pretrained_model.yaml

@@ -1,3 +1,4 @@
 type: file
+example: "model"
 label: "Pretrained model"
 summary: "A pretrained model for predicting the expression of one modality from another."

src/methods/novel/helper_functions.py

@@ -204,9 +204,9 @@ def forward(self, x):
 def rmse(y, y_pred):
     return np.sqrt(np.mean(np.square(y - y_pred)))
 
-def train_and_valid(model, optimizer, loss_fn, dataloader_train, dataloader_test, name_model, device):
+def train_and_valid(model, optimizer, loss_fn, dataloader_train, dataloader_test, name_model, device, n_epochs=100):
     best_score = 100000
-    for i in range(100):
+    for i in range(n_epochs):
         train_losses = []
         model.train()
 

src/methods/novel/novel_train/config.vsh.yaml

@@ -1,5 +1,12 @@
 __merge__: ../../../api/comp_method_train.yaml
 name: novel_train
+arguments:
+  - name: "--n_epochs"
+    type: integer
+    default: 100
+    description: Number of training epochs.
+    info:
+      test_default: 2
 resources:
   - path: script.py
     type: python_script

src/methods/novel/novel_train/script.py

@@ -23,7 +23,8 @@
 par = {
   'input_train_mod1': 'resources_test/task_predict_modality/openproblems_neurips2021/bmmc_multiome/normal/train_mod1.h5ad',
   'input_train_mod2': 'resources_test/task_predict_modality/openproblems_neurips2021/bmmc_multiome/normal/train_mod2.h5ad',
-  'output': 'resources_test/task_predict_modality/openproblems_neurips2021/bmmc_multiome/normal/models/novel'
+  'output': 'resources_test/task_predict_modality/openproblems_neurips2021/bmmc_multiome/normal/models/novel',
+  'n_epochs': 100
 }
 meta = {
   'resources_dir': 'src/methods/novel',
@@ -145,7 +146,7 @@
 output_transform = f"{par['output']}/transform.pkl"
 
 # train model
-train_and_valid(model, optimizer, loss_fn, dataloader_train, dataloader_test, output_model, device)
+train_and_valid(model, optimizer, loss_fn, dataloader_train, dataloader_test, output_model, device, n_epochs=par['n_epochs'])
 
 # Add model dim for use in predict part
 adata.uns["model_dim"] = {"mod1": n_vars_mod1, "mod2": n_vars_mod2}

src/methods/simple_mlp/simple_mlp_train/config.vsh.yaml

@@ -1,5 +1,12 @@
 __merge__: /src/api/comp_method_train.yaml
 name: simple_mlp_train
+arguments:
+  - name: "--n_epochs"
+    type: integer
+    required: false
+    description: Number of training epochs. Defaults to the value in the bundled model config.
+    info:
+      test_default: 2
 resources:
   - type: python_script
     path: script.py

src/methods/simple_mlp/simple_mlp_train/script.py

@@ -19,7 +19,8 @@
     'input_train_mod1': 'resources_test/task_predict_modality/openproblems_neurips2021/bmmc_multiome/swap/train_mod1.h5ad',
     'input_train_mod2': 'resources_test/task_predict_modality/openproblems_neurips2021/bmmc_multiome/swap/train_mod2.h5ad',
     'input_test_mod1': 'resources_test/task_predict_modality/openproblems_neurips2021/bmmc_multiome/swap/test_mod1.h5ad',
-    'output': 'output/model'
+    'output': 'output/model',
+    'n_epochs': None
 }
 meta = {
     'resources_dir': 'src/methods/simple_mlp',
@@ -138,6 +139,9 @@ def _train(X, y, Xt, yt, logger, config, num_workers):
 
     config = utils.load_yaml(yaml_path)
 
+    if par["n_epochs"] is not None:
+        config = config._replace(epochs=par["n_epochs"])
+
     if config.batch_size > X.shape[0]:
         config = config._replace(batch_size=math.ceil(X.shape[0] / 2))