Team prediction improvements

README.md

@@ -1130,7 +1130,8 @@ from metamon.backend.team_prediction.usage_stats import get_usage_stats
 from datetime import date
 usage_stats = get_usage_stats("gen1ou",
     start_date=date(2017, 12, 1),
-    end_date=date(2018, 3, 30)
+    end_date=date(2018, 3, 30),
+    rank="1500",
 )
 alakazam_info: dict = usage_stats["Alakazam"] # non alphanum chars and case are flexible
 ```

metamon/backend/team_prediction/predictor.py

@@ -17,14 +17,23 @@
 )
 from metamon.backend.team_prediction.usage_stats import (
     PreloadedSmogonUsageStats,
+    DEFAULT_USAGE_RANK,
+    RankLike,
 )
 from metamon.backend.replay_parser.str_parsing import pokemon_name
 from metamon.backend.team_prediction.team import TeamSet, PokemonSet, Roster
 
 
 class TeamPredictor(ABC):
-    def __init__(self, replay_stats_dir: Optional[str] = None):
+    def __init__(
+        self,
+        replay_stats_dir: Optional[str] = None,
+        usage_stats_rank: RankLike = DEFAULT_USAGE_RANK,
+    ):
+        if usage_stats_rank is None:
+            usage_stats_rank = DEFAULT_USAGE_RANK
         self.replay_stats_dir = replay_stats_dir
+        self.usage_stats_rank = usage_stats_rank
 
     def bin_usage_stats_dates(
         self, date: datetime.date
@@ -50,6 +59,7 @@ def get_legacy_team_builder(self, format: str, date: datetime.date) -> TeamBuild
             format=format,
             start_date=start_date,
             end_date=end_date,
+            rank=self.usage_stats_rank,
         )
 
     def get_usage_stats(
@@ -209,9 +219,13 @@ def __init__(
         top_k_scored_teams: int = 10,
         top_k_scored_movesets: int = 3,
         replay_stats_dir: Optional[str] = None,
+        usage_stats_rank: RankLike = DEFAULT_USAGE_RANK,
     ):
         assert not isinstance(top_k_consistent_teams, str)
-        super().__init__(replay_stats_dir)
+        super().__init__(
+            replay_stats_dir,
+            usage_stats_rank=usage_stats_rank,
+        )
         self.stat_format = None
         self.top_k_consistent_teams = top_k_consistent_teams
         self.top_k_consistent_movesets = top_k_consistent_movesets

metamon/backend/team_prediction/team.py

@@ -15,13 +15,13 @@
     unknown,
 )
 from metamon.backend.replay_parser.str_parsing import pokemon_name
-from metamon.backend.team_prediction.usage_stats import get_usage_stats
+from metamon.backend.team_prediction.usage_stats import get_usage_stats, DEFAULT_USAGE_RANK
 from metamon.backend.showdown_dex import Dex
 
 
 def moveset_size(pokemon_name: str, gen: int) -> int:
     # attempts to handle cases where we would expect a Pokemon to have less than 4 moves
-    stat = get_usage_stats(f"gen{gen}ubers")
+    stat = get_usage_stats(f"gen{gen}ubers", rank=DEFAULT_USAGE_RANK)
     try:
         moves = len(set(stat[pokemon_name]["moves"].keys()) - {"Nothing"})
     except KeyError:

metamon/backend/team_prediction/usage_stats/__init__.py

@@ -1,2 +1,8 @@
-from .stat_reader import get_usage_stats, PreloadedSmogonUsageStats
+from .stat_reader import (
+    get_usage_stats,
+    PreloadedSmogonUsageStats,
+    DEFAULT_USAGE_RANK,
+    RankLike,
+    list_available_usage_ranks,
+)
 from .legacy_team_builder import TeamBuilder, PokemonStatsLookupError

metamon/backend/team_prediction/usage_stats/create_usage_jsons.py

@@ -1,6 +1,7 @@
 import os
 import json
 import argparse
+from collections import defaultdict
 from tqdm import tqdm
 
 from metamon.backend.team_prediction.usage_stats.format_rules import (
@@ -20,54 +21,69 @@ def main(args):
         for gen in range(1, 10):
             for year in range(2014, 2026):
                 for month in range(1, 13):
+                    date = f"{year}-{month:02d}"
                     stat_dir = os.path.join(args.smogon_stat_dir)
-                    valid_movesets = []
+                    valid_movesets_by_rank = defaultdict(list)
                     for format in VALID_TIERS:
                         format_name = f"gen{gen}{format.name.lower()}"
-                        stat = SmogonStat(
+
+                        ranks = SmogonStat.available_ranks(
                             format_name,
                             raw_stats_dir=stat_dir,
-                            date=f"{year}-{month:02d}",
+                            date=date,
                         )
-                        if stat.movesets:
-                            # if we find data for this, save it
+
+                        for rank in ranks:
+                            stat = SmogonStat(
+                                format_name,
+                                raw_stats_dir=stat_dir,
+                                date=date,
+                                rank=rank,
+                                verbose=False,
+                            )
+                            if not stat.movesets:
+                                continue
+
                             path = os.path.join(
                                 args.save_dir,
                                 "movesets_data",
                                 f"gen{gen}",
                                 f"{format.name.lower()}",
-                                f"{year}-{month:02d}.json",
+                                rank,
+                                f"{date}.json",
                             )
                             os.makedirs(os.path.dirname(path), exist_ok=True)
                             with open(path, "w") as f:
                                 json.dump(stat.movesets, f)
-                            valid_movesets.append(stat.movesets)
 
-                            check_cheatsheet = {}
-                            for mon in stat.movesets.keys():
-                                checks = stat.movesets[mon]["checks"]
-                                check_cheatsheet[mon] = checks
+                            check_cheatsheet = {
+                                mon: stat.movesets[mon]["checks"]
+                                for mon in stat.movesets.keys()
+                            }
                             path = os.path.join(
                                 args.save_dir,
                                 "checks_data",
                                 f"gen{gen}",
                                 f"{format.name.lower()}",
-                                f"{year}-{month:02d}.json",
+                                rank,
+                                f"{date}.json",
                             )
                             os.makedirs(os.path.dirname(path), exist_ok=True)
                             with open(path, "w") as f:
                                 json.dump(check_cheatsheet, f)
+
+                            valid_movesets_by_rank[rank].append(stat.movesets)
                         pbar.update(1)
-                    if valid_movesets:
-                        # merge all the tiers. used to lookup rare Pokémon choices, i.e. fooling around
-                        # with low-tier Pokémon in OverUsed
-                        inclusive_movesets = merge_movesets(valid_movesets)
+
+                    for rank, tier_movesets in valid_movesets_by_rank.items():
+                        inclusive_movesets = merge_movesets(tier_movesets)
                         path = os.path.join(
                             args.save_dir,
                             "movesets_data",
                             f"gen{gen}",
                             "all_tiers",
-                            f"{year}-{month:02d}.json",
+                            rank,
+                            f"{date}.json",
                         )
                         os.makedirs(os.path.dirname(path), exist_ok=True)
                         with open(path, "w") as f:

metamon/backend/team_prediction/usage_stats/legacy_team_builder.py

@@ -7,7 +7,7 @@
 import numpy as np
 
 import metamon
-from metamon.backend.team_prediction.usage_stats import get_usage_stats
+from metamon.backend.team_prediction.usage_stats import get_usage_stats, RankLike
 from metamon.backend.team_prediction.usage_stats.constants import (
     HIDDEN_POWER_IVS,
     HIDDEN_POWER_DVS,
@@ -43,12 +43,18 @@ def __init__(
         format: str,
         start_date: datetime.date,
         end_date: datetime.date,
+        rank: RankLike = None,
         verbose: bool = False,
         remove_banned: bool = False,
     ):
         self.format = format
         self.gen = metamon.backend.format_to_gen(format)
-        self.stat = get_usage_stats(format, start_date, end_date)
+        self.stat = get_usage_stats(
+            format,
+            start_date,
+            end_date,
+            rank=rank,
+        )
         if remove_banned:
             self.stat.remove_banned_pm()
         self.verbose = verbose