112 lines
3.2 KiB
Python
Executable File
112 lines
3.2 KiB
Python
Executable File
import click
|
|
import logging
|
|
import tantri
|
|
import numpy
|
|
import tantri.cli.input_files.write_dipoles
|
|
import tantri.cli.file_importer
|
|
import json
|
|
import tantri.dipoles.generation
|
|
import pathlib
|
|
|
|
_logger = logging.getLogger(__name__)
|
|
|
|
LOG_PATTERN = "%(asctime)s | %(levelname)-7s | %(name)s | %(message)s"
|
|
|
|
|
|
def _set_up_logging(filename):
|
|
handlers = [logging.StreamHandler()]
|
|
if filename is not None:
|
|
handlers.append(logging.FileHandler(filename))
|
|
logging.basicConfig(
|
|
level=logging.DEBUG,
|
|
format=LOG_PATTERN,
|
|
handlers=handlers,
|
|
)
|
|
logging.getLogger("pdme").setLevel(logging.INFO)
|
|
logging.captureWarnings(True)
|
|
|
|
|
|
@click.group()
|
|
@click.option(
|
|
"--log", help="Enable logging to stream only", is_flag=True, default=False
|
|
)
|
|
@click.option("--log-file", help="A filename to use for logging (implies --log)")
|
|
@click.version_option(tantri.get_version())
|
|
def cli(log, log_file):
|
|
if log or (log_file is not None):
|
|
# log file has been provided, let's log
|
|
_set_up_logging(log_file)
|
|
|
|
|
|
@cli.command()
|
|
@click.option(
|
|
"--dipoles-file",
|
|
default="dipoles.json",
|
|
show_default=True,
|
|
type=click.Path(exists=True, path_type=pathlib.Path),
|
|
help="File with json array of dipoles",
|
|
)
|
|
@click.option(
|
|
"--dots-file",
|
|
default="dots.json",
|
|
show_default=True,
|
|
type=click.Path(exists=True, path_type=pathlib.Path),
|
|
help="File with json array of dots",
|
|
)
|
|
def read_files(dipoles_file, dots_file):
|
|
click.echo("in hello")
|
|
_logger.info("in hello")
|
|
_logger.info(
|
|
f"Received parameters [dipoles_file: {dipoles_file}] and [dots_file: {dots_file}]"
|
|
)
|
|
dipoles = tantri.cli.file_importer.read_dipoles_json_file(dipoles_file)
|
|
dots = tantri.cli.file_importer.read_dots_json_file(dots_file)
|
|
for dipole in dipoles:
|
|
_logger.info(dipole)
|
|
for dot in dots:
|
|
_logger.info(dot)
|
|
|
|
|
|
@cli.command()
|
|
@click.argument(
|
|
"generation_config",
|
|
type=click.Path(exists=True, path_type=pathlib.Path),
|
|
)
|
|
@click.argument(
|
|
"output_file",
|
|
type=click.File("w"),
|
|
)
|
|
@click.option(
|
|
"--override-rng-seed", type=int, help="Seed to override the generation config spec."
|
|
)
|
|
def generate_dipoles(generation_config, output_file, override_rng_seed):
|
|
"""Generate random dipoles as described by GENERATION_CONFIG and output to OUTPUT_FILE.
|
|
|
|
GENERATION_CONFIG should be a JSON file that matches the appropriate spec, and OUTPUT_FILE will contain JSON formatted contents.
|
|
OUTPUT_FILE will be overwritten, if it exists.
|
|
|
|
If the --override-rng-seed is set, it's better to keep logs of the generation!
|
|
"""
|
|
_logger.debug(
|
|
f"generate_dipoles was called, with config file {click.format_filename(generation_config)}"
|
|
)
|
|
_logger.debug(f"override_rng_seed: [{override_rng_seed}]")
|
|
|
|
with open(generation_config, "r") as config_file:
|
|
data = json.load(config_file)
|
|
config = tantri.dipoles.generation.DipoleGenerationConfig(**data)
|
|
|
|
override_rng = None
|
|
if override_rng_seed is not None:
|
|
_logger.info(f"Overriding the rng with a new one with seed {override_rng_seed}")
|
|
override_rng = numpy.random.default_rng(override_rng_seed)
|
|
_logger.debug(f"generating dipoles with config {config}...")
|
|
generated = tantri.dipoles.generation.make_dipoles(config, override_rng)
|
|
|
|
output_file.write(
|
|
json.dumps(
|
|
[g.as_dict() for g in generated],
|
|
cls=tantri.cli.input_files.write_dipoles.NumpyEncoder,
|
|
)
|
|
)
|