Initial pass on imat rules

rundetection/rules/imat_rules.py

@@ -8,6 +8,7 @@
 from pathlib import Path
 
 from rundetection.exceptions import RuleViolationError
+from rundetection.ingestion.ingest import load_h5py_dataset
 from rundetection.rules.rule import Rule
 
 if typing.TYPE_CHECKING:
@@ -78,12 +79,46 @@ def verify(self, job_request: JobRequest) -> None:
             # Find file with run number in it, often ending with .csv, then search for a dir in the same directory as it
             # called Tomo.
             imat_dir_path = find_correct_tomo_dir(exp_dir_path, str(job_request.run_number))
+            if imat_dir_path is None:
+                # We haven't found it yet, search the expected path for the correct directory
+                for child in exp_dir_path.iterdir():
+                    if child.is_dir():
+                        imat_dir_path = find_correct_tomo_dir(child, str(job_request.run_number))
+                        if imat_dir_path is not None:
+                            break
 
         if imat_dir_path is not None and imat_dir_path.exists():
+            job_request.additional_values["recon"] = "true"
+            job_request.additional_values["ngem"] = "false"
             job_request.additional_values["images_dir"] = str(imat_dir_path)
             job_request.additional_values["runno"] = job_request.run_number
         else:
-            logger.error("Images dir could not be found for experiment number: %s", job_request.experiment_number)
+            # Given there is no Images, let's check for an nGEM run.
+            # INES is temporary here and should be adjustable via env vars. Current technical limitation forces IMAT
+            # data here.
+            ngem_dir = os.environ.get("IMAT_NGEM_DIR", "/ngem/nGEM-INES")
+
+            # Grab the cycle from the .nxs file from the path /raw_data_1/run_cycle value in the form 26_1
+            cycle_str = load_h5py_dataset(job_request.filepath).get("run_cycle")[0].decode("utf-8")
+            cycle_year, cycle_num = cycle_str.split("_")
+
+            # Check if the correct dir exists in the nGEM dir for IMAT.
+            exp_dir_path = Path(ngem_dir) / "DATA" / f"IMAT_20{cycle_year}_0{cycle_num}"
+            if exp_dir_path.exists():
+                possible_path = exp_dir_path / f"IMAT{job_request.run_number:08d}"
+                if possible_path.exists() and possible_path.is_dir():
+                    # We found it
+                    job_request.additional_values["recon"] = "false"
+                    job_request.additional_values["ngem"] = "true"
+                    job_request.additional_values["ngem_path"] = str(possible_path)
+                    output_path = Path(str(exp_dir_path) + "_nxs") / "RUN"
+                    job_request.additional_values["ngem_output_path"] = str(output_path)
+
+        if "ngem" not in job_request.additional_values and "recon" not in job_request.additional_values:
+            # We did not find either an IMAT or nGEM detector run.
+            logger.error(
+                "Images dir and nGEM run could not be found for experiment number: %s", job_request.experiment_number
+            )
             raise RuleViolationError(
-                "Images dir could not be found for experiment number: %s", job_request.experiment_number
+                "Images dir and nGEM run could not be found for experiment number: %s", job_request.experiment_number
             )

test/rules/test_imat_rules.py

@@ -3,14 +3,15 @@
 import os
 import tempfile
 from pathlib import Path
+from unittest.mock import patch
 
 import pytest
 
 from rundetection.exceptions import RuleViolationError
 from rundetection.job_requests import JobRequest
 from rundetection.rules.imat_rules import IMATFindImagesRule
 
-EXPECTED_ADDITIONAL_VALUES_LEN = 2
+EXPECTED_ADDITIONAL_VALUES_LEN = 4
 RUN_NUMBER = 100
 
 
@@ -36,6 +37,31 @@ def job_request():
     )
 
 
+def test_delete_me(job_request):
+    """Test imat rules can find images successfully"""
+    with tempfile.TemporaryDirectory() as tmpdirname:
+        # Setup
+        os.environ["IMAT_DIR"] = tmpdirname
+        run_number = 36779
+        job_request.run_number = run_number
+        job_request.experiment_number = "2322026"
+        # Create required structure: a file with run number and a Tomo dir
+        exp_dir = Path(tmpdirname).joinpath("RB2322026")
+        exp_dir.mkdir(parents=True, exist_ok=True)
+        exp_dir.joinpath("run36779.csv").touch()
+        tomo_dir = exp_dir.joinpath("Tomo")
+        tomo_dir.mkdir()
+
+        # Test
+        rule = IMATFindImagesRule(True)
+        rule.verify(job_request)
+
+        # Assertions
+        assert len(job_request.additional_values) == EXPECTED_ADDITIONAL_VALUES_LEN
+        assert job_request.additional_values["images_dir"] == str(tomo_dir)
+        assert job_request.additional_values["runno"] == run_number
+
+
 def test_imat_find_images_success(job_request):
     """Test imat rules can find images successfully"""
     with tempfile.TemporaryDirectory() as tmpdirname:
@@ -56,6 +82,8 @@ def test_imat_find_images_success(job_request):
         assert len(job_request.additional_values) == EXPECTED_ADDITIONAL_VALUES_LEN
         assert job_request.additional_values["images_dir"] == str(tomo_dir)
         assert job_request.additional_values["runno"] == RUN_NUMBER
