Skip to content

Commit

Permalink
run: worked on automated mcnp runs
Browse files Browse the repository at this point in the history
  • Loading branch information
Arun Persaud committed Nov 10, 2024
1 parent caaed4a commit 59bb3b4
Showing 1 changed file with 149 additions and 136 deletions.
285 changes: 149 additions & 136 deletions src/pymcnp/cli/run.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -3,14 +3,11 @@
pymcnp run <file> [ --parallel=<threads> ] [ options ]
Options:
-p --parallel=<threads> Run files in parallel on <threads> threads.
-c --command=<command> Command to run.
-P --path=<path> Path to use.
-S --hosts=<hosts> Comma separated ist of hosts to run on.
-n --dry-run Don't run or create directories, just print what would happen
"""

import os
from pathlib import Path
import shutil
import subprocess
Expand All @@ -22,194 +19,216 @@
from ..files.inp import Inp, read_input
from . import _io

DEFAULT_NPP = 1000
DEFAULT_NPSMG = 1

def check_prog(name: str) -> None:
if shutil.which(name) is None:
print(f'[red]ERROR[/] Cannot find {name} program.')
sys.exit(3)


class Run:
"""
Encapsulates methods for running MCNP INP files.
"""Encapsulates methods for running a single PyMCNP INP instance.
``Run`` provides utilities for executing MCNP simulations concurrently or
in parallel. This class also stores the inputs and outputs of MNPC runs
in timestamped directories.
Can take either an Inp object or a filename that gets red in.
Always executes in a custom directory. By default this directory
doesn't get deleted.
Attributes:
inp: PyMCNP INP object to run.
path: Path to directory to store run inputs and outputs.
command: Terminal command to execute.
"""

def __init__(
self,
inp: Path | Inp,
inp: Inp | Path | str,
path: str | Path = Path('.'),
command: str = 'mcnp6',
dry_run: bool = False,
prefix: str = 'pymcnp',
run_name: str | None = None,
):
"""
Parameters:
path: Path to directory to store run inputs and outputs.
command: Terminal command to execute.
inp: The pymcnp input object or filename to run
path: Path to directory to store run inputs and outputs.
command: Terminal command to execute.
prefix: prefix of the run directory
run_name: name of the run, use a timestamp if not given
"""

self.command: str = command
self.dry_run: bool = dry_run
self.path: Path = Path(path)

self.inp = inp
self.prefix: str = prefix
self.run_name: str | None = run_name
self.run_path: Path | None = None # will be modified when we run
self.filename: Path | None = None # will be modified when we run

if isinstance(inp, Inp):
self.inp = inp
elif isinstance(inp, (str, Path)):
self.inp = read_input(inp)
else:
print('[red]Error[/] cannot parse input')

def prehook(self):
pass

def posthook(self):
pass

def parallel_prehook(self):
pass
def create_files(self, dry_run: bool = False):
"""Create directory and input files."""
if self.run_name is None:
self.run_name = _io.get_timestamp()

def parallel_posthook(self):
pass
self.run_path = self.path / f'{self.prefix}-{self.run_name}'
self.filename = self.run_path / f'input-{self.run_name}.i'

def check_prog(self, name: str) -> None:
if shutil.which(name) is None:
print(f'[red]ERROR[/] Cannot find {name} program.')
sys.exit(3)
if dry_run:
return

def run_single(self) -> Path:
"""
Runs a MCNP INP files.
self.run_path.mkdir(exist_ok=True, parents=True)
self.inp.to_mcnp_file(self.filename)

``run_single`` creates a directory to store a copy of the input file
passed to MCNP and the outputfiles generated by MCNP.
def run(self, *, dry_run: bool = False) -> Path:
"""Runs a PyMcnp INP instance.
Creates a custom directory and a copy of the input file. It
then calls MCNP. By default no cleanup is done.
Returns:
Path to run directory.
"""

self.check_prog(self.command)
if isinstance(self.inp, Inp):
timestamp = _io.get_timestamp()

directory_path = self.path / f'pymcnp-run-{timestamp}'
self.path = directory_path
inp_path = directory_path / f'pymcnp-inp-{timestamp}.inp'

directory_path.mkdir(exist_ok=True, parents=True)
self.inp.to_mcnp_file(inp_path)
"""
check_prog(self.command)

self.inp_path = inp_path
elif isinstance(self.inp, Path):
self.inp_path = self.inp
self.create_files(dry_run)

command_to_run = f'{self.command} i={self.inp_path.name}'
command_to_run = f'{self.command} i={self.filename.name}'

if self.dry_run:
if dry_run:
print('[yellow]INFO[/] Dry run:')
print(f' creating {self.run_path.absolute()}')
print(' executing prehook')
print(f' {command_to_run} in {self.path}')
print(f' {command_to_run}')
print(' executing posthook')

# some cleanup
self.run_path = None
self.filename = None
return self.path

self.prehook()
subprocess.run(command_to_run, cwd=self.path, shell=True)
subprocess.run(command_to_run, cwd=self.run_path, shell=True)
self.posthook()

return self.path

def run_parallel(self, count: int, hosts: str | None = None) -> Path:
"""
Runs MCNP INP files in parallel.
``run_parallel`` creates a directory to store a copy of the input file
passed to MCNP and directories containing each seperate run output
files. This method works by dividing the INP particles histories and
running the MCNP simulation many times concurrently.
Parameters:
count: Number of parallel threads to run.
Returns:
Path to run directory.
"""

