#!/bin/bash
#CFG --machine_type a3-highgpu-8g
#CFG --zone us-central1-a
#CFG --image_name jhm-macetools-openmm-2025-03-15
#CFG --user jhm
#CFG --project_id angstrom-genmd
#CFG --name ajc-testing
#CFG --provision SPOT
#CFG --restartable True
#CFG --efs_id 10.50.93.218:/ajc_test
source ~/.bashrc

source ~/efs/test/ajc_testing/miniconda3/bin/activate macetools-openmm
which python


mace_run_train\
    --name="sclj" \
    --train_file="../water_data/train_water.xyz" \
    --test_file="../water_data/test_water.xyz" \
    --enable_cueq="True" \
    --r_max=6.0 \
    --model="MACE" \
    --interaction_first="RealAgnosticDensityInteractionBlock" \
    --interaction="RealAgnosticDensityResidualInteractionBlock" \
    --scaling="rms_forces_scaling" \
    --E0s="{1: -13.571964772646918,8: -2043.933693071156}" \
    --num_workers=12 \
    --num_interactions=2 \
    --num_channels=128\
    --max_L=1 \
    --correlation=3 \
    --forces_weight=1000 \
    --energy_weight=40 \
    --weight_decay=1e-8 \
    --ema \
    --ema_decay=0.995 \
    --scheduler_patience=5 \
    --clip_grad=1.0 \
    --pair_repulsion="SCLJ" \
    --distance_transform="Soft" \
    --batch_size=32\
    --valid_batch_size=32\
    --enable_cueq=True \
    --max_num_epochs=190\
    --patience=50 \
    --lr 0.01 \
    --swa_lr 1e-3\
    --ema \
    --swa \
    --start_swa=150\
    --swa_energy_weight=1000 \
    --swa_forces_weight=10 \
    --error_table='PerAtomMAE' \
    --default_dtype="float64"\
    --device='cuda' \
    --seed=123 \
    --restart_latest \
    --save_cpu \