__main__.py 19.7 KB
Newer Older
Sean Leavey's avatar
Sean Leavey committed
1 2
"""Finesse command line interface"""

Sean Leavey's avatar
Sean Leavey committed
3
import sys
Sean Leavey's avatar
Sean Leavey committed
4 5
import logging
import datetime
Sean Leavey's avatar
Sean Leavey committed
6
import textwrap
7

Sean Leavey's avatar
Sean Leavey committed
8
import click
Sean Leavey's avatar
Sean Leavey committed
9 10 11
from click_default_group import (
    DefaultGroup,
)  # https://github.com/pallets/click/issues/430
12
import tabulate
Sean Leavey's avatar
Sean Leavey committed
13

Sean Leavey's avatar
Sean Leavey committed
14 15 16 17 18
from . import (
    __version__,
    PROGRAM,
    DESCRIPTION,
)
19 20
from .config import FinesseConfig
from .environment import show_tracebacks
Sean Leavey's avatar
Sean Leavey committed
21
from .utilities import graph_layouts
Sean Leavey's avatar
Sean Leavey committed
22 23


24 25 26
class ClickLogColorFormatter(logging.Formatter):
    """Stdout log formatter with colors."""

Sean Leavey's avatar
Sean Leavey committed
27 28 29
    # Log prefix format.
    # NOTE: this is NOT a :py:class:`logging.Formatter` compatible format string, but rather a
    # format specific to this class.
30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76
    PREFIX_FORMAT = "{short_name:>25s} [{record.levelname:>8s}]:"

    STYLES = {
        "critical": {"fg": "red",},
        "error": {"fg": "red",},
        "warning": {"fg": "yellow",},
        "info": {"fg": "green",},
        "debug": {"fg": "blue",},
    }

    def format(self, record):
        if record.exc_info:
            # Skip formatting of exception info.
            return super().format(record)

        level = record.levelname.casefold()
        msg = record.getMessage()

        if level in self.STYLES:
            # Remove "finesse." prefix.
            short_name = (
                record.name[8:] if record.name.startswith("finesse.") else record.name
            )

            prefix = click.style(
                self.PREFIX_FORMAT.format(short_name=short_name, record=record),
                **self.STYLES[level],
            )
            msg = "\n".join(f"{prefix} {line}" for line in msg.splitlines())

        return msg


class ClickLogHandler(logging.Handler):
    """Log handler that prints logs in a color-supporting platform-independent way."""

    # Use stderr instead of stdout.
    _USE_STDERR = True

    def emit(self, record):
        try:
            msg = self.format(record)
            click.echo(msg, err=self._USE_STDERR)
        except Exception:
            self.handleError(record)


Sean Leavey's avatar
Sean Leavey committed
77
class KatState:
Sean Leavey's avatar
Sean Leavey committed
78
    """Shared state for all CLI subcommands.
Samuel Rowlinson's avatar
Samuel Rowlinson committed
79

Sean Leavey's avatar
Sean Leavey committed
80 81 82
    This object gets built by Click when the CLI is called and encapsulates global settings for
    use by individual commands/groups.
    """
Sean Leavey's avatar
Sean Leavey committed
83

Sean Leavey's avatar
Sean Leavey committed
84
    VERBOSITY_DEFAULT = logging.WARNING
Sean Leavey's avatar
Sean Leavey committed
85 86

    def __init__(self):
Sean Leavey's avatar
Sean Leavey committed
87 88
        # Set up the Finesse logger.
        self.logger = logging.getLogger(PROGRAM)
89 90
        handler = ClickLogHandler()
        handler.setFormatter(ClickLogColorFormatter())
Sean Leavey's avatar
Sean Leavey committed
91
        self.logger.addHandler(handler)
Sean Leavey's avatar
Sean Leavey committed
92 93 94

        self.verbosity = self.VERBOSITY_DEFAULT

Sean Leavey's avatar
Sean Leavey committed
95 96 97
        # Log the version now that we've set everything up.
        self.logger.debug(f"{PROGRAM} {__version__}")

