From ec5212cef1199b2f473755617eff622d82673f91 Mon Sep 17 00:00:00 2001
From: cgohil8 <53237863+cgohil8@users.noreply.github.com>
Date: Thu, 15 Feb 2024 09:49:48 +0000
Subject: [PATCH] Changed reference brain + beamforming regularisation (#270)

* Don't pin down osl/osl-dynamics version in Windows env file.
* Use same standard brain as Matlab OSL.
* Option to specify regularisation in beamforming.
* Fixed bug generating the source report (without extra functions).
---
 envs/windows.yml                |  4 ++--
 osl/report/preproc_report.py    |  6 +++---
 osl/report/src_report.py        |  3 +--
 osl/source_recon/rhino/utils.py |  4 ++--
 osl/source_recon/wrappers.py    | 10 ++++++++++
 5 files changed, 18 insertions(+), 9 deletions(-)

diff --git a/envs/windows.yml b/envs/windows.yml
index b534ebe8..47a2193b 100644
--- a/envs/windows.yml
+++ b/envs/windows.yml
@@ -42,5 +42,5 @@ dependencies:
   - tensorflow_probability==0.17.0
   - tqdm==4.65.0
   - osfclient==0.0.5
-  - osl==0.5.1
-  - osl-dynamics==1.2.11
+  - osl
+  - osl-dynamics
diff --git a/osl/report/preproc_report.py b/osl/report/preproc_report.py
index 4fc5e4f1..ccbb63a2 100644
--- a/osl/report/preproc_report.py
+++ b/osl/report/preproc_report.py
@@ -441,9 +441,9 @@ def save_extra_funcs(raw, savebase=None):
     
     """
     extra_funcs = re.findall(
-    "%% extra_funcs start %%(.*?)%% extra_funcs end %%",
-    raw.info["description"],
-    flags=re.DOTALL,
+        "%% extra_funcs start %%(.*?)%% extra_funcs end %%",
+        raw.info["description"],
+        flags=re.DOTALL,
     )
     
     if savebase is not None:
diff --git a/osl/report/src_report.py b/osl/report/src_report.py
index cf43a1c4..6aa45b76 100644
--- a/osl/report/src_report.py
+++ b/osl/report/src_report.py
@@ -137,8 +137,7 @@ def save_extra_funcs(extra_funcs, reportdir):
     fpath : str
         Path to saved text file.
     """
-    
-    if reportdir is not None:
+    if reportdir is not None and extra_funcs is not None:
         fpath = reportdir / 'extra_funcs.txt'
         with(open(fpath, 'w')) as file:
             [print(f"{inspect.getsource(func)}\n\n", file=file) for func in extra_funcs]
diff --git a/osl/source_recon/rhino/utils.py b/osl/source_recon/rhino/utils.py
index ee6f53c7..dd36881a 100644
--- a/osl/source_recon/rhino/utils.py
+++ b/osl/source_recon/rhino/utils.py
@@ -80,7 +80,7 @@ def get_rhino_files(subjects_dir, subject):
         "bet_outskin_plus_nose_mesh_file": op.join(surfaces_dir, "outskin_plus_nose_mesh.nii.gz"),
         "bet_inskull_mesh_file": op.join(surfaces_dir, "inskull_mesh.nii.gz"),
         "bet_outskull_mesh_file": op.join(surfaces_dir, "outskull_mesh.nii.gz"),
-        "std_brain": op.join(os.environ["FSLDIR"], "data", "standard", "MNI152_T1_1mm_brain.nii.gz"),
+        "std_brain": op.join(os.environ["FSLDIR"], "data", "standard", "MNI152_T1_1mm.nii.gz"),
         "std_brain_bigfov": op.join(os.environ["FSLDIR"], "data", "standard", "MNI152_T1_1mm_BigFoV_facemask.nii.gz"),
         "completed": op.join(surfaces_dir, "completed.txt"),
     }
@@ -119,7 +119,7 @@ def get_rhino_files(subjects_dir, subject):
         "bet_outskin_plus_nose_mesh_file": op.join(coreg_dir, "scaled_outskin_plus_nose_mesh.nii.gz"),
         "bet_inskull_mesh_file": op.join(coreg_dir, "scaled_inskull_mesh.nii.gz"),
         "bet_outskull_mesh_file": op.join(coreg_dir, "scaled_outskull_mesh.nii.gz"),
-        "std_brain": op.join(os.environ["FSLDIR"], "data", "standard", "MNI152_T1_1mm_brain.nii.gz"),
+        "std_brain": op.join(os.environ["FSLDIR"], "data", "standard", "MNI152_T1_1mm.nii.gz"),
     }
 
     # All RHINO files
diff --git a/osl/source_recon/wrappers.py b/osl/source_recon/wrappers.py
index 5725c347..7010544e 100644
--- a/osl/source_recon/wrappers.py
+++ b/osl/source_recon/wrappers.py
@@ -411,6 +411,7 @@ def beamform(
     freq_range,
     chantypes,
     rank,
+    reg=0,
 ):
     """Wrapper function for beamforming.
 
@@ -433,6 +434,8 @@ def beamform(
         Channel types to use in beamforming.
     rank : dict
         Keys should be the channel types and the value should be the rank to use.
+    reg : float
+        The regularization for the whitened data covariance.
     """
     logger.info("beamforming")
 
@@ -470,6 +473,7 @@ def beamform(
         chantypes=chantypes,
         weight_norm="nai",
         rank=rank,
+        reg=reg,
         save_filters=True,
     )
 
@@ -485,6 +489,7 @@ def beamform(
             "beamform": True,
             "chantypes": chantypes,
             "rank": rank,
+            "reg": reg,
             "freq_range": freq_range,
             "filters_cov_plot": filters_cov_plot,
             "filters_svd_plot": filters_svd_plot,
@@ -665,6 +670,7 @@ def beamform_and_parcellate(
     spatial_resolution=None,
     reference_brain="mni",
     extra_chans="stim",
+    reg=0,
 ):
     """Wrapper function for beamforming and parcellation.
 
@@ -704,6 +710,8 @@ def beamform_and_parcellate(
     extra_chans : str or list of str
         Extra channels to include in the parc-raw.fif file. Defaults to 'stim'.
         Stim channels are always added to parc-raw.fif in addition to extra_chans.
+    reg : float
+        The regularization for the whitened data covariance in the beamforming stage.
     """
     logger.info("beamform_and_parcellate")
 
@@ -741,6 +749,7 @@ def beamform_and_parcellate(
         chantypes=chantypes,
         weight_norm="nai",
         rank=rank,
+        reg=reg,
         save_filters=True,
     )
 
@@ -823,6 +832,7 @@ def beamform_and_parcellate(
             "parcellate": True,
             "chantypes": chantypes,
             "rank": rank,
+            "reg": reg,
             "freq_range": freq_range,
             "filters_cov_plot": filters_cov_plot,
             "filters_svd_plot": filters_svd_plot,