add experimental model M-E01

configs/shimnet_600_M-E01.yaml

@@ -0,0 +1,134 @@
+model:
+  _target_: shimnet.models.ShimnetModular
+  encoder:
+    _target_: shimnet.models.ConvEncoder
+    hidden_dim: 64
+    output_dim: 128
+    activation: gelu
+    kernel_size: 7
+  local_feature_processor:
+    _target_: shimnet.models.ConvMLP
+    input_dim: 128
+    output_dim: 64
+    hidden_dims:
+    - 256
+    - 128
+    activation: gelu
+  attention_module:
+    _target_: shimnet.models.KVAttention
+    kv_dim: 32
+    num_heads: 8
+    k_processor:
+      _target_: shimnet.models.ConvMLP
+      input_dim: 128
+      output_dim: 256
+      hidden_dims:
+      - 512
+      - 256
+      activation: gelu
+    v_processor:
+      _target_: shimnet.models.MLP
+      input_dim: 128
+      output_dim: 256
+      hidden_dims:
+      - 512
+      - 256
+      activation: gelu
+  global_feature_processor:
+    _target_: shimnet.models.MLP
+    input_dim: 256
+    output_dim: 64
+    hidden_dims:
+    - 512
+    - 256
+    activation: gelu
+  response_head:
+    _target_: shimnet.models.MLP
+    input_dim: 256
+    output_dim: 81
+    hidden_dims:
+    - 512
+    - 256
+    activation: gelu
+  decoder:
+    _target_: shimnet.models.ConvDecoder
+    input_dim: 128
+    hidden_dim: 128
+    activation: gelu
+    kernel_size: 7
+    last_bias: false
+    last_activation: false
+training:
+#- batch_size: 64
+#learning_rate: 0.001
+#max_iters: 1600000
+- batch_size: 256
+  learning_rate: 0.0001
+  max_iters: 25600000
+- batch_size: 256
+  learning_rate: 0.00002
+  max_iters: 12800000
+losses:
+  clean:
+    function: mae
+    weight: 1.0
+  noised:
+    function: mae
+    weight: 1.0
+  response:
+    function: mae
+    weight: 1.0
+data:
+  _target_: shimnet.generators.Generator
+  include_response_function: true
+  seed: null # null means random seed
+  batch_size: null # to be set in training script
+  clean_spectra_generator:
+    _target_: shimnet.generators.TheoreticalMultipletSpectraGenerator
+    pixels: 2048
+    frq_step: ${metadata.frq_step}
+    peaks_parameter_generator:
+        _target_: shimnet.generators.PeaksParameterDataGenerator
+        atom_groups_data_file: data/multiplets_10000_parsed.txt
+        number_of_signals_min: 2
+        number_of_signals_max: 5
+        relative_frequency_min: -0.4
+        relative_frequency_max: 0.4
+        spectrum_width_min: 0.2
+        spectrum_width_max: 1.0
+        relative_width_min: 1.0
+        relative_width_max: 2.0
+        relative_height_min: 0.5
+        relative_height_max: 4
+        thf_min: 0.5
+        thf_max: 2
+        trf_min: 0.0
+        trf_max: 1.0
+        multiplicity_j1_min: 0.0
+        multiplicity_j1_max: 15
+        multiplicity_j2_min: 0.0
+        multiplicity_j2_max: 15
+  response_generator:
+    _target_: shimnet.generators.ResponseGenerator
+    response_function_library:
+      _target_: shimnet.generators.ResponseLibrary
+      response_files:
+        - data/smoothed_scrf_kernels/scrf_81_600MHz_smoothed_1-1-1.pt
+        - data/smoothed_scrf_kernels/scrf_81_600MHz_smoothed_1-2-1.pt
+        - data/smoothed_scrf_kernels/scrf_81_600MHz_smoothed_1-4-1.pt
+        - data/smoothed_scrf_kernels/scrf_81_600MHz_smoothed_1-3-3-1.pt
+    pad_to: 81
+    response_function_stretch_min: 0.8
+    response_function_stretch_max: 1.5
+    response_function_noise: 0.0
+    flip_response_function: false
+  noise_generator:
+    _target_: shimnet.generators.NoiseGenerator
+    spectrum_noise_min: 0.0
+    spectrum_noise_max: 0.015625
+logging:
+  step: 1000000
+  num_plots: 32
+metadata:
+  frq_step: 0.30048
+  spectrometer_frequency: 600.0

download_files.py

@@ -10,6 +10,10 @@ ALL_FILES_TO_DOWNLOAD = {
     {
         "url": "https://drive.google.com/uc?export=download&id=1_VxOpFGJcFsOa5DHOW2GJbP8RvHCmC1N",
         "destination": "weights/shimnet_600MHz.pt"
+    },
+    {
+        "url": "https://drive.google.com/uc?export=download&id=1643Il3qgCupY0n8Mar6WBc2WVuoQRzie",
+        "destination": "weights/shimnet_600MHz_M-E01.pt"
     }],
     "SCRF": [{
         "url": "https://drive.google.com/uc?export=download&id=113al7A__yYALx_2hkESuzFIDU3feVtNY",

predict-gui.py

@@ -140,7 +140,7 @@ with gr.Blocks() as app:
         with gr.Column():
             model_selection = gr.Radio(
                 label="Select Model",
-                choices=["600 MHz", "700 MHz", "Custom"],
+                choices=["600 MHz", "700 MHz", "M-E01", "Custom"],
                 value="600 MHz"
             )
             config_file = gr.File(label="Custom Config File (.yaml)", visible=False, height=120)
@@ -189,6 +189,9 @@ with gr.Blocks() as app:
         elif model_selection == "700 MHz":
             config_file = os.path.join(os.path.dirname(__file__), "configs/shimnet_700.yaml")
             weights_file = os.path.join(os.path.dirname(__file__), "weights/shimnet_700MHz.pt")
+        elif model_selection == "M-E01":
+            config_file = os.path.join(os.path.dirname(__file__), "configs/shimnet_600_M-E01.yaml")
+            weights_file = os.path.join(os.path.dirname(__file__), "weights/shimnet_600MHz_M-E01.pt")
         else:
             config_file = config_file.name
             weights_file = weights_file.name