Skip to content

Commit 664c0be

Browse files
michielkallenbergmichielkallenberg
committed
simple LSTM
1 parent 83e9fd3 commit 664c0be

1 file changed

Lines changed: 232 additions & 0 deletions

File tree

‎cybench/runs/run_de_agera5fpar_lstm.sh‎

Lines changed: 232 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,232 @@
1+
#!/usr/bin/env bash
2+
#
3+
# Germany-only colleague-style LSTM: AgERA5 + FPAR, avg_pool tokenizer, no static context.
4+
# Runs the same screening → walk-forward pipeline as the SLURM benchmark (other torch models).
5+
#
6+
# Usage (from repo root):
7+
# cybench/runs/run_de_agera5fpar_lstm.sh all
8+
# cybench/runs/run_de_agera5fpar_lstm.sh screening --batch my_de_lstm_v1
9+
# cybench/runs/run_de_agera5fpar_lstm.sh walk_forward --frozen-dir /path/to/.../2017_2018_2019_2020_2021
10+
# cybench/runs/run_de_agera5fpar_lstm.sh all --crop maize --device cpu --hp-trials 5 --dry-run
11+
#
12+
set -euo pipefail
13+
14+
usage() {
15+
cat <<'EOF'
16+
Usage: run_de_agera5fpar_lstm.sh <screening|walk_forward|all> [options]
17+
18+
Phases:
19+
screening HPO on train/val, final fit on train+val, evaluate on held-out test block
20+
walk_forward Rolling forecasts using frozen hyperparameters from screening
21+
all screening, then walk_forward (auto-discovers latest screening artifacts)
22+
23+
Options:
24+
--batch NAME experiment.name / output folder (default: baselines_de_agera5fpar_lstm_v1)
25+
--crop CROP wheat | maize (default: maize)
26+
--country CC ISO country code (default: DE)
27+
--horizon H end_of_sequence (default: middle-of-season → run dirs *_mid_season_*)
28+
--device DEV cuda | cpu (default: cuda)
29+
--hp-trials N Optuna trials in screening (default: 20)
30+
--frozen-dir PATH Screening split folder with optimal_model.yaml (walk_forward only)
31+
--collect After walk_forward, run collect_walk_forward_results.py
32+
--dry-run Print commands without executing
33+
-h, --help Show this help
34+
35+
Output (relative to repo root):
36+
../output/<batch>/maize_DE_lstm_lf_screening_mid_season_<timestamp>/
37+
../output/<batch>/maize_DE_lstm_lf_walk_forward_mid_season_<timestamp>/
38+
EOF
39+
}
40+
41+
PHASE=""
42+
BATCH="baselines_de_agera5fpar_lstm_v1"
43+
CROP="maize"
44+
COUNTRY="DE"
45+
HORIZON="middle-of-season"
46+
DEVICE="cuda"
47+
HP_TRIALS=20
48+
FROZEN_DIR=""
49+
COLLECT=false
50+
DRY_RUN=false
51+
52+
if [[ $# -lt 1 ]]; then
53+
usage
54+
exit 1
55+
fi
56+
57+
PHASE=$1
58+
shift
59+
60+
case "${PHASE}" in
61+
screening|walk_forward|all) ;;
62+
-h|--help)
63+
usage
64+
exit 0
65+
;;
66+
*)
67+
echo "Unknown phase: ${PHASE}" >&2
68+
usage
69+
exit 1
70+
;;
71+
esac
72+
73+
while [[ $# -gt 0 ]]; do
74+
case "$1" in
75+
--batch) BATCH=$2; shift 2 ;;
76+
--crop) CROP=$2; shift 2 ;;
77+
--country) COUNTRY=$2; shift 2 ;;
78+
--horizon) HORIZON=$2; shift 2 ;;
79+
--device) DEVICE=$2; shift 2 ;;
80+
--hp-trials) HP_TRIALS=$2; shift 2 ;;
81+
--frozen-dir) FROZEN_DIR=$2; shift 2 ;;
82+
--collect) COLLECT=true; shift ;;
83+
--dry-run) DRY_RUN=true; shift ;;
84+
-h|--help) usage; exit 0 ;;
85+
*)
86+
echo "Unknown option: $1" >&2
87+
usage
88+
exit 1
89+
;;
90+
esac
91+
done
92+
93+
SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
94+
REPO_ROOT="$(cd "${SCRIPT_DIR}/../.." && pwd)"
95+
cd "${REPO_ROOT}"
96+
97+
OUTPUT_ROOT="${REPO_ROOT}/../output"
98+
BATCH_DIR="${OUTPUT_ROOT}/${BATCH}"
99+
if [[ "${DRY_RUN}" != true ]]; then
100+
mkdir -p "${BATCH_DIR}"
101+
OUTPUT_ROOT="$(cd "${OUTPUT_ROOT}" && pwd)"
102+
BATCH_DIR="${OUTPUT_ROOT}/${BATCH}"
103+
fi
104+
105+
run_cmd() {
106+
if [[ "${DRY_RUN}" == true ]]; then
107+
printf ' '
108+
printf '%q ' "$@"
109+
printf '\n'
110+
else
111+
"$@"
112+
fi
113+
}
114+
115+
# Shared Hydra overrides (must match across screening and walk-forward).
116+
shared_overrides() {
117+
cat <<EOF
118+
dataset/crop=${CROP}
119+
dataset.country=${COUNTRY}
120+
dataset.framework=torch
121+
dataset.use_cache=true
122+
dataset/temporal=no_aggregate
123+
dataset.temporal.season.end_of_sequence=${HORIZON}
124+
model=lstm_lf
125+
model/torch_model/temporal_encoder/tokenizer=avg_pool
126+
model.torch_model.embed_dim=8
127+
~dataset.temporal.sources.ndvi
128+
~dataset.temporal.sources.soil_moisture
129+
dataset.temporal.sources.meteo.select=[tmin,tmax,tavg,prec,rad,et0,vpd]
130+
+process={name:select_context,drop:[year,sos_sin,sos_cos,eos_sin,eos_cos,loc_x,loc_y,loc_z,awc,bulk_density,drainage_class_1,drainage_class_2,drainage_class_3,drainage_class_4,drainage_class_5,drainage_class_6],keep:null}
131+
experiment.name=${BATCH}
132+
experiment.n_repetitions=1
133+
experiment.device=${DEVICE}
134+
EOF
135+
}
136+
137+
read_overrides() {
138+
mapfile -t OVERRIDES < <(shared_overrides)
139+
}
140+
141+
horizon_tag() {
142+
poetry run python -c \
143+
"from cybench.util.prediction_horizon import prediction_horizon_tag; print(prediction_horizon_tag('${HORIZON}'))"
144+
}
145+
146+
find_frozen_screening_dir() {
147+
local htag model_name run_dir frozen
148+
htag=$(horizon_tag)
149+
model_name="lstm_lf"
150+
run_dir=$(ls -td "${BATCH_DIR}/${CROP}_${COUNTRY}_${model_name}_screening_${htag}_"* 2>/dev/null | head -1 || true)
151+
if [[ -z "${run_dir}" ]]; then
152+
echo "No screening run under ${BATCH_DIR} for ${CROP}/${COUNTRY}/lstm_lf horizon=${HORIZON}" >&2
153+
return 1
154+
fi
155+
frozen=$(find "${run_dir}" -name optimal_model.yaml -printf '%h\n' 2>/dev/null | head -1)
156+
if [[ -n "${frozen}" ]]; then
157+
echo "${frozen}"
158+
return 0
159+
fi
160+
echo "No optimal_model.yaml under ${run_dir}" >&2
161+
return 1
162+
}
163+
164+
run_screening() {
165+
read_overrides
166+
local -a cmd=(poetry run python cybench/runs/run_experiments.py)
167+
cmd+=("${OVERRIDES[@]}")
168+
cmd+=(
169+
validation=screening
170+
+hp_search=bayesian
171+
"hp_search.n_trials=${HP_TRIALS}"
172+
)
173+
echo "== Screening | ${CROP}/${COUNTRY} | batch=${BATCH} | device=${DEVICE} | out=${BATCH_DIR}"
174+
run_cmd "${cmd[@]}"
175+
}
176+
177+
run_walk_forward() {
178+
local frozen=$1
179+
if [[ ! -f "${frozen}/optimal_model.yaml" ]]; then
180+
echo "Missing optimal_model.yaml in ${frozen}" >&2
181+
exit 1
182+
fi
183+
read_overrides
184+
local -a cmd=(poetry run python cybench/runs/run_experiments.py)
185+
cmd+=("${OVERRIDES[@]}")
186+
cmd+=(
187+
validation=walk_forward
188+
"validation.frozen_screening_dir=${frozen}"
189+
)
190+
echo "== Walk-forward | ${CROP}/${COUNTRY} | batch=${BATCH} | frozen=${frozen}"
191+
run_cmd "${cmd[@]}"
192+
}
193+
194+
run_collect() {
195+
local -a cmd=(
196+
poetry run python cybench/runs/analysis/collect_walk_forward_results.py
197+
--batch "${BATCH}"
198+
--output-dir "${OUTPUT_ROOT}/paper_${BATCH}"
199+
)
200+
echo "== Collect walk-forward results → ${OUTPUT_ROOT}/paper_${BATCH}"
201+
run_cmd "${cmd[@]}"
202+
}
203+
204+
case "${PHASE}" in
205+
screening)
206+
run_screening
207+
;;
208+
walk_forward)
209+
if [[ -z "${FROZEN_DIR}" ]]; then
210+
FROZEN_DIR=$(find_frozen_screening_dir)
211+
fi
212+
run_walk_forward "${FROZEN_DIR}"
213+
if [[ "${COLLECT}" == true ]]; then
214+
run_collect
215+
fi
216+
;;
217+
all)
218+
run_screening
219+
if [[ "${DRY_RUN}" == true ]]; then
220+
echo "(dry-run: would auto-discover frozen screening dir for walk-forward)"
221+
run_cmd echo walk_forward skipped in dry-run after screening
222+
else
223+
FROZEN_DIR=$(find_frozen_screening_dir)
224+
run_walk_forward "${FROZEN_DIR}"
225+
if [[ "${COLLECT}" == true ]]; then
226+
run_collect
227+
fi
228+
fi
229+
;;
230+
esac
231+
232+
echo "Done."

0 commit comments

Comments
 (0)