Sean Leavey's avatar
Sean Leavey committed
98 99 100
    @property
    def verbosity(self):
        """Log verbosity on stdout"""
Sean Leavey's avatar
Sean Leavey committed
101
        return self.logger.getEffectiveLevel()
Sean Leavey's avatar
Sean Leavey committed
102 103 104

    @verbosity.setter
    def verbosity(self, verbosity):
Sean Leavey's avatar
Sean Leavey committed
105 106
        verbosity = min(logging.CRITICAL, max(logging.DEBUG, verbosity))
        self.logger.setLevel(verbosity)
Sean Leavey's avatar
Sean Leavey committed
107 108 109 110 111 112 113

    @property
    def verbose(self):
        """Verbose output enabled

        Returns True if the verbosity is enough for INFO or DEBUG messages to be displayed.
        """
Sean Leavey's avatar
Sean Leavey committed
114 115
        return self.verbosity <= self.VERBOSITY_DEFAULT

116 117 118 119 120 121 122 123
    @property
    def debug(self):
        return show_tracebacks()

    @debug.setter
    def debug(self, value):
        show_tracebacks(value)

Sean Leavey's avatar
Sean Leavey committed
124 125 126 127 128 129 130 131 132 133 134 135 136
    def print(
        self, text, indent=0, error=False, exit_=False, exit_code=False, **kwargs
    ):
        text = str(text)

        if indent > 0:
            text = textwrap.indent(text, " " * 4)

        click.secho(text, err=error, **kwargs)

        if exit_:
            sys.exit(exit_code)

137 138 139 140
    def print_error(
        self, title, exception_or_msg=None, exit_=True, exit_code=1, **kwargs
    ):
        if exception_or_msg is not None:
Sean Leavey's avatar
Sean Leavey committed
141
            self.print(title, error=True, fg="red", bold=True)
142 143 144 145 146 147 148 149

            if self.debug and isinstance(exception_or_msg, BaseException):
                # Print the full traceback.
                import traceback

                # FIXME: use click stream.
                exception_or_msg = traceback.format_exc()

Sean Leavey's avatar
Sean Leavey committed
150
            self.print(
151
                exception_or_msg,
Sean Leavey's avatar
Sean Leavey committed
152 153 154 155 156 157 158 159 160
                indent=1,
                error=True,
                fg="red",
                exit_=True,
                exit_code=1,
                **kwargs,
            )
        else:
            self.print(title, error=True, fg="red", exit_=True, exit_code=1, **kwargs)
Sean Leavey's avatar
Sean Leavey committed
161

Sean Leavey's avatar
Sean Leavey committed
162
    def print_banner(self, kat_file=None, **kwargs):
Sean Leavey's avatar
Sean Leavey committed
163 164 165 166
        """Print Finesse banner."""
        if not self.verbose:
            return

Sean Leavey's avatar
Sean Leavey committed
167
        input_file = f"Input file: {kat_file.name}" if kat_file else ""
Sean Leavey's avatar
Sean Leavey committed
168
        timenow = datetime.datetime.now().strftime("%d %B %Y, %H:%M:%S")
Sean Leavey's avatar
Sean Leavey committed
169 170 171 172

        self.print(
            f"""------------------------------------------------------------------------
                     FINESSE {__version__}
Sean Leavey's avatar
Sean Leavey committed
173 174 175
        o_.-=.       Frequency domain INterferomEter Simulation SoftwarE
        (\\'\".\\|                         http://www.gwoptics.org/finesse/
        .>' (_--.
Sean Leavey's avatar
Sean Leavey committed
176
    _=/d   ,^\\       {input_file}
Sean Leavey's avatar
Sean Leavey committed
177 178
    ~~ \\)-'   '
    / |
Sean Leavey's avatar
Sean Leavey committed
179 180 181
    '  '                                       {timenow}
------------------------------------------------------------------------""",
            **kwargs,
Samuel Rowlinson's avatar
Samuel Rowlinson committed
182 183
        )

