From 053633ab91736e5927dcb4b417b818b892b99bd9 Mon Sep 17 00:00:00 2001
From: Chetan Gohil <chetan.gohil@psych.ox.ac.uk>
Date: Sat, 31 Aug 2024 03:11:52 +0100
Subject: [PATCH] Coreg with epoch files.

---
 osl/source_recon/rhino/coreg.py |  6 +++---
 osl/source_recon/wrappers.py    | 12 +++++++++---
 2 files changed, 12 insertions(+), 6 deletions(-)

diff --git a/osl/source_recon/rhino/coreg.py b/osl/source_recon/rhino/coreg.py
index 9b5805d..c3e8dfe 100644
--- a/osl/source_recon/rhino/coreg.py
+++ b/osl/source_recon/rhino/coreg.py
@@ -151,12 +151,12 @@ def coreg(
     # -------------------------------------------------------------------------------------------------------------------
     # Copy fif_file to new file for modification, and (optionally) changes dev_head_t to equal dev_ctf_t in fif file info
 
-    if fif_file[-7:] == "raw.fif":
+    if "raw.fif" in fif_file:
         raw = read_raw(fif_file)
-    elif fif_file[-10:] == "epochs.fif":
+    elif "epo.fif" in fif_file:
         raw = read_epochs(fif_file)
     else:
-        raise ValueError("Invalid fif file, needs to be a *raw.fif or a *epochs.fif file")
+        raise ValueError("Invalid fif file, needs to be a *-raw.fif or a *-epo.fif file")
     info = raw.info
 
     if use_dev_ctf_t:
diff --git a/osl/source_recon/wrappers.py b/osl/source_recon/wrappers.py
index 5de30de..1e63939 100644
--- a/osl/source_recon/wrappers.py
+++ b/osl/source_recon/wrappers.py
@@ -250,8 +250,9 @@ def compute_surfaces(
 def coregister(
     outdir,
     subject,
-    preproc_file,
     smri_file,
+    preproc_file=None,
+    epoch_file=None,
     use_nose=True,
     use_headshape=True,
     already_coregistered=False,
@@ -267,11 +268,13 @@ def coregister(
         Path to where to output the source reconstruction files.
     subject : str
         Subject name/id.
-    preproc_file : str
-        Path to the preprocessed fif file.
     smri_file : str
         Path to the T1 weighted structural MRI file to use in source
         reconstruction.
+    preproc_file : str, optional
+        Path to the preprocessed fif file.
+    epoch_file : str, optional
+        Path to the preprocessed epochs fif file.
     use_nose : bool, optional
         Should we use the nose in the coregistration?
     use_headshape : bool, optional
@@ -291,6 +294,9 @@ def coregister(
     reportdir : str, optional
         Path to report directory.
     """
+    if preproc_file is None:
+        preproc_file = epoch_file
+
     # Run coregistration
     rhino.coreg(
         fif_file=preproc_file,