SHIELDHIT support added

.travis.yml

@@ -63,7 +63,7 @@ install:
 
 # if scipy is used, then installation from pip last longer than 10min, thus we add travis_wait
 script:
-  - travis_wait tox --notest && tox -- -n 8
+  - travis_wait tox --notest && tox -- -n 1
 
 after_success:
   # deploy to normal if tag is present

MANIFEST.in

@@ -14,6 +14,7 @@
   include mcpartools/mcengine/data/*.sh
   include mcpartools/scheduler/data/*.sh
   include tests/res/*.inp
+  include tests/res/shieldhit/*.dat
 
   prune .travis
   exclude .gitkeep

mcpartools/generatemc.py

@@ -65,9 +65,9 @@ def main(args=sys.argv[1:]):
 
     opt = Options(args)
     generator = Generator(options=opt)
-    generator.run()
+    ret_code = generator.run()
 
-    return 0
+    return ret_code
 
 if __name__ == '__main__':
     sys.exit(main(sys.argv[1:]))

mcpartools/generator.py

@@ -60,7 +60,7 @@ def __init__(self, options):
     def run(self):
         if not self.options.valid:
             logger.error("Invalid options, aborting run")
-            return
+            return None
 
         # generate main dir according to date
         self.generate_main_dir()
@@ -77,6 +77,8 @@ def run(self):
         # save logs
         self.save_logs()
 
+        return 0
+
     def generate_main_dir(self):
         dir_name = time.strftime("run_%Y%m%d_%H%M%S")
         logger.debug("Generated main directory name: " + dir_name)
@@ -101,7 +103,7 @@ def generate_workspace(self):
             os.mkdir(jobdir_path)
             logger.debug("Generated job directory path: " + jobdir_path)
 
-            self.mc_engine.randomize(jobid + 1)
+            self.mc_engine.randomize(new_seed=jobid + 1)
             self.mc_engine.set_particle_no(self.options.particle_no)
             self.mc_engine.save_input(jobdir_path)
 

mcpartools/mcengine/data/collect_shieldhit.sh

@@ -0,0 +1,12 @@
+#!/usr/bin/env bash
+
+# Exit immediately if a simple command exits with a non-zero status.
+set -e
+
+# make output folder
+mkdir -p {output_dir:s}/output
+
+# copy all binary output files to the same folder
+for f in {output_dir:s}/workspace/job_*/*.bdo; do
+  cp $f {output_dir:s}/output/
+done

mcpartools/mcengine/data/run_shieldhit.sh

@@ -0,0 +1,26 @@
+#!/usr/bin/env bash
+
+# Exit immediately if a simple command exits with a non-zero status.
+set -e
+
+# location of SHIELDHIT binary file
+SHIELDHIT_BIN={shieldhit_bin:s}
+
+# working directory, output files will be saved here
+WORK_DIR={working_directory:s}
+
+# number of particles per job
+PARTICLE_NO={particle_no:d}
+
+# seed of RNG
+RNG_SEED={rnd_seed:d}
+
+# main SHIELDHIT12A input files
+BEAM_FILE={beam_file:s}
+GEO_FILE={geo_file:s}
+MAT_FILE={mat_file:s}
+DETECT_FILE={detect_file:s}
+
+# execute simulation
+$SHIELDHIT_BIN --beamfile=$BEAM_FILE --geofile=$GEO_FILE --matfile=$MAT_FILE --detectfile=$DETECT_FILE -n $PARTICLE_NO -N $RNG_SEED $WORK_DIR
+

mcpartools/mcengine/shieldhit.py

@@ -1,38 +1,80 @@
+import logging
 import os
+from pkg_resources import resource_string
 
 from mcpartools.mcengine.mcengine import Engine
 
+logger = logging.getLogger(__name__)
+
 
 class ShieldHit(Engine):
 
-    run_script = """#!/bin/bash
-shieldhit .... {:s}
-"""
+    default_run_script_path = os.path.join('data', 'run_shieldhit.sh')
+    collect_script = os.path.join('data', 'collect_shieldhit.sh')
+
+    def __init__(self, input_path, mc_run_script):
+        Engine.__init__(self, input_path, mc_run_script)
+
+        # user didn't provided path to input scripts, use default
+        if self.run_script_path is None:
+            tpl = resource_string(__name__, self.default_run_script_path)
+            self.run_script_content = tpl.decode('ascii')
+            logger.debug("Using default run script: " + self.default_run_script_path)
+        else:
+            tpl_fd = open(self.run_script_path, 'r')
+            self.run_script_content = tpl_fd.read()
+            tpl_fd.close()
+            logger.debug("Using user run script: " + self.run_script_path)
+
+        self.collect_script_content = resource_string(__name__, self.collect_script).decode('ascii')
+
+        self.particle_no = 1
+        self.rng_seed = 1
 
     @property
     def input_files(self):
