function for lar inside / outside NMS (l200)

docs/source/manual/meta.md

@@ -136,6 +136,13 @@ Supported fields per `simid`:
   - a list of the above strings to combine multiple volume patterns
   - `~vertices:NAME` to used the vertex positions similated by the `vtx` tier
     generator `NAME` (see {ref}`vtx-tier-meta`).
+  - `~function:NAME` to use a user-defined function to generate macro commands.
+    `NAME` should be in a format:
+  ```
+    module.function(<...>,*args,**kwargs)
+  ```
+  see {func}`legendsimflow.commands.get_confinement_from_function` for more
+  details. This function should return a list of the _remage_ macro commands.
 - `primaries_per_job` — integer, the number of primaries per job; becomes
   `N_EVENTS` in the macro file.
 - `number_of_jobs` — integer, how many jobs to split the simulation into.

tests/conftest.py

@@ -9,6 +9,7 @@
 from dbetto import AttrsDict
 from legendmeta import LegendMetadata
 from legendtestdata import LegendTestData
+from pygeoml200 import core
 
 testprod = Path(__file__).parent / "dummyprod"
 config_filename = testprod / "simflow-config.yaml"
@@ -26,6 +27,15 @@ def legend_test_metadata(legend_testdata):
     return LegendMetadata(legend_testdata["legend/metadata"])
 
 
+@pytest.fixture(scope="session")
+def test_generate_gdml(config):
+    geom_config = config.metadata.simprod.config.geom["l200p02-geom-config"]
+
+    return core.construct(
+        use_detailed_fiber_model=False, config=geom_config, public_geometry=True
+    )
+
+
 def make_config(legend_testdata):
     with config_filename.open() as f:
         config = yaml.safe_load(f)

tests/dummyprod/inputs/simprod/config/tier/stp/l200p02/simconfig.yaml

@@ -72,3 +72,17 @@ exotic_physics_hpge:
   confinement: ~defines:hpge_bulk
   primaries_per_job: 10_000
   number_of_jobs: 4
+
+lar_inside:
+  template: $_/templates/default.mac
+  generator: ~defines:K42
+  confinement: "~function:legendsimflow.confine.get_lar_minishroud_confine_commands(<...>,'minishroud_tube*',True,lar_name= 'liquid_argon')"
+  primaries_per_job: 10_000
+  number_of_jobs: 4
+
+lar_outside:
+  template: $_/templates/default.mac
+  generator: ~defines:K42
+  confinement: "~function:legendsimflow.confine.get_lar_minishroud_confine_commands(<...>,'minishroud_tube*',False,lar_name= 'liquid_argon')"
+  primaries_per_job: 10_000
+  number_of_jobs: 4

tests/test_aggregate.py

@@ -33,7 +33,7 @@ def test_simid_harvesting(config):
     simids = agg.gen_list_of_all_simids(config)
     assert isinstance(simids, type({}.keys()))
     assert all(isinstance(s, str) for s in simids)
-    assert len(simids) == 9
+    assert len(simids) == 11
 
 
 def test_simid_outputs(config):
@@ -165,4 +165,4 @@ def test_currmod_stuff(config):
 
 def test_tier_evt_stuff(config):
     files = agg.gen_list_of_all_tier_cvt_outputs(config)
-    assert len(files) == 9
+    assert len(files) == 11

tests/test_confine.py

@@ -0,0 +1,62 @@
+from __future__ import annotations
+
+import pytest
+
+from legendsimflow import commands, confine
+
+
+def test_get_lar_minishroud_confine_commands(test_generate_gdml):
+    lines = confine.get_lar_minishroud_confine_commands(test_generate_gdml, inside=True)
+
+    assert len(lines) > 0
+    assert isinstance(lines, list)
+
+    assert "/RMG/Generator/Confinement/Geometrical/AddSolid Cylinder" in lines
+
+    lines_outside = confine.get_lar_minishroud_confine_commands(
+        test_generate_gdml, inside=False
+    )
+
+    assert len(lines_outside) > 0
+    assert isinstance(lines_outside, list)
+
+    assert (
+        "/RMG/Generator/Confinement/Geometrical/AddExcludedSolid Cylinder"
+        in lines_outside
+    )
+
+    with pytest.raises(ValueError):
+        confine.get_lar_minishroud_confine_commands(
+            test_generate_gdml, pattern="non_existent_pattern*"
+        )
+    with pytest.raises(ValueError):
+        # exist pattern but not a nms
+        confine.get_lar_minishroud_confine_commands(test_generate_gdml, pattern="V**")
+
+    # test with eval
+
+    lines_eval = commands.get_confinement_from_function(
+        "legendsimflow.confine.get_lar_minishroud_confine_commands(<...>,inside=True)",
+        test_generate_gdml,
+    )
+    assert lines_eval == lines
+
+    lines_outside_eval = commands.get_confinement_from_function(
+        "legendsimflow.confine.get_lar_minishroud_confine_commands(<...>,inside=False)",
+        test_generate_gdml,
+    )
+    assert lines_outside_eval == lines_outside
+
+    # test with string
+    lines_eval_string = commands.get_confinement_from_function(
+        "legendsimflow.confine.get_lar_minishroud_confine_commands(<...>,lar_name= 'liquid_argon',inside=True)",
+        test_generate_gdml,
+    )
+    assert lines_eval_string == lines
+
+    # without any kwarg
+    lines_eval_args = commands.get_confinement_from_function(
+        "legendsimflow.confine.get_lar_minishroud_confine_commands(<...>,'minishroud_tube*',True,lar_name= 'liquid_argon')",
+        test_generate_gdml,
+    )
+    assert lines_eval_args == lines

workflow/src/legendsimflow/commands.py