Sean Leavey's avatar
Sean Leavey committed
184 185 186

def set_verbosity(ctx, option, value):
    """Callback to set stdout verbosity for root CLI command."""
Sean Leavey's avatar
Sean Leavey committed
187 188 189
    if not value:
        return

Sean Leavey's avatar
Sean Leavey committed
190 191
    state = ctx.ensure_object(KatState)

Sean Leavey's avatar
Sean Leavey committed
192 193 194 195 196 197
    if option.name == "verbose":
        offset = -10 * value
    elif option.name == "quiet":
        offset = 10 * value
    else:
        raise ValueError("Unsupported option.")
Sean Leavey's avatar
Sean Leavey committed
198

Sean Leavey's avatar
Sean Leavey committed
199
    state.verbosity = state.VERBOSITY_DEFAULT + offset
Samuel Rowlinson's avatar
Samuel Rowlinson committed
200

Sean Leavey's avatar
Sean Leavey committed
201

202 203 204 205 206 207 208 209
def set_debug(ctx, option, value):
    if not value:
        return

    state = ctx.ensure_object(KatState)
    state.debug = value


Sean Leavey's avatar
Sean Leavey committed
210
def print_banner(ctx, option, value):
Sean Leavey's avatar
Sean Leavey committed
211
    """Show the Finesse banner and exit."""
Sean Leavey's avatar
Sean Leavey committed
212 213 214 215 216 217 218
    if not value:
        return

    state = ctx.ensure_object(KatState)
    state.print_banner(exit_=True)


Sean Leavey's avatar
Sean Leavey committed
219 220 221 222 223 224 225 226 227 228 229 230
def parse_path(state, path, legacy=False, **error_kwargs):
    from finesse.script import parse_file, parse_legacy_file

    try:
        if legacy:
            return parse_legacy_file(path)
        else:
            return parse_file(path)
    except Exception as error:
        state.print_error(f"PARSING ERROR:", error, **error_kwargs)


231 232 233 234 235 236 237 238 239 240 241 242
def plot_graph(state, network, layout, graphviz=False):
    from .plotting.graph import plot_nx_graph, plot_graphviz

    if graphviz:
        try:
            plot_graphviz(network, layout)
        except ImportError as e:
            state.print_error("Dependency not installed", e)
    else:
        plot_nx_graph(network, layout)


Sean Leavey's avatar
Sean Leavey committed
243 244
input_file_argument = click.argument("input_file", type=click.File("r"))
output_file_argument = click.argument("output_file", type=click.File("w"), default="-")
Sean Leavey's avatar
Sean Leavey committed
245 246 247 248 249 250 251 252 253 254 255 256
plot_option = click.option(
    "--plot/--no-plot",
    default=True,
    show_default=True,
    help="Display results as figure.",
)
trace_option = click.option(
    "--trace/--no-trace",
    default=False,
    show_default=True,
    help="Displays the results of a beam trace of the model.",
)
257 258 259 260 261 262 263
graphviz_option = click.option(
    "--graphviz",
    is_flag=True,
    default=False,
    show_default=True,
    help="Generate layout and display using Graphviz.",
)
264 265 266
graph_layout_argument = click.option(
    "--layout",
    type=click.Choice(graph_layouts().keys()),
267
    default="spring",
268 269 270 271 272 273 274 275 276 277
    help="Graph layout algorithm to use.",
)
network_type_argument = click.option(
    "-t",
    "--type",
    "network_type",
    type=click.Choice(("full", "components", "optical")),
    default="full",
    help="Network to plot.",
)
278 279 280 281 282 283
table_format_argument = click.option(
    "--table-format",
    type=click.Choice(tabulate.tabulate_formats),
    default="fancy_grid",
    help="Table format to use.",
)
Sean Leavey's avatar
Sean Leavey committed
284 285 286 287
legacy_option = click.option(
    "--legacy",
    is_flag=True,
    default=False,
Sean Leavey's avatar
Sean Leavey committed
288
    help="Specify that the input file uses Finesse 2 syntax.",
Sean Leavey's avatar
Sean Leavey committed
289
)
290
verbose_option = click.option(
Sean Leavey's avatar
Sean Leavey committed
291 292 293 294 295 296 297
    "-v",
    "--verbose",
    count=True,
    callback=set_verbosity,
    expose_value=False,
    is_eager=True,
    help="Increase verbosity of log output (can be specified multiple times).",
Samuel Rowlinson's avatar
Samuel Rowlinson committed
298
)
299
quiet_option = click.option(
Sean Leavey's avatar
Sean Leavey committed
300 301 302 303 304 305 306
    "-q",
    "--quiet",
    count=True,
    callback=set_verbosity,
    expose_value=False,
    is_eager=True,
    help="Decrease verbosity of log output (can be specified multiple times).",
Samuel Rowlinson's avatar
Samuel Rowlinson committed
307
)
308 309 310 311 312 313 314 315 316 317
debug_option = click.option(
    "--debug",
    is_flag=True,
    default=False,
    callback=set_debug,
    expose_value=False,
    is_eager=True,
    help="Enable debug mode (print tracebacks).",
)

