# HG changeset patch # User Jeremy Thurgood # Date 1336759596 -7200 # Node ID 372d886f9e70c88138f428b66f0a7545de993210 # Parent 0090ecf08544d83b10e4ff18f5d0c32d859e3d0f New suggest_research() method on Lab. diff -r 0090ecf08544 -r 372d886f9e70 gamelib/lab.py --- a/gamelib/lab.py Fri May 11 17:17:44 2012 +0200 +++ b/gamelib/lab.py Fri May 11 20:06:36 2012 +0200 @@ -1,11 +1,123 @@ # -*- test-case-name: gamelib.tests.test_lab -*- -from random import random, choice +from random import random, choice, sample + +import networkx from gamelib import research, schematics from gamelib.game_base import get_subclasses +class ScienceGraph(object): + def __init__(self, all_science, known_science): + self.graph = networkx.DiGraph() + self.all_science = all_science + self.known_science = known_science + self.add_all_science() + self.tag_known_science() + assert networkx.is_directed_acyclic_graph(self.graph) + + def add_all_science(self): + # Add level 0 of everything to the graph. + for science in self.all_science: + self.graph.add_node((science, 0), known=False) + + # Walk dependencies and fill in intermediate nodes. + for science in self.all_science: + for dep in science.PREREQUISITES: + self.add_node_string(*dep) + self.graph.add_edge(dep, (science, 0)) + + def add_node_string(self, science, level): + node = (science, level) + if node in self.graph: + return node + + # We prepopulate with level 0 of eveything. + assert level >= 0 + + self.graph.add_node(node, known=False) + parent = self.add_node_string(science, level - 1) + self.graph.add_edge(parent, node) + return node + + def tag_known_science(self): + for science in self.known_science: + # We may know more of this than the graph has. + self.add_node_string(type(science), science.points + 1) + for i in range(science.points + 1): + self.graph.node[(type(science), i)]['known'] = True + + def is_known(self, node): + return self.graph.node[node]['known'] + + def distances_to_known(self, science): + nodes = set() + nodes_to_check = [(science, 0)] + + while nodes_to_check: + node = nodes_to_check.pop() + if self.is_known(node) or node in nodes: + continue + nodes.add(node) + nodes_to_check.extend(self.graph.predecessors(node)) + + distances = {} + for node in nodes: + distances.setdefault(node[0], 0) + distances[node[0]] += 1 + + return distances + + def count_unknown(self, sciences): + return len([s for s in sciences if not self.is_known((s, 0))]) + + def find_prospects(self): + prospects = {'research': [], 'schematic': []} + for science in self.all_science: + distances = self.distances_to_known(science) + if not distances: + # We already know this thing. + continue + # Remove the thing we're trying to get from the distance. + distances.pop(science) + if self.count_unknown(distances.keys()) > 0: + # We only want direct breakthroughs. + continue + prospects[science.SCIENCE_TYPE].append( + (sum(distances.values()), science, distances)) + return dict((k, sorted(v)) for k, v in prospects.items()) + + def find_promising_areas(self, size=3): + basic_science = False + areas_for_research = set() + areas_for_schematics = set() + prospects = self.find_prospects() + + for points, target, distances in prospects['schematic']: + if points > 0: + # We need nonzero points in these things. + areas_for_schematics.update(distances.keys()) + else: + # Any of these things qualify us. + areas_for_schematics.update(p for p, _ in target.PREREQUISITES) + + for points, target, distances in prospects['research']: + if points == 0: + basic_science = True + else: + areas_for_research.update(distances.keys()) + + suggestions = [] + k = min(size, len(areas_for_schematics)) + suggestions.extend(sample(areas_for_schematics, k)) + if len(suggestions) < size: + k = min(size - len(suggestions), len(areas_for_research)) + suggestions.extend(sample(areas_for_research, k)) + + return basic_science, suggestions + + class Lab(object): BASIC_RESEARCH_SUCCESS_RATE = 0.05 BASIC_RESEARCH_SUCCESS_MULTIPLIER = 2 @@ -14,6 +126,7 @@ self.science = [] self.new_research = get_subclasses(research.ResearchArea) self.new_schematics = get_subclasses(schematics.Schematic) + self.all_science = [s for s in self.new_research + self.new_schematics] if init_data is not None: # Load stored state. @@ -49,8 +162,8 @@ new_science = [physics] # We get two other random sciences with no prerequisites. - for _ in range(2): - science = choice(self.find_new_research())() + for science in sample(self.find_new_research(), 2): + science = science() self._gain_science(science) new_science.append(science) @@ -141,3 +254,13 @@ self._gain_science(breakthrough) breakthroughs = [breakthrough] return breakthroughs + + def suggest_research(self): + """Suggest research areas to pursue. + + Return value is a tuple of (bool, list), where the first element + indicates whether basic research might pay off and the second contains + Science classes that can be profitably pursued. + """ + graph = ScienceGraph(self.all_science, self.science) + return graph.find_promising_areas() diff -r 0090ecf08544 -r 372d886f9e70 gamelib/schematics.py --- a/gamelib/schematics.py Fri May 11 17:17:44 2012 +0200 +++ b/gamelib/schematics.py Fri May 11 20:06:36 2012 +0200 @@ -19,6 +19,7 @@ 'AQUATIC', 'INTELLIGENCE', 'AI', + 'COUNTERMEASURE', ) K = 1000 @@ -189,7 +190,7 @@ class LaserGun(Schematic): - NAME = "Laser Gun" + NAME = "laser gun" COST = 300 CATEGORIES = (cat.HAND_WEAPON,) PREREQUISITES = ( @@ -197,6 +198,16 @@ ) +class EmpMissile(Schematic): + NAME = "EMP missile" + COST = 1500 + CATEGORIES = (cat.COUNTERMEASURE,) + PREREQUISITES = ( + (research.Electrickery, 5), + (research.Rocketry, 2), + ) + + class DoomsdayVirus(Schematic): NAME = "doomsday virus" COST = 100 * K diff -r 0090ecf08544 -r 372d886f9e70 gamelib/tests/test_lab.py --- a/gamelib/tests/test_lab.py Fri May 11 17:17:44 2012 +0200 +++ b/gamelib/tests/test_lab.py Fri May 11 20:06:36 2012 +0200 @@ -1,6 +1,6 @@ from unittest import TestCase -from gamelib.lab import Lab +from gamelib.lab import Lab, ScienceGraph from gamelib import research, schematics @@ -9,6 +9,7 @@ 'research.Physics': 5, 'research.Rocketry': 2, 'research.Electrickery': 3, + 'schematic.MachineGun': 1, }, } @@ -46,3 +47,8 @@ new_science = dict( LAB_DATA['science'].items() + [('schematic.LightningGun', 0)]) self.assertEqual({'science': new_science}, lab.save_data()) + + def test_science_graph(self): + lab = Lab(LAB_DATA) + self.assertEqual((True, [research.Electrickery, research.Physics]), + lab.suggest_research())