-        base = self.input_path
+        base = os.path.abspath(self.input_path)
         # TODO add *.txt files with stopping power
         # TODO add *.ctx/.hed files with CT scans
-        files = ['beam.dat', 'geo.dat', 'mat.dat', 'detect.dat']
-        result = [os.path.join(base, f) for f in files]
+        files = ('beam.dat', 'geo.dat', 'mat.dat', 'detect.dat')
+        result = (os.path.join(base, f) for f in files)
         return result
 
-    def randomize(self, output_dir, new_seed):
-        pass
+    def randomize(self, new_seed, output_dir=None):
+        self.rng_seed = new_seed
 
     def set_particle_no(self, particle_no):
-        "in memory"
-        pass
+        self.particle_no = particle_no
 
     def save_input(self, output_dir):
-        "nothing, we will use switches -b, -g etc"
-        pass
+        logger.info("input files are not modified, we will used shieldhit switches instead")
 
-    def save_run_script(self, output_dir):
-        contents = self.run_script.format(output_dir)
+    def save_run_script(self, output_dir, job_id):
+        beam_file, geo_file, mat_file, detect_file = self.input_files
+        abs_output_dir = os.path.abspath(output_dir)
+        contents = self.run_script_content.format(
+            shieldhit_bin='shieldhit',
+            working_directory=abs_output_dir,
+            particle_no=self.particle_no,
+            rnd_seed=self.rng_seed,
+            beam_file=beam_file,
+            geo_file=geo_file,
+            mat_file=mat_file,
+            detect_file=detect_file
+        )
         out_file_name = "run.sh"
         out_file_path = os.path.join(output_dir, out_file_name)
         out_fd = open(out_file_path, 'w')
         out_fd.write(contents)
         out_fd.close()
+        os.chmod(out_file_path, 0o750)
+
+    def write_collect_script(self, output_dir):
+        output_dir_abs_path = os.path.abspath(output_dir)
+        contents = self.collect_script_content.format(output_dir=output_dir_abs_path)
+        out_file_name = "collect.sh"
+        out_file_path = os.path.join(output_dir, out_file_name)
+        out_fd = open(out_file_path, 'w')
+        out_fd.write(contents)
+        out_fd.close()
+        os.chmod(out_file_path, 0o750)

mcpartools/scheduler/data/run_slurm.sh

@@ -1,2 +1,7 @@
 #!/usr/bin/env bash
-{workspace_dir:s}/job_`printf %04d $SLURM_ARRAY_TASK_ID`/run.sh
+
+# Exit immediately if a simple command exits with a non-zero status.
+set -e
+
+# Run individual jobs
+{workspace_dir:s}/job_`printf %04d $SLURM_ARRAY_TASK_ID`/run.sh

tests/res/shieldhit/beam.dat

@@ -0,0 +1,10 @@
+*         Input file FOR023.DAT for the SHIELD Transport Code
+
+RNDSEED      	89736501     ! Random seed
+JPART0       	25           ! Incident particle type
+HIPROJ     	12.0    6.0  ! A and Z of heavy ion
+TMAX0      	391.0   0.0  ! Incident energy; (MeV/nucl)
+NSTAT           6000    3000 ! NSTAT, Step of saving
+STRAGG          2            ! Straggling: 0-Off 1-Gauss, 2-Vavilov
+MSCAT           2            ! Mult. scatt 0-Off 1-Gauss, 2-Moliere
+NUCRE           1            ! Nucl.Reac. switcher: 1-ON, 0-OFF

tests/res/shieldhit/detect.dat

