Refactor TestMassProfiler

-add analyse method
-decouple consume_figure from __init__
-cache profiler in subcommand profile and analyse

tidaltailsim/__main__.py

@@ -28,6 +28,8 @@ def parse_args(self, args=None, namespace=None):
  setattr(args, 'exiting', exiting)
  return args
 
+ # This method will be executed in the subparser context but parse_args will not
+ # therefore, we need to take care of both so that the exit mark can properly propagte upstream
  def parse_known_args(self, args=None, namespace=None):
  self.__mark_as_exiting = False
  namespace, args = super().parse_known_args(args, namespace)
@@ -123,6 +125,7 @@ def singleorbital_pickled_routine(args):
  pre_animate_func=lambda ax: [animator.plot_core_path(ax, core_index) for core_index in plotting_cores],
  animate_func=lambda fig, ax, speed, framerate: animator.animate(fig, ax, rate=speed, framerate=framerate, time_initial=args.timeinitial))
  else:
+ from pprint import pprint
  # Launch animation before launching the interactive shell
  fig, ax = plt.subplots(subplot_kw=dict(projection='3d') if not args.d2 else None)
 
@@ -165,8 +168,11 @@ def pick_handler(event):
  fig.show()
 
  # Animation has succesfully launched. Now launch the interactive shell
+
+ # initialise some variables
  continue_ = True
  test_mass_lines = []
+ profiler = None
 
  shell_parser = ShellArgumentParser(prog='', add_help=False)
  sub_shell_parsers = shell_parser.add_subparsers(title='command')
@@ -206,13 +212,17 @@ def pick_handler(event):
 
  profile_parser = sub_shell_parsers.add_parser('profile')
  profile_parser.set_defaults(action='profile')
- profile_parser.add_argument('test_mass_index', type=int)
+ profile_parser.add_argument('test_mass_index', type=int, nargs='?', default=None)
+
+ analyse_parser = sub_shell_parsers.add_parser('analyse')
+ analyse_parser.set_defaults(action='analyse')
+ analyse_parser.add_argument('test_mass_index', type=int, nargs='?', default=None)
 
