diff --git a/cubehandler/cli/commands/shrink.py b/cubehandler/cli/commands/shrink.py index 17344bb..1716567 100644 --- a/cubehandler/cli/commands/shrink.py +++ b/cubehandler/cli/commands/shrink.py @@ -1,13 +1,34 @@ +from pathlib import Path + import click from ...cube import Cube @click.command(help="Shrink a cube file.") -@click.argument("input_cube", type=click.Path(exists=True)) -@click.argument("output_cube", type=click.Path()) -def shrink(input_cube, output_cube): - cube = Cube.from_file(input_cube) - cube.reduce_data_density(points_per_angstrom=2) - cube.rescale_data() - cube.write_cube_file(output_cube, low_precision=True) +@click.argument( + "input_path", + type=click.Path(exists=True, file_okay=True, dir_okay=True), + required=True, +) +@click.argument( + "output_path", type=click.Path(file_okay=True, dir_okay=True), required=True +) +@click.option("-p", "--prefix", default="reduced_") +def shrink(input_path, output_path, prefix): + inp = Path(input_path) + out = Path(output_path) + + def run_reduction(inp, out): + cube = Cube.from_file(inp) + cube.reduce_data_density(points_per_angstrom=2) + # cube.rescale_data() # Rescaling happens below when low_precision is True + cube.write_cube_file(out, low_precision=True) + + if inp.is_file(): + run_reduction(inp, out) + elif inp.is_dir(): + out.mkdir(exist_ok=True) + for file in inp.glob("*cube"): + out_file = out / (prefix + file.name) + run_reduction(file.absolute(), out_file)