318

319
@click.group(cls=DefaultGroup, default="run", help=DESCRIPTION)
Samuel Rowlinson's avatar
Samuel Rowlinson committed
320
@click.option(
Sean Leavey's avatar
Sean Leavey committed
321
    "--banner",
Samuel Rowlinson's avatar
Samuel Rowlinson committed
322 323
    is_flag=True,
    default=False,
Sean Leavey's avatar
Sean Leavey committed
324 325 326
    callback=print_banner,
    expose_value=False,
    is_eager=True,
Sean Leavey's avatar
Sean Leavey committed
327
    help=print_banner.__doc__,
Samuel Rowlinson's avatar
Samuel Rowlinson committed
328
)
Sean Leavey's avatar
Sean Leavey committed
329 330 331 332 333 334 335
@click.version_option(version=__version__, prog_name=PROGRAM)
@click.pass_context
def cli(ctx):
    pass


@cli.command()
Sean Leavey's avatar
Sean Leavey committed
336
@input_file_argument
Sean Leavey's avatar
Sean Leavey committed
337 338
@plot_option
@trace_option
Sean Leavey's avatar
Sean Leavey committed
339
@legacy_option
340 341
@verbose_option
@quiet_option
342
@debug_option
Sean Leavey's avatar
Sean Leavey committed
343 344
@click.pass_context
def run(
Sean Leavey's avatar
Sean Leavey committed
345
    ctx, input_file, plot, trace, legacy,
Sean Leavey's avatar
Sean Leavey committed
346 347 348
):
    """Run a Finesse script."""
    state = ctx.ensure_object(KatState)
Sean Leavey's avatar
Sean Leavey committed
349
    state.print_banner(input_file)
Sean Leavey's avatar
Sean Leavey committed
350

Sean Leavey's avatar
Sean Leavey committed
351
    model = parse_path(state, input_file, legacy=legacy)
Sean Leavey's avatar
Sean Leavey committed
352 353 354 355 356 357 358 359 360 361 362 363 364 365 366 367 368 369 370

    try:
        solution = model.run()
    except Exception as error:
        state.print_error("ANALYSIS ERROR:", error)

    if trace:
        try:
            trace_solution = model.beam_trace()
        except Exception as error:
            state.print_error("TRACE ERROR:", error)

        state.print("Trace:", fg="green", bold=True)
        state.print(trace_solution)

    if solution is not None:
        state.print("Solution:", fg="green", bold=True)
        state.print(solution, indent=1)

371
        if plot:
372
            from finesse import configure
Sean Leavey's avatar
Sean Leavey committed
373

374
            configure(plotting=True)
375
            solution.plot(show=plot)
Sean Leavey's avatar
Sean Leavey committed
376 377


