Integrate prerun weight check

MEGnet/megnet_init.py

@@ -14,6 +14,7 @@
 
 _megnet_path = MEGnet.__path__[0]
 weights_path = op.join(_megnet_path, 'model_v2k3')
+model_path = op.join(weights_path, 'model_v2.keras')
 config_path = op.join(weights_path, 'config.json')
 min_model_version = 'v2.2'
 
@@ -26,7 +27,7 @@ def _version_tuple(version):
 
 
 def _check_weights():
-    if not op.exists(config_path):
+    if not op.isfile(model_path) or not op.isfile(config_path):
         return False
 
     try:
@@ -44,7 +45,7 @@ def _check_weights():
     if model_version_tuple is None or min_model_version_tuple is None:
         return False
 
-    return model_version_tuple > min_model_version_tuple
+    return model_version_tuple >= min_model_version_tuple
 
 
 def _download_weights():

MEGnet/prep_inputs/ICA.py

@@ -34,7 +34,7 @@
 from mne.viz.utils import _setup_vmin_vmax, _get_cmap, plt_show
 from scipy.io import savemat
 import PIL.Image
-import MEGnet
+from MEGnet import megnet_init
 from MEGnet.megnet_utilities import fPredictChunkAndVoting_parrallel
 import functools
 
@@ -49,6 +49,19 @@
 # Helper Functions
 # =============================================================================
 
+def _require_model_weights():
+    """Return the model path after checking its presence and version."""
+    if megnet_init._check_weights():
+        return megnet_init.model_path
+
+    raise RuntimeError(
+        'MEGnet model weights are missing or incompatible. Expected '
+        f'{megnet_init.model_path} and a config.json model_version at least '
+        f'{megnet_init.min_model_version}. Run `megnet_init` to download '
+        'compatible model weights.'
+    )
+
+
 # function to transform Cartesian coordinates to spherical coordinates
 # theta = azimuth
 # phi = elevation
@@ -656,9 +669,9 @@ def circle_plot(circle_pos=None, data=None, out_fname=None):
     #return matrix_out
 
 
