From a5a86e80981a36887f7fd0e550369e23e104aa10 Mon Sep 17 00:00:00 2001
From: Patrick Shriwise <pshriwise@gmail.com>
Date: Wed, 18 Sep 2024 22:42:01 -0500
Subject: [PATCH] Adding some simple cell by cell tests

---
 src/openmc_cad_adapter/to_cubit_journal.py | 40 +++++++++++++++++++---
 test/test_examples.py                      | 18 ++++++++++
 2 files changed, 53 insertions(+), 5 deletions(-)

diff --git a/src/openmc_cad_adapter/to_cubit_journal.py b/src/openmc_cad_adapter/to_cubit_journal.py
index 0055bee..cf92ae6 100644
--- a/src/openmc_cad_adapter/to_cubit_journal.py
+++ b/src/openmc_cad_adapter/to_cubit_journal.py
@@ -175,7 +175,31 @@ def vector_to_euler_xyz(v):
     return phi * oe, theta * oe, psi * oe
 
 
-def to_cubit_journal(geometry : openmc.Geometry, world : Iterable[Real] = None, cells: Iterable[int] = None, filename: str = "openmc.jou", to_cubit: bool = False, seen: set = set()):
+def to_cubit_journal(geometry : openmc.Geometry, world : Iterable[Real] = None,
+                     cells: Iterable[int, openmc.Cell] = None,
+                     filename: str = "openmc.jou",
+                     to_cubit: bool = False,
+                     seen: set = set()):
+    """Convert an OpenMC geometry to a Cubit journal.
+
+    Parameters
+    ----------
+        geometry : openmc.Geometry
+            The geometry to convert to a Cubit journal.
+        world : Iterable[Real], optional
+            Extents of the model in X, Y, and Z. Defaults to None.
+        cells : Iterable[int, openmc.Cell], optional
+            List of cells or cell IDs to write to individual journal files. If None,
+            all cells will be written to the same journal file. Defaults to None.
+        filename : str, optional
+            Output filename. Defaults to "openmc.jou".
+        to_cubit : bool, optional
+            Uses the cubit Python module to write the model as a .cub5 file.
+            Defaults to False.
+        seen : set, optional
+            Internal parameter.
+
+    """
 
     if not filename.endswith('.jou'):
         filename += '.jou'
@@ -183,6 +207,9 @@ def to_cubit_journal(geometry : openmc.Geometry, world : Iterable[Real] = None,
     if isinstance(geometry, openmc.Model):
         geometry = geometry.geometry
 
+    if cells is not None:
+        cells = [c if not isinstance(c, openmc.Cell) else c.id for c in cells]
+
     if to_cubit:
         try:
             import cubit
@@ -920,14 +947,17 @@ def do_cell(cell, cell_ids: Iterable[int] = None):
         after = len( cmds )
 
         if cell_ids is not None and cell.id in cell_ids:
-            with open( filename + f"_cell{cell.id}", "w" ) as f:
+            if filename.endswith(".jou"):
+                cell_filename = filename[:-4] + f"_cell{cell.id}.jou"
+            else:
+                cell_filename = filename + f"_cell{cell.id}"
+            with open(cell_filename, "w" ) as f:
                 for x in cmds[before:after]:
                     f.write( x + "\n" )
 
     for cell in geom.root_universe._cells.values():
-        if cells:
-            if cell.id in cells:
-                do_cell( cell, cell_ids=cells)
+        if cells is not None and cell.id in cells:
+            do_cell( cell, cell_ids=cells)
         else:
             do_cell( cell )
 
diff --git a/test/test_examples.py b/test/test_examples.py
index 1819366..491b2b2 100644
--- a/test/test_examples.py
+++ b/test/test_examples.py
@@ -49,6 +49,24 @@ def test_examples(example, request):
     diff_files(output, gold_file)
 
 
+def test_cell_by_cell_conversion(request):
+    openmc.reset_auto_ids()
+    exec(open(OPENMC_EXAMPLES_DIR / "pincell/build_xml.py").read())
+
+    openmc.reset_auto_ids()
+    model = openmc.Model.from_xml()
+
+    cell_ids = list(model.geometry.get_all_cells().keys())
+
+    world = [500, 500, 500]
+    output = 'pincell'
+    to_cubit_journal(model.geometry, world=world, cells=cell_ids, filename=output)
+
+    for cell_id in cell_ids:
+        output = f'pincell_cell{cell_id}.jou'
+        gold_file = request.path.parent / Path('gold') / Path(output)
+        diff_files(output, gold_file)
+
 @pytest.mark.parametrize("example", examples, ids=example_name)
 def test_examples_cli(example, request):