378 379
@cli.command()
@click.option(
380
    "--paths",
381
    is_flag=True,
382
    default=True,
383
    show_default=True,
384
    help="Show the paths used to build the configuration.",
385 386
)
@click.option(
387
    "--dump/--no-dump",
388 389 390
    is_flag=True,
    default=False,
    show_default=True,
391
    help="Show the current configuration in a table.",
392
)
393
@table_format_argument
394 395 396 397
@verbose_option
@quiet_option
@debug_option
@click.pass_context
398
def config(ctx, paths, dump, table_format):
399 400
    """Configuration information."""
    state = ctx.ensure_object(KatState)
401
    config = FinesseConfig()
402 403 404

    if paths:
        state.print(
Sean Leavey's avatar
Sean Leavey committed
405 406 407
            "Configuration paths (1 = highest priority; green = found, red = not found):",
            fg="green",
            bold=True,
408 409
        )

410 411 412
        configs = FinesseConfig.user_config_paths().items()

        for priority, (_, path) in enumerate(reversed(configs), start=1):
413 414 415 416
            fg = "green" if path.exists() else "red"
            pathstr = click.format_filename(str(path))
            state.print(f"{priority}: {pathstr}", indent=1, fg=fg)

417 418 419 420 421 422 423 424 425 426
    if dump:
        rows = []
        for section in config.sections():
            _section = click.style(section, fg="green", bold=True)
            for option, value in config.items(section):
                option = click.style(option, fg="blue", bold=True)
                rows.append((_section, option, value))

        state.print(tabulate.tabulate(rows, tablefmt=table_format))

427 428 429 430 431 432 433 434 435 436 437

@cli.command()
@verbose_option
@quiet_option
@debug_option
@click.pass_context
def config_edit(ctx):
    """Open the user configuration in the system default editor."""
    state = ctx.ensure_object(KatState)

    state.print("Opening user configuration in system editor...", fg="green", bold=True)
438
    click.edit(filename=FinesseConfig.user_config_path(), extension=".ini")
439 440 441 442 443


@cli.command()
@click.confirmation_option(
    prompt=(
444 445
        f"Are you sure you wish to reset the user configuration at "
        f"{FinesseConfig.user_config_path()}?"
446 447 448 449 450 451 452 453 454
    )
)
@verbose_option
@quiet_option
@debug_option
@click.pass_context
def config_reset(ctx):
    """Reset the user configuration."""
    state = ctx.ensure_object(KatState)
Sean Leavey's avatar
Sean Leavey committed
455
    FinesseConfig.write_user_config(force=True)
456 457 458 459

    state.print(f"Successfully reset user configuration.", fg="green", bold=True)


Sean Leavey's avatar
Sean Leavey committed
460
@cli.command()
Sean Leavey's avatar
Sean Leavey committed
461
@input_file_argument
Sean Leavey's avatar
Sean Leavey committed
462 463 464
@click.option(
    "--summary/--no-summary", is_flag=True, default=True, help="Show summary."
)
Sean Leavey's avatar
Sean Leavey committed
465 466 467
@legacy_option
@verbose_option
@quiet_option
468
@debug_option
Sean Leavey's avatar
Sean Leavey committed
469
@click.pass_context
Sean Leavey's avatar
Sean Leavey committed
470
def info(ctx, input_file, summary, legacy):
Sean Leavey's avatar
Sean Leavey committed
471 472
    """Print information about a model parsed from a script."""
    state = ctx.ensure_object(KatState)
Sean Leavey's avatar
Sean Leavey committed
473
    state.print_banner(input_file)
Sean Leavey's avatar
Sean Leavey committed
474

Sean Leavey's avatar
Sean Leavey committed
475
    model = parse_path(state, input_file, legacy=legacy)
Sean Leavey's avatar
Sean Leavey committed
476

Sean Leavey's avatar
Sean Leavey committed
477 478 479 480
    if summary:
        state.print("Summary:", fg="green", bold=True)
        state.print(model.info())

Sean Leavey's avatar
Sean Leavey committed
481 482 483 484 485 486 487 488 489 490 491 492 493 494 495 496 497 498 499 500 501 502