- for parser in (plot_path_parser, profile_parser):
- frame_slice_group = parser.add_mutually_exclusive_group(required=True)
- frame_slice_group.add_argument('-f', '--from', type=int, dest='from_', default=False,
+ for parser, required in zip((plot_path_parser, profile_parser, analyse_parser), (True, False, False)):
+ frame_slice_group = parser.add_mutually_exclusive_group(required=required)
+ frame_slice_group.add_argument('-f', '--from', type=int, dest='from_', default=None,
  metavar='begin_frame_index')
- frame_slice_group.add_argument('-ft', '--fromto', type=int, default=False, nargs=2,
+ frame_slice_group.add_argument('-ft', '--fromto', type=int, default=None, nargs=2,
  metavar=('begin_frame_index', 'end_frame_index'))
 
  print('Launching interactive shell\nTry \'help\' command to learn more')
@@ -263,10 +273,36 @@ def pick_handler(event):
  test_mass_lines = []
  fig.canvas.draw_idle()
  elif 'profile' == action:
- slice_ = slice(shell_args.fromto[0], shell_args.fromto[1]) if shell_args.fromto else slice(shell_args.from_, None)
+ if shell_args.fromto is not None or shell_args.from_ is not None:
+ slice_ = slice(shell_args.fromto[0], shell_args.fromto[1]) if shell_args.fromto else slice(shell_args.from_, None)
+ else:
+ slice_ = None
+
  fig = plt.figure() # so that pyplot keep reference to Figure created and hence not garbage-collected
- profiler = TestMassProfiler(problem, args.galaxy, args.orbital, shell_args.test_mass_index, figure=fig, frame_slice=slice_)
- fig.show()
+
+ if shell_args.test_mass_index is not None:
+ profiler = TestMassProfiler(problem, args.galaxy, args.orbital, shell_args.test_mass_index, figure=fig, frame_slice=slice_)
+ fig.show()
+ else:
+ if isinstance(profiler, TestMassProfiler):
+ profiler.consume_figure(fig, frame_slice=slice_)
+ print(f'Plotting from test_mass_index={profiler.test_mass_index}')
+ fig.show()
+ else:
+ print('Please either specify the test_mass_index or call \'analyse\' command first.')
+ elif 'analyse' == action:
+ if shell_args.fromto is not None or shell_args.from_ is not None:
+ slice_ = slice(shell_args.fromto[0], shell_args.fromto[1]) if shell_args.fromto else slice(shell_args.from_, None)
+ else:
+ slice_ = None
+
+ if shell_args.test_mass_index is not None:
+ profiler = TestMassProfiler(problem, args.galaxy, args.orbital, shell_args.test_mass_index, frame_slice=slice_)
+
+ if isinstance(profiler, TestMassProfiler):
+ pprint(profiler.analyse(frame_slice=slice_), sort_dicts=False)
+ else:
+ print('Please either specify the test_mass_index or call \'profile\' command first.')
 
 
 def parse_animation(args, dimension_is_3, animate_func, pre_animate_func=None):

tidaltailsim/galaxy_orbital_toolkit.py

@@ -202,6 +202,8 @@ def __init__(self, two_galaxy_problem: TwoGalaxyProblem,
  raise TypeError('two_galaxy_problem must be an instance of tidaltailsim.two_galaxy_problem.TwoGalaxyProblem')
  if galaxy_index not in (1, 2):
  raise ValueError('galaxy_index must be either 1 or 2')
+ if frame_slice is None:
+ frame_slice = slice(None)
  self._problem = two_galaxy_problem
  self._galaxy_index = galaxy_index
 
@@ -210,23 +212,12 @@ def __init__(self, two_galaxy_problem: TwoGalaxyProblem,
 
  self._orbital_index = orbital_index
  self._test_mass_index = test_mass_index
- if figure is None:
- figure = Figure()
 
- self._figure = figure # type: Figure
- self._axs = figure.subplots(2, 2, sharex=True)
- # rows: 0=>distance from core 1=>aceentricity
- # columns: 0=>rel to core1 1=>rel to core2
-
- self._figure.set_tight_layout(True)
- self._annotate_figure()
+ self.default_frame_slice = frame_slice # type: slice
 
  self._calculate_relative_states_from_cores()
 
- self._plot_distances_from_cores(frame_slice)
- self._plot_eccentricity(frame_slice)
-
- self.default_frame_slice = frame_slice # type: slice
+ self.consume_figure(figure)
 
  @property
  def two_galaxy_problem(self) -> TwoGalaxyProblem:
@@ -298,3 +289,54 @@ def _annotate_figure(self):
  ax.set_xlabel('Time')
  ax.set_ylabel(quantity)
  ax.set_title(f'{quantity} relative to\nthe core of galaxy {j+1}')
+
+ def consume_figure(self, figure: Figure, frame_slice: slice = None):
+ """plot distances from core and aceentricities and also annotate the figure"""
+ if frame_slice is None:
+ frame_slice = self.default_frame_slice
+
+ if not isinstance(figure, Figure):
+ self._figure = None # type: Figure
+ self._axs = None
+ else:
+ self._figure = figure # type: Figure
+ self._axs = figure.subplots(2, 2, sharex=True)
+ # rows: 0=>distance from core 1=>aceentricity
+ # columns: 0=>rel to core1 1=>rel to core2
+
+ self._figure.set_tight_layout(True)
+ self._annotate_figure()
+
+ self._plot_distances_from_cores(frame_slice)
+ self._plot_eccentricity(frame_slice)
+
+ return self._axs
+
+ def analyse(self, frame_slice: slice = None):
+ if frame_slice is None:
+ frame_slice = self.default_frame_slice
+ mean_distance = np.nanmean(self._distace_from_cores[:, frame_slice], axis=1)
+ var_distance = np.nanvar(self._distace_from_cores[:, frame_slice], axis=1, ddof=1)
+
+ mean_eccentricity = np.nanmean(self._eccentricity[:, frame_slice], axis=1)
+ var_eccentricity = np.nanvar(self._eccentricity[:, frame_slice], axis=1, ddof=1)
+
+ times = self._problem.time_domain[frame_slice]
+
+ result = [{
+ 'relative to galaxy': i + 1,
+ 'distance': {
+ 'mean': mean_distance[i],
+ 'variance': var_distance[i]
+ },
+ 'eccentricity': {
+ 'mean': mean_eccentricity[i],
+ 'variance': var_eccentricity[i]
+ }
+ } for i in range(2)]
+
+ return {
+ 'time limit': (times[0], times[-1]),
+ 'test_mass_index': self.test_mass_index,
+ 'result': result
+ }