@@ -0,0 +1,18 @@
+*----0---><----1---><----2---><----3---><----4---><----5---><----6--->
+CYL              0.0       0.0       0.0      10.0       7.0      30.0
+                   1         1       300        -1    ENERGY    ex_cyl
+*----0---><----1---><----2---><----3---><----4---><----5---><----6--->
+MSH             -5.0      -5.0       0.0       5.0       5.0      30.0
+                   1         1       300        -1    ENERGY   ex_zmsh
+*----0---><----1---><----2---><----3---><----4---><----5---><----6--->
+MSH             -1.0      -5.0      10.0       1.0       5.0      12.0
+                   1       100         1        -1    ENERGY   ex_ymsh
+*----0---><----1---><----2---><----3---><----4---><----5---><----6--->
+MSH             -5.0      -5.0       0.0       5.0       5.0      30.0
+                   1       100       300        -1    ENERGY  ex_yzmsh
+*----0---><----1---><----2---><----3---><----4---><----5---><----6--->
+GEOMAP          -1.0     -25.0     -15.0       1.0      25.0      35.0
+                   1        50        50                ZONE  ex_yzzon
+*----0---><----1---><----2---><----3---><----4---><----5---><----6--->
+GEOMAP         -25.0     -25.0      10.0      25.0      25.0      11.0
+                  50        50         1                ZONE  ex_xyzon

tests/res/shieldhit/geo.dat

@@ -0,0 +1,17 @@
+*---><---><--------><------------------------------------------------>
+    0    0           C12 200 MeV/A, H2O 30 cm cylinder, r=10, 1 zone
+*---><---><--------><--------><--------><--------><--------><-------->
+  RCC    1       0.0       0.0       0.0       0.0       0.0      30.0
+                10.0
+  RCC    2       0.0       0.0      -5.0       0.0       0.0      35.0
+                15.0
+  RCC    3       0.0       0.0     -10.0       0.0       0.0      40.0
+                20.0
+  END
+  001          +1
+  002          +2     -1
+  003          +3     -2      
+  END
+    1    2    3
+    1 1000    0
+

tests/res/shieldhit/mat.dat

@@ -0,0 +1,3 @@
+MEDIUM 1
+ICRU 276
+END

tests/test_generatemc.py

@@ -1,22 +1,55 @@
+import os
+import time
 import unittest
 
 from mcpartools import generatemc
 
 
-class TestFunMethod(unittest.TestCase):
+class TestRunGenerate(unittest.TestCase):
+
+    def setUp(self):
+        self.main_dir = os.path.join("tests", "res")
+
     def test_help(self):
-        self.assertRaises(SystemExit, generatemc.main, ["--help"])
+        time.sleep(1)  # TODO should be removed after fixing https://github.com/DataMedSci/mcpartools/issues/14
+        try:
+            generatemc.main(["--help"])
+        except SystemExit as e:
+            self.assertEqual(e.code, 0)
 
     def test_version(self):
-        self.assertRaises(SystemExit, generatemc.main, ["--version"])
-
-    def test_input_exit_code(self):
-        exit_code = generatemc.main(["-j", "2", "-p", "100", "tests/res/sample_fluka.inp"])
-        self.assertEqual(exit_code, 0)
+        time.sleep(1)
+        try:
+            generatemc.main(["--version"])
+        except SystemExit as e:
+            self.assertEqual(e.code, 0)
+
+    def test_fluka_input_ok(self):
+        time.sleep(1)
+        fluka_input = os.path.join(self.main_dir, "sample_fluka.inp")
+        ret_code = generatemc.main(["-j", "2", "-p", "100", fluka_input])
+        self.assertEqual(ret_code, 0)
+
+    def test_shieldhit_input_ok(self):
+        time.sleep(1)
+        shieldhit_input = os.path.join(self.main_dir, "shieldhit")
+        ret_code = generatemc.main(["-j", "2", "-p", "100", shieldhit_input])
+        self.assertEqual(ret_code, 0)
+
+    def test_shieldhit_scheduler_options(self):
+        time.sleep(1)
+        shieldhit_input = os.path.join(self.main_dir, "shieldhit")
+        ret_code = generatemc.main(["-j", "2", "-p", "100", "-s",
+                                    "--nodes=1 --ntasks-per-node=1 --mem=2000 --time=0:30:00",
+                                    shieldhit_input])
+        self.assertEqual(ret_code, 0)
 
     def test_input_bad(self):
-        self.assertRaises(SystemExit, generatemc.main, ["-j", "2"])
-
+        time.sleep(1)
+        try:
+            generatemc.main(["-j", "2"])
+        except SystemExit as e:
+            self.assertEqual(e.code, 2)
 
 if __name__ == '__main__':
     unittest.main()