@cli.command()
@input_file_argument
@click.option(
    "--from", "from_", help="Start node.",
)
@click.option(
    "--to", help="Stop node.",
)
@legacy_option
@verbose_option
@quiet_option
@debug_option
@click.pass_context
def path(ctx, input_file, from_, to, legacy):
    """Print information about a path through the parsed model."""
    state = ctx.ensure_object(KatState)
    state.print_banner(input_file)

    model = parse_path(state, input_file, legacy=legacy)

    state.print(model.path(from_, to))
Sean Leavey's avatar
Sean Leavey committed
503 504 505


@cli.command()
Sean Leavey's avatar
Sean Leavey committed
506
@input_file_argument
507 508
@network_type_argument
@graph_layout_argument
509
@graphviz_option
Sean Leavey's avatar
Sean Leavey committed
510 511 512 513 514
@legacy_option
@verbose_option
@quiet_option
@debug_option
@click.pass_context
515
def graph(ctx, input_file, network_type, layout, graphviz, legacy):
Sean Leavey's avatar
Sean Leavey committed
516 517
    """Show the parsed model graph."""
    state = ctx.ensure_object(KatState)
Sean Leavey's avatar
Sean Leavey committed
518
    state.print_banner(input_file)
Sean Leavey's avatar
Sean Leavey committed
519

Sean Leavey's avatar
Sean Leavey committed
520
    model = parse_path(state, input_file, legacy=legacy)
521

522 523 524 525
    try:
        plot_graph(state, model.get_network(network_type), layout, graphviz)
    except Exception as e:
        state.print_error(e)
Sean Leavey's avatar
Sean Leavey committed
526 527


528
@cli.command()
Sean Leavey's avatar
Sean Leavey committed
529
@input_file_argument
530 531 532 533 534
@click.option(
    "--direction",
    type=click.Choice(("x", "y")),
    help="Direction to compute mismatches in. If not specified, both 'x' and 'y' are computed.",
)
Sean Leavey's avatar
Sean Leavey committed
535
@legacy_option
536 537
@verbose_option
@quiet_option
538
@debug_option
539
@click.pass_context
Sean Leavey's avatar
Sean Leavey committed
540
def mismatches(ctx, input_file, direction, legacy):
541 542
    """Mismatch tools."""
    state = ctx.ensure_object(KatState)
Sean Leavey's avatar
Sean Leavey committed
543
    state.print_banner(input_file)
544

Sean Leavey's avatar
Sean Leavey committed
545
    model = parse_path(state, input_file, legacy=legacy)
546 547 548 549 550 551 552

    try:
        model.print_cavity_mismatches(direction=direction)
    except Exception as error:
        state.print_error("TRACE ERROR:", error)


Sean Leavey's avatar
Sean Leavey committed
553
@cli.command()
Sean Leavey's avatar
Sean Leavey committed
554
@input_file_argument
Samuel Rowlinson's avatar
Samuel Rowlinson committed
555
@click.option(
Sean Leavey's avatar
Sean Leavey committed
556
    "--from", "from_", help="Start node.",
Samuel Rowlinson's avatar
Samuel Rowlinson committed
557 558
)
@click.option(
Sean Leavey's avatar
Sean Leavey committed
559 560
    "--to", help="Stop node.",
)
561 562 563 564 565 566 567 568 569 570 571 572 573 574
@click.option(
    "--symbolic",
    is_flag=True,
    default=False,
    show_default=True,
    help="Perform symbolic trace computation.",
)
@click.option(
    "--print-distances",
    is_flag=True,
    default=False,
    show_default=True,
    help="Print component distances matrix.",
)
Sean Leavey's avatar
Sean Leavey committed
575 576
@click.option(
    "--plot/--no-plot",
Samuel Rowlinson's avatar
Samuel Rowlinson committed
577 578
    default=True,
    show_default=True,
Sean Leavey's avatar
Sean Leavey committed
579
    help="Display results as figure.",
Samuel Rowlinson's avatar
Samuel Rowlinson committed
580 581
)
@click.option(
Sean Leavey's avatar
Sean Leavey committed
582 583 584
    "--save-figure",
    type=click.File("wb", lazy=False),
    help="Save image of figure to file.",
Samuel Rowlinson's avatar
Samuel Rowlinson committed
585
)
Sean Leavey's avatar
Sean Leavey committed
586
@legacy_option
587 588
@verbose_option
@quiet_option
589
@debug_option
Sean Leavey's avatar
Sean Leavey committed
590
@click.pass_context
Sean Leavey's avatar
Sean Leavey committed
591 592 593
def trace(
    ctx, input_file, from_, to, symbolic, print_distances, plot, save_figure, legacy
):
Sean Leavey's avatar
Sean Leavey committed
594
    """Tracing tools."""
