Added some plots, use country colors in comparative plots, handle dup…

…licate countries better

src/stellarisdashboard/dash_server.py

@@ -22,10 +22,6 @@
 timeline_app.css.config.serve_locally = True
 timeline_app.scripts.config.serve_locally = True
 
-COLOR_PHYSICS = 'rgba(30,100,170,0.5)'
-COLOR_SOCIETY = 'rgba(60,150,90,0.5)'
-COLOR_ENGINEERING = 'rgba(190,150,30,0.5)'
-
 
 @flask_app.route("/")
 @flask_app.route("/checkversion/<version>/")
@@ -385,6 +381,7 @@ def get_galaxy(game_id, date):
             x=1.0,
             y=1.0,
         ),
+        height=720,
         hovermode='closest',
         plot_bgcolor=GALAXY_BG_COLOR,
         paper_bgcolor=STELLARIS_LIGHT_BG_COLOR,
@@ -504,16 +501,8 @@ def get_war_dicts(session, current_date):
 def get_country_color(country_name: str, alpha: float = 1.0) -> str:
     alpha = min(alpha, 1)
     alpha = max(alpha, 0)
-
-    if country_name == "physics":
-        return COLOR_PHYSICS
-    elif country_name == "society":
-        return COLOR_SOCIETY
-    elif country_name == "engineering":
-        return COLOR_ENGINEERING
-
-    random.seed(country_name)
-    r, g, b = [random.randint(20, 255) for _ in range(3)]
+    r, g, b = visualization_data.get_color_vals(country_name)
+    r, g, b = r * 255, g * 255, b * 255
     color = f"rgba({r},{g},{b},{alpha})"
     return color
 

src/stellarisdashboard/visualization_data.py

@@ -1,5 +1,6 @@
 import enum
 import logging
+import random
 from typing import List, Dict, Callable, Any, Tuple, Iterable, Set
 
 import dataclasses
@@ -10,6 +11,10 @@
 
 logger = logging.getLogger(__name__)
 
+COLOR_PHYSICS = (30, 100, 170)
+COLOR_SOCIETY = (60, 150, 90)
+COLOR_ENGINEERING = (190, 150, 30)
+
 
 @enum.unique
 class PlotStyle(enum.Enum):
@@ -51,6 +56,18 @@ class PlotSpecification:
     plot_data_function=lambda pd: pd.controlled_systems,
     style=PlotStyle.line,
 )
+NET_MINERAL_INCOME_GRAPH = PlotSpecification(
+    plot_id='net-mineral-income-graph',
+    title="Net Mineral Income (Warning: Might be inaccurate!)",
+    plot_data_function=lambda pd: pd.net_mineral_income,
+    style=PlotStyle.line,
+)
+NET_ENERGY_INCOME_GRAPH = PlotSpecification(
+    plot_id='net-energy-income-graph',
+    title="Net Energy Income (Warning: Might be inaccurate!)",
+    plot_data_function=lambda pd: pd.net_energy_income,
+    style=PlotStyle.line,
+)
 TECHNOLOGY_PROGRESS_GRAPH = PlotSpecification(
     plot_id='tech-count-graph',
     title="Researched Technologies",
@@ -70,6 +87,12 @@ class PlotSpecification:
     plot_data_function=lambda pd: pd.empire_research_output,
     style=PlotStyle.stacked,
 )