-def main(filename, outbasename=None, mains_freq=60.0, 
+def main(filename, results_dir, outbasename=None, mains_freq=60.0,
              save_preproc=False, save_ica=False, seedval=0,
-             results_dir=None, filename_raw=None, do_assess_bads=False,
+             filename_raw=None, do_assess_bads=False,
              bad_channels=[]):
     '''
         Perform all of the steps to preprocess the ica maps:
@@ -674,6 +687,8 @@ def main(filename, outbasename=None, mains_freq=60.0,
 
         filename : str or Raw MNE data object
             Path to file
+        results_dir : str / path
+            Path to output directory
         filename_raw : str
             Required for MEGIN datasets
             Path to file
@@ -688,12 +703,12 @@ def main(filename, outbasename=None, mains_freq=60.0,
             Save the ica output 
         seedval : Int
             Set the numpy random seed
-        results_dir : str / path
-            Path to output directory
         do_assess_bads : Bool
             Assess bad channels if not already done
 
     '''
+    _require_model_weights()
+
     if (type(filename) == str) | (type(filename) == PosixPath):
         raw = read_raw(filename)
     elif type(filename) in raw_typelist:
@@ -806,13 +821,14 @@ def classify_ica(results_dir=None, outbasename=None, filename=None):
 
     '''
     from scipy.io import loadmat
+    model_path = _require_model_weights()
+
     os.environ['TF_CPP_MIN_LOG_LEVEL'] = '2' 
     os.environ["CUDA_DEVICE_ORDER"] = "PCI_BUS_ID"   # see issue #152
     os.environ["CUDA_VISIBLE_DEVICES"] = ""
     os.environ["KERAS_BACKEND"] = "torch"
 
     import keras
-    model_path = op.join(MEGnet.__path__[0] ,  'model_v2k3/model_v2.keras') 
     # This is set to use CPU in initial import
     kModel=keras.models.load_model(model_path, compile=False)
 
@@ -830,6 +846,10 @@ def classify_ica(results_dir=None, outbasename=None, filename=None):
     arrSP = np.stack([loadmat(i)['array'] for i in arrSP_fnames])
     preds, probs = fPredictChunkAndVoting_parrallel(kModel, arrTS, arrSP)
     meg_rest_ica_classes = preds.argmax(axis=1)
+    np.save(
+        op.join(results_dir, 'megnet_classification.npy'),
+        meg_rest_ica_classes,
+    )
     ica_comps_toremove = [index for index, value in enumerate(meg_rest_ica_classes) if value in [1, 2, 3]]
     return {'classes':meg_rest_ica_classes,
             'bads_idx': ica_comps_toremove}
@@ -899,7 +919,8 @@ def cmdline():
 
     standard_args = parser.add_argument_group('standard')
     standard_args.add_argument('-filename', help='Path to MEG dataset')
-    standard_args.add_argument('-results_dir', help='Path to save the results')
+    standard_args.add_argument(
+        '-results_dir', required=True, help='Path to save the results')
     standard_args.add_argument('-line_freq', help='{60,50} Hz - AC electric frequency')
 
 
@@ -942,4 +963,3 @@ def cmdline():
 
 
 
-

MEGnet/tests/test_ica_model_weights.py

@@ -0,0 +1,89 @@
+import pytest
+import numpy as np
+
+from MEGnet.prep_inputs import ICA
+
+
+def test_main_requires_results_dir():
+    with pytest.raises(TypeError, match='results_dir'):
+        ICA.main('input.fif')
+
+
+def test_cmdline_requires_results_dir(monkeypatch):
+    monkeypatch.setattr(
+        'sys.argv',
+        ['ICA.py', '-filename', 'input.fif', '-line_freq', '60'],
+    )
+
+    with pytest.raises(SystemExit) as error:
+        ICA.cmdline()
+
+    assert error.value.code == 2
+
+
+def test_require_model_weights_returns_compatible_model(monkeypatch, tmp_path):
+    model_path = tmp_path / 'model_v2k3' / 'model_v2.keras'
+
+    monkeypatch.setattr(ICA.megnet_init, 'model_path', str(model_path))
+    monkeypatch.setattr(ICA.megnet_init, '_check_weights', lambda: True)
+
+    assert ICA._require_model_weights() == str(model_path)
+
+
+def test_require_model_weights_rejects_invalid_install(monkeypatch, tmp_path):
+    model_path = tmp_path / 'model_v2k3' / 'model_v2.keras'
+    monkeypatch.setattr(ICA.megnet_init, 'model_path', str(model_path))
+    monkeypatch.setattr(ICA.megnet_init, '_check_weights', lambda: False)
+
+    with pytest.raises(RuntimeError, match='Run `megnet_init`'):
+        ICA._require_model_weights()
+
+
+def test_classify_ica_saves_classification_vector(monkeypatch, tmp_path):
+    results_root = tmp_path / 'results'
+    output_dir = results_root / 'sample'
+    output_dir.mkdir(parents=True)
+    (output_dir / 'ICATimeSeries.mat').touch()
+    for index in range(1, 21):
+        (output_dir / f'component{index}.mat').touch()
+
+    class FakeKerasModels:
+        @staticmethod
+        def load_model(model_path, compile=False):
+            return object()
+
+    class FakeKeras:
+        models = FakeKerasModels
+
+    monkeypatch.setattr(ICA, '_require_model_weights', lambda: 'model_v2.keras')
+    monkeypatch.setitem(__import__('sys').modules, 'keras', FakeKeras)
+    monkeypatch.setattr(
+        ICA,
+        'fPredictChunkAndVoting_parrallel',
+        lambda model, arrTS, arrSP: (
+            np.array([
+                [1, 0, 0, 0],
+                [0, 1, 0, 0],
+                [0, 0, 1, 0],
+            ]),
+            None,
+        ),
+    )
+
+    def fake_loadmat(path):
+        if str(path).endswith('ICATimeSeries.mat'):
+            return {'arrICATimeSeries': np.ones((10, 3))}
+        return {'array': np.ones((2, 2, 3))}
+
+    monkeypatch.setattr('scipy.io.loadmat', fake_loadmat)
+
+    result = ICA.classify_ica(
+        results_dir=str(results_root),
+        outbasename='sample',
+        filename='ignored.fif',
+    )
+
+    classification_path = output_dir / 'megnet_classification.npy'
+    assert classification_path.is_file()
+    np.testing.assert_array_equal(np.load(classification_path), np.array([0, 1, 2]))
+    np.testing.assert_array_equal(result['classes'], np.array([0, 1, 2]))

MEGnet/tests/test_megnet_init.py

@@ -1,37 +1,55 @@
 import json
 import os.path as op
 
+import pytest
+
 from MEGnet import megnet_init
 
 
-def _set_config_path(monkeypatch, tmp_path):
+def _set_model_paths(monkeypatch, tmp_path):
     weights_path = tmp_path / 'model_v2k3'
     monkeypatch.setattr(megnet_init, 'weights_path', str(weights_path))
+    monkeypatch.setattr(megnet_init, 'model_path', str(weights_path / 'model_v2.keras'))
     monkeypatch.setattr(megnet_init, 'config_path', str(weights_path / 'config.json'))
     return weights_path
 
 
 def test_check_weights_requires_config_json(monkeypatch, tmp_path):
-    _set_config_path(monkeypatch, tmp_path).mkdir()
+    weights_path = _set_model_paths(monkeypatch, tmp_path)
+    weights_path.mkdir()
+    (weights_path / 'model_v2.keras').touch()
+
+    assert megnet_init._check_weights() is False
+
+
+def test_check_weights_requires_model_file(monkeypatch, tmp_path):
+    weights_path = _set_model_paths(monkeypatch, tmp_path)
+    weights_path.mkdir()
+    with open(op.join(weights_path, 'config.json'), 'w', encoding='utf-8') as fid:
+        json.dump({'model_version': 'v2.2'}, fid)
 
     assert megnet_init._check_weights() is False
 
 
-def test_check_weights_rejects_model_version_not_greater_than_min_version(monkeypatch, tmp_path):
-    weights_path = _set_config_path(monkeypatch, tmp_path)
+def test_check_weights_rejects_model_version_below_min_version(monkeypatch, tmp_path):
+    weights_path = _set_model_paths(monkeypatch, tmp_path)
     weights_path.mkdir()
-    monkeypatch.setattr(megnet_init, 'min_model_version', 'v2.1')
+    (weights_path / 'model_v2.keras').touch()
+    monkeypatch.setattr(megnet_init, 'min_model_version', 'v2.2')
     with open(op.join(weights_path, 'config.json'), 'w', encoding='utf-8') as fid:
         json.dump({'model_version': 'v2.1'}, fid)
 
     assert megnet_init._check_weights() is False
 
 
-def test_check_weights_accepts_model_version_greater_than_min_version(monkeypatch, tmp_path):
-    weights_path = _set_config_path(monkeypatch, tmp_path)
+@pytest.mark.parametrize('model_version', ['v2.2', 'v2.3'])
+def test_check_weights_accepts_compatible_model_version(
+        monkeypatch, tmp_path, model_version):
+    weights_path = _set_model_paths(monkeypatch, tmp_path)
     weights_path.mkdir()
-    monkeypatch.setattr(megnet_init, 'min_model_version', 'v2.1')
+    (weights_path / 'model_v2.keras').touch()
+    monkeypatch.setattr(megnet_init, 'min_model_version', 'v2.2')
     with open(op.join(weights_path, 'config.json'), 'w', encoding='utf-8') as fid:
-        json.dump({'model_version': 'v2.2'}, fid)
+        json.dump({'model_version': model_version}, fid)
 
     assert megnet_init._check_weights() is True