Sean Leavey's avatar
Sean Leavey committed
595
    state = ctx.ensure_object(KatState)
Sean Leavey's avatar
Sean Leavey committed
596
    state.print_banner(input_file)
Sean Leavey's avatar
Sean Leavey committed
597

Sean Leavey's avatar
Sean Leavey committed
598
    model = parse_path(state, input_file, legacy=legacy)
599

600
    try:
601
        trace_solution = model.propagate_beam(from_, to, symbolic=symbolic)
602
    except Exception as error:
Sean Leavey's avatar
Sean Leavey committed
603
        state.print_error("TRACE ERROR:", error)
604

Sean Leavey's avatar
Sean Leavey committed
605 606
    state.print("Trace:", fg="green", bold=True)
    state.print(trace_solution)
607

608 609 610 611
    if print_distances:
        # FIXME: print using state.
        trace_solution.print_distances_matrix()

Sean Leavey's avatar
Sean Leavey committed
612
    if plot or save_figure:
613
        from finesse import configure
614

615
        configure(plotting=True)
616
        trace_solution.plot(show=plot, filename=save_figure)
617

618

Sean Leavey's avatar
Sean Leavey committed
619
@cli.command()
Sean Leavey's avatar
Sean Leavey committed
620
@input_file_argument
Sean Leavey's avatar
Sean Leavey committed
621 622 623 624 625 626 627 628 629 630 631 632 633 634 635 636 637 638 639 640 641 642 643 644 645 646 647 648 649 650 651
@click.option(
    "--from", "from_", help="Start node.",
)
@click.option(
    "--to", help="Stop node.",
)
@click.option(
    "-p",
    "--parameter",
    "params",
    type=(str, click.Choice(("lin", "log")), float, float, int),
    multiple=True,
    help=(
        "Model parameter to animate in the form '<param> lin|log <start> <stop> <step>' (can be "
        "specified multiple times)."
    ),
)
@click.option(
    "--property",
    "props",
    type=click.Choice(("beamsize", "gouy", "curvature", "all")),
    multiple=True,
    help="Property to animate (can be specified multiple times).",
)
@click.option(
    "-i",
    "--interval",
    type=float,
    default=200,
    help="Delay between frames in milliseconds.",
)
Sean Leavey's avatar
Sean Leavey committed
652 653
@click.option(
    "--save-figure",
654
    type=click.Path(dir_okay=False, writable=True),
Sean Leavey's avatar
Sean Leavey committed
655 656
    help="Save image of figure to file.",
)
Sean Leavey's avatar
Sean Leavey committed
657 658 659 660 661
@legacy_option
@verbose_option
@quiet_option
@debug_option
@click.pass_context
Sean Leavey's avatar
Sean Leavey committed
662 663 664
def trace_animation(
    ctx, input_file, from_, to, params, props, interval, save_figure, legacy
):
Sean Leavey's avatar
Sean Leavey committed
665 666
    """Animate a trace."""
    state = ctx.ensure_object(KatState)
Sean Leavey's avatar
Sean Leavey committed
667
    state.print_banner(input_file)
Sean Leavey's avatar
Sean Leavey committed
668

Sean Leavey's avatar
Sean Leavey committed
669
    model = parse_path(state, input_file, legacy=legacy)
