fixup. PlotCce

sxs-collaboration · May 21, 2024 · c660ccf · c660ccf
1 parent 6756d2e
commit c660ccf
Show file tree

Hide file tree

Showing 2 changed files with 58 additions and 19 deletions.
diff --git a/src/Visualization/Python/PlotCce.py b/src/Visualization/Python/PlotCce.py
@@ -69,10 +69,12 @@ def _parse_modes(ctx, param, all_modes):
     type=int,
     help=(
         "Extraction radius of data to plot as an int. If there is only one Cce"
-        " subfile, that one will be used and this option will be ignored if"
-        " specified. If there are multiple, this option must be specified"
-        " otherwise an error will occur. The expected form of the Cce subfile"
-        " is 'SpectreRXXXX.cce' where XXXX is the zero-padded integer"
+        " subfile, that one will be used and this option does not need to be"
+        " specified. If there is only one Cce subfile and this option is"
+        " specified, the value will be checked against the extraction radius in"
+        " the subfile. If there are multiple subfiles, this option must be"
+        " specified otherwise an error will occur. The expected form of the Cce"
+        " subfile is 'SpectreRXXXX.cce' where XXXX is the zero-padded integer"
         " extraction radius. This option is ignored if the backwards"
         " compatibility option '--cce-group'/'-d' is specified."
     ),
@@ -170,6 +172,18 @@ def plot_cce_command(
                 " subfiles must end with the extension '.cce'"
             )
 
+        if (
+            len(cce_subfiles) == 1
+            and extraction_radius is not None
+            and f"SpectreR{extraction_radius:04}.cce" not in cce_subfiles[0]
+        ):
+            raise click.UsageError(
+                f"The extraction radius passed in ({extraction_radius}) does"
+                " not match the single Cce subfile that was found"
+                f" ({cce_subfiles[0]}). Either specify the correct extraction"
+                " radius, or remove the option altogether."
+            )
+
         if backward_cce_group is not None:
             cce_subfiles = [backward_cce_group]
 
@@ -182,6 +196,7 @@ def plot_cce_command(
                 " Please specify an extraction radius with"
                 " '--extraction-radius'/'-r'."
             )
+
         # If we didn't specify an extraction radius, the subfile name is just
         # the one listed in the file. If the extraction radius was specified
         # (and we've now guaranteed there is more than one subfile) use that as
@@ -242,21 +257,13 @@ def plot_cce_command(
             cycle_idx = (
                 j % len(cycle) if (real or imag) else (j // 2) % len(cycle)
             )
-            if i == 0:
-                ax.plot(
-                    data.index,
-                    data[f"{quantity}{mode}"],
-                    color=cycle[cycle_idx],
-                    linestyle=linestyle,
-                    label=mode,
-                )
-            else:
-                ax.plot(
-                    data.index,
-                    data[f"{quantity}{mode}"],
-                    color=cycle[cycle_idx],
-                    linestyle=linestyle,
-                )
+            ax.plot(
+                data.index,
+                data[f"{quantity}{mode}"],
+                color=cycle[cycle_idx],
+                linestyle=linestyle,
+                label=mode if i == 0 else None,
+            )
         ax.set_ylabel(quantity)
         # Legend only above top plot
         if i == 0:

diff --git a/tests/Unit/Visualization/Python/Test_PlotCce.py b/tests/Unit/Visualization/Python/Test_PlotCce.py
@@ -292,6 +292,38 @@ def test_plot_cce_errors(self):
         self.assertNotEqual(result.exit_code, 0, result.output)
         self.assertIn("Only specify one of '--real'/'--imag'.", result.output)
 
+        result = self.runner.invoke(
+            plot_cce_command,
+            [
+                os.path.join(
+                    self.work_dir, self.reduction_file_name_one_subfile
+                ),
+                "-o",
+                error_filename,
+                "-m",
+                "2,2",
+                "-r",
+                "1000",
+                "--x-bounds",
+                "0.0",
+                "3.0",
+                "--x-label",
+                "Retarded Time (M)",
+                "--title",
+                "Cce at Scri Error2",
+            ],
+            catch_exceptions=False,
+        )
+
+        self.assertNotEqual(result.exit_code, 0, result.output)
+        self.assertIn(
+            (
+                "Either specify the correct extraction radius, or remove the"
+                " option altogether."
+            ),
+            result.output,
+        )
+
         result = self.runner.invoke(
             plot_cce_command,
             [