+TOTAL_RESEARCH_OUTPUT_GRAPH = PlotSpecification(
+    plot_id='empire-research-output-comparison-graph',
+    title="Total Research Output",
+    plot_data_function=lambda pd: pd.total_research_output,
+    style=PlotStyle.line,
+)
 SURVEY_PROGRESS_GRAPH = PlotSpecification(
     plot_id='survey-count-graph',
     title="Exploration",
@@ -140,6 +163,8 @@ class PlotSpecification:
     "Economy": [
         PLANET_COUNT_GRAPH,
         SYSTEM_COUNT_GRAPH,
+        NET_ENERGY_INCOME_GRAPH,
+        NET_MINERAL_INCOME_GRAPH,
         EMPIRE_ENERGY_ECONOMY_GRAPH,
         EMPIRE_MINERAL_ECONOMY_GRAPH,
         EMPIRE_FOOD_ECONOMY_GRAPH,
@@ -155,6 +180,7 @@ class PlotSpecification:
     ],
     "Science": [
         TECHNOLOGY_PROGRESS_GRAPH,
+        TOTAL_RESEARCH_OUTPUT_GRAPH,
         SURVEY_PROGRESS_GRAPH,
         RESEARCH_OUTPUT_GRAPH,
         RESEARCH_ALLOCATION_GRAPH,
@@ -208,6 +234,19 @@ def show_military_info(country_data: models.CountryData):
             or country_data.has_federation_with_player)
 
 
+def get_color_vals(key_str: str, range_min: float = 0.1, range_max: float = 1.0):
+    if key_str == "physics":
+        r, g, b = COLOR_PHYSICS
+    elif key_str == "society":
+        r, g, b = COLOR_SOCIETY
+    elif key_str == "engineering":
+        r, g, b = COLOR_ENGINEERING
+    else:
+        random.seed(key_str)
+        r, g, b = [random.uniform(range_min, range_max) for _ in range(3)]
+    return r, g, b
+
+
 class EmpireProgressionPlotData:
     DEFAULT_VAL = float("nan")
 
@@ -218,6 +257,10 @@ def __init__(self, game_name):
         self.pop_count = None
         self.owned_planets = None
         self.controlled_systems = None
+        self.net_mineral_income = None
+        self.net_energy_income = None
+
+        self.total_research_output = None
         self.tech_count = None
         self.survey_count = None
         self.military_power = None
@@ -233,13 +276,19 @@ def __init__(self, game_name):
         self.empire_research_allocation = None
         self.show_everything = config.CONFIG.show_everything
 
+        self.data_dicts = []
+
     def initialize(self):
         self.dates: List[float] = []
         self.player_country: str = None
         self.pop_count: Dict[str, List[int]] = {}
         self.owned_planets: Dict[str, List[int]] = {}
         self.controlled_systems: Dict[str, List[int]] = {}
+        self.net_mineral_income: Dict[str, List[float]] = {}
+        self.net_energy_income: Dict[str, List[float]] = {}
+
         self.tech_count: Dict[str, List[int]] = {}
+        self.total_research_output: Dict[str, List[int]] = {}
         self.survey_count: Dict[str, List[int]] = {}
         self.military_power: Dict[str, List[float]] = {}
         self.fleet_size: Dict[str, List[float]] = {}
@@ -288,6 +337,18 @@ def initialize(self):
         )
         self.empire_research_output = dict(physics=[], society=[], engineering=[])
         self.empire_research_allocation = dict(physics=[], society=[], engineering=[])
+        self.data_dicts = [
+            self.pop_count,
+            self.owned_planets,
+            self.tech_count,
+            self.total_research_output,
+            self.survey_count,
+            self.military_power,
+            self.fleet_size,
+            self.controlled_systems,
+            self.net_mineral_income,
+            self.net_energy_income,
+        ]
 
     def update_with_new_gamestate(self):
         date_in_days = 360.0 * self.dates[-1] if self.dates else -1
@@ -304,7 +365,10 @@ def process_gamestate(self, gs: models.GameState):
             self._extract_pop_count(country_data)
             self._extract_planet_count(country_data)
             self._extract_system_count(country_data)
+            self._extract_energy_income(country_data)
+            self._extract_mineral_income(country_data)
             self._extract_tech_count(country_data)
+            self._extract_research_output(country_data)
             self._extract_exploration_progress(country_data)
             self._extract_military_strength(country_data)
             self._extract_fleet_size(country_data)
@@ -315,9 +379,9 @@ def process_gamestate(self, gs: models.GameState):
         self._extract_player_empire_budget_allocations(gs)
 
         # Pad every dict with the default value if no real value was added, to keep them consistent with the dates list
-        for data_dict in [self.pop_count, self.owned_planets, self.tech_count, self.survey_count, self.military_power, self.fleet_size]:
+        for data_dict in self.data_dicts:
             for key in data_dict:
-                if len(data_dict[key]) < len(self.dates):
+                while len(data_dict[key]) < len(self.dates):
                     data_dict[key].append(EmpireProgressionPlotData.DEFAULT_VAL)
 
     def _extract_player_empire_budget_allocations(self, gs: models.GameState):
@@ -399,13 +463,34 @@ def _extract_system_count(self, country_data: models.CountryData):
             new_val = EmpireProgressionPlotData.DEFAULT_VAL
         self._add_new_value_to_data_dict(self.controlled_systems, country_data.country.country_name, new_val)
 
+    def _extract_energy_income(self, country_data: models.CountryData):
+        if self.show_everything or show_economic_info(country_data):
+            new_val = country_data.energy_income - country_data.energy_spending
+        else:
+            new_val = EmpireProgressionPlotData.DEFAULT_VAL
+        self._add_new_value_to_data_dict(self.net_energy_income, country_data.country.country_name, new_val)
+
+    def _extract_mineral_income(self, country_data: models.CountryData):
+        if self.show_everything or show_economic_info(country_data):
+            new_val = country_data.mineral_income - country_data.mineral_spending
+        else:
+            new_val = EmpireProgressionPlotData.DEFAULT_VAL
+        self._add_new_value_to_data_dict(self.net_mineral_income, country_data.country.country_name, new_val)
+
     def _extract_tech_count(self, country_data: models.CountryData):
         if self.show_everything or show_tech_info(country_data):
             new_val = country_data.tech_progress
         else:
             new_val = EmpireProgressionPlotData.DEFAULT_VAL
         self._add_new_value_to_data_dict(self.tech_count, country_data.country.country_name, new_val)
 