Sean Leavey's avatar
Sean Leavey committed
670 671 672 673 674 675 676 677 678 679 680 681 682 683

    try:
        trace_solution = model.propagate_beam(from_, to, symbolic=True)
    except Exception as error:
        state.print_error("TRACE ERROR:", error)

    def parse_param_sweep(param, scale, start, stop, step):
        import numpy as np

        param = model.reduce_get_attr(param)
        func = np.linspace if scale == "lin" else np.logspace
        return param, func(start, stop, step)

    parameters = dict(parse_param_sweep(*param) for param in params)
684
    trace_solution.animate(parameters, *props, interval=interval, filename=save_figure)
Sean Leavey's avatar
Sean Leavey committed
685 686


Sean Leavey's avatar
Sean Leavey committed
687
@cli.command()
Sean Leavey's avatar
Sean Leavey committed
688 689
@input_file_argument
@output_file_argument
Sean Leavey's avatar
Sean Leavey committed
690
@legacy_option
691 692
@verbose_option
@quiet_option
693
@debug_option
Sean Leavey's avatar
Sean Leavey committed
694
@click.pass_context
695 696 697 698 699 700 701 702 703
def convert(ctx, input_file, output_file, legacy):
    """Convert and normalize kat script to canonical form.

    This can be used to convert Finesse 2 or Finesse 3 scripts to canonical Finesse 3 form. Note
    that some Finesse 2 instructions are not supported. The resulting script is arranged in standard
    order and with standard spacing.

    If OUTPUT_FILE is not specified, it defaults to stdout.
    """
Sean Leavey's avatar
Sean Leavey committed
704 705 706
    # Get state object to set verbosity, etc.
    state = ctx.ensure_object(KatState)

707 708
    model = parse_path(state, input_file, legacy=legacy)
    model.unparse_file(output_file)
709 710 711


@cli.group()
712 713
@click.pass_context
def debug(ctx):
714
    """Debug tools."""
715 716
    # The debug flag is implied.
    set_debug(ctx, None, True)
717 718


Sean Leavey's avatar
Sean Leavey committed
719 720 721 722 723 724 725 726 727 728 729 730
@debug.command(name="run")
@input_file_argument
@plot_option
@trace_option
@legacy_option
@verbose_option
@quiet_option
@debug_option
@click.pass_context
def run_debug(
    ctx, input_file, plot, trace, legacy,
):
731
    """Run a Finesse script with Python fault handler enabled.
Sean Leavey's avatar
Sean Leavey committed
732

733 734 735 736
    If the file extension is 'py', it is interpreted as Python code; otherwise it is parsed assuming
    it to be kat script.
    """
    import os
Sean Leavey's avatar
Sean Leavey committed
737 738
    import faulthandler

739
    click.secho("Enabling Python fault handler", fg="yellow")
Sean Leavey's avatar
Sean Leavey committed
740
    faulthandler.enable()
741 742 743 744 745 746 747

    root, ext = os.path.splitext(input_file.name)
    if ext.casefold() == ".py":
        exec(input_file.read())
    else:
        # Forward arguments to normal "run" command.
        ctx.forward(run)
Sean Leavey's avatar
Sean Leavey committed
748 749


750 751 752 753 754 755 756 757 758
@debug.command()
@input_file_argument
@graph_layout_argument
@graphviz_option
@verbose_option
@quiet_option
@debug_option
@click.pass_context
def ast(ctx, input_file, layout, graphviz):
Sean Leavey's avatar
Sean Leavey committed
759
    """Plot parsed syntax tree."""
760 761 762 763 764 765 766 767 768 769 770 771 772 773 774 775 776
    from finesse.script.builder import KatParser

    # Get state object to set verbosity, etc.
    state = ctx.ensure_object(KatState)

    script = input_file.read()
    builder = KatParser()

    try:
        parser = builder._parse(script)
    except Exception as e:
        state.print_error(f"PARSING ERROR:", e)

    try:
        plot_graph(state, parser.ast, layout, graphviz)
    except Exception as e:
        state.print_error(e)