Step 1: Prepare the Dataset

• Download the RadioMapSeer dataset from https://radiomapseer.github.io/
• Place the dataset folder under the EquUNet directory
• The folder structure should look like:

EquUNet/
└── RadioMapSeer/

Step 2: Train Group Equivariant UNet

The configuration file for the group equivariant UNet (GUNet) is located at EquUNet/configs/SingleGUNet.yaml

In EquUNet/configs/SingleGUNet.yaml, ensure that init_RConv is set to True during training, otherwise the training will not converge.

# EquUNet/configs/SingleGUNet.yaml
first_gunet:
  init_RConv: True  # <--- Ensure this is True
  group   : "D8"    # <--- Can be set to C4, D4, C8, D8, etc.
model_path: './data/GUNet' # <--- Folder to save trained model

Run the following command from the EquUNet/ directory to train the GUNet:

python train_SingleGUNet.py

Step 3: Train Variational Encoding Model

The configuration file for the variational autoencoder is located at EquUNet/DataCoupling/configs/first_radio.yaml

trainer:
  results_folder: './results/radio_Vae_3' # <--- Folder to save trained model

Run the following command from the EquUNet/DataCoupling/ directory to train the variational autoencoder:

python train_vae.py --cfg ./configs/first_radio.yaml

Step 4: Train Flow Matching Model

The configuration file for the flow matching model is located at EquUNet/DataCoupling/configs/radio_train.yaml

model:
  first_stage:
    ckpt_path: './results/radio_Vae_3/model-73.pt' # <--- the trained variational autoencoder
trainer:
  results_folder: "./results/Diffusion_fm/" # <--- folder to save trained flow matching model
GUNet:
  arch_param_path: "../configs/SingleGUNet.yaml"  # <--- configuration file for GUNet
  weigth_path: "../data/GUNet/C4best_model.mdl"   # <--- trained GUNet model

Run the following command from the EquUNet/DataCoupling/ directory to train the flow matching model:

python train_cond_ldm.py --cfg ./configs/radio_train.yaml

Step 5: Inference

The configuration file for inference is located at EquUNet/DataCoupling/configs/radio_sample.yaml

model:
  first_stage:
    ckpt_path: './results/radio_Vae_3/model-73.pt' # <--- the trained variational autoencoder
sampler:
  save_folder: './results/inference_fm'  # <--- folder to save inference result 
  ckpt_path: "./results/Diffusion_fm/model-200.pt" # <--- the trained flow matching model
GUNet:
  arch_param_path: "../configs/SingleGUNet.yaml"  # <--- configuration file for GUNet
  weigth_path: "../data/GUNet/C4best_model.mdl"   # <--- trained GUNet model

Run the following command from the EquUNet/DataCoupling/ directory to perform inference using the trained models:

python sample_cond_ldm.py --cfg ./configs/radio_sample.yaml