@@ -14,11 +14,14 @@
 # along with this program.  If not, see <https://www.gnu.org/licenses/>.
 from __future__ import annotations
 
+import ast
+import importlib
 import shlex
 from copy import copy
 from pathlib import Path
 
 import legenddataflowscripts as lds
+import pyg4ometry
 
 from . import SimflowConfig, nersc, patterns, utils
 from .exceptions import SimflowConfigError
@@ -98,12 +101,15 @@ def remage_run(
     if not isinstance(output, Path):
         output = Path(str(output))
 
+    # geometry is read only here
+    geom = nersc.dvs_ro(config, geom)
+
     # get the config block for this tier/simid
     block = f"simprod.config.tier.{tier}.{config.experiment}.simconfig.{simid}"
     sim_cfg = get_simconfig(config, tier, simid=simid)
 
     # get macro
-    macro_text, _ = make_remage_macro(config, simid, tier=tier)
+    macro_text, _ = make_remage_macro(config, simid, tier=tier, geom=geom)
 
     # need some modifications if this is a benchmark run
     try:
@@ -153,7 +159,7 @@ def remage_run(
         "--procs",
         str(procs),
         "--gdml-files",
-        str(nersc.dvs_ro(config, geom)),
+        str(geom),
         "--output-file",
         str(output),
     ]
@@ -202,9 +208,84 @@ def _confine_by_volume(
     return lines
 
 
+# Extract function path
+def _get_full_name(node: ast.AST) -> str:
+    """Get the name of the function being called, including the module path if it's an attribute access."""
+
+    if isinstance(node, ast.Name):
+        return node.id
+    if isinstance(node, ast.Attribute):
+        return _get_full_name(node.value) + "." + node.attr
+    msg = f"unsupported node type {type(node)} in function path!"
+    raise ValueError(msg)
+
+
+def get_confinement_from_function(
+    function_string: str, reg: pyg4ometry.gdml.Registry
+) -> list[str]:
+    """Get the confinement commands for a function defined in the GDML.
+
+    The function string must correspond to the following format:
+
+    .. code-block::
+
+        module.function(<...>, arg=...)
+
+    where ``<...>`` will be replaced with the
+    :class:`pyg4ometry.gdml.Registry` instance for the geometry.
+
+    Parameters
+    ----------
+    function_string
+        String describing the function to be used.
+    reg
+        The pyg4ometry registry containing the geometry information.
+
+    Returns
+    -------
+    list[str]
+        A list of remage confinement commands corresponding to the
+        function definition.
+    """
+
+    if "<...>" not in function_string:
+        msg = "the function string must contain the placeholder <...> for the registry argument!"
+        raise ValueError(msg)
+
+    function_string = function_string.replace("<...>", "___OBJ___")
+
+    tree = ast.parse(function_string, mode="eval")
+
+    if not isinstance(tree.body, ast.Call):
+        msg = f"{tree.body} is not a function call!"
+        raise ValueError(msg)
+
+    call_node = tree.body
+    full_name = _get_full_name(call_node.func)
+
+    # Resolve function object
+    module_path, func_name = full_name.rsplit(".", 1)
+    module = importlib.import_module(module_path)
+    func = getattr(module, func_name)
+
+    # Process arguments
+    args = []
+    for arg in call_node.args:
+        if isinstance(arg, ast.Name) and arg.id == "___OBJ___":
+            args.append(reg)
+        else:
+            args.append(ast.literal_eval(arg))
+
+    kwargs = {}
+    for kw in call_node.keywords:
+        kwargs[kw.arg] = ast.literal_eval(kw.value)
+
+    return func(*args, **kwargs)
+
+
 def make_remage_macro(
-    config: SimflowConfig, simid: str, tier: str = "stp"
-) -> (str, Path):
+    config: SimflowConfig, simid: str, tier: str = "stp", geom: str | None = None
+) -> tuple[str, Path]:
     """Render the remage macro for a given simulation and write it to disk.
 
     This function reads the simulation configuration for the provided tier/simid,
@@ -226,7 +307,8 @@ def make_remage_macro(
         Simulation identifier to select the simconfig.
     tier
         Simulation tier (e.g. "stp", "ver", ...). Default is "stp".
-
+    geom
+        Path to the geometry file.
     Returns
     -------
     A tuple with:
@@ -304,6 +386,20 @@ def make_remage_macro(
                     "/RMG/Generator/Confine FromFile",
                     f"/RMG/Generator/Confinement/FromFile/FileName {nersc.dvs_ro(config, vtx_file)}",
                 ]
+            elif sim_cfg.confinement.startswith("~function:"):
+                # in this case we need to parse the GDML to get the actual confinement commands
+                if not geom:
+                    msg = (
+                        "confinement uses '~function:' but no geometry (geom) was provided; "
+                        f"please set either {block}.geom or adjust {block}.confinement",
+                        f"{block}.confinement",
+                    )
+                    raise SimflowConfigError(*msg)
+
+                reg = pyg4ometry.gdml.Reader(str(geom)).getRegistry()
+
+                func_name = sim_cfg.confinement.removeprefix("~function:")
+                confinement = get_confinement_from_function(func_name, reg)
 
             elif sim_cfg.confinement.startswith("~defines:"):
                 key = sim_cfg.confinement.removeprefix("~defines:")
@@ -338,7 +434,7 @@ def make_remage_macro(
             msg = (
                 (
                     "the field must be a str or list[str] prefixed by "
-                    "~define: / ~volumes.surface: / ~volumes.bulk:"
+                    "~define: / ~volumes.surface: / ~volumes.bulk: / ~function: or ~vertices:"
                 ),
                 f"{block}.confinement",
             )