if count <= 0:
_io.error('Invalid Count.')

# some error checking
self.check_prog(self.command)
self.check_prog('parallel')

timestamp = _io.get_timestamp()

if isinstance(self.inp, Path):
self.inp = read_input(self.inp)
def cleanup(self):
"""Helper function that can be used in posthooks."""
if self.run_path:
for file_or_dir in self.run_path.rglob('*'):
if file_or_dir.is_file():
file_or_dir.unlink()
else:
file_or_dir.rmdir()
self.run_path.rmdir()

# create the main directory
directory_path = self.path / f'pymcnp-runs-{timestamp}'

L = len(str(count)) # number of digits
class Parallel:
"""Run multiple `Run` instancnes.
if hosts is None:
hosts = os.getenv('SLURM_JOB_NODELIST')
At the moment we only create the input directories with the correct files.
However, in the future we plan to automatically run those using Gnu parallel.
remote_option = ''
if hosts:
# we are doing remote, so add a working dir and also return data
remote_option = (
f"--workdir {Path.home()/ '.parallel'/'tmp'} "
# f"--return {WORKING_DIR}/{PREFIX}{{0#}} "
f"--transferfile {{}} "
f"--cleanup "
)

hosts = f'-S {hosts}' if hosts else ''
"""

command_to_run = f'parallel {hosts} {remote_option} pymcnp run {{}} :::'
def __init__(
self,
runs,
prefix: str = 'pymcnp',
run_name: str | None = None,
):
self.runs = runs
self.prefix = prefix

if run_name is None:
self.run_name = _io.get_timestamp()
else:
self.run_name = run_name

self.path = Path('.') / f'{self.prefix}-{self.run_name}'

def create_files(self):
if self.path.is_dir():
print(f'[red]Error[/] {self.path} already exists. Existing.')
return

self.path.mkdir()

N = len(str(len(self.runs))) # how many digits do we need?
for i, run in enumerate(self.runs):
run.path = self.path
run.prefix = self.prefix
run.run_name = f'{i:0{N}d}'
run.create_files()

def cleanup(self):
"""Helper function that can be used in posthooks."""
if self.path:
for file_or_dir in self.path.rglob('*'):
if file_or_dir.is_file():
file_or_dir.unlink()
else:
file_or_dir.rmdir()
self.path.rmdir()

if self.dry_run:
print('[yellow]INFO[/] Dry run:')
print(f' Create main directory {directory_path}')
print(' Run parallel pre-hook')
args = []
for n in range(0, count):
subdirectory_path = Path(f'pymcnp-run-{n:0{L}d}')
sub_inp_path = subdirectory_path / f'pymcnp-inp-{timestamp}-{n:0{L}d}.inp'
print(f' Create sub-directory: {subdirectory_path.name}')
print(' Reduce nps in input file')
args.append(str(sub_inp_path))
print(f" Running: {command_to_run} {' '.join(args)}")
print(' Run parallel post-hook')
return directory_path
def prehook(self):
pass

directory_path.mkdir(parents=True, exist_ok=False)
def posthook(self):
pass

if 'nps' not in self.inp.data:
print('ERROR: only support running files with a given nps at the moment')
sys.exit(1)
def run(self, *, dry_run: bool = False):
"""
Runs MCNP INP files in parallel.
self.parallel_prehook()
Creates a directory with subdirectories for all the different input files
nps = self.inp.data['nps'].npp.value
self.inp = self.inp.set_nps(nps // count)
Returns:
Path to run directory.
"""

args = []
for n in range(count):
subdirectory_path = Path(f'pymcnp-run-{n:0{L}d}')
(directory_path / subdirectory_path).mkdir(exist_ok=False)
# some error checking
self.check_prog('parallel')

inp_path = directory_path / subdirectory_path / f'pymcnp-inp-{timestamp}-{n:0{L}d}.inp'
return

self.inp = self.inp.set_seed()
self.inp.to_mcnp_file(inp_path)
args.append(str(inp_path))

subprocess.run(command_to_run, cwd=directory_path, shell=True)
class ParallelWithMaxNPS(Parallel):
"""Replace any run with nps > max_nps with multiple runs with a given nps."""

self.parallel_posthook()
def __init__(
self,
runs,
prefix: str = 'pymcnp',
run_name: str | None = None,
max_nps: float = 1e8,
):
super().__init__(runs, prefix, run_name)
self.max_nps = int(max_nps)

return directory_path
def prehook(self):
out = []
for run in self.runs:
nps = run.inp.get_nps()
if nps > self.max_nps:
while nps > 0:
if nps > self.max_nps:
run.inp = run.inp.set_nps(self.max_nps)
else:
run.inp = run.inp.set_nps(nps)
nps -= self.max_nps
run.inp = run.inp.set_seed()
out.append(run)
self.runs = out


def main() -> None:
Expand All @@ -227,17 +246,11 @@ def main() -> None:
command = args['--command'] if args['--command'] else 'mcnp6'
path = args['--path'] if args['--path'] else input_file.parent
dry_run = args['--dry-run']
nr_runs = int(args['--parallel']) if args['--parallel'] else None
hosts = args['--hosts']

run = Run(
input_file,
path=path,
command=command,
dry_run=dry_run,
)

if args['--parallel'] is not None:
run.run_parallel(nr_runs, hosts=hosts)
else:
run.run_single()
run.run(dry_run=dry_run)

0 comments on commit 59bb3b4

Please sign in to comment.