diff --git a/mccode_antlr/run/__init__.py b/mccode_antlr/run/__init__.py new file mode 100644 index 0000000..b54d650 --- /dev/null +++ b/mccode_antlr/run/__init__.py @@ -0,0 +1,6 @@ +from .runner import mccode_run_compiled, mccode_compile + +__all__ = [ + 'mccode_run_compiled', + 'mccode_compile', +] diff --git a/mccode_antlr/run/range.py b/mccode_antlr/run/range.py new file mode 100644 index 0000000..0e0869d --- /dev/null +++ b/mccode_antlr/run/range.py @@ -0,0 +1,231 @@ +from __future__ import annotations + +from typing import Union + + +class MRange: + """A range of values for a parameter in a MATLAB style. + The range is inclusive of the start and stop values, and the step is the difference between items in the range. + """ + def __init__(self, start, stop, step): + self.start = start + self.stop = stop + self.step = step + if self.start == self.stop: + raise ValueError(f'MRange start and stop values are equal: {self.start} ' + f'`list(MRange)` will be empty! Use a `Singular({self.start}, 1)` range instead.') + if self.step == 0: + raise ZeroDivisionError('MRange step cannot be zero') + + def __eq__(self, other): + return self.start == other.start and self.stop == other.stop and self.step == other.step + + @property + def min(self): + return self.start + + @property + def max(self): + return self.stop + + def __iter__(self): + def range_gen(start, stop, step): + v = start + i = 0 + while (step > 0 and v + step <= stop) or (step < 0 and v + step >= stop): + v = i * step + start + i += 1 + yield v + return range_gen(self.start, self.stop, self.step) + + def __getitem__(self, index: int): + if index < 0 or index >= len(self): + raise IndexError(f'Index {index} out of range') + return index * self.step + self.start + + def __str__(self): + return f'{self.start}:{self.step}:{self.stop}' + + def __repr__(self): + return f'MStyleRange({self})' + + def __len__(self): + return int((self.stop - self.start) / self.step) + 1 + + @classmethod + def from_str(cls, string): + """Parse a string in MATLAB style into a range. + The string should be of the form start:step:stop + """ + def float_or_int(s): + try: + return int(s) + except ValueError: + pass + return float(s) + + if string.count(':') > 2: + raise ValueError(f'Range string {string} contains more than two colons') + step = '1' + if ':' not in string: + start, stop = string, string + elif string.count(':') == 1: + start, stop = string.split(':') + else: + start, step, stop = string.split(':') + return cls(float_or_int(start), float_or_int(stop), float_or_int(step)) + + +class Singular: + """A singular range parameter for use with other range parameters in, e.g., a zip. + + Note: + The Singular range value will be repeated up to `maximum` times in an iterator. + If `maximum` is None, the Singular range will be repeated forever. + Therefore, care must be taken to ensure that the Singular range is used in a zip with a range that is + not infinite. + """ + def __init__(self, value, maximum=None): + self.value = value + self.maximum = maximum + + def __eq__(self, other): + return self.value == other.value and self.maximum == other.maximum + + def __str__(self): + return f'{self.value}(up to {self.maximum} times)' + + def __repr__(self): + return f'Singular({self.value}, {self.maximum})' + + @property + def min(self): + return self.value + + @property + def max(self): + return self.value + + def __iter__(self): + def forever(): + while True: + yield self.value + + def until(): + i = 0 + while i < self.maximum: + i += 1 + yield self.value + + return until() if self.maximum is not None else forever() + + def __len__(self): + return self.maximum + + @classmethod + def from_str(cls, string): + def float_or_int_or_str(s): + try: + return int(s) + except ValueError: + pass + try: + return float(s) + except ValueError: + return s + + if string.count(':') > 0: + raise ValueError(f'Singular string {string} contains a colon') + return cls(float_or_int_or_str(string)) + + +def parse_list(range_type, unparsed: list[str]): + ranges = {} + while len(unparsed): + if '=' in unparsed[0]: + k, v = unparsed[0].split('=', 1) + ranges[k.lower()] = range_type.from_str(v) + elif len(unparsed) > 1 and '=' not in unparsed[1]: + ranges[unparsed[0].lower()] = range_type.from_str(unparsed[1]) + del unparsed[1] + else: + raise ValueError(f'Invalid parameter: {unparsed[0]}') + del unparsed[0] + return ranges + + +def parameters_to_scan(parameters: dict[str, Union[list, MRange, Singular]], grid: bool = False): + """Convert a dictionary of ranged parameters to a list of parameter names and an iterable of parameter value tuples. + + The ranged parameters can be either MRange objects or lists of values. If a list of values is provided, it will be + iterated over directly. + + :parameter parameters: A dictionary of ranged parameters. + :parameter grid: Controls how the parameters are iterated; True implies a grid scan, False implies a linear scan. + """ + if grid: + for k, v in parameters.items(): + if isinstance(v, Singular): + parameters[k] = Singular(v.value, 1) + + names = [x.lower() for x in parameters.keys()] + values = [x if hasattr(x, '__iter__') else [x] for x in parameters.values()] + if not len(values): + return 0, names, [] + elif grid: + from itertools import product + from math import prod + # singular MRange objects *should* stop the grid along their axis: + n_pts = prod([len(v) for v in values]) + return n_pts, names, product(*values) + else: + # replace singular MRange entries with Singular iterators, to avoid stopping the zip early: + n_max = max([len(v) for v in values]) + for i, v in enumerate(values): + if len(v) > 1 and len(v) != n_max: + oth = [names[i] for i, n in enumerate(values) if len(n) == n_max] + par = 'parameters' if len(oth) > 1 else 'parameter' + have = 'have' if len(oth) > 1 else 'has' + raise ValueError(f'Parameter {names[i]} has {len(v)} values, but {par} {", ".join(oth)} {have} {n_max}') + return n_max, names, zip(*[v if len(v) > 1 else Singular(v[0] if isinstance(v, MRange) else v.value, n_max) for v in values]) + + +def _MRange_or_Singular(s: str): + if ':' in s: + return MRange.from_str(s) + return Singular.from_str(s) + + +def parse_command_line_parameters(unparsed: list[str]) -> dict[str, Union[Singular, MRange]]: + """Parse a list of input parameters into a dictionary of MRange objects. + + :parameter unparsed: A list of parameters. + """ + ranges = {} + index = 0 + while index < len(unparsed): + if '=' in unparsed[index]: + k, v = unparsed[index].split('=', 1) + ranges[k.lower()] = _MRange_or_Singular(v) + elif index + 1 < len(unparsed) and '=' not in unparsed[index + 1]: + ranges[unparsed[index].lower()] = _MRange_or_Singular(unparsed[index + 1]) + index += 1 + else: + raise ValueError(f'Invalid parameter: {unparsed[index]}') + index += 1 + return ranges + + +def parse_scan_parameters(unparsed: list[str]) -> dict[str, MRange | Singular]: + """Parse a list of input parameters into a dictionary of MRange or Singular objects. + + :parameter unparsed: A list of parameters. + :return: A dictionary of MRange or Singular objects. The Singular objects have their maximum length set to the + maximum iterations of all the ranges to avoid infinite iterations. + """ + ranges = parse_command_line_parameters(unparsed) + max_length = max(len(v) if isinstance(v, MRange) else 1 for v in ranges.values()) + for k, v in ranges.items(): + if isinstance(v, Singular) and v.maximum is None: + ranges[k] = Singular(v.value, max_length) + return ranges \ No newline at end of file diff --git a/mccode_antlr/run/runner.py b/mccode_antlr/run/runner.py new file mode 100644 index 0000000..e2b7f9d --- /dev/null +++ b/mccode_antlr/run/runner.py @@ -0,0 +1,256 @@ +from pathlib import Path +from mccode_antlr.reader import Registry + + +def regular_mccode_runtime_dict(args: dict) -> dict: + def insert_best_of(src: dict, snk: dict, names: tuple): + def get_best_of(): + for name in names: + if name in src: + return src[name] + raise RuntimeError(f"None of {names} found in {src}") + + if any(x in src for x in names): + snk[names[0]] = get_best_of() + return snk + + t = insert_best_of(args, {}, ('seed', 's')) + t = insert_best_of(args, t, ('ncount', 'n')) + t = insert_best_of(args, t, ('dir', 'out_dir', 'd')) + t = insert_best_of(args, t, ('trace', 't')) + t = insert_best_of(args, t, ('gravitation', 'g')) + t = insert_best_of(args, t, ('bufsiz',)) + t = insert_best_of(args, t, ('format',)) + return t + + +def mccode_runtime_dict_to_args_list(args: dict) -> list[str]: + """Convert a dictionary of McCode runtime arguments to a string. + + :parameter args: A dictionary of McCode runtime arguments. + :return: A list of arguments suitable for use in a command line call to a McCode compiled instrument. + """ + # convert to a standardized string: + out = [] + if 'seed' in args and args['seed'] is not None: + out.append(f'--seed={args["seed"]}') + if 'ncount' in args and args['ncount'] is not None: + out.append(f'--ncount={args["ncount"]}') + if 'dir' in args and args['dir'] is not None: + out.append(f'--dir={args["dir"]}') + if 'trace' in args and args['trace']: + out.append('--trace') + if 'gravitation' in args and args['gravitation']: + out.append('--gravitation') + if 'bufsiz' in args and args['bufsiz'] is not None: + out.append(f'--bufsiz={args["bufsiz"]}') + if 'format' in args and args['format'] is not None: + out.append(f'--format={args["format"]}') + return out + + +def mccode_runtime_parameters(args: dict, params: dict) -> str: + first = ' '.join(mccode_runtime_dict_to_args_list(args)) + second = ' '.join(f'{k}={v}' for k, v in params.items()) + return f'{first} {second}' + + +def sort_args(args: list[str]) -> list[str]: + """Take the list of arguments and sort them into the correct order for McCode run.""" + # TODO this is a bit of a hack, but it works for now + first, last = [], [] + k = 0 + while k < len(args): + if args[k].startswith('-'): + first.append(args[k]) + k += 1 + if '=' not in first[-1] and k < len(args) and not args[k].startswith('-') and '=' not in args[k]: + first.append(args[k]) + k += 1 + else: + last.append(args[k]) + k += 1 + return first + last + + +def mccode_run_script_parser(prog: str): + from argparse import ArgumentParser + from pathlib import Path + + def resolvable(name: str): + return None if name is None else Path(name).resolve() + + parser = ArgumentParser(prog=prog, description=f'Convert and run mccode_antlr-3 instr and comp files to {prog} runtime in C') + aa = parser.add_argument + + aa('filename', type=resolvable, nargs=1, help='.instr file name to be converted') + aa('parameters', nargs='*', help='Parameters to be passed to the instrument', type=str, default=None) + aa('-o', '--output-file', type=str, help='Output filename for C runtime binary', default=None) + aa('-d', '--directory', type=str, help='Output directory for C runtime artifacts') + aa('-I', '--search-dir', action='append', type=resolvable, help='Extra component search directory') + aa('-t', '--trace', action='store_true', help="Enable 'trace' mode for instrument display") + aa('-v', '--version', action='store_true', help='Print the McCode version') + aa('--source', action='store_true', help='Embed the instrument source code in the executable') + aa('--verbose', action='store_true', help='Verbose output') + + aa('-n', '--ncount', nargs=1, type=int, default=None, help='Number of neutrons to simulate') + aa('-m', '--mesh', action='store_true', default=False, help='N-dimensional mesh scan') + aa('-s', '--seed', nargs=1, type=int, default=None, help='Random number generator seed') + aa('-t', '--trace', action='store_true', default=False, help='Enable tracing') + aa('-g', '--gravitation', action='store_true', default=False, + help='Enable gravitation for all trajectories') + aa('--bufsiz', nargs=1, type=int, default=None, help='Monitor_nD list/buffer-size') + aa('--format', nargs=1, type=str, default=None, help='Output data files using FORMAT') + aa('--dryrun', action='store_true', default=False, + help='Do not run any simulations, just print the commands') + aa('--parallel', action='store_true', default=False, help='Use MPI multi-process parallelism') + aa('--gpu', action='store_true', default=False, help='Use GPU OpenACC parallelism') + aa('--process-count', nargs=1, type=int, default=0, help='MPI process count, 0 == System Default') + + return parser + + +def print_version_information(): + from mccode_antlr.version import version + print(f'mccode_antlr code generator version {version()}') + print(' Copyright (c) European Spallation Source ERIC, 2023-2024') + print('Based on McStas/McXtrace version 3') + print(' Copyright (c) DTU Physics and Risoe National Laboratory, 1997-2023') + print(' Additions (c) Institut Laue Langevin, 2003-2019') + print('All rights reserved\n\nComponents are (c) their authors, see component headers.') + + +def parse_mccode_run_script(prog: str): + import sys + from .range import parse_scan_parameters + sys.argv[1:] = sort_args(sys.argv[1:]) + args = mccode_run_script_parser(prog).parse_args() + parameters = parse_scan_parameters(args.parameters) + return args, parameters + + +def mccode_compile(instr, directory, generator, target: dict | None = None, config: dict | None = None, **kwargs): + from mccode_antlr.compiler.c import compile_instrument, CBinaryTarget + from loguru import logger + + def_target = CBinaryTarget(mpi=False, acc=False, count=1, nexus=False) + def_config = dict(default_main=True, enable_trace=False, portable=False, include_runtime=True, + embed_instrument_file=False, verbose=False) + def_config.update(config or {}) + def_target.update(target or {}) + + try: + binary = compile_instrument(instr, def_target, directory, generator=generator, config=def_config, **kwargs) + except RuntimeError as e: + logger.error(f'Failed to compile instrument: {e}') + raise e + # binary = Path(directory).joinpath(f'{instr.name}{module_config["ext"].get(str)}') + # if not binary.exists() or not binary.is_file() or not access(binary, R_OK): + # raise FileNotFoundError(f"No executable binary, {binary}, produced") + + return binary, def_target + + +def mccode_run_compiled(binary, target, directory: Path | str, parameters: str, capture: bool = True, dry_run: bool = False): + from mccode_antlr.compiler.c import run_compiled_instrument + from mccode_antlr.loader import read_mccode_dat + from pathlib import Path + + result = run_compiled_instrument(binary, target, f'--dir {directory} {parameters}', capture=capture, dry_run=dry_run) + sim_files = list(Path(directory).glob('**/*.dat')) + dats = {file.stem: read_mccode_dat(file) for file in sim_files} + return result, dats + + +def mccode_run_scan(name: str, binary, target, parameters, directory, grid: bool, capture: bool = True, dry_run: bool = False, **r_args): + from .range import parameters_to_scan + n_pts, names, scan = parameters_to_scan(parameters, grid=grid) + # n_zeros = len(str(n_pts)) + + args = regular_mccode_runtime_dict(r_args) + + if directory is None: + from datetime import datetime + directory = Path(f'{name}{datetime.now().strftime("%Y%m%d_%H%M%S")}') + + # if there is only one point, we don't need to scan + if n_pts > 1: + directory.mkdir(parents=True, exist_ok=True) + results = [] + for number, values in enumerate(scan): + # TODO Use the following line instead of the one after it when McCode is fixed to use zero-padded folder names + # # runtime_arguments['dir'] = args["dir"].joinpath(str(number).zfill(n_zeros)) + this_directory = directory.joinpath(str(number)) + pars = mccode_runtime_parameters(r_args, parameters) + r, d = mccode_run_compiled(binary, target, this_directory, pars, capture=capture, dry_run=dry_run) + results.append((r, d)) + return results + else: + directory.parent.mkdir(parents=True, exist_ok=True) + pars = mccode_runtime_parameters(args, parameters) + return mccode_run_compiled(binary, target, directory, pars, capture=capture, dry_run=dry_run) + + +def mccode_run(flavor: str, registry: Registry, generator: dict): + from pathlib import Path + from mccode_antlr.reader import Reader + from mccode_antlr.reader import LocalRegistry + from os import R_OK, access + + args, parameters = parse_mccode_run_script(flavor) + config = dict( + enable_trace=args.trace if args.trace is not None else False, + embed_instrument_file=args.source if args.source is not None else False, + verbose=args.verbose if args.verbose is not None else False, + output=args.output_file if args.output_file is not None else args.filename.with_suffix('.c') + ) + target = dict( + mpi=args.parallel, + acc=args.gpu, + count=args.process_count, + nexus=False + ) + runtime = dict( + seed=args.seed[0] if args.seed is not None else None, + ncount=args.ncount[0] if args.ncount is not None else None, + trace=args.trace, + gravitation=args.gravitation, + bufsiz=args.bufsiz[0] if args.bufsiz is not None else None, + format=args.format[0] if args.format is not None else None, + dry_run=args.dryrun, + capture=(not args.verbose) if args.verbose is not None else False, + ) + # check if the filename is actually a compiled instrument already: + if args.output_file is None and args.filename.exists() and access(args.filename, R_OK): + binary = args.filename + name = args.filename.stem + else: + # McCode always requires access to a remote Pooch repository: + registries = [registry] + # A user can specify extra (local) directories to search for included files using -I or --search-dir + if args.search_dir is not None and len(args.search_dir): + registries.extend([LocalRegistry(d.stem, d) for d in args.search_dir]) + # And McCode-3 users expect to always have access to files in the current working directory + registries.append(LocalRegistry('working_directory', f'{Path().resolve()}')) + # Construct the object which will read the instrument and component files, producing Python objects + reader = Reader(registries=registries) + # Read the provided .instr file, including all specified .instr and .comp files along the way + instrument = reader.get_instrument(args.filename) + name = instrument.name + # Generate the C binary for the instrument -- will output to, e.g., {instrument.name}.out, in the current directory + # unless if output_file was specified + binary, target = mccode_compile(instrument, args.output_file, generator, target=target, config=config) + + mccode_run_scan(name, binary, target, parameters, args.directory, args.mesh, **runtime) + + +def mcstas(): + from mccode_antlr.reader import MCSTAS_REGISTRY + from mccode_antlr.translators.target import MCSTAS_GENERATOR + mccode_run('mcstas', MCSTAS_REGISTRY, MCSTAS_GENERATOR) + + +def mcxtrace(): + from mccode_antlr.reader import MCXTRACE_REGISTRY + from mccode_antlr.translators.target import MCXTRACE_GENERATOR + mccode_run('mcxtrace', MCXTRACE_REGISTRY, MCXTRACE_GENERATOR) diff --git a/pyproject.toml b/pyproject.toml index d2556a7..6105ca9 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -33,6 +33,8 @@ hdf5 = ["h5py==3.10.0"] [project.scripts] mcstas-antlr = "mccode_antlr.commands:mcstas" mcxtrace-antlr = "mccode_antlr.commands:mcxtrace" +mcrun-antlr = "mccode_antlr.run.runner:mcstas" +mxrun-antlr = "mccode_antlr.run.runner:mcxtrace" [tool.setuptools_scm] diff --git a/tests/runtime/compiled.py b/tests/runtime/compiled.py index 6114d80..7ad90ce 100644 --- a/tests/runtime/compiled.py +++ b/tests/runtime/compiled.py @@ -82,39 +82,16 @@ def wrapper(*args, **kwargs): return wrapper -def compile_and_run(instr, parameters, run=True, dump_source=True, - target: dict | None = None, config: dict | None = None): - from mccode_antlr.compiler.c import compile_instrument, CBinaryTarget, run_compiled_instrument - from mccode_antlr.translators.target import MCSTAS_GENERATOR - from mccode_antlr.loader import read_mccode_dat - from mccode_antlr.config import config as module_config - from tempfile import TemporaryDirectory - from os import R_OK, access +def compile_and_run(instr, parameters, run=True, dump_source=True, target: dict | None = None, config: dict | None = None): from pathlib import Path - from loguru import logger + from tempfile import TemporaryDirectory + from mccode_antlr.translators.target import MCSTAS_GENERATOR + from mccode_antlr.run import mccode_compile, mccode_run_compiled - def_target = CBinaryTarget(mpi=False, acc=False, count=1, nexus=False) - def_config = dict(default_main=True, enable_trace=False, portable=False, include_runtime=True, - embed_instrument_file=False, verbose=False) - def_config.update(config or {}) - def_target.update(target or {}) + kwargs = dict(generator=MCSTAS_GENERATOR, target=target, config=config, dump_source=dump_source) with TemporaryDirectory() as directory: - try: - compile_instrument(instr, def_target, directory, - generator=MCSTAS_GENERATOR, config=def_config, dump_source=dump_source) - except RuntimeError as e: - logger.error(f'Failed to compile instrument: {e}') - raise e - binary = Path(directory).joinpath(f'{instr.name}{module_config["ext"].get(str)}') - - if not binary.exists() or not binary.is_file() or not access(binary, R_OK): - raise FileNotFoundError(f"No executable binary, {binary}, produced") - - if run: - result = run_compiled_instrument(binary, def_target, f"--dir {directory}/instr {parameters}", - capture=True) - sim_files = list(Path(directory).glob('**/*.dat')) - dats = {file.stem: read_mccode_dat(file) for file in sim_files} - return result, dats - return None, None + binary, target = mccode_compile(instr, directory, **kwargs) + # The runtime output directory used *can not* exist for McStas/McXtrace to work properly. + # So find a name inside this directory that doesn't exist (any name should work) + return mccode_run_compiled(binary, target, Path(directory).joinpath('t'), parameters) if run else (None, None)