+    def _extract_research_output(self, country_data: models.CountryData):
+        if self.show_everything or show_tech_info(country_data):
+            new_val = country_data.society_research + country_data.physics_research + country_data.engineering_research
+        else:
+            new_val = EmpireProgressionPlotData.DEFAULT_VAL
+        self._add_new_value_to_data_dict(self.total_research_output, country_data.country.country_name, new_val)
+
     def _extract_exploration_progress(self, country_data: models.CountryData):
         if self.show_everything or show_tech_info(country_data):
             new_val = country_data.exploration_progress
@@ -506,6 +591,9 @@ def _extract_player_empire_research(self, country_data: models.CountryData):
     def _add_new_value_to_data_dict(self, data_dict, key, new_val):
         if key not in data_dict:
             data_dict[key] = [EmpireProgressionPlotData.DEFAULT_VAL for _ in range(len(self.dates) - 1)]
+        if len(data_dict[key]) >= len(self.dates):
+            logger.info(f"Ignoring duplicate value for {key}.")
+            return
         data_dict[key].append(new_val)
 
 

src/stellarisdashboard/visualization_mpl.py

@@ -1,11 +1,12 @@
 import logging
 import math
-from typing import List, Dict
+from typing import List, Dict, Set
 
 import pathlib
 
 import itertools
 from matplotlib import pyplot as plt
+import matplotlib.lines
 
 from stellarisdashboard import models, visualization_data, config
 
@@ -143,9 +144,13 @@ def __init__(self, comparison_id: str):
         self.axes = None
         self.comparison_id = comparison_id
         self.plot_data: Dict[str, visualization_data.EmpireProgressionPlotData] = {}
+        self.countries: Set[str] = set()
+        self.countries_in_legend: Set[str] = set()
+        self.games_in_legend: Set[str] = set()
 
     def add_data(self, game_name: str, pd: visualization_data.EmpireProgressionPlotData):
         self.plot_data[game_name] = pd
+        self.countries |= pd.owned_planets.keys()
 
     def make_plots(self):
         for category, plot_specifications in visualization_data.THEMATICALLY_GROUPED_PLOTS.items():
@@ -154,6 +159,7 @@ def make_plots(self):
                 continue
             self._initialize_axes(category, plot_specifications)
             for plot_spec, ax in zip(plot_specifications, self.axes):
+                self.countries_in_legend = set()
                 self._make_line_plots(ax, plot_spec)
             self.save_plot(plot_id=category)
 
@@ -185,36 +191,35 @@ def _initialize_axes(self, category: str, plot_specifications: List[visualizatio
 
     def _make_line_plots(self, ax, plot_spec: visualization_data.PlotSpecification):
         ax.set_title(plot_spec.title)
+        game_handles = []
         for game_index, (game_name, style) in enumerate(zip(self.plot_data.keys(), itertools.cycle(MatplotLibComparativeVisualization.LINE_STYLES))):
             pd = self.plot_data[game_name]
-            for i, (key, x, y) in enumerate(pd.data_sorted_by_last_value(plot_spec)):
+            for i, (key, x, y) in enumerate(pd.iterate_data(plot_spec)):
                 if y:
                     plot_kwargs = self._get_country_plot_kwargs(
-                        plot_data=pd,
                         linestyle=style,
                         country_name=key,
-                        game_index=game_index,
                         game_name=game_name,
                     )
                     ax.plot(x, y, **plot_kwargs)
-        ax.legend()
+            game_handles.append(matplotlib.lines.Line2D([], [], linestyle=style, color='grey', label=game_name))
+        ax.legend(loc=2, prop={'size': 6})
+        self.fig.legend(handles=game_handles, loc=8)
 
     def _get_country_plot_kwargs(
-            self, plot_data: visualization_data.EmpireProgressionPlotData,
-            linestyle: str,
+            self, linestyle: str,
             country_name: str,
-            game_index: int,
             game_name: str,
     ):
-        linewidth = 0.25
-        color_index = game_index / max(1, len(self.plot_data) - 1)
-        c = MatplotLibVisualization.COLOR_MAP(color_index)
-        alpha = 0.5
+        linewidth = 0.75
+        alpha = 1
         label = None
-        if country_name == plot_data.player_country:
-            linewidth = 2
-            label = f"{game_name} ({country_name})"
-            alpha = 1.0
+        c = visualization_data.get_color_vals(country_name, range_min=0, range_max=0.9)
+        if country_name not in self.countries_in_legend:
+            self.countries_in_legend.add(country_name)
+            label = f"{country_name}"
+        if game_name not in self.games_in_legend:
+            self.games_in_legend.add(game_name)
         return dict(label=label, c=c, linewidth=linewidth, alpha=alpha, linestyle=linestyle)
 
     def save_plot(self, plot_id):