+        assert job_request.additional_values["recon"] == "true"
+        assert job_request.additional_values["ngem"] == "false"
 
 
 def test_imat_find_images_tomo_first(job_request):
@@ -105,41 +133,104 @@ def test_imat_find_images_missing_tomo(job_request):
     with tempfile.TemporaryDirectory() as tmpdirname:
         # Setup
         os.environ["IMAT_DIR"] = tmpdirname
+        os.environ["NGEM_DIR"] = tmpdirname  # Also set this to avoid loading nexus
         exp_dir = Path(tmpdirname).joinpath("RB12345")
         exp_dir.mkdir(parents=True, exist_ok=True)
         exp_dir.joinpath("run100.csv").touch()
 
         # Test
         rule = IMATFindImagesRule(True)
-        with pytest.raises(RuleViolationError):
-            rule.verify(job_request)
+        with patch("rundetection.rules.imat_rules.load_h5py_dataset") as mock_load:
+            mock_load.return_value.get.return_value = [b"26_1"]
+            with pytest.raises(RuleViolationError):
+                rule.verify(job_request)
 
 
 def test_imat_find_images_missing_run_file(job_request):
     """Test imat rules fail when run number file is missing"""
     with tempfile.TemporaryDirectory() as tmpdirname:
         # Setup
         os.environ["IMAT_DIR"] = tmpdirname
+        os.environ["NGEM_DIR"] = tmpdirname
         exp_dir = Path(tmpdirname).joinpath("RB12345")
         exp_dir.mkdir(parents=True, exist_ok=True)
         exp_dir.joinpath("Tomo").mkdir()
 
         # Test
         rule = IMATFindImagesRule(True)
-        with pytest.raises(RuleViolationError):
-            rule.verify(job_request)
+        with patch("rundetection.rules.imat_rules.load_h5py_dataset") as mock_load:
+            mock_load.return_value.get.return_value = [b"26_1"]
+            with pytest.raises(RuleViolationError):
+                rule.verify(job_request)
 
 
 def test_imat_find_images_failure(job_request):
     """Test imat rules act's correctly when it fails to find experiment dir"""
     with tempfile.TemporaryDirectory() as tmpdirname:
         # Setup
         os.environ["IMAT_DIR"] = tmpdirname
+        os.environ["NGEM_DIR"] = tmpdirname
 
         # Test
         rule = IMATFindImagesRule(True)
-        with pytest.raises(RuleViolationError):
-            rule.verify(job_request)
+        with patch("rundetection.rules.imat_rules.load_h5py_dataset") as mock_load:
+            mock_load.return_value.get.return_value = [b"26_1"]
+            with pytest.raises(RuleViolationError):
+                rule.verify(job_request)
 
         # Assertions
         assert len(job_request.additional_values) == 0
+
+
+def test_imat_find_images_ngem_success(job_request):
+    """Test imat rules can find nGEM images successfully when Tomo is missing"""
+    with tempfile.TemporaryDirectory() as tmpdirname:
+        # Setup
+        os.environ["IMAT_DIR"] = tmpdirname
+        os.environ["NGEM_DIR"] = tmpdirname
+
+        # Cycle information
+        cycle_year = "26"
+        cycle_num = "1"
+        cycle_str = f"{cycle_year}_{cycle_num}"
+
+        # Create nGEM directory structure
+        ngem_data_dir = Path(tmpdirname) / "DATA" / f"IMAT_20{cycle_year}_0{cycle_num}"
+        possible_path = ngem_data_dir / f"IMAT{job_request.run_number:08d}"
+        possible_path.mkdir(parents=True, exist_ok=True)
+
+        # Test
+        rule = IMATFindImagesRule(True)
+        with patch("rundetection.rules.imat_rules.load_h5py_dataset") as mock_load:
+            mock_load.return_value.get.return_value = [cycle_str.encode("utf-8")]
+            rule.verify(job_request)
+
+        # Assertions
+        assert job_request.additional_values["recon"] == "false"
+        assert job_request.additional_values["ngem"] == "true"
+        assert job_request.additional_values["ngem_path"] == str(possible_path)
+        expected_output_path = str(ngem_data_dir) + "_nxs/RUN"
+        assert job_request.additional_values["ngem_output_path"] == expected_output_path
+
+
+def test_imat_find_images_nested_success(job_request):
+    """Test imat rules can find images in a nested directory"""
+    with tempfile.TemporaryDirectory() as tmpdirname:
+        # Setup
+        os.environ["IMAT_DIR"] = tmpdirname
+        exp_dir = Path(tmpdirname).joinpath("RB12345")
+        nested_dir = exp_dir.joinpath("nested")
+        nested_dir.mkdir(parents=True, exist_ok=True)
+
+        # Create required structure in nested dir
+        nested_dir.joinpath("run100.csv").touch()
+        tomo_dir = nested_dir.joinpath("Tomo")
+        tomo_dir.mkdir()
+
+        # Test
+        rule = IMATFindImagesRule(True)
+        rule.verify(job_request)
+
+        # Assertions
+        assert job_request.additional_values["images_dir"] == str(tomo_dir)
+        assert job_request.additional_values["runno"] == RUN_NUMBER